From c8056dc9d70d27ebd8ed4cb101b8f157c3959991 Mon Sep 17 00:00:00 2001
From: Chris Mikkelson <cmikk@fsi.io>
Date: Sun, 28 Apr 2024 15:54:28 -0500
Subject: [PATCH] Add coalesce (merge function) support

---
 src/seekable.rs          |  8 ++++
 src/seekable/coalesce.rs | 88 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+)
 create mode 100644 src/seekable/coalesce.rs

diff --git a/src/seekable.rs b/src/seekable.rs
index 9749a16..d8bdde9 100644
--- a/src/seekable.rs
+++ b/src/seekable.rs
@@ -1,3 +1,4 @@
+mod coalesce;
 mod filter_map;
 mod merge;
 mod vec;
@@ -27,6 +28,13 @@ pub trait Seekable: Sized {
     {
         merge::merge(iter)
     }
+
+    fn coalesce<F>(self, cf: F) -> Coalesce<Self, F>
+    where
+        F: FnMut(&Self::Key, Self::Value, Self::Value) -> Self::Value,
+    {
+        coalesce::coalesce(self, cf)
+    }
 }
 
 #[derive(Debug)]
diff --git a/src/seekable/coalesce.rs b/src/seekable/coalesce.rs
new file mode 100644
index 0000000..527b9c1
--- /dev/null
+++ b/src/seekable/coalesce.rs
@@ -0,0 +1,88 @@
+use crate::seekable::{Iter, Seekable};
+use std::cell::RefCell;
+
+pub struct Coalesce<S, F>
+where
+    S: Seekable,
+    F: FnMut(&S::Key, S::Value, S::Value) -> S::Value,
+{
+    cur_item: RefCell<Option<(S::Key, S::Value)>>,
+    rest: S,
+    cfunc: F,
+}
+
+pub fn coalesce<S, F>(seekable: S, cf: F) -> Coalesce<S, F>
+where
+    S: Seekable,
+    F: FnMut(&S::Key, S::Value, S::Value) -> S::Value,
+{
+    Coalesce {
+        cur_item: RefCell::new(None),
+        rest: seekable,
+        cfunc: cf,
+    }
+}
+
+impl<S, F> IntoIterator for Coalesce<S, F>
+where
+    S: Seekable,
+    F: FnMut(&S::Key, S::Value, S::Value) -> S::Value,
+{
+    type Item = (S::Key, S::Value);
+    type IntoIter = Iter<Self>;
+    fn into_iter(self) -> Iter<Self> {
+        Iter(self)
+    }
+}
+
+impl<S, F> Seekable for Coalesce<S, F>
+where
+    S: Seekable,
+    F: FnMut(&S::Key, S::Value, S::Value) -> S::Value,
+{
+    type Key = S::Key;
+    type Value = S::Value;
+
+    fn next(&mut self) -> Option<(Self::Key, Self::Value)> {
+        if self.cur_item.borrow().is_none() {
+            let next = self.rest.next()?;
+            self.cur_item.replace(Some(next));
+        }
+        while let Some((k, v)) = self.rest.next() {
+            if let Some((cur_k, cur_v)) = self.cur_item.replace(None) {
+                if k != cur_k {
+                    self.cur_item.replace(Some((k, v)));
+                    return Some((cur_k, cur_v));
+                }
+                self.cur_item
+                    .replace(Some((k, ((self.cfunc)(&cur_k, cur_v, v)))));
+            }
+        }
+        self.cur_item.replace(None)
+    }
+
+    fn seek(&mut self, key: &Self::Key) {
+        self.cur_item.replace(None);
+        self.rest.seek(key);
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::seekable::SeekableVec;
+    use crate::Seekable;
+    use std::cmp::Ordering;
+
+    #[test]
+    fn test_coalesce() {
+        let v = vec![(1, 2), (1, 3), (2, 1), (2, 4)];
+        let vc = vec![(1, 2 + 3), (2, 1 + 4)];
+        let v: SeekableVec<i32, i32> = SeekableVec::from(v.clone());
+        assert_eq!(
+            v.coalesce(|_k, v1, v2| v1 + v2)
+                .into_iter()
+                .cmp(vc.into_iter()),
+            Ordering::Equal
+        )
+    }
+}
-- 
2.50.1