shadybug/
basic.rs

1/// Basic types useful with the Shader trait
2use std::ops::{Index, IndexMut};
3
4use glam::{UVec2, Vec3, Vec4, Vec4Swizzles};
5
6#[cfg(feature = "png")]
7use std::{fs::File, io::BufWriter, path::Path};
8
9#[derive(Clone, Debug)]
10pub struct Target<Pixel> {
11    pub size: UVec2,
12    pub pixels: Vec<Pixel>,
13}
14
15impl<Pixel> Target<Pixel> {
16    /// Create a new Target with the given `size` and fill it with `pixel`.
17    pub fn new(size: UVec2, pixel: Pixel) -> Self
18    where
19        Pixel: Clone,
20    {
21        Self {
22            size,
23            pixels: vec![pixel; size.element_product() as usize],
24        }
25    }
26}
27
28/// Convert from linear RGB to sRGB
29/// from https://gamedev.stackexchange.com/a/194038
30pub fn linear_to_srgb(color: Vec4) -> Vec4 {
31    let rgb = color.xyz();
32    let cutoff = rgb.cmplt(Vec3::splat(0.0031308));
33    let higher = 1.055 * rgb.powf(1. / 2.4) - 0.055;
34    let lower = rgb * 12.92;
35    Vec3::select(cutoff, lower, higher).extend(color.w)
36}
37
38#[cfg(feature = "png")]
39impl Target<Vec4> {
40    /// Write target data as a PNG
41    pub fn write_png<P: AsRef<Path>>(self, path: P) -> bool {
42        // from png crate documentation
43        let Ok(file) = File::create(path) else {
44            return false;
45        };
46        let ref mut w = BufWriter::new(file);
47        let mut encoder = png::Encoder::new(w, self.size.x, self.size.y);
48        encoder.set_color(png::ColorType::Rgba);
49        encoder.set_depth(png::BitDepth::Eight);
50        encoder.set_source_gamma(png::ScaledFloat::new(1. / 2.2));
51        let source_chromaticities = png::SourceChromaticities::new(
52            (0.31270, 0.32900),
53            (0.64000, 0.33000),
54            (0.30000, 0.60000),
55            (0.15000, 0.06000),
56        );
57        encoder.set_source_chromaticities(source_chromaticities);
58        let Ok(mut writer) = encoder.write_header() else {
59            return false;
60        };
61
62        // convert float pixels in linear RGBA to 8-bit sRGBA
63        let data: Vec<u8> = self
64            .pixels
65            .into_iter()
66            .flat_map(|color| {
67                linear_to_srgb(color)
68                    .to_array()
69                    .map(|x| (x * 256.).round().clamp(0., 255.) as u8)
70            })
71            .collect();
72
73        writer.write_image_data(&data).is_ok()
74    }
75}
76
77impl<Pixel> Index<UVec2> for Target<Pixel> {
78    type Output = Pixel;
79
80    fn index(&self, index: UVec2) -> &Self::Output {
81        &self.pixels[(index.y * self.size.x + index.x) as usize]
82    }
83}
84
85impl<Pixel> IndexMut<UVec2> for Target<Pixel> {
86    fn index_mut(&mut self, index: UVec2) -> &mut Self::Output {
87        &mut self.pixels[(index.y * self.size.x + index.x) as usize]
88    }
89}
90
91#[derive(Default, Clone, Copy, Debug)]
92pub struct Sample {
93    pub depth: f32,
94    pub color: Vec4,
95}