Make DeclNameRef able to store a module selector

beccadax · beccadax · commit 64d8b7094211 · 2025-10-24T16:23:33.000-07:00
In this commit, this change affects certain diagnostics but doesn’t actually alter name lookup behavior yet.
diff --git a/include/swift/AST/ASTBridging.h b/include/swift/AST/ASTBridging.h
@@ -164,10 +164,13 @@ BridgedDeclNameRef_createParsed(BridgedASTContext cContext,
 
 class BridgedDeclNameLoc {
   const void *_Nullable LocationInfo;
-  size_t NumArgumentLabels;
+  uint32_t NumArgumentLabels;
+  bool HasModuleSelectorLoc;
 
 public:
-  BridgedDeclNameLoc() : LocationInfo(nullptr), NumArgumentLabels(0) {}
+  BridgedDeclNameLoc()
+    : LocationInfo(nullptr), NumArgumentLabels(0), HasModuleSelectorLoc(false)
+  {}
 
   BRIDGED_INLINE BridgedDeclNameLoc(swift::DeclNameLoc loc);
 
diff --git a/include/swift/AST/ASTBridgingImpl.h b/include/swift/AST/ASTBridgingImpl.h
@@ -60,10 +60,12 @@ swift::DeclNameRef BridgedDeclNameRef::unbridged() const {
 
 BridgedDeclNameLoc::BridgedDeclNameLoc(swift::DeclNameLoc loc)
     : LocationInfo(loc.LocationInfo),
-      NumArgumentLabels(loc.NumArgumentLabels) {}
+      NumArgumentLabels(loc.NumArgumentLabels),
+      HasModuleSelectorLoc(loc.HasModuleSelectorLoc) {}
 
 swift::DeclNameLoc BridgedDeclNameLoc::unbridged() const {
-  return swift::DeclNameLoc(LocationInfo, NumArgumentLabels);
+  return swift::DeclNameLoc(LocationInfo, NumArgumentLabels,
+                            HasModuleSelectorLoc);
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/include/swift/AST/DeclNameLoc.h b/include/swift/AST/DeclNameLoc.h
@@ -35,61 +35,69 @@ class DeclNameLoc {
 
   /// Source location information.
   ///
-  /// If \c NumArgumentLabels == 0, this is the SourceLoc for the base name.
-  /// Otherwise, it points to an array of SourceLocs, which contains:
+  /// If \c NumArgumentLabels == 0 and \c !HasModuleSelectorLoc, this is the
+  /// SourceLoc for the base name. Otherwise, it points to an array of
+  /// SourceLocs, which contains:
   /// * The base name location
+  /// * The module selector location
   /// * The left parentheses location
   /// * The right parentheses location
   /// * The locations of each of the argument labels.
   const void *LocationInfo;
 
   /// The number of argument labels stored in the name.
-  unsigned NumArgumentLabels;
+  uint32_t NumArgumentLabels;
+  bool HasModuleSelectorLoc;
 
   enum {
     BaseNameIndex = 0,
-    LParenIndex = 1,
-    RParenIndex = 2,
-    FirstArgumentLabelIndex = 3,
+    ModuleSelectorIndex = 1,
+    LParenIndex = 2,
+    RParenIndex = 3,
+    FirstArgumentLabelIndex = 4,
   };
 
   /// Retrieve a pointer to either the only source location that was
   /// stored or to the array of source locations that was stored.
   SourceLoc const * getSourceLocs() const {
-    if (NumArgumentLabels == 0) 
+    if (NumArgumentLabels == 0 && !HasModuleSelectorLoc)
       return reinterpret_cast<SourceLoc const *>(&LocationInfo);
 
     return reinterpret_cast<SourceLoc const *>(LocationInfo);
   }
 
-  DeclNameLoc(const void *LocationInfo, unsigned NumArgumentLabels)
-      : LocationInfo(LocationInfo), NumArgumentLabels(NumArgumentLabels) {}
+  DeclNameLoc(const void *LocationInfo, unsigned NumArgumentLabels,
+              bool HasModuleSelectorLoc)
+      : LocationInfo(LocationInfo), NumArgumentLabels(NumArgumentLabels),
+        HasModuleSelectorLoc(HasModuleSelectorLoc) {}
 
 public:
   /// Create an invalid declaration name location.
-  DeclNameLoc() : DeclNameLoc(nullptr, 0) {}
+  DeclNameLoc() : DeclNameLoc(nullptr, 0, false) {}
 
   /// Create declaration name location information for a base name.
   explicit DeclNameLoc(SourceLoc baseNameLoc)
-      : DeclNameLoc(baseNameLoc.getOpaquePointerValue(), 0) {}
+      : DeclNameLoc(baseNameLoc.getOpaquePointerValue(), 0, false) {}
 
   explicit DeclNameLoc(ASTContext &ctx, SourceLoc moduleSelectorLoc,
                        SourceLoc baseNameLoc)
-    : DeclNameLoc(baseNameLoc) { }
+    : DeclNameLoc(ctx, moduleSelectorLoc, baseNameLoc,
+                  SourceLoc(), {}, SourceLoc()) { }
 
   /// Create declaration name location information for a compound
   /// name.
   DeclNameLoc(ASTContext &ctx, SourceLoc baseNameLoc,
               SourceLoc lParenLoc,
               ArrayRef<SourceLoc> argumentLabelLocs,
-              SourceLoc rParenLoc);
+              SourceLoc rParenLoc)
+    : DeclNameLoc(ctx, SourceLoc(), baseNameLoc,
+                  lParenLoc, argumentLabelLocs, rParenLoc) { }
 
   DeclNameLoc(ASTContext &ctx, SourceLoc moduleSelectorLoc,
               SourceLoc baseNameLoc,
               SourceLoc lParenLoc,
               ArrayRef<SourceLoc> argumentLabelLocs,
-              SourceLoc rParenLoc)
-    : DeclNameLoc(ctx, baseNameLoc, lParenLoc, argumentLabelLocs, rParenLoc) { }
+              SourceLoc rParenLoc);
 
   /// Whether the location information is valid.
   bool isValid() const { return getBaseNameLoc().isValid(); }
@@ -125,11 +133,12 @@ class DeclNameLoc {
   }
 
   SourceLoc getModuleSelectorLoc() const {
-    return SourceLoc();
+    if (!HasModuleSelectorLoc) return SourceLoc();
+    return getSourceLocs()[ModuleSelectorIndex];
   }
 
   SourceLoc getStartLoc() const {
-    return getBaseNameLoc();
+    return HasModuleSelectorLoc ? getModuleSelectorLoc() : getBaseNameLoc();
   }
 
   SourceLoc getEndLoc() const {
@@ -138,9 +147,7 @@ class DeclNameLoc {
   
   /// Retrieve the complete source range for this declaration name.
   SourceRange getSourceRange() const {
-    if (NumArgumentLabels == 0) return getBaseNameLoc();
-
-    return SourceRange(getBaseNameLoc(), getRParenLoc());
+    return SourceRange(getStartLoc(), getEndLoc());
   }
 };
 
diff --git a/include/swift/AST/Identifier.h b/include/swift/AST/Identifier.h
@@ -738,39 +738,75 @@ void simple_display(llvm::raw_ostream &out, DeclName name);
 /// An in-source reference to another declaration, including qualification
 /// information.
 class DeclNameRef {
-  DeclName FullName;
+  friend class ASTContext;
+  friend struct llvm::PointerLikeTypeTraits<DeclNameRef>;
+
+  /// Contains the name and module for a DeclNameRef with a module selector.
+  struct alignas(Identifier) SelectiveDeclNameRef : llvm::FoldingSetNode {
+    Identifier moduleSelector;   // Note: currently can never be empty().
+    DeclName fullName;
+
+    SelectiveDeclNameRef(Identifier moduleSelector, DeclName fullName)
+      : moduleSelector(moduleSelector), fullName(fullName) { }
+
+    /// Uniquing for the ASTContext.
+    static void Profile(llvm::FoldingSetNodeID &id, Identifier moduleSelector,
+                        DeclName fullName);
+
+    void Profile(llvm::FoldingSetNodeID &id) {
+      Profile(id, moduleSelector, fullName);
+    }
+  };
+
+  using Storage = llvm::PointerUnion<DeclName, SelectiveDeclNameRef *>;
+  Storage storage;
+
+  explicit DeclNameRef(void *_Nullable Opaque)
+    : storage(decltype(storage)::getFromOpaqueValue(Opaque)) { }
+
+  void initialize(ASTContext &C, Identifier moduleScope, DeclName fullName);
 
 public:
   static DeclNameRef createSubscript();
   static DeclNameRef createConstructor();
 
-  DeclNameRef() : FullName() { }
+  DeclNameRef() : storage(DeclName()) { }
 
-  void *_Nullable getOpaqueValue() const { return FullName.getOpaqueValue(); }
+  void *_Nullable getOpaqueValue() const {
+    return storage.getOpaqueValue();
+  }
   static DeclNameRef getFromOpaqueValue(void *_Nullable p);
 
   explicit DeclNameRef(ASTContext &C, Identifier moduleSelector,
-                       DeclName fullName)
-    : FullName(fullName) { }
+                       DeclName fullName) {
+    initialize(C, moduleSelector, fullName);
+  }
 
   explicit DeclNameRef(ASTContext &C, Identifier moduleSelector,
-                       DeclBaseName baseName, ArrayRef<Identifier> argLabels)
-    : FullName(C, baseName, argLabels) { }
+                       DeclBaseName baseName, ArrayRef<Identifier> argLabels) {
+    initialize(C, moduleSelector, DeclName(C, baseName, argLabels));
+  }
 
   explicit DeclNameRef(DeclName FullName)
-    : FullName(FullName) { }
+    : storage(FullName) { }
 
   bool hasModuleSelector() const {
-    return false;
+    return isa<SelectiveDeclNameRef *>(storage);
   }
 
   Identifier getModuleSelector() const {
-    return Identifier();
+    if (!hasModuleSelector())
+      return Identifier();
+
+    return cast<SelectiveDeclNameRef *>(storage)->moduleSelector;
   }
 
   /// The name of the declaration being referenced.
   DeclName getFullName() const {
-    return FullName;
+    if (!hasModuleSelector())
+      return cast<DeclName>(storage);
+
+    return cast<SelectiveDeclNameRef *>(storage)->fullName;
   }
 
   /// The base name of the declaration being referenced.
@@ -885,7 +921,7 @@ class DeclNameRef {
 };
 
 inline DeclNameRef DeclNameRef::getFromOpaqueValue(void *_Nullable p) {
-  return DeclNameRef(DeclName::getFromOpaqueValue(p));
+  return DeclNameRef(p);
 }
 
 inline DeclNameRef DeclNameRef::withoutArgumentLabels(ASTContext &C) const {
@@ -1059,7 +1095,6 @@ namespace llvm {
   };
 
   // A DeclNameRef is "pointer like" just like DeclNames.
-  template<typename T> struct PointerLikeTypeTraits;
   template<>
   struct PointerLikeTypeTraits<swift::DeclNameRef> {
   public:
@@ -1069,7 +1104,7 @@ namespace llvm {
     static inline swift::DeclNameRef getFromVoidPointer(void *_Nullable ptr) {
       return swift::DeclNameRef::getFromOpaqueValue(ptr);
     }
-    enum { NumLowBitsAvailable = PointerLikeTypeTraits<swift::DeclName>::NumLowBitsAvailable };
+    enum { NumLowBitsAvailable = PointerLikeTypeTraits<swift::DeclNameRef::Storage>::NumLowBitsAvailable };
   };
 
   // DeclNameRefs hash just like DeclNames.
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
@@ -641,6 +641,7 @@ struct ASTContext::Implementation {
   llvm::DenseMap<unsigned, BuiltinUnboundGenericType*> BuiltinUnboundGenericTypes;
   llvm::FoldingSet<BuiltinVectorType> BuiltinVectorTypes;
   llvm::FoldingSet<DeclName::CompoundDeclName> CompoundNames;
+  llvm::FoldingSet<DeclNameRef::SelectiveDeclNameRef> SelectiveNameRefs;
   llvm::DenseMap<UUID, GenericEnvironment *> OpenedElementEnvironments;
   llvm::FoldingSet<IndexSubset> IndexSubsets;
   llvm::FoldingSet<AutoDiffDerivativeFunctionIdentifier>
@@ -6233,6 +6234,39 @@ DeclName::DeclName(ASTContext &C, DeclBaseName baseName,
   initialize(C, baseName, names);
 }
 
+void DeclNameRef::SelectiveDeclNameRef::Profile(llvm::FoldingSetNodeID &id,
+                                                Identifier moduleSelector,
+                                                DeclName fullName) {
+  ASSERT(!moduleSelector.empty() &&
+         "Looking up SelectiveDeclNameRef with empty module?");
+  id.AddPointer(moduleSelector.getAsOpaquePointer());
+  id.AddPointer(fullName.getOpaqueValue());
+}
+
+void DeclNameRef::initialize(ASTContext &C, Identifier moduleSelector,
+                             DeclName fullName) {
+  if (moduleSelector.empty()) {
+    storage = fullName;
+    return;
+  }
+
+  llvm::FoldingSetNodeID id;
+  SelectiveDeclNameRef::Profile(id, moduleSelector, fullName);
+
+  void *insert = nullptr;
+  if (SelectiveDeclNameRef *selectiveRef
+        = C.getImpl().SelectiveNameRefs.FindNodeOrInsertPos(id, insert)) {
+    storage = selectiveRef;
+    return;
+  }
+
+  auto buf = C.Allocate(sizeof(SelectiveDeclNameRef),
+                        alignof(SelectiveDeclNameRef));
+  auto selectiveRef = new (buf) SelectiveDeclNameRef(moduleSelector, fullName);
+  storage = selectiveRef;
+  C.getImpl().SelectiveNameRefs.InsertNode(selectiveRef, insert);
+}
+
 /// Find the implementation of the named type in the given module.
 static NominalTypeDecl *findUnderlyingTypeInModule(ASTContext &ctx, 
                                                    Identifier name,
diff --git a/lib/AST/DeclNameLoc.cpp b/lib/AST/DeclNameLoc.cpp
@@ -20,16 +20,24 @@
 
 using namespace swift;
 
-DeclNameLoc::DeclNameLoc(ASTContext &ctx, SourceLoc baseNameLoc,
+DeclNameLoc::DeclNameLoc(ASTContext &ctx, SourceLoc moduleSelectorLoc,
+                         SourceLoc baseNameLoc,
                          SourceLoc lParenLoc,
                          ArrayRef<SourceLoc> argumentLabelLocs,
                          SourceLoc rParenLoc)
-  : NumArgumentLabels(argumentLabelLocs.size()) {
-  assert(NumArgumentLabels > 0 && "Use other constructor");
+    : NumArgumentLabels(argumentLabelLocs.size()),
+      HasModuleSelectorLoc(moduleSelectorLoc.isValid())
+{
+  if (NumArgumentLabels == 0 && !HasModuleSelectorLoc) {
+    LocationInfo = baseNameLoc.getOpaquePointerValue();
+    return;
+  }
 
   // Copy the location information into permanent storage.
-  auto storedLocs = ctx.Allocate<SourceLoc>(NumArgumentLabels + 3);
+  auto storedLocs =
+      ctx.Allocate<SourceLoc>(FirstArgumentLabelIndex + NumArgumentLabels);
   storedLocs[BaseNameIndex] = baseNameLoc;
+  storedLocs[ModuleSelectorIndex] = moduleSelectorLoc;
   storedLocs[LParenIndex] = lParenLoc;
   storedLocs[RParenIndex] = rParenLoc;
   std::memcpy(storedLocs.data() + FirstArgumentLabelIndex,
diff --git a/lib/Sema/CSSimplify.cpp b/lib/Sema/CSSimplify.cpp
@@ -10278,8 +10278,8 @@ performMemberLookup(ConstraintKind constraintKind, DeclNameRef memberName,
     // high.
     if (auto *args = getArgumentList(memberLocator)) {
       SmallVector<Identifier, 4> scratch;
-      memberName.getFullName() = DeclName(ctx, memberName.getBaseName(),
-                                          args->getArgumentLabels(scratch));
+      memberName = memberName.withArgumentLabels(
+                                  ctx, args->getArgumentLabels(scratch));
     }
   }
 
diff --git a/test/NameLookup/module_selector.swift b/test/NameLookup/module_selector.swift
@@ -204,7 +204,7 @@ let mog: Never = fatalError()
 
 func localVarsCantBeAccessedByModuleSelector() {
   let mag: Int.Swift::Magnitude = main::mag
-  // expected-error@-1 {{use of local variable 'mag' before its declaration}}
+  // expected-error@-1 {{use of local variable 'main::mag' before its declaration}}
   // expected-note@-2 {{'mag' declared here}}
 
   let mog: Never = main::mog
@@ -267,8 +267,9 @@ func badModuleNames() {
   _ = "foo".NonexistentModule::count
 
   let x: NonexistentModule::MyType = NonexistentModule::MyType()
-  // expected-error@-1 {{cannot find type 'MyType' in scope}}
+  // expected-error@-1 {{cannot find type 'NonexistentModule::MyType' in scope}}
+  // expected-error@-2 {{cannot find 'NonexistentModule::MyType' in scope}}
 
   let y: A.NonexistentModule::MyChildType = fatalError()
-  // expected-error@-1 {{'MyChildType' is not a member type of struct 'ModuleSelectorTestingKit.A'}}
+  // expected-error@-1 {{'NonexistentModule::MyChildType' is not a member type of struct 'ModuleSelectorTestingKit.A'}}
 }

Original file line number	Diff line number	Diff line change
`@@ -10278,8 +10278,8 @@ performMemberLookup(ConstraintKind constraintKind, DeclNameRef memberName,`
`10278`	`10278`	`// high.`
`10279`	`10279`	`if (auto *args = getArgumentList(memberLocator)) {`
`10280`	`10280`	`SmallVector<Identifier, 4> scratch;`
`10281`		`- memberName.getFullName() = DeclName(ctx, memberName.getBaseName(),`
`10282`		`- args->getArgumentLabels(scratch));`
	`10281`	`+ memberName = memberName.withArgumentLabels(`
	`10282`	`+ ctx, args->getArgumentLabels(scratch));`
`10283`	`10283`	`}`
`10284`	`10284`	`}`
`10285`	`10285`