Merge pull request #60497 from augusto2112/swift-meta-cache

Introduce ExternalTypeRefInfoProvider

Author: augusto2112

include/swift/Reflection/ReflectionContext.h

@@ -211,9 +211,10 @@ class ReflectionContext
     uint32_t ThreadPort;
   };
 
-  explicit ReflectionContext(std::shared_ptr<MemoryReader> reader)
-    : super(std::move(reader), *this)
-  {}
+  explicit ReflectionContext(
+      std::shared_ptr<MemoryReader> reader,
+      remote::ExternalTypeRefCache *externalCache = nullptr)
+      : super(std::move(reader), *this, externalCache) {}
 
   ReflectionContext(const ReflectionContext &other) = delete;
   ReflectionContext &operator=(const ReflectionContext &other) = delete;
@@ -227,8 +228,10 @@ class ReflectionContext
     return sizeof(StoredPointer) * 2;
   }
 
+  /// On success returns the ID of the newly registered Reflection Info.
   template <typename T>
-  bool readMachOSections(
+  llvm::Optional<uint32_t>
+  readMachOSections(
       RemoteAddress ImageStart,
       llvm::SmallVector<llvm::StringRef, 1> PotentialModuleNames = {}) {
     auto Buf =
@@ -350,7 +353,7 @@ class ReflectionContext
                            {MPEnumMdSec.first, MPEnumMdSec.second},
                            PotentialModuleNames};
 
-    this->addReflectionInfo(info);
+    auto InfoID = this->addReflectionInfo(info);
 
     // Find the __DATA segment.
     for (unsigned I = 0; I < NumCommands; ++I) {
@@ -374,10 +377,11 @@ class ReflectionContext
 
     savedBuffers.push_back(std::move(Buf));
     savedBuffers.push_back(std::move(Sections));
-    return true;
+    return InfoID;
   }
 
-  bool readPECOFFSections(
+  /// On success returns the ID of the newly registered Reflection Info.
+  llvm::Optional<uint32_t> readPECOFFSections(
       RemoteAddress ImageStart,
       llvm::SmallVector<llvm::StringRef, 1> PotentialModuleNames = {}) {
     auto DOSHdrBuf = this->getReader().readBytes(
@@ -477,11 +481,11 @@ class ReflectionContext
                            {ConformMdSec.first, ConformMdSec.second},
                            {MPEnumMdSec.first, MPEnumMdSec.second},
                            PotentialModuleNames};
-    this->addReflectionInfo(Info);
-    return true;
+    return this->addReflectionInfo(Info);
   }
 
-  bool readPECOFF(RemoteAddress ImageStart,
+  /// On success returns the ID of the newly registered Reflection Info.
+  llvm::Optional<uint32_t> readPECOFF(RemoteAddress ImageStart,
                   llvm::SmallVector<llvm::StringRef, 1> PotentialModuleNames = {}) {
     auto Buf = this->getReader().readBytes(ImageStart,
                                            sizeof(llvm::object::dos_header));
@@ -504,8 +508,9 @@ class ReflectionContext
     return readPECOFFSections(ImageStart, PotentialModuleNames);
   }
 
+  /// On success returns the ID of the newly registered Reflection Info.
   template <typename T>
-  bool readELFSections(
+  llvm::Optional<uint32_t> readELFSections(
       RemoteAddress ImageStart,
       llvm::Optional<llvm::sys::MemoryBlock> FileBuffer,
       llvm::SmallVector<llvm::StringRef, 1> PotentialModuleNames = {}) {
@@ -673,8 +678,7 @@ class ReflectionContext
                            {MPEnumMdSec.first, MPEnumMdSec.second},
                            PotentialModuleNames};
 
-    this->addReflectionInfo(info);
-    return true;
+    return this->addReflectionInfo(info);
   }
 
   /// Parses metadata information from an ELF image. Because the Section
@@ -693,22 +697,22 @@ class ReflectionContext
   ///     instance's memory reader.
   ///
   /// \return
-  ///     /b True if the metadata information was parsed successfully,
-  ///     /b false otherwise.
-  bool
+  ///     \b  The newly added reflection info ID if successful,
+  ///     \b llvm::None otherwise.
+  llvm::Optional<uint32_t>
   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)
-      return false;
+      return llvm::None;
 
     // Read the header.
     auto Hdr = reinterpret_cast<const llvm::ELF::Elf64_Ehdr *>(Buf.get());
 
     if (!Hdr->checkMagic())
-      return false;
+      return llvm::None;
 
     // Check if we have a ELFCLASS32 or ELFCLASS64
     unsigned char FileClass = Hdr->getFileClass();
@@ -719,11 +723,12 @@ class ReflectionContext
       return readELFSections<ELFTraits<llvm::ELF::ELFCLASS32>>(
           ImageStart, FileBuffer, PotentialModuleNames);
     } else {
-      return false;
+      return llvm::None;
     }
   }
 
