Skip to main content

lumen_engine/node/processing/
channel_shuffle.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!("channel_shuffle.wgsl");
9
10/// Remaps source raster color channels.
11#[derive(Debug, Clone, lumen_macros::Node)]
12#[node(
13    kind = "channel_shuffle",
14    name = "Channel Shuffle",
15    category = "processing"
16)]
17pub struct ChannelShuffle {
18    pub id: NodeId,
19    /// Source channel mapped into the red output channel.
20    #[property(kind = "string", format = "channel_selector")]
21    pub red: NodeProperty,
22    /// Source channel mapped into the green output channel.
23    #[property(kind = "string", format = "channel_selector")]
24    pub green: NodeProperty,
25    /// Source channel mapped into the blue output channel.
26    #[property(kind = "string", format = "channel_selector")]
27    pub blue: NodeProperty,
28    /// Source channel mapped into the alpha output channel.
29    #[property(kind = "string", format = "channel_selector")]
30    pub alpha: NodeProperty,
31    #[input()]
32    pub source: PortRef,
33}
34
35impl Default for ChannelShuffle {
36    fn default() -> Self {
37        Self {
38            id: NodeId::new(0),
39            red: NodeProperty::String("red".to_string()),
40            green: NodeProperty::String("green".to_string()),
41            blue: NodeProperty::String("blue".to_string()),
42            alpha: NodeProperty::String("alpha".to_string()),
43            source: PortRef::empty(),
44        }
45    }
46}
47
48impl GpuCompileNode for ChannelShuffle {
49    fn compile_gpu(
50        &self,
51        ctx: &mut crate::gpu::CompileContext<'_>,
52        port: &PortRef,
53    ) -> crate::Result<CompiledOutput> {
54        let (source, texture, params) = ctx.compile_unary_filter(
55            self.id,
56            &self.source,
57            port,
58            "channel-shuffle",
59            SHADER,
60            std::mem::size_of::<compiler::ChannelShuffleParams>() as u64,
61        )?;
62        ctx.push_frame_binding(FrameBinding::ChannelShuffle {
63            node_id: self.id,
64            red: self.red.clone(),
65            green: self.green.clone(),
66            blue: self.blue.clone(),
67            alpha: self.alpha.clone(),
68            buffer: params,
69        });
70        Ok(CompiledOutput::Raster(RasterHandle {
71            texture,
72            domain: source.domain,
73            metadata: source.metadata,
74        }))
75    }
76}
77
78impl GpuFrameBindNode for ChannelShuffle {
79    fn bind_gpu_frame(
80        &self,
81        ctx: &FrameBindContext<'_>,
82        binding: &FrameBinding,
83        bound: &mut BoundFrame,
84    ) -> crate::Result<()> {
85        let FrameBinding::ChannelShuffle {
86            node_id,
87            red,
88            green,
89            blue,
90            alpha,
91            buffer,
92        } = binding
93        else {
94            return Ok(());
95        };
96        let selectors = [
97            compiler::channel_selector(
98                *node_id,
99                "red",
100                &red.resolve_string(*node_id, "red", &ctx.expr_context(*node_id, "red"))?,
101            )?,
102            compiler::channel_selector(
103                *node_id,
104                "green",
105                &green.resolve_string(*node_id, "green", &ctx.expr_context(*node_id, "green"))?,
106            )?,
107            compiler::channel_selector(
108                *node_id,
109                "blue",
110                &blue.resolve_string(*node_id, "blue", &ctx.expr_context(*node_id, "blue"))?,
111            )?,
112            compiler::channel_selector(
113                *node_id,
114                "alpha",
115                &alpha.resolve_string(*node_id, "alpha", &ctx.expr_context(*node_id, "alpha"))?,
116            )?,
117        ];
118        let params = compiler::ChannelShuffleParams {
119            selector_indices: selectors.map(|selector| selector.index),
120            selector_values: selectors.map(|selector| selector.value),
121        };
122        bound.write_buffer(*buffer, 0, bytemuck::bytes_of(&params));
123        Ok(())
124    }
125}