1use std::cell::UnsafeCell;
2use std::collections::{BTreeMap, HashMap};
3use std::marker::PhantomData;
4
5use dfir_lang::graph::{DfirGraph, eliminate_extra_unions_tees, partition_graph};
6
7use super::compiled::CompiledFlow;
8use super::deploy::{DeployFlow, DeployResult};
9use super::deploy_provider::{ClusterSpec, Deploy, ExternalSpec, IntoProcessSpec};
10use super::ir::{HydroRoot, emit};
11#[cfg(feature = "viz")]
12use crate::graph::api::GraphApi;
13use crate::location::{Cluster, External, Process};
14#[cfg(feature = "sim")]
15#[cfg(stageleft_runtime)]
16use crate::sim::{flow::SimFlow, graph::SimNode};
17use crate::staging_util::Invariant;
18
19pub struct BuiltFlow<'a> {
20 pub(super) ir: Vec<HydroRoot>,
21 pub(super) process_id_name: Vec<(usize, String)>,
22 pub(super) cluster_id_name: Vec<(usize, String)>,
23 pub(super) external_id_name: Vec<(usize, String)>,
24
25 pub(super) _phantom: Invariant<'a>,
26}
27
28pub(crate) fn build_inner(ir: &mut Vec<HydroRoot>) -> BTreeMap<usize, DfirGraph> {
29 emit(ir)
30 .into_iter()
31 .map(|(k, v)| {
32 let (mut flat_graph, _, _) = v.build();
33 eliminate_extra_unions_tees(&mut flat_graph);
34 let partitioned_graph =
35 partition_graph(flat_graph).expect("Failed to partition (cycle detected).");
36 (k, partitioned_graph)
37 })
38 .collect()
39}
40
41impl<'a> BuiltFlow<'a> {
42 pub fn ir(&self) -> &Vec<HydroRoot> {
43 &self.ir
44 }
45
46 pub fn process_id_name(&self) -> &Vec<(usize, String)> {
47 &self.process_id_name
48 }
49
50 pub fn cluster_id_name(&self) -> &Vec<(usize, String)> {
51 &self.cluster_id_name
52 }
53
54 pub fn external_id_name(&self) -> &Vec<(usize, String)> {
55 &self.external_id_name
56 }
57
58 #[cfg(feature = "viz")]
60 pub fn graph_api(&self) -> GraphApi<'_> {
61 GraphApi::new(
62 &self.ir,
63 &self.process_id_name,
64 &self.cluster_id_name,
65 &self.external_id_name,
66 )
67 }
68
69 #[cfg(feature = "viz")]
71 pub fn mermaid_string(
72 &self,
73 show_metadata: bool,
74 show_location_groups: bool,
75 use_short_labels: bool,
76 ) -> String {
77 self.graph_api()
78 .mermaid_to_string(show_metadata, show_location_groups, use_short_labels)
79 }
80
81 #[cfg(feature = "viz")]
82 pub fn dot_string(
83 &self,
84 show_metadata: bool,
85 show_location_groups: bool,
86 use_short_labels: bool,
87 ) -> String {
88 self.graph_api()
89 .dot_to_string(show_metadata, show_location_groups, use_short_labels)
90 }
91
92 #[cfg(feature = "viz")]
93 pub fn reactflow_string(
94 &self,
95 show_metadata: bool,
96 show_location_groups: bool,
97 use_short_labels: bool,
98 ) -> String {
99 self.graph_api()
100 .reactflow_to_string(show_metadata, show_location_groups, use_short_labels)
101 }
102
103 #[cfg(feature = "viz")]
105 pub fn mermaid_to_file(
106 &self,
107 filename: &str,
108 show_metadata: bool,
109 show_location_groups: bool,
110 use_short_labels: bool,
111 ) -> Result<(), Box<dyn std::error::Error>> {
112 self.graph_api().mermaid_to_file(
113 filename,
114 show_metadata,
115 show_location_groups,
116 use_short_labels,
117 )
118 }
119
120 #[cfg(feature = "viz")]
121 pub fn dot_to_file(
122 &self,
123 filename: &str,
124 show_metadata: bool,
125 show_location_groups: bool,
126 use_short_labels: bool,
127 ) -> Result<(), Box<dyn std::error::Error>> {
128 self.graph_api().dot_to_file(
129 filename,
130 show_metadata,
131 show_location_groups,
132 use_short_labels,
133 )
134 }
135
136 #[cfg(feature = "viz")]
137 pub fn reactflow_to_file(
138 &self,
139 filename: &str,
140 show_metadata: bool,
141 show_location_groups: bool,
142 use_short_labels: bool,
143 ) -> Result<(), Box<dyn std::error::Error>> {
144 self.graph_api().reactflow_to_file(
145 filename,
146 show_metadata,
147 show_location_groups,
148 use_short_labels,
149 )
150 }
151
152 #[cfg(feature = "viz")]
154 pub fn mermaid_to_browser(
155 &self,
156 show_metadata: bool,
157 show_location_groups: bool,
158 use_short_labels: bool,
159 message_handler: Option<&dyn Fn(&str)>,
160 ) -> Result<(), Box<dyn std::error::Error>> {
161 self.graph_api().mermaid_to_browser(
162 show_metadata,
163 show_location_groups,
164 use_short_labels,
165 message_handler,
166 )
167 }
168
169 #[cfg(feature = "viz")]
170 pub fn dot_to_browser(
171 &self,
172 show_metadata: bool,
173 show_location_groups: bool,
174 use_short_labels: bool,
175 message_handler: Option<&dyn Fn(&str)>,
176 ) -> Result<(), Box<dyn std::error::Error>> {
177 self.graph_api().dot_to_browser(
178 show_metadata,
179 show_location_groups,
180 use_short_labels,
181 message_handler,
182 )
183 }
184
185 #[cfg(feature = "viz")]
186 pub fn reactflow_to_browser(
187 &self,
188 show_metadata: bool,
189 show_location_groups: bool,
190 use_short_labels: bool,
191 message_handler: Option<&dyn Fn(&str)>,
192 ) -> Result<(), Box<dyn std::error::Error>> {
193 self.graph_api().reactflow_to_browser(
194 show_metadata,
195 show_location_groups,
196 use_short_labels,
197 message_handler,
198 )
199 }
200
201 pub fn optimize_with(mut self, f: impl FnOnce(&mut [HydroRoot])) -> Self {
202 f(&mut self.ir);
203 BuiltFlow {
204 ir: std::mem::take(&mut self.ir),
205 process_id_name: std::mem::take(&mut self.process_id_name),
206 cluster_id_name: std::mem::take(&mut self.cluster_id_name),
207 external_id_name: std::mem::take(&mut self.external_id_name),
208 _phantom: PhantomData,
209 }
210 }
211
212 pub fn with_default_optimize<D: Deploy<'a>>(self) -> DeployFlow<'a, D> {
213 self.into_deploy()
214 }
215
216 #[cfg(feature = "sim")]
217 pub fn sim(mut self) -> SimFlow<'a> {
220 use std::cell::RefCell;
221 use std::rc::Rc;
222
223 use crate::sim::graph::SimExternal;
224
225 let external_ports = Rc::new(RefCell::new((vec![], 0)));
226 let processes = self
227 .process_id_name
228 .iter()
229 .map(|id| (id.0, SimNode {}))
230 .collect();
231
232 let clusters = self
233 .cluster_id_name
234 .iter()
235 .map(|id| (id.0, SimNode {}))
236 .collect();
237
238 let externals = self
239 .external_id_name
240 .iter()
241 .map(|id| {
242 (
243 id.0,
244 SimExternal {
245 external_ports: external_ports.clone(),
246 registered: RefCell::new(HashMap::new()),
247 },
248 )
249 })
250 .collect();
251
252 SimFlow {
253 ir: std::mem::take(&mut self.ir),
254 external_ports,
255 processes,
256 clusters,
257 externals,
258 _process_id_name: std::mem::take(&mut self.process_id_name),
259 _external_id_name: std::mem::take(&mut self.external_id_name),
260 _cluster_id_name: std::mem::take(&mut self.cluster_id_name),
261 _phantom: PhantomData,
262 }
263 }
264
265 pub fn into_deploy<D: Deploy<'a>>(mut self) -> DeployFlow<'a, D> {
266 let processes = if D::has_trivial_node() {
267 self.process_id_name
268 .iter()
269 .map(|id| (id.0, D::trivial_process(id.0)))
270 .collect()
271 } else {
272 HashMap::new()
273 };
274
275 let clusters = if D::has_trivial_node() {
276 self.cluster_id_name
277 .iter()
278 .map(|id| (id.0, D::trivial_cluster(id.0)))
279 .collect()
280 } else {
281 HashMap::new()
282 };
283
284 let externals = if D::has_trivial_node() {
285 self.external_id_name
286 .iter()
287 .map(|id| (id.0, D::trivial_external(id.0)))
288 .collect()
289 } else {
290 HashMap::new()
291 };
292
293 DeployFlow {
294 ir: UnsafeCell::new(std::mem::take(&mut self.ir)),
295 processes,
296 process_id_name: std::mem::take(&mut self.process_id_name),
297 clusters,
298 cluster_id_name: std::mem::take(&mut self.cluster_id_name),
299 externals,
300 external_id_name: std::mem::take(&mut self.external_id_name),
301 _phantom: PhantomData,
302 }
303 }
304
305 pub fn with_process<P, D: Deploy<'a>>(
306 self,
307 process: &Process<P>,
308 spec: impl IntoProcessSpec<'a, D>,
309 ) -> DeployFlow<'a, D> {
310 self.into_deploy().with_process(process, spec)
311 }
312
313 pub fn with_remaining_processes<D: Deploy<'a>, S: IntoProcessSpec<'a, D> + 'a>(
314 self,
315 spec: impl Fn() -> S,
316 ) -> DeployFlow<'a, D> {
317 self.into_deploy().with_remaining_processes(spec)
318 }
319
320 pub fn with_external<P, D: Deploy<'a>>(
321 self,
322 process: &External<P>,
323 spec: impl ExternalSpec<'a, D>,
324 ) -> DeployFlow<'a, D> {
325 self.into_deploy().with_external(process, spec)
326 }
327
328 pub fn with_remaining_externals<D: Deploy<'a>, S: ExternalSpec<'a, D> + 'a>(
329 self,
330 spec: impl Fn() -> S,
331 ) -> DeployFlow<'a, D> {
332 self.into_deploy().with_remaining_externals(spec)
333 }
334
335 pub fn with_cluster<C, D: Deploy<'a>>(
336 self,
337 cluster: &Cluster<C>,
338 spec: impl ClusterSpec<'a, D>,
339 ) -> DeployFlow<'a, D> {
340 self.into_deploy().with_cluster(cluster, spec)
341 }
342
343 pub fn with_remaining_clusters<D: Deploy<'a>, S: ClusterSpec<'a, D> + 'a>(
344 self,
345 spec: impl Fn() -> S,
346 ) -> DeployFlow<'a, D> {
347 self.into_deploy().with_remaining_clusters(spec)
348 }
349
350 pub fn compile<D: Deploy<'a>>(self, env: &D::CompileEnv) -> CompiledFlow<'a, D::GraphId> {
351 self.into_deploy::<D>().compile(env)
352 }
353
354 pub fn compile_no_network<D: Deploy<'a>>(self) -> CompiledFlow<'a, D::GraphId> {
355 self.into_deploy::<D>().compile_no_network()
356 }
357
358 pub fn deploy<D: Deploy<'a, CompileEnv = ()>>(
359 self,
360 env: &mut D::InstantiateEnv,
361 ) -> DeployResult<'a, D> {
362 self.into_deploy::<D>().deploy(env)
363 }
364
365 #[cfg(feature = "viz")]
366 pub fn generate_all_files(
367 &self,
368 prefix: &str,
369 show_metadata: bool,
370 show_location_groups: bool,
371 use_short_labels: bool,
372 ) -> Result<(), Box<dyn std::error::Error>> {
373 self.graph_api().generate_all_files(
374 prefix,
375 show_metadata,
376 show_location_groups,
377 use_short_labels,
378 )
379 }
380
381 #[cfg(feature = "viz")]
382 pub fn generate_graph_with_config(
383 &self,
384 config: &crate::graph::config::GraphConfig,
385 message_handler: Option<&dyn Fn(&str)>,
386 ) -> Result<(), Box<dyn std::error::Error>> {
387 self.graph_api()
388 .generate_graph_with_config(config, message_handler)
389 }
390
391 #[cfg(feature = "viz")]
392 pub fn generate_all_files_with_config(
393 &self,
394 config: &crate::graph::config::GraphConfig,
395 prefix: &str,
396 ) -> Result<(), Box<dyn std::error::Error>> {
397 self.graph_api()
398 .generate_all_files_with_config(config, prefix)
399 }
400}