Merge pull request #19486 from DougGregor/runtime-metadata-subst-same-type-extensions

 [Runtime] Handle non-key arguments when substituting from metadata.

Author: DougGregor

stdlib/public/runtime/MetadataLookup.cpp

@@ -798,6 +798,23 @@ Optional<unsigned> findAssociatedTypeByName(const ProtocolDescriptor *protocol,
   swift_runtime_unreachable("associated type names don't line up");
 }
 
+/// Retrieve the generic parameters introduced in this context.
+static ArrayRef<GenericParamDescriptor> getLocalGenericParams(
+    const ContextDescriptor *context) {
+  if (!context->isGeneric())
+    return { };
+
+  // Determine where to start looking at generic parameters.
+  unsigned startParamIndex;
+  if (auto parent = context->Parent.get())
+    startParamIndex = parent->getNumGenericParams();
+  else
+    startParamIndex = 0;
+
+  auto genericContext = context->getGenericContext();
+  return genericContext->getGenericParams().slice(startParamIndex);
+}
+
 /// Constructs metadata by decoding a mangled type name, for use with
 /// \c TypeDecoder.
 class DecodedMetadataBuilder {
@@ -1223,5 +1240,108 @@ swift_getTypeByMangledNameImpl(const char *typeNameStart, size_t typeNameLength,
   return swift_checkMetadataState(MetadataState::Complete, metadata).Value;
 }
 
+unsigned SubstGenericParametersFromMetadata::
+buildDescriptorPath(const ContextDescriptor *context) const {
+  // Terminating condition: we don't have a context.
+  if (!context)
+    return 0;
+
+  // Add the parent's contributino to the descriptor path.
+  unsigned numKeyGenericParamsInParent =
+    buildDescriptorPath(context->Parent.get());
+
+  // If this context is non-generic, we're done.
+  if (!context->isGeneric())
+    return numKeyGenericParamsInParent;
+
+  // Count the number of key generic params at this level.
+  unsigned numKeyGenericParamsHere = 0;
+  bool hasNonKeyGenericParams = false;
+  for (const auto &genericParam : getLocalGenericParams(context)) {
+    if (genericParam.hasKeyArgument())
+      ++numKeyGenericParamsHere;
+    else
+      hasNonKeyGenericParams = true;
+  }
+
+  // Form the path element.
+  descriptorPath.push_back(PathElement{context, numKeyGenericParamsInParent,
+                                       numKeyGenericParamsHere,
+                                       hasNonKeyGenericParams});
+  return numKeyGenericParamsInParent + numKeyGenericParamsHere;
+}
+
+void SubstGenericParametersFromMetadata::setup() const {
+  if (!descriptorPath.empty() || !base)
+    return;
+
+  buildDescriptorPath(base->getTypeContextDescriptor());
+}
+
+const Metadata *
+SubstGenericParametersFromMetadata::operator()(unsigned flatIndex) const {
+  // On first access, compute the descriptor path.
+  setup();
+
+  // Find the depth at which this parameter occurs.
+  unsigned depth = descriptorPath.size();
+  unsigned index = flatIndex;
+  for (const auto &pathElement : descriptorPath) {
+    // If the flat index is beyond the element at this position, we're done.
+    if (flatIndex >= pathElement.context->getNumGenericParams()) {
+      // Subtract off the number of parameters.
+      index -= pathElement.context->getNumGenericParams();
+      break;
+    }
+
+    --depth;
+  }
+
+  // Perform the access based on depth/index.
+  return (*this)(depth, index);
+}
+
+const Metadata *
+SubstGenericParametersFromMetadata::operator()(
+                                        unsigned depth, unsigned index) const {
+  // On first access, compute the descriptor path.
+  setup();
+
+  // If the depth is too great, there is nothing to do.
+  if (depth >= descriptorPath.size())
+    return nullptr;
+
+  /// Retrieve the descriptor path element at this depth.
+  auto &pathElement = descriptorPath[depth];
+  auto currentContext = pathElement.context;
+
+  // Check whether the index is clearly out of bounds.
+  if (index >= currentContext->getNumGenericParams())
+    return nullptr;
+
+  // Compute the flat index.
+  unsigned flatIndex = pathElement.numKeyGenericParamsInParent;
+  if (pathElement.hasNonKeyGenericParams > 0) {
+    // We have non-key generic parameters at this level, so the index needs to
+    // be checked more carefully.
+    auto genericParams = getLocalGenericParams(currentContext);
+
+    // Make sure that the requested parameter itself has a key argument.
+    if (!genericParams[index].hasKeyArgument())
+      return nullptr;
+
+    // Increase the flat index for each parameter with a key argument, up to
+    // the given index.
+    for (const auto &genericParam : genericParams.slice(0, index)) {
+      if (genericParam.hasKeyArgument())
+        ++flatIndex;
+    }
+  } else {
+    flatIndex += index;
+  }
+
+  return base->getGenericArgs()[flatIndex];
+}
+
 #define OVERRIDE_METADATALOOKUP COMPATIBILITY_OVERRIDE
 #include "CompatibilityOverride.def"

stdlib/public/runtime/Private.h

@@ -235,6 +235,51 @@ class TypeInfo {
   using SubstGenericParameterFn =
     llvm::function_ref<const Metadata *(unsigned depth, unsigned index)>;
 
+  /// Function object that produces substitutions for the generic parameters
+  /// that occur within a mangled name, using the generic arguments from
+  /// the given metadata.
+  ///
+  /// Use with \c _getTypeByMangledName to decode potentially-generic types.
+  class SWIFT_RUNTIME_LIBRARY_VISIBILITY SubstGenericParametersFromMetadata {
+    const Metadata *base;
+
+    /// An element in the descriptor path.
+    struct PathElement {
+      /// The context described by this path element.
+      const ContextDescriptor *context;
+
+      /// The number of key parameters in the parent.
+      unsigned numKeyGenericParamsInParent;
+
+      /// The number of key parameters locally introduced here.
+      unsigned numKeyGenericParamsHere;
+
+      /// Whether this context has any non-key generic parameters.
+      bool hasNonKeyGenericParams;
+    };
+
+    /// Information about the generic context descriptors that make up \c
+    /// descriptor, from the outermost to the innermost.
+    mutable std::vector<PathElement> descriptorPath;
+
+    /// Builds the descriptor path.
+    ///
+    /// \returns a pair containing the number of key generic parameters in
+    /// the path up to this point.
+    unsigned buildDescriptorPath(const ContextDescriptor *context) 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)
+      : base(base) { }
+
+    const Metadata *operator()(unsigned flatIndex) const;
+    const Metadata *operator()(unsigned depth, unsigned index) const;
+  };
+
   /// Retrieve the type metadata described by the given type name.
   ///
   /// \p substGenericParam Function that provides generic argument metadata

stdlib/public/runtime/ProtocolConformance.cpp

@@ -177,26 +177,11 @@ ProtocolConformanceDescriptor::getWitnessTable(const Metadata *type) const {
 
   case ConformanceFlags::ConformanceKind::ConditionalWitnessTableAccessor: {
     // Check the conditional requirements.
-    std::vector<unsigned> genericParamCounts;
-    (void)_gatherGenericParameterCounts(type->getTypeContextDescriptor(),
-                                        genericParamCounts);
-    auto genericArgs = type->getGenericArgs();
+    SubstGenericParametersFromMetadata substitutions(type);
     std::vector<const void *> conditionalArgs;
     bool failed =
       _checkGenericRequirements(getConditionalRequirements(), conditionalArgs,
-        [&](unsigned flatIndex) {
-          // FIXME: Adjust for non-key type parameters.
-          return genericArgs[flatIndex];
-        },
-        [&](unsigned depth, unsigned index) -> const Metadata *  {
-          // FIXME: Adjust for non-key type parameters.
-          if (auto flatIndex = _depthIndexToFlatIndex(depth, index,
-                                                      genericParamCounts)) {
-            return genericArgs[*flatIndex];
-          }
-
-          return nullptr;
-        });
+                                substitutions, substitutions);
     if (failed) return nullptr;
 
     return getWitnessTableAccessor()(

stdlib/public/runtime/ReflectionMirror.mm

@@ -332,44 +332,10 @@ static bool _shouldReportMissingReflectionMetadataWarnings() {
   if (!field.hasMangledTypeName())
     return {name, FieldType().withIndirect(field.isIndirectCase())};
 
-  std::vector<const ContextDescriptor *> descriptorPath;
-  {
-    const auto *parent = reinterpret_cast<
-                            const ContextDescriptor *>(baseDesc);
-    while (parent) {
-      if (parent->isGeneric())
-        descriptorPath.push_back(parent);
-
-      parent = parent->Parent.get();
-    }
-  }
-
   auto typeName = field.getMangledTypeName(0);
 
-  auto typeInfo = _getTypeByMangledName(
-      typeName,
-      [&](unsigned depth, unsigned index) -> const Metadata * {
-        if (depth >= descriptorPath.size())
-          return nullptr;
-
-        unsigned currentDepth = 0;
-        unsigned flatIndex = index;
-        const ContextDescriptor *currentContext = descriptorPath.back();
-
-        for (const auto *context : llvm::reverse(descriptorPath)) {
-          if (currentDepth >= depth)
-            break;
-
-          flatIndex += context->getNumGenericParams();
-          currentContext = context;
-          ++currentDepth;
-        }
-
-        if (index >= currentContext->getNumGenericParams())
-          return nullptr;
-
-        return base->getGenericArgs()[flatIndex];
-      });
+  SubstGenericParametersFromMetadata substitutions(base);
+  auto typeInfo = _getTypeByMangledName(typeName, substitutions);
 
   // Complete the type metadata before returning it to the caller.
   if (typeInfo) {

test/stdlib/Mirror.swift

@@ -1795,4 +1795,41 @@ mirrors.test("SymbolicReferenceInsideType") {
   expectEqual(expected, output)
 }
 
+protocol P1 { }
+protocol P2 { }
+protocol P3 { }
+
+struct ConformsToP1: P1 { }
+struct ConformsToP2: P2 { }
+struct ConformsToP3: P3 { }
+
+struct OuterTwoParams<T: P1, U: P2> {}
+
+struct ConformsToP1AndP2 : P1, P2 { }
+
+extension OuterTwoParams where U == T {
+  struct InnerEqualParams<V: P3> {
+    var x: T
+    var y: U
+    var z: V
+  }
+}
+
+mirrors.test("GenericNestedWithSameTypeConstraints") {
+  let value = OuterTwoParams.InnerEqualParams(x: ConformsToP1AndP2(),
+                                              y: ConformsToP1AndP2(),
+                                              z: ConformsToP3())
+  var output = ""
+  dump(value, to: &output)
+
+  // FIXME: The generic arguments are in the wrong place
+  let expected =
+    "▿ (extension in Mirror):Mirror.OuterTwoParams.InnerEqualParams<Mirror.ConformsToP1AndP2, Mirror.ConformsToP3>\n" +
+    "  - x: Mirror.ConformsToP1AndP2\n" +
+    "  - y: Mirror.ConformsToP1AndP2\n" +
+    "  - z: Mirror.ConformsToP3\n"
+
+  expectEqual(expected, output)
+}
+
 runAllTests()