long_screen/engine/
keyframes.rs

1//! The solve pass's keyframe index: which earlier keyframes a frame might be revisiting, by visual words.
2//!
3//! The index keeps only word postings in memory (per layer, per word, the keyframes that saw it, one per frame in
4//! frame order), as the TypeScript `KeyframeIndex` it replaces did. The keyframes themselves are rows the host
5//! stores (`keyframe/<id>`) and reads back. A revisit search is declare-then-provide: [`KeyframeIndex::retrieve`]
6//! names up to [`MAX_CANDIDATES`] candidate keyframes, the host loads those rows, keeps the ones
7//! [`Retrieval::admits`], and scores them with the kernels' candidate evaluation.
8
9use crate::kernels::features::{feature_words, DESCRIPTOR_WORDS};
10use std::collections::{HashMap, HashSet};
11
12/// A query needs at least this many features to search at all.
13pub const MIN_FEATURES: usize = 8;
14/// Words sampled from a query, spread evenly over all of its words.
15pub const QUERY_WORDS: usize = 48;
16/// Keyframes taken from each end of one word's posting.
17pub const POSTING_SPREAD: usize = 48;
18/// A posting longer than this overran the retrieval budget: [`BUDGET_DIAGNOSTIC`] is due for the layer.
19pub const RETRIEVAL_BUDGET: usize = 96;
20/// Candidates handed to evaluation, by vote count.
21pub const MAX_CANDIDATES: usize = 12;
22/// Votes a candidate needs.
23pub const MIN_VOTES: u32 = 3;
24/// Keyframes closer than this many frames to the query are odometry, not revisits, unless the query says otherwise.
25pub const DEFAULT_MIN_GAP: u32 = 3;
26/// The diagnostic code (severity warning) the host emits when [`Retrieval::budget_warning`] is set.
27pub const BUDGET_DIAGNOSTIC: &str = "RELOCALIZATION_BUDGET";
28
29/// One word's postings in one layer: interned keyframe ids in ascending frame order, one per frame.
30#[derive(Default)]
31struct Posting {
32    frames: Vec<u32>,
33    ids: Vec<u32>,
34}
35
36/// Visual words of the even-indexed descriptors (every other feature), as the index has always hashed them.
37fn words(descriptors: &[[u32; DESCRIPTOR_WORDS]]) -> Vec<u32> {
38    let even: Vec<[u32; DESCRIPTOR_WORDS]> = descriptors.iter().step_by(2).copied().collect();
39    feature_words(&even)
40}
41
42/// The outcome of [`KeyframeIndex::retrieve`].
43#[derive(Clone, Debug, PartialEq)]
44pub struct Retrieval {
45    /// Keyframe ids to load and evaluate, most votes first (ties in the order they first got a vote).
46    pub candidates: Vec<String>,
47    /// A posting overran [`RETRIEVAL_BUDGET`] and this layer has not been warned yet: emit [`BUDGET_DIAGNOSTIC`].
48    pub budget_warning: bool,
49    frame: u32,
50    min_gap: u32,
51}
52
53impl Retrieval {
54    /// Whether a loaded candidate keyframe at `keyframe_frame` is far enough from the query frame to be a revisit.
55    pub fn admits(&self, keyframe_frame: u32) -> bool {
56        keyframe_frame.abs_diff(self.frame) >= self.min_gap
57    }
58}
59
60#[derive(Default)]
61pub struct KeyframeIndex {
62    ids: Vec<String>,
63    interned: HashMap<String, u32>,
64    postings: HashMap<String, HashMap<u32, Posting>>,
65    warned_layers: HashSet<String>,
66}
67
68impl KeyframeIndex {
69    pub fn new() -> Self {
70        Self::default()
71    }
72
73    fn intern(&mut self, id: &str) -> u32 {
74        if let Some(&i) = self.interned.get(id) {
75            return i;
76        }
77        let i = self.ids.len() as u32;
78        self.ids.push(id.to_string());
79        self.interned.insert(id.to_string(), i);
80        i
81    }
82
83    /// Posts keyframe `id` (of `layer`, at `frame`) under the words of its features' descriptors. A keyframe at a
84    /// frame that already has one in a posting replaces it there.
85    pub fn add(
86        &mut self,
87        id: &str,
88        layer: &str,
89        frame: u32,
90        descriptors: &[[u32; DESCRIPTOR_WORDS]],
91    ) {
92        let interned = self.intern(id);
93        let postings = self.postings.entry(layer.to_string()).or_default();
94        for word in words(descriptors) {
95            let Posting { frames, ids } = postings.entry(word).or_default();
96            // Keyframes arrive in frame order, so this is almost always an append.
97            let mut at = frames.len();
98            while at > 0 && frames[at - 1] > frame {
99                at -= 1;
100            }
101            if at > 0 && frames[at - 1] == frame {
102                ids[at - 1] = interned;
103            } else {
104                frames.insert(at, frame);
105                ids.insert(at, interned);
106            }
107        }
108    }
109
110    /// Candidate keyframes of `layer` for a query at `frame` with these feature descriptors, by word votes. Each
111    /// sampled word votes for the earliest [`POSTING_SPREAD`] keyframes in its posting and, when the posting is
112    /// longer, the latest ones too (so recent keyframes still get found on a long recording); `exclude` never gets
113    /// a vote. `min_gap` (default [`DEFAULT_MIN_GAP`]) is what [`Retrieval::admits`] checks.
114    pub fn retrieve(
115        &mut self,
116        layer: &str,
117        frame: u32,
118        descriptors: &[[u32; DESCRIPTOR_WORDS]],
119        exclude: Option<&str>,
120        min_gap: Option<u32>,
121    ) -> Retrieval {
122        let mut retrieval = Retrieval {
123            candidates: Vec::new(),
124            budget_warning: false,
125            frame,
126            min_gap: min_gap.unwrap_or(DEFAULT_MIN_GAP),
127        };
128        if descriptors.len() < MIN_FEATURES {
129            return retrieval;
130        }
131        let all = words(descriptors);
132        let step = (all.len() / QUERY_WORDS).max(1);
133        let sampled = all.iter().step_by(step).take(QUERY_WORDS);
134        let exclude = exclude.and_then(|id| self.interned.get(id).copied());
135        // Insertion-ordered tally: the sort below is stable, so equal counts keep first-vote order.
136        let mut votes: Vec<(u32, u32)> = Vec::new();
137        let mut slot: HashMap<u32, usize> = HashMap::new();
138        let mut vote = |id: u32| match slot.get(&id) {
139            Some(&i) => votes[i].1 += 1,
140            None => {
141                slot.insert(id, votes.len());
142                votes.push((id, 1));
143            }
144        };
145        let mut truncated = false;
146        let postings = self.postings.get(layer);
147        for word in sampled {
148            let ids = postings
149                .and_then(|p| p.get(word))
150                .map_or(&[][..], |p| p.ids.as_slice());
151            if ids.len() > RETRIEVAL_BUDGET {
152                truncated = true;
153            }
154            let mut seen = HashSet::new();
155            for &id in ids.iter().take(POSTING_SPREAD) {
156                if Some(id) == exclude {
157                    continue;
158                }
159                seen.insert(id);
160                vote(id);
161            }
162            if ids.len() > POSTING_SPREAD {
163                for &id in ids[ids.len() - POSTING_SPREAD..].iter().rev() {
164                    if Some(id) == exclude || !seen.insert(id) {
165                        continue;
166                    }
167                    vote(id);
168                }
169            }
170        }
171        if truncated && self.warned_layers.insert(layer.to_string()) {
172            retrieval.budget_warning = true;
173        }
174        votes.sort_by(|a, b| b.1.cmp(&a.1));
175        retrieval.candidates = votes
176            .into_iter()
177            .take(MAX_CANDIDATES)
178            .filter(|&(_, count)| count >= MIN_VOTES)
179            .map(|(id, _)| self.ids[id as usize].clone())
180            .collect();
181        retrieval
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    /// A descriptor whose first band hashes to word `w` (< 4096) and whose other bands hash to `w + 1000k`-ish
190    /// words of their own; `d[k*2] ^ (d[k*2+1] >> 8)` is the band's word, so `d[k*2+1] = 0` makes it `d[k*2]`.
191    fn descriptor(w: u32) -> [u32; DESCRIPTOR_WORDS] {
192        let mut d = [0u32; DESCRIPTOR_WORDS];
193        for k in 0..4 {
194            d[k * 2] = (w + 1000 * k as u32) & 4095;
195        }
196        d
197    }
198
199    /// Features whose even-indexed descriptors carry `words` (odd ones are noise the index ignores).
200    fn features(words: &[u32]) -> Vec<[u32; DESCRIPTOR_WORDS]> {
201        words
202            .iter()
203            .flat_map(|&w| [descriptor(w), descriptor(4000)])
204            .collect()
205    }
206
207    #[test]
208    fn a_query_with_too_few_features_finds_nothing_and_never_warns() {
209        let mut index = KeyframeIndex::new();
210        index.add("L/1", "L", 1, &features(&[1, 2, 3, 4]));
211        let r = index.retrieve("L", 50, &features(&[1, 2, 3, 4])[..7], None, None);
212        assert_eq!(r.candidates, Vec::<String>::new());
213        assert!(!r.budget_warning);
214    }
215
216    #[test]
217    fn a_keyframe_needs_three_votes_and_only_its_own_layer_counts() {
218        let mut index = KeyframeIndex::new();
219        index.add("L/1", "L", 1, &features(&[1, 2, 3, 4]));
220        index.add("L/2", "L", 2, &features(&[1, 2, 90, 91]));
221        index.add("M/3", "M", 3, &features(&[1, 2, 3, 4]));
222        let r = index.retrieve("L", 50, &features(&[1, 2, 3, 4]), None, None);
223        // Each word's four bands vote: L/1 shares all bands of 1..4, L/2 those of 1 and 2.
224        assert_eq!(r.candidates, ["L/1", "L/2"]);
225        let none = index.retrieve("L", 50, &features(&[200, 201, 202, 203]), None, None);
226        assert!(none.candidates.is_empty());
227    }
228
229    #[test]
230    fn exclude_gets_no_vote() {
231        let mut index = KeyframeIndex::new();
232        index.add("L/1", "L", 1, &features(&[1, 2, 3, 4]));
233        index.add("L/2", "L", 2, &features(&[1, 2, 3, 4]));
234        let r = index.retrieve("L", 50, &features(&[1, 2, 3, 4]), Some("L/1"), None);
235        assert_eq!(r.candidates, ["L/2"]);
236    }
237
238    #[test]
239    fn ties_keep_first_vote_order_and_at_most_twelve_come_back() {
240        let mut index = KeyframeIndex::new();
241        for f in (0..20u32).rev() {
242            // Added out of order on purpose: postings stay in frame order regardless.
243            index.add(&format!("L/{f:02}"), "L", f, &features(&[1, 2, 3, 4]));
244        }
245        let r = index.retrieve("L", 500, &features(&[1, 2, 3, 4]), None, None);
246        let expected: Vec<String> = (0..12).map(|f| format!("L/{f:02}")).collect();
247        assert_eq!(r.candidates, expected);
248    }
249
250    /// The search reads a long posting as its earliest and its latest keyframes, so a posting must stay in frame
251    /// order. Failure case: a keyframe for an earlier frame arriving after later ones (two layers' solve steps
252    /// interleave, or a region re-minted from an older frame) appended at the end, taken for one of the latest.
253    #[test]
254    fn postings_stay_in_frame_order_whatever_order_keyframes_arrive_in() {
255        let mut index = KeyframeIndex::new();
256        for frame in [30u32, 10, 50, 20] {
257            index.add(
258                &format!("body/{frame}"),
259                "body",
260                frame,
261                &features(&[1, 2, 3, 4]),
262            );
263        }
264        let postings = &index.postings["body"];
265        assert!(!postings.is_empty());
266        for Posting { frames, ids } in postings.values() {
267            assert_eq!(frames, &[10, 20, 30, 50]);
268            let named: Vec<&str> = ids
269                .iter()
270                .map(|&i| index.ids[i as usize].as_str())
271                .collect();
272            assert_eq!(named, ["body/10", "body/20", "body/30", "body/50"]);
273        }
274    }
275
276    #[test]
277    fn a_later_keyframe_at_the_same_frame_replaces_the_earlier_one() {
278        let mut index = KeyframeIndex::new();
279        index.add("L/a", "L", 7, &features(&[1, 2, 3, 4]));
280        index.add("L/b", "L", 7, &features(&[1, 2, 3, 4]));
281        let r = index.retrieve("L", 50, &features(&[1, 2, 3, 4]), None, None);
282        assert_eq!(r.candidates, ["L/b"]);
283    }
284
285    #[test]
286    fn a_long_posting_votes_for_its_latest_keyframes_too_and_warns_once_per_layer() {
287        let mut index = KeyframeIndex::new();
288        for f in 0..120u32 {
289            index.add(&format!("L/{f:03}"), "L", f, &features(&[1, 2, 3, 4]));
290            index.add(&format!("M/{f:03}"), "M", f, &features(&[1, 2, 3, 4]));
291        }
292        let first = index.retrieve("L", 500, &features(&[1, 2, 3, 4]), None, None);
293        assert!(
294            first.budget_warning,
295            "a 120-long posting overran the budget of 96"
296        );
297        // Earliest and latest get the same votes; earliest voted first.
298        assert_eq!(first.candidates.len(), 12);
299        assert_eq!(first.candidates[0], "L/000");
300        let again = index.retrieve("L", 500, &features(&[1, 2, 3, 4]), None, None);
301        assert!(!again.budget_warning, "warned once per layer");
302        assert!(
303            index
304                .retrieve("M", 500, &features(&[1, 2, 3, 4]), None, None)
305                .budget_warning
306        );
307        // Only the latest end of a posting, past the first 48: still found.
308        let latest = index.retrieve("L", 500, &features(&[1, 2, 3, 4]), Some("L/000"), None);
309        assert!(!latest.candidates.contains(&"L/000".to_string()));
310    }
311
312    #[test]
313    fn the_latest_keyframes_of_a_long_posting_are_found() {
314        let mut index = KeyframeIndex::new();
315        for f in 0..60u32 {
316            index.add(&format!("L/{f:03}"), "L", f, &features(&[1, 2, 3, 4]));
317        }
318        // The latest keyframe also carries extra words, so it out-votes the earliest ones.
319        index.add("L/999", "L", 999, &features(&[1, 2, 3, 4, 5, 6]));
320        let r = index.retrieve("L", 2000, &features(&[1, 2, 3, 4, 5, 6]), None, None);
321        assert_eq!(r.candidates[0], "L/999");
322        assert!(!r.budget_warning);
323    }
324
325    #[test]
326    fn admits_keeps_odometry_neighbours_out() {
327        let mut index = KeyframeIndex::new();
328        let r = index.retrieve("L", 10, &[], None, None);
329        assert!(!r.admits(8) && !r.admits(12) && r.admits(7) && r.admits(13));
330        let wide = index.retrieve("L", 10, &[], None, Some(6));
331        assert!(!wide.admits(5) && wide.admits(4) && wide.admits(16));
332    }
333}
long_screen/engine/keyframes.rs

long_screen/engine/
keyframes.rs
