<commit_before>#[macro_use]
extern crate clap;

use petname::Petnames;

use std::collections::HashSet;
use std::fmt;
use std::fs;
use std::io;
use std::path;
use std::process;
use std::str::FromStr;

use clap::Arg;
use rand::seq::IteratorRandom;

fn main() {
    let matches = app().get_matches();
    match run(matches) {
        Err(Error::Disconnected) => {
            process::exit(0);
        }
        Err(e) => {
            eprintln!("Error: {}", e);
            process::exit(1);
        }
        Ok(()) => {
            process::exit(0);
        }
    }
}

enum Error {
    Io(io::Error),
    FileIo(path::PathBuf, io::Error),
    Cardinality(String),
    Alliteration(String),
    Disconnected,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Error::Io(ref e) => write!(f, "{}", e),
            Error::FileIo(ref path, ref e) => write!(f, "{}: {}", e, path.display()),
            Error::Cardinality(ref message) => write!(f, "cardinality is zero: {}", message),
            Error::Alliteration(ref message) => write!(f, "cannot alliterate: {}", message),
            Error::Disconnected => write!(f, "caller disconnected / stopped reading"),
        }
    }
}

impl From<io::Error> for Error {
    fn from(error: io::Error) -> Self {
        Error::Io(error)
    }
}

fn app<'a, 'b>() -> clap::App<'a, 'b> {
    clap::App::new("rust-petname")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Generate human readable random names.")
        .after_help(concat!(
            "Based on Dustin Kirkland's petname project ",
            "<https://github.com/dustinkirkland/petname>."
        ))
        .arg(
            Arg::with_name("words")
                .short("w")
                .long("words")
                .value_name("WORDS")
                .default_value("2")
                .help("Number of words in name")
                .takes_value(true)
                .validator(can_be_parsed::<u8>),
        )
        .arg(
            Arg::with_name("separator")
                .short("s")
                .long("separator")
                .value_name("SEP")
                .default_value("-")
                .help("Separator between words")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("complexity")
                .short("c")
                .long("complexity")
                .value_name("COM")
                .possible_values(&["0", "1", "2"])
                .hide_possible_values(true)
                .default_value("0")
                .help("Use small words (0), medium words (1), or large words (2)")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("directory")
                .short("d")
                .long("dir")
                .value_name("DIR")
                .help("Directory containing adjectives.txt, adverbs.txt, names.txt")
                .conflicts_with("complexity")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("count")
                .long("count")
                .value_name("COUNT")
                .default_value("1")
                .help(concat!(
                    "Generate multiple names; pass 0 to produce infinite ",
                    "names (--count=0 is deprecated; use --stream instead)"
                ))
                .takes_value(true)
                .validator(can_be_parsed::<usize>),
        )
        .arg(
            Arg::with_name("stream")
                .long("stream")
                .help("Stream names continuously")
                .conflicts_with("count")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("non-repeating")
                .long("non-repeating")
                .help("Do not generate the same name more than once")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("letters")
                .short("l")
                .long("letters")
                .value_name("LETTERS")
                .default_value("0")
                .help("Maxiumum number of letters in each word; 0 for unlimited")
                .takes_value(true)
                .validator(can_be_parsed::<usize>),
        )
        .arg(
            Arg::with_name("alliterate")
                .short("a")
                .long("alliterate")
                .help("Generate names where each word begins with the same letter")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("alliterate-with")
                .long("alliterate-with")
                .short("A")
                .value_name("LETTER")
                .help("Generate names where each word begins with the given letter")
                .takes_value(true)
                .validator(can_be_parsed::<char>),
        )
        .arg(
            // For compatibility with upstream.
            Arg::with_name("ubuntu")
                .short("u")
                .long("ubuntu")
                .help("Alias; see --alliterate")
                .takes_value(false),
        )
}

fn run(matches: clap::ArgMatches) -> Result<(), Error> {
    // Unwrapping is safe because these options have defaults.
    let opt_separator = matches.value_of("separator").unwrap();
    let opt_words = matches.value_of("words").unwrap();
    let opt_complexity = matches.value_of("complexity").unwrap();
    let opt_count = matches.value_of("count").unwrap();
    let opt_letters = matches.value_of("letters").unwrap();

    // Flags.
    let opt_stream = matches.is_present("stream");
    let opt_non_repeating = matches.is_present("non-repeating");
    let opt_alliterate = matches.is_present("alliterate")
        || matches.is_present("ubuntu")
        || matches.is_present("alliterate-with");

    // Optional arguments without defaults.
    let opt_directory = matches.value_of("directory");
    let opt_alliterate_char = matches
        .value_of("alliterate-with")
        .and_then(|s| s.parse::<char>().ok());

    // Parse numbers. Validated so unwrapping is okay.
    let opt_words: u8 = opt_words.parse().unwrap();
    let opt_count: usize = opt_count.parse().unwrap();
    let opt_letters: usize = opt_letters.parse().unwrap();

    // Load custom word lists, if specified.
    let words = match opt_directory {
        Some(dirname) => Words::load(dirname)?,
        None => Words::Builtin,
    };

    // Select the appropriate word list.
    let mut petnames = match words {
        Words::Custom(ref adjectives, ref adverbs, ref names) => {
            Petnames::init(adjectives, adverbs, names)
        }
        Words::Builtin => match opt_complexity {
            "0" => Petnames::small(),
            "1" => Petnames::medium(),
            "2" => Petnames::large(),
            _ => Petnames::small(),
        },
    };

    // If requested, limit the number of letters.
    if opt_letters != 0 {
        petnames.retain(|s| s.len() <= opt_letters);
    }

    // Check cardinality.
    if petnames.cardinality(opt_words) == 0 {
        return Err(Error::Cardinality(
            "no petnames to choose from; try relaxing constraints".to_string(),
        ));
    }

    // We're going to need a source of randomness.
    let mut rng = rand::thread_rng();

    // Handle alliteration, either by eliminating a specified
    // character, or using a random one.
    if opt_alliterate {
        // We choose the first letter from the intersection of the
        // first letters of each word list in `petnames`.
        let firsts =
            common_first_letters(&petnames.adjectives, &[&petnames.adverbs, &petnames.names]);
        // if a specific character was requested for alliteration,
        // attempt to use it.
        if let Some(c) = opt_alliterate_char {
            if firsts.contains(&c) {
                petnames.retain(|s| s.starts_with(c));
            } else {
                return Err(Error::Alliteration(
                    "no petnames begin with the choosen alliteration character".to_string(),
                ));
            }
        } else {
            // Otherwise choose the first letter at random; fails if
            // there are no letters.
            match firsts.iter().choose(&mut rng) {
                Some(c) => petnames.retain(|s| s.starts_with(*c)),
                None => {
                    return Err(Error::Alliteration(
                        "word lists have no initial letters in common".to_string(),
                    ))
                }
            };
        }
    }

    // Manage stdout.
    let stdout = io::stdout();
    let mut writer = io::BufWriter::new(stdout.lock());

    // Warn that --count=0 is deprecated.
    if opt_count == 0 {
        eprintln!(concat!(
            "Warning: specifying --count=0 to continuously produce petnames is ",
            "deprecated and its behaviour will change in a future version; ",
            "specify --stream instead.",
        ));
    }

    // Stream if count is 0. TODO: Only stream when --stream is specified.
    let count = if opt_stream || opt_count == 0 {
        None
    } else {
        Some(opt_count)
    };

    // Get an iterator for the names we want to print out.
    if opt_non_repeating {
        printer(
            &mut writer,
            petnames.iter_non_repeating(&mut rng, opt_words, &opt_separator),
            count,
        )
    } else {
        printer(
            &mut writer,
            petnames.iter(&mut rng, opt_words, &opt_separator),
            count,
        )
    }
}

fn printer<'a, OUT, NAMES>(
    writer: &mut OUT,
    names: NAMES,
    count: Option<usize>,
) -> Result<(), Error>
where
    OUT: io::Write,
    NAMES: Iterator<Item = String>,
{
    match count {
        None => {
            for name in names {
                writeln!(writer, "{}", name).map_err(suppress_disconnect)?;
            }
        }
        Some(n) => {
            for name in names.take(n) {
                writeln!(writer, "{}", name)?;
            }
        }
    }

    Ok(())
}

fn can_be_parsed<INTO>(value: String) -> Result<(), String>
where
    INTO: FromStr,
    <INTO as FromStr>::Err: std::fmt::Display,
{
    match value.parse::<INTO>() {
        Err(e) => Err(format!("{}", e)),
        Ok(_) => Ok(()),
    }
}

fn common_first_letters(init: &[&str], more: &[&[&str]]) -> HashSet<char> {
    let mut firsts = first_letters(init);
    let firsts_other: Vec<HashSet<char>> = more.iter().map(|list| first_letters(list)).collect();
    firsts.retain(|c| firsts_other.iter().all(|fs| fs.contains(c)));
    firsts
}

fn first_letters(names: &[&str]) -> HashSet<char> {
    names.iter().filter_map(|s| s.chars().next()).collect()
}

enum Words {
    Custom(String, String, String),
    Builtin,
}

impl Words {
    // Load word lists from the given directory. This function expects to find three
    // files in that directory: `adjectives.txt`, `adverbs.txt`, and `names.txt`.
    // Each should be valid UTF-8, and contain words separated by whitespace.
    fn load<T: AsRef<path::Path>>(dirname: T) -> Result<Self, Error> {
        let dirname = dirname.as_ref();
        Ok(Self::Custom(
            read_file_to_string(dirname.join("adjectives.txt"))?,
            read_file_to_string(dirname.join("adverbs.txt"))?,
            read_file_to_string(dirname.join("names.txt"))?,
        ))
    }
}

