csci5607/assignment-1c/src/image.rs

use std::{
  fs::File,
  io::{BufRead, BufReader, Read, Write},
  path::Path,
};

use anyhow::{Context, Result};
use generator::{done, Gn};
use itertools::Itertools;
use nalgebra::Vector3;

/// A pixel color represented by a red, green, and blue value in the range 0-1.
pub type Color = Vector3<f64>;

/// A representation of an image
#[derive(Derivative)]
#[derivative(Debug)]
pub struct Image {
  /// Width in pixels
  pub width: usize,

  /// Height in pixels
  pub height: usize,

  /// Pixel data in row-major form.
  #[derivative(Debug = "ignore")]
  pub data: Vec<Color>,
}

impl Image {
  pub fn from_file(path: impl AsRef<Path>) -> Result<Self> {
    let file = File::open(path.as_ref())?;
    Self::read(file)
  }

  /// Parse image from a Read
  pub fn read(r: impl Read + Send) -> Result<Self> {
    let mut line_reader = BufReader::new(r);

    let mut header = String::new();
    line_reader
      .read_line(&mut header)
      .context("Could not read line")?;
    let parts = header.trim().split(" ").collect::<Vec<_>>();

    let width = parts[1].parse::<usize>().context("Could not read width")?;
    let height = parts[2].parse::<usize>().context("Could not read height")?;
    let max_value = parts[3]
      .parse::<usize>()
      .context("Could not read max value")?;

    // Generator for reading numbers
    let numbers = Gn::<()>::new_scoped(move |mut s| {
      macro_rules! gen_try {
        ($expr:expr, $str:expr $(, $($arg:expr),* $(,)?)?) => {
          match $expr {
            Ok(v) => v,
            Err(e) => {
              s.yield_(
                Err(anyhow::Error::from(e))
                  .with_context(|| format!($str $(, $($arg,)*)?)),
              );
              done!();
            }
          }
        };
      }

      for line in line_reader.lines() {
        let line = gen_try!(line, "Could not read line");
        let parts = line.trim().split_whitespace();

        for part in parts {
          let int =
            gen_try!(part.parse::<u64>(), "Could not read int from: {}", part);
          s.yield_(Ok(int));
        }
      }

      done!()
    });

    let mut data = Vec::with_capacity(width * height);
    for mut chunk in &(numbers).chunks(3) {
      let (r, g, b) = match chunk.next_tuple() {
        Some(v) => v,
        None => bail!("Not enough elements"),
      };

      let r = r? as f64 / max_value as f64;
      let g = g? as f64 / max_value as f64;
      let b = b? as f64 / max_value as f64;

      let color = Color::new(r, g, b);
      data.push(color);
    }

    Ok(Image {
      width,
      height,
      data,
    })
  }

  /// Write the image in PPM format to a file.
  pub fn write(&self, mut w: impl Write) -> Result<()> {
    // Header
    let header = format!("P3 {} {} 255\n", self.width, self.height);
    w.write_all(header.as_bytes())?;

    // Pixel data
    assert_eq!(self.data.len(), self.width * self.height);

    for pixel in self.data.iter() {
      let pixel = pixel * 256.0;
      let red = pixel.x as u8;
      let green = pixel.y as u8;
      let blue = pixel.z as u8;
      let pixel = format!("{red} {green} {blue}\n");
      w.write_all(pixel.as_bytes())?;
    }

    Ok(())
  }
}