1use bevy_ecs::{prelude::Entity, world::World};
2use bevy_platform::collections::HashMap;
3#[cfg(feature = "trace")]
4use tracing::info_span;
5
6use alloc::{borrow::Cow, collections::VecDeque};
7use smallvec::{smallvec, SmallVec};
8use thiserror::Error;
9
10use crate::{
11 diagnostic::internal::{DiagnosticsRecorder, RenderDiagnosticsMutex},
12 render_graph::{
13 Edge, InternedRenderLabel, InternedRenderSubGraph, NodeRunError, NodeState, RenderGraph,
14 RenderGraphContext, SlotLabel, SlotType, SlotValue,
15 },
16 renderer::{RenderContext, RenderDevice},
17};
18
19pub(crate) struct RenderGraphRunner;
31
32#[derive(Error, Debug)]
33pub enum RenderGraphRunnerError {
34 #[error(transparent)]
35 NodeRunError(#[from] NodeRunError),
36 #[error("node output slot not set (index {slot_index}, name {slot_name})")]
37 EmptyNodeOutputSlot {
38 type_name: &'static str,
39 slot_index: usize,
40 slot_name: Cow<'static, str>,
41 },
42 #[error("graph '{sub_graph:?}' could not be run because slot '{slot_name}' at index {slot_index} has no value")]
43 MissingInput {
44 slot_index: usize,
45 slot_name: Cow<'static, str>,
46 sub_graph: Option<InternedRenderSubGraph>,
47 },
48 #[error("attempted to use the wrong type for input slot")]
49 MismatchedInputSlotType {
50 slot_index: usize,
51 label: SlotLabel,
52 expected: SlotType,
53 actual: SlotType,
54 },
55 #[error(
56 "node (name: '{node_name:?}') has {slot_count} input slots, but was provided {value_count} values"
57 )]
58 MismatchedInputCount {
59 node_name: InternedRenderLabel,
60 slot_count: usize,
61 value_count: usize,
62 },
63}
64
65impl RenderGraphRunner {
66 pub fn run(
67 graph: &RenderGraph,
68 render_device: RenderDevice,
69 mut diagnostics_recorder: Option<DiagnosticsRecorder>,
70 queue: &wgpu::Queue,
71 world: &World,
72 finalizer: impl FnOnce(&mut wgpu::CommandEncoder),
73 ) -> Result<Option<DiagnosticsRecorder>, RenderGraphRunnerError> {
74 if let Some(recorder) = &mut diagnostics_recorder {
75 recorder.begin_frame();
76 }
77
78 let mut render_context = RenderContext::new(render_device, diagnostics_recorder);
79 Self::run_graph(graph, None, &mut render_context, world, &[], None)?;
80 finalizer(render_context.command_encoder());
81
82 let (render_device, mut diagnostics_recorder) = {
83 let (commands, render_device, diagnostics_recorder) = render_context.finish();
84
85 #[cfg(feature = "trace")]
86 let _span = info_span!("submit_graph_commands").entered();
87 queue.submit(commands);
88
89 (render_device, diagnostics_recorder)
90 };
91
92 if let Some(recorder) = &mut diagnostics_recorder {
93 let render_diagnostics_mutex = world.resource::<RenderDiagnosticsMutex>().0.clone();
94 recorder.finish_frame(&render_device, move |diagnostics| {
95 *render_diagnostics_mutex.lock().expect("lock poisoned") = Some(diagnostics);
96 });
97 }
98
99 Ok(diagnostics_recorder)
100 }
101
102 fn run_graph<'w>(
105 graph: &RenderGraph,
106 sub_graph: Option<InternedRenderSubGraph>,
107 render_context: &mut RenderContext<'w>,
108 world: &'w World,
109 inputs: &[SlotValue],
110 view_entity: Option<Entity>,
111 ) -> Result<(), RenderGraphRunnerError> {
112 let mut node_outputs: HashMap<InternedRenderLabel, SmallVec<[SlotValue; 4]>> =
113 HashMap::default();
114 #[cfg(feature = "trace")]
115 let span = if let Some(label) = &sub_graph {
116 info_span!("run_graph", name = format!("{label:?}"))
117 } else {
118 info_span!("run_graph", name = "main_graph")
119 };
120 #[cfg(feature = "trace")]
121 let _guard = span.enter();
122
123 let mut node_queue: VecDeque<&NodeState> = graph
125 .iter_nodes()
126 .filter(|node| node.input_slots.is_empty())
127 .collect();
128
129 if let Some(input_node) = graph.get_input_node() {
131 let mut input_values: SmallVec<[SlotValue; 4]> = SmallVec::new();
132 for (i, input_slot) in input_node.input_slots.iter().enumerate() {
133 if let Some(input_value) = inputs.get(i) {
134 if input_slot.slot_type != input_value.slot_type() {
135 return Err(RenderGraphRunnerError::MismatchedInputSlotType {
136 slot_index: i,
137 actual: input_value.slot_type(),
138 expected: input_slot.slot_type,
139 label: input_slot.name.clone().into(),
140 });
141 }
142 input_values.push(input_value.clone());
143 } else {
144 return Err(RenderGraphRunnerError::MissingInput {
145 slot_index: i,
146 slot_name: input_slot.name.clone(),
147 sub_graph,
148 });
149 }
150 }
151
152 node_outputs.insert(input_node.label, input_values);
153
154 for (_, node_state) in graph
155 .iter_node_outputs(input_node.label)
156 .expect("node exists")
157 {
158 node_queue.push_front(node_state);
159 }
160 }
161
162 'handle_node: while let Some(node_state) = node_queue.pop_back() {
163 if node_outputs.contains_key(&node_state.label) {
165 continue;
166 }
167
168 let mut slot_indices_and_inputs: SmallVec<[(usize, SlotValue); 4]> = SmallVec::new();
169 for (edge, input_node) in graph
171 .iter_node_inputs(node_state.label)
172 .expect("node is in graph")
173 {
174 match edge {
175 Edge::SlotEdge {
176 output_index,
177 input_index,
178 ..
179 } => {
180 if let Some(outputs) = node_outputs.get(&input_node.label) {
181 slot_indices_and_inputs
182 .push((*input_index, outputs[*output_index].clone()));
183 } else {
184 node_queue.push_front(node_state);
185 continue 'handle_node;
186 }
187 }
188 Edge::NodeEdge { .. } => {
189 if !node_outputs.contains_key(&input_node.label) {
190 node_queue.push_front(node_state);
191 continue 'handle_node;
192 }
193 }
194 }
195 }
196
197 slot_indices_and_inputs.sort_by_key(|(index, _)| *index);
199 let inputs: SmallVec<[SlotValue; 4]> = slot_indices_and_inputs
200 .into_iter()
201 .map(|(_, value)| value)
202 .collect();
203
204 if inputs.len() != node_state.input_slots.len() {
205 return Err(RenderGraphRunnerError::MismatchedInputCount {
206 node_name: node_state.label,
207 slot_count: node_state.input_slots.len(),
208 value_count: inputs.len(),
209 });
210 }
211
212 let mut outputs: SmallVec<[Option<SlotValue>; 4]> =
213 smallvec![None; node_state.output_slots.len()];
214 {
215 let mut context = RenderGraphContext::new(graph, node_state, &inputs, &mut outputs);
216 if let Some(view_entity) = view_entity {
217 context.set_view_entity(view_entity);
218 }
219
220 {
221 #[cfg(feature = "trace")]
222 let _span = info_span!("node", name = node_state.type_name).entered();
223
224 node_state.node.run(&mut context, render_context, world)?;
225 }
226
227 for run_sub_graph in context.finish() {
228 let sub_graph = graph
229 .get_sub_graph(run_sub_graph.sub_graph)
230 .expect("sub graph exists because it was validated when queued.");
231 Self::run_graph(
232 sub_graph,
233 Some(run_sub_graph.sub_graph),
234 render_context,
235 world,
236 &run_sub_graph.inputs,
237 run_sub_graph.view_entity,
238 )?;
239 }
240 }
241
242 let mut values: SmallVec<[SlotValue; 4]> = SmallVec::new();
243 for (i, output) in outputs.into_iter().enumerate() {
244 if let Some(value) = output {
245 values.push(value);
246 } else {
247 let empty_slot = node_state.output_slots.get_slot(i).unwrap();
248 return Err(RenderGraphRunnerError::EmptyNodeOutputSlot {
249 type_name: node_state.type_name,
250 slot_index: i,
251 slot_name: empty_slot.name.clone(),
252 });
253 }
254 }
255 node_outputs.insert(node_state.label, values);
256
257 for (_, node_state) in graph
258 .iter_node_outputs(node_state.label)
259 .expect("node exists")
260 {
261 node_queue.push_front(node_state);
262 }
263 }
264
265 Ok(())
266 }
267}