<commit_before>#![allow(unused_must_use)] // because we of writeln!

//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Inokentiy Babushkin <inokentiy.babushkin@googlemail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE file
//  * that was distributed with this source code.

// spell-checker:ignore (ToDO) lstat

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg};
use std::fs;
use std::io::{ErrorKind, Write};
use uucore::InvalidEncodingHandling;

// operating mode
enum Mode {
    Default, // use filesystem to determine information and limits
    Basic,   // check basic compatibility with POSIX
    Extra,   // check for leading dashes and empty names
    Both,    // a combination of `Basic` and `Extra`
}

static NAME: &str = "pathchk";
static ABOUT: &str = "Check whether file names are valid or portable";

mod options {
    pub const POSIX: &str = "posix";
    pub const POSIX_SPECIAL: &str = "posix-special";
    pub const PORTABILITY: &str = "portability";
    pub const PATH: &str = "path";
}

// a few global constants as used in the GNU implementation
const POSIX_PATH_MAX: usize = 256;
const POSIX_NAME_MAX: usize = 14;

fn get_usage() -> String {
    format!("{0} [OPTION]... NAME...", executable!())
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = get_usage();
    let args = args
        .collect_str(InvalidEncodingHandling::ConvertLossy)
        .accept_any();

    let matches = App::new(executable!())
        .version(crate_version!())
        .about(ABOUT)
        .usage(&usage[..])
        .arg(
            Arg::with_name(options::POSIX)
                .short("p")
                .help("check for most POSIX systems"),
        )
        .arg(
            Arg::with_name(options::POSIX_SPECIAL)
                .short("P")
                .help(r#"check for empty names and leading "-""#),
        )
        .arg(
            Arg::with_name(options::PORTABILITY)
                .long(options::PORTABILITY)
                .help("check for all POSIX systems (equivalent to -p -P)"),
        )
        .arg(Arg::with_name(options::PATH).hidden(true).multiple(true))
        .get_matches_from(args);

    // set working mode
    let is_posix = matches.values_of(options::POSIX).is_some();
    let is_posix_special = matches.values_of(options::POSIX_SPECIAL).is_some();
    let is_portability = matches.values_of(options::PORTABILITY).is_some();

    let mode = if (is_posix && is_posix_special) || is_portability {
        Mode::Both
    } else if is_posix {
        Mode::Basic
    } else if is_posix_special {
        Mode::Extra
    } else {
        Mode::Default
    };

    // take necessary actions
    let paths = matches.values_of(options::PATH);
    let mut res = if paths.is_none() {
        show_error!("missing operand\nTry {} --help for more information", NAME);
        false
    } else {
        true
    };

    if res {
        // free strings are path operands
        // FIXME: TCS, seems inefficient and overly verbose (?)
        for p in paths.unwrap() {
            let mut path = Vec::new();
            for path_segment in p.split('/') {
                path.push(path_segment.to_string());
            }
            res &= check_path(&mode, &path);
        }
    }

    // determine error code
    if res {
        0
    } else {
        1
    }
}

// check a path, given as a slice of it's components and an operating mode
fn check_path(mode: &Mode, path: &[String]) -> bool {
    match *mode {
        Mode::Basic => check_basic(&path),
        Mode::Extra => check_default(&path) && check_extra(&path),
        Mode::Both => check_basic(&path) && check_extra(&path),
        _ => check_default(&path),
    }
}

// check a path in basic compatibility mode
fn check_basic(path: &[String]) -> bool {
    let joined_path = path.join("/");
    let total_len = joined_path.len();
    // path length
    if total_len > POSIX_PATH_MAX {
        writeln!(
            &mut std::io::stderr(),
            "limit {} exceeded by length {} of file name {}",
            POSIX_PATH_MAX,
            total_len,
            joined_path
        );
        return false;
    } else if total_len == 0 {
        writeln!(&mut std::io::stderr(), "empty file name");
        return false;
    }
    // components: character portability and length
    for p in path {
        let component_len = p.len();
        if component_len > POSIX_NAME_MAX {
            writeln!(
                &mut std::io::stderr(),
                "limit {} exceeded by length {} of file name component '{}'",
                POSIX_NAME_MAX,
                component_len,
                p
            );
            return false;
        }
        if !check_portable_chars(&p) {
            return false;
        }
    }
    // permission checks
    check_searchable(&joined_path)
}

// check a path in extra compatibility mode
fn check_extra(path: &[String]) -> bool {
    // components: leading hyphens
    for p in path {
        if !no_leading_hyphen(&p) {
            writeln!(
                &mut std::io::stderr(),
                "leading hyphen in file name component '{}'",
                p
            );
            return false;
        }
    }
    // path length
    if path.join("/").is_empty() {
        writeln!(&mut std::io::stderr(), "empty file name");
        return false;
    }
    true
}

// check a path in default mode (using the file system)
fn check_default(path: &[String]) -> bool {
    let joined_path = path.join("/");
    let total_len = joined_path.len();
    // path length
    if total_len > libc::PATH_MAX as usize {
        writeln!(
            &mut std::io::stderr(),
            "limit {} exceeded by length {} of file name '{}'",
            libc::PATH_MAX,
            total_len,
            joined_path
        );
        return false;
    }
    // components: length
    for p in path {
        let component_len = p.len();
        if component_len > libc::FILENAME_MAX as usize {
            writeln!(
                &mut std::io::stderr(),
                "limit {} exceeded by length {} of file name component '{}'",
                libc::FILENAME_MAX,
                component_len,
                p
            );
            return false;
        }
    }
    // permission checks
    check_searchable(&joined_path)
}

// check whether a path is or if other problems arise
fn check_searchable(path: &str) -> bool {
    // we use lstat, just like the original implementation
    match fs::symlink_metadata(path) {
        Ok(_) => true,
        Err(e) => {
            if e.kind() == ErrorKind::NotFound {
                true
            } else {
                writeln!(&mut std::io::stderr(), "{}", e);
                false
            }
        }
    }
}

// check for a hyphen at the beginning of a path segment
fn no_leading_hyphen(path_segment: &str) -> bool {
    !path_segment.starts_with('-')
}

// check whether a path segment contains only valid (read: portable) characters
fn check_portable_chars(path_segment: &str) -> bool {
    let valid_str = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789._-".to_string();
    for ch in path_segment.chars() {
        if !valid_str.contains(ch) {
            writeln!(
                &mut std::io::stderr(),
                "nonportable character '{}' in file name component '{}'",
                ch,
                path_segment
            );
            return false;
        }
    }
    true
}
<commit_msg>refactor/pathchk ~ fix `cargo clippy` complaint (clippy::needless_borrow)<commit_after>#![allow(unused_must_use)] // because we of writeln!

//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Inokentiy Babushkin <inokentiy.babushkin@googlemail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE file
//  * that was distributed with this source code.

// spell-checker:ignore (ToDO) lstat

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg};
use std::fs;
use std::io::{ErrorKind, Write};
use uucore::InvalidEncodingHandling;

// operating mode
enum Mode {
    Default, // use filesystem to determine information and limits
    Basic,   // check basic compatibility with POSIX
    Extra,   // check for leading dashes and empty names
    Both,    // a combination of `Basic` and `Extra`
}

static NAME: &str = "pathchk";
static ABOUT: &str = "Check whether file names are valid or portable";

mod options {
    pub const POSIX: &str = "posix";
    pub const POSIX_SPECIAL: &str = "posix-special";
    pub const PORTABILITY: &str = "portability";
    pub const PATH: &str = "path";
}

// a few global constants as used in the GNU implementation
const POSIX_PATH_MAX: usize = 256;
const POSIX_NAME_MAX: usize = 14;

