potlood/

lib.rs

#![feature(array_windows)]
use std::{fs::File, io::Write};

pub enum Color {
    Black,
    Grey,
    White,
    Red,
    Green,
    Blue,
    Yellow,
    Cyan,
    Magenta,
    Rgb(u8, u8, u8),
    Hex(u32),
    RgbFloat(f64, f64, f64),
}

#[derive(Clone, Copy)]
pub struct Pixel {
    r: u8,
    g: u8,
    b: u8,
}

impl Pixel {
    pub fn new(r: u8, g: u8, b: u8) -> Self {
        Self { r, g, b }
    }

    pub fn color(c: Color) -> Self {
        match c {
            Color::Black => Pixel::new(0, 0, 0),
            Color::Grey => Pixel::new(0xff / 2, 0xff / 2, 0xff / 2),
            Color::White => Pixel::new(0xff, 0xff, 0xff),
            Color::Red => Pixel::new(0xff, 0, 0),
            Color::Green => Pixel::new(0, 0xff, 0),
            Color::Blue => Pixel::new(0, 0, 0xff),
            Color::Yellow => Pixel::new(0xff, 0xff, 0),
            Color::Cyan => Pixel::new(0, 0xff, 0xff),
            Color::Magenta => Pixel::new(0xff, 0, 0xff),
            Color::Rgb(r, g, b) => Pixel::new(r, g, b),
            Color::RgbFloat(r, g, b) => {
                let convert = |v| (v * 255.0) as u8;
                Pixel::new(convert(r), convert(g), convert(b))
            }
            Color::Hex(h) => {
                let bytes = h.to_le_bytes();
                Pixel::new(bytes[2], bytes[1], bytes[0])
            }
        }
    }

    pub fn empty() -> Self {
        Self { r: 0, g: 0, b: 0 }
    }

    pub fn ppm(&self) -> String {
        format!("{} {} {}", self.r, self.g, self.b)
    }
}

#[derive(Debug)]
pub enum PaperError {
    OutOfBounds,
}

pub struct Paper {
    width: usize,
    height: usize,
    pixels: Vec<Pixel>,
}

impl Paper {
    pub fn new(width: usize, height: usize) -> Self {
        Self {
            width,
            height,
            pixels: vec![Pixel::empty(); width * height],
        }
    }

    fn ppm(&self) -> String {
        let header = format!("P3\n{} {}\n255", self.width, self.height);
        let pixels = self
            .pixels
            .iter()
            .map(|pixel| pixel.ppm())
            .collect::<Vec<String>>()
            .join("\n");
        return [header, pixels, String::new()].join("\n");
    }

    pub fn generate(&self, path: String) -> std::io::Result<()> {
        let mut file = File::create(path)?;
        file.write_all(self.ppm().as_bytes())?;
        Ok(())
    }

    pub fn set(&mut self, x: usize, y: usize, pixel: Pixel) {
        self.pixels[y * self.width + x] = pixel
    }

    pub fn set_checked(&mut self, x: usize, y: usize, pixel: Pixel) -> Result<(), PaperError> {
        if !(x < self.width && y < self.height) {
            return Err(PaperError::OutOfBounds);
        }

        self.set(x, y, pixel);
        Ok(())
    }

    pub fn set_horizontal(&mut self, y: usize, pixel: Pixel) {
        let index = y * self.width;
        let mut row = vec![pixel; self.width];
        self.pixels[index..index + self.width].swap_with_slice(&mut row)
    }

    pub fn set_vertical(&mut self, x: usize, pixel: Pixel) {
        self.pixels
            .chunks_mut(self.width)
            .for_each(|row| row[x] = pixel)
    }

    pub fn fill(&mut self, pixel: Pixel) {
        self.pixels.iter_mut().for_each(|p| *p = pixel)
    }

    pub fn rect(
        &mut self,
        x: usize,
        y: usize,
        width: usize,
        height: usize,
        pixel: Pixel,
    ) -> Result<(), PaperError> {
        if !(x + width < self.width && y + height < self.height) {
            return Err(PaperError::OutOfBounds);
        }

        for row in y..y + height {
            for col in x..x + width {
                self.set(col, row, pixel)
            }
        }

        Ok(())
    }

    pub fn circle(
        &mut self,
        x: usize,
        y: usize,
        radius: usize,
        pixel: Pixel,
    ) -> Result<(), PaperError> {
        if !(radius < x && radius < y && x + radius < self.width && y + radius < self.height) {
            return Err(PaperError::OutOfBounds);
        }

        let odd = 1 - radius % 2;

        for row in y - radius..y + radius {
            let v = y as isize - row as isize;
            for col in x - radius..x + radius {
                // From this square of pixels, we only set those which we calculate to be _inside_
                // the circle.
                //
                //        .(row, col)
                //      d/|
                //      /_|v
                //     M h
                //
                // With
                //     v = y - row
                // and
                //     h = x - col
                // we can see
                //     d^2 = h^2 + v^2.
                // In case the point lies in the circle,
                //     d < radius,
                // which is identical to
                //     d^2 < radius^2.
                let h = x as isize - col as isize;
                let d_squared = h * h + v * v;
                // TODO: Check out the effect of <= vs < and change code and above comment
                // accordingly.
                if (d_squared as usize) <= radius * radius + (1 * odd) {
                    self.set(col, row, pixel)
                }
            }
        }

        Ok(())
    }

    pub fn line(&mut self, x0: usize, y0: usize, x1: usize, y1: usize, pixel: Pixel) {
        let dx = x1.abs_diff(x0) as isize;
        let sx = if x0 < x1 { 1 } else { -1 };
        let dy = -(y1.abs_diff(y0) as isize);
        let sy = if y0 < y1 { 1 } else { -1 };
        let mut error = dx + dy;
        let mut x0 = x0 as isize;
        let mut y0 = y0 as isize;
        loop {
            if self.set_checked(x0 as usize, y0 as usize, pixel).is_err() {
                break;
            };
            if x0 == x1 as isize && y0 == y1 as isize {
                break;
            }
            if 2 * error >= dy {
                if x0 == x1 as isize {
                    break;
                }
                error = error + dy;
                x0 = x0 + sx;
            }
            if 2 * error <= dx {
                if y0 == y1 as isize {
                    break;
                }
                error = error + dx;
                y0 = y0 + sy;
            }
        }
    }

    pub fn stroke(&mut self, curve: Vec<(usize, usize)>, pixel: Pixel) {
        for &[(x0, y0), (x1, y1)] in curve.array_windows() {
            self.line(x0, y0, x1, y1, pixel)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn pixel_ppm() {
        let pixel = Pixel::new(20, 30, 40);
        let ppm = pixel.ppm();
        assert_eq!(ppm, "20 30 40");
    }

    #[test]
    fn ppm() {
        let mut paper = Paper::new(2, 3);
        paper.line(0, 0, 1, 2, Pixel::color(Color::Hex(0xffeeaa)));
        let ppm = "P3
2 3
255
255 238 170
0 0 0
0 0 0
255 238 170
0 0 0
255 238 170
";
        assert_eq!(paper.ppm(), ppm);
    }
}