<commit_before>#[feature(globs)];
#[feature(macro_rules)];

extern mod argparse;

use argparse::ArgumentParser;
use argparse::arg::*;
use argparse::value::*;

mod macros;

macro_rules! parse_args(
    ($parser:ident: $($arg:expr),*) => {
        match $parser.parse_args(&[$(~$arg),*]) {
            Ok(args) => args,
            Err(e) => fail!(e.to_str()),
        }
    };

    ($parser:ident) => (parse_args!($parser:));
)

#[test]
fn test_no_args()
{
    let parser = ArgumentParser::new();

    let args = parse_args!(parser);
    assert_eq!(args.len(), 0);
}

#[test]
fn test_bool_opt()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("-b", "--bool"; ty = ArgTyBool));

    let args = parse_args!(parser);
    assert_eq!(args.get::<bool>("bool"), false);

    let args = parse_args!(parser: "-b");
    assert_eq!(args.get::<bool>("bool"), true);

    let args = parse_args!(parser: "--bool");
    assert_eq!(args.get::<bool>("bool"), true);
}

#[test]
fn test_str_opt()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("-s", "--str"; default = Some(ArgValString(~"default"))));

    let args = parse_args!(parser);
    assert_eq!(args.get::<~str>("str"), ~"default");

    let args = parse_args!(parser: "-s", "toto");
    assert_eq!(args.get::<~str>("str"), ~"toto");

    let args = parse_args!(parser: "--str", "toto");
    assert_eq!(args.get::<~str>("str"), ~"toto");
}

#[test]
#[should_fail]
fn test_str_opt_missing_required()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("-s", "--str"; required = true));

    parse_args!(parser);
}

#[test]
fn test_int_opt()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("-i", "--int"; ty = ArgTyInteger, default = Some(ArgValInteger(42))));

    let args = parse_args!(parser);
    assert_eq!(args.get::<int>("int"), 42);

    let args = parse_args!(parser: "-i", "32");
    assert_eq!(args.get::<int>("int"), 32);

    let args = parse_args!(parser: "--int", "32");
    assert_eq!(args.get::<int>("int"), 32);
}

#[test]
#[should_fail]
fn test_int_opt_bad_format()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("-i", "--int"; ty = ArgTyInteger));

    parse_args!(parser: "-i", "toto");
}

#[test]
#[should_fail]
fn test_int_opt_missing_required()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("-i", "--int"; ty = ArgTyInteger, required = true));

    parse_args!(parser);
}

#[test]
fn test_positional_str()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("pos"; default = Some(ArgValString(~"default"))));

    let args = parse_args!(parser);
    assert_eq!(args.get::<~str>("pos"), ~"default");

    let args = parse_args!(parser: "toto");
    assert_eq!(args.get::<~str>("pos"), ~"toto");
}

#[test]
#[should_fail]
fn test_positional_str_missing_no_default()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("pos"));

    parse_args!(parser);
}

#[test]
fn test_positional_int()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("pos"; ty = ArgTyInteger, default = Some(ArgValInteger(42))));

    let args = parse_args!(parser);
    assert_eq!(args.get::<int>("pos"), 42);

    let args = parse_args!(parser: "32");
    assert_eq!(args.get::<int>("pos"), 32);
}

#[test]
#[should_fail]
fn test_positional_int_bad_format()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("pos"; ty = ArgTyInteger));

    parse_args!(parser: "toto");
}
<commit_msg>Add unit test.<commit_after>#[feature(globs)];
#[feature(macro_rules)];

extern mod argparse;

use argparse::ArgumentParser;
use argparse::arg::*;
use argparse::value::*;

mod macros;

macro_rules! parse_args(
    ($parser:ident: $($arg:expr),*) => {
        match $parser.parse_args(&[$(~$arg),*]) {
            Ok(args) => args,
            Err(e) => fail!(e.to_str()),
        }
    };

    ($parser:ident) => (parse_args!($parser:));
)

#[test]
fn test_no_args()
{
    let parser = ArgumentParser::new();

    let args = parse_args!(parser);
    assert_eq!(args.len(), 0);
}

#[test]
fn test_bool_opt()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("-b", "--bool"; ty = ArgTyBool));

    let args = parse_args!(parser);
    assert_eq!(args.get::<bool>("bool"), false);

    let args = parse_args!(parser: "-b");
    assert_eq!(args.get::<bool>("bool"), true);

    let args = parse_args!(parser: "--bool");
    assert_eq!(args.get::<bool>("bool"), true);
}

#[test]
fn test_str_opt()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("-s", "--str"; default = Some(ArgValString(~"default"))));

    let args = parse_args!(parser);
    assert_eq!(args.get::<~str>("str"), ~"default");

    let args = parse_args!(parser: "-s", "toto");
    assert_eq!(args.get::<~str>("str"), ~"toto");

    let args = parse_args!(parser: "--str", "toto");
    assert_eq!(args.get::<~str>("str"), ~"toto");
}

#[test]
#[should_fail]
fn test_str_opt_missing_required()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("-s", "--str"; required = true));

    parse_args!(parser);
}

#[test]
fn test_int_opt()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("-i", "--int"; ty = ArgTyInteger, default = Some(ArgValInteger(42))));

    let args = parse_args!(parser);
    assert_eq!(args.get::<int>("int"), 42);

    let args = parse_args!(parser: "-i", "32");
    assert_eq!(args.get::<int>("int"), 32);

    let args = parse_args!(parser: "--int", "32");
    assert_eq!(args.get::<int>("int"), 32);
}

#[test]
#[should_fail]
fn test_int_opt_bad_format()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("-i", "--int"; ty = ArgTyInteger));

    parse_args!(parser: "-i", "toto");
}

#[test]
#[should_fail]
fn test_int_opt_missing_required()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("-i", "--int"; ty = ArgTyInteger, required = true));

    parse_args!(parser);
}

#[test]
fn test_positional_str()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("pos"; default = Some(ArgValString(~"default"))));

    let args = parse_args!(parser);
    assert_eq!(args.get::<~str>("pos"), ~"default");

    let args = parse_args!(parser: "toto");
    assert_eq!(args.get::<~str>("pos"), ~"toto");
}

#[test]
#[should_fail]
fn test_positional_str_missing_no_default()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("pos"));

    parse_args!(parser);
}

#[test]
fn test_positional_int()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("pos"; ty = ArgTyInteger, default = Some(ArgValInteger(42))));

    let args = parse_args!(parser);
    assert_eq!(args.get::<int>("pos"), 42);

    let args = parse_args!(parser: "32");
    assert_eq!(args.get::<int>("pos"), 32);
}

#[test]
#[should_fail]
fn test_positional_int_bad_format()
{
    let mut parser = ArgumentParser::new();
    parser.add_argument(create_arg!("pos"; ty = ArgTyInteger));

    parse_args!(parser: "toto");
}

#[test]
fn test_multiple_small_opts()
{
    let mut parser = ArgumentParser::new();

    let opts = ~[
        create_arg!("-a"; ty = ArgTyBool),
        create_arg!("-b"; ty = ArgTyBool),
        create_arg!("-c"; ty = ArgTyBool),
        create_arg!("-d"; ty = ArgTyBool),
        create_arg!("-e"; ty = ArgTyBool),
    ];
    parser.add_arguments(opts);

    let args = parse_args!(parser);
    assert_eq!(args.get::<bool>("a"), false);
    assert_eq!(args.get::<bool>("b"), false);
    assert_eq!(args.get::<bool>("c"), false);
    assert_eq!(args.get::<bool>("d"), false);
    assert_eq!(args.get::<bool>("e"), false);

    let args = parse_args!(parser: "-a", "-b", "-c", "-d", "-e");
    assert_eq!(args.get::<bool>("a"), true);
    assert_eq!(args.get::<bool>("b"), true);
    assert_eq!(args.get::<bool>("c"), true);
    assert_eq!(args.get::<bool>("d"), true);
    assert_eq!(args.get::<bool>("e"), true);

    let args = parse_args!(parser: "-a", "-b", "-e");
    assert_eq!(args.get::<bool>("a"), true);
    assert_eq!(args.get::<bool>("b"), true);
    assert_eq!(args.get::<bool>("c"), false);
    assert_eq!(args.get::<bool>("d"), false);
    assert_eq!(args.get::<bool>("e"), true);

    let args = parse_args!(parser: "-abcde");
    assert_eq!(args.get::<bool>("a"), true);
    assert_eq!(args.get::<bool>("b"), true);
    assert_eq!(args.get::<bool>("c"), true);
    assert_eq!(args.get::<bool>("d"), true);
    assert_eq!(args.get::<bool>("e"), true);
}
<|endoftext|>


<commit_before><commit_msg>Name trait function variables<commit_after><|endoftext|>


<commit_before>use clap::crate_version;

use std::env;
use std::ffi::OsStr;
use std::io;
use std::process::{self, Command, ExitStatus};
use std::sync::{Arc, Condvar, Mutex};
use std::thread;
use std::time::{Duration, Instant};

