Added support for context-less functions.

jakubadamw · Oct 27, 2024 · f0f2c29 · f0f2c29
1 parent 191efc2
commit f0f2c29
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 12 deletions.
diff --git a/examples/fibonacci.rs b/examples/fibonacci.rs
@@ -33,7 +33,7 @@ struct Output {
     result: u32,
 }
 
-async fn execute(_context: Context, Input { n }: Input) -> anyhow::Result<Output> {
+async fn execute(Input { n }: Input) -> anyhow::Result<Output> {
     Ok(Output {
         result: fibonacci(n),
     })

diff --git a/src/workflow.rs b/src/workflow.rs
@@ -19,22 +19,59 @@ pub struct Step {
     pub(crate) timeout: std::time::Duration,
 }
 
-impl StepBuilder {
-    pub fn function<I, O, Fut, F>(mut self, function: &'static F) -> Self
-    where
-        I: serde::de::DeserializeOwned,
-        O: serde::ser::Serialize,
-        Fut: std::future::Future<Output = anyhow::Result<O>> + 'static,
-        F: Fn(Context, I) -> Fut,
-    {
+pub trait UserStepFunction<I, O, H> {
+    fn to_step_function(self) -> Arc<StepFunction>;
+}
+
+pub struct NoArguments;
+pub struct ContextArgument;
+
+impl<I, O, Fut, F> UserStepFunction<I, O, ContextArgument> for &'static F
+where
+    I: serde::de::DeserializeOwned,
+    O: serde::ser::Serialize,
+    Fut: std::future::Future<Output = anyhow::Result<O>> + 'static,
+    F: Fn(Context, I) -> Fut,
+{
+    fn to_step_function(self) -> Arc<StepFunction> {
         use futures_util::FutureExt;
-        self.function = Some(Arc::new(|context, value| {
-            let result = function(
+        Arc::new(|context, value| {
+            let result = (self)(
                 context,
                 serde_json::from_value(value).expect("must succeed"),
             );
             async { Ok(serde_json::to_value(result.await?).expect("must succeed")) }.boxed_local()
-        }));
+        })
+    }
+}
+
+impl<I, O, Fut, F> UserStepFunction<I, O, NoArguments> for &'static F
+where
+    I: serde::de::DeserializeOwned,
+    O: serde::ser::Serialize,
+    Fut: std::future::Future<Output = anyhow::Result<O>> + 'static,
+    F: Fn(I) -> Fut,
+{
+    fn to_step_function(self) -> Arc<StepFunction> {
+        use futures_util::FutureExt;
+        Arc::new(|_context, value| {
+            let result = (self)(serde_json::from_value(value).expect("must succeed"));
+            async { Ok(serde_json::to_value(result.await?).expect("must succeed")) }.boxed_local()
+        })
+    }
+}
+
+impl StepBuilder {
+    pub fn function<
+        AnyVariant,
+        I: serde::de::DeserializeOwned,
+        O: serde::ser::Serialize,
+        F: UserStepFunction<I, O, AnyVariant>,
+    >(
+        mut self,
+        function: F,
+    ) -> Self {
+        self.function = Some(function.to_step_function());
         self
     }
 }