Generalize Metadata Lookup for Shapes

Author: CodaFi

stdlib/public/runtime/MetadataLookup.cpp

@@ -2280,7 +2280,7 @@ static void installGetClassHook() {
 unsigned SubstGenericParametersFromMetadata::
 buildDescriptorPath(const ContextDescriptor *context,
                     Demangler &borrowFrom) const {
-  assert(sourceIsMetadata);
+  assert(sourceKind == SourceKind::Metadata);
 
   // Terminating condition: we don't have a context.
   if (!context)
@@ -2371,20 +2371,76 @@ buildEnvironmentPath(
   return totalKeyParamCount;
 }
 
+unsigned SubstGenericParametersFromMetadata::buildShapePath(
+    const TargetExtendedExistentialTypeShape<InProcess> *shape) const {
+  unsigned totalParamCount = 0;
+  unsigned totalKeyParamCount = 0;
+
+  auto genSig = shape->getGeneralizationSignature();
+  if (!genSig.getParams().empty()) {
+    bool hasNonKeyGenericParams = false;
+    unsigned numKeyGenericParamsHere = 0;
+    for (const auto &gp : genSig.getParams()) {
+      if (gp.hasKeyArgument())
+        numKeyGenericParamsHere++;
+      else
+        hasNonKeyGenericParams = true;
+    }
+    totalParamCount += numKeyGenericParamsHere;
+    totalKeyParamCount += numKeyGenericParamsHere;
+    descriptorPath.push_back(PathElement{genSig.getParams(), totalParamCount,
+                                         /*numKeyGenericParamsInParent*/ 0,
+                                         numKeyGenericParamsHere,
+                                         hasNonKeyGenericParams});
+  }
+
+  const unsigned genSigParamCount = genSig.getParams().size();
+  auto reqSig = shape->getRequirementSignature();
+  assert(reqSig.getParams().size() > genSig.getParams().size());
+  {
+    bool hasNonKeyGenericParams = false;
+    unsigned numKeyGenericParamsHere = 0;
+    auto remainingParams = reqSig.getParams().drop_front(genSig.getParams().size());
+    for (const auto &gp : remainingParams) {
+      if (gp.hasKeyArgument())
+        numKeyGenericParamsHere++;
+      else
+        hasNonKeyGenericParams = true;
+    }
+    totalParamCount += numKeyGenericParamsHere;
+    totalKeyParamCount += numKeyGenericParamsHere;
+    descriptorPath.push_back(PathElement{remainingParams,
+                                         totalParamCount,
+                                         genSigParamCount,
+                                         numKeyGenericParamsHere,
+                                         hasNonKeyGenericParams});
+  }
+
+  return totalKeyParamCount;
+}
+
 void SubstGenericParametersFromMetadata::setup() const {
   if (!descriptorPath.empty())
     return;
 
-  if (sourceIsMetadata && baseContext) {
+  switch (sourceKind) {
+  case SourceKind::Metadata: {
+    assert(baseContext);
     DemanglerForRuntimeTypeResolution<StackAllocatedDemangler<2048>> demangler;
     numKeyGenericParameters = buildDescriptorPath(baseContext, demangler);
     return;
   }
-
-  if (!sourceIsMetadata && environment) {
+  case SourceKind::Environment: {
+    assert(environment);
     numKeyGenericParameters = buildEnvironmentPath(environment);
     return;
   }
+  case SourceKind::Shape: {
+    assert(shape);
+    numKeyGenericParameters = buildShapePath(shape);
+    return;
+  }
+  }
 }
 
 const Metadata *

stdlib/public/runtime/Private.h

@@ -298,11 +298,17 @@ class TypeInfo {
   /// types.
   class SWIFT_RUNTIME_LIBRARY_VISIBILITY SubstGenericParametersFromMetadata {
     /// Whether the source is metadata (vs. a generic environment);
-    const bool sourceIsMetadata;
+    enum class SourceKind {
+      Metadata,
+      Environment,
+      Shape,
+    };
+    const SourceKind sourceKind;
 
     union {
       const TargetContextDescriptor<InProcess> *baseContext;
       const TargetGenericEnvironment<InProcess> *environment;
+      const TargetExtendedExistentialTypeShape<InProcess> *shape;
     };
 
     /// The generic arguments.
@@ -344,31 +350,40 @@ class TypeInfo {
     unsigned buildEnvironmentPath(
                const TargetGenericEnvironment<InProcess> *environment) const;
 
+    unsigned buildShapePath(
+        const TargetExtendedExistentialTypeShape<InProcess> *shape) const;
+
     // Set up the state we need to compute substitutions.
     void setup() const;
 
   public:
     /// Produce substitutions entirely from the given metadata.
     explicit SubstGenericParametersFromMetadata(const Metadata *base)
-      : sourceIsMetadata(true), baseContext(base->getTypeContextDescriptor()),
-        genericArgs(base ? (const void * const *)base->getGenericArgs()
-                         : nullptr) { }
-    
+        : sourceKind(SourceKind::Metadata),
+          baseContext(base->getTypeContextDescriptor()),
+          genericArgs(base ? (const void *const *)base->getGenericArgs()
+                           : nullptr) {}
+
     /// Produce substitutions from the given instantiation arguments for the
     /// given context.
     explicit SubstGenericParametersFromMetadata(const ContextDescriptor *base,
-                                                const void * const *args)
-      : sourceIsMetadata(true), baseContext(base), genericArgs(args)
-    {}
+                                                const void *const *args)
+        : sourceKind(SourceKind::Metadata), baseContext(base),
+          genericArgs(args) {}
 
     /// Produce substitutions from the given instantiation arguments for the
     /// given generic environment.
     explicit SubstGenericParametersFromMetadata(
-               const TargetGenericEnvironment<InProcess> *environment,
-               const void * const *arguments)
-      : sourceIsMetadata(false), environment(environment),
-        genericArgs(arguments) { }
-    
+        const TargetGenericEnvironment<InProcess> *environment,
+        const void *const *arguments)
+        : sourceKind(SourceKind::Environment), environment(environment),
+          genericArgs(arguments) {}
+
+    explicit SubstGenericParametersFromMetadata(
+        const TargetExtendedExistentialTypeShape<InProcess> *shape,
+        const void *const *arguments)
+        : sourceKind(SourceKind::Shape), shape(shape), genericArgs(arguments) {}
+
     const void * const *getGenericArgs() const { return genericArgs; }
 
     const Metadata *getMetadata(unsigned depth, unsigned index) const;