fn read_file_to_string<P: AsRef<path::Path>>(path: P) -> Result<String, Error> {
    fs::read_to_string(&path).map_err(|error| Error::FileIo(path.as_ref().to_path_buf(), error))
}

fn suppress_disconnect(err: io::Error) -> Error {
    match err.kind() {
        io::ErrorKind::BrokenPipe => Error::Disconnected,
        _ => err.into(),
    }
}
<commit_msg>Fix lint<commit_after>#[macro_use]
extern crate clap;

use petname::Petnames;

use std::collections::HashSet;
use std::fmt;
use std::fs;
use std::io;
use std::path;
use std::process;
use std::str::FromStr;

use clap::Arg;
use rand::seq::IteratorRandom;

fn main() {
    let matches = app().get_matches();
    match run(matches) {
        Err(Error::Disconnected) => {
            process::exit(0);
        }
        Err(e) => {
            eprintln!("Error: {}", e);
            process::exit(1);
        }
        Ok(()) => {
            process::exit(0);
        }
    }
}

enum Error {
    Io(io::Error),
    FileIo(path::PathBuf, io::Error),
    Cardinality(String),
    Alliteration(String),
    Disconnected,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Error::Io(ref e) => write!(f, "{}", e),
            Error::FileIo(ref path, ref e) => write!(f, "{}: {}", e, path.display()),
            Error::Cardinality(ref message) => write!(f, "cardinality is zero: {}", message),
            Error::Alliteration(ref message) => write!(f, "cannot alliterate: {}", message),
            Error::Disconnected => write!(f, "caller disconnected / stopped reading"),
        }
    }
}

impl From<io::Error> for Error {
    fn from(error: io::Error) -> Self {
        Error::Io(error)
    }
}

fn app<'a, 'b>() -> clap::App<'a, 'b> {
    clap::App::new("rust-petname")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Generate human readable random names.")
        .after_help(concat!(
            "Based on Dustin Kirkland's petname project ",
            "<https://github.com/dustinkirkland/petname>."
        ))
        .arg(
            Arg::with_name("words")
                .short("w")
                .long("words")
                .value_name("WORDS")
                .default_value("2")
                .help("Number of words in name")
                .takes_value(true)
                .validator(can_be_parsed::<u8>),
        )
        .arg(
            Arg::with_name("separator")
                .short("s")
                .long("separator")
                .value_name("SEP")
                .default_value("-")
                .help("Separator between words")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("complexity")
                .short("c")
                .long("complexity")
                .value_name("COM")
                .possible_values(&["0", "1", "2"])
                .hide_possible_values(true)
                .default_value("0")
                .help("Use small words (0), medium words (1), or large words (2)")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("directory")
                .short("d")
                .long("dir")
                .value_name("DIR")
                .help("Directory containing adjectives.txt, adverbs.txt, names.txt")
                .conflicts_with("complexity")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("count")
                .long("count")
                .value_name("COUNT")
                .default_value("1")
                .help(concat!(
                    "Generate multiple names; pass 0 to produce infinite ",
                    "names (--count=0 is deprecated; use --stream instead)"
                ))
                .takes_value(true)
                .validator(can_be_parsed::<usize>),
        )
        .arg(
            Arg::with_name("stream")
                .long("stream")
                .help("Stream names continuously")
                .conflicts_with("count")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("non-repeating")
                .long("non-repeating")
                .help("Do not generate the same name more than once")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("letters")
                .short("l")
                .long("letters")
                .value_name("LETTERS")
                .default_value("0")
                .help("Maxiumum number of letters in each word; 0 for unlimited")
                .takes_value(true)
                .validator(can_be_parsed::<usize>),
        )
        .arg(
            Arg::with_name("alliterate")
                .short("a")
                .long("alliterate")
                .help("Generate names where each word begins with the same letter")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("alliterate-with")
                .long("alliterate-with")
                .short("A")
                .value_name("LETTER")
                .help("Generate names where each word begins with the given letter")
                .takes_value(true)
                .validator(can_be_parsed::<char>),
        )
        .arg(
            // For compatibility with upstream.
            Arg::with_name("ubuntu")
                .short("u")
                .long("ubuntu")
                .help("Alias; see --alliterate")
                .takes_value(false),
        )
}

fn run(matches: clap::ArgMatches) -> Result<(), Error> {
    // Unwrapping is safe because these options have defaults.
    let opt_separator = matches.value_of("separator").unwrap();
    let opt_words = matches.value_of("words").unwrap();
    let opt_complexity = matches.value_of("complexity").unwrap();
    let opt_count = matches.value_of("count").unwrap();
    let opt_letters = matches.value_of("letters").unwrap();

    // Flags.
    let opt_stream = matches.is_present("stream");
    let opt_non_repeating = matches.is_present("non-repeating");
    let opt_alliterate = matches.is_present("alliterate")
        || matches.is_present("ubuntu")
        || matches.is_present("alliterate-with");

    // Optional arguments without defaults.
    let opt_directory = matches.value_of("directory");
    let opt_alliterate_char = matches
        .value_of("alliterate-with")
        .and_then(|s| s.parse::<char>().ok());

    // Parse numbers. Validated so unwrapping is okay.
    let opt_words: u8 = opt_words.parse().unwrap();
    let opt_count: usize = opt_count.parse().unwrap();
    let opt_letters: usize = opt_letters.parse().unwrap();

    // Load custom word lists, if specified.
    let words = match opt_directory {
        Some(dirname) => Words::load(dirname)?,
        None => Words::Builtin,
    };

    // Select the appropriate word list.
    let mut petnames = match words {
        Words::Custom(ref adjectives, ref adverbs, ref names) => {
            Petnames::init(adjectives, adverbs, names)
        }
        Words::Builtin => match opt_complexity {
            "0" => Petnames::small(),
            "1" => Petnames::medium(),
            "2" => Petnames::large(),
            _ => Petnames::small(),
        },
    };

    // If requested, limit the number of letters.
    if opt_letters != 0 {
        petnames.retain(|s| s.len() <= opt_letters);
    }

    // Check cardinality.
    if petnames.cardinality(opt_words) == 0 {
        return Err(Error::Cardinality(
            "no petnames to choose from; try relaxing constraints".to_string(),
        ));
    }

    // We're going to need a source of randomness.
    let mut rng = rand::thread_rng();

    // Handle alliteration, either by eliminating a specified
    // character, or using a random one.
    if opt_alliterate {
        // We choose the first letter from the intersection of the
        // first letters of each word list in `petnames`.
        let firsts =
            common_first_letters(&petnames.adjectives, &[&petnames.adverbs, &petnames.names]);
        // if a specific character was requested for alliteration,
        // attempt to use it.
        if let Some(c) = opt_alliterate_char {
            if firsts.contains(&c) {
                petnames.retain(|s| s.starts_with(c));
            } else {
                return Err(Error::Alliteration(
                    "no petnames begin with the choosen alliteration character".to_string(),
                ));
            }
        } else {
            // Otherwise choose the first letter at random; fails if
            // there are no letters.
            match firsts.iter().choose(&mut rng) {
                Some(c) => petnames.retain(|s| s.starts_with(*c)),
                None => {
                    return Err(Error::Alliteration(
                        "word lists have no initial letters in common".to_string(),
                    ))
                }
            };
        }
    }

    // Manage stdout.
    let stdout = io::stdout();
    let mut writer = io::BufWriter::new(stdout.lock());

    // Warn that --count=0 is deprecated.
    if opt_count == 0 {
        eprintln!(concat!(
            "Warning: specifying --count=0 to continuously produce petnames is ",
            "deprecated and its behaviour will change in a future version; ",
            "specify --stream instead.",
        ));
    }

    // Stream if count is 0. TODO: Only stream when --stream is specified.
    let count = if opt_stream || opt_count == 0 {
        None
    } else {
        Some(opt_count)
    };

    // Get an iterator for the names we want to print out.
    if opt_non_repeating {
        printer(
            &mut writer,
            petnames.iter_non_repeating(&mut rng, opt_words, opt_separator),
            count,
        )
    } else {
        printer(
            &mut writer,
            petnames.iter(&mut rng, opt_words, opt_separator),
            count,
        )
    }
}

fn printer<OUT, NAMES>(writer: &mut OUT, names: NAMES, count: Option<usize>) -> Result<(), Error>
where
    OUT: io::Write,
    NAMES: Iterator<Item = String>,
{
    match count {
        None => {
            for name in names {
                writeln!(writer, "{}", name).map_err(suppress_disconnect)?;
            }
        }
        Some(n) => {
            for name in names.take(n) {
                writeln!(writer, "{}", name)?;
            }
        }
    }

    Ok(())
}

fn can_be_parsed<INTO>(value: String) -> Result<(), String>
where
    INTO: FromStr,
    <INTO as FromStr>::Err: std::fmt::Display,
{
    match value.parse::<INTO>() {
        Err(e) => Err(format!("{}", e)),
        Ok(_) => Ok(()),
    }
}

fn common_first_letters(init: &[&str], more: &[&[&str]]) -> HashSet<char> {
    let mut firsts = first_letters(init);
    let firsts_other: Vec<HashSet<char>> = more.iter().map(|list| first_letters(list)).collect();
    firsts.retain(|c| firsts_other.iter().all(|fs| fs.contains(c)));
    firsts
}

fn first_letters(names: &[&str]) -> HashSet<char> {
    names.iter().filter_map(|s| s.chars().next()).collect()
}

enum Words {
    Custom(String, String, String),
    Builtin,
}

