Gather opaque type conformance requirements when scanning associated type infos from a binary

artemcm · artemcm · commit 92011f2f5a5d · 2022-06-29T11:24:22.000-07:00
When detecting that an associated type's substituted type is an opaque type, read out its opaque type descriptor to collect the names of protocols it must conform to.
diff --git a/include/swift/ABI/GenericContext.h b/include/swift/ABI/GenericContext.h
@@ -145,6 +145,18 @@ class TargetGenericRequirementDescriptor {
     return Protocol;
   }
 
+  /// Retreive the raw value of the Protocol requirement pointer.
+  int32_t getUnresolvedProtocolAddress() const {
+    assert(getKind() == GenericRequirementKind::Protocol);
+    return Protocol.getUnresolvedProtocolAddress();
+  }
+
+  /// Retreive the offset to the Protocol field
+  constexpr inline auto
+  getProtocolOffset() const -> typename Runtime::StoredSize {
+    return offsetof(typename std::remove_reference<decltype(*this)>::type, Protocol);
+  }
+
   /// Retrieve the right-hand type for a SameType or BaseClass requirement.
   llvm::StringRef getMangledTypeName() const {
     assert(getKind() == GenericRequirementKind::SameType ||
@@ -301,7 +313,8 @@ class TargetGenericEnvironment
   }
 };
 
-using GenericEnvironmentDescriptor = TargetGenericEnvironment<InProcess>;
+using GenericEnvironmentDescriptor =
+TargetGenericEnvironment<InProcess>;
 
 /// CRTP class for a context descriptor that includes trailing generic
 /// context description.
diff --git a/include/swift/ABI/Metadata.h b/include/swift/ABI/Metadata.h
@@ -2996,7 +2996,11 @@ struct TargetOpaqueTypeDescriptor final
     return cd->getKind() == ContextDescriptorKind::OpaqueType;
   }
 };
-  
+
+template <template <typename Runtime> class ObjCInteropKind,
+          unsigned PointerSize>
+using ExternalOpaqueTypeDescriptor = TargetOpaqueTypeDescriptor<
+    External<ObjCInteropKind<RuntimeTarget<PointerSize>>>>;
 using OpaqueTypeDescriptor = TargetOpaqueTypeDescriptor<InProcess>;
 
 /// The instantiation cache for generic metadata.  This must be guaranteed
diff --git a/include/swift/ABI/MetadataRef.h b/include/swift/ABI/MetadataRef.h
@@ -308,6 +308,16 @@ class RelativeTargetProtocolDescriptorPointer {
             swiftPointer.getPointer()));
   }
 
+  /// Retrieve a reference to the protocol.
+  int32_t getUnresolvedProtocolAddress() const {
+#if SWIFT_OBJC_INTEROP
+    if (isObjC()) {
+      return objcPointer.getUnresolvedOffset();
+    }
+#endif
+    return swiftPointer.getUnresolvedOffset();
+  }
+
   operator TargetProtocolDescriptorRef<Runtime>() const {
     return getProtocol();
   }
