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