lumen_engine/node/processing/
channel_shuffle.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!("channel_shuffle.wgsl");
9
10#[derive(
11    Debug, Clone, Copy, PartialEq, Eq, Default, lumen_macros::NodeEnum, lumen_macros::Delegate,
12)]
13#[repr(i64)]
14#[delegate(kind = "enum")]
15pub enum ChannelSelector {
16    #[default]
17    Red = 0,
18    Green = 1,
19    Blue = 2,
20    Alpha = 3,
21    Zero = 4,
22    One = 5,
23}
24
25impl ChannelSelector {
26    fn as_spec(self) -> &'static str {
27        match self {
28            Self::Red => "red",
29            Self::Green => "green",
30            Self::Blue => "blue",
31            Self::Alpha => "alpha",
32            Self::Zero => "zero",
33            Self::One => "one",
34        }
35    }
36}
37
38/// Remaps source raster color channels.
39#[derive(Debug, Clone, lumen_macros::Delegate)]
40pub struct ChannelShuffleParams {
41    /// Source channel mapped into the red output channel.
42    #[meta(kind = "enum", enum_type = ChannelSelector)]
43    pub red: ChannelSelector,
44    /// Source channel mapped into the green output channel.
45    #[meta(kind = "enum", enum_type = ChannelSelector)]
46    pub green: ChannelSelector,
47    /// Source channel mapped into the blue output channel.
48    #[meta(kind = "enum", enum_type = ChannelSelector)]
49    pub blue: ChannelSelector,
50    /// Source channel mapped into the alpha output channel.
51    #[meta(kind = "enum", enum_type = ChannelSelector)]
52    pub alpha: ChannelSelector,
53}
54
55impl Default for ChannelShuffleParams {
56    fn default() -> Self {
57        Self {
58            red: ChannelSelector::Red,
59            green: ChannelSelector::Green,
60            blue: ChannelSelector::Blue,
61            alpha: ChannelSelector::Alpha,
62        }
63    }
64}
65
66/// Remaps source raster color channels.
67#[derive(Debug, Clone, lumen_macros::Node)]
68#[node(
69    kind = "channel_shuffle",
70    name = "Channel Shuffle",
71    category = "processing"
72)]
73pub struct ChannelShuffle {
74    pub id: NodeId,
75    #[params]
76    pub params: ChannelShuffleParamsDelegate,
77
78    #[input()]
79    pub source: PortRef,
80}
81
82impl Default for ChannelShuffle {
83    fn default() -> Self {
84        Self {
85            id: NodeId::new(0),
86            params: ChannelShuffleParamsDelegate::default(),
87            source: PortRef::empty(),
88        }
89    }
90}
91
92#[derive(Debug, Clone)]
93struct CompiledChannelShuffle {
94    node_id: NodeId,
95    params: ChannelShuffleParamsDelegate,
96    buffer: lumen_gpu::BufferId,
97}
98
99impl GpuCompiledNode for CompiledChannelShuffle {
100    fn node_id(&self) -> NodeId {
101        self.node_id
102    }
103
104    fn bind(&self, ctx: &FrameBindContext<'_>, bound: &mut BoundFrame) -> crate::Result<()> {
105        let params = self.params.eval(&NodeParamEvalContext {
106            node_id: self.node_id,
107            expr: &ctx.expr_context(self.node_id, "params"),
108        })?;
109        let selectors = [
110            compiler::channel_selector(self.node_id, "red", params.red.as_spec())?,
111            compiler::channel_selector(self.node_id, "green", params.green.as_spec())?,
112            compiler::channel_selector(self.node_id, "blue", params.blue.as_spec())?,
113            compiler::channel_selector(self.node_id, "alpha", params.alpha.as_spec())?,
114        ];
115        let gpu_params = compiler::ChannelShuffleParams {
116            selector_indices: selectors.map(|selector| selector.index),
117            selector_values: selectors.map(|selector| selector.value),
118        };
119        bound.write_buffer(self.buffer, 0, bytemuck::bytes_of(&gpu_params));
120        Ok(())
121    }
122}
123
124impl GpuCompileNode for ChannelShuffle {
125    fn compile_gpu(
126        &self,
127        ctx: &mut crate::gpu::CompileContext<'_>,
128        port: &PortRef,
129    ) -> crate::Result<CompiledOutput> {
130        let (source, texture, params) = ctx.compile_unary_filter(
131            self.id,
132            &self.source,
133            port,
134            "channel-shuffle",
135            SHADER,
136            std::mem::size_of::<compiler::ChannelShuffleParams>() as u64,
137        )?;
138        ctx.register_compiled_node(CompiledChannelShuffle {
139            node_id: self.id,
140            params: self.params.clone(),
141            buffer: params,
142        });
143        Ok(CompiledOutput::Raster(RasterHandle {
144            texture,
145            domain: source.domain,
146            metadata: source.metadata,
147        }))
148    }
149}
lumen_engine/node/processing/channel_shuffle.rs

lumen_engine/node/processing/
channel_shuffle.rs
