Implement heuristic to prioritize in field descriptor search

augusto2112 · augusto2112 · commit 781feaa0bae2 · 2022-06-30T14:42:59.000-07:00
Currently when looking for field descriptors we parse the reflection
metadata in whatever order it was registered. This patch implements a
heuristic where we try to match a new optional Name field with the
module name of the type's field descriptor we're looking for.

rdar://87889973
diff --git a/include/swift/Reflection/ReflectionContext.h b/include/swift/Reflection/ReflectionContext.h
@@ -227,7 +227,9 @@ class ReflectionContext
     return sizeof(StoredPointer) * 2;
   }
 
-  template <typename T> bool readMachOSections(RemoteAddress ImageStart) {
+  template <typename T>
+  bool readMachOSections(RemoteAddress ImageStart,
+                         llvm::Optional<llvm::StringRef> Name = llvm::None) {
     auto Buf =
         this->getReader().readBytes(ImageStart, sizeof(typename T::Header));
     if (!Buf)
@@ -337,15 +339,15 @@ class ReflectionContext
         MPEnumMdSec.first == nullptr)
       return false;
 
-    ReflectionInfo info = {
-        {FieldMdSec.first, FieldMdSec.second},
-        {AssocTySec.first, AssocTySec.second},
-        {BuiltinTySec.first, BuiltinTySec.second},
-        {CaptureSec.first, CaptureSec.second},
-        {TypeRefMdSec.first, TypeRefMdSec.second},
-        {ReflStrMdSec.first, ReflStrMdSec.second},
-        {ConformMdSec.first, ConformMdSec.second},
-        {MPEnumMdSec.first, MPEnumMdSec.second}};
+    ReflectionInfo info = {{FieldMdSec.first, FieldMdSec.second},
+                           {AssocTySec.first, AssocTySec.second},
+                           {BuiltinTySec.first, BuiltinTySec.second},
+                           {CaptureSec.first, CaptureSec.second},
+                           {TypeRefMdSec.first, TypeRefMdSec.second},
+                           {ReflStrMdSec.first, ReflStrMdSec.second},
+                           {ConformMdSec.first, ConformMdSec.second},
+                           {MPEnumMdSec.first, MPEnumMdSec.second},
+                           Name};
 
     this->addReflectionInfo(info);
 
@@ -374,7 +376,8 @@ class ReflectionContext
     return true;
   }
 
-  bool readPECOFFSections(RemoteAddress ImageStart) {
+  bool readPECOFFSections(RemoteAddress ImageStart,
+                          llvm::Optional<llvm::StringRef> Name = llvm::None) {
     auto DOSHdrBuf = this->getReader().readBytes(
         ImageStart, sizeof(llvm::object::dos_header));
     if (!DOSHdrBuf)
@@ -463,20 +466,21 @@ class ReflectionContext
         MPEnumMdSec.first == nullptr)
       return false;
 
-    ReflectionInfo Info = {
-        {FieldMdSec.first, FieldMdSec.second},
-        {AssocTySec.first, AssocTySec.second},
-        {BuiltinTySec.first, BuiltinTySec.second},
-        {CaptureSec.first, CaptureSec.second},
-        {TypeRefMdSec.first, TypeRefMdSec.second},
-        {ReflStrMdSec.first, ReflStrMdSec.second},
-        {ConformMdSec.first, ConformMdSec.second},
-        {MPEnumMdSec.first, MPEnumMdSec.second}};
+    ReflectionInfo Info = {{FieldMdSec.first, FieldMdSec.second},
+                           {AssocTySec.first, AssocTySec.second},
+                           {BuiltinTySec.first, BuiltinTySec.second},
+                           {CaptureSec.first, CaptureSec.second},
+                           {TypeRefMdSec.first, TypeRefMdSec.second},
+                           {ReflStrMdSec.first, ReflStrMdSec.second},
+                           {ConformMdSec.first, ConformMdSec.second},
+                           {MPEnumMdSec.first, MPEnumMdSec.second},
+                           Name};
     this->addReflectionInfo(Info);
     return true;
   }
 
