[SourceKit] ⚡️Fast code completion within function bodies

include/swift/AST/Decl.h

@@ -5841,13 +5841,7 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// \sa hasBody()
   BraceStmt *getBody(bool canSynthesize = true) const;
 
-  void setBody(BraceStmt *S, BodyKind NewBodyKind = BodyKind::Parsed) {
-    assert(getBodyKind() != BodyKind::Skipped &&
-           "cannot set a body if it was skipped");
-
-    Body = S;
-    setBodyKind(NewBodyKind);
-  }
+  void setBody(BraceStmt *S, BodyKind NewBodyKind = BodyKind::Parsed);
 
   /// Note that the body was skipped for this function.  Function body
   /// cannot be attached after this call.
@@ -5866,7 +5860,8 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
 
   /// Note that parsing for the body was delayed.
   void setBodyDelayed(SourceRange bodyRange) {
-    assert(getBodyKind() == BodyKind::None);
+    assert(getBodyKind() == BodyKind::None ||
+           getBodyKind() == BodyKind::Skipped);
     assert(bodyRange.isValid());
     BodyRange = bodyRange;
     setBodyKind(BodyKind::Unparsed);
@@ -6657,6 +6652,9 @@ class ConstructorDecl : public AbstractFunctionDecl {
   /// initializer.
   BodyInitKind getDelegatingOrChainedInitKind(DiagnosticEngine *diags,
                                               ApplyExpr **init = nullptr) const;
+  void clearCachedDelegatingOrChainedInitKind() {
+    Bits.ConstructorDecl.ComputedBodyInitKind = 0;
+  }
 
   /// Whether this constructor is required.
   bool isRequired() const {

include/swift/AST/DiagnosticsFrontend.def

@@ -214,6 +214,9 @@ ERROR(repl_must_be_initialized,none,
 ERROR(error_doing_code_completion,none,
       "compiler is in code completion mode (benign diagnostic)", ())
 
+WARNING(completion_reusing_astcontext,none,
+        "completion reusing previous ASTContext (benign diagnostic)", ())
+
 ERROR(verify_encountered_fatal,none,
       "fatal error encountered while in -verify mode", ())
 

include/swift/AST/SourceFile.h

@@ -404,9 +404,11 @@ class SourceFile final : public FileUnit {
     InterfaceHash->update(a);
   }
 
-  void getInterfaceHash(llvm::SmallString<32> &str) {
+  void getInterfaceHash(llvm::SmallString<32> &str) const {
+    // Copy to preserve idempotence.
+    llvm::MD5 md5 = *InterfaceHash;
     llvm::MD5::MD5Result result;
-    InterfaceHash->final(result);
+    md5.final(result);
     llvm::MD5::stringifyResult(result, str);
   }
 

include/swift/Frontend/Frontend.h

@@ -318,20 +318,11 @@ class CompilerInvocation {
     return CodeCompletionOffset != ~0U;
   }
 
-  void setCodeCompletionFactory(CodeCompletionCallbacksFactory *Factory) {
-    CodeCompletionFactory = Factory;
-    disableASTScopeLookup();
-  }
-
   /// Called from lldb, see rdar://53971116
   void disableASTScopeLookup() {
     LangOpts.EnableASTScopeLookup = false;
   }
 
-  CodeCompletionCallbacksFactory *getCodeCompletionFactory() const {
-    return CodeCompletionFactory;
-  }
-
   /// Retrieve a module hash string that is suitable for uniquely
   /// identifying the conditions under which the module was built, for use
   /// in generating a cached PCH file for the bridging header.
@@ -483,6 +474,10 @@ class CompilerInstance {
     Diagnostics.addConsumer(*DC);
   }
 
+  void removeDiagnosticConsumer(DiagnosticConsumer *DC) {
+    Diagnostics.removeConsumer(*DC);
+  }
+
   void createDependencyTracker(bool TrackSystemDeps) {
     assert(!Context && "must be called before setup()");
     DepTracker = llvm::make_unique<DependencyTracker>(TrackSystemDeps);
@@ -547,6 +542,18 @@ class CompilerInstance {
   /// Returns true if there was an error during setup.
   bool setup(const CompilerInvocation &Invocation);
 
+  const CompilerInvocation &getInvocation() {
+    return Invocation;
+  }
+
+  bool hasPersistentParserState() const {
+    return bool(PersistentState);
+  }
+
+  PersistentParserState &getPersistentParserState() {
+    return *PersistentState.get();
+  }
+
 private:
   /// Set up the file system by loading and validating all VFS overlay YAML
   /// files. If the process of validating VFS files failed, or the overlay

include/swift/IDE/CompletionInstance.h

@@ -0,0 +1,93 @@
+//===--- CompletionInstance.h ---------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2019 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_COMPLETIONINSTANCE_H
+#define SWIFT_IDE_COMPLETIONINSTANCE_H
+
+#include "swift/Frontend/Frontend.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/VirtualFileSystem.h"
+
+namespace swift {
+
+class CompilerInstance;
+class CompilerInvocation;
+class DiagnosticConsumer;
+
+namespace ide {
+
+/// Copy a memory buffer inserting '\0' at the position of \c origBuf.
+std::unique_ptr<llvm::MemoryBuffer>
+makeCodeCompletionMemoryBuffer(const llvm::MemoryBuffer *origBuf,
+                               unsigned &Offset,
+                               llvm::StringRef bufferIdentifier);
+
+/// Manages \c CompilerInstance for completion like operations.
+class CompletionInstance {
+  unsigned MaxASTReuseCount = 100;
+
+  std::mutex mtx;
+
+  std::unique_ptr<CompilerInstance> CachedCI;
+  llvm::hash_code CachedArgHash;
+  unsigned CachedReuseCount = 0;
+
+  /// Calls \p Callback with cached \c CompilerInstance if it's usable for the
+  /// specified completion request.
+  /// Returns \c if the callback was called. Returns \c false if the compiler
+  /// argument has changed, primary file is not the same, the \c Offset is not
+  /// in function bodies, or the interface hash of the file has changed.
+  bool performCachedOperaitonIfPossible(
+      const swift::CompilerInvocation &Invocation, llvm::hash_code ArgsHash,
+      llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
+      DiagnosticConsumer *DiagC,
+      llvm::function_ref<void(CompilerInstance &)> Callback);
+
+  /// Calls \p Callback with new \c CompilerInstance for the completion
+  /// request. The \c CompilerInstace passed to the callback already performed
+  /// the first pass.
+  /// Returns \c false if it fails to setup the \c CompilerInstance.
+  bool performNewOperation(
+      llvm::Optional<llvm::hash_code> ArgsHash,
+      swift::CompilerInvocation &Invocation,
+      llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
+      llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
+      std::string &Error, DiagnosticConsumer *DiagC,
+      llvm::function_ref<void(CompilerInstance &)> Callback);
+
+public:
+  /// Calls \p Callback with a \c CompilerInstance which is prepared for the
+  /// second pass. \p Callback is resposible to perform the second pass on it.
+  /// The \c CompilerInstance may be reused from the previous completions,
+  /// and may be cached for the next completion.
+  /// Return \c true if \p is successfully called, \c it fails. In failure
+  /// cases \p Error is populated with an error message.
+  ///
+  /// NOTE: \p Args is only used for checking the equaity of the invocation.
+  /// Since this function assumes that it is already normalized, exact the same
+  /// arguments including their order is considered as the same invocation.
+  bool performOperation(
+      swift::CompilerInvocation &Invocation, llvm::ArrayRef<const char *> Args,
+      llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
+      llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
+      bool EnableASTCaching, std::string &Error, DiagnosticConsumer *DiagC,
+      llvm::function_ref<void(CompilerInstance &)> Callback);
+};
+
+} // namespace ide
+} // namespace swift
+
+#endif // SWIFT_IDE_COMPLETIONINSTANCE_H

include/swift/Parse/LocalContext.h

@@ -19,6 +19,7 @@
 #define SWIFT_PARSE_LOCALCONTEXT_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "swift/AST/Identifier.h"
 #include <cassert>
 
 namespace swift {

include/swift/Parse/Parser.h

@@ -1642,12 +1642,8 @@ class Parser {
   //===--------------------------------------------------------------------===//
   // Code completion second pass.
 
-  static void
-  performCodeCompletionSecondPass(PersistentParserState &ParserState,
-                                  CodeCompletionCallbacksFactory &Factory);
-
   void performCodeCompletionSecondPassImpl(
-      PersistentParserState::CodeCompletionDelayedDeclState &info);
+      CodeCompletionDelayedDeclState &info);
 };
 
 /// Describes a parsed declaration name.

include/swift/Parse/PersistentParserState.h

@@ -29,6 +29,33 @@ class SourceFile;
 class DeclContext;
 class IterableDeclContext;
 
+enum class CodeCompletionDelayedDeclKind {
+  TopLevelCodeDecl,
+  Decl,
+  FunctionBody,
+};
+
+class CodeCompletionDelayedDeclState {
+public:
+  CodeCompletionDelayedDeclKind Kind;
+  unsigned Flags;
+  DeclContext *ParentContext;
+  SavedScope Scope;
+  unsigned StartOffset;
+  unsigned EndOffset;
+  unsigned PrevOffset;
+
+  SavedScope takeScope() { return std::move(Scope); }
+
+  CodeCompletionDelayedDeclState(CodeCompletionDelayedDeclKind Kind,
+                                 unsigned Flags, DeclContext *ParentContext,
+                                 SavedScope &&Scope, unsigned StartOffset,
+                                 unsigned EndOffset, unsigned PrevOffset)
+      : Kind(Kind), Flags(Flags), ParentContext(ParentContext),
+        Scope(std::move(Scope)), StartOffset(StartOffset), EndOffset(EndOffset),
+        PrevOffset(PrevOffset) {}
+};
+
 /// Parser state persistent across multiple parses.
 class PersistentParserState {
 public:
@@ -39,36 +66,6 @@ class PersistentParserState {
     bool isValid() const { return Loc.isValid(); }
   };
 
-  enum class CodeCompletionDelayedDeclKind {
-    TopLevelCodeDecl,
-    Decl,
-    FunctionBody,
-  };
-
-  class CodeCompletionDelayedDeclState {
-    friend class PersistentParserState;
-    friend class Parser;
-    CodeCompletionDelayedDeclKind Kind;
-    unsigned Flags;
-    DeclContext *ParentContext;
-    ParserPos BodyPos;
-    SourceLoc BodyEnd;
-    SavedScope Scope;
-
-    SavedScope takeScope() {
-      return std::move(Scope);
-    }
-
-  public:
-    CodeCompletionDelayedDeclState(CodeCompletionDelayedDeclKind Kind,
-                                   unsigned Flags, DeclContext *ParentContext,
-                                   SourceRange BodyRange, SourceLoc PreviousLoc,
-                                   SavedScope &&Scope)
-        : Kind(Kind), Flags(Flags),
-          ParentContext(ParentContext), BodyPos{BodyRange.Start, PreviousLoc},
-          BodyEnd(BodyRange.End), Scope(std::move(Scope)) {}
-  };
-
   bool InPoundLineEnvironment = false;
   // FIXME: When condition evaluation moves to a later phase, remove this bit
   // and adjust the client call 'performParseOnly'.
@@ -92,16 +89,23 @@ class PersistentParserState {
   PersistentParserState(ASTContext &ctx) : PersistentParserState() { }
   ~PersistentParserState();
 
-  void setCodeCompletionDelayedDeclState(CodeCompletionDelayedDeclKind Kind,
+  void setCodeCompletionDelayedDeclState(SourceManager &SM, unsigned BufferID,
+                                         CodeCompletionDelayedDeclKind Kind,
                                          unsigned Flags,
                                          DeclContext *ParentContext,
                                          SourceRange BodyRange,
                                          SourceLoc PreviousLoc);
+  void restoreCodeCompletionDelayedDeclState(
+      const CodeCompletionDelayedDeclState &other);
 
   bool hasCodeCompletionDelayedDeclState() {
     return CodeCompletionDelayedDeclStat.get() != nullptr;
   }
 
+  CodeCompletionDelayedDeclState &getCodeCompletionDelayedDeclState() {
+    return *CodeCompletionDelayedDeclStat.get();
+  }
+
   std::unique_ptr<CodeCompletionDelayedDeclState>
   takeCodeCompletionDelayedDeclState() {
     assert(hasCodeCompletionDelayedDeclState());

include/swift/Parse/Scope.h

@@ -141,6 +141,8 @@ class Scope {
   bool isResolvable() const;
 
 public:
+  Scope(ScopeInfo &SI, ScopeKind SC, bool isInactiveConfigBlock = false);
+
   /// Create a lexical scope of the specified kind.
   Scope(Parser *P, ScopeKind SC, bool isInactiveConfigBlock = false);
 

lib/AST/Decl.cpp

@@ -6621,6 +6621,20 @@ BraceStmt *AbstractFunctionDecl::getBody(bool canSynthesize) const {
                            nullptr);
 }
 
+void AbstractFunctionDecl::setBody(BraceStmt *S, BodyKind NewBodyKind) {
+  assert(getBodyKind() != BodyKind::Skipped &&
+         "cannot set a body if it was skipped");
+
+  Body = S;
+  setBodyKind(NewBodyKind);
+
+  // Need to recompute init body kind.
+  if (NewBodyKind < BodyKind::TypeChecked) {
+    if (auto *ctor = dyn_cast<ConstructorDecl>(this))
+      ctor->clearCachedDelegatingOrChainedInitKind();
+  }
+}
+
 SourceRange AbstractFunctionDecl::getBodySourceRange() const {
   switch (getBodyKind()) {
   case BodyKind::None:
@@ -7397,8 +7411,11 @@ ConstructorDecl::getDelegatingOrChainedInitKind(DiagnosticEngine *diags,
 
   // If we already computed the result, return it.
   if (Bits.ConstructorDecl.ComputedBodyInitKind) {
-    return static_cast<BodyInitKind>(
-             Bits.ConstructorDecl.ComputedBodyInitKind - 1);
+    auto Kind = static_cast<BodyInitKind>(
+        Bits.ConstructorDecl.ComputedBodyInitKind - 1);
+    assert((Kind == BodyInitKind::None || !init) &&
+           "can't return cached result with the init expr");
+    return Kind;
   }
 
 

lib/Frontend/Frontend.cpp

@@ -894,12 +894,6 @@ void CompilerInstance::parseAndCheckTypesUpTo(
     }
   });
 
-  if (Invocation.isCodeCompletion()) {
-    assert(limitStage == SourceFile::NameBound);
-    performCodeCompletionSecondPass(*PersistentState.get(),
-                                    *Invocation.getCodeCompletionFactory());
-  }
-
   // If the limiting AST stage is name binding, we're done.
   if (limitStage <= SourceFile::NameBound) {
     return;

lib/IDE/CMakeLists.txt

@@ -2,6 +2,7 @@ add_swift_host_library(swiftIDE STATIC
   CodeCompletion.cpp
   CodeCompletionCache.cpp
   CommentConversion.cpp
+  CompletionInstance.cpp
   ConformingMethodList.cpp
   ExprContextAnalysis.cpp
   Formatting.cpp