tessera_ui/renderer/
composite.rs

1//! Composite command expansion for Tessera render graphs.
2//!
3//! ## Usage
4//!
5//! Expand composite commands into draw/compute ops before rendering.
6
7use 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
19/// Context provided to composite pipelines during expansion.
20pub struct CompositeContext<'a> {
21    /// Shared GPU resources for pipeline usage.
22    pub resources: RenderResources<'a>,
23    /// Registry for external persistent textures.
24    pub external_textures: ExternalTextureRegistry,
25    /// Pixel size of the current frame.
26    pub frame_size: PxSize,
27    /// Surface format for the current frame.
28    pub surface_format: wgpu::TextureFormat,
29    /// MSAA sample count for the current frame.
30    pub sample_count: u32,
31    /// Monotonic frame index.
32    pub frame_index: u64,
33}
34
35/// Type-erased metadata describing a composite command within a batch.
36pub struct ErasedCompositeBatchItem<'a> {
37    /// The composite command to expand.
38    pub command: &'a dyn CompositeCommand,
39    /// Measured size of the target region.
40    pub size: PxSize,
41    /// Absolute position of the target region.
42    pub position: PxPosition,
43    /// Opacity multiplier for the command.
44    pub opacity: f32,
45    /// Original sequence index for ordering.
46    pub sequence_index: usize,
47    /// Index of the op in the source render graph.
48    pub op_index: usize,
49}
50
51/// Strongly typed metadata describing a composite command within a batch.
52pub struct CompositeBatchItem<'a, C: CompositeCommand> {
53    /// The composite command to expand.
54    pub command: &'a C,
55    /// Measured size of the target region.
56    pub size: PxSize,
57    /// Absolute position of the target region.
58    pub position: PxPosition,
59    /// Opacity multiplier for the command.
60    pub opacity: f32,
61    /// Original sequence index for ordering.
62    pub sequence_index: usize,
63    /// Index of the op in the source render graph.
64    pub op_index: usize,
65}
66
67/// Replacement ops emitted for a composite command.
68pub struct CompositeReplacement {
69    /// Index of the original composite op to replace.
70    pub target_op: usize,
71    /// Ops that replace the composite command.
72    pub ops: Vec<RenderGraphOp>,
73}
74
75/// Composite pipeline output for a batch.
76pub struct CompositeOutput {
77    /// Local resources referenced by prelude and replacement ops.
78    pub resources: Vec<RenderResource>,
79    /// External resources referenced by prelude and replacement ops.
80    pub external_resources: Vec<ExternalTextureDesc>,
81    /// Prelude ops emitted before scene ops.
82    pub prelude_ops: Vec<RenderGraphOp>,
83    /// Replacement ops for composite commands.
84    pub replacements: Vec<CompositeReplacement>,
85}
86
87impl CompositeOutput {
88    /// Creates an empty composite output.
89    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    /// Adds an external texture reference and returns its resource id.
99    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
106/// Trait for pipelines that expand composite commands into graph ops.
107pub trait CompositePipeline<C: CompositeCommand>: Send + Sync + 'static {
108    /// Expands composite commands into draw/compute ops.
109    fn compile(
110        &mut self,
111        context: &CompositeContext<'_>,
112        items: &[CompositeBatchItem<'_, C>],
113    ) -> CompositeOutput;
114}
115
116/// Type-erased composite pipeline used by the registry.
117pub(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/// Registry for managing and dispatching composite pipelines.
163#[derive(Default)]
164pub struct CompositePipelineRegistry {
165    pipelines: HashMap<TypeId, Box<dyn ErasedCompositePipeline>>,
166}
167
168impl CompositePipelineRegistry {
169    /// Creates a new empty composite pipeline registry.
170    pub fn new() -> Self {
171        Self::default()
172    }
173
174    /// Registers a new composite pipeline for a specific command type.
175    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}