<commit_before>#![crate_name = "uu_tail"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Morten Olsen Lysgaard <morten@lysgaard.no>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 *
 */

extern crate getopts;

#[macro_use]
extern crate uucore;

use std::collections::VecDeque;
use std::error::Error;
use std::fmt;
use std::fs::File;
use std::io::{BufRead, BufReader, Read, Seek, SeekFrom, stdin, stdout, Write};
use std::path::Path;
use std::str::from_utf8;
use std::thread::sleep;
use std::time::Duration;

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

enum FilterMode {
    Bytes(u64),
    Lines(u64, u8), // (number of lines, delimiter)
}

struct Settings {
    mode: FilterMode,
    sleep_msec: u32,
    beginning: bool,
    follow: bool,
}

impl Default for Settings {
    fn default() -> Settings {
        Settings {
            mode: FilterMode::Lines(10, '\n' as u8),
            sleep_msec: 1000,
            beginning: false,
            follow: false,
        }
    }
}

pub fn uumain(args: Vec<String>) -> i32 {
    let mut settings: Settings = Default::default();

    // handle obsolete -number syntax
    let options = match obsolete(&args[1..]) {
        (args, Some(n)) => { settings.mode = FilterMode::Lines(n, '\n' as u8); args },
        (args, None) => args
    };

    let args = options;

    let mut opts = getopts::Options::new();

    opts.optopt("c", "bytes", "Number of bytes to print", "k");
    opts.optopt("n", "lines", "Number of lines to print", "k");
    opts.optflag("f", "follow", "Print the file as it grows");
    opts.optopt("s", "sleep-interval", "Number or seconds to sleep between polling the file when running with -f", "n");
    opts.optflag("z", "zero-terminated", "Line delimiter is NUL, not newline");
    opts.optflag("h", "help", "help");
    opts.optflag("V", "version", "version");

    let given_options = match opts.parse(&args) {
        Ok (m) => { m }
        Err(_) => {
            println!("{}", opts.usage(""));
            return 1;
        }
    };

    if given_options.opt_present("h") {
        println!("{}", opts.usage(""));
        return 0;
    }
    if given_options.opt_present("V") { version(); return 0 }

    settings.follow = given_options.opt_present("f");
    if settings.follow {
        match given_options.opt_str("s") {
            Some(n) => {
                let parsed: Option<u32> = n.parse().ok();
                match parsed {
                    Some(m) => { settings.sleep_msec = m * 1000 }
                    None => {}
                }
            }
            None => {}
        };
    }

    match given_options.opt_str("n") {
        Some(n) => {
            let mut slice: &str = n.as_ref();
            if slice.chars().next().unwrap_or('_') == '+' {
                settings.beginning = true;
                slice = &slice[1..];
            }
            match parse_size(slice) {
                Ok(m) => settings.mode = FilterMode::Lines(m, '\n' as u8),
                Err(e) => {
                    show_error!("{}", e.description());
                    return 1;
                }
            }
        }
        None => match given_options.opt_str("c") {
            Some(n) => {
                let mut slice: &str = n.as_ref();
                if slice.chars().next().unwrap_or('_') == '+' {
                    settings.beginning = true;
                    slice = &slice[1..];
                }
                match parse_size(slice) {
                    Ok(m) => settings.mode = FilterMode::Bytes(m),
                    Err(e) => {
                        show_error!("{}", e.description());
                        return 1;
                    }
                }
            }
            None => { }
        }
    };

    if given_options.opt_present("z") {
        if let FilterMode::Lines(count, _) = settings.mode {
            settings.mode = FilterMode::Lines(count, 0);
        }
    }

    let files = given_options.free;

    if files.is_empty() {
        let buffer = BufReader::new(stdin());
        unbounded_tail(buffer, &settings);
    } else {
        let mut multiple = false;
        let mut firstime = true;

        if files.len() > 1 {
            multiple = true;
        }

        for file in &files {
            if multiple {
                if !firstime { println!(""); }
                println!("==> {} <==", file);
            }
            firstime = false;

            let path = Path::new(file);
            let reader = File::open(&path).unwrap();
            bounded_tail(reader, &settings);
        }
    }

    0
}

#[derive(Debug, PartialEq, Eq)]
pub enum ParseSizeErr {
    ParseFailure(String),
    SizeTooBig(String),
}

impl Error for ParseSizeErr {
    fn description(&self) -> &str {
        match *self {
            ParseSizeErr::ParseFailure(ref s) => &*s,
            ParseSizeErr::SizeTooBig(ref s) => &*s,
        }
    }
}

impl fmt::Display for ParseSizeErr {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "{}", Error::description(self))
    }
}

impl ParseSizeErr {
    fn parse_failure(s: &str) -> ParseSizeErr {
        ParseSizeErr::ParseFailure(format!("invalid size: '{}'", s))
    }

    fn size_too_big(s: &str) -> ParseSizeErr {
        ParseSizeErr::SizeTooBig(
            format!("invalid size: '{}': Value too large to be stored in data type", s))
    }
}

pub type ParseSizeResult = Result<u64, ParseSizeErr>;

pub fn parse_size(mut size_slice: &str) -> Result<u64, ParseSizeErr> {
    let mut base =
        if size_slice.chars().last().unwrap_or('_') == 'B' {
            size_slice = &size_slice[..size_slice.len() - 1];
            1000u64
        } else {
            1024u64
        };

    let exponent =
        if size_slice.len() > 0 {
            let mut has_suffix = true;
            let exp = match size_slice.chars().last().unwrap_or('_') {
                'K' | 'k' => 1u64,
                'M' => 2u64,
                'G' => 3u64,
                'T' => 4u64,
                'P' => 5u64,
                'E' => 6u64,
                'Z' | 'Y' => {
                    return Err(ParseSizeErr::size_too_big(size_slice));
                },
                'b' => {
                    base = 512u64;
                    1u64
                }
                _ => {
                    has_suffix = false;
                    0u64
                }
            };
            if has_suffix {
                size_slice = &size_slice[..size_slice.len() - 1];
            }
            exp
        } else {
            0u64
        };

    let mut multiplier = 1u64;
    for _ in 0u64 .. exponent {
        multiplier *= base;
    }
    if base == 1000u64 && exponent == 0u64 {
        // sole B is not a valid suffix
        Err(ParseSizeErr::parse_failure(size_slice))
    } else {
        let value: Option<u64> = size_slice.parse().ok();
        value.map(|v| Ok(multiplier * v))
             .unwrap_or(Err(ParseSizeErr::parse_failure(size_slice)))
    }
}

// It searches for an option in the form of -123123
//
// In case is found, the options vector will get rid of that object so that
// getopts works correctly.
fn obsolete(options: &[String]) -> (Vec<String>, Option<u64>) {
    let mut options: Vec<String> = options.to_vec();
    let mut a = 0;
    let b = options.len();

    while a < b {
        let current = options[a].clone();
        let current = current.as_bytes();

        if current.len() > 1 && current[0] == '-' as u8 {
            let len = current.len();
            for pos in 1 .. len {
                // Ensure that the argument is only made out of digits
                if !(current[pos] as char).is_numeric() { break; }

                // If this is the last number
                if pos == len - 1 {
                    options.remove(a);
                    let number: Option<u64> = from_utf8(&current[1..len]).unwrap().parse().ok();
                    return (options, Some(number.unwrap()));
                }
            }
        }

        a += 1;
    };

    (options, None)
}

/// When reading files in reverse in `bounded_tail`, this is the size of each
/// block read at a time.
const BLOCK_SIZE: u64 = 1 << 16;

fn follow<T: Read>(mut reader: BufReader<T>, settings: &Settings) {
    assert!(settings.follow);
    loop {
        sleep(Duration::new(0, settings.sleep_msec*1000));
        loop {
            let mut datum = String::new();
            match reader.read_line(&mut datum) {
                Ok(0) => break,
                Ok(_) => print!("{}", datum),
                Err(err) => panic!(err)
            }
        }
    }
}

/// Iterate over bytes in the file, in reverse, until `should_stop` returns
/// true. The `file` is left seek'd to the position just after the byte that
/// `should_stop` returned true for.
fn backwards_thru_file<F>(file: &mut File, size: u64, buf: &mut Vec<u8>, delimiter: u8, should_stop: &mut F)
    where F: FnMut(u8) -> bool
{
    assert!(buf.len() >= BLOCK_SIZE as usize);

    let max_blocks_to_read = (size as f64 / BLOCK_SIZE as f64).ceil() as usize;

    for block_idx in 0..max_blocks_to_read {
        let block_size = if block_idx == max_blocks_to_read - 1 {
            size % BLOCK_SIZE
        } else {
            BLOCK_SIZE
        };

        // Seek backwards by the next block, read the full block into
        // `buf`, and then seek back to the start of the block again.
        let pos = file.seek(SeekFrom::Current(-(block_size as i64))).unwrap();
        file.read_exact(&mut buf[0..(block_size as usize)]).unwrap();
        let pos2 = file.seek(SeekFrom::Current(-(block_size as i64))).unwrap();
        assert_eq!(pos, pos2);

        // Iterate backwards through the bytes, calling `should_stop` on each
        // one.
        let slice = &buf[0..(block_size as usize)];
        for (i, ch) in slice.iter().enumerate().rev() {
            // Ignore one trailing newline.
            if block_idx == 0 && i as u64 == block_size - 1 && *ch == delimiter {
                continue;
            }

            if should_stop(*ch) {
                file.seek(SeekFrom::Current((i + 1) as i64)).unwrap();
                return;
            }
        }
    }
}

/// When tail'ing a file, we do not need to read the whole file from start to
/// finish just to find the last n lines or bytes. Instead, we can seek to the
/// end of the file, and then read the file "backwards" in blocks of size
/// `BLOCK_SIZE` until we find the location of the first line/byte. This ends up
/// being a nice performance win for very large files.
fn bounded_tail(mut file: File, settings: &Settings) {
    let size = file.seek(SeekFrom::End(0)).unwrap();
    let mut buf = vec![0; BLOCK_SIZE as usize];

    // Find the position in the file to start printing from.
    match settings.mode {
        FilterMode::Lines(mut count, delimiter) => {
            backwards_thru_file(&mut file, size, &mut buf, delimiter, &mut |byte| {
                if byte == delimiter {
                    count -= 1;
                    count == 0
                } else {
                    false
                }
            });
        },
        FilterMode::Bytes(count) => {
            file.seek(SeekFrom::End(-(count as i64))).unwrap();
        },
    }

    // Print the target section of the file.
    loop {
        let bytes_read = file.read(&mut buf).unwrap();

        let mut stdout = stdout();
        for b in &buf[0..bytes_read] {
            print_byte(&mut stdout, b);
        }

        if bytes_read == 0 {
            break;
        }
    }
}

