Skip to main content

hydro_lang/compile/
built.rs

1use std::marker::PhantomData;
2
3use dfir_lang::graph::{DfirGraph, eliminate_extra_unions_tees, partition_graph};
4use slotmap::{SecondaryMap, SlotMap, SparseSecondaryMap};
5
6use super::compiled::CompiledFlow;
7use super::deploy::{DeployFlow, DeployResult};
8use super::deploy_provider::{ClusterSpec, Deploy, ExternalSpec, IntoProcessSpec};
9use super::ir::{HydroRoot, emit};
10use crate::location::{Cluster, External, LocationKey, LocationType, Process};
11#[cfg(stageleft_runtime)]
12#[cfg(feature = "sim")]
13use crate::sim::{flow::SimFlow, graph::SimNode};
14use crate::staging_util::Invariant;
15#[cfg(stageleft_runtime)]
16#[cfg(feature = "viz")]
17use crate::viz::api::GraphApi;
18
19pub struct BuiltFlow<'a> {
20    pub(super) ir: Vec<HydroRoot>,
21    pub(super) locations: SlotMap<LocationKey, LocationType>,
22    pub(super) location_names: SecondaryMap<LocationKey, String>,
23
24    /// Application name used in telemetry.
25    pub(super) flow_name: String,
26
27    pub(super) _phantom: Invariant<'a>,
28}
29
30pub(crate) fn build_inner<'a, D: Deploy<'a>>(
31    ir: &mut Vec<HydroRoot>,
32) -> SecondaryMap<LocationKey, DfirGraph> {
33    emit::<D>(ir)
34        .into_iter()
35        .map(|(k, v)| {
36            let (mut flat_graph, _, _) = v.build();
37            eliminate_extra_unions_tees(&mut flat_graph);
38            let partitioned_graph =
39                partition_graph(flat_graph).expect("Failed to partition (cycle detected).");
40            (k, partitioned_graph)
41        })
42        .collect()
43}
44
45impl<'a> BuiltFlow<'a> {
46    /// Returns all [`HydroRoot`]s in the IR.
47    pub fn ir(&self) -> &[HydroRoot] {
48        &self.ir
49    }
50
51    /// Returns all raw location ID -> location name mappings.
52    pub fn location_names(&self) -> &SecondaryMap<LocationKey, String> {
53        &self.location_names
54    }
55
56    /// Get a GraphApi instance for this built flow
57    #[cfg(stageleft_runtime)]
58    #[cfg(feature = "viz")]
59    pub fn graph_api(&self) -> GraphApi<'_> {
60        GraphApi::new(&self.ir, self.location_names())
61    }
62
63    // String generation methods
64    #[cfg(feature = "viz")]
65    pub fn mermaid_string(
66        &self,
67        show_metadata: bool,
68        show_location_groups: bool,
69        use_short_labels: bool,
70    ) -> String {
71        self.graph_api()
72            .mermaid_to_string(show_metadata, show_location_groups, use_short_labels)
73    }
74
75    #[cfg(feature = "viz")]
76    pub fn dot_string(
77        &self,
78        show_metadata: bool,
79        show_location_groups: bool,
80        use_short_labels: bool,
81    ) -> String {
82        self.graph_api()
83            .dot_to_string(show_metadata, show_location_groups, use_short_labels)
84    }
85
86    #[cfg(feature = "viz")]
87    pub fn hydroscope_string(
88        &self,
89        show_metadata: bool,
90        show_location_groups: bool,
91        use_short_labels: bool,
92    ) -> String {
93        self.graph_api()
94            .hydroscope_to_string(show_metadata, show_location_groups, use_short_labels)
95    }
96
97    // File generation methods
98    #[cfg(feature = "viz")]
99    pub fn mermaid_to_file(
100        &self,
101        filename: &str,
102        show_metadata: bool,
103        show_location_groups: bool,
104        use_short_labels: bool,
105    ) -> Result<(), Box<dyn std::error::Error>> {
106        self.graph_api().mermaid_to_file(
107            filename,
108            show_metadata,
109            show_location_groups,
110            use_short_labels,
111        )
112    }
113
114    #[cfg(feature = "viz")]
115    pub fn dot_to_file(
116        &self,
117        filename: &str,
118        show_metadata: bool,
119        show_location_groups: bool,
120        use_short_labels: bool,
121    ) -> Result<(), Box<dyn std::error::Error>> {
122        self.graph_api().dot_to_file(
123            filename,
124            show_metadata,
125            show_location_groups,
126            use_short_labels,
127        )
128    }
129
130    #[cfg(feature = "viz")]
131    pub fn hydroscope_to_file(
132        &self,
133        filename: &str,
134        show_metadata: bool,
135        show_location_groups: bool,
136        use_short_labels: bool,
137    ) -> Result<(), Box<dyn std::error::Error>> {
138        self.graph_api().hydroscope_to_file(
139            filename,
140            show_metadata,
141            show_location_groups,
142            use_short_labels,
143        )
144    }
145
146    // Browser generation methods
147    #[cfg(feature = "viz")]
148    pub fn mermaid_to_browser(
149        &self,
150        show_metadata: bool,
151        show_location_groups: bool,
152        use_short_labels: bool,
153        message_handler: Option<&dyn Fn(&str)>,
154    ) -> Result<(), Box<dyn std::error::Error>> {
155        self.graph_api().mermaid_to_browser(
156            show_metadata,
157            show_location_groups,
158            use_short_labels,
159            message_handler,
160        )
161    }
162
163    #[cfg(feature = "viz")]
164    pub fn dot_to_browser(
165        &self,
166        show_metadata: bool,
167        show_location_groups: bool,
168        use_short_labels: bool,
169        message_handler: Option<&dyn Fn(&str)>,
170    ) -> Result<(), Box<dyn std::error::Error>> {
171        self.graph_api().dot_to_browser(
172            show_metadata,
173            show_location_groups,
174            use_short_labels,
175            message_handler,
176        )
177    }
178
179    #[cfg(feature = "viz")]
180    pub fn hydroscope_to_browser(
181        &self,
182        show_metadata: bool,
183        show_location_groups: bool,
184        use_short_labels: bool,
185        message_handler: Option<&dyn Fn(&str)>,
186    ) -> Result<(), Box<dyn std::error::Error>> {
187        self.graph_api().hydroscope_to_browser(
188            show_metadata,
189            show_location_groups,
190            use_short_labels,
191            message_handler,
192        )
193    }
194
195    pub fn optimize_with(mut self, f: impl FnOnce(&mut [HydroRoot])) -> Self {
196        f(&mut self.ir);
197        self
198    }
199
200    pub fn with_default_optimize<D: Deploy<'a>>(self) -> DeployFlow<'a, D> {
201        self.into_deploy()
202    }
203
204    #[cfg(feature = "sim")]
205    /// Creates a simulation for this builder, which can be used to run deterministic simulations
206    /// of the Hydro program.
207    pub fn sim(self) -> SimFlow<'a> {
208        use std::cell::RefCell;
209        use std::rc::Rc;
210
211        use slotmap::SparseSecondaryMap;
212
213        use crate::sim::graph::SimNodePort;
214
215        let shared_port_counter = Rc::new(RefCell::new(SimNodePort::default()));
216
217        let mut processes = SparseSecondaryMap::new();
218        let mut clusters = SparseSecondaryMap::new();
219        let externals = SparseSecondaryMap::new();
220
221        for (key, loc) in self.locations.iter() {
222            match loc {
223                LocationType::Process => {
224                    processes.insert(
225                        key,
226                        SimNode {
227                            shared_port_counter: shared_port_counter.clone(),
228                        },
229                    );
230                }
231                LocationType::Cluster => {
232                    clusters.insert(
233                        key,
234                        SimNode {
235                            shared_port_counter: shared_port_counter.clone(),
236                        },
237                    );
238                }
239                LocationType::External => {
240                    panic!("Sim cannot have externals");
241                }
242            }
243        }
244
245        SimFlow {
246            ir: self.ir,
247            processes,
248            clusters,
249            externals,
250            cluster_max_sizes: SparseSecondaryMap::new(),
251            externals_port_registry: Default::default(),
252            _phantom: PhantomData,
253        }
254    }
255
256    pub fn into_deploy<D: Deploy<'a>>(self) -> DeployFlow<'a, D> {
257        let (processes, clusters, externals) = Default::default();
258        DeployFlow {
259            ir: self.ir,
260            locations: self.locations,
261            location_names: self.location_names,
262            processes,
263            clusters,
264            externals,
265            sidecars: SparseSecondaryMap::new(),
266            flow_name: self.flow_name,
267            _phantom: PhantomData,
268        }
269    }
270
271    pub fn with_process<P, D: Deploy<'a>>(
272        self,
273        process: &Process<P>,
274        spec: impl IntoProcessSpec<'a, D>,
275    ) -> DeployFlow<'a, D> {
276        self.into_deploy().with_process(process, spec)
277    }
278
279    pub fn with_remaining_processes<D: Deploy<'a>, S: IntoProcessSpec<'a, D> + 'a>(
280        self,
281        spec: impl Fn() -> S,
282    ) -> DeployFlow<'a, D> {
283        self.into_deploy().with_remaining_processes(spec)
284    }
285
286    pub fn with_external<P, D: Deploy<'a>>(
287        self,
288        process: &External<P>,
289        spec: impl ExternalSpec<'a, D>,
290    ) -> DeployFlow<'a, D> {
291        self.into_deploy().with_external(process, spec)
292    }
293
294    pub fn with_remaining_externals<D: Deploy<'a>, S: ExternalSpec<'a, D> + 'a>(
295        self,
296        spec: impl Fn() -> S,
297    ) -> DeployFlow<'a, D> {
298        self.into_deploy().with_remaining_externals(spec)
299    }
300
301    pub fn with_cluster<C, D: Deploy<'a>>(
302        self,
303        cluster: &Cluster<C>,
304        spec: impl ClusterSpec<'a, D>,
305    ) -> DeployFlow<'a, D> {
306        self.into_deploy().with_cluster(cluster, spec)
307    }
308
309    pub fn with_remaining_clusters<D: Deploy<'a>, S: ClusterSpec<'a, D> + 'a>(
310        self,
311        spec: impl Fn() -> S,
312    ) -> DeployFlow<'a, D> {
313        self.into_deploy().with_remaining_clusters(spec)
314    }
315
316    pub fn compile<D: Deploy<'a>>(self) -> CompiledFlow<'a> {
317        self.into_deploy::<D>().compile()
318    }
319
320    pub fn deploy<D: Deploy<'a>>(self, env: &mut D::InstantiateEnv) -> DeployResult<'a, D> {
321        self.into_deploy::<D>().deploy(env)
322    }
323
324    #[cfg(feature = "viz")]
325    pub fn generate_all_files(
326        &self,
327        prefix: &str,
328        show_metadata: bool,
329        show_location_groups: bool,
330        use_short_labels: bool,
331    ) -> Result<(), Box<dyn std::error::Error>> {
332        self.graph_api().generate_all_files(
333            prefix,
334            show_metadata,
335            show_location_groups,
336            use_short_labels,
337        )
338    }
339
340    #[cfg(feature = "viz")]
341    pub fn generate_graph_with_config(
342        &self,
343        config: &crate::viz::config::GraphConfig,
344        message_handler: Option<&dyn Fn(&str)>,
345    ) -> Result<(), Box<dyn std::error::Error>> {
346        self.graph_api()
347            .generate_graph_with_config(config, message_handler)
348    }
349
350    #[cfg(feature = "viz")]
351    pub fn generate_all_files_with_config(
352        &self,
353        config: &crate::viz::config::GraphConfig,
354        prefix: &str,
355    ) -> Result<(), Box<dyn std::error::Error>> {
356        self.graph_api()
357            .generate_all_files_with_config(config, prefix)
358    }
359}