impl Words {
    // Load word lists from the given directory. This function expects to find three
    // files in that directory: `adjectives.txt`, `adverbs.txt`, and `names.txt`.
    // Each should be valid UTF-8, and contain words separated by whitespace.
    fn load<T: AsRef<path::Path>>(dirname: T) -> Result<Self, Error> {
        let dirname = dirname.as_ref();
        Ok(Self::Custom(
            read_file_to_string(dirname.join("adjectives.txt"))?,
            read_file_to_string(dirname.join("adverbs.txt"))?,
            read_file_to_string(dirname.join("names.txt"))?,
        ))
    }
}

fn read_file_to_string<P: AsRef<path::Path>>(path: P) -> Result<String, Error> {
    fs::read_to_string(&path).map_err(|error| Error::FileIo(path.as_ref().to_path_buf(), error))
}

fn suppress_disconnect(err: io::Error) -> Error {
    match err.kind() {
        io::ErrorKind::BrokenPipe => Error::Disconnected,
        _ => err.into(),
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::os::unix::fs::PermissionsExt;
use std::process::{self, Command, Stdio};

extern crate clap;
use clap::{App, Arg, ArgGroup};

extern crate isbfc;

fn main() {
    let matches = App::new("isbfc")
        .version("0.0.1")
        .author("Ian D. Scott <ian@iandouglasscott.com>")
        .about("Brainfuck compiler")
        .arg(
            Arg::with_name("output_asm")
                .short("S")
                .help("Assemble but do not link"),
        )
        .arg(
            Arg::with_name("dump_ir")
                .long("dumpir")
                .help("Dump intermediate representation; for debugging"),
        )
        .group(ArgGroup::with_name("actions").args(&["output_asm", "dump_ir"]))
        .arg(
            Arg::with_name("debugging_symbols")
                .short("g")
                .help("Generate debugging information"),
        )
        .arg(
            Arg::with_name("out_name")
                .short("o")
                .help("Output file name")
                .takes_value(true)
                .empty_values(false)
                .value_name("file"),
        )
        .arg(
            Arg::with_name("tape_size")
                .long("tape-size")
                .help("Size of tape")
                .takes_value(true)
                .empty_values(false)
                .default_value("8192")
                .value_name("bytes"),
        )
        .arg(
            Arg::with_name("minimal_elf")
                .long("minimal-elf")
                .help("Generate minimal ELF executable"),
        )
        .arg(
            Arg::with_name("level")
                .short("O")
                .help("Optimization level")
                .takes_value(true)
                .empty_values(false)
                .default_value("1"),
        )
        .arg(
            Arg::with_name("FILENAME")
                .help("Source file to compile")
                .required(true)
                .index(1),
        )
        .get_matches();

    let tape_size = matches
        .value_of("tape_size")
        .unwrap()
        .parse::<i32>()
        .unwrap();

    let level = matches.value_of("level").unwrap().parse::<u32>().unwrap();

    let path = matches.value_of("FILENAME").unwrap();
    let name = path.rsplitn(2, '.').last().unwrap();
    let mut file = File::open(&path).unwrap();
    let mut code = Vec::new();
    file.read_to_end(&mut code).unwrap();

    let mut ir = match isbfc::parse(&code) {
        Ok(ir) => ir,
        Err(err) => {
            println!("Parsing error: {}", err);
            process::exit(1);
        }
    };
    if level > 0 {
        ir = ir.optimize();
    }

    if matches.is_present("dump_ir") {
        if let Some(out_name) = matches.value_of("out_name") {
            let mut irfile = File::create(out_name).unwrap();
            writeln!(irfile, "{:#?}", ir).unwrap();
        } else {
            println!("{:#?}", ir);
        }
    } else if matches.is_present("output_asm") {
        println!("Compiling...");
        let output = ir.compile(tape_size);
        let def_name = format!("{}.s", name);
        let out_name = matches.value_of("out_name").unwrap_or(&def_name);
        let mut asmfile = File::create(out_name).unwrap();
        asmfile.write_all(&output.into_bytes()).unwrap();
    } else {
        println!("Compiling...");
        let output = ir.compile(tape_size);
        let out_name = matches.value_of("out_name").unwrap_or(name);
        let debug = matches.is_present("debugging_symbols");
        let minimal = matches.is_present("minimal_elf");
        asm_and_link(&output, &name, &out_name, debug, minimal);
    }
}

fn object_to_binary(o_name: &str) -> (Vec<u8>, u64) {
    let mut o_file = File::open(o_name).unwrap();
    let text = isbfc::elf64_get_section(&mut o_file, b".text")
        .unwrap()
        .unwrap();
    let bss = isbfc::elf64_get_section(&mut o_file, b".bss")
        .unwrap()
        .unwrap();
    let bss_offset = (text.sh_size + 0x1000 - 1) & !(0x1000 - 1);
    let bss_size = bss.sh_size;

    let bin = Command::new("ld")
        .arg("--oformat")
        .arg("binary")
        .arg("-Ttext")
        .arg("0x401000")
        .arg("-Tbss")
        .arg(format!("0x{:x}", 0x401000 + bss_offset))
        .arg("-o")
        .arg("/dev/stdout")
        .arg(o_name)
        .output()
        .unwrap()
        .stdout;

    (bin, bss_size)
}

fn assemble(code: &str, out_name: &str, debug: bool) -> io::Result<Option<i32>> {
    let mut command = Command::new("as");
    if debug {
        command.arg("-g");
    }
    let mut child = command
        .arg("-o")
        .arg(out_name)
        .arg("-") // Standard input
        .stdin(Stdio::piped())
        .spawn()?;

    child
        .stdin
        .take()
        .unwrap()
        .write_all(code.as_bytes())?;

    Ok(child.wait()?.code())
}

fn link(o_name: &str, out_name: &str, minimal: bool) -> io::Result<Option<i32>> {
    if minimal {
        let (bin, bss_size) = object_to_binary(&o_name);

        let mut file = File::create(out_name)?;
        isbfc::elf64_write(&mut file, &bin, bss_size)?;
        let mut permissions = file.metadata().unwrap().permissions();
        permissions.set_mode(permissions.mode() | 0o111);
        file.set_permissions(permissions)?;
        Ok(Some(0))
    } else {
        Ok(Command::new("ld")
            .arg("-o")
            .arg(out_name)
            .arg(o_name)
            .spawn()?
            .wait()?
            .code())
    }
}

fn asm_and_link(code: &str, name: &str, out_name: &str, debug: bool, minimal: bool) {
    let o_name = format!("{}.o", name);

    println!("Assembling...");

    let code = assemble(code, &o_name, debug).unwrap();
    if code != Some(0) {
        process::exit(1);
    }

    println!("Linking...");

    let code = link(&o_name, out_name, minimal).unwrap();
    if code != Some(0) {
        process::exit(1);
    }
}
<commit_msg>Make object_to_binary() return a Result<commit_after>use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::os::unix::fs::PermissionsExt;
use std::process::{self, Command, Stdio};

extern crate clap;
use clap::{App, Arg, ArgGroup};

extern crate isbfc;

fn main() {
    let matches = App::new("isbfc")
        .version("0.0.1")
        .author("Ian D. Scott <ian@iandouglasscott.com>")
        .about("Brainfuck compiler")
        .arg(
            Arg::with_name("output_asm")
                .short("S")
                .help("Assemble but do not link"),
        )
        .arg(
            Arg::with_name("dump_ir")
                .long("dumpir")
                .help("Dump intermediate representation; for debugging"),
        )
        .group(ArgGroup::with_name("actions").args(&["output_asm", "dump_ir"]))
        .arg(
            Arg::with_name("debugging_symbols")
                .short("g")
                .help("Generate debugging information"),
        )
        .arg(
            Arg::with_name("out_name")
                .short("o")
                .help("Output file name")
                .takes_value(true)
                .empty_values(false)
                .value_name("file"),
        )
        .arg(
            Arg::with_name("tape_size")
                .long("tape-size")
                .help("Size of tape")
                .takes_value(true)
                .empty_values(false)
                .default_value("8192")
                .value_name("bytes"),
        )
        .arg(
            Arg::with_name("minimal_elf")
                .long("minimal-elf")
                .help("Generate minimal ELF executable"),
        )
        .arg(
            Arg::with_name("level")
                .short("O")
                .help("Optimization level")
                .takes_value(true)
                .empty_values(false)
                .default_value("1"),
        )
        .arg(
            Arg::with_name("FILENAME")
                .help("Source file to compile")
                .required(true)
                .index(1),
        )
        .get_matches();

    let tape_size = matches
        .value_of("tape_size")
        .unwrap()
        .parse::<i32>()
        .unwrap();

    let level = matches.value_of("level").unwrap().parse::<u32>().unwrap();

    let path = matches.value_of("FILENAME").unwrap();
    let name = path.rsplitn(2, '.').last().unwrap();
    let mut file = File::open(&path).unwrap();
    let mut code = Vec::new();
    file.read_to_end(&mut code).unwrap();

    let mut ir = match isbfc::parse(&code) {
        Ok(ir) => ir,
        Err(err) => {
            println!("Parsing error: {}", err);
            process::exit(1);
        }
    };
    if level > 0 {
        ir = ir.optimize();
    }

    if matches.is_present("dump_ir") {
        if let Some(out_name) = matches.value_of("out_name") {
            let mut irfile = File::create(out_name).unwrap();
            writeln!(irfile, "{:#?}", ir).unwrap();
        } else {
            println!("{:#?}", ir);
        }
    } else if matches.is_present("output_asm") {
        println!("Compiling...");
        let output = ir.compile(tape_size);
        let def_name = format!("{}.s", name);
        let out_name = matches.value_of("out_name").unwrap_or(&def_name);
        let mut asmfile = File::create(out_name).unwrap();
        asmfile.write_all(&output.into_bytes()).unwrap();
    } else {
        println!("Compiling...");
        let output = ir.compile(tape_size);
        let out_name = matches.value_of("out_name").unwrap_or(name);
        let debug = matches.is_present("debugging_symbols");
        let minimal = matches.is_present("minimal_elf");
        asm_and_link(&output, &name, &out_name, debug, minimal);
    }
}

fn object_to_binary(o_name: &str) -> io::Result<(Vec<u8>, u64)> {
    let mut o_file = File::open(o_name)?;
    let text = isbfc::elf64_get_section(&mut o_file, b".text")?
        .unwrap();
    let bss = isbfc::elf64_get_section(&mut o_file, b".bss")?
        .unwrap();
    let bss_offset = (text.sh_size + 0x1000 - 1) & !(0x1000 - 1);
    let bss_size = bss.sh_size;

    let bin = Command::new("ld")
        .arg("--oformat")
        .arg("binary")
        .arg("-Ttext")
        .arg("0x401000")
        .arg("-Tbss")
        .arg(format!("0x{:x}", 0x401000 + bss_offset))
        .arg("-o")
        .arg("/dev/stdout")
        .arg(o_name)
        .output()?
        .stdout;

    Ok((bin, bss_size))
}

fn assemble(code: &str, out_name: &str, debug: bool) -> io::Result<Option<i32>> {
    let mut command = Command::new("as");
    if debug {
        command.arg("-g");
    }
    let mut child = command
        .arg("-o")
        .arg(out_name)
        .arg("-") // Standard input
        .stdin(Stdio::piped())
        .spawn()?;

    child
        .stdin
        .take()
        .unwrap()
        .write_all(code.as_bytes())?;

    Ok(child.wait()?.code())
}

fn link(o_name: &str, out_name: &str, minimal: bool) -> io::Result<Option<i32>> {
    if minimal {
        let (bin, bss_size) = object_to_binary(&o_name)?;

        let mut file = File::create(out_name)?;
        isbfc::elf64_write(&mut file, &bin, bss_size)?;
        let mut permissions = file.metadata().unwrap().permissions();
        permissions.set_mode(permissions.mode() | 0o111);
        file.set_permissions(permissions)?;
        Ok(Some(0))
    } else {
        Ok(Command::new("ld")
            .arg("-o")
            .arg(out_name)
            .arg(o_name)
            .spawn()?
            .wait()?
            .code())
    }
}

fn asm_and_link(code: &str, name: &str, out_name: &str, debug: bool, minimal: bool) {
    let o_name = format!("{}.o", name);

    println!("Assembling...");

    let code = assemble(code, &o_name, debug).unwrap();
    if code != Some(0) {
        process::exit(1);
    }

    println!("Linking...");

    let code = link(&o_name, out_name, minimal).unwrap();
    if code != Some(0) {
        process::exit(1);
    }
}
<|endoftext|>


<commit_before>use std::cell::Cell;

struct Coordinates {
    x: Cell<i32>,
    y: Cell<i32>,
}

fn main() {
    let origin = Coordinates { x: Cell::new(0), y: Cell::new(0) };

    move_up(&origin.y); // borrow the value
    move_up(&origin.y); // so that it can change

    println!("X is {:?} and Y is {:?}", origin.x, origin.y);
}

fn move_up(y_cord: &Cell<i32>) -> &Cell<i32> {
    let new_val = y_cord.get(); // get the value and store it
    y_cord.set(new_val + 1); // modify the stored value and set the incoming param
    y_cord // no semicolon here because it is the return value
}

fn move_down(y_cord: &Cell<i32>) -> &Cell<i32> {
    let new_val = y_cord.get(); // get the value and store it
    y_cord.set(new_val - 1); // modify the stored value and set the incoming param
    y_cord // no semicolon here because it is the return value
}

#[test]
fn it_can_move() {
    let origin = Coordinates { x: Cell::new(0), y: Cell::new(0) };

    assert_eq!(0, origin.y.get());

    move_up(&origin.y);

    assert_eq!(1, origin.y.get());

    move_up(&origin.y);
    move_up(&origin.y);
    move_up(&origin.y);

    assert_eq!(4, origin.y.get());

    for _ in 0..19 { // _ instead of x so that rust does not complain about x not being used
        move_up(&origin.y); // moves up 19 times
    }

    assert_eq!(23, origin.y.get());

    for _ in 0..19 {
        move_down(&origin.y);
    }

    assert_eq!(4, origin.y.get());
}
<commit_msg>can move right - fully tested<commit_after>use std::cell::Cell;

struct Coordinates {
    x: Cell<i32>,
    y: Cell<i32>,
}

fn main() {
    let origin = Coordinates { x: Cell::new(0), y: Cell::new(0) };

    move_up(&origin.y); // borrow the value
    move_up(&origin.y); // so that it can change

    println!("X is {:?} and Y is {:?}", origin.x, origin.y);
}

fn move_up(y_cord: &Cell<i32>) -> &Cell<i32> {
    let new_val = y_cord.get(); // get the value and store it
    y_cord.set(new_val + 1); // modify the stored value and set the incoming param
    y_cord // no semicolon here because it is the return value
}

fn move_down(y_cord: &Cell<i32>) -> &Cell<i32> {
    let new_val = y_cord.get(); // get the value and store it
    y_cord.set(new_val - 1); // modify the stored value and set the incoming param
    y_cord // no semicolon here because it is the return value
}

fn move_right(x_cord: &Cell<i32>) -> &Cell<i32> {
    let new_val = x_cord.get(); // get the value and store it
    x_cord.set(new_val + 1); // modify the stored value and set the incoming param
    x_cord // no semicolon here because it is the return value
}

#[test]
fn it_can_move() {
    let origin = Coordinates { x: Cell::new(0), y: Cell::new(0) };

    assert_eq!(0, origin.y.get());

    move_up(&origin.y);

    assert_eq!(1, origin.y.get());

    move_up(&origin.y);
    move_up(&origin.y);
    move_up(&origin.y);

    assert_eq!(4, origin.y.get());

    for _ in 0..19 { // _ instead of x so that rust does not complain about x not being used
        move_up(&origin.y); // moves up 19 times
    }

    assert_eq!(23, origin.y.get());

    for _ in 0..19 {
        move_down(&origin.y);
    }

    assert_eq!(4, origin.y.get());

    for _ in 0..19 {
        move_right(&origin.x);
    }

    assert_eq!(19, origin.x.get());
}
<|endoftext|>


<commit_before>#[cfg(windows)]
extern crate winapi;

#[macro_use]
extern crate clap;
#[macro_use]
extern crate failure;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate serde_json;

extern crate base64;
extern crate chrono;
extern crate directories;
extern crate regex;
extern crate reqwest;
extern crate serde;
extern crate toml;

mod backend;
mod config;
mod error;
mod input;
mod process;
mod series;

use backend::{anilist::Anilist, SyncBackend};
use config::Config;
use error::Error;
use series::Series;
use std::io::Read;
use std::path::PathBuf;

fn main() {
    match run() {
        Ok(_) => (),
        Err(e) => {
            let e: failure::Error = e.into();
            eprintln!("fatal error: {}", e.cause());

            for cause in e.causes().skip(1) {
                eprintln!("cause: {}", cause);
            }

            eprintln!("{}", e.backtrace());

            // This will allow errors to be displayed when the program
            // is launched within a newly created terminal
            let _ = std::io::stdin().bytes().next();
            std::process::exit(1);
        }
    }
}

fn run() -> Result<(), Error> {
    let args = clap_app!(anup =>
        (version: env!("CARGO_PKG_VERSION"))
        (author: env!("CARGO_PKG_AUTHORS"))
        (@arg NAME: "The name of the series to watch")
        (@arg PATH: -p --path +takes_value "Specifies the directory to look for video files in")
        (@arg SEASON: -s --season +takes_value "Specifies which season you want to watch")
        (@arg LIST: -l --list "Displays all saved series")
        (@arg DONT_SAVE_PASS: --dontsavepass "Disables saving of your account password")
    ).get_matches();

    if args.is_present("LIST") {
        return print_series_list();
    }

    watch_series(&args)
}

fn watch_series(args: &clap::ArgMatches) -> Result<(), Error> {
    let mut config = Config::load()?;
    config.remove_invalid_series();

    let path = get_series_path(&mut config, args)?;
    let sync_backend = Anilist::init(&mut config)?;

    config.save(!args.is_present("DONT_SAVE_PASS"))?;

    let season = args
        .value_of("SEASON")
        .and_then(|s| s.parse().ok())
        .unwrap_or(1);

    let mut series = Series::from_dir(&path, sync_backend)?;
    series.load_season(season)?.play_all_episodes()?;

    Ok(())
}

fn print_series_list() -> Result<(), Error> {
    let config = Config::load()?;

    println!("{} series found", config.series.len());
    println!(
        "note that any series marked as invalid will be removed the next time you watch a series"
    );
    println!();

    for (name, path) in config.series {
        let status_str = if path.exists() { "valid" } else { "invalid" };
        println!("[{}] {}: {}", status_str, name, path.to_string_lossy());
    }

    Ok(())
}

fn get_series_path(config: &mut Config, args: &clap::ArgMatches) -> Result<PathBuf, Error> {
    match args.value_of("PATH") {
        Some(path) => {
            if let Some(series_name) = args.value_of("NAME") {
                config.series.insert(series_name.into(), path.into());
            }

            Ok(path.into())
        }
        None => {
            let name = args.value_of("NAME").ok_or(Error::NoSeriesInfoProvided)?;

            config
                .series
                .get(name)
                .ok_or_else(|| Error::SeriesNotFound(name.into()))
                .map(|path| path.into())
        }
    }
}
<commit_msg>Rename --dontsavepass argument to --dontsavetoken<commit_after>#[cfg(windows)]
extern crate winapi;

#[macro_use]
extern crate clap;
#[macro_use]
extern crate failure;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate serde_json;

extern crate base64;
extern crate chrono;
extern crate directories;
extern crate regex;
extern crate reqwest;
extern crate serde;
extern crate toml;

mod backend;
mod config;
mod error;
mod input;
mod process;
mod series;

use backend::{anilist::Anilist, SyncBackend};
use config::Config;
use error::Error;
use series::Series;
use std::io::Read;
use std::path::PathBuf;

fn main() {
    match run() {
        Ok(_) => (),
        Err(e) => {
            let e: failure::Error = e.into();
            eprintln!("fatal error: {}", e.cause());

            for cause in e.causes().skip(1) {
                eprintln!("cause: {}", cause);
            }

            eprintln!("{}", e.backtrace());

            // This will allow errors to be displayed when the program
            // is launched within a newly created terminal
            let _ = std::io::stdin().bytes().next();
            std::process::exit(1);
        }
    }
}

fn run() -> Result<(), Error> {
    let args = clap_app!(anup =>
        (version: env!("CARGO_PKG_VERSION"))
        (author: env!("CARGO_PKG_AUTHORS"))
        (@arg NAME: "The name of the series to watch")
        (@arg PATH: -p --path +takes_value "Specifies the directory to look for video files in")
        (@arg SEASON: -s --season +takes_value "Specifies which season you want to watch")
        (@arg LIST: -l --list "Displays all saved series")
        (@arg DONT_SAVE_TOKEN: --dontsavetoken "Don't save your account access token")
    ).get_matches();

    if args.is_present("LIST") {
        return print_series_list();
    }

    watch_series(&args)
}

fn watch_series(args: &clap::ArgMatches) -> Result<(), Error> {
    let mut config = Config::load()?;
    config.remove_invalid_series();

    let path = get_series_path(&mut config, args)?;
    let sync_backend = Anilist::init(&mut config)?;

    config.save(!args.is_present("DONT_SAVE_TOKEN"))?;

    let season = args
        .value_of("SEASON")
        .and_then(|s| s.parse().ok())
        .unwrap_or(1);

    let mut series = Series::from_dir(&path, sync_backend)?;
    series.load_season(season)?.play_all_episodes()?;

    Ok(())
}

fn print_series_list() -> Result<(), Error> {
    let config = Config::load()?;

    println!("{} series found", config.series.len());
    println!(
        "note that any series marked as invalid will be removed the next time you watch a series"
    );
    println!();

    for (name, path) in config.series {
        let status_str = if path.exists() { "valid" } else { "invalid" };
        println!("[{}] {}: {}", status_str, name, path.to_string_lossy());
    }

    Ok(())
}

fn get_series_path(config: &mut Config, args: &clap::ArgMatches) -> Result<PathBuf, Error> {
    match args.value_of("PATH") {
        Some(path) => {
            if let Some(series_name) = args.value_of("NAME") {
                config.series.insert(series_name.into(), path.into());
            }

            Ok(path.into())
        }
        None => {
            let name = args.value_of("NAME").ok_or(Error::NoSeriesInfoProvided)?;

            config
                .series
                .get(name)
                .ok_or_else(|| Error::SeriesNotFound(name.into()))
                .map(|path| path.into())
        }
    }
}
<|endoftext|>


<commit_before>extern crate image;
extern crate tray_rust;

use std::iter;
use std::vec::Vec;
use std::sync::{Arc, TaskPool};
use std::sync::mpsc;
use std::sync::mpsc::{Sender, Receiver};
use std::time::duration::Duration;

use tray_rust::film;
use tray_rust::geometry::Geometry;
use tray_rust::sampler;
use tray_rust::sampler::{Sampler};
use tray_rust::scene;
use tray_rust::integrator::Integrator;

static WIDTH: usize = 800;
static HEIGHT: usize = 600;

/// Threads are each sent a sender end of the channel that is
/// read from by the render target thread which then saves the
/// values recieved to the render target
fn thread_work(tx: Sender<(f32, f32, film::Colorf)>, queue: Arc<sampler::BlockQueue>,
               scene: Arc<scene::Scene>) {
    let mut sampler = sampler::Uniform::new(queue.block_dim());
    let mut samples = Vec::with_capacity(sampler.max_spp());
    let mut sample_pos = Vec::with_capacity(sampler.max_spp());
    // Grab a block from the queue and start working on it, submitting samples
    // to the render target thread after each pixel
    for b in queue.iter() {
        sampler.select_block(&b);
        while sampler.has_samples() {
            // Get samples for a pixel and render them
            sampler.get_samples(&mut sample_pos);
            for s in sample_pos.iter() {
                let mut ray = scene.camera.generate_ray(s);
                if let Some(hit) = scene.intersect(&mut ray) {
                    let c = scene.integrator.illumination(&*scene, &ray, &hit);
                    samples.push((s.0, s.1, c));
                }
            }
            for s in samples.iter() {
                if let Err(e) = tx.send(*s) {
                    println!("Worker thread exiting with send error {:?}", e);
                    return;
                }
            }
            samples.clear();
        }
    }
}

/// Spawn `n` worker threads to render the scene in parallel. Returns the receive end
/// of the channel where the threads will write their samples so that the receiver
/// can write these samples to the render target
fn spawn_workers(pool: &TaskPool, n: usize, scene: Arc<scene::Scene>) -> Receiver<(f32, f32, film::Colorf)> {
    let (tx, rx) = mpsc::channel();
    let block_queue = Arc::new(sampler::BlockQueue::new((WIDTH as u32, HEIGHT as u32), (8, 8)));
    // TODO: the .. operator precedence is very low so we need this paren here at the moment
    // once (hopefully) it's raised we can remove the parens
    for _ in (0..n) {
        let q = block_queue.clone();
        let t = tx.clone();
        let s = scene.clone();
        pool.execute(move || {
            thread_work(t, q, s);
        });
    }
    rx
}

/// Render in the scene in parallel to the render target
fn render_parallel(rt: &mut film::RenderTarget){
    let scene = Arc::new(scene::Scene::new(WIDTH, HEIGHT));
    let n = 8;
    let pool = TaskPool::new(n);
    let rx = spawn_workers(&pool, n, scene);
    for m in rx.iter() {
        rt.write(m.0, m.1, &m.2);
    }
}

fn main() {
    let mut rt = film::RenderTarget::new(WIDTH, HEIGHT);
    let d = Duration::span(|| render_parallel(&mut rt));
    println!("Rendering took {}ms", d.num_milliseconds());
    let img = rt.get_render();
    image::save_buffer(&Path::new("out.png"), &img[], WIDTH as u32, HEIGHT as u32, image::RGB(8));
}

<commit_msg>Remove unused import from main<commit_after>extern crate image;
extern crate tray_rust;

use std::vec::Vec;
use std::sync::{Arc, TaskPool};
use std::sync::mpsc;
use std::sync::mpsc::{Sender, Receiver};
use std::time::duration::Duration;

use tray_rust::film;
use tray_rust::geometry::Geometry;
use tray_rust::sampler;
use tray_rust::sampler::{Sampler};
use tray_rust::scene;
use tray_rust::integrator::Integrator;

static WIDTH: usize = 800;
static HEIGHT: usize = 600;

/// Threads are each sent a sender end of the channel that is
/// read from by the render target thread which then saves the
/// values recieved to the render target
fn thread_work(tx: Sender<(f32, f32, film::Colorf)>, queue: Arc<sampler::BlockQueue>,
               scene: Arc<scene::Scene>) {
    let mut sampler = sampler::Uniform::new(queue.block_dim());
    let mut samples = Vec::with_capacity(sampler.max_spp());
    let mut sample_pos = Vec::with_capacity(sampler.max_spp());
    // Grab a block from the queue and start working on it, submitting samples
    // to the render target thread after each pixel
    for b in queue.iter() {
        sampler.select_block(&b);
        while sampler.has_samples() {
            // Get samples for a pixel and render them
            sampler.get_samples(&mut sample_pos);
            for s in sample_pos.iter() {
                let mut ray = scene.camera.generate_ray(s);
                if let Some(hit) = scene.intersect(&mut ray) {
                    let c = scene.integrator.illumination(&*scene, &ray, &hit);
                    samples.push((s.0, s.1, c));
                }
            }
            for s in samples.iter() {
                if let Err(e) = tx.send(*s) {
                    println!("Worker thread exiting with send error {:?}", e);
                    return;
                }
            }
            samples.clear();
        }
    }
}

/// Spawn `n` worker threads to render the scene in parallel. Returns the receive end
/// of the channel where the threads will write their samples so that the receiver
/// can write these samples to the render target
fn spawn_workers(pool: &TaskPool, n: usize, scene: Arc<scene::Scene>) -> Receiver<(f32, f32, film::Colorf)> {
    let (tx, rx) = mpsc::channel();
    let block_queue = Arc::new(sampler::BlockQueue::new((WIDTH as u32, HEIGHT as u32), (8, 8)));
    // TODO: the .. operator precedence is very low so we need this paren here at the moment
    // once (hopefully) it's raised we can remove the parens
    for _ in (0..n) {
        let q = block_queue.clone();
        let t = tx.clone();
        let s = scene.clone();
        pool.execute(move || {
            thread_work(t, q, s);
        });
    }
    rx
}

/// Render in the scene in parallel to the render target
fn render_parallel(rt: &mut film::RenderTarget){
    let scene = Arc::new(scene::Scene::new(WIDTH, HEIGHT));
    let n = 8;
    let pool = TaskPool::new(n);
    let rx = spawn_workers(&pool, n, scene);
    for m in rx.iter() {
        rt.write(m.0, m.1, &m.2);
    }
}

fn main() {
    let mut rt = film::RenderTarget::new(WIDTH, HEIGHT);
    let d = Duration::span(|| render_parallel(&mut rt));
    println!("Rendering took {}ms", d.num_milliseconds());
    let img = rt.get_render();
    image::save_buffer(&Path::new("out.png"), &img[], WIDTH as u32, HEIGHT as u32, image::RGB(8));
}

<|endoftext|>


<commit_before>extern crate riemann_client;
extern crate chrono;
extern crate docopt;

use std::process::Command;

use riemann_client::Client;
use riemann_client::proto::Event;

use std::io::prelude::*;
use std::net::TcpStream;

use docopt::Docopt;

use std::net::ToSocketAddrs;

// Write the Docopt usage string.
static USAGE: &'static str = "
Usage: dickedur -g GHOSTNAME -l HOSTNAME [-p PORT]
       dickedur (--help | -h)

Options:
    -h, --help     Show this screen.
    -l HOSTNAME   Hostname to advertise in graphite.
    -g GHOSTNAME   Graphite Hostname.
    -p PORT       Graphite port [default: 2003]
";

#[cfg(not(any(target_os="macos")))]
fn disk_free() -> std::process::Output {
	Command::new("df")
		.arg("--portability").output().unwrap_or_else(|e| { panic!("failed to execute process: {}", e) })
}

#[cfg(any(target_os="macos"))]
fn disk_free() -> std::process::Output {
	Command::new("df")
		.arg("-k").output().unwrap_or_else(|e| { panic!("failed to execute process: {}", e) })
}

fn main() {

	// Parse argv and exit the program with an error message if it fails.
	let args = Docopt::new(USAGE)
                  .and_then(|d| d.argv(std::env::args().into_iter()).parse())
                  .unwrap_or_else(|e| e.exit());

    let connect_string = format!("{}:{}", args.get_str("<graphite_hostname>"), args.get_str("graphite_port"));
    println!("Connect string is {}", connect_string);

    let o = disk_free();
	let stdout = String::from_utf8(o.stdout).
		ok().
		unwrap();

	let lines = stdout.split("\n");

	// graphite.proto.melvil.io 2003
	// timestamp = 1436278152
	// Graphite format
	// local.random.diceroll 4 `date +%s`
    let mut stream = TcpStream::connect(&*connect_string).unwrap();
    let dt = chrono::UTC::now();
    let timestamp = dt.timestamp();

	for line in lines {
		if line.starts_with("/") {
			let values: Vec<&str> = line.split(" ").filter(|e| {
				e.len() != 0
			}).collect();

			let disk = values[0];
			let all = values[2];
			let available = values[3];

	 	}
	}

    // let mut client = Client::connect(&("melvil.testing", 5555)).unwrap();
    // client.event({
    //     let mut event = Event::new();
    //     event.set_service("rust-riemann_client".to_string());
    //     event.set_state("ok".to_string());
    //     event.set_metric_d(128.128);
    //     event
    // }).unwrap();

    // client.event(riemann_client::Event {
    //     service: "rust-riemann_client",
    //     state: "ok",
    //     metric_d: 128.128
    //     ..Event::new()
    // }).unwrap()
}<commit_msg>First Working version<commit_after>extern crate riemann_client;
extern crate chrono;
extern crate docopt;

use std::process::Command;

use riemann_client::Client;
use riemann_client::proto::Event;

use std::io::prelude::*;
use std::net::TcpStream;

use docopt::Docopt;

use std::net::ToSocketAddrs;

// Write the Docopt usage string.
static USAGE: &'static str = "
Usage: dickedur -g GHOSTNAME -l HOSTNAME [-p PORT]
       dickedur (--help | -h)

Options:
    -h, --help     Show this screen.
    -l HOSTNAME    Hostname to advertise in graphite.
    -g GHOSTNAME   Graphite Hostname.
    -p PORT        Graphite port [default: 2003]
";

#[cfg(not(any(target_os="macos")))]
fn disk_free() -> std::process::Output {
	Command::new("df")
		.arg("--portability").output().unwrap_or_else(|e| { panic!("failed to execute process: {}", e) })
}

#[cfg(any(target_os="macos"))]
fn disk_free() -> std::process::Output {
	Command::new("df")
		.arg("-k").output().unwrap_or_else(|e| { panic!("failed to execute process: {}", e) })
}

fn main() {

	// Parse argv and exit the program with an error message if it fails.
	let args = Docopt::new(USAGE)
		.and_then(|d| d.argv(std::env::args().into_iter()).parse())
		.unwrap_or_else(|e| e.exit());

	println!("{:?}", args);

    let connect_string = format!("{}:{}", args.get_str("-g"), args.get_str("-p"));
    println!("Connect string is {}", connect_string);
    let my_hostname = args.get_str("-l");

	// graphite.proto.melvil.io 2003
	// timestamp = 1436278152
	// Graphite format
	// local.random.diceroll 4 `date +%s`

	while true {
	    let dt = chrono::UTC::now();
	    let timestamp = dt.timestamp();
	    let o = disk_free();
		let stdout = String::from_utf8(o.stdout).
			ok().
			unwrap();

		let lines = stdout.split("\n");

	    let mut stream = TcpStream::connect(&*connect_string).unwrap();

		for line in lines {
			if line.starts_with("/") {
				let values: Vec<&str> = line.split(" ").filter(|e| {
					e.len() != 0
				}).collect();

				let disk = values[0];
				let all = values[2];
				let available = values[3];

				let content = format!("test.{}.{}.all {} {}\n", my_hostname, disk, all, timestamp);
				println!("{}", content);
				let _ = stream.write(&content.as_bytes());

				let content = format!("test.{}.{}.available {} {}\n", my_hostname, disk, available, timestamp);
				println!("{}", content);
				let _ = stream.write(&content.as_bytes());
		 	}
		}
	 	std::thread::sleep_ms(10000);
	}

    // let mut client = Client::connect(&("melvil.testing", 5555)).unwrap();
    // client.event({
    //     let mut event = Event::new();
    //     event.set_service("rust-riemann_client".to_string());
    //     event.set_state("ok".to_string());
    //     event.set_metric_d(128.128);
    //     event
    // }).unwrap();

    // client.event(riemann_client::Event {
    //     service: "rust-riemann_client",
    //     state: "ok",
    //     metric_d: 128.128
    //     ..Event::new()
    // }).unwrap()
}<|endoftext|>


<commit_before>use clap::Clap;
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("#", "");
    if !col.chars().next().unwrap().is_alphabetic() {
        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.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("");
                    }
                }
                Ordering::Greater => {
                    panic!("Too many fields on row {}, fields: {:?}", num_rows + 1, row);
                }
                Ordering::Equal => {}
            }
            stmt.execute(&row).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 {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    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 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..r.column_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 mut writer = csv::Writer::from_path(destination_path).unwrap();
    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..r.column_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);
    }
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [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)]
struct Opts {
    #[clap(long, about = "Use ',' as the delimiter for the CSV")]
    comma: bool,
    #[clap(long, about = "Use '|' as the delimiter for the CSV")]
    pipe: bool,
    #[clap(long, about = "Use '\\t' as the delimiter for the CSV")]
    tab: 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) {
        (true, false, false) | (false, false, false) => b',',
        (false, true, false) => b'|',
        (false, false, true) => b'\t',
        _ => {
            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 = [
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
        ".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);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    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;
            }
        }
    }

    Ok(())
}

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

    #[test]
    fn test_normalize_col() {
        for (value, expected) in &[
            ("abc", "abc"),
            ("abc (123)", "abc"),
            ("2/6/2000", "c_2_6_2000"),
            ("COMBO#", "combo"),
        ] {
            assert_eq!(&&normalize_col(value), expected);
        }
    }
}
<commit_msg>handle empty col names (at least, don't panic)<commit_after>use clap::Clap;
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("#", "");
    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.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("");
                    }
                }
                Ordering::Greater => {
                    panic!("Too many fields on row {}, fields: {:?}", num_rows + 1, row);
                }
                Ordering::Equal => {}
            }
            stmt.execute(&row).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 {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    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 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..r.column_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 mut writer = csv::Writer::from_path(destination_path).unwrap();
    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..r.column_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);
    }
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [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)]
struct Opts {
    #[clap(long, about = "Use ',' as the delimiter for the CSV")]
    comma: bool,
    #[clap(long, about = "Use '|' as the delimiter for the CSV")]
    pipe: bool,
    #[clap(long, about = "Use '\\t' as the delimiter for the CSV")]
    tab: 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) {
        (true, false, false) | (false, false, false) => b',',
        (false, true, false) => b'|',
        (false, false, true) => b'\t',
        _ => {
            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 = [
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
        ".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);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    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;
            }
        }
    }

    Ok(())
}

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

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