fn unbounded_tail<T: Read>(mut reader: BufReader<T>, settings: &Settings) {
    // Read through each line/char and store them in a ringbuffer that always
    // contains count lines/chars. When reaching the end of file, output the
    // data in the ringbuf.
    match settings.mode {
        FilterMode::Lines(mut count, _delimiter) => {
            let mut ringbuf: VecDeque<String> = VecDeque::new();
            let mut skip = if settings.beginning {
                let temp = count;
                count = ::std::u64::MAX;
                temp - 1
            } else {
                0
            };
            loop {
                let mut datum = String::new();
                match reader.read_line(&mut datum) {
                    Ok(0) => break,
                    Ok(_) => {
                        if skip > 0 {
                            skip -= 1;
                        } else {
                            if count <= ringbuf.len() as u64 {
                                ringbuf.pop_front();
                            }
                            ringbuf.push_back(datum);
                        }
                    },
                    Err(err) => panic!(err)
                }
            }
            let mut stdout = stdout();
            for datum in &ringbuf {
                print_string(&mut stdout, datum);
            }
        },
        FilterMode::Bytes(mut count) => {
            let mut ringbuf: VecDeque<u8> = VecDeque::new();
            let mut skip = if settings.beginning {
                let temp = count;
                count = ::std::u64::MAX;
                temp - 1
            } else {
                0
            };
            loop {
                let mut datum = [0; 1];
                match reader.read(&mut datum) {
                    Ok(0) => break,
                    Ok(_) => {
                        if skip > 0 {
                            skip -= 1;
                        } else {
                            if count <= ringbuf.len() as u64 {
                                ringbuf.pop_front();
                            }
                            ringbuf.push_back(datum[0]);
                        }
                    },
                    Err(err) => panic!(err)
                }
            }
            let mut stdout = stdout();
            for datum in &ringbuf {
                print_byte(&mut stdout, datum);
            }
        }
    }

    if settings.follow {
        follow(reader, settings);
    }
}

#[inline]
fn print_byte<T: Write>(stdout: &mut T, ch: &u8) {
    if let Err(err) = stdout.write(&[*ch]) {
        crash!(1, "{}", err);
    }
}

#[inline]
fn print_string<T: Write>(_: &mut T, s: &str) {
    print!("{}", s);
}

fn version() {
    println!("{} {}", NAME, VERSION);
}
<commit_msg>tail: follow multiple files<commit_after>#![crate_name = "uu_tail"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Morten Olsen Lysgaard <morten@lysgaard.no>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 *
 */

extern crate getopts;

#[macro_use]
extern crate uucore;

use std::collections::VecDeque;
use std::error::Error;
use std::fmt;
use std::fs::File;
use std::io::{BufRead, BufReader, Read, Seek, SeekFrom, stdin, stdout, Write};
use std::path::Path;
use std::str::from_utf8;
use std::thread::sleep;
use std::time::Duration;

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

enum FilterMode {
    Bytes(u64),
    Lines(u64, u8), // (number of lines, delimiter)
}

struct Settings {
    mode: FilterMode,
    sleep_msec: u32,
    beginning: bool,
    follow: bool,
}

impl Default for Settings {
    fn default() -> Settings {
        Settings {
            mode: FilterMode::Lines(10, '\n' as u8),
            sleep_msec: 1000,
            beginning: false,
            follow: false,
        }
    }
}

pub fn uumain(args: Vec<String>) -> i32 {
    let mut settings: Settings = Default::default();

    // handle obsolete -number syntax
    let options = match obsolete(&args[1..]) {
        (args, Some(n)) => { settings.mode = FilterMode::Lines(n, '\n' as u8); args },
        (args, None) => args
    };

    let args = options;

    let mut opts = getopts::Options::new();

    opts.optopt("c", "bytes", "Number of bytes to print", "k");
    opts.optopt("n", "lines", "Number of lines to print", "k");
    opts.optflag("f", "follow", "Print the file as it grows");
    opts.optopt("s", "sleep-interval", "Number or seconds to sleep between polling the file when running with -f", "n");
    opts.optflag("z", "zero-terminated", "Line delimiter is NUL, not newline");
    opts.optflag("h", "help", "help");
    opts.optflag("V", "version", "version");

    let given_options = match opts.parse(&args) {
        Ok (m) => { m }
        Err(_) => {
            println!("{}", opts.usage(""));
            return 1;
        }
    };

    if given_options.opt_present("h") {
        println!("{}", opts.usage(""));
        return 0;
    }
    if given_options.opt_present("V") { version(); return 0 }

    settings.follow = given_options.opt_present("f");
    if settings.follow {
        match given_options.opt_str("s") {
            Some(n) => {
                let parsed: Option<u32> = n.parse().ok();
                match parsed {
                    Some(m) => { settings.sleep_msec = m * 1000 }
                    None => {}
                }
            }
            None => {}
        };
    }

    match given_options.opt_str("n") {
        Some(n) => {
            let mut slice: &str = n.as_ref();
            if slice.chars().next().unwrap_or('_') == '+' {
                settings.beginning = true;
                slice = &slice[1..];
            }
            match parse_size(slice) {
                Ok(m) => settings.mode = FilterMode::Lines(m, '\n' as u8),
                Err(e) => {
                    show_error!("{}", e.description());
                    return 1;
                }
            }
        }
        None => match given_options.opt_str("c") {
            Some(n) => {
                let mut slice: &str = n.as_ref();
                if slice.chars().next().unwrap_or('_') == '+' {
                    settings.beginning = true;
                    slice = &slice[1..];
                }
                match parse_size(slice) {
                    Ok(m) => settings.mode = FilterMode::Bytes(m),
                    Err(e) => {
                        show_error!("{}", e.description());
                        return 1;
                    }
                }
            }
            None => { }
        }
    };

    if given_options.opt_present("z") {
        if let FilterMode::Lines(count, _) = settings.mode {
            settings.mode = FilterMode::Lines(count, 0);
        }
    }

    let files = given_options.free;

    if files.is_empty() {
        let buffer = BufReader::new(stdin());
        unbounded_tail(buffer, &settings);
    } else {
        let mut multiple = false;
        let mut readers = Vec::new();

        if files.len() > 1 {
            multiple = true;
        }

        for filename in &files {
            if multiple {
                println!("==> {} <==", filename);
            }

            let path = Path::new(filename);
            let file = File::open(&path).unwrap();
            bounded_tail(&file, &settings);

            if settings.follow {
                let reader = BufReader::new(file);
                readers.push(reader);
            }
        }

        if settings.follow {
            follow(readers, &settings);
        }
    }

    0
}

#[derive(Debug, PartialEq, Eq)]
pub enum ParseSizeErr {
    ParseFailure(String),
    SizeTooBig(String),
}

impl Error for ParseSizeErr {
    fn description(&self) -> &str {
        match *self {
            ParseSizeErr::ParseFailure(ref s) => &*s,
            ParseSizeErr::SizeTooBig(ref s) => &*s,
        }
    }
}

impl fmt::Display for ParseSizeErr {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "{}", Error::description(self))
    }
}

impl ParseSizeErr {
    fn parse_failure(s: &str) -> ParseSizeErr {
        ParseSizeErr::ParseFailure(format!("invalid size: '{}'", s))
    }

    fn size_too_big(s: &str) -> ParseSizeErr {
        ParseSizeErr::SizeTooBig(
            format!("invalid size: '{}': Value too large to be stored in data type", s))
    }
}

pub type ParseSizeResult = Result<u64, ParseSizeErr>;

pub fn parse_size(mut size_slice: &str) -> Result<u64, ParseSizeErr> {
    let mut base =
        if size_slice.chars().last().unwrap_or('_') == 'B' {
            size_slice = &size_slice[..size_slice.len() - 1];
            1000u64
        } else {
            1024u64
        };

    let exponent =
        if size_slice.len() > 0 {
            let mut has_suffix = true;
            let exp = match size_slice.chars().last().unwrap_or('_') {
                'K' | 'k' => 1u64,
                'M' => 2u64,
                'G' => 3u64,
                'T' => 4u64,
                'P' => 5u64,
                'E' => 6u64,
                'Z' | 'Y' => {
                    return Err(ParseSizeErr::size_too_big(size_slice));
                },
                'b' => {
                    base = 512u64;
                    1u64
                }
                _ => {
                    has_suffix = false;
                    0u64
                }
            };
            if has_suffix {
                size_slice = &size_slice[..size_slice.len() - 1];
            }
            exp
        } else {
            0u64
        };

    let mut multiplier = 1u64;
    for _ in 0u64 .. exponent {
        multiplier *= base;
    }
    if base == 1000u64 && exponent == 0u64 {
        // sole B is not a valid suffix
        Err(ParseSizeErr::parse_failure(size_slice))
    } else {
        let value: Option<u64> = size_slice.parse().ok();
        value.map(|v| Ok(multiplier * v))
             .unwrap_or(Err(ParseSizeErr::parse_failure(size_slice)))
    }
}

// It searches for an option in the form of -123123
//
// In case is found, the options vector will get rid of that object so that
// getopts works correctly.
fn obsolete(options: &[String]) -> (Vec<String>, Option<u64>) {
    let mut options: Vec<String> = options.to_vec();
    let mut a = 0;
    let b = options.len();

    while a < b {
        let current = options[a].clone();
        let current = current.as_bytes();

        if current.len() > 1 && current[0] == '-' as u8 {
            let len = current.len();
            for pos in 1 .. len {
                // Ensure that the argument is only made out of digits
                if !(current[pos] as char).is_numeric() { break; }

                // If this is the last number
                if pos == len - 1 {
                    options.remove(a);
                    let number: Option<u64> = from_utf8(&current[1..len]).unwrap().parse().ok();
                    return (options, Some(number.unwrap()));
                }
            }
        }

        a += 1;
    };

    (options, None)
}

/// When reading files in reverse in `bounded_tail`, this is the size of each
/// block read at a time.
const BLOCK_SIZE: u64 = 1 << 16;

fn follow<T: Read>(mut readers: Vec<BufReader<T>>, settings: &Settings) {
    assert!(settings.follow);
    let mut last = readers.len();

    loop {
        sleep(Duration::new(0, settings.sleep_msec*1000));

        for reader in &mut readers {
            // Print all new content since the last pass
            loop {
                let mut datum = String::new();
                match reader.read_line(&mut datum) {
                    Ok(0) => break,
                    Ok(_) => {
                        // TODO: Print headers if i != last
                        print!("{}", datum);
                    },
                    Err(err) => panic!(err)
                }
            }
        }
    }
}

/// Iterate over bytes in the file, in reverse, until `should_stop` returns
/// true. The `file` is left seek'd to the position just after the byte that
/// `should_stop` returned true for.
fn backwards_thru_file<F>(mut file: &File, size: u64, buf: &mut Vec<u8>, delimiter: u8, should_stop: &mut F)
    where F: FnMut(u8) -> bool
{
    assert!(buf.len() >= BLOCK_SIZE as usize);

    let max_blocks_to_read = (size as f64 / BLOCK_SIZE as f64).ceil() as usize;

    for block_idx in 0..max_blocks_to_read {
        let block_size = if block_idx == max_blocks_to_read - 1 {
            size % BLOCK_SIZE
        } else {
            BLOCK_SIZE
        };

        // Seek backwards by the next block, read the full block into
        // `buf`, and then seek back to the start of the block again.
        let pos = file.seek(SeekFrom::Current(-(block_size as i64))).unwrap();
        file.read_exact(&mut buf[0..(block_size as usize)]).unwrap();
        let pos2 = file.seek(SeekFrom::Current(-(block_size as i64))).unwrap();
        assert_eq!(pos, pos2);

        // Iterate backwards through the bytes, calling `should_stop` on each
        // one.
        let slice = &buf[0..(block_size as usize)];
        for (i, ch) in slice.iter().enumerate().rev() {
            // Ignore one trailing newline.
            if block_idx == 0 && i as u64 == block_size - 1 && *ch == delimiter {
                continue;
            }

            if should_stop(*ch) {
                file.seek(SeekFrom::Current((i + 1) as i64)).unwrap();
                return;
            }
        }
    }
}

