<commit_before>#![feature(box_syntax)]
#![feature(core)]
#![feature(io)]
#![feature(collections)]
#![feature(path)]
#![feature(env)]

extern crate rdb;
extern crate getopts;
extern crate regex;
use std::env;
use std::old_io::{BufferedReader, File};
use getopts::Options;
use regex::Regex;

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

}

pub fn main() {
    let mut args = env::args();
    let program = args.next().unwrap();
    let mut opts = Options::new();

    opts.optopt("f", "format", "Format to output. Valid: json, plain, nil, protocol", "FORMAT");
    opts.optopt("k", "keys", "Keys to show. Can be a regular expression", "KEYS");
    opts.optmulti("d", "databases", "Database to show. Can be specified multiple times", "DB");
    opts.optmulti("t", "type", "Type to show. Can be specified multiple times", "TYPE");
    opts.optflag("h", "help", "print this help menu");

    let matches = match opts.parse(args) {
        Ok(m) => { m  }
        Err(e) => {
            println!("{}\n", e);
            print_usage(program.as_slice(), opts);
            return;
        }
    };

    if matches.opt_present("h") {
         print_usage(program.as_slice(), opts);
         return;
    }

    let mut filter = rdb::filter::Simple::new();

    for db in matches.opt_strs("d").iter() {
        filter.add_database(db.parse().unwrap());
    }

    for t in matches.opt_strs("t").iter() {
        let typ = match t.as_slice() {
            "string" => rdb::Type::String,
            "list" => rdb::Type::List,
            "set" => rdb::Type::Set,
            "sortedset" | "sorted-set" | "sorted_set" => rdb::Type::SortedSet,
            "hash" => rdb::Type::Hash,
            _ => {
                println!("Unknown type: {}\n", t);
                print_usage(program.as_slice(), opts);
                return;
            }
        };
        filter.add_type(typ);
    }

    if let Some(k) = matches.opt_str("k") {
        let re = match Regex::new(k.as_slice()) {
            Ok(re) => re,
            Err(err) => {
                println!("Incorrect regexp: {:?}\n", err);
                print_usage(program.as_slice(), opts);
                return;
            }
        };
        filter.add_keys(re);
    }

    if matches.free.is_empty() {
        print_usage(program.as_slice(), opts);
        return;
    }

    let path = matches.free[0].clone();
    let file = File::open(&Path::new(path));
    let reader = BufferedReader::new(file);

    let mut res = Ok(());

    if let Some(f) = matches.opt_str("f") {
        match f.as_slice() {
            "json" => {
                res = rdb::parse(reader, rdb::formatter::JSON::new(), filter);
            },
            "plain" => {
                res = rdb::parse(reader, rdb::formatter::Plain::new(), filter);
            },
            "nil" => {
                res = rdb::parse(reader, rdb::formatter::Nil::new(), filter);
            },
            "protocol" => {
                res = rdb::parse(reader, rdb::formatter::Protocol::new(), filter);

            },
            _ => {
                println!("Unknown format: {}\n", f);
                print_usage(program.as_slice(), opts);
            }
        }
    } else {
        res = rdb::parse(reader, rdb::formatter::JSON::new(), filter);
    }

    match res {
        Ok(()) => {},
        Err(e) => {
            println!("");
            let mut stderr = std::old_io::stderr();

            stderr.write_str(format!("Parsing failed: {}\n", e.desc).as_slice()).unwrap();

            if let Some(detail) = e.detail {
                stderr.write_str(format!("Details: {}\n", detail).as_slice()).unwrap();
            }
        }
    }
}
<commit_msg>Remove collections feature<commit_after>#![feature(box_syntax)]
#![feature(core)]
#![feature(io)]
#![feature(path)]
#![feature(env)]

extern crate rdb;
extern crate getopts;
extern crate regex;
use std::env;
use std::old_io::{BufferedReader, File};
use getopts::Options;
use regex::Regex;

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

}

