1use crate::db::Db;
31use crate::hll::{self, Encoding};
32use crate::keyspace::Keyspace;
33use crate::strings::check_len;
34use crate::value::{self, Str};
35use yo_common::{Code, Error, Result};
36use yo_index::RawMap;
37
38const NO_SUCH_KEY: &str = "The specified key does not exist";
40const NOT_SPARSE: &str = "HLL encoding is not sparse";
42
43impl Keyspace {
44 pub fn pfadd<'e, I>(&mut self, key: &[u8], eles: I) -> Result<bool>
51 where
52 I: Iterator<Item = &'e [u8]>,
53 {
54 self.reap(key);
55 self.string_only(key)?;
56 self.thaw(key)?;
57 check_len(key, hll::DENSE)?;
58
59 let mut buf = std::mem::take(&mut self.scratch);
60 buf.clear();
61 let (deadline, fresh) = match self.map.get(key) {
62 Some(rec) => {
63 value::read(rec).write_to(&mut buf);
64 (value::expire_at(rec), false)
65 }
66 None => (None, true),
67 };
68 if fresh {
69 hll::empty(&mut buf);
70 }
71
72 let outcome = 'work: {
73 if let Err(e) = hll::check(&buf) {
74 break 'work Err(e);
75 }
76 let mut changed = fresh;
77 for ele in eles {
78 let (index, count) = hll::place(ele);
79 match hll::set(&mut buf, index, count) {
80 Some(hit) => changed |= hit,
81 None => break 'work Err(hll::corrupt()),
82 }
83 }
84 Ok(changed)
85 };
86
87 if matches!(outcome, Ok(true)) {
93 hll::invalidate(&mut buf);
94 self.store_raw(key, &buf, deadline);
95 }
96 self.scratch = buf;
97 outcome
98 }
99
100 pub fn pfcount<'k, I>(&mut self, keys: I) -> Result<u64>
110 where
111 I: Iterator<Item = &'k [u8]> + Clone,
112 {
113 for key in keys.clone() {
114 self.hll_ready(key)?;
115 }
116 let mut one = keys.clone();
117 if let (Some(key), None) = (one.next(), one.next()) {
118 return self.count_one(key);
119 }
120
121 let mut max = [0u8; hll::REGISTERS];
126 for key in keys {
127 self.merge_sketch(key, &mut max)?;
128 }
129 Ok(estimate(&max))
130 }
131
132 pub fn pfmerge<'k, I>(&mut self, dest: &'k [u8], srcs: I) -> Result<()>
142 where
143 I: Iterator<Item = &'k [u8]> + Clone,
144 {
145 self.hll_ready(dest)?;
146 check_len(dest, hll::DENSE)?;
147 for src in srcs.clone() {
148 self.hll_ready(src)?;
149 }
150
151 let mut max = [0u8; hll::REGISTERS];
155 let mut dense = false;
156 for key in std::iter::once(dest).chain(srcs) {
157 dense |= self.merge_sketch(key, &mut max)?;
158 }
159 self.pfmerge_into(dest, &max, dense)
160 }
161
162 pub(crate) fn hll_ready(&mut self, key: &[u8]) -> Result<()> {
171 self.reap(key);
172 self.string_only(key)?;
173 self.thaw(key)?;
174 Ok(())
175 }
176
177 pub(crate) fn merge_sketch(&self, key: &[u8], max: &mut [u8; hll::REGISTERS]) -> Result<bool> {
183 let Some(bytes) = self.sketch(key)? else {
184 return Ok(false);
185 };
186 let enc = hll::check(bytes)?;
187 if !hll::merge(max, bytes, enc) {
188 return Err(hll::corrupt());
189 }
190 Ok(enc == Encoding::Dense)
191 }
192
193 pub(crate) fn pfmerge_into(
198 &mut self,
199 dest: &[u8],
200 max: &[u8; hll::REGISTERS],
201 dense: bool,
202 ) -> Result<()> {
203 let mut buf = std::mem::take(&mut self.scratch);
204 buf.clear();
205 let deadline = match self.map.get(dest) {
206 Some(rec) => {
207 value::read(rec).write_to(&mut buf);
208 value::expire_at(rec)
209 }
210 None => {
211 hll::empty(&mut buf);
212 None
213 }
214 };
215
216 let outcome = 'work: {
217 if dense && !hll::to_dense(&mut buf) {
222 break 'work Err(hll::corrupt());
223 }
224 for (i, &val) in max.iter().enumerate() {
225 if val != 0 && hll::set(&mut buf, i, val).is_none() {
226 break 'work Err(hll::corrupt());
227 }
228 }
229 Ok(())
230 };
231
232 if outcome.is_ok() {
233 hll::invalidate(&mut buf);
234 self.store_raw(dest, &buf, deadline);
235 }
236 self.scratch = buf;
237 outcome
238 }
239
240 pub fn pfgetreg(&mut self, key: &[u8], regs: &mut [u8; hll::REGISTERS]) -> Result<()> {
247 self.pftodense(key)?;
248 let bytes = self.sketch(key)?.ok_or_else(no_such_key)?;
249 let body = &bytes[hll::HDR..];
250 for (i, slot) in regs.iter_mut().enumerate() {
251 *slot = hll::dense_get(body, i);
252 }
253 Ok(())
254 }
255
256 pub fn pftodense(&mut self, key: &[u8]) -> Result<bool> {
258 self.reap(key);
259 self.string_only(key)?;
260 self.thaw(key)?;
261 let bytes = self.sketch(key)?.ok_or_else(no_such_key)?;
262 if hll::check(bytes)? == Encoding::Dense {
263 return Ok(false);
264 }
265
266 let mut buf = std::mem::take(&mut self.scratch);
267 buf.clear();
268 let deadline = match self.map.get(key) {
269 Some(rec) => {
270 value::read(rec).write_to(&mut buf);
271 value::expire_at(rec)
272 }
273 None => None,
274 };
275 let outcome = if hll::to_dense(&mut buf) {
276 self.store_raw(key, &buf, deadline);
277 Ok(true)
278 } else {
279 Err(hll::corrupt())
280 };
281 self.scratch = buf;
282 outcome
283 }
284
285 pub fn pfencoding(&mut self, key: &[u8]) -> Result<Encoding> {
287 self.reap(key);
288 self.string_only(key)?;
289 self.warm(key)?;
290 let bytes = self.sketch(key)?.ok_or_else(no_such_key)?;
291 hll::check(bytes)
292 }
293
294 pub fn pfdecode<T>(&mut self, key: &[u8], run: impl FnOnce(&[u8]) -> T) -> Result<T> {
300 self.reap(key);
301 self.string_only(key)?;
302 self.warm(key)?;
303
304 let mut buf = std::mem::take(&mut self.scratch);
308 buf.clear();
309 let outcome = 'work: {
310 let bytes = match self.sketch(key) {
311 Ok(Some(bytes)) => bytes,
312 Ok(None) => break 'work Err(no_such_key()),
313 Err(e) => break 'work Err(e),
314 };
315 match hll::check(bytes) {
316 Ok(Encoding::Sparse) => hll::decode(bytes, &mut buf),
317 Ok(Encoding::Dense) => break 'work Err(Error::new(Code::Invalid, NOT_SPARSE)),
318 Err(e) => break 'work Err(e),
319 }
320 Ok(())
321 };
322 let out = outcome.map(|()| run(&buf));
323 self.scratch = buf;
324 out
325 }
326
327 fn count_one(&mut self, key: &[u8]) -> Result<u64> {
329 let Some(bytes) = self.sketch(key)? else {
330 return Ok(0);
331 };
332 let enc = hll::check(bytes)?;
333 if let Some(n) = hll::cached(bytes) {
334 return Ok(n);
335 }
336 let n = hll::count(bytes, enc)?;
337
338 let hash = RawMap::hash_of(key);
342 if let Some(rec) = self.map.value_mut_hashed(hash, key)
343 && let Some(val) = value::raw_in_place(rec)
344 {
345 hll::cache(val, n);
346 }
347 Ok(n)
348 }
349
350 fn sketch(&self, key: &[u8]) -> Result<Option<&[u8]>> {
357 match self.peek(key) {
358 None => Ok(None),
359 Some(Str::Bytes(b)) => Ok(Some(b)),
360 Some(Str::Int(_)) => Err(hll::not_hll()),
361 }
362 }
363}
364
365impl Db {
366 pub fn pfcount<'k, I>(&self, keys: I) -> Result<u64>
380 where
381 I: Iterator<Item = &'k [u8]> + Clone,
382 {
383 if let Some(home) = self.one_stripe(keys.clone()) {
384 return self.hold_stripe(home).pfcount(keys);
385 }
386 let mut held = self.hold_many(keys.clone().map(|key| self.stripe_of(key)));
387 for key in keys.clone() {
388 held.stripe_mut(self.stripe_of(key)).hll_ready(key)?;
389 }
390 let mut max = [0u8; hll::REGISTERS];
391 for key in keys {
392 held.stripe(self.stripe_of(key))
393 .merge_sketch(key, &mut max)?;
394 }
395 Ok(estimate(&max))
396 }
397
398 pub fn pfmerge<'k, I>(&self, dest: &'k [u8], srcs: I) -> Result<()>
407 where
408 I: Iterator<Item = &'k [u8]> + Clone,
409 {
410 if let Some(home) = self.one_stripe(std::iter::once(dest).chain(srcs.clone())) {
411 return self.hold_stripe(home).pfmerge(dest, srcs);
412 }
413 let onto = self.stripe_of(dest);
414 let mut held = self
415 .hold_many(std::iter::once(onto).chain(srcs.clone().map(|src| self.stripe_of(src))));
416 held.stripe_mut(onto).hll_ready(dest)?;
417 check_len(dest, hll::DENSE)?;
418 for src in srcs.clone() {
419 held.stripe_mut(self.stripe_of(src)).hll_ready(src)?;
420 }
421
422 let mut max = [0u8; hll::REGISTERS];
423 let mut dense = false;
424 for key in std::iter::once(dest).chain(srcs) {
425 dense |= held
426 .stripe(self.stripe_of(key))
427 .merge_sketch(key, &mut max)?;
428 }
429 held.stripe_mut(onto).pfmerge_into(dest, &max, dense)
430 }
431}
432
433fn estimate(max: &[u8; hll::REGISTERS]) -> u64 {
435 let mut hist = [0u32; 64];
436 for &val in max {
437 hist[val as usize] += 1;
438 }
439 hll::estimate(&hist)
440}
441
442fn no_such_key() -> Error {
444 Error::new(Code::NotFound, NO_SUCH_KEY)
445}
446
447#[cfg(test)]
448mod tests {
449 use super::*;
450 use crate::keyspace::Keyspace;
451
452 fn db() -> Keyspace {
453 Keyspace::new()
454 }
455
456 fn names<'k>(list: &'k [&'k [u8]]) -> impl Iterator<Item = &'k [u8]> + Clone {
458 list.iter().copied()
459 }
460
461 fn bytes(db: &mut Keyspace, key: &[u8]) -> Vec<u8> {
463 db.get(key).expect("a value").expect("bytes").to_vec()
464 }
465
466 #[test]
468 fn a_sketch_a_client_reads_is_the_one_a_real_server_wrote() {
469 let mut db = db();
470 assert!(db.pfadd(b"h", names(&[b"a", b"b", b"c"])).expect("an add"));
471 assert_eq!(
472 bytes(&mut db, b"h"),
473 b"HYLL\x01\0\0\0\0\0\0\0\0\0\0\x80\x60\xf3\x80\x50\xb1\x84\x4b\xfb\x80\x42\x5a"
474 );
475 assert_eq!(db.pfcount(names(&[b"h"])).expect("a count"), 3);
476 assert_eq!(bytes(&mut db, b"h")[8..16], [3, 0, 0, 0, 0, 0, 0, 0]);
478 }
479
480 #[test]
482 fn adding_nothing_still_creates_the_key() {
483 let mut db = db();
484 assert!(db.pfadd(b"h", names(&[])).expect("an add"));
485 assert!(db.exists(b"h"));
486 assert_eq!(db.strlen(b"h").expect("a length"), 18);
487 assert!(!db.pfadd(b"h", names(&[])).expect("an add"));
488 assert_eq!(db.pfcount(names(&[b"h"])).expect("a count"), 0);
489 }
490
491 #[test]
493 fn adding_an_element_twice_leaves_the_value_alone() {
494 let mut db = db();
495 db.pfadd(b"h", names(&[b"a"])).expect("an add");
496 let before = bytes(&mut db, b"h");
497 assert!(!db.pfadd(b"h", names(&[b"a"])).expect("an add"));
498 assert_eq!(bytes(&mut db, b"h"), before);
499 }
500
501 #[test]
503 #[cfg_attr(
504 miri,
505 ignore = "the counts are the claim, so there is no smaller version of it"
506 )]
507 fn the_count_is_the_number_a_real_server_gives() {
508 for (n, want) in [(100usize, 100u64), (1000, 995), (10_000, 10_077)] {
509 let mut db = db();
510 for i in 0..n {
511 let ele = format!("e:{i}");
512 db.pfadd(b"h", names(&[ele.as_bytes()])).expect("an add");
513 }
514 assert_eq!(db.pfcount(names(&[b"h"])).expect("a count"), want, "{n}");
515 }
516 }
517
518 #[test]
520 #[cfg_attr(
521 miri,
522 ignore = "the sizes are the claim, and only these counts reach them"
523 )]
524 fn a_sketch_is_sparse_until_it_is_not() {
525 let mut db = db();
526 for i in 0..1000 {
527 let ele = format!("e:{i}");
528 db.pfadd(b"k1", names(&[ele.as_bytes()])).expect("an add");
529 }
530 assert_eq!(db.strlen(b"k1").expect("a length"), 1880);
531 assert_eq!(db.pfencoding(b"k1").expect("an encoding"), Encoding::Sparse);
532 const { assert!(1880 <= hll::SPARSE_MAX) };
533
534 for i in 0..10_000 {
535 let ele = format!("e:{i}");
536 db.pfadd(b"k2", names(&[ele.as_bytes()])).expect("an add");
537 }
538 assert_eq!(db.strlen(b"k2").expect("a length"), 12304);
539 assert_eq!(db.pfencoding(b"k2").expect("an encoding"), Encoding::Dense);
540 }
541
542 #[test]
549 fn counting_several_keys_counts_their_union() {
550 let mut db = db();
551 for i in 0..200 {
552 let ele = format!("e:{i}");
553 let key: &[u8] = if i < 150 { b"a" } else { b"b" };
554 db.pfadd(key, names(&[ele.as_bytes()])).expect("an add");
555 }
556 assert_eq!(db.pfcount(names(&[b"a"])).expect("a count"), 151);
557 assert_eq!(db.pfcount(names(&[b"b"])).expect("a count"), 49);
558 assert_eq!(db.pfcount(names(&[b"a", b"b"])).expect("a count"), 199);
559 assert_eq!(db.pfcount(names(&[b"gone"])).expect("a count"), 0);
561 assert_eq!(db.pfcount(names(&[b"a", b"gone"])).expect("a count"), 151);
562 assert!(!db.exists(b"gone"));
563 }
564
565 #[test]
567 fn a_merge_is_a_union_and_keeps_the_smaller_form() {
568 let mut db = db();
569 for i in 0..100 {
570 let ele = format!("e:{i}");
571 db.pfadd(b"s1", names(&[ele.as_bytes()])).expect("an add");
572 db.pfadd(b"s2", names(&[ele.as_bytes()])).expect("an add");
573 }
574 db.pfmerge(b"m", names(&[b"s1", b"s2"])).expect("a merge");
575 assert_eq!(db.strlen(b"m").expect("a length"), 279);
576 assert_eq!(db.pfencoding(b"m").expect("an encoding"), Encoding::Sparse);
577 assert_eq!(db.pfcount(names(&[b"m"])).expect("a count"), 100);
578
579 for i in 100..200 {
581 let ele = format!("e:{i}");
582 db.pfadd(b"s3", names(&[ele.as_bytes()])).expect("an add");
583 }
584 db.pfmerge(b"m", names(&[b"s3"])).expect("a merge");
585 assert_eq!(db.pfcount(names(&[b"m"])).expect("a count"), 199);
586 assert_eq!(db.strlen(b"m").expect("a length"), 499);
587
588 let before = bytes(&mut db, b"m");
590 db.pfmerge(b"m", names(&[])).expect("a merge");
591 assert_eq!(bytes(&mut db, b"m")[..8], before[..8]);
592 assert_eq!(db.pfcount(names(&[b"m"])).expect("a count"), 199);
593 }
594
595 #[test]
600 #[cfg_attr(
601 miri,
602 ignore = "nothing smaller is dense, and the merged count is a real server's"
603 )]
604 fn a_dense_source_makes_the_result_dense() {
605 let mut db = db();
606 for i in 0..100 {
607 let ele = format!("e:{i}");
608 db.pfadd(b"small", names(&[ele.as_bytes()]))
609 .expect("an add");
610 }
611 for i in 0..20_000 {
612 let ele = format!("e:{i}");
613 db.pfadd(b"big", names(&[ele.as_bytes()])).expect("an add");
614 }
615 db.pfmerge(b"m", names(&[b"small", b"big"]))
616 .expect("a merge");
617 assert_eq!(db.strlen(b"m").expect("a length"), 12304);
618 assert_eq!(db.pfencoding(b"m").expect("an encoding"), Encoding::Dense);
619 assert_eq!(db.pfcount(names(&[b"m"])).expect("a count"), 20096);
620 }
621
622 #[test]
624 fn a_write_keeps_the_deadline() {
625 let mut db = db();
626 let mut fresh = Vec::new();
627 hll::empty(&mut fresh);
628 db.setex(b"h", 100, &fresh).expect("a set");
629 let had = db.expire_at(b"h").expect("a deadline");
630 db.pfadd(b"h", names(&[b"a"])).expect("an add");
631 assert_eq!(db.expire_at(b"h"), Some(had));
632 db.pfmerge(b"h", names(&[])).expect("a merge");
633 assert_eq!(db.expire_at(b"h"), Some(had));
634 db.pftodense(b"h").expect("a conversion");
635 assert_eq!(db.expire_at(b"h"), Some(had));
636 assert_eq!(db.pfcount(names(&[b"h"])).expect("a count"), 1);
637 }
638
639 #[test]
641 fn the_debug_commands_say_what_a_real_server_says() {
642 let mut db = db();
643 db.pfadd(b"h", names(&[b"a", b"b", b"c"])).expect("an add");
644 let text = db.pfdecode(b"h", <[u8]>::to_vec).expect("a decode");
645 assert_eq!(text, b"Z:8436 v:1,1 Z:4274 v:2,1 Z:3068 v:1,1 Z:603");
646
647 let mut regs = [0u8; hll::REGISTERS];
650 db.pfgetreg(b"h", &mut regs).expect("the registers");
651 assert_eq!(regs[8436], 1);
652 assert_eq!(regs[12711], 2);
653 assert_eq!(regs[15780], 1);
654 assert_eq!(regs.iter().filter(|&&v| v != 0).count(), 3);
655 assert_eq!(db.pfencoding(b"h").expect("an encoding"), Encoding::Dense);
656 assert_eq!(db.strlen(b"h").expect("a length"), 12304);
657 assert_eq!(db.pfcount(names(&[b"h"])).expect("a count"), 3);
658
659 assert!(db.pfdecode(b"h", <[u8]>::to_vec).is_err());
661 assert!(!db.pftodense(b"h").expect("a conversion"));
662 }
663
664 #[test]
666 fn a_string_that_is_not_a_sketch_is_refused() {
667 let mut db = db();
668 db.set_plain(b"plain", b"not a sketch at all")
669 .expect("a set");
670 assert!(db.pfadd(b"plain", names(&[b"a"])).is_err());
671 assert!(db.pfcount(names(&[b"plain"])).is_err());
672 assert!(db.pfmerge(b"plain", names(&[])).is_err());
673 assert!(db.pfencoding(b"plain").is_err());
674 db.set_plain(b"n", b"12345").expect("a set");
676 assert!(db.pfcount(names(&[b"n"])).is_err());
677
678 assert!(db.pfencoding(b"gone").is_err());
681 assert!(db.pftodense(b"gone").is_err());
682 assert!(db.pfdecode(b"gone", <[u8]>::to_vec).is_err());
683 let mut regs = [0u8; hll::REGISTERS];
684 assert!(db.pfgetreg(b"gone", &mut regs).is_err());
685 }
686
687 #[test]
689 fn a_corrupted_sketch_is_reported() {
690 let mut db = db();
691 db.pfadd(b"h", names(&[b"a", b"b", b"c"])).expect("an add");
692 let mut short = bytes(&mut db, b"h");
693 short.pop();
694 db.set_plain(b"h", &short).expect("a set");
695 let err = db.pfcount(names(&[b"h"])).expect_err("a complaint");
696 assert_eq!(err.code(), Code::Corrupt);
697 assert!(db.pfcount(names(&[b"h", b"h"])).is_err());
698 assert!(db.pfmerge(b"m", names(&[b"h"])).is_err());
699 }
700}