Skip to main content

haystack_server/session/
peek.rs

1/// Result of a peek evaluation — check a small sample before full scan.
2#[derive(Debug)]
3pub struct PeekResult<T> {
4    /// Items found in the peek sample.
5    pub items: Vec<T>,
6    /// Whether the limit was satisfied by the peek alone.
7    pub satisfied: bool,
8    /// Number of items sampled.
9    pub sampled: usize,
10}
11
12/// Peek at the first N items from an iterator, checking if a limit is satisfied.
13pub fn peek_eval<T, I>(iter: I, peek_size: usize, limit: Option<usize>) -> PeekResult<T>
14where
15    I: Iterator<Item = T>,
16{
17    let mut items = Vec::new();
18    let mut count = 0;
19
20    for item in iter {
21        count += 1;
22        items.push(item);
23
24        if let Some(lim) = limit
25            && items.len() >= lim
26        {
27            return PeekResult {
28                items,
29                satisfied: true,
30                sampled: count,
31            };
32        }
33
34        if count >= peek_size {
35            break;
36        }
37    }
38
39    PeekResult {
40        satisfied: limit.is_some_and(|lim| items.len() >= lim),
41        sampled: count,
42        items,
43    }
44}
45
46/// Lazy evaluation wrapper that avoids full collection when limit is known.
47pub fn lazy_collect<T, I>(iter: I, limit: Option<usize>) -> Vec<T>
48where
49    I: Iterator<Item = T>,
50{
51    match limit {
52        Some(lim) => iter.take(lim).collect(),
53        None => iter.collect(),
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60
61    #[test]
62    fn peek_limit_satisfied() {
63        let data = vec![1, 2, 3, 4, 5];
64        let result = peek_eval(data.into_iter(), 10, Some(3));
65        assert!(result.satisfied);
66        assert_eq!(result.items.len(), 3);
67        assert_eq!(result.sampled, 3);
68    }
69
70    #[test]
71    fn peek_limit_not_satisfied() {
72        let data = vec![1, 2];
73        let result = peek_eval(data.into_iter(), 10, Some(5));
74        assert!(!result.satisfied);
75        assert_eq!(result.items.len(), 2);
76        assert_eq!(result.sampled, 2);
77    }
78
79    #[test]
80    fn peek_no_limit_stops_at_peek_size() {
81        let data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
82        let result = peek_eval(data.into_iter(), 3, None);
83        assert!(!result.satisfied);
84        assert_eq!(result.items.len(), 3);
85        assert_eq!(result.sampled, 3);
86    }
87
88    #[test]
89    fn peek_empty_iterator() {
90        let data: Vec<i32> = vec![];
91        let result = peek_eval(data.into_iter(), 10, Some(5));
92        assert!(!result.satisfied);
93        assert!(result.items.is_empty());
94        assert_eq!(result.sampled, 0);
95    }
96
97    #[test]
98    fn lazy_collect_with_limit() {
99        let data = vec![1, 2, 3, 4, 5];
100        let result = lazy_collect(data.into_iter(), Some(3));
101        assert_eq!(result, vec![1, 2, 3]);
102    }
103
104    #[test]
105    fn lazy_collect_without_limit() {
106        let data = vec![1, 2, 3, 4, 5];
107        let result = lazy_collect(data.into_iter(), None);
108        assert_eq!(result, vec![1, 2, 3, 4, 5]);
109    }
110
111    #[test]
112    fn lazy_collect_limit_exceeds_items() {
113        let data = vec![1, 2];
114        let result = lazy_collect(data.into_iter(), Some(10));
115        assert_eq!(result, vec![1, 2]);
116    }
117}