Skip to main content

lumen_engine/node/compositing/
merge.rs

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