-  bool readPECOFF(RemoteAddress ImageStart) {
+  bool readPECOFF(RemoteAddress ImageStart,
+                  llvm::Optional<llvm::StringRef> Name = llvm::None) {
     auto Buf = this->getReader().readBytes(ImageStart,
                                            sizeof(llvm::object::dos_header));
     if (!Buf)
@@ -495,12 +499,13 @@ class ReflectionContext
     if (memcmp(Buf.get(), llvm::COFF::PEMagic, sizeof(llvm::COFF::PEMagic)))
       return false;
 
-    return readPECOFFSections(ImageStart);
+    return readPECOFFSections(ImageStart, Name);
   }
 
   template <typename T>
   bool readELFSections(RemoteAddress ImageStart,
-                       llvm::Optional<llvm::sys::MemoryBlock> FileBuffer) {
+                       llvm::Optional<llvm::sys::MemoryBlock> FileBuffer,
+                       llvm::Optional<llvm::StringRef> Name = llvm::None) {
     // When reading from the FileBuffer we can simply return a pointer to
     // the underlying data.
     // When reading from the process, we need to keep the memory around
@@ -655,15 +660,15 @@ class ReflectionContext
         MPEnumMdSec.first == nullptr)
       return false;
 
-    ReflectionInfo info = {
-        {FieldMdSec.first, FieldMdSec.second},
-        {AssocTySec.first, AssocTySec.second},
-        {BuiltinTySec.first, BuiltinTySec.second},
-        {CaptureSec.first, CaptureSec.second},
-        {TypeRefMdSec.first, TypeRefMdSec.second},
-        {ReflStrMdSec.first, ReflStrMdSec.second},
-        {ConformMdSec.first, ConformMdSec.second},
-        {MPEnumMdSec.first, MPEnumMdSec.second}};
+    ReflectionInfo info = {{FieldMdSec.first, FieldMdSec.second},
+                           {AssocTySec.first, AssocTySec.second},
+                           {BuiltinTySec.first, BuiltinTySec.second},
+                           {CaptureSec.first, CaptureSec.second},
+                           {TypeRefMdSec.first, TypeRefMdSec.second},
+                           {ReflStrMdSec.first, ReflStrMdSec.second},
+                           {ConformMdSec.first, ConformMdSec.second},
+                           {MPEnumMdSec.first, MPEnumMdSec.second},
+                           Name};
 
     this->addReflectionInfo(info);
     return true;
@@ -687,7 +692,9 @@ class ReflectionContext
   /// \return
   ///     /b True if the metadata information was parsed successfully,
   ///     /b false otherwise.
-  bool readELF(RemoteAddress ImageStart, llvm::Optional<llvm::sys::MemoryBlock> FileBuffer) {
+  bool readELF(RemoteAddress ImageStart,
+               llvm::Optional<llvm::sys::MemoryBlock> FileBuffer,
+               llvm::Optional<llvm::StringRef> Name = llvm::None) {
     auto Buf =
         this->getReader().readBytes(ImageStart, sizeof(llvm::ELF::Elf64_Ehdr));
     if (!Buf)
@@ -703,16 +710,16 @@ class ReflectionContext
     unsigned char FileClass = Hdr->getFileClass();
     if (FileClass == llvm::ELF::ELFCLASS64) {
       return readELFSections<ELFTraits<llvm::ELF::ELFCLASS64>>(
-          ImageStart, FileBuffer);
+          ImageStart, FileBuffer, Name);
     } else if (FileClass == llvm::ELF::ELFCLASS32) {
       return readELFSections<ELFTraits<llvm::ELF::ELFCLASS32>>(
-          ImageStart, FileBuffer);
+          ImageStart, FileBuffer, Name);
     } else {
       return false;
     }
   }
 
