Skip to main content

proxima_core/
accept_rule.rs

1//! The accept decision for speculative rows as loadable data: [`AcceptRule::accepted_rows`] is the pure decision, and its count is the `accepted` argument of [`crate::serving_state::ServingState::accept_rows`].
2//!
3//! A similarity floor of 1.0 under exact equality, minimum run 0 and no cap reproduce [`crate::serving_state::ServingState::accept`].
4//!
5//! Under the `config` feature the rule loads from a builder, a TOML table and environment variables.
6
7/// Three scalars that decide how many leading drafted rows are accepted; copyable so it can ride in a config.
8#[derive(Debug, Clone, Copy, PartialEq)]
9#[cfg_attr(
10    feature = "config",
11    derive(
12        bon::Builder,
13        serde::Deserialize,
14        serde::Serialize,
15        conflaguration::Settings,
16        conflaguration::Validate
17    )
18)]
19#[cfg_attr(feature = "config", settings(prefix = "PROXIMA_ACCEPT_RULE"))]
20#[cfg_attr(feature = "config", builder(derive(Clone, Debug)))]
21pub struct AcceptRule {
22    /// A drafted row counts as accepted when its similarity to the verifier's row is at least this; 1.0 admits only rows the caller's similarity scores 1.0.
23    #[cfg_attr(feature = "config", setting(default = 1.0))]
24    #[cfg_attr(feature = "config", builder(default = 1.0))]
25    pub similarity_floor: f32,
26    /// A leading accepted run shorter than this is discarded, so the count is 0; 3 for a rule that only commits runs of three.
27    #[cfg_attr(feature = "config", setting(default = 0))]
28    #[cfg_attr(feature = "config", builder(default = 0))]
29    pub min_run: u16,
30    /// The leading run is cut at this many rows; 8.
31    #[cfg_attr(feature = "config", setting(default = 65535))]
32    #[cfg_attr(feature = "config", builder(default = u16::MAX))]
33    pub max_rows: u16,
34}
35
36impl Default for AcceptRule {
37    fn default() -> Self {
38        Self {
39            similarity_floor: 1.0,
40            min_run: 0,
41            max_rows: u16::MAX,
42        }
43    }
44}
45
46impl AcceptRule {
47    /// Count of leading rows accepted; the caller is [`crate::serving_state::ServingState::accept_rows`]. The similarity is an argument because the entry type is generic, so only the caller can score two entries.
48    #[must_use]
49    pub fn accepted_rows<Entry>(
50        &self,
51        draft: &[Entry],
52        choices: &[Entry],
53        similarity: impl Fn(&Entry, &Entry) -> f32,
54    ) -> usize {
55        let run = draft
56            .iter()
57            .zip(choices)
58            .take(usize::from(self.max_rows))
59            .take_while(|(drafted, chosen)| similarity(drafted, chosen) >= self.similarity_floor)
60            .count();
61        if run < usize::from(self.min_run) {
62            0
63        } else {
64            run
65        }
66    }
67}
68
69#[cfg(test)]
70// a failed transition in a test is a broken test; expect names it
71#[allow(clippy::unwrap_used, clippy::expect_used)]
72mod tests {
73    use super::*;
74    use crate::serving_state::ServingState;
75
76    fn exact(drafted: &u32, chosen: &u32) -> f32 {
77        if drafted == chosen { 1.0 } else { 0.0 }
78    }
79
80    fn rule(similarity_floor: f32, min_run: u16, max_rows: u16) -> AcceptRule {
81        AcceptRule {
82            similarity_floor,
83            min_run,
84            max_rows,
85        }
86    }
87
88    #[test]
89    fn accept_rule_default_matches_leading_equal_accept() {
90        let cases = [
91            ([3_u32, 4, 5], [3_u32, 4, 5]),
92            ([3, 4, 5], [3, 9, 5]),
93            ([3, 4, 5], [8, 4, 5]),
94        ];
95        let results: alloc::vec::Vec<_> = cases
96            .iter()
97            .map(|(draft, choices)| {
98                let state = ServingState::start(alloc::vec![1_u32], 0_usize)
99                    .advance_prefill(2, 1)
100                    .expect("prefill advances")
101                    .enter_verify(draft.to_vec())
102                    .expect("verify enters");
103                let caches = alloc::vec![10_usize, 11, 12];
104                let direct = state.clone().accept(choices, caches.clone());
105                let ruled = state.accept_rows(
106                    AcceptRule::default().accepted_rows(draft, choices, exact),
107                    choices,
108                    caches,
109                );
110                assert_eq!(direct, ruled);
111                ruled
112            })
113            .collect();
114
115        assert_eq!(
116            results,
117            alloc::vec![
118                Ok(ServingState::Accept {
119                    n: 3,
120                    next: 5,
121                    cache: 12
122                }),
123                Ok(ServingState::Rollback {
124                    snapshot: 11,
125                    to: 9
126                }),
127                Ok(ServingState::Rollback {
128                    snapshot: 10,
129                    to: 8
130                }),
131            ]
132        );
133    }
134
135    #[test]
136    fn accept_rule_similarity_floor_admits_near_matches() {
137        let draft = [10_u8, 20, 30];
138        let choices = [10_u8, 21, 35];
139        let similarity = |a: &u8, b: &u8| 1.0 - f32::from(a.abs_diff(*b)) / 10.0;
140
141        assert_eq!(
142            rule(1.0, 0, u16::MAX).accepted_rows(&draft, &choices, similarity),
143            1
144        );
145        assert_eq!(
146            rule(0.85, 0, u16::MAX).accepted_rows(&draft, &choices, similarity),
147            2
148        );
149        assert_eq!(
150            rule(0.4, 0, u16::MAX).accepted_rows(&draft, &choices, similarity),
151            3
152        );
153    }
154
155    #[test]
156    fn accept_rule_row_cap_limits_the_accepted_run() {
157        let rows = [1_u32, 2, 3, 4, 5];
158
159        assert_eq!(rule(1.0, 0, 3).accepted_rows(&rows, &rows, exact), 3);
160        assert_eq!(rule(1.0, 0, 9).accepted_rows(&rows, &rows, exact), 5);
161        assert_eq!(rule(1.0, 0, 0).accepted_rows(&rows, &rows, exact), 0);
162    }
163
164    #[test]
165    fn accept_rule_minimum_run_refuses_short_runs() {
166        let draft = [1_u32, 2, 3, 9];
167        let choices = [1_u32, 2, 3, 4];
168
169        assert_eq!(
170            rule(1.0, 0, u16::MAX).accepted_rows(&draft, &choices, exact),
171            3
172        );
173        assert_eq!(
174            rule(1.0, 3, u16::MAX).accepted_rows(&draft, &choices, exact),
175            3
176        );
177        assert_eq!(
178            rule(1.0, 4, u16::MAX).accepted_rows(&draft, &choices, exact),
179            0
180        );
181        assert_eq!(rule(1.0, 3, 2).accepted_rows(&draft, &choices, exact), 0);
182    }
183}
184
185#[cfg(all(test, feature = "config"))]
186// a failed load in a test is a broken test; expect names it
187#[allow(clippy::unwrap_used, clippy::expect_used)]
188mod config_tests {
189    use super::*;
190    use crate::serving_state::action_fixture::{Action, run_sequential, run_speculative};
191    use conflaguration::Settings;
192
193    const FLOOR_VAR: &str = "PROXIMA_ACCEPT_RULE_SIMILARITY_FLOOR";
194    const RUN_VAR: &str = "PROXIMA_ACCEPT_RULE_MIN_RUN";
195    const ROWS_VAR: &str = "PROXIMA_ACCEPT_RULE_MAX_ROWS";
196
197    fn field_agreement(drafted: &Action, chosen: &Action) -> f32 {
198        f32::from(
199            u8::from(drafted.key == chosen.key)
200                + u8::from(drafted.delta == chosen.delta)
201                + u8::from(drafted.version == chosen.version),
202        ) / 3.0
203    }
204
205    fn rule_from(table: &str) -> AcceptRule {
206        conflaguration::from_toml_str(table).expect("the rule table parses")
207    }
208
209    fn assert_rule_matches_sequential(table: &str, counts: (usize, usize, usize)) {
210        let rule = rule_from(table);
211        let sequential = run_sequential();
212        let speculative =
213            run_speculative(|draft, choices| rule.accepted_rows(draft, choices, field_agreement));
214
215        assert_eq!(speculative.history, sequential.history);
216        assert_eq!(speculative.environment, sequential.environment);
217        assert_eq!(
218            (
219                speculative.rounds,
220                speculative.rejected_rounds,
221                speculative.accepted_rows
222            ),
223            counts
224        );
225    }
226
227    fn assert_loaders_agree(expected: &AcceptRule, floor: &str, run: &str, rows: &str) {
228        let table = format!("similarity_floor = {floor}\nmin_run = {run}\nmax_rows = {rows}\n");
229        let from_toml: AcceptRule =
230            conflaguration::from_toml_str(&table).expect("the rule table parses");
231        assert_eq!(&from_toml, expected);
232
233        temp_env::with_vars(
234            [
235                (FLOOR_VAR, Some(floor)),
236                (RUN_VAR, Some(run)),
237                (ROWS_VAR, Some(rows)),
238            ],
239            || {
240                let from_env = AcceptRule::from_env().expect("the rule loads from env");
241                assert_eq!(&from_env, expected);
242            },
243        );
244    }
245
246    #[test]
247    fn accept_rule_builder_defaults_match_default_impl() {
248        assert_eq!(AcceptRule::builder().build(), AcceptRule::default());
249
250        temp_env::with_vars(
251            [(FLOOR_VAR, None::<&str>), (RUN_VAR, None), (ROWS_VAR, None)],
252            || {
253                let from_env = AcceptRule::from_env().expect("defaults load from an empty env");
254                assert_eq!(from_env, AcceptRule::default());
255            },
256        );
257    }
258
259    #[test]
260    fn accept_rule_round_trip_similarity_floor() {
261        let rule = AcceptRule::builder().similarity_floor(0.75).build();
262
263        assert_loaders_agree(&rule, "0.75", "0", "65535");
264    }
265
266    #[test]
267    fn accept_rule_round_trip_run_and_cap() {
268        let rule = AcceptRule::builder().min_run(3).max_rows(6).build();
269
270        assert_loaders_agree(&rule, "1.0", "3", "6");
271    }
272
273    #[test]
274    fn accept_rule_rejects_malformed_values() {
275        let bad_floor = conflaguration::from_toml_str::<AcceptRule>(
276            "similarity_floor = \"high\"\nmin_run = 0\nmax_rows = 8\n",
277        );
278        let bad_rows = conflaguration::from_toml_str::<AcceptRule>(
279            "similarity_floor = 1.0\nmin_run = 0\nmax_rows = 70000\n",
280        );
281        assert!(bad_floor.is_err());
282        assert!(bad_rows.is_err());
283
284        temp_env::with_vars([(ROWS_VAR, Some("lots"))], || {
285            assert!(AcceptRule::from_env().is_err());
286        });
287    }
288
289    #[test]
290    fn configured_rule_equality_matches_sequential() {
291        assert_rule_matches_sequential(
292            "similarity_floor = 1.0\nmin_run = 0\nmax_rows = 8\n",
293            (4, 4, 8),
294        );
295    }
296
297    #[test]
298    fn configured_rule_verifier_exact_matches_sequential() {
299        assert_rule_matches_sequential(
300            "similarity_floor = 1.0\nmin_run = 0\nmax_rows = 65535\n",
301            (4, 4, 8),
302        );
303    }
304
305    #[test]
306    fn configured_rule_anchor_run_matches_sequential() {
307        assert_rule_matches_sequential(
308            "similarity_floor = 1.0\nmin_run = 3\nmax_rows = 65535\n",
309            (6, 6, 6),
310        );
311    }
312}