<commit_before>extern crate rand;

use std::string::String;
use std::collections::HashMap;
use std::vec::Vec;
use std::fs::{File};
use std::env;
use std::path::{Path};
use std::io::{self, BufReader, BufRead};
use std::iter;

use rand::{random, Closed01, Rng};
use rand::distributions::{IndependentSample, Range, Sample, Gamma};

type Bag = HashMap<usize, usize>;

fn load_bags<P: AsRef<Path>>(path: P) -> io::Result<(Vec<usize>, Vec<Bag>)> {
    let mut bags = Vec::new();
    let mut labels = Vec::new();
    let file = try!(File::open(path));
    let file = BufReader::new(file);
    for line in file.lines() {
        let line = line.unwrap();
        let mut iter = line.split_whitespace();
        let label = iter.next().unwrap().parse::<usize>().unwrap();
        let mut bag = Bag::new();
        for elm in iter {
            let mut iter = elm.split(':');
            let index = iter.next().unwrap().parse::<usize>().unwrap();
            let value = iter.next().unwrap().parse::<usize>().unwrap();
            bag.insert(index - 1, value);
        }
        labels.push(label);
        bags.push(bag);
    }
    Ok((labels, bags))
}

fn load_text_vocabulary<P: AsRef<Path>>(path: P) -> io::Result<Vec<String>> {
    let mut vocab = Vec::new();
    let file = try!(File::open(path));
    let file = BufReader::new(file);
    for line in file.lines() {
        let line = line.unwrap();
        vocab.push(line);
    }
    Ok(vocab)
}

#[derive(Clone)]
struct Categorical {
    prop: Vec<f64>,
}

impl Categorical {
    fn new(weights: Vec<f64>) -> Categorical {
        let sum: f64 = weights.iter().sum();
        let prop: Vec<f64> = weights.iter().map(|w| w / sum).collect();
        Categorical {
            prop: prop,
        }
    }
}

impl Sample<usize> for Categorical {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> usize {
        self.ind_sample(rng)
    }
}

impl IndependentSample<usize> for Categorical {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> usize {
        let Closed01(x) = random::<Closed01<f64>>();

        let mut sum = 0.0;
        let mut ret = self.prop.len() - 1;
        for (k, &p) in self.prop.iter().enumerate() {
            sum += p;
            if x <= sum {
                ret = k;
                break;
            }
        }
        ret
    }
}

#[derive(Clone)]
struct Dirichlet {
    alpha: Vec<f64>,
}

impl Dirichlet {
    fn new(alpha: Vec<f64>) -> Dirichlet {
        Dirichlet {
            alpha: alpha,
        }
    }
}

impl Sample<Vec<f64>> for Dirichlet {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Vec<f64> {
        self.ind_sample(rng)
    }
}

impl IndependentSample<Vec<f64>> for Dirichlet {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Vec<f64> {
        let mut sum = 0.0;
        let mut xs = Vec::new();
        for &a in &self.alpha {
            let gamma = Gamma::new(a, 1.0);
            let y = gamma.ind_sample(rng);
            xs.push(y);
            sum += y;
        }
        for x in xs.iter_mut() {
            *x /= sum;
        }
        xs
    }
}