fn get_usage() -> String {
    format!("{0} [OPTION]... NAME...", executable!())
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = get_usage();
    let args = args
        .collect_str(InvalidEncodingHandling::ConvertLossy)
        .accept_any();

    let matches = App::new(executable!())
        .version(crate_version!())
        .about(ABOUT)
        .usage(&usage[..])
        .arg(
            Arg::with_name(options::POSIX)
                .short("p")
                .help("check for most POSIX systems"),
        )
        .arg(
            Arg::with_name(options::POSIX_SPECIAL)
                .short("P")
                .help(r#"check for empty names and leading "-""#),
        )
        .arg(
            Arg::with_name(options::PORTABILITY)
                .long(options::PORTABILITY)
                .help("check for all POSIX systems (equivalent to -p -P)"),
        )
        .arg(Arg::with_name(options::PATH).hidden(true).multiple(true))
        .get_matches_from(args);

    // set working mode
    let is_posix = matches.values_of(options::POSIX).is_some();
    let is_posix_special = matches.values_of(options::POSIX_SPECIAL).is_some();
    let is_portability = matches.values_of(options::PORTABILITY).is_some();

    let mode = if (is_posix && is_posix_special) || is_portability {
        Mode::Both
    } else if is_posix {
        Mode::Basic
    } else if is_posix_special {
        Mode::Extra
    } else {
        Mode::Default
    };

    // take necessary actions
    let paths = matches.values_of(options::PATH);
    let mut res = if paths.is_none() {
        show_error!("missing operand\nTry {} --help for more information", NAME);
        false
    } else {
        true
    };

    if res {
        // free strings are path operands
        // FIXME: TCS, seems inefficient and overly verbose (?)
        for p in paths.unwrap() {
            let mut path = Vec::new();
            for path_segment in p.split('/') {
                path.push(path_segment.to_string());
            }
            res &= check_path(&mode, &path);
        }
    }

    // determine error code
    if res {
        0
    } else {
        1
    }
}

// check a path, given as a slice of it's components and an operating mode
fn check_path(mode: &Mode, path: &[String]) -> bool {
    match *mode {
        Mode::Basic => check_basic(path),
        Mode::Extra => check_default(path) && check_extra(path),
        Mode::Both => check_basic(path) && check_extra(path),
        _ => check_default(path),
    }
}

// check a path in basic compatibility mode
fn check_basic(path: &[String]) -> bool {
    let joined_path = path.join("/");
    let total_len = joined_path.len();
    // path length
    if total_len > POSIX_PATH_MAX {
        writeln!(
            &mut std::io::stderr(),
            "limit {} exceeded by length {} of file name {}",
            POSIX_PATH_MAX,
            total_len,
            joined_path
        );
        return false;
    } else if total_len == 0 {
        writeln!(&mut std::io::stderr(), "empty file name");
        return false;
    }
    // components: character portability and length
    for p in path {
        let component_len = p.len();
        if component_len > POSIX_NAME_MAX {
            writeln!(
                &mut std::io::stderr(),
                "limit {} exceeded by length {} of file name component '{}'",
                POSIX_NAME_MAX,
                component_len,
                p
            );
            return false;
        }
        if !check_portable_chars(p) {
            return false;
        }
    }
    // permission checks
    check_searchable(&joined_path)
}

// check a path in extra compatibility mode
fn check_extra(path: &[String]) -> bool {
    // components: leading hyphens
    for p in path {
        if !no_leading_hyphen(p) {
            writeln!(
                &mut std::io::stderr(),
                "leading hyphen in file name component '{}'",
                p
            );
            return false;
        }
    }
    // path length
    if path.join("/").is_empty() {
        writeln!(&mut std::io::stderr(), "empty file name");
        return false;
    }
    true
}

// check a path in default mode (using the file system)
fn check_default(path: &[String]) -> bool {
    let joined_path = path.join("/");
    let total_len = joined_path.len();
    // path length
    if total_len > libc::PATH_MAX as usize {
        writeln!(
            &mut std::io::stderr(),
            "limit {} exceeded by length {} of file name '{}'",
            libc::PATH_MAX,
            total_len,
            joined_path
        );
        return false;
    }
    // components: length
    for p in path {
        let component_len = p.len();
        if component_len > libc::FILENAME_MAX as usize {
            writeln!(
                &mut std::io::stderr(),
                "limit {} exceeded by length {} of file name component '{}'",
                libc::FILENAME_MAX,
                component_len,
                p
            );
            return false;
        }
    }
    // permission checks
    check_searchable(&joined_path)
}

// check whether a path is or if other problems arise
fn check_searchable(path: &str) -> bool {
    // we use lstat, just like the original implementation
    match fs::symlink_metadata(path) {
        Ok(_) => true,
        Err(e) => {
            if e.kind() == ErrorKind::NotFound {
                true
            } else {
                writeln!(&mut std::io::stderr(), "{}", e);
                false
            }
        }
    }
}

// check for a hyphen at the beginning of a path segment
fn no_leading_hyphen(path_segment: &str) -> bool {
    !path_segment.starts_with('-')
}

// check whether a path segment contains only valid (read: portable) characters
fn check_portable_chars(path_segment: &str) -> bool {
    let valid_str = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789._-".to_string();
    for ch in path_segment.chars() {
        if !valid_str.contains(ch) {
            writeln!(
                &mut std::io::stderr(),
                "nonportable character '{}' in file name component '{}'",
                ch,
                path_segment
            );
            return false;
        }
    }
    true
}
<|endoftext|>


<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Alex Lyon <arcterus@mail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore (ToDO) tstr sigstr cmdname setpgid sigchld

#[macro_use]
extern crate uucore;

extern crate clap;

use clap::{crate_version, App, AppSettings, Arg};
use std::io::ErrorKind;
use std::process::{Command, Stdio};
use std::time::Duration;
use uucore::display::Quotable;
use uucore::error::{UResult, USimpleError};
use uucore::process::ChildExt;
use uucore::signals::{signal_by_name_or_value, signal_name_by_value};
use uucore::{format_usage, InvalidEncodingHandling};

static ABOUT: &str = "Start COMMAND, and kill it if still running after DURATION.";
const USAGE: &str = "{} [OPTION] DURATION COMMAND...";

const ERR_EXIT_STATUS: i32 = 125;

pub mod options {
    pub static FOREGROUND: &str = "foreground";
    pub static KILL_AFTER: &str = "kill-after";
    pub static SIGNAL: &str = "signal";
    pub static PRESERVE_STATUS: &str = "preserve-status";
    pub static VERBOSE: &str = "verbose";

    // Positional args.
    pub static DURATION: &str = "duration";
    pub static COMMAND: &str = "command";
}

struct Config {
    foreground: bool,
    kill_after: Option<Duration>,
    signal: usize,
    duration: Duration,
    preserve_status: bool,
    verbose: bool,

    command: Vec<String>,
}

impl Config {
    fn from(options: &clap::ArgMatches) -> Self {
        let signal = match options.value_of(options::SIGNAL) {
            Some(signal_) => {
                let signal_result = signal_by_name_or_value(signal_);
                match signal_result {
                    None => {
                        unreachable!("invalid signal {}", signal_.quote());
                    }
                    Some(signal_value) => signal_value,
                }
            }
            _ => uucore::signals::signal_by_name_or_value("TERM").unwrap(),
        };

        let kill_after = options
            .value_of(options::KILL_AFTER)
            .map(|time| uucore::parse_time::from_str(time).unwrap());

        let duration: Duration =
            uucore::parse_time::from_str(options.value_of(options::DURATION).unwrap()).unwrap();

        let preserve_status: bool = options.is_present(options::PRESERVE_STATUS);
        let foreground = options.is_present(options::FOREGROUND);
        let verbose = options.is_present(options::VERBOSE);

        let command = options
            .values_of(options::COMMAND)
            .unwrap()
            .map(String::from)
            .collect::<Vec<_>>();

        Self {
            foreground,
            kill_after,
            signal,
            duration,
            preserve_status,
            verbose,
            command,
        }
    }
}

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    let args = args
        .collect_str(InvalidEncodingHandling::ConvertLossy)
        .accept_any();

    let app = uu_app();

    let matches = app.get_matches_from(args);

    let config = Config::from(&matches);
    timeout(
        &config.command,
        config.duration,
        config.signal,
        config.kill_after,
        config.foreground,
        config.preserve_status,
        config.verbose,
    )
}

pub fn uu_app<'a>() -> App<'a> {
    App::new("timeout")
        .version(crate_version!())
        .about(ABOUT)
        .override_usage(format_usage(USAGE))
        .arg(
            Arg::new(options::FOREGROUND)
                .long(options::FOREGROUND)
                .help("when not running timeout directly from a shell prompt, allow COMMAND to read from the TTY and get TTY signals; in this mode, children of COMMAND will not be timed out")
        )
        .arg(
            Arg::new(options::KILL_AFTER)
                .short('k')
                .takes_value(true))
        .arg(
            Arg::new(options::PRESERVE_STATUS)
                .long(options::PRESERVE_STATUS)
                .help("exit with the same status as COMMAND, even when the command times out")
        )
        .arg(
            Arg::new(options::SIGNAL)
                .short('s')
                .long(options::SIGNAL)
                .help("specify the signal to be sent on timeout; SIGNAL may be a name like 'HUP' or a number; see 'kill -l' for a list of signals")
                .takes_value(true)
        )
        .arg(
            Arg::new(options::VERBOSE)
              .short('v')
              .long(options::VERBOSE)
              .help("diagnose to stderr any signal sent upon timeout")
        )
        .arg(
            Arg::new(options::DURATION)
                .index(1)
                .required(true)
        )
        .arg(
            Arg::new(options::COMMAND)
                .index(2)
                .required(true)
                .multiple_occurrences(true)
        )
        .setting(AppSettings::TrailingVarArg)
        .setting(AppSettings::InferLongArgs)
}

/// Remove pre-existing SIGCHLD handlers that would make waiting for the child's exit code fail.
fn unblock_sigchld() {
    unsafe {
        nix::sys::signal::signal(
            nix::sys::signal::Signal::SIGCHLD,
            nix::sys::signal::SigHandler::SigDfl,
        )
        .unwrap();
    }
}

/// TODO: Improve exit codes, and make them consistent with the GNU Coreutils exit codes.

fn timeout(
    cmd: &[String],
    duration: Duration,
    signal: usize,
    kill_after: Option<Duration>,
    foreground: bool,
    preserve_status: bool,
    verbose: bool,
) -> UResult<()> {
    if !foreground {
        unsafe { libc::setpgid(0, 0) };
    }
    let mut process = Command::new(&cmd[0])
        .args(&cmd[1..])
        .stdin(Stdio::inherit())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .spawn()
        .map_err(|err| {
            let status_code = if err.kind() == ErrorKind::NotFound {
                // FIXME: not sure which to use
                127
            } else {
                // FIXME: this may not be 100% correct...
                126
            };
            USimpleError::new(status_code, format!("failed to execute process: {}", err))
        })?;
    unblock_sigchld();
    match process.wait_or_timeout(duration) {
        Ok(Some(status)) => {
            let status_code = status.code().unwrap_or_else(|| status.signal().unwrap());
            if status_code == 0 {
                Ok(())
            } else {
                Err(status_code.into())
            }
        }
        Ok(None) => {
            if verbose {
                show_error!(
                    "sending signal {} to command {}",
                    signal_name_by_value(signal).unwrap(),
                    cmd[0].quote()
                );
            }
            process
                .send_signal(signal)
                .map_err(|e| USimpleError::new(ERR_EXIT_STATUS, format!("{}", e)))?;
            if let Some(kill_after) = kill_after {
                match process.wait_or_timeout(kill_after) {
                    Ok(Some(status)) => {
                        if preserve_status {
                            let status_code =
                                status.code().unwrap_or_else(|| status.signal().unwrap());
                            if status_code == 0 {
                                Ok(())
                            } else {
                                Err(status_code.into())
                            }
                        } else {
                            Err(124.into())
                        }
                    }
                    Ok(None) => {
                        if verbose {
                            show_error!("sending signal KILL to command {}", cmd[0].quote());
                        }
                        process
                            .send_signal(uucore::signals::signal_by_name_or_value("KILL").unwrap())
                            .map_err(|e| USimpleError::new(ERR_EXIT_STATUS, format!("{}", e)))?;
                        process
                            .wait()
                            .map_err(|e| USimpleError::new(ERR_EXIT_STATUS, format!("{}", e)))?;
                        Err(137.into())
                    }
                    Err(_) => Err(124.into()),
                }
            } else {
                Err(124.into())
            }
        }
        Err(_) => {
            // We're going to return ERR_EXIT_STATUS regardless of
            // whether `send_signal()` succeeds or fails, so just
            // ignore the return value.
            process
                .send_signal(signal)
                .map_err(|e| USimpleError::new(ERR_EXIT_STATUS, format!("{}", e)))?;
            Err(ERR_EXIT_STATUS.into())
        }
    }
}
<commit_msg>timeout: avoid panicking for empty string<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Alex Lyon <arcterus@mail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore (ToDO) tstr sigstr cmdname setpgid sigchld

#[macro_use]
extern crate uucore;

extern crate clap;

use clap::{crate_version, App, AppSettings, Arg};
use std::io::ErrorKind;
use std::process::{Command, Stdio};
use std::time::Duration;
use uucore::display::Quotable;
use uucore::error::{UResult, USimpleError};
use uucore::process::ChildExt;
use uucore::signals::{signal_by_name_or_value, signal_name_by_value};
use uucore::{format_usage, InvalidEncodingHandling};

static ABOUT: &str = "Start COMMAND, and kill it if still running after DURATION.";
const USAGE: &str = "{} [OPTION] DURATION COMMAND...";

const ERR_EXIT_STATUS: i32 = 125;

pub mod options {
    pub static FOREGROUND: &str = "foreground";
    pub static KILL_AFTER: &str = "kill-after";
    pub static SIGNAL: &str = "signal";
    pub static PRESERVE_STATUS: &str = "preserve-status";
    pub static VERBOSE: &str = "verbose";

    // Positional args.
    pub static DURATION: &str = "duration";
    pub static COMMAND: &str = "command";
}

struct Config {
    foreground: bool,
    kill_after: Option<Duration>,
    signal: usize,
    duration: Duration,
    preserve_status: bool,
    verbose: bool,

    command: Vec<String>,
}

