feat(hydroflow_plus): push down persists and implement Pi example (#1021

) Also fixes type inference issues with reduce the same way as we did for fold.
hydro-project · Jan 7, 2024 · af6e3be · af6e3be
1 parent 7e4d577
commit af6e3be
Show file tree

Hide file tree

Showing 21 changed files with 463 additions and 101 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/hydroflow_lang/src/graph/ops/reduce.rs b/hydroflow_lang/src/graph/ops/reduce.rs
@@ -86,10 +86,18 @@ pub const REDUCE: OperatorConstraints = OperatorConstraints {
                     let #ident = {
                         let mut #input = #input;
                         let #accumulator_ident = #input.next();
+
+                        #[inline(always)]
+                        /// A: accumulator type
+                        /// O: output type
+                        fn call_comb_type<A, O>(acc: &mut A, item: A, f: impl Fn(&mut A, A) -> O) -> O {
+                            f(acc, item)
+                        }
+
                         if let ::std::option::Option::Some(mut #accumulator_ident) = #accumulator_ident {
                             for #iterator_item_ident in #input {
                                 #[allow(clippy::redundant_closure_call)]
-                                (#func)(&mut #accumulator_ident, #iterator_item_ident);
+                                call_comb_type(&mut #accumulator_ident, #iterator_item_ident, #func);
                             }
 
                             ::std::option::Option::Some(#accumulator_ident)
@@ -115,10 +123,17 @@ pub const REDUCE: OperatorConstraints = OperatorConstraints {
                             #input.next()
                         };
 
+                        #[inline(always)]
+                        /// A: accumulator type
+                        /// O: output type
+                        fn call_comb_type<A, O>(acc: &mut A, item: A, f: impl Fn(&mut A, A) -> O) -> O {
+                            f(acc, item)
+                        }
+
                         let #ret_ident = if let ::std::option::Option::Some(mut #accumulator_ident) = #accumulator_ident {
                             for #iterator_item_ident in #input {
                                 #[allow(clippy::redundant_closure_call)]
-                                (#func)(&mut #accumulator_ident, #iterator_item_ident);
+                                call_comb_type(&mut #accumulator_ident, #iterator_item_ident, #func);
                             }
 
                             ::std::option::Option::Some(#accumulator_ident)

diff --git a/hydroflow_plus/Cargo.toml b/hydroflow_plus/Cargo.toml
@@ -25,3 +25,6 @@ hydroflow_lang = { path = "../hydroflow_lang", version = "^0.5.0" }
 serde = { version = "1", features = [ "derive" ] }
 bincode = "1.3"
 stageleft = { path = "../stageleft", version = "^0.1.0" }
+
+[build-dependencies]
+stageleft_tool = { path = "../stageleft_tool", version = "^0.1.0" }
diff --git a/hydroflow_plus/build.rs b/hydroflow_plus/build.rs
@@ -0,0 +1,3 @@
+fn main() {
+    stageleft_tool::gen_final!();
+}
diff --git a/hydroflow_plus/src/builder.rs b/hydroflow_plus/src/builder.rs
@@ -55,25 +55,25 @@ impl<'a, D: LocalDeploy<'a>> FlowBuilder<'a, D> {
         (&self.next_id, &self.builders)
     }
 
-    pub fn process(&'a self, builder: &impl ProcessSpec<'a, D>) -> D::Process {
+    pub fn process(&'a self, spec: &impl ProcessSpec<'a, D>) -> D::Process {
         let mut next_node_id = self.next_node_id.borrow_mut();
         let id = *next_node_id;
         *next_node_id += 1;
 
-        let node = builder.build(id, self, &mut self.meta.borrow_mut());
+        let node = spec.build(id, self, &mut self.meta.borrow_mut());
         self.nodes.borrow_mut().push(node.clone());
 
         self.update_metas();
 
         node
     }
 
-    pub fn cluster(&'a self, builder: &impl ClusterSpec<'a, D>) -> D::Cluster {
+    pub fn cluster(&'a self, spec: &impl ClusterSpec<'a, D>) -> D::Cluster {
         let mut next_node_id = self.next_node_id.borrow_mut();
         let id = *next_node_id;
         *next_node_id += 1;
 
-        let cluster = builder.build(id, self, &mut self.meta.borrow_mut());
+        let cluster = spec.build(id, self, &mut self.meta.borrow_mut());
         self.clusters.borrow_mut().push(cluster.clone());
 
         self.update_metas();

diff --git a/hydroflow_plus/src/cycle.rs b/hydroflow_plus/src/cycle.rs
@@ -20,7 +20,8 @@ pub struct HfCycle<'a, T, W, N: Location<'a>> {
 impl<'a, T, W, N: Location<'a>> HfCycle<'a, T, W, N> {
     pub fn complete(self, stream: &Stream<'a, T, W, N>) {
         let ident = self.ident;
-        let stream_ident = stream.ident.clone();
+        // TODO(shadaj): avoid having to concretize within cycles
+        let stream_ident = stream.ensure_concrete().ident.clone();
 
         self.builders
             .borrow_mut()

diff --git a/hydroflow_plus/src/lib.rs b/hydroflow_plus/src/lib.rs
@@ -1,3 +1,5 @@
+stageleft::stageleft_no_entry_crate!();
+
 use std::marker::PhantomData;
 
 use hydroflow::scheduled::context::Context;

diff --git a/hydroflow_plus/src/location/mod.rs b/hydroflow_plus/src/location/mod.rs
@@ -1,18 +1,19 @@
 use std::cell::RefCell;
 use std::io;
 use std::marker::PhantomData;
+use std::time::Duration;
 
 use hydroflow::bytes::BytesMut;
 use hydroflow::futures::stream::Stream as FuturesStream;
 use proc_macro2::Span;
-use stageleft::Quoted;
+use stageleft::{q, Quoted};
 use syn::parse_quote;
 
 use crate::builder::Builders;
 use crate::stream::{Async, Windowed};
 use crate::{FlowBuilder, HfCycle, Stream};
 
-mod graphs;
+pub mod graphs;
 pub use graphs::*;
 
 pub mod network;
@@ -72,6 +73,48 @@ pub trait Location<'a>: Clone {
 
     fn update_meta(&mut self, meta: &Self::Meta);
 
+    fn spin(&self) -> Stream<'a, (), Async, Self> {
+        let (next_id_cell, builders) = self.flow_builder();
+
+        let next_id = {
+            let mut next_id = next_id_cell.borrow_mut();
+            let id = *next_id;
+            *next_id += 1;
+            id
+        };
+
+        let ident = syn::Ident::new(&format!("stream_{}", next_id), Span::call_site());
+
+        builders
+            .borrow_mut()
+            .as_mut()
+            .unwrap()
+            .entry(self.id())
+            .or_default()
+            .add_statement(parse_quote! {
+                #ident = spin() -> tee();
+            });
+
+        Stream {
+            ident,
+            node: self.clone(),
+            next_id: next_id_cell,
+            builders,
+            is_delta: false,
+            _phantom: PhantomData,
+        }
+    }
+
+    fn spin_batch(
+        &self,
+        batch_size: impl Quoted<'a, usize> + Copy + 'a,
+    ) -> Stream<'a, (), Windowed, Self> {
+        self.spin()
+            .flat_map(q!(move |_| 0..batch_size))
+            .map(q!(|_| ()))
+            .tick_batch()
+    }
+
     fn source_stream<T, E: FuturesStream<Item = T> + Unpin>(
         &self,
         e: impl Quoted<'a, E>,
@@ -103,6 +146,7 @@ pub trait Location<'a>: Clone {
             node: self.clone(),
             next_id: next_id_cell,
             builders,
+            is_delta: false,
             _phantom: PhantomData,
         }
     }
@@ -146,6 +190,7 @@ pub trait Location<'a>: Clone {
                 node: self.clone(),
                 next_id: next_id_cell,
                 builders,
+                is_delta: false,
                 _phantom: PhantomData,
             },
         )
@@ -182,6 +227,43 @@ pub trait Location<'a>: Clone {
             node: self.clone(),
             next_id: next_id_cell,
             builders,
+            is_delta: false,
+            _phantom: PhantomData,
+        }
+    }
+
+    fn source_interval(
+        &self,
+        interval: impl Quoted<'a, Duration> + Copy + 'a,
+    ) -> Stream<'a, hydroflow::tokio::time::Instant, Async, Self> {
+        let (next_id_cell, builders) = self.flow_builder();
+
+        let next_id = {
+            let mut next_id = next_id_cell.borrow_mut();
+            let id = *next_id;
+            *next_id += 1;
+            id
+        };
+
+        let ident = syn::Ident::new(&format!("stream_{}", next_id), Span::call_site());
+        let interval = interval.splice();
+
+        builders
+            .borrow_mut()
+            .as_mut()
+            .unwrap()
+            .entry(self.id())
+            .or_default()
+            .add_statement(parse_quote! {
+                #ident = source_interval(#interval) -> tee();
+            });
+
+        Stream {
+            ident,
+            node: self.clone(),
+            next_id: next_id_cell,
+            builders,
+            is_delta: false,
             _phantom: PhantomData,
         }
     }
@@ -220,6 +302,7 @@ pub trait Location<'a>: Clone {
                 node: self.clone(),
                 next_id: next_id_cell,
                 builders,
+                is_delta: false,
                 _phantom: PhantomData,
             },
         )