Move implicit self lookup from resolveDeclRefExpr into ASTScope::unqualifiedLookup implementation

Author: calda

include/swift/AST/NameLookup.h

@@ -103,15 +103,18 @@ struct LookupResultEntry {
   /// extension (if it found something at that level).
   DeclContext *BaseDC;
 
+  /// The declaration that defines the base of the call to `Value`
+  ValueDecl *BaseDecl;
+  
   /// The declaration corresponds to the given name; i.e. the decl we are
   /// looking up.
   ValueDecl *Value;
 
 public:
-  LookupResultEntry(ValueDecl *value) : BaseDC(nullptr), Value(value) {}
+  LookupResultEntry(ValueDecl *value) : BaseDC(nullptr), BaseDecl(nullptr), Value(value) {}
 
-  LookupResultEntry(DeclContext *baseDC, ValueDecl *value)
-    : BaseDC(baseDC), Value(value) {}
+  LookupResultEntry(DeclContext *baseDC, ValueDecl *baseDecl, ValueDecl *value)
+    : BaseDC(baseDC), BaseDecl(baseDecl), Value(value) {}
 
   ValueDecl *getValueDecl() const { return Value; }
 

lib/AST/NameLookup.cpp

@@ -55,6 +55,9 @@ void VectorDeclConsumer::anchor() {}
 ValueDecl *LookupResultEntry::getBaseDecl() const {
   if (BaseDC == nullptr)
     return nullptr;
+  
+  if (BaseDecl)
+    return BaseDecl;
 
   if (auto *AFD = dyn_cast<AbstractFunctionDecl>(BaseDC))
     return AFD->getImplicitSelfDecl();

lib/AST/UnqualifiedLookup.cpp

@@ -343,8 +343,22 @@ void UnqualifiedLookupFactory::ResultFinderForTypeContext::findResults(
   SmallVector<ValueDecl *, 4> Lookup;
   contextForLookup->lookupQualified(selfBounds, Name, baseNLOptions, Lookup);
   for (auto Result : Lookup) {
-    results.emplace_back(const_cast<DeclContext *>(whereValueIsMember(Result)),
-                         Result);
+    auto baseDC = const_cast<DeclContext *>(whereValueIsMember(Result));
+    
+    // Perform an unqualified lookup for the base decl of this result. This handles cases
+    // where self was rebound (e.g. `guard let self = self`) earlier in the scope.
+    // Only do this in closures, since implicit self isn't allowed to be rebound
+    // in other contexts. Otherwise, we use the `ParamDecl` from `baseDC` if applicable.
+    ValueDecl* localBaseDecl = nullptr;
+    if (factory->DC->getContextKind() == DeclContextKind::AbstractClosureExpr && baseDC) {
+      auto *localDecl = ASTScope::lookupSingleLocalDecl(factory->DC->getParentSourceFile(),
+                                                       DeclName(factory->Ctx.Id_self),
+                                                       factory->Loc);
+      if (localDecl)
+        localBaseDecl = localDecl;
+    }
+    
+    results.emplace_back(baseDC, localBaseDecl, Result);
 #ifndef NDEBUG
     factory->addedResult(results.back());
 #endif

lib/Sema/PreCheckExpr.cpp

@@ -728,22 +728,7 @@ Expr *TypeChecker::resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE,
   ValueDecl *Base = nullptr;
   DeclContext *BaseDC = nullptr;
   for (auto Result : Lookup) {
-    // Perform an unqualified lookup for the base decl. This handles cases
-    // where self was rebound (e.g. `guard let self = self`) earlier in this scope.
-    // Only do this in closures, since implicit self isn't allowed to be rebound
-    // in other contexts.
-    ValueDecl* ThisBase = nullptr;
-    if (DC->getContextKind() == DeclContextKind::AbstractClosureExpr) {
-      auto &ctx = DC->getASTContext();
-      auto localDecl = ASTScope::lookupSingleLocalDecl(DC->getParentSourceFile(),
-                                                       DeclName(ctx.Id_self),
-                                                       UDRE->getLoc());
-      if (localDecl)
-        ThisBase = localDecl;
-    }
-    
-    if (!ThisBase)
-      ThisBase = Result.getBaseDecl();
+    auto ThisBase = Result.getBaseDecl();
 
     // Track the base for member declarations.
     if (ThisBase && !isa<ModuleDecl>(ThisBase)) {

lib/Sema/TypeCheckNameLookup.cpp

@@ -92,18 +92,21 @@ namespace {
     /// \param baseDC The declaration context through which we found the
     /// declaration.
     ///
+    /// \param baseDecl The declaration that defines the base of the
+    /// call to `found`
+    ///
     /// \param foundInType The type through which we found the
     /// declaration.
     ///
     /// \param isOuter Whether this is an outer result (i.e. a result that isn't
     /// from the innermost scope with results)
-    void add(ValueDecl *found, DeclContext *baseDC, Type foundInType,
+    void add(ValueDecl *found, DeclContext *baseDC, ValueDecl *baseDecl, Type foundInType,
              bool isOuter) {
       DeclContext *foundDC = found->getDeclContext();
 
       auto addResult = [&](ValueDecl *result) {
         if (Known.insert({{result, baseDC}, false}).second) {
-          Result.add(LookupResultEntry(baseDC, result), isOuter);
+          Result.add(LookupResultEntry(baseDC, baseDecl, result), isOuter);
           if (isOuter)
             FoundOuterDecls.push_back(result);
           else
@@ -267,7 +270,7 @@ LookupResult TypeChecker::lookupUnqualified(DeclContext *dc, DeclNameRef name,
       assert(foundInType && "bogus base declaration?");
     }
 
-    builder.add(found.getValueDecl(), found.getDeclContext(), foundInType,
+    builder.add(found.getValueDecl(), found.getDeclContext(), found.getBaseDecl(), foundInType,
                 /*isOuter=*/idx >= lookup.getIndexOfFirstOuterResult());
   }
   return result;
@@ -327,7 +330,7 @@ LookupResult TypeChecker::lookupMember(DeclContext *dc,
   dc->lookupQualified(type, name, subOptions, lookupResults);
 
   for (auto found : lookupResults)
-    builder.add(found, nullptr, type, /*isOuter=*/false);
+    builder.add(found, nullptr, nullptr, type, /*isOuter=*/false);
 
   return result;
 }