[CodeCompletion] Improve context type analysis

include/swift/AST/Expr.h

@@ -687,20 +687,13 @@ class ErrorExpr : public Expr {
 /// can help us preserve the context of the code completion position.
 class CodeCompletionExpr : public Expr {
   SourceRange Range;
-  bool Activated;
 
 public:
-  CodeCompletionExpr(SourceRange Range, Type Ty = Type()) :
-      Expr(ExprKind::CodeCompletion, /*Implicit=*/true, Ty),
-      Range(Range) {
-    Activated = false;
-  }
+  CodeCompletionExpr(SourceRange Range, Type Ty = Type())
+      : Expr(ExprKind::CodeCompletion, /*Implicit=*/true, Ty), Range(Range) {}
 
   SourceRange getSourceRange() const { return Range; }
 
-  bool isActivated() const { return Activated; }
-  void setActivated() { Activated = true; }
-
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::CodeCompletion;
   }

lib/IDE/CodeCompletion.cpp

@@ -1509,13 +1509,20 @@ static bool isTopLevelContext(const DeclContext *DC) {
 static Type getReturnTypeFromContext(const DeclContext *DC) {
   if (auto FD = dyn_cast<AbstractFunctionDecl>(DC)) {
     if (FD->hasInterfaceType()) {
-      if (auto FT = FD->getInterfaceType()->getAs<FunctionType>()) {
+      auto Ty = FD->getInterfaceType();
+      if (FD->getDeclContext()->isTypeContext())
+        Ty = FD->getMethodInterfaceType();
+      if (auto FT = Ty->getAs<AnyFunctionType>())
         return FT->getResult();
-      }
     }
-  } else if (auto CE = dyn_cast<AbstractClosureExpr>(DC)) {
-    if (CE->getType()) {
-      return CE->getResultType();
+  } else if (auto ACE = dyn_cast<AbstractClosureExpr>(DC)) {
+    if (ACE->getType())
+      return ACE->getResultType();
+    if (auto CE = dyn_cast<ClosureExpr>(ACE)) {
+      if (CE->hasExplicitResultType())
+        return const_cast<ClosureExpr *>(CE)
+            ->getExplicitResultTypeLoc()
+            .getType();
     }
   }
   return Type();
@@ -3893,7 +3900,7 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
   using FunctionParams = ArrayRef<AnyFunctionType::Param>;
 
   static bool
-  collectPossibleParamLists(DeclContext &DC, CallExpr *callExpr,
+  collectPossibleParamLists(DeclContext &DC, ApplyExpr *callExpr,
                             SmallVectorImpl<FunctionParams> &candidates) {
     auto *fnExpr = callExpr->getFn();
 
@@ -3940,27 +3947,20 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     if (!tuple)
       return false;
 
-    for (unsigned i = 0, n = tuple->getNumElements(); i != n; ++i) {
-      if (isa<CodeCompletionExpr>(tuple->getElement(i))) {
-        HasName = !tuple->getElementName(i).empty();
-        Position = i;
-        return true;
-      }
-    }
     auto &SM = DC.getASTContext().SourceMgr;
     for (unsigned i = 0, n = tuple->getNumElements(); i != n; ++i) {
       if (SM.isBeforeInBuffer(tuple->getElement(i)->getEndLoc(),
                               CCExpr->getStartLoc()))
         continue;
-      HasName = !tuple->getElementName(i).empty();
+      HasName = tuple->getElementNameLoc(i).isValid();
       Position = i;
       return true;
     }
     return false;
   }
 
   static bool
-  collectArgumentExpectation(DeclContext &DC, CallExpr *CallE, Expr *CCExpr,
+  collectArgumentExpectation(DeclContext &DC, ApplyExpr *CallE, Expr *CCExpr,
                              std::vector<Type> &ExpectedTypes,
                              std::vector<StringRef> &ExpectedNames) {
     // Collect parameter lists for possible func decls.
@@ -3974,15 +3974,17 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     if (!getPositionInArgs(DC, CallE->getArg(), CCExpr, Position, HasName))
       return false;
 
-    // Collect possible types at the position.
+    // Collect possible types (or labels) at the position.
     {
+      bool MayNeedName =
+          !HasName && isa<CallExpr>(CallE) && !CallE->isImplicit();
       SmallPtrSet<TypeBase *, 4> seenTypes;
       SmallPtrSet<Identifier, 4> seenNames;
       for (auto Params : Candidates) {
         if (Position >= Params.size())
           continue;
         const auto &Param = Params[Position];
-        if (Param.hasLabel() && !HasName) {
+        if (Param.hasLabel() && MayNeedName) {
           if (seenNames.insert(Param.getLabel()).second)
             ExpectedNames.push_back(Param.getLabel().str());
         } else {
@@ -4909,19 +4911,19 @@ namespace  {
   class ExprParentFinder : public ASTWalker {
     friend class CodeCompletionTypeContextAnalyzer;
     Expr *ChildExpr;
-    llvm::function_ref<bool(ASTNode)> Predicate;
+    llvm::function_ref<bool(ParentTy)> Predicate;
 
     bool arePositionsSame(Expr *E1, Expr *E2) {
       return E1->getSourceRange().Start == E2->getSourceRange().Start &&
         E1->getSourceRange().End == E2->getSourceRange().End;
     }
 
   public:
-    llvm::SmallVector<ASTNode, 5> Ancestors;
-    ASTNode ParentClosest;
-    ASTNode ParentFarthest;
+    llvm::SmallVector<ParentTy, 5> Ancestors;
+    ParentTy ParentClosest;
+    ParentTy ParentFarthest;
     ExprParentFinder(Expr* ChildExpr,
-                     llvm::function_ref<bool(ASTNode)> Predicate) :
+                     llvm::function_ref<bool(ParentTy)> Predicate) :
                      ChildExpr(ChildExpr), Predicate(Predicate) {}
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
@@ -4966,6 +4968,18 @@ namespace  {
         Ancestors.pop_back();
       return true;
     }
+
+    std::pair<bool, Pattern *> walkToPatternPre(Pattern *P) override {
+      if (Predicate(P))
+        Ancestors.push_back(P);
+      return { true, P };
+    }
+
+    Pattern *walkToPatternPost(Pattern *P) override {
+      if (Predicate(P))
+        Ancestors.pop_back();
+      return P;
+    }
   };
 } // end anonymous namespace
 
@@ -4981,16 +4995,19 @@ class CodeCompletionTypeContextAnalyzer {
 public:
   CodeCompletionTypeContextAnalyzer(DeclContext *DC, Expr *ParsedExpr) : DC(DC),
     ParsedExpr(ParsedExpr), SM(DC->getASTContext().SourceMgr),
-    Context(DC->getASTContext()), Finder(ParsedExpr, [](ASTNode Node) {
-      if (auto E = Node.dyn_cast<Expr *>()) {
+    Context(DC->getASTContext()),
+    Finder(ParsedExpr,  [](ASTWalker::ParentTy Node) {
+      if (auto E = Node.getAsExpr()) {
         switch(E->getKind()) {
         case ExprKind::Call:
+        case ExprKind::Binary:
+        case ExprKind::PrefixUnary:
         case ExprKind::Assign:
           return true;
         default:
           return false;
-      }
-      } else if (auto S = Node.dyn_cast<Stmt *>()) {
+        }
+      } else if (auto S = Node.getAsStmt()) {
         switch (S->getKind()) {
         case StmtKind::Return:
         case StmtKind::ForEach:
@@ -5002,25 +5019,34 @@ class CodeCompletionTypeContextAnalyzer {
         default:
           return false;
         }
-      } else if (auto D = Node.dyn_cast<Decl *>()) {
+      } else if (auto D = Node.getAsDecl()) {
         switch (D->getKind()) {
         case DeclKind::PatternBinding:
           return true;
         default:
           return false;
         }
+      } else if (auto P = Node.getAsPattern()) {
+        switch (P->getKind()) {
+        case PatternKind::Expr:
+          return true;
+        default:
+          return false;
+        }
       } else
         return false;
-  }) {}
+    }) {}
 
   void analyzeExpr(Expr *Parent, llvm::function_ref<void(Type)> Callback,
                    SmallVectorImpl<StringRef> &PossibleNames) {
     switch (Parent->getKind()) {
-      case ExprKind::Call: {
+      case ExprKind::Call:
+      case ExprKind::Binary:
+      case ExprKind::PrefixUnary: {
         std::vector<Type> PotentialTypes;
         std::vector<StringRef> ExpectedNames;
         CompletionLookup::collectArgumentExpectation(
-            *DC, cast<CallExpr>(Parent), ParsedExpr, PotentialTypes,
+            *DC, cast<ApplyExpr>(Parent), ParsedExpr, PotentialTypes,
             ExpectedNames);
         for (Type Ty : PotentialTypes)
           Callback(Ty);
@@ -5047,7 +5073,7 @@ class CodeCompletionTypeContextAnalyzer {
         break;
       }
       default:
-        llvm_unreachable("Unhandled expression kinds.");
+        llvm_unreachable("Unhandled expression kind.");
     }
   }
 
@@ -5072,7 +5098,7 @@ class CodeCompletionTypeContextAnalyzer {
       }
       break;
     default:
-      llvm_unreachable("Unhandled statement kinds.");
+      llvm_unreachable("Unhandled statement kind.");
     }
   }
 
@@ -5114,7 +5140,22 @@ class CodeCompletionTypeContextAnalyzer {
         break;
       }
       default:
-        llvm_unreachable("Unhandled decl kinds.");
+        llvm_unreachable("Unhandled decl kind.");
+    }
+  }
+
+  void analyzePattern(Pattern *P, llvm::function_ref<void(Type)> Callback) {
+    switch (P->getKind()) {
+    case PatternKind::Expr: {
+      auto ExprPat = cast<ExprPattern>(P);
+      if (auto D = ExprPat->getMatchVar()) {
+        if (D->hasInterfaceType())
+          Callback(D->getInterfaceType());
+      }
+      break;
+    }
+    default:
+      llvm_unreachable("Unhandled pattern kind.");
     }
   }
 
@@ -5136,12 +5177,14 @@ class CodeCompletionTypeContextAnalyzer {
 
     for (auto It = Finder.Ancestors.rbegin(); It != Finder.Ancestors.rend();
          ++ It) {
-      if (auto Parent = It->dyn_cast<Expr *>()) {
+      if (auto Parent = It->getAsExpr()) {
         analyzeExpr(Parent, Callback, PossibleNames);
-      } else if (auto Parent = It->dyn_cast<Stmt *>()) {
+      } else if (auto Parent = It->getAsStmt()) {
         analyzeStmt(Parent, Callback);
-      } else if (auto Parent = It->dyn_cast<Decl *>()) {
+      } else if (auto Parent = It->getAsDecl()) {
         analyzeDecl(Parent, Callback);
+      } else if (auto Parent = It->getAsPattern()) {
+        analyzePattern(Parent, Callback);
       }
       if (!PossibleTypes.empty() || !PossibleNames.empty())
         return true;
@@ -5417,12 +5460,12 @@ void CodeCompletionCallbacksImpl::doneParsing() {
   case CompletionKind::UnresolvedMember : {
     Lookup.setHaveDot(SourceLoc());
     SmallVector<Type, 1> PossibleTypes;
-    ExprParentFinder Walker(UnresolvedExpr, [&](ASTNode Node) {
-      return Node.is<Expr *>();
+    ExprParentFinder Walker(UnresolvedExpr, [&](ASTWalker::ParentTy Node) {
+      return Node.getAsExpr();
     });
     CurDeclContext->walkContext(Walker);
     bool Success = false;
-    if (auto PE = Walker.ParentFarthest.get<Expr *>()) {
+    if (auto PE = Walker.ParentFarthest.getAsExpr()) {
       prepareForRetypechecking(PE);
       Success = typeCheckUnresolvedExpr(*CurDeclContext, UnresolvedExpr, PE,
                                         PossibleTypes);
@@ -5452,11 +5495,7 @@ void CodeCompletionCallbacksImpl::doneParsing() {
 
   case CompletionKind::ReturnStmtExpr : {
     SourceLoc Loc = P.Context.SourceMgr.getCodeCompletionLoc();
-    if (auto FD = dyn_cast<AbstractFunctionDecl>(CurDeclContext)) {
-      if (auto FT = FD->getInterfaceType()->getAs<FunctionType>()) {
-        Lookup.setExpectedTypes(FT->getResult());
-      }
-    }
+    Lookup.setExpectedTypes(getReturnTypeFromContext(CurDeclContext));
     Lookup.getValueCompletionsInDeclContext(Loc);
     break;
   }

lib/Parse/ParseExpr.cpp

@@ -532,10 +532,9 @@ ParserResult<Expr> Parser::parseExprUnary(Diag<> Message, bool isExprBasic) {
   }
 
   ParserResult<Expr> SubExpr = parseExprUnary(Message, isExprBasic);
-  if (SubExpr.hasCodeCompletion())
-    return makeParserCodeCompletionResult<Expr>();
+  ParserStatus Status = SubExpr;
   if (SubExpr.isNull())
-    return nullptr;
+    return Status;
 
   // We are sure we can create a prefix prefix operator expr now.
   UnaryContext.setCreateSyntax(SyntaxKind::PrefixOperatorExpr);
@@ -545,12 +544,13 @@ ParserResult<Expr> Parser::parseExprUnary(Diag<> Message, bool isExprBasic) {
   if (auto *LE = dyn_cast<NumberLiteralExpr>(SubExpr.get())) {
     if (Operator->hasName() && Operator->getName().getBaseName() == "-") {
       LE->setNegative(Operator->getLoc());
-      return makeParserResult(LE);
+      return makeParserResult(Status, LE);
     }
   }
 
-  return makeParserResult(new (Context) PrefixUnaryExpr(
-      Operator, formUnaryArgument(Context, SubExpr.get())));
+  return makeParserResult(
+      Status, new (Context) PrefixUnaryExpr(
+                  Operator, formUnaryArgument(Context, SubExpr.get())));
 }
 
 /// expr-keypath-swift:
@@ -1667,9 +1667,6 @@ ParserResult<Expr> Parser::parseExprPrimary(Diag<> ID, bool isExprBasic) {
                                           rParenLoc,
                                           trailingClosure,
                                           SyntaxKind::FunctionCallArgumentList);
-      if (status.isError())
-        return nullptr;
-
       SyntaxContext->createNodeInPlace(SyntaxKind::FunctionCallExpr);
       return makeParserResult(
                  status,

lib/Parse/ParsePattern.cpp

@@ -1109,10 +1109,9 @@ ParserResult<Pattern> Parser::parseMatchingPattern(bool isExprBasic) {
   // disambiguate.
   ParserResult<Expr> subExpr =
     parseExprImpl(diag::expected_pattern, isExprBasic);
-  if (subExpr.hasCodeCompletion())
-    return makeParserCodeCompletionStatus();
+  ParserStatus status = subExpr;
   if (subExpr.isNull())
-    return nullptr;
+    return status;
 
   if (SyntaxContext->isEnabled()) {
     if (auto UPES = PatternCtx.popIf<UnresolvedPatternExprSyntax>()) {
@@ -1125,9 +1124,9 @@ ParserResult<Pattern> Parser::parseMatchingPattern(bool isExprBasic) {
   // obvious pattern, which will come back wrapped in an immediate
   // UnresolvedPatternExpr.  Transform this now to simplify later code.
   if (auto *UPE = dyn_cast<UnresolvedPatternExpr>(subExpr.get()))
-    return makeParserResult(UPE->getSubPattern());
+    return makeParserResult(status, UPE->getSubPattern());
 
-  return makeParserResult(new (Context) ExprPattern(subExpr.get()));
+  return makeParserResult(status, new (Context) ExprPattern(subExpr.get()));
 }
 
 ParserResult<Pattern> Parser::parseMatchingPatternAsLetOrVar(bool isLet,

lib/Sema/CSGen.cpp

@@ -1190,9 +1190,6 @@ namespace {
     }
 
     virtual Type visitCodeCompletionExpr(CodeCompletionExpr *E) {
-      if (!E->isActivated())
-        return Type();
-
       CS.Options |= ConstraintSystemFlags::SuppressDiagnostics;
       return CS.createTypeVariable(CS.getConstraintLocator(E),
                                    TVO_CanBindToLValue);

lib/Sema/TypeCheckConstraints.cpp

@@ -2231,9 +2231,6 @@ bool TypeChecker::typeCheckCompletionSequence(Expr *&expr, DeclContext *DC) {
   // Ensure the output expression is up to date.
   assert(exprAsBinOp == expr && isa<BinaryExpr>(expr) && "found wrong expr?");
 
-  // Add type variable for the code-completion expression.
-  CCE->setActivated();
-
   if (auto generated = CS.generateConstraints(expr)) {
     expr = generated;
   } else {

test/IDE/complete_dynamic_lookup.swift

@@ -463,7 +463,7 @@ func testAnyObject11_(_ dl: AnyObject) {
   dl.returnsObjcClass!(#^DL_FUNC_NAME_PAREN_1^#
 }
 // DL_FUNC_NAME_PAREN_1: Begin completions
-// DL_FUNC_NAME_PAREN_1-DAG: Pattern/CurrModule: ['(']{#Int#}[')'][#TopLevelObjcClass#]{{; name=.+$}}
+// DL_FUNC_NAME_PAREN_1-DAG: Pattern/CurrModule: ['(']{#(i): Int#}[')'][#TopLevelObjcClass#]{{; name=.+$}}
 // DL_FUNC_NAME_PAREN_1: End completions
 
 func testAnyObject12(_ dl: AnyObject) {