fn lda(dataset: Vec<Bag>, num_topics: usize, alpha: Vec<f64>, beta: Vec<f64>, burn_in: usize, num_samples: usize) {
    // Initialization
    let mut rng = rand::thread_rng();
    let topics = Range::new(0, num_topics);

    let vocab_size: usize = {
        let mut have_some = false;
        let mut max_index = 0;
        for bag in &dataset {
            for &index in bag.keys() {
                have_some = true;
                if index > max_index {
                    max_index = index;
                }
            }
        }
        if have_some {
            max_index + 1
        }
        else {
            0
        }
    };
    println!("K = {}", num_topics);
    println!("M = {}", &dataset.len());
    println!("V = {}", vocab_size);

    let mut w: Vec<Vec<usize>> = Vec::new();
    let mut z: Vec<Vec<usize>> = Vec::new();
    let mut z_samples: Vec<Vec<Vec<usize>>> = Vec::new();
    let mut theta: Vec<Vec<f64>> = Vec::new();
    let mut ndk: Vec<Vec<usize>> = Vec::new();
    for bag in &dataset {
        // w
        let mut w_d = Vec::new();
        for (&index, &count) in bag {
            for _ in 0..count {
                w_d.push(index);
            }
        }
        w.push(w_d);
        // z
        let mut z_d = Vec::new();
        let mut z_samples_d = Vec::new();
        for &count in bag.values() {
            for _ in 0..count {
                z_d.push(0);
                let zeros: Vec<usize> = iter::repeat(0).take(num_topics).collect();
                z_samples_d.push(zeros);
            }
        }
        let n_d = z_d.len();
        z.push(z_d);
        z_samples.push(z_samples_d);
        // theta
        let mut theta_d = Vec::with_capacity(num_topics);
        for _k in 0..num_topics {
            theta_d.push(1.0 / num_topics as f64);
        }
        theta.push(theta_d);
        // ndk
        let mut ndk_d = Vec::with_capacity(num_topics);
        ndk_d.push(n_d);
        for _k in 1..num_topics {
            ndk_d.push(0);
        }
        ndk.push(ndk_d);
    }
    println!("z = {:?}", z);
    println!("theta = {:?}", theta);
    println!("ndk = {:?}", ndk);

    let mut phi: Vec<Vec<f64>> = Vec::new();
    let mut nkv: Vec<Vec<usize>> = Vec::new();
    // phi
    for _k in 0..num_topics {
        let mut phi_k = Vec::new();
        for _v in 0..vocab_size {
            phi_k.push(1.0 / vocab_size as f64);
        }
        phi.push(phi_k);
    }
    // nkv
    let mut nkv_0 = Vec::new();
    for _v in 0..vocab_size {
        nkv_0.push(0);
    }
    for bag in &dataset {
        for (&index, &count) in bag {
            nkv_0[index] += count;
        }
    }
    nkv.push(nkv_0);
    for _k in 1..num_topics {
        let mut nkv_k = Vec::new();
        for _v in 0..vocab_size {
            nkv_k.push(0);
        }
        nkv.push(nkv_k);
    }
    println!("phi = {:?}", phi);
    println!("nkv = {:?}", nkv);

    // Sampling
    let mut rng = rand::thread_rng();
    for s in 0..(burn_in + num_samples) {
        for (d, w_d) in w.iter().enumerate() {
            for (i, &w_di) in w_d.iter().enumerate() {
                let v = w_di;
                // Sample z_di
                let mut weights: Vec<f64> = Vec::new();
                for k in 0..num_topics {
                    weights.push(theta[d][k] * phi[k][v]);
                }
                let cat = Categorical::new(weights);
                let old_z_di = z[d][i];
                let new_z_di = cat.ind_sample(&mut rng);
                z[d][i] = new_z_di;
                if s >= burn_in {
                    z_samples[d][i][new_z_di] += 1;
                }
                // Update ndk and nkv
                ndk[d][old_z_di] -= 1;
                ndk[d][new_z_di] += 1;
                nkv[old_z_di][v] -= 1;
                nkv[new_z_di][v] += 1;
            }
            // Sample theta_d
            let mut alpha_d: Vec<f64> = Vec::new();
            for k in 0..num_topics {
                alpha_d.push(ndk[d][k] as f64 + alpha[k]);
            }
            let dir = Dirichlet::new(alpha_d);
            theta[d] = dir.ind_sample(&mut rng);
        }
        for k in 0..num_topics {
            // Sample phi_k
            let mut beta_k: Vec<f64> = Vec::new();
            for v in 0..vocab_size {
                beta_k.push(nkv[k][v] as f64 + beta[v]);
            }
            let dir = Dirichlet::new(beta_k);
            phi[k] = dir.ind_sample(&mut rng);
        }
        // Evaluate the log-likelihood value for the current parameters
        let mut log_likelihood = 0.0;
        for k in 0..num_topics {
            for v in 0..vocab_size {
                log_likelihood += nkv[k][v] as f64 * f64::ln(phi[k][v]);
            }
        }
        println!("log_likelihood = {}", log_likelihood);
    }
    for (d, z_samples_d) in z_samples.iter().enumerate() {
        for (i, samples) in z_samples_d.iter().enumerate() {
            let mut sum = 0.0;
            for k in 0..num_topics {
                sum += samples[k] as f64;
            }
            print!("z[{:3}][{:3}] = ", d, i);
            for k in 0..num_topics {
                print!("{:3}:{:.2} ", k, samples[k] as f64 / sum);
            }
            println!("");
        }
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let (labels, dataset) = load_bags(&args[1]).unwrap();
    let vocab: Option<Vec<String>> = if args.len() < 3 {
        None
    }
    else {
        Some(load_text_vocabulary(&args[2]).unwrap())
    };
    println!("{:?}", labels);
    println!("{:?}", dataset);
    println!("{:?}", vocab);

    let num_topics = 3;
    let alpha: Vec<f64> = iter::repeat(1.0).take(num_topics).collect();
    let beta: Vec<f64> = iter::repeat(1.0).take(vocab.unwrap().len()).collect();

    lda(dataset, num_topics, alpha, beta, 1000, 10000);
}
<commit_msg>Add make_dataset() function<commit_after>extern crate rand;

use std::string::String;
use std::collections::HashMap;
use std::vec::Vec;
use std::fs::{File};
use std::env;
use std::path::{Path};
use std::io::{self, BufReader, BufRead};
use std::iter;

use rand::{random, Closed01, Rng};
use rand::distributions::{IndependentSample, Range, Sample, Gamma, LogNormal};

type Bag = HashMap<usize, usize>;

fn load_bags<P: AsRef<Path>>(path: P) -> io::Result<(Vec<usize>, Vec<Bag>)> {
    let mut bags = Vec::new();
    let mut labels = Vec::new();
    let file = try!(File::open(path));
    let file = BufReader::new(file);
    for line in file.lines() {
        let line = line.unwrap();
        let mut iter = line.split_whitespace();
        let label = iter.next().unwrap().parse::<usize>().unwrap();
        let mut bag = Bag::new();
        for elm in iter {
            let mut iter = elm.split(':');
            let index = iter.next().unwrap().parse::<usize>().unwrap();
            let value = iter.next().unwrap().parse::<usize>().unwrap();
            bag.insert(index - 1, value);
        }
        labels.push(label);
        bags.push(bag);
    }
    Ok((labels, bags))
}

fn load_text_vocabulary<P: AsRef<Path>>(path: P) -> io::Result<Vec<String>> {
    let mut vocab = Vec::new();
    let file = try!(File::open(path));
    let file = BufReader::new(file);
    for line in file.lines() {
        let line = line.unwrap();
        vocab.push(line);
    }
    Ok(vocab)
}

#[derive(Clone)]
struct Categorical {
    prop: Vec<f64>,
}

impl Categorical {
    fn new(weights: Vec<f64>) -> Categorical {
        let sum: f64 = weights.iter().sum();
        let prop: Vec<f64> = weights.iter().map(|w| w / sum).collect();
        Categorical {
            prop: prop,
        }
    }
}

impl Sample<usize> for Categorical {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> usize {
        self.ind_sample(rng)
    }
}

impl IndependentSample<usize> for Categorical {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> usize {
        let Closed01(x) = random::<Closed01<f64>>();

        let mut sum = 0.0;
        let mut ret = self.prop.len() - 1;
        for (k, &p) in self.prop.iter().enumerate() {
            sum += p;
            if x <= sum {
                ret = k;
                break;
            }
        }
        ret
    }
}

#[derive(Clone)]
struct Dirichlet {
    alpha: Vec<f64>,
}

impl Dirichlet {
    fn new(alpha: Vec<f64>) -> Dirichlet {
        Dirichlet {
            alpha: alpha,
        }
    }
}

impl Sample<Vec<f64>> for Dirichlet {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Vec<f64> {
        self.ind_sample(rng)
    }
}

impl IndependentSample<Vec<f64>> for Dirichlet {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Vec<f64> {
        let mut sum = 0.0;
        let mut xs = Vec::new();
        for &a in &self.alpha {
            let gamma = Gamma::new(a, 1.0);
            let y = gamma.ind_sample(rng);
            xs.push(y);
            sum += y;
        }
        for x in xs.iter_mut() {
            *x /= sum;
        }
        xs
    }
}

fn lda(dataset: Vec<Bag>, num_topics: usize, alpha: Vec<f64>, beta: Vec<f64>, burn_in: usize, num_samples: usize) {
    // Initialization
    let mut rng = rand::thread_rng();
    let topics = Range::new(0, num_topics);

    let vocab_size: usize = {
        let mut have_some = false;
        let mut max_index = 0;
        for bag in &dataset {
            for &index in bag.keys() {
                have_some = true;
                if index > max_index {
                    max_index = index;
                }
            }
        }
        if have_some {
            max_index + 1
        }
        else {
            0
        }
    };
    println!("K = {}", num_topics);
    println!("M = {}", &dataset.len());
    println!("V = {}", vocab_size);

    let mut w: Vec<Vec<usize>> = Vec::new();
    let mut z: Vec<Vec<usize>> = Vec::new();
    let mut z_samples: Vec<Vec<Vec<usize>>> = Vec::new();
    let mut theta: Vec<Vec<f64>> = Vec::new();
    let mut ndk: Vec<Vec<usize>> = Vec::new();
    for bag in &dataset {
        // w
        let mut w_d = Vec::new();
        for (&index, &count) in bag {
            for _ in 0..count {
                w_d.push(index);
            }
        }
        w.push(w_d);
        // z
        let mut z_d = Vec::new();
        let mut z_samples_d = Vec::new();
        for &count in bag.values() {
            for _ in 0..count {
                z_d.push(0);
                let zeros: Vec<usize> = iter::repeat(0).take(num_topics).collect();
                z_samples_d.push(zeros);
            }
        }
        let n_d = z_d.len();
        z.push(z_d);
        z_samples.push(z_samples_d);
        // theta
        let mut theta_d = Vec::with_capacity(num_topics);
        for _k in 0..num_topics {
            theta_d.push(1.0 / num_topics as f64);
        }
        theta.push(theta_d);
        // ndk
        let mut ndk_d = Vec::with_capacity(num_topics);
        ndk_d.push(n_d);
        for _k in 1..num_topics {
            ndk_d.push(0);
        }
        ndk.push(ndk_d);
    }
    println!("z = {:?}", z);
    println!("theta = {:?}", theta);
    println!("ndk = {:?}", ndk);

    let mut phi: Vec<Vec<f64>> = Vec::new();
    let mut nkv: Vec<Vec<usize>> = Vec::new();
    // phi
    for _k in 0..num_topics {
        let mut phi_k = Vec::new();
        for _v in 0..vocab_size {
            phi_k.push(1.0 / vocab_size as f64);
        }
        phi.push(phi_k);
    }
    // nkv
    let mut nkv_0 = Vec::new();
    for _v in 0..vocab_size {
        nkv_0.push(0);
    }
    for bag in &dataset {
        for (&index, &count) in bag {
            nkv_0[index] += count;
        }
    }
    nkv.push(nkv_0);
    for _k in 1..num_topics {
        let mut nkv_k = Vec::new();
        for _v in 0..vocab_size {
            nkv_k.push(0);
        }
        nkv.push(nkv_k);
    }
    println!("phi = {:?}", phi);
    println!("nkv = {:?}", nkv);

    // Sampling
    let mut rng = rand::thread_rng();
    for s in 0..(burn_in + num_samples) {
        for (d, w_d) in w.iter().enumerate() {
            for (i, &w_di) in w_d.iter().enumerate() {
                let v = w_di;
                // Sample z_di
                let mut weights: Vec<f64> = Vec::new();
                for k in 0..num_topics {
                    weights.push(theta[d][k] * phi[k][v]);
                }
                let cat = Categorical::new(weights);
                let old_z_di = z[d][i];
                let new_z_di = cat.ind_sample(&mut rng);
                z[d][i] = new_z_di;
                if s >= burn_in {
                    z_samples[d][i][new_z_di] += 1;
                }
                // Update ndk and nkv
                ndk[d][old_z_di] -= 1;
                ndk[d][new_z_di] += 1;
                nkv[old_z_di][v] -= 1;
                nkv[new_z_di][v] += 1;
            }
            // Sample theta_d
            let mut alpha_d: Vec<f64> = Vec::new();
            for k in 0..num_topics {
                alpha_d.push(ndk[d][k] as f64 + alpha[k]);
            }
            let dir = Dirichlet::new(alpha_d);
            theta[d] = dir.ind_sample(&mut rng);
        }
        for k in 0..num_topics {
            // Sample phi_k
            let mut beta_k: Vec<f64> = Vec::new();
            for v in 0..vocab_size {
                beta_k.push(nkv[k][v] as f64 + beta[v]);
            }
            let dir = Dirichlet::new(beta_k);
            phi[k] = dir.ind_sample(&mut rng);
        }
        // Evaluate the log-likelihood value for the current parameters
        let mut log_likelihood = 0.0;
        for k in 0..num_topics {
            for v in 0..vocab_size {
                log_likelihood += nkv[k][v] as f64 * f64::ln(phi[k][v]);
            }
        }
        println!("log_likelihood = {}", log_likelihood);
    }
    for (d, z_samples_d) in z_samples.iter().enumerate() {
        for (i, samples) in z_samples_d.iter().enumerate() {
            let mut sum = 0.0;
            for k in 0..num_topics {
                sum += samples[k] as f64;
            }
            print!("z[{:3}][{:3}] = ", d, i);
            for k in 0..num_topics {
                print!("{:3}:{:.2} ", k, samples[k] as f64 / sum);
            }
            println!("");
        }
    }
}

fn make_dataset(num_docs: usize, vocab_size: usize, num_topics: usize, alpha: Vec<f64>, beta: Vec<f64>) -> Vec<Bag> {
    let mut rng = rand::thread_rng();
    // phi
    let dir_beta = Dirichlet::new(beta);
    let mut phi: Vec<Vec<f64>> = Vec::with_capacity(num_topics);
    for k in 0..num_topics {
        // Sample phi_k
        phi.push(dir_beta.ind_sample(&mut rng));
    }
    // theta, nd, z, w
    let dir_alpha = Dirichlet::new(alpha);
    let lognorm = LogNormal::new(100.0, 30.0);
    let mut theta: Vec<Vec<f64>> = Vec::with_capacity(num_docs);
    let mut nd: Vec<usize> = Vec::with_capacity(num_docs);
    let mut z: Vec<Vec<usize>> = Vec::with_capacity(num_docs);
    let mut w: Vec<Vec<usize>> = Vec::with_capacity(num_docs);
    for d in 0..num_docs {
        theta.push(dir_alpha.ind_sample(&mut rng));
        nd.push(f64::floor(lognorm.ind_sample(&mut rng)) as usize);
        let cat_theta = Categorical::new(theta[d].clone());
        let mut z_d = Vec::with_capacity(nd[d]);
        let mut w_d = Vec::with_capacity(nd[d]);
        for i in 0..nd[d] {
            z_d.push(cat_theta.ind_sample(&mut rng));
            let cat_phi = Categorical::new(phi[z_d[i]].clone());
            w_d.push(cat_phi.ind_sample(&mut rng));
        }
        z.push(z_d);
        w.push(w_d);
    }
    // Make bags
    let mut bags: Vec<Bag> = Vec::new();
    for w_d in w {
        let mut bag: Bag = Bag::new();
        for v in w_d {
            let counter = bag.entry(v).or_insert(0);
            *counter += 1;
        }
        bags.push(bag);
    }
    bags
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let (labels, dataset) = load_bags(&args[1]).unwrap();
    let vocab: Option<Vec<String>> = if args.len() < 3 {
        None
    }
    else {
        Some(load_text_vocabulary(&args[2]).unwrap())
    };
    println!("{:?}", labels);
    println!("{:?}", dataset);
    println!("{:?}", vocab);

    let num_topics = 3;
    let alpha: Vec<f64> = iter::repeat(1.0).take(num_topics).collect();
    let beta: Vec<f64> = iter::repeat(1.0).take(vocab.unwrap().len()).collect();

    lda(dataset, num_topics, alpha, beta, 1000, 10000);
}
<|endoftext|>


<commit_before>extern crate raytracer;
extern crate rand;

use raytracer::vec3::Vec3;
use raytracer::ray::Ray;
use raytracer::hitable::*;
use raytracer::camera::Camera;

use rand::{thread_rng, Rng};

use std::io::{self, Write};

const NX: i32 = 200;
const NY: i32 = 100;
const NUM_SAMPLES: i32 = 64;
const MAX_DISTANCE: f64 = 1000.0;

fn random_in_unit_sphere() -> Vec3 {
    let mut rng = thread_rng();
    loop {
        let p = Vec3::new(rng.next_f64(), rng.next_f64(), rng.next_f64()).mul_scalar(2.0).sub_scalar(1.0);
        if p.dot(p) >=  1.0 {
            return p
        }
    }
}

fn color<T: Hitable>(r: &Ray, world: &HitableList<T>) -> Vec3 {
    match world.hit(r, 0.0, MAX_DISTANCE) {
        Some(h) => {
            // Sphere color
            let target = h.p + h.normal + random_in_unit_sphere();
            color(&Ray::new(h.p, target-h.p), world).mul_scalar(0.5)
        }
        None => {
            // Background
            let unit_dir = r.dir.normalized();
            let t = 0.5 * (unit_dir.y + 1.0);
            Vec3::new(1.0,1.0,1.0).mul_scalar(1.0-t) + Vec3::new(0.5, 0.7, 1.0).mul_scalar(t)
        }
    }
}

fn main() {
    io::stdout().write_fmt(format_args!("P3\n{} {}\n{}\n", NX, NY, 255)).unwrap();

    let camera = Camera {
        lower_left_corner: Vec3::new(-2.0, -1.0, -1.0),
        horizontal: Vec3::new(4.0, 0.0, 0.0),
        vertical: Vec3::new(0.0, 2.0, 0.0),
        origin: Vec3::new(0.0, 0.0, 0.0),
    };

    let world: HitableList<Sphere> = HitableList {
        items: vec![
            Sphere { center: Vec3::new(0.0,0.0,-1.0), radius: 0.5 },
            Sphere { center: Vec3::new(0.0,-100.5,-1.0), radius: 100.0 },
        ]
    };

    let mut rng = thread_rng();

    for j in (0..NY).rev() {
        for i in 0..NX {
            let mut c = Vec3::new(0.0, 0.0, 0.0);

            for s in 0..NUM_SAMPLES {
                // Get percent offset from bottom left corner
                let u = (i as f64 + rng.next_f64()) / (NX as f64);
                let v = (j as f64 + rng.next_f64()) / (NY as f64);

                // Make ray
                let r = camera.get_ray(u, v);

                // Get color
                c += color(&r, &world);
            }
            let c = c.div_scalar(NUM_SAMPLES as f64);

            // Output for PPM
            let c = c.mul_scalar(255.99);
            println!("{} {} {}", c.x as i32, c.y as i32, c.z as i32);
        }
    }
}
<commit_msg>Chapter 7: gamma fix<commit_after>extern crate raytracer;
extern crate rand;

use raytracer::vec3::Vec3;
use raytracer::ray::Ray;
use raytracer::hitable::*;
use raytracer::camera::Camera;

use rand::{thread_rng, Rng};

use std::io::{self, Write};

const NX: i32 = 200;
const NY: i32 = 100;
const NUM_SAMPLES: i32 = 300;
const MAX_DISTANCE: f64 = 100.0;

fn random_in_unit_sphere() -> Vec3 {
    let mut rng = thread_rng();
    loop {
        let p = Vec3::new(rng.next_f64(), rng.next_f64(), rng.next_f64()).mul_scalar(2.0).sub_scalar(1.0);
        if p.dot(p) >=  1.0 {
            return p
        }
    }
}

fn color<T: Hitable>(r: &Ray, world: &HitableList<T>) -> Vec3 {
    match world.hit(r, 0.0, MAX_DISTANCE) {
        Some(h) => {
            // Sphere color
            let target = h.p + h.normal + random_in_unit_sphere();
            color(&Ray::new(h.p, target-h.p), world).mul_scalar(0.5)
        }
        None => {
            // Background
            let unit_dir = r.dir.normalized();
            let t = 0.5 * (unit_dir.y + 1.0);
            Vec3::new(1.0,1.0,1.0).mul_scalar(1.0-t) + Vec3::new(0.5, 0.7, 1.0).mul_scalar(t)
        }
    }
}

fn main() {
    io::stdout().write_fmt(format_args!("P3\n{} {}\n{}\n", NX, NY, 255)).unwrap();

    let camera = Camera {
        lower_left_corner: Vec3::new(-2.0, -1.0, -1.0),
        horizontal: Vec3::new(4.0, 0.0, 0.0),
        vertical: Vec3::new(0.0, 2.0, 0.0),
        origin: Vec3::new(0.0, 0.0, 0.0),
    };

    let world: HitableList<Sphere> = HitableList {
        items: vec![
            Sphere { center: Vec3::new(0.0,0.0,-1.0), radius: 0.5 },
            Sphere { center: Vec3::new(0.0,-100.5,-1.0), radius: 100.0 },
        ]
    };

    let mut rng = thread_rng();

    for j in (0..NY).rev() {
        for i in 0..NX {
            let mut c = Vec3::new(0.0, 0.0, 0.0);

            for s in 0..NUM_SAMPLES {
                // Get percent offset from bottom left corner
                let u = (i as f64 + rng.next_f64()) / (NX as f64);
                let v = (j as f64 + rng.next_f64()) / (NY as f64);

                // Make ray
                let r = camera.get_ray(u, v);

                // Get color
                c += color(&r, &world);
            }
            let c = c.div_scalar(NUM_SAMPLES as f64);

            // Output for PPM
            // Gamma correct to 2: output color ^ (1/gamma) = x^(1/2) = sqrt
            let c = c.map(&|x: f64| x.sqrt());
            let c = c.mul_scalar(255.99);
            println!("{} {} {}", c.x as i32, c.y as i32, c.z as i32);
        }
    }
}
<|endoftext|>


<commit_before>#![feature(io, slice_patterns)]
use std::fs;
use std::io;
use std::io::{Read, Write};
use std::path::Path;

const MEMSIZE: usize = 4096;

#[derive(Debug)]
enum ParserResult {
    Something(Op),
    Nothing,
    EOF
}

use ParserResult::*;

#[derive(Debug)]
enum Op {
    Add(u8),
    Mov(usize),
    In,
    Out,
    Loop(OpStream)
}

use Op::*;

#[derive(Debug)]
struct OpStream {
    ops: Vec<Op>
}

impl OpStream {
    fn add(&mut self, op: Op) {
        self.ops.push(op);
    }

    fn optimize(&mut self) {
        let mut i = 0;
        while i < self.ops.len() {
            match &self.ops[i..] {
                [Add(a), Add(b), ..] => {
                    self.ops[i] = Add(a.wrapping_add(b));
                    self.ops.remove(i + 1);
                }
                [Mov(a), Mov(b), ..] => {
                    self.ops[i] = Mov(a.wrapping_add(b));
                    self.ops.remove(i + 1);
                }
                [Add(0), ..] | [Mov(0), ..] => {
                    self.ops.remove(i);
                }
                [Loop(_), ..] => {
                    if let &mut Loop(ref mut stream) = &mut self.ops[i] {
                        stream.optimize();
                    }
                    i += 1
                }
                _ => i += 1
            }
        }
    }

    fn get(&self) -> &[Op] {
        &self.ops[..]
    }

    fn new() -> OpStream {
        OpStream { ops: Vec::new() }
    }
}

struct State {
    index: usize,
    memory: [u8; MEMSIZE]
}

impl State {
    fn new() -> State {
        State { index: 0, memory: [0; MEMSIZE] }
    }

    fn peek(&self) -> u8 {
        self.memory[self.index % MEMSIZE]
    }

    fn poke(&mut self, value: u8) {
        self.memory[self.index % MEMSIZE] = value;
    }

    fn step(&mut self, op: &Op) -> bool {
        match op {
            &Add(i) => {
                // XXX: Ugly expression due to lexical scoping of borrow,
                //      cf. https://github.com/rust-lang/rust/issues/6393
                let x = self.peek();
                self.poke(x.wrapping_add(i));
            }
            &Mov(n) => {
                self.index = self.index.wrapping_add(n);
            }
            &In => {
                let mut c = vec![0u8];
                if io::stdin().read(&mut c).unwrap() == 0 {
                    return false;
                }
                self.poke(c[0]);
            }
            &Out => {
                io::stdout().write(&vec![self.peek()]).unwrap();
            }
            &Loop(ref ops) => {
                while self.peek() != 0 {
                    if !self.run(ops.get()) {
                        return false;
                    }
                }
            }
        }
        return true;
    }

    fn run(&mut self, ops: &[Op]) -> bool {
        for op in ops {
            if !self.step(op) {
                return false;
            }
        }
        return true;
    }
}

fn main() {
    let filenames: Vec<String> = std::iter::FromIterator::from_iter(std::env::args());

    for filename in &filenames[1..] {
        match reader(&Path::new(filename)) {
            Ok(reader) => {
                let mut chars = reader.chars().peekable();

                let mut opstream = OpStream::new();
                loop {
                    match parse(&mut chars) {
                        Something(op) => opstream.add(op),
                        EOF => break,
                        _ => {}
                    }
                }

                opstream.optimize();

                State::new().run(opstream.get());
            }
            Err(e) => {
                writeln!(&mut io::stderr(), "Error while processing {}: {}", filename, e).unwrap();
            }
        }
    }
}

fn parse<T: io::Read>(mut chars: &mut std::iter::Peekable<io::Chars<T>>) -> ParserResult {
    match chars.next() {
        Some(Ok(c)) => match c {
            '+' => Something(Add(1)),
            '-' => Something(Add(-1)),
            '<' => Something(Mov(-1)),
            '>' => Something(Mov(1)),
            ',' => Something(In),
            '.' => Something(Out),
            '[' => {
                let mut childstream = OpStream::new();
                while match chars.peek() {
                    Some(&Ok(c)) if c != ']' => true,
                    _ => false
                } {
                    match parse(&mut chars) {
                        Something(op) => childstream.add(op),
                        Nothing => {},
                        EOF => panic!()
                    }
                }
                chars.next();
                Something(Loop(childstream))
            },
            _ => Nothing  // other characters
        },
        _ => EOF  // XXX: really?
    }
}

fn reader(path: &Path) -> Result<io::BufReader<fs::File>, io::Error> {
    Ok(io::BufReader::new(try!(fs::File::open(path))))
}

#[cfg(test)]
mod tests {
    use std::io;
    use std::io::Read;
    use super::{parse, ParserResult, State, OpStream};
    use super::Op::*;

    macro_rules! assert_let(
        ($test:pat, $value:expr, $then:block) => (
            if let $test = $value {
                $then
            } else {
                panic!("{:?} doesn't match {}", $value, stringify!($test))
            }
        );
        ($test:pat, $value:expr) => (
            assert_let!($test, $value, {})
        );
    );

    #[test]
    fn test_state_peek() {
        let mut state = State::new();
        state.memory[state.index] = 23;
        assert_eq!(23, state.peek());
        state.index = 5;
        state.memory[state.index] = 42;
        assert_eq!(42, state.peek());
    }

    #[test]
    fn test_state_poke() {
        let mut state = State::new();
        state.poke(23);
        assert_eq!(23, state.memory[state.index]);
        state.index = 5;
        state.poke(42);
        assert_eq!(42, state.memory[state.index]);
    }

    #[test]
    fn test_state_step_add() {
        let mut state = State::new();
        state.step(&Add(23));
        assert_eq!(23, state.peek());
        state.step(&Add(42));
        assert_eq!(65, state.peek());
        state.step(&Add(-66));
        assert_eq!(-1, state.peek());
    }

    #[test]
    fn test_state_step_mov() {
        let mut state = State::new();
        state.step(&Mov(1));
        assert_eq!(1, state.index);
        state.step(&Mov(42));
        assert_eq!(43, state.index);
        state.step(&Mov(-1));
        assert_eq!(42, state.index);
    }

    #[test]
    fn test_state_step_loop() {
        let mut state = State::new();
        state.poke(23);
        state.step(&Loop(OpStream { ops: vec![Add(1)] }));
        assert_eq!(0, state.peek());
    }

    #[test]
    fn test_opstream_new_empty() {
        let opstream = OpStream::new();
        assert_eq!(0, opstream.ops.len());
    }

    #[test]
    fn test_opstream_optimize() {
        let mut opstream = OpStream::new();
        opstream.add(Mov(1));
        opstream.add(Mov(1));
        opstream.add(Add(1));
        opstream.add(Add(-1));
        opstream.add(Add(-1));
        opstream.add(Mov(1));
        opstream.add(Mov(-1));
        let mut opstream2 = OpStream::new();
        opstream2.add(Mov(2));
        opstream2.add(Mov(3));
        opstream.add(Loop(opstream2));
        opstream.optimize();

        assert_let!([Mov(2), Add(-1), Loop(ref s2)], &opstream.ops[..], {
            assert_let!([Mov(5)], &s2.ops[..]);
        });
    }

    #[test]
    fn test_parse() {
        let input = b"+-[+.,]+";
        let mut chars = io::BufReader::new(&input[..]).chars().peekable();
        assert_let!(ParserResult::Something(Add(1)), parse(&mut chars));
        assert_let!(ParserResult::Something(Add(-1)), parse(&mut chars));
        assert_let!(ParserResult::Something(Loop(loop_op)), parse(&mut chars), {
            assert_let!([Add(1), Out, In], &loop_op.ops[..]);
        });
        assert_let!(ParserResult::Something(Add(1)), parse(&mut chars));
    }
}
<commit_msg>refactor: simplify pattern<commit_after>#![feature(io, slice_patterns)]
use std::fs;
use std::io;
use std::io::{Read, Write};
use std::path::Path;

const MEMSIZE: usize = 4096;

#[derive(Debug)]
enum ParserResult {
    Something(Op),
    Nothing,
    EOF
}

use ParserResult::*;

#[derive(Debug)]
enum Op {
    Add(u8),
    Mov(usize),
    In,
    Out,
    Loop(OpStream)
}

use Op::*;

#[derive(Debug)]
struct OpStream {
    ops: Vec<Op>
}

impl OpStream {
    fn add(&mut self, op: Op) {
        self.ops.push(op);
    }

    fn optimize(&mut self) {
        let mut i = 0;
        while i < self.ops.len() {
            match &self.ops[i..] {
                [Add(a), Add(b), ..] => {
                    self.ops[i] = Add(a.wrapping_add(b));
                    self.ops.remove(i + 1);
                }
                [Mov(a), Mov(b), ..] => {
                    self.ops[i] = Mov(a.wrapping_add(b));
                    self.ops.remove(i + 1);
                }
                [Add(0), ..] | [Mov(0), ..] => {
                    self.ops.remove(i);
                }
                [Loop(_), ..] => {
                    if let &mut Loop(ref mut stream) = &mut self.ops[i] {
                        stream.optimize();
                    }
                    i += 1
                }
                _ => i += 1
            }
        }
    }

    fn get(&self) -> &[Op] {
        &self.ops[..]
    }

    fn new() -> OpStream {
        OpStream { ops: Vec::new() }
    }
}

struct State {
    index: usize,
    memory: [u8; MEMSIZE]
}

impl State {
    fn new() -> State {
        State { index: 0, memory: [0; MEMSIZE] }
    }

    fn peek(&self) -> u8 {
        self.memory[self.index % MEMSIZE]
    }

    fn poke(&mut self, value: u8) {
        self.memory[self.index % MEMSIZE] = value;
    }

    fn step(&mut self, op: &Op) -> bool {
        match op {
            &Add(i) => {
                // XXX: Ugly expression due to lexical scoping of borrow,
                //      cf. https://github.com/rust-lang/rust/issues/6393
                let x = self.peek();
                self.poke(x.wrapping_add(i));
            }
            &Mov(n) => {
                self.index = self.index.wrapping_add(n);
            }
            &In => {
                let mut c = vec![0u8];
                if io::stdin().read(&mut c).unwrap() == 0 {
                    return false;
                }
                self.poke(c[0]);
            }
            &Out => {
                io::stdout().write(&vec![self.peek()]).unwrap();
            }
            &Loop(ref ops) => {
                while self.peek() != 0 {
                    if !self.run(ops.get()) {
                        return false;
                    }
                }
            }
        }
        return true;
    }

    fn run(&mut self, ops: &[Op]) -> bool {
        for op in ops {
            if !self.step(op) {
                return false;
            }
        }
        return true;
    }
}

fn main() {
    let filenames: Vec<String> = std::iter::FromIterator::from_iter(std::env::args());

    for filename in &filenames[1..] {
        match reader(&Path::new(filename)) {
            Ok(reader) => {
                let mut chars = reader.chars().peekable();

                let mut opstream = OpStream::new();
                loop {
                    match parse(&mut chars) {
                        Something(op) => opstream.add(op),
                        EOF => break,
                        _ => {}
                    }
                }

                opstream.optimize();

                State::new().run(opstream.get());
            }
            Err(e) => {
                writeln!(&mut io::stderr(), "Error while processing {}: {}", filename, e).unwrap();
            }
        }
    }
}

fn parse<T: io::Read>(mut chars: &mut std::iter::Peekable<io::Chars<T>>) -> ParserResult {
    match chars.next() {
        Some(Ok(c)) => match c {
            '+' => Something(Add(1)),
            '-' => Something(Add(-1)),
            '<' => Something(Mov(-1)),
            '>' => Something(Mov(1)),
            ',' => Something(In),
            '.' => Something(Out),
            '[' => {
                let mut childstream = OpStream::new();
                while match chars.peek() {
                    Some(&Ok(c)) => c != ']',
                    _ => false
                } {
                    match parse(&mut chars) {
                        Something(op) => childstream.add(op),
                        Nothing => {},
                        EOF => panic!()
                    }
                }
                chars.next();
                Something(Loop(childstream))
            },
            _ => Nothing  // other characters
        },
        _ => EOF  // XXX: really?
    }
}

fn reader(path: &Path) -> Result<io::BufReader<fs::File>, io::Error> {
    Ok(io::BufReader::new(try!(fs::File::open(path))))
}

#[cfg(test)]
mod tests {
    use std::io;
    use std::io::Read;
    use super::{parse, ParserResult, State, OpStream};
    use super::Op::*;

    macro_rules! assert_let(
        ($test:pat, $value:expr, $then:block) => (
            if let $test = $value {
                $then
            } else {
                panic!("{:?} doesn't match {}", $value, stringify!($test))
            }
        );
        ($test:pat, $value:expr) => (
            assert_let!($test, $value, {})
        );
    );

    #[test]
    fn test_state_peek() {
        let mut state = State::new();
        state.memory[state.index] = 23;
        assert_eq!(23, state.peek());
        state.index = 5;
        state.memory[state.index] = 42;
        assert_eq!(42, state.peek());
    }

    #[test]
    fn test_state_poke() {
        let mut state = State::new();
        state.poke(23);
        assert_eq!(23, state.memory[state.index]);
        state.index = 5;
        state.poke(42);
        assert_eq!(42, state.memory[state.index]);
    }

    #[test]
    fn test_state_step_add() {
        let mut state = State::new();
        state.step(&Add(23));
        assert_eq!(23, state.peek());
        state.step(&Add(42));
        assert_eq!(65, state.peek());
        state.step(&Add(-66));
        assert_eq!(-1, state.peek());
    }

    #[test]
    fn test_state_step_mov() {
        let mut state = State::new();
        state.step(&Mov(1));
        assert_eq!(1, state.index);
        state.step(&Mov(42));
        assert_eq!(43, state.index);
        state.step(&Mov(-1));
        assert_eq!(42, state.index);
    }

    #[test]
    fn test_state_step_loop() {
        let mut state = State::new();
        state.poke(23);
        state.step(&Loop(OpStream { ops: vec![Add(1)] }));
        assert_eq!(0, state.peek());
    }

    #[test]
    fn test_opstream_new_empty() {
        let opstream = OpStream::new();
        assert_eq!(0, opstream.ops.len());
    }

    #[test]
    fn test_opstream_optimize() {
        let mut opstream = OpStream::new();
        opstream.add(Mov(1));
        opstream.add(Mov(1));
        opstream.add(Add(1));
        opstream.add(Add(-1));
        opstream.add(Add(-1));
        opstream.add(Mov(1));
        opstream.add(Mov(-1));
        let mut opstream2 = OpStream::new();
        opstream2.add(Mov(2));
        opstream2.add(Mov(3));
        opstream.add(Loop(opstream2));
        opstream.optimize();

        assert_let!([Mov(2), Add(-1), Loop(ref s2)], &opstream.ops[..], {
            assert_let!([Mov(5)], &s2.ops[..]);
        });
    }

    #[test]
    fn test_parse() {
        let input = b"+-[+.,]+";
        let mut chars = io::BufReader::new(&input[..]).chars().peekable();
        assert_let!(ParserResult::Something(Add(1)), parse(&mut chars));
        assert_let!(ParserResult::Something(Add(-1)), parse(&mut chars));
        assert_let!(ParserResult::Something(Loop(loop_op)), parse(&mut chars), {
            assert_let!([Add(1), Out, In], &loop_op.ops[..]);
        });
        assert_let!(ParserResult::Something(Add(1)), parse(&mut chars));
    }
}
<|endoftext|>


<commit_before>#![deny(
    absolute_paths_not_starting_with_crate,
    anonymous_parameters,
    async_idents,
    const_err,
    dead_code,
    deprecated,
    duplicate_associated_type_bindings,
    duplicate_macro_exports,
    ellipsis_inclusive_range_patterns,
    exceeding_bitshifts,
    illegal_floating_point_literal_pattern,
    improper_ctypes,
    incoherent_fundamental_impls,
    intra_doc_link_resolution_failure,
    invalid_type_param_default,
    irrefutable_let_patterns,
    late_bound_lifetime_arguments,
    legacy_constructor_visibility,
    legacy_directory_ownership,
    missing_copy_implementations,
    missing_docs,
    missing_fragment_specifier,
    mutable_transmutes,
    no_mangle_const_items,
    no_mangle_generic_items,
    non_camel_case_types,
    non_shorthand_field_patterns,
    non_snake_case,
    non_upper_case_globals,
    overflowing_literals,
    parenthesized_params_in_types_and_modules,
    path_statements,
    patterns_in_fns_without_body,
    plugin_as_library,
    private_in_public,
    private_no_mangle_fns,
    private_no_mangle_statics,
    proc_macro_derive_resolution_fallback,
    pub_use_of_private_extern_crate,
    question_mark_macro_sep,
    safe_extern_statics,
    safe_packed_borrows,
    stable_features,
    trivial_bounds,
    trivial_casts,
    trivial_numeric_casts,
    type_alias_bounds,
    tyvar_behind_raw_pointer,
    unconditional_recursion,
    unions_with_drop_fields,
    unknown_crate_types,
    unnameable_test_functions,
    unreachable_code,
    unreachable_patterns,
    unreachable_pub,
    unsafe_code,
    unstable_features,
    unstable_name_collisions,
    unused_allocation,
    unused_assignments,
    unused_attributes,
    unused_comparisons,
    unused_doc_comments,
    unused_extern_crates,
    unused_features,
    unused_import_braces,
    unused_imports,
    unused_labels,
    unused_lifetimes,
    unused_macros,
    unused_must_use,
    unused_mut,
    unused_parens,
    unused_qualifications,
    unused_unsafe,
    unused_variables,
    where_clauses_object_safety,
    while_true
)]
#![warn(macro_use_extern_crate, unknown_lints)]
#![allow(
    bare_trait_objects,
    box_pointers,
    elided_lifetimes_in_paths,
    missing_debug_implementations,
    single_use_lifetimes,
    unused_results,
    variant_size_differences,
    warnings,
    renamed_and_removed_lints
)]
#![cfg_attr(feature = "clippy", feature(plugin))]

