[NFC] Add representations for DeclNameRef/Loc with module selector

Author: beccadax

include/swift/AST/DeclNameLoc.h

@@ -29,59 +29,66 @@ class ASTContext;
 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);
   }
 
 public:
   /// Create an invalid declaration name location.
-  DeclNameLoc() : LocationInfo(0), NumArgumentLabels(0) { }
+  DeclNameLoc()
+      : LocationInfo(0), NumArgumentLabels(0), HasModuleSelectorLoc(false) { }
 
   /// Create declaration name location information for a base name.
   explicit DeclNameLoc(SourceLoc baseNameLoc)
     : LocationInfo(baseNameLoc.getOpaquePointerValue()),
-      NumArgumentLabels(0) { }
+      NumArgumentLabels(0), HasModuleSelectorLoc(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(); }
@@ -117,11 +124,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 {
@@ -130,9 +138,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());
   }
 };
 

include/swift/AST/Identifier.h

@@ -640,39 +640,86 @@ 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;
+
+  /// 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);
+    }
+
+    // Make vanilla new/delete illegal for SelectiveDeclNameRef.
+    void *operator new(size_t Bytes) = delete;
+    void operator delete(void *Data) = delete;
+
+    // Only allow allocation of SelectiveDeclNameRef using the allocator
+    // in ASTContext or by doing a placement new.
+    void *operator new(size_t Bytes, const ASTContext &C,
+                       unsigned Alignment = alignof(SelectiveDeclNameRef));
+    void *operator new(size_t Bytes, void *Mem) {
+      assert(Mem);
+      return Mem;
+    }
+  };
+
+  llvm::PointerUnion<DeclName, SelectiveDeclNameRef *> storage;
+
+  explicit DeclNameRef(void *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 *getOpaqueValue() const { return FullName.getOpaqueValue(); }
+  void *getOpaqueValue() const {
+    return storage.getOpaqueValue();
+  }
   static DeclNameRef getFromOpaqueValue(void *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(BaseName) { }
+    : storage(FullName) { }
 
   bool hasModuleSelector() const {
-    return false;
+    return storage.is<SelectiveDeclNameRef *>();
   }
 
   Identifier getModuleSelector() const {
-    return Identifier();
+    if (!hasModuleSelector())
+      return Identifier();
+
+    return storage.get<SelectiveDeclNameRef *>()->moduleSelector;
   }
 
   /// The name of the declaration being referenced.
   DeclName getFullName() const {
-    return FullName;
+    if (!hasModuleSelector())
+      return storage.get<DeclName>();
+
+    return storage.get<SelectiveDeclNameRef *>()->fullName;
   }
 
   /// The base name of the declaration being referenced.
@@ -784,7 +831,7 @@ class DeclNameRef {
 };
 
 inline DeclNameRef DeclNameRef::getFromOpaqueValue(void *p) {
-  return DeclNameRef(DeclName::getFromOpaqueValue(p));
+  return DeclNameRef(p);
 }
 
 inline DeclNameRef DeclNameRef::withoutArgumentLabels(ASTContext &C) const {
@@ -968,7 +1015,7 @@ namespace llvm {
     static inline swift::DeclNameRef getFromVoidPointer(void *ptr) {
       return swift::DeclNameRef::getFromOpaqueValue(ptr);
     }
-    enum { NumLowBitsAvailable = PointerLikeTypeTraits<swift::DeclName>::NumLowBitsAvailable };
+    enum { NumLowBitsAvailable = PointerLikeTypeTraits<swift::DeclName>::NumLowBitsAvailable - 1 };
   };
 
   // DeclNameRefs hash just like DeclNames.

lib/AST/ASTContext.cpp

@@ -469,6 +469,7 @@ struct ASTContext::Implementation {
   llvm::DenseMap<BuiltinIntegerWidth, BuiltinIntegerType*> IntegerTypes;
   llvm::FoldingSet<BuiltinVectorType> BuiltinVectorTypes;
   llvm::FoldingSet<DeclName::CompoundDeclName> CompoundNames;
+  llvm::FoldingSet<DeclNameRef::SelectiveDeclNameRef> SelectiveNameRefs;
   llvm::DenseMap<UUID, OpenedArchetypeType *> OpenedExistentialArchetypes;
   llvm::FoldingSet<IndexSubset> IndexSubsets;
   llvm::FoldingSet<AutoDiffDerivativeFunctionIdentifier>
@@ -4558,6 +4559,44 @@ DeclName::DeclName(ASTContext &C, DeclBaseName baseName,
   initialize(C, baseName, names);
 }
 
+// Only allow allocation of SelectiveDeclNameRefs using the allocator in
+// ASTContext.
+void *DeclNameRef::SelectiveDeclNameRef::
+operator new(size_t Bytes, const ASTContext &C, unsigned Alignment) {
+  return C.Allocate(Bytes, Alignment);
+}
+
+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 selectiveRef = new (C) 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,

lib/AST/DeclNameLoc.cpp

@@ -19,16 +19,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,

lib/Sema/CSSimplify.cpp

@@ -7192,8 +7192,7 @@ performMemberLookup(ConstraintKind constraintKind, DeclNameRef memberName,
     // anything else, because the cost of the general search is so
     // high.
     if (auto info = getArgumentInfo(memberLocator)) {
-      memberName.getFullName() = DeclName(ctx, memberName.getBaseName(),
-                                          info->Labels);
+      memberName = memberName.withArgumentLabels(ctx, info->Labels);
     }
   }