hermes_core/query/
docset.rs

1//! DocSet trait and concrete implementations for document iteration.
2//!
3//! `DocSet` is the base abstraction for forward-only cursors over sorted document IDs.
4//! Posting lists, filter results, and scorers all implement this trait.
5//! `IntersectionScorer` intersects a Scorer with a DocSet filter, driving from
6//! the smaller side by `size_hint`.
7
8use std::sync::Arc;
9
10use crate::DocId;
11use crate::structures::TERMINATED;
12
13// ── DocSet trait ─────────────────────────────────────────────────────────
14
15macro_rules! define_docset_trait {
16    ($($send_bounds:tt)*) => {
17        /// Forward-only cursor over sorted document IDs.
18        ///
19        /// This is the base iteration abstraction. Posting lists, filter cursors,
20        /// and scorers all implement this trait.
21        pub trait DocSet: $($send_bounds)* {
22            /// Current document ID, or [`TERMINATED`] if exhausted.
23            fn doc(&self) -> DocId;
24
25            /// Advance to the next document. Returns the new doc ID or [`TERMINATED`].
26            fn advance(&mut self) -> DocId;
27
28            /// Seek to the first document >= `target`. Returns doc ID or [`TERMINATED`].
29            fn seek(&mut self, target: DocId) -> DocId {
30                let mut doc = self.doc();
31                while doc < target {
32                    doc = self.advance();
33                }
34                doc
35            }
36
37            /// Estimated number of remaining documents.
38            fn size_hint(&self) -> u32;
39        }
40    };
41}
42
43#[cfg(not(target_arch = "wasm32"))]
44define_docset_trait!(Send + Sync);
45
46#[cfg(target_arch = "wasm32")]
47define_docset_trait!();
48
49// ── DocSet for Box<dyn DocSet> ───────────────────────────────────────────
50
51impl DocSet for Box<dyn DocSet + '_> {
52    #[inline]
53    fn doc(&self) -> DocId {
54        (**self).doc()
55    }
56    #[inline]
57    fn advance(&mut self) -> DocId {
58        (**self).advance()
59    }
60    #[inline]
61    fn seek(&mut self, target: DocId) -> DocId {
62        (**self).seek(target)
63    }
64    #[inline]
65    fn size_hint(&self) -> u32 {
66        (**self).size_hint()
67    }
68}
69
70// ── SortedVecDocSet ──────────────────────────────────────────────────────
71
72/// DocSet backed by a sorted `Vec<u32>`. Binary search for seek.
73pub struct SortedVecDocSet {
74    docs: Arc<Vec<u32>>,
75    pos: usize,
76}
77
78impl SortedVecDocSet {
79    pub fn new(docs: Arc<Vec<u32>>) -> Self {
80        Self { docs, pos: 0 }
81    }
82}
83
84impl DocSet for SortedVecDocSet {
85    #[inline]
86    fn doc(&self) -> DocId {
87        self.docs.get(self.pos).copied().unwrap_or(TERMINATED)
88    }
89
90    #[inline]
91    fn advance(&mut self) -> DocId {
92        if self.pos < self.docs.len() {
93            self.pos += 1;
94        }
95        self.doc()
96    }
97
98    fn seek(&mut self, target: DocId) -> DocId {
99        if self.pos >= self.docs.len() {
100            return TERMINATED;
101        }
102        let remaining = &self.docs[self.pos..];
103        match remaining.binary_search(&target) {
104            Ok(offset) => {
105                self.pos += offset;
106                self.docs[self.pos]
107            }
108            Err(offset) => {
109                self.pos += offset;
110                self.doc()
111            }
112        }
113    }
114
115    fn size_hint(&self) -> u32 {
116        self.docs.len().saturating_sub(self.pos) as u32
117    }
118}
119
120// ── IntersectionDocSet ───────────────────────────────────────────────────
121
122/// DocSet that yields the intersection of two DocSets.
123pub struct IntersectionDocSet<A: DocSet, B: DocSet> {
124    a: A,
125    b: B,
126}
127
128impl<A: DocSet, B: DocSet> IntersectionDocSet<A, B> {
129    pub fn new(mut a: A, mut b: B) -> Self {
130        // Align both on the first common doc
131        let mut da = a.doc();
132        let mut db = b.doc();
133        loop {
134            if da == TERMINATED || db == TERMINATED {
135                break;
136            }
137            if da == db {
138                break;
139            }
140            if da < db {
141                da = a.seek(db);
142            } else {
143                db = b.seek(da);
144            }
145        }
146        Self { a, b }
147    }
148}
149
150impl<A: DocSet, B: DocSet> DocSet for IntersectionDocSet<A, B> {
151    fn doc(&self) -> DocId {
152        let da = self.a.doc();
153        if da == TERMINATED || self.b.doc() == TERMINATED {
154            TERMINATED
155        } else {
156            da
157        }
158    }
159
160    fn advance(&mut self) -> DocId {
161        let mut da = self.a.advance();
162        let mut db = self.b.doc();
163        loop {
164            if da == TERMINATED || db == TERMINATED {
165                return TERMINATED;
166            }
167            if da == db {
168                return da;
169            }
170            if da < db {
171                da = self.a.seek(db);
172            } else {
173                db = self.b.seek(da);
174            }
175        }
176    }
177
178    fn seek(&mut self, target: DocId) -> DocId {
179        let mut da = self.a.seek(target);
180        let mut db = self.b.seek(target);
181        loop {
182            if da == TERMINATED || db == TERMINATED {
183                return TERMINATED;
184            }
185            if da == db {
186                return da;
187            }
188            if da < db {
189                da = self.a.seek(db);
190            } else {
191                db = self.b.seek(da);
192            }
193        }
194    }
195
196    fn size_hint(&self) -> u32 {
197        self.a.size_hint().min(self.b.size_hint())
198    }
199}
200
201// ── AllDocSet ────────────────────────────────────────────────────────────
202
203/// DocSet that yields all documents 0..num_docs.
204pub struct AllDocSet {
205    current: u32,
206    num_docs: u32,
207}
208
209impl AllDocSet {
210    pub fn new(num_docs: u32) -> Self {
211        Self {
212            current: 0,
213            num_docs,
214        }
215    }
216}
217
218impl DocSet for AllDocSet {
219    #[inline]
220    fn doc(&self) -> DocId {
221        if self.current >= self.num_docs {
222            TERMINATED
223        } else {
224            self.current
225        }
226    }
227
228    #[inline]
229    fn advance(&mut self) -> DocId {
230        self.current += 1;
231        self.doc()
232    }
233
234    #[inline]
235    fn seek(&mut self, target: DocId) -> DocId {
236        self.current = target;
237        self.doc()
238    }
239
240    fn size_hint(&self) -> u32 {
241        self.num_docs.saturating_sub(self.current)
242    }
243}
244
245// ── EmptyDocSet ──────────────────────────────────────────────────────────
246
247/// DocSet that is always empty.
248pub struct EmptyDocSet;
249
250impl DocSet for EmptyDocSet {
251    #[inline]
252    fn doc(&self) -> DocId {
253        TERMINATED
254    }
255    #[inline]
256    fn advance(&mut self) -> DocId {
257        TERMINATED
258    }
259    #[inline]
260    fn seek(&mut self, _target: DocId) -> DocId {
261        TERMINATED
262    }
263    fn size_hint(&self) -> u32 {
264        0
265    }
266}
267
268// ── IntersectionScorer ───────────────────────────────────────────────────
269
270/// Intersects a Scorer with a filter DocSet, driving from the smaller side.
271///
272/// This is the core composition primitive: filter queries create a filter DocSet,
273/// get the inner Scorer, and return `IntersectionScorer(scorer, filter)`.
274pub struct IntersectionScorer<'a> {
275    scorer: Box<dyn super::Scorer + 'a>,
276    filter: Box<dyn DocSet + 'a>,
277}
278
279impl<'a> IntersectionScorer<'a> {
280    pub fn new(mut scorer: Box<dyn super::Scorer + 'a>, mut filter: Box<dyn DocSet + 'a>) -> Self {
281        // Align both on first common doc
282        let mut ds = scorer.doc();
283        let mut df = filter.doc();
284        loop {
285            if ds == TERMINATED || df == TERMINATED {
286                break;
287            }
288            if ds == df {
289                break;
290            }
291            if ds < df {
292                ds = scorer.seek(df);
293            } else {
294                df = filter.seek(ds);
295            }
296        }
297        Self { scorer, filter }
298    }
299}
300
301impl DocSet for IntersectionScorer<'_> {
302    fn doc(&self) -> DocId {
303        let ds = self.scorer.doc();
304        if ds == TERMINATED || self.filter.doc() == TERMINATED {
305            TERMINATED
306        } else {
307            ds
308        }
309    }
310
311    fn advance(&mut self) -> DocId {
312        // Drive from the smaller side
313        let filter_smaller = self.filter.size_hint() < self.scorer.size_hint();
314
315        if filter_smaller {
316            // Filter drives
317            let mut df = self.filter.advance();
318            let mut ds = self.scorer.doc();
319            loop {
320                if df == TERMINATED || ds == TERMINATED {
321                    return TERMINATED;
322                }
323                if df == ds {
324                    return df;
325                }
326                if df < ds {
327                    df = self.filter.seek(ds);
328                } else {
329                    ds = self.scorer.seek(df);
330                }
331            }
332        } else {
333            // Scorer drives
334            let mut ds = self.scorer.advance();
335            let mut df = self.filter.doc();
336            loop {
337                if ds == TERMINATED || df == TERMINATED {
338                    return TERMINATED;
339                }
340                if ds == df {
341                    return ds;
342                }
343                if ds < df {
344                    ds = self.scorer.seek(df);
345                } else {
346                    df = self.filter.seek(ds);
347                }
348            }
349        }
350    }
351
352    fn seek(&mut self, target: DocId) -> DocId {
353        let mut ds = self.scorer.seek(target);
354        let mut df = self.filter.seek(target);
355        loop {
356            if ds == TERMINATED || df == TERMINATED {
357                return TERMINATED;
358            }
359            if ds == df {
360                return ds;
361            }
362            if ds < df {
363                ds = self.scorer.seek(df);
364            } else {
365                df = self.filter.seek(ds);
366            }
367        }
368    }
369
370    fn size_hint(&self) -> u32 {
371        self.scorer.size_hint().min(self.filter.size_hint())
372    }
373}
374
375impl super::Scorer for IntersectionScorer<'_> {
376    fn score(&self) -> crate::Score {
377        self.scorer.score()
378    }
379
380    fn matched_positions(&self) -> Option<super::MatchedPositions> {
381        self.scorer.matched_positions()
382    }
383}
384
385// ── Tests ────────────────────────────────────────────────────────────────
386
387#[cfg(test)]
388mod tests {
389    use super::*;
390
391    #[test]
392    fn test_sorted_vec_docset_basic() {
393        let docs = Arc::new(vec![1, 3, 5, 7, 9]);
394        let mut ds = SortedVecDocSet::new(docs);
395
396        assert_eq!(ds.doc(), 1);
397        assert_eq!(ds.advance(), 3);
398        assert_eq!(ds.advance(), 5);
399        assert_eq!(ds.seek(7), 7);
400        assert_eq!(ds.advance(), 9);
401        assert_eq!(ds.advance(), TERMINATED);
402        assert_eq!(ds.doc(), TERMINATED);
403    }
404
405    #[test]
406    fn test_sorted_vec_docset_seek_past() {
407        let docs = Arc::new(vec![1, 5, 10, 20]);
408        let mut ds = SortedVecDocSet::new(docs);
409
410        assert_eq!(ds.seek(3), 5);
411        assert_eq!(ds.seek(15), 20);
412        assert_eq!(ds.seek(21), TERMINATED);
413    }
414
415    #[test]
416    fn test_sorted_vec_docset_empty() {
417        let docs = Arc::new(vec![]);
418        let ds = SortedVecDocSet::new(docs);
419        assert_eq!(ds.doc(), TERMINATED);
420    }
421
422    #[test]
423    fn test_all_docset() {
424        let mut ds = AllDocSet::new(3);
425        assert_eq!(ds.doc(), 0);
426        assert_eq!(ds.advance(), 1);
427        assert_eq!(ds.advance(), 2);
428        assert_eq!(ds.advance(), TERMINATED);
429    }
430
431    #[test]
432    fn test_all_docset_seek() {
433        let mut ds = AllDocSet::new(10);
434        assert_eq!(ds.seek(5), 5);
435        assert_eq!(ds.seek(9), 9);
436        assert_eq!(ds.seek(10), TERMINATED);
437    }
438
439    #[test]
440    fn test_empty_docset() {
441        let mut ds = EmptyDocSet;
442        assert_eq!(ds.doc(), TERMINATED);
443        assert_eq!(ds.advance(), TERMINATED);
444        assert_eq!(ds.seek(5), TERMINATED);
445        assert_eq!(ds.size_hint(), 0);
446    }
447
448    #[test]
449    fn test_intersection_docset() {
450        let a = SortedVecDocSet::new(Arc::new(vec![1, 3, 5, 7, 9]));
451        let b = SortedVecDocSet::new(Arc::new(vec![2, 3, 5, 8, 9, 10]));
452        let mut isect = IntersectionDocSet::new(a, b);
453
454        assert_eq!(isect.doc(), 3);
455        assert_eq!(isect.advance(), 5);
456        assert_eq!(isect.advance(), 9);
457        assert_eq!(isect.advance(), TERMINATED);
458    }
459
460    #[test]
461    fn test_intersection_docset_empty() {
462        let a = SortedVecDocSet::new(Arc::new(vec![1, 3, 5]));
463        let b = SortedVecDocSet::new(Arc::new(vec![2, 4, 6]));
464        let isect = IntersectionDocSet::new(a, b);
465        assert_eq!(isect.doc(), TERMINATED);
466    }
467
468    #[test]
469    fn test_intersection_docset_seek() {
470        let a = SortedVecDocSet::new(Arc::new(vec![1, 5, 10, 20, 30]));
471        let b = SortedVecDocSet::new(Arc::new(vec![5, 10, 15, 20, 25, 30]));
472        let mut isect = IntersectionDocSet::new(a, b);
473
474        assert_eq!(isect.doc(), 5);
475        assert_eq!(isect.seek(15), 20);
476        assert_eq!(isect.advance(), 30);
477        assert_eq!(isect.advance(), TERMINATED);
478    }
479
480    #[test]
481    fn test_size_hint() {
482        let docs = Arc::new(vec![1, 2, 3, 4, 5]);
483        let mut ds = SortedVecDocSet::new(docs);
484        assert_eq!(ds.size_hint(), 5);
485        ds.advance();
486        assert_eq!(ds.size_hint(), 4);
487        ds.seek(4);
488        assert_eq!(ds.size_hint(), 2); // pos=3, remaining: [4, 5]
489    }
490}
hermes_core/query/docset.rs

hermes_core/query/
docset.rs
