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 crate::deploy::{ClusterSpec, Deploy, ExternalSpec, IntoProcessSpec};
10#[cfg(feature = "viz")]
11use crate::graph::api::GraphApi;
12use crate::ir::{HydroRoot, emit};
13use crate::location::{Cluster, External, Process};
14use crate::staging_util::Invariant;
15
16pub struct BuiltFlow<'a> {
17 pub(super) ir: Vec<HydroRoot>,
18 pub(super) process_id_name: Vec<(usize, String)>,
19 pub(super) cluster_id_name: Vec<(usize, String)>,
20 pub(super) external_id_name: Vec<(usize, String)>,
21
22 pub(super) _phantom: Invariant<'a>,
23}
24
25pub(crate) fn build_inner(ir: &mut Vec<HydroRoot>) -> BTreeMap<usize, DfirGraph> {
26 emit(ir)
27 .into_iter()
28 .map(|(k, v)| {
29 let (mut flat_graph, _, _) = v.build();
30 eliminate_extra_unions_tees(&mut flat_graph);
31 let partitioned_graph =
32 partition_graph(flat_graph).expect("Failed to partition (cycle detected).");
33 (k, partitioned_graph)
34 })
35 .collect()
36}
37
38impl<'a> BuiltFlow<'a> {
39 pub fn ir(&self) -> &Vec<HydroRoot> {
40 &self.ir
41 }
42
43 pub fn process_id_name(&self) -> &Vec<(usize, String)> {
44 &self.process_id_name
45 }
46
47 pub fn cluster_id_name(&self) -> &Vec<(usize, String)> {
48 &self.cluster_id_name
49 }
50
51 pub fn external_id_name(&self) -> &Vec<(usize, String)> {
52 &self.external_id_name
53 }
54
55 #[cfg(feature = "viz")]
57 pub fn graph_api(&self) -> GraphApi<'_> {
58 GraphApi::new(
59 &self.ir,
60 &self.process_id_name,
61 &self.cluster_id_name,
62 &self.external_id_name,
63 )
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 reactflow_string(
91 &self,
92 show_metadata: bool,
93 show_location_groups: bool,
94 use_short_labels: bool,
95 ) -> String {
96 self.graph_api()
97 .reactflow_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 reactflow_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().reactflow_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 reactflow_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().reactflow_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 BuiltFlow {
201 ir: std::mem::take(&mut self.ir),
202 process_id_name: std::mem::take(&mut self.process_id_name),
203 cluster_id_name: std::mem::take(&mut self.cluster_id_name),
204 external_id_name: std::mem::take(&mut self.external_id_name),
205 _phantom: PhantomData,
206 }
207 }
208
209 pub fn with_default_optimize<D: Deploy<'a>>(self) -> DeployFlow<'a, D> {
210 self.optimize_with(crate::rewrites::persist_pullup::persist_pullup)
211 .into_deploy()
212 }
213
214 pub fn into_deploy<D: Deploy<'a>>(mut self) -> DeployFlow<'a, D> {
215 let processes = if D::has_trivial_node() {
216 self.process_id_name
217 .iter()
218 .map(|id| (id.0, D::trivial_process(id.0)))
219 .collect()
220 } else {
221 HashMap::new()
222 };
223
224 let clusters = if D::has_trivial_node() {
225 self.cluster_id_name
226 .iter()
227 .map(|id| (id.0, D::trivial_cluster(id.0)))
228 .collect()
229 } else {
230 HashMap::new()
231 };
232
233 let externals = if D::has_trivial_node() {
234 self.external_id_name
235 .iter()
236 .map(|id| (id.0, D::trivial_external(id.0)))
237 .collect()
238 } else {
239 HashMap::new()
240 };
241
242 DeployFlow {
243 ir: UnsafeCell::new(std::mem::take(&mut self.ir)),
244 processes,
245 process_id_name: std::mem::take(&mut self.process_id_name),
246 clusters,
247 cluster_id_name: std::mem::take(&mut self.cluster_id_name),
248 externals,
249 external_id_name: std::mem::take(&mut self.external_id_name),
250 _phantom: PhantomData,
251 }
252 }
253
254 pub fn with_process<P, D: Deploy<'a>>(
255 self,
256 process: &Process<P>,
257 spec: impl IntoProcessSpec<'a, D>,
258 ) -> DeployFlow<'a, D> {
259 self.into_deploy().with_process(process, spec)
260 }
261
262 pub fn with_remaining_processes<D: Deploy<'a>, S: IntoProcessSpec<'a, D> + 'a>(
263 self,
264 spec: impl Fn() -> S,
265 ) -> DeployFlow<'a, D> {
266 self.into_deploy().with_remaining_processes(spec)
267 }
268
269 pub fn with_external<P, D: Deploy<'a>>(
270 self,
271 process: &External<P>,
272 spec: impl ExternalSpec<'a, D>,
273 ) -> DeployFlow<'a, D> {
274 self.into_deploy().with_external(process, spec)
275 }
276
277 pub fn with_remaining_externals<D: Deploy<'a>, S: ExternalSpec<'a, D> + 'a>(
278 self,
279 spec: impl Fn() -> S,
280 ) -> DeployFlow<'a, D> {
281 self.into_deploy().with_remaining_externals(spec)
282 }
283
284 pub fn with_cluster<C, D: Deploy<'a>>(
285 self,
286 cluster: &Cluster<C>,
287 spec: impl ClusterSpec<'a, D>,
288 ) -> DeployFlow<'a, D> {
289 self.into_deploy().with_cluster(cluster, spec)
290 }
291
292 pub fn with_remaining_clusters<D: Deploy<'a>, S: ClusterSpec<'a, D> + 'a>(
293 self,
294 spec: impl Fn() -> S,
295 ) -> DeployFlow<'a, D> {
296 self.into_deploy().with_remaining_clusters(spec)
297 }
298
299 pub fn compile<D: Deploy<'a>>(self, env: &D::CompileEnv) -> CompiledFlow<'a, D::GraphId> {
300 self.into_deploy::<D>().compile(env)
301 }
302
303 pub fn compile_no_network<D: Deploy<'a>>(self) -> CompiledFlow<'a, D::GraphId> {
304 self.into_deploy::<D>().compile_no_network()
305 }
306
307 pub fn deploy<D: Deploy<'a, CompileEnv = ()>>(
308 self,
309 env: &mut D::InstantiateEnv,
310 ) -> DeployResult<'a, D> {
311 self.into_deploy::<D>().deploy(env)
312 }
313
314 #[cfg(feature = "viz")]
315 pub fn generate_all_files(
316 &self,
317 prefix: &str,
318 show_metadata: bool,
319 show_location_groups: bool,
320 use_short_labels: bool,
321 ) -> Result<(), Box<dyn std::error::Error>> {
322 self.graph_api().generate_all_files(
323 prefix,
324 show_metadata,
325 show_location_groups,
326 use_short_labels,
327 )
328 }
329
330 #[cfg(feature = "viz")]
331 pub fn generate_graph_with_config(
332 &self,
333 config: &crate::graph_util::GraphConfig,
334 message_handler: Option<&dyn Fn(&str)>,
335 ) -> Result<(), Box<dyn std::error::Error>> {
336 self.graph_api()
337 .generate_graph_with_config(config, message_handler)
338 }
339
340 #[cfg(feature = "viz")]
341 pub fn generate_all_files_with_config(
342 &self,
343 config: &crate::graph_util::GraphConfig,
344 prefix: &str,
345 ) -> Result<(), Box<dyn std::error::Error>> {
346 self.graph_api()
347 .generate_all_files_with_config(config, prefix)
348 }
349}