<commit_before>// Based on smallpt, http://www.kevinbeason.com/smallpt/ which is also licensed under
// the MIT license.
use std::fs::File;
use std::io::BufWriter;
use std::io;
use std::io::prelude::*;
use std::sync::Arc;
use std::sync::mpsc::channel;

extern crate rand;
use rand::{Rng, XorShiftRng, SeedableRng};

extern crate threadpool;
use threadpool::ThreadPool;

extern crate num_cpus;

extern crate argparse;
use argparse::{ArgumentParser, Store};

mod path_tracer;
use path_tracer::*;

enum Material {
    Diffuse,
    Specular,
    Refractive
}

struct Sphere {
    material: Material,
    radius: f64,
    position: Vec3d,
    emission: Vec3d,
    colour: Vec3d
}

impl Sphere {
    fn new(material: Material, radius: f64, position: Vec3d, emission: Vec3d, colour: Vec3d) -> Sphere {
        Sphere {
          material: material,
          radius: radius,
          position: position,
          emission: emission,
          colour: colour
        }
    }
    fn intersect(&self, ray: &Ray) -> Option<f64> {
        let op = self.position - ray.origin;
        let b = op.dot(ray.direction);
        let determinant = b * b - op.dot(op) + self.radius * self.radius;
        if determinant < 0.0 { return None; }
        let determinant = determinant.sqrt();
        let t1 = b - determinant;
        let t2 = b + determinant;
        const EPSILON : f64 = 0.0001;
        if t1 > EPSILON {
            Some(t1)
        } else if t2 > EPSILON {
            Some(t2)
        } else {
            None
        }
    }
}

#[test]
fn intersection() {
    let sphere = Sphere::new(
        Material::Diffuse,
        100.0,
        Vec3d::new(0.0, 0.0, 200.0),
        Vec3d::zero(),
        Vec3d::zero());
    let ray = Ray::new(Vec3d::zero(), Vec3d::new(0.0, 0.0, 1.0));
    match sphere.intersect(&ray) {
        Some(x) => assert_eq!(x, 100.0),
        None => panic!("unexpected")
    }
    let ray = Ray::new(Vec3d::zero(), Vec3d::new(0.0, 1.0, 0.0));
    match sphere.intersect(&ray) {
        Some(_) => panic!("unexpected"),
        None => {}
    }
}

struct HitRecord<'a> {
    sphere: &'a Sphere,
    dist: f64
}

fn intersect<'a>(scene: &'a [Sphere], ray: &Ray) -> Option<HitRecord<'a>> {
    let mut result : Option<HitRecord<'a>> = None;
    for sph in scene {
        if let Some(dist) = sph.intersect(&ray) {
            if match result { None => true, Some(ref x) => dist < x.dist } {
                result = Some(HitRecord { sphere: &sph, dist: dist });
            }
        }
    }
    result
}

fn radiance<R: Rng>(scene: &[Sphere], ray: &Ray, depth: i32, rng: &mut R) -> Vec3d {
    if let Some(hit) = intersect(&scene, &ray) {
        let hit_pos = ray.origin + ray.direction * hit.dist;
        let hit_normal = (hit_pos - hit.sphere.position).normalized();
        let n1 = if hit_normal.dot(ray.direction) < 0.0 { hit_normal } else { hit_normal.neg() };
        let mut colour = hit.sphere.colour;
        let max_reflectance = colour.max_component();
        let depth = depth + 1;
        if depth > 5 {
            let rand = rng.gen::<f64>();
            if rand < max_reflectance && depth < 500 { // Rust's stack blows up ~600 on my machine
               colour = colour * (1.0 / max_reflectance);
            } else {
                return hit.sphere.emission;
            }
        }
        match hit.sphere.material {
            Material::Diffuse => {
                // Get a random polar coordinate
                let r1 = rng.gen::<f64>() * 2.0 * std::f64::consts::PI;
                let r2 = rng.gen::<f64>();
                let r2s = r2.sqrt();
                // Create a coordinate system u,v,w local to the point, where the w is the normal
                // pointing out of the sphere and the u and v are orthonormal to w.
                let w = n1;
                // Pick an arbitrary non-zero preferred axis for u
                let u = if n1.x.abs() > 0.1 { Vec3d::new(0.0, 1.0, 0.0) } else { Vec3d::new(1.0, 0.0, 0.0) }.cross(w);
                let v = w.cross(u);
                // construct the new direction
                let new_dir = u * r1.cos() * r2s + v * r1.sin() * r2s + w * (1.0 - r2).sqrt();
                colour = colour * radiance(scene, &Ray::new(hit_pos, new_dir.normalized()), depth, rng);
            },
            Material::Specular => {
                let reflection = ray.direction - hit_normal * 2.0 * hit_normal.dot(ray.direction);
                let reflected_ray = Ray::new(hit_pos, reflection);
                colour = colour * radiance(scene, &reflected_ray, depth, rng);
            },
            Material::Refractive => {
                let reflection = ray.direction - hit_normal * 2.0 * hit_normal.dot(ray.direction);
                let reflected_ray = Ray::new(hit_pos, reflection);
                let into = if hit_normal.dot(n1) > 0.0 { true } else { false };
                let nc = 1.0;
                let nt = 1.5;
                let nnt = if into { nc/nt } else { nt/nc };
                let ddn = ray.direction.dot(n1);
                let cos2t = 1.0 - nnt * nnt * (1.0 - ddn * ddn);
                if cos2t < 0.0 {
                    // Total internal reflection
                    colour = colour * radiance(scene, &reflected_ray, depth, rng);
                } else {
                    let tbd = ddn * nnt + cos2t.sqrt();
                    let tbd = if into { tbd } else { -tbd };
                    let tdir = (ray.direction * nnt - hit_normal * tbd).normalized();
                    let transmitted_ray = Ray::new(hit_pos, tdir);
                    let a = nt - nc;
                    let b = nt + nc;
                    let r0 = (a * a) / (b * b);
                    let c = 1.0 - if into { -ddn } else { tdir.dot(hit_normal) };
                    let re = r0 + (1.0 - r0) * c * c * c * c * c;
                    let tr = 1.0 - re;
                    let p = 0.25 + 0.5 * re;
                    let rp = re / p;
                    let tp = tr / (1.0 - p);
                    colour = colour * if depth > 2 {
                        if rng.gen::<f64>() < p {
                            radiance(scene, &reflected_ray, depth, rng) * rp
                        } else {
                            radiance(scene, &transmitted_ray, depth, rng) * tp
                        }
                    } else {
                        radiance(scene, &reflected_ray, depth, rng) * re +
                            radiance(scene, &transmitted_ray, depth, rng) * tr
                    }
                }
            }
        }
        hit.sphere.emission + colour
    } else {
        Vec3d::zero()
    }
}

fn random_samp<T: Rng>(rng: &mut T) -> f64 {
    let r = 2.0 * rng.gen::<f64>();
    if r < 1.0 { r.sqrt() - 1.0 } else { 1.0 - (2.0 - r).sqrt() }
}

fn to_int(v: f64) -> u8 {
    let ch = (v.powf(1.0/2.2) * 255.0 + 0.5) as i64;
    if ch < 0 { 0 } else if ch > 255 { 255 } else { ch as u8 }
}

