Rebase functionality via nice build pattern

Signed-off-by: Matthias Beyer <mail@beyermatthias.de>

3 files changed, 139 insertions, 0 deletions

diff --git a/src/async_dag.rs b/src/async_dag.rs
index fb3bf3d..187a9c6 100644
--- a/src/async_dag.rs
+++ b/src/async_dag.rs
@@ -158,6 +158,10 @@ where
         self.head = id.clone();
         Ok(id)
     }
+
+    pub fn rebase(&mut self) -> crate::Rebase<'_, Id, N, Backend> {
+        crate::rebase::Rebase(self)
+    }
 }
 
 pub trait Merger<Id, N>
diff --git a/src/lib.rs b/src/lib.rs
index bcb6ae0..c89237c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -16,5 +16,8 @@ pub use node::*;
 mod node_id;
 pub use node_id::*;
 
+mod rebase;
+pub use rebase::*;
+
 #[cfg(test)]
 mod test_impl;
diff --git a/src/rebase.rs b/src/rebase.rs
new file mode 100644
index 0000000..2bb11ee
--- /dev/null
+++ b/src/rebase.rs
@@ -0,0 +1,132 @@
+//! Module for the rebase types
+//!
+//! This module contains types for the `AsyncDag::rebase()` building of rebase functionality.
+//!
+
+use anyhow::Result;
+use futures::StreamExt;
+
+use crate::AsyncDag;
+use crate::DagBackend;
+use crate::Node;
+use crate::NodeId;
+
+pub struct Rebase<'a, Id, N, Backend>(pub(crate) &'a mut AsyncDag<Id, N, Backend>)
+    where Id: NodeId + Send,
+          N: Node<Id = Id>,
+          Backend: DagBackend<Id, N>;
+
+impl<'a, Id, N, Backend> Rebase<'a, Id, N, Backend>
+    where Id: NodeId + Send,
+          N: Node<Id = Id>,
+          Backend: DagBackend<Id, N>
+{
+    pub fn onto(self, id: Id) -> RebaseOnto<'a, Id, N, Backend, IdEqualSelector<Id>> {
+        RebaseOnto(self.0, id.clone(), IdEqualSelector(id))
+    }
+}
+
+pub trait RebaseNodeRewriterCreator<'a, Id, N, Backend, Selector>
+    where Id: NodeId + Send,
+          N: Node<Id = Id>,
+          Backend: DagBackend<Id, N>,
+          Selector: IdSelector<Id>
+{
+    fn doing<Rewriter>(self, rewriter: Rewriter) -> RebaseNodeRewriter<'a, Id, N, Backend, Selector, Rewriter>
+        where Rewriter: NodeRewriter<Id, N>;
+}
+
+pub struct RebaseOnto<'a, Id, N, Backend, Selector>(&'a mut AsyncDag<Id, N, Backend>, Id, Selector)
+    where Id: NodeId + Send,
+          N: Node<Id = Id>,
+          Backend: DagBackend<Id, N>,
+          Selector: IdSelector<Id>;
+
+impl<'a, Id, N, Backend, Selector> RebaseNodeRewriterCreator<'a, Id, N, Backend, Selector> for RebaseOnto<'a, Id, N, Backend, Selector>
+    where Id: NodeId + Send,
+          N: Node<Id = Id>,
+          Backend: DagBackend<Id, N>,
+          Selector: IdSelector<Id>
+{
+    fn doing<Rewriter>(self, rewriter: Rewriter) -> RebaseNodeRewriter<'a, Id, N, Backend, Selector, Rewriter>
+        where Rewriter: NodeRewriter<Id, N>
+    {
+        RebaseNodeRewriter {
+            dag: self.0,
+            base: self.1,
+            selector: self.2,
+            rewriter,
+        }
+    }
+}
+
+pub trait IdSelector<Id: NodeId> {
+    fn select(&self, id: &Id) -> bool;
+}
+
+pub struct IdEqualSelector<Id: NodeId>(Id);
+impl<Id: NodeId> IdSelector<Id> for IdEqualSelector<Id> {
+    fn select(&self, id: &Id) -> bool {
+        self.0 == *id
+    }
+}
+
+pub trait NodeRewriter<Id: NodeId, N: Node<Id = Id>> {
+    fn rewrite_node(&mut self, previous_id: Id, id: Id, node: N) -> Result<N>;
+}
+
+pub struct RebaseNodeRewriter<'a, Id, N, Backend, Selector, Rewriter>
+    where Id: NodeId + Send,
+          N: Node<Id = Id>,
+          Backend: DagBackend<Id, N>,
+          Selector: IdSelector<Id>,
+          Rewriter: NodeRewriter<Id, N>,
+{
+    dag: &'a mut AsyncDag<Id, N, Backend>,
+    base: Id,
+    selector: Selector,
+    rewriter: Rewriter,
+}
+
+impl<'a, Id, N, Backend, Selector, Rewriter> RebaseNodeRewriter<'a, Id, N, Backend, Selector, Rewriter>
+    where Id: NodeId + Send,
+          N: Node<Id = Id>,
+          Backend: DagBackend<Id, N>,
+          Selector: IdSelector<Id>,
+          Rewriter: NodeRewriter<Id, N>
+{
+    pub async fn run(self) -> Result<Id> {
+        let dag = self.dag;
+        let sel = self.selector;
+        let rew = self.rewriter;
+
+        let v = dag.stream()
+            .map(|node| {
+                match node {
+                    Ok((id, node)) => (!sel.select(&id), Ok((id, node))),
+                    Err(e) => (true, Err(e)),
+                }
+            })
+            .take_while(|tpl| futures::future::ready(tpl.0))
+            .map(|tpl| tpl.1)
+            .collect::<Vec<_>>()
+            .await;
+
+        futures::stream::iter(v.into_iter().rev())
+            .fold(Ok((dag, rew, self.base)), |previous, r| async {
+                match (previous, r) {
+                    (Ok((dag, mut rew, prev)), Ok((id, node))) => {
+                        let node = rew.rewrite_node(prev, id, node)?;
+                        dag.update_head_unchecked(node)
+                            .await
+                            .map(|id| (dag, rew, id))
+                    },
+                    (_, Err(e)) => Err(e),
+                    (Err(e), _) => Err(e),
+                }
+            })
+            .await
+            .map(|tpl| tpl.2)
+    }
+}
+