impl Config {
    fn from(options: &clap::ArgMatches) -> UResult<Self> {
        let signal = match options.value_of(options::SIGNAL) {
            Some(signal_) => {
                let signal_result = signal_by_name_or_value(signal_);
                match signal_result {
                    None => {
                        unreachable!("invalid signal {}", signal_.quote());
                    }
                    Some(signal_value) => signal_value,
                }
            }
            _ => uucore::signals::signal_by_name_or_value("TERM").unwrap(),
        };

        let kill_after = options
            .value_of(options::KILL_AFTER)
            .map(|time| uucore::parse_time::from_str(time).unwrap());

        let duration =
            match uucore::parse_time::from_str(options.value_of(options::DURATION).unwrap()) {
                Ok(duration) => duration,
                Err(err) => return Err(USimpleError::new(1, err)),
            };

        let preserve_status: bool = options.is_present(options::PRESERVE_STATUS);
        let foreground = options.is_present(options::FOREGROUND);
        let verbose = options.is_present(options::VERBOSE);

        let command = options
            .values_of(options::COMMAND)
            .unwrap()
            .map(String::from)
            .collect::<Vec<_>>();

        Ok(Self {
            foreground,
            kill_after,
            signal,
            duration,
            preserve_status,
            verbose,
            command,
        })
    }
}

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    let args = args
        .collect_str(InvalidEncodingHandling::ConvertLossy)
        .accept_any();

    let app = uu_app();

    let matches = app.get_matches_from(args);

    let config = Config::from(&matches)?;
    timeout(
        &config.command,
        config.duration,
        config.signal,
        config.kill_after,
        config.foreground,
        config.preserve_status,
        config.verbose,
    )
}

pub fn uu_app<'a>() -> App<'a> {
    App::new("timeout")
        .version(crate_version!())
        .about(ABOUT)
        .override_usage(format_usage(USAGE))
        .arg(
            Arg::new(options::FOREGROUND)
                .long(options::FOREGROUND)
                .help("when not running timeout directly from a shell prompt, allow COMMAND to read from the TTY and get TTY signals; in this mode, children of COMMAND will not be timed out")
        )
        .arg(
            Arg::new(options::KILL_AFTER)
                .short('k')
                .takes_value(true))
        .arg(
            Arg::new(options::PRESERVE_STATUS)
                .long(options::PRESERVE_STATUS)
                .help("exit with the same status as COMMAND, even when the command times out")
        )
        .arg(
            Arg::new(options::SIGNAL)
                .short('s')
                .long(options::SIGNAL)
                .help("specify the signal to be sent on timeout; SIGNAL may be a name like 'HUP' or a number; see 'kill -l' for a list of signals")
                .takes_value(true)
        )
        .arg(
            Arg::new(options::VERBOSE)
              .short('v')
              .long(options::VERBOSE)
              .help("diagnose to stderr any signal sent upon timeout")
        )
        .arg(
            Arg::new(options::DURATION)
                .index(1)
                .required(true)
        )
        .arg(
            Arg::new(options::COMMAND)
                .index(2)
                .required(true)
                .multiple_occurrences(true)
        )
        .setting(AppSettings::TrailingVarArg)
        .setting(AppSettings::InferLongArgs)
}

/// Remove pre-existing SIGCHLD handlers that would make waiting for the child's exit code fail.
fn unblock_sigchld() {
    unsafe {
        nix::sys::signal::signal(
            nix::sys::signal::Signal::SIGCHLD,
            nix::sys::signal::SigHandler::SigDfl,
        )
        .unwrap();
    }
}

/// TODO: Improve exit codes, and make them consistent with the GNU Coreutils exit codes.