-  bool addImage(RemoteAddress ImageStart) {
+  bool addImage(RemoteAddress ImageStart, llvm::Optional<llvm::StringRef> Name = llvm::None) {
     // Read the first few bytes to look for a magic header.
     auto Magic = this->getReader().readBytes(ImageStart, sizeof(uint32_t));
     if (!Magic)
@@ -723,18 +730,18 @@ class ReflectionContext
 
     // 32- and 64-bit Mach-O.
     if (MagicWord == llvm::MachO::MH_MAGIC) {
-      return readMachOSections<MachOTraits<4>>(ImageStart);
+      return readMachOSections<MachOTraits<4>>(ImageStart, Name);
     }
 
     if (MagicWord == llvm::MachO::MH_MAGIC_64) {
-      return readMachOSections<MachOTraits<8>>(ImageStart);
+      return readMachOSections<MachOTraits<8>>(ImageStart, Name);
     }
 
     // PE. (This just checks for the DOS header; `readPECOFF` will further
     // validate the existence of the PE header.)
     auto MagicBytes = (const char*)Magic.get();
     if (MagicBytes[0] == 'M' && MagicBytes[1] == 'Z') {
-      return readPECOFF(ImageStart);
+      return readPECOFF(ImageStart, Name);
     }
 
 
@@ -743,7 +750,7 @@ class ReflectionContext
         && MagicBytes[1] == llvm::ELF::ElfMagic[1]
         && MagicBytes[2] == llvm::ELF::ElfMagic[2]
         && MagicBytes[3] == llvm::ELF::ElfMagic[3]) {
-      return readELF(ImageStart, llvm::Optional<llvm::sys::MemoryBlock>());
+      return readELF(ImageStart, llvm::Optional<llvm::sys::MemoryBlock>(), Name);
     }
 
     // We don't recognize the format.
@@ -759,8 +766,9 @@ class ReflectionContext
   ///     \b True if any of the reflection sections were registered,
   ///     \b false otherwise.
   bool addImage(llvm::function_ref<
-                std::pair<RemoteRef<void>, uint64_t>(ReflectionSectionKind)>
-                    FindSection) {
+                    std::pair<RemoteRef<void>, uint64_t>(ReflectionSectionKind)>
+                    FindSection,
+                llvm::Optional<llvm::StringRef> Name) {
     auto Sections = {
         ReflectionSectionKind::fieldmd, ReflectionSectionKind::assocty,
         ReflectionSectionKind::builtin, ReflectionSectionKind::capture,
@@ -784,11 +792,15 @@ class ReflectionContext
     if (llvm::all_of(Pairs, [](const auto &Pair) { return !Pair.first; }))
       return false;
 
-    ReflectionInfo Info = {
-        {Pairs[0].first, Pairs[0].second}, {Pairs[1].first, Pairs[1].second},
-        {Pairs[2].first, Pairs[2].second}, {Pairs[3].first, Pairs[3].second},
-        {Pairs[4].first, Pairs[4].second}, {Pairs[5].first, Pairs[5].second},
-        {Pairs[6].first, Pairs[6].second}, {Pairs[7].first, Pairs[7].second}};
+    ReflectionInfo Info = {{Pairs[0].first, Pairs[0].second},
+                           {Pairs[1].first, Pairs[1].second},
+                           {Pairs[2].first, Pairs[2].second},
+                           {Pairs[3].first, Pairs[3].second},
+                           {Pairs[4].first, Pairs[4].second},
+                           {Pairs[5].first, Pairs[5].second},
+                           {Pairs[6].first, Pairs[6].second},
+                           {Pairs[7].first, Pairs[7].second},
+                           Name};
     this->addReflectionInfo(Info);
     return true;
   }
diff --git a/include/swift/Reflection/TypeRefBuilder.h b/include/swift/Reflection/TypeRefBuilder.h
@@ -29,6 +29,7 @@
 #include <ostream>
 #include <sstream>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 namespace swift {
@@ -264,6 +265,7 @@ struct ReflectionInfo {
   GenericSection ReflectionString;
   GenericSection Conformance;
   MultiPayloadEnumSection MultiPayloadEnum;
+  llvm::Optional<llvm::StringRef> Name;
 };
 
 struct ClosureContextInfo {
@@ -827,10 +829,9 @@ class TypeRefBuilder {
 private:
   std::vector<ReflectionInfo> ReflectionInfos;
 
-  /// Index of the next Reflection Info that should be processed.
-  /// This assumes that Reflection Infos are never removed from the vector.
-  size_t FirstUnprocessedReflectionInfoIndex = 0;
-    
+  /// Reflection Infos we've already processed.
+  std::unordered_set<const ReflectionInfo *> ProcessedReflectionInfos;
+
   llvm::Optional<std::string> normalizeReflectionName(RemoteRef<char> name);
   bool reflectionNameMatches(RemoteRef<char> reflectionName,
                              StringRef searchName);
@@ -1548,7 +1549,7 @@ class TypeRefBuilder {
                                      mangledTypeName};
     }
   };
-
+  void PopulateFieldTypeInfoCacheWithReflectionInfo(ReflectionInfo &Info);
 public:
   template <template <typename Runtime> class ObjCInteropKind,
             unsigned PointerSize>
diff --git a/stdlib/public/Reflection/TypeRefBuilder.cpp b/stdlib/public/Reflection/TypeRefBuilder.cpp
@@ -193,36 +193,85 @@ const TypeRef *TypeRefBuilder::lookupSuperclass(const TypeRef *TR) {
   return Unsubstituted->subst(*this, *SubstMap);
 }
 
+static llvm::Optional<StringRef> FindOutermostModuleName(NodePointer Node) {
+  // Breadth first search until we find the module name so we find the outermost
+  // one.
+  llvm::SmallVector<NodePointer, 8> Queue;
+  Queue.push_back(Node);
+  // Instead of removing items from the front of the queue we just iterate over
+  // them.
+  for (size_t i = 0; i < Queue.size(); ++i) {
+    NodePointer Current = Queue[i];
+    if (Current->getKind() == Node::Kind::Module) {
+      if (Current->hasText())
+        return Current->getText();
+      else
+        return {};
+    }
+    for (auto Child : *Current)
+      Queue.push_back(Child);
+  }
+  return {};
+}
+
+void TypeRefBuilder::PopulateFieldTypeInfoCacheWithReflectionInfo(
+    ReflectionInfo &Info) {
+  if (ProcessedReflectionInfos.find(&Info) != ProcessedReflectionInfos.end())
+    return;
+
+  for (auto FD : Info.Field) {
+    if (!FD->hasMangledTypeName())
+      continue;
+    auto CandidateMangledName = readTypeRef(FD, FD->MangledTypeName);
+    if (auto NormalizedName = normalizeReflectionName(CandidateMangledName)) {
+      FieldTypeInfoCache[std::move(*NormalizedName)] = FD;
+    }
+  }
+
+  ProcessedReflectionInfos.emplace(&Info);
+}
+
 RemoteRef<FieldDescriptor>
 TypeRefBuilder::getFieldTypeInfo(const TypeRef *TR) {
   const std::string *MangledName;
-  if (auto N = dyn_cast<NominalTypeRef>(TR))
+  NodePointer Node;
+  Demangler Dem;
+  if (auto N = dyn_cast<NominalTypeRef>(TR)) {
+    Node = N->getDemangling(Dem);
     MangledName = &N->getMangledName();
-  else if (auto BG = dyn_cast<BoundGenericTypeRef>(TR))
+  } else if (auto BG = dyn_cast<BoundGenericTypeRef>(TR)) {
+    Node = BG->getDemangling(Dem);
     MangledName = &BG->getMangledName();
+  }
   else
     return nullptr;
 
   // Try the cache.
   auto Found = FieldTypeInfoCache.find(*MangledName);
   if (Found != FieldTypeInfoCache.end())
     return Found->second;
+  
+  // Heuristic: find the outermost Module node available, and try to parse the 
+  // ReflectionInfos with a matching name first.
+  auto ModuleName = FindOutermostModuleName(Node);
+  // If we couldn't find a module name or the type is imported (__C module) we
+  // don't any useful information on which image to look for the type.
+  if (ModuleName && ModuleName != llvm::StringRef("__C")) {
+    for (auto &Info : ReflectionInfos) {
+      if (ModuleName == Info.Name) {
+        PopulateFieldTypeInfoCacheWithReflectionInfo(Info);
+        Found = FieldTypeInfoCache.find(*MangledName);
+        if (Found != FieldTypeInfoCache.end()) {
+          return Found->second;
+        }
+      }
+    }
+  }
 
   // On failure, fill out the cache, ReflectionInfo by ReflectionInfo,
   // until we find the field descriptor we're looking for.
-  while (FirstUnprocessedReflectionInfoIndex < ReflectionInfos.size()) {
-    auto &Info = ReflectionInfos[FirstUnprocessedReflectionInfoIndex];
-    for (auto FD : Info.Field) {
-      if (!FD->hasMangledTypeName())
-        continue;
-      auto CandidateMangledName = readTypeRef(FD, FD->MangledTypeName);
-      if (auto NormalizedName = normalizeReflectionName(CandidateMangledName))
-        FieldTypeInfoCache[std::move(*NormalizedName)] = FD;
-    }
-
-    // Since we're done with the current ReflectionInfo, increment early in
-    // case we get a cache hit.
-    ++FirstUnprocessedReflectionInfoIndex;
+  for (auto &Info : ReflectionInfos) {
+    PopulateFieldTypeInfoCacheWithReflectionInfo(Info);
     Found = FieldTypeInfoCache.find(*MangledName);
     if (Found != FieldTypeInfoCache.end())
       return Found->second;
diff --git a/stdlib/public/SwiftRemoteMirror/SwiftRemoteMirror.cpp b/stdlib/public/SwiftRemoteMirror/SwiftRemoteMirror.cpp
@@ -238,14 +238,15 @@ swift_reflection_addReflectionInfo(SwiftReflectionContextRef ContextRef,
   }
 
   ReflectionInfo ContextInfo{
-    sectionFromInfo<FieldDescriptorIterator>(Info, Info.field),
-    sectionFromInfo<AssociatedTypeIterator>(Info, Info.associated_types),
-    sectionFromInfo<BuiltinTypeDescriptorIterator>(Info, Info.builtin_types),
-    sectionFromInfo<CaptureDescriptorIterator>(Info, Info.capture),
-    sectionFromInfo<const void *>(Info, Info.type_references),
-    sectionFromInfo<const void *>(Info, Info.reflection_strings),
-    ReflectionSection<const void *>(nullptr, 0),
-    ReflectionSection<MultiPayloadEnumDescriptorIterator>(0, 0)};
+      sectionFromInfo<FieldDescriptorIterator>(Info, Info.field),
+      sectionFromInfo<AssociatedTypeIterator>(Info, Info.associated_types),
+      sectionFromInfo<BuiltinTypeDescriptorIterator>(Info, Info.builtin_types),
+      sectionFromInfo<CaptureDescriptorIterator>(Info, Info.capture),
+      sectionFromInfo<const void *>(Info, Info.type_references),
+      sectionFromInfo<const void *>(Info, Info.reflection_strings),
+      ReflectionSection<const void *>(nullptr, 0),
+      ReflectionSection<MultiPayloadEnumDescriptorIterator>(0, 0),
+      llvm::None};
 
   Context->addReflectionInfo(ContextInfo);
 }
@@ -264,10 +265,11 @@ void swift_reflection_addReflectionMappingInfo(
       reflectionSectionFromLocalAndRemote<CaptureDescriptorIterator>(
           Info.capture),
       reflectionSectionFromLocalAndRemote<const void *>(Info.type_references),
-      reflectionSectionFromLocalAndRemote<const void *>(Info.reflection_strings),
+      reflectionSectionFromLocalAndRemote<const void *>(
+          Info.reflection_strings),
       ReflectionSection<const void *>(nullptr, 0),
-      MultiPayloadEnumSection(0, 0)
-  };
+      MultiPayloadEnumSection(0, 0),
+      llvm::None};
 
   Context->addReflectionInfo(ContextInfo);
 }
