hydro_lang/compile/
builder.rs

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

hydro_lang/compile/
builder.rs
