[Runtime] Fix key argument indexing when checking invertible protocols.

Track the key argument index separately from the generic parameter index when performing the invertible protocol checking in _checkGenericRequirements. This keeps the indexing correct when a non-key argument is followed by a key argument.

rdar://128774651

Author: mikeash

stdlib/public/runtime/DynamicCast.cpp

@@ -1864,8 +1864,8 @@ static DynamicCastResult tryCastToExtendedExistential(
         [&substitutions](unsigned depth, unsigned index) {
           return substitutions.getMetadata(depth, index).Ptr;
         },
-        [&substitutions](unsigned ordinal) {
-          return substitutions.getMetadataOrdinal(ordinal).Ptr;
+        [&substitutions](unsigned fullOrdinal, unsigned keyOrdinal) {
+          return substitutions.getMetadataKeyArgOrdinal(keyOrdinal).Ptr;
         },
         [](const Metadata *type, unsigned index) -> const WitnessTable * {
           swift_unreachable("Resolution of witness tables is not supported");

stdlib/public/runtime/MetadataLookup.cpp

@@ -1195,7 +1195,7 @@ class SubstGenericParametersFromWrittenArgs {
         genericParamCounts(genericParamCounts) {}
 
   MetadataOrPack getMetadata(unsigned depth, unsigned index) const;
-  MetadataOrPack getMetadataOrdinal(unsigned ordinal) const;
+  MetadataOrPack getMetadataFullOrdinal(unsigned ordinal) const;
   const WitnessTable *getWitnessTable(const Metadata *type,
                                       unsigned index) const;
 };
@@ -1415,8 +1415,8 @@ _gatherGenericParameters(const ContextDescriptor *context,
         [&substitutions](unsigned depth, unsigned index) {
           return substitutions.getMetadata(depth, index).Ptr;
         },
-        [&substitutions](unsigned ordinal) {
-          return substitutions.getMetadataOrdinal(ordinal).Ptr;
+        [&substitutions](unsigned fullOrdinal, unsigned keyOrdinal) {
+          return substitutions.getMetadataFullOrdinal(fullOrdinal).Ptr;
         },
         [&substitutions](const Metadata *type, unsigned index) {
           return substitutions.getWitnessTable(type, index);
@@ -1849,12 +1849,12 @@ class DecodedMetadataBuilder {
           // FIXME: variadic generics
           return genArgs[index].getMetadata();
         },
-        [genArgs](unsigned ordinal) {
-          if (ordinal >= genArgs.size())
+        [genArgs](unsigned fullOrdinal, unsigned keyOrdinal) {
+          if (fullOrdinal >= genArgs.size())
             return (const Metadata*)nullptr;
 
           // FIXME: variadic generics
-          return genArgs[ordinal].getMetadata();
+          return genArgs[fullOrdinal].getMetadata();
         },
         [](const Metadata *type, unsigned index) -> const WitnessTable * {
           swift_unreachable("never called");
@@ -2814,8 +2814,8 @@ swift_distributed_getWitnessTables(GenericEnvironmentDescriptor *genericEnv,
       [&substFn](unsigned depth, unsigned index) {
         return substFn.getMetadata(depth, index).Ptr;
       },
-      [&substFn](unsigned ordinal) {
-        return substFn.getMetadataOrdinal(ordinal).Ptr;
+      [&substFn](unsigned fullOrdinal, unsigned keyOrdinal) {
+        return substFn.getMetadataKeyArgOrdinal(keyOrdinal).Ptr;
       },
       [&substFn](const Metadata *type, unsigned index) {
         return substFn.getWitnessTable(type, index);
@@ -3243,8 +3243,8 @@ SubstGenericParametersFromMetadata::getMetadata(
   return MetadataOrPack(genericArgs[flatIndex]);
 }
 
-MetadataOrPack
-SubstGenericParametersFromMetadata::getMetadataOrdinal(unsigned ordinal) const {
+MetadataOrPack SubstGenericParametersFromMetadata::getMetadataKeyArgOrdinal(
+    unsigned ordinal) const {
   // Don't attempt anything if we have no generic parameters.
   if (genericArgs == nullptr)
     return MetadataOrPack();
@@ -3281,8 +3281,8 @@ MetadataOrPack SubstGenericParametersFromWrittenArgs::getMetadata(
   return MetadataOrPack();
 }
 
-MetadataOrPack SubstGenericParametersFromWrittenArgs::getMetadataOrdinal(
-                                        unsigned ordinal) const {
+MetadataOrPack SubstGenericParametersFromWrittenArgs::getMetadataFullOrdinal(
+    unsigned ordinal) const {
   if (ordinal < allGenericArgs.size()) {
     return MetadataOrPack(allGenericArgs[ordinal]);
   }

stdlib/public/runtime/Private.h

@@ -295,11 +295,13 @@ class TypeInfo {
     std::function<const void *(unsigned depth, unsigned index)>;
 
   /// Callback used to provide the substitution of a generic parameter
-  /// (described by the ordinal, or "flat index") to its metadata.
+  /// (described by the ordinal, or "flat index") to its metadata. The index may
+  /// be "full" or it may be only relative to key arguments. The call is
+  /// provided both indexes and may use the one it requires.
   ///
   /// The return type here is a lie; it's actually a MetadataOrPack.
   using SubstGenericParameterOrdinalFn =
-    std::function<const void *(unsigned ordinal)>;
+    std::function<const void *(unsigned fullOrdinal, unsigned keyOrdinal)>;
 
   /// Callback used to provide the substitution of a witness table based on
   /// its index into the enclosing generic environment.
@@ -460,7 +462,7 @@ class TypeInfo {
     const void * const *getGenericArgs() const { return genericArgs; }
 
     MetadataOrPack getMetadata(unsigned depth, unsigned index) const;
-    MetadataOrPack getMetadataOrdinal(unsigned ordinal) const;
+    MetadataOrPack getMetadataKeyArgOrdinal(unsigned ordinal) const;
     const WitnessTable *getWitnessTable(const Metadata *type,
                                         unsigned index) const;
   };

stdlib/public/runtime/ProtocolConformance.cpp

@@ -340,8 +340,8 @@ ProtocolConformanceDescriptor::getWitnessTable(const Metadata *type) const {
         [&substitutions](unsigned depth, unsigned index) {
           return substitutions.getMetadata(depth, index).Ptr;
         },
-        [&substitutions](unsigned ordinal) {
-          return substitutions.getMetadataOrdinal(ordinal).Ptr;
+        [&substitutions](unsigned fullOrdinal, unsigned keyOrdinal) {
+          return substitutions.getMetadataKeyArgOrdinal(keyOrdinal).Ptr;
         },
         [&substitutions](const Metadata *type, unsigned index) {
           return substitutions.getWitnessTable(type, index);
@@ -1843,8 +1843,8 @@ checkInvertibleRequirements(const Metadata *type,
         [&substFn](unsigned depth, unsigned index) {
           return substFn.getMetadata(depth, index).Ptr;
         },
-        [&substFn](unsigned ordinal) {
-          return substFn.getMetadataOrdinal(ordinal).Ptr;
+        [&substFn](unsigned fullOrdinal, unsigned keyOrdinal) {
+          return substFn.getMetadataKeyArgOrdinal(keyOrdinal).Ptr;
         },
         [&substFn](const Metadata *type, unsigned index) {
           return substFn.getWitnessTable(type, index);
@@ -1886,6 +1886,7 @@ std::optional<TypeLookupError> swift::_checkGenericRequirements(
 
   // Now, check all of the generic arguments for invertible protocols.
   unsigned numGenericParams = genericParams.size();
+  unsigned keyIndex = 0;
   for (unsigned index = 0; index != numGenericParams; ++index) {
     // Non-key arguments don't need to be checked, because they are
     // aliased to another type.
@@ -1896,7 +1897,7 @@ std::optional<TypeLookupError> swift::_checkGenericRequirements(
     if (index < allSuppressed.size())
       suppressed = allSuppressed[index];
 
-    MetadataOrPack metadataOrPack(substGenericParamOrdinal(index));
+    MetadataOrPack metadataOrPack(substGenericParamOrdinal(index, keyIndex));
     switch (genericParams[index].getKind()) {
     case GenericParamKind::Type: {
       if (!metadataOrPack || metadataOrPack.isMetadataPack()) {
@@ -1937,6 +1938,7 @@ std::optional<TypeLookupError> swift::_checkGenericRequirements(
       return TYPE_LOOKUP_ERROR_FMT("unknown generic parameter kind %u",
                                    index);
     }
+    keyIndex++;
   }
 
   // Success!

test/Interpreter/conditional_conformances_runtime.swift

@@ -33,11 +33,35 @@ extension Dictionary: P where Value == (Key) -> Bool {
 
 
 let yx = Y(wrapped: X())
-assert(tryAsP(yx) == 11)
+precondition(tryAsP(yx) == 11)
 
 let dict: [Int : (Int) -> Bool] = [:]
-assert(tryAsP(dict) == 2)
+precondition(tryAsP(dict) == 2)
 
 let yDict = Y(wrapped: dict)
-assert(tryAsP(yDict) == 12)
+precondition(tryAsP(yDict) == 12)
 
+// <rdar://128774651> Generic argument indexing runs off the end when a key
+// argument comes after a non-key argument.
+struct Outer<T> {}
+
+extension Outer where T == Int {
+  struct Inner<U> {
+    // Add some stored properties that will have field offsets in the metadata.
+    // If we read past the end of the generic arguments, we'll find these and
+    // crash on something that's definitely not a valid pointer.
+    var a: T?
+    var b: T?
+    var c: T?
+    var d: T?
+  }
+}
+
+extension Outer.Inner: P where U == String {
+  func foo() -> Int { return 1 }
+}
+
+let conformingInner: Any = Outer.Inner<String>()
+let nonconformingInner: Any = Outer.Inner<Int>()
+precondition(conformingInner is P)
+precondition(!(nonconformingInner is P))