fn timeout(
    cmd: &[String],
    duration: Duration,
    signal: usize,
    kill_after: Option<Duration>,
    foreground: bool,
    preserve_status: bool,
    verbose: bool,
) -> UResult<()> {
    if !foreground {
        unsafe { libc::setpgid(0, 0) };
    }
    let mut process = Command::new(&cmd[0])
        .args(&cmd[1..])
        .stdin(Stdio::inherit())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .spawn()
        .map_err(|err| {
            let status_code = if err.kind() == ErrorKind::NotFound {
                // FIXME: not sure which to use
                127
            } else {
                // FIXME: this may not be 100% correct...
                126
            };
            USimpleError::new(status_code, format!("failed to execute process: {}", err))
        })?;
    unblock_sigchld();
    match process.wait_or_timeout(duration) {
        Ok(Some(status)) => {
            let status_code = status.code().unwrap_or_else(|| status.signal().unwrap());
            if status_code == 0 {
                Ok(())
            } else {
                Err(status_code.into())
            }
        }
        Ok(None) => {
            if verbose {
                show_error!(
                    "sending signal {} to command {}",
                    signal_name_by_value(signal).unwrap(),
                    cmd[0].quote()
                );
            }
            process
                .send_signal(signal)
                .map_err(|e| USimpleError::new(ERR_EXIT_STATUS, format!("{}", e)))?;
            if let Some(kill_after) = kill_after {
                match process.wait_or_timeout(kill_after) {
                    Ok(Some(status)) => {
                        if preserve_status {
                            let status_code =
                                status.code().unwrap_or_else(|| status.signal().unwrap());
                            if status_code == 0 {
                                Ok(())
                            } else {
                                Err(status_code.into())
                            }
                        } else {
                            Err(124.into())
                        }
                    }
                    Ok(None) => {
                        if verbose {
                            show_error!("sending signal KILL to command {}", cmd[0].quote());
                        }
                        process
                            .send_signal(uucore::signals::signal_by_name_or_value("KILL").unwrap())
                            .map_err(|e| USimpleError::new(ERR_EXIT_STATUS, format!("{}", e)))?;
                        process
                            .wait()
                            .map_err(|e| USimpleError::new(ERR_EXIT_STATUS, format!("{}", e)))?;
                        Err(137.into())
                    }
                    Err(_) => Err(124.into()),
                }
            } else {
                Err(124.into())
            }
        }
        Err(_) => {
            // We're going to return ERR_EXIT_STATUS regardless of
            // whether `send_signal()` succeeds or fails, so just
            // ignore the return value.
            process
                .send_signal(signal)
                .map_err(|e| USimpleError::new(ERR_EXIT_STATUS, format!("{}", e)))?;
            Err(ERR_EXIT_STATUS.into())
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2013-2014 The gl-rs developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// 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.

#![crate_name = "gl_generator"]
#![comment = "OpenGL function loader generator."]
#![license = "ASL2"]
#![crate_type = "dylib"]

#![feature(phase)]
#![feature(globs)]
#![feature(macro_rules)]
#![feature(plugin_registrar)]
#![feature(quote)]

#[phase(plugin, link)]
extern crate log;

extern crate rustc;
extern crate syntax;

use std::path::Path;
use std::io::{File, Reader};

use syntax::parse::token;
use syntax::ast::{ Item, TokenTree };
use syntax::ext::base::{expr_to_string, get_exprs_from_tts, DummyResult, ExtCtxt, MacResult};
use syntax::codemap::Span;

use registry::*;
use static_gen::StaticGenerator;
use struct_gen::StructGenerator;

mod common;
pub mod static_gen;
pub mod struct_gen;
pub mod registry;
pub mod ty;

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut ::rustc::plugin::Registry) {
    reg.register_macro("generate_gl_bindings", macro_handler);
}

// this is the object that we will return from the "generate_gl_bindings" macro expansion
struct MacroResult {
    content: Vec<::std::gc::Gc<Item>>
}
impl MacResult for MacroResult {
    fn make_def(&self) -> Option<::syntax::ext::base::MacroDef> { None }
    fn make_expr(&self) -> Option<::std::gc::Gc<::syntax::ast::Expr>> { None }
    fn make_pat(&self) -> Option<::std::gc::Gc<::syntax::ast::Pat>> { None }
    fn make_stmt(&self) -> Option<::std::gc::Gc<::syntax::ast::Stmt>> { None }

    fn make_items(&self) -> Option<::syntax::util::small_vector::SmallVector<::std::gc::Gc<Item>>> {
        Some(::syntax::util::small_vector::SmallVector::many(self.content.clone()))
    }
}

// handler for generate_gl_bindings!
fn macro_handler(ecx: &mut ExtCtxt, span: Span, token_tree: &[TokenTree]) -> Box<MacResult> {
    // getting the arguments from the macro
    let (ns_literal, api, profile, version, generator, exts) = match parse_macro_arguments(ecx, span.clone(), token_tree) {
        Some(t) => t,
        None => return DummyResult::any(span)
    };

    let ns = match ns_literal.as_slice() {
        "gl"  => Gl,
        "glx" => {
            ecx.span_err(span, "glx generation unimplemented");
            return DummyResult::any(span)
        },
        "wgl" => {
            ecx.span_err(span, "wgl generation unimplemented");
            return DummyResult::any(span)
        }
        ns => {
            ecx.span_err(span, format!("Unexpected opengl namespace '{}'", ns).as_slice());
            return DummyResult::any(span)
        }
    };

    let path = {
        use std::os;
        let p = Path::new(ecx.codemap().span_to_filename(span));
        let mut p = os::make_absolute(&p);
        p.pop(); p.pop(); p.pop();
        p.push(format!("deps/khronos-api/{}.xml", ns_literal));
        p
    };

    let filter = Some(Filter {
        extensions: exts.unwrap_or("".to_string()).as_slice().split(',').map(|s| s.to_string()).collect(),
        profile: profile,
        version: version,
        api: api,
    });

    // generating the registry of all bindings
    let reg = Registry::from_xml(
        match
            match File::open(&path) {
                                Ok(f) => f,
                                Err(err) => {
                                    ecx.span_err(span, format!("{}", err).as_slice());
                                    return DummyResult::any(span)
                                }
                            }.read_to_string()
        {
            Ok(s) => s,
            Err(_) => {
                ecx.span_err(span, "registry source not utf8!");
                return DummyResult::any(span)
            }

        }.as_slice(), ns, filter
    );

    // generating the Rust bindings as a source code into "buffer"
    let mut buffer = ::std::io::MemWriter::new();
    match generator.as_slice() {
        "static" => StaticGenerator::write(&mut buffer, &reg, ns, false),
        "struct" => StructGenerator::write(&mut buffer, &reg, ns, false),
        generator => {
            ecx.span_err(span, format!("unknown generator type: {}", generator).as_slice());
            return DummyResult::any(span)
        },
    }

    // creating a new Rust parser from these bindings
    let content = match String::from_utf8(buffer.unwrap()) {
        Ok(s) => s,
        Err(err) => {
            ecx.span_err(span, format!("{}", err).as_slice());
            return DummyResult::any(span)
        }
    };
    let mut parser = ::syntax::parse::new_parser_from_source_str(ecx.parse_sess(), ecx.cfg(), path.display().to_string(), content);

    // getting all the items defined by the bindings
    let mut items = Vec::new();
    loop {
        match parser.parse_item_with_outer_attributes() {
            None => break,
            Some(i) => items.push(i)
        }
    }
    if !parser.eat(&token::EOF) {
        ecx.span_err(span, "the rust parser failed to compile all the generated bindings (meaning there is a bug in this library!)");
        return DummyResult::any(span)
    }
    box MacroResult { content: items } as Box<MacResult>
}

fn parse_macro_arguments(ecx: &mut ExtCtxt, span: Span, tts: &[syntax::ast::TokenTree]) -> Option<(String, String, String, String, String, Option<String>)> {
    let values = match get_exprs_from_tts(ecx, span, tts) {
        Some(v) => v,
        None => return None
    };

    if values.len() != 5 && values.len() != 6 {
        ecx.span_err(span, format!("expected 5 or 6 arguments but only got {}", values.len()).as_slice());
        return None;
    }

    match (
        expr_to_string(ecx, values[0].clone(), "expected string literal").map(|e| match e { (s, _) => s.get().to_string() }),
        expr_to_string(ecx, values[1].clone(), "expected string literal").map(|e| match e { (s, _) => s.get().to_string() }),
        expr_to_string(ecx, values[2].clone(), "expected string literal").map(|e| match e { (s, _) => s.get().to_string() }),
        expr_to_string(ecx, values[3].clone(), "expected string literal").map(|e| match e { (s, _) => s.get().to_string() }),
        expr_to_string(ecx, values[4].clone(), "expected string literal").map(|e| match e { (s, _) => s.get().to_string() }),
        values.as_slice().get(5).and_then(|e| expr_to_string(ecx, *e, "expected string literal")).map(|e| match e { (s, _) => s.get().to_string() }),
    ) {
        (Some(a), Some(b), Some(c), Some(d), Some(e), f) => Some((a, b, c, d, e, f)),
        _ => None
    }
}
<commit_msg>Isolate registry parsing and bindings generating in tasks<commit_after>// Copyright 2013-2014 The gl-rs developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// 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.

#![crate_name = "gl_generator"]
#![comment = "OpenGL function loader generator."]
#![license = "ASL2"]
#![crate_type = "dylib"]

#![feature(phase)]
#![feature(globs)]
#![feature(macro_rules)]
#![feature(plugin_registrar)]
#![feature(quote)]

#[phase(plugin, link)]
extern crate log;

extern crate rustc;
extern crate syntax;

use std::path::Path;
use std::io::{File, Reader};

use syntax::parse::token;
use syntax::ast::{ Item, TokenTree };
use syntax::ext::base::{expr_to_string, get_exprs_from_tts, DummyResult, ExtCtxt, MacResult};
use syntax::codemap::Span;

use registry::*;
use static_gen::StaticGenerator;
use struct_gen::StructGenerator;

mod common;
pub mod static_gen;
pub mod struct_gen;
pub mod registry;
pub mod ty;

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut ::rustc::plugin::Registry) {
    reg.register_macro("generate_gl_bindings", macro_handler);
}

// this is the object that we will return from the "generate_gl_bindings" macro expansion
struct MacroResult {
    content: Vec<::std::gc::Gc<Item>>
}
impl MacResult for MacroResult {
    fn make_def(&self) -> Option<::syntax::ext::base::MacroDef> { None }
    fn make_expr(&self) -> Option<::std::gc::Gc<::syntax::ast::Expr>> { None }
    fn make_pat(&self) -> Option<::std::gc::Gc<::syntax::ast::Pat>> { None }
    fn make_stmt(&self) -> Option<::std::gc::Gc<::syntax::ast::Stmt>> { None }

    fn make_items(&self) -> Option<::syntax::util::small_vector::SmallVector<::std::gc::Gc<Item>>> {
        Some(::syntax::util::small_vector::SmallVector::many(self.content.clone()))
    }
}

// handler for generate_gl_bindings!
fn macro_handler(ecx: &mut ExtCtxt, span: Span, token_tree: &[TokenTree]) -> Box<MacResult> {
    // getting the arguments from the macro
    let (ns_literal, api, profile, version, generator, exts) = match parse_macro_arguments(ecx, span.clone(), token_tree) {
        Some(t) => t,
        None => return DummyResult::any(span)
    };

    let ns = match ns_literal.as_slice() {
        "gl"  => Gl,
        "glx" => {
            ecx.span_err(span, "glx generation unimplemented");
            return DummyResult::any(span)
        },
        "wgl" => {
            ecx.span_err(span, "wgl generation unimplemented");
            return DummyResult::any(span)
        }
        ns => {
            ecx.span_err(span, format!("Unexpected opengl namespace '{}'", ns).as_slice());
            return DummyResult::any(span)
        }
    };

    let path = {
        use std::os;
        let p = Path::new(ecx.codemap().span_to_filename(span));
        let mut p = os::make_absolute(&p);
        p.pop(); p.pop(); p.pop();
        p.push(format!("deps/khronos-api/{}.xml", ns_literal));
        p
    };

    let filter = Some(Filter {
        extensions: exts.unwrap_or("".to_string()).as_slice().split(',').map(|s| s.to_string()).collect(),
        profile: profile,
        version: version,
        api: api,
    });

    // generating the registry of all bindings
    let reg = {
        use std::task;

        let mut file = match File::open(&path) {
            Ok(f) => f,
            Err(err) => {
                ecx.span_err(span, format!("{}", err).as_slice());
                return DummyResult::any(span)
            }
        };

        let file_content = match file.read_to_string() {
            Ok(s) => s,
            Err(_) => {
                ecx.span_err(span, "registry source not utf8!");
                return DummyResult::any(span)
            }
        };

        match task::try(proc() Registry::from_xml(file_content.as_slice(), ns, filter)) {
            Ok(reg) => reg,
            Err(err) => {
                use std::any::{Any, AnyRefExt};
                let err: &Any = err;

                match err {
                    err if err.is::<String>() => {
                        ecx.span_err(span, "error while parsing the registry");
                        ecx.span_err(span, err.downcast_ref::<String>().unwrap().as_slice());
                    },
                    err if err.is::<&'static str>() => {
                        ecx.span_err(span, "error while parsing the registry");
                        ecx.span_err(span, err.downcast_ref::<&'static str>().unwrap().as_slice());
                    },
                    _ => {
                        ecx.span_err(span, "unknown error while parsing the registry");
                    }
                }

                return DummyResult::any(span);
            }
        }
    };

    // generating the Rust bindings as a source code into "buffer"
    let buffer = {
        use std::io::MemWriter;
        use std::task;

        // calling the generator
        let result = match generator.as_slice() {
            "static" => task::try(proc() {
                let mut buffer = MemWriter::new();
                StaticGenerator::write(&mut buffer, &reg, ns, false);
                buffer
            }),

            "struct" => task::try(proc() {
                let mut buffer = MemWriter::new();
                StructGenerator::write(&mut buffer, &reg, ns, false);
                buffer
            }),

            generator => {
                ecx.span_err(span, format!("unknown generator type: {}", generator).as_slice());
                return DummyResult::any(span);
            },
        };

        // processing the result
        match result {
            Ok(buffer) => buffer.unwrap(),
            Err(err) => {
                use std::any::{Any, AnyRefExt};
                let err: &Any = err;

                match err {
                    err if err.is::<String>() => {
                        ecx.span_err(span, "error while generating the bindings");
                        ecx.span_err(span, err.downcast_ref::<String>().unwrap().as_slice());
                    },
                    err if err.is::<&'static str>() => {
                        ecx.span_err(span, "error while generating the bindings");
                        ecx.span_err(span, err.downcast_ref::<&'static str>().unwrap().as_slice());
                    },
                    _ => {
                        ecx.span_err(span, "unknown error while generating the bindings");
                    }
                }

                return DummyResult::any(span);
            }
        }
    };

    // creating a new Rust parser from these bindings
    let content = match String::from_utf8(buffer) {
        Ok(s) => s,
        Err(err) => {
            ecx.span_err(span, format!("{}", err).as_slice());
            return DummyResult::any(span)
        }
    };
    let mut parser = ::syntax::parse::new_parser_from_source_str(ecx.parse_sess(), ecx.cfg(), path.display().to_string(), content);

    // getting all the items defined by the bindings
    let mut items = Vec::new();
    loop {
        match parser.parse_item_with_outer_attributes() {
            None => break,
            Some(i) => items.push(i)
        }
    }
    if !parser.eat(&token::EOF) {
        ecx.span_err(span, "the rust parser failed to compile all the generated bindings (meaning there is a bug in this library!)");
        return DummyResult::any(span)
    }
    box MacroResult { content: items } as Box<MacResult>
}

fn parse_macro_arguments(ecx: &mut ExtCtxt, span: Span, tts: &[syntax::ast::TokenTree]) -> Option<(String, String, String, String, String, Option<String>)> {
    let values = match get_exprs_from_tts(ecx, span, tts) {
        Some(v) => v,
        None => return None
    };

    if values.len() != 5 && values.len() != 6 {
        ecx.span_err(span, format!("expected 5 or 6 arguments but only got {}", values.len()).as_slice());
        return None;
    }

    match (
        expr_to_string(ecx, values[0].clone(), "expected string literal").map(|e| match e { (s, _) => s.get().to_string() }),
        expr_to_string(ecx, values[1].clone(), "expected string literal").map(|e| match e { (s, _) => s.get().to_string() }),
        expr_to_string(ecx, values[2].clone(), "expected string literal").map(|e| match e { (s, _) => s.get().to_string() }),
        expr_to_string(ecx, values[3].clone(), "expected string literal").map(|e| match e { (s, _) => s.get().to_string() }),
        expr_to_string(ecx, values[4].clone(), "expected string literal").map(|e| match e { (s, _) => s.get().to_string() }),
        values.as_slice().get(5).and_then(|e| expr_to_string(ecx, *e, "expected string literal")).map(|e| match e { (s, _) => s.get().to_string() }),
    ) {
        (Some(a), Some(b), Some(c), Some(d), Some(e), f) => Some((a, b, c, d, e, f)),
        _ => None
    }
}
<|endoftext|>


<commit_before>use ::bitreader::BitReader;
use std::io::Read;

// For Huffman codes used in the Deflate spec, is seems that the length of a
// code is at most 10 bits (max alphabet size is 704).
// For this simple use case, we don't need/want to deal with type parameters.
pub type Symbol = u16;

#[derive(Debug, Clone, PartialEq)]
pub struct Node {
	left: Option<Box<Tree>>,
	right: Option<Box<Tree>>,
}

#[derive(Debug, Clone, PartialEq)]
pub enum Tree {
	Leaf(Symbol),
	Inner(Node),
}

const EMPTY_NODE: Tree = Tree::Inner(Node{
	left: None,
	right: None,
});

impl Tree {
	pub fn new() -> Tree {
		EMPTY_NODE
	}

	pub fn is_empty(&self) -> bool {
		*self == EMPTY_NODE
	}

	pub fn insert(&mut self, code: Vec<bool>, symbol: Symbol) {
		if code.len() == 1 {
			if code[0] {
				*self = Tree::Inner(Node{
					left: match *self {
						Tree::Inner(Node{
							ref left,
							right: _,
						}) => left.clone(),
						Tree::Leaf(_) => unreachable!(),
					},
					right: Some(Box::new(Tree::Leaf(symbol))),
				});
			} else {
				*self = Tree::Inner(Node{
					left: Some(Box::new(Tree::Leaf(symbol))),
					right: match *self {
						Tree::Inner(Node{
							left: _,
							ref right,
						}) => right.clone(),
						Tree::Leaf(_) => unreachable!(),
					},
				});
			}
		} else {
			if code[0] {
				match *self {
					Tree::Inner(Node{
						left: _,
						ref mut right,
					}) => {
						match *right {
							None => *right = Some(Box::new(Tree::new())),
							Some(_) => {
								// Nothing to do
							},
						};
						// Now we can be sure that right is a subtree. So we can delegate to it.

						match *right {
							Some(ref mut boxed_tree) => (*boxed_tree).insert(code[1..].to_vec(), symbol),
							_ => unreachable!(),
						};
					},
					Tree::Leaf(_) => unreachable!(),
				}
			} else {
				match *self {
					Tree::Inner(Node{
						ref mut left,
						right: _,
					}) => {
						match *left {
							None => *left = Some(Box::new(Tree::new())),
							Some(_) => {
								// Nothing to do
							},
						};
						// Now we can be sure that left is a subtree. So we can delegate to it.

						match *left {
							Some(ref mut boxed_tree) => (*boxed_tree).insert(code[1..].to_vec(), symbol),
							_ => unreachable!(),
						};
					},
					Tree::Leaf(_) => unreachable!(),
				}
			}
		}
	}

	pub fn lookup(&self, c: bool) -> Option<Tree> {
		match *self {
			Tree::Leaf(_) => None,
			Tree::Inner(Node{
				ref left,
				ref right
			}) =>
				if c {
					match *right {
						Some(ref boxed_tree) => Some((**boxed_tree).clone()),
						None => None,
					}
				} else {
					match *left {
						Some(ref boxed_tree) => Some((**boxed_tree).clone()),
						None => None,
					}
				}
		}
	}

	pub fn lookup_symbol<R: Read>(&self, r: &mut BitReader<R>) -> Option<Symbol> {
		loop {
			match r.read_bit() {
				Ok(bit) =>
					match self.lookup(bit) {
						Some(Tree::Leaf(symbol)) => return Some(symbol),
						Some(inner) => return inner.lookup_symbol(r),
						None => unreachable!(),
					},
				Err(_) => return None,
			}
		}
	}
}

// These are not very good tests, because they rely
// on implementation details, for the tree data structure
mod tests {
	#[test]
	fn should_create_empty_tree() {
		use super::Tree;
		assert!(Tree::new().is_empty());
	}

	#[test]
	fn should_create_different_instances() {
		use super::Tree;
		let mut tree_0 = Tree::new();
		let tree_1 = Tree::new();

		tree_0.insert(vec![false], 666);
		assert!(!tree_0.is_empty());
		assert!(tree_1.is_empty());
	}

	#[test]
	fn should_insert_first_level_leaf_on_left() {
		use super::{Tree, Node};
		use super::Tree::{ Inner, Leaf };

		let mut tree = Tree::new();
		tree.insert(vec![false], 666);

		assert_eq!(Inner(Node{
			left: Some(Box::new(Leaf(666))),
			right: None,
		}), tree);
	}

	#[test]
	fn should_insert_first_level_leaf_on_right() {
		use super::{Tree, Node};
		use super::Tree::{ Inner, Leaf };

		let mut tree = Tree::new();
		tree.insert(vec![true], 666);

		assert_eq!(Inner(Node{
			left: None,
			right: Some(Box::new(Leaf(666))),
		}), tree);
	}

	#[test]
	fn should_insert_first_level_leaf_on_left_then_on_right() {
		use super::{Tree, Node};
		use super::Tree::{ Inner, Leaf };

		let mut tree = Tree::new();
		tree.insert(vec![false], 667);
		tree.insert(vec![true], 666);

		assert_eq!(Inner(Node{
			left: Some(Box::new(Leaf(667))),
			right: Some(Box::new(Leaf(666))),
		}), tree);
	}

	#[test]
	fn should_insert_first_level_leaf_on_right_then_on_left() {
		use super::{Tree, Node};
		use super::Tree::{ Inner, Leaf };

		let mut tree = Tree::new();
		tree.insert(vec![true], 666);
		tree.insert(vec![false], 667);

		assert_eq!(Inner(Node{
			left: Some(Box::new(Leaf(667))),
			right: Some(Box::new(Leaf(666))),
		}), tree);
	}

	#[test]
	fn should_insert_second_level_leaf_left_right() {
		use super::{Tree, Node};
		use super::Tree::{ Inner, Leaf };

		let mut tree = Tree::new();
		tree.insert(vec![false, true], 6666);

		assert_eq!(Inner(Node{
			left: Some(Box::new(Inner(Node{
				left: None,
				right: Some(Box::new(Leaf(6666))),
			}))),
			right: None
		}), tree);
	}

	#[test]
	fn should_insert_second_level_leaf_right_left() {
		use super::{Tree, Node};
		use super::Tree::{ Inner, Leaf };

		let mut tree = Tree::new();
		tree.insert(vec![true, false], 6666);

		assert_eq!(Inner(Node{
			left: None,
			right: Some(Box::new(Inner(Node{
				left: Some(Box::new(Leaf(6666))),
				right: None,
			}))),
		}), tree);
	}

	#[test]
	fn should_lookup_first_level_leaf_left() {
		use super::Tree;
		use super::Tree::Leaf;

		let mut tree = Tree::new();
		tree.insert(vec![true, false], 6666);
		tree.insert(vec![false], 666);
		tree.insert(vec![true, true], 6667);

		assert_eq!(Some(Leaf(666)), tree.lookup(false));
	}

	#[test]
	fn should_lookup_first_level_leaf_right() {
		use super::Tree;
		use super::Tree::Leaf;

		let mut tree = Tree::new();
		tree.insert(vec![false, false], 6666);
		tree.insert(vec![true], 666);
		tree.insert(vec![false, true], 6667);

		assert_eq!(Some(Leaf(666)), tree.lookup(true));
	}

	#[test]
	fn should_lookup_first_level_node_left() {
		use super::{ Tree, Node };
		use super::Tree::{ Inner, Leaf };

		let mut tree = Tree::new();
		tree.insert(vec![false, false], 6666);
		tree.insert(vec![true], 666);
		tree.insert(vec![false, true], 6667);

		assert_eq!(Some(Inner(Node{
			left: Some(Box::new(Leaf(6666))),
			right: Some(Box::new(Leaf(6667))),
		})), tree.lookup(false));
	}

	#[test]
	fn should_result_in_none() {
		use super::Tree;

		let mut tree = Tree::new();
		tree.insert(vec![true], 666);

		assert_eq!(None, tree.lookup(false));
	}
}
<commit_msg>MOD: improved tree tests to avoid relying on implementation details<commit_after>use ::bitreader::BitReader;
use std::io::Read;

// For Huffman codes used in the Deflate spec, is seems that the length of a
// code is at most 10 bits (max alphabet size is 704).
// For this simple use case, we don't need/want to deal with type parameters.
pub type Symbol = u16;

#[derive(Debug, Clone, PartialEq)]
pub struct Node {
	left: Option<Box<Tree>>,
	right: Option<Box<Tree>>,
}

#[derive(Debug, Clone, PartialEq)]
pub enum Tree {
	Leaf(Symbol),
	Inner(Node),
}

const EMPTY_NODE: Tree = Tree::Inner(Node{
	left: None,
	right: None,
});

impl Tree {
	pub fn new() -> Tree {
		EMPTY_NODE
	}

	pub fn is_empty(&self) -> bool {
		*self == EMPTY_NODE
	}

	pub fn insert(&mut self, code: Vec<bool>, symbol: Symbol) {
		if code.len() == 1 {
			if code[0] {
				*self = Tree::Inner(Node{
					left: match *self {
						Tree::Inner(Node{
							ref left,
							right: _,
						}) => left.clone(),
						Tree::Leaf(_) => unreachable!(),
					},
					right: Some(Box::new(Tree::Leaf(symbol))),
				});
			} else {
				*self = Tree::Inner(Node{
					left: Some(Box::new(Tree::Leaf(symbol))),
					right: match *self {
						Tree::Inner(Node{
							left: _,
							ref right,
						}) => right.clone(),
						Tree::Leaf(_) => unreachable!(),
					},
				});
			}
		} else {
			if code[0] {
				match *self {
					Tree::Inner(Node{
						left: _,
						ref mut right,
					}) => {
						match *right {
							None => *right = Some(Box::new(Tree::new())),
							Some(_) => {
								// Nothing to do
							},
						};
						// Now we can be sure that right is a subtree. So we can delegate to it.

						match *right {
							Some(ref mut boxed_tree) => (*boxed_tree).insert(code[1..].to_vec(), symbol),
							_ => unreachable!(),
						};
					},
					Tree::Leaf(_) => unreachable!(),
				}
			} else {
				match *self {
					Tree::Inner(Node{
						ref mut left,
						right: _,
					}) => {
						match *left {
							None => *left = Some(Box::new(Tree::new())),
							Some(_) => {
								// Nothing to do
							},
						};
						// Now we can be sure that left is a subtree. So we can delegate to it.

						match *left {
							Some(ref mut boxed_tree) => (*boxed_tree).insert(code[1..].to_vec(), symbol),
							_ => unreachable!(),
						};
					},
					Tree::Leaf(_) => unreachable!(),
				}
			}
		}
	}

	pub fn lookup(&self, c: bool) -> Option<Tree> {
		match *self {
			Tree::Leaf(_) => None,
			Tree::Inner(Node{
				ref left,
				ref right
			}) =>
				if c {
					match *right {
						Some(ref boxed_tree) => Some((**boxed_tree).clone()),
						None => None,
					}
				} else {
					match *left {
						Some(ref boxed_tree) => Some((**boxed_tree).clone()),
						None => None,
					}
				}
		}
	}

	pub fn lookup_symbol<R: Read>(&self, r: &mut BitReader<R>) -> Option<Symbol> {
		loop {
			match r.read_bit() {
				Ok(bit) =>
					match self.lookup(bit) {
						Some(Tree::Leaf(symbol)) => return Some(symbol),
						Some(inner) => return inner.lookup_symbol(r),
						None => unreachable!(),
					},
				Err(_) => return None,
			}
		}
	}
}

mod tests {
	#[test]
	fn should_create_empty_tree() {
		use super::Tree;
		assert!(Tree::new().is_empty());
	}

	#[test]
	fn should_create_different_instances() {
		use super::Tree;
		let mut tree_0 = Tree::new();
		let tree_1 = Tree::new();

		tree_0.insert(vec![false], 666);
		assert!(!tree_0.is_empty());
		assert!(tree_1.is_empty());
	}

	#[test]
	fn should_insert_and_lookup_first_level_leaf_on_left() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![0]));
		let mut tree = Tree::new();
		tree.insert(vec![false], 666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
	}

	#[test]
	fn should_insert_and_lookup_first_level_leaf_on_right() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![1]));
		let mut tree = Tree::new();
		tree.insert(vec![true], 666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
	}

	#[test]
	fn should_insert_first_level_leaf_on_left_then_on_right() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![2]));
		let mut tree = Tree::new();
		tree.insert(vec![false], 667);
		tree.insert(vec![true], 666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(667));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
	}

	#[test]
	fn should_insert_first_level_leaf_on_right_then_on_left() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![1]));
		let mut tree = Tree::new();
		tree.insert(vec![true], 666);
		tree.insert(vec![false], 667);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(667));
	}

	#[test]
	fn should_insert_second_level_leaf_left_right() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![2]));
		let mut tree = Tree::new();
		tree.insert(vec![false, true], 6666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6666));
	}

	#[test]
	fn should_insert_second_level_leaf_right_left() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![1]));
		let mut tree = Tree::new();
		tree.insert(vec![true, false], 6666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6666));
	}

	#[test]
	fn should_lookup_first_level_leaf_left() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![0b11001]));
		let mut tree = Tree::new();
		tree.insert(vec![true, false], 6666);
		tree.insert(vec![false], 666);
		tree.insert(vec![true, true], 6667);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6667));
	}

	#[test]
	fn should_lookup_first_level_leaf_right() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![0b10100]));
		let mut tree = Tree::new();
		tree.insert(vec![false, false], 6666);
		tree.insert(vec![true], 666);
		tree.insert(vec![false, true], 6667);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6667));
	}

	#[test]
	fn should_result_in_none() {
		use super::Tree;

		let mut tree = Tree::new();
		tree.insert(vec![true], 666);

		assert_eq!(None, tree.lookup(false));
	}
}
<|endoftext|>


