<commit_before>extern crate mmap;
extern crate time;
extern crate timely;
extern crate getopts;

use timely::progress::timestamp::RootTimestamp;
// use timely::progress::nested::Summary::Local;
use timely::dataflow::*;
use timely::dataflow::operators::*;
// use timely::communication::*;
use timely::dataflow::channels::pact::Exchange;
use timely::drain::DrainExt;

mod typedrw;
mod graphmap;
mod sorting;
use graphmap::GraphMMap;
use sorting::{SegmentList, radix_sort_32};

fn main () {

    let filename = std::env::args().skip(1).next().unwrap();
    let strategy = std::env::args().skip(2).next().unwrap() == "process";

    // currently need timely's full option set to parse args
    let mut opts = getopts::Options::new();
    opts.optopt("w", "workers", "", "");
    opts.optopt("p", "process", "", "");
    opts.optopt("n", "processes", "", "");
    opts.optopt("h", "hostfile", "", "");

    if let Ok(matches) = opts.parse(std::env::args().skip(3)) {

        let workers: usize = matches.opt_str("w").map(|x| x.parse().unwrap_or(1)).unwrap_or(1);

        timely::execute_from_args(std::env::args().skip(2), move |root| {

            let index = root.index() as usize;
            let peers = root.peers() as usize;

            let start = time::precise_time_s();

            let nodes = GraphMMap::new(&filename).nodes();

            let mut segments = SegmentList::new(1024); // list of edge segments

            let mut src = vec![];   // holds ranks
            let mut deg = vec![];   // holds source degrees
            let mut rev = vec![];   // holds (dst, deg) pairs
            let mut trn = vec![];   // holds transposed sources

            let mut going = start;

            let mut input = root.scoped(|builder| {

                let (input, edges) = builder.new_input::<(u32, u32)>();
                let (cycle, ranks) = builder.loop_variable::<(u32, f32)>(20, 1);

                let mut ranks = edges.binary_notify(&ranks,
                                    Exchange::new(|x: &(u32,u32)| x.0 as u64),
                                    Exchange::new(|x: &(u32,f32)| x.0 as u64),
                                    "pagerank",
                                    vec![RootTimestamp::new(0)],
                                    move |input1, input2, output, notificator| {

                    // receive incoming edges (should only be iter 0)
                    while let Some((_index, data)) = input1.next() {
                        segments.push(data.drain_temp());
                    }

                    // all inputs received for iter, commence multiplication
                    while let Some((iter, _)) = notificator.next() {

                        let now = time::now();

                        if index == 0 { println!("{}:{}:{}.{} starting iteration {}", now.tm_hour, now.tm_min, now.tm_sec, now.tm_nsec, iter.inner); }

                        // if the very first iteration, prepare some stuff.
                        // specifically, transpose edges and sort by destination.
                        if iter.inner == 0 {
                            let (a, b, c) = transpose(segments.finalize(), peers, nodes);
                            deg = a; rev = b; trn = c;
                            src = vec![0.0f32; deg.len()];
                        }

                        // record some timings in order to estimate per-iteration times
                        if iter.inner == 0  && index == 0 { println!("src: {}, dst: {}, edges: {}", src.len(), rev.len(), trn.len()); }
                        if iter.inner == 10 && index == 0 { going = time::precise_time_s(); }
                        if iter.inner == 20 && index == 0 { println!("average: {}", (time::precise_time_s() - going) / 10.0 ); }

                        // prepare src for transmitting to destinations
                        for s in 0..src.len() { src[s] = 0.15 + 0.85 * src[s] / deg[s] as f32; }

                        // wander through destinations
                        let mut trn_slice = &trn[..];
                        let mut session = output.session(&iter);
                        for &(dst, deg) in &rev {
                            let mut accum = 0.0;
                            for &s in &trn_slice[..deg as usize] {
                                // accum += src[s as usize];
                                unsafe { accum += *src.get_unchecked(s as usize); }
                            }
                            trn_slice = &trn_slice[deg as usize..];
                            session.give((dst, accum));
                        }

                        for s in &mut src { *s = 0.0; }
                    }

                    // receive data from workers, accumulate in src
                    while let Some((iter, data)) = input2.next() {
                        notificator.notify_at(&iter);
                        for &(node, rank) in data.iter() {
                            src[node as usize / peers] += rank;
                        }
                    }
                });

                // optionally, do process-local accumulation
                if strategy {
                    let local_base = workers * (index / workers);
                    let local_index = index % workers;
                    let mut acc = vec![0.0; (nodes / workers) + 1];   // holds ranks
                    ranks = ranks.unary_notify(
                        Exchange::new(move |x: &(u32,f32)| (local_base as u64 + (x.0 as u64 % workers as u64))),
                        "aggregation",
                        vec![],
                        move |input, output, iterator| {
                            while let Some((iter, data)) = input.next() {
                                iterator.notify_at(&iter);
                                for &(node, rank) in data.iter() {
                                    acc[node as usize / workers] += rank;
                                }
                            }

                            while let Some((item, _)) = iterator.next() {
                                output.session(&item)
                                      .give_iterator(acc.drain_temp()
                                                        .enumerate()
                                                        .filter(|x| x.1 != 0.0)
                                                        .map(|(u,f)| ((u * workers + local_index) as u32, f)));

                                for _ in 0..(1 + (nodes/workers)) { acc.push(0.0); }
                            }
                        }
                    );
                }

                ranks.connect_loop(cycle);

                input
            });

            // introduce edges into the computation;
            // allow mmaped file to drop
            {
                let graph = GraphMMap::new(&filename);
                for node in 0..graph.nodes() {
                    if node % peers == index {
                        for dst in graph.edges(node) {
                            input.send((node as u32, *dst as u32));
                        }
                    }
                }
            }
            input.close();
            while root.step() { }

            if index == 0 { println!("elapsed: {}", time::precise_time_s() - start); }
        });
    }
    else {
        println!("error parsing arguments");
        println!("usage:\tpagerank <source> (worker|process) [timely options]");
    }
}

// returns [src/peers] degrees, (dst, deg) pairs, and a list of [src/peers] endpoints
fn transpose(mut edges: Vec<Vec<(u32, u32)>>, peers: usize, nodes: usize) -> (Vec<u32>, Vec<(u32, u32)>, Vec<u32>)  {

    let mut deg = vec![0; (nodes as usize / peers) + 1];
    for list in &edges {
        for &(s, _) in list {
            deg[s as usize / peers] += 1;
        }
    }

    radix_sort_32(&mut edges, &mut Vec::new(), &|&(_,d)| d);

    let mut rev = Vec::<(u32,u32)>::with_capacity(deg.len());
    let mut trn = Vec::with_capacity(edges.len() * 1024);
    for list in edges {
        for (s,d) in list {
            let len = rev.len();
            if (len == 0) || (rev[len-1].0 < d) {
                rev.push((d, 0u32));
            }

            let len = rev.len();
            rev[len-1].1 += 1;
            trn.push(s / peers as u32);
        }
    }

    return (deg, rev, trn);
}
<commit_msg>print shape information for each worker<commit_after>extern crate mmap;
extern crate time;
extern crate timely;
extern crate getopts;

use timely::progress::timestamp::RootTimestamp;
// use timely::progress::nested::Summary::Local;
use timely::dataflow::*;
use timely::dataflow::operators::*;
// use timely::communication::*;
use timely::dataflow::channels::pact::Exchange;
use timely::drain::DrainExt;

mod typedrw;
mod graphmap;
mod sorting;
use graphmap::GraphMMap;
use sorting::{SegmentList, radix_sort_32};

fn main () {

    let filename = std::env::args().skip(1).next().unwrap();
    let strategy = std::env::args().skip(2).next().unwrap() == "process";

    // currently need timely's full option set to parse args
    let mut opts = getopts::Options::new();
    opts.optopt("w", "workers", "", "");
    opts.optopt("p", "process", "", "");
    opts.optopt("n", "processes", "", "");
    opts.optopt("h", "hostfile", "", "");

    if let Ok(matches) = opts.parse(std::env::args().skip(3)) {

        let workers: usize = matches.opt_str("w").map(|x| x.parse().unwrap_or(1)).unwrap_or(1);

        timely::execute_from_args(std::env::args().skip(2), move |root| {

            let index = root.index() as usize;
            let peers = root.peers() as usize;

            let start = time::precise_time_s();

            let nodes = GraphMMap::new(&filename).nodes();

            let mut segments = SegmentList::new(1024); // list of edge segments

            let mut src = vec![];   // holds ranks
            let mut deg = vec![];   // holds source degrees
            let mut rev = vec![];   // holds (dst, deg) pairs
            let mut trn = vec![];   // holds transposed sources

            let mut going = start;

            let mut input = root.scoped(|builder| {

                let (input, edges) = builder.new_input::<(u32, u32)>();
                let (cycle, ranks) = builder.loop_variable::<(u32, f32)>(20, 1);

                let mut ranks = edges.binary_notify(&ranks,
                                    Exchange::new(|x: &(u32,u32)| x.0 as u64),
                                    Exchange::new(|x: &(u32,f32)| x.0 as u64),
                                    "pagerank",
                                    vec![RootTimestamp::new(0)],
                                    move |input1, input2, output, notificator| {

                    // receive incoming edges (should only be iter 0)
                    while let Some((_index, data)) = input1.next() {
                        segments.push(data.drain_temp());
                    }

                    // all inputs received for iter, commence multiplication
                    while let Some((iter, _)) = notificator.next() {

                        let now = time::now();

                        if index == 0 { println!("{}:{}:{}.{} starting iteration {}", now.tm_hour, now.tm_min, now.tm_sec, now.tm_nsec, iter.inner); }

                        // if the very first iteration, prepare some stuff.
                        // specifically, transpose edges and sort by destination.
                        if iter.inner == 0 {
                            let (a, b, c) = transpose(segments.finalize(), peers, nodes);
                            deg = a; rev = b; trn = c;
                            src = vec![0.0f32; deg.len()];
                        }

                        // record some timings in order to estimate per-iteration times
                        if iter.inner == 0  { println!("src: {}, dst: {}, edges: {}", src.len(), rev.len(), trn.len()); }
                        if iter.inner == 10 && index == 0 { going = time::precise_time_s(); }
                        if iter.inner == 20 && index == 0 { println!("average: {}", (time::precise_time_s() - going) / 10.0 ); }

                        // prepare src for transmitting to destinations
                        for s in 0..src.len() { src[s] = 0.15 + 0.85 * src[s] / deg[s] as f32; }

                        // wander through destinations
                        let mut trn_slice = &trn[..];
                        let mut session = output.session(&iter);
                        for &(dst, deg) in &rev {
                            let mut accum = 0.0;
                            for &s in &trn_slice[..deg as usize] {
                                // accum += src[s as usize];
                                unsafe { accum += *src.get_unchecked(s as usize); }
                            }
                            trn_slice = &trn_slice[deg as usize..];
                            session.give((dst, accum));
                        }

                        for s in &mut src { *s = 0.0; }
                    }

                    // receive data from workers, accumulate in src
                    while let Some((iter, data)) = input2.next() {
                        notificator.notify_at(&iter);
                        for &(node, rank) in data.iter() {
                            src[node as usize / peers] += rank;
                        }
                    }
                });

                // optionally, do process-local accumulation
                if strategy {
                    let local_base = workers * (index / workers);
                    let local_index = index % workers;
                    let mut acc = vec![0.0; (nodes / workers) + 1];   // holds ranks
                    ranks = ranks.unary_notify(
                        Exchange::new(move |x: &(u32,f32)| (local_base as u64 + (x.0 as u64 % workers as u64))),
                        "aggregation",
                        vec![],
                        move |input, output, iterator| {
                            while let Some((iter, data)) = input.next() {
                                iterator.notify_at(&iter);
                                for &(node, rank) in data.iter() {
                                    acc[node as usize / workers] += rank;
                                }
                            }

                            while let Some((item, _)) = iterator.next() {
                                output.session(&item)
                                      .give_iterator(acc.drain_temp()
                                                        .enumerate()
                                                        .filter(|x| x.1 != 0.0)
                                                        .map(|(u,f)| ((u * workers + local_index) as u32, f)));

                                for _ in 0..(1 + (nodes/workers)) { acc.push(0.0); }
                            }
                        }
                    );
                }

                ranks.connect_loop(cycle);

                input
            });

            // introduce edges into the computation;
            // allow mmaped file to drop
            {
                let graph = GraphMMap::new(&filename);
                for node in 0..graph.nodes() {
                    if node % peers == index {
                        for dst in graph.edges(node) {
                            input.send((node as u32, *dst as u32));
                        }
                    }
                }
            }
            input.close();
            while root.step() { }

            if index == 0 { println!("elapsed: {}", time::precise_time_s() - start); }
        });
    }
    else {
        println!("error parsing arguments");
        println!("usage:\tpagerank <source> (worker|process) [timely options]");
    }
}

