haystack_server/session/
peek.rs1#[derive(Debug)]
3pub struct PeekResult<T> {
4 pub items: Vec<T>,
6 pub satisfied: bool,
8 pub sampled: usize,
10}
11
12pub 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
46pub 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}