fn main() {
    let mut samps = 1;
    let mut width = 1024;
    let mut height = 768;
    let mut output_filename = "image.ppm".to_string();
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Render a simple image");
        ap.refer(&mut samps).add_option(&["-s", "--samples"], Store, "Number of samples");
        ap.refer(&mut height).add_option(&["-h", "--height"], Store, "Height");
        ap.refer(&mut width).add_option(&["-w", "--width"], Store, "Width");
        ap.refer(&mut output_filename).add_option(&["-o", "--output"], Store,
                                                  "Filename to output to");
        ap.parse_args_or_exit();
    }
    samps = samps / 4;
    if samps < 1 { samps = 1; }
    const BLACK : Vec3d = Vec3d { x: 0.0, y: 0.0, z: 0.0 };
    const RED : Vec3d = Vec3d { x: 0.75, y: 0.25, z: 0.25 };
    const BLUE : Vec3d = Vec3d { x: 0.25, y: 0.25, z: 0.75 };
    const GREY : Vec3d = Vec3d { x: 0.75, y: 0.75, z: 0.75 };
    const WHITE : Vec3d = Vec3d { x: 0.999, y: 0.999, z: 0.999 };
    let scene = Arc::new(vec!{
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(1e5+1.0, 40.8, 81.6),
            BLACK, RED),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(-1e5+99.0, 40.8, 81.6),
            BLACK, BLUE),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, 40.8, 1e5),
            BLACK, GREY),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, 40.8, -1e5 + 170.0),
            BLACK, BLACK),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, 1e5, 81.6),
            BLACK, GREY),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, -1e5 + 81.6, 81.6),
            BLACK, GREY),
        Sphere::new(Material::Specular, 16.5,
            Vec3d::new(27.0, 16.5, 47.0),
            BLACK, WHITE),
        Sphere::new(Material::Refractive, 16.5,
            Vec3d::new(73.0, 16.5, 78.0),
            BLACK, WHITE),
        Sphere::new(Material::Diffuse, 600.0,
            Vec3d::new(50.0, 681.6 - 0.27, 81.6),
            Vec3d::new(12.0, 12.0, 12.0), BLACK),
    });

    let camera_pos = Vec3d::new(50.0, 52.0, 295.6);
    let camera_dir = Vec3d::new(0.0, -0.042612, -1.0).normalized();
    let camera_x = Vec3d::new(width as f64 * 0.5135 / height as f64, 0.0, 0.0);
    let camera_y = camera_x.cross(camera_dir).normalized() * 0.5135;

    let num_threads = num_cpus::get();
    println!("Using {} threads", num_threads);
    let pool = ThreadPool::new(num_threads);
    let (tx, rx) = channel();

    for y in 0..height {
        let tx = tx.clone();
        let scene = scene.clone();
        pool.execute(move || {
            let mut line = Vec::with_capacity(width);
            let mut rng = XorShiftRng::from_seed([1 + (y * y) as u32, 0x193a6754, 0x15aac60d, 0xb017f00d]);
            for x in 0..width {
                let mut sum = Vec3d::zero();
                for sx in 0..2 {
                    for sy in 0..2 {
                        let mut r = Vec3d::zero();
                        for _samp in 0..samps {
                            let dx = random_samp(&mut rng);
                            let dy = random_samp(&mut rng);
                            let sub_x = (sx as f64 + 0.5 + dx) / 2.0;
                            let dir_x = (sub_x + x as f64) / width as f64 - 0.5;
                            let sub_y = (sy as f64 + 0.5 + dy) / 2.0;
                            let dir_y = (sub_y + (height -y - 1) as f64) / height as f64 - 0.5;
                            let dir = (camera_x * dir_x + camera_y * dir_y + camera_dir).normalized();
                            let jittered_ray = Ray::new(camera_pos + dir * 140.0, dir);
                            let sample = radiance(&scene, &jittered_ray, 0, &mut rng);
                            r = r + (sample / (samps as f64));
                        }
                        sum = sum + (r.clamp() * 0.25);
                    }
                }
                line.push(sum);
            }
            tx.send((y, line)).unwrap();
        });
    }
    let mut left = height;
    let mut screen : Vec<Vec<Vec3d>> = Vec::new();
    for _y in 0..height {
        screen.push(Vec::new());
    }
    while left > 0 {
        print!("Rendering ({} spp) {:.4}%...\r", samps * 4, 100.0 * (height - left) as f64 / height as f64);
        io::stdout().flush().ok().expect("Could not flush stdout");
        let (y, line) = rx.recv().unwrap();
        screen[y] = line;
        left -= 1;
    }
    println!("\nWriting output to '{}'", output_filename);
    let output_file = File::create(output_filename).unwrap();
    let mut writer = BufWriter::new(output_file);
    write!(&mut writer, "P3\n{} {}\n255\n", width, height).unwrap();
    for y in 0..height {
        for x in 0..width {
            let sum = screen[y][x];
            write!(&mut writer, "{} {} {} ", to_int(sum.x), to_int(sum.y), to_int(sum.z)).unwrap();
        }
    }
}
<commit_msg>Add more idiomatic rust usage of Option<> (thanks to @llogiq)<commit_after>// Based on smallpt, http://www.kevinbeason.com/smallpt/ which is also licensed under
// the MIT license.
use std::fs::File;
use std::io::BufWriter;
use std::io;
use std::io::prelude::*;
use std::sync::Arc;
use std::sync::mpsc::channel;

extern crate rand;
use rand::{Rng, XorShiftRng, SeedableRng};

extern crate threadpool;
use threadpool::ThreadPool;

extern crate num_cpus;

extern crate argparse;
use argparse::{ArgumentParser, Store};

mod path_tracer;
use path_tracer::*;

enum Material {
    Diffuse,
    Specular,
    Refractive
}

struct Sphere {
    material: Material,
    radius: f64,
    position: Vec3d,
    emission: Vec3d,
    colour: Vec3d
}

impl Sphere {
    fn new(material: Material, radius: f64, position: Vec3d, emission: Vec3d, colour: Vec3d) -> Sphere {
        Sphere {
          material: material,
          radius: radius,
          position: position,
          emission: emission,
          colour: colour
        }
    }
    fn intersect(&self, ray: &Ray) -> Option<f64> {
        let op = self.position - ray.origin;
        let b = op.dot(ray.direction);
        let determinant = b * b - op.dot(op) + self.radius * self.radius;
        if determinant < 0.0 { return None; }
        let determinant = determinant.sqrt();
        let t1 = b - determinant;
        let t2 = b + determinant;
        const EPSILON : f64 = 0.0001;
        if t1 > EPSILON {
            Some(t1)
        } else if t2 > EPSILON {
            Some(t2)
        } else {
            None
        }
    }
}

#[test]
fn intersection() {
    let sphere = Sphere::new(
        Material::Diffuse,
        100.0,
        Vec3d::new(0.0, 0.0, 200.0),
        Vec3d::zero(),
        Vec3d::zero());
    let ray = Ray::new(Vec3d::zero(), Vec3d::new(0.0, 0.0, 1.0));
    match sphere.intersect(&ray) {
        Some(x) => assert_eq!(x, 100.0),
        None => panic!("unexpected")
    }
    let ray = Ray::new(Vec3d::zero(), Vec3d::new(0.0, 1.0, 0.0));
    match sphere.intersect(&ray) {
        Some(_) => panic!("unexpected"),
        None => {}
    }
}

struct HitRecord<'a> {
    sphere: &'a Sphere,
    dist: f64
}

fn intersect<'a>(scene: &'a [Sphere], ray: &Ray) -> Option<HitRecord<'a>> {
    let mut result : Option<HitRecord<'a>> = None;
    for sph in scene {
        if let Some(dist) = sph.intersect(&ray) {
            if result.as_ref().map_or(true, |x| dist < x.dist) {
                result = Some(HitRecord { sphere: &sph, dist: dist });
            }
        }
    }
    result
}

fn radiance<R: Rng>(scene: &[Sphere], ray: &Ray, depth: i32, rng: &mut R) -> Vec3d {
    if let Some(hit) = intersect(&scene, &ray) {
        let hit_pos = ray.origin + ray.direction * hit.dist;
        let hit_normal = (hit_pos - hit.sphere.position).normalized();
        let n1 = if hit_normal.dot(ray.direction) < 0.0 { hit_normal } else { hit_normal.neg() };
        let mut colour = hit.sphere.colour;
        let max_reflectance = colour.max_component();
        let depth = depth + 1;
        if depth > 5 {
            let rand = rng.gen::<f64>();
            if rand < max_reflectance && depth < 500 { // Rust's stack blows up ~600 on my machine
               colour = colour * (1.0 / max_reflectance);
            } else {
                return hit.sphere.emission;
            }
        }
        match hit.sphere.material {
            Material::Diffuse => {
                // Get a random polar coordinate
                let r1 = rng.gen::<f64>() * 2.0 * std::f64::consts::PI;
                let r2 = rng.gen::<f64>();
                let r2s = r2.sqrt();
                // Create a coordinate system u,v,w local to the point, where the w is the normal
                // pointing out of the sphere and the u and v are orthonormal to w.
                let w = n1;
                // Pick an arbitrary non-zero preferred axis for u
                let u = if n1.x.abs() > 0.1 { Vec3d::new(0.0, 1.0, 0.0) } else { Vec3d::new(1.0, 0.0, 0.0) }.cross(w);
                let v = w.cross(u);
                // construct the new direction
                let new_dir = u * r1.cos() * r2s + v * r1.sin() * r2s + w * (1.0 - r2).sqrt();
                colour = colour * radiance(scene, &Ray::new(hit_pos, new_dir.normalized()), depth, rng);
            },
            Material::Specular => {
                let reflection = ray.direction - hit_normal * 2.0 * hit_normal.dot(ray.direction);
                let reflected_ray = Ray::new(hit_pos, reflection);
                colour = colour * radiance(scene, &reflected_ray, depth, rng);
            },
            Material::Refractive => {
                let reflection = ray.direction - hit_normal * 2.0 * hit_normal.dot(ray.direction);
                let reflected_ray = Ray::new(hit_pos, reflection);
                let into = if hit_normal.dot(n1) > 0.0 { true } else { false };
                let nc = 1.0;
                let nt = 1.5;
                let nnt = if into { nc/nt } else { nt/nc };
                let ddn = ray.direction.dot(n1);
                let cos2t = 1.0 - nnt * nnt * (1.0 - ddn * ddn);
                if cos2t < 0.0 {
                    // Total internal reflection
                    colour = colour * radiance(scene, &reflected_ray, depth, rng);
                } else {
                    let tbd = ddn * nnt + cos2t.sqrt();
                    let tbd = if into { tbd } else { -tbd };
                    let tdir = (ray.direction * nnt - hit_normal * tbd).normalized();
                    let transmitted_ray = Ray::new(hit_pos, tdir);
                    let a = nt - nc;
                    let b = nt + nc;
                    let r0 = (a * a) / (b * b);
                    let c = 1.0 - if into { -ddn } else { tdir.dot(hit_normal) };
                    let re = r0 + (1.0 - r0) * c * c * c * c * c;
                    let tr = 1.0 - re;
                    let p = 0.25 + 0.5 * re;
                    let rp = re / p;
                    let tp = tr / (1.0 - p);
                    colour = colour * if depth > 2 {
                        if rng.gen::<f64>() < p {
                            radiance(scene, &reflected_ray, depth, rng) * rp
                        } else {
                            radiance(scene, &transmitted_ray, depth, rng) * tp
                        }
                    } else {
                        radiance(scene, &reflected_ray, depth, rng) * re +
                            radiance(scene, &transmitted_ray, depth, rng) * tr
                    }
                }
            }
        }
        hit.sphere.emission + colour
    } else {
        Vec3d::zero()
    }
}