<commit_before>#![allow(dead_code)]
// #![allow(unused_features)]
// #![feature(test)]
#[macro_use] extern crate clap;
extern crate cgmath;
use cgmath::{Matrix4, Point3, Deg, perspective};
// use cgmath::prelude::*;

extern crate gl;

extern crate glutin;
use glutin::{
    CursorState,
    ElementState,
    MouseScrollDelta,
    GlContext,
    VirtualKeyCode,
    WindowEvent,
};

extern crate gltf;
extern crate image;

extern crate futures;

use clap::{Arg, App};

// use std::sync::mpsc::Receiver;
use std::time::Instant;

mod shader;
use shader::Shader;
mod camera;
use camera::Camera;
use camera::CameraMovement::*;
mod macros;
mod http_source;
use http_source::HttpSource;
mod utils;
use utils::{print_elapsed, FrameTimer};

mod render;
use render::*;

const SCR_WIDTH: u32 = 800;
const SCR_HEIGHT: u32 = 600;

pub fn main() {
    let args = App::new("gltf-viewer")
        .version(crate_version!())
        .arg(Arg::with_name("FILE/URL")
            .required(true)
            .takes_value(true)
            .help("glTF file name or URL"))
        .arg(Arg::with_name("screenshot")
            .long("screenshot")
            .short("s")
            .help("Create screenshot (NOT WORKING YET)"))
        .get_matches();
    let source = args.value_of("FILE/URL").unwrap();
    let screenshot = args.is_present("screenshot");

    let mut camera = Camera {
        // TODO!: position.z - bounding box length
        position: Point3::new(0.0, 0.0, 1.0),
        zoom: 60.0,
        ..Camera::default()
    };

    let mut first_mouse = true;
    let mut last_x: f32 = SCR_WIDTH as f32 / 2.0;
    let mut last_y: f32 = SCR_HEIGHT as f32 / 2.0;

    // glutin: initialize and configure
    let mut events_loop = glutin::EventsLoop::new();

    // TODO?: hints for 4.1, core profile, forward compat
    let window = glutin::WindowBuilder::new()
            .with_title("gltf-viewer")
            // TODO: configurable initial dimensions
            .with_dimensions(SCR_WIDTH, SCR_HEIGHT);

    let context = glutin::ContextBuilder::new()
        .with_vsync(true);
    let gl_window = glutin::GlWindow::new(window, context, &events_loop).unwrap();

    unsafe {
        gl_window.make_current().unwrap();
    }

    // TODO!: capturing - on click or uncapture somehow?
    let _ = gl_window.set_cursor_state(CursorState::Hide);

    // gl: load all OpenGL function pointers
    gl::load_with(|symbol| gl_window.get_proc_address(symbol) as *const _);

    let (mut shader, mut scene, loc_projection, loc_view) = unsafe {
            gl::ClearColor(0.1, 1.0, 0.3, 1.0);
            gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);

        gl::Enable(gl::DEPTH_TEST);

        let mut shader = Shader::from_source(
            include_str!("shaders/simple.vs"),
            include_str!("shaders/simple.fs"));

        // NOTE: shader debug version
        // let shader = Shader::new(
        //     "src/shaders/simple.vs",
        //     "src/shaders/simple.fs");

        shader.use_program();
        let loc_projection = shader.uniform_location("projection");
        let loc_view = shader.uniform_location("view");

        let mut start_time = Instant::now();
        let gltf =
            if source.starts_with("http") {
                let http_source = HttpSource::new(source);
                let import = gltf::Import::custom(http_source, Default::default());
                let gltf = import_gltf(import);
                println!(); // to end the "progress dots"
                gltf
            }
            else {
                let import = gltf::Import::from_path(source);
                import_gltf(import)
            };
        print_elapsed("Imported glTF in ", &start_time);
        start_time = Instant::now();

        // load first scene
        let scene = Scene::from_gltf(gltf.scenes().nth(0).unwrap());
        print_elapsed("Loaded scene in", &start_time);
        println!("Nodes: {:<2}\nMeshes: {:<2}",
            gltf.nodes().count(),
            scene.meshes.len());

        // TODO: keyboard switch?
        // draw in wireframe
        // gl::PolygonMode(gl::FRONT_AND_BACK, gl::LINE);

        (shader, scene, loc_projection, loc_view)
    };

    // timing
    let mut delta_time: f64; // seconds
    let mut last_frame: Instant = Instant::now();

    let mut render_timer = FrameTimer::new("rendering", 300);

    // render loop
    let mut running = true;
    while running {
        // per-frame time logic
        // NOTE: Deliberately ignoring the seconds of `elapsed()`
        delta_time = (last_frame.elapsed().subsec_nanos() as f64) / 1_000_000_000.0;
        last_frame = Instant::now();

        // TODO!!: refactor into struct...
        running = process_events(&mut events_loop, &gl_window, &mut first_mouse, &mut last_x, &mut last_y, &mut camera);
        camera.update(delta_time); // navigation

        // render
        unsafe {
            render_timer.start();

            gl::ClearColor(0.1, 0.2, 0.3, 1.0);
            gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);

            shader.use_program();

            // view/projection transformations
            // TODO!: only re-compute/set perspective on Zoom changes (also view?)
            let projection: Matrix4<f32> = perspective(Deg(camera.zoom), SCR_WIDTH as f32 / SCR_HEIGHT as f32, 0.01, 1000.0);
            let view = camera.get_view_matrix();
            shader.set_mat4(loc_projection, &projection);
            shader.set_mat4(loc_view, &view);

            scene.draw(&mut shader);

            render_timer.end();
        }

        gl_window.swap_buffers().unwrap();

        // TODO!: implement screenshotting
        if screenshot { return }
    }
}