/// When tail'ing a file, we do not need to read the whole file from start to
/// finish just to find the last n lines or bytes. Instead, we can seek to the
/// end of the file, and then read the file "backwards" in blocks of size
/// `BLOCK_SIZE` until we find the location of the first line/byte. This ends up
/// being a nice performance win for very large files.
fn bounded_tail(mut file: &File, settings: &Settings) {
    let size = file.seek(SeekFrom::End(0)).unwrap();
    let mut buf = vec![0; BLOCK_SIZE as usize];

    // Find the position in the file to start printing from.
    match settings.mode {
        FilterMode::Lines(mut count, delimiter) => {
            backwards_thru_file(&mut file, size, &mut buf, delimiter, &mut |byte| {
                if byte == delimiter {
                    count -= 1;
                    count == 0
                } else {
                    false
                }
            });
        },
        FilterMode::Bytes(count) => {
            file.seek(SeekFrom::End(-(count as i64))).unwrap();
        },
    }

    // Print the target section of the file.
    loop {
        let bytes_read = file.read(&mut buf).unwrap();

        let mut stdout = stdout();
        for b in &buf[0..bytes_read] {
            print_byte(&mut stdout, b);
        }

        if bytes_read == 0 {
            break;
        }
    }
}

fn unbounded_tail<T: Read>(mut reader: BufReader<T>, settings: &Settings) {
    // Read through each line/char and store them in a ringbuffer that always
    // contains count lines/chars. When reaching the end of file, output the
    // data in the ringbuf.
    match settings.mode {
        FilterMode::Lines(mut count, _delimiter) => {
            let mut ringbuf: VecDeque<String> = VecDeque::new();
            let mut skip = if settings.beginning {
                let temp = count;
                count = ::std::u64::MAX;
                temp - 1
            } else {
                0
            };
            loop {
                let mut datum = String::new();
                match reader.read_line(&mut datum) {
                    Ok(0) => break,
                    Ok(_) => {
                        if skip > 0 {
                            skip -= 1;
                        } else {
                            if count <= ringbuf.len() as u64 {
                                ringbuf.pop_front();
                            }
                            ringbuf.push_back(datum);
                        }
                    },
                    Err(err) => panic!(err)
                }
            }
            let mut stdout = stdout();
            for datum in &ringbuf {
                print_string(&mut stdout, datum);
            }
        },
        FilterMode::Bytes(mut count) => {
            let mut ringbuf: VecDeque<u8> = VecDeque::new();
            let mut skip = if settings.beginning {
                let temp = count;
                count = ::std::u64::MAX;
                temp - 1
            } else {
                0
            };
            loop {
                let mut datum = [0; 1];
                match reader.read(&mut datum) {
                    Ok(0) => break,
                    Ok(_) => {
                        if skip > 0 {
                            skip -= 1;
                        } else {
                            if count <= ringbuf.len() as u64 {
                                ringbuf.pop_front();
                            }
                            ringbuf.push_back(datum[0]);
                        }
                    },
                    Err(err) => panic!(err)
                }
            }
            let mut stdout = stdout();
            for datum in &ringbuf {
                print_byte(&mut stdout, datum);
            }
        }
    }

    // TODO: make following stdin work with the new follow() signature
    // maybe wrap stdin in a 1-element vec?
    if settings.follow {
        //follow(reader, settings);
    }
}

#[inline]
fn print_byte<T: Write>(stdout: &mut T, ch: &u8) {
    if let Err(err) = stdout.write(&[*ch]) {
        crash!(1, "{}", err);
    }
}

#[inline]
fn print_string<T: Write>(_: &mut T, s: &str) {
    print!("{}", s);
}

fn version() {
    println!("{} {}", NAME, VERSION);
}
<|endoftext|>


<commit_before>#![deny(warnings)]
#![feature(io, core, env, path, std_misc)]

extern crate "cargo-registry" as cargo_registry;
extern crate "conduit-middleware" as conduit_middleware;
extern crate "conduit-test" as conduit_test;
extern crate "rustc-serialize" as rustc_serialize;
extern crate conduit;
extern crate curl;
extern crate git2;
extern crate time;
extern crate url;
extern crate semver;

use std::collections::HashMap;
use std::error::Error as StdError;
use std::old_io::Command;
use std::old_io::process::InheritFd;
use std::env;
use std::sync::{Once, ONCE_INIT, Arc};
use rustc_serialize::json::{self, Json};

use conduit::Request;
use conduit_test::MockRequest;
use cargo_registry::app::App;
use cargo_registry::db::{self, RequestTransaction};
use cargo_registry::dependency::Kind;
use cargo_registry::{User, Crate, Version, Keyword, Dependency};

macro_rules! t{ ($e:expr) => (
    match $e {
        Ok(e) => e,
        Err(m) => panic!("{} failed with: {}", stringify!($e), m),
    }
) }

macro_rules! t_resp{ ($e:expr) => ({
    t!($e)
}) }

macro_rules! ok_resp{ ($e:expr) => ({
    let resp = t_resp!($e);
    if !::ok_resp(&resp) { panic!("bad response: {:?}", resp.status); }
    resp
}) }

macro_rules! bad_resp{ ($e:expr) => ({
    let mut resp = t_resp!($e);
    match ::bad_resp(&mut resp) {
        None => panic!("ok response: {:?}", resp.status),
        Some(b) => b,
    }
}) }

#[derive(RustcDecodable, Debug)]
struct Error { detail: String }
#[derive(RustcDecodable)]
struct Bad { errors: Vec<Error> }

mod middleware;
mod keyword;
mod krate;
mod user;
mod record;
mod git;
mod version;

fn app() -> (record::Bomb, Arc<App>, conduit_middleware::MiddlewareBuilder) {
    struct NoCommit;
    static INIT: Once = ONCE_INIT;
    git::init();

    let (proxy, bomb) = record::proxy();
    let config = cargo_registry::Config {
        s3_bucket: env::var_string("S3_BUCKET").unwrap_or(String::new()),
        s3_access_key: env::var_string("S3_ACCESS_KEY").unwrap_or(String::new()),
        s3_secret_key: env::var_string("S3_SECRET_KEY").unwrap_or(String::new()),
        s3_region: env::var_string("S3_REGION").ok(),
        s3_proxy: Some(proxy),
        session_key: "test".to_string(),
        git_repo_checkout: git::checkout(),
        gh_client_id: "".to_string(),
        gh_client_secret: "".to_string(),
        db_url: env("TEST_DATABASE_URL"),
        env: cargo_registry::Env::Test,
        max_upload_size: 1000,
    };
    INIT.call_once(|| db_setup(config.db_url.as_slice()));
    let app = App::new(&config);
    let app = Arc::new(app);
    let mut middleware = cargo_registry::middleware(app.clone());
    middleware.add(NoCommit);
    return (bomb, app, middleware);

    fn env(s: &str) -> String {
        match env::var_string(s).ok() {
            Some(s) => s,
            None => panic!("must have `{}` defined", s),
        }
    }

    fn db_setup(db: &str) {
        let migrate = env::current_exe().unwrap().dir_path().join("../migrate");
        assert!(Command::new(migrate).env("DATABASE_URL", db)
                        .stdout(InheritFd(1))
                        .stderr(InheritFd(2))
                        .status().unwrap().success());
    }

    impl conduit_middleware::Middleware for NoCommit {
        fn after(&self, req: &mut Request,
                 res: Result<conduit::Response, Box<StdError+Send>>)
                 -> Result<conduit::Response, Box<StdError+Send>> {
            req.extensions().find::<db::Transaction>()
               .expect("Transaction not present in request")
               .rollback();
            return res;
        }
    }
}

fn req(app: Arc<App>, method: conduit::Method, path: &str) -> MockRequest {
    let mut req = MockRequest::new(method, path);
    req.mut_extensions().insert(db::Transaction::new(app));
    return req;
}

fn ok_resp(r: &conduit::Response) -> bool {
    r.status.0 == 200
}

fn bad_resp(r: &mut conduit::Response) -> Option<Bad> {
    let bad = json::<Bad>(r);
    if bad.errors.len() == 0 { return None }
    Some(bad)
}

fn json<T: rustc_serialize::Decodable>(r: &mut conduit::Response) -> T {
    let data = r.body.read_to_end().unwrap();
    let s = std::str::from_utf8(data.as_slice()).unwrap();
    let j = match Json::from_str(s) {
        Ok(t) => t,
        Err(e) => panic!("failed to decode: {:?}\n{}", e, s),
    };
    let j = fixup(j);
    let s = j.to_string();
    return match json::decode(s.as_slice()) {
        Ok(t) => t,
        Err(e) => panic!("failed to decode: {:?}\n{}", e, s),
    };

    fn fixup(json: Json) -> Json {
        match json {
            Json::Object(object) => {
                Json::Object(object.into_iter().map(|(k, v)| {
                    let k = if k.as_slice() == "crate" {
                        "krate".to_string()
                    } else {
                        k
                    };
                    (k, fixup(v))
                }).collect())
            }
            Json::Array(list) => {
                Json::Array(list.into_iter().map(fixup).collect())
            }
            j => j,
        }
    }
}

fn user(login: &str) -> User {
    User {
        id: 10000,
        gh_login: login.to_string(),
        email: None,
        name: None,
        avatar: None,
        gh_access_token: User::new_api_token(), // just randomize it
        api_token: User::new_api_token(),
    }
}

fn krate(name: &str) -> Crate {
    cargo_registry::krate::Crate {
        id: 10000,
        name: name.to_string(),
        user_id: 100,
        updated_at: time::now().to_timespec(),
        created_at: time::now().to_timespec(),
        downloads: 10,
        max_version: semver::Version::parse("0.0.0").unwrap(),
        documentation: None,
        homepage: None,
        description: None,
        readme: None,
        keywords: Vec::new(),
        license: None,
        repository: None,
    }
}

fn mock_user(req: &mut Request, u: User) -> User {
    let u = User::find_or_insert(req.tx().unwrap(),
                                 u.gh_login.as_slice(),
                                 u.email.as_ref().map(|s| s.as_slice()),
                                 u.name.as_ref().map(|s| s.as_slice()),
                                 u.avatar.as_ref().map(|s| s.as_slice()),
                                 u.gh_access_token.as_slice(),
                                 u.api_token.as_slice()).unwrap();
    req.mut_extensions().insert(u.clone());
    return u;
}

fn mock_crate(req: &mut Request, krate: Crate) -> (Crate, Version) {
    mock_crate_vers(req, krate, &semver::Version::parse("1.0.0").unwrap())
}
fn mock_crate_vers(req: &mut Request, krate: Crate, v: &semver::Version)
                   -> (Crate, Version) {
    let user = req.extensions().find::<User>().unwrap();
    let mut krate = Crate::find_or_insert(req.tx().unwrap(), krate.name.as_slice(),
                                      user.id, &krate.description,
                                      &krate.homepage,
                                      &krate.documentation,
                                      &krate.readme,
                                      krate.keywords.as_slice(),
                                      &krate.repository,
                                      &krate.license,
                                      &None).unwrap();
    Keyword::update_crate(req.tx().unwrap(), &krate,
                          krate.keywords.as_slice()).unwrap();
    let v = krate.add_version(req.tx().unwrap(), v, &HashMap::new(), &[]);
    (krate, v.unwrap())
}

fn mock_dep(req: &mut Request, version: &Version, krate: &Crate,
            target: Option<&str>) -> Dependency {
    Dependency::insert(req.tx().unwrap(),
                       version.id,
                       krate.id,
                       &semver::VersionReq::parse(">= 0").unwrap(),
                       Kind::Normal,
                       false, true, &[],
                       &target.map(|s| s.to_string())).unwrap()
}

fn mock_keyword(req: &mut Request, name: &str) -> Keyword {
    Keyword::find_or_insert(req.tx().unwrap(), name).unwrap()
}

fn logout(req: &mut Request) {
    req.mut_extensions().pop::<User>();
}
<commit_msg>Fix tests<commit_after>#![deny(warnings)]
#![feature(io, core, env, path, std_misc)]

