Merge pull request #41948 from ahoppen/pr/migrate-afterpoundexpr-solver-based

[CodeCompletion] Migrate AfterPoundExprCompletion to solver-based

Author: ahoppen

include/swift/IDE/AfterPoundExprCompletion.h

@@ -0,0 +1,57 @@
+//===--- AfterPoundExprCompletion.h ---------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_IDE_AFTERPOUNDEXPRCOMPLETION_H
+#define SWIFT_IDE_AFTERPOUNDEXPRCOMPLETION_H
+
+#include "swift/IDE/CodeCompletionConsumer.h"
+#include "swift/IDE/CodeCompletionContext.h"
+#include "swift/IDE/TypeCheckCompletionCallback.h"
+
+namespace swift {
+namespace ide {
+
+/// Used to collect and store information needed to perform unresolved member
+/// completion (\c CompletionKind::UnresolvedMember ) from the solutions
+/// formed during expression type-checking.
+class AfterPoundExprCompletion : public TypeCheckCompletionCallback {
+  struct Result {
+    Type ExpectedTy;
+    bool IsImplicitSingleExpressionReturn;
+
+    /// Whether the surrounding context is async and thus calling async
+    /// functions is supported.
+    bool IsInAsyncContext;
+  };
+
+  CodeCompletionExpr *CompletionExpr;
+  DeclContext *DC;
+  Optional<StmtKind> ParentStmtKind;
+
+  SmallVector<Result, 4> Results;
+
+  void sawSolutionImpl(const constraints::Solution &solution) override;
+
+public:
+  AfterPoundExprCompletion(CodeCompletionExpr *CompletionExpr, DeclContext *DC,
+                           Optional<StmtKind> ParentStmtKind)
+      : CompletionExpr(CompletionExpr), DC(DC), ParentStmtKind(ParentStmtKind) {
+  }
+
+  void deliverResults(ide::CodeCompletionContext &CompletionCtx,
+                      CodeCompletionConsumer &Consumer);
+};
+
+} // end namespace ide
+} // end namespace swift
+
+#endif // SWIFT_IDE_AFTERPOUNDEXPRCOMPLETION_H

lib/IDE/AfterPoundExprCompletion.cpp

@@ -0,0 +1,59 @@
+//===--- AfterPoundExprCompletion.cpp -------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "swift/IDE/AfterPoundExprCompletion.h"
+#include "swift/IDE/CodeCompletion.h"
+#include "swift/IDE/CompletionLookup.h"
+#include "swift/Sema/CompletionContextFinder.h"
+#include "swift/Sema/ConstraintSystem.h"
+#include "swift/Sema/IDETypeChecking.h"
+
+using namespace swift;
+using namespace swift::constraints;
+using namespace swift::ide;
+
+void AfterPoundExprCompletion::sawSolutionImpl(const constraints::Solution &S) {
+  auto &CS = S.getConstraintSystem();
+  Type ExpectedTy = getTypeForCompletion(S, CompletionExpr);
+
+  bool IsAsync = isContextAsync(S, DC);
+
+  // If ExpectedTy is a duplicate of any other result, ignore this solution.
+  auto IsEqual = [&](const Result &R) {
+    return R.ExpectedTy->isEqual(ExpectedTy);
+  };
+  if (!llvm::any_of(Results, IsEqual)) {
+    bool SingleExprBody = isImplicitSingleExpressionReturn(CS, CompletionExpr);
+    Results.push_back({ExpectedTy, SingleExprBody, IsAsync});
+  }
+}
+
+void AfterPoundExprCompletion::deliverResults(
+    ide::CodeCompletionContext &CompletionCtx,
+    CodeCompletionConsumer &Consumer) {
+  ASTContext &Ctx = DC->getASTContext();
+  CompletionLookup Lookup(CompletionCtx.getResultSink(), Ctx, DC,
+                          &CompletionCtx);
+
+  Lookup.shouldCheckForDuplicates(Results.size() > 1);
+
+  for (auto &Result : Results) {
+    Lookup.setExpectedTypes({Result.ExpectedTy},
+                            Result.IsImplicitSingleExpressionReturn,
+                            /*expectsNonVoid=*/true);
+    Lookup.addPoundAvailable(ParentStmtKind);
+    Lookup.addPoundLiteralCompletions(/*needPound=*/false);
+    Lookup.addObjCPoundKeywordCompletions(/*needPound=*/false);
+  }
+
+  deliverCompletionResults(CompletionCtx, Lookup, DC, Consumer);
+}

lib/IDE/CMakeLists.txt

@@ -1,5 +1,6 @@
 
 add_swift_host_library(swiftIDE STATIC
+  AfterPoundExprCompletion.cpp
   ArgumentCompletion.cpp
   CodeCompletion.cpp
   CodeCompletionCache.cpp

lib/IDE/CodeCompletion.cpp

@@ -32,6 +32,7 @@
 #include "swift/ClangImporter/ClangImporter.h"
 #include "swift/ClangImporter/ClangModule.h"
 #include "swift/Frontend/FrontendOptions.h"
+#include "swift/IDE/AfterPoundExprCompletion.h"
 #include "swift/IDE/ArgumentCompletion.h"
 #include "swift/IDE/CodeCompletionCache.h"
 #include "swift/IDE/CodeCompletionConsumer.h"
@@ -1418,6 +1419,25 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
     Lookup.deliverResults(CCLoc, CompletionContext, Consumer);
     return true;
   }
+  case CompletionKind::AfterPoundExpr: {
+    assert(CodeCompleteTokenExpr);
+    assert(CurDeclContext);
+
+    AfterPoundExprCompletion Lookup(CodeCompleteTokenExpr, CurDeclContext,
+                                    ParentStmtKind);
+    llvm::SaveAndRestore<TypeCheckCompletionCallback *> CompletionCollector(
+        Context.CompletionCallback, &Lookup);
+    typeCheckContextAt(CurDeclContext, CompletionLoc);
+
+    if (!Lookup.gotCallback()) {
+      Lookup.fallbackTypeCheck(CurDeclContext);
+    }
+
+    addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);
+
+    Lookup.deliverResults(CompletionContext, Consumer);
+    return true;
+  }
   default:
     return false;
   }
@@ -1543,6 +1563,7 @@ void CodeCompletionCallbacksImpl::doneParsing() {
   case CompletionKind::StmtOrExpr:
   case CompletionKind::ForEachSequence:
   case CompletionKind::PostfixExprBeginning:
+  case CompletionKind::AfterPoundExpr:
     llvm_unreachable("should be already handled");
     return;
 
@@ -1829,17 +1850,6 @@ void CodeCompletionCallbacksImpl::doneParsing() {
     break;
   }
 
-  case CompletionKind::AfterPoundExpr: {
-    ExprContextInfo ContextInfo(CurDeclContext, CodeCompleteTokenExpr);
-    Lookup.setExpectedTypes(ContextInfo.getPossibleTypes(),
-                            ContextInfo.isImplicitSingleExpressionReturn());
-
-    Lookup.addPoundAvailable(ParentStmtKind);
-    Lookup.addPoundLiteralCompletions(/*needPound=*/false);
-    Lookup.addObjCPoundKeywordCompletions(/*needPound=*/false);
-    break;
-  }
-
   case CompletionKind::AfterPoundDirective: {
     addPoundDirectives(CompletionContext.getResultSink());
     // FIXME: Add pound expressions (e.g. '#selector()') if it's at statements