fn process_events(
    events_loop: &mut glutin::EventsLoop,
    gl_window: &glutin::GlWindow,
    first_mouse: &mut bool,
    last_x: &mut f32,
    last_y: &mut f32,
    mut camera: &mut Camera) -> bool
{
    let mut keep_running = true;
    events_loop.poll_events(|event| {
        match event {
            glutin::Event::WindowEvent{ event, .. } => match event {
                WindowEvent::Closed => keep_running = false,
                WindowEvent::Resized(w, h) => gl_window.resize(w, h), // TODO!: handle aspect ratio changes
                WindowEvent::DroppedFile(_path_buf) => (), // TODO!: drag file in
                WindowEvent::MouseMoved { position: (xpos, ypos), .. } => {
                    let (xpos, ypos) = (xpos as f32, ypos as f32);
                    if *first_mouse {
                        *last_x = xpos;
                        *last_y = ypos;
                        *first_mouse = false;
                    }

                    let xoffset = xpos - *last_x;
                    let yoffset = *last_y - ypos; // reversed since y-coordinates go from bottom to top

                    *last_x = xpos;
                    *last_y = ypos;

                    camera.process_mouse_movement(xoffset, yoffset, true);
                },
                WindowEvent::MouseWheel { delta: MouseScrollDelta::PixelDelta(_xoffset, yoffset), .. } => {
                    // TODO: need to handle LineDelta case too?
                    camera.process_mouse_scroll(yoffset);
                }
                WindowEvent::KeyboardInput { input, .. } => {
                    keep_running = process_input(input, &mut camera);
                }
                _ => ()
            },
            _ => ()
        }
    });

    keep_running
}

