1use std::hash::{Hash, Hasher};
8
9use casefold::index_fold_char;
10
11use crate::MAX_SPARSE_GRAM_SIZE;
12use crate::ngram::NGram;
13use crate::table::{bigram_h, bigram_priority_rolling};
14
15const RING_MASK: u32 = MAX_SPARSE_GRAM_SIZE as u32 - 1;
17
18#[derive(Clone, Copy, Debug)]
19struct PosState {
20 index: u32,
21 value: u32,
22}
23
24#[derive(Clone, Debug, Default)]
25struct Queue {
26 idx_buf: [u32; MAX_SPARSE_GRAM_SIZE],
27 val_buf: [u32; MAX_SPARSE_GRAM_SIZE],
28 head: u32,
29 len: u32,
30}
31
32impl Queue {
33 fn new() -> Self {
34 Self {
35 idx_buf: [0; MAX_SPARSE_GRAM_SIZE],
36 val_buf: [0; MAX_SPARSE_GRAM_SIZE],
37 head: 0,
38 len: 0,
39 }
40 }
41
42 fn clear(&mut self) {
43 self.len = 0;
44 }
45
46 fn is_empty(&self) -> bool {
47 self.len == 0
48 }
49
50 fn front_idx(&self) -> u32 {
51 debug_assert!(!self.is_empty());
54 self.idx_buf[self.head as usize]
55 }
56
57 fn front_value(&self) -> u32 {
58 debug_assert!(!self.is_empty());
61 self.val_buf[self.head as usize]
62 }
63
64 fn pop_front_idx(&mut self) -> u32 {
65 debug_assert!(!self.is_empty());
66 let first = self.idx_buf[self.head as usize];
67 self.head = (self.head + 1) & RING_MASK;
68 self.len -= 1;
69 first
70 }
71
72 fn push(&mut self, state: PosState) {
73 while !self.is_empty() {
76 let slot = ((self.head + self.len - 1) & RING_MASK) as usize;
77 if self.val_buf[slot] <= state.value {
78 break;
79 }
80 self.len -= 1;
81 }
82 debug_assert!(self.len < MAX_SPARSE_GRAM_SIZE as u32);
83 let slot = ((self.head + self.len) & RING_MASK) as usize;
84 self.idx_buf[slot] = state.index;
85 self.val_buf[slot] = state.value;
86 self.len += 1;
87 }
88}
89
90#[derive(Clone)]
98pub struct QueryGrams {
99 queue: Queue,
101 content: u64,
103 content_end_idx: u32,
105 h: u32,
107}
108
109impl Eq for QueryGrams {}
110
111impl PartialEq for QueryGrams {
112 fn eq(&self, other: &Self) -> bool {
113 self.state() == other.state()
114 }
115}
116
117impl Hash for QueryGrams {
118 fn hash<H: Hasher>(&self, state: &mut H) {
119 self.state().hash(state);
120 }
121}
122
123impl Default for QueryGrams {
124 fn default() -> Self {
125 Self {
126 content: 0,
127 queue: Queue::new(),
128 content_end_idx: 0,
129 h: 0,
130 }
131 }
132}
133
134impl QueryGrams {
135 #[inline]
154 pub fn state(&self) -> (u32, u64) {
155 let begin = if self.queue.is_empty() {
156 0
157 } else {
158 self.queue.front_idx() - 2
162 };
163 let content_len = self.content_end_idx - begin;
164 debug_assert!(content_len < MAX_SPARSE_GRAM_SIZE as u32);
165 let mask = (1u64 << (content_len * 8)) - 1;
166 (content_len, self.content & mask)
167 }
168
169 pub fn min_priority(&self) -> u32 {
178 if self.queue.is_empty() {
179 u32::MAX
180 } else {
181 self.queue.front_value()
182 }
183 }
184
185 #[inline]
193 fn follow_byte(&self, end: u32) -> Option<u8> {
194 if end >= self.content_end_idx {
195 None
196 } else {
197 let shift = self.content_end_idx - 1 - end;
198 debug_assert!(shift < MAX_SPARSE_GRAM_SIZE as u32);
199 Some((self.content >> (shift * 8)) as u8)
200 }
201 }
202
203 fn extract_gram<F>(&mut self, begin_index: u32, end_index: u32, consumer: &mut F)
204 where
205 F: FnMut(NGram, u32, Option<u8>),
206 {
207 debug_assert!(end_index >= begin_index);
208 debug_assert!(end_index <= self.content_end_idx);
209
210 let len = (end_index - begin_index + 2) as usize;
211 let dist = self.content_end_idx - end_index;
212 let shifted = self.content >> (dist * 8);
213 let follow = self.follow_byte(end_index);
216 consumer(NGram::from_window_masked(shifted, len), end_index, follow);
217 }
218
219 pub fn append_char<F>(&mut self, c: char, consumer: F)
223 where
224 F: FnMut(NGram, u32, Option<u8>),
225 {
226 self.append_byte(index_fold_char(c), consumer);
227 }
228
229 pub fn append_byte<F>(&mut self, right: u8, mut consumer: F)
235 where
236 F: FnMut(NGram, u32, Option<u8>),
237 {
238 let left = (self.content & 0xFF) as u8;
239 self.content_end_idx += 1;
240 self.content = (self.content << 8) | right as u64;
241
242 let idx = self.content_end_idx;
243
244 if idx == 1 {
247 self.h = bigram_h(right);
248 } else {
249 let (value, h_b) = bigram_priority_rolling(left, right, self.h);
250 self.h = h_b;
251 if !self.queue.is_empty()
252 && let priority = self.queue.front_value()
253 && value < priority
254 {
255 let mut begin = self.queue.pop_front_idx();
256 while !self.queue.is_empty() && self.queue.front_value() == priority {
257 let end = self.queue.pop_front_idx();
258 self.extract_gram(begin, end, &mut consumer);
259 begin = end;
260 }
261 self.queue.clear();
262 self.queue.push(PosState { index: idx, value });
263 self.extract_gram(begin, idx, &mut consumer);
264 } else {
265 self.queue.push(PosState { index: idx, value });
266 }
267 if idx - self.queue.front_idx() + 2 >= MAX_SPARSE_GRAM_SIZE as u32 && self.queue.len > 1
268 {
269 let begin = self.queue.pop_front_idx();
270 let end = self.queue.front_idx();
271 self.extract_gram(begin, end, &mut consumer);
272 }
273 }
274 }
275
276 pub fn flush<F>(mut self, mut consumer: F)
278 where
279 F: FnMut(NGram, u32, Option<u8>),
280 {
281 if self.content_end_idx == 2 {
282 self.extract_gram(2, 2, &mut consumer);
283 } else {
284 while self.queue.len > 1 {
285 let begin = self.queue.pop_front_idx();
286 let end = self.queue.front_idx();
287 self.extract_gram(begin, end, &mut consumer);
288 }
289 }
290 }
291
292 pub fn consume_first<F>(&mut self, mut consumer: F)
294 where
295 F: FnMut(NGram, u32, Option<u8>),
296 {
297 if self.queue.len > 1 {
298 let begin = self.queue.pop_front_idx();
300 let end = self.queue.front_idx();
301 self.extract_gram(begin, end, &mut consumer);
302 } else if self.content_end_idx == 2 {
303 self.extract_gram(2, 2, &mut consumer);
305 }
306
307 if self.queue.len <= 1 {
311 self.queue.clear();
312 if self.content_end_idx > 1 {
313 self.content_end_idx = 1;
318 } else {
319 self.content_end_idx = 0;
325 }
326 }
327 }
328}
329
330#[cfg(test)]
331mod tests {
332 use super::*;
333 use crate::collect_sparse_grams_deque;
334 use std::collections::hash_map::DefaultHasher;
335 use std::ops::Range;
336
337 fn query_intervals(input: &str) -> Vec<Range<u32>> {
343 let mut q = QueryGrams::default();
344 let mut out = Vec::new();
345 for c in input.chars() {
346 q.append_char(c, |gram, end, _follow| {
347 let begin = end + 1 - gram.len() as u32;
348 out.push(begin..end - 1);
349 });
350 }
351 q.flush(|gram, end, _follow| {
352 let begin = end + 1 - gram.len() as u32;
353 out.push(begin..end - 1);
354 });
355 out
356 }
357
358 fn candidate_intervals(bytes: &[u8]) -> Vec<Range<u32>> {
363 let mut out = Vec::new();
364 collect_sparse_grams_deque(bytes, |gram, idx| {
365 let begin = idx + 1 - gram.len() as u32;
366 out.push(begin..idx - 1);
367 });
368 out
369 }
370
371 fn min_cover_size_query_semantics(m: u32, intervals: &[Range<u32>]) -> usize {
372 let inf = usize::MAX / 4;
373 let mut dp = vec![0usize; m as usize + 1];
374 for pos in (1..m).rev() {
375 let mut best = inf;
376 for iv in intervals {
377 if iv.start <= pos && iv.end > pos {
378 let next = iv.end.min(m) as usize;
379 best = best.min(1 + dp[next]);
380 }
381 }
382 dp[pos as usize] = best;
383 }
384 dp[1]
385 }
386
387 fn is_valid_query_cover_chain(m: u32, intervals: &[Range<u32>]) -> bool {
388 if m <= 1 {
389 return true;
390 }
391 let mut produced = 1u32;
392 for iv in intervals {
393 if iv.start > produced || iv.end <= produced {
394 return false;
395 }
396 produced = iv.end;
397 if produced >= m {
398 return true;
399 }
400 }
401 false
402 }
403
404 #[derive(Clone, Copy, Debug)]
405 struct Rng64(u64);
406
407 impl Rng64 {
408 fn new(seed: u64) -> Self {
409 Self(seed)
410 }
411
412 fn next_u64(&mut self) -> u64 {
413 let mut x = self.0;
415 x ^= x >> 12;
416 x ^= x << 25;
417 x ^= x >> 27;
418 self.0 = x;
419 x.wrapping_mul(0x2545_F491_4F6C_DD1D)
420 }
421
422 fn gen_range(&mut self, upper: usize) -> usize {
423 (self.next_u64() % upper as u64) as usize
424 }
425 }
426
427 #[test]
428 fn append_and_flush_emit_grams() {
429 let mut q = QueryGrams::default();
430 for c in "hello world".chars() {
431 q.append_char(c, |_gram, _idx, _follow| {});
432 }
433 let mut out = Vec::new();
434 q.flush(|gram, idx, _follow| out.push((gram, idx)));
435 assert!(!out.is_empty());
436 assert!(
437 out.iter()
438 .all(|(g, _)| (2..=MAX_SPARSE_GRAM_SIZE).contains(&g.len()))
439 );
440 }
441
442 #[test]
443 fn single_bigram_has_no_follow() {
444 let bytes: Vec<u8> = "ab".chars().map(crate::index_fold_char).collect();
447 let mut q = QueryGrams::default();
448 for &b in &bytes {
449 q.append_byte(b, |_gram, _end, _follow| {});
450 }
451 let mut emitted = Vec::new();
452 q.flush(|gram, end, follow| emitted.push((gram.len(), end, follow)));
453 assert_eq!(emitted, vec![(2usize, 2u32, None)]);
454 }
455
456 #[test]
457 fn emitted_follow_is_the_actual_next_char() {
458 for input in [
461 "ab",
462 "abc",
463 "hello",
464 "querystream",
465 "aaaaaaaaaa",
466 "lardeee",
467 "sssdk",
468 ] {
469 let bytes: Vec<u8> = input.chars().map(crate::index_fold_char).collect();
470 let n = bytes.len() as u32;
471 let mut q = QueryGrams::default();
472 let mut check = |_gram: NGram, end: u32, follow: Option<u8>| {
473 if let Some(b) = follow {
474 assert!(
475 end < n,
476 "follow reported past input for {input:?}: end={end}"
477 );
478 assert_eq!(
479 b, bytes[end as usize],
480 "wrong follow byte for {input:?} at end={end}"
481 );
482 }
483 };
484 for &b in &bytes {
485 q.append_byte(b, &mut check);
486 }
487 q.flush(&mut check);
488 }
489 }
490
491 #[test]
492 fn state_tracks_active_window_each_step() {
493 let input = b"hello world";
502 let expected_len = [1u32, 2, 3, 2, 3, 2, 3, 2, 3, 4, 5];
504
505 let pack = |bytes: &[u8]| bytes.iter().fold(0u64, |acc, &b| (acc << 8) | b as u64);
506
507 let mut q = QueryGrams::default();
508 let mut prev = q.state();
509 assert_eq!(prev, (0, 0), "a fresh state must be empty");
510
511 for (i, &byte) in input.iter().enumerate() {
512 assert_eq!(
514 q.state(),
515 prev,
516 "state moved without an append before byte {i}"
517 );
518
519 q.append_byte(byte, |_, _, _| {});
520
521 let len = expected_len[i];
522 let window = &input[i + 1 - len as usize..=i];
523 let after = q.state();
524 assert_eq!(
525 after,
526 (len, pack(window)),
527 "unexpected state after byte {i} ({:?})",
528 char::from(byte),
529 );
530 prev = after;
531 }
532 }
533
534 #[test]
535 fn state_eq_and_hash_ignore_absolute_history() {
536 let mut a = QueryGrams::default();
537 let mut b = QueryGrams::default();
538
539 for c in "abc".chars() {
540 a.append_char(c, |_gram, _idx, _follow| {});
541 }
542 for c in "zabc".chars() {
543 b.append_char(c, |_gram, _idx, _follow| {});
544 }
545 b.consume_first(|_gram, _idx, _follow| {});
546
547 if a == b {
549 let mut ha = DefaultHasher::new();
550 let mut hb = DefaultHasher::new();
551 a.hash(&mut ha);
552 b.hash(&mut hb);
553 assert_eq!(ha.finish(), hb.finish());
554 }
555 }
556
557 #[test]
558 fn query_flush_is_minimum_cover_on_small_inputs() {
559 for input in [
560 "abc",
561 "abcd",
562 "abcdef",
563 "hello",
564 "hello world",
565 "ababababab",
566 ] {
567 let bytes = input.as_bytes();
568 if bytes.len() < 3 {
569 continue;
570 }
571 let produced = query_intervals(input);
572 let candidates = candidate_intervals(bytes);
573 let m = bytes.len() as u32 - 1;
574
575 assert!(
576 is_valid_query_cover_chain(m, &produced),
577 "produced set is not a valid cover chain for {input:?}: {:?}",
578 produced
579 );
580 let optimum = min_cover_size_query_semantics(m, &candidates);
581 assert_eq!(
582 produced.len(),
583 optimum,
584 "produced set is not minimum-size query cover for {input:?}; produced={:?}; optimum={optimum}",
585 produced
586 );
587 }
588 }
589
590 #[test]
591 fn query_lardeee_diagnostic() {
592 let input = "lardeee";
593 let bytes = input.as_bytes();
594 let produced = query_intervals(input);
595 let candidates = candidate_intervals(bytes);
596 let m = bytes.len() as u32 - 1;
597
598 assert!(
599 is_valid_query_cover_chain(m, &produced),
600 "produced set is not a valid cover chain for {input:?}: {:?}",
601 produced
602 );
603
604 let optimum = min_cover_size_query_semantics(m, &candidates);
605
606 for iv in &produced {
609 assert!(
610 candidates
611 .iter()
612 .any(|c| c.start == iv.start && c.end == iv.end),
613 "produced interval not present in oracle candidates for {input:?}: {:?}; candidates={:?}",
614 iv,
615 candidates
616 );
617 }
618
619 assert_eq!(
621 produced.len(),
622 optimum,
623 "minimum-cover mismatch for {input:?}; produced={:?}; optimum={optimum}; candidates={:?}",
624 produced,
625 candidates
626 );
627 }
628
629 #[test]
630 fn query_sssdk_diagnostic() {
631 let input = "sssdk";
632 let bytes = input.as_bytes();
633 let produced = query_intervals(input);
634 let candidates = candidate_intervals(bytes);
635 let m = bytes.len() as u32 - 1;
636
637 assert!(
638 is_valid_query_cover_chain(m, &produced),
639 "produced set is not a valid cover chain for {input:?}: {:?}",
640 produced
641 );
642
643 let optimum = min_cover_size_query_semantics(m, &candidates);
644
645 for iv in &produced {
646 assert!(
647 candidates
648 .iter()
649 .any(|c| c.start == iv.start && c.end == iv.end),
650 "produced interval not present in oracle candidates for {input:?}: {:?}; candidates={:?}",
651 iv,
652 candidates
653 );
654 }
655
656 assert_eq!(
657 produced.len(),
658 optimum,
659 "minimum-cover mismatch for {input:?}; produced={:?}; optimum={optimum}; candidates={:?}",
660 produced,
661 candidates
662 );
663 }
664
665 #[test]
666 fn query_flush_is_minimum_cover_on_randomized_inputs() {
667 let mut rng = Rng64::new(0xA5A5_0123_89AB_CDEF);
668
669 for _ in 0..2000 {
670 let len = 3 + rng.gen_range(14); let mut bytes = vec![0u8; len];
672 for b in &mut bytes {
673 *b = b'a' + rng.gen_range(26) as u8; }
675 let input = std::str::from_utf8(&bytes).expect("ASCII should be valid UTF-8");
676
677 let produced = query_intervals(input);
678 let candidates = candidate_intervals(&bytes);
679 let m = bytes.len() as u32 - 1;
680
681 assert!(
682 is_valid_query_cover_chain(m, &produced),
683 "produced set is not a valid cover chain for randomized input {:?}: {:?}",
684 input,
685 produced
686 );
687 assert!(
688 produced.iter().all(|iv| {
689 let gram_len = (iv.end - iv.start + 2) as usize;
690 (2..=MAX_SPARSE_GRAM_SIZE).contains(&gram_len)
691 }),
692 "produced set contains out-of-range gram length for randomized input {:?}: {:?}",
693 input,
694 produced
695 );
696 let optimum = min_cover_size_query_semantics(m, &candidates);
697 assert_eq!(
698 produced.len(),
699 optimum,
700 "produced set is not minimum-size query cover for randomized input {:?}; produced={:?}; optimum={optimum}",
701 input,
702 produced
703 );
704 }
705 }
706
707 #[test]
708 fn query_consume_first_on_randomized_inputs() {
709 let mut rng = Rng64::new(0xC0DE_CAFE_1234_5678);
710
711 for _ in 0..2000 {
712 let len = 4 + rng.gen_range(13); let mut bytes = vec![0u8; len];
714 for b in &mut bytes {
715 *b = b'a' + rng.gen_range(26) as u8;
716 }
717 let input = std::str::from_utf8(&bytes).expect("ASCII should be valid UTF-8");
718
719 let split = 2 + rng.gen_range(len - 2);
720 let (prefix, suffix) = input.split_at(split);
721 let mut q = QueryGrams::default();
722
723 let mut first_half = Vec::new();
726 for c in prefix.chars() {
727 q.append_char(c, |gram, end, _follow| {
728 let begin = end + 1 - gram.len() as u32;
729 first_half.push(begin..end - 1);
730 });
731 }
732
733 while !q.queue.is_empty() {
734 q.consume_first(|gram, end, _follow| {
735 let begin = end + 1 - gram.len() as u32;
736 first_half.push(begin..end - 1);
737 });
738 }
739
740 assert!(first_half.iter().all(|iv| {
741 let gram_len = (iv.end - iv.start + 2) as usize;
742 (2..=MAX_SPARSE_GRAM_SIZE).contains(&gram_len)
743 }));
744
745 let prefix_bytes = prefix.as_bytes();
749 if prefix_bytes.len() >= 3 {
750 let prefix_m = prefix_bytes.len() as u32 - 1;
751 let prefix_candidates = candidate_intervals(prefix_bytes);
752 assert!(
753 is_valid_query_cover_chain(prefix_m, &first_half),
754 "first half is not a valid cover chain for prefix {:?}: {:?}",
755 prefix,
756 first_half
757 );
758 let prefix_optimum = min_cover_size_query_semantics(prefix_m, &prefix_candidates);
759 assert_eq!(
760 first_half.len(),
761 prefix_optimum,
762 "first half is not a minimum-size query cover for prefix {:?}; produced={:?}; optimum={prefix_optimum}",
763 prefix,
764 first_half
765 );
766 }
767
768 let retained = char::from((q.content & 0xFF) as u8);
769 let local_input = format!("{retained}{suffix}");
770
771 let mut remaining = Vec::new();
772 for c in suffix.chars() {
773 q.append_char(c, |gram, end, _follow| {
774 let begin = end + 1 - gram.len() as u32;
775 remaining.push(begin..end - 1);
776 });
777 }
778 q.flush(|gram, end, _follow| {
779 let begin = end + 1 - gram.len() as u32;
780 remaining.push(begin..end - 1);
781 });
782
783 assert_eq!(
784 remaining,
785 query_intervals(&local_input),
786 "post-drain continuation mismatch for randomized input {:?}; prefix={:?}; suffix={:?}; first_half={:?}; remaining={:?}; local_input={:?}",
787 input,
788 prefix,
789 suffix,
790 first_half,
791 remaining,
792 local_input
793 );
794
795 let local_bytes = local_input.as_bytes();
797 if local_bytes.len() >= 3 {
798 let local_m = local_bytes.len() as u32 - 1;
799 let local_candidates = candidate_intervals(local_bytes);
800 assert!(
801 is_valid_query_cover_chain(local_m, &remaining),
802 "second half is not a valid cover chain for {:?}: {:?}",
803 local_input,
804 remaining
805 );
806 let local_optimum = min_cover_size_query_semantics(local_m, &local_candidates);
807 assert_eq!(
808 remaining.len(),
809 local_optimum,
810 "second half is not a minimum-size query cover for {:?}; produced={:?}; optimum={local_optimum}",
811 local_input,
812 remaining
813 );
814 }
815 }
816 }
817
818 #[test]
824 fn consume_first_converges_to_default() {
825 for input in ["abc", "abcdef", "hello world", "ababababab", "mississippi"] {
826 let mut q = QueryGrams::default();
827 for c in input.chars() {
828 q.append_char(c, |_gram, _idx, _follow| {});
829 }
830 for _ in 0..(input.len() + 8) {
833 q.consume_first(|_gram, _idx, _follow| {});
834 }
835 assert_eq!(
836 q.state(),
837 QueryGrams::default().state(),
838 "consume_first did not converge to the default state for input {input:?}",
839 );
840 }
841 }
842}