feat(hydro_lang): support in-memory inputs for embedded Hydro programs

Author: shadaj

hydro_lang/src/compile/builder.rs

@@ -231,7 +231,7 @@ impl<'a> FlowBuilder<'a> {
         self.with_default_optimize().with_remaining_clusters(spec)
     }
 
-    pub fn compile<D: Deploy<'a>>(self) -> CompiledFlow<'a> {
+    pub fn compile<D: Deploy<'a, InstantiateEnv = ()>>(self) -> CompiledFlow<'a> {
         self.with_default_optimize::<D>().compile()
     }
 

hydro_lang/src/compile/built.rs

@@ -316,7 +316,7 @@ impl<'a> BuiltFlow<'a> {
         self.into_deploy().with_remaining_clusters(spec)
     }
 
-    pub fn compile<D: Deploy<'a>>(self) -> CompiledFlow<'a> {
+    pub fn compile<D: Deploy<'a, InstantiateEnv = ()>>(self) -> CompiledFlow<'a> {
         self.into_deploy::<D>().compile()
     }
 

hydro_lang/src/compile/deploy.rs

@@ -262,14 +262,17 @@ impl<'a, D: Deploy<'a>> DeployFlow<'a, D> {
     /// Compiles the flow into DFIR ([`dfir_lang::graph::DfirGraph`]) including networking.
     ///
     /// (This does not compile the DFIR itself, instead use [`Self::deploy`] to compile & deploy the DFIR).
-    pub fn compile(mut self) -> CompiledFlow<'a> {
-        self.compile_internal()
+    pub fn compile(mut self) -> CompiledFlow<'a>
+    where
+        D: Deploy<'a, InstantiateEnv = ()>,
+    {
+        self.compile_internal(&mut ())
     }
 
     /// Same as [`Self::compile`] but does not invalidate `self`, for internal use.
     ///
     /// Empties `self.sidecars` and modifies `self.ir`, leaving `self` in a partial state.
-    pub(super) fn compile_internal(&mut self) -> CompiledFlow<'a> {
+    pub(super) fn compile_internal(&mut self, env: &mut D::InstantiateEnv) -> CompiledFlow<'a> {
         let mut seen_tees: HashMap<_, _> = HashMap::new();
         let mut extra_stmts = SparseSecondaryMap::new();
         for leaf in self.ir.iter_mut() {
@@ -279,6 +282,7 @@ impl<'a, D: Deploy<'a>> DeployFlow<'a, D> {
                 &self.processes,
                 &self.clusters,
                 &self.externals,
+                env,
             );
         }
 
@@ -348,7 +352,7 @@ impl<'a, D: Deploy<'a>> DeployFlow<'a, D> {
             mut extra_stmts,
             mut sidecars,
             _phantom,
-        } = self.compile_internal();
+        } = self.compile_internal(env);
 
         let mut compiled = dfir;
         self.cluster_id_stmts(&mut extra_stmts);

hydro_lang/src/compile/deploy_provider.rs

@@ -118,6 +118,18 @@ pub trait Deploy<'a> {
     fn cluster_membership_stream(
         location_id: &LocationId,
     ) -> impl QuotedWithContext<'a, Box<dyn Stream<Item = (TaglessMemberId, MembershipEvent)> + Unpin>, ()>;
+
+    /// Registers an embedded input for the given ident and element type.
+    ///
+    /// Only meaningful for the embedded deployment backend. The default
+    /// implementation panics.
+    fn register_embedded_input(
+        _env: &mut Self::InstantiateEnv,
+        _ident: &syn::Ident,
+        _element_type: &syn::Type,
+    ) {
+        panic!("register_embedded_input is only supported by EmbeddedDeploy");
+    }
 }
 
 pub trait ProcessSpec<'a, D>

hydro_lang/src/compile/embedded.rs

@@ -48,7 +48,7 @@ pub struct EmbeddedNode {
 impl Node for EmbeddedNode {
     type Port = ();
     type Meta = ();
-    type InstantiateEnv = ();
+    type InstantiateEnv = EmbeddedInstantiateEnv;
 
     fn next_port(&self) -> Self::Port {}
 
@@ -140,9 +140,20 @@ impl<S: Into<String>> ExternalSpec<'_, EmbeddedDeploy> for S {
     }
 }
 
+/// Collected embedded input registrations.
+///
+/// During `compile_network`, each `EmbeddedInput` IR node registers its ident
+/// and element type here. `generate_embedded` then uses this to add parameters
+/// to the generated functions.
+#[derive(Default)]
+pub struct EmbeddedInstantiateEnv {
+    /// (ident name, element type) pairs collected during compilation.
+    pub inputs: Vec<(syn::Ident, syn::Type)>,
+}
+
 impl<'a> Deploy<'a> for EmbeddedDeploy {
     type Meta = ();
-    type InstantiateEnv = ();
+    type InstantiateEnv = EmbeddedInstantiateEnv;
 
     type Process = EmbeddedNode;
     type Cluster = EmbeddedNode;
@@ -302,6 +313,14 @@ impl<'a> Deploy<'a> for EmbeddedDeploy {
             "EmbeddedDeploy does not support cluster membership streams"
         ))
     }