// returns [src/peers] degrees, (dst, deg) pairs, and a list of [src/peers] endpoints
fn transpose(mut edges: Vec<Vec<(u32, u32)>>, peers: usize, nodes: usize) -> (Vec<u32>, Vec<(u32, u32)>, Vec<u32>)  {

    let mut deg = vec![0; (nodes as usize / peers) + 1];
    for list in &edges {
        for &(s, _) in list {
            deg[s as usize / peers] += 1;
        }
    }

    radix_sort_32(&mut edges, &mut Vec::new(), &|&(_,d)| d);

    let mut rev = Vec::<(u32,u32)>::with_capacity(deg.len());
    let mut trn = Vec::with_capacity(edges.len() * 1024);
    for list in edges {
        for (s,d) in list {
            let len = rev.len();
            if (len == 0) || (rev[len-1].0 < d) {
                rev.push((d, 0u32));
            }

            let len = rev.len();
            rev[len-1].1 += 1;
            trn.push(s / peers as u32);
        }
    }

    return (deg, rev, trn);
}
<|endoftext|>


<commit_before>
#[derive(Debug,Copy,Clone)]
enum Raw {
  // Numbers in this language can be simple integers or rational.
  // They are the same to the user and automatically switch to the best
  // type on every calculation.
  Integer(i64),
  Rational(i64, i64),
}

type Value = Result<Raw, &'static str>;

// Calculate the greatest common divisor using Euclid's algorithm.
fn gcd(mut a: i64, mut b: i64) -> i64 {
    let mut c;
    while a != 0 {
        c = a;
        a = b % a;
        b = c;
    }
    b
}

#[inline(always)]
fn new_integer(n: i64) -> Value {
    Ok(Raw::Integer(n))
}

// This normalizes inputs and sometimes even converts to integers.
fn new_rational(mut n: i64, mut d: i64) -> Value {
    let g = gcd(n, d);
    n = n / g;
    d = d / g;
    if d < 0 {
      n = -n;
      d = -d;
    }
    if d == 1 {
        Ok(Raw::Integer(n))
    }
    else {
        Ok(Raw::Rational(n, d))
    }
}

// We could just let new_rational handle the gcd here, but this
// avoids overflows by dividing the gcd in first.
fn rational_add(n1: i64, d1: i64, n2: i64, d2: i64) -> Value {
  if d1 == d2 {
    // Fast path for common divisors.
    new_rational(n1 + n2, d1)
  }
  else {
    let g = gcd(d1, d2);
    new_rational(n1 * (d2 / g) + n2 * (d1 / g), d1 / g * d2)
  }
}

fn add(left: Value, right: Value) -> Value {
    match (try!(left), try!(right)) {
        (Raw::Integer(a), Raw::Integer(b)) => new_integer(a + b),
        (Raw::Rational(n, d), Raw::Integer(b)) => rational_add(n, d, b, 1),
        (Raw::Integer(a), Raw::Rational(n, d)) => rational_add(a, 1, n, d),
        (Raw::Rational(n1, d1), Raw::Rational(n2, d2)) => rational_add(n1, d1, n2, d2)
    }
}

fn subtract(left: Value, right: Value) -> Value {
    match (try!(left), try!(right)) {
        (Raw::Integer(a), Raw::Integer(b)) => new_integer(a - b),
        (Raw::Rational(n, d), Raw::Integer(b)) => rational_add(n, d, -b, 1),
        (Raw::Integer(a), Raw::Rational(n, d)) => rational_add(a, 1, -n, d),
        (Raw::Rational(n1, d1), Raw::Rational(n2, d2)) => rational_add(n1, d1, -n2, d2)
    }
}

fn divide(left: Value, right: Value) -> Value {
    match (try!(left), try!(right)) {
        (Raw::Integer(a), Raw::Integer(b)) => new_rational(a, b),
        (Raw::Rational(n, d), Raw::Integer(b)) => new_rational(n, d * b),
        (Raw::Integer(a), Raw::Rational(n, d)) => new_rational(a * d, n),
        (Raw::Rational(n1, d1), Raw::Rational(n2, d2)) => new_rational(n1 * d2, d1 * n2),
    }
}

fn multiply(left: Value, right: Value) -> Value {
    match (try!(left), try!(right)) {
        (Raw::Integer(a), Raw::Integer(b)) => new_integer(a * b),
        (Raw::Rational(n, d), Raw::Integer(b)) => new_rational(n * b, d),
        (Raw::Integer(a), Raw::Rational(n, d)) => new_rational(a * n, d),
        (Raw::Rational(n1, d1), Raw::Rational(n2, d2)) => new_rational(n1 * n2, d1 * d2),
    }
}

fn test(a: Value, b: Value) {
    println!("{:?} / {:?} = {:?}", a, b, divide(a, b));
    println!("{:?} * {:?} = {:?}", a, b, multiply(a, b));
    println!("{:?} + {:?} = {:?}", a, b, add(a, b));
    println!("{:?} - {:?} = {:?}", a, b, subtract(a, b));
}

fn main() {
    test(new_integer(44), new_integer(14));
    test(new_integer(14), new_integer(44));
    test(new_integer(44), new_integer(-14));
    test(new_integer(14), new_integer(-44));
    test(new_integer(-44), new_integer(14));
    test(new_integer(-14), new_integer(44));
    test(new_integer(-44), new_integer(-14));
    test(new_integer(-14), new_integer(-44));
    test(new_rational(1, 2), new_rational(2, 1));
    test(new_rational(2, 1), new_rational(1, 2));
    test(new_rational(1, 2), new_rational(1, 2));
    test(new_rational(2, 2), new_rational(2, 1));
}
<commit_msg>Use borrowed values instead of copied values<commit_after>
#[derive(Debug)]
enum Value {
  // Numbers in this language can be simple integers or rational.
  // They are the same to the user and automatically switch to the best
  // type on every calculation.
  Integer(i64),
  Rational(i64, i64),
  String(String)
}
type Error = &'static str;

// Calculate the greatest common divisor using Euclid's algorithm.
fn gcd(mut a: i64, mut b: i64) -> i64 {
    let mut c;
    while a != 0 {
        c = a;
        a = b % a;
        b = c;
    }
    b
}

#[inline(always)]
fn new_integer(n: i64) -> Value {
    Value::Integer(n)
}

// This normalizes inputs and sometimes even converts to integers.
fn new_rational(mut n: i64, mut d: i64) -> Value {
    let g = gcd(n, d);
    n = n / g;
    d = d / g;
    if d < 0 {
      n = -n;
      d = -d;
    }
    if d == 1 {
        Value::Integer(n)
    }
    else {
        Value::Rational(n, d)
    }
}

// We could just let new_rational handle the gcd here, but this
// avoids overflows by dividing the gcd in first.
fn rational_add(n1: i64, d1: i64, n2: i64, d2: i64) -> Value {
  if d1 == d2 {
    // Fast path for common divisors.
    new_rational(n1 + n2, d1)
  }
  else {
    let g = gcd(d1, d2);
    new_rational(n1 * (d2 / g) + n2 * (d1 / g), d1 / g * d2)
  }
}

fn add(left: &Value, right: &Value) -> Result<Value, Error> {
    match (left, right) {
        (&Value::Integer(a), &Value::Integer(b)) => Ok(new_integer(a + b)),
        (&Value::Rational(n, d), &Value::Integer(b)) => Ok(rational_add(n, d, b, 1)),
        (&Value::Integer(a), &Value::Rational(n, d)) => Ok(rational_add(a, 1, n, d)),
        (&Value::Rational(n1, d1), &Value::Rational(n2, d2)) => Ok(rational_add(n1, d1, n2, d2)),
        _ => Err("Add requires two numbers"),
    }
}

