[Keypaths] Encode generic environment in the key-path pattern.

DougGregor · DougGregor · commit f759296cc89e · 2018-11-16T10:13:06.000-08:00
Extend the key-path pattern with a representation of the generic environment
of the key-path, which includes the generic parameters and generic
requirements of the environment.
diff --git a/include/swift/ABI/Metadata.h b/include/swift/ABI/Metadata.h
@@ -2611,6 +2611,57 @@ class TargetGenericRequirementDescriptor {
 using GenericRequirementDescriptor =
   TargetGenericRequirementDescriptor<InProcess>;
 
+template<typename Runtime>
+class TargetGenericEnvironment
+    : public swift::ABI::TrailingObjects<TargetGenericEnvironment<Runtime>,
+               uint16_t, GenericParamDescriptor,
+               TargetGenericRequirementDescriptor<Runtime>> {
+  using GenericRequirementDescriptor =
+    TargetGenericRequirementDescriptor<Runtime>;
+  using TrailingObjects =
+     swift::ABI::TrailingObjects<TargetGenericEnvironment<Runtime>,
+       uint16_t, GenericParamDescriptor, GenericRequirementDescriptor>;
+  friend TrailingObjects;
+
+  template<typename T>
+  using OverloadToken = typename TrailingObjects::template OverloadToken<T>;
+
+  size_t numTrailingObjects(OverloadToken<uint16_t>) const {
+    return Flags.getNumGenericParameterLevels();
+  }
+
+  size_t numTrailingObjects(OverloadToken<GenericParamDescriptor>) const {
+    return getGenericParameterCounts().back();
+  }
+
+  size_t numTrailingObjects(OverloadToken<GenericRequirementDescriptor>) const {
+    return Flags.getNumGenericRequirements();
+  }
+
+  GenericEnvironmentFlags Flags;
+
+public:
+  /// Retrieve the cumulative generic parameter counts at each level of genericity.
+  ArrayRef<uint16_t> getGenericParameterCounts() const {
+    return ArrayRef<uint16_t>(this->template getTrailingObjects<uint16_t>(),
+                              Flags.getNumGenericParameterLevels());
+  }
+
+  /// Retrieve the generic parameters descriptors.
+  ArrayRef<GenericParamDescriptor> getGenericParameters() const {
+    return ArrayRef<GenericParamDescriptor>(
+             this->template getTrailingObjects<GenericParamDescriptor>(),
+             getGenericParameterCounts().back());
+  }
+
+  /// Retrieve the generic requirements.
+  ArrayRef<GenericRequirementDescriptor> getGenericRequirements() const {
+    return ArrayRef<GenericRequirementDescriptor>(
+             this->template getTrailingObjects<GenericRequirementDescriptor>(),
+             Flags.getNumGenericRequirements());
+  }
+};
+
 /// CRTP class for a context descriptor that includes trailing generic
 /// context description.
 template<class Self,
diff --git a/include/swift/ABI/MetadataValues.h b/include/swift/ABI/MetadataValues.h
@@ -1508,6 +1508,45 @@ enum class GenericRequirementLayoutKind : uint32_t {
   Class = 0,
 };
 
+class GenericEnvironmentFlags {
+  uint32_t Value;
+
+  enum : uint32_t {
+    NumGenericParameterLevelsMask = 0xFFF,
+    NumGenericRequirementsShift = 12,
+    NumGenericRequirementsMask = 0xFFFF << NumGenericRequirementsShift,
+  };
+
+  constexpr explicit GenericEnvironmentFlags(uint32_t value) : Value(value) { }
+
+public:
+  constexpr GenericEnvironmentFlags() : Value(0) { }
+
+  constexpr GenericEnvironmentFlags
+  withNumGenericParameterLevels(uint16_t numGenericParameterLevels) const {
+    return GenericEnvironmentFlags((Value &~ NumGenericParameterLevelsMask)
+                                   | numGenericParameterLevels);
+  }
+
+  constexpr GenericEnvironmentFlags
+  withNumGenericRequirements(uint16_t numGenericRequirements) const {
+    return GenericEnvironmentFlags((Value &~ NumGenericParameterLevelsMask)
+             | (numGenericRequirements << NumGenericRequirementsShift));
+  }
+
+  constexpr unsigned getNumGenericParameterLevels() const {
+    return Value & NumGenericParameterLevelsMask;
+  }
+
+  constexpr unsigned getNumGenericRequirements() const {
+    return (Value & NumGenericRequirementsMask) >> NumGenericRequirementsShift;
+  }
+
+  constexpr uint32_t getIntValue() const {
+    return Value;
+  }
+};
+
 /// Flags used by generic metadata patterns.
 class GenericMetadataPatternFlags : public FlagSet<uint32_t> {
   enum {
diff --git a/lib/IRGen/GenKeyPath.cpp b/lib/IRGen/GenKeyPath.cpp
@@ -657,7 +657,7 @@ emitGeneratorForKeyPath(IRGenModule &IGM,
                         ArrayRef<GenericRequirement> requirements,
                         llvm::function_ref<void(IRGenFunction&,CanType)> emit) {
 
-  return IGM.getAddrOfStringForMetadataRef(name,
+  return IGM.getAddrOfStringForMetadataRef(name, /*alignment=*/2,
       /*shouldSetLowBit=*/true,
       [&](ConstantInitBuilder &B) {
         // Build a stub that loads the necessary bindings from the key path's
@@ -1171,6 +1171,9 @@ IRGenModule::getAddrOfKeyPathPattern(KeyPathPattern *pattern,
     fields.addInt32(0);
   }
 
+  // Add the generic environment.
+  fields.addRelativeAddressOrNull(
+    getAddrOfGenericEnvironment(pattern->getGenericSignature()));
   // Store type references for the root and leaf.
   fields.addRelativeAddress(
     emitMetadataGeneratorForKeyPath(*this, rootTy, genericEnv, requirements));
diff --git a/lib/IRGen/GenProto.cpp b/lib/IRGen/GenProto.cpp
@@ -67,6 +67,7 @@
 #include "GenericRequirement.h"
 #include "IRGenDebugInfo.h"
 #include "IRGenFunction.h"
+#include "IRGenMangler.h"
 #include "IRGenModule.h"
 #include "MetadataPath.h"
 #include "MetadataRequest.h"
@@ -3419,3 +3420,58 @@ llvm::Value *irgen::emitProtocolDescriptorRef(IRGenFunction &IGF,
 
   return val;
 }
+
+llvm::Constant *IRGenModule::getAddrOfGenericEnvironment(
+                                                CanGenericSignature signature) {
+  if (!signature)
+    return nullptr;
+
+  IRGenMangler mangler;
+  auto symbolName = mangler.mangleSymbolNameForGenericEnvironment(signature);
+  return getAddrOfStringForMetadataRef(symbolName, /*alignment=*/0, false,
+      [&] (ConstantInitBuilder &builder) -> ConstantInitFuture {
+        /// Collect the cumulative count of parameters at each level.
+        llvm::SmallVector<uint16_t, 4> genericParamCounts;
+        unsigned curDepth = 0;
+        unsigned genericParamCount = 0;
+        for (const auto gp : signature->getGenericParams()) {
+          if (curDepth != gp->getDepth()) {
+            genericParamCounts.push_back(genericParamCount);
+            curDepth = gp->getDepth();
+          }
+
+          ++genericParamCount;
+        }
+        genericParamCounts.push_back(genericParamCount);
+
+        auto flags = GenericEnvironmentFlags()
+          .withNumGenericParameterLevels(genericParamCounts.size())
+          .withNumGenericRequirements(signature->getRequirements().size());
+
+        ConstantStructBuilder fields = builder.beginStruct();
+        fields.setPacked(true);
+
+        // Flags
+        fields.addInt32(flags.getIntValue());
+
+        // Parameter counts.
+        for (auto count : genericParamCounts) {
+          fields.addInt16(count);
+        }
+
+        // Generic parameters.
+        signature->forEachParam([&](GenericTypeParamType *param,
+                                    bool canonical) {
+          fields.addInt(Int8Ty,
+                        GenericParamDescriptor(GenericParamKind::Type,
+                                               canonical,
+                                               false)
+                          .getIntValue());
+        });
+
+        // Generic requirements
+        irgen::addGenericRequirements(*this, fields, signature,
+                                      signature->getRequirements());
+        return fields.finishAndCreateFuture();
+      });
+}
diff --git a/lib/IRGen/IRGenMangler.cpp b/lib/IRGen/IRGenMangler.cpp
@@ -294,3 +294,11 @@ std::string IRGenMangler::mangleSymbolNameForKeyPathMetadata(
     assert(conformance.isInvalid() && "Unknown protocol conformance");
   return finalize();
 }
+
+std::string IRGenMangler::mangleSymbolNameForGenericEnvironment(
+                                              CanGenericSignature genericSig) {
+  beginManglingWithoutPrefix();
+  Buffer << "generic environment ";
+  appendGenericSignature(genericSig);
+  return finalize();
+}
diff --git a/lib/IRGen/IRGenMangler.h b/lib/IRGen/IRGenMangler.h
@@ -416,6 +416,8 @@ class IRGenMangler : public Mangle::ASTMangler {
                                            CanType type,
                                            ProtocolConformanceRef conformance);
 
+  std::string mangleSymbolNameForGenericEnvironment(
+                                                CanGenericSignature genericSig);
 protected:
   SymbolicMangling
   withSymbolicReferences(IRGenModule &IGM,
diff --git a/lib/IRGen/IRGenModule.h b/lib/IRGen/IRGenModule.h
@@ -1032,6 +1032,7 @@ class IRGenModule {
   ///
   /// \param symbolName The name of the symbol that describes the metadata
   /// being referenced.
+  /// \param alignment If non-zero, the alignment of the requested variable.
   /// \param shouldSetLowBit Whether to set the low bit of the result
   /// constant, which is used by some clients to indicate that the result is
   /// a mangled name.
@@ -1041,6 +1042,7 @@ class IRGenModule {
   /// \returns the address of the global variable describing this metadata.
   llvm::Constant *getAddrOfStringForMetadataRef(
       StringRef symbolName,
+      unsigned alignment,
       bool shouldSetLowBit,
       llvm::function_ref<ConstantInitFuture(ConstantInitBuilder &)> body);
 
@@ -1276,6 +1278,7 @@ private:                            \
   llvm::Constant *getAddrOfObjCModuleContextDescriptor();
   llvm::Constant *getAddrOfClangImporterModuleContextDescriptor();
   ConstantReference getAddrOfParentContextDescriptor(DeclContext *from);
+  llvm::Constant *getAddrOfGenericEnvironment(CanGenericSignature signature);
   llvm::Constant *getAddrOfProtocolRequirementsBaseDescriptor(
                                                   ProtocolDecl *proto);
   llvm::GlobalValue *defineProtocolRequirementsBaseDescriptor(
diff --git a/lib/IRGen/MetadataRequest.cpp b/lib/IRGen/MetadataRequest.cpp
@@ -220,6 +220,7 @@ MetadataDependency MetadataDependencyCollector::finish(IRGenFunction &IGF) {
 
 llvm::Constant *IRGenModule::getAddrOfStringForMetadataRef(
     StringRef symbolName,
+    unsigned alignment,
     bool shouldSetLowBit,
     llvm::function_ref<ConstantInitFuture (ConstantInitBuilder &)> body) {
   // Call this to form the return value.
@@ -251,7 +252,8 @@ llvm::Constant *IRGenModule::getAddrOfStringForMetadataRef(
     llvm::GlobalValue::HiddenVisibility,
     llvm::GlobalValue::DefaultStorageClass})
   .to(var);
-  var->setAlignment(2);
+  if (alignment)
+    var->setAlignment(alignment);
   setTrueConstGlobal(var);
   var->setSection(getReflectionTypeRefSectionName());
 
diff --git a/stdlib/public/core/KeyPath.swift b/stdlib/public/core/KeyPath.swift
@@ -2245,7 +2245,7 @@ internal var keyPathObjectHeaderSize: Int {
 }
 
 internal var keyPathPatternHeaderSize: Int {
-  return 12
+  return 16
 }
 
 // Runtime entry point to instantiate a key path object.
@@ -2323,7 +2323,7 @@ public func _swift_getKeyPath(pattern: UnsafeMutableRawPointer,
   }
 
   // Adopt the KVC string from the pattern.
-  let kvcStringBase = patternPtr.advanced(by: 8)
+  let kvcStringBase = patternPtr.advanced(by: 12)
   let kvcStringOffset = kvcStringBase.load(as: Int32.self)
 
   if kvcStringOffset == 0 {
@@ -2465,9 +2465,10 @@ internal struct KeyPathPatternComputedArguments {
 }
 
 internal protocol KeyPathPatternVisitor {
-  mutating func visitHeader(rootMetadataRef: MetadataReference,
+  mutating func visitHeader(genericEnvironment: UnsafeRawPointer?,
+                            rootMetadataRef: MetadataReference,
                             leafMetadataRef: MetadataReference,
-                            kvcCompatibilityString: UnsafeRawPointer)
+                            kvcCompatibilityString: UnsafeRawPointer?)
   mutating func visitStoredComponent(kind: KeyPathStructOrClass,
                                      mutable: Bool,
                                      offset: KeyPathPatternStoredOffset)
@@ -2517,14 +2518,17 @@ internal func _walkKeyPathPattern<W: KeyPathPatternVisitor>(
                                   _ pattern: UnsafeRawPointer,
                                   walker: inout W) {
   // Visit the header.
-  let rootMetadataRef = _loadRelativeAddress(at: pattern,
+  let genericEnvironment = _loadRelativeAddress(at: pattern,
+                                                as: UnsafeRawPointer.self)
+  let rootMetadataRef = _loadRelativeAddress(at: pattern, fromByteOffset: 4,
                                              as: MetadataReference.self)
-  let leafMetadataRef = _loadRelativeAddress(at: pattern, fromByteOffset: 4,
+  let leafMetadataRef = _loadRelativeAddress(at: pattern, fromByteOffset: 8,
                                              as: MetadataReference.self)
-  let kvcString = _loadRelativeAddress(at: pattern, fromByteOffset: 8,
+  let kvcString = _loadRelativeAddress(at: pattern, fromByteOffset: 12,
                                        as: UnsafeRawPointer.self)
 
-  walker.visitHeader(rootMetadataRef: rootMetadataRef,
+  walker.visitHeader(genericEnvironment: genericEnvironment,
+                     rootMetadataRef: rootMetadataRef,
                      leafMetadataRef: leafMetadataRef,
                      kvcCompatibilityString: kvcString)
 
@@ -2803,9 +2807,10 @@ internal struct GetKeyPathClassAndInstanceSizeFromPattern
     size = MemoryLayout<Int>._roundingUpToAlignment(size)
   }
 
-  mutating func visitHeader(rootMetadataRef: MetadataReference,
+  mutating func visitHeader(genericEnvironment: UnsafeRawPointer?,
+                            rootMetadataRef: MetadataReference,
                             leafMetadataRef: MetadataReference,
-                            kvcCompatibilityString: UnsafeRawPointer) {
+                            kvcCompatibilityString: UnsafeRawPointer?) {
     // Get the root and leaf type metadata so we can form the class type
     // for the entire key path.
     root = _resolveKeyPathMetadataReference(rootMetadataRef,
@@ -3035,9 +3040,10 @@ internal struct InstantiateKeyPathBuffer : KeyPathPatternVisitor {
     return oldValue
   }
 
-  mutating func visitHeader(rootMetadataRef: MetadataReference,
+  mutating func visitHeader(genericEnvironment: UnsafeRawPointer?,
+                            rootMetadataRef: MetadataReference,
                             leafMetadataRef: MetadataReference,
-                            kvcCompatibilityString: UnsafeRawPointer) {
+                            kvcCompatibilityString: UnsafeRawPointer?) {
   }
 
   mutating func visitStoredComponent(kind: KeyPathStructOrClass,
@@ -3261,13 +3267,16 @@ internal struct ValidatingInstantiateKeyPathBuffer: KeyPathPatternVisitor {
     origDest = self.instantiateVisitor.destData.baseAddress.unsafelyUnwrapped
   }
 
-  mutating func visitHeader(rootMetadataRef: MetadataReference,
+  mutating func visitHeader(genericEnvironment: UnsafeRawPointer?,
+                            rootMetadataRef: MetadataReference,
                             leafMetadataRef: MetadataReference,
-                            kvcCompatibilityString: UnsafeRawPointer) {
-    sizeVisitor.visitHeader(rootMetadataRef: rootMetadataRef,
+                            kvcCompatibilityString: UnsafeRawPointer?) {
+    sizeVisitor.visitHeader(genericEnvironment: genericEnvironment,
+                            rootMetadataRef: rootMetadataRef,
                             leafMetadataRef: leafMetadataRef,
                             kvcCompatibilityString: kvcCompatibilityString)
-    instantiateVisitor.visitHeader(rootMetadataRef: rootMetadataRef,
+    instantiateVisitor.visitHeader(genericEnvironment: genericEnvironment,
+                                 rootMetadataRef: rootMetadataRef,
                                  leafMetadataRef: leafMetadataRef,
                                  kvcCompatibilityString: kvcCompatibilityString)
   }
diff --git a/test/IRGen/keypaths.sil b/test/IRGen/keypaths.sil
@@ -195,6 +195,7 @@ sil_vtable C2 {}
 // -- %j: Gen<A>.y
 // CHECK: [[KP_J:@keypath(\..*)?]] = private global <{ {{.*}} }> <{
 // CHECK-SAME: i32 0
+// CHECK-SAME: @"generic environment l"
 // CHECK-SAME: @"keypath_get_type
 // CHECK-SAME: @"keypath_get_type
 //             -- size 8
@@ -204,6 +205,9 @@ sil_vtable C2 {}
 // CHECK-32-SAME: i32 20 }>
 // CHECK-64-SAME: i32 36 }>
 
+// CHECK-LABEL: @"generic environment SHRzSHR_r0_l" = linkonce_odr hidden constant
+// CHECK-SAME: i32 8192, i16 2, i8 -128, i8 -128, i32 128
+
 // -- %t
 // CHECK: [[KP_T:@keypath(\..*)?]] = private global <{ {{.*}} }> <{ {{.*}} i32 1, {{.*}} @"got.$s8keypaths1GV1xxvpMV"
 // CHECK-SAME:   @"keypath_get_type