<commit_before>use ::bitreader::BitReader;
use std::io::Read;

// For Huffman codes used in the Deflate spec, is seems that the length of a
// code is at most 10 bits (max alphabet size is 704).
// For this simple use case, we don't need/want to deal with type parameters.
pub type Symbol = u16;

#[derive(Debug, Clone, PartialEq)]
pub struct Node {
	left: Option<Box<Tree>>,
	right: Option<Box<Tree>>,
}

#[derive(Debug, Clone, PartialEq)]
pub enum Tree {
	Leaf(Symbol),
	Inner(Node),
}

const EMPTY_NODE: Tree = Tree::Inner(Node{
	left: None,
	right: None,
});

impl Tree {
	pub fn new() -> Tree {
		EMPTY_NODE
	}

	pub fn insert(&mut self, code: Vec<bool>, symbol: Symbol) {
		if code.len() == 1 {
			if code[0] {
				*self = Tree::Inner(Node{
					left: match *self {
						Tree::Inner(Node{
							ref left,
							right: _,
						}) => left.clone(),
						Tree::Leaf(_) => unreachable!(),
					},
					right: Some(Box::new(Tree::Leaf(symbol))),
				});
			} else {
				*self = Tree::Inner(Node{
					left: Some(Box::new(Tree::Leaf(symbol))),
					right: match *self {
						Tree::Inner(Node{
							left: _,
							ref right,
						}) => right.clone(),
						Tree::Leaf(_) => unreachable!(),
					},
				});
			}
		} else {
			if code[0] {
				match *self {
					Tree::Inner(Node{
						left: _,
						ref mut right,
					}) => {
						match *right {
							None => *right = Some(Box::new(Tree::new())),
							Some(_) => {
								// Nothing to do
							},
						};
						// Now we can be sure that right is a subtree. So we can delegate to it.

						match *right {
							Some(ref mut boxed_tree) => (*boxed_tree).insert(code[1..].to_vec(), symbol),
							_ => unreachable!(),
						};
					},
					Tree::Leaf(_) => unreachable!(),
				}
			} else {
				match *self {
					Tree::Inner(Node{
						ref mut left,
						right: _,
					}) => {
						match *left {
							None => *left = Some(Box::new(Tree::new())),
							Some(_) => {
								// Nothing to do
							},
						};
						// Now we can be sure that left is a subtree. So we can delegate to it.

						match *left {
							Some(ref mut boxed_tree) => (*boxed_tree).insert(code[1..].to_vec(), symbol),
							_ => unreachable!(),
						};
					},
					Tree::Leaf(_) => unreachable!(),
				}
			}
		}
	}