//! # cargo-make
//!
//! Rust task runner and build tool.

//! The cargo-make task runner enables to define and configure sets of tasks and run them as a flow.

//! A task is a command or a script to execute.

//! Tasks can have dependencies which are also tasks that will be executed before the task itself.

//! With a simple toml based configuration file, you can define a multi platform build script that can run build, test,
//! documentation generation, bench tests execution, security validations and more by running a single command.
//!
//! ## Installation
//! In order to install, just run the following command
//!
//! ```sh
//! cargo install cargo-make
//! ```
//!
//! This will install cargo-make in your ~/.cargo/bin.

//! Make sure to add ~/.cargo/bin directory to your PATH variable.
//!
//! # Contributing
//! See [contributing guide](https://github.com/sagiegurari/simple_redis/blob/master/.github/CONTRIBUTING.md)
//!
//! # License
//! Developed by Sagie Gur-Ari and licensed under the
//! [Apache 2](https://github.com/sagiegurari/simple_redis/blob/master/LICENSE) open source license.
//!

extern crate ci_info;
extern crate clap;
extern crate dirs;
extern crate fern;
extern crate glob;
extern crate indexmap;
#[macro_use]
extern crate log;
extern crate rand;
extern crate run_script;
extern crate rust_info;
extern crate semver;
#[macro_use]
extern crate serde_derive;
extern crate shell2batch;
extern crate toml;