fn process_input(input: glutin::KeyboardInput, camera: &mut Camera) -> bool {
    let pressed = match input.state {
        ElementState::Pressed => true,
        ElementState::Released => false
    };
    if let Some(code) = input.virtual_keycode {
        match code {
            VirtualKeyCode::Escape if pressed => return false,
            VirtualKeyCode::W => camera.process_keyboard(FORWARD, pressed),
            VirtualKeyCode::S => camera.process_keyboard(BACKWARD, pressed),
            VirtualKeyCode::A => camera.process_keyboard(LEFT, pressed),
            VirtualKeyCode::D => camera.process_keyboard(RIGHT, pressed),
            _ => ()
        }
    }
    true
}

fn import_gltf<S: gltf::import::Source>(import: gltf::Import<S>) -> gltf::Gltf {
    match import.sync() {
        Ok(gltf) => gltf,
        Err(err) => {
            println!("glTF import failed: {:?}", err);
            std::process::exit(1);
        }
    }
}

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

    #[test]
    fn print_struct_sizes() {
        // run with `cargo test -- --nocapture`
        println!("Sizes in bytes:");
        println!("Scene:     {:>3}", std::mem::size_of::<Scene>());
        println!("Node:      {:>3}", std::mem::size_of::<Node>());
        println!("Mesh:      {:>3}", std::mem::size_of::<Mesh>());
        println!("Primitive: {:>3}", std::mem::size_of::<Primitive>());
        println!("Vertex:    {:>3}", std::mem::size_of::<Vertex>());
        println!();
        println!("Option<String>: {:>3}", std::mem::size_of::<Option<String>>());
        println!("String:         {:>3}", std::mem::size_of::<String>());
        println!("Vec<f32>:       {:>3}", std::mem::size_of::<Vec<f32>>());
        println!("Vec<Node>:      {:>3}", std::mem::size_of::<Vec<Node>>());
    }

