Merge pull request #59802 from augusto2112/heuristic-get-field-descriptor

augusto2112 · web-flow · commit 027797121d9b · 2022-07-06T13:24:32.000-03:00
Implement heuristic to prioritize reflection info in field descriptor search
diff --git a/include/swift/Reflection/ReflectionContext.h b/include/swift/Reflection/ReflectionContext.h
@@ -227,7 +227,10 @@ class ReflectionContext
     return sizeof(StoredPointer) * 2;
   }
 
-  template <typename T> bool readMachOSections(RemoteAddress ImageStart) {
+  template <typename T>
+  bool readMachOSections(
+      RemoteAddress ImageStart,
+      llvm::SmallVector<llvm::StringRef, 1> PotentialModuleNames = {}) {
     auto Buf =
         this->getReader().readBytes(ImageStart, sizeof(typename T::Header));
     if (!Buf)
@@ -337,15 +340,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},
+                           PotentialModuleNames};
 
     this->addReflectionInfo(info);
 
@@ -374,7 +377,9 @@ class ReflectionContext
     return true;
   }
 
-  bool readPECOFFSections(RemoteAddress ImageStart) {
+  bool readPECOFFSections(
+      RemoteAddress ImageStart,
+      llvm::SmallVector<llvm::StringRef, 1> PotentialModuleNames = {}) {
     auto DOSHdrBuf = this->getReader().readBytes(
         ImageStart, sizeof(llvm::object::dos_header));
     if (!DOSHdrBuf)
@@ -463,20 +468,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},
+                           PotentialModuleNames};
     this->addReflectionInfo(Info);
     return true;
   }
 
-  bool readPECOFF(RemoteAddress ImageStart) {
+  bool readPECOFF(RemoteAddress ImageStart,
+                  llvm::SmallVector<llvm::StringRef, 1> PotentialModuleNames = {}) {
     auto Buf = this->getReader().readBytes(ImageStart,
                                            sizeof(llvm::object::dos_header));
     if (!Buf)
@@ -495,12 +501,14 @@ class ReflectionContext
     if (memcmp(Buf.get(), llvm::COFF::PEMagic, sizeof(llvm::COFF::PEMagic)))
       return false;
 
-    return readPECOFFSections(ImageStart);
+    return readPECOFFSections(ImageStart, PotentialModuleNames);
   }
 
   template <typename T>
-  bool readELFSections(RemoteAddress ImageStart,
-                       llvm::Optional<llvm::sys::MemoryBlock> FileBuffer) {
+  bool readELFSections(
+      RemoteAddress ImageStart,
+      llvm::Optional<llvm::sys::MemoryBlock> FileBuffer,
+      llvm::SmallVector<llvm::StringRef, 1> PotentialModuleNames = {}) {
     // 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 +663,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},
+                           PotentialModuleNames};
 
     this->addReflectionInfo(info);
     return true;
@@ -687,7 +695,10 @@ 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::SmallVector<llvm::StringRef, 1> PotentialModuleNames = {}) {
     auto Buf =
         this->getReader().readBytes(ImageStart, sizeof(llvm::ELF::Elf64_Ehdr));
     if (!Buf)
@@ -703,16 +714,18 @@ class ReflectionContext
     unsigned char FileClass = Hdr->getFileClass();
     if (FileClass == llvm::ELF::ELFCLASS64) {
       return readELFSections<ELFTraits<llvm::ELF::ELFCLASS64>>(
-          ImageStart, FileBuffer);
+          ImageStart, FileBuffer, PotentialModuleNames);
     } else if (FileClass == llvm::ELF::ELFCLASS32) {
       return readELFSections<ELFTraits<llvm::ELF::ELFCLASS32>>(
-          ImageStart, FileBuffer);
+          ImageStart, FileBuffer, PotentialModuleNames);
     } else {
       return false;
     }
   }
 
-  bool addImage(RemoteAddress ImageStart) {
+  bool
+  addImage(RemoteAddress ImageStart,
+           llvm::SmallVector<llvm::StringRef, 1> PotentialModuleNames = {}) {
     // Read the first few bytes to look for a magic header.
     auto Magic = this->getReader().readBytes(ImageStart, sizeof(uint32_t));
     if (!Magic)
@@ -723,18 +736,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, PotentialModuleNames);
     }
 
     if (MagicWord == llvm::MachO::MH_MAGIC_64) {
-      return readMachOSections<MachOTraits<8>>(ImageStart);
+      return readMachOSections<MachOTraits<8>>(ImageStart, PotentialModuleNames);
     }
 
     // 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, PotentialModuleNames);
     }
 
 
@@ -743,7 +756,8 @@ 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>(),
+                     PotentialModuleNames);
     }
 
     // We don't recognize the format.
