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 fn the_count_is_the_number_a_real_server_gives() {
504 for (n, want) in [(100usize, 100u64), (1000, 995), (10_000, 10_077)] {
505 let mut db = db();
506 for i in 0..n {
507 let ele = format!("e:{i}");
508 db.pfadd(b"h", names(&[ele.as_bytes()])).expect("an add");
509 }
510 assert_eq!(db.pfcount(names(&[b"h"])).expect("a count"), want, "{n}");
511 }
512 }
513
514 #[test]
516 fn a_sketch_is_sparse_until_it_is_not() {
517 let mut db = db();
518 for i in 0..1000 {
519 let ele = format!("e:{i}");
520 db.pfadd(b"k1", names(&[ele.as_bytes()])).expect("an add");
521 }
522 assert_eq!(db.strlen(b"k1").expect("a length"), 1880);
523 assert_eq!(db.pfencoding(b"k1").expect("an encoding"), Encoding::Sparse);
524 const { assert!(1880 <= hll::SPARSE_MAX) };
525
526 for i in 0..10_000 {
527 let ele = format!("e:{i}");
528 db.pfadd(b"k2", names(&[ele.as_bytes()])).expect("an add");
529 }
530 assert_eq!(db.strlen(b"k2").expect("a length"), 12304);
531 assert_eq!(db.pfencoding(b"k2").expect("an encoding"), Encoding::Dense);
532 }
533
534 #[test]
541 fn counting_several_keys_counts_their_union() {
542 let mut db = db();
543 for i in 0..200 {
544 let ele = format!("e:{i}");
545 let key: &[u8] = if i < 150 { b"a" } else { b"b" };
546 db.pfadd(key, names(&[ele.as_bytes()])).expect("an add");
547 }
548 assert_eq!(db.pfcount(names(&[b"a"])).expect("a count"), 151);
549 assert_eq!(db.pfcount(names(&[b"b"])).expect("a count"), 49);
550 assert_eq!(db.pfcount(names(&[b"a", b"b"])).expect("a count"), 199);
551 assert_eq!(db.pfcount(names(&[b"gone"])).expect("a count"), 0);
553 assert_eq!(db.pfcount(names(&[b"a", b"gone"])).expect("a count"), 151);
554 assert!(!db.exists(b"gone"));
555 }
556
557 #[test]
559 fn a_merge_is_a_union_and_keeps_the_smaller_form() {
560 let mut db = db();
561 for i in 0..100 {
562 let ele = format!("e:{i}");
563 db.pfadd(b"s1", names(&[ele.as_bytes()])).expect("an add");
564 db.pfadd(b"s2", names(&[ele.as_bytes()])).expect("an add");
565 }
566 db.pfmerge(b"m", names(&[b"s1", b"s2"])).expect("a merge");
567 assert_eq!(db.strlen(b"m").expect("a length"), 279);
568 assert_eq!(db.pfencoding(b"m").expect("an encoding"), Encoding::Sparse);
569 assert_eq!(db.pfcount(names(&[b"m"])).expect("a count"), 100);
570
571 for i in 100..200 {
573 let ele = format!("e:{i}");
574 db.pfadd(b"s3", names(&[ele.as_bytes()])).expect("an add");
575 }
576 db.pfmerge(b"m", names(&[b"s3"])).expect("a merge");
577 assert_eq!(db.pfcount(names(&[b"m"])).expect("a count"), 199);
578 assert_eq!(db.strlen(b"m").expect("a length"), 499);
579
580 let before = bytes(&mut db, b"m");
582 db.pfmerge(b"m", names(&[])).expect("a merge");
583 assert_eq!(bytes(&mut db, b"m")[..8], before[..8]);
584 assert_eq!(db.pfcount(names(&[b"m"])).expect("a count"), 199);
585 }
586
587 #[test]
589 fn a_dense_source_makes_the_result_dense() {
590 let mut db = db();
591 for i in 0..100 {
592 let ele = format!("e:{i}");
593 db.pfadd(b"small", names(&[ele.as_bytes()]))
594 .expect("an add");
595 }
596 for i in 0..20_000 {
597 let ele = format!("e:{i}");
598 db.pfadd(b"big", names(&[ele.as_bytes()])).expect("an add");
599 }
600 db.pfmerge(b"m", names(&[b"small", b"big"]))
601 .expect("a merge");
602 assert_eq!(db.strlen(b"m").expect("a length"), 12304);
603 assert_eq!(db.pfencoding(b"m").expect("an encoding"), Encoding::Dense);
604 assert_eq!(db.pfcount(names(&[b"m"])).expect("a count"), 20096);
605 }
606
607 #[test]
609 fn a_write_keeps_the_deadline() {
610 let mut db = db();
611 let mut fresh = Vec::new();
612 hll::empty(&mut fresh);
613 db.setex(b"h", 100, &fresh).expect("a set");
614 let had = db.expire_at(b"h").expect("a deadline");
615 db.pfadd(b"h", names(&[b"a"])).expect("an add");
616 assert_eq!(db.expire_at(b"h"), Some(had));
617 db.pfmerge(b"h", names(&[])).expect("a merge");
618 assert_eq!(db.expire_at(b"h"), Some(had));
619 db.pftodense(b"h").expect("a conversion");
620 assert_eq!(db.expire_at(b"h"), Some(had));
621 assert_eq!(db.pfcount(names(&[b"h"])).expect("a count"), 1);
622 }
623
624 #[test]
626 fn the_debug_commands_say_what_a_real_server_says() {
627 let mut db = db();
628 db.pfadd(b"h", names(&[b"a", b"b", b"c"])).expect("an add");
629 let text = db.pfdecode(b"h", <[u8]>::to_vec).expect("a decode");
630 assert_eq!(text, b"Z:8436 v:1,1 Z:4274 v:2,1 Z:3068 v:1,1 Z:603");
631
632 let mut regs = [0u8; hll::REGISTERS];
635 db.pfgetreg(b"h", &mut regs).expect("the registers");
636 assert_eq!(regs[8436], 1);
637 assert_eq!(regs[12711], 2);
638 assert_eq!(regs[15780], 1);
639 assert_eq!(regs.iter().filter(|&&v| v != 0).count(), 3);
640 assert_eq!(db.pfencoding(b"h").expect("an encoding"), Encoding::Dense);
641 assert_eq!(db.strlen(b"h").expect("a length"), 12304);
642 assert_eq!(db.pfcount(names(&[b"h"])).expect("a count"), 3);
643
644 assert!(db.pfdecode(b"h", <[u8]>::to_vec).is_err());
646 assert!(!db.pftodense(b"h").expect("a conversion"));
647 }
648
649 #[test]
651 fn a_string_that_is_not_a_sketch_is_refused() {
652 let mut db = db();
653 db.set_plain(b"plain", b"not a sketch at all")
654 .expect("a set");
655 assert!(db.pfadd(b"plain", names(&[b"a"])).is_err());
656 assert!(db.pfcount(names(&[b"plain"])).is_err());
657 assert!(db.pfmerge(b"plain", names(&[])).is_err());
658 assert!(db.pfencoding(b"plain").is_err());
659 db.set_plain(b"n", b"12345").expect("a set");
661 assert!(db.pfcount(names(&[b"n"])).is_err());
662
663 assert!(db.pfencoding(b"gone").is_err());
666 assert!(db.pftodense(b"gone").is_err());
667 assert!(db.pfdecode(b"gone", <[u8]>::to_vec).is_err());
668 let mut regs = [0u8; hll::REGISTERS];
669 assert!(db.pfgetreg(b"gone", &mut regs).is_err());
670 }
671
672 #[test]
674 fn a_corrupted_sketch_is_reported() {
675 let mut db = db();
676 db.pfadd(b"h", names(&[b"a", b"b", b"c"])).expect("an add");
677 let mut short = bytes(&mut db, b"h");
678 short.pop();
679 db.set_plain(b"h", &short).expect("a set");
680 let err = db.pfcount(names(&[b"h"])).expect_err("a complaint");
681 assert_eq!(err.code(), Code::Corrupt);
682 assert!(db.pfcount(names(&[b"h", b"h"])).is_err());
683 assert!(db.pfmerge(b"m", names(&[b"h"])).is_err());
684 }
685}