//     extern crate test;
//     use self::test::Bencher;
//     #[bench]
//     fn bench_frame_timer(b: &mut Bencher) {
//         let mut timer = FrameTimer::new("Foobar", 60);
//         b.iter(|| {
//             for _ in 0..60 {
//                 timer.start();
//                 timer.end();
//             }
//         })
//     }
}
<commit_msg>refactor main function into struct<commit_after>#![allow(dead_code)]
// #![allow(unused_features)]
// #![feature(test)]
#[macro_use] extern crate clap;
extern crate cgmath;
use cgmath::{Matrix4, Point3, Deg, perspective};
// use cgmath::prelude::*;

extern crate gl;

extern crate glutin;
use glutin::{
    CursorState,
    ElementState,
    MouseScrollDelta,
    GlContext,
    VirtualKeyCode,
    WindowEvent,
};

extern crate gltf;
extern crate image;

extern crate futures;

use clap::{Arg, App};

use std::time::Instant;

mod shader;
use shader::Shader;
mod camera;
use camera::Camera;
use camera::CameraMovement::*;
mod macros;
mod http_source;
use http_source::HttpSource;
mod utils;
use utils::{print_elapsed, FrameTimer};

mod render;
use render::*;

const SCR_WIDTH: u32 = 800;
const SCR_HEIGHT: u32 = 600;