+
+    fn register_embedded_input(
+        env: &mut Self::InstantiateEnv,
+        ident: &syn::Ident,
+        element_type: &syn::Type,
+    ) {
+        env.inputs.push((ident.clone(), element_type.clone()));
+    }
 }
 
 impl super::deploy::DeployFlow<'_, EmbeddedDeploy> {
@@ -336,7 +355,12 @@ impl super::deploy::DeployFlow<'_, EmbeddedDeploy> {
     /// include!(concat!(env!("OUT_DIR"), "/embedded.rs"));
     /// ```
     pub fn generate_embedded(mut self, crate_name: &str) -> syn::File {
-        let compiled = self.compile_internal();
+        let mut env = EmbeddedInstantiateEnv::default();
+        let compiled = self.compile_internal(&mut env);
+
+        // Sort inputs by name for deterministic output.
+        let mut inputs = env.inputs;
+        inputs.sort_by(|a, b| a.0.to_string().cmp(&b.0.to_string()));
 
         let root = crate::staging_util::get_this_crate();
         let orig_crate_name = quote::format_ident!("{}", crate_name.replace('-', "_"));
@@ -347,6 +371,14 @@ impl super::deploy::DeployFlow<'_, EmbeddedDeploy> {
         let mut location_keys: Vec<_> = compiled.all_dfir().keys().collect();
         location_keys.sort();
 
+        // Build the input parameters for each generated function.
+        let input_params: Vec<proc_macro2::TokenStream> = inputs
+            .iter()
+            .map(|(ident, element_type)| {
+                quote! { #ident: impl __root_dfir_rs::futures::Stream<Item = #element_type> + Unpin + 'a }
+            })
+            .collect();
+
         for location_key in location_keys {
             let graph = &compiled.all_dfir()[location_key];
 
@@ -368,7 +400,7 @@ impl super::deploy::DeployFlow<'_, EmbeddedDeploy> {
 
             let func: syn::Item = syn::parse_quote! {
                 #[allow(unused, non_snake_case, clippy::suspicious_else_formatting)]
-                pub fn #fn_ident<'a>() -> #root::runtime_support::dfir_rs::scheduled::graph::Dfir<'a> {
+                pub fn #fn_ident<'a>(#(#input_params),*) -> #root::runtime_support::dfir_rs::scheduled::graph::Dfir<'a> {
                     #dfir_tokens
                 }
             };

hydro_lang/src/compile/ir/mod.rs

@@ -685,10 +685,12 @@ impl HydroRoot {
         processes: &SparseSecondaryMap<LocationKey, D::Process>,
         clusters: &SparseSecondaryMap<LocationKey, D::Cluster>,
         externals: &SparseSecondaryMap<LocationKey, D::External>,
+        env: &mut D::InstantiateEnv,
     ) where
         D: Deploy<'a>,
     {
         let refcell_extra_stmts = RefCell::new(extra_stmts);
+        let refcell_env = RefCell::new(env);
         self.transform_bottom_up(
             &mut |l| {
                 if let HydroRoot::SendExternal {
@@ -887,6 +889,16 @@ impl HydroRoot {
                         connect_fn: Some(connect_fn),
                     }
                     .into();
+                } else if let HydroNode::EmbeddedInput { ident, metadata } = n {
+                    let element_type = match &metadata.collection_kind {
+                        CollectionKind::Stream { element_type, .. } => element_type.0.as_ref().clone(),
+                        _ => panic!("EmbeddedInput must have Stream collection kind"),
+                    };
+                    D::register_embedded_input(
+                        &mut refcell_env.borrow_mut(),
+                        ident,
+                        &element_type,
+                    );
                 }
             },
             seen_tees,
@@ -1721,6 +1733,15 @@ pub enum HydroNode {
         input: Box<HydroNode>,
         metadata: HydroIrMetadata,
     },
+
+    /// An external input for embedded deployment mode.
+    ///
+    /// This node compiles to `source_stream(ident)` where `ident` is a parameter
+    /// added to the generated function signature.
+    EmbeddedInput {
+        ident: syn::Ident,
+        metadata: HydroIrMetadata,
+    },
 }
 
 pub type SeenTees = HashMap<*const RefCell<HydroNode>, Rc<RefCell<HydroNode>>>;
@@ -1776,7 +1797,8 @@ impl HydroNode {
             HydroNode::Source { .. }
             | HydroNode::SingletonSource { .. }
             | HydroNode::CycleSource { .. }
-            | HydroNode::ExternalInput { .. } => {}
+            | HydroNode::ExternalInput { .. }
+            | HydroNode::EmbeddedInput { .. } => {}
 
             HydroNode::Tee { inner, .. } => {
                 if let Some(transformed) = seen_tees.get(&inner.as_ptr()) {
@@ -2120,6 +2142,10 @@ impl HydroNode {
                 input: Box::new(input.deep_clone(seen_tees)),
                 metadata: metadata.clone(),
             },
+            HydroNode::EmbeddedInput { ident, metadata } => HydroNode::EmbeddedInput {
+                ident: ident.clone(),
+                metadata: metadata.clone(),
+            },
         }
     }
 
@@ -3430,6 +3456,33 @@ impl HydroNode {
 
                         ident_stack.push(counter_ident);
                     }
+
+                    HydroNode::EmbeddedInput { ident, .. } => {
+                        let source_ident =
+                            syn::Ident::new(&format!("stream_{}", *next_stmt_id), Span::call_site());
+
+                        let ident = ident.clone();
+
+                        match builders_or_callback {
+                            BuildersOrCallback::Builders(graph_builders) => {
+                                let builder = graph_builders.get_dfir_mut(&out_location);
+                                builder.add_dfir(
+                                    parse_quote! {
+                                        #source_ident = source_stream(#ident);
+                                    },
+                                    None,
+                                    Some(&next_stmt_id.to_string()),
+                                );
+                            }
+                            BuildersOrCallback::Callback(_, node_callback) => {
+                                node_callback(node, next_stmt_id);
+                            }
+                        }
+
+                        *next_stmt_id += 1;
+
+                        ident_stack.push(source_ident);
+                    }
                 }
             },
             seen_tees,
@@ -3508,6 +3561,7 @@ impl HydroNode {
                     transform(deserialize_fn);
                 }
             }
+            HydroNode::EmbeddedInput { .. } => {}
             HydroNode::Counter { duration, .. } => {
                 transform(duration);
             }
@@ -3560,6 +3614,7 @@ impl HydroNode {
             HydroNode::ExternalInput { metadata, .. } => metadata,
             HydroNode::Network { metadata, .. } => metadata,
             HydroNode::Counter { metadata, .. } => metadata,
+            HydroNode::EmbeddedInput { metadata, .. } => metadata,
         }
     }
 
