nodedb_fts/backend/
memory.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! In-memory FTS backend for Lite and WASM deployments.
4//!
5//! All data lives in HashMaps behind `RefCell` for interior mutability,
6//! matching the `&self` trait signature. Rebuilt from documents on cold
7//! start — acceptable for edge-scale datasets.
8//!
9//! Keys are fully structural tuples `(database_id, tid, collection, field, …)`
10//! — database, tenant, and index isolation never depends on lexical-prefix
11//! ordering.
12
13use std::cell::RefCell;
14use std::collections::{BTreeMap, HashMap};
15use std::fmt;
16
17use nodedb_types::Surrogate;
18
19use crate::backend::FtsBackend;
20use crate::posting::Posting;
21use crate::scope::IndexScope;
22
23/// In-memory backend error (infallible in practice, but trait requires it).
24#[derive(Debug)]
25pub struct MemoryError(String);
26
27impl fmt::Display for MemoryError {
28    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
29        write!(f, "memory backend: {}", self.0)
30    }
31}
32
33/// `(database_id, tid, collection, field)`: one index.
34type IndexKey = (u64, u64, String, String);
35/// An index key plus a term, meta subkey, or segment id.
36type SubKey = (IndexKey, String);
37/// An index key plus a document.
38type DocLenKey = (IndexKey, Surrogate);
39/// One term's postings keyed by document, so one document's posting is
40/// replaced or removed without touching the rest of the list.
41type PostingList = BTreeMap<Surrogate, Posting>;
42
43/// In-memory FTS backend backed by HashMaps keyed by
44/// `(database_id, tid, collection, field, …)` tuples.
45///
46/// Uses `RefCell` for interior mutability so the `FtsBackend` trait
47/// can use `&self` uniformly (redb has its own transactional isolation).
48#[derive(Debug, Default)]
49pub struct MemoryBackend {
50    /// `(index, term) → posting list`, ordered by document.
51    postings: RefCell<HashMap<SubKey, PostingList>>,
52    /// `(index, doc_id) → token count`.
53    doc_lengths: RefCell<HashMap<DocLenKey, u32>>,
54    /// `index → (doc_count, total_token_sum)`.
55    stats: RefCell<HashMap<IndexKey, (u32, u64)>>,
56    /// `(index, subkey) → blob` for fieldnorms, analyzer, language.
57    meta: RefCell<HashMap<SubKey, Vec<u8>>>,
58    /// `(index, segment_id) → compressed segment bytes`.
59    segments: RefCell<HashMap<SubKey, Vec<u8>>>,
60}
61
62impl MemoryBackend {
63    pub fn new() -> Self {
64        Self::default()
65    }
66}
67
68fn index_key(database_id: u64, tid: u64, index: IndexScope<'_>) -> IndexKey {
69    (
70        database_id,
71        tid,
72        index.collection().to_string(),
73        index.field_key().to_string(),
74    )
75}
76
77fn sub_key(database_id: u64, tid: u64, index: IndexScope<'_>, sub: &str) -> SubKey {
78    (index_key(database_id, tid, index), sub.to_string())
79}
80
81fn doc_len_key(database_id: u64, tid: u64, index: IndexScope<'_>, doc_id: Surrogate) -> DocLenKey {
82    (index_key(database_id, tid, index), doc_id)
83}
84
85/// Whether `key` belongs to `(database_id, tid)`, and to `collection` when given.
86fn owned_by(key: &IndexKey, database_id: u64, tid: u64, collection: Option<&str>) -> bool {
87    key.0 == database_id && key.1 == tid && collection.is_none_or(|c| key.2 == c)
88}
89
90impl MemoryBackend {
91    /// Drop every entry owned by `(database_id, tid[, collection])`.
92    fn purge_owned(&self, database_id: u64, tid: u64, collection: Option<&str>) -> usize {
93        let mut postings = self.postings.borrow_mut();
94        let mut doc_lengths = self.doc_lengths.borrow_mut();
95        let before = postings.len() + doc_lengths.len();
96        postings.retain(|k, _| !owned_by(&k.0, database_id, tid, collection));
97        doc_lengths.retain(|k, _| !owned_by(&k.0, database_id, tid, collection));
98        self.stats
99            .borrow_mut()
100            .retain(|k, _| !owned_by(k, database_id, tid, collection));
101        self.meta
102            .borrow_mut()
103            .retain(|k, _| !owned_by(&k.0, database_id, tid, collection));
104        self.segments
105            .borrow_mut()
106            .retain(|k, _| !owned_by(&k.0, database_id, tid, collection));
107        let after = postings.len() + doc_lengths.len();
108        before - after
109    }
110}
111
112impl FtsBackend for MemoryBackend {
113    type Error = MemoryError;
114
115    fn read_postings(
116        &self,
117        database_id: u64,
118        tid: u64,
119        index: IndexScope<'_>,
120        term: &str,
121    ) -> Result<Vec<Posting>, Self::Error> {
122        Ok(self
123            .postings
124            .borrow()
125            .get(&sub_key(database_id, tid, index, term))
126            .map(|list| list.values().cloned().collect())
127            .unwrap_or_default())
128    }
129
130    fn write_postings(
131        &self,
132        database_id: u64,
133        tid: u64,
134        index: IndexScope<'_>,
135        term: &str,
136        postings: &[Posting],
137    ) -> Result<(), Self::Error> {
138        let key = sub_key(database_id, tid, index, term);
139        let mut map = self.postings.borrow_mut();
140        if postings.is_empty() {
141            map.remove(&key);
142        } else {
143            let list: PostingList = postings.iter().map(|p| (p.doc_id, p.clone())).collect();
144            map.insert(key, list);
145        }
146        Ok(())
147    }
148
149    fn upsert_posting(
150        &self,
151        database_id: u64,
152        tid: u64,
153        index: IndexScope<'_>,
154        term: &str,
155        posting: &Posting,
156    ) -> Result<(), Self::Error> {
157        self.postings
158            .borrow_mut()
159            .entry(sub_key(database_id, tid, index, term))
160            .or_default()
161            .insert(posting.doc_id, posting.clone());
162        Ok(())
163    }
164
165    fn remove_posting(
166        &self,
167        database_id: u64,
168        tid: u64,
169        index: IndexScope<'_>,
170        term: &str,
171        doc_id: Surrogate,
172    ) -> Result<(), Self::Error> {
173        let key = sub_key(database_id, tid, index, term);
174        let mut map = self.postings.borrow_mut();
175        let emptied = match map.get_mut(&key) {
176            Some(list) => {
177                list.remove(&doc_id);
178                list.is_empty()
179            }
180            None => false,
181        };
182        if emptied {
183            map.remove(&key);
184        }
185        Ok(())
186    }
187
188    fn remove_postings(
189        &self,
190        database_id: u64,
191        tid: u64,
192        index: IndexScope<'_>,
193        term: &str,
194    ) -> Result<(), Self::Error> {
195        self.postings
196            .borrow_mut()
197            .remove(&sub_key(database_id, tid, index, term));
198        Ok(())
199    }
200
201    fn read_doc_length(
202        &self,
203        database_id: u64,
204        tid: u64,
205        index: IndexScope<'_>,
206        doc_id: Surrogate,
207    ) -> Result<Option<u32>, Self::Error> {
208        Ok(self
209            .doc_lengths
210            .borrow()
211            .get(&doc_len_key(database_id, tid, index, doc_id))
212            .copied())
213    }
214
215    fn write_doc_length(
216        &self,
217        database_id: u64,
218        tid: u64,
219        index: IndexScope<'_>,
220        doc_id: Surrogate,
221        length: u32,
222    ) -> Result<(), Self::Error> {
223        self.doc_lengths
224            .borrow_mut()
225            .insert(doc_len_key(database_id, tid, index, doc_id), length);
226        Ok(())
227    }
228
229    fn remove_doc_length(
230        &self,
231        database_id: u64,
232        tid: u64,
233        index: IndexScope<'_>,
234        doc_id: Surrogate,
235    ) -> Result<(), Self::Error> {
236        self.doc_lengths
237            .borrow_mut()
238            .remove(&doc_len_key(database_id, tid, index, doc_id));
239        Ok(())
240    }
241
242    fn collection_terms(
243        &self,
244        database_id: u64,
245        tid: u64,
246        index: IndexScope<'_>,
247    ) -> Result<Vec<String>, Self::Error> {
248        let key = index_key(database_id, tid, index);
249        Ok(self
250            .postings
251            .borrow()
252            .keys()
253            .filter(|(k, _)| *k == key)
254            .map(|(_, term)| term.clone())
255            .collect())
256    }
257
258    fn collection_stats(
259        &self,
260        database_id: u64,
261        tid: u64,
262        index: IndexScope<'_>,
263    ) -> Result<(u32, u64), Self::Error> {
264        Ok(self
265            .stats
266            .borrow()
267            .get(&index_key(database_id, tid, index))
268            .copied()
269            .unwrap_or((0, 0)))
270    }
271
272    fn increment_stats(
273        &self,
274        database_id: u64,
275        tid: u64,
276        index: IndexScope<'_>,
277        doc_len: u32,
278    ) -> Result<(), Self::Error> {
279        let mut stats = self.stats.borrow_mut();
280        let entry = stats
281            .entry(index_key(database_id, tid, index))
282            .or_insert((0, 0));
283        entry.0 += 1;
284        entry.1 += doc_len as u64;
285        Ok(())
286    }
287
288    fn decrement_stats(
289        &self,
290        database_id: u64,
291        tid: u64,
292        index: IndexScope<'_>,
293        doc_len: u32,
294    ) -> Result<(), Self::Error> {
295        let mut stats = self.stats.borrow_mut();
296        let entry = stats
297            .entry(index_key(database_id, tid, index))
298            .or_insert((0, 0));
299        entry.0 = entry.0.saturating_sub(1);
300        entry.1 = entry.1.saturating_sub(doc_len as u64);
301        Ok(())
302    }
303
304    fn read_meta(
305        &self,
306        database_id: u64,
307        tid: u64,
308        index: IndexScope<'_>,
309        subkey: &str,
310    ) -> Result<Option<Vec<u8>>, Self::Error> {
311        Ok(self
312            .meta
313            .borrow()
314            .get(&sub_key(database_id, tid, index, subkey))
315            .cloned())
316    }
317
318    fn write_meta(
319        &self,
320        database_id: u64,
321        tid: u64,
322        index: IndexScope<'_>,
323        subkey: &str,
324        value: &[u8],
325    ) -> Result<(), Self::Error> {
326        self.meta
327            .borrow_mut()
328            .insert(sub_key(database_id, tid, index, subkey), value.to_vec());
329        Ok(())
330    }
331
332    fn write_segment(
333        &self,
334        database_id: u64,
335        tid: u64,
336        index: IndexScope<'_>,
337        segment_id: &str,
338        data: &[u8],
339    ) -> Result<(), Self::Error> {
340        self.segments
341            .borrow_mut()
342            .insert(sub_key(database_id, tid, index, segment_id), data.to_vec());
343        Ok(())
344    }
345
346    fn read_segment(
347        &self,
348        database_id: u64,
349        tid: u64,
350        index: IndexScope<'_>,
351        segment_id: &str,
352    ) -> Result<Option<Vec<u8>>, Self::Error> {
353        Ok(self
354            .segments
355            .borrow()
356            .get(&sub_key(database_id, tid, index, segment_id))
357            .cloned())
358    }
359
360    fn list_segments(
361        &self,
362        database_id: u64,
363        tid: u64,
364        index: IndexScope<'_>,
365    ) -> Result<Vec<String>, Self::Error> {
366        let key = index_key(database_id, tid, index);
367        Ok(self
368            .segments
369            .borrow()
370            .keys()
371            .filter(|(k, _)| *k == key)
372            .map(|(_, seg)| seg.clone())
373            .collect())
374    }
375
376    fn remove_segment(
377        &self,
378        database_id: u64,
379        tid: u64,
380        index: IndexScope<'_>,
381        segment_id: &str,
382    ) -> Result<(), Self::Error> {
383        self.segments
384            .borrow_mut()
385            .remove(&sub_key(database_id, tid, index, segment_id));
386        Ok(())
387    }
388
389    fn purge_collection(
390        &self,
391        database_id: u64,
392        tid: u64,
393        collection: &str,
394    ) -> Result<usize, Self::Error> {
395        Ok(self.purge_owned(database_id, tid, Some(collection)))
396    }
397
398    fn purge_tenant(&self, database_id: u64, tid: u64) -> Result<usize, Self::Error> {
399        Ok(self.purge_owned(database_id, tid, None))
400    }
401}
402
403#[cfg(test)]
404mod tests {
405    use super::*;
406
407    const DB: u64 = 0;
408    const T: u64 = 1;
409    const COL: IndexScope<'static> = IndexScope::document("col");
410    const OTHER: IndexScope<'static> = IndexScope::document("other");
411
412    fn posting(position: u32) -> Vec<Posting> {
413        vec![Posting {
414            doc_id: Surrogate(1),
415            term_freq: 1,
416            positions: vec![position],
417        }]
418    }
419
420    #[test]
421    fn roundtrip_postings() {
422        let backend = MemoryBackend::new();
423        let postings = vec![Posting {
424            doc_id: Surrogate(1),
425            term_freq: 2,
426            positions: vec![0, 5],
427        }];
428        backend
429            .write_postings(DB, T, COL, "hello", &postings)
430            .unwrap();
431
432        let read = backend.read_postings(DB, T, COL, "hello").unwrap();
433        assert_eq!(read.len(), 1);
434        assert_eq!(read[0].doc_id, Surrogate(1));
435    }
436
437    #[test]
438    fn upsert_replaces_one_document_and_remove_empties_the_list() {
439        let backend = MemoryBackend::new();
440        let doc = |id: u32, freq: u32| Posting {
441            doc_id: Surrogate(id),
442            term_freq: freq,
443            positions: vec![0],
444        };
445        for (id, freq) in [(3, 1), (1, 1), (2, 1), (1, 4)] {
446            backend
447                .upsert_posting(DB, T, COL, "ab", &doc(id, freq))
448                .unwrap();
449        }
450        backend
451            .upsert_posting(DB, T, COL, "abc", &doc(9, 1))
452            .unwrap();
453
454        let ab = backend.read_postings(DB, T, COL, "ab").unwrap();
455        let ids: Vec<(u32, u32)> = ab.iter().map(|p| (p.doc_id.0, p.term_freq)).collect();
456        assert_eq!(ids, vec![(1, 4), (2, 1), (3, 1)]);
457        assert_eq!(backend.read_postings(DB, T, COL, "abc").unwrap().len(), 1);
458
459        for id in [1, 2, 3] {
460            backend
461                .remove_posting(DB, T, COL, "ab", Surrogate(id))
462                .unwrap();
463        }
464        assert!(backend.read_postings(DB, T, COL, "ab").unwrap().is_empty());
465        assert_eq!(backend.collection_terms(DB, T, COL).unwrap(), vec!["abc"]);
466        backend
467            .remove_posting(DB, T, COL, "missing", Surrogate(1))
468            .unwrap();
469    }
470
471    #[test]
472    fn roundtrip_doc_lengths() {
473        let backend = MemoryBackend::new();
474        backend
475            .write_doc_length(DB, T, COL, Surrogate(1), 42)
476            .unwrap();
477        assert_eq!(
478            backend.read_doc_length(DB, T, COL, Surrogate(1)).unwrap(),
479            Some(42)
480        );
481
482        backend.remove_doc_length(DB, T, COL, Surrogate(1)).unwrap();
483        assert_eq!(
484            backend.read_doc_length(DB, T, COL, Surrogate(1)).unwrap(),
485            None
486        );
487    }
488
489    #[test]
490    fn incremental_stats() {
491        let backend = MemoryBackend::new();
492        backend.increment_stats(DB, T, COL, 10).unwrap();
493        backend.increment_stats(DB, T, COL, 20).unwrap();
494        assert_eq!(backend.collection_stats(DB, T, COL).unwrap(), (2, 30));
495
496        backend.decrement_stats(DB, T, COL, 10).unwrap();
497        assert_eq!(backend.collection_stats(DB, T, COL).unwrap(), (1, 20));
498    }
499
500    #[test]
501    fn stats_saturating_sub() {
502        let backend = MemoryBackend::new();
503        backend.decrement_stats(DB, T, COL, 100).unwrap();
504        assert_eq!(backend.collection_stats(DB, T, COL).unwrap(), (0, 0));
505    }
506
507    #[test]
508    fn field_scopes_are_isolated_from_the_document_scope() {
509        let backend = MemoryBackend::new();
510        let title = IndexScope::field("col", "title").unwrap();
511        backend
512            .write_postings(DB, T, title, "rust", &posting(0))
513            .unwrap();
514        backend.increment_stats(DB, T, title, 3).unwrap();
515        backend
516            .write_doc_length(DB, T, title, Surrogate(1), 3)
517            .unwrap();
518
519        assert!(
520            backend
521                .read_postings(DB, T, COL, "rust")
522                .unwrap()
523                .is_empty()
524        );
525        assert_eq!(backend.collection_stats(DB, T, COL).unwrap(), (0, 0));
526        assert_eq!(
527            backend.read_doc_length(DB, T, COL, Surrogate(1)).unwrap(),
528            None
529        );
530        assert_eq!(
531            backend.collection_terms(DB, T, title).unwrap(),
532            vec!["rust"]
533        );
534        assert!(backend.collection_terms(DB, T, COL).unwrap().is_empty());
535    }
536
537    #[test]
538    fn purge_clears_stats_and_isolates_collections() {
539        let backend = MemoryBackend::new();
540        let title = IndexScope::field("col", "title").unwrap();
541        backend.increment_stats(DB, T, COL, 10).unwrap();
542        backend.increment_stats(DB, T, title, 2).unwrap();
543        backend
544            .write_doc_length(DB, T, COL, Surrogate(1), 10)
545            .unwrap();
546        backend
547            .write_postings(DB, T, COL, "hello", &posting(0))
548            .unwrap();
549        backend
550            .write_postings(DB, T, title, "hello", &posting(0))
551            .unwrap();
552
553        backend.increment_stats(DB, T, OTHER, 7).unwrap();
554        backend
555            .write_doc_length(DB, T, OTHER, Surrogate(1), 7)
556            .unwrap();
557        backend
558            .write_postings(DB, T, OTHER, "world", &posting(0))
559            .unwrap();
560
561        backend.purge_collection(DB, T, "col").unwrap();
562        assert_eq!(backend.collection_stats(DB, T, COL).unwrap(), (0, 0));
563        assert_eq!(backend.collection_stats(DB, T, title).unwrap(), (0, 0));
564        assert!(
565            backend
566                .read_postings(DB, T, COL, "hello")
567                .unwrap()
568                .is_empty()
569        );
570        assert!(
571            backend
572                .read_postings(DB, T, title, "hello")
573                .unwrap()
574                .is_empty()
575        );
576        assert_eq!(
577            backend.read_doc_length(DB, T, COL, Surrogate(1)).unwrap(),
578            None
579        );
580
581        assert_eq!(backend.collection_stats(DB, T, OTHER).unwrap(), (1, 7));
582        assert_eq!(
583            backend.read_postings(DB, T, OTHER, "world").unwrap().len(),
584            1
585        );
586        assert_eq!(
587            backend.read_doc_length(DB, T, OTHER, Surrogate(1)).unwrap(),
588            Some(7)
589        );
590    }
591
592    #[test]
593    fn collection_terms() {
594        let backend = MemoryBackend::new();
595        backend
596            .write_postings(DB, T, COL, "hello", &posting(0))
597            .unwrap();
598        backend
599            .write_postings(DB, T, COL, "world", &posting(1))
600            .unwrap();
601
602        let mut terms = backend.collection_terms(DB, T, COL).unwrap();
603        terms.sort();
604        assert_eq!(terms, vec!["hello", "world"]);
605    }
606
607    #[test]
608    fn segment_roundtrip() {
609        let backend = MemoryBackend::new();
610        let data = b"compressed segment bytes";
611        backend.write_segment(DB, T, COL, "id1", data).unwrap();
612        assert_eq!(
613            backend.read_segment(DB, T, COL, "id1").unwrap(),
614            Some(data.to_vec())
615        );
616        assert_eq!(backend.read_segment(DB, T, COL, "missing").unwrap(), None);
617    }
618
619    #[test]
620    fn segment_list_filters_by_index() {
621        let backend = MemoryBackend::new();
622        let title = IndexScope::field("col", "title").unwrap();
623        backend.write_segment(DB, T, COL, "a", b"a").unwrap();
624        backend.write_segment(DB, T, COL, "b", b"b").unwrap();
625        backend.write_segment(DB, T, OTHER, "c", b"c").unwrap();
626        backend.write_segment(DB, T, title, "d", b"d").unwrap();
627
628        let mut segs = backend.list_segments(DB, T, COL).unwrap();
629        segs.sort();
630        assert_eq!(segs, vec!["a", "b"]);
631
632        assert_eq!(backend.list_segments(DB, T, OTHER).unwrap(), vec!["c"]);
633        assert_eq!(backend.list_segments(DB, T, title).unwrap(), vec!["d"]);
634    }
635
636    #[test]
637    fn segment_remove() {
638        let backend = MemoryBackend::new();
639        backend.write_segment(DB, T, COL, "id1", b"data").unwrap();
640        backend.remove_segment(DB, T, COL, "id1").unwrap();
641        assert_eq!(backend.read_segment(DB, T, COL, "id1").unwrap(), None);
642    }
643
644    #[test]
645    fn purge_clears_segments() {
646        let backend = MemoryBackend::new();
647        backend.write_segment(DB, T, COL, "a", b"a").unwrap();
648        backend.write_segment(DB, T, OTHER, "b", b"b").unwrap();
649
650        backend.purge_collection(DB, T, "col").unwrap();
651        assert!(backend.list_segments(DB, T, COL).unwrap().is_empty());
652        assert_eq!(backend.list_segments(DB, T, OTHER).unwrap().len(), 1);
653    }
654
655    #[test]
656    fn purge_tenant_isolates_tenants() {
657        let backend = MemoryBackend::new();
658        backend.increment_stats(DB, 1, COL, 5).unwrap();
659        backend.increment_stats(DB, 2, COL, 7).unwrap();
660        backend
661            .write_postings(DB, 1, COL, "t", &posting(0))
662            .unwrap();
663        backend
664            .write_postings(DB, 2, COL, "t", &posting(0))
665            .unwrap();
666
667        backend.purge_tenant(DB, 1).unwrap();
668        assert_eq!(backend.collection_stats(DB, 1, COL).unwrap(), (0, 0));
669        assert!(backend.read_postings(DB, 1, COL, "t").unwrap().is_empty());
670        assert_eq!(backend.collection_stats(DB, 2, COL).unwrap(), (1, 7));
671        assert_eq!(backend.read_postings(DB, 2, COL, "t").unwrap().len(), 1);
672    }
673
674    #[test]
675    fn databases_isolated() {
676        let backend = MemoryBackend::new();
677        backend.increment_stats(0, T, COL, 5).unwrap();
678        backend.increment_stats(9, T, COL, 7).unwrap();
679        backend.write_postings(0, T, COL, "t", &posting(0)).unwrap();
680        backend.write_postings(9, T, COL, "t", &posting(0)).unwrap();
681
682        backend.purge_tenant(0, T).unwrap();
683        assert_eq!(backend.collection_stats(0, T, COL).unwrap(), (0, 0));
684        assert!(backend.read_postings(0, T, COL, "t").unwrap().is_empty());
685        // Same tenant in a different database must be unaffected.
686        assert_eq!(backend.collection_stats(9, T, COL).unwrap(), (1, 7));
687        assert_eq!(backend.read_postings(9, T, COL, "t").unwrap().len(), 1);
688    }
689}
nodedb_fts/backend/memory.rs

nodedb_fts/backend/
memory.rs
