Skip to main content

lumen_engine/node/compositing/
merge.rs

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