extern crate "cargo-registry" as cargo_registry;
extern crate "conduit-middleware" as conduit_middleware;
extern crate "conduit-test" as conduit_test;
extern crate "rustc-serialize" as rustc_serialize;
extern crate conduit;
extern crate curl;
extern crate git2;
extern crate time;
extern crate url;
extern crate semver;

use std::collections::HashMap;
use std::error::Error as StdError;
use std::old_io::Command;
use std::old_io::process::InheritFd;
use std::env;
use std::sync::{Once, ONCE_INIT, Arc};
use rustc_serialize::json::{self, Json};

use conduit::Request;
use conduit_test::MockRequest;
use cargo_registry::app::App;
use cargo_registry::db::{self, RequestTransaction};
use cargo_registry::dependency::Kind;
use cargo_registry::{User, Crate, Version, Keyword, Dependency};

macro_rules! t{ ($e:expr) => (
    match $e {
        Ok(e) => e,
        Err(m) => panic!("{} failed with: {}", stringify!($e), m),
    }
) }

macro_rules! t_resp{ ($e:expr) => ({
    t!($e)
}) }

macro_rules! ok_resp{ ($e:expr) => ({
    let resp = t_resp!($e);
    if !::ok_resp(&resp) { panic!("bad response: {:?}", resp.status); }
    resp
}) }

macro_rules! bad_resp{ ($e:expr) => ({
    let mut resp = t_resp!($e);
    match ::bad_resp(&mut resp) {
        None => panic!("ok response: {:?}", resp.status),
        Some(b) => b,
    }
}) }

#[derive(RustcDecodable, Debug)]
struct Error { detail: String }
#[derive(RustcDecodable)]
struct Bad { errors: Vec<Error> }

mod middleware;
mod keyword;
mod krate;
mod user;
mod record;
mod git;
mod version;

fn app() -> (record::Bomb, Arc<App>, conduit_middleware::MiddlewareBuilder) {
    struct NoCommit;
    static INIT: Once = ONCE_INIT;
    git::init();

    let (proxy, bomb) = record::proxy();
    let config = cargo_registry::Config {
        s3_bucket: env::var_string("S3_BUCKET").unwrap_or(String::new()),
        s3_access_key: env::var_string("S3_ACCESS_KEY").unwrap_or(String::new()),
        s3_secret_key: env::var_string("S3_SECRET_KEY").unwrap_or(String::new()),
        s3_region: env::var_string("S3_REGION").ok(),
        s3_proxy: Some(proxy),
        session_key: "test".to_string(),
        git_repo_checkout: git::checkout(),
        gh_client_id: "".to_string(),
        gh_client_secret: "".to_string(),
        db_url: env("TEST_DATABASE_URL"),
        env: cargo_registry::Env::Test,
        max_upload_size: 1000,
    };
    INIT.call_once(|| db_setup(config.db_url.as_slice()));
    let app = App::new(&config);
    let app = Arc::new(app);
    let mut middleware = cargo_registry::middleware(app.clone());
    middleware.add(NoCommit);
    return (bomb, app, middleware);

    fn env(s: &str) -> String {
        match env::var_string(s).ok() {
            Some(s) => s,
            None => panic!("must have `{}` defined", s),
        }
    }

    fn db_setup(db: &str) {
        let migrate = env::current_exe().unwrap().join("../migrate");
        assert!(Command::new(migrate).env("DATABASE_URL", db)
                        .stdout(InheritFd(1))
                        .stderr(InheritFd(2))
                        .status().unwrap().success());
    }

    impl conduit_middleware::Middleware for NoCommit {
        fn after(&self, req: &mut Request,
                 res: Result<conduit::Response, Box<StdError+Send>>)
                 -> Result<conduit::Response, Box<StdError+Send>> {
            req.extensions().find::<db::Transaction>()
               .expect("Transaction not present in request")
               .rollback();
            return res;
        }
    }
}

fn req(app: Arc<App>, method: conduit::Method, path: &str) -> MockRequest {
    let mut req = MockRequest::new(method, path);
    req.mut_extensions().insert(db::Transaction::new(app));
    return req;
}

fn ok_resp(r: &conduit::Response) -> bool {
    r.status.0 == 200
}

fn bad_resp(r: &mut conduit::Response) -> Option<Bad> {
    let bad = json::<Bad>(r);
    if bad.errors.len() == 0 { return None }
    Some(bad)
}

fn json<T: rustc_serialize::Decodable>(r: &mut conduit::Response) -> T {
    let data = r.body.read_to_end().unwrap();
    let s = std::str::from_utf8(data.as_slice()).unwrap();
    let j = match Json::from_str(s) {
        Ok(t) => t,
        Err(e) => panic!("failed to decode: {:?}\n{}", e, s),
    };
    let j = fixup(j);
    let s = j.to_string();
    return match json::decode(s.as_slice()) {
        Ok(t) => t,
        Err(e) => panic!("failed to decode: {:?}\n{}", e, s),
    };

    fn fixup(json: Json) -> Json {
        match json {
            Json::Object(object) => {
                Json::Object(object.into_iter().map(|(k, v)| {
                    let k = if k.as_slice() == "crate" {
                        "krate".to_string()
                    } else {
                        k
                    };
                    (k, fixup(v))
                }).collect())
            }
            Json::Array(list) => {
                Json::Array(list.into_iter().map(fixup).collect())
            }
            j => j,
        }
    }
}

fn user(login: &str) -> User {
    User {
        id: 10000,
        gh_login: login.to_string(),
        email: None,
        name: None,
        avatar: None,
        gh_access_token: User::new_api_token(), // just randomize it
        api_token: User::new_api_token(),
    }
}

fn krate(name: &str) -> Crate {
    cargo_registry::krate::Crate {
        id: 10000,
        name: name.to_string(),
        user_id: 100,
        updated_at: time::now().to_timespec(),
        created_at: time::now().to_timespec(),
        downloads: 10,
        max_version: semver::Version::parse("0.0.0").unwrap(),
        documentation: None,
        homepage: None,
        description: None,
        readme: None,
        keywords: Vec::new(),
        license: None,
        repository: None,
    }
}

fn mock_user(req: &mut Request, u: User) -> User {
    let u = User::find_or_insert(req.tx().unwrap(),
                                 u.gh_login.as_slice(),
                                 u.email.as_ref().map(|s| s.as_slice()),
                                 u.name.as_ref().map(|s| s.as_slice()),
                                 u.avatar.as_ref().map(|s| s.as_slice()),
                                 u.gh_access_token.as_slice(),
                                 u.api_token.as_slice()).unwrap();
    req.mut_extensions().insert(u.clone());
    return u;
}

fn mock_crate(req: &mut Request, krate: Crate) -> (Crate, Version) {
    mock_crate_vers(req, krate, &semver::Version::parse("1.0.0").unwrap())
}
fn mock_crate_vers(req: &mut Request, krate: Crate, v: &semver::Version)
                   -> (Crate, Version) {
    let user = req.extensions().find::<User>().unwrap();
    let mut krate = Crate::find_or_insert(req.tx().unwrap(), krate.name.as_slice(),
                                      user.id, &krate.description,
                                      &krate.homepage,
                                      &krate.documentation,
                                      &krate.readme,
                                      krate.keywords.as_slice(),
                                      &krate.repository,
                                      &krate.license,
                                      &None).unwrap();
    Keyword::update_crate(req.tx().unwrap(), &krate,
                          krate.keywords.as_slice()).unwrap();
    let v = krate.add_version(req.tx().unwrap(), v, &HashMap::new(), &[]);
    (krate, v.unwrap())
}

fn mock_dep(req: &mut Request, version: &Version, krate: &Crate,
            target: Option<&str>) -> Dependency {
    Dependency::insert(req.tx().unwrap(),
                       version.id,
                       krate.id,
                       &semver::VersionReq::parse(">= 0").unwrap(),
                       Kind::Normal,
                       false, true, &[],
                       &target.map(|s| s.to_string())).unwrap()
}

fn mock_keyword(req: &mut Request, name: &str) -> Keyword {
    Keyword::find_or_insert(req.tx().unwrap(), name).unwrap()
}

fn logout(req: &mut Request) {
    req.mut_extensions().pop::<User>();
}
<|endoftext|>


<commit_before>use std::iter::Peekable;
use std::str::FromStr;
use std::time::Duration;

use game::{Position, Move, White, Black, fen_to_position, NumPlies, NumMoves};
use super::types::{CmdVal, RegisterParam, GoParam};
use super::types::options;
use types::NumNodes;

pub fn parse(s: &str) -> Option<CmdVal> {
    let mut words = s.words().peekable();

    let mut cmd_val: Option<CmdVal> = None;

    while cmd_val.is_none() {
        let next_word_opt = words.next();
        let next_word = match next_word_opt {
            Some(val) => val,
            None => break,
        };

        cmd_val = match next_word {
            "uci" => Some(CmdVal::Uci),
            "debug" => parse_on_off(&mut words).map(|val| CmdVal::Debug(val)),
            "isready" => Some(CmdVal::IsReady),
            "setoption" => parse_option_val(&mut words).map(|(name, val)| CmdVal::SetOption(name, val)),
            "register" => parse_register_vec(&mut words).map(|val| CmdVal::Register(val)),
            "ucinewgame" => Some(CmdVal::UciNewGame),
            "position" => {
                let pos_opt = parse_position(&mut words);
                match parse_move_vec(&mut words) {
                    Some(moves) => Some(CmdVal::SetupPosition(pos_opt, moves)),
                    None => None,
                }
            },
            "go" => parse_go_param_vec(&mut words).map(|val| CmdVal::Go(val)),
            "stop" => Some(CmdVal::Stop),
            "ponderhit" => Some(CmdVal::PonderHit),
            "quit" => Some(CmdVal::Quit),
            _ => None,
        }
    }

    cmd_val
}

fn parse_on_off(words: &mut Iterator<&str>) -> Option<bool> {
    let mut ans = None;
    while ans.is_none() {
        let next_word_opt = words.next();
        let next_word = match next_word_opt {
            Some(val) => val,
            None => break,
        };
        ans = match next_word {
            "on" => Some(true),
            "off" => Some(false),
            _ => None,
        }
    }
    ans
}

fn parse_option_val<'a,T>(words: &mut T) -> Option<options::NameAndVal>
where T: Iterator<&'a str> {
    // TODO parse options.
    Some((options::Name::Dummy, options::Val::Spin(1)))
}

fn parse_register_vec<'a,T>(words: &mut Peekable<&'a str,T>) -> Option<Vec<RegisterParam>>
where T: Iterator<&'a str> {
    let mut res = Vec::<RegisterParam>::new();
    while let Some(next_word) = words.next() {
        let register_val = match next_word {
            "later" => Some(RegisterParam::Later),
            "name" => {
                let mut name_vec = Vec::<&str>::new();
                loop {
                    let curr_name = words.peek().and_then( |s| {
                        if ["late","name","code"].contains(s) { None }
                        else { Some(*s) }
                    });
                    match curr_name {
                        Some(val) => {
                            name_vec.push(val);
                            words.next();
                        },
                        None => break,
                    }
                }
                Some(RegisterParam::Name(name_vec.connect(" ")))
            },
            "code" => words.next().map( |x| RegisterParam::Code(x.to_string())),
            _ => None,
        };
        if let Some(val) = register_val {
            res.push(val);
        }
    }
    if res.is_empty() { None } else { Some(res) }
}

fn parse_position<'a,T>(words: &mut T) -> Option<Position>
where T: Iterator<&'a str> {
    let six_words: Vec<_> = words.by_ref().take(6).collect();
    fen_to_position(&*six_words.connect(" ")).ok()
}

fn parse_move_vec<'a,T>(words: &mut Peekable<&'a str,T>) -> Option<Vec<Move>>
where T: Iterator<&'a str> {
    let mut res = Vec::<Move>::new();
    while let Some(val) = words.peek().and_then(|val| FromStr::from_str(*val)) {
        res.push(val);
        words.next();
    }
    if res.is_empty() { None } else { Some(res) }
}

