[CodeCompletion] Migrate conforming methods list to solver-based

Author: ahoppen

lib/IDE/ConformingMethodList.cpp

@@ -16,8 +16,10 @@
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/USRGeneration.h"
+#include "swift/IDE/TypeCheckCompletionCallback.h"
 #include "swift/Parse/IDEInspectionCallbacks.h"
 #include "swift/Sema/IDETypeChecking.h"
+#include "swift/Sema/ConstraintSystem.h"
 #include "clang/AST/Attr.h"
 #include "clang/AST/Decl.h"
 
@@ -30,7 +32,7 @@ class ConformingMethodListCallbacks : public CodeCompletionCallbacks,
   ArrayRef<const char *> ExpectedTypeNames;
   ConformingMethodListConsumer &Consumer;
   SourceLoc Loc;
-  Expr *ParsedExpr = nullptr;
+  CodeCompletionExpr *CCExpr = nullptr;
   DeclContext *CurDeclContext = nullptr;
 
   void getMatchingMethods(Type T,
@@ -56,34 +58,67 @@ class ConformingMethodListCallbacks : public CodeCompletionCallbacks,
 void ConformingMethodListCallbacks::completeDotExpr(CodeCompletionExpr *E,
                                                     SourceLoc DotLoc) {
   CurDeclContext = P.CurDeclContext;
-  ParsedExpr = E->getBase();
+  CCExpr = E;
 }
 
 void ConformingMethodListCallbacks::completePostfixExpr(CodeCompletionExpr *E,
                                                         bool hasSpace) {
   CurDeclContext = P.CurDeclContext;
-  ParsedExpr = E->getBase();
+  CCExpr = E;
 }
 
-void ConformingMethodListCallbacks::doneParsing(SourceFile *SrcFile) {
-  if (!ParsedExpr)
-    return;
+class ConformingMethodListCallback : public TypeCheckCompletionCallback {
+public:
+  struct Result {
+    Type Ty;
 
-  typeCheckContextAt(TypeCheckASTNodeAtLocContext::declContext(CurDeclContext),
-                     ParsedExpr->getLoc());
+    /// Types of variables that were determined in the solution that produced
+    /// this result. This in particular includes parameters of closures that
+    /// were type-checked with the code completion expression.
+    llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes;
+  };
+private:
+  CodeCompletionExpr *CCExpr;
 
-  Type T = ParsedExpr->getType();
+  SmallVector<Result> Results;
 
-  // Type check the expression if needed.
-  if (!T || T->is<ErrorType>()) {
-    ConcreteDeclRef ReferencedDecl = nullptr;
-    auto optT = getTypeOfCompletionContextExpr(P.Context, CurDeclContext,
-                                               CompletionTypeCheckKind::Normal,
-                                               ParsedExpr, ReferencedDecl);
-    if (!optT)
+  void sawSolutionImpl(const constraints::Solution &S) override {
+    if (!S.hasType(CCExpr->getBase())) {
       return;
-    T = *optT;
+    }
+    if (Type T = getTypeForCompletion(S, CCExpr->getBase())) {
+      llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes;
+      getSolutionSpecificVarTypes(S, SolutionSpecificVarTypes);
+      Results.push_back({T, SolutionSpecificVarTypes});
+    }
+  }
+
+public:
+  ConformingMethodListCallback(CodeCompletionExpr *CCExpr) : CCExpr(CCExpr) {}
+
+  ArrayRef<Result> getResults() const { return Results; }
+};
+
+void ConformingMethodListCallbacks::doneParsing(SourceFile *SrcFile) {
+  if (!CCExpr || !CCExpr->getBase())
+    return;
+
+  ConformingMethodListCallback TypeCheckCallback(CCExpr);
+  {
+    llvm::SaveAndRestore<TypeCheckCompletionCallback *> CompletionCollector(
+        Context.CompletionCallback, &TypeCheckCallback);
+    typeCheckContextAt(
+        TypeCheckASTNodeAtLocContext::declContext(CurDeclContext),
+        CCExpr->getLoc());
+  }
+
+  if (TypeCheckCallback.getResults().size() != 1) {
+    // Either no results or results were ambiguous, which we cannot handle.
+    return;
   }
+  auto Res = TypeCheckCallback.getResults()[0];
+  Type T = Res.Ty;
+  WithSolutionSpecificVarTypesRAII VarType(Res.SolutionSpecificVarTypes);
 
   if (!T || T->is<ErrorType>() || T->is<UnresolvedType>())
     return;

lib/Sema/TypeCheckStmt.cpp

@@ -2586,8 +2586,8 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
       auto optBody = TypeChecker::applyResultBuilderBodyTransform(
           func, builderType,
           /*ClosuresInResultBuilderDontParticipateInInference=*/
-              ctx.CompletionCallback == nullptr && ctx.SolutionCallback == nullptr);
-      if (ctx.CompletionCallback && ctx.CompletionCallback->gotCallback()) {
+          ctx.CompletionCallback == nullptr && ctx.SolutionCallback == nullptr);
+      if ((ctx.CompletionCallback && ctx.CompletionCallback->gotCallback()) || (ctx.SolutionCallback && ctx.SolutionCallback->gotCallback())) {
         // We already informed the completion callback of solutions found by
         // type checking the entire result builder from
         // applyResultBuilderBodyTransform. No need to typecheck the requested

validation-test/IDE/crashers_2_fixed/0036-conformingmethodlist-already-typechecked.swift

@@ -0,0 +1,17 @@
+// RUN: %target-swift-ide-test -conforming-methods -source-filename %s -code-completion-token COMPLETE
+
+// This test used to crash while PostfixExpr completion was migrated to solver-based.
+
+struct MyPublisher {
+  func removeDuplicates() {}
+}
+
+func handleEvents(receiveOutput: ((String) -> Void)? = nil) -> MyPublisher {}
+
+protocol AnyCancellable {}
+
+class CategoriesSearchViewModel {
+  func foo() {
+    var searchCancellable: AnyCancellable = handleEvents(receiveOutput: { [weak self] _ in }).removeDuplicates #^COMPLETE^#
+  }
+}