Expose BlobReader::Keys()

PiperOrigin-RevId: 694166186
google · Nov 7, 2024 · d4050a2 · d4050a2
1 parent 868b016
commit d4050a2
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 0 deletions.
diff --git a/compression/blob_store.cc b/compression/blob_store.cc
@@ -201,6 +201,10 @@ class BlobStore {
     return false;
   }
 
+  hwy::Span<const hwy::uint128_t> Keys() const {
+    return hwy::Span<const hwy::uint128_t>(keys_, num_blobs_);
+  }
+
  private:
   uint32_t magic_;
   uint32_t num_blobs_;      // never 0
@@ -303,6 +307,10 @@ BlobError BlobReader::ReadOne(hwy::uint128_t key, void* data,
   return 0;
 }
 
+hwy::Span<const hwy::uint128_t> BlobReader::Keys() const {
+  return blob_store_->Keys();
+}
+
 BlobError BlobWriter::WriteAll(hwy::ThreadPool& pool, const Path& filename) {
   HWY_ASSERT(keys_.size() == blobs_.size());
 

diff --git a/compression/blob_store.h b/compression/blob_store.h
@@ -84,6 +84,9 @@ class BlobReader {
   // Reads one blob directly.
   BlobError ReadOne(hwy::uint128_t key, void* data, size_t size) const;
 
+  // Returns all available blob keys.
+  hwy::Span<const hwy::uint128_t> Keys() const;
+
  private:
   BlobStorePtr blob_store_;  // holds header, not the entire file
   std::vector<BlobIO> requests_;

diff --git a/compression/blob_store_test.cc b/compression/blob_store_test.cc
@@ -70,6 +70,11 @@ TEST(BlobStoreTest, TestReadWrite) {
     HWY_ASSERT_STRING_EQ("DATA", buffer.data());
   }
 
+  const hwy::Span<const hwy::uint128_t> keys = reader.Keys();
+  HWY_ASSERT_EQ(keys.size(), 2);
+  HWY_ASSERT_EQ(keys[0], keyA);
+  HWY_ASSERT_EQ(keys[1], keyB);
+
   close(fd);
   unlink(path_str);
 }