From d4050a2917345232867f5e3dd35d8799269b2d6e Mon Sep 17 00:00:00 2001
From: Paul Chang <paulchang@google.com>
Date: Thu, 7 Nov 2024 10:27:57 -0800
Subject: [PATCH] Expose BlobReader::Keys()

PiperOrigin-RevId: 694166186
---
 compression/blob_store.cc      | 8 ++++++++
 compression/blob_store.h       | 3 +++
 compression/blob_store_test.cc | 5 +++++
 3 files changed, 16 insertions(+)

diff --git a/compression/blob_store.cc b/compression/blob_store.cc
index 659946cd..06bcb56f 100644
--- 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
index 4aba006f..94bbacef 100644
--- 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
index 64647567..dbba55fc 100644
--- 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);
 }