metrale_speculative/
ngram.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Prompt-lookup draft proposer: finds the longest suffix of the
4//! token history (prompt and generated tokens) that occurred earlier, and
5//! proposes the token that followed it. Host code only.
6//!
7//! Owner: speculative.
8//! Invariants: none beyond the types.
9
10/// 2026-09-25: The proposer. It keeps no index: `propose` searches the
11/// history it is given.
12pub struct NgramProposer {
13    /// 2026-09-25: Shortest suffix length tried.
14    min_match: usize,
15    /// 2026-09-25: Longest suffix length tried.
16    max_match: usize,
17    /// 2026-09-25: Counts added to by `record_accept` and `record_reject`.
18    pub accepts: u64,
19    pub rejects: u64,
20}
21
22#[allow(clippy::len_without_is_empty)]
23impl NgramProposer {
24    pub fn new(_order: usize) -> Self {
25        Self {
26            min_match: 2,
27            max_match: 16,
28            accepts: 0,
29            rejects: 0,
30        }
31    }
32
33    /// 2026-09-25: Counts an accepted draft in `accepts`.
34    pub fn record_accept(&mut self) {
35        self.accepts += 1;
36    }
37
38    /// 2026-09-25: Counts a rejected draft in `rejects`.
39    pub fn record_reject(&mut self) {
40        self.rejects += 1;
41    }
42
43    /// 2026-09-25: For the longest suffix of `all_tokens` (from `min_match`
44    /// to `max_match` tokens) that also occurs earlier, returns the token after
45    /// its earliest occurrence; `None` when no suffix repeats.
46    pub fn propose(&self, all_tokens: &[u32]) -> Option<u32> {
47        let len = all_tokens.len();
48        if len < self.min_match + 1 {
49            return None;
50        }
51
52        let max_n = self.max_match.min(len - 1);
53        let mut best_next: Option<u32> = None;
54        let mut best_match_len: usize = 0;
55
56        // 2026-09-25: Suffix = all_tokens[len-n..len], searched at positions
57        // 0..=len-n-1 so that a token follows the match.
58        for n in self.min_match..=max_n {
59            let suffix = &all_tokens[len - n..len];
60            for start in 0..=(len - n - 1) {
61                if all_tokens[start..start + n] == *suffix && n > best_match_len {
62                    best_match_len = n;
63                    best_next = Some(all_tokens[start + n]);
64                    break;
65                }
66            }
67        }
68
69        best_next
70    }
71
72    /// 2026-09-25: No-op: the proposer keeps no index.
73    pub fn observe(&mut self, _history: &[u32], _next: u32) {}
74
75    pub fn len(&self) -> usize {
76        0
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn test_prompt_lookup_basic() {
86        let p = NgramProposer::new(4);
87
88        // 2026-09-25: The suffix [1, 2, 3] occurs at position 0, followed by 4.
89        let tokens = vec![1, 2, 3, 4, 5, 1, 2, 3];
90        assert_eq!(p.propose(&tokens), Some(4));
91    }
92
93    #[test]
94    fn test_prompt_lookup_no_match() {
95        let p = NgramProposer::new(4);
96
97        let tokens = vec![1, 2, 3, 4, 5, 6, 7, 8];
98        assert_eq!(p.propose(&tokens), None);
99    }
100
101    #[test]
102    fn test_prompt_lookup_short() {
103        let p = NgramProposer::new(4);
104
105        let tokens = vec![1, 2];
106        assert_eq!(p.propose(&tokens), None);
107    }
108
109    #[test]
110    fn test_prompt_lookup_repetitive() {
111        let p = NgramProposer::new(4);
112
113        // 2026-09-25: [A, B, C, A, B, C, A, B]: the longest repeated suffix is
114        // [A, B, C, A, B], at position 0, followed by C.
115        let tokens = vec![10, 20, 30, 10, 20, 30, 10, 20];
116        assert_eq!(p.propose(&tokens), Some(30));
117    }
118}
metrale_speculative/ngram.rs

metrale_speculative/
ngram.rs
