Handle module selectors with local vars right

Author: beccadax

include/swift/AST/Identifier.h

@@ -769,6 +769,7 @@ class DeclNameRef {
 public:
   static DeclNameRef createSubscript();
   static DeclNameRef createConstructor();
+  static DeclNameRef createSelf(const ASTContext &ctx);
 
   DeclNameRef() : storage(DeclName()) { }
 

include/swift/AST/NameLookup.h

@@ -33,7 +33,7 @@
 
 namespace swift {
 class ASTContext;
-class DeclName;
+class DeclNameRef;
 class Type;
 class TypeDecl;
 class ValueDecl;
@@ -773,20 +773,20 @@ class ASTScope : public ASTAllocated<ASTScope> {
   ///
   /// \param stopAfterInnermostBraceStmt If lookup should consider
   /// local declarations inside the innermost syntactic scope only.
-  static void lookupLocalDecls(SourceFile *, DeclName, SourceLoc,
+  static void lookupLocalDecls(SourceFile *, DeclNameRef, SourceLoc,
                                bool stopAfterInnermostBraceStmt,
                                ABIRole roleFilter,
                                SmallVectorImpl<ValueDecl *> &);
 
-  static void lookupLocalDecls(SourceFile *sf, DeclName name, SourceLoc loc,
+  static void lookupLocalDecls(SourceFile *sf, DeclNameRef name, SourceLoc loc,
                                bool stopAfterInnermostBraceStmt,
                                SmallVectorImpl<ValueDecl *> &results) {
     lookupLocalDecls(sf, name, loc, stopAfterInnermostBraceStmt,
                      ABIRole::ProvidesAPI, results);
   }
 
   /// Returns the result if there is exactly one, nullptr otherwise.
-  static ValueDecl *lookupSingleLocalDecl(SourceFile *, DeclName, SourceLoc);
+  static ValueDecl *lookupSingleLocalDecl(SourceFile *, DeclNameRef, SourceLoc);
 
   /// Entry point to record the visible statement labels from the given
   /// point.

lib/AST/Identifier.cpp

@@ -220,6 +220,10 @@ llvm::raw_ostream &DeclName::printPretty(llvm::raw_ostream &os) const {
   return print(os, /*skipEmptyArgumentNames=*/!isSpecial());
 }
 
+DeclNameRef DeclNameRef::createSelf(const ASTContext &ctx) {
+  return DeclNameRef(ctx.Id_self);
+}
+
 void DeclNameRef::dump() const {
   llvm::errs() << *this << "\n";
 }

lib/AST/UnqualifiedLookup.cpp

@@ -383,7 +383,8 @@ ValueDecl *UnqualifiedLookupFactory::lookupBaseDecl(const DeclContext *baseDC) c
     return nullptr;
 
   auto selfDecl = ASTScope::lookupSingleLocalDecl(DC->getParentSourceFile(),
-                                                  DeclName(Ctx.Id_self), Loc);
+                                                  DeclNameRef::createSelf(Ctx),
+                                                  Loc);
   if (!selfDecl) {
     return nullptr;
   }
@@ -910,14 +911,14 @@ namespace {
 
 class ASTScopeDeclConsumerForLocalLookup
     : public AbstractASTScopeDeclConsumer {
-  DeclName name;
+  DeclNameRef name;
   bool stopAfterInnermostBraceStmt;
   ABIRole roleFilter;
   SmallVectorImpl<ValueDecl *> &results;
 
 public:
   ASTScopeDeclConsumerForLocalLookup(
-      DeclName name, bool stopAfterInnermostBraceStmt,
+      DeclNameRef name, bool stopAfterInnermostBraceStmt,
       ABIRole roleFilter, SmallVectorImpl<ValueDecl *> &results)
     : name(name), stopAfterInnermostBraceStmt(stopAfterInnermostBraceStmt),
       roleFilter(roleFilter), results(results) {}
@@ -928,6 +929,8 @@ class ASTScopeDeclConsumerForLocalLookup
 
   bool consume(ArrayRef<ValueDecl *> values,
                NullablePtr<DeclContext> baseDC) override {
+    if (name.hasModuleSelector()) return false;
+    
     for (auto *value: values) {
       bool foundMatch = false;
       if (auto *varDecl = dyn_cast<VarDecl>(value)) {
@@ -941,7 +944,7 @@ class ASTScopeDeclConsumerForLocalLookup
         });
       }
 
-      if (!foundMatch && value->getName().matchesRef(name)
+      if (!foundMatch && value->getName().matchesRef(name.getFullName())
               && hasCorrectABIRole(value))
         results.push_back(value);
     }
@@ -968,7 +971,7 @@ class ASTScopeDeclConsumerForLocalLookup
 
 /// Lookup that only finds local declarations and does not trigger
 /// interface type computation.
-void ASTScope::lookupLocalDecls(SourceFile *sf, DeclName name, SourceLoc loc,
+void ASTScope::lookupLocalDecls(SourceFile *sf, DeclNameRef name, SourceLoc loc,
                                 bool stopAfterInnermostBraceStmt,
                                 ABIRole roleFilter,
                                 SmallVectorImpl<ValueDecl *> &results) {
@@ -977,7 +980,7 @@ void ASTScope::lookupLocalDecls(SourceFile *sf, DeclName name, SourceLoc loc,
   ASTScope::unqualifiedLookup(sf, loc, consumer);
 }
 
-ValueDecl *ASTScope::lookupSingleLocalDecl(SourceFile *sf, DeclName name,
+ValueDecl *ASTScope::lookupSingleLocalDecl(SourceFile *sf, DeclNameRef name,
                                            SourceLoc loc) {
   SmallVector<ValueDecl *, 1> result;
   ASTScope::lookupLocalDecls(sf, name, loc,

lib/Sema/CSDiagnostics.cpp

@@ -1758,7 +1758,7 @@ class VarDeclMultipleReferencesChecker : public ASTWalker {
       if (name.isSimpleName(varDecl->getName()) && loc.isValid()) {
         auto *otherDecl =
             ASTScope::lookupSingleLocalDecl(DC->getParentSourceFile(),
-                                            name.getFullName(), loc);
+                                            name, loc);
         if (otherDecl == varDecl)
           ++count;
       }

lib/Sema/CSGen.cpp

@@ -897,7 +897,7 @@ TypeVarRefCollector::walkToExprPre(Expr *expr) {
     auto loc = declRef->getLoc();
     if (name.isSimpleName() && loc.isValid()) {
       auto *SF = CS.DC->getParentSourceFile();
-      auto *D = ASTScope::lookupSingleLocalDecl(SF, name.getFullName(), loc);
+      auto *D = ASTScope::lookupSingleLocalDecl(SF, name, loc);
       inferTypeVars(D);
     }
   }

lib/Sema/PreCheckTarget.cpp

@@ -560,8 +560,7 @@ static Expr *resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE, DeclContext *DC,
       }
     }
 
-    DeclName lookupName(context, Name.getBaseName(), lookupLabels);
-    LookupName = DeclNameRef(lookupName);
+    LookupName = Name.withArgumentLabels(context, lookupLabels);
   }
 
   LookupResult Lookup;
@@ -571,8 +570,7 @@ static Expr *resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE, DeclContext *DC,
 
   // First, look for a local binding in scope.
   if (Loc.isValid() && !Name.isOperator()) {
-    ASTScope::lookupLocalDecls(DC->getParentSourceFile(),
-                               LookupName.getFullName(), Loc,
+    ASTScope::lookupLocalDecls(DC->getParentSourceFile(), LookupName, Loc,
                                /*stopAfterInnermostBraceStmt=*/false,
                                ResultValues);
     for (auto *localDecl : ResultValues) {
@@ -2177,8 +2175,9 @@ VarDecl *PreCheckTarget::getImplicitSelfDeclForSuperContext(SourceLoc Loc) {
 
   // Do an actual lookup for 'self' in case it shows up in a capture list.
   auto *methodSelf = methodContext->getImplicitSelfDecl();
-  auto *lookupSelf = ASTScope::lookupSingleLocalDecl(DC->getParentSourceFile(),
-                                                     Ctx.Id_self, Loc);
+  auto *lookupSelf = ASTScope::lookupSingleLocalDecl(
+                                   DC->getParentSourceFile(),
+                                   DeclNameRef::createSelf(Ctx), Loc);
   if (lookupSelf && lookupSelf != methodSelf) {
     // FIXME: This is the wrong diagnostic for if someone manually declares a
     // variable named 'self' using backticks.

lib/Sema/TypeCheckDecl.cpp

@@ -577,8 +577,8 @@ BodyInitKindRequest::evaluate(Evaluator &evaluator,
         auto loc = declRef->getLoc();
         if (name.isSimpleName(ctx.Id_self)) {
           auto *otherSelfDecl =
-            ASTScope::lookupSingleLocalDecl(Decl->getParentSourceFile(),
-                                            name.getFullName(), loc);
+            ASTScope::lookupSingleLocalDecl(Decl->getParentSourceFile(), name,
+                                            loc);
           if (otherSelfDecl == Decl->getImplicitSelfDecl())
             myKind = BodyInitKind::Delegating;
         }

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -757,7 +757,8 @@ CheckRedeclarationRequest::evaluate(Evaluator &eval, ValueDecl *current,
       else
         roleFilter = currentIsABIOnly ? ABIRole::ProvidesABI
                                       : ABIRole::ProvidesAPI;
-      ASTScope::lookupLocalDecls(currentFile, current->getBaseName(),
+      ASTScope::lookupLocalDecls(currentFile,
+                                 DeclNameRef(current->getBaseName()),
                                  current->getLoc(),
                                  /*stopAfterInnermostBraceStmt=*/true,
                                  roleFilter, otherDefinitions);

test/NameLookup/module_selector.swift

@@ -309,20 +309,23 @@ func decl1(
   label p3: @escaping () -> A
 ) {
   switch Optional(main::p2) {
+  // expected-error@-1 {{cannot find 'main::p2' in scope}}
   case Optional.some(let decl1i):
     break
   case .none:
     break
   }
 
   switch Optional(main::p2) {
+  // expected-error@-1 {{cannot find 'main::p2' in scope}}
   case let Optional.some(decl1j):
     break
   case .none:
     break
   }
 
   switch Optional(main::p2) {
+  // expected-error@-1 {{cannot find 'main::p2' in scope}}
  case let decl1k?:
     break
   case .none: