[CodeCompletion] Fast completion for top-level code in single file script

e.g. Playground.
A single file script is like a single function body; the interface of
the file does not affect any other files.
So when a completion happens in a single file script, re-parse the whole
file. But we are still be able to reuse imported modules.

rdar://problem/58378157

Author: rintaro

include/swift/Frontend/Frontend.h

@@ -653,6 +653,9 @@ class CompilerInstance {
     explicit ImplicitImports(CompilerInstance &compiler);
   };
 
+  static void addAdditionalInitialImportsTo(
+    SourceFile *SF, const ImplicitImports &implicitImports);
+
 private:
   void addMainFileToModule(const ImplicitImports &implicitImports);
 

lib/Frontend/Frontend.cpp

@@ -640,7 +640,7 @@ ModuleDecl *CompilerInstance::getMainModule() const {
   return MainModule;
 }
 
-static void addAdditionalInitialImportsTo(
+void CompilerInstance::addAdditionalInitialImportsTo(
     SourceFile *SF, const CompilerInstance::ImplicitImports &implicitImports) {
   SmallVector<SourceFile::ImportedModuleDesc, 4> additionalImports;
 

lib/IDE/CompletionInstance.cpp

@@ -189,10 +189,7 @@ bool CompletionInstance::performCachedOperaitonIfPossible(
     return false;
 
   auto &oldInfo = oldState.getCodeCompletionDelayedDeclState();
-
-  // Currently, only completions within a function body is supported.
-  if (oldInfo.Kind != CodeCompletionDelayedDeclKind::FunctionBody)
-    return false;
+  auto *oldSF = oldInfo.ParentContext->getParentSourceFile();
 
   // Parse the new buffer into temporary SourceFile.
   SourceManager tmpSM;
@@ -204,94 +201,141 @@ bool CompletionInstance::performCachedOperaitonIfPossible(
   TypeCheckerOptions typeckOpts;
   SearchPathOptions searchPathOpts;
   DiagnosticEngine tmpDiags(tmpSM);
-  std::unique_ptr<ASTContext> Ctx(
+  std::unique_ptr<ASTContext> tmpCtx(
       ASTContext::get(langOpts, typeckOpts, searchPathOpts, tmpSM, tmpDiags));
-  registerIDERequestFunctions(Ctx->evaluator);
-  registerTypeCheckerRequestFunctions(Ctx->evaluator);
-  registerSILGenRequestFunctions(Ctx->evaluator);
-  ModuleDecl *M = ModuleDecl::create(Identifier(), *Ctx);
+  registerIDERequestFunctions(tmpCtx->evaluator);
+  registerTypeCheckerRequestFunctions(tmpCtx->evaluator);
+  registerSILGenRequestFunctions(tmpCtx->evaluator);
+  ModuleDecl *tmpM = ModuleDecl::create(Identifier(), *tmpCtx);
   PersistentParserState newState;
-  SourceFile *newSF =
-      new (*Ctx) SourceFile(*M, SourceFileKind::Library, tmpBufferID,
-                            SourceFile::ImplicitModuleImportKind::None);
-  newSF->enableInterfaceHash();
+  SourceFile *tmpSF =
+      new (*tmpCtx) SourceFile(*tmpM, oldSF->Kind, tmpBufferID,
+                               SourceFile::ImplicitModuleImportKind::None);
+  tmpSF->enableInterfaceHash();
   // Ensure all non-function-body tokens are hashed into the interface hash
-  Ctx->LangOpts.EnableTypeFingerprints = false;
-  parseIntoSourceFile(*newSF, tmpBufferID, &newState);
+  tmpCtx->LangOpts.EnableTypeFingerprints = false;
+  parseIntoSourceFile(*tmpSF, tmpBufferID, &newState);
   // Couldn't find any completion token?
   if (!newState.hasCodeCompletionDelayedDeclState())
     return false;
 
   auto &newInfo = newState.getCodeCompletionDelayedDeclState();
+  unsigned newBufferID;
+
+  switch (newInfo.Kind) {
+  case CodeCompletionDelayedDeclKind::FunctionBody: {
+    // If the interface has changed, AST must be refreshed.
+    llvm::SmallString<32> oldInterfaceHash{};
+    llvm::SmallString<32> newInterfaceHash{};
+    oldSF->getInterfaceHash(oldInterfaceHash);
+    tmpSF->getInterfaceHash(newInterfaceHash);
+    if (oldInterfaceHash != newInterfaceHash)
+      return false;
+
+    DeclContext *DC =
+        getEquivalentDeclContextFromSourceFile(newInfo.ParentContext, oldSF);
+    if (!DC)
+      return false;
+
+    // OK, we can perform fast completion for this. Update the orignal delayed
+    // decl state.
+
+    // Fast completion keeps the buffer in memory for multiple completions.
+    // To reduce the consumption, slice the source buffer so it only holds
+    // the portion that is needed for the second pass.
+    auto startOffset = newInfo.StartOffset;
+    if (newInfo.PrevOffset != ~0u)
+      startOffset = newInfo.PrevOffset;
+    auto startLoc = tmpSM.getLocForOffset(tmpBufferID, startOffset);
+    startLoc = Lexer::getLocForStartOfLine(tmpSM, startLoc);
+    startOffset = tmpSM.getLocOffsetInBuffer(startLoc, tmpBufferID);
+
+    auto endOffset = newInfo.EndOffset;
+    auto endLoc = tmpSM.getLocForOffset(tmpBufferID, endOffset);
+    endLoc = Lexer::getLocForEndOfToken(tmpSM, endLoc);
+    endOffset = tmpSM.getLocOffsetInBuffer(endLoc, tmpBufferID);
+
+    newInfo.StartOffset -= startOffset;
+    newInfo.EndOffset -= startOffset;
+    if (newInfo.PrevOffset != ~0u)
+      newInfo.PrevOffset -= startOffset;
+
+    auto sourceText =
+        completionBuffer->getBuffer().slice(startOffset, endOffset);
+    auto newOffset = Offset - startOffset;
+
+    newBufferID = SM.addMemBufferCopy(sourceText,
+                                      completionBuffer->getBufferIdentifier());
+    SM.openVirtualFile(SM.getLocForBufferStart(newBufferID),
+                       tmpSM.getDisplayNameForLoc(startLoc),
+                       tmpSM.getLineAndColumn(startLoc).first - 1);
+    SM.setCodeCompletionPoint(newBufferID, newOffset);
+
+    // Construct dummy scopes. We don't need to restore the original scope
+    // because they are probably not 'isResolvable()' anyway.
+    auto &SI = oldState.getScopeInfo();
+    assert(SI.getCurrentScope() == nullptr);
+    Scope Top(SI, ScopeKind::TopLevel);
+    Scope Body(SI, ScopeKind::FunctionBody);
+
+    oldInfo.ParentContext = DC;
+    oldInfo.StartOffset = newInfo.StartOffset;
+    oldInfo.EndOffset = newInfo.EndOffset;
+    oldInfo.PrevOffset = newInfo.PrevOffset;
+    oldState.restoreCodeCompletionDelayedDeclState(oldInfo);
+
+    auto *AFD = cast<AbstractFunctionDecl>(DC);
+    if (AFD->isBodySkipped())
+      AFD->setBodyDelayed(AFD->getBodySourceRange());
+
+    break;
+  }
+  case CodeCompletionDelayedDeclKind::Decl:
+  case CodeCompletionDelayedDeclKind::TopLevelCodeDecl: {
+    // Support decl/top-level code only if the completion happens in a single
+    // file 'main' script (e.g. playground).
+    auto *oldM = oldInfo.ParentContext->getParentModule();
+    if (oldM->getFiles().size() != 1 || oldSF->Kind != SourceFileKind::Main)
+      return false;
+
+    // Perform fast completion.
+
+    // Prepare the new buffer in the source manager.
+    auto sourceText = completionBuffer->getBuffer();
+    if (newInfo.Kind == CodeCompletionDelayedDeclKind::TopLevelCodeDecl) {
+      // We don't need the source text after the top-level code.
+      auto endOffset = newInfo.EndOffset;
+      auto endLoc = tmpSM.getLocForOffset(tmpBufferID, endOffset);
+      endLoc = Lexer::getLocForEndOfToken(tmpSM, endLoc);
+      endOffset = tmpSM.getLocOffsetInBuffer(endLoc, tmpBufferID);
+      sourceText = sourceText.slice(0, endOffset);
+    }
+    newBufferID = SM.addMemBufferCopy(sourceText,
+                                      completionBuffer->getBufferIdentifier());
+    SM.setCodeCompletionPoint(newBufferID, Offset);
+
+    // Create a new module and a source file using the current AST context.
+    auto &Ctx = oldM->getASTContext();
+    auto newM = ModuleDecl::create(oldM->getName(), Ctx);
+    CompilerInstance::ImplicitImports implicitImport(CI);
+    SourceFile *newSF = new (Ctx) SourceFile(*newM, SourceFileKind::Main,
+                                             newBufferID, implicitImport.kind);
+    newM->addFile(*newSF);
+    CompilerInstance::addAdditionalInitialImportsTo(newSF, implicitImport);
+    newSF->enableInterfaceHash();
+
+    // Re-parse the whole file. Still re-use imported modules.
+    (void)oldState.takeCodeCompletionDelayedDeclState();
+    parseIntoSourceFile(*newSF, newBufferID, &oldState);
+    performNameBinding(*newSF);
+    bindExtensions(*newSF);
+
+    assert(oldState.hasCodeCompletionDelayedDeclState() &&
+           oldState.getCodeCompletionDelayedDeclState().Kind == newInfo.Kind);
+    break;
+  }
+  }
 
-  // The new completion must happens in function body too.
-  if (newInfo.Kind != CodeCompletionDelayedDeclKind::FunctionBody)
-    return false;
-
-  auto *oldSF = oldInfo.ParentContext->getParentSourceFile();
-
-  // If the interface has changed, AST must be refreshed.
-  llvm::SmallString<32> oldInterfaceHash{};
-  llvm::SmallString<32> newInterfaceHash{};
-  oldSF->getInterfaceHash(oldInterfaceHash);
-  newSF->getInterfaceHash(newInterfaceHash);
-  if (oldInterfaceHash != newInterfaceHash)
-    return false;
-
-  DeclContext *DC =
-      getEquivalentDeclContextFromSourceFile(newInfo.ParentContext, oldSF);
-  if (!DC)
-    return false;
-
-  // OK, we can perform fast completion for this. Update the orignal delayed
-  // decl state.
-
-  // Fast completion keeps the buffer in memory for multiple completions.
-  // To reduce the consumption, slice the source buffer so it only holds
-  // the portion that is needed for the second pass.
-  auto startOffset = newInfo.StartOffset;
-  if (newInfo.PrevOffset != ~0u)
-    startOffset = newInfo.PrevOffset;
-  auto startLoc = tmpSM.getLocForOffset(tmpBufferID, startOffset);
-  startLoc = Lexer::getLocForStartOfLine(tmpSM, startLoc);
-  startOffset = tmpSM.getLocOffsetInBuffer(startLoc, tmpBufferID);
-
-  auto endOffset = newInfo.EndOffset;
-  auto endLoc = tmpSM.getLocForOffset(tmpBufferID, endOffset);
-  endLoc = Lexer::getLocForEndOfToken(tmpSM, endLoc);
-  endOffset = tmpSM.getLocOffsetInBuffer(endLoc, tmpBufferID);
-
-  newInfo.StartOffset -= startOffset;
-  newInfo.EndOffset -= startOffset;
-  if (newInfo.PrevOffset != ~0u)
-    newInfo.PrevOffset -= startOffset;
-
-  auto sourceText = completionBuffer->getBuffer().slice(startOffset, endOffset);
-  auto newOffset = Offset - startOffset;
-
-  auto newBufferID =
-      SM.addMemBufferCopy(sourceText, completionBuffer->getBufferIdentifier());
-  SM.openVirtualFile(SM.getLocForBufferStart(newBufferID),
-                     tmpSM.getDisplayNameForLoc(startLoc),
-                     tmpSM.getLineAndColumn(startLoc).first - 1);
-  SM.setCodeCompletionPoint(newBufferID, newOffset);
-
-  // Construct dummy scopes. We don't need to restore the original scope
-  // because they are probably not 'isResolvable()' anyway.
-  auto &SI = oldState.getScopeInfo();
-  assert(SI.getCurrentScope() == nullptr);
-  Scope Top(SI, ScopeKind::TopLevel);
-  Scope Body(SI, ScopeKind::FunctionBody);
-
-  oldInfo.ParentContext = DC;
-  oldInfo.StartOffset = newInfo.StartOffset;
-  oldInfo.EndOffset = newInfo.EndOffset;
-  oldInfo.PrevOffset = newInfo.PrevOffset;
-  oldState.restoreCodeCompletionDelayedDeclState(oldInfo);
-
-  auto *AFD = cast<AbstractFunctionDecl>(DC);
-  if (AFD->isBodySkipped())
-    AFD->setBodyDelayed(AFD->getBodySourceRange());
   if (DiagC)
     CI.addDiagnosticConsumer(DiagC);
 
@@ -366,6 +410,9 @@ bool swift::ide::CompletionInstance::performOperation(
   // weaken that hash, disable them here:
   Invocation.getLangOptions().EnableTypeFingerprints = false;
 
+  // We don't need token list.
+  Invocation.getLangOptions().CollectParsedToken = false;
+
   // FIXME: ASTScopeLookup doesn't support code completion yet.
   Invocation.disableASTScopeLookup();
 

test/SourceKit/CodeComplete/complete_sequence_accessor.swift

@@ -39,18 +39,18 @@ enum S {
 // Enabled.
 // RUN: %sourcekitd-test \
 // RUN:   -req=track-compiles == \
-// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=12:9 %s -- %s == \
-// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=15:15 %s -- %s == \
-// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=16:15 %s -- %s == \
-// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=23:9 %s -- %s == \
-// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=26:15 %s -- %s == \
-// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=27:15 %s -- %s == \
-// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=30:9 %s -- %s == \
-// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=33:15 %s -- %s == \
-// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=34:15 %s -- %s == \
-// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=12:1 %s -- %s == \
-// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=23:1 %s -- %s == \
-// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=16:1 %s -- %s > %t.response
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=12:9 %s -- %s -parse-as-library == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=15:15 %s -- %s -parse-as-library == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=16:15 %s -- %s -parse-as-library == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=23:9 %s -- %s -parse-as-library == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=26:15 %s -- %s -parse-as-library == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=27:15 %s -- %s -parse-as-library == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=30:9 %s -- %s -parse-as-library == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=33:15 %s -- %s -parse-as-library == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=34:15 %s -- %s -parse-as-library == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=12:1 %s -- %s -parse-as-library == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=23:1 %s -- %s -parse-as-library == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=16:1 %s -- %s -parse-as-library > %t.response
 // RUN: %FileCheck --check-prefix=RESULT  %s < %t.response
 // RUN: %FileCheck --check-prefix=TRACE  %s < %t.response
 

test/SourceKit/CodeComplete/complete_sequence_toplevel.swift

@@ -0,0 +1,49 @@
+class Foo {
+  var x: Int
+  var y: Int
+  func fooMethod() {}
+}
+struct Bar {
+  var a: Int
+  var b: Int
+}
+extension Bar {
+  func barMethod() {}
+}
+func foo(arg: Foo) {
+  _ = arg.
+}
+_ = Bar(a: 12, b: 42)
+
+// Enabled.
+// RUN: %sourcekitd-test \
+// RUN:   -req=track-compiles == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=14:11 %s -- %s == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=13:15 %s -- %s == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=16:22 %s -- %s > %t.response
+// RUN: %FileCheck --check-prefix=RESULT %s < %t.response
+// RUN: %FileCheck --check-prefix=TRACE %s < %t.response
+
+// RESULT-LABEL: key.results: [
+// RESULT-DAG: key.name: "fooMethod()"
+// RESULT-DAG: key.name: "self"
+// RESULT-DAG: key.name: "x"
+// RESULT-DAG: key.name: "y"
+// RESULT: ]
+// RESULT-LABEL: key.results: [
+// RESULT-DAG: key.name: "Foo"
+// RESULT-DAG: key.name: "Bar"
+// RESULT: ]
+// RESULT-LABEL: key.results: [
+// RESULT-DAG: key.name: "barMethod()"
+// RESULT-DAG: key.name: "self"
+// RESULT-DAG: key.name: "a"
+// RESULT-DAG: key.name: "b"
+// RESULT: ]
+
+// TRACE-LABEL: key.notification: source.notification.compile-did-finish,
+// TRACE-NOT: key.description: "completion reusing previous ASTContext (benign diagnostic)"
+// TRACE-LABEL: key.notification: source.notification.compile-did-finish,
+// TRACE: key.description: "completion reusing previous ASTContext (benign diagnostic)"
+// TRACE-LABEL: key.notification: source.notification.compile-did-finish,
+// TRACE: key.description: "completion reusing previous ASTContext (benign diagnostic)"