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(
132 path: impl AsRef<Path>, precision: u8,
133 ) -> Result<Self, HLLError> {
134 if !(MIN_PRECISION..=MAX_PRECISION).contains(&precision) {
135 return Err(HLLError::InvalidPrecision);
136 }
137 let (file, mmap) = crate::mmf_attach::create_or_attach(
138 path.as_ref(),
139 hll_file_size(precision),
140 |ptr| unsafe { Self::init_region(ptr, precision) },
141 |ptr| unsafe { (*(ptr as *const HLLHeader)).magic == HLL_MAGIC },
142 )?;
143 Self::from_region(file, mmap, precision)
144 }
145
146 unsafe fn init_region(ptr: *mut u8, precision: u8) {
154 let hdr = ptr as *mut HLLHeader;
155 unsafe {
156 (*hdr).precision = precision as u32;
157 (*hdr).m = 1u32 << precision;
158 std::ptr::write_volatile(&raw mut (*hdr).magic, HLL_MAGIC);
159 }
160 }
161
162 fn from_region(file: File, mmap: MmapMut, precision: u8) -> Result<Self, HLLError> {
165 let hdr = unsafe { &*(mmap.as_ptr() as *const HLLHeader) };
166 if hdr.magic != HLL_MAGIC || hdr.precision != precision as u32 {
167 return Err(HLLError::LayoutMismatch);
168 }
169 let m = hdr.m;
170 Ok(Self {
171 _file: file, mmap, precision, m,
172 header_sidecar: subetha_core::HandshakeHeader::new(),
173 ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
174 })
175 }
176
177 pub fn open(
178 path: impl AsRef<Path>, expected_precision: u8,
179 ) -> Result<Self, HLLError> {
180 let total = hll_file_size(expected_precision);
181 let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
182 if file.metadata()?.len() < total as u64 {
183 return Err(HLLError::LayoutMismatch);
184 }
185 let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
186 Self::from_region(file, mmap, expected_precision)
187 }
188
189 #[inline]
190 pub fn precision(&self) -> u8 { self.precision }
191 #[inline]
192 pub fn n_registers(&self) -> u32 { self.m }
193
194 fn register(&self, idx: usize) -> &AtomicU8 {
195 let base = unsafe { self.mmap.as_ptr().add(size_of::<HLLHeader>()) };
196 unsafe { &*(base.add(idx) as *const AtomicU8) }
197 }
198
199 pub fn insert(&self, item: &[u8]) {
201 let h = hash_for_hll(item);
202 let p = self.precision as u32;
203 let reg_idx = (h >> (64 - p)) as usize;
204 let w = (h << p) | (1u64 << (p.saturating_sub(1)));
208 let rank = (w.leading_zeros() as u8) + 1;
209 self.register(reg_idx).fetch_max(rank, Ordering::AcqRel);
210 self.ring_sidecar
211 .push_op(crate::sidecar_ops::sketch::OP_INSERT, 0);
212 }
213
214 pub fn estimate(&self) -> u64 {
217 let m = self.m as f64;
218 let alpha = match self.m {
219 16 => 0.673,
220 32 => 0.697,
221 64 => 0.709,
222 _ => 0.7213 / (1.0 + 1.079 / m),
223 };
224 let mut sum = 0.0f64;
225 let mut zeros = 0u32;
226 for i in 0..self.m as usize {
227 let r = self.register(i).load(Ordering::Acquire);
228 if r == 0 { zeros += 1; }
229 sum += 2f64.powi(-(r as i32));
230 }
231 let raw = alpha * m * m / sum;
232 let v = if raw <= 2.5 * m && zeros > 0 {
235 let z = zeros as f64;
236 (m * (m / z).ln()).round() as u64
237 } else {
238 raw.round() as u64
239 };
240 self.ring_sidecar
241 .push_op(crate::sidecar_ops::sketch::OP_QUERY, 0);
242 v
243 }
244
245 pub fn reset(&self) {
247 for i in 0..self.m as usize {
248 self.register(i).store(0, Ordering::Release);
249 }
250 self.ring_sidecar
251 .push_op(crate::sidecar_ops::sketch::OP_CLEAR, 0);
252 }
253
254 pub fn flush(&self) -> Result<(), HLLError> {
255 self.mmap.flush()?;
256 Ok(())
257 }
258 pub fn flush_async(&self) -> Result<(), HLLError> {
259 self.mmap.flush_async()?;
260 Ok(())
261 }
262}
263
264#[cfg(test)]
265mod tests {
266 use super::*;
267 use std::sync::Arc;
268 use std::thread;
269
270 fn tmp(name: &str) -> std::path::PathBuf {
271 let mut p = std::env::temp_dir();
272 let pid = std::process::id();
273 p.push(format!("subetha-hll-{name}-{pid}.bin"));
274 p
275 }
276
277 #[test]
280 fn second_create_attaches_and_keeps_registers() {
281 let p = tmp("attach");
282 std::fs::remove_file(&p).ok();
283 let h = SharedHyperLogLog::create(&p, 12).unwrap();
284 h.insert(b"one");
285
286 let h2 = SharedHyperLogLog::create(&p, 12).unwrap();
287 assert_eq!(h2.estimate(), 1, "attach zeroed live registers");
288 assert!(matches!(
289 SharedHyperLogLog::create(&p, 10),
290 Err(HLLError::LayoutMismatch),
291 ));
292
293 h2.reset();
294 assert_eq!(h.estimate(), 0, "reset did not zero for every handle");
295 drop(h);
296 drop(h2);
297 std::fs::remove_file(&p).ok();
298 }
299
300 #[test]
301 fn create_initial_empty_estimate_is_zero() {
302 let p = tmp("init");
303 let h = SharedHyperLogLog::create(&p, 12).unwrap();
304 assert_eq!(h.precision(), 12);
305 assert_eq!(h.n_registers(), 4096);
306 assert_eq!(h.estimate(), 0);
307 std::fs::remove_file(&p).ok();
308 }
309
310 #[test]
311 fn invalid_precision_rejected() {
312 let p = tmp("invalid");
313 assert_eq!(
314 SharedHyperLogLog::create(&p, 3).err(),
315 Some(HLLError::InvalidPrecision)
316 );
317 assert_eq!(
318 SharedHyperLogLog::create(&p, 17).err(),
319 Some(HLLError::InvalidPrecision)
320 );
321 std::fs::remove_file(&p).ok();
322 }
323
324 #[test]
325 fn single_insert_estimate_is_one() {
326 let p = tmp("single");
327 let h = SharedHyperLogLog::create(&p, 12).unwrap();
328 h.insert(b"hello");
329 let est = h.estimate();
330 assert_eq!(est, 1, "single insert should estimate 1, got {est}");
333 std::fs::remove_file(&p).ok();
334 }
335
336 #[test]
337 fn idempotent_inserts_same_item() {
338 let p = tmp("idempotent");
339 let h = SharedHyperLogLog::create(&p, 12).unwrap();
340 for _ in 0..100 { h.insert(b"same-item"); }
341 let est = h.estimate();
342 assert_eq!(est, 1, "100 inserts of same item should estimate 1, got {est}");
343 std::fs::remove_file(&p).ok();
344 }
345
346 #[test]
347 fn distinct_count_within_error_bound_p12() {
348 let p = tmp("distinct-1k");
351 let h = SharedHyperLogLog::create(&p, 12).unwrap();
352 for i in 0..1000u32 {
353 h.insert(format!("item-{i:05}").as_bytes());
354 }
355 let est = h.estimate();
356 assert!(
357 (900..=1100).contains(&est),
358 "estimate {est} should be within 10% of 1000 (got {est})",
359 );
360 std::fs::remove_file(&p).ok();
361 }
362
363 #[test]
364 fn distinct_count_within_error_bound_p14() {
365 let p = tmp("distinct-10k");
368 let h = SharedHyperLogLog::create(&p, 14).unwrap();
369 for i in 0..10_000u32 {
370 h.insert(format!("k{i:06}").as_bytes());
371 }
372 let est = h.estimate();
373 assert!(
374 (9600..=10400).contains(&est),
375 "estimate {est} should be within 4% of 10000",
376 );
377 std::fs::remove_file(&p).ok();
378 }
379
380 #[test]
381 fn reset_clears_registers() {
382 let p = tmp("reset");
383 let h = SharedHyperLogLog::create(&p, 10).unwrap();
384 for i in 0..100u32 { h.insert(format!("k{i}").as_bytes()); }
385 assert!(h.estimate() > 0);
386 h.reset();
387 assert_eq!(h.estimate(), 0);
388 std::fs::remove_file(&p).ok();
389 }
390
391 #[test]
392 fn cross_handle_visibility() {
393 let p = tmp("cross-handle");
394 let w = SharedHyperLogLog::create(&p, 10).unwrap();
395 let r = SharedHyperLogLog::open(&p, 10).unwrap();
396 for i in 0..50u32 { w.insert(format!("k{i}").as_bytes()); }
397 let est_w = w.estimate();
398 let est_r = r.estimate();
399 assert_eq!(est_w, est_r);
400 assert!((40..=60).contains(&est_r), "estimate {est_r} should be near 50");
402 std::fs::remove_file(&p).ok();
403 }
404
405 #[test]
406 fn config_mismatch_at_open_rejected() {
407 let p = tmp("mismatch");
408 let _w = SharedHyperLogLog::create(&p, 10).unwrap();
409 assert!(matches!(
410 SharedHyperLogLog::open(&p, 12),
411 Err(HLLError::LayoutMismatch)
412 ));
413 std::fs::remove_file(&p).ok();
414 }
415
416 #[test]
417 fn concurrent_inserters_correct_estimate() {
418 let p = tmp("concurrent");
421 let h: Arc<SharedHyperLogLog>
422 = Arc::new(SharedHyperLogLog::create(&p, 12).unwrap());
423 let mut handles = vec![];
424 for t in 0..4u32 {
425 let h = h.clone();
426 handles.push(thread::spawn(move || {
427 for i in 0..1000u32 {
428 h.insert(format!("t{t}-i{i:05}").as_bytes());
429 }
430 }));
431 }
432 for h in handles { h.join().unwrap(); }
433 let est = h.estimate();
434 assert!(
435 (3700..=4300).contains(&est),
436 "concurrent inserts: estimate {est} should be within ~7% of 4000",
437 );
438 std::fs::remove_file(&p).ok();
439 }
440
441 #[test]
442 fn disk_persistence_survives_reopen() {
443 let p = tmp("disk");
444 {
445 let h = SharedHyperLogLog::create(&p, 10).unwrap();
446 for i in 0..100u32 { h.insert(format!("k{i}").as_bytes()); }
447 h.flush().unwrap();
448 }
449 let h2 = SharedHyperLogLog::open(&p, 10).unwrap();
450 let est = h2.estimate();
451 assert!((80..=120).contains(&est),
452 "reopened estimate {est} should be near 100");
453 std::fs::remove_file(&p).ok();
454 }
455}