	pub fn lookup(&self, c: bool) -> Option<Tree> {
		match *self {
			Tree::Leaf(_) => None,
			Tree::Inner(Node{
				ref left,
				ref right
			}) =>
				if c {
					match *right {
						Some(ref boxed_tree) => Some((**boxed_tree).clone()),
						None => None,
					}
				} else {
					match *left {
						Some(ref boxed_tree) => Some((**boxed_tree).clone()),
						None => None,
					}
				}
		}
	}

	pub fn lookup_symbol<R: Read>(&self, r: &mut BitReader<R>) -> Option<Symbol> {
		loop {
			match r.read_bit() {
				Ok(bit) =>
					match self.lookup(bit) {
						Some(Tree::Leaf(symbol)) => return Some(symbol),
						Some(inner) => return inner.lookup_symbol(r),
						None => unreachable!(),
					},
				Err(_) => return None,
			}
		}
	}
}

mod tests {
	#[test]
	fn should_create_empty_tree() {
		use super::Tree;
		assert_eq!(Tree::new(), super::EMPTY_NODE);
	}

	#[test]
	fn should_create_different_instances() {
		use super::Tree;
		let mut tree_0 = Tree::new();
		let tree_1 = Tree::new();

		tree_0.insert(vec![false], 666);
		assert!(tree_0 != super::EMPTY_NODE);
		assert!(tree_1 == super::EMPTY_NODE);
	}

	#[test]
	fn should_insert_and_lookup_first_level_leaf_on_left() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![0]));
		let mut tree = Tree::new();
		tree.insert(vec![false], 666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
	}

	#[test]
	fn should_insert_and_lookup_first_level_leaf_on_right() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![1]));
		let mut tree = Tree::new();
		tree.insert(vec![true], 666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
	}

	#[test]
	fn should_insert_first_level_leaf_on_left_then_on_right() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![2]));
		let mut tree = Tree::new();
		tree.insert(vec![false], 667);
		tree.insert(vec![true], 666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(667));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
	}

	#[test]
	fn should_insert_first_level_leaf_on_right_then_on_left() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![1]));
		let mut tree = Tree::new();
		tree.insert(vec![true], 666);
		tree.insert(vec![false], 667);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(667));
	}

	#[test]
	fn should_insert_second_level_leaf_left_right() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![2]));
		let mut tree = Tree::new();
		tree.insert(vec![false, true], 6666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6666));
	}

	#[test]
	fn should_insert_second_level_leaf_right_left() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![1]));
		let mut tree = Tree::new();
		tree.insert(vec![true, false], 6666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6666));
	}

	#[test]
	fn should_lookup_first_level_leaf_left() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![0b11001]));
		let mut tree = Tree::new();
		tree.insert(vec![true, false], 6666);
		tree.insert(vec![false], 666);
		tree.insert(vec![true, true], 6667);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6667));
	}

	#[test]
	fn should_lookup_first_level_leaf_right() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![0b10100]));
		let mut tree = Tree::new();
		tree.insert(vec![false, false], 6666);
		tree.insert(vec![true], 666);
		tree.insert(vec![false, true], 6667);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6667));
	}

	#[test]
	fn should_result_in_none() {
		use super::Tree;

		let mut tree = Tree::new();
		tree.insert(vec![true], 666);

		assert_eq!(None, tree.lookup(false));
	}
}
<commit_msg>MOD: lookup helper function does not need to be public<commit_after>use ::bitreader::BitReader;
use std::io::Read;

// For Huffman codes used in the Deflate spec, is seems that the length of a
// code is at most 10 bits (max alphabet size is 704).
// For this simple use case, we don't need/want to deal with type parameters.
pub type Symbol = u16;

#[derive(Debug, Clone, PartialEq)]
pub struct Node {
	left: Option<Box<Tree>>,
	right: Option<Box<Tree>>,
}

#[derive(Debug, Clone, PartialEq)]
pub enum Tree {
	Leaf(Symbol),
	Inner(Node),
}

const EMPTY_NODE: Tree = Tree::Inner(Node{
	left: None,
	right: None,
});

impl Tree {
	pub fn new() -> Tree {
		EMPTY_NODE
	}

	pub fn insert(&mut self, code: Vec<bool>, symbol: Symbol) {
		if code.len() == 1 {
			if code[0] {
				*self = Tree::Inner(Node{
					left: match *self {
						Tree::Inner(Node{
							ref left,
							right: _,
						}) => left.clone(),
						Tree::Leaf(_) => unreachable!(),
					},
					right: Some(Box::new(Tree::Leaf(symbol))),
				});
			} else {
				*self = Tree::Inner(Node{
					left: Some(Box::new(Tree::Leaf(symbol))),
					right: match *self {
						Tree::Inner(Node{
							left: _,
							ref right,
						}) => right.clone(),
						Tree::Leaf(_) => unreachable!(),
					},
				});
			}
		} else {
			if code[0] {
				match *self {
					Tree::Inner(Node{
						left: _,
						ref mut right,
					}) => {
						match *right {
							None => *right = Some(Box::new(Tree::new())),
							Some(_) => {
								// Nothing to do
							},
						};
						// Now we can be sure that right is a subtree. So we can delegate to it.

						match *right {
							Some(ref mut boxed_tree) => (*boxed_tree).insert(code[1..].to_vec(), symbol),
							_ => unreachable!(),
						};
					},
					Tree::Leaf(_) => unreachable!(),
				}
			} else {
				match *self {
					Tree::Inner(Node{
						ref mut left,
						right: _,
					}) => {
						match *left {
							None => *left = Some(Box::new(Tree::new())),
							Some(_) => {
								// Nothing to do
							},
						};
						// Now we can be sure that left is a subtree. So we can delegate to it.

						match *left {
							Some(ref mut boxed_tree) => (*boxed_tree).insert(code[1..].to_vec(), symbol),
							_ => unreachable!(),
						};
					},
					Tree::Leaf(_) => unreachable!(),
				}
			}
		}
	}

	fn lookup(&self, c: bool) -> Option<Tree> {
		match *self {
			Tree::Leaf(_) => None,
			Tree::Inner(Node{
				ref left,
				ref right
			}) =>
				if c {
					match *right {
						Some(ref boxed_tree) => Some((**boxed_tree).clone()),
						None => None,
					}
				} else {
					match *left {
						Some(ref boxed_tree) => Some((**boxed_tree).clone()),
						None => None,
					}
				}
		}
	}

	pub fn lookup_symbol<R: Read>(&self, r: &mut BitReader<R>) -> Option<Symbol> {
		loop {
			match r.read_bit() {
				Ok(bit) =>
					match self.lookup(bit) {
						Some(Tree::Leaf(symbol)) => return Some(symbol),
						Some(inner) => return inner.lookup_symbol(r),
						None => unreachable!(),
					},
				Err(_) => return None,
			}
		}
	}
}

mod tests {
	#[test]
	fn should_create_empty_tree() {
		use super::Tree;
		assert_eq!(Tree::new(), super::EMPTY_NODE);
	}