#[path = "../common.rs"]
mod common;

fn main() {
    if !common::archive_file_path().exists() {
        let version = common::afl_rustc_version();
        eprintln!(
            "AFL LLVM runtime is not built with Rust {}, run `cargo \
             install --force afl` to build it.",
            version
        );
        process::exit(1);
    }

    let app_matches = clap_app().get_matches();
    // This unwrap is okay because we set SubcommandRequiredElseHelp at the top level, and afl is
    // the only subcommand
    let afl_matches = app_matches.subcommand_matches("afl").unwrap();

    match afl_matches.subcommand() {
        Some(("analyze", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-analyze args")
                .unwrap_or_default();
            run_afl(args, "afl-analyze", None);
        }
        Some(("cmin", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-cmin args")
                .unwrap_or_default();
            run_afl(args, "afl-cmin", None);
        }
        Some(("fuzz", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-fuzz args")
                .unwrap_or_default();
            let timeout = sub_matches.value_of("max_total_time").map(|_| {
                sub_matches
                    .value_of_t::<u64>("max_total_time")
                    .unwrap_or_else(|e| e.exit())
            });
            run_afl(args, "afl-fuzz", timeout);
        }
        Some(("gotcpu", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-gotcpu args")
                .unwrap_or_default();
            run_afl(args, "afl-gotcpu", None);
        }
        Some(("plot", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-plot args")
                .unwrap_or_default();
            run_afl(args, "afl-plot", None);
        }
        Some(("showmap", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-showmap args")
                .unwrap_or_default();
            run_afl(args, "afl-showmap", None);
        }
        Some(("tmin", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-tmin args")
                .unwrap_or_default();
            run_afl(args, "afl-tmin", None);
        }
        Some(("whatsup", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-whatsup args")
                .unwrap_or_default();
            run_afl(args, "afl-whatsup", None);
        }
        Some((subcommand, sub_matches)) => {
            let args = sub_matches.values_of_os("").unwrap_or_default();
            run_cargo(subcommand, args);
        }
        // unreachable due to SubcommandRequiredElseHelp on "afl" subcommand
        None => unreachable!(),
    }
}

fn clap_app() -> clap::App<'static> {
    use clap::{
        App,
        AppSettings::{
            AllowExternalSubcommands, AllowHyphenValues, AllowInvalidUtf8ForExternalSubcommands,
            DisableHelpFlag, DisableHelpSubcommand, DisableVersionFlag, SubcommandRequiredElseHelp,
        },
        Arg,
    };

    App::new("cargo afl")
        .bin_name("cargo")
        .setting(SubcommandRequiredElseHelp)
        .subcommand(
            App::new("afl")
                .version(crate_version!())
                .setting(SubcommandRequiredElseHelp)
                .setting(AllowExternalSubcommands)
                .setting(AllowInvalidUtf8ForExternalSubcommands)
                .override_usage("cargo afl [SUBCOMMAND or Cargo SUBCOMMAND]")
                .after_help(
                    "In addition to the subcommands above, Cargo subcommands are also \
                 supported (see `cargo help` for a list of all Cargo subcommands).",
                )
                .subcommand(
                    App::new("analyze")
                        .about("Invoke afl-analyze")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-analyze args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("cmin")
                        .about("Invoke afl-cmin")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-cmin args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("fuzz")
                        .about("Invoke afl-fuzz")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("max_total_time")
                                .long("max_total_time")
                                .takes_value(true)
                                .help("Maximum amount of time to run the fuzzer"),
                        )
                        .arg(
                            Arg::new("afl-fuzz args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("gotcpu")
                        .about("Invoke afl-gotcpu")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-gotcpu args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("plot")
                        .about("Invoke afl-plot")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-plot args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("showmap")
                        .about("Invoke afl-showmap")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-showmap args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("tmin")
                        .about("Invoke afl-tmin")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-tmin args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("whatsup")
                        .about("Invoke afl-whatsup")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-whatsup args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                ),
        )
}

fn run_timeout_terminate(mut cmd: Command, timeout: Option<u64>) -> Result<ExitStatus, io::Error> {
    let timeout = match timeout {
        Some(timeout) => Duration::from_secs(timeout),
        None => return cmd.status(),
    };
    let start_time = Instant::now();

    let mut child = cmd.spawn()?;
    let pid = child.id();

    let pair = Arc::new((Mutex::new(false), Condvar::new()));
    let (stop_mutex, condvar) = &*pair;
    let thread_handle = {
        let pair = pair.clone();
        thread::spawn(move || -> Result<(), io::Error> {
            // This thread will wait until the child process has exited, or the
            // timeout has elapsed, whichever comes first. If the timeout
            // elapses, and the process is still running, it will send SIGTERM
            // to the child process.

            let (stop_mutex, condvar) = &*pair;
            let mut stop = stop_mutex.lock().unwrap();
            loop {
                let elapsed = Instant::now() - start_time;
                if elapsed >= timeout {
                    break;
                }

                let dur = timeout - elapsed;
                let results = condvar.wait_timeout(stop, dur).unwrap();
                stop = results.0;
                if *stop {
                    // Blocking waitid call on the main thread has returned,
                    // thus the child process has terminated
                    return Ok(());
                }
                if results.1.timed_out() {
                    break;
                }
            }

            // Since the waitid call on the main thread is using WNOWAIT, the
            // child process won't be cleaned up (until after this thread
            // exits and the main thread calls wait) and thus its PID won't be
            // reused by another, unrelated process.
            unsafe {
                let ret = libc::kill(pid as i32, libc::SIGTERM);
                if ret == -1 {
                    Err(io::Error::last_os_error())?
                }
            }

            Ok(())
        })
    };

    unsafe {
        // Block until the child process terminates, but leave it in a waitable
        // state still
        let ret = libc::waitid(
            libc::P_PID,
            pid,
            std::ptr::null_mut(),
            libc::WEXITED | libc::WNOWAIT,
        );
        if ret == -1 {
            Err(io::Error::last_os_error())?
        }
    }
    {
        let mut stop = stop_mutex.lock().unwrap();
        *stop = true;
    }
    // Tell the timeout thread to stop, wake it, and wait for it to exit
    condvar.notify_one();
    thread_handle.join().unwrap()?;

    // Clean up zombie and get exit status (this won't block, because the child
    // process has terminated and is still waitable)
    child.wait()
}

fn run_afl<I, S>(args: I, cmd: &str, timeout: Option<u64>)
where
    I: IntoIterator<Item = S>,
    S: AsRef<OsStr>,
{
    let cmd_path = common::afl_dir().join("bin").join(cmd);
    let mut cmd = Command::new(cmd_path);
    cmd.args(args);
    let status = run_timeout_terminate(cmd, timeout).unwrap();
    process::exit(status.code().unwrap_or(1));
}

fn run_cargo<I, S>(subcommand: &str, args: I)
where
    I: IntoIterator<Item = S>,
    S: AsRef<OsStr>,
{
    let cargo_path = env::var("CARGO").expect("Could not determine `cargo` path");

    // add some flags to sanitizers to make them work with Rust code
    let asan_options = env::var("ASAN_OPTIONS").unwrap_or_default();
    let asan_options = format!(
        "detect_odr_violation=0:abort_on_error=1:symbolize=0:{}",
        asan_options
    );

    let tsan_options = env::var("TSAN_OPTIONS").unwrap_or_default();
    let tsan_options = format!("report_signal_unsafe=0:{}", tsan_options);

    // The new LLVM pass manager was not enabled in rustc 1.57 as expected:
    // https://github.com/rust-lang/rust/pull/91263
    // The fix for now is to pass `-C passes=sancov-module` only to nightly
    // compilers for which the LLVM version is >= 13.

    let version_meta = rustc_version::version_meta().unwrap();
    let passes = if is_nightly() && version_meta.llvm_version.map_or(true, |v| v.major >= 13) {
        "sancov-module"
    } else {
        "sancov"
    };

    // `-C codegen-units=1` is needed to work around link errors
    // https://github.com/rust-fuzz/afl.rs/pull/193#issuecomment-933550430
    let mut rustflags = format!(
        "--cfg fuzzing \
         -C debug-assertions \
         -C overflow_checks \
         -C passes={} \
         -C codegen-units=1 \
         -C llvm-args=-sanitizer-coverage-level=3 \
         -C llvm-args=-sanitizer-coverage-trace-pc-guard \
         -C llvm-args=-sanitizer-coverage-prune-blocks=0 \
         -C opt-level=3 \
         -C target-cpu=native \
         -C debuginfo=0 \
         -l afl-llvm-rt \
         -L {} ",
        passes,
        common::afl_llvm_rt_dir().display()
    );

    // RUSTFLAGS are not used by rustdoc, instead RUSTDOCFLAGS are used. Since
    // doctests will try to link against afl-llvm-rt, set up RUSTDOCFLAGS to
    // have doctests built the same as other code to avoid issues with doctests.
    let mut rustdocflags = format!(
        "--cfg fuzzing \
         -C debug-assertions \
         -C overflow_checks \
         -C passes={} \
         -C codegen-units=1 \
         -C llvm-args=-sanitizer-coverage-level=3 \
         -C llvm-args=-sanitizer-coverage-trace-pc-guard \
         -C llvm-args=-sanitizer-coverage-prune-blocks=0 \
         -C opt-level=3 \
         -C target-cpu=native \
         -C debuginfo=0 \
         -L {} ",
        passes,
        common::afl_llvm_rt_dir().display()
    );

    if cfg!(target_os = "linux") {
        // work around https://github.com/rust-fuzz/afl.rs/issues/141 /
        // https://github.com/rust-lang/rust/issues/53945, can be removed once
        // those are fixed.
        rustflags.push_str("-Clink-arg=-fuse-ld=gold ");
        rustdocflags.push_str("-Clink-arg=-fuse-ld=gold ");
    }

    // add user provided flags
    rustflags.push_str(&env::var("RUSTFLAGS").unwrap_or_default());
    rustdocflags.push_str(&env::var("RUSTDOCFLAGS").unwrap_or_default());

    let status = Command::new(cargo_path)
        .arg(subcommand)
        .args(args)
        .env("RUSTFLAGS", &rustflags)
        .env("RUSTDOCFLAGS", &rustdocflags)
        .env("ASAN_OPTIONS", asan_options)
        .env("TSAN_OPTIONS", tsan_options)
        .status()
        .unwrap();
    process::exit(status.code().unwrap_or(1));
}

fn is_nightly() -> bool {
    Command::new("rustc")
        .args(&["-Z", "help"])
        .status()
        .unwrap()
        .success()
}
<commit_msg>Add `clap::App` unit test<commit_after>use clap::crate_version;

use std::env;
use std::ffi::OsStr;
use std::io;
use std::process::{self, Command, ExitStatus};
use std::sync::{Arc, Condvar, Mutex};
use std::thread;
use std::time::{Duration, Instant};

#[path = "../common.rs"]
mod common;

fn main() {
    if !common::archive_file_path().exists() {
        let version = common::afl_rustc_version();
        eprintln!(
            "AFL LLVM runtime is not built with Rust {}, run `cargo \
             install --force afl` to build it.",
            version
        );
        process::exit(1);
    }

    let app_matches = clap_app().get_matches();
    // This unwrap is okay because we set SubcommandRequiredElseHelp at the top level, and afl is
    // the only subcommand
    let afl_matches = app_matches.subcommand_matches("afl").unwrap();

    match afl_matches.subcommand() {
        Some(("analyze", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-analyze args")
                .unwrap_or_default();
            run_afl(args, "afl-analyze", None);
        }
        Some(("cmin", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-cmin args")
                .unwrap_or_default();
            run_afl(args, "afl-cmin", None);
        }
        Some(("fuzz", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-fuzz args")
                .unwrap_or_default();
            let timeout = sub_matches.value_of("max_total_time").map(|_| {
                sub_matches
                    .value_of_t::<u64>("max_total_time")
                    .unwrap_or_else(|e| e.exit())
            });
            run_afl(args, "afl-fuzz", timeout);
        }
        Some(("gotcpu", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-gotcpu args")
                .unwrap_or_default();
            run_afl(args, "afl-gotcpu", None);
        }
        Some(("plot", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-plot args")
                .unwrap_or_default();
            run_afl(args, "afl-plot", None);
        }
        Some(("showmap", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-showmap args")
                .unwrap_or_default();
            run_afl(args, "afl-showmap", None);
        }
        Some(("tmin", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-tmin args")
                .unwrap_or_default();
            run_afl(args, "afl-tmin", None);
        }
        Some(("whatsup", sub_matches)) => {
            let args = sub_matches
                .values_of_os("afl-whatsup args")
                .unwrap_or_default();
            run_afl(args, "afl-whatsup", None);
        }
        Some((subcommand, sub_matches)) => {
            let args = sub_matches.values_of_os("").unwrap_or_default();
            run_cargo(subcommand, args);
        }
        // unreachable due to SubcommandRequiredElseHelp on "afl" subcommand
        None => unreachable!(),
    }
}

fn clap_app() -> clap::App<'static> {
    use clap::{
        App,
        AppSettings::{
            AllowExternalSubcommands, AllowHyphenValues, AllowInvalidUtf8ForExternalSubcommands,
            DisableHelpFlag, DisableHelpSubcommand, DisableVersionFlag, SubcommandRequiredElseHelp,
        },
        Arg,
    };

    App::new("cargo afl")
        .bin_name("cargo")
        .setting(SubcommandRequiredElseHelp)
        .subcommand(
            App::new("afl")
                .version(crate_version!())
                .setting(SubcommandRequiredElseHelp)
                .setting(AllowExternalSubcommands)
                .setting(AllowInvalidUtf8ForExternalSubcommands)
                .override_usage("cargo afl [SUBCOMMAND or Cargo SUBCOMMAND]")
                .after_help(
                    "In addition to the subcommands above, Cargo subcommands are also \
                 supported (see `cargo help` for a list of all Cargo subcommands).",
                )
                .subcommand(
                    App::new("analyze")
                        .about("Invoke afl-analyze")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-analyze args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("cmin")
                        .about("Invoke afl-cmin")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-cmin args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("fuzz")
                        .about("Invoke afl-fuzz")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("max_total_time")
                                .long("max_total_time")
                                .takes_value(true)
                                .help("Maximum amount of time to run the fuzzer"),
                        )
                        .arg(
                            Arg::new("afl-fuzz args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("gotcpu")
                        .about("Invoke afl-gotcpu")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-gotcpu args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("plot")
                        .about("Invoke afl-plot")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-plot args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("showmap")
                        .about("Invoke afl-showmap")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-showmap args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("tmin")
                        .about("Invoke afl-tmin")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-tmin args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                )
                .subcommand(
                    App::new("whatsup")
                        .about("Invoke afl-whatsup")
                        .setting(AllowHyphenValues)
                        .setting(DisableHelpSubcommand)
                        .setting(DisableHelpFlag)
                        .setting(DisableVersionFlag)
                        .arg(
                            Arg::new("afl-whatsup args")
                                .allow_invalid_utf8(true)
                                .multiple_values(true),
                        ),
                ),
        )
}

fn run_timeout_terminate(mut cmd: Command, timeout: Option<u64>) -> Result<ExitStatus, io::Error> {
    let timeout = match timeout {
        Some(timeout) => Duration::from_secs(timeout),
        None => return cmd.status(),
    };
    let start_time = Instant::now();

    let mut child = cmd.spawn()?;
    let pid = child.id();

    let pair = Arc::new((Mutex::new(false), Condvar::new()));
    let (stop_mutex, condvar) = &*pair;
    let thread_handle = {
        let pair = pair.clone();
        thread::spawn(move || -> Result<(), io::Error> {
            // This thread will wait until the child process has exited, or the
            // timeout has elapsed, whichever comes first. If the timeout
            // elapses, and the process is still running, it will send SIGTERM
            // to the child process.

            let (stop_mutex, condvar) = &*pair;
            let mut stop = stop_mutex.lock().unwrap();
            loop {
                let elapsed = Instant::now() - start_time;
                if elapsed >= timeout {
                    break;
                }

                let dur = timeout - elapsed;
                let results = condvar.wait_timeout(stop, dur).unwrap();
                stop = results.0;
                if *stop {
                    // Blocking waitid call on the main thread has returned,
                    // thus the child process has terminated
                    return Ok(());
                }
                if results.1.timed_out() {
                    break;
                }
            }

            // Since the waitid call on the main thread is using WNOWAIT, the
            // child process won't be cleaned up (until after this thread
            // exits and the main thread calls wait) and thus its PID won't be
            // reused by another, unrelated process.
            unsafe {
                let ret = libc::kill(pid as i32, libc::SIGTERM);
                if ret == -1 {
                    Err(io::Error::last_os_error())?
                }
            }

            Ok(())
        })
    };

    unsafe {
        // Block until the child process terminates, but leave it in a waitable
        // state still
        let ret = libc::waitid(
            libc::P_PID,
            pid,
            std::ptr::null_mut(),
            libc::WEXITED | libc::WNOWAIT,
        );
        if ret == -1 {
            Err(io::Error::last_os_error())?
        }
    }
    {
        let mut stop = stop_mutex.lock().unwrap();
        *stop = true;
    }
    // Tell the timeout thread to stop, wake it, and wait for it to exit
    condvar.notify_one();
    thread_handle.join().unwrap()?;

    // Clean up zombie and get exit status (this won't block, because the child
    // process has terminated and is still waitable)
    child.wait()
}

fn run_afl<I, S>(args: I, cmd: &str, timeout: Option<u64>)
where
    I: IntoIterator<Item = S>,
    S: AsRef<OsStr>,
{
    let cmd_path = common::afl_dir().join("bin").join(cmd);
    let mut cmd = Command::new(cmd_path);
    cmd.args(args);
    let status = run_timeout_terminate(cmd, timeout).unwrap();
    process::exit(status.code().unwrap_or(1));
}

fn run_cargo<I, S>(subcommand: &str, args: I)
where
    I: IntoIterator<Item = S>,
    S: AsRef<OsStr>,
{
    let cargo_path = env::var("CARGO").expect("Could not determine `cargo` path");

    // add some flags to sanitizers to make them work with Rust code
    let asan_options = env::var("ASAN_OPTIONS").unwrap_or_default();
    let asan_options = format!(
        "detect_odr_violation=0:abort_on_error=1:symbolize=0:{}",
        asan_options
    );

    let tsan_options = env::var("TSAN_OPTIONS").unwrap_or_default();
    let tsan_options = format!("report_signal_unsafe=0:{}", tsan_options);

    // The new LLVM pass manager was not enabled in rustc 1.57 as expected:
    // https://github.com/rust-lang/rust/pull/91263
    // The fix for now is to pass `-C passes=sancov-module` only to nightly
    // compilers for which the LLVM version is >= 13.

    let version_meta = rustc_version::version_meta().unwrap();
    let passes = if is_nightly() && version_meta.llvm_version.map_or(true, |v| v.major >= 13) {
        "sancov-module"
    } else {
        "sancov"
    };

    // `-C codegen-units=1` is needed to work around link errors
    // https://github.com/rust-fuzz/afl.rs/pull/193#issuecomment-933550430
    let mut rustflags = format!(
        "--cfg fuzzing \
         -C debug-assertions \
         -C overflow_checks \
         -C passes={} \
         -C codegen-units=1 \
         -C llvm-args=-sanitizer-coverage-level=3 \
         -C llvm-args=-sanitizer-coverage-trace-pc-guard \
         -C llvm-args=-sanitizer-coverage-prune-blocks=0 \
         -C opt-level=3 \
         -C target-cpu=native \
         -C debuginfo=0 \
         -l afl-llvm-rt \
         -L {} ",
        passes,
        common::afl_llvm_rt_dir().display()
    );

    // RUSTFLAGS are not used by rustdoc, instead RUSTDOCFLAGS are used. Since
    // doctests will try to link against afl-llvm-rt, set up RUSTDOCFLAGS to
    // have doctests built the same as other code to avoid issues with doctests.
    let mut rustdocflags = format!(
        "--cfg fuzzing \
         -C debug-assertions \
         -C overflow_checks \
         -C passes={} \
         -C codegen-units=1 \
         -C llvm-args=-sanitizer-coverage-level=3 \
         -C llvm-args=-sanitizer-coverage-trace-pc-guard \
         -C llvm-args=-sanitizer-coverage-prune-blocks=0 \
         -C opt-level=3 \
         -C target-cpu=native \
         -C debuginfo=0 \
         -L {} ",
        passes,
        common::afl_llvm_rt_dir().display()
    );

    if cfg!(target_os = "linux") {
        // work around https://github.com/rust-fuzz/afl.rs/issues/141 /
        // https://github.com/rust-lang/rust/issues/53945, can be removed once
        // those are fixed.
        rustflags.push_str("-Clink-arg=-fuse-ld=gold ");
        rustdocflags.push_str("-Clink-arg=-fuse-ld=gold ");
    }

    // add user provided flags
    rustflags.push_str(&env::var("RUSTFLAGS").unwrap_or_default());
    rustdocflags.push_str(&env::var("RUSTDOCFLAGS").unwrap_or_default());

    let status = Command::new(cargo_path)
        .arg(subcommand)
        .args(args)
        .env("RUSTFLAGS", &rustflags)
        .env("RUSTDOCFLAGS", &rustdocflags)
        .env("ASAN_OPTIONS", asan_options)
        .env("TSAN_OPTIONS", tsan_options)
        .status()
        .unwrap();
    process::exit(status.code().unwrap_or(1));
}

fn is_nightly() -> bool {
    Command::new("rustc")
        .args(&["-Z", "help"])
        .status()
        .unwrap()
        .success()
}

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

    #[test]
    fn test_app() {
        clap_app().debug_assert();
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

#![feature(plugin)]
#![cfg_attr(feature = "dev", plugin(clippy))]

extern crate tikv;
extern crate getopts;
extern crate protobuf;
extern crate kvproto;
extern crate rocksdb;

use std::{env, str, u64};
use getopts::Options;
use protobuf::Message;
use kvproto::raft_cmdpb::RaftCmdRequest;
use kvproto::metapb::Region;
use kvproto::eraftpb::Entry;
use rocksdb::DB;
use tikv::util::{self, escape, unescape};
use tikv::raftstore::store::keys;
use tikv::raftstore::store::engine::{Peekable, Iterable};
use tikv::storage::DEFAULT_CFS;

/// # Message dump tool
///
/// A simple tool that dump the message from rocksdb directory. Very useful when you want
/// to take a deep look in the data.

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

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optopt("", "db", "set rocksdb path, required", "");
    opts.optopt("r",
                "region-id",
                "set the region id",
                "required when getting raft message");
    opts.optflag("h", "help", "print this help menu");
    opts.optflag("", "info", "print the region info");
    opts.optopt("i", "index", "set the raft log index", "");
    opts.optopt("k", "key", "set the query raw key, in escape format", "");
    opts.optopt("f",
                "from",
                "set the scan from raw key, in escaped format",
                "");
    opts.optopt("t", "to", "set the scan end raw key, in escaped format", "");
    opts.optopt("l", "limit", "set the scan limit", "");
    opts.optopt("c",
                "cf",
                "column family name, only avialbe for dump-range",
                "");
    let matches = opts.parse(&args[1..]).expect("opts parse failed");
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }

    let db_str = matches.opt_str("db").unwrap();
    let db = util::rocksdb::open(&db_str, DEFAULT_CFS).unwrap();
    let key = matches.opt_str("k");
    let from = matches.opt_str("f");
    let to = matches.opt_str("t");
    let limit = matches.opt_str("l").map(|s| s.parse().unwrap());
    let idx = matches.opt_str("i");
    let region = matches.opt_str("r");
    let cf = matches.opt_str("c");
    let cf_name = cf.as_ref().map_or("default", |s| s.as_str());
    if let Some(key) = key {
        dump_raw_value(db, cf_name, key);
    } else if let Some(idx) = idx {
        dump_raft_log_entry(db, region.unwrap(), idx);
    } else if matches.opt_present("info") {
        dump_region_info(db, region.unwrap());
    } else if let Some(from) = from {
        dump_range(db, from, to, limit, cf_name);
    } else {
        panic!("currently only random key-value and raft log entry query are supported.");
    }
}

fn dump_raw_value(db: DB, cf: &str, key: String) {
    let key = unescape(&key);
    let value = db.get_value_cf(cf, &key).unwrap();
    println!("value: {}", value.map_or("None".to_owned(), |v| escape(&v)));
}

fn dump_raft_log_entry(db: DB, region_id_str: String, idx_str: String) {
    let region_id = u64::from_str_radix(&region_id_str, 10).unwrap();
    let idx = u64::from_str_radix(&idx_str, 10).unwrap();

    let idx_key = keys::raft_log_key(region_id, idx);
    println!("idx_key: {}", escape(&idx_key));
    let mut ent: Entry = db.get_msg(&idx_key).unwrap().unwrap();
    let data = ent.take_data();
    println!("entry {:?}", ent);
    let mut msg = RaftCmdRequest::new();
    msg.merge_from_bytes(&data).unwrap();
    println!("msg {:?}", msg);
}

fn dump_region_info(db: DB, region_id_str: String) {
    let region_id = u64::from_str_radix(&region_id_str, 10).unwrap();
    let region_state_key = keys::region_state_key(region_id);
    println!("info_key: {}", escape(&region_state_key));
    let region: Option<Region> = db.get_msg(&region_state_key).unwrap();
    println!("info: {:?}", region);
}

fn dump_range(db: DB, from: String, to: Option<String>, limit: Option<u64>, cf: &str) {
    let from = unescape(&from);
    let to = to.map_or_else(|| vec![0xff], |s| unescape(&s));
    let limit = limit.unwrap_or(u64::MAX);

    if limit == 0 {
        return;
    }

    let mut cnt = 0;
    db.scan_cf(cf,
                 &from,
                 &to,
                 &mut |k, v| {
                     println!("key: {}, value: {}", escape(k), escape(v));
                     cnt += 1;
                     Ok(cnt < limit)
                 })
        .unwrap()
}
<commit_msg>bin: make region info dump all state (#882)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

#![feature(plugin)]
#![cfg_attr(feature = "dev", plugin(clippy))]

extern crate tikv;
extern crate getopts;
extern crate protobuf;
extern crate kvproto;
extern crate rocksdb;

use std::{env, str, u64};
use getopts::Options;
use protobuf::Message;
use kvproto::raft_cmdpb::RaftCmdRequest;
use kvproto::raft_serverpb::{RaftLocalState, RegionLocalState, RaftApplyState};
use kvproto::eraftpb::Entry;
use rocksdb::DB;
use tikv::util::{self, escape, unescape};
use tikv::raftstore::store::keys;
use tikv::raftstore::store::engine::{Peekable, Iterable};
use tikv::storage::DEFAULT_CFS;

/// # Message dump tool
///
/// A simple tool that dump the message from rocksdb directory. Very useful when you want
/// to take a deep look in the data.

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

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optopt("", "db", "set rocksdb path, required", "");
    opts.optopt("r",
                "region-id",
                "set the region id",
                "required when getting raft message");
    opts.optflag("h", "help", "print this help menu");
    opts.optflag("", "info", "print the region info");
    opts.optopt("i", "index", "set the raft log index", "");
    opts.optopt("k", "key", "set the query raw key, in escape format", "");
    opts.optopt("f",
                "from",
                "set the scan from raw key, in escaped format",
                "");
    opts.optopt("t", "to", "set the scan end raw key, in escaped format", "");
    opts.optopt("l", "limit", "set the scan limit", "");
    opts.optopt("c",
                "cf",
                "column family name, only avialbe for dump-range",
                "");
    let matches = opts.parse(&args[1..]).expect("opts parse failed");
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }

    let db_str = matches.opt_str("db").unwrap();
    let db = util::rocksdb::open(&db_str, DEFAULT_CFS).unwrap();
    let key = matches.opt_str("k");
    let from = matches.opt_str("f");
    let to = matches.opt_str("t");
    let limit = matches.opt_str("l").map(|s| s.parse().unwrap());
    let idx = matches.opt_str("i");
    let region = matches.opt_str("r");
    let cf = matches.opt_str("c");
    let cf_name = cf.as_ref().map_or("default", |s| s.as_str());
    if let Some(key) = key {
        dump_raw_value(db, cf_name, key);
    } else if let Some(idx) = idx {
        dump_raft_log_entry(db, region.unwrap(), idx);
    } else if matches.opt_present("info") {
        dump_region_info(db, region.unwrap());
    } else if let Some(from) = from {
        dump_range(db, from, to, limit, cf_name);
    } else {
        panic!("currently only random key-value and raft log entry query are supported.");
    }
}

fn dump_raw_value(db: DB, cf: &str, key: String) {
    let key = unescape(&key);
    let value = db.get_value_cf(cf, &key).unwrap();
    println!("value: {}", value.map_or("None".to_owned(), |v| escape(&v)));
}

fn dump_raft_log_entry(db: DB, region_id_str: String, idx_str: String) {
    let region_id = u64::from_str_radix(&region_id_str, 10).unwrap();
    let idx = u64::from_str_radix(&idx_str, 10).unwrap();

    let idx_key = keys::raft_log_key(region_id, idx);
    println!("idx_key: {}", escape(&idx_key));
    let mut ent: Entry = db.get_msg(&idx_key).unwrap().unwrap();
    let data = ent.take_data();
    println!("entry {:?}", ent);
    let mut msg = RaftCmdRequest::new();
    msg.merge_from_bytes(&data).unwrap();
    println!("msg {:?}", msg);
}

fn dump_region_info(db: DB, region_id_str: String) {
    let region_id = u64::from_str_radix(&region_id_str, 10).unwrap();

    let region_state_key = keys::region_state_key(region_id);
    println!("region state key: {}", escape(&region_state_key));
    let region_state: Option<RegionLocalState> = db.get_msg(&region_state_key).unwrap();
    println!("region state: {:?}", region_state);

    let raft_state_key = keys::raft_state_key(region_id);
    println!("raft state key: {}", escape(&raft_state_key));
    let raft_state: Option<RaftLocalState> = db.get_msg(&raft_state_key).unwrap();
    println!("raft state: {:?}", raft_state);

    let apply_state_key = keys::apply_state_key(region_id);
    println!("apply state key: {}", escape(&apply_state_key));
    let apply_state: Option<RaftApplyState> = db.get_msg(&apply_state_key).unwrap();
    println!("apply state: {:?}", apply_state);
}

fn dump_range(db: DB, from: String, to: Option<String>, limit: Option<u64>, cf: &str) {
    let from = unescape(&from);
    let to = to.map_or_else(|| vec![0xff], |s| unescape(&s));
    let limit = limit.unwrap_or(u64::MAX);

    if limit == 0 {
        return;
    }

    let mut cnt = 0;
    db.scan_cf(cf,
                 &from,
                 &to,
                 &mut |k, v| {
                     println!("key: {}, value: {}", escape(k), escape(v));
                     cnt += 1;
                     Ok(cnt < limit)
                 })
        .unwrap()
}
<|endoftext|>


<commit_before>extern crate dcpu16;
extern crate getopts;

use std::io::prelude::*;
use std::io::BufReader;
use std::fs::File;
use std::path::Path;
use std::env;
use getopts::Options;
use dcpu16::assembler;

fn main() {
    let opts = Options::new();
    let args: Vec<String> = env::args().collect();
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m },
        Err(f) => { panic!(f.to_string()) },
    };

    if matches.free.len() != 1 {
        println!("Please input file");
        return;
    }
    let ref filename = matches.free[0];
    let mut lines: Vec<String> = Vec::new();

    let x: &[_] = &[' ', '\n', '\t'];
    let path = Path::new(filename);
    let file = match File::open(&path) {
        Err(_) => panic!(""),//Could not open file {}: {}", path.display(), Exception::description(&why)),
        Ok(file) => file,
    };
    for line in BufReader::new(&file).lines() {
        match line {
            Ok(s) => {lines.push(s.trim_matches(x).to_string())},
            Err(_) => {},
        }
    }
    let mut cpu = assembler::PCPU::new();
    let mut line_no = 0usize;
    for line in lines.iter() {
        let l = &line[..];

        if l.len() == 0 {
            line_no += 1;
            continue;
        }

        println!("\x1b[1;37m{}\x1b[0m", l);
        match assembler::tokenize(line_no, l, &mut cpu) {
            Ok(tokens) => {
                if tokens.len() > 0 {
                    for t in tokens.iter() {
                        println!("{} \x1b[1;30m{}:{}\x1b[0m", t.ttype, t.col, t.col+t.len);
                    }
                    println!("");
                }
            },
            Err(err) => {
                assembler::print_parse_error(&cpu, l, err);
            },
        };

        line_no += 1;
    }
}
<commit_msg>Better error handling in tokenizer.<commit_after>extern crate dcpu16;
extern crate getopts;

use std::io::prelude::*;
use std::io::BufReader;
use std::fs::File;
use std::path::Path;
use std::env;
use getopts::Options;
use dcpu16::assembler;

fn main() {
    let opts = Options::new();
    let args: Vec<String> = env::args().collect();
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m },
        Err(why) => {
            println!("{}", why);
            return;
        },
    };

    if matches.free.len() != 1 {
        println!("Please input file");
        return;
    }
    let ref filename = matches.free[0];
    let mut lines: Vec<String> = Vec::new();

    let x: &[_] = &[' ', '\n', '\t'];
    let path = Path::new(filename);
    let file = match File::open(&path) {
        Err(why) => {
            println!("Could load file {}: {}", path.display(), why);
            return;
        },
        Ok(file) => file,
    };
    for line in BufReader::new(&file).lines() {
        match line {
            Ok(s) => {lines.push(s.trim_matches(x).to_string())},
            Err(_) => {},
        }
    }
    let mut cpu = assembler::PCPU::new();
    let mut line_no = 0usize;
    for line in lines.iter() {
        let l = &line[..];

        if l.len() == 0 {
            line_no += 1;
            continue;
        }

        println!("\x1b[1;37m{}\x1b[0m", l);
        match assembler::tokenize(line_no, l, &mut cpu) {
            Ok(tokens) => {
                if tokens.len() > 0 {
                    for t in tokens.iter() {
                        println!("{} \x1b[1;30m{}:{}\x1b[0m", t.ttype, t.col, t.col+t.len);
                    }
                    println!("");
                }
            },
            Err(err) => {
                assembler::print_parse_error(&cpu, l, err);
            },
        };

        line_no += 1;
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! TUI peer display

use std::cmp::Ordering;

use crate::servers::{PeerStats, ServerStats};

use crate::tui::humansize::{file_size_opts::CONVENTIONAL, FileSize};
use chrono::prelude::*;

use cursive::direction::Orientation;
use cursive::event::Key;
use cursive::traits::{Boxable, Identifiable};
use cursive::view::View;
use cursive::views::{BoxView, Dialog, LinearLayout, OnEventView, TextView};
use cursive::Cursive;

use crate::tui::constants::{MAIN_MENU, TABLE_PEER_STATUS, VIEW_PEER_SYNC};
use crate::tui::table::{TableView, TableViewItem};
use crate::tui::types::TUIStatusListener;

#[derive(Copy, Clone, PartialEq, Eq, Hash)]
enum PeerColumn {
	Address,
	State,
	UsedBandwidth,
	TotalDifficulty,
	Direction,
	Version,
	UserAgent,
}

impl PeerColumn {
	fn _as_str(&self) -> &str {
		match *self {
			PeerColumn::Address => "Address",
			PeerColumn::State => "State",
			PeerColumn::UsedBandwidth => "Used bandwidth",
			PeerColumn::Version => "Version",
			PeerColumn::TotalDifficulty => "Total Difficulty",
			PeerColumn::Direction => "Direction",
			PeerColumn::UserAgent => "User Agent",
		}
	}
}

impl TableViewItem<PeerColumn> for PeerStats {
	fn to_column(&self, column: PeerColumn) -> String {
		// Converts optional size to human readable size
		fn size_to_string(size: u64) -> String {
			size.file_size(CONVENTIONAL).unwrap_or("-".to_string())
		}

		match column {
			PeerColumn::Address => self.addr.clone(),
			PeerColumn::State => self.state.clone(),
			PeerColumn::UsedBandwidth => format!(
				"S: {}, R: {}",
				size_to_string(self.sent_bytes_per_sec),
				size_to_string(self.received_bytes_per_sec),
			)
			.to_string(),
			PeerColumn::TotalDifficulty => format!(
				"{} D @ {} H ({}s)",
				self.total_difficulty,
				self.height,
				(Utc::now() - self.last_seen).num_seconds(),
			)
			.to_string(),
			PeerColumn::Direction => self.direction.clone(),
			PeerColumn::Version => format!("{}", self.version),
			PeerColumn::UserAgent => self.user_agent.clone(),
		}
	}

	fn cmp(&self, other: &Self, column: PeerColumn) -> Ordering
	where
		Self: Sized,
	{
		// Compares used bandwidth of two peers
		fn cmp_used_bandwidth(curr: &PeerStats, other: &PeerStats) -> Ordering {
			let curr_recv_bytes = curr.received_bytes_per_sec;
			let curr_sent_bytes = curr.sent_bytes_per_sec;
			let other_recv_bytes = other.received_bytes_per_sec;
			let other_sent_bytes = other.sent_bytes_per_sec;

			let curr_sum = curr_recv_bytes + curr_sent_bytes;
			let other_sum = other_recv_bytes + other_sent_bytes;

			curr_sum.cmp(&other_sum)
		};

		match column {
			PeerColumn::Address => self.addr.cmp(&other.addr),
			PeerColumn::State => self.state.cmp(&other.state),
			PeerColumn::UsedBandwidth => cmp_used_bandwidth(&self, &other),
			PeerColumn::TotalDifficulty => self.total_difficulty.cmp(&other.total_difficulty),
			PeerColumn::Direction => self.direction.cmp(&other.direction),
			PeerColumn::Version => self.version.cmp(&other.version),
			PeerColumn::UserAgent => self.user_agent.cmp(&other.user_agent),
		}
	}
}

pub struct TUIPeerView;

impl TUIStatusListener for TUIPeerView {
	fn create() -> Box<dyn View> {
		let table_view = TableView::<PeerStats, PeerColumn>::new()
			.column(PeerColumn::Address, "Address", |c| c.width_percent(16))
			.column(PeerColumn::State, "State", |c| c.width_percent(8))
			.column(PeerColumn::UsedBandwidth, "Used bandwidth", |c| {
				c.width_percent(16)
			})
			.column(PeerColumn::Direction, "Direction", |c| c.width_percent(8))
			.column(PeerColumn::TotalDifficulty, "Total Difficulty", |c| {
				c.width_percent(24)
			})
			.column(PeerColumn::Version, "Proto", |c| c.width_percent(6))
			.column(PeerColumn::UserAgent, "User Agent", |c| c.width_percent(18));
		let peer_status_view = BoxView::with_full_screen(
			LinearLayout::new(Orientation::Vertical)
				.child(
					LinearLayout::new(Orientation::Horizontal)
						.child(TextView::new("  ").with_id("peers_total")),
				)
				.child(
					LinearLayout::new(Orientation::Horizontal)
						.child(TextView::new("Longest Chain: "))
						.child(TextView::new("  ").with_id("longest_work_peer")),
				)
				.child(TextView::new("   "))
				.child(
					Dialog::around(table_view.with_id(TABLE_PEER_STATUS).min_size((50, 20)))
						.title("Connected Peers"),
				),
		)
		.with_id(VIEW_PEER_SYNC);

		let peer_status_view =
			OnEventView::new(peer_status_view).on_pre_event(Key::Esc, move |c| {
				let _ = c.focus_id(MAIN_MENU);
			});

		Box::new(peer_status_view)
	}

	fn update(c: &mut Cursive, stats: &ServerStats) {
		let lp = stats
			.peer_stats
			.iter()
			.max_by(|x, y| x.total_difficulty.cmp(&y.total_difficulty));
		let lp_str = match lp {
			Some(l) => format!(
				"{} D @ {} H vs Us: {} D @ {} H",
				l.total_difficulty, l.height, stats.head.total_difficulty, stats.head.height
			)
			.to_string(),
			None => "".to_string(),
		};
		let _ = c.call_on_id(
			TABLE_PEER_STATUS,
			|t: &mut TableView<PeerStats, PeerColumn>| {
				t.set_items(stats.peer_stats.clone());
			},
		);
		let _ = c.call_on_id("peers_total", |t: &mut TextView| {
			t.set_content(format!(
				"Total Peers: {} (Outbound: {})",
				stats.peer_stats.len(),
				stats
					.peer_stats
					.iter()
					.filter(|x| x.direction == "Outbound")
					.count(),
			));
		});
		let _ = c.call_on_id("longest_work_peer", |t: &mut TextView| {
			t.set_content(lp_str);
		});
	}
}
<commit_msg>Replace S and R with corresponding ASCII arrows in bandwidth column (#2836)<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! TUI peer display

use std::cmp::Ordering;

use crate::servers::{PeerStats, ServerStats};

use crate::tui::humansize::{file_size_opts::CONVENTIONAL, FileSize};
use chrono::prelude::*;

use cursive::direction::Orientation;
use cursive::event::Key;
use cursive::traits::{Boxable, Identifiable};
use cursive::view::View;
use cursive::views::{BoxView, Dialog, LinearLayout, OnEventView, TextView};
use cursive::Cursive;

use crate::tui::constants::{MAIN_MENU, TABLE_PEER_STATUS, VIEW_PEER_SYNC};
use crate::tui::table::{TableView, TableViewItem};
use crate::tui::types::TUIStatusListener;

#[derive(Copy, Clone, PartialEq, Eq, Hash)]
enum PeerColumn {
	Address,
	State,
	UsedBandwidth,
	TotalDifficulty,
	Direction,
	Version,
	UserAgent,
}

impl PeerColumn {
	fn _as_str(&self) -> &str {
		match *self {
			PeerColumn::Address => "Address",
			PeerColumn::State => "State",
			PeerColumn::UsedBandwidth => "Used bandwidth",
			PeerColumn::Version => "Version",
			PeerColumn::TotalDifficulty => "Total Difficulty",
			PeerColumn::Direction => "Direction",
			PeerColumn::UserAgent => "User Agent",
		}
	}
}

impl TableViewItem<PeerColumn> for PeerStats {
	fn to_column(&self, column: PeerColumn) -> String {
		// Converts optional size to human readable size
		fn size_to_string(size: u64) -> String {
			size.file_size(CONVENTIONAL).unwrap_or("-".to_string())
		}

		match column {
			PeerColumn::Address => self.addr.clone(),
			PeerColumn::State => self.state.clone(),
			PeerColumn::UsedBandwidth => format!(
				"↑: {}, ↓: {}",
				size_to_string(self.sent_bytes_per_sec),
				size_to_string(self.received_bytes_per_sec),
			)
			.to_string(),
			PeerColumn::TotalDifficulty => format!(
				"{} D @ {} H ({}s)",
				self.total_difficulty,
				self.height,
				(Utc::now() - self.last_seen).num_seconds(),
			)
			.to_string(),
			PeerColumn::Direction => self.direction.clone(),
			PeerColumn::Version => format!("{}", self.version),
			PeerColumn::UserAgent => self.user_agent.clone(),
		}
	}

	fn cmp(&self, other: &Self, column: PeerColumn) -> Ordering
	where
		Self: Sized,
	{
		// Compares used bandwidth of two peers
		fn cmp_used_bandwidth(curr: &PeerStats, other: &PeerStats) -> Ordering {
			let curr_recv_bytes = curr.received_bytes_per_sec;
			let curr_sent_bytes = curr.sent_bytes_per_sec;
			let other_recv_bytes = other.received_bytes_per_sec;
			let other_sent_bytes = other.sent_bytes_per_sec;

			let curr_sum = curr_recv_bytes + curr_sent_bytes;
			let other_sum = other_recv_bytes + other_sent_bytes;

			curr_sum.cmp(&other_sum)
		};

		match column {
			PeerColumn::Address => self.addr.cmp(&other.addr),
			PeerColumn::State => self.state.cmp(&other.state),
			PeerColumn::UsedBandwidth => cmp_used_bandwidth(&self, &other),
			PeerColumn::TotalDifficulty => self.total_difficulty.cmp(&other.total_difficulty),
			PeerColumn::Direction => self.direction.cmp(&other.direction),
			PeerColumn::Version => self.version.cmp(&other.version),
			PeerColumn::UserAgent => self.user_agent.cmp(&other.user_agent),
		}
	}
}

pub struct TUIPeerView;

impl TUIStatusListener for TUIPeerView {
	fn create() -> Box<dyn View> {
		let table_view = TableView::<PeerStats, PeerColumn>::new()
			.column(PeerColumn::Address, "Address", |c| c.width_percent(16))
			.column(PeerColumn::State, "State", |c| c.width_percent(8))
			.column(PeerColumn::UsedBandwidth, "Used bandwidth", |c| {
				c.width_percent(16)
			})
			.column(PeerColumn::Direction, "Direction", |c| c.width_percent(8))
			.column(PeerColumn::TotalDifficulty, "Total Difficulty", |c| {
				c.width_percent(24)
			})
			.column(PeerColumn::Version, "Proto", |c| c.width_percent(6))
			.column(PeerColumn::UserAgent, "User Agent", |c| c.width_percent(18));
		let peer_status_view = BoxView::with_full_screen(
			LinearLayout::new(Orientation::Vertical)
				.child(
					LinearLayout::new(Orientation::Horizontal)
						.child(TextView::new("  ").with_id("peers_total")),
				)
				.child(
					LinearLayout::new(Orientation::Horizontal)
						.child(TextView::new("Longest Chain: "))
						.child(TextView::new("  ").with_id("longest_work_peer")),
				)
				.child(TextView::new("   "))
				.child(
					Dialog::around(table_view.with_id(TABLE_PEER_STATUS).min_size((50, 20)))
						.title("Connected Peers"),
				),
		)
		.with_id(VIEW_PEER_SYNC);

		let peer_status_view =
			OnEventView::new(peer_status_view).on_pre_event(Key::Esc, move |c| {
				let _ = c.focus_id(MAIN_MENU);
			});

		Box::new(peer_status_view)
	}

	fn update(c: &mut Cursive, stats: &ServerStats) {
		let lp = stats
			.peer_stats
			.iter()
			.max_by(|x, y| x.total_difficulty.cmp(&y.total_difficulty));
		let lp_str = match lp {
			Some(l) => format!(
				"{} D @ {} H vs Us: {} D @ {} H",
				l.total_difficulty, l.height, stats.head.total_difficulty, stats.head.height
			)
			.to_string(),
			None => "".to_string(),
		};
		let _ = c.call_on_id(
			TABLE_PEER_STATUS,
			|t: &mut TableView<PeerStats, PeerColumn>| {
				t.set_items(stats.peer_stats.clone());
			},
		);
		let _ = c.call_on_id("peers_total", |t: &mut TextView| {
			t.set_content(format!(
				"Total Peers: {} (Outbound: {})",
				stats.peer_stats.len(),
				stats
					.peer_stats
					.iter()
					.filter(|x| x.direction == "Outbound")
					.count(),
			));
		});
		let _ = c.call_on_id("longest_work_peer", |t: &mut TextView| {
			t.set_content(lp_str);
		});
	}
}
<|endoftext|>


<commit_before>use std::ops;
use std::ptr;
use std::io;

use std::cell::Cell;

use buffer::DataSource;

const DEFAULT_BUFFER_SIZE: usize = 6 * 1024;

/// Trait all parser buffers implement.
///
/// Enables the consumer to request specific amounts of data and only consume partial parts of the
/// buffer.
pub trait Buffer<I>: ops::Deref<Target=[I]> {
    /// Attempt to fill the buffer using the closure `F`.
    ///
    /// The successful return from `F` should contain the number of items successfully written to
    /// the slice.
    ///
    /// # Notes
    ///
    /// * The returned value must *NOT* be larger than the length of the given slice.
    ///
    /// * Return `0` if no more data is available or if the slice is of zero length.
    ///
    /// * The slice might contain uninitialized memory, do not read from the slice.
    fn fill<S: DataSource<Item=I>>(&mut self, &mut S) -> io::Result<usize>;

    /// Buffer attempts to clear space for additional items.
    fn request_space(&mut self, usize);

    /// Consumes the given amount of bytes, must be less than or equal to `len()`.
    ///
    /// Does not invalidate any borrow of data from self.
    fn consume(&self, items: usize);

    /// Returns the number of bytes left in the buffer.
    fn len(&self) -> usize;

    /// Returns the maximum amount of data which can be stored
    fn capacity(&self) -> usize;
}

/// A fixed size buffer.
///
/// Only allocates when created.
// TODO: Tests
#[derive(Debug, Eq, PartialEq)]
pub struct FixedSizeBuffer<I: Default + Clone> {
    /// Backing memory.
    buffer:    Vec<I>,
    /// Number of items of `buffer` which contain actual data.
    populated: usize,
    /// The number of bytes from the start of the buffer which are used.
    ///
    /// As long as used <= populated it is safe.
    used:      Cell<usize>,
}

impl<I: Default + Clone> FixedSizeBuffer<I> {
    /// Creates a fixed-size buffer with the default buffer size.
    pub fn new() -> Self {
        Self::with_size(DEFAULT_BUFFER_SIZE)
    }

    /// Creates a fixed-size buffer with the supplied buffer size.
    pub fn with_size(size: usize) -> Self {
        assert!(size > 0);

        FixedSizeBuffer {
            buffer:    vec![Default::default(); size],
            populated: 0,
            used:      Cell::new(0),
        }
    }
}

impl<I: Default + Clone> ops::Deref for FixedSizeBuffer<I> {
    type Target = [I];

    fn deref(&self) -> &[I] {
        &self.buffer[self.used.get()..self.populated]
    }
}

impl<I: Default + Clone> ops::DerefMut for FixedSizeBuffer<I> {
    fn deref_mut(&mut self) -> &mut [I] {
        &mut self.buffer[self.used.get()..self.populated]
    }
}

impl<I: Default + Clone> Buffer<I> for FixedSizeBuffer<I> {
    fn fill<S: DataSource<Item=I>>(&mut self, s: &mut S) -> io::Result<usize> {
        s.read(&mut self.buffer[self.populated..]).map(|n| {
            debug_assert!(self.populated + n <= self.buffer.len());

            self.populated += n;

            n
        })
    }

    fn request_space(&mut self, items: usize) {
        use std::ptr;

        assert!(self.populated >= self.used.get());

        // Only copy if we actually need to free the space
        if self.buffer.len() - self.populated < items {
            unsafe {
                ptr::copy(self.buffer.as_ptr().offset(self.used.get() as isize), self.buffer.as_mut_ptr(), self.populated - self.used.get());
            }

            self.populated -= self.used.get();
            self.used.set(0);
        }
    }

    fn consume(&self, items: usize) {
        debug_assert!(self.used.get() + items <= self.populated);

        self.used.set(self.used.get() + items)
    }

    fn len(&self) -> usize {
        self.populated - self.used.get()
    }

    fn capacity(&self) -> usize {
        self.buffer.len()
    }
}

/// A buffer which will reallocate to fit the requested amount of data.
///
/// # Note:
///
/// Will not decrease in size.
// TODO: Tests
pub struct GrowingBuffer<I> {
    /// Backing memory.
    buffer:    Vec<I>,
    /// Number of items of `buffer` which contain actual data.
    populated: usize,
    /// Maximal size of the buffer, 0 means infinity.
    limit:     usize,
    /// The number of bytes from the start of the buffer which are used.
    ///
    /// As long as used <= populated it is safe.
    used:      Cell<usize>,
}

impl<I> GrowingBuffer<I> {
    /// Creates a new unlimited `GrowingBuffer`.
    pub fn new() -> Self {
        Self::with_limit(0)
    }

    /// Creates a new `GrowingBuffer` with the specified limit.
    ///
    /// # Note
    ///
    /// The actual amount of allocated memory might be larger than the specified limit, depends on
    /// the allocator.
    pub fn with_limit(limit: usize) -> Self {
        GrowingBuffer {
            buffer:    Vec::new(),
            populated: 0,
            limit:     limit,
            used:      Cell::new(0),
        }
    }
}

impl<I> ops::Deref for GrowingBuffer<I> {
    type Target = [I];

    fn deref(&self) -> &[I] {
        &self.buffer[self.used.get()..self.populated]
    }
}

impl<I> ops::DerefMut for GrowingBuffer<I> {
    fn deref_mut(&mut self) -> &mut [I] {
        &mut self.buffer[self.used.get()..self.populated]
    }
}

impl<I> Buffer<I> for GrowingBuffer<I> {
    fn fill<S: DataSource<Item=I>>(&mut self, s: &mut S) -> io::Result<usize> {
        s.read(&mut self.buffer[self.populated..]).map(|n| {
            debug_assert!(self.populated + n <= self.buffer.len());

            self.populated += n;

            n
        })
    }

    fn request_space(&mut self, items: usize) {
        // If we are over the limit, refuse
        if self.limit != 0 && self.buffer.capacity() > self.limit {
            return;
        }

        if items + self.len() > self.buffer.capacity() {
            // We do not have enough space for the new items, reallocate
            self.buffer.reserve(items);

            let cap = self.buffer.capacity();

            // TODO: Would it be better with a Clone and Default requirement on I?
            // We set the length here to allow fill() to hand out a slice of uninitialized memory
            // to be populated.
            // NOTE: We cannot actually expose this memory to the parser since self.populated will
            // be the upper limit for the deref to slice.
            unsafe {
                self.buffer.set_len(cap);
            }
        }

        // Only copy if we actually need to free the space
        if self.buffer.len() - self.populated < items {
            unsafe {
                ptr::copy(self.buffer.as_ptr().offset(self.used.get() as isize), self.buffer.as_mut_ptr(), self.populated - self.used.get());
            }

            self.populated -= self.used.get();
            self.used.set(0);
        }
    }

    fn consume(&self, items: usize) {
        debug_assert!(self.used.get() + items <= self.populated);

        self.used.set(self.used.get() + items)
    }

    fn len(&self) -> usize {
        self.populated - self.used.get()
    }

    fn capacity(&self) -> usize {
        self.buffer.len()
    }
}
<commit_msg>Buffers: Made the interal items require Copy, since we use ptr::copy<commit_after>use std::ops;
use std::ptr;
use std::io;

use std::cell::Cell;

use buffer::DataSource;

const DEFAULT_BUFFER_SIZE: usize = 6 * 1024;

/// Trait all parser buffers implement.
///
/// Enables the consumer to request specific amounts of data and only consume partial parts of the
/// buffer.
pub trait Buffer<I>: ops::Deref<Target=[I]> {
    /// Attempt to fill the buffer using the closure `F`.
    ///
    /// The successful return from `F` should contain the number of items successfully written to
    /// the slice.
    ///
    /// # Notes
    ///
    /// * The returned value must *NOT* be larger than the length of the given slice.
    ///
    /// * Return `0` if no more data is available or if the slice is of zero length.
    ///
    /// * The slice might contain uninitialized memory, do not read from the slice.
    fn fill<S: DataSource<Item=I>>(&mut self, &mut S) -> io::Result<usize>;

    /// Buffer attempts to clear space for additional items.
    fn request_space(&mut self, usize);

    /// Consumes the given amount of bytes, must be less than or equal to `len()`.
    ///
    /// Does not invalidate any borrow of data from self.
    fn consume(&self, items: usize);

    /// Returns the number of bytes left in the buffer.
    fn len(&self) -> usize;

    /// Returns the maximum amount of data which can be stored
    fn capacity(&self) -> usize;
}

/// A fixed size buffer.
///
/// Only allocates when created.
// TODO: Tests
#[derive(Debug, Eq, PartialEq)]
pub struct FixedSizeBuffer<I: Default + Copy> {
    /// Backing memory.
    buffer:    Vec<I>,
    /// Number of items of `buffer` which contain actual data.
    populated: usize,
    /// The number of bytes from the start of the buffer which are used.
    ///
    /// As long as used <= populated it is safe.
    used:      Cell<usize>,
}

impl<I: Default + Copy> FixedSizeBuffer<I> {
    /// Creates a fixed-size buffer with the default buffer size.
    pub fn new() -> Self {
        Self::with_size(DEFAULT_BUFFER_SIZE)
    }

    /// Creates a fixed-size buffer with the supplied buffer size.
    pub fn with_size(size: usize) -> Self {
        assert!(size > 0);

        FixedSizeBuffer {
            buffer:    vec![Default::default(); size],
            populated: 0,
            used:      Cell::new(0),
        }
    }
}

impl<I: Default + Copy> ops::Deref for FixedSizeBuffer<I> {
    type Target = [I];

    fn deref(&self) -> &[I] {
        &self.buffer[self.used.get()..self.populated]
    }
}

impl<I: Default + Copy> ops::DerefMut for FixedSizeBuffer<I> {
    fn deref_mut(&mut self) -> &mut [I] {
        &mut self.buffer[self.used.get()..self.populated]
    }
}

impl<I: Default + Copy> Buffer<I> for FixedSizeBuffer<I> {
    fn fill<S: DataSource<Item=I>>(&mut self, s: &mut S) -> io::Result<usize> {
        s.read(&mut self.buffer[self.populated..]).map(|n| {
            debug_assert!(self.populated + n <= self.buffer.len());

            self.populated += n;

            n
        })
    }

    fn request_space(&mut self, items: usize) {
        use std::ptr;

        assert!(self.populated >= self.used.get());

        // Only copy if we actually need to free the space
        if self.buffer.len() - self.populated < items {
            unsafe {
                ptr::copy(self.buffer.as_ptr().offset(self.used.get() as isize), self.buffer.as_mut_ptr(), self.populated - self.used.get());
            }

            self.populated -= self.used.get();
            self.used.set(0);
        }
    }

    fn consume(&self, items: usize) {
        debug_assert!(self.used.get() + items <= self.populated);

        self.used.set(self.used.get() + items)
    }

    fn len(&self) -> usize {
        self.populated - self.used.get()
    }

    fn capacity(&self) -> usize {
        self.buffer.len()
    }
}

/// A buffer which will reallocate to fit the requested amount of data.
///
/// # Note:
///
/// Will not decrease in size.
// TODO: Tests
pub struct GrowingBuffer<I: Copy> {
    /// Backing memory.
    buffer:    Vec<I>,
    /// Number of items of `buffer` which contain actual data.
    populated: usize,
    /// Maximal size of the buffer, 0 means infinity.
    limit:     usize,
    /// The number of bytes from the start of the buffer which are used.
    ///
    /// As long as used <= populated it is safe.
    used:      Cell<usize>,
}

impl<I: Copy> GrowingBuffer<I> {
    /// Creates a new unlimited `GrowingBuffer`.
    pub fn new() -> Self {
        Self::with_limit(0)
    }

    /// Creates a new `GrowingBuffer` with the specified limit.
    ///
    /// # Note
    ///
    /// The actual amount of allocated memory might be larger than the specified limit, depends on
    /// the allocator.
    pub fn with_limit(limit: usize) -> Self {
        GrowingBuffer {
            buffer:    Vec::new(),
            populated: 0,
            limit:     limit,
            used:      Cell::new(0),
        }
    }
}

impl<I: Copy> ops::Deref for GrowingBuffer<I> {
    type Target = [I];

    fn deref(&self) -> &[I] {
        &self.buffer[self.used.get()..self.populated]
    }
}

impl<I: Copy> ops::DerefMut for GrowingBuffer<I> {
    fn deref_mut(&mut self) -> &mut [I] {
        &mut self.buffer[self.used.get()..self.populated]
    }
}

impl<I: Copy> Buffer<I> for GrowingBuffer<I> {
    fn fill<S: DataSource<Item=I>>(&mut self, s: &mut S) -> io::Result<usize> {
        s.read(&mut self.buffer[self.populated..]).map(|n| {
            debug_assert!(self.populated + n <= self.buffer.len());

            self.populated += n;

            n
        })
    }

    fn request_space(&mut self, items: usize) {
        // If we are over the limit, refuse
        if self.limit != 0 && self.buffer.capacity() > self.limit {
            return;
        }

        if items + self.len() > self.buffer.capacity() {
            // We do not have enough space for the new items, reallocate
            self.buffer.reserve(items);

            let cap = self.buffer.capacity();

            // TODO: Would it be better with a Default requirement on I?
            // We set the length here to allow fill() to hand out a slice of uninitialized memory
            // to be populated.
            // NOTE: We cannot actually expose this memory to the parser since self.populated will
            // be the upper limit for the deref to slice.
            unsafe {
                self.buffer.set_len(cap);
            }
        }

        // Only copy if we actually need to free the space
        if self.buffer.len() - self.populated < items {
            unsafe {
                ptr::copy(self.buffer.as_ptr().offset(self.used.get() as isize), self.buffer.as_mut_ptr(), self.populated - self.used.get());
            }

            self.populated -= self.used.get();
            self.used.set(0);
        }
    }

    fn consume(&self, items: usize) {
        debug_assert!(self.used.get() + items <= self.populated);

        self.used.set(self.used.get() + items)
    }

    fn len(&self) -> usize {
        self.populated - self.used.get()
    }

    fn capacity(&self) -> usize {
        self.buffer.len()
    }
}
<|endoftext|>