fn parse_go_param_vec<'a,T>(words: &mut Peekable<&'a str,T>) -> Option<Vec<GoParam>>
where T: Iterator<&'a str> {
    let mut res = Vec::<GoParam>::new();
    while let Some(next_word) = words.next() {
        match next_word {
            "search moves" => parse_move_vec(words).map(|x| GoParam::SearchMoves(x)),
            "ponder" => Some(GoParam::Ponder),
            "wtime"     => words.next().and_then(|s| s.parse::<i64>())
                                       .map(|x| GoParam::Time(White,Duration::milliseconds(x))),
            "btime"     => words.next().and_then(|s| s.parse::<i64>())
                                       .map(|x| GoParam::Time(Black,Duration::milliseconds(x))),
            "movestogo" => words.next().and_then(|s| s.parse::<u16>())
                                       .map(|x| GoParam::MovesToGo(NumMoves(x))),
            "depth"     => words.next().and_then(|s| s.parse::<u16>())
                                       .map(|x| GoParam::Depth(NumPlies(x))),
            "nodes"     => words.next().and_then(|s| s.parse::<u64>())
                                       .map(|x| GoParam::Nodes(NumNodes(x))),
            "mate"      => words.next().and_then(|s| s.parse::<u16>())
                                       .map(|x| GoParam::Mate(NumMoves(x))),
            "movetime"  => words.next().and_then(|s| s.parse::<i64>())
                                       .map(|x| GoParam::MoveTime(Duration::milliseconds(x))),
            "infinite" => Some(GoParam::Infinite),
            _ => None,
        }.map(|val| res.push(val));
    }
    if res.is_empty() { None } else { Some(res) }
}
<commit_msg>Fix lack of parsing "winc" and "binc".<commit_after>use std::iter::Peekable;
use std::str::FromStr;
use std::time::Duration;

use game::{Position, Move, White, Black, fen_to_position, NumPlies, NumMoves};
use super::types::{CmdVal, RegisterParam, GoParam};
use super::types::options;
use types::NumNodes;

pub fn parse(s: &str) -> Option<CmdVal> {
    let mut words = s.words().peekable();

    let mut cmd_val: Option<CmdVal> = None;

    while cmd_val.is_none() {
        let next_word_opt = words.next();
        let next_word = match next_word_opt {
            Some(val) => val,
            None => break,
        };

        cmd_val = match next_word {
            "uci" => Some(CmdVal::Uci),
            "debug" => parse_on_off(&mut words).map(|val| CmdVal::Debug(val)),
            "isready" => Some(CmdVal::IsReady),
            "setoption" => parse_option_val(&mut words).map(|(name, val)| CmdVal::SetOption(name, val)),
            "register" => parse_register_vec(&mut words).map(|val| CmdVal::Register(val)),
            "ucinewgame" => Some(CmdVal::UciNewGame),
            "position" => {
                let pos_opt = parse_position(&mut words);
                match parse_move_vec(&mut words) {
                    Some(moves) => Some(CmdVal::SetupPosition(pos_opt, moves)),
                    None => None,
                }
            },
            "go" => parse_go_param_vec(&mut words).map(|val| CmdVal::Go(val)),
            "stop" => Some(CmdVal::Stop),
            "ponderhit" => Some(CmdVal::PonderHit),
            "quit" => Some(CmdVal::Quit),
            _ => None,
        }
    }

    cmd_val
}

fn parse_on_off(words: &mut Iterator<&str>) -> Option<bool> {
    let mut ans = None;
    while ans.is_none() {
        let next_word_opt = words.next();
        let next_word = match next_word_opt {
            Some(val) => val,
            None => break,
        };
        ans = match next_word {
            "on" => Some(true),
            "off" => Some(false),
            _ => None,
        }
    }
    ans
}

fn parse_option_val<'a,T>(words: &mut T) -> Option<options::NameAndVal>
where T: Iterator<&'a str> {
    // TODO parse options.
    Some((options::Name::Dummy, options::Val::Spin(1)))
}

fn parse_register_vec<'a,T>(words: &mut Peekable<&'a str,T>) -> Option<Vec<RegisterParam>>
where T: Iterator<&'a str> {
    let mut res = Vec::<RegisterParam>::new();
    while let Some(next_word) = words.next() {
        let register_val = match next_word {
            "later" => Some(RegisterParam::Later),
            "name" => {
                let mut name_vec = Vec::<&str>::new();
                loop {
                    let curr_name = words.peek().and_then( |s| {
                        if ["late","name","code"].contains(s) { None }
                        else { Some(*s) }
                    });
                    match curr_name {
                        Some(val) => {
                            name_vec.push(val);
                            words.next();
                        },
                        None => break,
                    }
                }
                Some(RegisterParam::Name(name_vec.connect(" ")))
            },
            "code" => words.next().map( |x| RegisterParam::Code(x.to_string())),
            _ => None,
        };
        if let Some(val) = register_val {
            res.push(val);
        }
    }
    if res.is_empty() { None } else { Some(res) }
}

fn parse_position<'a,T>(words: &mut T) -> Option<Position>
where T: Iterator<&'a str> {
    let six_words: Vec<_> = words.by_ref().take(6).collect();
    fen_to_position(&*six_words.connect(" ")).ok()
}

fn parse_move_vec<'a,T>(words: &mut Peekable<&'a str,T>) -> Option<Vec<Move>>
where T: Iterator<&'a str> {
    let mut res = Vec::<Move>::new();
    while let Some(val) = words.peek().and_then(|val| FromStr::from_str(*val)) {
        res.push(val);
        words.next();
    }
    if res.is_empty() { None } else { Some(res) }
}

fn parse_go_param_vec<'a,T>(words: &mut Peekable<&'a str,T>) -> Option<Vec<GoParam>>
where T: Iterator<&'a str> {
    let mut res = Vec::<GoParam>::new();
    while let Some(next_word) = words.next() {
        match next_word {
            "search moves" => parse_move_vec(words).map(|x| GoParam::SearchMoves(x)),
            "ponder" => Some(GoParam::Ponder),
            "wtime"     => words.next().and_then(|s| s.parse::<i64>())
                                       .map(|x| GoParam::Time(White,Duration::milliseconds(x))),
            "btime"     => words.next().and_then(|s| s.parse::<i64>())
                                       .map(|x| GoParam::Time(Black,Duration::milliseconds(x))),
            "winc"      => words.next().and_then(|s| s.parse::<i64>())
                                       .map(|x| GoParam::IncTime(White,Duration::milliseconds(x))),
            "binc"      => words.next().and_then(|s| s.parse::<i64>())
                                       .map(|x| GoParam::IncTime(Black,Duration::milliseconds(x))),
            "movestogo" => words.next().and_then(|s| s.parse::<u16>())
                                       .map(|x| GoParam::MovesToGo(NumMoves(x))),
            "depth"     => words.next().and_then(|s| s.parse::<u16>())
                                       .map(|x| GoParam::Depth(NumPlies(x))),
            "nodes"     => words.next().and_then(|s| s.parse::<u64>())
                                       .map(|x| GoParam::Nodes(NumNodes(x))),
            "mate"      => words.next().and_then(|s| s.parse::<u16>())
                                       .map(|x| GoParam::Mate(NumMoves(x))),
            "movetime"  => words.next().and_then(|s| s.parse::<i64>())
                                       .map(|x| GoParam::MoveTime(Duration::milliseconds(x))),
            "infinite" => Some(GoParam::Infinite),
            _ => None,
        }.map(|val| res.push(val));
    }
    if res.is_empty() { None } else { Some(res) }
}
<|endoftext|>


<commit_before>// Copyright (c) 2015, The Radare Project. All rights reserved.
// See the COPYING file at the top-level directory of this distribution.
// Licensed under the BSD 3-Clause License:
// <http://opensource.org/licenses/BSD-3-Clause>
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Module implements Pipeline and Runner structs to make radeco-lib easier to use.

use ::frontend::parser::{Parser};
use ::frontend::structs::{LOpInfo, LRegInfo};
use ::frontend::r2::R2;
use ::frontend::esilssa::SSAConstruction;
use ::middle::ir::{MInst};
use ::middle::cfg::CFG;
use ::middle::dot;
use ::middle::dce;
use ::middle::ssa::SSAStorage;
use ::analysis::constant_propagation::constant;
use ::middle::ssa::verifier;

use std::io::prelude::*;
use std::fs;
use std::fs::{File};
use std::path::{PathBuf};

macro_rules! out {
	($str: expr, $m: expr) => { if $m { println!($str) } }
}

#[derive(Clone, Copy, Debug)]
pub enum Analysis {
	ConstProp,
}

#[derive(Clone, Copy, Debug)]
pub enum Pipeline {
	ReadFromR2,
	ParseEsil,
	CFG,
	SSA,
	AnalyzeSSA(Analysis),
	DCE,
	Verify,
	CWriter,
}

// Enum that represents the output of a stage in the Pipeline.
// Every stage has a corresponding Pipeout type.
#[derive(Clone)]
pub enum Pipeout {
	Esil(Vec<String>),
	LOpInfo(Vec<LOpInfo>),
	Instructions {i: Vec<MInst>},
	CFG {cfg: CFG},
	SSA {ssa: SSAStorage},
}

impl Pipeout {
	fn to_string(&self) -> String {
		let s = match *self {
			Pipeout::Esil(_) => "esil",
			Pipeout::LOpInfo(_) => "esil",
			Pipeout::Instructions {i: _} => "ir",
			Pipeout::CFG {cfg: _} => "cfg",
			Pipeout::SSA {ssa: _} => "ssa",
		};
		s.to_string()
	}
}

// States all the vars important to various stages of the pipeline together.
// Also acts as the result struct.
#[allow(dead_code)]
pub struct State {
	r2: Option<R2>,
	esil: Option<Vec<String>>,
	pub reg_info: Option<LRegInfo>,
	p: Option<Parser>,
	cfg: Option<CFG>,
	ssa: Option<SSAStorage>,
	pub pipeout: Option<Pipeout>,
}

impl State {
	fn new() -> State {
		State {
			r2: None,
			esil: None,
			reg_info: None,
			p: None,
			cfg: None,
			ssa: None,
			pipeout: None,
		}
	}
}

pub struct Runner {
	name: String,
	bin_name: Option<String>,
	addr: Option<String>,
	verbose: bool,
	pipeline: Vec<Pipeline>,
	results: Vec<Pipeout>,
	outpath: String,
	pub state: State,
}

impl Runner {

	pub fn new(name: String, bin_name: Option<String>,
               addr: Option<String>, verbose: bool,
               pipeline: Vec<Pipeline>, outpath: Option<String>) -> Runner
	{
		Runner {
			name: name,
			bin_name: bin_name,
			addr: addr,
			verbose: verbose,
			pipeline: pipeline,
			results: Vec::new(),
			state: State::new(),
			outpath: outpath.unwrap_or("./outputs".to_owned()),
		}
	}

	fn set_reg_info(&mut self, reg_info: &LRegInfo) {
		self.state.reg_info = Some(reg_info.clone());
	}

	fn set_pipeout(&mut self, pipeout: &Pipeout) {
		self.state.pipeout = Some(pipeout.clone());
	}