fn subtract(left: &Value, right: &Value) -> Result<Value, Error> {
    match (left, right) {
        (&Value::Integer(a), &Value::Integer(b)) => Ok(new_integer(a - b)),
        (&Value::Rational(n, d), &Value::Integer(b)) => Ok(rational_add(n, d, -b, 1)),
        (&Value::Integer(a), &Value::Rational(n, d)) => Ok(rational_add(a, 1, -n, d)),
        (&Value::Rational(n1, d1), &Value::Rational(n2, d2)) => Ok(rational_add(n1, d1, -n2, d2)),
        _ => Err("Subtract requires two numbers"),
    }
}

fn divide(left: &Value, right: &Value) -> Result<Value, Error> {
    match (left, right) {
        (&Value::Integer(a), &Value::Integer(b)) => Ok(new_rational(a, b)),
        (&Value::Rational(n, d), &Value::Integer(b)) => Ok(new_rational(n, d * b)),
        (&Value::Integer(a), &Value::Rational(n, d)) => Ok(new_rational(a * d, n)),
        (&Value::Rational(n1, d1), &Value::Rational(n2, d2)) => Ok(new_rational(n1 * d2, d1 * n2)),
        _ => Err("Divide requires two numbers"),
    }
}

fn multiply(left: &Value, right: &Value) -> Result<Value, Error> {
    match (left, right) {
        (&Value::Integer(a), &Value::Integer(b)) => Ok(new_integer(a * b)),
        (&Value::Rational(n, d), &Value::Integer(b)) => Ok(new_rational(n * b, d)),
        (&Value::Integer(a), &Value::Rational(n, d)) => Ok(new_rational(a * n, d)),
        (&Value::Rational(n1, d1), &Value::Rational(n2, d2)) => Ok(new_rational(n1 * n2, d1 * d2)),
        _ => Err("Multiply requires two numbers"),
    }
}

fn test(a: Value, b: Value) {
    println!("{:?} / {:?} = {:?}", a, b, divide(&a, &b));
    println!("{:?} * {:?} = {:?}", a, b, multiply(&a, &b));
    println!("{:?} + {:?} = {:?}", a, b, add(&a, &b));
    println!("{:?} - {:?} = {:?}", a, b, subtract(&a, &b));
}

fn main() {
    test(new_integer(44), new_integer(14));
    test(new_integer(14), new_integer(44));
    test(new_integer(44), new_integer(-14));
    test(new_integer(14), new_integer(-44));
    test(new_integer(-44), new_integer(14));
    test(new_integer(-14), new_integer(44));
    test(new_integer(-44), new_integer(-14));
    test(new_integer(-14), new_integer(-44));
    test(new_rational(1, 2), new_rational(2, 1));
    test(new_rational(2, 1), new_rational(1, 2));
    test(new_rational(1, 2), new_rational(1, 2));
    test(new_rational(2, 2), new_rational(2, 1));
    test(Value::String(str::to_string("Input string")),
         new_rational(2, 1));
}
<|endoftext|>


<commit_before>use std::io::prelude::*;
use std::fs::File;
mod configuration;
mod shell;
mod command;
mod output;

fn main() {
    println!("sokan");
    let mut fileResult :std::fs::File;
    match File::open("default.yaml") {
        Ok(result) => fileResult = result,
        Err(error) => {
            println!("Error occured while accessing default.yaml: {:?}", error);
            return;
        }
    }

    let mut yaml_file = String::new();
    match fileResult.read_to_string(&mut yaml_file) {
        Err(error) => println!("fail"),
        _ => { }
    }

    let mut commands = configuration::from_yaml(yaml_file.to_string());
    commands.reverse();
    while let Some(cmd) = commands.pop() {
        let shellout = shell::run_command(&cmd);
        output::print_shellout(&cmd, &shellout);
    }
}
<commit_msg>end provisioning if one command fails<commit_after>use std::io::prelude::*;
use std::fs::File;
use std::process;
mod configuration;
mod shell;
mod command;
mod output;

fn main() {
    println!("sokan");
    let mut fileResult :std::fs::File;
    match File::open("default.yaml") {
        Ok(result) => fileResult = result,
        Err(error) => {
            println!("Error occured while accessing default.yaml: {:?}", error);
            return;
        }
    }

    let mut yaml_file = String::new();
    match fileResult.read_to_string(&mut yaml_file) {
        Err(error) => println!("fail"),
        _ => { }
    }

    let mut commands = configuration::from_yaml(yaml_file.to_string());
    commands.reverse();

    while let Some(cmd) = commands.pop() {
        let shellout = shell::run_command(&cmd);
        output::print_shellout(&cmd, &shellout);
        let exit_status = shellout.status.clone();
        if exit_status.success() == false {
            let code = exit_status.code().unwrap();
            process::exit(code);
        }
    }
}
<|endoftext|>


<commit_before>//!
//! # WebSocket Serial Server
//!
//! WebSocket Serial Server is a program that allows
//! that allows browsers to access serial ports
//! on localhost
//!
//! ## Running
//!
//! ```./wsss```
//!
//! You can also specift a port:
//!
//! ```./wsss -p PORT_NUM``` or ```./wsss --port PORT_NUM```
//!
//! When you specify a port, a simple http server will
//! be bound at "/" on PORT_NUM, and a websocket bound at "/"
//! on PORT_NUM + 1
//!
//! The webpage served at "/" on PORT_NUM provides a convenient
//! way for testing serial port connectivity
//!
//! You can open/close ports, read and write data, and see the
//! responses. All messages are in JSON format

extern crate argparse;
extern crate base64;
extern crate env_logger;
extern crate hyper;
#[macro_use]
extern crate log;
extern crate rand;
extern crate serde_json;
extern crate serialport;
extern crate websocket;

extern crate serial_support;

use std::io::Write;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender};
use std::thread;

use argparse::{ArgumentParser, Store};
use websocket::client::Writer;
use websocket::result::WebSocketError;
use websocket::message::Type;
use websocket::{Server, Message};
use websocket::server::upgrade::WsUpgrade;
use hyper::Server as HttpServer;
use hyper::net::Fresh;
use hyper::server::request::Request;
use hyper::server::response::Response;
use rand::{Rng, thread_rng};

use serial_support::dynamic_sleep::DynamicSleep;
use serial_support::errors as e;
use serial_support::manage::Manager;
use serial_support::messages::*;

/// Max number of failures we allow when trying to send
/// data to client before exiting
/// TODO: Make configurable
pub const MAX_SEND_ERROR_COUNT: u32 = 5;

/// Launches wsss
pub fn main() {

  // Init logger
  env_logger::init().unwrap();

  // Default port number
  let mut port = 8080;
  // Parse cmdline args
  {
    let mut ap = ArgumentParser::new();
    ap.set_description("Provide access to serial ports over JSON Websockets");
    ap.refer(&mut port)
      .add_option(&["-p", "--port"], Store, "Http Port");
    ap.parse_args_or_exit();
  }

  // websocket port
  let ws_port = port + 1;

  // html file for landing page
  let websocket_html = include_str!("websockets.html").replace(
    "__WS_PORT__ = 8081",
    &format!("__WS_PORT__ = {}", ws_port),
  );

  // The HTTP server handler
  let http_handler = move |_: Request, response: Response<Fresh>| {
    let mut response = response.start().expect(&"Could not start response");
    // Send a client webpage
    response
      .write_all(websocket_html.as_bytes())
      .expect(&"Could not get template as bytes");
    response.end().expect(&"Send response failed");
  };

  info!("Using ports {} {}", port, ws_port);

  // Set up channels and Manager
  let (sub_tx, sub_rx) = channel::<SubscriptionRequest>();
  let (sreq_tx, sreq_rx) = channel::<(String, SerialRequest)>();
  Manager::spawn(sreq_rx, sub_rx);

  // Start listening for http connections
  thread::spawn(
    move || {
      let http_server = HttpServer::http(format!("127.0.0.1:{}", port)).expect(
        &format!(
          "Failed to create http server on port {}",
          port
        ),
      );
      http_server
        .handle(http_handler)
        .expect(&"Failed to listen");
    },
  );

  // Start listening for WebSocket connections
  let ws_server = Server::bind(format!("127.0.0.1:{}", ws_port)).expect(
    &format!(
      "Failed bind on websocket port {}",
      ws_port
    ),
  );

  // Continuously iterate over connections,
  // spawning handlers
  for connection in ws_server.filter_map(Result::ok) {
    // Set up subscription id
    // let ts = SystemTime::now() - UNIX_EPOCH
    let prefix: String = thread_rng().gen_ascii_chars().take(8).collect();
    let sub_id = format!("thread-{}-{}", prefix, rand::random::<u16>());
    debug!("{}: spawned.", sub_id);

    // Spawn a new thread for each connection.
    let sub_tx_clone = sub_tx.clone();
    let sreq_tx_clone = sreq_tx.clone();
    spawn_ws_handler(sub_id, sub_tx_clone, sreq_tx_clone, connection);
  }
}

/// Spawn a websocket handler into its own thread
fn spawn_ws_handler(
  sub_id: String,
  sub_tx_clone: Sender<SubscriptionRequest>,
  sreq_tx_clone: Sender<(String, SerialRequest)>,
  connection: WsUpgrade<TcpStream>,
) {
  thread::spawn(move || ws_handler(sub_id, &sub_tx_clone, &sreq_tx_clone, connection));
}