pub fn main() {
    let mut args = env::args();
    let program = args.next().unwrap();
    let mut opts = Options::new();

    opts.optopt("f", "format", "Format to output. Valid: json, plain, nil, protocol", "FORMAT");
    opts.optopt("k", "keys", "Keys to show. Can be a regular expression", "KEYS");
    opts.optmulti("d", "databases", "Database to show. Can be specified multiple times", "DB");
    opts.optmulti("t", "type", "Type to show. Can be specified multiple times", "TYPE");
    opts.optflag("h", "help", "print this help menu");

    let matches = match opts.parse(args) {
        Ok(m) => { m  }
        Err(e) => {
            println!("{}\n", e);
            print_usage(program.as_slice(), opts);
            return;
        }
    };

    if matches.opt_present("h") {
         print_usage(program.as_slice(), opts);
         return;
    }

    let mut filter = rdb::filter::Simple::new();

    for db in matches.opt_strs("d").iter() {
        filter.add_database(db.parse().unwrap());
    }

    for t in matches.opt_strs("t").iter() {
        let typ = match t.as_slice() {
            "string" => rdb::Type::String,
            "list" => rdb::Type::List,
            "set" => rdb::Type::Set,
            "sortedset" | "sorted-set" | "sorted_set" => rdb::Type::SortedSet,
            "hash" => rdb::Type::Hash,
            _ => {
                println!("Unknown type: {}\n", t);
                print_usage(program.as_slice(), opts);
                return;
            }
        };
        filter.add_type(typ);
    }

    if let Some(k) = matches.opt_str("k") {
        let re = match Regex::new(k.as_slice()) {
            Ok(re) => re,
            Err(err) => {
                println!("Incorrect regexp: {:?}\n", err);
                print_usage(program.as_slice(), opts);
                return;
            }
        };
        filter.add_keys(re);
    }

    if matches.free.is_empty() {
        print_usage(program.as_slice(), opts);
        return;
    }

    let path = matches.free[0].clone();
    let file = File::open(&Path::new(path));
    let reader = BufferedReader::new(file);

    let mut res = Ok(());

    if let Some(f) = matches.opt_str("f") {
        match f.as_slice() {
            "json" => {
                res = rdb::parse(reader, rdb::formatter::JSON::new(), filter);
            },
            "plain" => {
                res = rdb::parse(reader, rdb::formatter::Plain::new(), filter);
            },
            "nil" => {
                res = rdb::parse(reader, rdb::formatter::Nil::new(), filter);
            },
            "protocol" => {
                res = rdb::parse(reader, rdb::formatter::Protocol::new(), filter);

            },
            _ => {
                println!("Unknown format: {}\n", f);
                print_usage(program.as_slice(), opts);
            }
        }
    } else {
        res = rdb::parse(reader, rdb::formatter::JSON::new(), filter);
    }

    match res {
        Ok(()) => {},
        Err(e) => {
            println!("");
            let mut stderr = std::old_io::stderr();

            stderr.write_str(format!("Parsing failed: {}\n", e.desc).as_slice()).unwrap();

            if let Some(detail) = e.detail {
                stderr.write_str(format!("Details: {}\n", detail).as_slice()).unwrap();
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::fs;
use std::io::{self, ErrorKind, Read, BufWriter};
use std::path::Path;
use std::process::exit;
use std::time::{Duration, Instant};

use getopts::Options;

use ppbert::prelude::*;
use ppbert::parser::*;
use ppbert::pp::*;

const PROG_NAME: &str = "ppbert";

fn opt_usize(m: &getopts::Matches, opt: &str, default: usize) -> usize {
    match m.opt_get_default(opt, default) {
        Ok(n) => n,
        Err(_) => {
            eprintln!("'{}' must be a number", opt);
            exit(1);
        }
    }
}

fn main() {
    let mut opts = Options::new();
    opts.optflag("V", "version", "display version");
    opts.optflag("h", "help", "display this help");
    opts.optopt("i", "indent", "indent with NUM spaces", "NUM");
    opts.optopt("m", "per-line", "print at most NUM basic terms per line", "NUM");
    opts.optflag("p", "parse", "parse only, not pretty print");
    opts.optflag("1", "bert1", "force ppbert to use regular BERT parser");
    opts.optflag("2", "bert2", "force ppbert to use BERT2 parser");
    opts.optflag("d", "disk-log", "force ppbert to use DiskLog parser");
    opts.optflag("v", "verbose", "show diagnostics on stderr");
    opts.optflag("j", "json", "print as JSON");
    opts.optflag("t", "transform-proplists", "convert proplists to JSON objects");
    opts.optflag("b", "as-bert", "print as BERT");

    let mut matches = match opts.parse(env::args().skip(1)) {
        Ok(m) => m,
        Err(e) => {
            eprintln!("{}: {}", PROG_NAME, e);
            eprintln!("{}", opts.usage(&format!("{} {}", PROG_NAME, VERSION)));
            exit(1);
        }
    };

    if matches.opt_present("help") {
        println!("{}", opts.usage(&format!("{} {}", PROG_NAME, VERSION)));
        exit(0);
    }

    if matches.opt_present("version") {
        println!("{} {}", PROG_NAME, VERSION);
        exit(0);
    }

    // If no files to process, use stdin.
    if matches.free.is_empty() {
        matches.free.push("-".to_owned());
    }

    let indent_width = opt_usize(&matches, "indent", 2);
    let max_per_line = opt_usize(&matches, "per-line", 6);
    let parse_only = matches.opt_present("parse");
    let json = matches.opt_present("json");
    let as_bert = matches.opt_present("as-bert");
    let transform_proplists = matches.opt_present("transform-proplists");
    let verbose = matches.opt_present("verbose");

    let parser_choice: Option<ParserNext> =
        if matches.opt_present("bert1") {
            Some(BertParser::bert1_next)
        } else if matches.opt_present("bert2") {
            Some(BertParser::bert2_next)
        } else if matches.opt_present("disk-log") {
            Some(BertParser::disk_log_next)
        } else {
            None
        };

    let pp: Box<dyn PrettyPrinter> =
        if json {
            Box::new(JsonPrettyPrinter::new(transform_proplists))
        } else if as_bert {
            Box::new(BertWriter::new())
        } else {
            Box::new(ErlangPrettyPrinter::new(indent_width, max_per_line))
        };

    let mut return_code = 0;
    for file in &matches.free {
        let res = handle_file(file, parse_only, verbose, parser_choice, &pp);
        match res {
            Ok(()) => (),
            Err(ref e) => {
                if broken_pipe(e) {
                    break;
                }
                return_code = 1;
                eprintln!("{}: {}: {}", PROG_NAME, file, e);
            }
        }
    }
    exit(return_code);
}

fn broken_pipe(err: &BertError) -> bool {
    match *err {
        BertError::IoError(ref ioerr) =>
            ioerr.kind() == ErrorKind::BrokenPipe,
        _ => false
    }
}

fn read_bytes(filename: &str) -> Result<Vec<u8>> {
    if filename == "-" {
        let stdin = io::stdin();
        let mut stdin = stdin.lock();
        let mut buf: Vec<u8> = Vec::with_capacity(4096);
        stdin.read_to_end(&mut buf)?;
        return Ok(buf);
    } else {
        let buf = fs::read(filename)?;
        return Ok(buf);
    }
}

fn parser_from_ext(filename: &str) -> ParserNext {
    let ext: Option<&str> =
        Path::new(filename)
        .extension()
        .and_then(|x| x.to_str());
    match ext {
        Some("bert") | Some("bert1") => BertParser::bert1_next,
        Some("bert2") => BertParser::bert2_next,
        Some("log") => BertParser::disk_log_next,
        _ => {
            eprintln!("{}: cannot find an appropriate parser for {}; using BERT",
                      PROG_NAME, filename);
            BertParser::bert1_next
        },
    }
}

fn handle_file(
    filename: &str,
    parse_only: bool,
    verbose: bool,
    parser_choice: Option<ParserNext>,
    pp: &Box<dyn PrettyPrinter>,
) -> Result<()> {
    // Read file or stdin into buffer
    let now = Instant::now();
    let bytes = read_bytes(filename)?;
    let read_dur = now.elapsed();
    let mut parser = BertParser::new(bytes);

    let parser_next: ParserNext = match parser_choice {
        Some(f) => f,
        None => parser_from_ext(filename),
    };

    let mut parse_dur = Duration::new(0, 0);
    let mut pp_dur = Duration::new(0, 0);

    loop {
        let now = Instant::now();
        let term = match parser_next(&mut parser) {
            Some(term) => term?,
            None => break,
        };
        parse_dur += now.elapsed();
        if !parse_only {
            let now = Instant::now();
            let stdout = BufWriter::new(io::stdout());
            pp.write(&term, Box::new(stdout))?;
            pp_dur += now.elapsed();
        }
    }

    if verbose {
        eprintln!("{}: {} read time: {:?}", PROG_NAME, filename, read_dur);
        eprintln!("{}: {} parse time: {:?}", PROG_NAME, filename, parse_dur);
        if !parse_only {
            eprintln!("{}: {} print time: {:?}", PROG_NAME, filename, pp_dur);
        }
    }

    return Ok(());
}
<commit_msg>Reduce verbiage for error handling in main()<commit_after>use std::env;
use std::fs;
use std::io::{self, ErrorKind, Read, BufWriter};
use std::path::Path;
use std::process::exit;
use std::time::{Duration, Instant};

use getopts::Options;

use ppbert::prelude::*;
use ppbert::parser::*;
use ppbert::pp::*;

const PROG_NAME: &str = "ppbert";

fn opt_usize(m: &getopts::Matches, opt: &str, default: usize) -> usize {
    match m.opt_get_default(opt, default) {
        Ok(n) => n,
        Err(_) => {
            eprintln!("'{}' must be a number", opt);
            exit(1);
        }
    }
}

fn main() {
    let mut opts = Options::new();
    opts.optflag("V", "version", "display version");
    opts.optflag("h", "help", "display this help");
    opts.optopt("i", "indent", "indent with NUM spaces", "NUM");
    opts.optopt("m", "per-line", "print at most NUM basic terms per line", "NUM");
    opts.optflag("p", "parse", "parse only, not pretty print");
    opts.optflag("1", "bert1", "force ppbert to use regular BERT parser");
    opts.optflag("2", "bert2", "force ppbert to use BERT2 parser");
    opts.optflag("d", "disk-log", "force ppbert to use DiskLog parser");
    opts.optflag("v", "verbose", "show diagnostics on stderr");
    opts.optflag("j", "json", "print as JSON");
    opts.optflag("t", "transform-proplists", "convert proplists to JSON objects");
    opts.optflag("b", "as-bert", "print as BERT");

    let mut matches = match opts.parse(env::args().skip(1)) {
        Ok(m) => m,
        Err(e) => {
            eprintln!("{}: {}", PROG_NAME, e);
            eprintln!("{}", opts.usage(&format!("{} {}", PROG_NAME, VERSION)));
            exit(1);
        }
    };

    if matches.opt_present("help") {
        println!("{}", opts.usage(&format!("{} {}", PROG_NAME, VERSION)));
        exit(0);
    }

    if matches.opt_present("version") {
        println!("{} {}", PROG_NAME, VERSION);
        exit(0);
    }

    // If no files to process, use stdin.
    if matches.free.is_empty() {
        matches.free.push("-".to_owned());
    }

    let indent_width = opt_usize(&matches, "indent", 2);
    let max_per_line = opt_usize(&matches, "per-line", 6);
    let parse_only = matches.opt_present("parse");
    let json = matches.opt_present("json");
    let as_bert = matches.opt_present("as-bert");
    let transform_proplists = matches.opt_present("transform-proplists");
    let verbose = matches.opt_present("verbose");

    let parser_choice: Option<ParserNext> =
        if matches.opt_present("bert1") {
            Some(BertParser::bert1_next)
        } else if matches.opt_present("bert2") {
            Some(BertParser::bert2_next)
        } else if matches.opt_present("disk-log") {
            Some(BertParser::disk_log_next)
        } else {
            None
        };

    let pp: Box<dyn PrettyPrinter> =
        if json {
            Box::new(JsonPrettyPrinter::new(transform_proplists))
        } else if as_bert {
            Box::new(BertWriter::new())
        } else {
            Box::new(ErlangPrettyPrinter::new(indent_width, max_per_line))
        };

    let mut return_code = 0;
    for file in &matches.free {
        if let Err(ref e) = handle_file(file, parse_only, verbose, parser_choice, &pp) {
            if broken_pipe(e) {
                break;
            }
            return_code = 1;
            eprintln!("{}: {}: {}", PROG_NAME, file, e);
        }
    }
    exit(return_code);
}

fn broken_pipe(err: &BertError) -> bool {
    match *err {
        BertError::IoError(ref ioerr) =>
            ioerr.kind() == ErrorKind::BrokenPipe,
        _ => false
    }
}

fn read_bytes(filename: &str) -> Result<Vec<u8>> {
    if filename == "-" {
        let stdin = io::stdin();
        let mut stdin = stdin.lock();
        let mut buf: Vec<u8> = Vec::with_capacity(4096);
        stdin.read_to_end(&mut buf)?;
        return Ok(buf);
    } else {
        let buf = fs::read(filename)?;
        return Ok(buf);
    }
}

fn parser_from_ext(filename: &str) -> ParserNext {
    let ext: Option<&str> =
        Path::new(filename)
        .extension()
        .and_then(|x| x.to_str());
    match ext {
        Some("bert") | Some("bert1") => BertParser::bert1_next,
        Some("bert2") => BertParser::bert2_next,
        Some("log") => BertParser::disk_log_next,
        _ => {
            eprintln!("{}: cannot find an appropriate parser for {}; using BERT",
                      PROG_NAME, filename);
            BertParser::bert1_next
        },
    }
}

fn handle_file(
    filename: &str,
    parse_only: bool,
    verbose: bool,
    parser_choice: Option<ParserNext>,
    pp: &Box<dyn PrettyPrinter>,
) -> Result<()> {
    // Read file or stdin into buffer
    let now = Instant::now();
    let bytes = read_bytes(filename)?;
    let read_dur = now.elapsed();
    let mut parser = BertParser::new(bytes);

    let parser_next: ParserNext = match parser_choice {
        Some(f) => f,
        None => parser_from_ext(filename),
    };

    let mut parse_dur = Duration::new(0, 0);
    let mut pp_dur = Duration::new(0, 0);

    loop {
        let now = Instant::now();
        let term = match parser_next(&mut parser) {
            Some(term) => term?,
            None => break,
        };
        parse_dur += now.elapsed();
        if !parse_only {
            let now = Instant::now();
            let stdout = BufWriter::new(io::stdout());
            pp.write(&term, Box::new(stdout))?;
            pp_dur += now.elapsed();
        }
    }

    if verbose {
        eprintln!("{}: {} read time: {:?}", PROG_NAME, filename, read_dur);
        eprintln!("{}: {} parse time: {:?}", PROG_NAME, filename, parse_dur);
        if !parse_only {
            eprintln!("{}: {} print time: {:?}", PROG_NAME, filename, pp_dur);
        }
    }

    return Ok(());
}
<|endoftext|>


<commit_before>use clap::Parser;
use std::cmp::Ordering;
use std::fs::File;

fn normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)$").unwrap();
    }
    let mut col = RE
        .replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace('(', "")
        .replace(')', "")
        .replace(' ', "_")
        .replace('.', "_")
        .replace('-', "_")
        .replace('/', "_")
        .replace('?', "")
        .replace(',', "_")
        .replace('&', "_")
        .replace(':', "")
        .replace('#', "");
    if !col.chars().next().map(char::is_alphabetic).unwrap_or(true) {
        col = format!("c_{}", col)
    }
    col
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("\"{}\" varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: &str,
    delimiter: u8,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::ReaderBuilder::new()
        .flexible(true)
        .delimiter(delimiter)
        .from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.byte_records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).expect("tx.prepare() failed");
        while let Some(row) = records.next() {
            let mut row = row?;
            match row.len().cmp(&normalized_cols.len()) {
                Ordering::Less => {
                    for _ in 0..normalized_cols.len() - row.len() {
                        row.push_field(b"");
                    }
                }
                Ordering::Greater => {
                    panic!("Too many fields on row {}, fields: {:?}", num_rows + 1, row);
                }
                Ordering::Equal => {}
            }
            stmt.execute(rusqlite::params_from_iter(
                row.iter().map(String::from_utf8_lossy),
            ))
            .unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({}) from {:?}",
        num_rows,
        table_name,
        normalized_cols.join(", "),
        path,
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(query).map_err(|e| e.to_string())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;
    let col_count = stmt.column_count();

    let mut table = comfy_table::Table::new();
    table.load_preset("││──╞═╪╡┆    ┬┴┌┐└┘");
    table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
    let mut title_row = comfy_table::Row::new();
    for col in stmt.column_names() {
        title_row.add_cell(comfy_table::Cell::new(col));
    }
    table.set_header(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = comfy_table::Row::new();
        for i in 0..col_count {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(comfy_table::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    println!("{}", table);
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;
    let col_count = stmt.column_count();

    let mut writer = csv::Writer::from_path(destination_path).map_err(|e| format!("{:?}", e))?;
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..col_count).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else if line.starts_with(".schema") {
        _handle_query(
            conn,
            "SELECT sql AS schema FROM sqlite_master WHERE name like 't%'",
        )
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

fn install_udfs(c: &mut rusqlite::Connection) -> Result<(), Box<dyn std::error::Error>> {
    c.create_scalar_function(
        "regexp_extract",
        3,
        rusqlite::functions::FunctionFlags::SQLITE_UTF8
            | rusqlite::functions::FunctionFlags::SQLITE_DETERMINISTIC,
        |ctx| {
            let re = ctx.get_or_create_aux(
                0,
                |vr| -> Result<_, Box<dyn std::error::Error + Send + Sync + 'static>> {
                    Ok(regex::Regex::new(vr.as_str()?)?)
                },
            )?;
            let value = ctx.get::<Box<str>>(1)?;
            let replacement = ctx.get::<Box<str>>(2)?;

            let caps = match re.captures(&value) {
                Some(caps) => caps,
                None => return Ok("".to_string()),
            };
            let mut dest = String::new();
            caps.expand(&replacement, &mut dest);
            Ok(dest)
        },
    )?;
    c.create_scalar_function(
        "regexp",
        2,
        rusqlite::functions::FunctionFlags::SQLITE_UTF8
            | rusqlite::functions::FunctionFlags::SQLITE_DETERMINISTIC,
        |ctx| {
            let re = ctx.get_or_create_aux(
                0,
                |vr| -> Result<_, Box<dyn std::error::Error + Send + Sync + 'static>> {
                    Ok(regex::Regex::new(vr.as_str()?)?)
                },
            )?;
            let value = ctx.get::<Box<str>>(1)?;
            Ok(re.is_match(&value))
        },
    )?;

    Ok(())
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 5] = [b' ', b'(', b')', b',', b'.'];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    type Hint = String;

    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::validate::Validator for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

#[derive(clap::Parser)]
struct Opts {
    #[clap(long, help = "Use ',' as the delimiter for the CSV")]
    comma: bool,
    #[clap(long, help = "Use '|' as the delimiter for the CSV")]
    pipe: bool,
    #[clap(long, help = "Use '\\t' as the delimiter for the CSV")]
    tab: bool,
    #[clap(long, help = "Use ';' as the delimiter for the CSV")]
    semicolon: bool,

    #[clap()]
    paths: Vec<String>,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let opts = Opts::parse();

    let delim = match (opts.comma, opts.pipe, opts.tab, opts.semicolon) {
        (true, false, false, false) | (false, false, false, false) => b',',
        (false, true, false, false) => b'|',
        (false, false, true, false) => b'\t',
        (false, false, false, true) => b';',
        _ => {
            eprintln!("Can't pass more than one of --comma, --pipe, and --tab");
            std::process::exit(1);
        }
    };

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        // keywords
        "distinct",
        "select",
        "from",
        "group",
        "by",
        "order",
        "where",
        "count",
        "limit",
        "offset",
        // functions
        "length",
        "coalesce",
        "regexp_extract",
        // csv-sql commands
        "export",
        "schema",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if opts.paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", &opts.paths[0], delim)?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in opts.paths.iter().enumerate() {
            let mut col_names =
                _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path, delim)?;
            base_words.append(&mut col_names);
        }
    }

    install_udfs(&mut conn)?;

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    let history_path = dirs::home_dir().unwrap().join(".csv-sql-history");
    let _ = rl.load_history(&history_path);
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }
    rl.save_history(&history_path).unwrap();

    Ok(())
}

