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        if !self.finalized && !std::thread::panicking() {
192            panic!(
193                "Dropped FlowBuilder without finalizing, you may have forgotten to call `with_default_optimize`, `optimize_with`, or `finalize`."
194            );
195        }
196    }
197}
198
199#[expect(missing_docs, reason = "TODO")]
200impl<'a> FlowBuilder<'a> {
201    /// Creates a new `FlowBuilder` to construct a Hydro program, using the Cargo package name as the program name.
202    #[expect(
203        clippy::new_without_default,
204        reason = "call `new` explicitly, not `default`"
205    )]
206    pub fn new() -> Self {
207        let mut name = std::env::var("CARGO_PKG_NAME").unwrap_or_else(|_| "unknown".to_owned());
208        if let Ok(bin_path) = std::env::current_exe()
209            && let Some(bin_name) = bin_path.file_stem()
210        {
211            name = format!("{}/{}", name, bin_name.display());
212        }
213        Self::with_name(name)
214    }
215
216    /// Creates a new `FlowBuilder` to construct a Hydro program, with the given program name.
217    pub fn with_name(name: impl Into<String>) -> Self {
218        Self {
219            flow_state: Rc::new(RefCell::new(FlowStateInner {
220                roots: Some(vec![]),
221                next_external_port: crate::Counter::default(),
222                next_cycle_id: crate::Counter::default(),
223                next_clock_id: crate::Counter::default(),
224                next_sidecar_id: crate::Counter::default(),
225                next_sim_hook_id: 0,
226                sidecars: Vec::new(),
227                live_collection_nodes: Vec::new(),
228            })),
229            locations: SlotMap::with_key(),
230            location_names: SecondaryMap::new(),
231            #[cfg(feature = "sim")]
232            location_version: SecondaryMap::new(),
233            #[cfg(feature = "sim")]
234            location_version_group_root: SecondaryMap::new(),
235            flow_name: name.into(),
236            finalized: false,
237            _phantom: PhantomData,
238        }
239    }
240
241    pub(crate) fn flow_state(&self) -> &FlowState {
242        &self.flow_state
243    }
244
245    fn insert_location(&mut self, ty: LocationType, name: String) -> LocationKey {
246        let key = self.locations.insert(ty);
247        self.location_names.insert(key, name);
248        #[cfg(feature = "sim")]
249        {
250            self.location_version.insert(key, 0);
251            self.location_version_group_root.insert(key, key);
252        }
253        key
254    }
255
256    pub fn process<P>(&mut self) -> Process<'a, P> {
257        let key = self.insert_location(LocationType::Process, type_name::<P>().to_owned());
258        Process {
259            key,
260            flow_state: self.flow_state().clone(),
261            _phantom: PhantomData,
262        }
263    }
264
265    pub fn cluster<C>(&mut self) -> Cluster<'a, C> {
266        let key = self.insert_location(LocationType::Cluster, type_name::<C>().to_owned());
267        Cluster {
268            key,
269            flow_state: self.flow_state().clone(),
270            _phantom: PhantomData,
271        }
272    }
273
274    pub fn external<E>(&mut self) -> External<'a, E> {
275        let key = self.insert_location(LocationType::External, type_name::<E>().to_owned());
276        External {
277            key,
278            flow_state: self.flow_state().clone(),
279            _phantom: PhantomData,
280        }
281    }
282
283    /// Creates a fresh set of **simulator hook handles** (see [`crate::sim_hooks`]).
284    ///
285    /// A handle is attached to a specific unsafe operator with the
286    /// `nondet!(/** reason */ hook = handle)` syntax, and lets a simulation test script the
287    /// non-deterministic decisions of that operator. `B` may be a single handle type (e.g.
288    /// `BatchHook<u32>`) or a struct of handles implementing
289    /// [`SimHook`](crate::sim_hooks::SimHook).
290    ///
291    /// Creating a handle and never binding it is allowed (a bundle may be only partially
292    /// used by a particular program configuration), but scripting a decision on an unbound
293    /// handle panics at that call. Binding the same handle to two different operators is an
294    /// error at flow build time.
295    pub fn sim_hook<B: crate::sim_hooks::SimHook>(&mut self) -> B {
296        let flow_state = self.flow_state.clone();
297        B::create(&mut move || flow_state.borrow_mut().next_sim_hook_id())
298    }
299
300    #[cfg(feature = "sim")]
301    pub fn next_version<C>(&mut self, cluster: &Cluster<'a, C>) -> Cluster<'a, C> {
302        let group_root = self.location_version_group_root[cluster.key];
303        let version = self
304            .location_version_group_root
305            .values()
306            .filter(|&&r| r == group_root)
307            .count() as u32;
308        let key = self.insert_location(LocationType::Cluster, type_name::<C>().to_owned());
309        self.location_version.insert(key, version);
310        self.location_version_group_root.insert(key, group_root);
311        Cluster {
312            key,
313            flow_state: self.flow_state().clone(),
314            _phantom: PhantomData,
315        }
316    }
317}
318
319#[cfg(feature = "build")]
320#[cfg_attr(docsrs, doc(cfg(feature = "build")))]
321#[expect(missing_docs, reason = "TODO")]
322impl<'a> FlowBuilder<'a> {
323    pub fn finalize(mut self) -> super::built::BuiltFlow<'a> {
324        self.finalized = true;
325
326        let mut flow_state = self.flow_state.borrow_mut();
327
328        // Yank the IR from any live collections (streams, singletons, ...) that are still
329        // alive, since their `Drop` will run after finalization and would otherwise
330        // silently fail to register their dataflow as a root.
331        let live_collection_nodes = std::mem::take(&mut flow_state.live_collection_nodes);
332        for node_cell in live_collection_nodes {
333            if let Some(node_cell) = node_cell.upgrade() {
334                let ir_node = node_cell.replace(HydroNode::Placeholder);
335                if !matches!(ir_node, HydroNode::Placeholder) && !ir_node.is_shared_with_others() {
336                    flow_state.push_root(HydroRoot::Null {
337                        input: Box::new(ir_node),
338                        op_metadata: HydroIrOpMetadata::new(),
339                    });
340                }
341            }
342        }
343
344        let mut ir = flow_state.roots.take().unwrap();
345        let sidecars = std::mem::take(&mut flow_state.sidecars);
346        drop(flow_state);
347
348        super::ir::unify_atomic_ticks(&mut ir);
349
350        super::built::BuiltFlow {
351            ir,
352            locations: std::mem::take(&mut self.locations),
353            location_names: std::mem::take(&mut self.location_names),
354            sidecars,
355            flow_name: std::mem::take(&mut self.flow_name),
356            #[cfg(feature = "sim")]
357            location_version: std::mem::take(&mut self.location_version),
358            #[cfg(feature = "sim")]
359            location_version_group_root: std::mem::take(&mut self.location_version_group_root),
360            _phantom: PhantomData,
361        }
362    }
363
364    pub fn with_default_optimize<D: Deploy<'a>>(self) -> DeployFlow<'a, D> {
365        self.finalize().with_default_optimize()
366    }
367
368    pub fn optimize_with(self, f: impl FnOnce(&mut [HydroRoot])) -> super::built::BuiltFlow<'a> {
369        self.finalize().optimize_with(f)
370    }
371
372    pub fn with_process<P, D: Deploy<'a>>(
373        self,
374        process: &Process<'_, P>,
375        spec: impl IntoProcessSpec<'a, D>,
376    ) -> DeployFlow<'a, D> {
377        self.with_default_optimize().with_process(process, spec)
378    }
379
380    pub fn with_remaining_processes<D: Deploy<'a>, S: IntoProcessSpec<'a, D> + 'a>(
381        self,
382        spec: impl Fn() -> S,
383    ) -> DeployFlow<'a, D> {
384        self.with_default_optimize().with_remaining_processes(spec)
385    }
386
387    pub fn with_external<P, D: Deploy<'a>>(
388        self,
389        process: &External<'_, P>,
390        spec: impl ExternalSpec<'a, D>,
391    ) -> DeployFlow<'a, D> {
392        self.with_default_optimize().with_external(process, spec)
393    }
394
395    pub fn with_remaining_externals<D: Deploy<'a>, S: ExternalSpec<'a, D> + 'a>(
396        self,
397        spec: impl Fn() -> S,
398    ) -> DeployFlow<'a, D> {
399        self.with_default_optimize().with_remaining_externals(spec)
400    }
401
402    pub fn with_cluster<C, D: Deploy<'a>>(
403        self,
404        cluster: &Cluster<'_, C>,
405        spec: impl ClusterSpec<'a, D>,
406    ) -> DeployFlow<'a, D> {
407        self.with_default_optimize().with_cluster(cluster, spec)
408    }
409
410    pub fn with_remaining_clusters<D: Deploy<'a>, S: ClusterSpec<'a, D> + 'a>(
411        self,
412        spec: impl Fn() -> S,
413    ) -> DeployFlow<'a, D> {
414        self.with_default_optimize().with_remaining_clusters(spec)
415    }
416
417    pub fn compile<D: Deploy<'a, InstantiateEnv = ()>>(self) -> CompiledFlow<'a> {
418        self.with_default_optimize::<D>().compile()
419    }
420
421    pub fn deploy<D: Deploy<'a>>(self, env: &mut D::InstantiateEnv) -> DeployResult<'a, D> {
422        self.with_default_optimize().deploy(env)
423    }
424
425    #[cfg(feature = "sim")]
426    /// Creates a simulation for this builder, which can be used to run deterministic simulations
427    /// of the Hydro program.
428    pub fn sim(self) -> SimFlow<'a> {
429        self.finalize().sim()
430    }
431
432    pub fn from_built<'b>(built: &super::built::BuiltFlow<'_>) -> FlowBuilder<'b> {
433        FlowBuilder {
434            flow_state: Rc::new(RefCell::new(FlowStateInner {
435                roots: None,
436                next_external_port: crate::Counter::default(),
437                next_cycle_id: crate::Counter::default(),
438                next_clock_id: crate::Counter::default(),
439                next_sidecar_id: crate::Counter::default(),
440                next_sim_hook_id: 0,
441                sidecars: Vec::new(),
442                live_collection_nodes: Vec::new(),
443            })),
444            locations: built.locations.clone(),
445            location_names: built.location_names.clone(),
446            #[cfg(feature = "sim")]
447            location_version: built.location_version.clone(),
448            #[cfg(feature = "sim")]
449            location_version_group_root: built.location_version_group_root.clone(),
450            flow_name: built.flow_name.clone(),
451            finalized: false,
452            _phantom: PhantomData,
453        }
454    }
455
456    #[doc(hidden)] // TODO(mingwei): This is an unstable API for now
457    pub fn replace_ir(&mut self, roots: Vec<HydroRoot>) {
458        self.flow_state.borrow_mut().roots = Some(roots);
459    }
460
461    #[doc(hidden)] // TODO(mingwei): This is an unstable API for now
462    pub fn next_clock_id(&mut self) -> ClockId {
463        self.flow_state.borrow_mut().next_clock_id()
464    }
465
466    #[doc(hidden)] // TODO(mingwei): This is an unstable API for now
467    pub fn next_cycle_id(&mut self) -> CycleId {
468        self.flow_state.borrow_mut().next_cycle_id()
469    }
470}
hydro_lang/compile/builder.rs

hydro_lang/compile/
builder.rs
