fidget_wgpu/
color.rs

1//! Data types for evaluating shape color
2//!
3//! In this module, only [`ShapeColor`] and [`ShapeColorBuffers`] should be used
4//! directly (through direct construction and [`ShapeColorBuffers::new`]
5//! respectively).  Other types are generic and are re-exported with a specific
6//! paramerization in an `effects` module; e.g. the generic
7//! [`ColorWorkspace<C>`] is specialized to
8//! [`voxel::effects::ColorWorkspace`](crate::voxel::effects::ColorWorkspace)
9use crate::{
10    CopyVarsChanged, CopyVarsError, Gpu,
11    buf::{BufferSizeError, ConfigBufferWrite, FlexBuffer, FlexConfigBuffer},
12    tag,
13    voxel::VarsBufferTag,
14};
15use fidget_bytecode::{Bytecode, ReservedRegister};
16use fidget_core::{
17    eval::Function,
18    shape::ShapeVars,
19    var::{Var, VarMap},
20    vm::VmShape,
21};
22
23use std::collections::HashMap;
24use zerocopy::IntoBytes;
25
26tag!(ShapeStartBufferTag, u32, usize, STORAGE | COPY_DST);
27
28/// Generic workspace for evaluating diffuse color
29pub struct ColorWorkspace<C> {
30    /// Config and serialized tapes, appended end to end
31    config: FlexConfigBuffer<C, u32>,
32
33    /// Start of the channel tapes for each shape
34    ///
35    /// The high bit indicates whether this is RGB (0) or HSL (1)
36    shape_start: FlexBuffer<ShapeStartBufferTag>,
37
38    /// Scratch space to upload variable values
39    vars_buf: FlexBuffer<VarsBufferTag>,
40
41    /// Lazily-constructed bind group for the config buffers
42    bind_group: std::cell::OnceCell<wgpu::BindGroup>,
43}
44
45impl<C> ColorWorkspace<C>
46where
47    C: zerocopy::IntoBytes + zerocopy::Immutable + Copy,
48{
49    pub(crate) fn new(device: &wgpu::Device) -> Self {
50        let config = FlexBuffer::new(device, "color config", 4.into()).unwrap();
51        let shape_start =
52            FlexBuffer::new(device, "shape start", 4usize).unwrap();
53        let vars_buf = FlexBuffer::new(device, "color vars", 4usize).unwrap();
54        Self {
55            config,
56            shape_start,
57            vars_buf,
58            bind_group: Default::default(),
59        }
60    }
61
62    pub(crate) fn config_bind_group(
63        &self,
64        device: &wgpu::Device,
65        layout: &wgpu::BindGroupLayout,
66    ) -> &wgpu::BindGroup {
67        self.bind_group.get_or_init(|| {
68            device.create_bind_group(&wgpu::BindGroupDescriptor {
69                label: Some("config bind group"),
70                layout,
71                entries: &[
72                    wgpu::BindGroupEntry {
73                        binding: 0,
74                        resource: self.config.bind_active(),
75                    },
76                    wgpu::BindGroupEntry {
77                        binding: 1,
78                        resource: self.shape_start.bind_active(),
79                    },
80                    wgpu::BindGroupEntry {
81                        binding: 2,
82                        resource: self.vars_buf.bind_active(),
83                    },
84                ],
85            })
86        })
87    }
88
89    pub(crate) fn copy_vars(
90        &mut self,
91        gpu: &Gpu,
92        vs: &VarMap,
93        vars: &ShapeVars<f32>,
94    ) -> Result<(), CopyVarsError> {
95        match crate::copy_vars(gpu, vs, vars, &mut self.vars_buf)? {
96            CopyVarsChanged::BufferChanged => {
97                self.bind_group = Default::default()
98            }
99            CopyVarsChanged::BufferUnchanged => (),
100        }
101        Ok(())
102    }
103
104    pub(crate) fn copy_config_and_tape(
105        &mut self,
106        gpu: &Gpu,
107        config: &C,
108        bytecode: &[u32],
109    ) -> Result<(), BufferSizeError> {
110        match self.config.write(gpu, config, bytecode)? {
111            ConfigBufferWrite::BufferChanged => {
112                self.bind_group = Default::default()
113            }
114            ConfigBufferWrite::BufferUnchanged => (),
115        }
116        Ok(())
117    }
118
119    pub(crate) fn copy_shape_starts(
120        &mut self,
121        gpu: &Gpu,
122        shape_starts: &[u32],
123    ) -> Result<(), BufferSizeError> {
124        let r = self
125            .shape_start
126            .grow_to_fit(&gpu.device, shape_starts.len())?;
127        if !matches!(r, std::cmp::Ordering::Equal) {
128            self.bind_group = Default::default();
129        }
130        let Ok(byte_count) =
131            (std::mem::size_of_val(shape_starts) as u64).try_into()
132        else {
133            return Ok(());
134        };
135        let mut writer = gpu
136            .queue
137            .write_buffer_with(self.shape_start.data(), 0, byte_count)
138            .unwrap();
139        writer.copy_from_slice(shape_starts.as_bytes());
140        Ok(())
141    }
142}
143
144/// Error type when constructing a [`ShapeColorBuffers`]
145#[derive(Debug, thiserror::Error)]
146pub enum ShapeColorError {
147    /// The shape uses a reserved register
148    #[error(transparent)]
149    RegisterError(#[from] ReservedRegister),
150}
151
152/// Color buffers for rendering a shape's diffuse color
153///
154/// These buffers live purely on the CPU, so they can be shared between threads.
155/// The object may be heavy (and therefore does not implement `Clone`); to share
156/// it, wrap it in an `Arc` or similar.
157pub struct ShapeColorBuffers {
158    /// Unified [`VarMap`] object
159    var_map: VarMap,
160
161    /// Number of shapes available
162    shape_count: usize,
163
164    /// Index of each shape's start position in the bytecode data
165    shape_start: Vec<u32>,
166
167    /// Maximum number of registers used by any tape
168    reg_count: u8,
169
170    /// Raw bytecode, which should be copied to the GPU config buffer tail
171    bytecode: Vec<u32>,
172}
173
174// Accessor methods
175impl ShapeColorBuffers {
176    pub(crate) fn var_map(&self) -> &VarMap {
177        &self.var_map
178    }
179    pub(crate) fn shape_count(&self) -> usize {
180        self.shape_count
181    }
182    pub(crate) fn shape_start(&self) -> &[u32] {
183        &self.shape_start
184    }
185    pub(crate) fn reg_count(&self) -> u8 {
186        self.reg_count
187    }
188    pub(crate) fn bytecode(&self) -> &[u32] {
189        &self.bytecode
190    }
191}
192
193/// Generic shape color
194pub enum ShapeColor<T> {
195    /// Red / green / blue channels, in the 0-1 range
196    Rgb {
197        /// Red component
198        r: T,
199        /// Green component
200        g: T,
201        /// Blue component
202        b: T,
203    },
204    /// Hue / saturation / lightness channels, in the 0-1 range
205    Hsl {
206        /// Hue
207        h: T,
208        /// Saturation
209        s: T,
210        /// Lightness
211        l: T,
212    },
213}
214
215impl<T> ShapeColor<T> {
216    fn channels(&self) -> [&T; 3] {
217        match self {
218            ShapeColor::Rgb { r, g, b } => [r, g, b],
219            ShapeColor::Hsl { h, s, l } => [h, s, l],
220        }
221    }
222}
223
224impl ShapeColorBuffers {
225    /// Builds a new set of buffers for evaluating shape color
226    pub fn new(
227        colors: &[ShapeColor<VmShape>],
228    ) -> Result<Self, ShapeColorError> {
229        // Build a single unified variable map, used across all tapes
230        let mut var_map = VarMap::new();
231        for c in colors {
232            for channel in c.channels() {
233                let vars = channel.inner().vars();
234                for (v, _index) in vars.iter() {
235                    var_map.insert(v);
236                }
237            }
238        }
239        let mut reg_count = 0;
240        let mut shape_start = Vec::with_capacity(colors.len());
241        let mut bytecode_data: Vec<u32> = Vec::new();
242        let mut local_var_map = HashMap::new();
243        for c in colors {
244            // Divide by 2 to convert from `u32` to `TapeWord`
245            let kind = match c {
246                ShapeColor::Rgb { .. } => 0,
247                ShapeColor::Hsl { .. } => 1 << 31,
248            };
249            let index = u32::try_from(bytecode_data.len() / 2).unwrap();
250            assert!(
251                index & (1 << 31) == 0,
252                "you have built more than 2 GiB of shape tapes?!"
253            );
254            shape_start.push(index | kind);
255            for channel in c.channels() {
256                // Build a local variable remapping array, reusing allocations
257                local_var_map.clear();
258                local_var_map.extend(channel.inner().vars().iter().map(
259                    |(v, i)| {
260                        (
261                            u32::try_from(i).unwrap(),
262                            u32::try_from(var_map.get(&v).unwrap()).unwrap(),
263                        )
264                    },
265                ));
266
267                // Generate bytecode for the root tape
268                let bytecode = Bytecode::new_with_input_map(
269                    channel.inner().data(),
270                    &local_var_map,
271                )?;
272                bytecode_data.extend(bytecode.data());
273                reg_count = reg_count.max(bytecode.reg_count());
274            }
275        }
276
277        Ok(Self {
278            shape_count: colors.len(),
279            shape_start,
280            bytecode: bytecode_data,
281            var_map,
282            reg_count,
283        })
284    }
285
286    /// Helper function to return XYZ variable indices
287    pub(crate) fn axes(&self) -> [u32; 3] {
288        [Var::X, Var::Y, Var::Z]
289            .map(|a| self.var_map.get(&a).map(|v| v as u32).unwrap_or(u32::MAX))
290    }
291}
fidget_wgpu/color.rs

fidget_wgpu/
color.rs
