Skip to main content

lumen_engine/node/compositing/
raster_multimerge.rs

1use crate::node::{NodeId, NodeProperty, PortRef};
2
3use crate::gpu::{
4    BoundFrame, CompiledOutput, FrameBindContext, FrameBinding, GpuCompileNode, GpuFrameBindNode,
5    RasterHandle, compiler,
6};
7
8pub(crate) const SHADER: &str = include_str!("raster_multimerge.wgsl");
9
10/// Composites a variadic stack of raster layers in order.
11#[derive(Debug, Clone, lumen_macros::Node)]
12#[node(
13    kind = "raster_multimerge",
14    name = "Raster Multi Merge",
15    category = "compositing"
16)]
17pub struct RasterMultiMerge {
18    pub id: NodeId,
19    /// Opacity applied to each layer as it is composited.
20    #[property(kind = "float", min = 0, max = 1, step = 0.05)]
21    pub opacity: NodeProperty,
22    /// Blend mode used for each layer in the stack.
23    #[property(kind = "enum", enum_type = crate::node::compositing::BlendMode)]
24    pub blend_mode: NodeProperty,
25    #[input(optional, variadic)]
26    pub layers: Vec<PortRef>,
27}
28
29impl Default for RasterMultiMerge {
30    fn default() -> Self {
31        Self {
32            id: NodeId::new(0),
33            opacity: NodeProperty::Float(1.0),
34            blend_mode: NodeProperty::Int(0),
35            layers: Vec::new(),
36        }
37    }
38}
39
40impl GpuCompileNode for RasterMultiMerge {
41    fn compile_gpu(
42        &self,
43        ctx: &mut crate::gpu::CompileContext<'_>,
44        port: &PortRef,
45    ) -> crate::Result<CompiledOutput> {
46        if port.port != "output" {
47            return Err(ctx.missing_output(self.id, &port.port));
48        }
49
50        let mut layers = self
51            .layers
52            .iter()
53            .filter(|layer| !layer.is_empty())
54            .map(|layer| {
55                ctx.compile_port(layer)
56                    .and_then(|output| output.into_raster(layer.id, &layer.port))
57            })
58            .collect::<crate::Result<Vec<_>>>()?
59            .into_iter();
60        let Some(first) = layers.next() else {
61            return Ok(ctx.compile_transparent(self.id));
62        };
63
64        let size = first.domain.storage_size;
65        let params = ctx.builder_mut().buffer_for(
66            lumen_gpu::NodeKey(self.id.0),
67            Some(format!("raster-multimerge:{}:params", self.id.0)),
68            lumen_gpu::BufferDesc::uniform(
69                std::mem::size_of::<compiler::RasterMultiMergeParams>() as u64
70            ),
71        );
72        let program = ctx.builder_mut().program_for(
73            lumen_gpu::NodeKey(self.id.0),
74            lumen_gpu::ProgramDesc::Compute(lumen_gpu::ComputeProgramDesc {
75                label: Some("raster-multimerge".to_string()),
76                shader: SHADER.to_string(),
77                entry: "cs_main".to_string(),
78                bind_groups: lumen_gpu::BindGroupLayoutSpec::single(vec![
79                    lumen_gpu::BindingLayoutEntry::texture(
80                        0,
81                        lumen_gpu::wgpu::ShaderStages::COMPUTE,
82                    ),
83                    lumen_gpu::BindingLayoutEntry::texture(
84                        1,
85                        lumen_gpu::wgpu::ShaderStages::COMPUTE,
86                    ),
87                    lumen_gpu::BindingLayoutEntry::uniform(
88                        2,
89                        lumen_gpu::wgpu::ShaderStages::COMPUTE,
90                    ),
91                    lumen_gpu::BindingLayoutEntry::storage_texture(
92                        3,
93                        lumen_gpu::wgpu::ShaderStages::COMPUTE,
94                        lumen_gpu::wgpu::TextureFormat::Rgba8Unorm,
95                        lumen_gpu::wgpu::StorageTextureAccess::WriteOnly,
96                    ),
97                ]),
98            }),
99        );
100
101        let mut current = first;
102        for (index, overlay) in layers.enumerate() {
103            let texture = ctx.builder_mut().texture_for(
104                lumen_gpu::NodeKey(self.id.0),
105                Some(format!("raster-multimerge:{}:layer-{index}", self.id.0)),
106                lumen_gpu::TextureDesc::storage(size, lumen_gpu::wgpu::TextureFormat::Rgba8Unorm),
107            );
108            ctx.builder_mut().compute_pass(lumen_gpu::ComputePassDesc {
109                label: Some(format!("raster-multimerge:{}:layer-{index}", self.id.0)),
110                owner: Some(lumen_gpu::NodeKey(self.id.0)),
111                program,
112                bindings: vec![
113                    lumen_gpu::Binding::sampled_texture(0, 0, current.texture),
114                    lumen_gpu::Binding::sampled_texture(0, 1, overlay.texture),
115                    lumen_gpu::Binding::uniform(0, 2, params),
116                    lumen_gpu::Binding::storage_texture(0, 3, texture),
117                ],
118                dispatch: compiler::dispatch_for(size).into(),
119            });
120            current = RasterHandle {
121                texture,
122                domain: lumen_gpu::TextureDomain::full_frame(size),
123                metadata: current.metadata,
124            };
125        }
126        ctx.builder_mut().param(
127            lumen_gpu::ParamKey {
128                owner: lumen_gpu::NodeKey(self.id.0),
129                slot: 0,
130            },
131            lumen_gpu::ParamTarget::Buffer(params),
132        );
133        ctx.push_frame_binding(FrameBinding::RasterMultiMerge {
134            node_id: self.id,
135            opacity: self.opacity.clone(),
136            blend_mode: self.blend_mode.clone(),
137            buffer: params,
138        });
139
140        Ok(CompiledOutput::Raster(current))
141    }
142}
143
144impl GpuFrameBindNode for RasterMultiMerge {
145    fn bind_gpu_frame(
146        &self,
147        ctx: &FrameBindContext<'_>,
148        binding: &FrameBinding,
149        bound: &mut BoundFrame,
150    ) -> crate::Result<()> {
151        let FrameBinding::RasterMultiMerge {
152            node_id,
153            opacity,
154            blend_mode,
155            buffer,
156        } = binding
157        else {
158            return Ok(());
159        };
160        let params = compiler::RasterMultiMergeParams {
161            values: [
162                opacity.resolve_float(
163                    *node_id,
164                    "opacity",
165                    &ctx.expr_context(*node_id, "opacity"),
166                )? as f32,
167                blend_mode.resolve_int(
168                    *node_id,
169                    "blend_mode",
170                    &ctx.expr_context(*node_id, "blend_mode"),
171                )? as f32,
172                0.0,
173                0.0,
174            ],
175        };
176        bound.write_buffer(*buffer, 0, bytemuck::bytes_of(&params));
177        Ok(())
178    }
179}