Merge pull request #41637 from artemcm/StaticMirrorAssociatedTypes

[Static Mirror] Refactor gathering of associated types from a binary and add SwiftStaticMirror entry-point for it.

Author: artemcm

include/swift-c/StaticMirror/BinaryScan.h

@@ -25,7 +25,7 @@
 /// SWIFTSTATICMIRROR_VERSION_MINOR should increase when there are API additions.
 /// SWIFTSTATICMIRROR_VERSION_MAJOR is intended for "major" source/ABI breaking changes.
 #define SWIFTSTATICMIRROR_VERSION_MAJOR 0
-#define SWIFTSTATICMIRROR_VERSION_MINOR 1
+#define SWIFTSTATICMIRROR_VERSION_MINOR 2 // Added associated-type gather
 
 SWIFTSTATICMIRROR_BEGIN_DECLS
 
@@ -42,11 +42,33 @@ typedef void *swift_static_mirror_t;
 typedef struct swift_static_mirror_conformance_info_s
     *swift_static_mirror_conformance_info_t;
 
+/// Opaque container to details of an associated type specification.
+typedef struct swift_static_mirror_type_alias_s
+    *swift_static_mirror_type_alias_t;
+
+/// Opaque container to an associated type (typealias) info of a given type.
+typedef struct swift_static_mirror_associated_type_info_s
+    *swift_static_mirror_associated_type_info_t;
+
+typedef struct {
+  swift_static_mirror_type_alias_t *type_aliases;
+  size_t count;
+} swift_static_mirror_type_alias_set_t;
+
+
+typedef struct {
+  swift_static_mirror_associated_type_info_t *associated_type_infos;
+  size_t count;
+} swift_static_mirror_associated_type_info_set_t;
+
 typedef struct {
   swift_static_mirror_conformance_info_t *conformances;
   size_t count;
 } swift_static_mirror_conformances_set_t;
 
+
+// swift_static_mirror_conformance_info query methods
+
 SWIFTSTATICMIRROR_PUBLIC swift_static_mirror_string_ref_t
     swift_static_mirror_conformance_info_get_type_name(
         swift_static_mirror_conformance_info_t);
@@ -63,6 +85,25 @@ SWIFTSTATICMIRROR_PUBLIC void
     swift_static_mirror_conformance_info_dispose(
         swift_static_mirror_conformance_info_t);
 
+// swift_static_mirror_associated_type query methods
+SWIFTSTATICMIRROR_PUBLIC swift_static_mirror_string_ref_t
+    swift_static_mirror_type_alias_get_type_alias_name(
+        swift_static_mirror_type_alias_t);
+SWIFTSTATICMIRROR_PUBLIC swift_static_mirror_string_ref_t
+    swift_static_mirror_type_alias_get_substituted_type_name(
+        swift_static_mirror_type_alias_t);
+SWIFTSTATICMIRROR_PUBLIC swift_static_mirror_string_ref_t
+    swift_static_mirror_type_alias_get_substituted_type_mangled_name(
+        swift_static_mirror_type_alias_t);
+
+// swift_static_mirror_associated_type_info query methods
+SWIFTSTATICMIRROR_PUBLIC swift_static_mirror_string_ref_t
+    swift_static_mirror_associated_type_info_get_mangled_type_name(
+        swift_static_mirror_associated_type_info_t);
+SWIFTSTATICMIRROR_PUBLIC swift_static_mirror_type_alias_set_t*
+    swift_static_mirror_associated_type_info_get_type_alias_set(
+        swift_static_mirror_associated_type_info_t);
+
 /// Create an \c swift_static_mirror_t instance.
 /// The returned \c swift_static_mirror_t is owned by the caller and must be
 /// disposed of using \c swift_static_mirror_dispose .
