From 0a9d8fb2836a886b29a6b7610ee0503c974178c3 Mon Sep 17 00:00:00 2001
From: Chris Mikkelson <cmikk@fsi.io>
Date: Fri, 13 Sep 2024 12:58:13 -0500
Subject: [PATCH] Add filter with seek capability

---
 src/filter.rs | 113 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 113 insertions(+)
 create mode 100644 src/filter.rs

diff --git a/src/filter.rs b/src/filter.rs
new file mode 100644
index 0000000..ddcdac2
--- /dev/null
+++ b/src/filter.rs
@@ -0,0 +1,113 @@
+use crate::{Entry, Iter, Source};
+
+pub struct FilterIter<'a, I, F> {
+    iter: I,
+    filter_func: &'a F,
+}
+
+impl<'a, I, F> FilterIter<'a, I, F>
+where
+    F: Fn(&Entry, &mut dyn Iter) -> bool,
+{
+    pub fn new(iter: I, filter_func: &'a F) -> Self {
+        Self { iter, filter_func }
+    }
+}
+
+impl<'a, I, F> Iterator for FilterIter<'a, I, F>
+where
+    F: Fn(&Entry, &mut dyn Iter) -> bool,
+    I: Iter,
+{
+    type Item = Entry;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        while let Some(e) = self.iter.next() {
+            if (self.filter_func)(&e, &mut self.iter) {
+                return Some(e);
+            }
+        }
+        None
+    }
+}
+
+impl<'a, I, F> Iter for FilterIter<'a, I, F>
+where
+    I: Iter,
+    F: Fn(&Entry, &mut dyn Iter) -> bool,
+{
+    fn seek(&mut self, key: &[u8]) {
+        self.iter.seek(key);
+    }
+}
+
+pub struct FilterSource<S, F> {
+    source: S,
+    filter_func: F,
+}
+
+impl<S, F> FilterSource<S, F>
+where
+    S: Source,
+    F: Fn(&Entry, &mut dyn Iter) -> bool,
+{
+    pub fn new(source: S, filter_func: F) -> Self {
+        Self {
+            source,
+            filter_func,
+        }
+    }
+}
+
+impl<S, F> Source for FilterSource<S, F>
+where
+    S: Source,
+    F: Fn(&Entry, &mut dyn Iter) -> bool,
+{
+    fn iter(&self) -> impl Iter {
+        self.source.iter().filter_func(&self.filter_func)
+    }
+
+    fn get(&self, key: &[u8]) -> impl Iter {
+        self.source.get(key).filter_func(&self.filter_func)
+    }
+
+    fn get_prefix(&self, prefix: &[u8]) -> impl Iter {
+        self.source
+            .get_prefix(prefix)
+            .filter_func(&self.filter_func)
+    }
+
+    fn get_range(&self, start: &[u8], end: &[u8]) -> impl Iter {
+        self.source
+            .get_range(start, end)
+            .filter_func(&self.filter_func)
+    }
+}
+
+#[test]
+fn test_filter() {
+    use crate::source::test_source::TestSource;
+
+    let ts = TestSource(
+        (0u8..10)
+            .into_iter()
+            .map(|n| Entry::new(vec![n], vec![]))
+            .collect(),
+    )
+    .filter(|e, it| {
+        // pass only even values
+        if e.key()[0] % 2 != 0 {
+            return false;
+        } else if e.key()[0] == 4 {
+            // at key 4, seek to 8
+            it.seek(vec![8].as_slice());
+        }
+        true
+    });
+
+    assert_eq!(
+        vec![0, 2, 4, 8],
+        ts.iter().map(|e| e.key()[0]).collect::<Vec<u8>>()
+    );
+}
-- 
2.50.1