fn random_samp<T: Rng>(rng: &mut T) -> f64 {
    let r = 2.0 * rng.gen::<f64>();
    if r < 1.0 { r.sqrt() - 1.0 } else { 1.0 - (2.0 - r).sqrt() }
}

fn to_int(v: f64) -> u8 {
    let ch = (v.powf(1.0/2.2) * 255.0 + 0.5) as i64;
    if ch < 0 { 0 } else if ch > 255 { 255 } else { ch as u8 }
}

fn main() {
    let mut samps = 1;
    let mut width = 1024;
    let mut height = 768;
    let mut output_filename = "image.ppm".to_string();
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Render a simple image");
        ap.refer(&mut samps).add_option(&["-s", "--samples"], Store, "Number of samples");
        ap.refer(&mut height).add_option(&["-h", "--height"], Store, "Height");
        ap.refer(&mut width).add_option(&["-w", "--width"], Store, "Width");
        ap.refer(&mut output_filename).add_option(&["-o", "--output"], Store,
                                                  "Filename to output to");
        ap.parse_args_or_exit();
    }
    samps = samps / 4;
    if samps < 1 { samps = 1; }
    const BLACK : Vec3d = Vec3d { x: 0.0, y: 0.0, z: 0.0 };
    const RED : Vec3d = Vec3d { x: 0.75, y: 0.25, z: 0.25 };
    const BLUE : Vec3d = Vec3d { x: 0.25, y: 0.25, z: 0.75 };
    const GREY : Vec3d = Vec3d { x: 0.75, y: 0.75, z: 0.75 };
    const WHITE : Vec3d = Vec3d { x: 0.999, y: 0.999, z: 0.999 };
    let scene = Arc::new(vec!{
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(1e5+1.0, 40.8, 81.6),
            BLACK, RED),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(-1e5+99.0, 40.8, 81.6),
            BLACK, BLUE),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, 40.8, 1e5),
            BLACK, GREY),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, 40.8, -1e5 + 170.0),
            BLACK, BLACK),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, 1e5, 81.6),
            BLACK, GREY),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, -1e5 + 81.6, 81.6),
            BLACK, GREY),
        Sphere::new(Material::Specular, 16.5,
            Vec3d::new(27.0, 16.5, 47.0),
            BLACK, WHITE),
        Sphere::new(Material::Refractive, 16.5,
            Vec3d::new(73.0, 16.5, 78.0),
            BLACK, WHITE),
        Sphere::new(Material::Diffuse, 600.0,
            Vec3d::new(50.0, 681.6 - 0.27, 81.6),
            Vec3d::new(12.0, 12.0, 12.0), BLACK),
    });

    let camera_pos = Vec3d::new(50.0, 52.0, 295.6);
    let camera_dir = Vec3d::new(0.0, -0.042612, -1.0).normalized();
    let camera_x = Vec3d::new(width as f64 * 0.5135 / height as f64, 0.0, 0.0);
    let camera_y = camera_x.cross(camera_dir).normalized() * 0.5135;

    let num_threads = num_cpus::get();
    println!("Using {} threads", num_threads);
    let pool = ThreadPool::new(num_threads);
    let (tx, rx) = channel();

    for y in 0..height {
        let tx = tx.clone();
        let scene = scene.clone();
        pool.execute(move || {
            let mut line = Vec::with_capacity(width);
            let mut rng = XorShiftRng::from_seed([1 + (y * y) as u32, 0x193a6754, 0x15aac60d, 0xb017f00d]);
            for x in 0..width {
                let mut sum = Vec3d::zero();
                for sx in 0..2 {
                    for sy in 0..2 {
                        let mut r = Vec3d::zero();
                        for _samp in 0..samps {
                            let dx = random_samp(&mut rng);
                            let dy = random_samp(&mut rng);
                            let sub_x = (sx as f64 + 0.5 + dx) / 2.0;
                            let dir_x = (sub_x + x as f64) / width as f64 - 0.5;
                            let sub_y = (sy as f64 + 0.5 + dy) / 2.0;
                            let dir_y = (sub_y + (height -y - 1) as f64) / height as f64 - 0.5;
                            let dir = (camera_x * dir_x + camera_y * dir_y + camera_dir).normalized();
                            let jittered_ray = Ray::new(camera_pos + dir * 140.0, dir);
                            let sample = radiance(&scene, &jittered_ray, 0, &mut rng);
                            r = r + (sample / (samps as f64));
                        }
                        sum = sum + (r.clamp() * 0.25);
                    }
                }
                line.push(sum);
            }
            tx.send((y, line)).unwrap();
        });
    }
    let mut left = height;
    let mut screen : Vec<Vec<Vec3d>> = Vec::new();
    for _y in 0..height {
        screen.push(Vec::new());
    }
    while left > 0 {
        print!("Rendering ({} spp) {:.4}%...\r", samps * 4, 100.0 * (height - left) as f64 / height as f64);
        io::stdout().flush().ok().expect("Could not flush stdout");
        let (y, line) = rx.recv().unwrap();
        screen[y] = line;
        left -= 1;
    }
    println!("\nWriting output to '{}'", output_filename);
    let output_file = File::create(output_filename).unwrap();
    let mut writer = BufWriter::new(output_file);
    write!(&mut writer, "P3\n{} {}\n255\n", width, height).unwrap();
    for y in 0..height {
        for x in 0..width {
            let sum = screen[y][x];
            write!(&mut writer, "{} {} {} ", to_int(sum.x), to_int(sum.y), to_int(sum.z)).unwrap();
        }
    }
}
<|endoftext|>


<commit_before>#![feature(plugin)]
#![plugin(clippy)]

extern crate avsser;
extern crate getopts;

use getopts::Options;
use std::env;
use std::path::PathBuf;
use avsser::generic::input::get_list_of_files;
use avsser::generic::output::create_avs_script;
use avsser::generic::output::AvsOptions;

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options] INPUT", program);
    print!("{}", opts.usage(&brief));
}

fn main() {
    let program = env::args().next().unwrap();
    let args: Vec<String> = env::args().skip(1).collect();

    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help menu");
    let matches = match opts.parse(&args) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }
    let input = if !matches.free.is_empty() {
        matches.free[0].clone()
    } else {
        print_usage(&program, opts);
        return;
    };

    let input = get_list_of_files(PathBuf::from(input), false).ok().expect("Unable to read input file(s)");
    for file in input {
        let path = PathBuf::from(file);
        create_avs_script(
            path.clone(),
            path.with_extension("avs"),
            AvsOptions {
                remove_grain: Some(1)
            }).ok();
    }
}
<commit_msg>Skip files that aren't videos<commit_after>#![feature(plugin)]
#![plugin(clippy)]

extern crate avsser;
extern crate getopts;

use getopts::Options;
use std::env;
use std::path::PathBuf;
use avsser::generic::input::get_list_of_files;
use avsser::generic::input::determine_input_type;
use avsser::generic::output::create_avs_script;
use avsser::generic::output::AvsOptions;

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options] INPUT", program);
    print!("{}", opts.usage(&brief));
}