	fn read_from_r2(&mut self) {
		assert!(!self.bin_name.is_none());
		assert!(!self.addr.is_none());
		out!("[*] Reading from R2", self.verbose);

		if self.state.r2.is_none() {
			let bin_name = self.bin_name.clone();
			// TODO: Error Handling
			let mut _r2 = R2::new(bin_name).unwrap();
			_r2.init();
			self.state.r2 = Some(_r2);
		}
		let func_info;
		let reg_info;
		match self.state.r2.as_mut() {
			Some(r2) => {
				reg_info = r2.get_reg_info().unwrap();
				let addr = self.addr.clone().unwrap();
				func_info = r2.get_function(&*addr).unwrap();
			},
			None => panic!("Unable to Initialize r2. Something is wrong!"),
		}

		self.set_reg_info(&reg_info);
		self.set_pipeout(&Pipeout::LOpInfo(func_info.ops.unwrap()));
	}

	fn parse_esil(&mut self) {
		let pipein = self.state.pipeout.clone().unwrap();
		out!("[*] Parsing ESIL", self.verbose);
		let mut p = Parser::new(None);
		if let Some(ref r) = self.state.reg_info {
			p.set_register_profile(r);
		}
		match pipein {
			Pipeout::Esil(strs) => {
				for _str in strs {
					p.parse_str(&*_str).ok();
				}
			},
			Pipeout::LOpInfo(mut ops) => {
				for op in ops.iter_mut() {
					p.parse_opinfo(op).ok();
				}
			},
			_  => panic!("Incompatible type found in the pipeline!"),
		}

		let insts = p.emit_insts();
		let pipeout = Pipeout::Instructions {i: insts};
		self.set_pipeout(&pipeout);
		self.state.p = Some(p);
	}

	fn construct_cfg(&mut self) {
		let mut pipein = self.state.pipeout.clone().unwrap();
		out!("[*] Starting CFG Construction", self.verbose);
		match pipein {
			Pipeout::Instructions {i: ref mut insts} => {
				let mut cfg = CFG::new();
				cfg.build(insts);
				let pipeout = Pipeout::CFG {cfg: cfg.clone()};
				self.set_pipeout(&pipeout);
				self.state.cfg = Some(cfg);
			},
			_ => panic!("Incompatible type found in the pipeline!"),
		}
	}

	fn construct_ssa(&mut self) {
		// TODO: Relax this condition.
		assert!(self.state.reg_info.is_some());
		out!("[*] Starting SSA Construction", self.verbose);
		let pipein = self.state.pipeout.clone().unwrap();
		let r = self.state.reg_info.clone().unwrap();
		match pipein {
			Pipeout::CFG {ref cfg} => {
				let mut ssa = SSAStorage::new();
				{
					let mut con = SSAConstruction::new(&mut ssa, &r);
					con.run(cfg);
				}
				let pipeout = Pipeout::SSA {ssa: ssa.clone()};
				self.set_pipeout(&pipeout);
				self.state.ssa = Some(ssa);
			},
			_ => panic!("Incompatible type found in the pipeline!"),
		}
	}

	fn analyze(&mut self, analysis: &Analysis) {
		out!("[*] Starting Analysis", self.verbose);
		let pipein = self.state.pipeout.clone().unwrap();
		let mut ssa = if let Pipeout::SSA { ssa } = pipein {
			ssa
		} else {
			panic!("Incompatible type found in the pipeline!");
		};
		let ssa = match *analysis {
			Analysis::ConstProp => {
				let mut analyzer = constant::Analyzer::new(&mut ssa);
				analyzer.analyze();
				analyzer.emit_ssa()
			},
		};

		self.set_pipeout(&Pipeout::SSA { ssa: ssa.clone() });
	}

	fn dce(&mut self) {
		out!("[*] Running DCE", self.verbose);
		let pipein = self.state.pipeout.clone().unwrap();
		let mut ssa = if let Pipeout::SSA { ssa } = pipein {
			ssa
		} else {
			panic!("Incompatible type found in the pipeline!");
		};

		{
			dce::collect(&mut ssa);
		}

		self.set_pipeout(&Pipeout::SSA { ssa: ssa.clone() });
	}

	fn verify(&mut self) {
		out!("[*] Verifying the Integrity of SSA.", self.verbose);
		let pipein = self.state.pipeout.clone().unwrap();
		let ssa = if let Pipeout::SSA { ssa } = pipein {
			ssa
		} else {
			panic!("Incompatible type found in the pipeline!");
		};

		{
			verifier::verify(&ssa).unwrap();
		}
	}

	// TODO: Return Error. Never panic!()
	pub fn run(&mut self) {
		let pipe_iter = self.pipeline.clone();
		for stage in pipe_iter.iter() {
			match *stage {
				Pipeline::ReadFromR2 => self.read_from_r2(),
				Pipeline::ParseEsil => self.parse_esil(),
				Pipeline::CFG => self.construct_cfg(),
				Pipeline::SSA => self.construct_ssa(),
				Pipeline::AnalyzeSSA(ref a) => self.analyze(a),
				Pipeline::DCE => self.dce(),
				Pipeline::Verify => self.verify(),
				Pipeline::CWriter => unimplemented!(),
			}
			self.results.push(self.state.pipeout.clone().unwrap());
		}
	}

	fn write_file(&self, fname: PathBuf, res: String) {
		let mut file = File::create(fname).ok().expect("Error. Cannot create file!\n");
		file.write_all(res.as_bytes()).ok().expect("Error. Cannot write file!\n");
	}

	pub fn dump(&self) {
		let mut phase_num = 1;
		for res in self.results.iter() {
			let mut write_out = String::new();
			let mut ext;
			match *res {
				Pipeout::Esil(ref s) => {
					ext = "esil";
					for esil in s {
						let tmp = format!("{}\n",esil);
						write_out.push_str(&*tmp);
					}
				},
				Pipeout::LOpInfo(ref ops) => {
					ext = "esil";
					for op in ops {
						let tmp = format!("{}:\t{}\n", op.offset.unwrap(), op.esil.clone().unwrap());
						write_out.push_str(&*tmp);
					}
				},
				Pipeout::Instructions {i: ref insts} => {
					ext = "insts";
					for inst in insts {
						let tmp = format!("0x{:08X}:\t{}\n", inst.addr.val, inst);
						write_out.push_str(&*tmp);
					}
				},
				Pipeout::CFG {ref cfg} => {
					ext = "dot";
					let tmp = dot::emit_dot(cfg);
					write_out.push_str(&*tmp);
				},
				Pipeout::SSA {ref ssa} => {
					ext = "dot";
					let tmp = dot::emit_dot(ssa);
					write_out.push_str(&*tmp);
				},
			}

			// Format of output file name:
			// test_name-type-phase_num.ext
			let mut p = PathBuf::from(&self.outpath);
			p.push(self.name.clone());
			let pstr = format!("{}-res_{}-phase{}", self.name, res.to_string(), phase_num.to_string());
			fs::create_dir_all(&p).ok();
			p.push(pstr);
			p.set_extension(ext);
			self.write_file(p, write_out);
			phase_num += 1;
		}
	}
}
<commit_msg>Adds supporting features to utils for radeco<commit_after>// Copyright (c) 2015, The Radare Project. All rights reserved.
// See the COPYING file at the top-level directory of this distribution.
// Licensed under the BSD 3-Clause License:
// <http://opensource.org/licenses/BSD-3-Clause>
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Module implements Pipeline and Runner structs to make radeco-lib easier to use.

use ::frontend::parser::{Parser};
use ::frontend::structs::{LOpInfo, LRegInfo};
use ::frontend::r2::R2;
use ::frontend::esilssa::SSAConstruction;
use ::middle::ir::{MInst};
use ::middle::cfg::CFG;
use ::middle::dot;
use ::middle::dce;
use ::middle::ssa::SSAStorage;
use ::analysis::constant_propagation::constant;
use ::middle::ssa::verifier;

use std::io::prelude::*;
use std::fs;
use std::fs::{File};
use std::fmt;
use std::path::{PathBuf};

macro_rules! out {
	($str: expr, $m: expr) => { if $m { println!($str) } }
}

#[derive(Clone, Copy, Debug)]
pub enum Analysis {
	ConstProp,
}

#[derive(Clone, Copy, Debug)]
pub enum Pipeline {
	ReadFromR2,
	ParseEsil,
	CFG,
	SSA,
	AnalyzeSSA(Analysis),
	DCE,
	Verify,
	CWriter,
}

// Enum that represents the output of a stage in the Pipeline.
// Every stage has a corresponding Pipeout type.
#[derive(Clone)]
pub enum Pipeout {
	Esil(Vec<String>),
	LOpInfo(Vec<LOpInfo>),
	Instructions {i: Vec<MInst>},
	CFG {cfg: CFG},
	SSA {ssa: SSAStorage},
}

impl Pipeout {
	fn to_string(&self) -> String {
		let s = match *self {
			Pipeout::Esil(_) => "esil",
			Pipeout::LOpInfo(_) => "esil",
			Pipeout::Instructions {i: _} => "ir",
			Pipeout::CFG {cfg: _} => "cfg",
			Pipeout::SSA {ssa: _} => "ssa",
		};
		s.to_string()
	}
}

// States all the vars important to various stages of the pipeline together.
// Also acts as the result struct.
#[allow(dead_code)]
pub struct State {
	r2: Option<R2>,
	esil: Option<Vec<String>>,
	pub reg_info: Option<LRegInfo>,
	p: Option<Parser>,
	cfg: Option<CFG>,
	ssa: Option<SSAStorage>,
	pub pipeout: Option<Pipeout>,
}

impl State {
	fn new() -> State {
		State {
			r2: None,
			esil: None,
			reg_info: None,
			p: None,
			cfg: None,
			ssa: None,
			pipeout: None,
		}
	}
}

pub struct Runner {
	name: String,
	bin_name: Option<String>,
	addr: Option<String>,
	verbose: bool,
	pipeline: Vec<Pipeline>,
	results: Vec<Pipeout>,
	outpath: String,
	pub state: State,
}

impl fmt::Display for Pipeline {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		match *self {
			Pipeline::ReadFromR2 => write!(f, "{}", "r2"),
			Pipeline::ParseEsil => write!(f, "{}", "Parse ESIL"),
			Pipeline::CFG => write!(f, "{}", "Control Flow Graph Construction"),
			Pipeline::SSA => write!(f, "{}", "SSA Construction"),
			Pipeline::AnalyzeSSA(_) => write!(f, "{}", "Constant Propagation"),
			Pipeline::DCE => write!(f, "{}", "Dead Code Elimination"),
			Pipeline::Verify => write!(f, "{}", "Verify SSA"),
			Pipeline::CWriter => write!(f, "{}", "C Writer"),
		}
	}
}

impl fmt::Display for Runner {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		let mut res;
		res = format!("Name: {}\n", self.name);
		res = format!("{}Binary: {}\n", res, self.bin_name.as_ref().unwrap_or(&"-".to_owned()));
		res = format!("{}Address: {}\n", res, self.addr.as_ref().unwrap_or(&"-".to_owned()));
		res = format!("{}Pipeline:\n", res);
		for p in &self.pipeline {
			res = format!("{}    - {}\n", res, p);
		}
		res = format!("{}Output: {}", res, self.outpath);
		write!(f, "{}", res)
	}
}

impl Runner {

	pub fn new(name: String, bin_name: Option<String>,
               addr: Option<String>, verbose: bool,
               pipeline: Vec<Pipeline>, outpath: Option<String>) -> Runner
	{
		Runner {
			name: name,
			bin_name: bin_name,
			addr: addr,
			verbose: verbose,
			pipeline: pipeline,
			results: Vec::new(),
			state: State::new(),
			outpath: outpath.unwrap_or("./outputs".to_owned()),
		}
	}

	pub fn is_verbose(&self) -> bool {
		self.verbose
	}

	fn set_reg_info(&mut self, reg_info: &LRegInfo) {
		self.state.reg_info = Some(reg_info.clone());
	}

