Merge two fields into a PointerUnion in SILDeclRef to save space

Author: aschwaighofer

include/swift/SIL/SILDeclRef.h

@@ -150,15 +150,28 @@ struct SILDeclRef {
   unsigned isForeign : 1;
   /// The default argument index for a default argument getter.
   unsigned defaultArgIndex : 10;
+
+  PointerUnion<AutoDiffDerivativeFunctionIdentifier *,
+               const GenericSignatureImpl *>
+      pointer;
+
   /// The derivative function identifier.
-  AutoDiffDerivativeFunctionIdentifier *derivativeFunctionIdentifier = nullptr;
+  AutoDiffDerivativeFunctionIdentifier * getDerivativeFunctionIdentifier() const {
+    if (!pointer.is<AutoDiffDerivativeFunctionIdentifier *>())
+      return nullptr;
+    return pointer.get<AutoDiffDerivativeFunctionIdentifier *>();
+  }
 
-  GenericSignature specializedSignature;
+  GenericSignature getSpecializedSignature() const {
+    if (!pointer.is<const GenericSignatureImpl *>())
+      return GenericSignature();
+    else
+      return GenericSignature(pointer.get<const GenericSignatureImpl *>());
+  }
 
   /// Produces a null SILDeclRef.
   SILDeclRef()
-      : loc(), kind(Kind::Func), isForeign(0), defaultArgIndex(0),
-        derivativeFunctionIdentifier(nullptr) {}
+      : loc(), kind(Kind::Func), isForeign(0), defaultArgIndex(0) {}
 
   /// Produces a SILDeclRef of the given kind for the given decl.
   explicit SILDeclRef(
@@ -294,7 +307,7 @@ struct SILDeclRef {
     return loc.getOpaqueValue() == rhs.loc.getOpaqueValue() &&
            kind == rhs.kind && isForeign == rhs.isForeign &&
            defaultArgIndex == rhs.defaultArgIndex &&
-           derivativeFunctionIdentifier == rhs.derivativeFunctionIdentifier;
+           pointer == rhs.pointer;
   }
   bool operator!=(SILDeclRef rhs) const {
     return !(*this == rhs);
@@ -309,7 +322,7 @@ struct SILDeclRef {
   /// decl.
   SILDeclRef asForeign(bool foreign = true) const {
     return SILDeclRef(loc.getOpaqueValue(), kind, foreign, defaultArgIndex,
-                      derivativeFunctionIdentifier);
+                      pointer.get<AutoDiffDerivativeFunctionIdentifier *>());
   }
 
   /// Returns the entry point for the corresponding autodiff derivative
@@ -318,16 +331,16 @@ struct SILDeclRef {
       AutoDiffDerivativeFunctionIdentifier *derivativeId) const {
     assert(derivativeId);
     SILDeclRef declRef = *this;
-    declRef.derivativeFunctionIdentifier = derivativeId;
+    declRef.pointer = derivativeId;
     return declRef;
   }
 
   /// Returns the entry point for the original function corresponding to an
   /// autodiff derivative function.
   SILDeclRef asAutoDiffOriginalFunction() const {
-    assert(derivativeFunctionIdentifier);
+    assert(pointer.get<AutoDiffDerivativeFunctionIdentifier *>());
     SILDeclRef declRef = *this;
-    declRef.derivativeFunctionIdentifier = nullptr;
+    declRef.pointer = (AutoDiffDerivativeFunctionIdentifier *)nullptr;
     return declRef;
   }
 
@@ -405,13 +418,14 @@ struct SILDeclRef {
   bool canBeDynamicReplacement() const;
 
   bool isAutoDiffDerivativeFunction() const {
-    return derivativeFunctionIdentifier != nullptr;
+    return pointer.is<AutoDiffDerivativeFunctionIdentifier *>() &&
+           pointer.get<AutoDiffDerivativeFunctionIdentifier *>() != nullptr;
   }
 
   AutoDiffDerivativeFunctionIdentifier *
   getAutoDiffDerivativeFunctionIdentifier() const {
     assert(isAutoDiffDerivativeFunction());
-    return derivativeFunctionIdentifier;
+    return pointer.get<AutoDiffDerivativeFunctionIdentifier *>();
   }
 
 private:
@@ -422,7 +436,7 @@ struct SILDeclRef {
                       AutoDiffDerivativeFunctionIdentifier *derivativeId)
       : loc(Loc::getFromOpaqueValue(opaqueLoc)), kind(kind),
         isForeign(isForeign), defaultArgIndex(defaultArgIndex),
-        derivativeFunctionIdentifier(derivativeId) {}
+        pointer(derivativeId) {}
 };
 
 inline llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, SILDeclRef C) {
@@ -457,7 +471,7 @@ template<> struct DenseMapInfo<swift::SILDeclRef> {
                     ? UnsignedInfo::getHashValue(Val.defaultArgIndex)
                     : 0;
     unsigned h4 = UnsignedInfo::getHashValue(Val.isForeign);
-    unsigned h5 = PointerInfo::getHashValue(Val.derivativeFunctionIdentifier);
+    unsigned h5 = PointerInfo::getHashValue(Val.pointer.getOpaqueValue());
     return h1 ^ (h2 << 4) ^ (h3 << 9) ^ (h4 << 7) ^ (h5 << 11);
   }
   static bool isEqual(swift::SILDeclRef const &LHS,

lib/IRGen/GenMeta.cpp

@@ -1683,7 +1683,7 @@ namespace {
     void emitNonoverriddenMethod(SILDeclRef fn) {
       // TODO: Derivative functions do not distinguish themselves in the mangled
       // names of method descriptor symbols yet, causing symbol name collisions.
-      if (fn.derivativeFunctionIdentifier)
+      if (fn.getDerivativeFunctionIdentifier())
         return;
 
      HasNonoverriddenMethods = true;

lib/SIL/IR/SILDeclRef.cpp

@@ -119,10 +119,11 @@ bool swift::requiresForeignEntryPoint(ValueDecl *vd) {
 SILDeclRef::SILDeclRef(ValueDecl *vd, SILDeclRef::Kind kind, bool isForeign,
                        AutoDiffDerivativeFunctionIdentifier *derivativeId)
     : loc(vd), kind(kind), isForeign(isForeign), defaultArgIndex(0),
-      derivativeFunctionIdentifier(derivativeId) {}
+      pointer(derivativeId) {}
 
 SILDeclRef::SILDeclRef(SILDeclRef::Loc baseLoc, bool asForeign)
-    : defaultArgIndex(0), derivativeFunctionIdentifier(nullptr) {
+    : defaultArgIndex(0),
+      pointer((AutoDiffDerivativeFunctionIdentifier *)nullptr) {
   if (auto *vd = baseLoc.dyn_cast<ValueDecl*>()) {
     if (auto *fd = dyn_cast<FuncDecl>(vd)) {
       // Map FuncDecls directly to Func SILDeclRefs.
@@ -164,7 +165,7 @@ SILDeclRef::SILDeclRef(SILDeclRef::Loc baseLoc, bool asForeign)
 SILDeclRef::SILDeclRef(SILDeclRef::Loc baseLoc,
                        GenericSignature prespecializedSig)
     : SILDeclRef(baseLoc, false) {
-  specializedSignature = prespecializedSig;
+  pointer = prespecializedSig.getPointer();
 }
 
 Optional<AnyFunctionRef> SILDeclRef::getAnyFunctionRef() const {
@@ -232,7 +233,7 @@ bool SILDeclRef::isImplicit() const {
 SILLinkage SILDeclRef::getLinkage(ForDefinition_t forDefinition) const {
 
   // Prespecializations are public.
-  if (specializedSignature) {
+  if (getSpecializedSignature()) {
     return SILLinkage::Public;
   }
 
@@ -678,6 +679,7 @@ std::string SILDeclRef::mangle(ManglingKind MKind) const {
   using namespace Mangle;
   ASTMangler mangler;
 
+  auto *derivativeFunctionIdentifier = getDerivativeFunctionIdentifier();
   if (derivativeFunctionIdentifier) {
     std::string originalMangled = asAutoDiffOriginalFunction().mangle(MKind);
     auto *silParameterIndices = autodiff::getLoweredParameterIndices(
@@ -716,14 +718,15 @@ std::string SILDeclRef::mangle(ManglingKind MKind) const {
   }
 
   // Mangle prespecializations.
-  if (specializedSignature) {
+  if (getSpecializedSignature()) {
     SILDeclRef nonSpecializedDeclRef = *this;
-    nonSpecializedDeclRef.specializedSignature = GenericSignature();
+    nonSpecializedDeclRef.pointer =
+        (AutoDiffDerivativeFunctionIdentifier *)nullptr;
     auto mangledNonSpecializedString = nonSpecializedDeclRef.mangle();
     auto *funcDecl = cast<AbstractFunctionDecl>(getDecl());
     auto genericSig = funcDecl->getGenericSignature();
     return GenericSpecializationMangler::manglePrespecialization(
-        mangledNonSpecializedString, genericSig, specializedSignature);
+        mangledNonSpecializedString, genericSig, getSpecializedSignature());
   }
 
   ASTMangler::SymbolKind SKind = ASTMangler::SymbolKind::Default;
@@ -818,7 +821,7 @@ std::string SILDeclRef::mangle(ManglingKind MKind) const {
 // Returns true if the given JVP/VJP SILDeclRef requires a new vtable entry.
 // FIXME(TF-1213): Also consider derived declaration `@derivative` attributes.
 static bool derivativeFunctionRequiresNewVTableEntry(SILDeclRef declRef) {
-  assert(declRef.derivativeFunctionIdentifier &&
+  assert(declRef.getDerivativeFunctionIdentifier() &&
          "Expected a derivative function SILDeclRef");
   auto overridden = declRef.getOverridden();
   if (!overridden)
@@ -828,7 +831,7 @@ static bool derivativeFunctionRequiresNewVTableEntry(SILDeclRef declRef) {
       declRef.getDecl()->getAttrs().getAttributes<DifferentiableAttr>(),
       [&](const DifferentiableAttr *derivedDiffAttr) {
         return derivedDiffAttr->getParameterIndices() ==
-               declRef.derivativeFunctionIdentifier->getParameterIndices();
+               declRef.getDerivativeFunctionIdentifier()->getParameterIndices();
       });
   assert(derivedDiffAttr && "Expected `@differentiable` attribute");
   // Otherwise, if the base `@differentiable` attribute specifies a derivative
@@ -838,7 +841,7 @@ static bool derivativeFunctionRequiresNewVTableEntry(SILDeclRef declRef) {
       overridden.getDecl()->getAttrs().getAttributes<DifferentiableAttr>();
   for (auto *baseDiffAttr : baseDiffAttrs) {
     if (baseDiffAttr->getParameterIndices() ==
-        declRef.derivativeFunctionIdentifier->getParameterIndices())
+        declRef.getDerivativeFunctionIdentifier()->getParameterIndices())
       return false;
   }
   // Otherwise, if there is no base `@differentiable` attribute exists, then a
@@ -847,7 +850,7 @@ static bool derivativeFunctionRequiresNewVTableEntry(SILDeclRef declRef) {
 }
 
 bool SILDeclRef::requiresNewVTableEntry() const {
-  if (derivativeFunctionIdentifier)
+  if (getDerivativeFunctionIdentifier())
     if (derivativeFunctionRequiresNewVTableEntry(*this))
       return true;
   if (!hasDecl())
@@ -928,15 +931,16 @@ SILDeclRef SILDeclRef::getNextOverriddenVTableEntry() const {
 
     // JVPs/VJPs are overridden only if the base declaration has a
     // `@differentiable` attribute with the same parameter indices.
-    if (derivativeFunctionIdentifier) {
+    if (getDerivativeFunctionIdentifier()) {
       auto overriddenAttrs =
           overridden.getDecl()->getAttrs().getAttributes<DifferentiableAttr>();
       for (const auto *attr : overriddenAttrs) {
         if (attr->getParameterIndices() !=
-            derivativeFunctionIdentifier->getParameterIndices())
+            getDerivativeFunctionIdentifier()->getParameterIndices())
           continue;
-        auto *overriddenDerivativeId = overridden.derivativeFunctionIdentifier;
-        overridden.derivativeFunctionIdentifier =
+        auto *overriddenDerivativeId =
+            overridden.getDerivativeFunctionIdentifier();
+        overridden.pointer =
             AutoDiffDerivativeFunctionIdentifier::get(
                 overriddenDerivativeId->getKind(),
                 overriddenDerivativeId->getParameterIndices(),

lib/SIL/IR/SILFunctionType.cpp

@@ -3266,7 +3266,7 @@ TypeConverter::getConstantInfo(TypeExpansionContext expansion,
   //   preserving SIL typing invariants.
   //
   // Always use (ad) to compute lowered derivative function types.
-  if (auto *derivativeId = constant.derivativeFunctionIdentifier) {
+  if (auto *derivativeId = constant.getDerivativeFunctionIdentifier()) {
     // Get lowered original function type.
     auto origFnConstantInfo = getConstantInfo(
         TypeExpansionContext::minimal(), constant.asAutoDiffOriginalFunction());

lib/SIL/IR/SILPrinter.cpp

@@ -349,19 +349,19 @@ void SILDeclRef::print(raw_ostream &OS) const {
   if (isForeign)
     OS << (isDot ? '.' : '!')  << "foreign";
 
-  if (derivativeFunctionIdentifier) {
+  if (getDerivativeFunctionIdentifier()) {
     OS << ((isDot || isForeign) ? '.' : '!');
-    switch (derivativeFunctionIdentifier->getKind()) {
+    switch (getDerivativeFunctionIdentifier()->getKind()) {
     case AutoDiffDerivativeFunctionKind::JVP:
       OS << "jvp.";
       break;
     case AutoDiffDerivativeFunctionKind::VJP:
       OS << "vjp.";
       break;
     }
-    OS << derivativeFunctionIdentifier->getParameterIndices()->getString();
+    OS << getDerivativeFunctionIdentifier()->getParameterIndices()->getString();
     if (auto derivativeGenSig =
-            derivativeFunctionIdentifier->getDerivativeGenericSignature()) {
+            getDerivativeFunctionIdentifier()->getDerivativeGenericSignature()) {
       OS << "." << derivativeGenSig;
     }
   }

lib/SIL/IR/TypeLowering.cpp

@@ -2520,7 +2520,7 @@ getFunctionInterfaceTypeWithCaptures(TypeConverter &TC,
 }
 
 CanAnyFunctionType TypeConverter::makeConstantInterfaceType(SILDeclRef c) {
-  if (auto *derivativeId = c.derivativeFunctionIdentifier) {
+  if (auto *derivativeId = c.getDerivativeFunctionIdentifier()) {
     auto originalFnTy =
         makeConstantInterfaceType(c.asAutoDiffOriginalFunction());
     auto *derivativeFnTy = originalFnTy->getAutoDiffDerivativeFunctionType(

lib/SILGen/SILGenPoly.cpp

@@ -4577,7 +4577,7 @@ getWitnessFunctionRef(SILGenFunction &SGF,
                       SILLocation loc) {
   switch (witnessKind) {
   case WitnessDispatchKind::Static:
-    if (auto *derivativeId = witness.derivativeFunctionIdentifier) {
+    if (auto *derivativeId = witness.getDerivativeFunctionIdentifier()) {
       auto originalFn =
           SGF.emitGlobalFunctionRef(loc, witness.asAutoDiffOriginalFunction());
       auto *loweredParamIndices = autodiff::getLoweredParameterIndices(
@@ -4594,7 +4594,7 @@ getWitnessFunctionRef(SILGenFunction &SGF,
     }
     return SGF.emitGlobalFunctionRef(loc, witness);
   case WitnessDispatchKind::Dynamic:
-    assert(!witness.derivativeFunctionIdentifier);
+    assert(!witness.getDerivativeFunctionIdentifier());
     return SGF.emitDynamicMethodRef(loc, witness, witnessFTy).getValue();
   case WitnessDispatchKind::Witness: {
     auto typeAndConf =
@@ -4609,7 +4609,7 @@ getWitnessFunctionRef(SILGenFunction &SGF,
     // If `witness` is a derivative function `SILDeclRef`, replace the
     // derivative function identifier's generic signature with the witness thunk
     // substitution map's generic signature.
-    if (auto *derivativeId = witness.derivativeFunctionIdentifier) {
+    if (auto *derivativeId = witness.getDerivativeFunctionIdentifier()) {
       auto *newDerivativeId = AutoDiffDerivativeFunctionIdentifier::get(
           derivativeId->getKind(), derivativeId->getParameterIndices(),
           witnessSubs.getGenericSignature(), SGF.getASTContext());

lib/SILGen/SILGenThunk.cpp

@@ -175,7 +175,7 @@ getOrCreateReabstractionThunk(CanSILFunctionType thunkType,
 
 SILFunction *SILGenModule::getOrCreateAutoDiffClassMethodThunk(
     SILDeclRef derivativeFnDeclRef, CanSILFunctionType constantTy) {
-  auto *derivativeId = derivativeFnDeclRef.derivativeFunctionIdentifier;
+  auto *derivativeId = derivativeFnDeclRef.getDerivativeFunctionIdentifier();
   assert(derivativeId);
   auto *derivativeFnDecl = derivativeFnDeclRef.getDecl();
 

lib/SILGen/SILGenType.cpp

@@ -90,7 +90,7 @@ SILGenModule::emitVTableMethod(ClassDecl *theClass,
     implFn = getDynamicThunk(
         derived, Types.getConstantInfo(TypeExpansionContext::minimal(), derived)
                      .SILFnType);
-  } else if (auto *derivativeId = derived.derivativeFunctionIdentifier) {
+  } else if (auto *derivativeId = derived.getDerivativeFunctionIdentifier()) {
     // For JVP/VJP methods, create a vtable entry thunk. The thunk contains an
     // `differentiable_function` instruction, which is later filled during the
     // differentiation transform.
@@ -168,7 +168,7 @@ SILGenModule::emitVTableMethod(ClassDecl *theClass,
         base.kind == SILDeclRef::Kind::Allocator);
     }
     // TODO(TF-685): Use proper autodiff thunk mangling.
-    if (auto *derivativeId = derived.derivativeFunctionIdentifier) {
+    if (auto *derivativeId = derived.getDerivativeFunctionIdentifier()) {
       switch (derivativeId->getKind()) {
       case AutoDiffDerivativeFunctionKind::JVP:
         name += "_jvp";
@@ -743,7 +743,7 @@ SILFunction *SILGenModule::emitProtocolWitness(
   std::string nameBuffer =
       NewMangler.mangleWitnessThunk(manglingConformance, requirement.getDecl());
   // TODO(TF-685): Proper mangling for derivative witness thunks.
-  if (auto *derivativeId = requirement.derivativeFunctionIdentifier) {
+  if (auto *derivativeId = requirement.getDerivativeFunctionIdentifier()) {
     std::string kindString;
     switch (derivativeId->getKind()) {
     case AutoDiffDerivativeFunctionKind::JVP: