[CodeCompletion] Enable ASTScope in code completion

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() {

include/swift/Basic/SourceManager.h

@@ -38,6 +38,20 @@ class SourceManager {
   /// to speed up stats.
   mutable llvm::DenseMap<StringRef, llvm::vfs::Status> StatusCache;
 
+  struct ReplacedRangeType {
+    SourceRange Original;
+    SourceRange New;
+    ReplacedRangeType() {}
+    ReplacedRangeType(NoneType) {}
+    ReplacedRangeType(SourceRange Original, SourceRange New)
+        : Original(Original), New(New) {
+      assert(Original.isValid() && New.isValid());
+    }
+
+    explicit operator bool() const { return Original.isValid(); }
+  };
+  ReplacedRangeType ReplacedRange;
+
   // \c #sourceLocation directive handling.
   struct VirtualFile {
     CharSourceRange Range;
@@ -89,6 +103,9 @@ class SourceManager {
 
   SourceLoc getCodeCompletionLoc() const;
 
+  const ReplacedRangeType &getReplacedRange() const { return ReplacedRange; }
+  void setReplacedRange(const ReplacedRangeType &val) { ReplacedRange = val; }
+
   /// Returns true if \c LHS is before \c RHS in the source buffer.
   bool isBeforeInBuffer(SourceLoc LHS, SourceLoc RHS) const {
     return LHS.Value.getPointer() < RHS.Value.getPointer();

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

@@ -30,7 +30,6 @@
 #include "swift/AST/TypeRepr.h"
 #include "swift/Basic/STLExtras.h"
 #include "llvm/Support/Compiler.h"
-#include <algorithm>
 
 using namespace swift;
 using namespace namelookup;
@@ -82,6 +81,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,33 +151,40 @@ 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 {
   // Use binary search to find the child that contains this location.
-  struct CompareLocs {
-    SourceManager &sourceMgr;
-
-    bool operator()(const ASTScopeImpl *scope, SourceLoc loc) {
-      ASTScopeAssert(scope->checkSourceRangeOfThisASTNode(), "Bad range.");
-      return -1 == ASTScopeImpl::compare(scope->getSourceRangeOfScope(), 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);
-    }
-  };
-  auto *const *child = std::lower_bound(
-      getChildren().begin(), getChildren().end(), loc, CompareLocs{sourceMgr});
-
-  if (child != getChildren().end() &&
-      sourceMgr.rangeContainsTokenLoc((*child)->getSourceRangeOfScope(), loc))
-    return *child;
+  auto *const *child = llvm::lower_bound(
+      getChildren(), loc,
+      [&sourceMgr](const ASTScopeImpl *scope, SourceLoc loc) {
+        ASTScopeAssert(scope->checkSourceRangeOfThisASTNode(), "Bad range.");
+        auto rangeOfScope = scope->getSourceRangeOfScope();
+        loc = translateLocForReplacedRange(sourceMgr, rangeOfScope, loc);
+        return -1 == ASTScopeImpl::compare(rangeOfScope, loc, sourceMgr,
+                                           /*ensureDisjoint=*/false);
+      });
+
+  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

@@ -417,6 +417,8 @@ bool CompletionInstance::performCachedOperationIfPossible(
 
     auto *AFD = cast<AbstractFunctionDecl>(DC);
     AFD->setBodyToBeReparsed(newBodyRange);
+    SM.setReplacedRange({AFD->getOriginalBodySourceRange(), newBodyRange});
+    oldSF->clearScope();
 
     traceDC = AFD;
     break;
@@ -602,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^# })
+}

test/decl/ext/extensions.swift

@@ -349,3 +349,10 @@ struct SR_10466<T> {
 extension SR_10466 where T == Never { // expected-note {{requirement specified as 'T' == 'Never' [with T = T]}}
   typealias A = Int
 }
+
+#if true
+protocol Rdar66943328 {
+  associatedtype Assoc
+}
+extension Rdar66943328 where Assoc == Int // expected-error {{expected '{' in extension}}
+#endif