[Sema] Walk SyntacticElementTarget for completion

Instead of walking the single ASTNode from the
target, walk all AST nodes associated with the
target to find the completion expr. This is needed
to find the completion expr in a pattern for an
initialization target.

Author: hamishknight

include/swift/Sema/CompletionContextFinder.h

@@ -20,6 +20,10 @@
 
 namespace swift {
 
+namespace constraints {
+class SyntacticElementTarget;
+}
+
 class CompletionContextFinder : public ASTWalker {
   enum class ContextKind {
     FallbackExpression,
@@ -53,12 +57,9 @@ class CompletionContextFinder : public ASTWalker {
     return MacroWalking::Arguments;
   }
 
-  /// Finder for completion contexts within the provided initial expression.
-  CompletionContextFinder(ASTNode initialNode, DeclContext *DC)
-      : InitialExpr(initialNode.dyn_cast<Expr *>()), InitialDC(DC) {
-    assert(DC);
-    initialNode.walk(*this);
-  };
+  /// Finder for completion contexts within the provided SyntacticElementTarget.
+  CompletionContextFinder(constraints::SyntacticElementTarget target,
+                          DeclContext *DC);
 
   /// Finder for completion contexts within the outermost non-closure context of
   /// the code completion expression's direct context.

lib/Sema/BuilderTransform.cpp

@@ -984,7 +984,8 @@ Optional<BraceStmt *> TypeChecker::applyResultBuilderBodyTransform(
     SmallVector<Solution, 4> solutions;
     cs.solveForCodeCompletion(solutions);
 
-    CompletionContextFinder analyzer(func, func->getDeclContext());
+    SyntacticElementTarget funcTarget(func);
+    CompletionContextFinder analyzer(funcTarget, func->getDeclContext());
     if (analyzer.hasCompletion()) {
       filterSolutionsForCodeCompletion(solutions, analyzer);
       for (const auto &solution : solutions) {

lib/Sema/CompletionContextFinder.cpp

@@ -12,10 +12,19 @@
 
 #include "swift/Sema/CompletionContextFinder.h"
 #include "swift/Parse/Lexer.h"
+#include "swift/Sema/SyntacticElementTarget.h"
 
 using namespace swift;
+using namespace constraints;
 using Fallback = CompletionContextFinder::Fallback;
 
+CompletionContextFinder::CompletionContextFinder(
+    SyntacticElementTarget target, DeclContext *DC)
+    : InitialExpr(target.getAsExpr()), InitialDC(DC) {
+  assert(DC);
+  target.walk(*this);
+}
+
 ASTWalker::PreWalkResult<Expr *>
 CompletionContextFinder::walkToExprPre(Expr *E) {
   if (auto *closure = dyn_cast<ClosureExpr>(E)) {

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -574,15 +574,12 @@ bool TypeChecker::typeCheckForCodeCompletion(
       return false;
   }
 
-  auto node = target.getAsASTNode();
-  if (!node)
-    return false;
-
-  if (auto *expr = getAsExpr(node)) {
-    node = expr->walk(SanitizeExpr(Context));
+  if (getAsExpr(target.getAsASTNode())) {
+    SanitizeExpr sanitizer(Context);
+    target = *target.walk(sanitizer);
   }
 
-  CompletionContextFinder contextAnalyzer(node, DC);
+  CompletionContextFinder contextAnalyzer(target, DC);
 
   // If there was no completion expr (e.g. if the code completion location was
   // among tokens that were skipped over during parser error recovery) bail.