]> git.mikk.net Git - mtbl-rs/commitdiff
Add filter with seek capability
authorChris Mikkelson <cmikk@fsi.io>
Fri, 13 Sep 2024 17:58:13 +0000 (12:58 -0500)
committerChris Mikkelson <cmikk@fsi.io>
Fri, 13 Sep 2024 17:58:13 +0000 (12:58 -0500)
src/filter.rs [new file with mode: 0644]

diff --git a/src/filter.rs b/src/filter.rs
new file mode 100644 (file)
index 0000000..ddcdac2
--- /dev/null
@@ -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>>()
+    );
+}