hydro_lang/compile/
builder.rs

1use std::any::type_name;
2use std::cell::RefCell;
3use std::marker::PhantomData;
4use std::rc::{Rc, Weak};
5
6use slotmap::{SecondaryMap, SlotMap};
7
8#[cfg(feature = "build")]
9use super::compiled::CompiledFlow;
10#[cfg(feature = "build")]
11use super::deploy::{DeployFlow, DeployResult};
12#[cfg(feature = "build")]
13use super::deploy_provider::{ClusterSpec, Deploy, ExternalSpec, IntoProcessSpec};
14#[cfg(feature = "build")]
15use super::ir::HydroIrOpMetadata;
16use super::ir::{HydroNode, HydroRoot};
17use crate::location::{Cluster, External, LocationKey, LocationType, Process};
18
19/// A compile-time directive to spawn a future on a location's `LocalSet`
20/// alongside the DFIR scheduler.
21pub enum Sidecar {
22    /// A ready-to-go future expression (e.g. telemetry metrics collection).
23    Simple {
24        location_key: LocationKey,
25        future_expr: Box<syn::Expr>,
26    },
27    /// A user-owned sidecar that returns a `(Stream, Sink)` pair to the framework.
28    /// The closure is called at startup; the returned stream feeds items into the
29    /// dataflow and the returned sink receives items from the dataflow.
30    Bidi {
31        location_key: LocationKey,
32        sidecar_id: SidecarId,
33        sidecar_closure: Box<syn::Expr>,
34    },
35}
36#[cfg(feature = "sim")]
37#[cfg(stageleft_runtime)]
38use crate::sim::flow::SimFlow;
39use crate::staging_util::Invariant;
40
41#[stageleft::export(ExternalPortId, CycleId, ClockId, SidecarId, StmtId, HandoffId)]
42crate::newtype_counter! {
43    /// ID for an external output.
44    pub struct ExternalPortId(usize);
45
46    /// ID for a [`crate::location::Location::forward_ref`] cycle.
47    pub struct CycleId(usize);
48
49    /// ID for clocks (ticks).
50    pub struct ClockId(usize);
51
52    /// ID for user-owned sidecars.
53    pub struct SidecarId(usize);
54
55    /// ID for a statement in the emitted DFIR graph.
56    pub struct StmtId(usize);
57
58    /// ID for a handoff channel in the simulator.
59    pub struct HandoffId(usize);
60}
61
62impl CycleId {
63    #[cfg(feature = "build")]
64    pub(crate) fn as_ident(&self) -> syn::Ident {
65        syn::Ident::new(&format!("cycle_{}", self), proc_macro2::Span::call_site())
66    }
67}
68
69impl SidecarId {
70    /// Derives the two idents for a bidi sidecar: `(stream, sink)`.
71    pub fn idents(&self) -> (syn::Ident, syn::Ident) {
72        let span = proc_macro2::Span::call_site();
73        (
74            syn::Ident::new(&format!("__hydro_sidecar_{}_stream", self), span),
75            syn::Ident::new(&format!("__hydro_sidecar_{}_sink", self), span),
76        )
77    }
78}
79
80pub(crate) type FlowState = Rc<RefCell<FlowStateInner>>;
81
82pub(crate) struct FlowStateInner {
83    /// Tracks the roots of the dataflow IR. This is referenced by
84    /// `Stream` and `HfCycle` to build the IR. The inner option will
85    /// be set to `None` when this builder is finalized.
86    roots: Option<Vec<HydroRoot>>,
87
88    /// Counter for generating unique external output identifiers.
89    next_external_port: crate::Counter<ExternalPortId>,
90
91    /// Counters for generating identifiers for cycles.
92    next_cycle_id: crate::Counter<CycleId>,
93
94    /// Counters for clock IDs.
95    next_clock_id: crate::Counter<ClockId>,
96
97    /// Counter for generating unique sidecar identifiers, not used for anything else.
98    next_sidecar_id: crate::Counter<SidecarId>,
99
100    /// Compile-time sidecar directives. Processed during compilation,
101    /// not part of the dataflow IR.
102    pub sidecars: Vec<Sidecar>,
103
104    /// Weak references to the IR nodes of all live collections (streams, singletons, ...)
105    /// created against this flow. When the flow is finalized, any collection that is still
106    /// alive (its `Rc` has not been dropped) has its IR yanked and registered as a root,
107    /// so that dataflow with side effects (e.g. `inspect`) is not silently lost just
108    /// because the collection was dropped after finalization.
109    pub(crate) live_collection_nodes: Vec<Weak<RefCell<HydroNode>>>,
110}
111
112impl FlowStateInner {
113    pub fn next_external_port(&mut self) -> ExternalPortId {
114        self.next_external_port.get_and_increment()
115    }
116
117    pub fn next_cycle_id(&mut self) -> CycleId {
118        self.next_cycle_id.get_and_increment()
119    }
120
121    pub fn next_clock_id(&mut self) -> ClockId {
122        self.next_clock_id.get_and_increment()
123    }
124
125    pub fn next_sidecar_id(&mut self) -> SidecarId {
126        self.next_sidecar_id.get_and_increment()
127    }
128
129    pub fn push_root(&mut self, root: HydroRoot) {
130        self.roots
131            .as_mut()
132            .expect("Attempted to add a root to a flow that has already been finalized. No roots can be added after the flow has been compiled.")
133            .push(root);
134    }
135
136    pub fn try_push_root(&mut self, root: HydroRoot) {
137        if let Some(roots) = self.roots.as_mut() {
138            roots.push(root);
139        }
140    }
141}
142
143pub struct FlowBuilder<'a> {
144    /// Hydro IR and associated counters
145    flow_state: FlowState,
146
147    /// Locations and their type.
148    locations: SlotMap<LocationKey, LocationType>,
149    /// Map from raw location ID to name (including externals).
150    location_names: SecondaryMap<LocationKey, String>,
151    /// The program version each location belongs to. Every location has an entry (0 unless it is a
152    /// `next_version` successor); populated eagerly at location creation.
153    #[cfg(feature = "sim")]
154    location_version: SecondaryMap<LocationKey, u32>,
155    /// Maps each location to the root key of its cross-version correspondence group: version 0 of
156    /// the same logical location. Every location has an entry (its own key unless it is a
157    /// `next_version` successor); populated eagerly at location creation.
158    #[cfg(feature = "sim")]
159    location_version_group_root: SecondaryMap<LocationKey, LocationKey>,
160
161    /// Application name used in telemetry.
162    #[cfg_attr(
163        not(feature = "build"),
164        expect(dead_code, reason = "unused without build")
165    )]
166    flow_name: String,
167
168    /// Tracks whether this flow has been finalized; it is an error to
169    /// drop without finalizing.
170    finalized: bool,
171
172    /// 'a on a FlowBuilder is used to ensure that staged code does not
173    /// capture more data that it is allowed to; 'a is generated at the
174    /// entrypoint of the staged code and we keep it invariant here
175    /// to enforce the appropriate constraints
176    _phantom: Invariant<'a>,
177}
178
179impl Drop for FlowBuilder<'_> {
180    fn drop(&mut self) {
181        if !self.finalized && !std::thread::panicking() {
182            panic!(
183                "Dropped FlowBuilder without finalizing, you may have forgotten to call `with_default_optimize`, `optimize_with`, or `finalize`."
184            );
185        }
186    }
187}
188
189#[expect(missing_docs, reason = "TODO")]
190impl<'a> FlowBuilder<'a> {
191    /// Creates a new `FlowBuilder` to construct a Hydro program, using the Cargo package name as the program name.
192    #[expect(
193        clippy::new_without_default,
194        reason = "call `new` explicitly, not `default`"
195    )]
196    pub fn new() -> Self {
197        let mut name = std::env::var("CARGO_PKG_NAME").unwrap_or_else(|_| "unknown".to_owned());
198        if let Ok(bin_path) = std::env::current_exe()
199            && let Some(bin_name) = bin_path.file_stem()
200        {
201            name = format!("{}/{}", name, bin_name.display());
202        }
203        Self::with_name(name)
204    }
205
206    /// Creates a new `FlowBuilder` to construct a Hydro program, with the given program name.
207    pub fn with_name(name: impl Into<String>) -> Self {
208        Self {
209            flow_state: Rc::new(RefCell::new(FlowStateInner {
210                roots: Some(vec![]),
211                next_external_port: crate::Counter::default(),
212                next_cycle_id: crate::Counter::default(),
213                next_clock_id: crate::Counter::default(),
214                next_sidecar_id: crate::Counter::default(),
215                sidecars: Vec::new(),
216                live_collection_nodes: Vec::new(),
217            })),
218            locations: SlotMap::with_key(),
219            location_names: SecondaryMap::new(),
220            #[cfg(feature = "sim")]
221            location_version: SecondaryMap::new(),
222            #[cfg(feature = "sim")]
223            location_version_group_root: SecondaryMap::new(),
224            flow_name: name.into(),
225            finalized: false,
226            _phantom: PhantomData,
227        }
228    }
229
230    pub(crate) fn flow_state(&self) -> &FlowState {
231        &self.flow_state
232    }
233
234    fn insert_location(&mut self, ty: LocationType, name: String) -> LocationKey {
235        let key = self.locations.insert(ty);
236        self.location_names.insert(key, name);
237        #[cfg(feature = "sim")]
238        {
239            self.location_version.insert(key, 0);
240            self.location_version_group_root.insert(key, key);
241        }
242        key
243    }
244
245    pub fn process<P>(&mut self) -> Process<'a, P> {
246        let key = self.insert_location(LocationType::Process, type_name::<P>().to_owned());
247        Process {
248            key,
249            flow_state: self.flow_state().clone(),
250            _phantom: PhantomData,
251        }
252    }
253
254    pub fn cluster<C>(&mut self) -> Cluster<'a, C> {
255        let key = self.insert_location(LocationType::Cluster, type_name::<C>().to_owned());
256        Cluster {
257            key,
258            flow_state: self.flow_state().clone(),
259            _phantom: PhantomData,
260        }
261    }
262
263    pub fn external<E>(&mut self) -> External<'a, E> {
264        let key = self.insert_location(LocationType::External, type_name::<E>().to_owned());
265        External {
266            key,
267            flow_state: self.flow_state().clone(),
268            _phantom: PhantomData,
269        }
270    }
271
272    #[cfg(feature = "sim")]
273    pub fn next_version<C>(&mut self, cluster: &Cluster<'a, C>) -> Cluster<'a, C> {
274        let group_root = self.location_version_group_root[cluster.key];
275        let version = self
276            .location_version_group_root
277            .values()
278            .filter(|&&r| r == group_root)
279            .count() as u32;
280        let key = self.insert_location(LocationType::Cluster, type_name::<C>().to_owned());
281        self.location_version.insert(key, version);
282        self.location_version_group_root.insert(key, group_root);
283        Cluster {
284            key,
285            flow_state: self.flow_state().clone(),
286            _phantom: PhantomData,
287        }
288    }
289}
290
291#[cfg(feature = "build")]
292#[cfg_attr(docsrs, doc(cfg(feature = "build")))]
293#[expect(missing_docs, reason = "TODO")]
294impl<'a> FlowBuilder<'a> {
295    pub fn finalize(mut self) -> super::built::BuiltFlow<'a> {
296        self.finalized = true;
297
298        let mut flow_state = self.flow_state.borrow_mut();
299
300        // Yank the IR from any live collections (streams, singletons, ...) that are still
301        // alive, since their `Drop` will run after finalization and would otherwise
302        // silently fail to register their dataflow as a root.
303        let live_collection_nodes = std::mem::take(&mut flow_state.live_collection_nodes);
304        for node_cell in live_collection_nodes {
305            if let Some(node_cell) = node_cell.upgrade() {
306                let ir_node = node_cell.replace(HydroNode::Placeholder);
307                if !matches!(ir_node, HydroNode::Placeholder) && !ir_node.is_shared_with_others() {
308                    flow_state.push_root(HydroRoot::Null {
309                        input: Box::new(ir_node),
310                        op_metadata: HydroIrOpMetadata::new(),
311                    });
312                }
313            }
314        }
315
316        let mut ir = flow_state.roots.take().unwrap();
317        let sidecars = std::mem::take(&mut flow_state.sidecars);
318        drop(flow_state);
319
320        super::ir::unify_atomic_ticks(&mut ir);
321
322        super::built::BuiltFlow {
323            ir,
324            locations: std::mem::take(&mut self.locations),
325            location_names: std::mem::take(&mut self.location_names),
326            sidecars,
327            flow_name: std::mem::take(&mut self.flow_name),
328            #[cfg(feature = "sim")]
329            location_version: std::mem::take(&mut self.location_version),
330            #[cfg(feature = "sim")]
331            location_version_group_root: std::mem::take(&mut self.location_version_group_root),
332            _phantom: PhantomData,
333        }
334    }
335
336    pub fn with_default_optimize<D: Deploy<'a>>(self) -> DeployFlow<'a, D> {
337        self.finalize().with_default_optimize()
338    }
339
340    pub fn optimize_with(self, f: impl FnOnce(&mut [HydroRoot])) -> super::built::BuiltFlow<'a> {
341        self.finalize().optimize_with(f)
342    }
343
344    pub fn with_process<P, D: Deploy<'a>>(
345        self,
346        process: &Process<P>,
347        spec: impl IntoProcessSpec<'a, D>,
348    ) -> DeployFlow<'a, D> {
349        self.with_default_optimize().with_process(process, spec)
350    }
351
352    pub fn with_remaining_processes<D: Deploy<'a>, S: IntoProcessSpec<'a, D> + 'a>(
353        self,
354        spec: impl Fn() -> S,
355    ) -> DeployFlow<'a, D> {
356        self.with_default_optimize().with_remaining_processes(spec)
357    }
358
359    pub fn with_external<P, D: Deploy<'a>>(
360        self,
361        process: &External<P>,
362        spec: impl ExternalSpec<'a, D>,
363    ) -> DeployFlow<'a, D> {
364        self.with_default_optimize().with_external(process, spec)
365    }
366
367    pub fn with_remaining_externals<D: Deploy<'a>, S: ExternalSpec<'a, D> + 'a>(
368        self,
369        spec: impl Fn() -> S,
370    ) -> DeployFlow<'a, D> {
371        self.with_default_optimize().with_remaining_externals(spec)
372    }
373
374    pub fn with_cluster<C, D: Deploy<'a>>(
375        self,
376        cluster: &Cluster<C>,
377        spec: impl ClusterSpec<'a, D>,
378    ) -> DeployFlow<'a, D> {
379        self.with_default_optimize().with_cluster(cluster, spec)
380    }
381
382    pub fn with_remaining_clusters<D: Deploy<'a>, S: ClusterSpec<'a, D> + 'a>(
383        self,
384        spec: impl Fn() -> S,
385    ) -> DeployFlow<'a, D> {
386        self.with_default_optimize().with_remaining_clusters(spec)
387    }
388
389    pub fn compile<D: Deploy<'a, InstantiateEnv = ()>>(self) -> CompiledFlow<'a> {
390        self.with_default_optimize::<D>().compile()
391    }
392
393    pub fn deploy<D: Deploy<'a>>(self, env: &mut D::InstantiateEnv) -> DeployResult<'a, D> {
394        self.with_default_optimize().deploy(env)
395    }
396
397    #[cfg(feature = "sim")]
398    /// Creates a simulation for this builder, which can be used to run deterministic simulations
399    /// of the Hydro program.
400    pub fn sim(self) -> SimFlow<'a> {
401        self.finalize().sim()
402    }
403
404    pub fn from_built<'b>(built: &super::built::BuiltFlow) -> FlowBuilder<'b> {
405        FlowBuilder {
406            flow_state: Rc::new(RefCell::new(FlowStateInner {
407                roots: None,
408                next_external_port: crate::Counter::default(),
409                next_cycle_id: crate::Counter::default(),
410                next_clock_id: crate::Counter::default(),
411                next_sidecar_id: crate::Counter::default(),
412                sidecars: Vec::new(),
413                live_collection_nodes: Vec::new(),
414            })),
415            locations: built.locations.clone(),
416            location_names: built.location_names.clone(),
417            #[cfg(feature = "sim")]
418            location_version: built.location_version.clone(),
419            #[cfg(feature = "sim")]
420            location_version_group_root: built.location_version_group_root.clone(),
421            flow_name: built.flow_name.clone(),
422            finalized: false,
423            _phantom: PhantomData,
424        }
425    }
426
427    #[doc(hidden)] // TODO(mingwei): This is an unstable API for now
428    pub fn replace_ir(&mut self, roots: Vec<HydroRoot>) {
429        self.flow_state.borrow_mut().roots = Some(roots);
430    }
431
432    #[doc(hidden)] // TODO(mingwei): This is an unstable API for now
433    pub fn next_clock_id(&mut self) -> ClockId {
434        self.flow_state.borrow_mut().next_clock_id()
435    }
436
437    #[doc(hidden)] // TODO(mingwei): This is an unstable API for now
438    pub fn next_cycle_id(&mut self) -> CycleId {
439        self.flow_state.borrow_mut().next_cycle_id()
440    }
441}
hydro_lang/compile/builder.rs

hydro_lang/compile/
builder.rs
