fix: hll_union kept memory of input sketches' allocators

hll/include/HllUnion-internal.hpp

@@ -54,7 +54,9 @@ void hll_union_alloc<A>::update(const hll_sketch_alloc<A>& sketch) {
 template<typename A>
 void hll_union_alloc<A>::update(hll_sketch_alloc<A>&& sketch) {
   if (sketch.is_empty()) { return; }
-  if (gadget_.is_empty() && sketch.get_target_type() == HLL_8 && sketch.get_lg_config_k() <= lg_max_k_) {
+  // taking over the sketch is only possible if its memory can be released with the union's allocator
+  if (gadget_.is_empty() && sketch.get_target_type() == HLL_8 && sketch.get_lg_config_k() <= lg_max_k_
+      && sketch.sketch_impl->getAllocator() == gadget_.sketch_impl->getAllocator()) {
     if (sketch.get_current_mode() == HLL || sketch.get_lg_config_k() == lg_max_k_) {
       gadget_ = std::move(sketch);
     }
@@ -208,17 +210,20 @@ double hll_union_alloc<A>::get_rel_err(bool upper_bound, bool unioned,
 }
 
 template<typename A>
-HllSketchImpl<A>* hll_union_alloc<A>::copy_or_downsample(const HllSketchImpl<A>* src_impl, uint8_t tgt_lg_k) {
+HllSketchImpl<A>* hll_union_alloc<A>::copy_or_downsample(const HllSketchImpl<A>* src_impl, uint8_t tgt_lg_k,
+    const A& allocator) {
   if (src_impl->getCurMode() != HLL) {
     throw std::logic_error("Attempt to downsample non-HLL sketch");
   }
   const HllArray<A>* src = static_cast<const HllArray<A>*>(src_impl);
-  const int src_lg_k = src->getLgConfigK();
-  if (src_lg_k <= tgt_lg_k) {
+  const uint8_t src_lg_k = src->getLgConfigK();
+  // a copy keeps the source's allocator, so it is only made if that is the union's allocator
+  if (src_lg_k <= tgt_lg_k && src->getAllocator() == allocator) {
     return src->copyAs(HLL_8);
   }
   typedef typename std::allocator_traits<A>::template rebind_alloc<Hll8Array<A>> hll8Alloc;
-  Hll8Array<A>* tgtHllArr = new (hll8Alloc(src->getAllocator()).allocate(1)) Hll8Array<A>(tgt_lg_k, src->getAllocator());
+  const uint8_t lg_k = std::min(src_lg_k, tgt_lg_k);
+  Hll8Array<A>* tgtHllArr = new (hll8Alloc(allocator).allocate(1)) Hll8Array<A>(lg_k, allocator);
   tgtHllArr->mergeHll(*src);
   //both of these are required for isomorphism
   tgtHllArr->putHipAccum(src->getHipAccum());
@@ -239,8 +244,12 @@ template<typename A>
 void hll_union_alloc<A>::union_impl(const hll_sketch_alloc<A>& sketch, uint8_t lg_max_k) {
   const HllSketchImpl<A>* src_impl = sketch.sketch_impl; //default
   HllSketchImpl<A>* dst_impl = gadget_.sketch_impl; //default
+  // the gadget allocates with the union's allocator: the input sketch may use another allocator
+  // instance (e.g. a stateful one), and copies of it would keep that allocator
+  const A allocator = dst_impl->getAllocator();
   if (src_impl->getCurMode() == LIST || src_impl->getCurMode() == SET) {
-    if (dst_impl->isEmpty() && src_impl->getLgConfigK() == dst_impl->getLgConfigK()) {
+    if (dst_impl->isEmpty() && src_impl->getLgConfigK() == dst_impl->getLgConfigK()
+        && src_impl->getAllocator() == allocator) {
       dst_impl = src_impl->copyAs(HLL_8);
       gadget_.sketch_impl->get_deleter()(gadget_.sketch_impl); // gadget to be replaced
     } else {
@@ -254,12 +263,12 @@ void hll_union_alloc<A>::union_impl(const hll_sketch_alloc<A>& sketch, uint8_t l
       // swap so that src is LIST or SET, tgt is HLL
       // use lg_max_k because LIST has effective K of 2^26
       const CouponList<A>* src = static_cast<const CouponList<A>*>(dst_impl);
-      dst_impl = copy_or_downsample(src_impl, lg_max_k);
+      dst_impl = copy_or_downsample(src_impl, lg_max_k, allocator);
       static_cast<Hll8Array<A>*>(dst_impl)->mergeList(*src);
       gadget_.sketch_impl->get_deleter()(gadget_.sketch_impl); // gadget to be replaced
     } else { // gadget is HLL
       if (src_impl->getLgConfigK() < dst_impl->getLgConfigK()) {
-        dst_impl = copy_or_downsample(dst_impl, sketch.get_lg_config_k());
+        dst_impl = copy_or_downsample(dst_impl, sketch.get_lg_config_k(), allocator);
         gadget_.sketch_impl->get_deleter()(gadget_.sketch_impl); // gadget to be replaced
       }
       const HllArray<A>* src = static_cast<const HllArray<A>*>(src_impl);
@@ -268,7 +277,7 @@ void hll_union_alloc<A>::union_impl(const hll_sketch_alloc<A>& sketch, uint8_t l
       static_cast<Hll8Array<A>*>(dst_impl)->putHipAccum(0);
     }
   } else { // src is HLL, gadget is empty
-    dst_impl = copy_or_downsample(src_impl, lg_max_k);
+    dst_impl = copy_or_downsample(src_impl, lg_max_k, allocator);
     gadget_.sketch_impl->get_deleter()(gadget_.sketch_impl); // gadget to be replaced
   }
   gadget_.sketch_impl = dst_impl; // gadget replaced

hll/include/hll.hpp

@@ -640,7 +640,7 @@ class hll_union_alloc {
     */
     inline void union_impl(const hll_sketch_alloc<A>& sketch, uint8_t lg_max_k);
 
-    static HllSketchImpl<A>* copy_or_downsample(const HllSketchImpl<A>* src_impl, uint8_t tgt_lg_k);
+    static HllSketchImpl<A>* copy_or_downsample(const HllSketchImpl<A>* src_impl, uint8_t tgt_lg_k, const A& allocator);
 
     void coupon_update(uint32_t coupon);
 

hll/test/HllUnionTest.cpp

@@ -384,4 +384,77 @@ TEST_CASE("hll union: serialization is grouping independent", "[hll_union]") {
   REQUIRE(serialize_nested_union(a, b, c) == serialize_flat_union(a, b, c));
 }
 
+// A stateful allocator: instances with different ids are different allocators, and the bytes each one
+// has outstanding are counted.
+static long long pool_bytes[2] = {0, 0};
+
+template<typename T>
+struct pool_allocator {
+  using value_type = T;
+  int id;
+  explicit pool_allocator(int id): id(id) {}
+  template<typename U>
+  pool_allocator(const pool_allocator<U>& other): id(other.id) {}
+  T* allocate(std::size_t n) {
+    pool_bytes[id] += n * sizeof(T);
+    return static_cast<T*>(::operator new(n * sizeof(T)));
+  }
+  void deallocate(T* p, std::size_t n) {
+    pool_bytes[id] -= n * sizeof(T);
+    ::operator delete(p);
+  }
+};
+template<typename T, typename U>
+bool operator==(const pool_allocator<T>& a, const pool_allocator<U>& b) { return a.id == b.id; }
+template<typename T, typename U>
+bool operator!=(const pool_allocator<T>& a, const pool_allocator<U>& b) { return a.id != b.id; }
+
+using pool_hll_sketch = hll_sketch_alloc<pool_allocator<uint8_t>>;
+using pool_hll_union = hll_union_alloc<pool_allocator<uint8_t>>;
+
+static pool_hll_sketch make_pool_sketch(int id, uint8_t lg_k, target_hll_type type, uint64_t n) {
+  pool_hll_sketch sketch(lg_k, type, false, pool_allocator<uint8_t>(id));
+  for (uint64_t i = 0; i < n; ++i) sketch.update(i);
+  return sketch;
+}
+
+TEST_CASE("hll union: keeps no memory of an input sketch's allocator", "[hll_union]") {
+  const uint8_t lg_max_k = 10;
+  // LIST, SET and HLL inputs, with a smaller, the same and a larger lg_k than the union
+  const uint64_t sizes[] = {5, 100, 10000};
+  const uint8_t lg_ks[] = {8, 10, 12};
+  for (uint64_t n: sizes) {
+    for (uint8_t lg_k: lg_ks) {
+      for (target_hll_type type: {HLL_4, HLL_6, HLL_8}) {
+        for (bool gadget_empty: {true, false}) {
+          for (bool move: {false, true}) {
+            hll_union u(lg_max_k);
+            pool_hll_union pool_union(lg_max_k, pool_allocator<uint8_t>(0));
+            if (!gadget_empty) {
+              for (uint64_t i = 0; i < 3 * n; ++i) {
+                u.update(i + n);
+                pool_union.update(i + n);
+              }
+            }
+            {
+              hll_sketch control(lg_k, type);
+              for (uint64_t i = 0; i < n; ++i) control.update(i);
+              u.update(control);
+              pool_hll_sketch input = make_pool_sketch(1, lg_k, type, n);
+              if (move) {
+                pool_union.update(std::move(input));
+              } else {
+                pool_union.update(input);
+              }
+            }
+            REQUIRE(pool_bytes[1] == 0);
+            REQUIRE(pool_union.get_estimate() == u.get_estimate());
+          }
+        }
+      }
+    }
+  }
+  REQUIRE(pool_bytes[0] == 0);
+}
+
 } /* namespace datasketches */