/// Websocket handler
fn ws_handler(
  sub_id: String,
  sub_tx: &Sender<SubscriptionRequest>,
  sreq_tx: &Sender<(String, SerialRequest)>,
  connection: WsUpgrade<TcpStream>,
) {

  if !connection
        .protocols()
        .contains(&"websocket-serial-json".to_string()) {
    connection
      .reject()
      .expect(&"Connection rejection failed.");
    return;
  }

  connection
    .tcp_stream()
    .set_nonblocking(true)
    .expect(&"Setting stream non-blocking failed.");

  // Create response channel
  let (sub_resp_tx, sub_resp_rx) = channel::<SerialResponse>();

  // Register sub_id with manager
  sub_tx
    .send(
      SubscriptionRequest {
        sub_id: sub_id.clone(),
        subscriber: sub_resp_tx,
      },
    )
    .expect(&format!("{}: Registering with manager failed.", sub_id));

  let client = connection
    .use_protocol(format!("websocket-serial-json"))
    .accept()
    .expect(&format!("{}: Accept protocol failed.", sub_id));

  let ip = client
    .peer_addr()
    .expect(&format!("{}: Could not get peer address", sub_id));

  info!("{}: Connection from {}", sub_id, ip);

  let (mut receiver, mut sender) = client
    .split()
    .expect(&format!("{}: WS client error", sub_id));

  let mut send_error_count = 0;

  let mut dynamic_sleep = DynamicSleep::new("main");

  'msg_loop: loop {

    dynamic_sleep.sleep();

    // Try and read a WS message
    match receiver.recv_message::<Message, _, _>() {
      Ok(message) => {
        match message.opcode {

          Type::Close => {
            let message = Message::close();
            sender
              .send_message(&message)
              .unwrap_or(info!("{}: Client {} hung up!", sub_id, ip));
            // Send close request to cleanup resources
            sreq_tx
              .send((sub_id.clone(), SerialRequest::Close { port: None }))
              .unwrap_or_else(
                |e| {
                  warn!(
                    "Client exit cleanup failed for sub_id '{}', cause '{}'",
                    sub_id,
                    e
                  )
                },
              );
            info!("{}: Client {} disconnected", sub_id, ip);
            break 'msg_loop;
          }

          Type::Ping => {
            let message = Message::pong(message.payload);
            sender
              .send_message(&message)
              .unwrap_or(info!("{}:  Could not ping client {}!", sub_id, ip));
          }

          _ => {

            // Get the payload, in a lossy manner
            let msg = String::from_utf8_lossy(&message.payload);

            // So we will get a result <SerialRequest::*,SerialResponse::Error> back
            match serde_json::from_str(&msg) {
              Ok(req) => {
                match sreq_tx.send((sub_id.clone(), req)) {
                  Err(err) => {
                    let error = e::ErrorKind::SendRequest(err).into();
                    send_serial_response_error(&sub_id, &mut sender, error);

                  }
                  _ => {}
                };
              }
              Err(err) => {
                let error = e::ErrorKind::Json(err).into();
                send_serial_response_error(&sub_id, &mut sender, error);
              }
            };
          }
        };
      }
      Err(e) => {
        match e {
          WebSocketError::NoDataAvailable => { /*Logging?*/ }
          _ => { /*Logging?*/ }
        };
      }
    }

    // Send on any serial responses
    match sub_resp_rx.try_recv() {
      Ok(resp) => {
        match serde_json::to_string(&resp) {
          Ok(json) => {
            let reply = Message::text(json.clone());
            sender
              .send_message(&reply)
              .unwrap_or_else(
                |e| {
                  send_error_count += 1;
                  info!(
                    "{}: Could not send message '{}' to client '{}', cause '{}'",
                    sub_id,
                    json,
                    ip,
                    e
                  )
                },
              );
          }
          Err(_) => {}
        }
      }
      _ => { /*Logging*/ }
    };

    if send_error_count > MAX_SEND_ERROR_COUNT {
      warn!(
        "{}: Client send error count exceeded! Shutting down msg loop.",
        &sub_id
      );
      break 'msg_loop;
    }
  }

  info!("{}: Shutting down!", sub_id);

}

/// Send an error to the given subscriber
/// Log a warning if the message can't be sent
/// This is usually ok as it means the client
/// has simply disconnected
fn send_serial_response_error(sub_id: &String, sender: &mut Writer<TcpStream>, error: e::Error) {
  let error = e::to_serial_response_error(error);
  serde_json::to_string(&error)
    .map_err(|err| e::ErrorKind::Json(err))
    .map(|json| Message::text(json))
    .map(
      |msg| {
        sender
          .send_message::<Message, _>(&msg)
          .map_err::<e::Error, _>(|err| e::ErrorKind::SendWsMessage(err).into())
      },
    )
    .unwrap_or_else(
      |_| {
        warn!("{}: Problem sending bad json error response", sub_id);
        Ok(())
      },
    )
    .is_ok(); // This shouldn't be needed?

}
<commit_msg>Move allocation of http server out of spawn<commit_after>//!
//! # WebSocket Serial Server
//!
//! WebSocket Serial Server is a program that allows
//! that allows browsers to access serial ports
//! on localhost
//!
//! ## Running
//!
//! ```./wsss```
//!
//! You can also specift a port:
//!
//! ```./wsss -p PORT_NUM``` or ```./wsss --port PORT_NUM```
//!
//! When you specify a port, a simple http server will
//! be bound at "/" on PORT_NUM, and a websocket bound at "/"
//! on PORT_NUM + 1
//!
//! The webpage served at "/" on PORT_NUM provides a convenient
//! way for testing serial port connectivity
//!
//! You can open/close ports, read and write data, and see the
//! responses. All messages are in JSON format

extern crate argparse;
extern crate base64;
extern crate env_logger;
extern crate hyper;
#[macro_use]
extern crate log;
extern crate rand;
extern crate serde_json;
extern crate serialport;
extern crate websocket;

extern crate serial_support;

use std::io::Write;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender};
use std::thread;

use argparse::{ArgumentParser, Store};
use websocket::client::Writer;
use websocket::result::WebSocketError;
use websocket::message::Type;
use websocket::{Server, Message};
use websocket::server::upgrade::WsUpgrade;
use hyper::Server as HttpServer;
use hyper::net::Fresh;
use hyper::server::request::Request;
use hyper::server::response::Response;
use rand::{Rng, thread_rng};

use serial_support::dynamic_sleep::DynamicSleep;
use serial_support::errors as e;
use serial_support::manage::Manager;
use serial_support::messages::*;

/// Max number of failures we allow when trying to send
/// data to client before exiting
/// TODO: Make configurable
pub const MAX_SEND_ERROR_COUNT: u32 = 5;

/// Launches wsss
pub fn main() {

  // Init logger
  env_logger::init().unwrap();

  // Default port number
  let mut port = 8080;
  // Parse cmdline args
  {
    let mut ap = ArgumentParser::new();
    ap.set_description("Provide access to serial ports over JSON Websockets");
    ap.refer(&mut port)
      .add_option(&["-p", "--port"], Store, "Http Port");
    ap.parse_args_or_exit();
  }

  // websocket port
  let ws_port = port + 1;

  // html file for landing page
  let websocket_html = include_str!("websockets.html").replace(
    "__WS_PORT__ = 8081",
    &format!("__WS_PORT__ = {}", ws_port),
  );

  // The HTTP server handler
  let http_handler = move |_: Request, response: Response<Fresh>| {
    let mut response = response.start().expect(&"Could not start response");
    // Send a client webpage
    response
      .write_all(websocket_html.as_bytes())
      .expect(&"Could not get template as bytes");
    response.end().expect(&"Send response failed");
  };

  info!("Using ports {} {}", port, ws_port);

  // Set up channels and Manager
  let (sub_tx, sub_rx) = channel::<SubscriptionRequest>();
  let (sreq_tx, sreq_rx) = channel::<(String, SerialRequest)>();
  Manager::spawn(sreq_rx, sub_rx);

  // Start listening for http connections
  let http_server = HttpServer::http(format!("127.0.0.1:{}", port)).expect(
    &format!(
      "Failed to create http server on port {}",
      port
    ),
  );

  thread::spawn(
    move || {
      http_server
        .handle(http_handler)
        .expect(&"Failed to listen");
    },
  );

  // Start listening for WebSocket connections
  let ws_server = Server::bind(format!("127.0.0.1:{}", ws_port)).expect(
    &format!(
      "Failed bind on websocket port {}",
      ws_port
    ),
  );

  // Continuously iterate over connections,
  // spawning handlers
  for connection in ws_server.filter_map(Result::ok) {
    // Set up subscription id
    // let ts = SystemTime::now() - UNIX_EPOCH
    let prefix: String = thread_rng().gen_ascii_chars().take(8).collect();
    let sub_id = format!("thread-{}-{}", prefix, rand::random::<u16>());
    debug!("{}: spawned.", sub_id);

    // Spawn a new thread for each connection.
    let sub_tx_clone = sub_tx.clone();
    let sreq_tx_clone = sreq_tx.clone();
    spawn_ws_handler(sub_id, sub_tx_clone, sreq_tx_clone, connection);
  }
}

/// Spawn a websocket handler into its own thread
fn spawn_ws_handler(
  sub_id: String,
  sub_tx_clone: Sender<SubscriptionRequest>,
  sreq_tx_clone: Sender<(String, SerialRequest)>,
  connection: WsUpgrade<TcpStream>,
) {
  thread::spawn(move || ws_handler(sub_id, &sub_tx_clone, &sreq_tx_clone, connection));
}