-  bool
+  /// On success returns the ID of the newly registered Reflection Info.
+  llvm::Optional<uint32_t>
   addImage(RemoteAddress ImageStart,
            llvm::SmallVector<llvm::StringRef, 1> PotentialModuleNames = {}) {
     // Read the first few bytes to look for a magic header.
@@ -761,7 +766,7 @@ class ReflectionContext
     }
 
     // We don't recognize the format.
-    return false;
+    return llvm::None;
   }
 
   /// Adds an image using the FindSection closure to find the swift metadata
@@ -770,9 +775,9 @@ class ReflectionContext
   ///     of freeing the memory buffer in the RemoteRef return value.
   ///     process.
   /// \return
-  ///     \b True if any of the reflection sections were registered,
-  ///     \b false otherwise.
-  bool
+  ///     \b  The newly added reflection info ID if successful,
+  ///     \b llvm::None otherwise.
+  llvm::Optional<uint32_t>
   addImage(llvm::function_ref<
                std::pair<RemoteRef<void>, uint64_t>(ReflectionSectionKind)>
                FindSection,
@@ -798,7 +803,7 @@ class ReflectionContext
 
     // If we didn't find any sections, return.
     if (llvm::all_of(Pairs, [](const auto &Pair) { return !Pair.first; }))
-      return false;
+      return {};
 
     ReflectionInfo Info = {{Pairs[0].first, Pairs[0].second},
                            {Pairs[1].first, Pairs[1].second},
@@ -809,12 +814,12 @@ class ReflectionContext
                            {Pairs[6].first, Pairs[6].second},
                            {Pairs[7].first, Pairs[7].second},
                            PotentialModuleNames};
-    this->addReflectionInfo(Info);
-    return true;
+    return addReflectionInfo(Info);
   }
 
