1use std::{any::TypeId, collections::HashMap};
8
9use crate::{
10 Command, CompositeCommand, PxPosition, PxSize,
11 render_graph::{
12 ExternalTextureDesc, RenderGraph, RenderGraphOp, RenderGraphParts, RenderResource,
13 RenderResourceId,
14 },
15};
16
17use super::{core::RenderResources, external::ExternalTextureRegistry};
18
19pub struct CompositeContext<'a> {
21 pub resources: RenderResources<'a>,
23 pub external_textures: ExternalTextureRegistry,
25 pub frame_size: PxSize,
27 pub surface_format: wgpu::TextureFormat,
29 pub sample_count: u32,
31 pub frame_index: u64,
33}
34
35pub struct ErasedCompositeBatchItem<'a> {
37 pub command: &'a dyn CompositeCommand,
39 pub size: PxSize,
41 pub position: PxPosition,
43 pub opacity: f32,
45 pub sequence_index: usize,
47 pub op_index: usize,
49}
50
51pub struct CompositeBatchItem<'a, C: CompositeCommand> {
53 pub command: &'a C,
55 pub size: PxSize,
57 pub position: PxPosition,
59 pub opacity: f32,
61 pub sequence_index: usize,
63 pub op_index: usize,
65}
66
67pub struct CompositeReplacement {
69 pub target_op: usize,
71 pub ops: Vec<RenderGraphOp>,
73}
74
75pub struct CompositeOutput {
77 pub resources: Vec<RenderResource>,
79 pub external_resources: Vec<ExternalTextureDesc>,
81 pub prelude_ops: Vec<RenderGraphOp>,
83 pub replacements: Vec<CompositeReplacement>,
85}
86
87impl CompositeOutput {
88 pub fn empty() -> Self {
90 Self {
91 resources: Vec::new(),
92 external_resources: Vec::new(),
93 prelude_ops: Vec::new(),
94 replacements: Vec::new(),
95 }
96 }
97
98 pub fn add_external_texture(&mut self, desc: ExternalTextureDesc) -> RenderResourceId {
100 let index = self.external_resources.len() as u32;
101 self.external_resources.push(desc);
102 RenderResourceId::External(index)
103 }
104}
105
106pub trait CompositePipeline<C: CompositeCommand>: Send + Sync + 'static {
108 fn compile(
110 &mut self,
111 context: &CompositeContext<'_>,
112 items: &[CompositeBatchItem<'_, C>],
113 ) -> CompositeOutput;
114}
115
116pub(crate) trait ErasedCompositePipeline: Send + Sync {
118 fn compile_erased(
119 &mut self,
120 context: &CompositeContext<'_>,
121 items: &[ErasedCompositeBatchItem<'_>],
122 ) -> CompositeOutput;
123}
124
125struct CompositePipelineImpl<C: CompositeCommand, P: CompositePipeline<C>> {
126 pipeline: P,
127 _command: std::marker::PhantomData<C>,
128}
129
130impl<C: CompositeCommand + 'static, P: CompositePipeline<C>> ErasedCompositePipeline
131 for CompositePipelineImpl<C, P>
132{
133 fn compile_erased(
134 &mut self,
135 context: &CompositeContext<'_>,
136 items: &[ErasedCompositeBatchItem<'_>],
137 ) -> CompositeOutput {
138 if items.is_empty() {
139 return CompositeOutput::empty();
140 }
141
142 let mut typed_items: Vec<CompositeBatchItem<'_, C>> = Vec::with_capacity(items.len());
143 for item in items {
144 let command = item
145 .command
146 .downcast_ref::<C>()
147 .expect("Composite batch contained command of unexpected type");
148 typed_items.push(CompositeBatchItem {
149 command,
150 size: item.size,
151 position: item.position,
152 opacity: item.opacity,
153 sequence_index: item.sequence_index,
154 op_index: item.op_index,
155 });
156 }
157
158 self.pipeline.compile(context, &typed_items)
159 }
160}
161
162#[derive(Default)]
164pub struct CompositePipelineRegistry {
165 pipelines: HashMap<TypeId, Box<dyn ErasedCompositePipeline>>,
166}
167
168impl CompositePipelineRegistry {
169 pub fn new() -> Self {
171 Self::default()
172 }
173
174 pub fn register<C: CompositeCommand + 'static>(
176 &mut self,
177 pipeline: impl CompositePipeline<C> + 'static,
178 ) {
179 let erased = Box::new(CompositePipelineImpl {
180 pipeline,
181 _command: std::marker::PhantomData,
182 });
183 self.pipelines.insert(TypeId::of::<C>(), erased);
184 }
185
186 pub(crate) fn compile_erased(
187 &mut self,
188 context: &CompositeContext<'_>,
189 items: &[ErasedCompositeBatchItem<'_>],
190 ) -> CompositeOutput {
191 if items.is_empty() {
192 return CompositeOutput::empty();
193 }
194
195 let command_type_id = items[0].command.as_any().type_id();
196 if let Some(pipeline) = self.pipelines.get_mut(&command_type_id) {
197 pipeline.compile_erased(context, items)
198 } else {
199 panic!(
200 "No composite pipeline found for command {:?}",
201 std::any::type_name_of_val(items[0].command)
202 );
203 }
204 }
205}
206
207pub(crate) fn expand_composites(
208 scene: RenderGraph,
209 context: CompositeContext<'_>,
210 registry: &mut CompositePipelineRegistry,
211) -> RenderGraph {
212 let RenderGraphParts {
213 ops,
214 resources,
215 external_resources,
216 } = scene.into_parts();
217
218 let mut type_order: Vec<TypeId> = Vec::new();
219 let mut batches: HashMap<TypeId, Vec<ErasedCompositeBatchItem<'_>>> = HashMap::new();
220
221 for (index, op) in ops.iter().enumerate() {
222 if let Command::Composite(command) = &op.command {
223 let type_id = command.as_any().type_id();
224 let entry = batches.entry(type_id).or_insert_with(|| {
225 type_order.push(type_id);
226 Vec::new()
227 });
228 entry.push(ErasedCompositeBatchItem {
229 command: command.as_ref(),
230 size: op.size,
231 position: op.position,
232 opacity: op.opacity,
233 sequence_index: op.sequence_index,
234 op_index: index,
235 });
236 }
237 }
238
239 if type_order.is_empty() {
240 return RenderGraph::from_parts(RenderGraphParts {
241 ops,
242 resources,
243 external_resources,
244 });
245 }
246
247 let mut new_resources = resources;
248 let mut new_external_resources = external_resources;
249 let mut prelude_ops: Vec<RenderGraphOp> = Vec::new();
250 let mut replacements: HashMap<usize, Vec<Vec<RenderGraphOp>>> = HashMap::new();
251
252 for type_id in type_order {
253 let items = batches
254 .get(&type_id)
255 .expect("composite batch missing type entry");
256 let output = registry.compile_erased(&context, items);
257
258 let resource_map = map_resources(&mut new_resources, &output.resources);
259 let external_map =
260 map_external_resources(&mut new_external_resources, &output.external_resources);
261
262 let prelude_base = prelude_ops.len();
263 let mut mapped_prelude = output.prelude_ops;
264 for op in &mut mapped_prelude {
265 remap_resources(op, &resource_map, &external_map);
266 remap_deps(op, prelude_base);
267 }
268 prelude_ops.extend(mapped_prelude);
269
270 for replacement in output.replacements {
271 let mut mapped_ops = replacement.ops;
272 for op in &mut mapped_ops {
273 remap_resources(op, &resource_map, &external_map);
274 }
275 replacements
276 .entry(replacement.target_op)
277 .or_default()
278 .push(mapped_ops);
279 }
280 }
281
282 let mut new_ops: Vec<RenderGraphOp> = Vec::with_capacity(prelude_ops.len() + ops.len());
283 new_ops.extend(prelude_ops);
284
285 for (index, op) in ops.into_iter().enumerate() {
286 match op.command {
287 Command::Composite(_) => {
288 if let Some(mut fragments) = replacements.remove(&index) {
289 for mut fragment_ops in fragments.drain(..) {
290 let base = new_ops.len();
291 for op in &mut fragment_ops {
292 remap_deps(op, base);
293 }
294 new_ops.extend(fragment_ops);
295 }
296 }
297 }
298 _ => new_ops.push(op),
299 }
300 }
301
302 if !replacements.is_empty() {
303 for mut fragments in replacements.into_values() {
304 for mut fragment_ops in fragments.drain(..) {
305 let base = new_ops.len();
306 for op in &mut fragment_ops {
307 remap_deps(op, base);
308 }
309 new_ops.extend(fragment_ops);
310 }
311 }
312 }
313
314 for (seq, op) in new_ops.iter_mut().enumerate() {
315 op.sequence_index = seq;
316 }
317
318 RenderGraph::from_parts(RenderGraphParts {
319 ops: new_ops,
320 resources: new_resources,
321 external_resources: new_external_resources,
322 })
323}
324
325fn map_resources(
326 resources: &mut Vec<RenderResource>,
327 locals: &[RenderResource],
328) -> Vec<RenderResourceId> {
329 let mut map = Vec::with_capacity(locals.len());
330 for resource in locals {
331 let index = resources.len() as u32;
332 resources.push(resource.clone());
333 map.push(RenderResourceId::Local(index));
334 }
335 map
336}
337
338fn map_external_resources(
339 resources: &mut Vec<ExternalTextureDesc>,
340 externals: &[ExternalTextureDesc],
341) -> Vec<RenderResourceId> {
342 let mut map = Vec::with_capacity(externals.len());
343 for resource in externals {
344 if let Some(index) = resources
345 .iter()
346 .position(|existing| existing.handle_id == resource.handle_id)
347 {
348 map.push(RenderResourceId::External(index as u32));
349 continue;
350 }
351 let index = resources.len() as u32;
352 resources.push(resource.clone());
353 map.push(RenderResourceId::External(index));
354 }
355 map
356}
357
358fn remap_resources(
359 op: &mut RenderGraphOp,
360 local_map: &[RenderResourceId],
361 external_map: &[RenderResourceId],
362) {
363 op.read = op
364 .read
365 .map(|resource| map_resource(resource, local_map, external_map));
366 op.write = op
367 .write
368 .map(|resource| map_resource(resource, local_map, external_map));
369}
370
371fn map_resource(
372 resource: RenderResourceId,
373 local_map: &[RenderResourceId],
374 external_map: &[RenderResourceId],
375) -> RenderResourceId {
376 match resource {
377 RenderResourceId::Local(index) => local_map
378 .get(index as usize)
379 .copied()
380 .unwrap_or(RenderResourceId::Local(index)),
381 RenderResourceId::External(index) => external_map
382 .get(index as usize)
383 .copied()
384 .unwrap_or(RenderResourceId::External(index)),
385 other => other,
386 }
387}
388
389fn remap_deps(op: &mut RenderGraphOp, base: usize) {
390 if base == 0 {
391 return;
392 }
393 for dep in op.deps.iter_mut() {
394 *dep += base;
395 }
396}