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 {
171 if self.queue.is_empty() {
172 0
173 } else {
174 self.queue.front_value()
175 }
176 }
177
178 #[inline]
186 fn follow_byte(&self, end: u32) -> Option<u8> {
187 if end >= self.content_end_idx {
188 None
189 } else {
190 let shift = self.content_end_idx - 1 - end;
191 debug_assert!(shift < MAX_SPARSE_GRAM_SIZE as u32);
192 Some((self.content >> (shift * 8)) as u8)
193 }
194 }
195
196 fn extract_gram<F>(&mut self, begin_index: u32, end_index: u32, consumer: &mut F)
197 where
198 F: FnMut(NGram, u32, Option<u8>),
199 {
200 debug_assert!(end_index >= begin_index);
201 debug_assert!(end_index <= self.content_end_idx);
202
203 let len = (end_index - begin_index + 2) as usize;
204 let dist = self.content_end_idx - end_index;
205 let shifted = self.content >> (dist * 8);
206 let follow = self.follow_byte(end_index);
209 consumer(NGram::from_window_masked(shifted, len), end_index, follow);
210 }
211
212 pub fn append_char<F>(&mut self, c: char, consumer: F)
216 where
217 F: FnMut(NGram, u32, Option<u8>),
218 {
219 self.append_byte(index_fold_char(c), consumer);
220 }
221
222 pub fn append_byte<F>(&mut self, right: u8, mut consumer: F)
228 where
229 F: FnMut(NGram, u32, Option<u8>),
230 {
231 let left = (self.content & 0xFF) as u8;
232 self.content_end_idx += 1;
233 self.content = (self.content << 8) | right as u64;
234
235 let idx = self.content_end_idx;
236
237 if idx == 1 {
240 self.h = bigram_h(right);
241 } else {
242 let (value, h_b) = bigram_priority_rolling(left, right, self.h);
243 self.h = h_b;
244 if !self.queue.is_empty()
245 && let priority = self.queue.front_value()
246 && value < priority
247 {
248 let mut begin = self.queue.pop_front_idx();
249 while !self.queue.is_empty() && self.queue.front_value() == priority {
250 let end = self.queue.pop_front_idx();
251 self.extract_gram(begin, end, &mut consumer);
252 begin = end;
253 }
254 self.queue.clear();
255 self.queue.push(PosState { index: idx, value });
256 self.extract_gram(begin, idx, &mut consumer);
257 } else {
258 self.queue.push(PosState { index: idx, value });
259 }
260 if idx - self.queue.front_idx() + 2 >= MAX_SPARSE_GRAM_SIZE as u32 && self.queue.len > 1
261 {
262 let begin = self.queue.pop_front_idx();
263 let end = self.queue.front_idx();
264 self.extract_gram(begin, end, &mut consumer);
265 }
266 }
267 }
268
269 pub fn flush<F>(mut self, mut consumer: F)
271 where
272 F: FnMut(NGram, u32, Option<u8>),
273 {
274 if self.content_end_idx == 2 {
275 self.extract_gram(2, 2, &mut consumer);
276 } else {
277 while self.queue.len > 1 {
278 let begin = self.queue.pop_front_idx();
279 let end = self.queue.front_idx();
280 self.extract_gram(begin, end, &mut consumer);
281 }
282 }
283 }
284
285 pub fn consume_first<F>(&mut self, mut consumer: F)
287 where
288 F: FnMut(NGram, u32, Option<u8>),
289 {
290 if self.queue.len > 1 {
291 let begin = self.queue.pop_front_idx();
293 let end = self.queue.front_idx();
294 self.extract_gram(begin, end, &mut consumer);
295 } else if self.content_end_idx == 2 {
296 self.extract_gram(2, 2, &mut consumer);
298 }
299
300 if self.queue.len <= 1 {
304 self.queue.clear();
305 if self.content_end_idx > 1 {
306 let last = (self.content & 0xFF) as u8;
308 self.content = last as u64;
309 self.content_end_idx = 1;
310 }
311 }
312 }
313}
314
315#[cfg(test)]
316mod tests {
317 use super::*;
318 use crate::collect_sparse_grams_deque;
319 use std::collections::hash_map::DefaultHasher;
320 use std::ops::Range;
321
322 fn query_intervals(input: &str) -> Vec<Range<u32>> {
328 let mut q = QueryGrams::default();
329 let mut out = Vec::new();
330 for c in input.chars() {
331 q.append_char(c, |gram, end, _follow| {
332 let begin = end + 1 - gram.len() as u32;
333 out.push(begin..end - 1);
334 });
335 }
336 q.flush(|gram, end, _follow| {
337 let begin = end + 1 - gram.len() as u32;
338 out.push(begin..end - 1);
339 });
340 out
341 }
342
343 fn candidate_intervals(bytes: &[u8]) -> Vec<Range<u32>> {
348 let mut out = Vec::new();
349 collect_sparse_grams_deque(bytes, |gram, idx| {
350 let begin = idx + 1 - gram.len() as u32;
351 out.push(begin..idx - 1);
352 });
353 out
354 }
355
356 fn min_cover_size_query_semantics(m: u32, intervals: &[Range<u32>]) -> usize {
357 let inf = usize::MAX / 4;
358 let mut dp = vec![0usize; m as usize + 1];
359 for pos in (1..m).rev() {
360 let mut best = inf;
361 for iv in intervals {
362 if iv.start <= pos && iv.end > pos {
363 let next = iv.end.min(m) as usize;
364 best = best.min(1 + dp[next]);
365 }
366 }
367 dp[pos as usize] = best;
368 }
369 dp[1]
370 }
371
372 fn is_valid_query_cover_chain(m: u32, intervals: &[Range<u32>]) -> bool {
373 if m <= 1 {
374 return true;
375 }
376 let mut produced = 1u32;
377 for iv in intervals {
378 if iv.start > produced || iv.end <= produced {
379 return false;
380 }
381 produced = iv.end;
382 if produced >= m {
383 return true;
384 }
385 }
386 false
387 }
388
389 #[derive(Clone, Copy, Debug)]
390 struct Rng64(u64);
391
392 impl Rng64 {
393 fn new(seed: u64) -> Self {
394 Self(seed)
395 }
396
397 fn next_u64(&mut self) -> u64 {
398 let mut x = self.0;
400 x ^= x >> 12;
401 x ^= x << 25;
402 x ^= x >> 27;
403 self.0 = x;
404 x.wrapping_mul(0x2545_F491_4F6C_DD1D)
405 }
406
407 fn gen_range(&mut self, upper: usize) -> usize {
408 (self.next_u64() % upper as u64) as usize
409 }
410 }
411
412 #[test]
413 fn append_and_flush_emit_grams() {
414 let mut q = QueryGrams::default();
415 for c in "hello world".chars() {
416 q.append_char(c, |_gram, _idx, _follow| {});
417 }
418 let mut out = Vec::new();
419 q.flush(|gram, idx, _follow| out.push((gram, idx)));
420 assert!(!out.is_empty());
421 assert!(
422 out.iter()
423 .all(|(g, _)| (2..=MAX_SPARSE_GRAM_SIZE).contains(&g.len()))
424 );
425 }
426
427 #[test]
428 fn single_bigram_has_no_follow() {
429 let bytes: Vec<u8> = "ab".chars().map(crate::index_fold_char).collect();
432 let mut q = QueryGrams::default();
433 for &b in &bytes {
434 q.append_byte(b, |_gram, _end, _follow| {});
435 }
436 let mut emitted = Vec::new();
437 q.flush(|gram, end, follow| emitted.push((gram.len(), end, follow)));
438 assert_eq!(emitted, vec![(2usize, 2u32, None)]);
439 }
440
441 #[test]
442 fn emitted_follow_is_the_actual_next_char() {
443 for input in [
446 "ab",
447 "abc",
448 "hello",
449 "querystream",
450 "aaaaaaaaaa",
451 "lardeee",
452 "sssdk",
453 ] {
454 let bytes: Vec<u8> = input.chars().map(crate::index_fold_char).collect();
455 let n = bytes.len() as u32;
456 let mut q = QueryGrams::default();
457 let mut check = |_gram: NGram, end: u32, follow: Option<u8>| {
458 if let Some(b) = follow {
459 assert!(
460 end < n,
461 "follow reported past input for {input:?}: end={end}"
462 );
463 assert_eq!(
464 b, bytes[end as usize],
465 "wrong follow byte for {input:?} at end={end}"
466 );
467 }
468 };
469 for &b in &bytes {
470 q.append_byte(b, &mut check);
471 }
472 q.flush(&mut check);
473 }
474 }
475
476 #[test]
477 fn state_tracks_active_window_each_step() {
478 let input = b"hello world";
487 let expected_len = [1u32, 2, 3, 2, 3, 2, 3, 2, 3, 4, 5];
489
490 let pack = |bytes: &[u8]| bytes.iter().fold(0u64, |acc, &b| (acc << 8) | b as u64);
491
492 let mut q = QueryGrams::default();
493 let mut prev = q.state();
494 assert_eq!(prev, (0, 0), "a fresh state must be empty");
495
496 for (i, &byte) in input.iter().enumerate() {
497 assert_eq!(
499 q.state(),
500 prev,
501 "state moved without an append before byte {i}"
502 );
503
504 q.append_byte(byte, |_, _, _| {});
505
506 let len = expected_len[i];
507 let window = &input[i + 1 - len as usize..=i];
508 let after = q.state();
509 assert_eq!(
510 after,
511 (len, pack(window)),
512 "unexpected state after byte {i} ({:?})",
513 char::from(byte),
514 );
515 prev = after;
516 }
517 }
518
519 #[test]
520 fn state_eq_and_hash_ignore_absolute_history() {
521 let mut a = QueryGrams::default();
522 let mut b = QueryGrams::default();
523
524 for c in "abc".chars() {
525 a.append_char(c, |_gram, _idx, _follow| {});
526 }
527 for c in "zabc".chars() {
528 b.append_char(c, |_gram, _idx, _follow| {});
529 }
530 b.consume_first(|_gram, _idx, _follow| {});
531
532 if a == b {
534 let mut ha = DefaultHasher::new();
535 let mut hb = DefaultHasher::new();
536 a.hash(&mut ha);
537 b.hash(&mut hb);
538 assert_eq!(ha.finish(), hb.finish());
539 }
540 }
541
542 #[test]
543 fn query_flush_is_minimum_cover_on_small_inputs() {
544 for input in [
545 "abc",
546 "abcd",
547 "abcdef",
548 "hello",
549 "hello world",
550 "ababababab",
551 ] {
552 let bytes = input.as_bytes();
553 if bytes.len() < 3 {
554 continue;
555 }
556 let produced = query_intervals(input);
557 let candidates = candidate_intervals(bytes);
558 let m = bytes.len() as u32 - 1;
559
560 assert!(
561 is_valid_query_cover_chain(m, &produced),
562 "produced set is not a valid cover chain for {input:?}: {:?}",
563 produced
564 );
565 let optimum = min_cover_size_query_semantics(m, &candidates);
566 assert_eq!(
567 produced.len(),
568 optimum,
569 "produced set is not minimum-size query cover for {input:?}; produced={:?}; optimum={optimum}",
570 produced
571 );
572 }
573 }
574
575 #[test]
576 fn query_lardeee_diagnostic() {
577 let input = "lardeee";
578 let bytes = input.as_bytes();
579 let produced = query_intervals(input);
580 let candidates = candidate_intervals(bytes);
581 let m = bytes.len() as u32 - 1;
582
583 assert!(
584 is_valid_query_cover_chain(m, &produced),
585 "produced set is not a valid cover chain for {input:?}: {:?}",
586 produced
587 );
588
589 let optimum = min_cover_size_query_semantics(m, &candidates);
590
591 for iv in &produced {
594 assert!(
595 candidates
596 .iter()
597 .any(|c| c.start == iv.start && c.end == iv.end),
598 "produced interval not present in oracle candidates for {input:?}: {:?}; candidates={:?}",
599 iv,
600 candidates
601 );
602 }
603
604 assert_eq!(
606 produced.len(),
607 optimum,
608 "minimum-cover mismatch for {input:?}; produced={:?}; optimum={optimum}; candidates={:?}",
609 produced,
610 candidates
611 );
612 }
613
614 #[test]
615 fn query_sssdk_diagnostic() {
616 let input = "sssdk";
617 let bytes = input.as_bytes();
618 let produced = query_intervals(input);
619 let candidates = candidate_intervals(bytes);
620 let m = bytes.len() as u32 - 1;
621
622 assert!(
623 is_valid_query_cover_chain(m, &produced),
624 "produced set is not a valid cover chain for {input:?}: {:?}",
625 produced
626 );
627
628 let optimum = min_cover_size_query_semantics(m, &candidates);
629
630 for iv in &produced {
631 assert!(
632 candidates
633 .iter()
634 .any(|c| c.start == iv.start && c.end == iv.end),
635 "produced interval not present in oracle candidates for {input:?}: {:?}; candidates={:?}",
636 iv,
637 candidates
638 );
639 }
640
641 assert_eq!(
642 produced.len(),
643 optimum,
644 "minimum-cover mismatch for {input:?}; produced={:?}; optimum={optimum}; candidates={:?}",
645 produced,
646 candidates
647 );
648 }
649
650 #[test]
651 fn query_flush_is_minimum_cover_on_randomized_inputs() {
652 let mut rng = Rng64::new(0xA5A5_0123_89AB_CDEF);
653
654 for _ in 0..2000 {
655 let len = 3 + rng.gen_range(14); let mut bytes = vec![0u8; len];
657 for b in &mut bytes {
658 *b = b'a' + rng.gen_range(26) as u8; }
660 let input = std::str::from_utf8(&bytes).expect("ASCII should be valid UTF-8");
661
662 let produced = query_intervals(input);
663 let candidates = candidate_intervals(&bytes);
664 let m = bytes.len() as u32 - 1;
665
666 assert!(
667 is_valid_query_cover_chain(m, &produced),
668 "produced set is not a valid cover chain for randomized input {:?}: {:?}",
669 input,
670 produced
671 );
672 assert!(
673 produced.iter().all(|iv| {
674 let gram_len = (iv.end - iv.start + 2) as usize;
675 (2..=MAX_SPARSE_GRAM_SIZE).contains(&gram_len)
676 }),
677 "produced set contains out-of-range gram length for randomized input {:?}: {:?}",
678 input,
679 produced
680 );
681 let optimum = min_cover_size_query_semantics(m, &candidates);
682 assert_eq!(
683 produced.len(),
684 optimum,
685 "produced set is not minimum-size query cover for randomized input {:?}; produced={:?}; optimum={optimum}",
686 input,
687 produced
688 );
689 }
690 }
691
692 #[test]
693 fn query_consume_first_on_randomized_inputs() {
694 let mut rng = Rng64::new(0xC0DE_CAFE_1234_5678);
695
696 for _ in 0..2000 {
697 let len = 4 + rng.gen_range(13); let mut bytes = vec![0u8; len];
699 for b in &mut bytes {
700 *b = b'a' + rng.gen_range(26) as u8;
701 }
702 let input = std::str::from_utf8(&bytes).expect("ASCII should be valid UTF-8");
703
704 let split = 2 + rng.gen_range(len - 2);
705 let (prefix, suffix) = input.split_at(split);
706 let mut q = QueryGrams::default();
707
708 let mut first_half = Vec::new();
711 for c in prefix.chars() {
712 q.append_char(c, |gram, end, _follow| {
713 let begin = end + 1 - gram.len() as u32;
714 first_half.push(begin..end - 1);
715 });
716 }
717
718 while !q.queue.is_empty() {
719 q.consume_first(|gram, end, _follow| {
720 let begin = end + 1 - gram.len() as u32;
721 first_half.push(begin..end - 1);
722 });
723 }
724
725 assert!(first_half.iter().all(|iv| {
726 let gram_len = (iv.end - iv.start + 2) as usize;
727 (2..=MAX_SPARSE_GRAM_SIZE).contains(&gram_len)
728 }));
729
730 let prefix_bytes = prefix.as_bytes();
734 if prefix_bytes.len() >= 3 {
735 let prefix_m = prefix_bytes.len() as u32 - 1;
736 let prefix_candidates = candidate_intervals(prefix_bytes);
737 assert!(
738 is_valid_query_cover_chain(prefix_m, &first_half),
739 "first half is not a valid cover chain for prefix {:?}: {:?}",
740 prefix,
741 first_half
742 );
743 let prefix_optimum = min_cover_size_query_semantics(prefix_m, &prefix_candidates);
744 assert_eq!(
745 first_half.len(),
746 prefix_optimum,
747 "first half is not a minimum-size query cover for prefix {:?}; produced={:?}; optimum={prefix_optimum}",
748 prefix,
749 first_half
750 );
751 }
752
753 let retained = char::from((q.content & 0xFF) as u8);
754 let local_input = format!("{retained}{suffix}");
755
756 let mut remaining = Vec::new();
757 for c in suffix.chars() {
758 q.append_char(c, |gram, end, _follow| {
759 let begin = end + 1 - gram.len() as u32;
760 remaining.push(begin..end - 1);
761 });
762 }
763 q.flush(|gram, end, _follow| {
764 let begin = end + 1 - gram.len() as u32;
765 remaining.push(begin..end - 1);
766 });
767
768 assert_eq!(
769 remaining,
770 query_intervals(&local_input),
771 "post-drain continuation mismatch for randomized input {:?}; prefix={:?}; suffix={:?}; first_half={:?}; remaining={:?}; local_input={:?}",
772 input,
773 prefix,
774 suffix,
775 first_half,
776 remaining,
777 local_input
778 );
779
780 let local_bytes = local_input.as_bytes();
782 if local_bytes.len() >= 3 {
783 let local_m = local_bytes.len() as u32 - 1;
784 let local_candidates = candidate_intervals(local_bytes);
785 assert!(
786 is_valid_query_cover_chain(local_m, &remaining),
787 "second half is not a valid cover chain for {:?}: {:?}",
788 local_input,
789 remaining
790 );
791 let local_optimum = min_cover_size_query_semantics(local_m, &local_candidates);
792 assert_eq!(
793 remaining.len(),
794 local_optimum,
795 "second half is not a minimum-size query cover for {:?}; produced={:?}; optimum={local_optimum}",
796 local_input,
797 remaining
798 );
799 }
800 }
801 }
802}