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 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 pub fn ir(&self) -> &[HydroRoot] {
48 &self.ir
49 }
50
51 pub fn location_names(&self) -> &SecondaryMap<LocationKey, String> {
53 &self.location_names
54 }
55
56 #[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 #[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 #[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 #[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 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}