#[cfg(test)]
mod test {
    use super::normalize_col;

    #[test]
    fn test_normalize_col() {
        for (value, expected) in &[
            ("", ""),
            ("abc", "abc"),
            ("(S)AO", "sao"),
            ("abc (123)", "abc"),
            ("2/6/2000", "c_2_6_2000"),
            ("COMBO#", "combo"),
        ] {
            assert_eq!(&&normalize_col(value), expected);
        }
    }
}
<commit_msg>clippy fix<commit_after>use clap::Parser;
use std::cmp::Ordering;
use std::fs::File;

fn normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)$").unwrap();
    }
    let mut col = RE
        .replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace(['(', ')'], "")
        .replace([' ', '.', '-', '/'], "_")
        .replace('?', "")
        .replace([',', '&'], "_")
        .replace([':', '#'], "");
    if !col.chars().next().map(char::is_alphabetic).unwrap_or(true) {
        col = format!("c_{}", col)
    }
    col
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("\"{}\" varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: &str,
    delimiter: u8,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::ReaderBuilder::new()
        .flexible(true)
        .delimiter(delimiter)
        .from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.byte_records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).expect("tx.prepare() failed");
        while let Some(row) = records.next() {
            let mut row = row?;
            match row.len().cmp(&normalized_cols.len()) {
                Ordering::Less => {
                    for _ in 0..normalized_cols.len() - row.len() {
                        row.push_field(b"");
                    }
                }
                Ordering::Greater => {
                    panic!("Too many fields on row {}, fields: {:?}", num_rows + 1, row);
                }
                Ordering::Equal => {}
            }
            stmt.execute(rusqlite::params_from_iter(
                row.iter().map(String::from_utf8_lossy),
            ))
            .unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({}) from {:?}",
        num_rows,
        table_name,
        normalized_cols.join(", "),
        path,
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(query).map_err(|e| e.to_string())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;
    let col_count = stmt.column_count();

    let mut table = comfy_table::Table::new();
    table.load_preset("││──╞═╪╡┆    ┬┴┌┐└┘");
    table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
    let mut title_row = comfy_table::Row::new();
    for col in stmt.column_names() {
        title_row.add_cell(comfy_table::Cell::new(col));
    }
    table.set_header(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = comfy_table::Row::new();
        for i in 0..col_count {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(comfy_table::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    println!("{}", table);
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;
    let col_count = stmt.column_count();

    let mut writer = csv::Writer::from_path(destination_path).map_err(|e| format!("{:?}", e))?;
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..col_count).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else if line.starts_with(".schema") {
        _handle_query(
            conn,
            "SELECT sql AS schema FROM sqlite_master WHERE name like 't%'",
        )
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

fn install_udfs(c: &mut rusqlite::Connection) -> Result<(), Box<dyn std::error::Error>> {
    c.create_scalar_function(
        "regexp_extract",
        3,
        rusqlite::functions::FunctionFlags::SQLITE_UTF8
            | rusqlite::functions::FunctionFlags::SQLITE_DETERMINISTIC,
        |ctx| {
            let re = ctx.get_or_create_aux(
                0,
                |vr| -> Result<_, Box<dyn std::error::Error + Send + Sync + 'static>> {
                    Ok(regex::Regex::new(vr.as_str()?)?)
                },
            )?;
            let value = ctx.get::<Box<str>>(1)?;
            let replacement = ctx.get::<Box<str>>(2)?;

            let caps = match re.captures(&value) {
                Some(caps) => caps,
                None => return Ok("".to_string()),
            };
            let mut dest = String::new();
            caps.expand(&replacement, &mut dest);
            Ok(dest)
        },
    )?;
    c.create_scalar_function(
        "regexp",
        2,
        rusqlite::functions::FunctionFlags::SQLITE_UTF8
            | rusqlite::functions::FunctionFlags::SQLITE_DETERMINISTIC,
        |ctx| {
            let re = ctx.get_or_create_aux(
                0,
                |vr| -> Result<_, Box<dyn std::error::Error + Send + Sync + 'static>> {
                    Ok(regex::Regex::new(vr.as_str()?)?)
                },
            )?;
            let value = ctx.get::<Box<str>>(1)?;
            Ok(re.is_match(&value))
        },
    )?;

    Ok(())
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 5] = [b' ', b'(', b')', b',', b'.'];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    type Hint = String;

    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::validate::Validator for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

#[derive(clap::Parser)]
struct Opts {
    #[clap(long, help = "Use ',' as the delimiter for the CSV")]
    comma: bool,
    #[clap(long, help = "Use '|' as the delimiter for the CSV")]
    pipe: bool,
    #[clap(long, help = "Use '\\t' as the delimiter for the CSV")]
    tab: bool,
    #[clap(long, help = "Use ';' as the delimiter for the CSV")]
    semicolon: bool,

    #[clap()]
    paths: Vec<String>,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let opts = Opts::parse();

    let delim = match (opts.comma, opts.pipe, opts.tab, opts.semicolon) {
        (true, false, false, false) | (false, false, false, false) => b',',
        (false, true, false, false) => b'|',
        (false, false, true, false) => b'\t',
        (false, false, false, true) => b';',
        _ => {
            eprintln!("Can't pass more than one of --comma, --pipe, and --tab");
            std::process::exit(1);
        }
    };

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        // keywords
        "distinct",
        "select",
        "from",
        "group",
        "by",
        "order",
        "where",
        "count",
        "limit",
        "offset",
        // functions
        "length",
        "coalesce",
        "regexp_extract",
        // csv-sql commands
        "export",
        "schema",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if opts.paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", &opts.paths[0], delim)?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in opts.paths.iter().enumerate() {
            let mut col_names =
                _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path, delim)?;
            base_words.append(&mut col_names);
        }
    }

    install_udfs(&mut conn)?;

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    let history_path = dirs::home_dir().unwrap().join(".csv-sql-history");
    let _ = rl.load_history(&history_path);
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }
    rl.save_history(&history_path).unwrap();

    Ok(())
}

