[CodeCompletion] Enable ASTScope in code completion

* Re-create `ASTScope` for each completion
* Add generic params and where clause scope even without missing body
* Use `getOriginalBodySourceRange()` for `AbstractFunctionBodyScope`
* Source range translations for replaced ranges when finding scopes
* Bypass source range checks when the completion happens in the replaced
  range
* Be lenient with ASTScope / DeclContext mismatch in code completion

Author: rintaro

include/swift/AST/SourceFile.h

@@ -172,7 +172,7 @@ class SourceFile final : public FileUnit {
   bool IsPrimary;
 
   /// The scope map that describes this source file.
-  std::unique_ptr<ASTScope> Scope;
+  NullablePtr<ASTScope> Scope = nullptr;
 
   /// The set of validated opaque return type decls in the source file.
   llvm::SmallVector<OpaqueTypeDecl *, 4> OpaqueReturnTypes;
@@ -468,6 +468,10 @@ class SourceFile final : public FileUnit {
   /// Retrieve the scope that describes this source file.
   ASTScope &getScope();
 
+  void clearScope() {
+    Scope = nullptr;
+  }
+
   /// Retrieves the previously set delayed parser state, asserting that it
   /// exists.
   PersistentParserState *getDelayedParserState() {

lib/AST/ASTScopeCreation.cpp

@@ -795,7 +795,7 @@ void ASTSourceFileScope::
 void ASTSourceFileScope::expandFunctionBody(AbstractFunctionDecl *AFD) {
   if (!AFD)
     return;
-  auto sr = AFD->getBodySourceRange();
+  auto sr = AFD->getOriginalBodySourceRange();
   if (sr.isInvalid())
     return;
   ASTScopeImpl *bodyScope = findInnermostEnclosingScope(sr.Start, nullptr);
@@ -1596,11 +1596,6 @@ ASTScopeImpl *GenericTypeOrExtensionWholePortion::expandScope(
   // Get now in case recursion emancipates scope
   auto *const ip = scope->getParent().get();
 
-  // Prevent circular request bugs caused by illegal input and
-  // doing lookups that getExtendedNominal in the midst of getExtendedNominal.
-  if (scope->shouldHaveABody() && !scope->doesDeclHaveABody())
-    return ip;
-
   auto *context = scope->getGenericContext();
   auto *genericParams = (isa<TypeAliasDecl>(context)
                          ? context->getParsedGenericParams()
@@ -1609,6 +1604,12 @@ ASTScopeImpl *GenericTypeOrExtensionWholePortion::expandScope(
       scope->getDecl(), genericParams, scope);
   if (context->getTrailingWhereClause())
     scope->createTrailingWhereClauseScope(deepestScope, scopeCreator);
+
+  // Prevent circular request bugs caused by illegal input and
+  // doing lookups that getExtendedNominal in the midst of getExtendedNominal.
+  if (scope->shouldHaveABody() && !scope->doesDeclHaveABody())
+    return ip;
+
   scope->createBodyScope(deepestScope, scopeCreator);
   return ip;
 }

lib/AST/ASTScopeLookup.cpp

@@ -82,6 +82,15 @@ const ASTScopeImpl *ASTScopeImpl::findStartingScopeForLookup(
   // Someday, just use the assertion below. For now, print out lots of info for
   // debugging.
   if (!startingScope) {
+
+    // Be lenient in code completion mode. There are cases where the decl
+    // context doesn't match with the ASTScope. e.g. dangling attributes.
+    // FIXME: Create ASTScope tree even for invalid code.
+    if (innermost &&
+        startingContext->getASTContext().SourceMgr.hasCodeCompletionBuffer()) {
+      return innermost;
+    }
+
     llvm::errs() << "ASTScopeImpl: resorting to startingScope hack, file: "
                  << sourceFile->getFilename() << "\n";
     // The check is costly, and inactive lookups will end up here, so don't
@@ -143,6 +152,20 @@ bool ASTScopeImpl::checkSourceRangeOfThisASTNode() const {
   return true;
 }
 
+/// If the \p loc is in a new buffer but \p range is not, consider the location
+/// is at the start of replaced range. Otherwise, returns \p loc as is.
+static SourceLoc translateLocForReplacedRange(SourceManager &sourceMgr,
+                                              SourceRange range,
+                                              SourceLoc loc) {
+  if (const auto &replacedRange = sourceMgr.getReplacedRange()) {
+    if (sourceMgr.rangeContainsTokenLoc(replacedRange.New, loc) &&
+        !sourceMgr.rangeContains(replacedRange.New, range)) {
+      return replacedRange.Original.Start;
+    }
+  }
+  return loc;
+}
+
 NullablePtr<ASTScopeImpl>
 ASTScopeImpl::findChildContaining(SourceLoc loc,
                                   SourceManager &sourceMgr) const {
@@ -152,24 +175,24 @@ ASTScopeImpl::findChildContaining(SourceLoc loc,
 
     bool operator()(const ASTScopeImpl *scope, SourceLoc loc) {
       ASTScopeAssert(scope->checkSourceRangeOfThisASTNode(), "Bad range.");
-      return -1 == ASTScopeImpl::compare(scope->getSourceRangeOfScope(), loc,
-                                         sourceMgr,
+      auto rangeOfScope = scope->getSourceRangeOfScope();
+      loc = translateLocForReplacedRange(sourceMgr, rangeOfScope, loc);
+      return -1 == ASTScopeImpl::compare(rangeOfScope, loc, sourceMgr,
                                          /*ensureDisjoint=*/false);
     }
     bool operator()(SourceLoc loc, const ASTScopeImpl *scope) {
-      ASTScopeAssert(scope->checkSourceRangeOfThisASTNode(), "Bad range.");
-      // Alternatively, we could check that loc < start-of-scope
-      return 0 >= ASTScopeImpl::compare(loc, scope->getSourceRangeOfScope(),
-                                        sourceMgr,
-                                        /*ensureDisjoint=*/false);
+      return !(*this)(scope, loc);
     }
   };
   auto *const *child = std::lower_bound(
       getChildren().begin(), getChildren().end(), loc, CompareLocs{sourceMgr});
 
-  if (child != getChildren().end() &&
-      sourceMgr.rangeContainsTokenLoc((*child)->getSourceRangeOfScope(), loc))
-    return *child;
+  if (child != getChildren().end()) {
+    auto rangeOfScope = (*child)->getSourceRangeOfScope();
+    loc = translateLocForReplacedRange(sourceMgr, rangeOfScope, loc);
+    if (sourceMgr.rangeContainsTokenLoc(rangeOfScope, loc))
+      return *child;
+  }
 
   return nullptr;
 }

lib/AST/ASTScopeSourceRange.cpp

@@ -67,6 +67,14 @@ ASTScopeImpl::widenSourceRangeForChildren(const SourceRange range,
   if (range.isInvalid())
     return childRange;
   auto r = range;
+
+  // HACK: For code completion. If the range of the child is from another
+  // source buffer, don't widen using that range.
+  if (const auto &replacedRange = getSourceManager().getReplacedRange()) {
+    if (getSourceManager().rangeContains(replacedRange.Original, range) &&
+        getSourceManager().rangeContains(replacedRange.New, childRange))
+      return r;
+  }
   r.widen(childRange);
   return r;
 }
@@ -114,6 +122,14 @@ bool ASTScopeImpl::verifyThatChildrenAreContainedWithin(
                   getChildren().back()->getSourceRangeOfScope().End);
   if (getSourceManager().rangeContains(range, rangeOfChildren))
     return true;
+
+  // HACK: For code completion. Handle replaced range.
+  if (const auto &replacedRange = getSourceManager().getReplacedRange()) {
+    if (getSourceManager().rangeContains(replacedRange.Original, range) &&
+        getSourceManager().rangeContains(replacedRange.New, rangeOfChildren))
+      return true;
+  }
+
   auto &out = verificationError() << "children not contained in its parent\n";
   if (getChildren().size() == 1) {
     out << "\n***Only Child node***\n";
@@ -200,7 +216,7 @@ SourceRange DifferentiableAttributeScope::getSourceRangeOfThisASTNode(
 
 SourceRange AbstractFunctionBodyScope::getSourceRangeOfThisASTNode(
     const bool omitAssertions) const {
-  return decl->getBodySourceRange();
+  return decl->getOriginalBodySourceRange();
 }
 
 SourceRange TopLevelCodeScope::getSourceRangeOfThisASTNode(
@@ -357,7 +373,7 @@ SourceRange AbstractFunctionDeclScope::getSourceRangeOfThisASTNode(
     ASTScopeAssert(r.End.isValid(), "Start valid imples end valid.");
     return r;
   }
-  return decl->getBodySourceRange();
+  return decl->getOriginalBodySourceRange();
 }
 
 SourceRange ParameterListScope::getSourceRangeOfThisASTNode(
@@ -609,7 +625,7 @@ SourceRange IterableTypeScope::sourceRangeForDeferredExpansion() const {
   return portion->sourceRangeForDeferredExpansion(this);
 }
 SourceRange AbstractFunctionBodyScope::sourceRangeForDeferredExpansion() const {
-  const auto bsr = decl->getBodySourceRange();
+  const auto bsr = decl->getOriginalBodySourceRange();
   const SourceLoc endEvenIfNoCloseBraceAndEndsWithInterpolatedStringLiteral =
       getLocEncompassingPotentialLookups(getSourceManager(), bsr.End);
   return SourceRange(bsr.Start,

lib/AST/Module.cpp

@@ -2387,7 +2387,7 @@ StringRef SourceFile::getFilename() const {
 
 ASTScope &SourceFile::getScope() {
   if (!Scope)
-    Scope = std::unique_ptr<ASTScope>(new (getASTContext()) ASTScope(this));
+    Scope = new (getASTContext()) ASTScope(this);
   return *Scope.get();
 }
 

lib/IDE/CompletionInstance.cpp

@@ -418,6 +418,7 @@ bool CompletionInstance::performCachedOperationIfPossible(
     auto *AFD = cast<AbstractFunctionDecl>(DC);
     AFD->setBodyToBeReparsed(newBodyRange);
     SM.setReplacedRange({AFD->getOriginalBodySourceRange(), newBodyRange});
+    oldSF->clearScope();
 
     traceDC = AFD;
     break;
@@ -603,9 +604,6 @@ bool swift::ide::CompletionInstance::performOperation(
   // We don't need token list.
   Invocation.getLangOptions().CollectParsedToken = false;
 
-  // FIXME: ASTScopeLookup doesn't support code completion yet.
-  Invocation.disableASTScopeLookup();
-
   if (EnableASTCaching) {
     // Compute the signature of the invocation.
     llvm::hash_code ArgsHash(0);

test/IDE/complete_attributes.swift

@@ -1,5 +1,7 @@
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=TOP_LEVEL_ATTR_1 -code-completion-keywords=false | %FileCheck %s -check-prefix=ERROR_COMMON
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=MEMBER_DECL_ATTR_1 -code-completion-keywords=false | %FileCheck %s -check-prefix=ERROR_COMMON
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=ATTRARG_MEMBER | %FileCheck %s -check-prefix=MEMBER_MyValue
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=ATTRARG_MEMBER_IN_CLOSURE | %FileCheck %s -check-prefix=MEMBER_MyValue
 
 // ERROR_COMMON: found code completion token
 // ERROR_COMMON-NOT: Keyword/
@@ -9,3 +11,22 @@
 class MemberDeclAttribute {
   @#^MEMBER_DECL_ATTR_1^# func memberDeclAttr1() {}
 }
+
+struct MyValue {
+  init() {}
+  static var val: Int
+}
+
+// MEMBER_MyValue: Begin completions, 4 items
+// MEMBER_MyValue-DAG: Keyword[self]/CurrNominal:          self[#MyValue.Type#];
+// MEMBER_MyValue-DAG: Keyword/CurrNominal:                Type[#MyValue.Type#];
+// MEMBER_MyValue-DAG: Decl[Constructor]/CurrNominal:      init()[#MyValue#];
+// MEMBER_MyValue-DAG: Decl[StaticVar]/CurrNominal:        val[#Int#];
+// MEMBER_MyValue: End completions
+
+class TestUknownDanglingAttr1 {
+  @UknownAttr(arg: MyValue.#^ATTRARG_MEMBER^#)
+}
+class TestUknownDanglingAttr2 {
+  @UknownAttr(arg: { MyValue.#^ATTRARG_MEMBER_IN_CLOSURE^# })
+}