pub fn main() {
    let args = App::new("gltf-viewer")
        .version(crate_version!())
        .arg(Arg::with_name("FILE/URL")
            .required(true)
            .takes_value(true)
            .help("glTF file name or URL"))
        .arg(Arg::with_name("screenshot")
            .long("screenshot")
            .short("s")
            .help("Create screenshot (NOT WORKING YET)"))
        .get_matches();
    let source = args.value_of("FILE/URL").unwrap();
    let screenshot = args.is_present("screenshot");

    let mut viewer = GltfViewer::new(source);
    viewer.start_render_loop(screenshot);
}

struct GltfViewer {
    camera: Camera,
    first_mouse: bool,
    last_x: f32,
    last_y: f32,
    events_loop: glutin::EventsLoop,
    gl_window: glutin::GlWindow,

    shader: Shader,
    loc_projection: i32,
    loc_view: i32,

    scene: Scene,

    delta_time: f64, // seconds
    last_frame: Instant,
}

impl GltfViewer {
    pub fn new(source: &str) -> GltfViewer {
        let camera = Camera {
            // TODO!: position.z - bounding box length
            position: Point3::new(0.0, 0.0, 1.0),
            zoom: 60.0,
            ..Camera::default()
        };

        let first_mouse = true;
        let last_x: f32 = SCR_WIDTH as f32 / 2.0;
        let last_y: f32 = SCR_HEIGHT as f32 / 2.0;

        // glutin: initialize and configure
        let events_loop = glutin::EventsLoop::new();

        // TODO?: hints for 4.1, core profile, forward compat
        let window = glutin::WindowBuilder::new()
                .with_title("gltf-viewer")
                // TODO: configurable initial dimensions
                .with_dimensions(SCR_WIDTH, SCR_HEIGHT);

        let context = glutin::ContextBuilder::new()
            .with_vsync(true);
        let gl_window = glutin::GlWindow::new(window, context, &events_loop).unwrap();

        unsafe {
            gl_window.make_current().unwrap();
        }

        // TODO!: capturing - on click or uncapture somehow?
        // TODO!!: find solution for macOS - see https://github.com/tomaka/glutin/issues/226
         #[cfg(target_os = "macos")]
        let _ = gl_window.set_cursor_state(CursorState::Hide);
         #[cfg(not(target_os = "macos"))]
        let _ = gl_window.set_cursor_state(CursorState::Grab);

        // gl: load all OpenGL function pointers
        gl::load_with(|symbol| gl_window.get_proc_address(symbol) as *const _);

        let (shader, loc_projection, loc_view) = unsafe {
            gl::ClearColor(0.1, 1.0, 0.3, 1.0);
            gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);

            gl::Enable(gl::DEPTH_TEST);

            let mut shader = Shader::from_source(
                include_str!("shaders/simple.vs"),
                include_str!("shaders/simple.fs"));

            // NOTE: shader debug version
            // let shader = Shader::new(
            //     "src/shaders/simple.vs",
            //     "src/shaders/simple.fs");

            shader.use_program();
            let loc_projection = shader.uniform_location("projection");
            let loc_view = shader.uniform_location("view");

            // TODO: keyboard switch?
            // draw in wireframe
            gl::PolygonMode(gl::FRONT_AND_BACK, gl::LINE);

            (shader, loc_projection, loc_view)
        };

        GltfViewer {
            camera,
            first_mouse, last_x, last_y,

            events_loop,
            gl_window,

            shader, loc_projection, loc_view,

            scene: Self::load(source),

            delta_time: 0.0, // seconds
            last_frame: Instant::now(),
        }
    }

    pub fn load(source: &str) -> Scene {
        let mut start_time = Instant::now();
        let gltf =
            if source.starts_with("http") {
                let http_source = HttpSource::new(source);
                let import = gltf::Import::custom(http_source, Default::default());
                let gltf = import_gltf(import);
                println!(); // to end the "progress dots"
                gltf
            }
            else {
                let import = gltf::Import::from_path(source);
                import_gltf(import)
            };
        print_elapsed("Imported glTF in ", &start_time);
        start_time = Instant::now();

        // load first scene
        let scene = Scene::from_gltf(gltf.scenes().nth(0).unwrap());
        print_elapsed("Loaded scene in", &start_time);
        println!("Nodes: {:<2}\nMeshes: {:<2}",
            gltf.nodes().count(),
            scene.meshes.len());

        scene
    }

    pub fn start_render_loop(&mut self, screenshot: bool) {
        let mut render_timer = FrameTimer::new("rendering", 300);
        loop {
            // per-frame time logic
            // NOTE: Deliberately ignoring the seconds of `elapsed()`
            self.delta_time = (self.last_frame.elapsed().subsec_nanos() as f64) / 1_000_000_000.0;
            self.last_frame = Instant::now();

            let keep_running = process_events(
                &mut self.events_loop, &self.gl_window,
                &mut self.first_mouse, &mut self.last_x, &mut self.last_y,
                &mut self.camera);
            if !keep_running { break }
            self.camera.update(self.delta_time); // navigation

            // render
            unsafe {
                render_timer.start();

                gl::ClearColor(0.1, 0.2, 0.3, 1.0);
                gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);

                self.shader.use_program();

                // view/projection transformations
                // TODO!: only re-compute/set perspective on Zoom changes (also view?)
                let projection: Matrix4<f32> = perspective(Deg(self.camera.zoom), SCR_WIDTH as f32 / SCR_HEIGHT as f32, 0.01, 1000.0);
                let view = self.camera.get_view_matrix();
                self.shader.set_mat4(self.loc_projection, &projection);
                self.shader.set_mat4(self.loc_view, &view);

                self.scene.draw(&mut self.shader);

                render_timer.end();
            }

            self.gl_window.swap_buffers().unwrap();

            // TODO!: implement screenshotting
            if screenshot { return }
        }
    }
}

fn import_gltf<S: gltf::import::Source>(import: gltf::Import<S>) -> gltf::Gltf {
    match import.sync() {
        Ok(gltf) => gltf,
        Err(err) => {
            println!("glTF import failed: {:?}", err);
            std::process::exit(1);
        }
    }
}

fn process_events(
    events_loop: &mut glutin::EventsLoop,
    gl_window: &glutin::GlWindow,
    first_mouse: &mut bool,
    last_x: &mut f32,
    last_y: &mut f32,
    mut camera: &mut Camera) -> bool
{
    let mut keep_running = true;
    events_loop.poll_events(|event| {
        match event {
            glutin::Event::WindowEvent{ event, .. } => match event {
                WindowEvent::Closed => keep_running = false,
                WindowEvent::Resized(w, h) => gl_window.resize(w, h), // TODO!: handle aspect ratio changes
                WindowEvent::DroppedFile(_path_buf) => (), // TODO!: drag file in
                WindowEvent::MouseMoved { position: (xpos, ypos), .. } => {
                    let (xpos, ypos) = (xpos as f32, ypos as f32);
                    if *first_mouse {
                        *last_x = xpos;
                        *last_y = ypos;
                        *first_mouse = false;
                    }

                    let xoffset = xpos - *last_x;
                    let yoffset = *last_y - ypos; // reversed since y-coordinates go from bottom to top

                    *last_x = xpos;
                    *last_y = ypos;

                    camera.process_mouse_movement(xoffset, yoffset, true);
                },
                WindowEvent::MouseWheel { delta: MouseScrollDelta::PixelDelta(_xoffset, yoffset), .. } => {
                    // TODO: need to handle LineDelta case too?
                    camera.process_mouse_scroll(yoffset);
                }
                WindowEvent::KeyboardInput { input, .. } => {
                    keep_running = process_input(input, &mut camera);
                }
                _ => ()
            },
            _ => ()
        }
    });

    keep_running
}

fn process_input(input: glutin::KeyboardInput, camera: &mut Camera) -> bool {
    let pressed = match input.state {
        ElementState::Pressed => true,
        ElementState::Released => false
    };
    if let Some(code) = input.virtual_keycode {
        match code {
            VirtualKeyCode::Escape if pressed => return false,
            VirtualKeyCode::W => camera.process_keyboard(FORWARD, pressed),
            VirtualKeyCode::S => camera.process_keyboard(BACKWARD, pressed),
            VirtualKeyCode::A => camera.process_keyboard(LEFT, pressed),
            VirtualKeyCode::D => camera.process_keyboard(RIGHT, pressed),
            _ => ()
        }
    }
    true
}

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

    #[test]
    fn print_struct_sizes() {
        // run with `cargo test -- --nocapture`
        println!("Sizes in bytes:");
        println!("Scene:     {:>3}", std::mem::size_of::<Scene>());
        println!("Node:      {:>3}", std::mem::size_of::<Node>());
        println!("Mesh:      {:>3}", std::mem::size_of::<Mesh>());
        println!("Primitive: {:>3}", std::mem::size_of::<Primitive>());
        println!("Vertex:    {:>3}", std::mem::size_of::<Vertex>());
        println!();
        println!("Option<String>: {:>3}", std::mem::size_of::<Option<String>>());
        println!("String:         {:>3}", std::mem::size_of::<String>());
        println!("Vec<f32>:       {:>3}", std::mem::size_of::<Vec<f32>>());
        println!("Vec<Node>:      {:>3}", std::mem::size_of::<Vec<Node>>());
    }

//     extern crate test;
//     use self::test::Bencher;
//     #[bench]
//     fn bench_frame_timer(b: &mut Bencher) {
//         let mut timer = FrameTimer::new("Foobar", 60);
//         b.iter(|| {
//             for _ in 0..60 {
//                 timer.start();
//                 timer.end();
//             }
//         })
//     }
}
<|endoftext|>