#[cfg(test)]
mod test {
    use super::normalize_col;

    #[test]
    fn test_normalize_col() {
        for (value, expected) in &[
            ("", ""),
            ("abc", "abc"),
            ("(S)AO", "sao"),
            ("abc (123)", "abc"),
            ("2/6/2000", "c_2_6_2000"),
            ("COMBO#", "combo"),
        ] {
            assert_eq!(&&normalize_col(value), expected);
        }
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::{Read, Write, Result};

mod buffer;
#[cfg(test)]
mod parser_tests;
#[cfg(test)]
mod interp_tests;
#[cfg(test)]
mod opt_test;

#[derive(Debug, Clone, Copy, PartialEq)]
enum Instruction {
    MoveLeft(usize),
    MoveRight(usize),
    Inc(u8),
    Dec(u8),
    Input,
    Output,
    LoopEntry(usize),
    LoopExit(usize),
}

use Instruction::*;

#[derive(Debug, PartialEq)]
struct Program {
    instructions: Vec<Instruction>,
}

impl Program {
    pub fn new(text: String) -> Program {
        // open brackets while parsing
        let mut bracket_stack = Vec::new();
        // idx to patch with index of LoopEntry to patch in
        let mut relocs = Vec::new();
        // "lex" token stream - aka skip whitespace
        let token = text.chars()
            .filter(move |c| {
                        *c == '>' || *c == '<' || *c == '+' || *c == '-' || *c == '.' ||
                        *c == ',' || *c == '[' || *c == ']'
                    });
        let mut instructions = token
            .enumerate()
            .map(|(idx, c)| {
                match c {
                    '<' => MoveLeft(1),
                    '>' => MoveRight(1),
                    '+' => Inc(1),
                    '-' => Dec(1),
                    '.' => Output,
                    ',' => Input,
                    '[' => {
                        bracket_stack.push(idx);
                        // value must be patched later
                        LoopEntry(std::usize::MAX)
                    }
                    ']' => {
                        if let Some(loop_entry) = bracket_stack.pop() {
                            relocs.push((loop_entry, idx));
                            LoopExit(loop_entry)
                        } else {
                            panic!("Unbalanced {:?} at pc={:}", c, idx);
                        }
                    }
                    c => panic!("Unknown instruction {:?} at pc={:}", c, idx),
                }
            })
            .collect::<Vec<_>>();
        if let Some(unbalanced_idx) = bracket_stack.pop() {
            panic!("Unbalanced {:?} at pc={:}", '[', unbalanced_idx);
        }
        for (idx, value) in relocs {
            if idx >= instructions.len() || LoopEntry(std::usize::MAX) != instructions[idx] {
                panic!("Unexpected instruction {:?} at pc={:} for reloc",
                       instructions[idx],
                       idx);
            } else {
                instructions[idx] = LoopEntry(value);
            }
        }
        Program { instructions: instructions }
    }

    pub fn optimize(&self) -> Program {
        // optimized program
        let mut instructions: Vec<Instruction> = Vec::new();
        // open brackets while parsing
        let mut bracket_stack = Vec::new();
        // idx to patch with index of LoopEntry to patch in
        let mut relocs = Vec::new();

        for (idx, instr) in self.instructions.iter().enumerate() {
            match *instr {
                MoveLeft(offset) => {
                    let last_instr = instructions.last().cloned();
                    if let Some(MoveLeft(old_offset)) = last_instr {
                        let last_idx = instructions.len() - 1;
                        instructions[last_idx] = MoveLeft(old_offset + offset);
                    } else {
                        instructions.push(*instr);
                    }
                }
                MoveRight(offset) => {
                    let last_instr = instructions.last().cloned();
                    if let Some(MoveRight(old_offset)) = last_instr {
                        let last_idx = instructions.len() - 1;
                        instructions[last_idx] = MoveRight(old_offset + offset);
                    } else {
                        instructions.push(*instr);
                    }
                }
                Inc(val) => {
                    let last_instr = instructions.last().cloned();
                    if let Some(Inc(old_val)) = last_instr {
                        let last_idx = instructions.len() - 1;
                        instructions[last_idx] = Inc(old_val.wrapping_add(val));
                    } else {
                        instructions.push(*instr);
                    }
                }
                Dec(val) => {
                    let last_instr = instructions.last().cloned();
                    if let Some(Dec(old_val)) = last_instr {
                        let last_idx = instructions.len() - 1;
                        instructions[last_idx] = Dec(old_val.wrapping_add(val));
                    } else {
                        instructions.push(*instr);
                    }
                }
                LoopEntry(_) => {
                    bracket_stack.push(instructions.len());
                    instructions.push(LoopEntry(std::usize::MAX));
                }
                LoopExit(_) => {
                    if let Some(loop_entry) = bracket_stack.pop() {
                        relocs.push((loop_entry, instructions.len()));
                        instructions.push(LoopExit(loop_entry));
                    } else {
                        panic!("Unbalanced {:?} at pc={:}", instr, idx);
                    }
                }
                instr => instructions.push(instr),
            }
        }
        if let Some(unbalanced_idx) = bracket_stack.pop() {
            panic!("Unbalanced {:?} at pc={:}", '[', unbalanced_idx);
        }
        for (idx, value) in relocs {
            if idx >= instructions.len() || LoopEntry(std::usize::MAX) != instructions[idx] {
                panic!("Unexpected instruction {:?} at pc={:} for reloc",
                       instructions[idx],
                       idx);
            } else {
                instructions[idx] = LoopEntry(value);
            }
        }
        Program { instructions: instructions }
    }

    pub fn interp(&self, input: &mut Read, output: &mut Write) {
        let mut memory = vec![0u8; 30000];

        let mut pc: usize = 0;
        let mut dataptr: usize = 0;

        fn get_char(input: &mut Read) -> u8 {
            let mut buf = [0u8; 1];
            match input.read(&mut buf) {
                Err(_) => panic!("Cannot read from stdin"),
                Ok(_) => buf[0],
            }
        }

        while pc < self.instructions.len() {
            match self.instructions[pc] {
                MoveLeft(offset) => dataptr -= offset,
                MoveRight(offset) => dataptr += offset,
                Inc(increment) => memory[dataptr] = memory[dataptr].wrapping_add(increment),
                Dec(decrement) => memory[dataptr] = memory[dataptr].wrapping_sub(decrement),
                Output => write!(output, "{:}", memory[dataptr] as char).expect("Output error"),
                Input => memory[dataptr] = get_char(input),
                LoopEntry(target) => {
                    if 0 == memory[dataptr] {
                        pc = target;
                    }
                }
                LoopExit(target) => {
                    if 0 != memory[dataptr] {
                        pc = target;
                    }
                }
                // ref c => panic!("Unknown instruction {:?} at pc={:}", c, pc),
            }
            pc += 1;
        }
    }
}

fn load_program(fname: String) -> Result<Program> {
    let mut file = File::open(fname)?;
    let mut contents: Vec<u8> = Vec::new();
    file.read_to_end(&mut contents)?;
    let text = String::from_utf8(contents).unwrap();

    Ok(Program::new(text))
}

fn main() {
    match load_program("examples/mandelbrot.bf".to_string()) {
        Ok(ref mut p) => {
            let p = p.optimize();
            //println!("{:?}", p);
            p.interp(&mut std::io::stdin(), &mut std::io::stdout());
        }
        Err(err) => panic!("Cannot read file because {:?}", err),
    }
}
<commit_msg>generalized loop parsing and patching<commit_after>use std::fs::File;
use std::io::{Read, Write, Result};

mod buffer;
#[cfg(test)]
mod parser_tests;
#[cfg(test)]
mod interp_tests;
#[cfg(test)]
mod opt_test;

#[derive(Debug, Clone, Copy, PartialEq)]
enum Instruction {
    MoveLeft(usize),
    MoveRight(usize),
    Inc(u8),
    Dec(u8),
    Input,
    Output,
    LoopEntry(usize),
    LoopExit(usize),
}

use Instruction::*;

#[derive(Debug)]
struct LoopHelper {
    bracket_stack: Vec<usize>,
    relocs: Vec<(usize, usize)>,
}

impl LoopHelper {
    fn new() -> LoopHelper {
        LoopHelper {
            bracket_stack: Vec::new(),
            relocs: Vec::new(),
        }
    }

    fn loop_entry(&mut self, idx: usize) -> Instruction {
        self.bracket_stack.push(idx);
        // value must be patched later
        LoopEntry(std::usize::MAX)
    }

    fn loop_exit(&mut self, idx: usize) -> Instruction {
        if let Some(loop_entry) = self.bracket_stack.pop() {
            self.relocs.push((loop_entry, idx));
            LoopExit(loop_entry)
        } else {
            panic!("Unbalanced {:?} at pc={:}", ']', idx);
        }
    }

    fn relocate(mut self, instructions: &mut Vec<Instruction>) {
        if let Some(unbalanced_idx) = self.bracket_stack.pop() {
            panic!("Unbalanced {:?} at pc={:}", '[', unbalanced_idx);
        }
        for (idx, value) in self.relocs {
            if idx >= instructions.len() || LoopEntry(std::usize::MAX) != instructions[idx] {
                panic!("Unexpected instruction {:?} at pc={:} for reloc",
                       instructions[idx],
                       idx);
            } else {
                instructions[idx] = LoopEntry(value);
            }
        }
    }
}

#[derive(Debug, PartialEq)]
struct Program {
    instructions: Vec<Instruction>,
}

impl Program {
    pub fn new(text: String) -> Program {
        let mut loop_helper = LoopHelper::new();
        // "lex" token stream - aka skip whitespace
        let token = text.chars()
            .filter(move |c| {
                        *c == '>' || *c == '<' || *c == '+' || *c == '-' || *c == '.' ||
                        *c == ',' || *c == '[' || *c == ']'
                    });
        let mut instructions = token
            .enumerate()
            .map(|(idx, c)| match c {
                     '<' => MoveLeft(1),
                     '>' => MoveRight(1),
                     '+' => Inc(1),
                     '-' => Dec(1),
                     '.' => Output,
                     ',' => Input,
                     '[' => loop_helper.loop_entry(idx),
                     ']' => loop_helper.loop_exit(idx),
                     c => panic!("Unknown instruction {:?} at pc={:}", c, idx),
                 })
            .collect::<Vec<_>>();
        loop_helper.relocate(&mut instructions);
        Program { instructions: instructions }
    }

    pub fn optimize(&self) -> Program {
        let mut loop_helper = LoopHelper::new();
        // optimized program
        let mut instructions: Vec<Instruction> = Vec::new();

        for instr in self.instructions.iter() {
            match *instr {
                MoveLeft(offset) => {
                    let last_instr = instructions.last().cloned();
                    if let Some(MoveLeft(old_offset)) = last_instr {
                        let last_idx = instructions.len() - 1;
                        instructions[last_idx] = MoveLeft(old_offset + offset);
                    } else {
                        instructions.push(*instr);
                    }
                }
                MoveRight(offset) => {
                    let last_instr = instructions.last().cloned();
                    if let Some(MoveRight(old_offset)) = last_instr {
                        let last_idx = instructions.len() - 1;
                        instructions[last_idx] = MoveRight(old_offset + offset);
                    } else {
                        instructions.push(*instr);
                    }
                }
                Inc(val) => {
                    let last_instr = instructions.last().cloned();
                    if let Some(Inc(old_val)) = last_instr {
                        let last_idx = instructions.len() - 1;
                        instructions[last_idx] = Inc(old_val.wrapping_add(val));
                    } else {
                        instructions.push(*instr);
                    }
                }
                Dec(val) => {
                    let last_instr = instructions.last().cloned();
                    if let Some(Dec(old_val)) = last_instr {
                        let last_idx = instructions.len() - 1;
                        instructions[last_idx] = Dec(old_val.wrapping_add(val));
                    } else {
                        instructions.push(*instr);
                    }
                }
                LoopEntry(_) => {
                    let idx = instructions.len();
                    instructions.push(loop_helper.loop_entry(idx));
                }
                LoopExit(_) => {
                    let idx = instructions.len();
                    instructions.push(loop_helper.loop_exit(idx));
                }
                instr => instructions.push(instr),
            }
        }
        loop_helper.relocate(&mut instructions);
        Program { instructions: instructions }
    }

    pub fn interp(&self, input: &mut Read, output: &mut Write) {
        let mut memory = vec![0u8; 30000];

        let mut pc: usize = 0;
        let mut dataptr: usize = 0;

        fn get_char(input: &mut Read) -> u8 {
            let mut buf = [0u8; 1];
            match input.read(&mut buf) {
                Err(_) => panic!("Cannot read from stdin"),
                Ok(_) => buf[0],
            }
        }

        while pc < self.instructions.len() {
            match self.instructions[pc] {
                MoveLeft(offset) => dataptr -= offset,
                MoveRight(offset) => dataptr += offset,
                Inc(increment) => memory[dataptr] = memory[dataptr].wrapping_add(increment),
                Dec(decrement) => memory[dataptr] = memory[dataptr].wrapping_sub(decrement),
                Output => write!(output, "{:}", memory[dataptr] as char).expect("Output error"),
                Input => memory[dataptr] = get_char(input),
                LoopEntry(target) => {
                    if 0 == memory[dataptr] {
                        pc = target;
                    }
                }
                LoopExit(target) => {
                    if 0 != memory[dataptr] {
                        pc = target;
                    }
                }
                // ref c => panic!("Unknown instruction {:?} at pc={:}", c, pc),
            }
            pc += 1;
        }
    }
}

fn load_program(fname: String) -> Result<Program> {
    let mut file = File::open(fname)?;
    let mut contents: Vec<u8> = Vec::new();
    file.read_to_end(&mut contents)?;
    let text = String::from_utf8(contents).unwrap();

    Ok(Program::new(text))
}

fn main() {
    match load_program("examples/mandelbrot.bf".to_string()) {
        Ok(ref mut p) => {
            let p = p.optimize();
            //println!("{:?}", p);
            p.interp(&mut std::io::stdin(), &mut std::io::stdout());
        }
        Err(err) => panic!("Cannot read file because {:?}", err),
    }
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate aluminum;

use std::io;
use std::io::prelude::*;
use std::error::Error;
use clap::{App, Arg, AppSettings, SubCommand};
use aluminum::commands;

const VERSION_NUMBER: &'static str = "0.1.0";

fn main() {
    let matches = App::new("Rukyll")
        .setting(AppSettings::ArgRequiredElseHelp)
        .version(VERSION_NUMBER)
        .about("Static site generator")
        .subcommand(SubCommand::with_name("new").arg(Arg::with_name("project name")
                                                         .index(1)
                                                         .required(true)))
        .subcommand(SubCommand::with_name("build"))
        .subcommand(SubCommand::with_name("clean"))
        .get_matches();

    if let ("new", Some(new)) = matches.subcommand() {
        let project_name = if new.is_present("project name") {
            new.value_of("project name").unwrap()
        } else {
            ""
        };

        commands::new_project(project_name).unwrap();
    } else if matches.is_present("build") {
        println!("Building project...");
        match commands::build_project() {
            Ok(_) => {},
            Err(what) => {
                io::stderr()
                    .write(
                        format!("Error: {}", Error::description(&what))
                            .as_bytes()
                    )
                    .unwrap();
            }
        }
    } else if matches.is_present("clean") {
        println!("Cleaning project...");
        match commands::clean_project() {
            Ok(_) => {},
            Err(what) => {
                io::stderr()
                    .write(
                        format!("Error: {}", Error::description(&what))
                            .as_bytes()
                    )
                    .unwrap();
            }
        }
    }
}
<commit_msg>chore: make clean fail silently<commit_after>extern crate clap;
extern crate aluminum;

use std::io;
use std::io::prelude::*;
use std::error::Error;
use clap::{App, Arg, AppSettings, SubCommand};
use aluminum::commands;

const VERSION_NUMBER: &'static str = "0.1.0";

fn main() {
    let matches = App::new("Rukyll")
        .setting(AppSettings::ArgRequiredElseHelp)
        .version(VERSION_NUMBER)
        .about("Static site generator")
        .subcommand(SubCommand::with_name("new").arg(Arg::with_name("project name")
                                                         .index(1)
                                                         .required(true)))
        .subcommand(SubCommand::with_name("build"))
        .subcommand(SubCommand::with_name("clean"))
        .get_matches();

    if let ("new", Some(new)) = matches.subcommand() {
        let project_name = if new.is_present("project name") {
            new.value_of("project name").unwrap()
        } else {
            ""
        };

        commands::new_project(project_name).unwrap();
    } else if matches.is_present("build") {
        println!("Building project...");
        match commands::build_project() {
            Ok(_) => {},
            Err(what) => {
                io::stderr()
                    .write(
                        format!("Error: {}", Error::description(&what))
                            .as_bytes()
                    )
                    .unwrap();
            }
        }
    } else if matches.is_present("clean") {
        println!("Cleaning project...");
        match commands::clean_project() {
            Ok(_) => {},
            Err(_) => {}
        }
    }
}
<|endoftext|>


<commit_before>extern crate libusb;
use std::time::Duration;
use std::thread;
use std::process;
use std::io::{self, Write};

fn main() {
    let context = libusb::Context::new().unwrap();
    let device = find_device(&context);

    match device {
        Ok(handle) => {
            if let Err(e) = run(handle) {
                println!("{}", e);
                process::exit(2);
            }
        },
        Err(libusb::Error::NotFound) => {
            println!("Device not found");
            process::exit(1);
        }
        Err(e) => {
            println!("{}", e);
            process::exit(1);
        }
    }
}

fn find_device<'a>(context: &'a libusb::Context) -> Result<libusb::DeviceHandle<'a>, libusb::Error> {
    let devices = try!(context.devices());
    devices.iter().find(|device| {
        device.device_descriptor().ok().map_or(false, |desc| desc.vendor_id() == 0x05c6)
    }).map_or(Err(libusb::Error::NotFound), |device| device.open())
}

fn run(mut handle: libusb::DeviceHandle) -> Result<(), libusb::Error> {
    let mut stdout = io::stdout();
    println!("Found device");
    try!(handle.claim_interface(0));

    let mut read_buf = [0; 4096];

    loop {
        println!("Connecting DMSS");

        try!(handle.write_bulk(0x01, &[0x0c, 0x00, 0x00, 0x00,  0x00, 0x47, 0xb8, 0x7e], Duration::from_secs(10)));
        let len = try!(handle.read_bulk(0x81, &mut read_buf, Duration::from_secs(10)));
        let rx = &read_buf[0..len];

        if rx == &[0x13, 0x0c, 0x00, 0x00,  0x00, 0x00, 0x72, 0xce, 0x7e] {
            break;
        } else {
            thread::sleep(Duration::from_millis(100));
        }
    }

    try!(handle.write_bulk(0x01, &[0x7d, 0x5d, 0x04, 0x34,  0x21, 0x34, 0x21, 0x00,  0x00, 0x1f, 0x00, 0x00,  0x00, 0xab, 0xdf, 0x7e], Duration::from_secs(1)));
    println!("DMSS is connected");

    loop {
        let len = try!(handle.read_bulk(0x81, &mut read_buf, Duration::from_secs(0)));

        for packet in Hdlc::new(read_buf[0..len].iter().map(|&x| x)) {
            if packet.len() > 0 {
                parse_packet(&mut stdout, packet).unwrap();
            }
        }
    }
}

fn split_byte(s: &[u8], sep: u8) -> (&[u8], &[u8]) {
    if let Some(pos) = s.iter().position(|&i| i == sep) {
        let (a, b) = s.split_at(pos);
        (a, &b[1..])
    } else {
        (s, &[])
    }
}

fn strip_trailing_newlines(mut s: &[u8]) -> &[u8] {
    while s.ends_with(b"\n") {
        s = &s[0..s.len()-1];
    }
    s
}

fn le32(s: &[u8]) -> u32 {
    (s[0] as u32) | (s[1] as u32) << 8 | (s[2] as u32) << 16 | (s[3] as u32) << 24
}

fn le16(s: &[u8]) -> u16 {
    (s[0] as u16) | (s[1] as u16) << 8
}

fn parse_packet(stdout: &mut Write, packet: Vec<u8>) -> io::Result<()> {
    if packet.len() < 24 || packet[0] != 0x79 || packet[2] != 0 {
        print!("[unknown packet: ");
        for b in &packet { print!("{:02x}", b); }
        println!("]");
        return Ok(());
    }

    let (header, rest) = packet.split_at(20);
    let timestamp = le32(&header[6..10]);
    let lineno = le16(&header[12..14]);
    let (msg, rest) = split_byte(rest, 0);
    let (file, _) = split_byte(rest, 0);

    //for b in &packet { print!("{:02x} ", b); }

    try!(write!(stdout, "{:10}: ", timestamp));
    try!(stdout.write_all(strip_trailing_newlines(msg)));
    try!(write!(stdout, " ("));
    try!(stdout.write_all(file));
    try!(write!(stdout, ":{})\n", lineno));

    Ok(())
}

/// Iterator over packets in a HDLC-framed sequence of bytes
struct Hdlc<I> {
    input: I,
    end: bool,
}

impl<I> Hdlc<I> where I:Iterator<Item=u8> {
    fn new(input: I) -> Hdlc<I> {
        Hdlc {
            input: input,
            end: false,
        }
    }

