1use std::fs::{File, OpenOptions};
33use std::mem::size_of;
34use std::path::Path;
35use std::sync::atomic::{AtomicU8, Ordering};
36
37use memmap2::{MmapMut, MmapOptions};
38
39pub const HLL_MAGIC: u64 = 0x4150_484C_4C56_3031;
40pub const MIN_PRECISION: u8 = 4;
42pub const MAX_PRECISION: u8 = 16;
44
45#[repr(C, align(64))]
46pub struct HLLHeader {
47 pub magic: u64,
48 pub precision: u32,
49 pub m: u32,
50 _pad: [u8; 48],
51}
52
53const _: () = {
54 assert!(size_of::<HLLHeader>() == 64);
55};
56
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub enum HLLError {
59 InvalidPrecision,
60 LayoutMismatch,
61 IoError(std::io::ErrorKind),
62}
63
64impl From<std::io::Error> for HLLError {
65 fn from(e: std::io::Error) -> Self { Self::IoError(e.kind()) }
66}
67
68const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
69const FNV_PRIME: u64 = 0x100_0000_01b3;
70
71#[inline]
72fn fnv1a_64(bytes: &[u8]) -> u64 {
73 let mut h = FNV_OFFSET_BASIS;
74 for &b in bytes {
75 h ^= b as u64;
76 h = h.wrapping_mul(FNV_PRIME);
77 }
78 h
79}
80
81#[inline]
87fn fmix64(mut h: u64) -> u64 {
88 h ^= h >> 33;
89 h = h.wrapping_mul(0xff51_afd7_ed55_8ccd);
90 h ^= h >> 33;
91 h = h.wrapping_mul(0xc4ce_b9fe_1a85_ec53);
92 h ^= h >> 33;
93 h
94}
95
96#[inline]
97fn hash_for_hll(item: &[u8]) -> u64 {
98 fmix64(fnv1a_64(item))
99}
100
101pub fn hll_file_size(precision: u8) -> usize {
102 size_of::<HLLHeader>() + (1usize << precision)
103}
104
105pub struct SharedHyperLogLog {
106 _file: File,
107 mmap: MmapMut,
108 precision: u8,
109 m: u32,
110 header_sidecar: subetha_core::HandshakeHeader,
111 ring_sidecar: Box<subetha_core::ObservationRing>,
112}
113
114unsafe impl Send for SharedHyperLogLog {}
115unsafe impl Sync for SharedHyperLogLog {}
116
117impl subetha_sidecar::AdaptiveInstance for SharedHyperLogLog {
118 fn header(&self) -> &subetha_core::HandshakeHeader { &self.header_sidecar }
119 fn ring(&self) -> &subetha_core::ObservationRing { &self.ring_sidecar }
120 fn make_policy(&self) -> Box<dyn subetha_sidecar::Policy> {
121 Box::new(subetha_sidecar::NoMigrationPolicy)
122 }
123}
124
125impl SharedHyperLogLog {
126 pub fn create(
127 path: impl AsRef<Path>, precision: u8,
128 ) -> Result<Self, HLLError> {
129 if !(MIN_PRECISION..=MAX_PRECISION).contains(&precision) {
130 return Err(HLLError::InvalidPrecision);
131 }
132 let m = 1u32 << precision;
133 let total = hll_file_size(precision);
134 let file = OpenOptions::new()
135 .read(true).write(true).create(true).truncate(true)
136 .open(path.as_ref())?;
137 file.set_len(total as u64)?;
138 let mut mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
139 let hdr = mmap.as_mut_ptr() as *mut HLLHeader;
140 unsafe {
141 std::ptr::write_bytes(hdr as *mut u8, 0, size_of::<HLLHeader>());
142 (*hdr).magic = HLL_MAGIC;
143 (*hdr).precision = precision as u32;
144 (*hdr).m = m;
145 }
146 Ok(Self {
148 _file: file, mmap, precision, m,
149 header_sidecar: subetha_core::HandshakeHeader::new(),
150 ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
151 })
152 }
153
154 pub fn open(
155 path: impl AsRef<Path>, expected_precision: u8,
156 ) -> Result<Self, HLLError> {
157 let total = hll_file_size(expected_precision);
158 let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
159 if file.metadata()?.len() < total as u64 {
160 return Err(HLLError::LayoutMismatch);
161 }
162 let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
163 let hdr = unsafe { &*(mmap.as_ptr() as *const HLLHeader) };
164 if hdr.magic != HLL_MAGIC || hdr.precision != expected_precision as u32 {
165 return Err(HLLError::LayoutMismatch);
166 }
167 let m = hdr.m;
168 Ok(Self {
169 _file: file, mmap, precision: expected_precision, m,
170 header_sidecar: subetha_core::HandshakeHeader::new(),
171 ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
172 })
173 }
174
175 #[inline]
176 pub fn precision(&self) -> u8 { self.precision }
177 #[inline]
178 pub fn n_registers(&self) -> u32 { self.m }
179
180 fn register(&self, idx: usize) -> &AtomicU8 {
181 let base = unsafe { self.mmap.as_ptr().add(size_of::<HLLHeader>()) };
182 unsafe { &*(base.add(idx) as *const AtomicU8) }
183 }
184
185 pub fn insert(&self, item: &[u8]) {
187 let h = hash_for_hll(item);
188 let p = self.precision as u32;
189 let reg_idx = (h >> (64 - p)) as usize;
190 let w = (h << p) | (1u64 << (p.saturating_sub(1)));
194 let rank = (w.leading_zeros() as u8) + 1;
195 self.register(reg_idx).fetch_max(rank, Ordering::AcqRel);
196 self.ring_sidecar
197 .push_op(crate::sidecar_ops::sketch::OP_INSERT, 0);
198 }
199
200 pub fn estimate(&self) -> u64 {
203 let m = self.m as f64;
204 let alpha = match self.m {
205 16 => 0.673,
206 32 => 0.697,
207 64 => 0.709,
208 _ => 0.7213 / (1.0 + 1.079 / m),
209 };
210 let mut sum = 0.0f64;
211 let mut zeros = 0u32;
212 for i in 0..self.m as usize {
213 let r = self.register(i).load(Ordering::Acquire);
214 if r == 0 { zeros += 1; }
215 sum += 2f64.powi(-(r as i32));
216 }
217 let raw = alpha * m * m / sum;
218 let v = if raw <= 2.5 * m && zeros > 0 {
221 let z = zeros as f64;
222 (m * (m / z).ln()).round() as u64
223 } else {
224 raw.round() as u64
225 };
226 self.ring_sidecar
227 .push_op(crate::sidecar_ops::sketch::OP_QUERY, 0);
228 v
229 }
230
231 pub fn reset(&self) {
233 for i in 0..self.m as usize {
234 self.register(i).store(0, Ordering::Release);
235 }
236 self.ring_sidecar
237 .push_op(crate::sidecar_ops::sketch::OP_CLEAR, 0);
238 }
239
240 pub fn flush(&self) -> Result<(), HLLError> {
241 self.mmap.flush()?;
242 Ok(())
243 }
244 pub fn flush_async(&self) -> Result<(), HLLError> {
245 self.mmap.flush_async()?;
246 Ok(())
247 }
248}
249
250#[cfg(test)]
251mod tests {
252 use super::*;
253 use std::sync::Arc;
254 use std::thread;
255
256 fn tmp(name: &str) -> std::path::PathBuf {
257 let mut p = std::env::temp_dir();
258 let pid = std::process::id();
259 p.push(format!("subetha-hll-{name}-{pid}.bin"));
260 p
261 }
262
263 #[test]
264 fn create_initial_empty_estimate_is_zero() {
265 let p = tmp("init");
266 let h = SharedHyperLogLog::create(&p, 12).unwrap();
267 assert_eq!(h.precision(), 12);
268 assert_eq!(h.n_registers(), 4096);
269 assert_eq!(h.estimate(), 0);
270 std::fs::remove_file(&p).ok();
271 }
272
273 #[test]
274 fn invalid_precision_rejected() {
275 let p = tmp("invalid");
276 assert_eq!(
277 SharedHyperLogLog::create(&p, 3).err(),
278 Some(HLLError::InvalidPrecision)
279 );
280 assert_eq!(
281 SharedHyperLogLog::create(&p, 17).err(),
282 Some(HLLError::InvalidPrecision)
283 );
284 std::fs::remove_file(&p).ok();
285 }
286
287 #[test]
288 fn single_insert_estimate_is_one() {
289 let p = tmp("single");
290 let h = SharedHyperLogLog::create(&p, 12).unwrap();
291 h.insert(b"hello");
292 let est = h.estimate();
293 assert_eq!(est, 1, "single insert should estimate 1, got {est}");
296 std::fs::remove_file(&p).ok();
297 }
298
299 #[test]
300 fn idempotent_inserts_same_item() {
301 let p = tmp("idempotent");
302 let h = SharedHyperLogLog::create(&p, 12).unwrap();
303 for _ in 0..100 { h.insert(b"same-item"); }
304 let est = h.estimate();
305 assert_eq!(est, 1, "100 inserts of same item should estimate 1, got {est}");
306 std::fs::remove_file(&p).ok();
307 }
308
309 #[test]
310 fn distinct_count_within_error_bound_p12() {
311 let p = tmp("distinct-1k");
314 let h = SharedHyperLogLog::create(&p, 12).unwrap();
315 for i in 0..1000u32 {
316 h.insert(format!("item-{i:05}").as_bytes());
317 }
318 let est = h.estimate();
319 assert!(
320 (900..=1100).contains(&est),
321 "estimate {est} should be within 10% of 1000 (got {est})",
322 );
323 std::fs::remove_file(&p).ok();
324 }
325
326 #[test]
327 fn distinct_count_within_error_bound_p14() {
328 let p = tmp("distinct-10k");
331 let h = SharedHyperLogLog::create(&p, 14).unwrap();
332 for i in 0..10_000u32 {
333 h.insert(format!("k{i:06}").as_bytes());
334 }
335 let est = h.estimate();
336 assert!(
337 (9600..=10400).contains(&est),
338 "estimate {est} should be within 4% of 10000",
339 );
340 std::fs::remove_file(&p).ok();
341 }
342
343 #[test]
344 fn reset_clears_registers() {
345 let p = tmp("reset");
346 let h = SharedHyperLogLog::create(&p, 10).unwrap();
347 for i in 0..100u32 { h.insert(format!("k{i}").as_bytes()); }
348 assert!(h.estimate() > 0);
349 h.reset();
350 assert_eq!(h.estimate(), 0);
351 std::fs::remove_file(&p).ok();
352 }
353
354 #[test]
355 fn cross_handle_visibility() {
356 let p = tmp("cross-handle");
357 let w = SharedHyperLogLog::create(&p, 10).unwrap();
358 let r = SharedHyperLogLog::open(&p, 10).unwrap();
359 for i in 0..50u32 { w.insert(format!("k{i}").as_bytes()); }
360 let est_w = w.estimate();
361 let est_r = r.estimate();
362 assert_eq!(est_w, est_r);
363 assert!((40..=60).contains(&est_r), "estimate {est_r} should be near 50");
365 std::fs::remove_file(&p).ok();
366 }
367
368 #[test]
369 fn config_mismatch_at_open_rejected() {
370 let p = tmp("mismatch");
371 let _w = SharedHyperLogLog::create(&p, 10).unwrap();
372 assert!(matches!(
373 SharedHyperLogLog::open(&p, 12),
374 Err(HLLError::LayoutMismatch)
375 ));
376 std::fs::remove_file(&p).ok();
377 }
378
379 #[test]
380 fn concurrent_inserters_correct_estimate() {
381 let p = tmp("concurrent");
384 let h: Arc<SharedHyperLogLog>
385 = Arc::new(SharedHyperLogLog::create(&p, 12).unwrap());
386 let mut handles = vec![];
387 for t in 0..4u32 {
388 let h = h.clone();
389 handles.push(thread::spawn(move || {
390 for i in 0..1000u32 {
391 h.insert(format!("t{t}-i{i:05}").as_bytes());
392 }
393 }));
394 }
395 for h in handles { h.join().unwrap(); }
396 let est = h.estimate();
397 assert!(
398 (3700..=4300).contains(&est),
399 "concurrent inserts: estimate {est} should be within ~7% of 4000",
400 );
401 std::fs::remove_file(&p).ok();
402 }
403
404 #[test]
405 fn disk_persistence_survives_reopen() {
406 let p = tmp("disk");
407 {
408 let h = SharedHyperLogLog::create(&p, 10).unwrap();
409 for i in 0..100u32 { h.insert(format!("k{i}").as_bytes()); }
410 h.flush().unwrap();
411 }
412 let h2 = SharedHyperLogLog::open(&p, 10).unwrap();
413 let est = h2.estimate();
414 assert!((80..=120).contains(&est),
415 "reopened estimate {est} should be near 100");
416 std::fs::remove_file(&p).ok();
417 }
418}