	fn set_pipeout(&mut self, pipeout: &Pipeout) {
		self.state.pipeout = Some(pipeout.clone());
	}

	fn read_from_r2(&mut self) {
		//assert!(!self.bin_name.is_none());
		//assert!(!self.addr.is_none());
		out!("[*] Reading from R2", self.verbose);

		if self.state.r2.is_none() {
			let bin_name = self.bin_name.clone();
			// TODO: Error Handling
			let mut _r2 = R2::new(bin_name).unwrap();
			_r2.init();
			self.state.r2 = Some(_r2);
		}
		let func_info;
		let reg_info;
		match self.state.r2.as_mut() {
			Some(r2) => {
				reg_info = r2.get_reg_info().unwrap();
				let addr = self.addr.clone().unwrap();
				func_info = r2.get_function(&*addr).unwrap();
			},
			None => panic!("Unable to Initialize r2. Something is wrong!"),
		}

		self.set_reg_info(&reg_info);
		self.set_pipeout(&Pipeout::LOpInfo(func_info.ops.unwrap()));
	}

	fn parse_esil(&mut self) {
		let pipein = self.state.pipeout.clone().unwrap();
		out!("[*] Parsing ESIL", self.verbose);
		let mut p = Parser::new(None);
		if let Some(ref r) = self.state.reg_info {
			p.set_register_profile(r);
		}
		match pipein {
			Pipeout::Esil(strs) => {
				for _str in strs {
					p.parse_str(&*_str).ok();
				}
			},
			Pipeout::LOpInfo(mut ops) => {
				for op in ops.iter_mut() {
					p.parse_opinfo(op).ok();
				}
			},
			_  => panic!("Incompatible type found in the pipeline!"),
		}

		let insts = p.emit_insts();
		let pipeout = Pipeout::Instructions {i: insts};
		self.set_pipeout(&pipeout);
		self.state.p = Some(p);
	}

	fn construct_cfg(&mut self) {
		let mut pipein = self.state.pipeout.clone().unwrap();
		out!("[*] Starting CFG Construction", self.verbose);
		match pipein {
			Pipeout::Instructions {i: ref mut insts} => {
				let mut cfg = CFG::new();
				cfg.build(insts);
				let pipeout = Pipeout::CFG {cfg: cfg.clone()};
				self.set_pipeout(&pipeout);
				self.state.cfg = Some(cfg);
			},
			_ => panic!("Incompatible type found in the pipeline!"),
		}
	}

	fn construct_ssa(&mut self) {
		// TODO: Relax this condition.
		assert!(self.state.reg_info.is_some());
		out!("[*] Starting SSA Construction", self.verbose);
		let pipein = self.state.pipeout.clone().unwrap();
		let r = self.state.reg_info.clone().unwrap();
		match pipein {
			Pipeout::CFG {ref cfg} => {
				let mut ssa = SSAStorage::new();
				{
					let mut con = SSAConstruction::new(&mut ssa, &r);
					con.run(cfg);
				}
				let pipeout = Pipeout::SSA {ssa: ssa.clone()};
				self.set_pipeout(&pipeout);
				self.state.ssa = Some(ssa);
			},
			_ => panic!("Incompatible type found in the pipeline!"),
		}
	}

	fn analyze(&mut self, analysis: &Analysis) {
		out!("[*] Starting Analysis", self.verbose);
		let pipein = self.state.pipeout.clone().unwrap();
		let mut ssa = if let Pipeout::SSA { ssa } = pipein {
			ssa
		} else {
			panic!("Incompatible type found in the pipeline!");
		};
		let ssa = match *analysis {
			Analysis::ConstProp => {
				let mut analyzer = constant::Analyzer::new(&mut ssa);
				analyzer.analyze();
				analyzer.emit_ssa()
			},
		};

		self.set_pipeout(&Pipeout::SSA { ssa: ssa.clone() });
	}

	fn dce(&mut self) {
		out!("[*] Running DCE", self.verbose);
		let pipein = self.state.pipeout.clone().unwrap();
		let mut ssa = if let Pipeout::SSA { ssa } = pipein {
			ssa
		} else {
			panic!("Incompatible type found in the pipeline!");
		};

		{
			dce::collect(&mut ssa);
		}

		self.set_pipeout(&Pipeout::SSA { ssa: ssa.clone() });
	}

	fn verify(&mut self) {
		out!("[*] Verifying the Integrity of SSA.", self.verbose);
		let pipein = self.state.pipeout.clone().unwrap();
		let ssa = if let Pipeout::SSA { ssa } = pipein {
			ssa
		} else {
			panic!("Incompatible type found in the pipeline!");
		};

		{
			verifier::verify(&ssa).unwrap();
		}
	}

	// TODO: Return Error. Never panic!()
	pub fn run(&mut self) {
		let pipe_iter = self.pipeline.clone();
		for stage in pipe_iter.iter() {
			match *stage {
				Pipeline::ReadFromR2 => self.read_from_r2(),
				Pipeline::ParseEsil => self.parse_esil(),
				Pipeline::CFG => self.construct_cfg(),
				Pipeline::SSA => self.construct_ssa(),
				Pipeline::AnalyzeSSA(ref a) => self.analyze(a),
				Pipeline::DCE => self.dce(),
				Pipeline::Verify => self.verify(),
				Pipeline::CWriter => unimplemented!(),
			}
			self.results.push(self.state.pipeout.clone().unwrap());
		}
	}

	fn write_file(&self, fname: PathBuf, res: String) {
		let mut file = File::create(fname).ok().expect("Error. Cannot create file!\n");
		file.write_all(res.as_bytes()).ok().expect("Error. Cannot write file!\n");
	}

	pub fn output(&self, phases: Option<Vec<u16>>) {
		let mut phase_num = 1;
		let count = self.pipeline.len() - 1;
		let phases = phases.unwrap_or((0..count as u16).collect::<Vec<_>>());
		for res in self.results.iter() {
			if !phases.contains(&phase_num) {
				phase_num += 1;
				continue;
			}
			let mut write_out = String::new();
			let mut ext;
			match *res {
				Pipeout::Esil(ref s) => {
					ext = "esil";
					for esil in s {
						let tmp = format!("{}\n",esil);
						write_out.push_str(&*tmp);
					}
				},
				Pipeout::LOpInfo(ref ops) => {
					ext = "esil";
					for op in ops {
						let tmp = format!("{}:\t{}\n", op.offset.unwrap(), op.esil.clone().unwrap());
						write_out.push_str(&*tmp);
					}
				},
				Pipeout::Instructions {i: ref insts} => {
					ext = "insts";
					for inst in insts {
						let tmp = format!("0x{:08X}:\t{}\n", inst.addr.val, inst);
						write_out.push_str(&*tmp);
					}
				},
				Pipeout::CFG {ref cfg} => {
					ext = "dot";
					let tmp = dot::emit_dot(cfg);
					write_out.push_str(&*tmp);
				},
				Pipeout::SSA {ref ssa} => {
					ext = "dot";
					let tmp = dot::emit_dot(ssa);
					write_out.push_str(&*tmp);
				},
			}

			// Format of output file name:
			// test_name-type-phase_num.ext
			let mut p = PathBuf::from(&self.outpath);
			p.push(self.name.clone());
			let pstr = format!("{}-res_{}-phase{}", self.name, res.to_string(), phase_num.to_string());
			fs::create_dir_all(&p).ok();
			p.push(pstr);
			p.set_extension(ext);
			self.write_file(p, write_out);
			phase_num += 1;
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use serde::Deserialize;
use serde::de::{self, Deserializer, Visitor};

use std::any::{Any, TypeId};
use std::borrow::Borrow;
use std::cell::RefCell;
use std::collections::HashMap;
use std::convert::AsRef;
use std::ffi::OsStr;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;
use std::mem;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::sync::Mutex;

use builder::Step;

pub struct Interned<T>(usize, PhantomData<*const T>);

impl<'de> Deserialize<'de> for Interned<String> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
        where D: Deserializer<'de>
    {
        struct StrVisitor;

        impl<'de> Visitor<'de> for StrVisitor {
            type Value = &'de str;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("a str")
            }

            fn visit_borrowed_str<E>(self, value: &'de str) -> Result<&'de str, E>
                where E: de::Error
            {
                Ok(value)
            }
        }

        Ok(INTERNER.intern_str(deserializer.deserialize_str(StrVisitor)?))
    }
}
impl Default for Interned<String> {
    fn default() -> Self {
        INTERNER.intern_string(String::default())
    }
}

impl Default for Interned<PathBuf> {
    fn default() -> Self {
        INTERNER.intern_path(PathBuf::default())
    }
}

impl<T> Copy for Interned<T> {}
impl<T> Clone for Interned<T> {
    fn clone(&self) -> Interned<T> {
        *self
    }
}

impl<T> PartialEq for Interned<T> {
    fn eq(&self, other: &Self) -> bool {
        self.0 == other.0
    }
}
impl<T> Eq for Interned<T> {}

impl PartialEq<str> for Interned<String> {
    fn eq(&self, other: &str) -> bool {
       *self == other
    }
}
impl<'a> PartialEq<&'a str> for Interned<String> {
    fn eq(&self, other: &&str) -> bool {
        **self == **other
    }
}
impl<'a, T> PartialEq<&'a Interned<T>> for Interned<T> {
    fn eq(&self, other: &&Self) -> bool {
        self.0 == other.0
    }
}
impl<'a, T> PartialEq<Interned<T>> for &'a Interned<T> {
    fn eq(&self, other: &Interned<T>) -> bool {
        self.0 == other.0
    }
}

unsafe impl<T> Send for Interned<T> {}
unsafe impl<T> Sync for Interned<T> {}

impl fmt::Display for Interned<String> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let s: &str = &*self;
        f.write_str(s)
    }
}

impl fmt::Debug for Interned<String> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let s: &str = &*self;
        f.write_fmt(format_args!("{:?}", s))
    }
}
impl fmt::Debug for Interned<PathBuf> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let s: &Path = &*self;
        f.write_fmt(format_args!("{:?}", s))
    }
}

impl Hash for Interned<String> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        let l = INTERNER.strs.lock().unwrap();
        l.get(*self).hash(state)
    }
}

impl Hash for Interned<PathBuf> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        let l = INTERNER.paths.lock().unwrap();
        l.get(*self).hash(state)
    }
}

impl Deref for Interned<String> {
    type Target = str;
    fn deref(&self) -> &'static str {
        let l = INTERNER.strs.lock().unwrap();
        unsafe { mem::transmute::<&str, &'static str>(l.get(*self)) }
    }
}

impl Deref for Interned<PathBuf> {
    type Target = Path;
    fn deref(&self) -> &'static Path {
        let l = INTERNER.paths.lock().unwrap();
        unsafe { mem::transmute::<&Path, &'static Path>(l.get(*self)) }
    }
}

impl AsRef<Path> for Interned<PathBuf> {
    fn as_ref(&self) -> &'static Path {
        let l = INTERNER.paths.lock().unwrap();
        unsafe { mem::transmute::<&Path, &'static Path>(l.get(*self)) }
    }
}

impl AsRef<Path> for Interned<String> {
    fn as_ref(&self) -> &'static Path {
        let l = INTERNER.strs.lock().unwrap();
        unsafe { mem::transmute::<&Path, &'static Path>(l.get(*self).as_ref()) }
    }
}

impl AsRef<OsStr> for Interned<PathBuf> {
    fn as_ref(&self) -> &'static OsStr {
        let l = INTERNER.paths.lock().unwrap();
        unsafe { mem::transmute::<&OsStr, &'static OsStr>(l.get(*self).as_ref()) }
    }
}