/// Websocket handler
fn ws_handler(
  sub_id: String,
  sub_tx: &Sender<SubscriptionRequest>,
  sreq_tx: &Sender<(String, SerialRequest)>,
  connection: WsUpgrade<TcpStream>,
) {

  if !connection
        .protocols()
        .contains(&"websocket-serial-json".to_string()) {
    connection
      .reject()
      .expect(&"Connection rejection failed.");
    return;
  }

  connection
    .tcp_stream()
    .set_nonblocking(true)
    .expect(&"Setting stream non-blocking failed.");

  // Create response channel
  let (sub_resp_tx, sub_resp_rx) = channel::<SerialResponse>();

  // Register sub_id with manager
  sub_tx
    .send(
      SubscriptionRequest {
        sub_id: sub_id.clone(),
        subscriber: sub_resp_tx,
      },
    )
    .expect(&format!("{}: Registering with manager failed.", sub_id));

  let client = connection
    .use_protocol(format!("websocket-serial-json"))
    .accept()
    .expect(&format!("{}: Accept protocol failed.", sub_id));

  let ip = client
    .peer_addr()
    .expect(&format!("{}: Could not get peer address", sub_id));

  info!("{}: Connection from {}", sub_id, ip);

  let (mut receiver, mut sender) = client
    .split()
    .expect(&format!("{}: WS client error", sub_id));

  let mut send_error_count = 0;

  let mut dynamic_sleep = DynamicSleep::new("main");

  'msg_loop: loop {

    dynamic_sleep.sleep();

    // Try and read a WS message
    match receiver.recv_message::<Message, _, _>() {
      Ok(message) => {
        match message.opcode {

          Type::Close => {
            let message = Message::close();
            sender
              .send_message(&message)
              .unwrap_or(info!("{}: Client {} hung up!", sub_id, ip));
            // Send close request to cleanup resources
            sreq_tx
              .send((sub_id.clone(), SerialRequest::Close { port: None }))
              .unwrap_or_else(
                |e| {
                  warn!(
                    "Client exit cleanup failed for sub_id '{}', cause '{}'",
                    sub_id,
                    e
                  )
                },
              );
            info!("{}: Client {} disconnected", sub_id, ip);
            break 'msg_loop;
          }

          Type::Ping => {
            let message = Message::pong(message.payload);
            sender
              .send_message(&message)
              .unwrap_or(info!("{}:  Could not ping client {}!", sub_id, ip));
          }

          _ => {

            // Get the payload, in a lossy manner
            let msg = String::from_utf8_lossy(&message.payload);

            // So we will get a result <SerialRequest::*,SerialResponse::Error> back
            match serde_json::from_str(&msg) {
              Ok(req) => {
                match sreq_tx.send((sub_id.clone(), req)) {
                  Err(err) => {
                    let error = e::ErrorKind::SendRequest(err).into();
                    send_serial_response_error(&sub_id, &mut sender, error);

                  }
                  _ => {}
                };
              }
              Err(err) => {
                let error = e::ErrorKind::Json(err).into();
                send_serial_response_error(&sub_id, &mut sender, error);
              }
            };
          }
        };
      }
      Err(e) => {
        match e {
          WebSocketError::NoDataAvailable => { /*Logging?*/ }
          _ => { /*Logging?*/ }
        };
      }
    }

    // Send on any serial responses
    match sub_resp_rx.try_recv() {
      Ok(resp) => {
        match serde_json::to_string(&resp) {
          Ok(json) => {
            let reply = Message::text(json.clone());
            sender
              .send_message(&reply)
              .unwrap_or_else(
                |e| {
                  send_error_count += 1;
                  info!(
                    "{}: Could not send message '{}' to client '{}', cause '{}'",
                    sub_id,
                    json,
                    ip,
                    e
                  )
                },
              );
          }
          Err(_) => {}
        }
      }
      _ => { /*Logging*/ }
    };

    if send_error_count > MAX_SEND_ERROR_COUNT {
      warn!(
        "{}: Client send error count exceeded! Shutting down msg loop.",
        &sub_id
      );
      break 'msg_loop;
    }
  }

  info!("{}: Shutting down!", sub_id);

}

/// Send an error to the given subscriber
/// Log a warning if the message can't be sent
/// This is usually ok as it means the client
/// has simply disconnected
fn send_serial_response_error(sub_id: &String, sender: &mut Writer<TcpStream>, error: e::Error) {
  let error = e::to_serial_response_error(error);
  serde_json::to_string(&error)
    .map_err(|err| e::ErrorKind::Json(err))
    .map(|json| Message::text(json))
    .map(
      |msg| {
        sender
          .send_message::<Message, _>(&msg)
          .map_err::<e::Error, _>(|err| e::ErrorKind::SendWsMessage(err).into())
      },
    )
    .unwrap_or_else(
      |_| {
        warn!("{}: Problem sending bad json error response", sub_id);
        Ok(())
      },
    )
    .is_ok(); // This shouldn't be needed?

}
<|endoftext|>


<commit_before>extern crate ansi_term;
extern crate clap;
extern crate config;
extern crate git2;

mod modules;

use ansi_term::{ANSIString, ANSIStrings};
use clap::{App, Arg};
use config::{Config, File, FileFormat};
use std::path::PathBuf;

fn main() {
    let matches = App::new("contrail")
        .version("0.1.0")
        .about("Fast and configurable shell prompter")
        .arg(Arg::with_name("exit_code")
                 .short("e")
                 .long("exit_code")
                 .value_name("CODE")
                 .help("The exit code of the last-executed command")
                 .takes_value(true))
        .arg(Arg::with_name("config")
                 .short("c")
                 .long("config")
                 .value_name("FILE")
                 .help("The configuration file")
                 .takes_value(true))
        .arg(Arg::with_name("zsh").short("z").long("zsh").help("Enable ZSH mode"))
        .get_matches();

    let mut c = Config::new();

    // Merge in the defaults
    modules::merge_defaults(&mut c);

    // Merge in the config file
    if let Some(f) = matches.value_of("config") {
        c.merge(File::new(f, FileFormat::Toml).required(false))
            .expect("Unable to read the config file!");
    } else if let Some(f) = find_config_file() {
        c.merge(File::new(f.to_str().unwrap(), FileFormat::Toml).required(false))
            .expect("Unable to read the config file!");
    }

    // Merge in the command-line arguments
    if matches.is_present("zsh") {
        c.set("global.shell", "zsh").unwrap();
    }

    // A few things to note:
    // - Calling `c.get()` is a workaround, `c.get_array()` consumes `self`
    // - We do `rev()` because it's easier to handle things like finding
    //   the next module's background color when you go in backwards order
    let module_names: Vec<String> = c.get("global.modules")
        .unwrap()
        .into_array()
        .unwrap_or_default()
        .into_iter()
        .map(|m| m.into_str().unwrap())
        .rev()
        .collect();

    let mut ansi_strings: Vec<ANSIString<'static>> = Vec::new();

    // The name of the last successfully formatted module
    let mut last_successful_format: Option<&str> = None;

    for name in &module_names {
        let formatted_module: Option<ANSIString<'static>>;

        let result = match name.as_ref() {
            "directory" => modules::format_module_directory(&mut c, last_successful_format),
            "exit_code" => {
                let exit_code = matches.value_of("exit_code").unwrap_or("255");
                modules::format_module_exit_code(&mut c, last_successful_format, exit_code)
            }
            "git" => modules::format_module_git(&mut c, last_successful_format),
            "prompt" => {
                let exit_code = matches.value_of("exit_code").unwrap_or("255");
                modules::format_module_prompt(&mut c, last_successful_format, exit_code)
            }
            _ => {
                // The output will be found in the config file
                modules::format_module(&mut c, name, None, last_successful_format)
            }
        };

        // Updates the last formatted module variable and prints it
        // out only if it was successful
        formatted_module = result.1;
        if result.0.is_some() {
            last_successful_format = result.0;
            ansi_strings.insert(0, formatted_module.unwrap());
        }

    }

    print!("{}", ANSIStrings(ansi_strings.as_slice()));
}

/// Looks through typical directories for a `contrail/config.toml`
// TODO Implement and test this
fn find_config_file() -> Option<PathBuf> {
    // This is NOT the desired behavior, this is just for testing
    let pb = PathBuf::from("config.toml");

    if pb.exists() { Some(pb) } else { None }
}
<commit_msg>refactor: automatically detect program version number<commit_after>extern crate ansi_term;
extern crate clap;
extern crate config;
extern crate git2;

mod modules;

use ansi_term::{ANSIString, ANSIStrings};
use clap::{App, Arg};
use config::{Config, File, FileFormat};
use std::path::PathBuf;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

fn main() {
    let matches = App::new("contrail")
        .version(VERSION)
        .about("Fast and configurable shell prompter")
        .arg(Arg::with_name("exit_code")
                 .short("e")
                 .long("exit_code")
                 .value_name("CODE")
                 .help("The exit code of the last-executed command")
                 .takes_value(true))
        .arg(Arg::with_name("config")
                 .short("c")
                 .long("config")
                 .value_name("FILE")
                 .help("The configuration file")
                 .takes_value(true))
        .arg(Arg::with_name("zsh").short("z").long("zsh").help("Enables ZSH mode"))
        .get_matches();

    let mut c = Config::new();

    // Merge in the defaults
    modules::merge_defaults(&mut c);

    // Merge in the config file
    if let Some(f) = matches.value_of("config") {
        c.merge(File::new(f, FileFormat::Toml).required(false))
            .expect("Unable to read the config file!");
    } else if let Some(f) = find_config_file() {
        c.merge(File::new(f.to_str().unwrap(), FileFormat::Toml).required(false))
            .expect("Unable to read the config file!");
    }

    // Merge in the command-line arguments
    if matches.is_present("zsh") {
        c.set("global.shell", "zsh").unwrap();
    }

    // A few things to note:
    // - Calling `c.get()` is a workaround, `c.get_array()` consumes `self`
    // - We do `rev()` because it's easier to handle things like finding
    //   the next module's background color when you go in backwards order
    let module_names: Vec<String> = c.get("global.modules")
        .unwrap()
        .into_array()
        .unwrap_or_default()
        .into_iter()
        .map(|m| m.into_str().unwrap())
        .rev()
        .collect();

    let mut ansi_strings: Vec<ANSIString<'static>> = Vec::new();

    // The name of the last successfully formatted module
    let mut last_successful_format: Option<&str> = None;

    for name in &module_names {
        let formatted_module: Option<ANSIString<'static>>;

        let result = match name.as_ref() {
            "directory" => modules::format_module_directory(&mut c, last_successful_format),
            "exit_code" => {
                let exit_code = matches.value_of("exit_code").unwrap_or("255");
                modules::format_module_exit_code(&mut c, last_successful_format, exit_code)
            }
            "git" => modules::format_module_git(&mut c, last_successful_format),
            "prompt" => {
                let exit_code = matches.value_of("exit_code").unwrap_or("255");
                modules::format_module_prompt(&mut c, last_successful_format, exit_code)
            }
            _ => {
                // The output will be found in the config file
                modules::format_module(&mut c, name, None, last_successful_format)
            }
        };

        // Updates the last formatted module variable and prints it
        // out only if it was successful
        formatted_module = result.1;
        if result.0.is_some() {
            last_successful_format = result.0;
            ansi_strings.insert(0, formatted_module.unwrap());
        }

    }

    print!("{}", ANSIStrings(ansi_strings.as_slice()));
}

/// Looks through typical directories for a `contrail/config.toml`
// TODO Implement and test this
fn find_config_file() -> Option<PathBuf> {
    // This is NOT the desired behavior, this is just for testing
    let pb = PathBuf::from("config.toml");

    if pb.exists() { Some(pb) } else { None }
}
<|endoftext|>


<commit_before>#![feature(duration_span)]

extern crate image;
extern crate rand;
extern crate docopt;
extern crate rustc_serialize;
extern crate num_cpus;
extern crate scoped_threadpool;
extern crate tray_rust;

