Skip to main content

hydro_lang/compile/
built.rs

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