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