[runtime] Handle same-type constraints when resolving generic params

Generic parameters for a context are normally classified as "key",
meaning they have actual metadata provided at runtime, or non-key,
meaning they're derivable from somewhere else. However, a nested
context or constrained extension can take what would be a "key"
parameter in a parent context and make it non-key in a child context.
This messes with the mapping between the (depth, index) representation
of generic parameters and the flat list of generic arguments. Fix this
by (1) consistently substituting out extension contexts with the
contexts of the extended types, and (2) using the most nested context
to decide which parameters are key, instead of the context a parameter
was originally introduced in.

Note that (1) may have problems if/when extensions start introducing
their /own/ generic parameters. For now I tried to be consistent with
what was there.

rdar://problem/52364601

Author: jrose-apple

stdlib/public/runtime/MetadataLookup.cpp

@@ -246,11 +246,17 @@ _findContextDescriptor(Demangle::NodePointer node,
                            Demangle::Demangler &Dem);
 
 /// Find the context descriptor for the type extended by the given extension.
+///
+/// If \p maybeExtension isn't actually an extension context, returns nullptr.
 static const ContextDescriptor *
-_findExtendedTypeContextDescriptor(const ExtensionContextDescriptor *extension,
+_findExtendedTypeContextDescriptor(const ContextDescriptor *maybeExtension,
                                    Demangler &demangler,
                                    Demangle::NodePointer *demangledNode
                                      = nullptr) {
+  auto extension = dyn_cast<ExtensionContextDescriptor>(maybeExtension);
+  if (!extension)
+    return nullptr;
+
   Demangle::NodePointer localNode;
   Demangle::NodePointer &node = demangledNode ? *demangledNode : localNode;
 
@@ -844,28 +850,26 @@ bool swift::_gatherGenericParameterCounts(
   DemanglerForRuntimeTypeResolution<> demangler;
   demangler.providePreallocatedMemory(BorrowFrom);
 
-  if (auto extension = dyn_cast<ExtensionContextDescriptor>(descriptor)) {
-    // If we have an nominal type extension descriptor, extract the extended type
+  if (auto extension = _findExtendedTypeContextDescriptor(descriptor,
+                                                          demangler)) {
+    // If we have a nominal type extension descriptor, extract the extended type
     // and use that. If the extension is not nominal, then we can use the
     // extension's own signature.
-    if (auto extendedType =
-          _findExtendedTypeContextDescriptor(extension, demangler)) {
-      descriptor = extendedType;
-    }
+    descriptor = extension;
   }
 
   // Once we hit a non-generic descriptor, we're done.
   if (!descriptor->isGeneric()) return false;
 
   // Recurse to record the parent context's generic parameters.
-  if (auto parent = descriptor->Parent.get())
-    (void)_gatherGenericParameterCounts(parent, genericParamCounts, demangler);
+  auto parent = descriptor->Parent.get();
+  (void)_gatherGenericParameterCounts(parent, genericParamCounts, demangler);
 
   // Record a new level of generic parameters if the count exceeds the
   // previous count.
-  auto myCount =
-    descriptor->getGenericContext()->getGenericContextHeader().NumParams;
-  if (genericParamCounts.empty() || myCount > genericParamCounts.back()) {
+  unsigned parentCount = parent->getNumGenericParams();
+  unsigned myCount = descriptor->getNumGenericParams();
+  if (myCount > parentCount) {
     genericParamCounts.push_back(myCount);
     return true;
   }
@@ -1658,32 +1662,50 @@ static void installGetClassHook() {
 #endif
 
 unsigned SubstGenericParametersFromMetadata::
-buildDescriptorPath(const ContextDescriptor *context) const {
+buildDescriptorPath(const ContextDescriptor *context,
+                    Demangler &borrowFrom) const {
+  assert(sourceIsMetadata);
+
   // Terminating condition: we don't have a context.
   if (!context)
     return 0;
 
+  DemanglerForRuntimeTypeResolution<> demangler;
+  demangler.providePreallocatedMemory(borrowFrom);
+
+  if (auto extension = _findExtendedTypeContextDescriptor(context, demangler)) {
+    // If we have a nominal type extension descriptor, extract the extended type
+    // and use that. If the extension is not nominal, then we can use the
+    // extension's own signature.
+    context = extension;
+  }
+
   // Add the parent's contribution to the descriptor path.
-  unsigned numKeyGenericParamsInParent =
-    buildDescriptorPath(context->Parent.get());
+  const ContextDescriptor *parent = context->Parent.get();
+  unsigned numKeyGenericParamsInParent = buildDescriptorPath(parent, demangler);
 
   // If this context is non-generic, we're done.
   if (!context->isGeneric())
     return numKeyGenericParamsInParent;
 
   // Count the number of key generic params at this level.
+  auto allGenericParams = baseContext->getGenericContext()->getGenericParams();
+  unsigned parentCount = parent->getNumGenericParams();
+  unsigned localCount = context->getNumGenericParams();
+  auto localGenericParams = allGenericParams.slice(parentCount,
+                                                   localCount - parentCount);
+
   unsigned numKeyGenericParamsHere = 0;
   bool hasNonKeyGenericParams = false;
-  auto localGenericParams = getLocalGenericParams(context);
   for (const auto &genericParam : localGenericParams) {
     if (genericParam.hasKeyArgument())
       ++numKeyGenericParamsHere;
     else
       hasNonKeyGenericParams = true;
   }
 
-  // Form the path element if there are any generic parameters to be found.
-  if (numKeyGenericParamsHere != 0)
+  // Form the path element if there are any new generic parameters.
+  if (localCount > parentCount)
     descriptorPath.push_back(PathElement{localGenericParams,
                                          context->getNumGenericParams(),
                                          numKeyGenericParamsInParent,
@@ -1738,7 +1760,8 @@ void SubstGenericParametersFromMetadata::setup() const {
     return;
 
   if (sourceIsMetadata && baseContext) {
-    numKeyGenericParameters = buildDescriptorPath(baseContext);
+    DemanglerForRuntimeTypeResolution<StackAllocatedDemangler<2048>> demangler;
+    numKeyGenericParameters = buildDescriptorPath(baseContext, demangler);
     return;
   }
 

stdlib/public/runtime/Private.h

@@ -309,7 +309,8 @@ class TypeInfo {
     ///
     /// \returns a pair containing the number of key generic parameters in
     /// the path up to this point.
-    unsigned buildDescriptorPath(const ContextDescriptor *context) const;
+    unsigned buildDescriptorPath(const ContextDescriptor *context,
+                                 Demangler &demangler) const;
 
     /// Builds a path from the generic environment.
     unsigned buildEnvironmentPath(

test/IRGen/nested_generics.swift

@@ -105,6 +105,8 @@ protocol HasAssoc {
   associatedtype Assoc
 }
 
+class SeparateGenericSuperclass<T> {}
+
 enum Container<T : TagProtocol> {
   class _Superclass {}
   // CHECK-CONSTANTS-LABEL: @"$s15nested_generics9ContainerO9_SubclassCMn" =
@@ -127,27 +129,30 @@ enum Container<T : TagProtocol> {
   // CHECK-CONSTANTS-SAME: @"symbolic _____yx_G 15nested_generics9ContainerO11_SuperclassC"
   class _Subclass2<U: Collection>: _Superclass where T == U.Element {}
 
-
   // CHECK-CONSTANTS-LABEL: @"$s15nested_generics9ContainerO10_Subclass3CMn" =
-  // FIXME: That "qd__" still causes problems: it's (depth: 1, index: 0), but
-  // the runtime doesn't count the parameters at depth 0 correctly.
   // CHECK-CONSTANTS-SAME: @"symbolic _____y______qd__G 15nested_generics9ContainerO18_GenericSuperclassC AA5OuterO"
   class _GenericSuperclass<U> {}
   class _Subclass3<U>: _GenericSuperclass<U> where T == Outer {}
 
+  class MoreNesting {
+    // CHECK-CONSTANTS-LABEL: @"$s15nested_generics9ContainerO11MoreNestingC9_SubclassCMn" =
+    // CHECK-CONSTANTS-SAME: @"symbolic _____y______G 15nested_generics9ContainerO11_SuperclassC AA5OuterO"
+    class _Subclass<U>: _Superclass where T == Outer {}
+  }
+
+  // CHECK-CONSTANTS-LABEL: @"$s15nested_generics9ContainerO24_SeparateGenericSubclassCMn" =
+  // CHECK-CONSTANTS-SAME: @"symbolic _____yxSgG 15nested_generics25SeparateGenericSuperclassC"
+  class _SeparateGenericSubclass: SeparateGenericSuperclass<T?> {}
+
   // CHECK-CONSTANTS-LABEL: @"$s15nested_generics9ContainerO6FieldsVMF" =
   // CHECK-CONSTANTS-SAME: @"symbolic _____ 15nested_generics5OuterO"
-  // FIXME: This still causes problems: it's (depth: 1, index: 0), but
-  // the runtime doesn't count the parameters at depth 0 correctly.
   // CHECK-CONSTANTS-SAME: @"symbolic qd__"
   struct Fields<U> where T == Outer {
     var x: T
     var y: U
   }
 
   // CHECK-CONSTANTS-LABEL: @"$s15nested_generics9ContainerO5CasesOMF" =
-  // FIXME: This still causes problems: it's (depth: 1, index: 0), but
-  // the runtime doesn't count the parameters at depth 0 correctly.
   // CHECK-CONSTANTS-SAME: @"symbolic qd__"
   enum Cases<U> where T == Outer {
     case a(T)

test/stdlib/Mirror.swift

@@ -1851,4 +1851,158 @@ mirrors.test("CustomAlignment") {
   expectEqual(expected, output)
 }
 
+protocol STSTagProtocol {}
+struct STSOuter : STSTagProtocol {}
+
+enum STSContainer<T : STSTagProtocol> {
+  class Superclass {}
+  class Subclass<U>: Superclass where T == STSOuter {
+    class ExtraNested: Superclass {}
+  }
+
+  class GenericSuperclass<U> {}
+  class Subclass2<U>: GenericSuperclass<U> where T == STSOuter {}
+
+  class Subclass3<U: Collection>: Superclass where T == U.Element {}
+
+  class MoreNesting<X> {
+    class Subclass<U>: Superclass where T == STSOuter {}
+  }
+
+  struct Fields<U> where T == STSOuter {
+    var x: T?
+    var y: U?
+  }
+
+  enum Cases<U> where T == STSOuter {
+    case a(T)
+    case b(U)
+  }
+}
+
+// A new type with an easily-recognizable, easily-strippable suffix character.
+enum STSContainer℠<T : STSTagProtocol> {
+  class Superclass {}
+  class GenericSuperclass<U> {}
+}
+extension STSContainer℠ where T == STSOuter {
+  class Subclass<U>: Superclass {
+    class ExtraNested: Superclass {}
+  }
+
+  class Subclass2<U>: GenericSuperclass<U> {}
+
+  class MoreNesting<X> {
+    class Subclass<U>: Superclass {}
+  }
+
+  struct Fields<U> {
+    var x: T?
+    var y: U?
+  }
+
+  enum Cases<U> {
+    case a(T)
+    case b(U)
+  }
+}
+
+var sameTypeSuite = TestSuite("Mirrors.SameTypeConstraints")
+
+func testSTSDump<RegularValue, ExtensionValue>(
+  _ regular: RegularValue,
+  _ ext: ExtensionValue,
+  _ expected: String
+) {
+  var output = ""
+  dump(regular, to: &output)
+  expectEqual(expected, output)
+
+  func clean(_ dumpOutput: String) -> String {
+    // This isn't efficient but it doesn't have to be.
+    var result = dumpOutput
+    result.removeAll { $0 == "℠" }
+    if let openParenIndex = result.firstIndex(of: "(") {
+      if result[openParenIndex...].hasPrefix("(extension in Mirror):") {
+        let colonIndex = result[openParenIndex...].firstIndex(of: ":")!
+        result.removeSubrange(openParenIndex...colonIndex)
+      }
+    }
+    return result
+  }
+
+  var extensionOutput = ""
+  dump(ext, to: &extensionOutput)
+  expectEqual(expected, clean(extensionOutput))
+}
+
+func testSTSDump<RegularValue>(_ regular: RegularValue, _ expected: String) {
+  var output = ""
+  dump(regular, to: &output)
+  expectEqual(expected, output)
+}
+
+
+sameTypeSuite.test("Subclass") {
+  testSTSDump(STSContainer.Subclass<Int>(),
+              STSContainer℠.Subclass<Int>(),
+              """
+    - Mirror.STSContainer<Mirror.STSOuter>.Subclass<Swift.Int> #0
+      - super: Mirror.STSContainer<Mirror.STSOuter>.Superclass\n
+    """)
+
+  testSTSDump(STSContainer.Subclass2<Int>(),
+              STSContainer℠.Subclass2<Int>(),
+              """
+    - Mirror.STSContainer<Mirror.STSOuter>.Subclass2<Swift.Int> #0
+      - super: Mirror.STSContainer<Mirror.STSOuter>.GenericSuperclass<Swift.Int>\n
+    """)
+
+  testSTSDump(STSContainer.Subclass3<[STSOuter]>(),
+              """
+    - Mirror.STSContainer<Mirror.STSOuter>.Subclass3<Swift.Array<Mirror.STSOuter>> #0
+      - super: Mirror.STSContainer<Mirror.STSOuter>.Superclass\n
+    """)
+
+  testSTSDump(STSContainer.Subclass<Int>.ExtraNested(),
+              STSContainer℠.Subclass<Int>.ExtraNested(),
+              """
+    - Mirror.STSContainer<Mirror.STSOuter>.Subclass<Swift.Int>.ExtraNested #0
+      - super: Mirror.STSContainer<Mirror.STSOuter>.Superclass\n
+    """)
+
+  testSTSDump(STSContainer.MoreNesting<Bool>.Subclass<Int>(),
+              STSContainer℠.MoreNesting<Bool>.Subclass<Int>(),
+              """
+    - Mirror.STSContainer<Mirror.STSOuter>.MoreNesting<Swift.Bool>.Subclass<Swift.Int> #0
+      - super: Mirror.STSContainer<Mirror.STSOuter>.Superclass\n
+    """)
+}
+
+sameTypeSuite.test("Fields") {
+  testSTSDump(STSContainer.Fields<Int>(),
+              STSContainer℠.Fields<Int>(),
+              """
+    ▿ Mirror.STSContainer<Mirror.STSOuter>.Fields<Swift.Int>
+      - x: nil
+      - y: nil\n
+    """)
+}
+
+sameTypeSuite.test("Cases") {
+  testSTSDump(STSContainer.Cases<Int>.a(.init()),
+              STSContainer℠.Cases<Int>.a(.init()),
+              """
+    - Mirror.STSContainer<Mirror.STSOuter>.Cases<Swift.Int>.a\n
+    """)
+
+  testSTSDump(STSContainer.Cases<Int>.b(.init()),
+              STSContainer℠.Cases<Int>.b(.init()),
+              """
+    ▿ Mirror.STSContainer<Mirror.STSOuter>.Cases<Swift.Int>.b
+      - b: 0\n
+    """)
+}
+
+
 runAllTests()