csci5607/assignment-1d/src/main.rs

#[macro_use]
extern crate anyhow;
#[macro_use]
extern crate tracing;

use std::path::PathBuf;
use std::{fs::File, str::FromStr};

use anyhow::Result;
use assignment_1d::{image::Image, ray::Ray, scene::Scene};

use clap::{ArgAction, Parser};

use rayon::prelude::{IntoParallelIterator, ParallelIterator};
use tracing::metadata::LevelFilter;
use tracing_appender::non_blocking::WorkerGuard;
use tracing_subscriber::{
  fmt::Layer, prelude::__tracing_subscriber_SubscriberExt,
  util::SubscriberInitExt,
};

/// Simple raytracer with Blinn-Phong illumination and shadowing.
#[derive(Parser)]
#[clap(author, version, about, long_about = None)]
struct Opt {
  /// Path to the input file to use.
  #[clap()]
  input_path: PathBuf,

  /// Path to the output (defaults to the same file name as the input except
  /// with an extension of .ppm)
  #[clap(short = 'o', long = "output")]
  output_path: Option<PathBuf>,

  /// Log output in json
  #[clap(long = "json")]
  use_json: bool,

  /// Which file to send logs to (stderr by default)
  #[clap(long = "log-output")]
  log_output: Option<PathBuf>,

  /// Force parallel projection to be used
  #[clap(long = "parallel")]
  force_parallel: bool,

  /// Override distance from eye
  #[clap(long = "distance", default_value = "1.0")]
  distance: f64,

  /// Verbosity
  #[clap(short, long, action = ArgAction::Count)]
  verbosity: u8,

  /// Evaluate at a single pixel
  #[clap(short, long = "render-pixel")]
  render_pixel: Option<RenderPixel>,
}

fn main() -> Result<()> {
  let opt = Opt::parse();

  let _guard = setup_logging(&opt);

  // Rename the output file if it's not provided
  let out_file = opt
    .output_path
    .unwrap_or_else(|| opt.input_path.with_extension("ppm"));

  let mut scene = Scene::from_input_file(&opt.input_path)?;
  let distance = opt.distance;

  // Force-override parallel projection
  if opt.force_parallel {
    scene.parallel_projection = true;
  }

  // Translate image pixels to real-world 3d coords
  let translate_pixel = scene.pixel_translation_function(distance);

  let evaluate_at_pixel = |px, py| {
    let span = trace_span!("main_loop", px = px, py = py);
    let _enter = span.enter();

    let pixel_in_space = translate_pixel(px, py);

    let ray_start = if scene.parallel_projection {
      // For a parallel projection, we'll just take the view direction and
      // subtract it from the target point. This means every single
      // ray will be viewed from a point at infinity, rather than a single eye
      // position.
      let n = scene.view_dir.normalize();
      let view_dir = n * distance;
      pixel_in_space - view_dir
    } else {
      scene.eye_pos
    };

    let ray = Ray::from_endpoints(ray_start, pixel_in_space);

    // let res= rayon::spawn(|| scene.trace_single_ray(ray, 0));
    scene.trace_single_ray(scene.eye_pos, ray, 0)
  };

  // For debugging purposes!
  if let Some(RenderPixel(px, py)) = opt.render_pixel {
    let pixel_color = evaluate_at_pixel(px, py)?;
    println!("Pixel color: {pixel_color}");
    return Ok(());
  }

  // Generate a parallel iterator for pixels
  // The iterator preserves order and uses row-major order
  let pixels_iter = (0..scene.image_height)
    .into_par_iter()
    .flat_map(|y| (0..scene.image_width).into_par_iter().map(move |x| (x, y)));

  // Loop through every single pixel of the output file
  let pixels = pixels_iter
    .map(|(px, py)| evaluate_at_pixel(px, py))
    .collect::<Result<Vec<_>>>()?;

  // Construct and emit image
  let image = Image {
    width: scene.image_width,
    height: scene.image_height,
    data: pixels,
  };

  {
    let file = File::create(out_file)?;
    image.write(file)?;
  }

  Ok(())
}

#[derive(Clone)]
struct RenderPixel(usize, usize);

impl FromStr for RenderPixel {
  type Err = anyhow::Error;

  fn from_str(s: &str) -> Result<Self, Self::Err> {
    let parts = s.split(",").collect::<Vec<_>>();
    ensure!(parts.len() == 2, "must be a pair");

    let x = parts[0].parse::<usize>()?;
    let y = parts[1].parse::<usize>()?;

    Ok(RenderPixel(x, y))
  }
}

/// A little bit of engineering to make it easy to write conditional builders
/// for logging setup because the tracing-subscriber crate for some reason
/// decided it would be a good idea to have all of its builders be polymorphic?
macro_rules! logsetup_if {
  ($ident:ident , $cond:expr , $iftrue:expr , $iffalse:expr => { $($body:tt)* }) => {
    if ($cond) {
      let $ident = $iftrue;
      $($body)*
    } else {
      let $ident = $iffalse;
      $($body)*
    }
  };
}

fn setup_logging(opt: &Opt) -> Option<WorkerGuard> {
  let mut result = None;

  let level_filter = match opt.verbosity {
    0 => LevelFilter::ERROR,
    1 => LevelFilter::WARN,
    2 => LevelFilter::INFO,
    3 => LevelFilter::DEBUG,
    _ => LevelFilter::TRACE,
  };

  let layer = Layer::default();

  logsetup_if! (layer, opt.use_json, layer.json(), layer => {
    let layer = layer
      .with_target(false)
      .with_timer(tracing_subscriber::fmt::time::uptime())
      .with_level(true);

    logsetup_if! (layer, opt.log_output.is_some(), {
      let log_output = opt.log_output.clone().unwrap();
      let file_appender = tracing_appender::rolling::never(".", log_output);
      let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
      result = Some(guard);
      layer.with_writer(non_blocking)
    }, layer => {
      tracing_subscriber::registry()
        .with(layer)
        .with(level_filter)
        .init();
    });
  });

  result
}