use std::vec::Vec;
use std::iter;
use std::path::PathBuf;
use std::time::Duration;
use std::io::ErrorKind;

use rand::StdRng;
use docopt::Docopt;

use tray_rust::film::{self, ImageSample};
use tray_rust::geometry::{Geometry, Instance, Emitter};
use tray_rust::sampler::{self, Sampler};
use tray_rust::scene;
use tray_rust::integrator::Integrator;

static USAGE: &'static str = "
Usage: tray_rust <scenefile> [options]

Options:
  -o <file>     Specify the output file or directory to save the image or frames. Supported formats are
                PNG, JPG and PPM. Default is 'frame<#>.png'.
  -n <number>   Specify the number of threads to use for rendering. Defaults to the number of cores
                on the system.
  -h, --help    Show this message.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_scenefile: String,
    flag_o: Option<String>,
    flag_n: Option<u32>,
}

/// Threads are each sent a sender end of the channel that is
/// read from by the render target thread which then saves the
/// values recieved to the render target
fn thread_work(spp: usize, queue: &sampler::BlockQueue, scene: &scene::Scene,
               target: &film::RenderTarget, light_list: &Vec<&Emitter>) {
    let mut sampler = sampler::LowDiscrepancy::new(queue.block_dim(), spp);
    let mut sample_pos = Vec::with_capacity(sampler.max_spp());
    let mut time_samples: Vec<_> = iter::repeat(0.0).take(sampler.max_spp()).collect();
    let block_dim = queue.block_dim();
    let mut block_samples = Vec::with_capacity(sampler.max_spp() * (block_dim.0 * block_dim.1) as usize);
    let mut rng = match StdRng::new() {
        Ok(r) => r,
        Err(e) => { println!("Failed to get StdRng, {}", e); return }
    };
    // Grab a block from the queue and start working on it, submitting samples
    // to the render target thread after each pixel
    for b in queue.iter() {
        sampler.select_block(b);
        let mut pixel_samples = 0;
        while sampler.has_samples() {
            // Get samples for a pixel and render them
            sampler.get_samples(&mut sample_pos, &mut rng);
            sampler.get_samples_1d(&mut time_samples[..], &mut rng);
            for (s, t) in sample_pos.iter().zip(time_samples.iter()) {
                let mut ray = scene.camera.generate_ray(s, *t);
                if let Some(hit) = scene.intersect(&mut ray) {
                    let c = scene.integrator.illumination(scene, light_list, &ray,
                                                          &hit, &mut sampler, &mut rng).clamp();
                    block_samples.push(ImageSample::new(s.0, s.1, c));
                }
            }
            // If the samples are ok the samples for the next pixel start at the end of the current
            // pixel's samples
            if sampler.report_results(&block_samples[pixel_samples..]) {
                pixel_samples = block_samples.len();
            }
        }
        target.write(&block_samples, sampler.get_region());
        block_samples.clear();
    }
}

/// Render the scene in parallel using `n` threads and write the result to the render target
fn render_parallel(rt: &mut film::RenderTarget, scene: &scene::Scene, n: u32, spp: usize){
    let mut pool = scoped_threadpool::Pool::new(n);
    let dim = rt.dimensions();
    let block_queue = sampler::BlockQueue::new((dim.0 as u32, dim.1 as u32), (8, 8));
    let light_list: Vec<_> = scene.bvh.into_iter().filter_map(|x| {
        match x {
            &Instance::Emitter(ref e) => Some(e),
            _ => None,
        }
    }).collect();
    assert!(!light_list.is_empty(), "At least one light is required");
    pool.scoped(|scope| {
        for _ in 0..n {
            let b = &block_queue;
            let r = &*rt;
            let l = &light_list;
            scope.execute(move || {
                thread_work(spp, b, scene, r, l);
            });
        }
    });
}

fn main() {
    let args: Args = Docopt::new(USAGE).and_then(|d| d.decode()).unwrap_or_else(|e| e.exit());
    let n = match args.flag_n {
        Some(n) => n,
        None => num_cpus::get() as u32,
    };
    let out_path = match &args.flag_o {
        &Some(ref f) => {
            let p = PathBuf::from(f);
            // If we're writing to a directory make sure it exists
            if p.extension() == None {
                match std::fs::create_dir(p.as_path()) {
                    Err(e) => {
                        if e.kind() != ErrorKind::AlreadyExists {
                            panic!("Failed to create output directory");
                        }
                    },
                    Ok(_) => {},
                }
            }
            p
        },
        &None => PathBuf::from("./"),
    };

    let (mut scene, mut rt, spp, frame_info) = scene::Scene::load_file(&args.arg_scenefile[..]);
    let image_dim = rt.dimensions();
    println!("Rendering using {} threads\n--------------------", n);

    let time_step = frame_info.time / frame_info.frames as f32;
    for i in frame_info.start..frame_info.end + 1 {
        let frame_start = i as f32 * time_step;
        let frame_end = (i as f32 + 1.0) * time_step;
        scene.camera.update_shutter(frame_start, frame_end);
        // TODO: How often to re-build the BVH?
        println!("Frame {}: re-building bvh for {} to {}", i, frame_start, frame_end);
        scene.bvh.rebuild(frame_start, frame_end);
        println!("Frame {}: rendering for {} to {}", i, frame_start, frame_end);
        let d = Duration::span(|| render_parallel(&mut rt, &scene, n, spp));
        let time = d.as_secs() as f64 + (d.subsec_nanos() as f64) / 1_000_000_000.0;
        println!("Frame {}: Rendering took {}s", i, time);

        let img = rt.get_render();
        let out_file = match out_path.extension() {
            Some(_) => out_path.clone(),
            None => out_path.join(PathBuf::from(format!("frame{:05}.png", i))),
        };
        match image::save_buffer(&out_file.as_path(), &img[..], image_dim.0 as u32, image_dim.1 as u32, image::RGB(8)) {
            Ok(_) => {},
            Err(e) => println!("Error saving image, {}", e),
        };
        rt.clear();
        println!("Frame {}: Rendered to '{}'\n--------------------", i, out_file.display());
    }
}

<commit_msg>We need to push black samples when we miss for the filtering to work<commit_after>#![feature(duration_span)]

extern crate image;
extern crate rand;
extern crate docopt;
extern crate rustc_serialize;
extern crate num_cpus;
extern crate scoped_threadpool;
extern crate tray_rust;

use std::vec::Vec;
use std::iter;
use std::path::PathBuf;
use std::time::Duration;
use std::io::ErrorKind;

use rand::StdRng;
use docopt::Docopt;

use tray_rust::film::{self, ImageSample, Colorf};
use tray_rust::geometry::{Geometry, Instance, Emitter};
use tray_rust::sampler::{self, Sampler};
use tray_rust::scene;
use tray_rust::integrator::Integrator;

static USAGE: &'static str = "
Usage: tray_rust <scenefile> [options]

Options:
  -o <file>     Specify the output file or directory to save the image or frames. Supported formats are
                PNG, JPG and PPM. Default is 'frame<#>.png'.
  -n <number>   Specify the number of threads to use for rendering. Defaults to the number of cores
                on the system.
  -h, --help    Show this message.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_scenefile: String,
    flag_o: Option<String>,
    flag_n: Option<u32>,
}

/// Threads are each sent a sender end of the channel that is
/// read from by the render target thread which then saves the
/// values recieved to the render target
fn thread_work(spp: usize, queue: &sampler::BlockQueue, scene: &scene::Scene,
               target: &film::RenderTarget, light_list: &Vec<&Emitter>) {
    let mut sampler = sampler::LowDiscrepancy::new(queue.block_dim(), spp);
    let mut sample_pos = Vec::with_capacity(sampler.max_spp());
    let mut time_samples: Vec<_> = iter::repeat(0.0).take(sampler.max_spp()).collect();
    let block_dim = queue.block_dim();
    let mut block_samples = Vec::with_capacity(sampler.max_spp() * (block_dim.0 * block_dim.1) as usize);
    let mut rng = match StdRng::new() {
        Ok(r) => r,
        Err(e) => { println!("Failed to get StdRng, {}", e); return }
    };
    // Grab a block from the queue and start working on it, submitting samples
    // to the render target thread after each pixel
    for b in queue.iter() {
        sampler.select_block(b);
        let mut pixel_samples = 0;
        while sampler.has_samples() {
            // Get samples for a pixel and render them
            sampler.get_samples(&mut sample_pos, &mut rng);
            sampler.get_samples_1d(&mut time_samples[..], &mut rng);
            for (s, t) in sample_pos.iter().zip(time_samples.iter()) {
                let mut ray = scene.camera.generate_ray(s, *t);
                if let Some(hit) = scene.intersect(&mut ray) {
                    let c = scene.integrator.illumination(scene, light_list, &ray,
                                                          &hit, &mut sampler, &mut rng).clamp();
                    block_samples.push(ImageSample::new(s.0, s.1, c));
                } else {
                    block_samples.push(ImageSample::new(s.0, s.1, Colorf::black()));
                }
            }
            // If the samples are ok the samples for the next pixel start at the end of the current
            // pixel's samples
            if sampler.report_results(&block_samples[pixel_samples..]) {
                pixel_samples = block_samples.len();
            }
        }
        target.write(&block_samples, sampler.get_region());
        block_samples.clear();
    }
}

/// Render the scene in parallel using `n` threads and write the result to the render target
fn render_parallel(rt: &mut film::RenderTarget, scene: &scene::Scene, n: u32, spp: usize){
    let mut pool = scoped_threadpool::Pool::new(n);
    let dim = rt.dimensions();
    let block_queue = sampler::BlockQueue::new((dim.0 as u32, dim.1 as u32), (8, 8));
    let light_list: Vec<_> = scene.bvh.into_iter().filter_map(|x| {
        match x {
            &Instance::Emitter(ref e) => Some(e),
            _ => None,
        }
    }).collect();
    assert!(!light_list.is_empty(), "At least one light is required");
    pool.scoped(|scope| {
        for _ in 0..n {
            let b = &block_queue;
            let r = &*rt;
            let l = &light_list;
            scope.execute(move || {
                thread_work(spp, b, scene, r, l);
            });
        }
    });
}