@@ -3609,6 +3664,7 @@ impl HydroNode {
             HydroNode::ExternalInput { metadata, .. } => metadata,
             HydroNode::Network { metadata, .. } => metadata,
             HydroNode::Counter { metadata, .. } => metadata,
+            HydroNode::EmbeddedInput { metadata, .. } => metadata,
         }
     }
 
@@ -3621,6 +3677,7 @@ impl HydroNode {
             | HydroNode::SingletonSource { .. }
             | HydroNode::ExternalInput { .. }
             | HydroNode::CycleSource { .. }
+            | HydroNode::EmbeddedInput { .. }
             | HydroNode::Tee { .. } => {
                 // Tee should find its input in separate special ways
                 vec![]
@@ -3752,6 +3809,7 @@ impl HydroNode {
             HydroNode::Counter { tag, duration, .. } => {
                 format!("Counter({:?}, {:?})", tag, duration)
             }
+            HydroNode::EmbeddedInput { ident, .. } => format!("EmbeddedInput({})", ident),
         }
     }
 }

hydro_lang/src/location/mod.rs

@@ -520,6 +520,35 @@ pub trait Location<'a>: dynamic::DynLocation {
         (SimSender(external.port_id, PhantomData), stream)
     }
 
+    /// Creates an external input stream for embedded deployment mode.
+    ///
+    /// The `name` parameter specifies the name of the generated function parameter
+    /// that will supply data to this stream at runtime. The generated function will
+    /// accept an `impl Stream<Item = T> + Unpin` argument with this name.
+    fn embedded_input<T>(
+        &self,
+        name: impl Into<String>,
+    ) -> Stream<T, Self, Unbounded, TotalOrder, ExactlyOnce>
+    where
+        Self: Sized + NoTick,
+    {
+        let ident = syn::Ident::new(&name.into(), Span::call_site());
+
+        Stream::new(
+            self.clone(),
+            HydroNode::EmbeddedInput {
+                ident,
+                metadata: self.new_node_metadata(Stream::<
+                    T,
+                    Self,
+                    Unbounded,
+                    TotalOrder,
+                    ExactlyOnce,
+                >::collection_kind()),
+            },
+        )
+    }
+
     /// Establishes a server on this location to receive a bidirectional connection from a single
     /// client, identified by the given `External` handle. Returns a port handle for the external
     /// process to connect to, a stream of incoming messages, and a handle to send outgoing

hydro_lang/src/sim/flow.rs

@@ -110,6 +110,7 @@ impl<'a> SimFlow<'a> {
                 &self.processes,
                 &self.clusters,
                 &self.externals,
+                &mut (),
             );
         });
 

hydro_lang/src/viz/render.rs

@@ -1550,6 +1550,10 @@ impl HydroNode {
                 },
                 duration,
             ),
+
+            HydroNode::EmbeddedInput { ident, metadata } => {
+                build_source_node(structure, metadata, format!("embedded_input({})", ident))
+            }
         }
     }
 }

hydro_test/src/local/capitalize.rs

@@ -0,0 +1,7 @@
+use hydro_lang::prelude::*;
+
+pub fn capitalize<'a>(input: Stream<String, Process<'a, ()>>) {
+    input
+        .map(q!(|s| s.to_uppercase()))
+        .for_each(q!(|s| println!("{}", s)));
+}