diff --git a/include/swift/Basic/RelativePointer.h b/include/swift/Basic/RelativePointer.h
@@ -347,11 +347,7 @@ class RelativeIndirectablePointerIntPair {
 
 public:
   const ValueTy *getPointer() const & {
-    static_assert(alignof(ValueTy) >= 2 && alignof(Offset) >= 2,
-                  "alignment of value and offset must be at least 2 to "
-                  "make room for indirectable flag");
-
-    Offset offset = (RelativeOffsetPlusIndirectAndInt & ~getIntMask());
+    Offset offset = getUnresolvedOffset();
 
     // Check for null.
     if (Nullable && offset == 0)
@@ -370,10 +366,17 @@ class RelativeIndirectablePointerIntPair {
     }
   }
 
+  Offset getUnresolvedOffset() const & {
+    static_assert(alignof(ValueTy) >= 2 && alignof(Offset) >= 2,
+                  "alignment of value and offset must be at least 2 to "
+                  "make room for indirectable flag");
+    Offset offset = (RelativeOffsetPlusIndirectAndInt & ~getIntMask());
+    return offset;
+  }
+
   /// A zero relative offset encodes a null reference.
   bool isNull() const & {
-    Offset offset = (RelativeOffsetPlusIndirectAndInt & ~getIntMask());
-    return offset == 0;
+    return getUnresolvedOffset() == 0;
   }
 
   IntTy getInt() const & {
diff --git a/include/swift/Reflection/TypeRefBuilder.h b/include/swift/Reflection/TypeRefBuilder.h
@@ -315,6 +315,7 @@ struct AssociatedType {
   std::string SubstitutedTypeMangledName;
   std::string SubstitutedTypeFullyQualifiedName;
   std::string SubstitutedTypeDiagnosticPrintName;
+  std::vector<std::string> OpaqueTypeProtocolConformanceRequirements;
 };
 
 /// Info about all of a given type's associated types, as read out from an Image
@@ -523,10 +524,10 @@ class TypeRefBuilder {
                             Node::Kind::OpaqueTypeDescriptorSymbolicReference) {
       auto underlyingTy = OpaqueUnderlyingTypeReader(
                                          opaqueDescriptor->getIndex(), ordinal);
-      
+
       if (!underlyingTy)
         return nullptr;
-      
+
       GenericArgumentMap subs;
       for (unsigned d = 0, de = genericArgs.size(); d < de; ++d) {
         auto argsForDepth = genericArgs[d];
@@ -1005,12 +1006,115 @@ class TypeRefBuilder {
                    bool printTypeName = false);
   FieldTypeCollectionResult collectFieldTypes(llvm::Optional<std::string> forMangledTypeName);
   void dumpFieldSection(std::ostream &stream);
-  AssociatedTypeCollectionResult collectAssociatedTypes(llvm::Optional<std::string> forMangledTypeName);
-  void dumpAssociatedTypeSection(std::ostream &stream);
   void dumpBuiltinTypeSection(std::ostream &stream);
   void dumpCaptureSection(std::ostream &stream);
   void dumpMultiPayloadEnumSection(std::ostream &stream);
 
+  ///
+  /// Extraction of associated types
+  ///
+public:
+  template <template <typename Runtime> class ObjCInteropKind,
+            unsigned PointerSize>
+  AssociatedTypeCollectionResult
+  collectAssociatedTypes(llvm::Optional<std::string> forMangledTypeName) {
+    AssociatedTypeCollectionResult result;
+    for (const auto &sections : ReflectionInfos) {
+      for (auto descriptor : sections.AssociatedType) {
+        // Read out the relevant info from the associated type descriptor:
+        // The type's name and which protocol conformance it corresponds to
+        auto typeRef = readTypeRef(descriptor, descriptor->ConformingTypeName);
+        auto typeName = nodeToString(demangleTypeRef(typeRef));
+        auto optionalMangledTypeName = normalizeReflectionName(typeRef);
+        auto protocolNode = demangleTypeRef(
+            readTypeRef(descriptor, descriptor->ProtocolTypeName));
+        auto protocolName = nodeToString(protocolNode);
+        clearNodeFactory();
+        if (optionalMangledTypeName.hasValue()) {
+          auto mangledTypeName = optionalMangledTypeName.getValue();
+          if (forMangledTypeName.hasValue()) {
+            if (mangledTypeName != forMangledTypeName.getValue())
+              continue;
+          }
+
+          // For each associated type, gather its typealias name,
+          // the substituted type info, and if the substituted type is opaque -
+          // gather its protocol conformance requirements
+          std::vector<AssociatedType> associatedTypes;
+          for (const auto &associatedTypeRef : *descriptor.getLocalBuffer()) {
+            auto associatedType = descriptor.getField(associatedTypeRef);
+            std::string typealiasTypeName =
+                getTypeRefString(
+                    readTypeRef(associatedType, associatedType->Name))
+                    .str();
+
+            std::string mangledSubstitutedTypeName =
+                std::string(associatedType->SubstitutedTypeName);
+            auto substitutedTypeRef = readTypeRef(
+                associatedType, associatedType->SubstitutedTypeName);
+            auto optionalMangledSubstitutedTypeName =
+                normalizeReflectionName(substitutedTypeRef);
+            if (optionalMangledSubstitutedTypeName.hasValue()) {
+              mangledSubstitutedTypeName =
+                  optionalMangledSubstitutedTypeName.getValue();
+            }
+
+            // We intentionally do not want to resolve opaque type
+            // references, because if the substituted type is opaque, we
+            // would like to get at its OpaqueTypeDescriptor address, which
+            // is stored on the OpaqueTypeDescriptorSymbolicReference typeRef.
+            auto substitutedDemangleTree =
+                demangleTypeRef(substitutedTypeRef,
+                                /* useOpaqueTypeSymbolicReferences */ true);
+
+            // If the substituted type is an opaque type, also gather info
+            // about which protocols it is required to conform to
+            std::vector<std::string> OpaqueTypeConformanceRequirements;
+            gatherConformanceRequirementsIfOpaque<ObjCInteropKind, PointerSize>(
+                substitutedDemangleTree, OpaqueTypeConformanceRequirements);
+
+            auto substitutedTypeName = nodeToString(substitutedDemangleTree);
+            std::stringstream OS;
+            dumpTypeRef(substitutedTypeRef, OS);
+            associatedTypes.emplace_back(
+                AssociatedType{typealiasTypeName, mangledSubstitutedTypeName,
+                               substitutedTypeName, OS.str(),
+                               OpaqueTypeConformanceRequirements});
+          }
+          result.AssociatedTypeInfos.emplace_back(AssociatedTypeInfo{
+              mangledTypeName, typeName, protocolName, associatedTypes});
+        }
+      }
+    }
+    return result;
+  }
+
+  template <template <typename Runtime> class ObjCInteropKind,
+            unsigned PointerSize>
+  void gatherConformanceRequirementsIfOpaque(
+      Demangle::Node *substitutedTypeDemangleTree,
+      std::vector<std::string> &OpaqueTypeConformanceRequirements) {
+    // With unresolved opaque symbolic references, the demangle tree we
+    // extract the opaque type descriptor's address from is of the form:
+    // kind=Type
+    //  kind=OpaqueType
+    //    kind=OpaqueTypeDescriptorSymbolicReference, index={{1-9+}}
+    // Where the `index` value is the descriptor's address
+    //
+    if (substitutedTypeDemangleTree->getKind() == Node::Kind::Type) {
+      auto childDemangleTree = substitutedTypeDemangleTree->getFirstChild();
+      if (childDemangleTree->getKind() == Node::Kind::OpaqueType) {
+        auto opaqueTypeChildDemangleTree = childDemangleTree->getFirstChild();
+        if (opaqueTypeChildDemangleTree->getKind() ==
+            Node::Kind::OpaqueTypeDescriptorSymbolicReference) {
+          OpaqueTypeConformanceRequirements =
+              collectOpaqueTypeConformanceNames<ObjCInteropKind, PointerSize>(
+                  opaqueTypeChildDemangleTree->getIndex());
+        }
+      }
+    }
+  }
+
 private:
   struct ContextNameInfo {
     std::string name;
@@ -1142,7 +1246,6 @@ class TypeRefBuilder {
           remote::RemoteAddress(protocolDescriptorAddress),
           sizeof(ExternalContextDescriptor<ObjCInteropKind, PointerSize>));
       if (!protocolContextDescriptorBytes.get()) {
-        // Error = "Failed to read context (protocol) descriptor.";
         return llvm::None;
       }
       const ExternalContextDescriptor<ObjCInteropKind, PointerSize>
@@ -1332,6 +1435,83 @@ class TypeRefBuilder {
     }
   };
 
+  template <template <typename Runtime> class ObjCInteropKind,
+            unsigned PointerSize>
+  void dumpAssociatedTypeSection(std::ostream &stream) {
+    auto associatedTypeCollectionResult =
+        collectAssociatedTypes<ObjCInteropKind, PointerSize>(
+            llvm::Optional<std::string>());
+    for (const auto &info :
+         associatedTypeCollectionResult.AssociatedTypeInfos) {
+      stream << "- " << info.FullyQualifiedName << " : "
+             << info.ProtocolFullyQualifiedName << "\n";
+      for (const auto &typeAlias : info.AssociatedTypes) {
+        stream << "typealias " << typeAlias.TypeAliasName << " = "
+               << typeAlias.SubstitutedTypeFullyQualifiedName << "\n";
+        stream << typeAlias.SubstitutedTypeDiagnosticPrintName;
+        if (!typeAlias.OpaqueTypeProtocolConformanceRequirements.empty()) {
+          stream << "opaque type conformance requirements: \n";
+          for (const auto &protocolName :
+               typeAlias.OpaqueTypeProtocolConformanceRequirements) {
+            stream << protocolName << "\n";
+          }
+        }
+      }
+      stream << "\n";
+    }
+  }
+
+  template <template <typename Runtime> class ObjCInteropKind,
+            unsigned PointerSize>
+  std::vector<std::string>
+  collectOpaqueTypeConformanceNames(uintptr_t opaqueTypeDescriptorAddress) {
+    std::vector<std::string> result;
+    auto opaqueTypeDescriptorBytes = OpaqueByteReader(
+        remote::RemoteAddress(opaqueTypeDescriptorAddress),
+        sizeof(ExternalOpaqueTypeDescriptor<ObjCInteropKind, PointerSize>));
+    if (!opaqueTypeDescriptorBytes.get()) {
+      return result;
+    }
+    const ExternalOpaqueTypeDescriptor<ObjCInteropKind, PointerSize>
+        *opaqueTypeDescriptor =
+            (const ExternalOpaqueTypeDescriptor<ObjCInteropKind, PointerSize> *)
+                opaqueTypeDescriptorBytes.get();
+
+    if (!opaqueTypeDescriptor) {
+      return result;
+    }
+
+    for (const auto &req : opaqueTypeDescriptor->getGenericRequirements()) {
+      if (req.getKind() == GenericRequirementKind::Protocol) {
+        // Compute the address of the protocol descriptor offset as:
+        // opaqueTypeDescriptorAddress + offset of the protocol descriptor
+        // offset in the descriptor
+        auto protocolDescriptorOffsetOffset = (uintptr_t)(&req) +
+                                              req.getProtocolOffset() -
+                                              (uintptr_t)opaqueTypeDescriptor;
+        auto protocolDescriptorOffsetAddress =
+            opaqueTypeDescriptorAddress + protocolDescriptorOffsetOffset;
+        auto protocolDescriptorOffsetValue = req.getUnresolvedProtocolAddress();
+
+        // Compute the address of the protocol descriptor by following the
+        // offset
+        auto protocolDescriptorAddress = detail::applyRelativeOffset(
+            (const char *)protocolDescriptorOffsetAddress,
+            protocolDescriptorOffsetValue);
+
+        auto nameReader =
+            QualifiedContextNameReader<ObjCInteropKind, PointerSize>(
+                OpaqueByteReader, OpaqueStringReader, OpaquePointerReader,
+                OpaqueDynamicSymbolResolver);
+        auto conformanceRequirementProtocolName =
+            nameReader.readFullyQualifiedProtocolNameFromProtocolDescriptor(
+                protocolDescriptorAddress);
+        result.push_back(*conformanceRequirementProtocolName);
+      }
+    }
+    return result;
+  }
+
   ///
   /// Extraction of protocol conformances
   ///
@@ -1687,7 +1867,7 @@ class TypeRefBuilder {
     stream << "\n";
     stream << "ASSOCIATED TYPES:\n";
     stream << "=================\n";
-    dumpAssociatedTypeSection(stream);
+    dumpAssociatedTypeSection<ObjCInteropKind, PointerSize>(stream);
     stream << "\n";
     stream << "BUILTIN TYPES:\n";
     stream << "==============\n";
diff --git a/lib/StaticMirror/BinaryScanningTool.cpp b/lib/StaticMirror/BinaryScanningTool.cpp
@@ -80,12 +80,46 @@ BinaryScanningTool::collectConformances(const std::vector<std::string> &protocol
 
 AssociatedTypeCollectionResult
 BinaryScanningTool::collectAssociatedTypes(const std::string &mangledTypeName) {
-  return Context->Builder.collectAssociatedTypes(mangledTypeName);
+  switch (PointerSize) {
+    case 4:
+      // FIXME: This could/should be configurable.
+#if SWIFT_OBJC_INTEROP
+      return Context->Builder.collectAssociatedTypes<WithObjCInterop, 4>(mangledTypeName);
+#else
+      return Context->Builder.collectAssociatedTypes<NoObjCInterop, 4>(mangledTypeName);
+#endif
+    case 8:
+#if SWIFT_OBJC_INTEROP
+      return Context->Builder.collectAssociatedTypes<WithObjCInterop, 8>(mangledTypeName);
+#else
+      return Context->Builder.collectAssociatedTypes<NoObjCInterop, 8>(mangledTypeName);
+#endif
+    default:
+      fputs("unsupported word size in object file\n", stderr);
+      abort();
+  }
 }
 
 AssociatedTypeCollectionResult
 BinaryScanningTool::collectAllAssociatedTypes() {
-  return Context->Builder.collectAssociatedTypes(llvm::Optional<std::string>());
+  switch (PointerSize) {
+    case 4:
+      // FIXME: This could/should be configurable.
+#if SWIFT_OBJC_INTEROP
+      return Context->Builder.collectAssociatedTypes<WithObjCInterop, 4>(llvm::Optional<std::string>());
+#else
+      return Context->Builder.collectAssociatedTypes<NoObjCInterop, 4>(llvm::Optional<std::string>());
+#endif
+    case 8:
+#if SWIFT_OBJC_INTEROP
+      return Context->Builder.collectAssociatedTypes<WithObjCInterop, 8>(llvm::Optional<std::string>());
+#else
+      return Context->Builder.collectAssociatedTypes<NoObjCInterop, 8>(llvm::Optional<std::string>());
+#endif
+    default:
+      fputs("unsupported word size in object file\n", stderr);
+      abort();
+  }
 }
 
 FieldTypeCollectionResult
diff --git a/stdlib/public/Reflection/TypeRefBuilder.cpp b/stdlib/public/Reflection/TypeRefBuilder.cpp
diff --git a/test/Reflection/Inputs/swiftmodules/testModB.swift b/test/Reflection/Inputs/swiftmodules/testModB.swift
diff --git a/test/Reflection/opaque_associated_type_requirements.swift b/test/Reflection/opaque_associated_type_requirements.swift