// make types public for docs
pub mod types;

mod cache;
mod cli;
mod command;
mod condition;
mod config;
mod descriptor;
mod environment;
mod installer;
mod legacy;
mod logger;
mod runner;
mod scriptengine;
mod storage;
mod version;

#[cfg(test)]
#[path = "./main_test.rs"]
mod main_test;

#[cfg(test)]
#[path = "./test.rs"]
mod test;

fn main() {
    cli::run_cli();
}
<commit_msg>update lint rules<commit_after>#![deny(
    absolute_paths_not_starting_with_crate,
    anonymous_parameters,
    const_err,
    dead_code,
    deprecated,
    duplicate_associated_type_bindings,
    duplicate_macro_exports,
    ellipsis_inclusive_range_patterns,
    exceeding_bitshifts,
    explicit_outlives_requirements,
    illegal_floating_point_literal_pattern,
    improper_ctypes,
    incoherent_fundamental_impls,
    intra_doc_link_resolution_failure,
    invalid_type_param_default,
    irrefutable_let_patterns,
    keyword_idents,
    late_bound_lifetime_arguments,
    legacy_constructor_visibility,
    legacy_directory_ownership,
    macro_expanded_macro_exports_accessed_by_absolute_paths,
    missing_copy_implementations,
    missing_doc_code_examples,
    missing_docs,
    missing_fragment_specifier,
    mutable_transmutes,
    no_mangle_const_items,
    no_mangle_generic_items,
    non_camel_case_types,
    non_shorthand_field_patterns,
    non_snake_case,
    non_upper_case_globals,
    overflowing_literals,
    parenthesized_params_in_types_and_modules,
    path_statements,
    patterns_in_fns_without_body,
    plugin_as_library,
    private_in_public,
    proc_macro_derive_resolution_fallback,
    pub_use_of_private_extern_crate,
    question_mark_macro_sep,
    safe_extern_statics,
    safe_packed_borrows,
    stable_features,
    trivial_bounds,
    trivial_casts,
    trivial_numeric_casts,
    type_alias_bounds,
    tyvar_behind_raw_pointer,
    unconditional_recursion,
    unions_with_drop_fields,
    unknown_crate_types,
    unnameable_test_items,
    unreachable_code,
    unreachable_patterns,
    unreachable_pub,
    unsafe_code,
    unstable_features,
    unstable_name_collisions,
    unused_allocation,
    unused_assignments,
    unused_attributes,
    unused_comparisons,
    unused_doc_comments,
    unused_extern_crates,
    unused_features,
    unused_import_braces,
    unused_imports,
    unused_labels,
    unused_lifetimes,
    unused_macros,
    unused_must_use,
    unused_mut,
    unused_parens,
    unused_qualifications,
    unused_unsafe,
    unused_variables,
    where_clauses_object_safety,
    while_true
)]
#![warn(macro_use_extern_crate, unknown_lints)]
#![allow(
    bare_trait_objects,
    box_pointers,
    elided_lifetimes_in_paths,
    missing_debug_implementations,
    single_use_lifetimes,
    unused_results,
    variant_size_differences,
    warnings,
    renamed_and_removed_lints
)]
#![cfg_attr(feature = "clippy", feature(plugin))]

