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

nodedb_fts/backend/
memory.rs
