Add Self- and SelfWitnessTable metadata sources to capture descriptors

These can show up in partial applications of functions with the
witness_method calling convention.

Author: bitjammer

include/swift/Reflection/MetadataSource.h

@@ -182,6 +182,9 @@ class MetadataSource {
         return decodeGenericArgument(A, it, end);
       case 'P':
         return decodeParent(A, it, end);
+      case 'S':
+        ++it;
+        return A.template createSelf();
       default:
         return nullptr;
     }
@@ -315,6 +318,7 @@ class GenericArgumentMetadataSource final : public MetadataSource {
   }
 };
 
+/// Metadata gotten through the parent of a nominal type's metadata.
 class ParentMetadataSource final : public MetadataSource {
   const MetadataSource *Child;
 public:
@@ -337,6 +341,41 @@ class ParentMetadataSource final : public MetadataSource {
   }
 };
 
+/// A source of metadata from the Self metadata parameter passed via
+/// a witness_method convention function.
+class SelfMetadataSource final : public MetadataSource {
+public:
+  SelfMetadataSource() : MetadataSource(MetadataSourceKind::Self) {}
+
+  template <typename Allocator>
+  static const SelfMetadataSource*
+  create(Allocator &A) {
+    return A.template make_source<SelfMetadataSource>();
+  }
+
+  static bool classof(const MetadataSource *MS) {
+    return MS->getKind() == MetadataSourceKind::Self;
+  }
+};
+
+/// A source of metadata from the Self witness table parameter passed via
+/// a witness_method convention function.
+class SelfWitnessTableMetadataSource final : public MetadataSource {
+public:
+  SelfWitnessTableMetadataSource()
+    : MetadataSource(MetadataSourceKind::SelfWitnessTable) {}
+
+  template <typename Allocator>
+  static const SelfWitnessTableMetadataSource*
+  create(Allocator &A) {
+    return A.template make_source<SelfWitnessTableMetadataSource>();
+  }
+
+  static bool classof(const MetadataSource *MS) {
+    return MS->getKind() == MetadataSourceKind::SelfWitnessTable;
+  }
+};
+
 template <typename ImplClass, typename RetTy = void, typename... Args>
 class MetadataSourceVisitor {
 public:

include/swift/Reflection/MetadataSourceBuilder.h

@@ -62,6 +62,16 @@ class MetadataSourceBuilder {
   createParent(const MetadataSource *Child) {
     return ParentMetadataSource::create(*this, Child);
   }
+
+  const SelfMetadataSource *
+  createSelf() {
+    return SelfMetadataSource::create(*this);
+  }
+
+  const SelfWitnessTableMetadataSource *
+  createSelfWitnessTable() {
+    return SelfWitnessTableMetadataSource::create(*this);
+  }
 };
 
 } // end namespace reflection

include/swift/Reflection/MetadataSources.def

@@ -21,3 +21,5 @@ METADATA_SOURCE(ReferenceCapture, MetadataSource)
 METADATA_SOURCE(MetadataCapture, MetadataSource)
 METADATA_SOURCE(GenericArgument, MetadataSource)
 METADATA_SOURCE(Parent, MetadataSource)
+METADATA_SOURCE(Self, MetadataSource)
+METADATA_SOURCE(SelfWitnessTable, MetadataSource)

lib/IRGen/GenReflection.cpp

@@ -70,6 +70,15 @@ class MetadataSourceEncoder
     visit(P->getChild());
     OS << '_';
   }
+
+  void visitSelfMetadataSource(const SelfMetadataSource *S) {
+    OS << 'S';
+  }
+
+  void
+  visitSelfWitnessTableMetadataSource(const SelfWitnessTableMetadataSource *S) {
+    OS << 'W';
+  }
 };
 
 class PrintMetadataSource
@@ -148,6 +157,18 @@ class PrintMetadataSource
     printRec(P->getChild());
     closeForm();
   }
+
+  void
+  visitSelfMetadataSource(const SelfMetadataSource *S) {
+    printHeader("self");
+    closeForm();
+  }
+
+  void
+  visitSelfWitnessTableMetadataSource(const SelfWitnessTableMetadataSource *S) {
+    printHeader("self-witness-table");
+    closeForm();
+  }
 };
 
 class ReflectionMetadataBuilder : public ConstantBuilder<> {
@@ -511,6 +532,8 @@ class CaptureDescriptorBuilder : public ReflectionMetadataBuilder {
 
     PolymorphicConvention Convention(IGM, CalleeType);
 
+    using SourceKind = PolymorphicConvention::SourceKind;
+
     auto Generics = Callee.getContextGenericParams()->getNestedGenericParams();
     for (auto GenericParam : Generics) {
       // The generic type parameter (depth, index) serves as the key to the
@@ -527,11 +550,22 @@ class CaptureDescriptorBuilder : public ReflectionMetadataBuilder {
         // The convention fulfills the requirement, so record how to get
         // to the metadata.
 
+        auto ConventionSource = Convention.getSource(Fulfillment->SourceIndex);
+        if (ConventionSource.getKind() == SourceKind::SelfMetadata) {
+          SourceMap.insert({GenericParamType, SourceBuilder.createSelf()});
+          continue;
+        } else if (ConventionSource.getKind() == SourceKind::SelfWitnessTable) {
+          SourceMap.insert({
+            GenericParamType,
+            SourceBuilder.createSelfWitnessTable()
+          });
+          continue;
+        }
+
         // Since captures are created via partial_apply instructions, we need
         // to see which function parameter fulfilled this metadata need and
         // grab its type.
-        auto FnParameterIndex
-          = Convention.getSource(Fulfillment->SourceIndex).getParamIndex();
+        auto FnParameterIndex = ConventionSource.getParamIndex();
         auto FnParameter = CalleeType->getParameters()[FnParameterIndex];
         auto ParameterType = FnParameter.getType()->getCanonicalType();
 

stdlib/public/Reflection/MetadataSource.cpp

@@ -93,6 +93,17 @@ class PrintMetadataSource
     printRec(P->getChild());
     closeForm();
   }
+
+  void visitSelfMetadataSource(const SelfMetadataSource *S) {
+    printHeader("self");
+    closeForm();
+  }
+
+  void
+  visitSelfWitnessTableMetadataSource(const SelfWitnessTableMetadataSource *W) {
+    printHeader("self-witness-table");
+    closeForm();
+  }
 };
 
 void MetadataSource::dump() const {