1use std::marker::PhantomData;
2
3use dfir_lang::graph::{
4 DfirGraph, FlatGraphBuilderOutput, PartitionError, eliminate_extra_unions_tees, partition_graph,
5};
6use slotmap::{SecondaryMap, SlotMap};
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) sidecars: Vec<super::builder::Sidecar>,
28
29 pub(super) flow_name: String,
31
32 #[cfg(feature = "sim")]
35 pub(super) location_version: SecondaryMap<LocationKey, u32>,
36
37 #[cfg(feature = "sim")]
40 pub(super) location_version_group_root: SecondaryMap<LocationKey, LocationKey>,
41
42 pub(super) _phantom: Invariant<'a>,
43}
44
45pub(crate) fn build_inner(
56 ir: &mut Vec<HydroRoot>,
57) -> SecondaryMap<LocationKey, Result<DfirGraph, PartitionError>> {
58 emit(ir)
59 .into_iter()
60 .map(|(k, v)| {
61 let FlatGraphBuilderOutput { mut flat_graph, .. } =
62 v.build().expect("Failed to build DFIR flat graph.");
63 eliminate_extra_unions_tees(&mut flat_graph);
64 (k, partition_graph(flat_graph))
65 })
66 .collect()
67}
68
69impl<'a> BuiltFlow<'a> {
70 pub fn ir(&self) -> &[HydroRoot] {
72 &self.ir
73 }
74
75 #[cfg(feature = "viz")]
77 pub fn ir_json(&self) -> Result<String, serde_json::Error> {
78 super::ir::serialize_dedup_shared(|| serde_json::to_string_pretty(&self.ir))
79 }
80
81 pub fn location_names(&self) -> &SecondaryMap<LocationKey, String> {
83 &self.location_names
84 }
85
86 #[cfg(stageleft_runtime)]
88 #[cfg(feature = "viz")]
89 pub fn graph_api(&self) -> GraphApi<'_> {
90 GraphApi::new(&self.ir, self.location_names())
91 }
92
93 #[cfg(feature = "viz")]
95 pub fn render_graph(
96 &self,
97 format: crate::viz::config::GraphType,
98 use_short_labels: bool,
99 show_metadata: bool,
100 ) -> String {
101 self.graph_api()
102 .render(format, use_short_labels, show_metadata)
103 }
104
105 #[cfg(feature = "viz")]
107 pub fn write_graph_to_file(
108 &self,
109 format: crate::viz::config::GraphType,
110 filename: &str,
111 use_short_labels: bool,
112 show_metadata: bool,
113 ) -> Result<(), Box<dyn std::error::Error>> {
114 self.graph_api()
115 .write_to_file(format, filename, use_short_labels, show_metadata)
116 }
117
118 #[cfg(feature = "viz")]
120 pub fn generate_graph(
121 &self,
122 config: &crate::viz::config::GraphConfig,
123 ) -> Result<Option<String>, Box<dyn std::error::Error>> {
124 self.graph_api().generate_graph(config)
125 }
126
127 pub fn optimize_with(mut self, f: impl FnOnce(&mut [HydroRoot])) -> Self {
128 f(&mut self.ir);
129 self
130 }
131
132 pub fn with_default_optimize<D: Deploy<'a>>(self) -> DeployFlow<'a, D> {
133 self.into_deploy()
134 }
135
136 #[cfg(feature = "sim")]
137 pub fn sim(self) -> SimFlow<'a> {
140 use std::cell::RefCell;
141 use std::rc::Rc;
142
143 use slotmap::SparseSecondaryMap;
144
145 use crate::sim::graph::SimNodePort;
146
147 let shared_port_counter = Rc::new(RefCell::new(crate::Counter::<SimNodePort>::default()));
148
149 let mut processes = SparseSecondaryMap::new();
150 let mut clusters = SparseSecondaryMap::new();
151 let externals = SparseSecondaryMap::new();
152
153 for (key, loc) in self.locations.iter() {
154 match loc {
155 LocationType::Process => {
156 processes.insert(
157 key,
158 SimNode {
159 shared_port_counter: shared_port_counter.clone(),
160 },
161 );
162 }
163 LocationType::Cluster => {
164 clusters.insert(
165 key,
166 SimNode {
167 shared_port_counter: shared_port_counter.clone(),
168 },
169 );
170 }
171 LocationType::External => {
172 panic!("Sim cannot have externals");
173 }
174 }
175 }
176
177 SimFlow {
178 ir: self.ir,
179 processes,
180 clusters,
181 externals,
182 cluster_max_sizes: SparseSecondaryMap::new(),
183 externals_port_registry: Default::default(),
184 location_version: self.location_version,
185 location_version_group_root: self.location_version_group_root,
186 test_safety_only: false,
187 skip_consistency_assertions: false,
188 unit_test_fuzz_iterations: 8192,
189 _phantom: PhantomData,
190 }
191 }
192
193 pub fn into_deploy<D: Deploy<'a>>(self) -> DeployFlow<'a, D> {
194 let (processes, clusters, externals) = Default::default();
195 DeployFlow {
196 ir: self.ir,
197 locations: self.locations,
198 location_names: self.location_names,
199 processes,
200 clusters,
201 externals,
202 sidecars: self.sidecars,
203 flow_name: self.flow_name,
204 _phantom: PhantomData,
205 }
206 }
207
208 pub fn with_process<P, D: Deploy<'a>>(
209 self,
210 process: &Process<'_, P>,
211 spec: impl IntoProcessSpec<'a, D>,
212 ) -> DeployFlow<'a, D> {
213 self.into_deploy().with_process(process, spec)
214 }
215
216 pub fn with_remaining_processes<D: Deploy<'a>, S: IntoProcessSpec<'a, D> + 'a>(
217 self,
218 spec: impl Fn() -> S,
219 ) -> DeployFlow<'a, D> {
220 self.into_deploy().with_remaining_processes(spec)
221 }
222
223 pub fn with_external<P, D: Deploy<'a>>(
224 self,
225 process: &External<'_, P>,
226 spec: impl ExternalSpec<'a, D>,
227 ) -> DeployFlow<'a, D> {
228 self.into_deploy().with_external(process, spec)
229 }
230
231 pub fn with_remaining_externals<D: Deploy<'a>, S: ExternalSpec<'a, D> + 'a>(
232 self,
233 spec: impl Fn() -> S,
234 ) -> DeployFlow<'a, D> {
235 self.into_deploy().with_remaining_externals(spec)
236 }
237
238 pub fn with_cluster<C, D: Deploy<'a>>(
239 self,
240 cluster: &Cluster<'_, C>,
241 spec: impl ClusterSpec<'a, D>,
242 ) -> DeployFlow<'a, D> {
243 self.into_deploy().with_cluster(cluster, spec)
244 }
245
246 pub fn with_remaining_clusters<D: Deploy<'a>, S: ClusterSpec<'a, D> + 'a>(
247 self,
248 spec: impl Fn() -> S,
249 ) -> DeployFlow<'a, D> {
250 self.into_deploy().with_remaining_clusters(spec)
251 }
252
253 pub fn compile<D: Deploy<'a, InstantiateEnv = ()>>(self) -> CompiledFlow<'a> {
254 self.into_deploy::<D>().compile()
255 }
256
257 pub fn deploy<D: Deploy<'a>>(self, env: &mut D::InstantiateEnv) -> DeployResult<'a, D> {
258 self.into_deploy::<D>().deploy(env)
259 }
260}