Merge pull request #5154 from augusto2112/swift-persistent-meta-cache

[lldb] Use on disk hash table as the binary format of SwiftMetadataCache

Author: augusto2112

lldb/source/Core/CoreProperties.td

@@ -32,7 +32,7 @@ let Definition = "modulelist" in {
     Desc<"The module loading mode to use when loading modules for Swift.">;
   def EnableSwiftMetadataCache: Property<"enable-swift-metadata-cache", "Boolean">,
     Global,
-    DefaultFalse,
+    DefaultTrue,
     Desc<"Enable caching for Swift reflection metadata in LLDB.">;
   def SwiftMetadataCachePath: Property<"swift-metadata-cache-path", "FileSpec">,
     Global,

lldb/source/Plugins/LanguageRuntime/Swift/SwiftMetadataCache.cpp

@@ -2,13 +2,9 @@
 
 #include "lldb/Utility/DataEncoder.h"
 #include "lldb/Utility/LLDBLog.h"
-#include "lldb/Utility/Log.h"
 #include "lldb/Version/Version.h"
-#include "llvm/CodeGen/AccelTable.h"
-#include "llvm/CodeGen/AsmPrinter.h"
 #include "llvm/Support/BLAKE3.h"
 #include "llvm/Support/CachePruning.h"
-#include "llvm/Support/Compression.h"
 
 using namespace lldb;
 using namespace lldb_private;
@@ -32,9 +28,7 @@ SwiftMetadataCache::SwiftMetadataCache() {
   }
 }
 
