From f114b67fd7ddfcc4aadb65e72975fb5e32e9f7fb Mon Sep 17 00:00:00 2001
From: Matt Hunzinger <matthunz2@gmail.com>
Date: Sun, 18 Feb 2024 08:52:45 -0500
Subject: [PATCH] Add Then combinator and refactor

---
 src/lib.rs       | 21 +-------------------
 src/task/mod.rs  | 51 ++++++++++++++++++++++++++++++++++++++++++++++++
 src/task/then.rs | 48 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 100 insertions(+), 20 deletions(-)
 create mode 100644 src/task/mod.rs
 create mode 100644 src/task/then.rs

diff --git a/src/lib.rs b/src/lib.rs
index 043cc6b..cdafe4a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,20 +1 @@
-use std::rc::Rc;
-
-pub struct Context<M, A> {
-    waker: Rc<dyn Fn(Rc<dyn Fn(&mut M) -> Option<A>>)>,
-}
-
-pub trait Task<M, A = ()> {
-    type Output;
-
-    type State;
-
-    fn build(&mut self, cx: &Context<M, A>, model: &mut M) -> (Self::Output, Self::State);
-
-    fn rebuild(
-        &mut self,
-        cx: &Context<M, A>,
-        model: &mut M,
-        state: &mut Self::State,
-    ) -> Self::Output;
-}
+pub mod task;
diff --git a/src/task/mod.rs b/src/task/mod.rs
new file mode 100644
index 0000000..a445957
--- /dev/null
+++ b/src/task/mod.rs
@@ -0,0 +1,51 @@
+use std::rc::Rc;
+
+mod then;
+pub use self::then::Then;
+
+pub struct Context<M, A> {
+    waker: Rc<dyn Fn(Rc<dyn Fn(&mut M) -> Option<A>>)>,
+}
+
+pub trait Task<M, A = ()> {
+    type Output;
+
+    type State;
+
+    fn build(&mut self, cx: &Context<M, A>, model: &mut M) -> (Self::Output, Self::State);
+
+    fn rebuild(
+        &mut self,
+        cx: &Context<M, A>,
+        model: &mut M,
+        state: &mut Self::State,
+    ) -> Self::Output;
+
+    fn then<F, T>(self, f: F) -> Then<Self, F, T, M>
+    where
+        Self: Sized + 'static,
+        F: FnMut(&mut M, Self::Output) -> T + 'static,
+        T: Task<M> + 'static,
+        M: 'static,
+    {
+        Then::new(self, f)
+    }
+}
+
+impl<M> Task<M> for () {
+    type Output = ();
+
+    type State = ();
+
+    fn build(&mut self, _cx: &Context<M, ()>, _model: &mut M) -> (Self::Output, Self::State) {
+        ((), ())
+    }
+
+    fn rebuild(
+        &mut self,
+        _cx: &Context<M, ()>,
+        _model: &mut M,
+        _state: &mut Self::State,
+    ) -> Self::Output {
+    }
+}
diff --git a/src/task/then.rs b/src/task/then.rs
new file mode 100644
index 0000000..67b0e24
--- /dev/null
+++ b/src/task/then.rs
@@ -0,0 +1,48 @@
+use super::{Context, Task};
+use std::marker::PhantomData;
+
+pub struct Then<T1, F, T2, M> {
+    task: T1,
+    f: F,
+    _marker: PhantomData<(T2, M)>,
+}
+
+impl<T1, F, T2, M> Then<T1, F, T2, M> {
+    pub(super) fn new(task: T1, f: F) -> Self {
+        Self {
+            task,
+            f,
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<T1, F, T2, M> Task<M> for Then<T1, F, T2, M>
+where
+    T1: Task<M> + 'static,
+    F: FnMut(&mut M, T1::Output) -> T2 + 'static,
+    T2: Task<M> + 'static,
+    M: 'static,
+{
+    type Output = T2::Output;
+
+    type State = (T1::State, T2::State);
+
+    fn build(&mut self, cx: &Context<M, ()>, model: &mut M) -> (Self::Output, Self::State) {
+        let (output, state) = self.task.build(cx, model);
+        let mut next = (self.f)(model, output);
+        let (next_output, next_state) = next.build(cx, model);
+        (next_output, (state, next_state))
+    }
+
+    fn rebuild(
+        &mut self,
+        cx: &Context<M, ()>,
+        model: &mut M,
+        state: &mut Self::State,
+    ) -> Self::Output {
+        let output = self.task.rebuild(cx, model, &mut state.0);
+        let mut next = (self.f)(model, output);
+        next.rebuild(cx, model, &mut state.1)
+    }
+}