//! # cargo-make
//!
//! Rust task runner and build tool.

//! The cargo-make task runner enables to define and configure sets of tasks and run them as a flow.

//! A task is a command or a script to execute.

//! Tasks can have dependencies which are also tasks that will be executed before the task itself.

//! With a simple toml based configuration file, you can define a multi platform build script that can run build, test,
//! documentation generation, bench tests execution, security validations and more by running a single command.
//!
//! ## Installation
//! In order to install, just run the following command
//!
//! ```sh
//! cargo install cargo-make
//! ```
//!
//! This will install cargo-make in your ~/.cargo/bin.

//! Make sure to add ~/.cargo/bin directory to your PATH variable.
//!
//! # Contributing
//! See [contributing guide](https://github.com/sagiegurari/simple_redis/blob/master/.github/CONTRIBUTING.md)
//!
//! # License
//! Developed by Sagie Gur-Ari and licensed under the
//! [Apache 2](https://github.com/sagiegurari/simple_redis/blob/master/LICENSE) open source license.
//!

extern crate ci_info;
extern crate clap;
extern crate dirs;
extern crate fern;
extern crate glob;
extern crate indexmap;
#[macro_use]
extern crate log;
extern crate rand;
extern crate run_script;
extern crate rust_info;
extern crate semver;
#[macro_use]
extern crate serde_derive;
extern crate shell2batch;
extern crate toml;

// make types public for docs
pub mod types;

mod cache;
mod cli;
mod command;
mod condition;
mod config;
mod descriptor;
mod environment;
mod installer;
mod legacy;
mod logger;
mod runner;
mod scriptengine;
mod storage;
mod version;

#[cfg(test)]
#[path = "./main_test.rs"]
mod main_test;

#[cfg(test)]
#[path = "./test.rs"]
mod test;

fn main() {
    cli::run_cli();
}
<|endoftext|>


<commit_before>extern crate rand;
extern crate rustc_serialize;
extern crate docopt;

use std::collections::HashMap;
use std::fmt;
use std::fs::File;
use std::io::prelude::*;

use docopt::*;
use rand::{Rng, ThreadRng};

struct MarkovGenerator<T> {
    map: HashMap<T, HashMap<T, u32>>,
    count: u32,
    rng: ThreadRng
}

impl<T:Eq+std::hash::Hash+Clone> MarkovGenerator<T> {
    fn new() -> Self {
        MarkovGenerator{
            map: HashMap::new(),
            count: 0,
            rng: rand::thread_rng()
        }
    }

    fn process_all(&mut self, input: &Vec<T>) -> Result<(),()> {
        let mut prev_item: Option<&T> = None;
        for item in input {
            if prev_item.is_some() {
                let mut submap = self.map.entry(prev_item.unwrap().clone()).or_insert(HashMap::new());
                let mut prob = submap.entry(item.clone()).or_insert(0);
                *prob += 1;
                self.count += 1;

            }
            prev_item = Some(&item);
        }
        self.map.insert(prev_item.unwrap().clone(), HashMap::new());
        Ok(())
    }

    fn get_random(&mut self) -> T {
        let keys = &self.map.keys().collect::<Vec<&T>>();
        (*self.rng.choose(keys).unwrap()).clone()
    }

    fn get(&mut self, key: &T) -> Result<T, ()> {
        match self.map.get(key) {
            Some(submap) => {
                loop {
                    let random = self.rng.gen_range(1u32, self.count);
                    for (k, v) in submap.iter() {
                        if *v >= random {
                            let r:T = k.clone();
                            return Ok(r);
                        }
                    }
                }
            },
            None => Err(())
        }
    }
}

impl<T: fmt::Display+Eq+std::hash::Hash+Clone> fmt::Display for MarkovGenerator<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut result = String::new();
        result.push_str("MarkovGenerator: {\n");
        for (k,v) in &self.map {
            result.push_str(&format!("  \"{}\" => {{\n", &k));
            for (kk, vv) in v {
                result.push_str(&format!("    \"{}\" : {},\n", &kk, &vv));
            }
            result.push_str("  },\n");
        }
        result.push_str("}");
        write!(f, "{}", result)
    }
}

const USAGE: &'static str = "
Random text generator.

Usage:
  markov-text --input <INPUT> --size <SIZE> --output <OUTPUT>
  markov-text (-h | --help)
  markov-text --version

Options:
  -h --help                Show this screen.
  --version                Show version.
  --input <INPUT>          Path to the file with full text.
  --output <OUTPUT>        File with result of generation.
  --size <SIZE>            Size of the generated text in sentences.
";

#[derive(RustcDecodable)]
struct Args {
    flag_version: bool,
    flag_input: Option<String>,
    flag_output: Option<String>,
    flag_size: Option<u32>,
}

fn main() {
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());
    if args.flag_version {
        println!("markov-text v0.1.0");
        return;
    }
    let input_path = args.flag_input.unwrap_or_else(|| {panic!("Input file is not specified!")});
    let output_path = args.flag_output.unwrap_or_else(|| {panic!("Output file is not specified!")});
    let gen_size = args.flag_size.unwrap_or_else(|| {panic!("Text size is not specified!")});

    let mut f = File::open(input_path).unwrap();
    let mut s = String::new();
    f.read_to_string(&mut s).unwrap();

    let mut markov_net = MarkovGenerator::new();
    let _ = markov_net.process_all(&s.split(|c| c==' ' || c=='\n').collect());

    let mut count = 0;
    let mut word = markov_net.get_random();
    let mut gen_text = String::new();
    while count < gen_size {
        gen_text.push_str(word);
        gen_text.push(' ');
        word = markov_net.get(&word).unwrap();
        count += 1;
    }

    f = File::create(output_path).unwrap();
    f.write_all(gen_text.as_bytes()).unwrap();
}
<commit_msg>Fixed typos. Triggering first travis-ci build.<commit_after>extern crate rand;
extern crate rustc_serialize;
extern crate docopt;