impl AsRef<OsStr> for Interned<String> {
    fn as_ref(&self) -> &'static OsStr {
        let l = INTERNER.strs.lock().unwrap();
        unsafe { mem::transmute::<&OsStr, &'static OsStr>(l.get(*self).as_ref()) }
    }
}

struct TyIntern<T> {
    items: Vec<T>,
    set: HashMap<T, Interned<T>>,
}

impl<T: Hash + Clone + Eq> TyIntern<T> {
    fn new() -> TyIntern<T> {
        TyIntern {
            items: Vec::new(),
            set: HashMap::new(),
        }
    }

    fn intern_borrow<B>(&mut self, item: &B) -> Interned<T>
    where
        B: Eq + Hash + ToOwned<Owned=T> + ?Sized,
        T: Borrow<B>,
    {
        if let Some(i) = self.set.get(&item) {
            return *i;
        }
        let item = item.to_owned();
        let interned =  Interned(self.items.len(), PhantomData::<*const T>);
        self.set.insert(item.clone(), interned);
        self.items.push(item);
        interned
    }

    fn intern(&mut self, item: T) -> Interned<T> {
        if let Some(i) = self.set.get(&item) {
            return *i;
        }
        let interned =  Interned(self.items.len(), PhantomData::<*const T>);
        self.set.insert(item.clone(), interned);
        self.items.push(item);
        interned
    }

    fn get(&self, i: Interned<T>) -> &T {
        &self.items[i.0]
    }
}

pub struct Interner {
    strs: Mutex<TyIntern<String>>,
    paths: Mutex<TyIntern<PathBuf>>,
}

impl Interner {
    fn new() -> Interner {
        Interner {
            strs: Mutex::new(TyIntern::new()),
            paths: Mutex::new(TyIntern::new()),
        }
    }

    pub fn intern_str(&self, s: &str) -> Interned<String> {
        self.strs.lock().unwrap().intern_borrow(s)
    }
    pub fn intern_string(&self, s: String) -> Interned<String> {
        self.strs.lock().unwrap().intern(s)
    }

    pub fn intern_path(&self, s: PathBuf) -> Interned<PathBuf> {
        self.paths.lock().unwrap().intern(s)
    }
}

lazy_static! {
    pub static ref INTERNER: Interner = Interner::new();
}

/// This is essentially a HashMap which allows storing any type in its input and
/// any type in its output. It is a write-once cache; values are never evicted,
/// which means that references to the value can safely be returned from the
/// get() method.
//
// FIXME: This type does not permit retrieving &Path from a PathBuf, primarily
// due to a lack of any obvious way to ensure that this is safe, but also not
#[derive(Debug)]
pub struct Cache(
    RefCell<HashMap<
        TypeId,
        Box<Any>, // actually a HashMap<Step, Interned<Step::Output>>
    >>
);

impl Cache {
    pub fn new() -> Cache {
        Cache(RefCell::new(HashMap::new()))
    }

    pub fn put<S: Step>(&self, step: S, value: S::Output) {
        let mut cache = self.0.borrow_mut();
        let type_id = TypeId::of::<S>();
        let stepcache = cache.entry(type_id)
                        .or_insert_with(|| Box::new(HashMap::<S, S::Output>::new()))
                        .downcast_mut::<HashMap<S, S::Output>>()
                        .expect("invalid type mapped");
        assert!(!stepcache.contains_key(&step), "processing {:?} a second time", step);
        stepcache.insert(step, value);
    }

    pub fn get<S: Step>(&self, step: &S) -> Option<S::Output> {
        let mut cache = self.0.borrow_mut();
        let type_id = TypeId::of::<S>();
        let stepcache = cache.entry(type_id)
                        .or_insert_with(|| Box::new(HashMap::<S, S::Output>::new()))
                        .downcast_mut::<HashMap<S, S::Output>>()
                        .expect("invalid type mapped");
        stepcache.get(step).cloned()
    }
}
<commit_msg>Remove outdated FIXME from cache<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use serde::Deserialize;
use serde::de::{self, Deserializer, Visitor};

use std::any::{Any, TypeId};
use std::borrow::Borrow;
use std::cell::RefCell;
use std::collections::HashMap;
use std::convert::AsRef;
use std::ffi::OsStr;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;
use std::mem;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::sync::Mutex;

use builder::Step;

pub struct Interned<T>(usize, PhantomData<*const T>);

impl<'de> Deserialize<'de> for Interned<String> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
        where D: Deserializer<'de>
    {
        struct StrVisitor;

        impl<'de> Visitor<'de> for StrVisitor {
            type Value = &'de str;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("a str")
            }

            fn visit_borrowed_str<E>(self, value: &'de str) -> Result<&'de str, E>
                where E: de::Error
            {
                Ok(value)
            }
        }

        Ok(INTERNER.intern_str(deserializer.deserialize_str(StrVisitor)?))
    }
}
impl Default for Interned<String> {
    fn default() -> Self {
        INTERNER.intern_string(String::default())
    }
}

impl Default for Interned<PathBuf> {
    fn default() -> Self {
        INTERNER.intern_path(PathBuf::default())
    }
}

impl<T> Copy for Interned<T> {}
impl<T> Clone for Interned<T> {
    fn clone(&self) -> Interned<T> {
        *self
    }
}

impl<T> PartialEq for Interned<T> {
    fn eq(&self, other: &Self) -> bool {
        self.0 == other.0
    }
}
impl<T> Eq for Interned<T> {}

impl PartialEq<str> for Interned<String> {
    fn eq(&self, other: &str) -> bool {
       *self == other
    }
}
impl<'a> PartialEq<&'a str> for Interned<String> {
    fn eq(&self, other: &&str) -> bool {
        **self == **other
    }
}
impl<'a, T> PartialEq<&'a Interned<T>> for Interned<T> {
    fn eq(&self, other: &&Self) -> bool {
        self.0 == other.0
    }
}
impl<'a, T> PartialEq<Interned<T>> for &'a Interned<T> {
    fn eq(&self, other: &Interned<T>) -> bool {
        self.0 == other.0
    }
}

unsafe impl<T> Send for Interned<T> {}
unsafe impl<T> Sync for Interned<T> {}

impl fmt::Display for Interned<String> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let s: &str = &*self;
        f.write_str(s)
    }
}

impl fmt::Debug for Interned<String> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let s: &str = &*self;
        f.write_fmt(format_args!("{:?}", s))
    }
}
impl fmt::Debug for Interned<PathBuf> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let s: &Path = &*self;
        f.write_fmt(format_args!("{:?}", s))
    }
}

impl Hash for Interned<String> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        let l = INTERNER.strs.lock().unwrap();
        l.get(*self).hash(state)
    }
}

impl Hash for Interned<PathBuf> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        let l = INTERNER.paths.lock().unwrap();
        l.get(*self).hash(state)
    }
}

impl Deref for Interned<String> {
    type Target = str;
    fn deref(&self) -> &'static str {
        let l = INTERNER.strs.lock().unwrap();
        unsafe { mem::transmute::<&str, &'static str>(l.get(*self)) }
    }
}

impl Deref for Interned<PathBuf> {
    type Target = Path;
    fn deref(&self) -> &'static Path {
        let l = INTERNER.paths.lock().unwrap();
        unsafe { mem::transmute::<&Path, &'static Path>(l.get(*self)) }
    }
}

impl AsRef<Path> for Interned<PathBuf> {
    fn as_ref(&self) -> &'static Path {
        let l = INTERNER.paths.lock().unwrap();
        unsafe { mem::transmute::<&Path, &'static Path>(l.get(*self)) }
    }
}

impl AsRef<Path> for Interned<String> {
    fn as_ref(&self) -> &'static Path {
        let l = INTERNER.strs.lock().unwrap();
        unsafe { mem::transmute::<&Path, &'static Path>(l.get(*self).as_ref()) }
    }
}

impl AsRef<OsStr> for Interned<PathBuf> {
    fn as_ref(&self) -> &'static OsStr {
        let l = INTERNER.paths.lock().unwrap();
        unsafe { mem::transmute::<&OsStr, &'static OsStr>(l.get(*self).as_ref()) }
    }
}

impl AsRef<OsStr> for Interned<String> {
    fn as_ref(&self) -> &'static OsStr {
        let l = INTERNER.strs.lock().unwrap();
        unsafe { mem::transmute::<&OsStr, &'static OsStr>(l.get(*self).as_ref()) }
    }
}

struct TyIntern<T> {
    items: Vec<T>,
    set: HashMap<T, Interned<T>>,
}

impl<T: Hash + Clone + Eq> TyIntern<T> {
    fn new() -> TyIntern<T> {
        TyIntern {
            items: Vec::new(),
            set: HashMap::new(),
        }
    }

    fn intern_borrow<B>(&mut self, item: &B) -> Interned<T>
    where
        B: Eq + Hash + ToOwned<Owned=T> + ?Sized,
        T: Borrow<B>,
    {
        if let Some(i) = self.set.get(&item) {
            return *i;
        }
        let item = item.to_owned();
        let interned =  Interned(self.items.len(), PhantomData::<*const T>);
        self.set.insert(item.clone(), interned);
        self.items.push(item);
        interned
    }

    fn intern(&mut self, item: T) -> Interned<T> {
        if let Some(i) = self.set.get(&item) {
            return *i;
        }
        let interned =  Interned(self.items.len(), PhantomData::<*const T>);
        self.set.insert(item.clone(), interned);
        self.items.push(item);
        interned
    }

    fn get(&self, i: Interned<T>) -> &T {
        &self.items[i.0]
    }
}

pub struct Interner {
    strs: Mutex<TyIntern<String>>,
    paths: Mutex<TyIntern<PathBuf>>,
}

impl Interner {
    fn new() -> Interner {
        Interner {
            strs: Mutex::new(TyIntern::new()),
            paths: Mutex::new(TyIntern::new()),
        }
    }

    pub fn intern_str(&self, s: &str) -> Interned<String> {
        self.strs.lock().unwrap().intern_borrow(s)
    }
    pub fn intern_string(&self, s: String) -> Interned<String> {
        self.strs.lock().unwrap().intern(s)
    }

    pub fn intern_path(&self, s: PathBuf) -> Interned<PathBuf> {
        self.paths.lock().unwrap().intern(s)
    }
}

lazy_static! {
    pub static ref INTERNER: Interner = Interner::new();
}

/// This is essentially a HashMap which allows storing any type in its input and
/// any type in its output. It is a write-once cache; values are never evicted,
/// which means that references to the value can safely be returned from the
/// get() method.
#[derive(Debug)]
pub struct Cache(
    RefCell<HashMap<
        TypeId,
        Box<Any>, // actually a HashMap<Step, Interned<Step::Output>>
    >>
);

impl Cache {
    pub fn new() -> Cache {
        Cache(RefCell::new(HashMap::new()))
    }

    pub fn put<S: Step>(&self, step: S, value: S::Output) {
        let mut cache = self.0.borrow_mut();
        let type_id = TypeId::of::<S>();
        let stepcache = cache.entry(type_id)
                        .or_insert_with(|| Box::new(HashMap::<S, S::Output>::new()))
                        .downcast_mut::<HashMap<S, S::Output>>()
                        .expect("invalid type mapped");
        assert!(!stepcache.contains_key(&step), "processing {:?} a second time", step);
        stepcache.insert(step, value);
    }

    pub fn get<S: Step>(&self, step: &S) -> Option<S::Output> {
        let mut cache = self.0.borrow_mut();
        let type_id = TypeId::of::<S>();
        let stepcache = cache.entry(type_id)
                        .or_insert_with(|| Box::new(HashMap::<S, S::Output>::new()))
                        .downcast_mut::<HashMap<S, S::Output>>()
                        .expect("invalid type mapped");
        stepcache.get(step).cloned()
    }
}
<|endoftext|>