    fn input_byte(&mut self) -> Option<u8> {
        let x = self.input.next();
        self.end |= x.is_none();
        x
    }
}

impl<I> Iterator for Hdlc<I> where I:Iterator<Item=u8> {
    type Item = Vec<u8>;
    fn next(&mut self) -> Option<Vec<u8>> {
        if self.end {
            return None;
        }

        let mut out = vec![];

        loop {
            match self.input_byte() {
                Some(0x7e) | None => break,
                Some(0x7d) => {
                    match self.input_byte() {
                        Some(c) => out.push(c ^ 0x20),
                        None => break,
                    }
                }
                Some(c) => out.push(c),
            };
        }

        Some(out)
    }
}
<commit_msg>main: don't print "unknown packet" warnings<commit_after>extern crate libusb;
use std::time::Duration;
use std::thread;
use std::process;
use std::io::{self, Write};

fn main() {
    let context = libusb::Context::new().unwrap();
    let device = find_device(&context);

    match device {
        Ok(handle) => {
            if let Err(e) = run(handle) {
                println!("{}", e);
                process::exit(2);
            }
        },
        Err(libusb::Error::NotFound) => {
            println!("Device not found");
            process::exit(1);
        }
        Err(e) => {
            println!("{}", e);
            process::exit(1);
        }
    }
}

fn find_device<'a>(context: &'a libusb::Context) -> Result<libusb::DeviceHandle<'a>, libusb::Error> {
    let devices = try!(context.devices());
    devices.iter().find(|device| {
        device.device_descriptor().ok().map_or(false, |desc| desc.vendor_id() == 0x05c6)
    }).map_or(Err(libusb::Error::NotFound), |device| device.open())
}