fn main() {
    let args: Args = Docopt::new(USAGE).and_then(|d| d.decode()).unwrap_or_else(|e| e.exit());
    let n = match args.flag_n {
        Some(n) => n,
        None => num_cpus::get() as u32,
    };
    let out_path = match &args.flag_o {
        &Some(ref f) => {
            let p = PathBuf::from(f);
            // If we're writing to a directory make sure it exists
            if p.extension() == None {
                match std::fs::create_dir(p.as_path()) {
                    Err(e) => {
                        if e.kind() != ErrorKind::AlreadyExists {
                            panic!("Failed to create output directory");
                        }
                    },
                    Ok(_) => {},
                }
            }
            p
        },
        &None => PathBuf::from("./"),
    };

    let (mut scene, mut rt, spp, frame_info) = scene::Scene::load_file(&args.arg_scenefile[..]);
    let image_dim = rt.dimensions();
    println!("Rendering using {} threads\n--------------------", n);

    let time_step = frame_info.time / frame_info.frames as f32;
    for i in frame_info.start..frame_info.end + 1 {
        let frame_start = i as f32 * time_step;
        let frame_end = (i as f32 + 1.0) * time_step;
        scene.camera.update_shutter(frame_start, frame_end);
        // TODO: How often to re-build the BVH?
        println!("Frame {}: re-building bvh for {} to {}", i, frame_start, frame_end);
        scene.bvh.rebuild(frame_start, frame_end);
        println!("Frame {}: rendering for {} to {}", i, frame_start, frame_end);
        let d = Duration::span(|| render_parallel(&mut rt, &scene, n, spp));
        let time = d.as_secs() as f64 + (d.subsec_nanos() as f64) / 1_000_000_000.0;
        println!("Frame {}: Rendering took {}s", i, time);

        let img = rt.get_render();
        let out_file = match out_path.extension() {
            Some(_) => out_path.clone(),
            None => out_path.join(PathBuf::from(format!("frame{:05}.png", i))),
        };
        match image::save_buffer(&out_file.as_path(), &img[..], image_dim.0 as u32, image_dim.1 as u32, image::RGB(8)) {
            Ok(_) => {},
            Err(e) => println!("Error saving image, {}", e),
        };
        rt.clear();
        println!("Frame {}: Rendered to '{}'\n--------------------", i, out_file.display());
    }
}

<|endoftext|>


<commit_before>#![allow(unknown_lints)]
#![warn(clippy)]

pub mod lexer;
pub mod parser;

#[macro_use]
extern crate log;
extern crate simplelog;
extern crate structopt;
#[macro_use]
extern crate structopt_derive;

use simplelog::Config;
use simplelog::LogLevelFilter;
use simplelog::TermLogger;
use structopt::StructOpt;

use lexer::lexeme_iterator::LexemeIterator;
use parser::interpreter::Interpreter;
use parser::interpreter::InterpreterState;
use parser::token::Token;
use parser::token_iterator::TokenIterator;

use std::env;
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::io::BufReader;
use std::path::Path;

#[derive(StructOpt, Debug)]
#[structopt(name = "MiniPLInterpreter")]
struct Cli {
    #[structopt(
        long = "inputpath",
        short = "i",
        help = "Path to file to be interpreted. If left empty, the program will wait for stdin to be interpreted."
    )]
    input_path: Option<String>,

    #[structopt(long = "loglevel", short = "l")]
    log_level: Option<LogLevel>,
}

#[derive(StructOpt, Debug)]
#[structopt(name = "loglevel")]
enum LogLevel {
    #[structopt(name = "error")]
    Error,
    #[structopt(name = "info")]
    Info,
    #[structopt(name = "warning")]
    Warning,
    #[structopt(name = "debug")]
    Debug,
    #[structopt(name = "trace")]
    Trace,
}

impl std::str::FromStr for LogLevel {
    type Err = std::string::ParseError;

    fn from_str(text: &str) -> std::result::Result<Self, Self::Err> {
        match text {
            "info" => Result::Ok(LogLevel::Info),
            "warning" => Result::Ok(LogLevel::Warning),
            "debug" => Result::Ok(LogLevel::Debug),
            "trace" => Result::Ok(LogLevel::Trace),
            "error" => Result::Ok(LogLevel::Error),
            _ => Result::Ok(LogLevel::Error),
        }
    }
}

impl Into<LogLevelFilter> for LogLevel {
    fn into(self) -> LogLevelFilter {
        match self {
            LogLevel::Error => LogLevelFilter::Error,
            LogLevel::Info => LogLevelFilter::Info,
            LogLevel::Warning => LogLevelFilter::Warn,
            LogLevel::Debug => LogLevelFilter::Debug,
            LogLevel::Trace => LogLevelFilter::Trace,
        }
    }
}

fn main() {
    let _example = 1230000000;
    let args = Cli::from_args();
    let ll = args.log_level;
    let level_filter: LogLevelFilter = match ll {
        Some(level) => level.into(),
        _ => LogLevelFilter::Error,
    };
    let _ = TermLogger::init(level_filter, Config::default());

    info!("MiniPL Interpreter starting!");

    let path_str: String = args.input_path.unwrap_or("".to_string());
    let path = Path::new(&path_str);

    let mut absolute_path = env::current_dir().unwrap();
    let mut state = InterpreterState::new();

    debug!("The current directory is {}", absolute_path.display());
    absolute_path.push(path);
    let exit_keyword = "exit";
    debug!("The file is {}", absolute_path.display());

    if path.exists() && path.is_file() {
        let f = File::open(path).unwrap();
        let mut reader = BufReader::new(f);
        let mut buffer = String::new();
        reader.read_to_string(&mut buffer).unwrap();
        info!("File read");
        buffer.trim_left_matches("\u{feff}");

        let _ = eval_file(&buffer);
    } else {
        println!("Type commands");
        println!("Type {:?} to exit", exit_keyword);

        loop {
            print!("MiniPL> ");
            io::stdout().flush().unwrap();
            let mut input_text: String = String::new();
            io::stdin().read_line(&mut input_text).expect("failed to read from stdin");
            trace!("Read: {}", input_text);
            io::stdout().flush().unwrap();
            if input_text.to_string().trim() == exit_keyword {
                break;
            }
            eval_line(&input_text.to_string().trim(), &mut state);
        }
    };
}

fn eval_file(file_contents: &str) -> Result<(), String> {
    let lexeme_it = LexemeIterator::new(file_contents);
    let mut token_iterator: TokenIterator<LexemeIterator> = TokenIterator::new(lexeme_it);
    let mut state = InterpreterState::new();
    let mut interpreter = Interpreter::new(&mut token_iterator as &mut Iterator<Item = Token>, &mut state);
    interpreter.interpret()
}

fn eval_line(line: &str, mut state: &mut InterpreterState) {
    debug!("{:?}", line);
    let it = LexemeIterator::new(line);
    let mut token_iterator: TokenIterator<LexemeIterator> = TokenIterator::new(it);
    let mut interpreter = Interpreter::new(&mut token_iterator, &mut state);
    match interpreter.interpret() {
        Err(msg) => error!("{}", msg),
        _ => {}
    }
}

#[test]
fn sample1_var_definition_expression_print() {
    let code = r#"
var X : int := 4 + (6 * 2);
print X;
"#;
    eval_file(code).unwrap();
}

#[test]
fn sample2_loop_print() {
    let code = r#"
var nTimes : int := 0;
print nTimes;
print "How many times?";
nTimes := 3;
print nTimes;
var x : int;
for x in 0..nTimes-1 do
	print x;
	print " : Hello, World!\n";
end for;
print x;
print nTimes;
assert (x = (nTimes-1));
"#;
    eval_file(code).unwrap();
}

#[test]
fn sample4_decl_assign_print() {
    let code = r#"
var X : int;
X := 15;
print X;
"#;
    eval_file(code).unwrap();
}

#[test]
fn test_decl_assign_boolean() {
    let code = r#"
var X : bool;
X := 15;
print X;
"#;
    eval_file(code).unwrap();
}
<commit_msg>2018 preview<commit_after>#![feature(rust_2018_preview, uniform_paths)]

#![allow(unknown_lints)]
#![warn(clippy)]

pub mod lexer;
pub mod parser;

#[macro_use]
extern crate log;
extern crate simplelog;
extern crate structopt;
#[macro_use]
extern crate structopt_derive;

use simplelog::Config;
use simplelog::LogLevelFilter;
use simplelog::TermLogger;
use structopt::StructOpt;

use lexer::lexeme_iterator::LexemeIterator;
use parser::interpreter::Interpreter;
use parser::interpreter::InterpreterState;
use parser::token::Token;
use parser::token_iterator::TokenIterator;

use std::env;
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::io::BufReader;
use std::path::Path;

#[derive(StructOpt, Debug)]
#[structopt(name = "MiniPLInterpreter")]
struct Cli {
    #[structopt(
        long = "inputpath",
        short = "i",
        help = "Path to file to be interpreted. If left empty, the program will wait for stdin to be interpreted."
    )]
    input_path: Option<String>,

    #[structopt(long = "loglevel", short = "l")]
    log_level: Option<LogLevel>,
}

#[derive(StructOpt, Debug)]
#[structopt(name = "loglevel")]
enum LogLevel {
    #[structopt(name = "error")]
    Error,
    #[structopt(name = "info")]
    Info,
    #[structopt(name = "warning")]
    Warning,
    #[structopt(name = "debug")]
    Debug,
    #[structopt(name = "trace")]
    Trace,
}

impl std::str::FromStr for LogLevel {
    type Err = std::string::ParseError;

    fn from_str(text: &str) -> std::result::Result<Self, Self::Err> {
        match text {
            "info" => Result::Ok(LogLevel::Info),
            "warning" => Result::Ok(LogLevel::Warning),
            "debug" => Result::Ok(LogLevel::Debug),
            "trace" => Result::Ok(LogLevel::Trace),
            "error" => Result::Ok(LogLevel::Error),
            _ => Result::Ok(LogLevel::Error),
        }
    }
}

impl Into<LogLevelFilter> for LogLevel {
    fn into(self) -> LogLevelFilter {
        match self {
            LogLevel::Error => LogLevelFilter::Error,
            LogLevel::Info => LogLevelFilter::Info,
            LogLevel::Warning => LogLevelFilter::Warn,
            LogLevel::Debug => LogLevelFilter::Debug,
            LogLevel::Trace => LogLevelFilter::Trace,
        }
    }
}

