chore: expose direct API on Bloom objects (#2845)

Needed to be able to (de)serialize SBF objects.

Signed-off-by: Roman Gershman <roman@dragonflydb.io>

src/core/bloom.cc

@@ -69,6 +69,14 @@ void Bloom::Init(uint64_t entries, double fp_prob, PMR_NS::memory_resource* heap
   bit_log_ = absl::countr_zero(bits);
 }
 
+void Bloom::Init(uint8_t* blob, size_t len, unsigned hash_cnt) {
+  DCHECK_EQ(len * 8, absl::bit_ceil(len * 8));  // must be power of two.
+  CHECK(bf_ == nullptr);
+  hash_cnt_ = hash_cnt;
+  bf_ = blob;
+  bit_log_ = absl::countr_zero(len * 8);
+}
+
 void Bloom::Destroy(PMR_NS::memory_resource* resource) {
   resource->deallocate(CHECK_NOTNULL(bf_), bitlen() / 8);
   bf_ = nullptr;
@@ -141,6 +149,16 @@ SBF::SBF(uint64_t initial_capacity, double fp_prob, double grow_factor, PMR_NS::
   max_capacity_ = filters_.front().Capacity(fp_prob_);
 }
 
+SBF::SBF(double grow_factor, double fp_prob, size_t max_capacity, size_t prev_size,
+         size_t current_size, PMR_NS::memory_resource* mr)
+    : filters_(mr),
+      grow_factor_(grow_factor),
+      fp_prob_(fp_prob),
+      prev_size_(prev_size),
+      current_size_(current_size),
+      max_capacity_(max_capacity) {
+}
+
 SBF::~SBF() {
   PMR_NS::memory_resource* mr = filters_.get_allocator().resource();
   for (auto& f : filters_)
@@ -158,6 +176,13 @@ SBF& SBF::operator=(SBF&& src) {
   return *this;
 }
 
+void SBF::AddFilter(const std::string& blob, unsigned hash_cnt) {
+  PMR_NS::memory_resource* mr = filters_.get_allocator().resource();
+  uint8_t* ptr = (uint8_t*)mr->allocate(blob.size(), 1);
+  memcpy(ptr, blob.data(), blob.size());
+  filters_.emplace_back().Init(ptr, blob.size(), hash_cnt);
+}
+
 bool SBF::Add(std::string_view str) {
   DCHECK_LT(current_size_, max_capacity_);
 

src/core/bloom.h

@@ -29,6 +29,9 @@ class Bloom {
   // heap
   void Init(uint64_t entries, double fp_prob, PMR_NS::memory_resource* resource);
 
+  // Direct initializer. len*8 must be power of 2.
+  void Init(uint8_t* blob, size_t len, unsigned hash_cnt);
+
   // Destroys the object, must be called before destructing the object.
   // resource - resource with which the object was initialized.
   void Destroy(PMR_NS::memory_resource* resource);
@@ -55,6 +58,14 @@ class Bloom {
   // derived from fp_prob.
   size_t Capacity(double fp_prob) const;
 
+  std::string_view data() const {
+    return std::string_view{reinterpret_cast<const char*>(bf_), bitlen() / 8};
+  }
+
+  unsigned hash_cnt() const {
+    return hash_cnt_;
+  }
+
  private:
   bool IsSet(size_t index) const;
   bool Set(size_t index);  // return true if bit was set (i.e was 0 before)
@@ -77,30 +88,63 @@ class SBF {
 
  public:
   SBF(uint64_t initial_capacity, double fp_prob, double grow_factor, PMR_NS::memory_resource* mr);
+
+  // C'tor used for loading persisted filters into SBF.
+  // Should be followed by AddFilter.
+  SBF(double grow_factor, double fp_prob, size_t max_capacity, size_t prev_size,
+      size_t current_size, PMR_NS::memory_resource* mr);
   ~SBF();
 
   SBF& operator=(SBF&& src);
 
+  void AddFilter(const std::string& blob, unsigned hash_cnt);
+
   bool Add(std::string_view str);
   bool Exists(std::string_view str) const;
 
-  size_t GetSize() const {
-    return prev_size_ + current_size_;
+  size_t current_size() const {
+    return current_size_;
   }
 
-  size_t MallocUsed() const;
+  size_t prev_size() const {
+    return prev_size_;
+  }
 
   double grow_factor() const {
     return grow_factor_;
   }
 
+  // expected fp probability for the current filter.
+  double fp_probability() const {
+    return fp_prob_;
+  }
+
+  uint32_t num_filters() const {
+    return filters_.size();
+  }
+
+  std::string_view data(size_t idx) const {
+    return filters_[idx].data();
+  }
+
+  unsigned hashfunc_cnt(size_t idx) const {
+    return filters_[idx].hash_cnt();
+  }
+
+  // max capacity of the current filter.
+  size_t max_capacity() const {
+    return max_capacity_;
+  }
+
+  size_t MallocUsed() const;
+
  private:
   // multiple filters from the smallest to the largest.
   std::vector<Bloom, PMR_NS::polymorphic_allocator<Bloom>> filters_;
   double grow_factor_;
   double fp_prob_;
-  size_t current_size_ = 0;
   size_t prev_size_ = 0;
+  size_t current_size_ = 0;
   size_t max_capacity_;
 };
 

src/core/compact_object.h

@@ -303,6 +303,11 @@ class CompactObj {
   // pre condition - the type here is OBJ_JSON and was set with SetJson
   JsonType* GetJson() const;
 
+  void SetSBF(SBF* sbf) {
+    SetMeta(SBF_TAG);
+    u_.sbf = sbf;
+  }
+
   void SetSBF(uint64_t initial_capacity, double fp_prob, double grow_factor);
   SBF* GetSBF() const;