fn run(mut handle: libusb::DeviceHandle) -> Result<(), libusb::Error> {
    let mut stdout = io::stdout();
    println!("Found device");
    try!(handle.claim_interface(0));

    let mut read_buf = [0; 4096];

    loop {
        println!("Connecting DMSS");

        try!(handle.write_bulk(0x01, &[0x0c, 0x00, 0x00, 0x00,  0x00, 0x47, 0xb8, 0x7e], Duration::from_secs(10)));
        let len = try!(handle.read_bulk(0x81, &mut read_buf, Duration::from_secs(10)));
        let rx = &read_buf[0..len];

        if rx == &[0x13, 0x0c, 0x00, 0x00,  0x00, 0x00, 0x72, 0xce, 0x7e] {
            break;
        } else {
            thread::sleep(Duration::from_millis(100));
        }
    }

    try!(handle.write_bulk(0x01, &[0x7d, 0x5d, 0x04, 0x34,  0x21, 0x34, 0x21, 0x00,  0x00, 0x1f, 0x00, 0x00,  0x00, 0xab, 0xdf, 0x7e], Duration::from_secs(1)));
    println!("DMSS is connected");

    loop {
        let len = try!(handle.read_bulk(0x81, &mut read_buf, Duration::from_secs(0)));

        for packet in Hdlc::new(read_buf[0..len].iter().map(|&x| x)) {
            if packet.len() > 0 {
                parse_packet(&mut stdout, packet).unwrap();
            }
        }
    }
}

fn split_byte(s: &[u8], sep: u8) -> (&[u8], &[u8]) {
    if let Some(pos) = s.iter().position(|&i| i == sep) {
        let (a, b) = s.split_at(pos);
        (a, &b[1..])
    } else {
        (s, &[])
    }
}

fn strip_trailing_newlines(mut s: &[u8]) -> &[u8] {
    while s.ends_with(b"\n") {
        s = &s[0..s.len()-1];
    }
    s
}

fn le32(s: &[u8]) -> u32 {
    (s[0] as u32) | (s[1] as u32) << 8 | (s[2] as u32) << 16 | (s[3] as u32) << 24
}

fn le16(s: &[u8]) -> u16 {
    (s[0] as u16) | (s[1] as u16) << 8
}

fn parse_packet(stdout: &mut Write, packet: Vec<u8>) -> io::Result<()> {
    if packet.len() < 24 || packet[0] != 0x79 || packet[2] != 0 {
        //print!("[unknown packet: ");
        //for b in &packet { print!("{:02x}", b); }
        //println!("]");
        return Ok(());
    }

    let (header, rest) = packet.split_at(20);
    let timestamp = le32(&header[6..10]);
    let lineno = le16(&header[12..14]);
    let (msg, rest) = split_byte(rest, 0);
    let (file, _) = split_byte(rest, 0);

    //for b in &packet { print!("{:02x} ", b); }

    try!(write!(stdout, "{:10}: ", timestamp));
    try!(stdout.write_all(strip_trailing_newlines(msg)));
    try!(write!(stdout, " ("));
    try!(stdout.write_all(file));
    try!(write!(stdout, ":{})\n", lineno));

    Ok(())
}

/// Iterator over packets in a HDLC-framed sequence of bytes
struct Hdlc<I> {
    input: I,
    end: bool,
}

impl<I> Hdlc<I> where I:Iterator<Item=u8> {
    fn new(input: I) -> Hdlc<I> {
        Hdlc {
            input: input,
            end: false,
        }
    }

    fn input_byte(&mut self) -> Option<u8> {
        let x = self.input.next();
        self.end |= x.is_none();
        x
    }
}

impl<I> Iterator for Hdlc<I> where I:Iterator<Item=u8> {
    type Item = Vec<u8>;
    fn next(&mut self) -> Option<Vec<u8>> {
        if self.end {
            return None;
        }

        let mut out = vec![];

        loop {
            match self.input_byte() {
                Some(0x7e) | None => break,
                Some(0x7d) => {
                    match self.input_byte() {
                        Some(c) => out.push(c ^ 0x20),
                        None => break,
                    }
                }
                Some(c) => out.push(c),
            };
        }

        Some(out)
    }
}
<|endoftext|>


<commit_before>extern crate image;
extern crate toml;
extern crate rustc_serialize;

use std::env;
use std::thread;
use std::path::Path;
use image::ImageBuffer;

mod color;
mod math;
mod mesh;
mod primitive;
mod ray_tracer;
mod scene;

use ray_tracer::RayTracer;
use scene::Scene;

fn main() {
    let args: Vec<_> = env::args().collect();
    if args.len() < 2 {
        println!("Usage: command <scene.toml>");
        return;
    }

    let path = Path::new(&args[1]);
    let scene = Scene::read(path).prepare();

    let threads = scene.threads;
    let handles: Vec<_> = (0..threads).map(|t| {
        let scene = scene.clone();
        thread::spawn(move || {
            let ray_tracer = RayTracer::new(&scene);
            let extra = if t == threads - 1 { scene.height % threads } else { 0 };
            let (h, w) = (scene.height / threads + extra, scene.width);
            let mut imgbuf = ImageBuffer::new(w, h);
            for y in (0..h) {
                for x in (0..w) {
                    let yy = t * h + y;
                    imgbuf.put_pixel(x, y, ray_tracer.trace_pixel(x, yy).rgb());
                }
            }
            imgbuf
        })
    }).collect();

    let chunk_bufs: Vec<_> = handles.into_iter().map(|h| { h.join().unwrap() }).collect();

    let mut imgbuf = ImageBuffer::new(scene.width, scene.height);
    let mut y_off = 0;
    for buf in chunk_bufs {
        for (x, y, pixel) in buf.enumerate_pixels() {
            imgbuf.put_pixel(x, y_off + y, *pixel);
        }
        y_off += buf.height();
    }

    let fout = Path::new(&scene.image);
    let _ = imgbuf.save(fout);
    println!("Wrote file {:?}", fout);
}
<commit_msg>Use Arc to share scene<commit_after>extern crate image;
extern crate toml;
extern crate rustc_serialize;

use std::env;
use std::path::Path;
use std::sync::Arc;
use std::thread;
use image::ImageBuffer;

mod color;
mod math;
mod mesh;
mod primitive;
mod ray_tracer;
mod scene;

use ray_tracer::RayTracer;
use scene::Scene;

