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

#[macro_use]
extern crate tracing;

use std::fs::File;
use std::path::PathBuf;

use anyhow::Result;
use assignment_1b::image::Image;
use assignment_1b::ray::Ray;
use assignment_1b::scene::Scene;

use clap::Parser;

use rayon::prelude::{IntoParallelIterator, ParallelIterator};

/// 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>,

  /// 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,
}

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

  // Set up logging
  tracing_subscriber::fmt()
    .with_target(false)
    .with_timer(tracing_subscriber::fmt::time::uptime())
    .with_level(true)
    .init();

  // 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);

  // 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)| {
      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 intersections = scene
        .objects
        .iter()
        .enumerate()
        .filter_map(|(i, object)| {
          match object.kind.intersects_ray_at(&ray) {
            Ok(Some(t)) => {
              // Return both the t and the sphere, because we want to sort on
              // the t but later retrieve attributes from the sphere
              Some(Ok((i, t, object)))
            }
            Ok(None) => None,
            Err(err) => {
              error!("Error: {err}");
              Some(Err(err))
            }
          }
        })
        .collect::<Result<Vec<_>>>()?;

      // Sort the list of intersection times by the lowest one.
      let earliest_intersection =
        intersections.into_iter().min_by_key(|(_, t, _)| t.time);

      Ok(match earliest_intersection {
        // Take the object's material color
        Some((obj_idx, intersection_context, object)) => scene
          .compute_pixel_color(obj_idx, object.material, intersection_context),

        // There was no intersection, so this should default to the scene's
        // background color
        None => scene.bkg_color,
      })
    })
    .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(())
}