@@ -81,6 +122,19 @@ swift_static_mirror_conformances_set_create(
 SWIFTSTATICMIRROR_PUBLIC void swift_static_mirror_conformances_set_dispose(
     swift_static_mirror_conformances_set_t *);
 
+/// Identify and collect all associated types of a given input type (by mangled name)
+SWIFTSTATICMIRROR_PUBLIC swift_static_mirror_associated_type_info_set_t *
+swift_static_mirror_associated_type_info_set_create(
+    swift_static_mirror_t, const char *);
+
+/// Identify and collect associated types of all discovered types.
+SWIFTSTATICMIRROR_PUBLIC swift_static_mirror_associated_type_info_set_t *
+swift_static_mirror_all_associated_type_info_set_create(
+    swift_static_mirror_t);
+
+SWIFTSTATICMIRROR_PUBLIC void swift_static_mirror_associated_type_info_set_dispose(
+    swift_static_mirror_associated_type_info_set_t *);
+
 SWIFTSTATICMIRROR_END_DECLS
 
 #endif // SWIFT_C_BINARY_SCAN_H

include/swift/Reflection/TypeRefBuilder.h

@@ -297,6 +297,26 @@ struct ConformanceCollectionResult {
   std::vector<std::string> Errors;
 };
 
+struct AssociatedType {
+  std::string TypeAliasName;
+  std::string SubstitutedTypeMangledName;
+  std::string SubstitutedTypeFullyQualifiedName;
+  std::string SubstitutedTypeDiagnosticPrintName;
+};
+
+/// Info about a given type's associated type, as read out from an Image
+struct AssociatedTypeInfo {
+  std::string MangledTypeName;
+  std::string FullyQualifiedName;
+  std::string ProtocolFullyQualifiedName;
+  std::vector<AssociatedType> AssociatedTypes;
+};
+
+struct AssociatedTypeCollectionResult {
+  std::vector<AssociatedTypeInfo> AssociatedTypeInfos;
+  std::vector<std::string> Errors;
+};
+
 /// An implementation of MetadataReader's BuilderType concept for
 /// building TypeRefs, and parsing field metadata from any images
 /// it has been made aware of.
@@ -880,6 +900,7 @@ class TypeRefBuilder {
   void dumpTypeRef(RemoteRef<char> MangledName, std::ostream &stream,
                    bool printTypeName = false);
   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);

include/swift/StaticMirror/BinaryScanImpl.h

@@ -30,4 +30,15 @@ struct swift_static_mirror_conformance_info_s {
   swift_static_mirror_string_ref_t protocol_name;
 };
 
+struct swift_static_mirror_type_alias_s {
+  swift_static_mirror_string_ref_t type_alias_name;
+  swift_static_mirror_string_ref_t substituted_type_name;
+  swift_static_mirror_string_ref_t substituted_type_mangled_name;
+};
+
+struct swift_static_mirror_associated_type_info_s {
+  swift_static_mirror_string_ref_t mangled_type_name;
+  swift_static_mirror_type_alias_set_t *type_alias_set;
+};
+
 #endif // SWIFT_C_BINARY_SCAN_IMPL_H

include/swift/StaticMirror/BinaryScanningTool.h

@@ -46,6 +46,14 @@ class BinaryScanningTool {
   reflection::ConformanceCollectionResult
   collectConformances(const std::vector<std::string> &protocolNames);
 
+  /// Collect all associated types of the specified type
+  reflection::AssociatedTypeCollectionResult
+  collectAssociatedTypes(const std::string &mangledTypeName);
+
+  /// Collect associated types of all discovered types
+  reflection::AssociatedTypeCollectionResult
+  collectAllAssociatedTypes();
+
 private:
   // Note: binaryOrError and objectOrError own the memory for our ObjectFile;
   // once they go out of scope, we can no longer do anything.

lib/StaticMirror/BinaryScanningTool.cpp

@@ -78,5 +78,14 @@ BinaryScanningTool::collectConformances(const std::vector<std::string> &protocol
   }
 }
 
+AssociatedTypeCollectionResult
+BinaryScanningTool::collectAssociatedTypes(const std::string &mangledTypeName) {
+  return Context->Builder.collectAssociatedTypes(mangledTypeName);
+}
+
+AssociatedTypeCollectionResult
+BinaryScanningTool::collectAllAssociatedTypes() {
+  return Context->Builder.collectAssociatedTypes(llvm::Optional<std::string>());
+}
 } // end namespace static_mirror
 } // end namespace swift

stdlib/public/Reflection/TypeRefBuilder.cpp

@@ -25,6 +25,7 @@
 #include "swift/Remote/MetadataReader.h"
 #include <iomanip>
 #include <iostream>
+#include <sstream>
 
 using namespace swift;
 using namespace reflection;
@@ -457,32 +458,72 @@ void TypeRefBuilder::dumpFieldSection(std::ostream &stream) {
   }
 }
 
