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 #[cfg(feature = "runtime_support")]
54 pub fn ir_json(&self) -> Result<String, serde_json::Error> {
55 super::ir::serialize_dedup_shared(|| serde_json::to_string_pretty(&self.ir))
56 }
57
58 pub fn location_names(&self) -> &SecondaryMap<LocationKey, String> {
60 &self.location_names
61 }
62
63 #[cfg(stageleft_runtime)]
65 #[cfg(feature = "viz")]
66 pub fn graph_api(&self) -> GraphApi<'_> {
67 GraphApi::new(&self.ir, self.location_names())
68 }
69
70 #[cfg(feature = "viz")]
72 pub fn render_graph(
73 &self,
74 format: crate::viz::config::GraphType,
75 use_short_labels: bool,
76 show_metadata: bool,
77 ) -> String {
78 self.graph_api()
79 .render(format, use_short_labels, show_metadata)
80 }
81
82 #[cfg(feature = "viz")]
84 pub fn write_graph_to_file(
85 &self,
86 format: crate::viz::config::GraphType,
87 filename: &str,
88 use_short_labels: bool,
89 show_metadata: bool,
90 ) -> Result<(), Box<dyn std::error::Error>> {
91 self.graph_api()
92 .write_to_file(format, filename, use_short_labels, show_metadata)
93 }
94
95 #[cfg(feature = "viz")]
97 pub fn generate_graph(
98 &self,
99 config: &crate::viz::config::GraphConfig,
100 ) -> Result<Option<String>, Box<dyn std::error::Error>> {
101 self.graph_api().generate_graph(config)
102 }
103
104 pub fn optimize_with(mut self, f: impl FnOnce(&mut [HydroRoot])) -> Self {
105 f(&mut self.ir);
106 self
107 }
108
109 pub fn with_default_optimize<D: Deploy<'a>>(self) -> DeployFlow<'a, D> {
110 self.into_deploy()
111 }
112
113 #[cfg(feature = "sim")]
114 pub fn sim(self) -> SimFlow<'a> {
117 use std::cell::RefCell;
118 use std::rc::Rc;
119
120 use slotmap::SparseSecondaryMap;
121
122 use crate::sim::graph::SimNodePort;
123
124 let shared_port_counter = Rc::new(RefCell::new(SimNodePort::default()));
125
126 let mut processes = SparseSecondaryMap::new();
127 let mut clusters = SparseSecondaryMap::new();
128 let externals = SparseSecondaryMap::new();
129
130 for (key, loc) in self.locations.iter() {
131 match loc {
132 LocationType::Process => {
133 processes.insert(
134 key,
135 SimNode {
136 shared_port_counter: shared_port_counter.clone(),
137 },
138 );
139 }
140 LocationType::Cluster => {
141 clusters.insert(
142 key,
143 SimNode {
144 shared_port_counter: shared_port_counter.clone(),
145 },
146 );
147 }
148 LocationType::External => {
149 panic!("Sim cannot have externals");
150 }
151 }
152 }
153
154 SimFlow {
155 ir: self.ir,
156 processes,
157 clusters,
158 externals,
159 cluster_max_sizes: SparseSecondaryMap::new(),
160 externals_port_registry: Default::default(),
161 test_safety_only: false,
162 skip_consistency_assertions: false,
163 unit_test_fuzz_iterations: 8192,
164 _phantom: PhantomData,
165 }
166 }
167
168 pub fn into_deploy<D: Deploy<'a>>(self) -> DeployFlow<'a, D> {
169 let (processes, clusters, externals) = Default::default();
170 DeployFlow {
171 ir: self.ir,
172 locations: self.locations,
173 location_names: self.location_names,
174 processes,
175 clusters,
176 externals,
177 sidecars: SparseSecondaryMap::new(),
178 flow_name: self.flow_name,
179 _phantom: PhantomData,
180 }
181 }
182
183 pub fn with_process<P, D: Deploy<'a>>(
184 self,
185 process: &Process<P>,
186 spec: impl IntoProcessSpec<'a, D>,
187 ) -> DeployFlow<'a, D> {
188 self.into_deploy().with_process(process, spec)
189 }
190
191 pub fn with_remaining_processes<D: Deploy<'a>, S: IntoProcessSpec<'a, D> + 'a>(
192 self,
193 spec: impl Fn() -> S,
194 ) -> DeployFlow<'a, D> {
195 self.into_deploy().with_remaining_processes(spec)
196 }
197
198 pub fn with_external<P, D: Deploy<'a>>(
199 self,
200 process: &External<P>,
201 spec: impl ExternalSpec<'a, D>,
202 ) -> DeployFlow<'a, D> {
203 self.into_deploy().with_external(process, spec)
204 }
205
206 pub fn with_remaining_externals<D: Deploy<'a>, S: ExternalSpec<'a, D> + 'a>(
207 self,
208 spec: impl Fn() -> S,
209 ) -> DeployFlow<'a, D> {
210 self.into_deploy().with_remaining_externals(spec)
211 }
212
213 pub fn with_cluster<C, D: Deploy<'a>>(
214 self,
215 cluster: &Cluster<C>,
216 spec: impl ClusterSpec<'a, D>,
217 ) -> DeployFlow<'a, D> {
218 self.into_deploy().with_cluster(cluster, spec)
219 }
220
221 pub fn with_remaining_clusters<D: Deploy<'a>, S: ClusterSpec<'a, D> + 'a>(
222 self,
223 spec: impl Fn() -> S,
224 ) -> DeployFlow<'a, D> {
225 self.into_deploy().with_remaining_clusters(spec)
226 }
227
228 pub fn compile<D: Deploy<'a, InstantiateEnv = ()>>(self) -> CompiledFlow<'a> {
229 self.into_deploy::<D>().compile()
230 }
231
232 pub fn deploy<D: Deploy<'a>>(self, env: &mut D::InstantiateEnv) -> DeployResult<'a, D> {
233 self.into_deploy::<D>().deploy(env)
234 }
235}