use std::collections::HashMap;
use std::fmt;
use std::fs::File;
use std::io::prelude::*;

use docopt::*;
use rand::{Rng, ThreadRng};

struct MarkovGenerator<T> {
    map: HashMap<T, HashMap<T, u32>>,
    count: u32,
    rng: ThreadRng
}

impl<T:Eq+std::hash::Hash+Clone> MarkovGenerator<T> {
    fn new() -> Self {
        MarkovGenerator{
            map: HashMap::new(),
            count: 0,
            rng: rand::thread_rng()
        }
    }

    fn process_all(&mut self, input: &Vec<T>) -> Result<(),()> {
        let mut prev_item: Option<&T> = None;
        for item in input {
            if prev_item.is_some() {
                let mut submap = self.map.entry(prev_item.unwrap().clone()).or_insert(HashMap::new());
                let mut prob = submap.entry(item.clone()).or_insert(0);
                *prob += 1;
                self.count += 1;

            }
            prev_item = Some(&item);
        }
        self.map.insert(prev_item.unwrap().clone(), HashMap::new());
        Ok(())
    }

    fn get_random(&mut self) -> T {
        let keys = &self.map.keys().collect::<Vec<&T>>();
        (*self.rng.choose(keys).unwrap()).clone()
    }

    fn get(&mut self, key: &T) -> Result<T, ()> {
        match self.map.get(key) {
            Some(submap) => {
                loop {
                    let random = self.rng.gen_range(1u32, self.count);
                    for (k, v) in submap.iter() {
                        if *v >= random {
                            let r:T = k.clone();
                            return Ok(r);
                        }
                    }
                }
            },
            None => Err(())
        }
    }
}

impl<T: fmt::Display+Eq+std::hash::Hash+Clone> fmt::Display for MarkovGenerator<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut result = String::new();
        result.push_str("MarkovGenerator: {\n");
        for (k,v) in &self.map {
            result.push_str(&format!("  \"{}\" => {{\n", &k));
            for (kk, vv) in v {
                result.push_str(&format!("    \"{}\" : {},\n", &kk, &vv));
            }
            result.push_str("  },\n");
        }
        result.push_str("}");
        write!(f, "{}", result)
    }
}

const USAGE: &'static str = "
Random text generator.

Usage:
  markov-text --input <INPUT> --size <SIZE> --output <OUTPUT>
  markov-text (-h | --help)
  markov-text --version

Options:
  -h --help                Show this screen.
  --version                Show version.
  --input <INPUT>          Path to the file with full text.
  --output <OUTPUT>        File with result of generation.
  --size <SIZE>            Size of the generated text in words.
";

#[derive(RustcDecodable)]
struct Args {
    flag_version: bool,
    flag_input: Option<String>,
    flag_output: Option<String>,
    flag_size: Option<u32>,
}

fn main() {
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());
    if args.flag_version {
        println!("markov-text v0.1.0");
        return;
    }
    let input_path = args.flag_input.unwrap_or_else(|| {panic!("Input file is not specified!")});
    let output_path = args.flag_output.unwrap_or_else(|| {panic!("Output file is not specified!")});
    let gen_size = args.flag_size.unwrap_or_else(|| {panic!("Text size is not specified!")});

    let mut f = File::open(input_path).unwrap();
    let mut s = String::new();
    f.read_to_string(&mut s).unwrap();

    let mut markov_net = MarkovGenerator::new();
    let _ = markov_net.process_all(&s.split(|c| c==' ' || c=='\n').collect());

    let mut count = 0;
    let mut word = markov_net.get_random();
    let mut gen_text = String::new();
    while count < gen_size {
        gen_text.push_str(word);
        gen_text.push(' ');
        word = markov_net.get(&word).unwrap();
        count += 1;
    }

    f = File::create(output_path).unwrap();
    f.write_all(gen_text.as_bytes()).unwrap();
}
<|endoftext|>


<commit_before>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

extern crate backtrace;
#[macro_use(crate_version, crate_authors)]
extern crate clap;

extern crate api_server;
extern crate fc_util;
extern crate jailer;
#[macro_use]
extern crate logger;
extern crate mmds;
extern crate seccomp;
extern crate vmm;

use backtrace::Backtrace;
use clap::{App, Arg};

use std::io::ErrorKind;
use std::panic;
use std::path::PathBuf;
use std::process;
use std::sync::mpsc::channel;
use std::sync::{Arc, RwLock};

use api_server::{ApiServer, Error};
use fc_util::validators::validate_instance_id;
use logger::{Metric, LOGGER, METRICS};
use mmds::MMDS;
use vmm::signal_handler::register_signal_handlers;
use vmm::vmm_config::instance_info::{InstanceInfo, InstanceState};

const DEFAULT_API_SOCK_PATH: &str = "/tmp/firecracker.socket";
const DEFAULT_INSTANCE_ID: &str = "anonymous-instance";