	#[test]
	fn should_create_different_instances() {
		use super::Tree;
		let mut tree_0 = Tree::new();
		let tree_1 = Tree::new();

		tree_0.insert(vec![false], 666);
		assert!(tree_0 != super::EMPTY_NODE);
		assert!(tree_1 == super::EMPTY_NODE);
	}

	#[test]
	fn should_insert_and_lookup_first_level_leaf_on_left() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![0]));
		let mut tree = Tree::new();
		tree.insert(vec![false], 666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
	}

	#[test]
	fn should_insert_and_lookup_first_level_leaf_on_right() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![1]));
		let mut tree = Tree::new();
		tree.insert(vec![true], 666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
	}

	#[test]
	fn should_insert_first_level_leaf_on_left_then_on_right() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![2]));
		let mut tree = Tree::new();
		tree.insert(vec![false], 667);
		tree.insert(vec![true], 666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(667));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
	}

	#[test]
	fn should_insert_first_level_leaf_on_right_then_on_left() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![1]));
		let mut tree = Tree::new();
		tree.insert(vec![true], 666);
		tree.insert(vec![false], 667);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(667));
	}

	#[test]
	fn should_insert_second_level_leaf_left_right() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![2]));
		let mut tree = Tree::new();
		tree.insert(vec![false, true], 6666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6666));
	}

	#[test]
	fn should_insert_second_level_leaf_right_left() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![1]));
		let mut tree = Tree::new();
		tree.insert(vec![true, false], 6666);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6666));
	}

	#[test]
	fn should_lookup_first_level_leaf_left() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![0b11001]));
		let mut tree = Tree::new();
		tree.insert(vec![true, false], 6666);
		tree.insert(vec![false], 666);
		tree.insert(vec![true, true], 6667);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6667));
	}

	#[test]
	fn should_lookup_first_level_leaf_right() {
		use ::bitreader::BitReader;
		use super::Tree;
		use std::io::Cursor;

		let mut lookup_stream = BitReader::new(Cursor::new(vec![0b10100]));
		let mut tree = Tree::new();
		tree.insert(vec![false, false], 6666);
		tree.insert(vec![true], 666);
		tree.insert(vec![false, true], 6667);

		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(666));
		assert_eq!(tree.lookup_symbol(&mut lookup_stream), Some(6667));
	}

	#[test]
	fn should_result_in_none() {
		use super::Tree;

		let mut tree = Tree::new();
		tree.insert(vec![true], 666);

		assert_eq!(None, tree.lookup(false));
	}
}
<|endoftext|>


<commit_before>#![cfg(feature = "instructions")]

//! Special x86_64 instructions.

pub mod interrupts;
pub mod port;
pub mod random;
pub mod segmentation;
pub mod tables;
pub mod tlb;

/// Halts the CPU until the next interrupt arrives.
#[inline]
pub fn hlt() {
    #[cfg(feature = "inline_asm")]
    unsafe {
        asm!("hlt", options(nomem, nostack));
    }

    #[cfg(not(feature = "inline_asm"))]
    unsafe {
        crate::asm::x86_64_asm_hlt();
    }
}

/// Executes the `nop` instructions, which performs no operation (i.e. does nothing).
///
/// This operation is useful to work around the LLVM bug that endless loops are illegally
/// optimized away (see [the issue](https://github.com/rust-lang/rust/issues/28728)). By invoking this
/// instruction (which is marked as volatile), the compiler should no longer optimize the
/// endless loop away.
#[inline]
pub fn nop() {
    #[cfg(feature = "inline_asm")]
    unsafe {
        asm!("nop", options(nomem, nostack, preserves_flags));
    }

    #[cfg(not(feature = "inline_asm"))]
    unsafe {
        crate::asm::x86_64_asm_nop();
    }
}

/// Emits a '[magic breakpoint](https://wiki.osdev.org/Bochs#Magic_Breakpoint)' instruction for the [Bochs](http://bochs.sourceforge.net/) CPU
/// emulator. Make sure to set `magic_break: enabled=1` in your `.bochsrc` file.
#[cfg(feature = "inline_asm")]
#[inline]
pub fn bochs_breakpoint() {
    unsafe {
        asm!("xchgw bx, bx", options(nomem, nostack));
    }
}