fn main() {
    let args: Vec<_> = env::args().collect();
    if args.len() < 2 {
        println!("Usage: command <scene.toml>");
        return;
    }

    let path = Path::new(&args[1]);
    let scene = Scene::read(path).prepare();
    let scene = Arc::new(scene);

    let threads = scene.threads;
    let handles: Vec<_> = (0..threads).map(|t| {
        let scene = scene.clone();
        thread::spawn(move || {
            let ray_tracer = RayTracer::new(&scene);
            let extra = if t == threads - 1 { scene.height % threads } else { 0 };
            let (h, w) = (scene.height / threads + extra, scene.width);
            let mut imgbuf = ImageBuffer::new(w, h);
            for y in (0..h) {
                for x in (0..w) {
                    let yy = t * h + y;
                    imgbuf.put_pixel(x, y, ray_tracer.trace_pixel(x, yy).rgb());
                }
            }
            imgbuf
        })
    }).collect();

    let chunk_bufs: Vec<_> = handles.into_iter().map(|h| { h.join().unwrap() }).collect();

    let mut imgbuf = ImageBuffer::new(scene.width, scene.height);
    let mut y_off = 0;
    for buf in chunk_bufs {
        for (x, y, pixel) in buf.enumerate_pixels() {
            imgbuf.put_pixel(x, y_off + y, *pixel);
        }
        y_off += buf.height();
    }

    let fout = Path::new(&scene.image);
    let _ = imgbuf.save(fout);
    println!("Wrote file {:?}", fout);
}
<|endoftext|>


<commit_before>extern crate rlisp;

use std::error::Error;
use rlisp::parser::Parser;
use rlisp::scope::GlobalScope;

fn main()
{
    let mut stdin = std::io::stdin();
    let mut env = GlobalScope::new();
    env.load_stdlib();

    loop
    {
        let mut text = String::new();
        match stdin.read_line(&mut text) {
            Ok(0) => break,
            Ok(_) => {},
            Err(e) => panic!("Error reading input: {}", e.description()),
        }

        match Parser::new(&text).parse() {
            Ok(vs) => for val in vs
            {
                match val.eval(&mut env) {
                    Ok(v) => println!("Result: {}", v),
                    Err(e) => println!("Error: {}", e),
                }
            },
            Err(e) => println!("Error: {}", e),
        }
    }
}
<commit_msg>simplify stuff<commit_after>extern crate rlisp;

use rlisp::parser::Parser;
use rlisp::scope::GlobalScope;

fn main()
{
    let mut stdin = std::io::stdin();
    let mut env = GlobalScope::new();
    env.load_stdlib();

    loop
    {
        let mut text = String::new();
        match stdin.read_line(&mut text) {
            Ok(0) => break,
            Ok(_) => {},
            Err(e) => panic!("Error reading input: {}", e),
        }

        match Parser::new(&text).parse() {
            Ok(vs) => for val in vs
            {
                match val.eval(&mut env) {
                    Ok(v) => println!("Result: {}", v),
                    Err(e) => println!("Error: {}", e),
                }
            },
            Err(e) => println!("Error: {}", e),
        }
    }
}
<|endoftext|>


<commit_before>extern crate byteorder;
extern crate crypto;
extern crate futures;
extern crate hyper;
extern crate rayon;
extern crate rustc_serialize;
extern crate tokio_core;
extern crate tokio_io;
extern crate url;

use byteorder::{ByteOrder, LittleEndian};
use crypto::digest::Digest;
use crypto::sha2::Sha256;
use futures::{Future, future, Stream};
use hyper::StatusCode;
use hyper::server::{Service, Request, Response, Http};
use rayon::current_num_threads;
use rayon::prelude::*;
use rustc_serialize::hex::{FromHex, ToHex};
use std::fmt::Write;
use std::io;
use std::sync::{Arc, atomic};
use tokio_core::net::TcpListener;
use tokio_core::reactor::Core;
use tokio_io::{AsyncRead, AsyncWrite};
use url::Url;

fn detect_hup<I: AsyncRead + AsyncWrite>(io: I) -> (DetectHUP<I>, futures::sync::oneshot::Receiver<()>) {
    let (sender, receiver) = futures::sync::oneshot::channel();
    (DetectHUP { io: io, sender: Some(sender) }, receiver)
}
struct DetectHUP<I> {
    io: I,
    sender: Option<futures::sync::oneshot::Sender<()>>
}
impl<I> DetectHUP<I> {
    fn abortlogic<T: std::fmt::Debug>(&mut self, x: Result<T, io::Error>) -> Result<T, io::Error> {
        println!("{:?}", x);
        if let Err(ref e) = x {
            println!("e.kind {:?}", e.kind());
            if e.kind() == io::ErrorKind::ConnectionReset {
                let _ = self.sender.take().map(|sender| sender.send(()));
            }
        }
        x
    }
}
// TODO: more delegation (all Read+Write+AsyncRead+AsyncWrite methods) for efficiency
impl<I: AsyncRead> io::Read for DetectHUP<I> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let tmp = self.io.read(buf);
        self.abortlogic(tmp)
    }
}
impl<I: AsyncRead> AsyncRead for DetectHUP<I> {}

impl<I: AsyncWrite> io::Write for DetectHUP<I> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let tmp = self.io.write(buf);
        self.abortlogic(tmp)
    }
    fn flush(&mut self) -> io::Result<()> {
        self.io.flush()
    }
}
impl<I: AsyncWrite> AsyncWrite for DetectHUP<I> {
    fn shutdown(&mut self) -> futures::Poll<(), io::Error> {
        self.io.shutdown()
    }
}

fn proofofwork(mask: &[u8], goal: &[u8], done: Arc<atomic::AtomicBool>) -> (String, String) {
    assert_eq!(mask.len(), goal.len());
    assert_eq!(mask.len(), 32);
    let f = |x: u64| {
        if done.load(atomic::Ordering::Relaxed) {
            return None;
        }
        let mut tmpinput = [0; 8];
        let mut tmpoutput = [0; 32];
        LittleEndian::write_u64(&mut tmpinput, x);
        // TODO: efficient alphanumeric hashes
        //if !tmpinput.iter().all(|&c| (c as char).is_alphanumeric()) { return None; }
        let mut hasher = Sha256::new();
        hasher.input(&tmpinput);
        hasher.result(&mut tmpoutput);
        if mask.iter().zip(goal).zip(tmpoutput.iter()).all(|((&m, &g), o)| (m & o) == (m & g)) {
            //return Some((std::str::from_utf8(&tmpinput).unwrap().into(), hasher.result_str()));
            return Some((tmpinput.to_hex(), hasher.result_str()));
        }
        None
    };
    (0..2u64.pow(63)).into_par_iter().filter_map(|x| f(x)).find_any(|_| true).unwrap()
}

const HELP_MSG: &'static str = "Usage examples:
$ time curl localhost:3000/sha256?mask=$(python -c 'print \"00\"*29+\"ff\"*3')\\&goal=$(python -c 'print \"00\"*29+\"deadbe\"')
159a360000000000 has hash 9d0ec7b3dd909e1d7ee64186e13b8065c32e88695c8aa938bdbb9a9c6ddeadbe

real    0m3.049s
user    0m0.040s
sys     0m0.008s
$ python -c 'import sys; sys.stdout.write(\"159a360000000000\".decode(\"hex\"))' | sha256sum
9d0ec7b3dd909e1d7ee64186e13b8065c32e88695c8aa938bdbb9a9c6ddeadbe  -
$ python
>>> import requests
>>> requests.get('http://localhost:3000/sha256', params={'mask': '00'+'ff'*3+'00'*28, 'goal': '00badc0d'+'00'*28}).text
u'78a3170000000000 has hash fcbadc0d5856bb6eea467a236218eb5e16017a1636e335e2946618feb0aae620\\n'
>>> requests.get('http://localhost:3000/sha256', params={'mask': 'ff'*4+'00'*28, 'goal': '00abcdef'+'00'*28}).text
u'c72b530200000040 has hash 00abcdef83801fd557e1740187560ac4fdc557645e175f5faeb407e54a2d9958\\n'

Intended general usage (more algos will be added later):
GET /sha256?mask=<some 32 byte hex encoded mask>&goal=<some 32 byte hex encoded goal>";

struct POWService(Arc<atomic::AtomicBool>);
impl Service for POWService {
    type Request = Request;
    type Response = Response;
    type Error = hyper::Error;
    type Future = Box<Future<Item=Response, Error=hyper::Error>>;

    fn call(&self, req: Request) -> Self::Future {
        match req.method() {
            &hyper::Get => (),
            _ => return Box::new(future::ok(Response::new().with_status(StatusCode::NotFound))),
        }
        match req.path() {
            "/sha256" => { powserver(req, self.0.clone()) }
            _ => Box::new(future::ok(Response::new().with_body(HELP_MSG.as_bytes())))
        }
    }
}

fn fmt_error_to_hyper_error(e: std::fmt::Error) -> hyper::Error {
    let tmp: std::io::Error = std::io::Error::new(std::io::ErrorKind::Other, e);
    tmp.into()
}

fn powserver(req: Request, done: Arc<atomic::AtomicBool>) -> Box<Future<Item=Response, Error=hyper::Error>> {
    let base_url = Url::parse("http://foo").unwrap();
    println!("{:?}", req.uri());
    println!("{:?}", req.headers());
    if let Ok(url) = Url::options().base_url(Some(&base_url)).parse(&req.uri().as_ref()) {
        let mut mask = None;
        let mut goal = None;
        for (k, v) in url.query_pairs() {
            if k == "mask" && v.len() == 32*2 {
                mask = v.from_hex().ok();
            }
            if k == "goal" && v.len() == 32*2 {
                goal = v.from_hex().ok();
            }
        }
        println!("mask: {:?}\ngoal: {:?}", mask.clone().map(|x| x.to_hex()), goal.clone().map(|x| x.to_hex()));
        if let (Some(mask), Some(goal)) = (mask, goal) {
            let mut resp = String::new();
            let (x, hash) = proofofwork(&mask, &goal, done.clone());
            println!("sending preimage {}", x);
            if let Err(e) = writeln!(resp, "{} has hash {}", x, hash) {
                return Box::new(future::err(fmt_error_to_hyper_error(e)));
            }
            return Box::new(future::ok(Response::new().with_body(resp)));
        }
    }
    Box::new(future::ok(Response::new().with_body(HELP_MSG.as_bytes())))
}