fn main() {
    LOGGER
        .preinit(Some(DEFAULT_INSTANCE_ID.to_string()))
        .expect("Failed to register logger");

    if let Err(e) = register_signal_handlers() {
        error!("Failed to register signal handlers: {}", e);
        process::exit(i32::from(vmm::FC_EXIT_CODE_GENERIC_ERROR));
    }
    // Start firecracker by setting up a panic hook, which will be called before
    // terminating as we're building with panic = "abort".
    // It's worth noting that the abort is caused by sending a SIG_ABORT signal to the process.
    panic::set_hook(Box::new(move |info| {
        // We're currently using the closure parameter, which is a &PanicInfo, for printing the
        // origin of the panic, including the payload passed to panic! and the source code location
        // from which the panic originated.
        error!("Firecracker {}", info);
        METRICS.vmm.panic_count.inc();
        let bt = Backtrace::new();
        error!("{:?}", bt);

        // Log the metrics before aborting.
        if let Err(e) = LOGGER.log_metrics() {
            error!("Failed to log metrics while panicking: {}", e);
        }
    }));

    let cmd_arguments = App::new("firecracker")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Launch a microvm.")
        .arg(
            Arg::with_name("api_sock")
                .long("api-sock")
                .help("Path to unix domain socket used by the API")
                .takes_value(true)
                .default_value(DEFAULT_API_SOCK_PATH),
        )
        .arg(
            Arg::with_name("id")
                .long("id")
                .help("MicroVM unique identifier")
                .takes_value(true)
                .default_value(DEFAULT_INSTANCE_ID)
                .validator(|s: String| -> Result<(), String> {
                    validate_instance_id(&s).map_err(|e| format!("{}", e))
                }),
        )
        .arg(
            Arg::with_name("seccomp-level")
                .long("seccomp-level")
                .help(
                    "Level of seccomp filtering.\n
                            - Level 0: No filtering.\n
                            - Level 1: Seccomp filtering by syscall number.\n
                            - Level 2: Seccomp filtering by syscall number and argument values.\n
                        ",
                )
                .takes_value(true)
                .default_value("2")
                .possible_values(&["0", "1", "2"]),
        )
        .arg(
            Arg::with_name("start-time-us")
                .long("start-time-us")
                .takes_value(true)
                .hidden(true),
        )
        .arg(
            Arg::with_name("start-time-cpu-us")
                .long("start-time-cpu-us")
                .takes_value(true)
                .hidden(true),
        )
        .get_matches();

    let bind_path = cmd_arguments
        .value_of("api_sock")
        .map(PathBuf::from)
        .expect("Missing argument: api_sock");

    // It's safe to unwrap here because clap's been provided with a default value
    let instance_id = cmd_arguments.value_of("id").unwrap().to_string();

    // We disable seccomp filtering when testing, because when running the test_gnutests
    // integration test from test_unittests.py, an invalid syscall is issued, and we crash
    // otherwise.
    #[cfg(test)]
    let seccomp_level = seccomp::SECCOMP_LEVEL_NONE;
    #[cfg(not(test))]
    // It's safe to unwrap here because clap's been provided with a default value,
    // and allowed values are guaranteed to parse to u32.
    let seccomp_level = cmd_arguments
        .value_of("seccomp-level")
        .unwrap()
        .parse::<u32>()
        .unwrap();

    let start_time_us = cmd_arguments.value_of("start-time-us").map(|s| {
        s.parse::<u64>()
            .expect("'start-time-us' parameter expected to be of 'u64' type.")
    });

    let start_time_cpu_us = cmd_arguments.value_of("start-time-cpu-us").map(|s| {
        s.parse::<u64>()
            .expect("'start-time-cpu_us' parameter expected to be of 'u64' type.")
    });

    let shared_info = Arc::new(RwLock::new(InstanceInfo {
        state: InstanceState::Uninitialized,
        id: instance_id,
        vmm_version: crate_version!().to_string(),
    }));
    let mmds_info = MMDS.clone();
    let (to_vmm, from_api) = channel();
    let server =
        ApiServer::new(mmds_info, shared_info.clone(), to_vmm).expect("Cannot create API server");

    let api_event_fd = server
        .get_event_fd_clone()
        .expect("Cannot clone API eventFD.");

    let _vmm_thread_handle =
        vmm::start_vmm_thread(shared_info, api_event_fd, from_api, seccomp_level);

    match server.bind_and_run(bind_path, start_time_us, start_time_cpu_us, seccomp_level) {
        Ok(_) => (),
        Err(Error::Io(inner)) => match inner.kind() {
            ErrorKind::AddrInUse => panic!("Failed to open the API socket: {:?}", Error::Io(inner)),
            _ => panic!(
                "Failed to communicate with the API socket: {:?}",
                Error::Io(inner)
            ),
        },
        Err(eventfd_err @ Error::Eventfd(_)) => {
            panic!("Failed to open the API socket: {:?}", eventfd_err)
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate tempfile;

    use self::tempfile::NamedTempFile;
    use super::*;

    use logger::AppInfo;
    use std::fs::File;
    use std::io::BufRead;
    use std::io::BufReader;
    use std::path::Path;
    use std::time::Duration;
    use std::{fs, thread};

    /// Look through the log for lines that match expectations.
    /// expectations is a list of tuples of words we're looking for.
    /// A tuple matches a line if all the words in the tuple can be found on that line.
    /// For this test to pass, every tuple must match at least one line.
    fn validate_backtrace(
        log_path: &str,
        expectations: &[(&'static str, &'static str, &'static str)],
    ) -> bool {
        let f = File::open(log_path).unwrap();
        let reader = BufReader::new(f);
        let mut expectation_iter = expectations.iter();
        let mut expected_words = expectation_iter.next().unwrap();

        for ln_res in reader.lines() {
            let line = ln_res.unwrap();
            if !(line.contains(expected_words.0)
                && line.contains(expected_words.1)
                && line.contains(expected_words.2))
            {
                continue;
            }
            if let Some(w) = expectation_iter.next() {
                expected_words = w;
                continue;
            }
            return true;
        }
        false
    }

    #[test]
    #[allow(clippy::unit_cmp)]
    fn test_main() {
        const FIRECRACKER_INIT_TIMEOUT_MILLIS: u64 = 100;

        // There is no reason to run this test if the default socket path exists.
        assert!(!Path::new(DEFAULT_API_SOCK_PATH).exists());

        let log_file_temp =
            NamedTempFile::new().expect("Failed to create temporary output logging file.");
        let metrics_file_temp =
            NamedTempFile::new().expect("Failed to create temporary metrics logging file.");
        let log_file = String::from(log_file_temp.path().to_path_buf().to_str().unwrap());

        // Start Firecracker in a separate thread
        thread::spawn(|| {
            main();
        });

        // Wait around for a bit, so Firecracker has time to initialize and create the
        // API socket.
        thread::sleep(Duration::from_millis(FIRECRACKER_INIT_TIMEOUT_MILLIS));

        // If Firecracker hasn't finished initializing yet, something is really wrong!
        assert!(Path::new(DEFAULT_API_SOCK_PATH).exists());

        // Initialize the logger
        LOGGER
            .init(
                &AppInfo::new("Firecracker", "1.0"),
                "TEST-ID",
                log_file_temp.path().to_str().unwrap().to_string(),
                metrics_file_temp.path().to_str().unwrap().to_string(),
                &[],
            )
            .expect("Could not initialize logger.");

        // Cause some controlled panic and see if a backtrace shows up in the log,
        // as it's supposed to.
        let _ = panic::catch_unwind(|| {
            panic!("Oh, noes!");
        });
        // Look for the expected backtrace inside the log
        assert!(
            validate_backtrace(
                log_file.as_str(),
                &[
                    // Lines containing these words should have appeared in the log, in this order
                    ("ERROR", "main.rs", "Firecracker panicked at"),
                    ("ERROR", "main.rs", "stack backtrace:"),
                    ("0:", "0x", "firecracker::main::"),
                ],
            ) || println!("Could not validate backtrace!\n {:?}", Backtrace::new()) != ()
        );

        // Clean up
        fs::remove_file(DEFAULT_API_SOCK_PATH).expect("failure in removing socket file");
    }
}
<commit_msg>src: fix main.rs test<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

extern crate backtrace;
#[macro_use(crate_version, crate_authors)]
extern crate clap;

extern crate api_server;
extern crate fc_util;
extern crate jailer;
#[macro_use]
extern crate logger;
extern crate mmds;
extern crate seccomp;
extern crate vmm;

use backtrace::Backtrace;
use clap::{App, Arg};

use std::io::ErrorKind;
use std::panic;
use std::path::PathBuf;
use std::process;
use std::sync::mpsc::channel;
use std::sync::{Arc, RwLock};

use api_server::{ApiServer, Error};
use fc_util::validators::validate_instance_id;
use logger::{Metric, LOGGER, METRICS};
use mmds::MMDS;
use vmm::signal_handler::register_signal_handlers;
use vmm::vmm_config::instance_info::{InstanceInfo, InstanceState};

const DEFAULT_API_SOCK_PATH: &str = "/tmp/firecracker.socket";
const DEFAULT_INSTANCE_ID: &str = "anonymous-instance";

fn main() {
    LOGGER
        .preinit(Some(DEFAULT_INSTANCE_ID.to_string()))
        .expect("Failed to register logger");

    if let Err(e) = register_signal_handlers() {
        error!("Failed to register signal handlers: {}", e);
        process::exit(i32::from(vmm::FC_EXIT_CODE_GENERIC_ERROR));
    }
    // Start firecracker by setting up a panic hook, which will be called before
    // terminating as we're building with panic = "abort".
    // It's worth noting that the abort is caused by sending a SIG_ABORT signal to the process.
    panic::set_hook(Box::new(move |info| {
        // We're currently using the closure parameter, which is a &PanicInfo, for printing the
        // origin of the panic, including the payload passed to panic! and the source code location
        // from which the panic originated.
        error!("Firecracker {}", info);
        METRICS.vmm.panic_count.inc();
        let bt = Backtrace::new();
        error!("{:?}", bt);

        // Log the metrics before aborting.
        if let Err(e) = LOGGER.log_metrics() {
            error!("Failed to log metrics while panicking: {}", e);
        }
    }));

    let cmd_arguments = App::new("firecracker")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Launch a microvm.")
        .arg(
            Arg::with_name("api_sock")
                .long("api-sock")
                .help("Path to unix domain socket used by the API")
                .takes_value(true)
                .default_value(DEFAULT_API_SOCK_PATH),
        )
        .arg(
            Arg::with_name("id")
                .long("id")
                .help("MicroVM unique identifier")
                .takes_value(true)
                .default_value(DEFAULT_INSTANCE_ID)
                .validator(|s: String| -> Result<(), String> {
                    validate_instance_id(&s).map_err(|e| format!("{}", e))
                }),
        )
        .arg(
            Arg::with_name("seccomp-level")
                .long("seccomp-level")
                .help(
                    "Level of seccomp filtering.\n
                            - Level 0: No filtering.\n
                            - Level 1: Seccomp filtering by syscall number.\n
                            - Level 2: Seccomp filtering by syscall number and argument values.\n
                        ",
                )
                .takes_value(true)
                .default_value("2")
                .possible_values(&["0", "1", "2"]),
        )
        .arg(
            Arg::with_name("start-time-us")
                .long("start-time-us")
                .takes_value(true)
                .hidden(true),
        )
        .arg(
            Arg::with_name("start-time-cpu-us")
                .long("start-time-cpu-us")
                .takes_value(true)
                .hidden(true),
        )
        .get_matches();

    let bind_path = cmd_arguments
        .value_of("api_sock")
        .map(PathBuf::from)
        .expect("Missing argument: api_sock");

    // It's safe to unwrap here because clap's been provided with a default value
    let instance_id = cmd_arguments.value_of("id").unwrap().to_string();

    // We disable seccomp filtering when testing, because when running the test_gnutests
    // integration test from test_unittests.py, an invalid syscall is issued, and we crash
    // otherwise.
    #[cfg(test)]
    let seccomp_level = seccomp::SECCOMP_LEVEL_NONE;
    #[cfg(not(test))]
    // It's safe to unwrap here because clap's been provided with a default value,
    // and allowed values are guaranteed to parse to u32.
    let seccomp_level = cmd_arguments
        .value_of("seccomp-level")
        .unwrap()
        .parse::<u32>()
        .unwrap();

    let start_time_us = cmd_arguments.value_of("start-time-us").map(|s| {
        s.parse::<u64>()
            .expect("'start-time-us' parameter expected to be of 'u64' type.")
    });

    let start_time_cpu_us = cmd_arguments.value_of("start-time-cpu-us").map(|s| {
        s.parse::<u64>()
            .expect("'start-time-cpu_us' parameter expected to be of 'u64' type.")
    });

    let shared_info = Arc::new(RwLock::new(InstanceInfo {
        state: InstanceState::Uninitialized,
        id: instance_id,
        vmm_version: crate_version!().to_string(),
    }));
    let mmds_info = MMDS.clone();
    let (to_vmm, from_api) = channel();
    let server =
        ApiServer::new(mmds_info, shared_info.clone(), to_vmm).expect("Cannot create API server");

    let api_event_fd = server
        .get_event_fd_clone()
        .expect("Cannot clone API eventFD.");

    let _vmm_thread_handle =
        vmm::start_vmm_thread(shared_info, api_event_fd, from_api, seccomp_level);

    match server.bind_and_run(bind_path, start_time_us, start_time_cpu_us, seccomp_level) {
        Ok(_) => (),
        Err(Error::Io(inner)) => match inner.kind() {
            ErrorKind::AddrInUse => panic!("Failed to open the API socket: {:?}", Error::Io(inner)),
            _ => panic!(
                "Failed to communicate with the API socket: {:?}",
                Error::Io(inner)
            ),
        },
        Err(eventfd_err @ Error::Eventfd(_)) => {
            panic!("Failed to open the API socket: {:?}", eventfd_err)
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate tempfile;

    use self::tempfile::NamedTempFile;
    use super::*;

    use logger::AppInfo;
    use std::fs::File;
    use std::io::BufRead;
    use std::io::BufReader;
    use std::path::Path;
    use std::time::Duration;
    use std::{fs, thread};

    /// Look through the log for lines that match expectations.
    /// expectations is a list of tuples of words we're looking for.
    /// A tuple matches a line if all the words in the tuple can be found on that line.
    /// For this test to pass, every tuple must match at least one line.
    fn validate_backtrace(
        log_path: &str,
        expectations: &[(&'static str, &'static str, &'static str)],
    ) -> bool {
        let f = File::open(log_path).unwrap();
        let reader = BufReader::new(f);
        let mut expectation_iter = expectations.iter();
        let mut expected_words = expectation_iter.next().unwrap();

        for ln_res in reader.lines() {
            let line = ln_res.unwrap();
            if !(line.contains(expected_words.0)
                && line.contains(expected_words.1)
                && line.contains(expected_words.2))
            {
                continue;
            }
            if let Some(w) = expectation_iter.next() {
                expected_words = w;
                continue;
            }
            return true;
        }
        false
    }

    #[allow(clippy::unit_cmp)]
    #[test]
    fn test_main() {
        const FIRECRACKER_INIT_TIMEOUT_MILLIS: u64 = 150;

        // If the default api socket path exist, remove it so we can continue running the test.
        if Path::new(DEFAULT_API_SOCK_PATH).exists() {
            fs::remove_file(DEFAULT_API_SOCK_PATH).expect("failure in removing socket file");
        }

        let log_file_temp =
            NamedTempFile::new().expect("Failed to create temporary output logging file.");
        let metrics_file_temp =
            NamedTempFile::new().expect("Failed to create temporary metrics logging file.");
        let log_file = String::from(log_file_temp.path().to_path_buf().to_str().unwrap());

        // Start Firecracker in a separate thread
        thread::spawn(|| {
            main();
        });

        // Wait around for a bit, so Firecracker has time to initialize and create the
        // API socket.
        thread::sleep(Duration::from_millis(FIRECRACKER_INIT_TIMEOUT_MILLIS));

        // If Firecracker hasn't finished initializing yet, something is really wrong!
        assert!(Path::new(DEFAULT_API_SOCK_PATH).exists());

        // Initialize the logger
        LOGGER
            .init(
                &AppInfo::new("Firecracker", "1.0"),
                "TEST-ID",
                log_file_temp.path().to_str().unwrap().to_string(),
                metrics_file_temp.path().to_str().unwrap().to_string(),
                &[],
            )
            .expect("Could not initialize logger.");

        // Cause some controlled panic and see if a backtrace shows up in the log,
        // as it's supposed to.
        let _ = panic::catch_unwind(|| {
            panic!("Oh, noes!");
        });

        // Before checking the backtrace let's remove the API socket to make sure we don't
        // leave it on the host in case the assert fails.
        fs::remove_file(DEFAULT_API_SOCK_PATH).expect("failure in removing socket file");

        // Look for the expected backtrace inside the log
        assert!(
            validate_backtrace(
                log_file.as_str(),
                &[
                    // Lines containing these words should have appeared in the log, in this order
                    ("ERROR", "main.rs", "Firecracker panicked at"),
                    ("ERROR", "main.rs", "stack backtrace:"),
                    ("0:", "0x", "firecracker::main::"),
                ],
            ) || println!("Could not validate backtrace!\n {:?}", Backtrace::new()) != ()
        );
    }
}
<|endoftext|>


<commit_before>//! Implementation of a simple uTP client and server.
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate getopts;
extern crate daemonize;
extern crate udt;

use daemonize::{Daemonize};
use std::process;
use std::process::Command;
use std::thread;
use std::net::{SocketAddr, SocketAddrV4, Ipv4Addr};
use std::net;
use std::str;
use std::env;
use std::fs::File;
use std::str::FromStr;
use std::path::Path;
use getopts::Options;
use udt::*;

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

fn main() {
    use std::io::{stdin, stdout, stderr, Read, Write};

    // This example may run in either server or client mode.
    // Using an enum tends to make the code cleaner and easier to read.
    enum Mode {Server, Client}

    // Start logging
    env_logger::init().expect("Error starting logger");

    // Fetch arguments
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optopt("o", "output", "set output file name", "NAME");
    opts.optflag("s", "server", "server mode");
    opts.optflag("p", "port-range", "server listening port range");
    opts.optflag("h", "help", "print this help menu");

    let matches = match opts.parse(&args[1..]) {
        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;
    };

    // Parse the mode argument
    let mode: Mode = match matches.opt_present("s") {
        true => Mode::Server,
        false => Mode::Client
    };

    match mode {
        Mode::Server => {
            udt::init();
            let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Stream).unwrap();
            sock.setsockopt(UdtOpts::UDP_RCVBUF, 5590000i32);
            sock.setsockopt(UdtOpts::UDP_SNDBUF, 5590000i32);
            sock.bind(SocketAddr::V4(SocketAddrV4::from_str("0.0.0.0:55000").unwrap())).unwrap();
            let my_addr = sock.getsockname().unwrap();
            println!("Server bound to {:?}", my_addr);

            sock.listen(5).unwrap();

            let (mut stream, peer) = sock.accept().unwrap();
            println!("Received new connection from peer {:?}", peer);

            let mut f = File::open(input).unwrap();

            let daemonize = Daemonize::new();

            let mut writer = stdout();
            // let _ = writeln!(&mut stderr(), "Serving on {}", addr);

            let mut startmagic = vec![0; 1];
            let amount_read = stream.recv(&mut startmagic, 1).unwrap();
            println!("got {} magic byte(s).", amount_read);

            // match daemonize.start() {
            //     Ok(_) => { let _ = writeln!(&mut stderr(), "daemonized"); }
            //     Err(e) => { let _ = writeln!(&mut stderr(), "RWRWARWARARRR"); }
            // }

            let mut total = 0;
            // Create a reasonably sized buffer
            let mut payload = vec![0; 1300];
            // loop {
            //     match stream.send(&payload) {
            //         Ok(written) => { print!("."); }
            //         Err(e) => { panic!("{:?}", e); }
            //     }
            // }
            loop {
                match f.read(&mut payload) {
                    Ok(read) => {
                        println!("file read {}", read);
                        match stream.send(&payload[0..read]) {
                            Ok(written) => {
                                total += written;
                                print!("\rwritten {}", total);
                            },
                            Err(e) => {
                                stream.close().expect("Error closing stream");
                                panic!("{:?}", e);
                            }
                        }
                    },
                    Err(e) => {
                        stream.close().expect("Error closing stream");
                        panic!("{}", e);
                    }
                }
            }
            println!("yay");
        }
        Mode::Client => {
            let sections: Vec<&str> = input.split(":").collect();
            let addr: String = sections[0].to_owned();
            let path: String = sections[1].to_owned();
            let cmd = format!("~/bin/shoop -s {}", path);
            println!("addr: {}, path: {}, cmd: {}", addr, path, cmd);

            // let output = Command::new("ssh")
            //                      .arg(addr.to_owned())
            //                      .arg(cmd)
            //                      .output()
            //                      .unwrap_or_else(|e| {
            //                          panic!("failed to execute process: {}", e);
            //                      });
            // let udp_addr = String::from_utf8_lossy(&output.stdout).trim().to_owned();
            //
            // // Create a stream and try to connect to the remote address
            // println!("shoop server told us to connect to {}", udp_addr);
            udt::init();
            let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Stream).unwrap();
            sock.setsockopt(UdtOpts::UDP_RCVBUF, 5590000i32);
            sock.setsockopt(UdtOpts::UDP_SNDBUF, 5590000i32);
            let addr: SocketAddr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::from_str(&addr).unwrap(), 55000));
            sock.bind(SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::from_str("0.0.0.0").unwrap(), 0)));
            match sock.connect(addr) {
                Ok(()) => {
                    println!("connected!");
                },
                Err(e) => {
                    panic!("errrrrrrr {:?}", e);
                }
            }

            let mut f = File::create("outfile").unwrap();
            println!("created file.");
            let mut payload = vec![0; 1024 * 1024];
            sock.send(b"\x01");
            println!("write magic byte.");

            let mut total = 0;
            loop {
                match sock.recv(&mut payload, 1024 * 1024) {
                    Ok(0) => {
                        println!("EOF");
                        break
                    },
                    Ok(read) => {
                        total += read;
                        println!("read {}", total);
                        f.write_all(&payload[0..(read-1) as usize]);
                    },
                    Err(e) => {
                        panic!("{:?}", e);
                    }
                }
            }
            sock.close();
        }
    }
}

<commit_msg>EOF handling<commit_after>//! Implementation of a simple uTP client and server.
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate getopts;
extern crate daemonize;
extern crate udt;

use daemonize::{Daemonize};
use std::process;
use std::process::Command;
use std::thread;
use std::net::{SocketAddr, SocketAddrV4, Ipv4Addr};
use std::net;
use std::str;
use std::env;
use std::fs::File;
use std::str::FromStr;
use std::path::Path;
use getopts::Options;
use udt::*;

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

fn main() {
    use std::io::{stdin, stdout, stderr, Read, Write};

    // This example may run in either server or client mode.
    // Using an enum tends to make the code cleaner and easier to read.
    enum Mode {Server, Client}

    // Start logging
    env_logger::init().expect("Error starting logger");

    // Fetch arguments
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optopt("o", "output", "set output file name", "NAME");
    opts.optflag("s", "server", "server mode");
    opts.optflag("p", "port-range", "server listening port range");
    opts.optflag("h", "help", "print this help menu");

    let matches = match opts.parse(&args[1..]) {
        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;
    };

    // Parse the mode argument
    let mode: Mode = match matches.opt_present("s") {
        true => Mode::Server,
        false => Mode::Client
    };

    match mode {
        Mode::Server => {
            udt::init();
            let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Stream).unwrap();
            sock.setsockopt(UdtOpts::UDP_RCVBUF, 5590000i32);
            sock.setsockopt(UdtOpts::UDP_SNDBUF, 5590000i32);
            sock.bind(SocketAddr::V4(SocketAddrV4::from_str("0.0.0.0:55000").unwrap())).unwrap();
            let my_addr = sock.getsockname().unwrap();
            println!("Server bound to {:?}", my_addr);

            sock.listen(5).unwrap();

            let (mut stream, peer) = sock.accept().unwrap();
            println!("Received new connection from peer {:?}", peer);

            let mut f = File::open(input).unwrap();

            let daemonize = Daemonize::new();

            let mut writer = stdout();
            // let _ = writeln!(&mut stderr(), "Serving on {}", addr);

            let mut startmagic = vec![0; 1];
            let amount_read = stream.recv(&mut startmagic, 1).unwrap();
            println!("got {} magic byte(s).", amount_read);

            // match daemonize.start() {
            //     Ok(_) => { let _ = writeln!(&mut stderr(), "daemonized"); }
            //     Err(e) => { let _ = writeln!(&mut stderr(), "RWRWARWARARRR"); }
            // }

            let mut total = 0;
            // Create a reasonably sized buffer
            let mut payload = vec![0; 1300];
            // loop {
            //     match stream.send(&payload) {
            //         Ok(written) => { print!("."); }
            //         Err(e) => { panic!("{:?}", e); }
            //     }
            // }
            loop {
                match f.read(&mut payload) {
                    Ok(read) => {
                        match stream.send(&payload[0..read]) {
                            Ok(0) => {
                                println!("EOF.");
                                break;
                            }
                            Ok(written) => {
                                total += written;
                                print!("\rwritten {}", total);
                            },
                            Err(e) => {
                                stream.close().expect("Error closing stream");
                                panic!("{:?}", e);
                            }
                        }
                    },
                    Err(e) => {
                        stream.close().expect("Error closing stream");
                        panic!("{}", e);
                    }
                }
            }
            stream.close().expect("Error closing stream.");
            println!("yay");
        }
        Mode::Client => {
            let sections: Vec<&str> = input.split(":").collect();
            let addr: String = sections[0].to_owned();
            let path: String = sections[1].to_owned();
            let cmd = format!("~/bin/shoop -s {}", path);
            println!("addr: {}, path: {}, cmd: {}", addr, path, cmd);

            // let output = Command::new("ssh")
            //                      .arg(addr.to_owned())
            //                      .arg(cmd)
            //                      .output()
            //                      .unwrap_or_else(|e| {
            //                          panic!("failed to execute process: {}", e);
            //                      });
            // let udp_addr = String::from_utf8_lossy(&output.stdout).trim().to_owned();
            //
            // // Create a stream and try to connect to the remote address
            // println!("shoop server told us to connect to {}", udp_addr);
            udt::init();
            let sock = UdtSocket::new(SocketFamily::AFInet, SocketType::Stream).unwrap();
            sock.setsockopt(UdtOpts::UDP_RCVBUF, 5590000i32);
            sock.setsockopt(UdtOpts::UDP_SNDBUF, 5590000i32);
            let addr: SocketAddr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::from_str(&addr).unwrap(), 55000));
            sock.bind(SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::from_str("0.0.0.0").unwrap(), 0)));
            match sock.connect(addr) {
                Ok(()) => {
                    println!("connected!");
                },
                Err(e) => {
                    panic!("errrrrrrr {:?}", e);
                }
            }

            let mut f = File::create("outfile").unwrap();
            println!("created file.");
            let mut payload = vec![0; 1024 * 1024];
            sock.send(b"\x01");
            println!("write magic byte.");

            let mut total = 0;
            loop {
                match sock.recv(&mut payload, 1024 * 1024) {
                    Ok(0) => {
                        println!("EOF");
                        break
                    },
                    Ok(read) => {
                        total += read;
                        println!("read {}", total);
                        f.write_all(&payload[0..(read-1) as usize]);
                    },
                    Err(e) => {
                        panic!("{:?}", e);
                    }
                }
            }
            sock.close().expect("Error closing socket.");
        }
    }
}

<|endoftext|>