-  void addReflectionInfo(ReflectionInfo I) {
-    getBuilder().addReflectionInfo(I);
+  /// Adds the reflection info and returns it's id.
+  uint32_t addReflectionInfo(ReflectionInfo I) {
+    return getBuilder().addReflectionInfo(I);
   }
 
   bool ownsObject(RemoteAddress ObjectAddress) {

include/swift/Reflection/TypeRefBuilder.h

@@ -18,6 +18,7 @@
 #ifndef SWIFT_REFLECTION_TYPEREFBUILDER_H
 #define SWIFT_REFLECTION_TYPEREFBUILDER_H
 
+#include "swift/Remote/ExternalTypeRefCache.h"
 #include "swift/Remote/MetadataReader.h"
 #include "swift/Reflection/MetadataSourceBuilder.h"
 #include "swift/Reflection/Records.h"
@@ -426,6 +427,8 @@ class TypeRefBuilder {
   TypeConverter TC;
   MetadataSourceBuilder MSB;
 
+  remote::ExternalTypeRefCache *ExternalTypeRefCache = nullptr;
+
 #define TYPEREF(Id, Parent) \
   std::unordered_map<TypeRefID, const Id##TypeRef *, \
                      TypeRefID::Hash, TypeRefID::Equal> Id##TypeRefs;
@@ -915,10 +918,17 @@ class TypeRefBuilder {
   /// Parsing reflection metadata
   ///
 
-  void addReflectionInfo(ReflectionInfo I) {
+  /// Add the ReflectionInfo and return a unique ID for the reflection image
+  /// added. Since we only add reflection infos, the ID can be its index.
+  /// We return a uint32_t since it's extremely unlikely we'll run out of
+  /// indexes.
+  uint32_t addReflectionInfo(ReflectionInfo I) {
     ReflectionInfos.push_back(I);
+    auto InfoID = ReflectionInfos.size() - 1;
+    assert(InfoID <= UINT32_MAX && "ReflectionInfo ID overflow");
+    return InfoID;
   }
-  
+
   const std::vector<ReflectionInfo> &getReflectionInfos() {
     return ReflectionInfos;
   }
@@ -943,6 +953,9 @@ class TypeRefBuilder {
   llvm::Optional<RemoteRef<FieldDescriptor>>
   findFieldDescriptorAtIndex(size_t Index, const std::string &MangledName);
 
+  llvm::Optional<RemoteRef<FieldDescriptor>>
+  getFieldDescriptorFromExternalCache(const std::string &MangledName);
+
 public:
   RemoteRef<char> readTypeRef(uint64_t remoteAddr);
 
@@ -983,8 +996,9 @@ class TypeRefBuilder {
 
 public:
   template<typename Runtime>
-  TypeRefBuilder(remote::MetadataReader<Runtime, TypeRefBuilder> &reader)
-    : TC(*this),
+  TypeRefBuilder(remote::MetadataReader<Runtime, TypeRefBuilder> &reader,
+                 remote::ExternalTypeRefCache *externalCache = nullptr)
+      : TC(*this), ExternalTypeRefCache(externalCache),
       PointerSize(sizeof(typename Runtime::StoredPointer)),
       TypeRefDemangler(
       [this, &reader](RemoteRef<char> string, bool useOpaqueTypeSymbolicReferences) -> Demangle::Node * {

include/swift/Remote/ExternalTypeRefCache.h

@@ -0,0 +1,75 @@
+//===--- ExternalTypeRefCache.h - Abstract access to external caches of
+//typeref ------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+/// @file
+/// This file declares an abstract interface for external caches of
+/// typeref information.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_REMOTE_EXTERNALTYPEREFCACHE_H
+#define SWIFT_REMOTE_EXTERNALTYPEREFCACHE_H
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/Optional.h"
+
+#include <string>
+
+namespace swift {
+namespace reflection {
+
+template <typename T>
+class ReflectionSection;
+class FieldDescriptorIterator;
+using FieldSection = ReflectionSection<FieldDescriptorIterator>;
+}
+
+namespace remote {
+/// A struct with the information required to locate a specific field
+/// descriptor.
+struct FieldDescriptorLocator {
+  /// The reflection info ID the field descriptor belongs to.
+  uint64_t InfoID;
+
+  /// The offset of the field descriptor in the FieldSection buffer.
+  uint64_t Offset;
+};
+
+/// An abstract interface for providing external type layout information.
+struct ExternalTypeRefCache {
+  virtual ~ExternalTypeRefCache() = default;
+
+  /// Cache the field descriptors of a reflection info with a given id with
+  /// their corresponding mangled names. The amount of field descriptors and
+  /// mangled names must be the same. If a field descriptor does not have a
+  /// mangled name a corresponding empty string must be in the mangled_names
+  /// array.
+  virtual void
+  cacheFieldDescriptors(uint64_t InfoID,
+                        const swift::reflection::FieldSection &FieldDescriptors,
+                        llvm::ArrayRef<std::string> MangledNames) = 0;
+
+  /// Retrieve a pair representing the reflection info id and the offset of a
+  /// field descriptor in the field section buffer, if available.
+  virtual llvm::Optional<FieldDescriptorLocator>
+  getFieldDescriptorLocator(const std::string &Name) = 0;
+
+  /// Returns whether the reflection info with the corresponding ID has been
+  /// cached already.
+  virtual bool isReflectionInfoCached(uint64_t InfoID) = 0;
+};
+
+} // namespace remote
+} // namespace swift
+#endif