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 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 pub fn ir(&self) -> &[HydroRoot] {
51 &self.ir
52 }
53
54 pub fn location_names(&self) -> &SecondaryMap<LocationKey, String> {
56 &self.location_names
57 }
58
59 #[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 #[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 #[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 #[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 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}