Skip to main content

lumen_engine/node/compositing/
raster_multimerge.rs

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