fn main() {
    let _example = 1230000000;
    let args = Cli::from_args();
    let ll = args.log_level;
    let level_filter: LogLevelFilter = match ll {
        Some(level) => level.into(),
        _ => LogLevelFilter::Error,
    };
    let _ = TermLogger::init(level_filter, Config::default());

    info!("MiniPL Interpreter starting!");

    let path_str: String = args.input_path.unwrap_or("".to_string());
    let path = Path::new(&path_str);

    let mut absolute_path = env::current_dir().unwrap();
    let mut state = InterpreterState::new();

    debug!("The current directory is {}", absolute_path.display());
    absolute_path.push(path);
    let exit_keyword = "exit";
    debug!("The file is {}", absolute_path.display());

    if path.exists() && path.is_file() {
        let f = File::open(path).unwrap();
        let mut reader = BufReader::new(f);
        let mut buffer = String::new();
        reader.read_to_string(&mut buffer).unwrap();
        info!("File read");
        buffer.trim_left_matches("\u{feff}");

        let _ = eval_file(&buffer);
    } else {
        println!("Type commands");
        println!("Type {:?} to exit", exit_keyword);

        loop {
            print!("MiniPL> ");
            io::stdout().flush().unwrap();
            let mut input_text: String = String::new();
            io::stdin().read_line(&mut input_text).expect("failed to read from stdin");
            trace!("Read: {}", input_text);
            io::stdout().flush().unwrap();
            if input_text.to_string().trim() == exit_keyword {
                break;
            }
            eval_line(&input_text.to_string().trim(), &mut state);
        }
    };
}

fn eval_file(file_contents: &str) -> Result<(), String> {
    let lexeme_it = LexemeIterator::new(file_contents);
    let mut token_iterator: TokenIterator<LexemeIterator> = TokenIterator::new(lexeme_it);
    let mut state = InterpreterState::new();
    let mut interpreter = Interpreter::new(&mut token_iterator as &mut Iterator<Item = Token>, &mut state);
    interpreter.interpret()
}

fn eval_line(line: &str, mut state: &mut InterpreterState) {
    debug!("{:?}", line);
    let it = LexemeIterator::new(line);
    let mut token_iterator: TokenIterator<LexemeIterator> = TokenIterator::new(it);
    let mut interpreter = Interpreter::new(&mut token_iterator, &mut state);
    match interpreter.interpret() {
        Err(msg) => error!("{}", msg),
        _ => {}
    }
}

#[test]
fn sample1_var_definition_expression_print() {
    let code = r#"
var X : int := 4 + (6 * 2);
print X;
"#;
    eval_file(code).unwrap();
}

#[test]
fn sample2_loop_print() {
    let code = r#"
var nTimes : int := 0;
print nTimes;
print "How many times?";
nTimes := 3;
print nTimes;
var x : int;
for x in 0..nTimes-1 do
	print x;
	print " : Hello, World!\n";
end for;
print x;
print nTimes;
assert (x = (nTimes-1));
"#;
    eval_file(code).unwrap();
}

#[test]
fn sample4_decl_assign_print() {
    let code = r#"
var X : int;
X := 15;
print X;
"#;
    eval_file(code).unwrap();
}

#[test]
fn test_decl_assign_boolean() {
    let code = r#"
var X : bool;
X := 15;
print X;
"#;
    eval_file(code).unwrap();
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate walkdir;
extern crate twox_hash;

use std::collections::HashMap;
use std::collections::HashSet;
use std::fs::File;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::os::linux::fs::MetadataExt;

use clap::{App, Arg, ArgMatches};
use walkdir::WalkDir;
use twox_hash::XxHash;
use std::hash::Hasher;

#[derive(Default)]
struct Stats {
    file_count: u64,
    total_size: u64,
    duplicate_count: u64,
}

fn parse_args<'a>() -> ArgMatches<'a> {
    App::new("dupr")
        .version("0.1.0")
        .author("Mike Sampson")
        .about("Duplicate file finder")
        .arg(
            Arg::with_name("DIR")
                .help("Directory to process")
                .required(true),
        )
        .arg(Arg::with_name("summary").short("s").long("summary").help(
            "Print out summary information after duplicates",
        ))
        .get_matches()
}

fn collect_paths(path: &str, stats: &mut Stats) -> HashMap<u64, Vec<PathBuf>> {

    let mut length_paths = HashMap::new();

    for entry in WalkDir::new(path) {
        let entry = match entry {
            Ok(e) => e,
            Err(err) => {
                println!("dupr: {}", err); // stderr
                continue;
            }
        };

        if !entry.file_type().is_file() {
            continue;
        }

        let metadata = match entry.metadata() {
            Ok(e) => e,
            Err(err) => {
                println!("dupr: {}", err); //std err
                continue;
            }
        };

        if metadata.len() == 0 {
            continue;
        }

        stats.file_count += 1;
        stats.total_size += metadata.len();

        let path = entry.path().to_path_buf();

        length_paths
            .entry(metadata.len())
            .or_insert(Vec::new())
            .push(path);
    }

    length_paths
}

fn main() {
    let matches = parse_args();

    let mut stats: Stats = Default::default();

    let files = collect_paths(matches.value_of("DIR").unwrap(), &mut stats);

    for f_paths in files.values().filter(|x| x.len() > 1) {

        // Avoid hardlinks to same inode
        let mut hardlinks = Vec::new();
        let mut di = HashSet::new();
        for f_path in f_paths {
            let metadata = f_path.metadata().unwrap();
            let dev_inode = (metadata.st_dev(), metadata.st_ino());
            if !di.contains(&dev_inode) {
                hardlinks.push(f_path);
                di.insert(dev_inode);
            }
        }

        let mut hashes = HashMap::new();
        for f_path in hardlinks {
            hashes
                .entry(hash_file(&f_path))
                .or_insert(Vec::new())
                .push(f_path);
        }

        for (hash, paths) in hashes.iter().filter(|&(_, v)| v.len() > 1) {
            stats.duplicate_count += paths.len() as u64;
            for p in paths {
                println!("{:016x} - {}", hash, p.display());
            }
        }
    }

    if matches.is_present("summary") {
        println!(
            "Processed {} files with a total size of {} bytes. {} duplicates found.",
            stats.file_count,
            stats.total_size,
            stats.duplicate_count
        );
    }
}

fn hash_file<P: AsRef<Path>>(path: P) -> u64 {
    let mut file = File::open(path).unwrap();

    let mut hash = XxHash::with_seed(0);

    let mut buffer = [0; 32 * 1024];

    while let Ok(n) = file.read(&mut buffer) {
        if n == 0 {
            break;
        }
        hash.write(&buffer[0..n]);
    }

    hash.finish()
}
<commit_msg>Remove unwrap<commit_after>extern crate clap;
extern crate walkdir;
extern crate twox_hash;

use std::collections::HashMap;
use std::collections::HashSet;
use std::fs::File;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::os::linux::fs::MetadataExt;

use clap::{App, Arg, ArgMatches};
use walkdir::WalkDir;
use twox_hash::XxHash;
use std::hash::Hasher;

#[derive(Default)]
struct Stats {
    file_count: u64,
    total_size: u64,
    duplicate_count: u64,
}

fn parse_args<'a>() -> ArgMatches<'a> {
    App::new("dupr")
        .version("0.1.0")
        .author("Mike Sampson")
        .about("Duplicate file finder")
        .arg(
            Arg::with_name("DIR")
                .help("Directory to process")
                .required(true),
        )
        .arg(Arg::with_name("summary").short("s").long("summary").help(
            "Print out summary information after duplicates",
        ))
        .get_matches()
}

fn collect_paths(path: &str, stats: &mut Stats) -> HashMap<u64, Vec<PathBuf>> {

    let mut length_paths = HashMap::new();

    for entry in WalkDir::new(path) {
        let entry = match entry {
            Ok(e) => e,
            Err(err) => {
                println!("dupr: {}", err); // stderr
                continue;
            }
        };

        if !entry.file_type().is_file() {
            continue;
        }

        let metadata = match entry.metadata() {
            Ok(e) => e,
            Err(err) => {
                println!("dupr: {}", err); //std err
                continue;
            }
        };

        if metadata.len() == 0 {
            continue;
        }

        stats.file_count += 1;
        stats.total_size += metadata.len();

        let path = entry.path().to_path_buf();

        length_paths
            .entry(metadata.len())
            .or_insert(Vec::new())
            .push(path);
    }

    length_paths
}

fn main() {
    let matches = parse_args();

    let mut stats: Stats = Default::default();

    let files = collect_paths(matches.value_of("DIR").unwrap(), &mut stats);

    for f_paths in files.values().filter(|x| x.len() > 1) {

        // Avoid hardlinks to same inode
        let mut hardlinks = Vec::new();
        let mut di = HashSet::new();
        for f_path in f_paths {
            let metadata = match f_path.metadata() {
                Ok(m) => m,
                Err(err) => {
                    println!("dupr: {}", err);
                    continue;
                }
            };
            let dev_inode = (metadata.st_dev(), metadata.st_ino());
            if !di.contains(&dev_inode) {
                hardlinks.push(f_path);
                di.insert(dev_inode);
            }
        }

        let mut hashes = HashMap::new();
        for f_path in hardlinks {
            hashes
                .entry(hash_file(&f_path))
                .or_insert(Vec::new())
                .push(f_path);
        }

        for (hash, paths) in hashes.iter().filter(|&(_, v)| v.len() > 1) {
            stats.duplicate_count += paths.len() as u64;
            for p in paths {
                println!("{:016x} - {}", hash, p.display());
            }
        }
    }

    if matches.is_present("summary") {
        println!(
            "Processed {} files with a total size of {} bytes. {} duplicates found.",
            stats.file_count,
            stats.total_size,
            stats.duplicate_count
        );
    }
}

fn hash_file<P: AsRef<Path>>(path: P) -> u64 {
    let mut file = File::open(path).unwrap();

    let mut hash = XxHash::with_seed(0);

    let mut buffer = [0; 32 * 1024];

    while let Ok(n) = file.read(&mut buffer) {
        if n == 0 {
            break;
        }
        hash.write(&buffer[0..n]);
    }

    hash.finish()
}
<|endoftext|>