[CodeCompletion] Don't record solution results twice in ExprCompletion

This fixes a performance problem where we were sometimes reporting >1000 solutions but all the relevant parameters of the result were equal.

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

@@ -90,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/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/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,12 +84,11 @@ void ExprTypeCheckCompletionCallback::sawSolutionImpl(
     }
   }
 
-  Results.push_back(
-      {ExpectedTy, ImplicitReturn, IsAsync, SolutionSpecificVarTypes});
+  addResult(ImplicitReturn, IsAsync, SolutionSpecificVarTypes);
+  addExpectedType(ExpectedTy);
 
   if (auto PatternMatchType = getPatternMatchType(S, CompletionExpr)) {
-    Results.push_back(
-        {PatternMatchType, ImplicitReturn, IsAsync, SolutionSpecificVarTypes});
+    addExpectedType(PatternMatchType);
   }
 }
 
@@ -53,10 +100,6 @@ void ExprTypeCheckCompletionCallback::deliverResults(
                           &CompletionCtx);
   Lookup.shouldCheckForDuplicates(Results.size() > 1);
 
-  SmallVector<Type, 2> ExpectedTypes;
-  for (auto &Result : Results) {
-    ExpectedTypes.push_back(Result.ExpectedType);
-  }
   for (auto &Result : Results) {
     Lookup.setExpectedTypes(ExpectedTypes,
                             Result.IsImplicitSingleExpressionReturn);

lib/IDE/TypeCheckCompletionCallback.cpp

@@ -208,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);
+  }
+}