@@ -758,9 +772,11 @@ class ReflectionContext
   /// \return
   ///     \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) {
+  bool
+  addImage(llvm::function_ref<
+               std::pair<RemoteRef<void>, uint64_t>(ReflectionSectionKind)>
+               FindSection,
+           llvm::SmallVector<llvm::StringRef, 1> PotentialModuleNames = {}) {
     auto Sections = {
         ReflectionSectionKind::fieldmd, ReflectionSectionKind::assocty,
         ReflectionSectionKind::builtin, ReflectionSectionKind::capture,
@@ -784,11 +800,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},
+                           PotentialModuleNames};
     this->addReflectionInfo(Info);
     return true;
   }
diff --git a/include/swift/Reflection/TypeRefBuilder.h b/include/swift/Reflection/TypeRefBuilder.h
@@ -24,6 +24,7 @@
 #include "swift/Reflection/TypeLowering.h"
 #include "swift/Reflection/TypeRef.h"
 #include "llvm/ADT/Optional.h"
+#include "llvm/ADT/SmallVector.h"
 #include <iomanip>
 #include <iostream>
 #include <ostream>
@@ -264,6 +265,7 @@ struct ReflectionInfo {
   GenericSection ReflectionString;
   GenericSection Conformance;
   MultiPayloadEnumSection MultiPayloadEnum;
+  llvm::SmallVector<llvm::StringRef, 1> PotentialModuleNames;
 };
 
 struct ClosureContextInfo {
@@ -926,13 +928,15 @@ 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;
-    
+  /// Indexes of Reflection Infos we've already processed.
+  llvm::DenseSet<size_t> ProcessedReflectionInfoIndexes;
+
   llvm::Optional<std::string> normalizeReflectionName(RemoteRef<char> name);
   bool reflectionNameMatches(RemoteRef<char> reflectionName,
                              StringRef searchName);
+  void populateFieldTypeInfoCacheWithReflectionAtIndex(size_t Index);
+  llvm::Optional<RemoteRef<FieldDescriptor>>
+  findFieldDescriptorAtIndex(size_t Index, const std::string &MangledName);
 
 public:
   RemoteRef<char> readTypeRef(uint64_t remoteAddr);
@@ -1871,7 +1875,6 @@ class TypeRefBuilder {
                                      mangledTypeName};
     }
   };
-
 public:
   template <template <typename Runtime> class ObjCInteropKind,
             unsigned PointerSize>
diff --git a/stdlib/public/Reflection/TypeRefBuilder.cpp b/stdlib/public/Reflection/TypeRefBuilder.cpp
@@ -193,40 +193,93 @@ const TypeRef *TypeRefBuilder::lookupSuperclass(const TypeRef *TR) {
   return Unsubstituted->subst(*this, *SubstMap);
 }
 
-RemoteRef<FieldDescriptor>
-TypeRefBuilder::getFieldTypeInfo(const TypeRef *TR) {
+static llvm::Optional<StringRef> FindOutermostModuleName(NodePointer Node) {
+  if (!Node)
+    return {};
+  // 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::populateFieldTypeInfoCacheWithReflectionAtIndex(
+    size_t Index) {
+  if (ProcessedReflectionInfoIndexes.contains(Index))
+    return;
+
+  const auto &Info = ReflectionInfos[Index];
+  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;
+    }
+  }
+
+  ProcessedReflectionInfoIndexes.insert(Index);
+}
+
+llvm::Optional<RemoteRef<FieldDescriptor>>
+TypeRefBuilder::findFieldDescriptorAtIndex(size_t Index,
+                                           const std::string &MangledName) {
+  populateFieldTypeInfoCacheWithReflectionAtIndex(Index);
+  auto Found = FieldTypeInfoCache.find(MangledName);
+  if (Found != FieldTypeInfoCache.end()) {
+    return Found->second;
+  }
+  return llvm::None;
+}
+
+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
+  } 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 (size_t i = 0; i < ReflectionInfos.size(); ++i)
+      if (llvm::is_contained(ReflectionInfos[i].PotentialModuleNames,
+                             ModuleName))
+        if (auto FD = findFieldDescriptorAtIndex(i, *MangledName))
+          return *FD;
+
   // 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;
-    Found = FieldTypeInfoCache.find(*MangledName);
-    if (Found != FieldTypeInfoCache.end())
-      return Found->second;
-  }
+  for (size_t i = 0; i < ReflectionInfos.size(); ++i)
+    if (auto FD = findFieldDescriptorAtIndex(i, *MangledName))
+      return *FD;
 
   return nullptr;
 }
diff --git a/stdlib/public/SwiftRemoteMirror/SwiftRemoteMirror.cpp b/stdlib/public/SwiftRemoteMirror/SwiftRemoteMirror.cpp