-void TypeRefBuilder::dumpAssociatedTypeSection(std::ostream &stream) {
+AssociatedTypeCollectionResult TypeRefBuilder::collectAssociatedTypes(
+    llvm::Optional<std::string> forMangledTypeName) {
+  AssociatedTypeCollectionResult result;
   for (const auto &sections : ReflectionInfos) {
     for (auto descriptor : sections.AssociatedType) {
-      auto conformingTypeNode = demangleTypeRef(
-          readTypeRef(descriptor, descriptor->ConformingTypeName));
-      auto conformingTypeName = nodeToString(conformingTypeNode);
+      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();
-
-      stream << "- " << conformingTypeName << " : " << protocolName << "\n";
-
-      for (const auto &associatedTypeRef : *descriptor.getLocalBuffer()) {
-        auto associatedType = descriptor.getField(associatedTypeRef);
-
-        std::string name =
-            getTypeRefString(readTypeRef(associatedType, associatedType->Name))
-                .str();
-        stream << "typealias " << name << " = ";
-        dumpTypeRef(
-            readTypeRef(associatedType, associatedType->SubstitutedTypeName),
-            stream);
+      if (optionalMangledTypeName.hasValue()) {
+        auto mangledTypeName =
+            optionalMangledTypeName.getValue().insert(0, "$s");
+        if (forMangledTypeName.hasValue()) {
+          if (mangledTypeName != forMangledTypeName.getValue())
+            continue;
+        }
+        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 =
+                "$s" + optionalMangledSubstitutedTypeName.getValue();
+          }
+          auto substitutedDemangleTree = demangleTypeRef(substitutedTypeRef);
+          auto substitutedTypeName = nodeToString(substitutedDemangleTree);
+          std::stringstream OS;
+          dumpTypeRef(substitutedTypeRef, OS);
+          associatedTypes.emplace_back(
+              AssociatedType{typealiasTypeName, mangledSubstitutedTypeName,
+                             substitutedTypeName, OS.str()});
+        }
+        result.AssociatedTypeInfos.emplace_back(AssociatedTypeInfo{
+            mangledTypeName, typeName, protocolName, associatedTypes});
       }
     }
   }
+  return result;
+}
+
+void TypeRefBuilder::dumpAssociatedTypeSection(std::ostream &stream) {
+  auto associatedTypeCollectionResult =
+      collectAssociatedTypes(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;
+    }
+    stream << "\n";
+  }
 }
 
 void TypeRefBuilder::dumpBuiltinTypeSection(std::ostream &stream) {

tools/libStaticMirror/libStaticMirror.cpp

@@ -18,7 +18,6 @@
 #include "swift/StaticMirror/BinaryScanningTool.h"
 #include "swift/StaticMirror/BinaryScanImpl.h"
 #include "swift/DependencyScan/StringUtils.h"
-//#include "swift/Option/Options.h"
 
 // FIXME: Code duplication with StringUtils.cpp
 namespace swift {
@@ -136,3 +135,119 @@ void swift_static_mirror_conformances_set_dispose(
   delete[] set->conformances;
   delete set;
 }
+
+static swift_static_mirror_associated_type_info_set_t *
+convertAssociatedTypeQueryResult(
+    const swift::reflection::AssociatedTypeCollectionResult &scanResult) {
+  swift_static_mirror_associated_type_info_set_t *result =
+      new swift_static_mirror_associated_type_info_set_t;
+  result->count = scanResult.AssociatedTypeInfos.size();
+  result->associated_type_infos = new swift_static_mirror_associated_type_info_t
+      [scanResult.AssociatedTypeInfos.size()];
+
+  int associatedTypeInfoIndex = 0;
+  for (const auto &assocTypeInfo : scanResult.AssociatedTypeInfos) {
+    swift_static_mirror_associated_type_info_s *info =
+        new swift_static_mirror_associated_type_info_s;
+    info->mangled_type_name = swift::c_string_utils::create_clone(
+        assocTypeInfo.MangledTypeName.c_str());
+    info->type_alias_set = new swift_static_mirror_type_alias_set_t;
+    info->type_alias_set->count = assocTypeInfo.AssociatedTypes.size();
+    info->type_alias_set->type_aliases =
+        new swift_static_mirror_type_alias_t[assocTypeInfo.AssociatedTypes
+                                                      .size()];
+    int typealiasIndex = 0;
+    for (const auto &typeAliasInfo : assocTypeInfo.AssociatedTypes) {
+      swift_static_mirror_type_alias_s *typealiasDetails =
+          new swift_static_mirror_type_alias_s;
+      typealiasDetails->type_alias_name = swift::c_string_utils::create_clone(
+          typeAliasInfo.TypeAliasName.c_str());
+      typealiasDetails->substituted_type_name =
+          swift::c_string_utils::create_clone(
+              typeAliasInfo.SubstitutedTypeFullyQualifiedName.c_str());
+      typealiasDetails->substituted_type_mangled_name =
+          swift::c_string_utils::create_clone(
+              typeAliasInfo.SubstitutedTypeMangledName.c_str());
+      info->type_alias_set->type_aliases[typealiasIndex] =
+          typealiasDetails;
+      typealiasIndex += 1;
+    }
+    result->associated_type_infos[associatedTypeInfoIndex] = info;
+    associatedTypeInfoIndex += 1;
+  }
+  return result;
+}
+
+swift_static_mirror_associated_type_info_set_t *
+swift_static_mirror_associated_type_info_set_create(
+    swift_static_mirror_t static_mirror, const char *forTypeName) {
+  BinaryScanningTool *scanTool = unwrap(static_mirror);
+  auto scanResult = scanTool->collectAssociatedTypes(forTypeName);
+  return convertAssociatedTypeQueryResult(scanResult);
+}
+
+/// Identify and collect associated types of all discovered types.
+swift_static_mirror_associated_type_info_set_t *
+swift_static_mirror_all_associated_type_info_set_create(
+    swift_static_mirror_t static_mirror) {
+  BinaryScanningTool *scanTool = unwrap(static_mirror);
+  auto scanResult = scanTool->collectAllAssociatedTypes();
+  return convertAssociatedTypeQueryResult(scanResult);
+}
+
+// swift_static_mirror_associated_type query methods
+swift_static_mirror_string_ref_t
+swift_static_mirror_type_alias_get_type_alias_name(
+    swift_static_mirror_type_alias_t type_alias) {
+  return type_alias->type_alias_name;
+}
+swift_static_mirror_string_ref_t
+swift_static_mirror_type_alias_get_substituted_type_name(
+    swift_static_mirror_type_alias_t type_alias) {
+  return type_alias->substituted_type_name;
+}
+swift_static_mirror_string_ref_t
+swift_static_mirror_type_alias_get_substituted_type_mangled_name(
+    swift_static_mirror_type_alias_t type_alias) {
+  return type_alias->substituted_type_mangled_name;
+}
+
+// swift_static_mirror_associated_type_info query methods
+swift_static_mirror_string_ref_t
+swift_static_mirror_associated_type_info_get_mangled_type_name(
+    swift_static_mirror_associated_type_info_t associated_type_info) {
+  return associated_type_info->mangled_type_name;
+}
+swift_static_mirror_type_alias_set_t*
+swift_static_mirror_associated_type_info_get_type_alias_set(
+    swift_static_mirror_associated_type_info_t associated_type_info) {
+  return associated_type_info->type_alias_set;
+}
+
+void swift_static_mirror_type_alias_dispose(swift_static_mirror_type_alias_t type_alias) {
+  swift_static_mirror_string_dispose(type_alias->substituted_type_mangled_name);
+  swift_static_mirror_string_dispose(type_alias->substituted_type_name);
+  swift_static_mirror_string_dispose(type_alias->type_alias_name);
+  delete type_alias;
+}
+
+void swift_static_mirror_type_alias_set_dispose(swift_static_mirror_type_alias_set_t *type_alias_set) {
+  for (size_t i = 0; i < type_alias_set->count; ++i) {
+    swift_static_mirror_type_alias_dispose(type_alias_set->type_aliases[i]);
+  }
+  delete[] type_alias_set->type_aliases;
+  delete type_alias_set;
+}
+
+void swift_static_mirror_associated_type_info_dispose(swift_static_mirror_associated_type_info_t associated_type_info) {
+  swift_static_mirror_string_dispose(associated_type_info->mangled_type_name);
+  swift_static_mirror_type_alias_set_dispose(associated_type_info->type_alias_set);
+}
+
+void swift_static_mirror_associated_type_info_set_dispose(swift_static_mirror_associated_type_info_set_t *associated_type_info_set) {
+  for (size_t i = 0; i < associated_type_info_set->count; ++i) {
+    swift_static_mirror_associated_type_info_dispose(associated_type_info_set->associated_type_infos[i]);
+  }
+  delete[] associated_type_info_set->associated_type_infos;
+  delete associated_type_info_set;
+}