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};
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.into_deploy()
211 }
212
213 pub fn into_deploy<D: Deploy<'a>>(mut self) -> DeployFlow<'a, D> {
214 let processes = if D::has_trivial_node() {
215 self.process_id_name
216 .iter()
217 .map(|id| (id.0, D::trivial_process(id.0)))
218 .collect()
219 } else {
220 HashMap::new()
221 };
222
223 let clusters = if D::has_trivial_node() {
224 self.cluster_id_name
225 .iter()
226 .map(|id| (id.0, D::trivial_cluster(id.0)))
227 .collect()
228 } else {
229 HashMap::new()
230 };
231
232 let externals = if D::has_trivial_node() {
233 self.external_id_name
234 .iter()
235 .map(|id| (id.0, D::trivial_external(id.0)))
236 .collect()
237 } else {
238 HashMap::new()
239 };
240
241 DeployFlow {
242 ir: UnsafeCell::new(std::mem::take(&mut self.ir)),
243 processes,
244 process_id_name: std::mem::take(&mut self.process_id_name),
245 clusters,
246 cluster_id_name: std::mem::take(&mut self.cluster_id_name),
247 externals,
248 external_id_name: std::mem::take(&mut self.external_id_name),
249 _phantom: PhantomData,
250 }
251 }
252
253 pub fn with_process<P, D: Deploy<'a>>(
254 self,
255 process: &Process<P>,
256 spec: impl IntoProcessSpec<'a, D>,
257 ) -> DeployFlow<'a, D> {
258 self.into_deploy().with_process(process, spec)
259 }
260
261 pub fn with_remaining_processes<D: Deploy<'a>, S: IntoProcessSpec<'a, D> + 'a>(
262 self,
263 spec: impl Fn() -> S,
264 ) -> DeployFlow<'a, D> {
265 self.into_deploy().with_remaining_processes(spec)
266 }
267
268 pub fn with_external<P, D: Deploy<'a>>(
269 self,
270 process: &External<P>,
271 spec: impl ExternalSpec<'a, D>,
272 ) -> DeployFlow<'a, D> {
273 self.into_deploy().with_external(process, spec)
274 }
275
276 pub fn with_remaining_externals<D: Deploy<'a>, S: ExternalSpec<'a, D> + 'a>(
277 self,
278 spec: impl Fn() -> S,
279 ) -> DeployFlow<'a, D> {
280 self.into_deploy().with_remaining_externals(spec)
281 }
282
283 pub fn with_cluster<C, D: Deploy<'a>>(
284 self,
285 cluster: &Cluster<C>,
286 spec: impl ClusterSpec<'a, D>,
287 ) -> DeployFlow<'a, D> {
288 self.into_deploy().with_cluster(cluster, spec)
289 }
290
291 pub fn with_remaining_clusters<D: Deploy<'a>, S: ClusterSpec<'a, D> + 'a>(
292 self,
293 spec: impl Fn() -> S,
294 ) -> DeployFlow<'a, D> {
295 self.into_deploy().with_remaining_clusters(spec)
296 }
297
298 pub fn compile<D: Deploy<'a>>(self, env: &D::CompileEnv) -> CompiledFlow<'a, D::GraphId> {
299 self.into_deploy::<D>().compile(env)
300 }
301
302 pub fn compile_no_network<D: Deploy<'a>>(self) -> CompiledFlow<'a, D::GraphId> {
303 self.into_deploy::<D>().compile_no_network()
304 }
305
306 pub fn deploy<D: Deploy<'a, CompileEnv = ()>>(
307 self,
308 env: &mut D::InstantiateEnv,
309 ) -> DeployResult<'a, D> {
310 self.into_deploy::<D>().deploy(env)
311 }
312
313 #[cfg(feature = "viz")]
314 pub fn generate_all_files(
315 &self,
316 prefix: &str,
317 show_metadata: bool,
318 show_location_groups: bool,
319 use_short_labels: bool,
320 ) -> Result<(), Box<dyn std::error::Error>> {
321 self.graph_api().generate_all_files(
322 prefix,
323 show_metadata,
324 show_location_groups,
325 use_short_labels,
326 )
327 }
328
329 #[cfg(feature = "viz")]
330 pub fn generate_graph_with_config(
331 &self,
332 config: &crate::graph::config::GraphConfig,
333 message_handler: Option<&dyn Fn(&str)>,
334 ) -> Result<(), Box<dyn std::error::Error>> {
335 self.graph_api()
336 .generate_graph_with_config(config, message_handler)
337 }
338
339 #[cfg(feature = "viz")]
340 pub fn generate_all_files_with_config(
341 &self,
342 config: &crate::graph::config::GraphConfig,
343 prefix: &str,
344 ) -> Result<(), Box<dyn std::error::Error>> {
345 self.graph_api()
346 .generate_all_files_with_config(config, prefix)
347 }
348}