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    /// Counter for generating unique simulator hook handle IDs (see
101    /// [`FlowBuilder::sim_hook`]).
102    next_sim_hook_id: usize,
103
104    /// Compile-time sidecar directives. Processed during compilation,
105    /// not part of the dataflow IR.
106    pub sidecars: Vec<Sidecar>,
107
108    /// Weak references to the IR nodes of all live collections (streams, singletons, ...)
109    /// created against this flow. When the flow is finalized, any collection that is still
110    /// alive (its `Rc` has not been dropped) has its IR yanked and registered as a root,
111    /// so that dataflow with side effects (e.g. `inspect`) is not silently lost just
112    /// because the collection was dropped after finalization.
113    pub(crate) live_collection_nodes: Vec<Weak<RefCell<HydroNode>>>,
114}
115
116impl FlowStateInner {
117    pub fn next_external_port(&mut self) -> ExternalPortId {
118        self.next_external_port.get_and_increment()
119    }
120
121    pub fn next_cycle_id(&mut self) -> CycleId {
122        self.next_cycle_id.get_and_increment()
123    }
124
125    pub fn next_clock_id(&mut self) -> ClockId {
126        self.next_clock_id.get_and_increment()
127    }
128
129    pub fn next_sidecar_id(&mut self) -> SidecarId {
130        self.next_sidecar_id.get_and_increment()
131    }
132
133    pub fn next_sim_hook_id(&mut self) -> usize {
134        let id = self.next_sim_hook_id;
135        self.next_sim_hook_id += 1;
136        id
137    }
138
139    pub fn push_root(&mut self, root: HydroRoot) {
140        self.roots
141            .as_mut()
142            .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.")
143            .push(root);
144    }
145
146    pub fn try_push_root(&mut self, root: HydroRoot) {
147        if let Some(roots) = self.roots.as_mut() {
148            roots.push(root);
149        }
150    }
151}
152
153pub struct FlowBuilder<'a> {
154    /// Hydro IR and associated counters
155    flow_state: FlowState,
156
157    /// Locations and their type.
158    locations: SlotMap<LocationKey, LocationType>,
159    /// Map from raw location ID to name (including externals).
160    location_names: SecondaryMap<LocationKey, String>,
161    /// The program version each location belongs to. Every location has an entry (0 unless it is a
162    /// `next_version` successor); populated eagerly at location creation.
163    #[cfg(feature = "sim")]
164    location_version: SecondaryMap<LocationKey, u32>,
165    /// Maps each location to the root key of its cross-version correspondence group: version 0 of
166    /// the same logical location. Every location has an entry (its own key unless it is a
167    /// `next_version` successor); populated eagerly at location creation.
168    #[cfg(feature = "sim")]
169    location_version_group_root: SecondaryMap<LocationKey, LocationKey>,
170
171    /// Application name used in telemetry.
172    #[cfg_attr(
173        not(feature = "build"),
174        expect(dead_code, reason = "unused without build")
175    )]
176    flow_name: String,
177
178    /// Tracks whether this flow has been finalized; it is an error to
179    /// drop without finalizing.
180    finalized: bool,
181
182    /// 'a on a `FlowBuilder` is used to ensure that staged code does not
183    /// capture more data that it is allowed to; 'a is generated at the
184    /// entrypoint of the staged code and we keep it invariant here
185    /// to enforce the appropriate constraints
186    _phantom: Invariant<'a>,
187}
188
189impl Drop for FlowBuilder<'_> {
190    fn drop(&mut self) {
191        assert!(
192            self.finalized || std::thread::panicking(),
193            "Dropped FlowBuilder without finalizing, you may have forgotten to call `with_default_optimize`, `optimize_with`, or `finalize`."
194        );
195    }
196}
197
198#[expect(missing_docs, reason = "TODO")]
199impl<'a> FlowBuilder<'a> {
200    /// Creates a new `FlowBuilder` to construct a Hydro program, using the Cargo package name as the program name.
201    #[expect(
202        clippy::new_without_default,
203        reason = "call `new` explicitly, not `default`"
204    )]
205    pub fn new() -> Self {
206        let mut name = std::env::var("CARGO_PKG_NAME").unwrap_or_else(|_| "unknown".to_owned());
207        if let Ok(bin_path) = std::env::current_exe()
208            && let Some(bin_name) = bin_path.file_stem()
209        {
210            name = format!("{}/{}", name, bin_name.display());
211        }
212        Self::with_name(name)
213    }
214
215    /// Creates a new `FlowBuilder` to construct a Hydro program, with the given program name.
216    pub fn with_name(name: impl Into<String>) -> Self {
217        Self {
218            flow_state: Rc::new(RefCell::new(FlowStateInner {
219                roots: Some(vec![]),
220                next_external_port: crate::Counter::default(),
221                next_cycle_id: crate::Counter::default(),
222                next_clock_id: crate::Counter::default(),
223                next_sidecar_id: crate::Counter::default(),
224                next_sim_hook_id: 0,
225                sidecars: Vec::new(),
226                live_collection_nodes: Vec::new(),
227            })),
228            locations: SlotMap::with_key(),
229            location_names: SecondaryMap::new(),
230            #[cfg(feature = "sim")]
231            location_version: SecondaryMap::new(),
232            #[cfg(feature = "sim")]
233            location_version_group_root: SecondaryMap::new(),
234            flow_name: name.into(),
235            finalized: false,
236            _phantom: PhantomData,
237        }
238    }
239
240    pub(crate) fn flow_state(&self) -> &FlowState {
241        &self.flow_state
242    }
243
244    fn insert_location(&mut self, ty: LocationType, name: String) -> LocationKey {
245        let key = self.locations.insert(ty);
246        self.location_names.insert(key, name);
247        #[cfg(feature = "sim")]
248        {
249            self.location_version.insert(key, 0);
250            self.location_version_group_root.insert(key, key);
251        }
252        key
253    }
254
255    pub fn process<P>(&mut self) -> Process<'a, P> {
256        let key = self.insert_location(LocationType::Process, type_name::<P>().to_owned());
257        Process {
258            key,
259            flow_state: self.flow_state().clone(),
260            _phantom: PhantomData,
261        }
262    }
263
264    pub fn cluster<C>(&mut self) -> Cluster<'a, C> {
265        let key = self.insert_location(LocationType::Cluster, type_name::<C>().to_owned());
266        Cluster {
267            key,
268            flow_state: self.flow_state().clone(),
269            _phantom: PhantomData,
270        }
271    }
272
273    pub fn external<E>(&mut self) -> External<'a, E> {
274        let key = self.insert_location(LocationType::External, type_name::<E>().to_owned());
275        External {
276            key,
277            flow_state: self.flow_state().clone(),
278            _phantom: PhantomData,
279        }
280    }
281
282    /// Creates a fresh set of **simulator hook handles** (see [`crate::sim_hooks`]).
283    ///
284    /// A handle is attached to a specific unsafe operator with the
285    /// `nondet!(/** reason */ hook = handle)` syntax, and lets a simulation test script the
286    /// non-deterministic decisions of that operator. `B` may be a single handle type (e.g.
287    /// `BatchHook<u32>`) or a struct of handles implementing
288    /// [`SimHook`](crate::sim_hooks::SimHook).
289    ///
290    /// Creating a handle and never binding it is allowed (a bundle may be only partially
291    /// used by a particular program configuration), but scripting a decision on an unbound
292    /// handle panics at that call. Binding the same handle to two different operators is an
293    /// error at flow build time.
294    pub fn sim_hook<B: crate::sim_hooks::SimHook>(&mut self) -> B {
295        let flow_state = self.flow_state.clone();
296        B::create(&mut move || flow_state.borrow_mut().next_sim_hook_id())
297    }
298
299    #[cfg(feature = "sim")]
300    pub fn next_version<C>(&mut self, cluster: &Cluster<'a, C>) -> Cluster<'a, C> {
301        let group_root = self.location_version_group_root[cluster.key];
302        let version = self
303            .location_version_group_root
304            .values()
305            .filter(|&&r| r == group_root)
306            .count() as u32;
307        let key = self.insert_location(LocationType::Cluster, type_name::<C>().to_owned());
308        self.location_version.insert(key, version);
309        self.location_version_group_root.insert(key, group_root);
310        Cluster {
311            key,
312            flow_state: self.flow_state().clone(),
313            _phantom: PhantomData,
314        }
315    }
316}
317
318#[cfg(feature = "build")]
319#[cfg_attr(docsrs, doc(cfg(feature = "build")))]
320#[expect(missing_docs, reason = "TODO")]
321impl<'a> FlowBuilder<'a> {
322    pub fn finalize(mut self) -> super::built::BuiltFlow<'a> {
323        self.finalized = true;
324
325        let mut flow_state = self.flow_state.borrow_mut();
326
327        // Yank the IR from any live collections (streams, singletons, ...) that are still
328        // alive, since their `Drop` will run after finalization and would otherwise
329        // silently fail to register their dataflow as a root.
330        let live_collection_nodes = std::mem::take(&mut flow_state.live_collection_nodes);
331        for node_cell in live_collection_nodes {
332            if let Some(node_cell) = node_cell.upgrade() {
333                let ir_node = node_cell.replace(HydroNode::Placeholder);
334                if !matches!(ir_node, HydroNode::Placeholder) && !ir_node.is_shared_with_others() {
335                    flow_state.push_root(HydroRoot::Null {
336                        input: Box::new(ir_node),
337                        op_metadata: HydroIrOpMetadata::new(),
338                    });
339                }
340            }
341        }
342
343        let mut ir = flow_state.roots.take().unwrap();
344        let sidecars = std::mem::take(&mut flow_state.sidecars);
345        drop(flow_state);
346
347        super::ir::unify_atomic_ticks(&mut ir);
348
349        super::built::BuiltFlow {
350            ir,
351            locations: std::mem::take(&mut self.locations),
352            location_names: std::mem::take(&mut self.location_names),
353            sidecars,
354            flow_name: std::mem::take(&mut self.flow_name),
355            #[cfg(feature = "sim")]
356            location_version: std::mem::take(&mut self.location_version),
357            #[cfg(feature = "sim")]
358            location_version_group_root: std::mem::take(&mut self.location_version_group_root),
359            _phantom: PhantomData,
360        }
361    }
362
363    pub fn with_default_optimize<D: Deploy<'a>>(self) -> DeployFlow<'a, D> {
364        self.finalize().with_default_optimize()
365    }
366
367    pub fn optimize_with(self, f: impl FnOnce(&mut [HydroRoot])) -> super::built::BuiltFlow<'a> {
368        self.finalize().optimize_with(f)
369    }
370
371    pub fn with_process<P, D: Deploy<'a>>(
372        self,
373        process: &Process<'_, P>,
374        spec: impl IntoProcessSpec<'a, D>,
375    ) -> DeployFlow<'a, D> {
376        self.with_default_optimize().with_process(process, spec)
377    }
378
379    pub fn with_remaining_processes<D: Deploy<'a>, S: IntoProcessSpec<'a, D> + 'a>(
380        self,
381        spec: impl Fn() -> S,
382    ) -> DeployFlow<'a, D> {
383        self.with_default_optimize().with_remaining_processes(spec)
384    }
385
386    pub fn with_external<P, D: Deploy<'a>>(
387        self,
388        process: &External<'_, P>,
389        spec: impl ExternalSpec<'a, D>,
390    ) -> DeployFlow<'a, D> {
391        self.with_default_optimize().with_external(process, spec)
392    }
393
394    pub fn with_remaining_externals<D: Deploy<'a>, S: ExternalSpec<'a, D> + 'a>(
395        self,
396        spec: impl Fn() -> S,
397    ) -> DeployFlow<'a, D> {
398        self.with_default_optimize().with_remaining_externals(spec)
399    }
400
401    pub fn with_cluster<C, D: Deploy<'a>>(
402        self,
403        cluster: &Cluster<'_, C>,
404        spec: impl ClusterSpec<'a, D>,
405    ) -> DeployFlow<'a, D> {
406        self.with_default_optimize().with_cluster(cluster, spec)
407    }
408
409    pub fn with_remaining_clusters<D: Deploy<'a>, S: ClusterSpec<'a, D> + 'a>(
410        self,
411        spec: impl Fn() -> S,
412    ) -> DeployFlow<'a, D> {
413        self.with_default_optimize().with_remaining_clusters(spec)
414    }
415
416    pub fn compile<D: Deploy<'a, InstantiateEnv = ()>>(self) -> CompiledFlow<'a> {
417        self.with_default_optimize::<D>().compile()
418    }
419
420    pub fn deploy<D: Deploy<'a>>(self, env: &mut D::InstantiateEnv) -> DeployResult<'a, D> {
421        self.with_default_optimize().deploy(env)
422    }
423
424    #[cfg(feature = "sim")]
425    /// Creates a simulation for this builder, which can be used to run deterministic simulations
426    /// of the Hydro program.
427    pub fn sim(self) -> SimFlow<'a> {
428        self.finalize().sim()
429    }
430
431    pub fn from_built<'b>(built: &super::built::BuiltFlow<'_>) -> FlowBuilder<'b> {
432        FlowBuilder {
433            flow_state: Rc::new(RefCell::new(FlowStateInner {
434                roots: None,
435                next_external_port: crate::Counter::default(),
436                next_cycle_id: crate::Counter::default(),
437                next_clock_id: crate::Counter::default(),
438                next_sidecar_id: crate::Counter::default(),
439                next_sim_hook_id: 0,
440                sidecars: Vec::new(),
441                live_collection_nodes: Vec::new(),
442            })),
443            locations: built.locations.clone(),
444            location_names: built.location_names.clone(),
445            #[cfg(feature = "sim")]
446            location_version: built.location_version.clone(),
447            #[cfg(feature = "sim")]
448            location_version_group_root: built.location_version_group_root.clone(),
449            flow_name: built.flow_name.clone(),
450            finalized: false,
451            _phantom: PhantomData,
452        }
453    }
454
455    #[doc(hidden)] // TODO(mingwei): This is an unstable API for now
456    pub fn replace_ir(&mut self, roots: Vec<HydroRoot>) {
457        self.flow_state.borrow_mut().roots = Some(roots);
458    }
459
460    #[doc(hidden)] // TODO(mingwei): This is an unstable API for now
461    pub fn next_clock_id(&mut self) -> ClockId {
462        self.flow_state.borrow_mut().next_clock_id()
463    }
464
465    #[doc(hidden)] // TODO(mingwei): This is an unstable API for now
466    pub fn next_cycle_id(&mut self) -> CycleId {
467        self.flow_state.borrow_mut().next_cycle_id()
468    }
469}
hydro_lang/compile/builder.rs

hydro_lang/compile/
builder.rs