/// Gets the current instruction pointer. Note that this is only approximate as it requires a few
/// instructions to execute.
#[cfg(feature = "inline_asm")]
#[inline(always)]
pub fn read_rip() -> u64 {
    let rip: u64;
    unsafe {
        asm!(
            "lea {}, [rip]", out(reg) rip, options(nostack, nomem)
        );
    }
    rip
}
<commit_msg>Fix bad conversion from llvm_asm! to asm! (#218)<commit_after>#![cfg(feature = "instructions")]

//! Special x86_64 instructions.

pub mod interrupts;
pub mod port;
pub mod random;
pub mod segmentation;
pub mod tables;
pub mod tlb;

/// Halts the CPU until the next interrupt arrives.
#[inline]
pub fn hlt() {
    #[cfg(feature = "inline_asm")]
    unsafe {
        asm!("hlt", options(nomem, nostack));
    }

    #[cfg(not(feature = "inline_asm"))]
    unsafe {
        crate::asm::x86_64_asm_hlt();
    }
}

/// Executes the `nop` instructions, which performs no operation (i.e. does nothing).
///
/// This operation is useful to work around the LLVM bug that endless loops are illegally
/// optimized away (see [the issue](https://github.com/rust-lang/rust/issues/28728)). By invoking this
/// instruction (which is marked as volatile), the compiler should no longer optimize the
/// endless loop away.
#[inline]
pub fn nop() {
    #[cfg(feature = "inline_asm")]
    unsafe {
        asm!("nop", options(nomem, nostack, preserves_flags));
    }

    #[cfg(not(feature = "inline_asm"))]
    unsafe {
        crate::asm::x86_64_asm_nop();
    }
}

/// Emits a '[magic breakpoint](https://wiki.osdev.org/Bochs#Magic_Breakpoint)' instruction for the [Bochs](http://bochs.sourceforge.net/) CPU
/// emulator. Make sure to set `magic_break: enabled=1` in your `.bochsrc` file.
#[cfg(feature = "inline_asm")]
#[inline]
pub fn bochs_breakpoint() {
    unsafe {
        asm!("xchg bx, bx", options(nomem, nostack));
    }
}

/// Gets the current instruction pointer. Note that this is only approximate as it requires a few
/// instructions to execute.
#[cfg(feature = "inline_asm")]
#[inline(always)]
pub fn read_rip() -> u64 {
    let rip: u64;
    unsafe {
        asm!(
            "lea {}, [rip]", out(reg) rip, options(nostack, nomem)
        );
    }
    rip
}
<|endoftext|>


<commit_before>use flate2::read::ZlibDecoder;
use serde_json;
use websocket::message::OwnedMessage;
use websocket::sync::stream::{TcpStream, TlsStream};
use websocket::sync::Client as WsClient;
use gateway::GatewayError;
use internal::prelude::*;

pub trait ReceiverExt {
    fn recv_json<F, T>(&mut self, decode: F) -> Result<T> where F: FnOnce(Value) -> Result<T>;
}

pub trait SenderExt {
    fn send_json(&mut self, value: &Value) -> Result<()>;
}

impl ReceiverExt for WsClient<TlsStream<TcpStream>> {
    fn recv_json<F, T>(&mut self, decode: F) -> Result<T>
        where F: FnOnce(Value) -> Result<T> {
        let message = self.recv_message()?;

        let res = match message {
            OwnedMessage::Binary(bytes) => {
                let value = serde_json::from_reader(ZlibDecoder::new(&bytes[..]))?;

                Some(decode(value).map_err(|why| {
                    let s = String::from_utf8_lossy(&bytes);

                    warn!("(╯°□°）╯︵ ┻━┻ Error decoding: {}", s);

                    why
                }))
            },
            OwnedMessage::Close(data) => Some(Err(Error::Gateway(GatewayError::Closed(data)))),
            OwnedMessage::Text(payload) => {
                let value = serde_json::from_str(&payload)?;

                Some(decode(value).map_err(|why| {
                    warn!("(╯°□°）╯︵ ┻━┻ Error decoding: {}", payload);

                    why
                }))
            },
            OwnedMessage::Ping(x) => {
                self.send_message(&OwnedMessage::Pong(x))
                    .map_err(Error::from)?;

                None
            },
            OwnedMessage::Pong(_) => None,
        };

        res.unwrap()
    }
}

impl SenderExt for WsClient<TlsStream<TcpStream>> {
    fn send_json(&mut self, value: &Value) -> Result<()> {
        serde_json::to_string(value)
            .map(OwnedMessage::Text)
            .map_err(Error::from)
            .and_then(|m| self.send_message(&m).map_err(Error::from))
    }
}
<commit_msg>Handle the `None`s better<commit_after>use flate2::read::ZlibDecoder;
use serde_json;
use websocket::message::OwnedMessage;
use websocket::sync::stream::{TcpStream, TlsStream};
use websocket::sync::Client as WsClient;
use gateway::GatewayError;
use internal::prelude::*;

pub trait ReceiverExt {
    fn recv_json<F, T>(&mut self, decode: F) -> Result<T> where F: Fn(Value) -> Result<T>;
}

pub trait SenderExt {
    fn send_json(&mut self, value: &Value) -> Result<()>;
}

impl ReceiverExt for WsClient<TlsStream<TcpStream>> {
    fn recv_json<F, T>(&mut self, decode: F) -> Result<T>
        where F: Fn(Value) -> Result<T> {
        let message = self.recv_message()?;

        let res = match message {
            OwnedMessage::Binary(bytes) => {
                let value = serde_json::from_reader(ZlibDecoder::new(&bytes[..]))?;

                Some(decode(value).map_err(|why| {
                    let s = String::from_utf8_lossy(&bytes);

                    warn!("(╯°□°）╯︵ ┻━┻ Error decoding: {}", s);

                    why
                }))
            },
            OwnedMessage::Close(data) => Some(Err(Error::Gateway(GatewayError::Closed(data)))),
            OwnedMessage::Text(payload) => {
                let value = serde_json::from_str(&payload)?;

                Some(decode(value).map_err(|why| {
                    warn!("(╯°□°）╯︵ ┻━┻ Error decoding: {}", payload);

                    why
                }))
            },
            OwnedMessage::Ping(x) => {
                self.send_message(&OwnedMessage::Pong(x))
                    .map_err(Error::from)?;

                None
            },
            OwnedMessage::Pong(_) => None,
        };

        // As to ignore the `None`s returned from `Ping` and `Pong`.
        // Since they're essentially useless to us anyway.
        match res {
            Some(data) => data,
            None => self.recv_json(decode),
        }
    }
}

impl SenderExt for WsClient<TlsStream<TcpStream>> {
    fn send_json(&mut self, value: &Value) -> Result<()> {
        serde_json::to_string(value)
            .map(OwnedMessage::Text)
            .map_err(Error::from)
            .and_then(|m| self.send_message(&m).map_err(Error::from))
    }
}
<|endoftext|>


<commit_before>use crate::result::*;
use chrono::{DateTime, Utc};
use serde_json::Value;
use serde_urlencoded;

mod types;

pub use self::types::*;
use crate::client::{Client, Method};
use crate::response::{parse_list_elements_json, parse_object_json, ListResponse};

/// Query parameters for `get_all_released_orders`

#[derive(Debug, Serialize, Default)]
#[allow(non_snake_case)]
pub struct ReleasedQueryParams {
  pub limit: Option<i32>,
  pub createdStartDate: Option<DateTime<Utc>>,
  pub nextCursor: Option<String>,
}

/// Query parameters for `get_all_orders`
#[derive(Debug, Serialize, Default)]
#[allow(non_snake_case)]
pub struct QueryParams {
  pub sku: Option<String>,
  pub customerOrderId: Option<String>,
  pub purchaseOrderId: Option<String>,
  pub status: Option<String>,
  pub createdStartDate: Option<DateTime<Utc>>,
  pub createdEndDate: Option<DateTime<Utc>>,
  pub fromExpectedShipDate: Option<DateTime<Utc>>,
  pub toExpectedShipDate: Option<DateTime<Utc>>,
  pub limit: Option<i32>,
  pub nextCursor: Option<String>,
  pub shipNodeType: Option<String>,
}

#[derive(Debug, Clone)]
#[allow(non_snake_case)]
pub struct ShipParams {
  pub lineNumber: String,
  pub shipDateTime: DateTime<Utc>,
  pub carrierName: Option<String>,
  pub methodCode: String,
  pub trackingNumber: String,
  pub trackingURL: String,
  pub otherCarrier: Option<String>,
  pub unitOfMeasurement: Option<String>,
  pub amount: Option<String>,
  pub shipFromCountry: String,
}

impl ShipParams {
  pub fn to_value(&self) -> Value {
    let timestamp: i64 =
      self.shipDateTime.timestamp() * 1000 + self.shipDateTime.timestamp_subsec_millis() as i64;
    json!({
      "lineNumber": self.lineNumber,
      "orderLineStatuses": {
        "shipFromCountry": self.shipFromCountry,
        "orderLineStatus": [
          {
            "status": "Shipped",
            "statusQuantity": {
              "unitOfMeasurement": self.unitOfMeasurement.clone().unwrap_or_else(|| "EACH".to_owned()),
              "amount": self.amount.clone().unwrap_or_else(|| "1".to_owned()),
            },
            "trackingInfo": {
              "shipDateTime": timestamp,
              "carrierName": {
                "otherCarrier": self.otherCarrier,
                "carrier": self.carrierName,
              },
              "methodCode": self.methodCode,
              "trackingNumber": self.trackingNumber,
              "trackingURL": self.trackingURL
            }
          }
        ]
      }
    })
  }
}

pub type OrderList = ListResponse<Order>;

impl Client {
  pub fn get_all_released_orders(&self, params: &ReleasedQueryParams) -> WalmartResult<OrderList> {
    let qs = serde_urlencoded::to_string(params)?;
    let mut res = self.send(self.request_json(Method::GET, "/v3/orders/released", qs)?)?;
    parse_list_elements_json(res.status(), &mut res, "order").map_err(Into::into)
  }

  pub fn get_all_orders(&self, params: &QueryParams) -> WalmartResult<OrderList> {
    let mut res = self.send(self.request_json(
      Method::GET,
      "/v3/orders",
      serde_urlencoded::to_string(params)?,
    )?)?;
    parse_list_elements_json(res.status(), &mut res, "order").map_err(Into::into)
  }

  pub fn get_all_orders_by_next_cursor(&self, next_cursor: &str) -> WalmartResult<OrderList> {
    use url::form_urlencoded;
    let mut res = self.send(
      self.request_json(
        Method::GET,
        "/v3/orders",
        form_urlencoded::parse((&next_cursor[1..]).as_bytes())
          .into_owned()
          .collect::<Vec<_>>(),
      )?,
    )?;
    parse_list_elements_json(res.status(), &mut res, "order").map_err(Into::into)
  }

  pub fn get_order(&self, purchase_order_id: &str) -> WalmartResult<Order> {
    let path = format!("/v3/orders/{}", purchase_order_id);
    let mut res = self.send(self.request_json(Method::GET, &path, ())?)?;
    parse_object_json(res.status(), &mut res, "order").map_err(Into::into)
  }

  pub fn ack_order(&self, purchase_order_id: &str) -> WalmartResult<Order> {
    let path = format!("/v3/orders/{}/acknowledge", purchase_order_id);
    let mut res = self.send(
      self
        .request_json(Method::POST, &path, ())?
        .json(&Vec::<i32>::new()),
    )?;
    parse_object_json(res.status(), &mut res, "order").map_err(Into::into)
  }

  pub fn ship_order_line(
    &self,
    purchase_order_id: &str,
    line: &ShipParams,
  ) -> WalmartResult<Order> {
    self.ship_order(purchase_order_id, &[line.clone()])
  }

  pub fn ship_order(&self, purchase_order_id: &str, lines: &[ShipParams]) -> WalmartResult<Order> {
    let line_values: Vec<_> = lines.into_iter().map(ShipParams::to_value).collect();
    let body = json!({
      "orderShipment": {
        "orderLines": {
          "orderLine": line_values,
        }
      }
    });
    let path = format!("/v3/orders/{}/shipping", purchase_order_id);
    let mut res = self.send(
      self
        .request_json(
          Method::POST,
          &path,
          vec![("purchaseOrderId", purchase_order_id)],
        )?
        .json(&body),
    )?;
    parse_object_json(res.status(), &mut res, "order").map_err(Into::into)
  }
}
<commit_msg>fix shipFromCountry<commit_after>use crate::result::*;
use chrono::{DateTime, Utc};
use serde_json::Value;
use serde_urlencoded;

mod types;

pub use self::types::*;
use crate::client::{Client, Method};
use crate::response::{parse_list_elements_json, parse_object_json, ListResponse};

/// Query parameters for `get_all_released_orders`

#[derive(Debug, Serialize, Default)]
#[allow(non_snake_case)]
pub struct ReleasedQueryParams {
  pub limit: Option<i32>,
  pub createdStartDate: Option<DateTime<Utc>>,
  pub nextCursor: Option<String>,
}

/// Query parameters for `get_all_orders`
#[derive(Debug, Serialize, Default)]
#[allow(non_snake_case)]
pub struct QueryParams {
  pub sku: Option<String>,
  pub customerOrderId: Option<String>,
  pub purchaseOrderId: Option<String>,
  pub status: Option<String>,
  pub createdStartDate: Option<DateTime<Utc>>,
  pub createdEndDate: Option<DateTime<Utc>>,
  pub fromExpectedShipDate: Option<DateTime<Utc>>,
  pub toExpectedShipDate: Option<DateTime<Utc>>,
  pub limit: Option<i32>,
  pub nextCursor: Option<String>,
  pub shipNodeType: Option<String>,
}

#[derive(Debug, Clone)]
#[allow(non_snake_case)]
pub struct ShipParams {
  pub lineNumber: String,
  pub shipDateTime: DateTime<Utc>,
  pub carrierName: Option<String>,
  pub methodCode: String,
  pub trackingNumber: String,
  pub trackingURL: String,
  pub otherCarrier: Option<String>,
  pub unitOfMeasurement: Option<String>,
  pub amount: Option<String>,
  pub shipFromCountry: String,
}

impl ShipParams {
  pub fn to_value(&self) -> Value {
    let timestamp: i64 =
      self.shipDateTime.timestamp() * 1000 + self.shipDateTime.timestamp_subsec_millis() as i64;
    json!({
      "lineNumber": self.lineNumber,
      "shipFromCountry": self.shipFromCountry,
      "orderLineStatuses": {
        "orderLineStatus": [
          {
            "status": "Shipped",
            "statusQuantity": {
              "unitOfMeasurement": self.unitOfMeasurement.clone().unwrap_or_else(|| "EACH".to_owned()),
              "amount": self.amount.clone().unwrap_or_else(|| "1".to_owned()),
            },
            "trackingInfo": {
              "shipDateTime": timestamp,
              "carrierName": {
                "otherCarrier": self.otherCarrier,
                "carrier": self.carrierName,
              },
              "methodCode": self.methodCode,
              "trackingNumber": self.trackingNumber,
              "trackingURL": self.trackingURL
            }
          }
        ]
      }
    })
  }
}

pub type OrderList = ListResponse<Order>;

impl Client {
  pub fn get_all_released_orders(&self, params: &ReleasedQueryParams) -> WalmartResult<OrderList> {
    let qs = serde_urlencoded::to_string(params)?;
    let mut res = self.send(self.request_json(Method::GET, "/v3/orders/released", qs)?)?;
    parse_list_elements_json(res.status(), &mut res, "order").map_err(Into::into)
  }

  pub fn get_all_orders(&self, params: &QueryParams) -> WalmartResult<OrderList> {
    let mut res = self.send(self.request_json(
      Method::GET,
      "/v3/orders",
      serde_urlencoded::to_string(params)?,
    )?)?;
    parse_list_elements_json(res.status(), &mut res, "order").map_err(Into::into)
  }

  pub fn get_all_orders_by_next_cursor(&self, next_cursor: &str) -> WalmartResult<OrderList> {
    use url::form_urlencoded;
    let mut res = self.send(
      self.request_json(
        Method::GET,
        "/v3/orders",
        form_urlencoded::parse((&next_cursor[1..]).as_bytes())
          .into_owned()
          .collect::<Vec<_>>(),
      )?,
    )?;
    parse_list_elements_json(res.status(), &mut res, "order").map_err(Into::into)
  }

  pub fn get_order(&self, purchase_order_id: &str) -> WalmartResult<Order> {
    let path = format!("/v3/orders/{}", purchase_order_id);
    let mut res = self.send(self.request_json(Method::GET, &path, ())?)?;
    parse_object_json(res.status(), &mut res, "order").map_err(Into::into)
  }

  pub fn ack_order(&self, purchase_order_id: &str) -> WalmartResult<Order> {
    let path = format!("/v3/orders/{}/acknowledge", purchase_order_id);
    let mut res = self.send(
      self
        .request_json(Method::POST, &path, ())?
        .json(&Vec::<i32>::new()),
    )?;
    parse_object_json(res.status(), &mut res, "order").map_err(Into::into)
  }

  pub fn ship_order_line(
    &self,
    purchase_order_id: &str,
    line: &ShipParams,
  ) -> WalmartResult<Order> {
    self.ship_order(purchase_order_id, &[line.clone()])
  }

  pub fn ship_order(&self, purchase_order_id: &str, lines: &[ShipParams]) -> WalmartResult<Order> {
    let line_values: Vec<_> = lines.into_iter().map(ShipParams::to_value).collect();
    let body = json!({
      "orderShipment": {
        "orderLines": {
          "orderLine": line_values,
        }
      }
    });
    let path = format!("/v3/orders/{}/shipping", purchase_order_id);
    let mut res = self.send(
      self
        .request_json(
          Method::POST,
          &path,
          vec![("purchaseOrderId", purchase_order_id)],
        )?
        .json(&body),
    )?;
    parse_object_json(res.status(), &mut res, "order").map_err(Into::into)
  }
}
<|endoftext|>