-bool SwiftMetadataCache::is_enabled() {
-  return llvm::zlib::isAvailable() && m_data_file_cache.hasValue();
-}
+bool SwiftMetadataCache::is_enabled() { return m_data_file_cache.hasValue(); }
 
 void SwiftMetadataCache::registerModuleWithReflectionInfoID(ModuleSP module,
                                                             uint64_t info_id) {
@@ -44,7 +38,8 @@ void SwiftMetadataCache::registerModuleWithReflectionInfoID(ModuleSP module,
   if (!is_enabled())
     return;
 
-  m_info_to_module[info_id] = {module, false};
+  /// Insert the module cache info as not processed.
+  m_reflection_info_to_module.insert({info_id, module});
 
   // Attempt to load the cached file.
   auto module_name = getTyperefCacheFileNameForModule(module);
@@ -57,17 +52,32 @@ void SwiftMetadataCache::registerModuleWithReflectionInfoID(ModuleSP module,
     return;
   }
 
+  // Move it to the instance variable so references to this data don't go
+  // out of scope.
+  m_hash_table_buffers.emplace_back(std::move(mem_buffer_up));
+  auto &mem_buffer = m_hash_table_buffers.back();
+
   // Extractor used to extract the header information (see the .h file for
   // details on the format).
-  DataExtractor header_extractor(mem_buffer_up->getBufferStart(),
-                                 mem_buffer_up->getBufferSize(),
+  DataExtractor header_extractor(mem_buffer->getBufferStart(),
+                                 mem_buffer->getBufferSize(),
                                  module->GetObjectFile()->GetByteOrder(),
                                  module->GetObjectFile()->GetAddressByteSize());
 
   lldb::offset_t read_offset = 0;
+  uint16_t cached_UUID_size = 0;
+  if (!header_extractor.GetU16(&read_offset, &cached_UUID_size, 1)) {
+    LLDB_LOG(log,
+             "[SwiftMetadataCache] Failed to read cached UUID size for module {0}.",
+             module->GetFileSpec().GetFilename());
+    m_data_file_cache->RemoveCacheFile(module_name);
+    return;
+  }
 
-  std::string UUID = module->GetUUID().GetAsString();
-  std::string cached_UUID = header_extractor.GetCStr(&read_offset);
+  const auto *cached_UUID_data = reinterpret_cast<const uint8_t *>(
+      header_extractor.GetData(&read_offset, cached_UUID_size));
+
+  llvm::ArrayRef<uint8_t> cached_UUID(cached_UUID_data, cached_UUID_size);
   // If no uuid in the file something is wrong with the cache.
   if (cached_UUID.empty()) {
     LLDB_LOG(log,
@@ -77,6 +87,7 @@ void SwiftMetadataCache::registerModuleWithReflectionInfoID(ModuleSP module,
     return;
   }
 
+  auto UUID = module->GetUUID().GetBytes();
   // If the UUIDs don't match this is most likely a stale cache.
   if (cached_UUID != UUID) {
     LLDB_LOGV(log, "[SwiftMetadataCache] Module UUID mismatch for {0}.",
@@ -85,73 +96,31 @@ void SwiftMetadataCache::registerModuleWithReflectionInfoID(ModuleSP module,
     return;
   }
 
-  uint64_t expanded_size = 0;
-  if (!header_extractor.GetU64(&read_offset, &expanded_size, 1)) {
+  // The on disk hash table must have a 4-byte alignment, skip
+  // the padding when reading.
+  read_offset = llvm::alignTo(read_offset, 4);
+
+  // The offset of the hash table control structure, which follows the payload.
+  uint32_t table_control_offset = 0;
+  if (!header_extractor.GetU32(&read_offset, &table_control_offset, 1)) {
     LLDB_LOGV(log,
-              "[SwiftMetadataCache] Failed to read decompressed cache size for "
+              "[SwiftMetadataCache] Failed to read table offset for "
               "module {0}.",
               module->GetFileSpec().GetFilename());
     m_data_file_cache->RemoveCacheFile(module_name);
     return;
   }
 
-  const auto *start = (const char *)header_extractor.GetData(&read_offset, 0);
-  // Create a reference to the compressed data.
-  llvm::StringRef string_buffer(start, (uint64_t)mem_buffer_up->getBufferEnd() -
-                                           (uint64_t)start);
-
-  llvm::SmallString<0> decompressed;
-  auto error =
-      llvm::zlib::uncompress(string_buffer, decompressed, expanded_size);
-  if (error) {
-    auto error_string = llvm::toString(std::move(error));
-    LLDB_LOG(log,
-             "[SwiftMetadataCache] Cache decompression failed with error: {0}. "
-             "Deleting cached file.",
-             error_string);
-    m_data_file_cache->RemoveCacheFile(module_name);
-    return;
-  }
+  const auto *table_contents = reinterpret_cast<const uint8_t *>(
+      header_extractor.GetData(&read_offset, 0));
 
-  // Extractor to extract the body of the cached file (see SwiftMetadataCache.h
-  // for more details of the format).
-  DataExtractor body_extractor(decompressed.data(), decompressed.size(),
-                               module->GetObjectFile()->GetByteOrder(),
-                               module->GetObjectFile()->GetAddressByteSize());
-  read_offset = 0;
-  auto num_entries = body_extractor.GetU64(&read_offset);
-
-  // Map to extract the encoded data to. Since extraction can fail we don't want
-  // to insert values into the final map in case we have to abort midway.
-  llvm::StringMap<swift::remote::FieldDescriptorLocator> temp_map;
-  for (size_t i = 0; i < num_entries; i++) {
-    const auto *mangled_name = body_extractor.GetCStr(&read_offset);
-    if (!mangled_name) {
-      LLDB_LOG(log,
-               "[SwiftMetadataCache] Failed to read mangled name {0} at offset "
-               "{1} for module {2}.",
-               i, read_offset, module->GetFileSpec().GetFilename());
-      m_data_file_cache->RemoveCacheFile(module_name);
-      return;
-    }
-    uint64_t offset = 0;
-    if (!body_extractor.GetU64(&read_offset, &offset, 1)) {
-      LLDB_LOG(log,
-               "[SwiftMetadataCache] Failed to read mangled name {0} at offset "
-               "{1} for module {2}.",
-               i, read_offset, module->GetFileSpec().GetFilename());
-      m_data_file_cache->RemoveCacheFile(module_name);
-      return;
-    }
-    temp_map[mangled_name] = {info_id, offset};
-  }
+  const auto *table_control = table_contents + table_control_offset;
 
-  // Move the values to the actual map now that we know that it's safe.
-  for (auto &p : temp_map)
-    m_mangled_name_to_offset.try_emplace(p.getKey(), p.second);
+  // Store the hash table.
+  m_reflection_info_to_module.find(info_id)->second.cache_hash_table.reset(
+      llvm::OnDiskChainedHashTable<TypeRefInfo>::Create(
+          table_control, table_contents, m_info));
 
-  // Mark this reflection info as processed.
-  m_info_to_module[info_id] = {module, true};
   LLDB_LOGV(log, "[SwiftMetadataCache] Loaded cache for module {0}.",
             module->GetFileSpec().GetFilename());
 }
@@ -167,33 +136,38 @@ static bool areMangledNamesAndFieldSectionSameSize(
   return field_descriptors_size == mangled_names.size();
 }
 
-bool SwiftMetadataCache::writeMangledNamesAndOffsetsToEncoder(
+llvm::Optional<std::pair<uint32_t, llvm::SmallString<32>>>
+SwiftMetadataCache::generateHashTableBlob(
     uint64_t info_id, const swift::reflection::FieldSection &field_descriptors,
-    const std::vector<std::string> &mangled_names, DataEncoder &encoder) {
+    const std::vector<std::string> &mangled_names) {
   Log *log = GetLog(LLDBLog::Types);
-  auto num_entries = mangled_names.size();
-  encoder.AppendU64(num_entries);
+  llvm::SmallString<32> hash_table_blob;
+  llvm::raw_svector_ostream blobStream(hash_table_blob);
 
   // If the amount of mangled names and field descriptors don't match something
   // unexpected happened.
   if (!areMangledNamesAndFieldSectionSameSize(field_descriptors,
                                               mangled_names)) {
     LLDB_LOG(log, "[SwiftMetadataCache] Mismatch between number of mangled "
                   "names and field descriptors passed in.");
-    return false;
+    return {};
   }
 
+  llvm::OnDiskChainedHashTableGenerator<TypeRefInfo> table_generator;
   for (auto pair : llvm::zip(field_descriptors, mangled_names)) {
     auto field_descriptor = std::get<0>(pair);
     auto &mangled_name = std::get<1>(pair);
     if (mangled_name.empty())
       continue;
     auto offset = field_descriptor.getAddressData() -
                   field_descriptors.startAddress().getAddressData();
-    encoder.AppendCString(mangled_name.data());
-    encoder.AppendU64(offset);
+    table_generator.insert(mangled_name, offset, m_info);
   }
-  return true;
+
+  // Make sure that no bucket is at offset 0.
+  llvm::support::endian::write<uint32_t>(blobStream, 0, llvm::support::little);
+  uint32_t table_control_offset = table_generator.Emit(blobStream, m_info);
+  return {{std::move(table_control_offset), std::move(hash_table_blob)}};
 }
 
 void SwiftMetadataCache::cacheFieldDescriptors(
@@ -205,33 +179,40 @@ void SwiftMetadataCache::cacheFieldDescriptors(
   if (!is_enabled())
     return;
 
-  auto it = m_info_to_module.find(info_id);
-  if (it == m_info_to_module.end()) {
+  auto it = m_reflection_info_to_module.find(info_id);
+  if (it == m_reflection_info_to_module.end()) {
     LLDB_LOGV(log, "[SwiftMetadataCache] No module found with module id {0}.",
               info_id);
     return;
   }
 
-  auto module = std::get<ModuleSP>(it->second);
-  // Write the data to the body encoder with the format expected by the current
-  // cache version.
-  DataEncoder body_encoder;
-  if (!writeMangledNamesAndOffsetsToEncoder(info_id, field_descriptors,
-                                            mangled_names, body_encoder))
-    return;
+  auto &module = it->second.module;
 
-  uint64_t typeref_buffer_size = body_encoder.GetData().size();
-  llvm::StringRef typeref_buffer((const char *)body_encoder.GetData().data(),
-                                 typeref_buffer_size);
+  auto maybe_pair =
+      generateHashTableBlob(info_id, field_descriptors, mangled_names);
+  if (!maybe_pair)
+    return;
 
-  llvm::SmallString<0> compressed_buffer;
-  llvm::zlib::compress(typeref_buffer, compressed_buffer);
+  auto &table_offset = maybe_pair->first;
+  auto &hash_table_blob = maybe_pair->second;
 
   // Write the header followed by the body.
   DataEncoder encoder;
-  encoder.AppendCString(module->GetUUID().GetAsString());
-  encoder.AppendU64(typeref_buffer_size);
-  encoder.AppendData(compressed_buffer);
+  auto uuid = module->GetUUID().GetBytes();
+  // Append the uuid size followed by the uuid itself.
+  encoder.AppendU16(uuid.size());
+  encoder.AppendData(uuid);
+
+
+  auto size_so_far = encoder.GetByteSize();
+  // The on disk hash table must have a 4-byte alignment, so
+  // write 0 bytes until we get to the required alignemnt.
+  auto padding = llvm::alignTo(size_so_far, 4) - size_so_far;
+  while (padding-- > 0)
+    encoder.AppendU8(0);
+
+  encoder.AppendU32(table_offset);
+  encoder.AppendData(hash_table_blob);
 
   auto filename = getTyperefCacheFileNameForModule(module);
 
@@ -244,24 +225,33 @@ llvm::Optional<swift::remote::FieldDescriptorLocator>
 SwiftMetadataCache::getFieldDescriptorLocator(const std::string &Name) {
   std::lock_guard<std::recursive_mutex> guard(m_mutex);
   Log *log = GetLog(LLDBLog::Types);
-  auto it = m_mangled_name_to_offset.find(Name);
-  if (it != m_mangled_name_to_offset.end()) {
-    LLDB_LOGV(
-        log,
-        "[SwiftMetadataCache] Returning field descriptor for mangled name {0}",
-        Name);
-    return it->second;
+  // Compute hash outside of loop as an optimization.
+  auto hash = m_info.ComputeHash(Name);
+  for (auto &pair : m_reflection_info_to_module) {
+    auto &cache_hash_table = pair.second.cache_hash_table;
+    // No cache for this reflection module.
+    if (!cache_hash_table)
+      continue;
+    auto it = cache_hash_table->find_hashed(Name, hash, &m_info);
+    if (it != cache_hash_table->end()) {
+      LLDB_LOGV(log,
+                "[SwiftMetadataCache] Returning field descriptor for mangled "
+                "name {0}",
+                Name);
+      auto info_id = pair.first;
+      return {{info_id, *it}};
+    }
   }
   return {};
 }
 
 bool SwiftMetadataCache::isReflectionInfoCached(uint64_t info_id) {
   std::lock_guard<std::recursive_mutex> guard(m_mutex);
-  auto it = m_info_to_module.find(info_id);
+  auto it = m_reflection_info_to_module.find(info_id);
   // First check if we've registered the reflection info with that id.
-  if (it != m_info_to_module.end())
-    // Then check whether we've already parsed it or not.
-    return std::get<bool>(it->second);
+  if (it != m_reflection_info_to_module.end())
+    // Then check whether we have a cache for it or not.
+    return it->second.cache_hash_table.get() != nullptr;
   return false;
 }