Merge pull request #41886 from ahoppen/pr/migrate-foreachsequence-postfixexprbeginning-solver-based

[CodeCompletion] Migrate ForEachSequence and PostfixExprBeginning to solver-based

Author: ahoppen

include/swift/IDE/ExprCompletion.h

@@ -23,9 +23,6 @@ namespace ide {
 class ExprTypeCheckCompletionCallback : public TypeCheckCompletionCallback {
 public:
   struct Result {
-    /// The contextual type that the code completion expression should produce.
-    Type ExpectedType;
-
     /// If the code completion expression is an implicit return in a
     /// single-expression closure.
     bool IsImplicitSingleExpressionReturn;
@@ -38,14 +35,37 @@ class ExprTypeCheckCompletionCallback : public TypeCheckCompletionCallback {
     /// this result. This in particular includes parameters of closures that
     /// were type-checked with the code completion expression.
     llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes;
+
+    bool operator==(const Result &Other) const;
   };
 
 private:
   CodeCompletionExpr *CompletionExpr;
   DeclContext *DC;
 
+  /// The contextual types to which the code completion results should be
+  /// convertible.
+  /// Technically, each result should have its own expected type because some
+  /// expected types may only be available e.g. for certain
+  /// \c SolutionSpecificVarTypes. But that means that we need to do a separate
+  /// completion lookup for each expected type and de-duplicate the results,
+  /// which can have huge performance implications (>5mins instead of <2secs).
+  /// In practice sharing ExpectedTypes between results yields identical results
+  /// in almost all cases and acceptable results in the other cases.
+  SmallVector<Type, 4> ExpectedTypes;
+
   SmallVector<Result, 4> Results;
 
+  /// Adds the given type to \c ExpectedTypes unless \c ExpectedTypes already
+  /// contains the type.
+  void addExpectedType(Type ExpectedType);
+
+  /// Adds the result with the given parameters to \c Results unless \c Results
+  /// already contains an entry with exactly the same values.
+  void addResult(
+      bool IsImplicitSingleExpressionReturn, bool IsInAsyncContext,
+      llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes);
+
   void sawSolutionImpl(const constraints::Solution &solution) override;
 
 public:

include/swift/IDE/TypeCheckCompletionCallback.h

@@ -67,7 +67,13 @@ class TypeCheckCompletionCallback {
 
 // MARK: - Utility functions for subclasses of TypeCheckCompletionCallback
 
-Type getTypeForCompletion(const constraints::Solution &S, Expr *E);
+Type getTypeForCompletion(const constraints::Solution &S, ASTNode Node);
+
+/// If \p E occurs in a pattern matching position, returns the type that it is
+/// being pattern-matched against.
+/// If that type is an enum, it allows us to suggest the enum cases for the code
+/// completion expression \p E.
+Type getPatternMatchType(const constraints::Solution &S, Expr *E);
 
 /// Whether the given completion expression is the only expression in its
 /// containing closure or function body and its value is implicitly returned.
@@ -84,6 +90,9 @@ bool isImplicitSingleExpressionReturn(constraints::ConstraintSystem &CS,
 /// Returns \c true iff the decl context \p DC allows calling async functions.
 bool isContextAsync(const constraints::Solution &S, DeclContext *DC);
 
+/// Returns true if both types are null or if they are equal.
+bool nullableTypesEqual(Type LHS, Type RHS);
+
 } // namespace ide
 } // namespace swift
 

lib/AST/Type.cpp

@@ -4308,11 +4308,11 @@ operator()(CanType dependentType, Type conformingReplacementType,
 ProtocolConformanceRef MakeAbstractConformanceForGenericType::
 operator()(CanType dependentType, Type conformingReplacementType,
            ProtocolDecl *conformedProtocol) const {
-  assert((conformingReplacementType->is<ErrorType>()
-          || conformingReplacementType->is<SubstitutableType>()
-          || conformingReplacementType->is<DependentMemberType>()
-          || conformingReplacementType->is<TypeVariableType>())
-         && "replacement requires looking up a concrete conformance");
+  assert((conformingReplacementType->is<ErrorType>() ||
+          conformingReplacementType->is<SubstitutableType>() ||
+          conformingReplacementType->is<DependentMemberType>() ||
+          conformingReplacementType->hasTypeVariable()) &&
+         "replacement requires looking up a concrete conformance");
   // A class-constrained archetype might conform to the protocol
   // concretely.
   if (auto *archetypeType = conformingReplacementType->getAs<ArchetypeType>()) {

lib/IDE/ArgumentCompletion.cpp

@@ -20,18 +20,6 @@ using namespace swift;
 using namespace swift::ide;
 using namespace swift::constraints;
 
-/// Returns true if both types are null or if they are equal.
-static bool nullableTypesEqual(Type LHS, Type RHS) {
-  if (LHS.isNull() && RHS.isNull()) {
-    return true;
-  } else if (LHS.isNull() || RHS.isNull()) {
-    // One type is null but the other is not.
-    return false;
-  } else {
-    return LHS->isEqual(RHS);
-  }
-}
-
 bool ArgumentTypeCheckCompletionCallback::addPossibleParams(
     const ArgumentTypeCheckCompletionCallback::Result &Res,
     SmallVectorImpl<PossibleParamInfo> &Params, SmallVectorImpl<Type> &Types) {

lib/IDE/CodeCompletion.cpp

@@ -1396,7 +1396,9 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
                           CurDeclContext, CompletionContext, Consumer);
     return true;
   }
-  case CompletionKind::StmtOrExpr: {
+  case CompletionKind::StmtOrExpr:
+  case CompletionKind::ForEachSequence:
+  case CompletionKind::PostfixExprBeginning: {
     assert(CodeCompleteTokenExpr);
     assert(CurDeclContext);
 
@@ -1539,18 +1541,11 @@ void CodeCompletionCallbacksImpl::doneParsing() {
   case CompletionKind::KeyPathExprSwift:
   case CompletionKind::CallArg:
   case CompletionKind::StmtOrExpr:
+  case CompletionKind::ForEachSequence:
+  case CompletionKind::PostfixExprBeginning:
     llvm_unreachable("should be already handled");
     return;
 
-  case CompletionKind::ForEachSequence:
-  case CompletionKind::PostfixExprBeginning: {
-    ExprContextInfo ContextInfo(CurDeclContext, CodeCompleteTokenExpr);
-    Lookup.setExpectedTypes(ContextInfo.getPossibleTypes(),
-                            ContextInfo.isImplicitSingleExpressionReturn());
-    DoPostfixExprBeginning();
-    break;
-  }
-
   case CompletionKind::PostfixExpr: {
     Lookup.setHaveLeadingSpace(HasSpace);
     if (isDynamicLookup(*ExprType))

lib/IDE/ExprCompletion.cpp

@@ -19,6 +19,54 @@ using namespace swift;
 using namespace swift::ide;
 using namespace swift::constraints;
 
+static bool solutionSpecificVarTypesEqual(
+    const llvm::SmallDenseMap<const VarDecl *, Type> &LHS,
+    const llvm::SmallDenseMap<const VarDecl *, Type> &RHS) {
+  if (LHS.size() != RHS.size()) {
+    return false;
+  }
+  for (auto LHSEntry : LHS) {
+    auto RHSEntry = RHS.find(LHSEntry.first);
+    if (RHSEntry == RHS.end()) {
+      // Entry of the LHS doesn't exist in RHS
+      return false;
+    } else if (!nullableTypesEqual(LHSEntry.second, RHSEntry->second)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool ExprTypeCheckCompletionCallback::Result::operator==(
+    const Result &Other) const {
+  return IsImplicitSingleExpressionReturn ==
+             Other.IsImplicitSingleExpressionReturn &&
+         IsInAsyncContext == Other.IsInAsyncContext &&
+         solutionSpecificVarTypesEqual(SolutionSpecificVarTypes,
+                                       Other.SolutionSpecificVarTypes);
+}
+
+void ExprTypeCheckCompletionCallback::addExpectedType(Type ExpectedType) {
+  auto IsEqual = [&ExpectedType](Type Other) {
+    return nullableTypesEqual(ExpectedType, Other);
+  };
+  if (llvm::any_of(ExpectedTypes, IsEqual)) {
+    return;
+  }
+  ExpectedTypes.push_back(ExpectedType);
+}
+
+void ExprTypeCheckCompletionCallback::addResult(
+    bool IsImplicitSingleExpressionReturn, bool IsInAsyncContext,
+    llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes) {
+  Result NewResult = {IsImplicitSingleExpressionReturn, IsInAsyncContext,
+                      SolutionSpecificVarTypes};
+  if (llvm::is_contained(Results, NewResult)) {
+    return;
+  }
+  Results.push_back(NewResult);
+}
+
 void ExprTypeCheckCompletionCallback::sawSolutionImpl(
     const constraints::Solution &S) {
   auto &CS = S.getConstraintSystem();
@@ -36,8 +84,12 @@ void ExprTypeCheckCompletionCallback::sawSolutionImpl(
     }
   }
 
-  Results.push_back(
-      {ExpectedTy, ImplicitReturn, IsAsync, SolutionSpecificVarTypes});
+  addResult(ImplicitReturn, IsAsync, SolutionSpecificVarTypes);
+  addExpectedType(ExpectedTy);
+
+  if (auto PatternMatchType = getPatternMatchType(S, CompletionExpr)) {
+    addExpectedType(PatternMatchType);
+  }
 }
 
 void ExprTypeCheckCompletionCallback::deliverResults(
@@ -46,9 +98,10 @@ void ExprTypeCheckCompletionCallback::deliverResults(
   ASTContext &Ctx = DC->getASTContext();
   CompletionLookup Lookup(CompletionCtx.getResultSink(), Ctx, DC,
                           &CompletionCtx);
+  Lookup.shouldCheckForDuplicates(Results.size() > 1);
 
   for (auto &Result : Results) {
-    Lookup.setExpectedTypes(Result.ExpectedType,
+    Lookup.setExpectedTypes(ExpectedTypes,
                             Result.IsImplicitSingleExpressionReturn);
     Lookup.setCanCurrDeclContextHandleAsync(Result.IsInAsyncContext);
     Lookup.setSolutionSpecificVarTypes(Result.SolutionSpecificVarTypes);

lib/IDE/ExprContextAnalysis.cpp

@@ -107,6 +107,10 @@ void swift::ide::typeCheckContextAt(DeclContext *DC, SourceLoc Loc) {
         auto *Param = AFD->getParameters()->get(defaultArg->getIndex());
         (void)Param->getTypeCheckedDefaultExpr();
       }
+      if (auto *SD = dyn_cast<SubscriptDecl>(defaultArg->getParent())) {
+        auto *Param = SD->getIndices()->get(defaultArg->getIndex());
+        (void)Param->getTypeCheckedDefaultExpr();
+      }
     }
     break;
 

lib/IDE/TypeCheckCompletionCallback.cpp

@@ -40,8 +40,9 @@ void TypeCheckCompletionCallback::fallbackTypeCheck(DeclContext *DC) {
 
 // MARK: - Utility functions for subclasses of TypeCheckCompletionCallback
 
-Type swift::ide::getTypeForCompletion(const constraints::Solution &S, Expr *E) {
-  if (!S.hasType(E)) {
+Type swift::ide::getTypeForCompletion(const constraints::Solution &S,
+                                      ASTNode Node) {
+  if (!S.hasType(Node)) {
     assert(false && "Expression wasn't type checked?");
     return nullptr;
   }
@@ -66,8 +67,8 @@ Type swift::ide::getTypeForCompletion(const constraints::Solution &S, Expr *E) {
   // expression should still have a type of `[T]` instead of `[<<hole>>]`
   // because it helps to produce correct contextual member list based on
   // a conformance requirement associated with generic parameter `T`.
-  if (isa<CodeCompletionExpr>(E)) {
-    auto completionTy = S.getType(E).transform([&](Type type) -> Type {
+  if (isExpr<CodeCompletionExpr>(Node)) {
+    auto completionTy = S.getType(Node).transform([&](Type type) -> Type {
       if (auto *typeVar = type->getAs<TypeVariableType>())
         return S.getFixedType(typeVar);
       return type;
@@ -92,18 +93,77 @@ Type swift::ide::getTypeForCompletion(const constraints::Solution &S, Expr *E) {
       return type;
     }));
   } else {
-    Result = S.getResolvedType(E);
+    Result = S.getResolvedType(Node);
   }
 
   if (!Result || Result->is<UnresolvedType>()) {
-    Result = CS.getContextualType(E, /*forConstraint=*/false);
+    Result = CS.getContextualType(Node, /*forConstraint=*/false);
   }
   if (Result && Result->is<UnresolvedType>()) {
     Result = Type();
   }
   return Result;
 }
 
+/// If the code completion expression \p E occurs in a pattern matching
+/// position, we have an AST that looks like this.
+/// \code
+/// (binary_expr implicit type='$T3'
+///   (overloaded_decl_ref_expr function_ref=compound decls=[
+///     Swift.(file).~=,
+///     Swift.(file).Optional extension.~=])
+///   (argument_list implicit
+///     (argument
+///       (code_completion_expr implicit type='$T1'))
+///     (argument
+///       (declref_expr implicit decl=swift_ide_test.(file).foo(x:).$match))))
+/// \endcode
+/// If the code completion expression occurs in such an AST, return the
+/// declaration of the \c $match variable, otherwise return \c nullptr.
+static VarDecl *getMatchVarIfInPatternMatch(Expr *E, ConstraintSystem &CS) {
+  auto &Context = CS.getASTContext();
+
+  auto *Binary = dyn_cast_or_null<BinaryExpr>(CS.getParentExpr(E));
+  if (!Binary || !Binary->isImplicit() || Binary->getLHS() != E) {
+    return nullptr;
+  }
+
+  auto CalledOperator = Binary->getFn();
+  if (!isPatternMatchingOperator(CalledOperator)) {
+    return nullptr;
+  }
+
+  auto MatchArg = dyn_cast_or_null<DeclRefExpr>(Binary->getRHS());
+  if (!MatchArg || !MatchArg->isImplicit()) {
+    return nullptr;
+  }
+
+  auto MatchVar = MatchArg->getDecl();
+  if (MatchVar && MatchVar->isImplicit() &&
+      MatchVar->getBaseName() == Context.Id_PatternMatchVar) {
+    return dyn_cast<VarDecl>(MatchVar);
+  } else {
+    return nullptr;
+  }
+}
+
+Type swift::ide::getPatternMatchType(const constraints::Solution &S, Expr *E) {
+  if (auto MatchVar = getMatchVarIfInPatternMatch(E, S.getConstraintSystem())) {
+    Type MatchVarType;
+    // If the MatchVar has an explicit type, it's not part of the solution. But
+    // we can look it up in the constraint system directly.
+    if (auto T = S.getConstraintSystem().getVarType(MatchVar)) {
+      MatchVarType = T;
+    } else {
+      MatchVarType = getTypeForCompletion(S, MatchVar);
+    }
+    if (MatchVarType) {
+      return MatchVarType;
+    }
+  }
+  return nullptr;
+}
+
 bool swift::ide::isImplicitSingleExpressionReturn(ConstraintSystem &CS,
                                                   Expr *CompletionExpr) {
   Expr *ParentExpr = CS.getParentExpr(CompletionExpr);
@@ -148,3 +208,14 @@ bool swift::ide::isContextAsync(const constraints::Solution &S,
   //    async.
   return canDeclContextHandleAsync(DC);
 }
+
+bool swift::ide::nullableTypesEqual(Type LHS, Type RHS) {
+  if (LHS.isNull() && RHS.isNull()) {
+    return true;
+  } else if (LHS.isNull() || RHS.isNull()) {
+    // One type is null but the other is not.
+    return false;
+  } else {
+    return LHS->isEqual(RHS);
+  }
+}