avenger_wgpu/marks/
gradient.rs

1use crate::marks::multi::GRADIENT_TEXTURE_CODE;
2use avenger::marks::value::{ColorOrGradient, Gradient};
3use colorgrad::Color;
4use image::{DynamicImage, Rgba};
5use wgpu::Extent3d;
6
7const GRADIENT_WIDTH: u32 = 256;
8const GRADIENT_HEIGH: u32 = 32;
9pub const GRADIENT_LINEAR: f32 = 0.0;
10pub const GRADIENT_RADIAL: f32 = 1.0;
11pub const COLORWAY_LENGTH: u32 = 250;
12
13pub struct GradientAtlasBuilder {
14    extent: Extent3d,
15    next_image: image::RgbaImage,
16    images: Vec<DynamicImage>,
17    next_grad_row: usize,
18    initialized: bool,
19}
20
21impl Default for GradientAtlasBuilder {
22    fn default() -> Self {
23        Self::new()
24    }
25}
26
27impl GradientAtlasBuilder {
28    pub fn new() -> Self {
29        // Initialize with single pixel image
30        Self {
31            extent: Extent3d {
32                width: 1,
33                height: 1,
34                depth_or_array_layers: 1,
35            },
36            next_image: image::RgbaImage::new(1, 1),
37            images: vec![],
38            next_grad_row: 0,
39            initialized: false,
40        }
41    }
42
43    pub fn register_gradients(&mut self, gradients: &[Gradient]) -> (Option<usize>, Vec<f32>) {
44        if gradients.is_empty() {
45            return (None, Vec::new());
46        }
47
48        // Handle initialization
49        if !self.initialized {
50            // Initialze next_image that we can write to
51            self.next_image = image::RgbaImage::new(GRADIENT_WIDTH, GRADIENT_HEIGH);
52            self.extent = Extent3d {
53                width: GRADIENT_WIDTH,
54                height: GRADIENT_HEIGH,
55                depth_or_array_layers: 1,
56            };
57            self.initialized = true;
58        }
59
60        // Handle creation of new image when current image is full
61        if self.next_grad_row + gradients.len() > GRADIENT_HEIGH as usize {
62            let full_image = std::mem::take(&mut self.next_image);
63            self.next_image = image::RgbaImage::new(GRADIENT_WIDTH, GRADIENT_HEIGH);
64            self.images
65                .push(image::DynamicImage::ImageRgba8(full_image));
66            self.next_grad_row = 0;
67        }
68
69        // Write gradient values
70        for (pos, grad) in gradients.iter().enumerate() {
71            let row = (pos + self.next_grad_row) as u32;
72
73            // Build gradient colorway using colorgrad
74            let s = grad.stops();
75            let mut binding = colorgrad::CustomGradient::new();
76            let offsets = s.iter().map(|stop| stop.offset as f64).collect::<Vec<_>>();
77            let colors = s
78                .iter()
79                .map(|stop| {
80                    Color::new(
81                        stop.color[0] as f64,
82                        stop.color[1] as f64,
83                        stop.color[2] as f64,
84                        stop.color[3] as f64,
85                    )
86                })
87                .collect::<Vec<_>>();
88
89            let builder = binding.domain(offsets.as_slice()).colors(colors.as_slice());
90            let b = builder.build().unwrap();
91
92            // Store 250-bin colorway in pixels 6 through 255
93            let col_offset = GRADIENT_WIDTH - COLORWAY_LENGTH;
94            for i in 0..COLORWAY_LENGTH {
95                let p = (i as f64) / COLORWAY_LENGTH as f64;
96                let c = b.at(p).to_rgba8();
97                self.next_image
98                    .put_pixel(i + col_offset, row, Rgba::from(c));
99            }
100
101            // Encode the gradient control points in the first two or three pixels of the texture
102            match grad {
103                Gradient::LinearGradient(grad) => {
104                    // Write gradient type to column 0
105                    let control_color0 = Rgba::from([(GRADIENT_LINEAR * 255.0) as u8, 0, 0, 0]);
106                    self.next_image.put_pixel(0, row, control_color0);
107
108                    // Write x/y control points to column 1
109                    let control_color1 = Rgba::from([
110                        (grad.x0 * 255.0) as u8,
111                        (grad.y0 * 255.0) as u8,
112                        (grad.x1 * 255.0) as u8,
113                        (grad.y1 * 255.0) as u8,
114                    ]);
115                    self.next_image.put_pixel(1, row, control_color1);
116                }
117                Gradient::RadialGradient(grad) => {
118                    // Write gradient type to column 0
119                    let control_color0 = Rgba::from([(GRADIENT_RADIAL * 255.0) as u8, 0, 0, 0]);
120                    self.next_image.put_pixel(0, row, control_color0);
121
122                    // Write x/y control points to column 1
123                    let control_color1 = Rgba::from([
124                        (grad.x0 * 255.0) as u8,
125                        (grad.y0 * 255.0) as u8,
126                        (grad.x1 * 255.0) as u8,
127                        (grad.y1 * 255.0) as u8,
128                    ]);
129                    self.next_image.put_pixel(1, row, control_color1);
130
131                    // Write radius control points to column 2
132                    let control_color2 =
133                        Rgba::from([(grad.r0 * 255.0) as u8, (grad.r1 * 255.0) as u8, 0, 0]);
134                    self.next_image.put_pixel(2, row, control_color2);
135                }
136            }
137        }
138
139        // Compute texture coords of gradient rows.
140        // Add 0.1 of pixel offset so that we don't land on the edge
141        let coords = (self.next_grad_row..(self.next_grad_row + gradients.len()))
142            .map(|i| (i as f32 + 0.1) / GRADIENT_HEIGH as f32)
143            .collect::<Vec<_>>();
144
145        // Update next gradient row (Could be greater than GRADIENT_HEIGH, this will be
146        // handled on the next call to register_gradients.
147        self.next_grad_row += gradients.len();
148
149        // Compute gradient atlas index (index into the vector if images that will be
150        // returned by build).
151        let atlas_index = self.images.len();
152
153        (Some(atlas_index), coords)
154    }
155
156    pub fn build(&self) -> (Extent3d, Vec<DynamicImage>) {
157        let mut images = self.images.clone();
158        images.push(image::DynamicImage::ImageRgba8(self.next_image.clone()));
159        (self.extent, images)
160    }
161}
162
163pub fn to_color_or_gradient_coord(
164    color_or_gradient: &ColorOrGradient,
165    grad_coords: &[f32],
166) -> [f32; 4] {
167    match color_or_gradient {
168        ColorOrGradient::Color(c) => *c,
169        ColorOrGradient::GradientIndex(grad_idx) => [
170            GRADIENT_TEXTURE_CODE,
171            grad_coords[*grad_idx as usize],
172            0.0,
173            0.0,
174        ],
175    }
176}