fn main() {
    let program = env::args().next().unwrap();
    let args: Vec<String> = env::args().skip(1).collect();

    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help menu");
    let matches = match opts.parse(&args) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }
    let input = if !matches.free.is_empty() {
        matches.free[0].clone()
    } else {
        print_usage(&program, opts);
        return;
    };

    let input = get_list_of_files(PathBuf::from(input), false).ok().expect("Unable to read input file(s)");
    for file in input {
        if determine_input_type(file.clone()).is_none() {
            continue;
        }
        let path = PathBuf::from(file);
        create_avs_script(
            path.clone(),
            path.with_extension("avs"),
            AvsOptions {
                remove_grain: Some(1)
            }).ok();
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::prelude::*;
use std::net::TcpStream;
use std::fs::File;
use std::env;

extern crate chrono;
use chrono::*;

#[macro_use]
extern crate clap;

extern crate env_logger;
#[macro_use]
extern crate log;

extern crate stats;

extern crate prometheus;
use prometheus::Encoder;
extern crate hyper;

mod analyzer;
mod args;
mod filter;
mod log_parser;
use log_parser::log_events::*;
mod render;
use render::Renderer;
mod request_response_matcher;
use request_response_matcher::*;

fn main() {
    env_logger::init().expect("Failed to initialize logging.");

    let args = args::parse_args(env::args()).expect("Failed to parse arguments.");

    let result = run(&args);

    let mut stream;
    let mut renderer: Box<render::Renderer>;

    renderer = match args.graphite_server {
        Some(ref graphite_server) => {
            stream = TcpStream::connect((graphite_server.as_ref(), args.graphite_port.unwrap()))
                .expect("Could not connect to the Graphite server");

            Box::new(render::GraphiteRenderer::new(UTC::now(),
                                                   args.graphite_prefix.clone(),
                                                   &mut stream))
        }
        None => Box::new(render::TerminalRenderer::new()),
    };

    match result {
        Some(result) => {
            renderer.render(result);
        }
        None => warn!("No matching log lines in file."),
    }

    let addr = "127.0.0.1:9898";
    println!("listening addr {:?}", addr);
    hyper::server::Server::http(addr)
        .unwrap()
        .handle(move |_: hyper::server::Request, mut res: hyper::server::Response| {

            let result = run(&args);

            let mut renderer = render::prometheus::PrometheusRenderer::new();
            renderer.render(result.unwrap());
            res.headers_mut()
                .set(hyper::header::ContentType(renderer.encoder
                    .format_type()
                    .parse::<hyper::mime::Mime>()
                    .unwrap()));
            res.send(&renderer.buffer).unwrap();
        })
        .unwrap();
}

fn run(args: &args::RequestLogAnalyzerArgs) -> Option<analyzer::RequestLogAnalyzerResult> {
    let input: Box<io::Read> = match args.filename.as_ref() {
        "-" => Box::new(io::stdin()),
        _ => Box::new(File::open(&args.filename).expect("Failed to open file.")),
    };

    let timings = extract_timings(input, &args.conditions);
    let result = analyzer::analyze(&timings);

    result
}

fn extract_timings(input: Box<io::Read>, conditions: &filter::FilterConditions) -> Vec<i64> {
    let reader = io::BufReader::new(input);

    let mut requests: Vec<Request> = Vec::new();
    let mut responses: Vec<Response> = Vec::new();
    let mut timings: Vec<i64> = Vec::new();

    for line in reader.lines() {
        let parsed_line = log_parser::parse_line(&line.expect("Failed to read line."));

        match parsed_line {
            Ok(event) => {
                match event {
                    LogEvent::Request(request) => requests.push(request),
                    LogEvent::Response(response) => responses.push(response),
                }

                let pairs = extract_matching_request_response_pairs(&mut requests, &mut responses);

                let mut new_timings: Vec<i64> = pairs.iter()
                    .filter(|pair| filter::matches_filter(&pair, conditions))
                    .map(|pair| pair.response.response_time.num_milliseconds())
                    .collect();

                timings.append(&mut new_timings);
            }
            Err(err) => warn!("{}", err),
        }
    }
    timings
}

#[test]
fn test_extract_timings() {
    let conditions = filter::FilterConditions {
        include_terms: None,
        exclude_terms: None,
        latest_time: None,
    };

    let timings = extract_timings(Box::new(File::open("src/test/simple-1.log").unwrap()),
                                  &conditions);

    assert_eq!(timings, vec![7, 10]);
}

#[test]
fn test_extract_timings_ignore_broken_lines() {
    let conditions = filter::FilterConditions {
        include_terms: None,
        exclude_terms: None,
        latest_time: None,
    };

    let timings = extract_timings(Box::new(File::open("src/test/broken.log").unwrap()),
                                  &conditions);

    assert_eq!(timings, vec![7]);
}

#[test]
fn test_run() {
    let args = args::RequestLogAnalyzerArgs {
        filename: String::from("src/test/simple-1.log"),
        conditions: filter::FilterConditions {
            include_terms: None,
            exclude_terms: None,
            latest_time: None,
        },
        graphite_server: None,
        graphite_port: Some(2003),
        graphite_prefix: None,
    };

    let result = run(&args).unwrap();
    assert_eq!(result.count, 2);
}
<commit_msg>Extract Prometheus HTTP server setup<commit_after>use std::io;
use std::io::prelude::*;
use std::net::TcpStream;
use std::fs::File;
use std::env;

extern crate chrono;
use chrono::*;

#[macro_use]
extern crate clap;

extern crate env_logger;
#[macro_use]
extern crate log;

extern crate stats;

extern crate prometheus;
use prometheus::Encoder;
extern crate hyper;

mod analyzer;
mod args;
mod filter;
mod log_parser;
use log_parser::log_events::*;
mod render;
use render::Renderer;
mod request_response_matcher;
use request_response_matcher::*;

fn main() {
    env_logger::init().expect("Failed to initialize logging.");

    let args = args::parse_args(env::args()).expect("Failed to parse arguments.");

    let result = run(&args);

    let mut stream;
    let mut renderer: Box<render::Renderer>;

    renderer = match args.graphite_server {
        Some(ref graphite_server) => {
            stream = TcpStream::connect((graphite_server.as_ref(), args.graphite_port.unwrap()))
                .expect("Could not connect to the Graphite server");

            Box::new(render::GraphiteRenderer::new(UTC::now(),
                                                   args.graphite_prefix.clone(),
                                                   &mut stream))
        }
        None => Box::new(render::TerminalRenderer::new()),
    };

    match result {
        Some(result) => {
            renderer.render(result);
        }
        None => warn!("No matching log lines in file."),
    }

    listen_http(args, "127.0.0.1:9898");
}

fn listen_http(args: args::RequestLogAnalyzerArgs, binding_address: &str) {
    info!("listening addr {:?}", binding_address);
    hyper::server::Server::http(binding_address)
        .unwrap()
        .handle(move |_: hyper::server::Request, mut res: hyper::server::Response| {

            let result = run(&args);

            let mut renderer = render::prometheus::PrometheusRenderer::new();
            renderer.render(result.unwrap());
            res.headers_mut()
                .set(hyper::header::ContentType(renderer.encoder
                    .format_type()
                    .parse::<hyper::mime::Mime>()
                    .unwrap()));
            res.send(&renderer.buffer).unwrap();
        })
        .unwrap();
}

fn run(args: &args::RequestLogAnalyzerArgs) -> Option<analyzer::RequestLogAnalyzerResult> {
    let input: Box<io::Read> = match args.filename.as_ref() {
        "-" => Box::new(io::stdin()),
        _ => Box::new(File::open(&args.filename).expect("Failed to open file.")),
    };

    let timings = extract_timings(input, &args.conditions);
    let result = analyzer::analyze(&timings);

    result
}

fn extract_timings(input: Box<io::Read>, conditions: &filter::FilterConditions) -> Vec<i64> {
    let reader = io::BufReader::new(input);

    let mut requests: Vec<Request> = Vec::new();
    let mut responses: Vec<Response> = Vec::new();
    let mut timings: Vec<i64> = Vec::new();

    for line in reader.lines() {
        let parsed_line = log_parser::parse_line(&line.expect("Failed to read line."));

        match parsed_line {
            Ok(event) => {
                match event {
                    LogEvent::Request(request) => requests.push(request),
                    LogEvent::Response(response) => responses.push(response),
                }

                let pairs = extract_matching_request_response_pairs(&mut requests, &mut responses);

                let mut new_timings: Vec<i64> = pairs.iter()
                    .filter(|pair| filter::matches_filter(&pair, conditions))
                    .map(|pair| pair.response.response_time.num_milliseconds())
                    .collect();

                timings.append(&mut new_timings);
            }
            Err(err) => warn!("{}", err),
        }
    }
    timings
}

#[test]
fn test_extract_timings() {
    let conditions = filter::FilterConditions {
        include_terms: None,
        exclude_terms: None,
        latest_time: None,
    };

    let timings = extract_timings(Box::new(File::open("src/test/simple-1.log").unwrap()),
                                  &conditions);

    assert_eq!(timings, vec![7, 10]);
}

#[test]
fn test_extract_timings_ignore_broken_lines() {
    let conditions = filter::FilterConditions {
        include_terms: None,
        exclude_terms: None,
        latest_time: None,
    };

    let timings = extract_timings(Box::new(File::open("src/test/broken.log").unwrap()),
                                  &conditions);

    assert_eq!(timings, vec![7]);
}

#[test]
fn test_run() {
    let args = args::RequestLogAnalyzerArgs {
        filename: String::from("src/test/simple-1.log"),
        conditions: filter::FilterConditions {
            include_terms: None,
            exclude_terms: None,
            latest_time: None,
        },
        graphite_server: None,
        graphite_port: Some(2003),
        graphite_prefix: None,
    };

    let result = run(&args).unwrap();
    assert_eq!(result.count, 2);
}
<|endoftext|>


<commit_before>pub mod bookitem;
pub mod bookconfig;

pub mod bookconfig_test;

pub use self::bookitem::{BookItem, BookItems};
pub use self::bookconfig::BookConfig;

use std::path::{Path, PathBuf};
use std::fs::{self, File};
use std::error::Error;
use std::io;
use std::io::Write;
use std::io::ErrorKind;
use std::process::Command;

use {theme, parse, utils};
use renderer::{Renderer, HtmlHandlebars};

pub struct MDBook {
    root: PathBuf,
    dest: PathBuf,
    src: PathBuf,
    theme_path: PathBuf,

    pub title: String,
    pub author: String,
    pub description: String,

    pub content: Vec<BookItem>,
    renderer: Box<Renderer>,

    livereload: Option<String>,

    /// Should `mdbook build` create files referenced from SUMMARY.md if they
    /// don't exist
    pub create_missing: bool,
}

impl MDBook {
    /// Create a new `MDBook` struct with root directory `root`
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate mdbook;
    /// # use mdbook::MDBook;
    /// # use std::path::Path;
    /// # fn main() {
    /// let book = MDBook::new(Path::new("root_dir"));
    /// # }
    /// ```
    ///
    /// In this example, `root_dir` will be the root directory of our book and is specified in function
    /// of the current working directory by using a relative path instead of an absolute path.
    ///
    /// Default directory paths:
    ///
    /// - source: `root/src`
    /// - output: `root/book`
    /// - theme: `root/theme`
    ///
    /// They can both be changed by using [`set_src()`](#method.set_src) and [`set_dest()`](#method.set_dest)

    pub fn new(root: &Path) -> MDBook {

        if !root.exists() || !root.is_dir() {
            warn!("{:?} No directory with that name", root);
        }

        MDBook {
            root: root.to_owned(),
            dest: root.join("book"),
            src: root.join("src"),
            theme_path: root.join("theme"),

            title: String::new(),
            author: String::new(),
            description: String::new(),

            content: vec![],
            renderer: Box::new(HtmlHandlebars::new()),

            livereload: None,
            create_missing: true,
        }
    }

    /// Returns a flat depth-first iterator over the elements of the book, it returns an [BookItem enum](bookitem.html):
    /// `(section: String, bookitem: &BookItem)`
    ///
    /// ```no_run
    /// # extern crate mdbook;
    /// # use mdbook::MDBook;
    /// # use mdbook::BookItem;
    /// # use std::path::Path;
    /// # fn main() {
    /// # let mut book = MDBook::new(Path::new("mybook"));
    /// for item in book.iter() {
    ///     match item {
    ///         &BookItem::Chapter(ref section, ref chapter) => {},
    ///         &BookItem::Affix(ref chapter) => {},
    ///         &BookItem::Spacer => {},
    ///     }
    /// }
    ///
    /// // would print something like this:
    /// // 1. Chapter 1
    /// // 1.1 Sub Chapter
    /// // 1.2 Sub Chapter
    /// // 2. Chapter 2
    /// //
    /// // etc.
    /// # }
    /// ```

    pub fn iter(&self) -> BookItems {
        BookItems {
            items: &self.content[..],
            current_index: 0,
            stack: Vec::new(),
        }
    }

    /// `init()` creates some boilerplate files and directories to get you started with your book.
    ///
    /// ```text
    /// book-test/
    /// ├── book
    /// └── src
    ///     ├── chapter_1.md
    ///     └── SUMMARY.md
    /// ```
    ///
    /// It uses the paths given as source and output directories and adds a `SUMMARY.md` and a
    /// `chapter_1.md` to the source directory.

    pub fn init(&mut self) -> Result<(), Box<Error>> {

        debug!("[fn]: init");

        if !self.root.exists() {
            fs::create_dir_all(&self.root).unwrap();
            info!("{:?} created", &self.root);
        }

        {

            if !self.dest.exists() {
                debug!("[*]: {:?} does not exist, trying to create directory", self.dest);
                try!(fs::create_dir_all(&self.dest));
            }

            if !self.src.exists() {
                debug!("[*]: {:?} does not exist, trying to create directory", self.src);
                try!(fs::create_dir_all(&self.src));
            }

            let summary = self.src.join("SUMMARY.md");

            if !summary.exists() {

                // Summary does not exist, create it

                debug!("[*]: {:?} does not exist, trying to create SUMMARY.md", self.src.join("SUMMARY.md"));
                let mut f = try!(File::create(&self.src.join("SUMMARY.md")));

                debug!("[*]: Writing to SUMMARY.md");

                try!(writeln!(f, "# Summary"));
                try!(writeln!(f, ""));
                try!(writeln!(f, "- [Chapter 1](./chapter_1.md)"));
            }
        }

        // parse SUMMARY.md, and create the missing item related file
        try!(self.parse_summary());

        debug!("[*]: constructing paths for missing files");
        for item in self.iter() {
            debug!("[*]: item: {:?}", item);
            let ch = match *item {
                BookItem::Spacer => continue,
                BookItem::Chapter(_, ref ch) |
                BookItem::Affix(ref ch) => ch,
            };
            if ch.path.as_os_str().is_empty() {
                let path = self.src.join(&ch.path);

                if !path.exists() {
                    if !self.create_missing {
                        return Err(format!(
                            "'{}' referenced from SUMMARY.md does not exist.",
                            path.to_string_lossy()).into());
                    }
                    debug!("[*]: {:?} does not exist, trying to create file", path);
                    try!(::std::fs::create_dir_all(path.parent().unwrap()));
                    let mut f = try!(File::create(path));

                    // debug!("[*]: Writing to {:?}", path);
                    try!(writeln!(f, "# {}", ch.name));
                }
            }
        }

        debug!("[*]: init done");
        Ok(())
    }

    pub fn create_gitignore(&self) {
        let gitignore = self.get_gitignore();

        if !gitignore.exists() {
            // Gitignore does not exist, create it

            // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`. If it
            // is not, `strip_prefix` will return an Error.
            if !self.get_dest().starts_with(&self.root) {
                return;
            }

            let relative = self.get_dest()
                               .strip_prefix(&self.root)
                               .expect("Destination is not relative to root.");
            let relative = relative.to_str()
                                   .expect("Path could not be yielded into a string slice.");

            debug!("[*]: {:?} does not exist, trying to create .gitignore", gitignore);

            let mut f = File::create(&gitignore).expect("Could not create file.");

            debug!("[*]: Writing to .gitignore");

            writeln!(f, "{}", relative).expect("Could not write to file.");
        }
    }

    /// The `build()` method is the one where everything happens. First it parses `SUMMARY.md` to
    /// construct the book's structure in the form of a `Vec<BookItem>` and then calls `render()`
    /// method of the current renderer.
    ///
    /// It is the renderer who generates all the output files.
    pub fn build(&mut self) -> Result<(), Box<Error>> {
        debug!("[fn]: build");

        try!(self.init());

        // Clean output directory
        try!(utils::fs::remove_dir_content(&self.dest));

        try!(self.renderer.render(&self));

        Ok(())
    }

    pub fn get_gitignore(&self) -> PathBuf {
        self.root.join(".gitignore")
    }

    pub fn copy_theme(&self) -> Result<(), Box<Error>> {
        debug!("[fn]: copy_theme");

        let theme_dir = self.src.join("theme");

        if !theme_dir.exists() {
            debug!("[*]: {:?} does not exist, trying to create directory", theme_dir);
            try!(fs::create_dir(&theme_dir));
        }

        // index.hbs
        let mut index = try!(File::create(&theme_dir.join("index.hbs")));
        try!(index.write_all(theme::INDEX));

        // book.css
        let mut css = try!(File::create(&theme_dir.join("book.css")));
        try!(css.write_all(theme::CSS));

        // favicon.png
        let mut favicon = try!(File::create(&theme_dir.join("favicon.png")));
        try!(favicon.write_all(theme::FAVICON));

        // book.js
        let mut js = try!(File::create(&theme_dir.join("book.js")));
        try!(js.write_all(theme::JS));

        // highlight.css
        let mut highlight_css = try!(File::create(&theme_dir.join("highlight.css")));
        try!(highlight_css.write_all(theme::HIGHLIGHT_CSS));

        // highlight.js
        let mut highlight_js = try!(File::create(&theme_dir.join("highlight.js")));
        try!(highlight_js.write_all(theme::HIGHLIGHT_JS));

        Ok(())
    }

    pub fn write_file<P: AsRef<Path>>(&self, filename: P, content: &[u8]) -> Result<(), Box<Error>> {
        let path = self.get_dest().join(filename);
        try!(utils::fs::create_file(&path).and_then(|mut file| {
            file.write_all(content)
        }).map_err(|e| {
            io::Error::new(io::ErrorKind::Other, format!("Could not create {}: {}", path.display(), e))
        }));
        Ok(())
    }

    /// Parses the `book.json` file (if it exists) to extract the configuration parameters.
    /// The `book.json` file should be in the root directory of the book.
    /// The root directory is the one specified when creating a new `MDBook`

    pub fn read_config(mut self) -> Self {

        let config = BookConfig::new(&self.root)
                         .read_config(&self.root)
                         .to_owned();

        self.title = config.title;
        self.description = config.description;
        self.author = config.author;

        self.dest = config.dest;
        self.src = config.src;
        self.theme_path = config.theme_path;

        self
    }

    /// You can change the default renderer to another one by using this method. The only requirement
    /// is for your renderer to implement the [Renderer trait](../../renderer/renderer/trait.Renderer.html)
    ///
    /// ```no_run
    /// extern crate mdbook;
    /// use mdbook::MDBook;
    /// use mdbook::renderer::HtmlHandlebars;
    /// # use std::path::Path;
    ///
    /// fn main() {
    ///     let mut book = MDBook::new(Path::new("mybook"))
    ///                         .set_renderer(Box::new(HtmlHandlebars::new()));
    ///
    ///     // In this example we replace the default renderer by the default renderer...
    ///     // Don't forget to put your renderer in a Box
    /// }
    /// ```
    ///
    /// **note:** Don't forget to put your renderer in a `Box` before passing it to `set_renderer()`

    pub fn set_renderer(mut self, renderer: Box<Renderer>) -> Self {
        self.renderer = renderer;
        self
    }

    pub fn test(&mut self) -> Result<(), Box<Error>> {
        // read in the chapters
        try!(self.parse_summary());
        for item in self.iter() {

            if let BookItem::Chapter(_, ref ch) = *item {
                if ch.path != PathBuf::new() {

                    let path = self.get_src().join(&ch.path);

                    println!("[*]: Testing file: {:?}", path);

                    let output_result = Command::new("rustdoc")
                                            .arg(&path)
                                            .arg("--test")
                                            .output();
                    let output = try!(output_result);

                    if !output.status.success() {
                        return Err(Box::new(io::Error::new(ErrorKind::Other, format!(
                                        "{}\n{}",
                                        String::from_utf8_lossy(&output.stdout),
                                        String::from_utf8_lossy(&output.stderr)))) as Box<Error>);
                    }
                }
            }
        }
        Ok(())
    }

    pub fn get_root(&self) -> &Path {
        &self.root
    }

    pub fn set_dest(mut self, dest: &Path) -> Self {

        // Handle absolute and relative paths
        if dest.is_absolute() {
            self.dest = dest.to_owned();
        } else {
            let dest = self.root.join(dest).to_owned();
            self.dest = dest;
        }

        self
    }

    pub fn get_dest(&self) -> &Path {
        &self.dest
    }

    pub fn set_src(mut self, src: &Path) -> Self {

        // Handle absolute and relative paths
        if src.is_absolute() {
            self.src = src.to_owned();
        } else {
            let src = self.root.join(src).to_owned();
            self.src = src;
        }

        self
    }

    pub fn get_src(&self) -> &Path {
        &self.src
    }

    pub fn set_title(mut self, title: &str) -> Self {
        self.title = title.to_owned();
        self
    }

    pub fn get_title(&self) -> &str {
        &self.title
    }

    pub fn set_author(mut self, author: &str) -> Self {
        self.author = author.to_owned();
        self
    }

    pub fn get_author(&self) -> &str {
        &self.author
    }

    pub fn set_description(mut self, description: &str) -> Self {
        self.description = description.to_owned();
        self
    }

    pub fn get_description(&self) -> &str {
        &self.description
    }

    pub fn set_livereload(&mut self, livereload: String) -> &mut Self {
        self.livereload = Some(livereload);
        self
    }

    pub fn unset_livereload(&mut self) -> &Self {
        self.livereload = None;
        self
    }

    pub fn get_livereload(&self) -> Option<&String> {
        self.livereload.as_ref()
    }

    pub fn set_theme_path(mut self, theme_path: &Path) -> Self {
        self.theme_path = if theme_path.is_absolute() {
            theme_path.to_owned()
        } else {
            self.root.join(theme_path).to_owned()
        };
        self
    }

    pub fn get_theme_path(&self) -> &Path {
        &self.theme_path
    }

    // Construct book
    fn parse_summary(&mut self) -> Result<(), Box<Error>> {
        // When append becomes stable, use self.content.append() ...
        self.content = try!(parse::construct_bookitems(&self.src.join("SUMMARY.md")));
        Ok(())
    }
}
<commit_msg>Fix incorrect conditional.<commit_after>pub mod bookitem;
pub mod bookconfig;

pub mod bookconfig_test;

pub use self::bookitem::{BookItem, BookItems};
pub use self::bookconfig::BookConfig;

use std::path::{Path, PathBuf};
use std::fs::{self, File};
use std::error::Error;
use std::io;
use std::io::Write;
use std::io::ErrorKind;
use std::process::Command;

use {theme, parse, utils};
use renderer::{Renderer, HtmlHandlebars};

pub struct MDBook {
    root: PathBuf,
    dest: PathBuf,
    src: PathBuf,
    theme_path: PathBuf,

    pub title: String,
    pub author: String,
    pub description: String,

    pub content: Vec<BookItem>,
    renderer: Box<Renderer>,

    livereload: Option<String>,

    /// Should `mdbook build` create files referenced from SUMMARY.md if they
    /// don't exist
    pub create_missing: bool,
}

impl MDBook {
    /// Create a new `MDBook` struct with root directory `root`
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate mdbook;
    /// # use mdbook::MDBook;
    /// # use std::path::Path;
    /// # fn main() {
    /// let book = MDBook::new(Path::new("root_dir"));
    /// # }
    /// ```
    ///
    /// In this example, `root_dir` will be the root directory of our book and is specified in function
    /// of the current working directory by using a relative path instead of an absolute path.
    ///
    /// Default directory paths:
    ///
    /// - source: `root/src`
    /// - output: `root/book`
    /// - theme: `root/theme`
    ///
    /// They can both be changed by using [`set_src()`](#method.set_src) and [`set_dest()`](#method.set_dest)

    pub fn new(root: &Path) -> MDBook {

        if !root.exists() || !root.is_dir() {
            warn!("{:?} No directory with that name", root);
        }

        MDBook {
            root: root.to_owned(),
            dest: root.join("book"),
            src: root.join("src"),
            theme_path: root.join("theme"),

            title: String::new(),
            author: String::new(),
            description: String::new(),

            content: vec![],
            renderer: Box::new(HtmlHandlebars::new()),

            livereload: None,
            create_missing: true,
        }
    }

    /// Returns a flat depth-first iterator over the elements of the book, it returns an [BookItem enum](bookitem.html):
    /// `(section: String, bookitem: &BookItem)`
    ///
    /// ```no_run
    /// # extern crate mdbook;
    /// # use mdbook::MDBook;
    /// # use mdbook::BookItem;
    /// # use std::path::Path;
    /// # fn main() {
    /// # let mut book = MDBook::new(Path::new("mybook"));
    /// for item in book.iter() {
    ///     match item {
    ///         &BookItem::Chapter(ref section, ref chapter) => {},
    ///         &BookItem::Affix(ref chapter) => {},
    ///         &BookItem::Spacer => {},
    ///     }
    /// }
    ///
    /// // would print something like this:
    /// // 1. Chapter 1
    /// // 1.1 Sub Chapter
    /// // 1.2 Sub Chapter
    /// // 2. Chapter 2
    /// //
    /// // etc.
    /// # }
    /// ```

    pub fn iter(&self) -> BookItems {
        BookItems {
            items: &self.content[..],
            current_index: 0,
            stack: Vec::new(),
        }
    }

    /// `init()` creates some boilerplate files and directories to get you started with your book.
    ///
    /// ```text
    /// book-test/
    /// ├── book
    /// └── src
    ///     ├── chapter_1.md
    ///     └── SUMMARY.md
    /// ```
    ///
    /// It uses the paths given as source and output directories and adds a `SUMMARY.md` and a
    /// `chapter_1.md` to the source directory.

    pub fn init(&mut self) -> Result<(), Box<Error>> {

        debug!("[fn]: init");

        if !self.root.exists() {
            fs::create_dir_all(&self.root).unwrap();
            info!("{:?} created", &self.root);
        }

        {

            if !self.dest.exists() {
                debug!("[*]: {:?} does not exist, trying to create directory", self.dest);
                try!(fs::create_dir_all(&self.dest));
            }

            if !self.src.exists() {
                debug!("[*]: {:?} does not exist, trying to create directory", self.src);
                try!(fs::create_dir_all(&self.src));
            }

            let summary = self.src.join("SUMMARY.md");

            if !summary.exists() {

                // Summary does not exist, create it

                debug!("[*]: {:?} does not exist, trying to create SUMMARY.md", self.src.join("SUMMARY.md"));
                let mut f = try!(File::create(&self.src.join("SUMMARY.md")));

                debug!("[*]: Writing to SUMMARY.md");

                try!(writeln!(f, "# Summary"));
                try!(writeln!(f, ""));
                try!(writeln!(f, "- [Chapter 1](./chapter_1.md)"));
            }
        }

        // parse SUMMARY.md, and create the missing item related file
        try!(self.parse_summary());

        debug!("[*]: constructing paths for missing files");
        for item in self.iter() {
            debug!("[*]: item: {:?}", item);
            let ch = match *item {
                BookItem::Spacer => continue,
                BookItem::Chapter(_, ref ch) |
                BookItem::Affix(ref ch) => ch,
            };
            if !ch.path.as_os_str().is_empty() {
                let path = self.src.join(&ch.path);

                if !path.exists() {
                    if !self.create_missing {
                        return Err(format!(
                            "'{}' referenced from SUMMARY.md does not exist.",
                            path.to_string_lossy()).into());
                    }
                    debug!("[*]: {:?} does not exist, trying to create file", path);
                    try!(::std::fs::create_dir_all(path.parent().unwrap()));
                    let mut f = try!(File::create(path));

                    // debug!("[*]: Writing to {:?}", path);
                    try!(writeln!(f, "# {}", ch.name));
                }
            }
        }

        debug!("[*]: init done");
        Ok(())
    }

    pub fn create_gitignore(&self) {
        let gitignore = self.get_gitignore();

        if !gitignore.exists() {
            // Gitignore does not exist, create it

            // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`. If it
            // is not, `strip_prefix` will return an Error.
            if !self.get_dest().starts_with(&self.root) {
                return;
            }

            let relative = self.get_dest()
                               .strip_prefix(&self.root)
                               .expect("Destination is not relative to root.");
            let relative = relative.to_str()
                                   .expect("Path could not be yielded into a string slice.");

            debug!("[*]: {:?} does not exist, trying to create .gitignore", gitignore);

            let mut f = File::create(&gitignore).expect("Could not create file.");

            debug!("[*]: Writing to .gitignore");

            writeln!(f, "{}", relative).expect("Could not write to file.");
        }
    }

    /// The `build()` method is the one where everything happens. First it parses `SUMMARY.md` to
    /// construct the book's structure in the form of a `Vec<BookItem>` and then calls `render()`
    /// method of the current renderer.
    ///
    /// It is the renderer who generates all the output files.
    pub fn build(&mut self) -> Result<(), Box<Error>> {
        debug!("[fn]: build");

        try!(self.init());

        // Clean output directory
        try!(utils::fs::remove_dir_content(&self.dest));

        try!(self.renderer.render(&self));

        Ok(())
    }

    pub fn get_gitignore(&self) -> PathBuf {
        self.root.join(".gitignore")
    }

    pub fn copy_theme(&self) -> Result<(), Box<Error>> {
        debug!("[fn]: copy_theme");

        let theme_dir = self.src.join("theme");

        if !theme_dir.exists() {
            debug!("[*]: {:?} does not exist, trying to create directory", theme_dir);
            try!(fs::create_dir(&theme_dir));
        }

        // index.hbs
        let mut index = try!(File::create(&theme_dir.join("index.hbs")));
        try!(index.write_all(theme::INDEX));

        // book.css
        let mut css = try!(File::create(&theme_dir.join("book.css")));
        try!(css.write_all(theme::CSS));

        // favicon.png
        let mut favicon = try!(File::create(&theme_dir.join("favicon.png")));
        try!(favicon.write_all(theme::FAVICON));

        // book.js
        let mut js = try!(File::create(&theme_dir.join("book.js")));
        try!(js.write_all(theme::JS));

        // highlight.css
        let mut highlight_css = try!(File::create(&theme_dir.join("highlight.css")));
        try!(highlight_css.write_all(theme::HIGHLIGHT_CSS));

        // highlight.js
        let mut highlight_js = try!(File::create(&theme_dir.join("highlight.js")));
        try!(highlight_js.write_all(theme::HIGHLIGHT_JS));

        Ok(())
    }

    pub fn write_file<P: AsRef<Path>>(&self, filename: P, content: &[u8]) -> Result<(), Box<Error>> {
        let path = self.get_dest().join(filename);
        try!(utils::fs::create_file(&path).and_then(|mut file| {
            file.write_all(content)
        }).map_err(|e| {
            io::Error::new(io::ErrorKind::Other, format!("Could not create {}: {}", path.display(), e))
        }));
        Ok(())
    }

    /// Parses the `book.json` file (if it exists) to extract the configuration parameters.
    /// The `book.json` file should be in the root directory of the book.
    /// The root directory is the one specified when creating a new `MDBook`

    pub fn read_config(mut self) -> Self {

        let config = BookConfig::new(&self.root)
                         .read_config(&self.root)
                         .to_owned();

        self.title = config.title;
        self.description = config.description;
        self.author = config.author;

        self.dest = config.dest;
        self.src = config.src;
        self.theme_path = config.theme_path;

        self
    }

    /// You can change the default renderer to another one by using this method. The only requirement
    /// is for your renderer to implement the [Renderer trait](../../renderer/renderer/trait.Renderer.html)
    ///
    /// ```no_run
    /// extern crate mdbook;
    /// use mdbook::MDBook;
    /// use mdbook::renderer::HtmlHandlebars;
    /// # use std::path::Path;
    ///
    /// fn main() {
    ///     let mut book = MDBook::new(Path::new("mybook"))
    ///                         .set_renderer(Box::new(HtmlHandlebars::new()));
    ///
    ///     // In this example we replace the default renderer by the default renderer...
    ///     // Don't forget to put your renderer in a Box
    /// }
    /// ```
    ///
    /// **note:** Don't forget to put your renderer in a `Box` before passing it to `set_renderer()`

    pub fn set_renderer(mut self, renderer: Box<Renderer>) -> Self {
        self.renderer = renderer;
        self
    }

    pub fn test(&mut self) -> Result<(), Box<Error>> {
        // read in the chapters
        try!(self.parse_summary());
        for item in self.iter() {

            if let BookItem::Chapter(_, ref ch) = *item {
                if ch.path != PathBuf::new() {

                    let path = self.get_src().join(&ch.path);

                    println!("[*]: Testing file: {:?}", path);

                    let output_result = Command::new("rustdoc")
                                            .arg(&path)
                                            .arg("--test")
                                            .output();
                    let output = try!(output_result);

                    if !output.status.success() {
                        return Err(Box::new(io::Error::new(ErrorKind::Other, format!(
                                        "{}\n{}",
                                        String::from_utf8_lossy(&output.stdout),
                                        String::from_utf8_lossy(&output.stderr)))) as Box<Error>);
                    }
                }
            }
        }
        Ok(())
    }

    pub fn get_root(&self) -> &Path {
        &self.root
    }

    pub fn set_dest(mut self, dest: &Path) -> Self {

        // Handle absolute and relative paths
        if dest.is_absolute() {
            self.dest = dest.to_owned();
        } else {
            let dest = self.root.join(dest).to_owned();
            self.dest = dest;
        }

        self
    }

    pub fn get_dest(&self) -> &Path {
        &self.dest
    }

    pub fn set_src(mut self, src: &Path) -> Self {

        // Handle absolute and relative paths
        if src.is_absolute() {
            self.src = src.to_owned();
        } else {
            let src = self.root.join(src).to_owned();
            self.src = src;
        }

        self
    }

    pub fn get_src(&self) -> &Path {
        &self.src
    }

    pub fn set_title(mut self, title: &str) -> Self {
        self.title = title.to_owned();
        self
    }

    pub fn get_title(&self) -> &str {
        &self.title
    }

    pub fn set_author(mut self, author: &str) -> Self {
        self.author = author.to_owned();
        self
    }

    pub fn get_author(&self) -> &str {
        &self.author
    }

    pub fn set_description(mut self, description: &str) -> Self {
        self.description = description.to_owned();
        self
    }

    pub fn get_description(&self) -> &str {
        &self.description
    }

    pub fn set_livereload(&mut self, livereload: String) -> &mut Self {
        self.livereload = Some(livereload);
        self
    }

    pub fn unset_livereload(&mut self) -> &Self {
        self.livereload = None;
        self
    }

    pub fn get_livereload(&self) -> Option<&String> {
        self.livereload.as_ref()
    }

    pub fn set_theme_path(mut self, theme_path: &Path) -> Self {
        self.theme_path = if theme_path.is_absolute() {
            theme_path.to_owned()
        } else {
            self.root.join(theme_path).to_owned()
        };
        self
    }

    pub fn get_theme_path(&self) -> &Path {
        &self.theme_path
    }

    // Construct book
    fn parse_summary(&mut self) -> Result<(), Box<Error>> {
        // When append becomes stable, use self.content.append() ...
        self.content = try!(parse::construct_bookitems(&self.src.join("SUMMARY.md")));
        Ok(())
    }
}
<|endoftext|>


<commit_before>use crate::exercise::{Exercise, ExerciseList};
use crate::run::run;
use crate::verify::verify;
use clap::{crate_version, App, Arg, SubCommand};
use notify::DebouncedEvent;
use notify::{RecommendedWatcher, RecursiveMode, Watcher};
use std::ffi::OsStr;
use std::fs;
use std::path::Path;
use std::sync::mpsc::channel;
use std::time::Duration;

mod exercise;
mod run;
mod verify;

fn main() {
    let matches = App::new("rustlings")
        .version(crate_version!())
        .author("Olivia Hugger, Carol Nichols")
        .about("Rustlings is a collection of small exercises to get you used to writing and reading Rust code")
        .subcommand(SubCommand::with_name("verify").alias("v").about("Verifies all exercises according to the recommended order"))
        .subcommand(SubCommand::with_name("watch").alias("w").about("Reruns `verify` when files were edited"))
        .subcommand(
            SubCommand::with_name("run")
                .alias("r")
                .about("Runs/Tests a single exercise")
                .arg(Arg::with_name("file").required(true).index(1))
                .arg(Arg::with_name("test").short("t").long("test").help("Run the file as a test")),
        )
        .get_matches();

    if None == matches.subcommand_name() {
        println!();
        println!(r#"       welcome to...                      "#);
        println!(r#"                 _   _ _                  "#);
        println!(r#"  _ __ _   _ ___| |_| (_)_ __   __ _ ___  "#);
        println!(r#" | '__| | | / __| __| | | '_ \ / _` / __| "#);
        println!(r#" | |  | |_| \__ \ |_| | | | | | (_| \__ \ "#);
        println!(r#" |_|   \__,_|___/\__|_|_|_| |_|\__, |___/ "#);
        println!(r#"                               |___/      "#);
        println!();
    }

    if !Path::new("info.toml").exists() {
        println!(
            "{} must be run from the rustlings directory",
            std::env::current_exe().unwrap().to_str().unwrap()
        );
        println!("Try `cd rustlings/`!");
        std::process::exit(1);
    }

    let toml_str = &fs::read_to_string("info.toml").unwrap();
    let exercises = toml::from_str::<ExerciseList>(toml_str).unwrap().exercises;

    if let Some(ref matches) = matches.subcommand_matches("run") {
        let filename = matches.value_of("file").unwrap_or_else(|| {
            println!("Please supply a file name!");
            std::process::exit(1);
        });

        let matching_exercise = |e: &&Exercise| {
            Path::new(filename)
                .canonicalize()
                .map(|p| p.ends_with(&e.path))
                .unwrap_or(false)
        };

        let exercise = exercises.iter().find(matching_exercise).unwrap_or_else(|| {
            println!("No exercise found for your file name!");
            std::process::exit(1)
        });

        run(&exercise).unwrap_or_else(|_| std::process::exit(1));
    }

    if matches.subcommand_matches("verify").is_some() {
        verify(&exercises).unwrap_or_else(|_| std::process::exit(1));
    }

    if matches.subcommand_matches("watch").is_some() {
        watch(&exercises).unwrap();
    }

    if matches.subcommand_name().is_none() {
        let text = fs::read_to_string("default_out.txt").unwrap();
        println!("{}", text);
    }
}

fn watch(exercises: &[Exercise]) -> notify::Result<()> {
    let (tx, rx) = channel();

    let mut watcher: RecommendedWatcher = Watcher::new(tx, Duration::from_secs(2))?;
    watcher.watch(Path::new("./exercises"), RecursiveMode::Recursive)?;

    let _ignored = verify(exercises.iter());

    loop {
        match rx.recv() {
            Ok(event) => match event {
                DebouncedEvent::Create(b) | DebouncedEvent::Chmod(b) | DebouncedEvent::Write(b) => {
                    if b.extension() == Some(OsStr::new("rs")) {
                        println!("----------**********----------\n");
                        let filepath = b.as_path().canonicalize().unwrap();
                        let exercise = exercises
                            .iter()
                            .skip_while(|e| !filepath.ends_with(&e.path));
                        let _ignored = verify(exercise);
                    }
                }
                _ => {}
            },
            Err(e) => println!("watch error: {:?}", e),
        }
    }
}
<commit_msg>Check if changed exercise file exists before calling verify.<commit_after>use crate::exercise::{Exercise, ExerciseList};
use crate::run::run;
use crate::verify::verify;
use clap::{crate_version, App, Arg, SubCommand};
use notify::DebouncedEvent;
use notify::{RecommendedWatcher, RecursiveMode, Watcher};
use std::ffi::OsStr;
use std::fs;
use std::path::Path;
use std::sync::mpsc::channel;
use std::time::Duration;

mod exercise;
mod run;
mod verify;

fn main() {
    let matches = App::new("rustlings")
        .version(crate_version!())
        .author("Olivia Hugger, Carol Nichols")
        .about("Rustlings is a collection of small exercises to get you used to writing and reading Rust code")
        .subcommand(SubCommand::with_name("verify").alias("v").about("Verifies all exercises according to the recommended order"))
        .subcommand(SubCommand::with_name("watch").alias("w").about("Reruns `verify` when files were edited"))
        .subcommand(
            SubCommand::with_name("run")
                .alias("r")
                .about("Runs/Tests a single exercise")
                .arg(Arg::with_name("file").required(true).index(1))
                .arg(Arg::with_name("test").short("t").long("test").help("Run the file as a test")),
        )
        .get_matches();

    if None == matches.subcommand_name() {
        println!();
        println!(r#"       welcome to...                      "#);
        println!(r#"                 _   _ _                  "#);
        println!(r#"  _ __ _   _ ___| |_| (_)_ __   __ _ ___  "#);
        println!(r#" | '__| | | / __| __| | | '_ \ / _` / __| "#);
        println!(r#" | |  | |_| \__ \ |_| | | | | | (_| \__ \ "#);
        println!(r#" |_|   \__,_|___/\__|_|_|_| |_|\__, |___/ "#);
        println!(r#"                               |___/      "#);
        println!();
    }

    if !Path::new("info.toml").exists() {
        println!(
            "{} must be run from the rustlings directory",
            std::env::current_exe().unwrap().to_str().unwrap()
        );
        println!("Try `cd rustlings/`!");
        std::process::exit(1);
    }

    let toml_str = &fs::read_to_string("info.toml").unwrap();
    let exercises = toml::from_str::<ExerciseList>(toml_str).unwrap().exercises;

    if let Some(ref matches) = matches.subcommand_matches("run") {
        let filename = matches.value_of("file").unwrap_or_else(|| {
            println!("Please supply a file name!");
            std::process::exit(1);
        });

        let matching_exercise = |e: &&Exercise| {
            Path::new(filename)
                .canonicalize()
                .map(|p| p.ends_with(&e.path))
                .unwrap_or(false)
        };

        let exercise = exercises.iter().find(matching_exercise).unwrap_or_else(|| {
            println!("No exercise found for your file name!");
            std::process::exit(1)
        });

        run(&exercise).unwrap_or_else(|_| std::process::exit(1));
    }

    if matches.subcommand_matches("verify").is_some() {
        verify(&exercises).unwrap_or_else(|_| std::process::exit(1));
    }

    if matches.subcommand_matches("watch").is_some() {
        watch(&exercises).unwrap();
    }

    if matches.subcommand_name().is_none() {
        let text = fs::read_to_string("default_out.txt").unwrap();
        println!("{}", text);
    }
}

fn watch(exercises: &[Exercise]) -> notify::Result<()> {
    let (tx, rx) = channel();

    let mut watcher: RecommendedWatcher = Watcher::new(tx, Duration::from_secs(2))?;
    watcher.watch(Path::new("./exercises"), RecursiveMode::Recursive)?;

    let _ignored = verify(exercises.iter());

    loop {
        match rx.recv() {
            Ok(event) => match event {
                DebouncedEvent::Create(b) | DebouncedEvent::Chmod(b) | DebouncedEvent::Write(b) => {
                    if b.extension() == Some(OsStr::new("rs")) && b.exists() {
                        println!("----------**********----------\n");
                        let filepath = b.as_path().canonicalize().unwrap();
                        let exercise = exercises
                            .iter()
                            .skip_while(|e| !filepath.ends_with(&e.path));
                        let _ignored = verify(exercise);
                    }
                }
                _ => {}
            },
            Err(e) => println!("watch error: {:?}", e),
        }
    }
}
<|endoftext|>