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

Replace the previous, simple, compressed list of string offset pairs
representation with an OnDiskChainedHashTable. This format can be
mmaped and is faster loading up.

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;
 }