fn main() {
    if let Some(Ok(port)) = std::env::args().nth(1).map(|x| x.parse::<u16>()) {
        println!("current_num_threads: {}", current_num_threads());
        let bindaddr = ("0.0.0.0".parse::<std::net::IpAddr>().unwrap(), port);
        /*Http::new().bind(&bindaddr.into(), || Ok(POWService)).expect("Failed to bind server")
            .run().expect("Fatal error while running the server");*/
        let http = Http::new();
        let mut core = Core::new().unwrap();
        let handle = core.handle();
        let listener = TcpListener::bind(&bindaddr.into(), &handle).unwrap();
        core.run(listener.incoming().for_each(|(sock, addr)| {
            let (sock, hup) = detect_hup(sock);
            let done = Arc::new(atomic::AtomicBool::new(false));
            http.bind_connection(&handle, sock, addr, POWService(done.clone()));
            handle.spawn(hup.and_then(move |()| { done.store(true, atomic::Ordering::Relaxed); Ok(()) }).map_err(|_| ()));
            Ok(())
        })).unwrap();
    } else {
        println!("Expecting the port as the first argument.");
    }
}

#[test]
fn test_pow() {
    println!("current_num_threads: {}", current_num_threads());
    println!("POW: {:?}", proofofwork(b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff\xff\xff",
                                      b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x1a\x85\xbd"));
}
<commit_msg>Make proofofwork spawn a future on Rayon's eventloop.<commit_after>extern crate byteorder;
extern crate crypto;
extern crate futures;
extern crate hyper;
extern crate rayon;
extern crate rustc_serialize;
extern crate tokio_core;
extern crate tokio_io;
extern crate url;

use byteorder::{ByteOrder, LittleEndian};
use crypto::digest::Digest;
use crypto::sha2::Sha256;
use futures::{Future, future, Stream};
use hyper::StatusCode;
use hyper::server::{Service, Request, Response, Http};
use rayon::{current_num_threads, RayonFuture};
use rayon::prelude::*;
use rustc_serialize::hex::{FromHex, ToHex};
use std::fmt::Write;
use std::io;
use std::sync::{Arc, atomic};
use tokio_core::net::TcpListener;
use tokio_core::reactor::Core;
use tokio_io::{AsyncRead, AsyncWrite};
use url::Url;

fn detect_hup<I: AsyncRead + AsyncWrite>(io: I) -> (DetectHUP<I>, futures::sync::oneshot::Receiver<()>) {
    let (sender, receiver) = futures::sync::oneshot::channel();
    (DetectHUP { io: io, sender: Some(sender) }, receiver)
}
struct DetectHUP<I> {
    io: I,
    sender: Option<futures::sync::oneshot::Sender<()>>
}
impl<I> DetectHUP<I> {
    fn abortlogic<T: std::fmt::Debug>(&mut self, x: Result<T, io::Error>) -> Result<T, io::Error> {
        println!("{:?}", x);
        if let Err(ref e) = x {
            println!("e.kind {:?}", e.kind());
            if e.kind() == io::ErrorKind::ConnectionReset {
                let _ = self.sender.take().map(|sender| sender.send(()));
            }
        }
        x
    }
}
// TODO: more delegation (all Read+Write+AsyncRead+AsyncWrite methods) for efficiency
impl<I: AsyncRead> io::Read for DetectHUP<I> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let tmp = self.io.read(buf);
        self.abortlogic(tmp)
    }
}
impl<I: AsyncRead> AsyncRead for DetectHUP<I> {}

impl<I: AsyncWrite> io::Write for DetectHUP<I> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let tmp = self.io.write(buf);
        self.abortlogic(tmp)
    }
    fn flush(&mut self) -> io::Result<()> {
        self.io.flush()
    }
}
impl<I: AsyncWrite> AsyncWrite for DetectHUP<I> {
    fn shutdown(&mut self) -> futures::Poll<(), io::Error> {
        self.io.shutdown()
    }
}

fn proofofwork(mask: Vec<u8>, goal: Vec<u8>, done: Arc<atomic::AtomicBool>) -> RayonFuture<(String, String), hyper::Error> {
    assert_eq!(mask.len(), goal.len());
    assert_eq!(mask.len(), 32);
    rayon::spawn_future_async(future::lazy(move || {
        let mask = mask;
        let mask = &mask[..];
        let goal = goal;
        let goal = &goal[..];
        let f = |x: u64| {
            if done.load(atomic::Ordering::Relaxed) {
                return None;
            }
            let mut tmpinput = [0; 8];
            let mut tmpoutput = [0; 32];
            LittleEndian::write_u64(&mut tmpinput, x);
            // TODO: efficient alphanumeric hashes
            //if !tmpinput.iter().all(|&c| (c as char).is_alphanumeric()) { return None; }
            let mut hasher = Sha256::new();
            hasher.input(&tmpinput);
            hasher.result(&mut tmpoutput);
            if mask.iter().zip(goal).zip(tmpoutput.iter()).all(|((m, g), o)| (m & o) == (m & g)) {
                //return Some((std::str::from_utf8(&tmpinput).unwrap().into(), hasher.result_str()));
                return Some((tmpinput.to_hex(), hasher.result_str()));
            }
            None
        };
        Ok((0..2u64.pow(63)).into_par_iter().filter_map(f).find_any(|_| true).unwrap())
    }))
}

const HELP_MSG: &'static str = "Usage examples:
$ time curl localhost:3000/sha256?mask=$(python -c 'print \"00\"*29+\"ff\"*3')\\&goal=$(python -c 'print \"00\"*29+\"deadbe\"')
159a360000000000 has hash 9d0ec7b3dd909e1d7ee64186e13b8065c32e88695c8aa938bdbb9a9c6ddeadbe

real    0m3.049s
user    0m0.040s
sys     0m0.008s
$ python -c 'import sys; sys.stdout.write(\"159a360000000000\".decode(\"hex\"))' | sha256sum
9d0ec7b3dd909e1d7ee64186e13b8065c32e88695c8aa938bdbb9a9c6ddeadbe  -
$ python
>>> import requests
>>> requests.get('http://localhost:3000/sha256', params={'mask': '00'+'ff'*3+'00'*28, 'goal': '00badc0d'+'00'*28}).text
u'78a3170000000000 has hash fcbadc0d5856bb6eea467a236218eb5e16017a1636e335e2946618feb0aae620\\n'
>>> requests.get('http://localhost:3000/sha256', params={'mask': 'ff'*4+'00'*28, 'goal': '00abcdef'+'00'*28}).text
u'c72b530200000040 has hash 00abcdef83801fd557e1740187560ac4fdc557645e175f5faeb407e54a2d9958\\n'

Intended general usage (more algos will be added later):
GET /sha256?mask=<some 32 byte hex encoded mask>&goal=<some 32 byte hex encoded goal>";

struct POWService(Arc<atomic::AtomicBool>);
impl Service for POWService {
    type Request = Request;
    type Response = Response;
    type Error = hyper::Error;
    type Future = Box<Future<Item=Response, Error=hyper::Error>>;

    fn call(&self, req: Request) -> Self::Future {
        match req.method() {
            &hyper::Get => (),
            _ => return Box::new(future::ok(Response::new().with_status(StatusCode::NotFound))),
        }
        match req.path() {
            "/sha256" => { powserver(req, self.0.clone()) }
            _ => Box::new(future::ok(Response::new().with_body(HELP_MSG.as_bytes())))
        }
    }
}

fn fmt_error_to_hyper_error(e: std::fmt::Error) -> hyper::Error {
    let tmp: std::io::Error = std::io::Error::new(std::io::ErrorKind::Other, e);
    tmp.into()
}

fn powserver(req: Request, done: Arc<atomic::AtomicBool>) -> Box<Future<Item=Response, Error=hyper::Error>> {
    let base_url = Url::parse("http://foo").unwrap();
    println!("{:?}", req.uri());
    println!("{:?}", req.headers());
    if let Ok(url) = Url::options().base_url(Some(&base_url)).parse(&req.uri().as_ref()) {
        let mut mask = None;
        let mut goal = None;
        for (k, v) in url.query_pairs() {
            if k == "mask" && v.len() == 32*2 {
                mask = v.from_hex().ok();
            }
            if k == "goal" && v.len() == 32*2 {
                goal = v.from_hex().ok();
            }
        }
        println!("mask: {:?}\ngoal: {:?}", mask.clone().map(|x| x.to_hex()), goal.clone().map(|x| x.to_hex()));
        if let (Some(mask), Some(goal)) = (mask, goal) {
            let mut resp = String::new();
            return Box::new(proofofwork(mask, goal, done.clone()).and_then(|(x, hash)| {
                println!("sending preimage {}", x);
                if let Err(e) = writeln!(resp, "{} has hash {}", x, hash) {
                    return Err(fmt_error_to_hyper_error(e));
                }
                Ok(Response::new().with_body(resp))
            }));
        }
    }
    Box::new(future::ok(Response::new().with_body(HELP_MSG.as_bytes())))
}

fn main() {
    if let Some(Ok(port)) = std::env::args().nth(1).map(|x| x.parse::<u16>()) {
        println!("current_num_threads: {}", current_num_threads());
        let bindaddr = ("0.0.0.0".parse::<std::net::IpAddr>().unwrap(), port);
        /*Http::new().bind(&bindaddr.into(), || Ok(POWService)).expect("Failed to bind server")
            .run().expect("Fatal error while running the server");*/
        let http = Http::new();
        let mut core = Core::new().unwrap();
        let handle = core.handle();
        let listener = TcpListener::bind(&bindaddr.into(), &handle).unwrap();
        core.run(listener.incoming().for_each(|(sock, addr)| {
            let (sock, hup) = detect_hup(sock);
            let done = Arc::new(atomic::AtomicBool::new(false));
            http.bind_connection(&handle, sock, addr, POWService(done.clone()));
            handle.spawn(hup.and_then(move |()| { done.store(true, atomic::Ordering::Relaxed); Ok(()) }).map_err(|_| ()));
            Ok(())
        })).unwrap();
    } else {
        println!("Expecting the port as the first argument.");
    }
}

#[test]
fn test_pow() {
    println!("current_num_threads: {}", current_num_threads());
    println!("POW: {:?}", proofofwork(b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff\xff\xff",
                                      b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x1a\x85\xbd"));
}
<|endoftext|>