<commit_before>#![crate_name = "mv"]
#![feature(slice_patterns, str_char)]
#![allow(deprecated)]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Orvar Segerström <orvarsegerstrom@gmail.com>
 * (c) Sokovikov Evgeniy  <skv-headless@yandex.ru>
 *
 * For the full copyright and license information, please view the LICENSE file
 * that was distributed with this source code.
 */

extern crate getopts;
extern crate libc;

use std::fs;
use std::io::{BufRead, BufReader, Result, stdin, Write};
use std::os::unix::fs::MetadataExt;
use std::path::{Path, PathBuf};

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

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

use filesystem::UUPathExt;

static NAME: &'static str = "mv";
static VERSION: &'static str = "0.0.1";

pub struct Behaviour {
    overwrite: OverwriteMode,
    backup: BackupMode,
    suffix: String,
    update: bool,
    target_dir: Option<String>,
    no_target_dir: bool,
    verbose: bool,
}

#[derive(Clone, Eq, PartialEq)]
pub enum OverwriteMode {
    NoClobber,
    Interactive,
    Force,
}

#[derive(Clone, Eq, PartialEq)]
pub enum BackupMode {
    NoBackup,
    SimpleBackup,
    NumberedBackup,
    ExistingBackup,
}

pub fn uumain(args: Vec<String>) -> i32 {
    let mut opts = getopts::Options::new();

    opts.optflagopt("",  "backup", "make a backup of each existing destination file", "CONTROL");
    opts.optflag("b", "", "like --backup but does not accept an argument");
    opts.optflag("f", "force", "do not prompt before overwriting");
    opts.optflag("i", "interactive", "prompt before override");
    opts.optflag("n", "no-clobber", "do not overwrite an existing file");
    opts.optflag("",  "strip-trailing-slashes", "remove any trailing slashes from each SOURCE\n \
                                                 argument");
    opts.optopt("S", "suffix", "override the usual backup suffix", "SUFFIX");
    opts.optopt("t", "target-directory", "move all SOURCE arguments into DIRECTORY", "DIRECTORY");
    opts.optflag("T", "no-target-directory", "treat DEST as a normal file");
    opts.optflag("u", "update", "move only when the SOURCE file is newer\n \
                                than the destination file or when the\n \
                                destination file is missing");
    opts.optflag("v", "verbose", "explain what is being done");
    opts.optflag("h", "help", "display this help and exit");
    opts.optflag("V", "version", "output version information and exit");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => {
            show_error!("Invalid options\n{}", f);
            return 1;
        }
    };
    let usage = opts.usage("Move SOURCE to DEST, or multiple SOURCE(s) to DIRECTORY.");

    /* This does not exactly match the GNU implementation:
     * The GNU mv defaults to Force, but if more than one of the
     * overwrite options are supplied, only the last takes effect.
     * To default to no-clobber in that situation seems safer:
     */
    let overwrite_mode = if matches.opt_present("no-clobber") {
        OverwriteMode::NoClobber
    } else if matches.opt_present("interactive") {
        OverwriteMode::Interactive
    } else {
        OverwriteMode::Force
    };

    let backup_mode = if matches.opt_present("b") {
        BackupMode::SimpleBackup
    } else if matches.opt_present("backup") {
        match matches.opt_str("backup") {
            None => BackupMode::SimpleBackup,
            Some(mode) => match &mode[..] {
                "simple" | "never" => BackupMode::SimpleBackup,
                "numbered" | "t"   => BackupMode::NumberedBackup,
                "existing" | "nil" => BackupMode::ExistingBackup,
                "none" | "off"     => BackupMode::NoBackup,
                x => {
                    show_error!("invalid argument ‘{}’ for ‘backup type’\n\
                                Try '{} --help' for more information.", x, NAME);
                    return 1;
                }
            }
        }
    } else {
        BackupMode::NoBackup
    };

    if overwrite_mode == OverwriteMode::NoClobber && backup_mode != BackupMode::NoBackup {
        show_error!("options --backup and --no-clobber are mutually exclusive\n\
                    Try '{} --help' for more information.", NAME);
        return 1;
    }

    let backup_suffix = if matches.opt_present("suffix") {
        match matches.opt_str("suffix") {
            Some(x) => x,
            None => {
                show_error!("option '--suffix' requires an argument\n\
                            Try '{} --help' for more information.", NAME);
                return 1;
            }
        }
    } else {
        "~".to_string()
    };

    if matches.opt_present("T") && matches.opt_present("t") {
        show_error!("cannot combine --target-directory (-t) and --no-target-directory (-T)");
        return 1;
    }

    let behaviour = Behaviour {
        overwrite: overwrite_mode,
        backup: backup_mode,
        suffix: backup_suffix,
        update: matches.opt_present("u"),
        target_dir: matches.opt_str("t"),
        no_target_dir: matches.opt_present("T"),
        verbose: matches.opt_present("v"),
    };

    let paths: Vec<PathBuf> = {
        fn string_to_path<'a>(s: &'a String) -> &'a Path {
            Path::new(s)
        };
        fn strip_slashes<'a>(p: &'a Path) -> &'a Path {
            p.components().as_path()
        }
        let to_owned = |p: &Path| p.to_owned();
        let arguments = matches.free.iter().map(string_to_path);
        if matches.opt_present("strip-trailing-slashes") {
            arguments.map(strip_slashes).map(to_owned).collect()
        } else {
            arguments.map(to_owned).collect()
        }
    };

    if matches.opt_present("version") {
        println!("{} {}", NAME, VERSION);
        0
    } else if matches.opt_present("help") {
        help(&usage);
        0
    } else {
        exec(&paths[..], behaviour)
    }
}

fn help(usage: &str) {
    println!("{0} {1}\n\n\
    Usage: {0} SOURCE DEST\n   \
       or: {0} SOURCE... DIRECTORY\n\n\
    {2}", NAME, VERSION, usage);
}

fn exec(files: &[PathBuf], b: Behaviour) -> i32 {
    match b.target_dir {
        Some(ref name) => return move_files_into_dir(files, &PathBuf::from(name), &b),
        None => {}
    }
    match files {
        [] | [_] => {
            show_error!("missing file operand\n\
                        Try '{} --help' for more information.", NAME);
            return 1;
        },
        [ref source, ref target] => {
            if !source.uu_exists() {
                show_error!("cannot stat ‘{}’: No such file or directory", source.display());
                return 1;
            }

            if target.uu_is_dir() {
                if b.no_target_dir {
                    if !source.uu_is_dir() {
                        show_error!("cannot overwrite directory ‘{}’ with non-directory",
                            target.display());
                        return 1;
                    }

                    return match rename(source, target, &b) {
                        Err(e) => {
                            show_error!("{}", e);
                            1
                        },
                        _ => 0
                    }
                }

                return move_files_into_dir(&[source.clone()], target, &b);
            }

            match rename(source, target, &b) {
                Err(e) => {
                    show_error!("{}", e);
                    return 1;
                },
                _ => {}
            }
        }
        fs => {
            if b.no_target_dir {
                show_error!("mv: extra operand ‘{}’\n\
                            Try '{} --help' for more information.", fs[2].display(), NAME);
                return 1;
            }
            let target_dir = fs.last().unwrap();
            move_files_into_dir(fs.init(), target_dir, &b);
        }
    }
    0
}

fn move_files_into_dir(files: &[PathBuf], target_dir: &PathBuf, b: &Behaviour) -> i32 {
    if !target_dir.uu_is_dir() {
        show_error!("target ‘{}’ is not a directory", target_dir.display());
        return 1;
    }

    let mut all_successful = true;
    for sourcepath in files.iter() {
        let targetpath = match sourcepath.as_os_str().to_str() {
            Some(name) => target_dir.join(name),
            None => {
                show_error!("cannot stat ‘{}’: No such file or directory",
                            sourcepath.display());

                all_successful = false;
                continue;
            }
        };

        match rename(sourcepath, &targetpath, b) {
            Err(e) => {
                show_error!("mv: cannot move ‘{}’ to ‘{}’: {}",
                            sourcepath.display(), targetpath.display(), e);
                all_successful = false;
            },
            _ => {}
        }
    };
    if all_successful { 0 } else { 1 }
}

fn rename(from: &PathBuf, to: &PathBuf, b: &Behaviour) -> Result<()> {
    let mut backup_path = None;

    if to.uu_exists() {
        match b.overwrite {
            OverwriteMode::NoClobber => return Ok(()),
            OverwriteMode::Interactive => {
                print!("{}: overwrite ‘{}’? ", NAME, to.display());
                if !read_yes() {
                    return Ok(());
                }
            },
            OverwriteMode::Force => {}
        };

        backup_path = match b.backup {
            BackupMode::NoBackup => None,
            BackupMode::SimpleBackup => Some(simple_backup_path(to, &b.suffix)),
            BackupMode::NumberedBackup => Some(numbered_backup_path(to)),
            BackupMode::ExistingBackup => Some(existing_backup_path(to, &b.suffix))
        };
        if let Some(ref p) = backup_path {
            try!(fs::rename(to, p));
        }

        if b.update {
            if try!(fs::metadata(from)).mtime() <= try!(fs::metadata(to)).mtime() {
                return Ok(());
            }
        }
    }

    try!(fs::rename(from, to));

    if b.verbose {
        print!("‘{}’ -> ‘{}’", from.display(), to.display());
        match backup_path {
            Some(path) => println!(" (backup: ‘{}’)", path.display()),
            None => println!("")
        }
    }
    Ok(())
}

fn read_yes() -> bool {
    let mut s = String::new();
    match BufReader::new(stdin()).read_line(&mut s) {
        Ok(_) => match s.slice_shift_char() {
            Some((x, _)) => x == 'y' || x == 'Y',
            _ => false
        },
        _ => false
    }
}

fn simple_backup_path(path: &PathBuf, suffix: &String) -> PathBuf {
    let mut p = path.as_os_str().to_str().unwrap().to_string();
    p.push_str(suffix);
    return PathBuf::from(p);
}

fn numbered_backup_path(path: &PathBuf) -> PathBuf {
    let mut i: u64 = 1;
    loop {
        let new_path = simple_backup_path(path, &format!(".~{}~", i));
        if !new_path.uu_exists() {
            return new_path;
        }
        i = i + 1;
    }
}

fn existing_backup_path(path: &PathBuf, suffix: &String) -> PathBuf {
    let test_path = simple_backup_path(path, &".~1~".to_string());
    if test_path.uu_exists() {
        return numbered_backup_path(path);
    }
    simple_backup_path(path, suffix)
}
<commit_msg>Replace unstable init() method.<commit_after>#![crate_name = "mv"]
#![feature(slice_patterns, str_char)]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Orvar Segerström <orvarsegerstrom@gmail.com>
 * (c) Sokovikov Evgeniy  <skv-headless@yandex.ru>
 *
 * For the full copyright and license information, please view the LICENSE file
 * that was distributed with this source code.
 */

extern crate getopts;
extern crate libc;

use std::fs;
use std::io::{BufRead, BufReader, Result, stdin, Write};
use std::os::unix::fs::MetadataExt;
use std::path::{Path, PathBuf};

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

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

use filesystem::UUPathExt;

static NAME: &'static str = "mv";
static VERSION: &'static str = "0.0.1";

pub struct Behaviour {
    overwrite: OverwriteMode,
    backup: BackupMode,
    suffix: String,
    update: bool,
    target_dir: Option<String>,
    no_target_dir: bool,
    verbose: bool,
}

#[derive(Clone, Eq, PartialEq)]
pub enum OverwriteMode {
    NoClobber,
    Interactive,
    Force,
}

#[derive(Clone, Eq, PartialEq)]
pub enum BackupMode {
    NoBackup,
    SimpleBackup,
    NumberedBackup,
    ExistingBackup,
}

pub fn uumain(args: Vec<String>) -> i32 {
    let mut opts = getopts::Options::new();

    opts.optflagopt("",  "backup", "make a backup of each existing destination file", "CONTROL");
    opts.optflag("b", "", "like --backup but does not accept an argument");
    opts.optflag("f", "force", "do not prompt before overwriting");
    opts.optflag("i", "interactive", "prompt before override");
    opts.optflag("n", "no-clobber", "do not overwrite an existing file");
    opts.optflag("",  "strip-trailing-slashes", "remove any trailing slashes from each SOURCE\n \
                                                 argument");
    opts.optopt("S", "suffix", "override the usual backup suffix", "SUFFIX");
    opts.optopt("t", "target-directory", "move all SOURCE arguments into DIRECTORY", "DIRECTORY");
    opts.optflag("T", "no-target-directory", "treat DEST as a normal file");
    opts.optflag("u", "update", "move only when the SOURCE file is newer\n \
                                than the destination file or when the\n \
                                destination file is missing");
    opts.optflag("v", "verbose", "explain what is being done");
    opts.optflag("h", "help", "display this help and exit");
    opts.optflag("V", "version", "output version information and exit");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => {
            show_error!("Invalid options\n{}", f);
            return 1;
        }
    };
    let usage = opts.usage("Move SOURCE to DEST, or multiple SOURCE(s) to DIRECTORY.");

    /* This does not exactly match the GNU implementation:
     * The GNU mv defaults to Force, but if more than one of the
     * overwrite options are supplied, only the last takes effect.
     * To default to no-clobber in that situation seems safer:
     */
    let overwrite_mode = if matches.opt_present("no-clobber") {
        OverwriteMode::NoClobber
    } else if matches.opt_present("interactive") {
        OverwriteMode::Interactive
    } else {
        OverwriteMode::Force
    };

    let backup_mode = if matches.opt_present("b") {
        BackupMode::SimpleBackup
    } else if matches.opt_present("backup") {
        match matches.opt_str("backup") {
            None => BackupMode::SimpleBackup,
            Some(mode) => match &mode[..] {
                "simple" | "never" => BackupMode::SimpleBackup,
                "numbered" | "t"   => BackupMode::NumberedBackup,
                "existing" | "nil" => BackupMode::ExistingBackup,
                "none" | "off"     => BackupMode::NoBackup,
                x => {
                    show_error!("invalid argument ‘{}’ for ‘backup type’\n\
                                Try '{} --help' for more information.", x, NAME);
                    return 1;
                }
            }
        }
    } else {
        BackupMode::NoBackup
    };

    if overwrite_mode == OverwriteMode::NoClobber && backup_mode != BackupMode::NoBackup {
        show_error!("options --backup and --no-clobber are mutually exclusive\n\
                    Try '{} --help' for more information.", NAME);
        return 1;
    }

    let backup_suffix = if matches.opt_present("suffix") {
        match matches.opt_str("suffix") {
            Some(x) => x,
            None => {
                show_error!("option '--suffix' requires an argument\n\
                            Try '{} --help' for more information.", NAME);
                return 1;
            }
        }
    } else {
        "~".to_string()
    };

    if matches.opt_present("T") && matches.opt_present("t") {
        show_error!("cannot combine --target-directory (-t) and --no-target-directory (-T)");
        return 1;
    }

    let behaviour = Behaviour {
        overwrite: overwrite_mode,
        backup: backup_mode,
        suffix: backup_suffix,
        update: matches.opt_present("u"),
        target_dir: matches.opt_str("t"),
        no_target_dir: matches.opt_present("T"),
        verbose: matches.opt_present("v"),
    };

    let paths: Vec<PathBuf> = {
        fn string_to_path<'a>(s: &'a String) -> &'a Path {
            Path::new(s)
        };
        fn strip_slashes<'a>(p: &'a Path) -> &'a Path {
            p.components().as_path()
        }
        let to_owned = |p: &Path| p.to_owned();
        let arguments = matches.free.iter().map(string_to_path);
        if matches.opt_present("strip-trailing-slashes") {
            arguments.map(strip_slashes).map(to_owned).collect()
        } else {
            arguments.map(to_owned).collect()
        }
    };

    if matches.opt_present("version") {
        println!("{} {}", NAME, VERSION);
        0
    } else if matches.opt_present("help") {
        help(&usage);
        0
    } else {
        exec(&paths[..], behaviour)
    }
}

fn help(usage: &str) {
    println!("{0} {1}\n\n\
    Usage: {0} SOURCE DEST\n   \
       or: {0} SOURCE... DIRECTORY\n\n\
    {2}", NAME, VERSION, usage);
}

fn exec(files: &[PathBuf], b: Behaviour) -> i32 {
    match b.target_dir {
        Some(ref name) => return move_files_into_dir(files, &PathBuf::from(name), &b),
        None => {}
    }
    match files {
        [] | [_] => {
            show_error!("missing file operand\n\
                        Try '{} --help' for more information.", NAME);
            return 1;
        },
        [ref source, ref target] => {
            if !source.uu_exists() {
                show_error!("cannot stat ‘{}’: No such file or directory", source.display());
                return 1;
            }

            if target.uu_is_dir() {
                if b.no_target_dir {
                    if !source.uu_is_dir() {
                        show_error!("cannot overwrite directory ‘{}’ with non-directory",
                            target.display());
                        return 1;
                    }

                    return match rename(source, target, &b) {
                        Err(e) => {
                            show_error!("{}", e);
                            1
                        },
                        _ => 0
                    }
                }

                return move_files_into_dir(&[source.clone()], target, &b);
            }

            match rename(source, target, &b) {
                Err(e) => {
                    show_error!("{}", e);
                    return 1;
                },
                _ => {}
            }
        }
        fs => {
            if b.no_target_dir {
                show_error!("mv: extra operand ‘{}’\n\
                            Try '{} --help' for more information.", fs[2].display(), NAME);
                return 1;
            }
            let target_dir = fs.last().unwrap();
            move_files_into_dir(&fs[0..fs.len()-1], target_dir, &b);
        }
    }
    0
}

fn move_files_into_dir(files: &[PathBuf], target_dir: &PathBuf, b: &Behaviour) -> i32 {
    if !target_dir.uu_is_dir() {
        show_error!("target ‘{}’ is not a directory", target_dir.display());
        return 1;
    }

    let mut all_successful = true;
    for sourcepath in files.iter() {
        let targetpath = match sourcepath.as_os_str().to_str() {
            Some(name) => target_dir.join(name),
            None => {
                show_error!("cannot stat ‘{}’: No such file or directory",
                            sourcepath.display());

                all_successful = false;
                continue;
            }
        };

        match rename(sourcepath, &targetpath, b) {
            Err(e) => {
                show_error!("mv: cannot move ‘{}’ to ‘{}’: {}",
                            sourcepath.display(), targetpath.display(), e);
                all_successful = false;
            },
            _ => {}
        }
    };
    if all_successful { 0 } else { 1 }
}

fn rename(from: &PathBuf, to: &PathBuf, b: &Behaviour) -> Result<()> {
    let mut backup_path = None;

    if to.uu_exists() {
        match b.overwrite {
            OverwriteMode::NoClobber => return Ok(()),
            OverwriteMode::Interactive => {
                print!("{}: overwrite ‘{}’? ", NAME, to.display());
                if !read_yes() {
                    return Ok(());
                }
            },
            OverwriteMode::Force => {}
        };

        backup_path = match b.backup {
            BackupMode::NoBackup => None,
            BackupMode::SimpleBackup => Some(simple_backup_path(to, &b.suffix)),
            BackupMode::NumberedBackup => Some(numbered_backup_path(to)),
            BackupMode::ExistingBackup => Some(existing_backup_path(to, &b.suffix))
        };
        if let Some(ref p) = backup_path {
            try!(fs::rename(to, p));
        }

        if b.update {
            if try!(fs::metadata(from)).mtime() <= try!(fs::metadata(to)).mtime() {
                return Ok(());
            }
        }
    }

    try!(fs::rename(from, to));

    if b.verbose {
        print!("‘{}’ -> ‘{}’", from.display(), to.display());
        match backup_path {
            Some(path) => println!(" (backup: ‘{}’)", path.display()),
            None => println!("")
        }
    }
    Ok(())
}

fn read_yes() -> bool {
    let mut s = String::new();
    match BufReader::new(stdin()).read_line(&mut s) {
        Ok(_) => match s.slice_shift_char() {
            Some((x, _)) => x == 'y' || x == 'Y',
            _ => false
        },
        _ => false
    }
}

fn simple_backup_path(path: &PathBuf, suffix: &String) -> PathBuf {
    let mut p = path.as_os_str().to_str().unwrap().to_string();
    p.push_str(suffix);
    return PathBuf::from(p);
}

fn numbered_backup_path(path: &PathBuf) -> PathBuf {
    let mut i: u64 = 1;
    loop {
        let new_path = simple_backup_path(path, &format!(".~{}~", i));
        if !new_path.uu_exists() {
            return new_path;
        }
        i = i + 1;
    }
}

fn existing_backup_path(path: &PathBuf, suffix: &String) -> PathBuf {
    let test_path = simple_backup_path(path, &".~1~".to_string());
    if test_path.uu_exists() {
        return numbered_backup_path(path);
    }
    simple_backup_path(path, suffix)
}
<|endoftext|>


<commit_before>#![crate_name = "od"]
#![allow(unstable)]

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

extern crate getopts;
extern crate collections;

use collections::string::String;
use std::io::File;

#[derive(Show)]
enum Radix { Decimal, Hexadecimal, Octal, Binary }

pub fn uumain(args: Vec<String>) -> isize {
    let opts = [
        getopts::optopt("A", "address-radix",
                        "Select the base in which file offsets are printed.", "RADIX"),
        getopts::optopt("j", "skip-bytes",
                        "Skip bytes input bytes before formatting and writing.", "BYTES"),
        getopts::optopt("N", "read-bytes",
                        "limit dump to BYTES input bytes", "BYTES"),
        getopts::optopt("S", "strings",
                        ("output strings of at least BYTES graphic chars. 3 is assumed when \
                          BYTES is not specified."),
                        "BYTES"),
        getopts::optopt("t", "format", "select output format or formats", "TYPE"),
        getopts::optflag("v", "output-duplicates", "do not use * to mark line suppression"),
        getopts::optopt("w", "width",
                        ("output BYTES bytes per output line. 32 is implied when BYTES is not \
                          specified."),
                        "BYTES"),
        getopts::optflag("h", "help", "display this help and exit."),
        getopts::optflag("v", "version", "output version information and exit."),
    ];

    let matches = match getopts::getopts(args.tail(), &opts) {
        Ok(m) => m,
        Err(f) => panic!("Invalid options\n{}", f)
    };

    let input_offset_base = match parse_radix(matches.opt_str("A")) {
        Ok(r) => r,
        Err(f) => { panic!("Invalid -A/--address-radix\n{}", f) }
    };

    let mut fname;
    match args.last() {
        Some(n) => fname = n,
        None    => { panic!("Need fname for now") ; }
    };

    main(input_offset_base, fname.clone());

    0
}

fn main(input_offset_base: Radix, fname: String) {
    let mut f = match File::open(&Path::new(fname)) {
        Ok(f) => f,
        Err(e) => panic!("file error: {}", e)
    };

    let mut addr = 0;
    let bytes = &mut [b'\x00'; 16];
    loop {
        match f.read(bytes) {
            Ok(n) => {
                print!("{:07o}", addr);
                match input_offset_base {
                    Radix::Decimal => {},
                    Radix::Octal => {
                        for b in range(0, n / std::u16::BYTES) {
                            let bs = &bytes[(2 * b) .. (2 * b + 2)];
                            let p: u16 = (bs[1] as u16) << 8 | bs[0] as u16;
                            print!(" {:06o}", p);
                        }
                        if n % std::u16::BYTES == 1 {
                            print!(" {:06o}", bytes[n - 1]);
                        }
                    }
                    _ => { }
                };
                print!("\n");
                addr += n;
            },
            Err(_) => { println!("{:07o}", addr); break; }
        };
    };
}

fn parse_radix(radix_str: Option<String>) -> Result<Radix, &'static str> {
    match radix_str {
        None => Ok(Radix::Octal),
        Some(s) => {
            let st = s.into_bytes();
            if st.len() != 1 {
                Err("Radix must be one of [d, o, b, x]\n")
            } else {
                let radix: char = *(st.get(0)
                                      .expect("byte string of length 1 lacks a 0th elem")) as char;
                match radix {
                    'd' => Ok(Radix::Decimal),
                    'x' => Ok(Radix::Hexadecimal),
                    'o' => Ok(Radix::Octal),
                    'b' => Ok(Radix::Binary),
                    _ => Err("Radix must be one of [d, o, b, x]\n")
                }
            }
        }
    }
}
<commit_msg>od: rm -v alias for --version<commit_after>#![crate_name = "od"]
#![allow(unstable)]

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

extern crate getopts;
extern crate collections;

use collections::string::String;
use std::io::File;

#[derive(Show)]
enum Radix { Decimal, Hexadecimal, Octal, Binary }

pub fn uumain(args: Vec<String>) -> isize {
    let opts = [
        getopts::optopt("A", "address-radix",
                        "Select the base in which file offsets are printed.", "RADIX"),
        getopts::optopt("j", "skip-bytes",
                        "Skip bytes input bytes before formatting and writing.", "BYTES"),
        getopts::optopt("N", "read-bytes",
                        "limit dump to BYTES input bytes", "BYTES"),
        getopts::optopt("S", "strings",
                        ("output strings of at least BYTES graphic chars. 3 is assumed when \
                          BYTES is not specified."),
                        "BYTES"),
        getopts::optopt("t", "format", "select output format or formats", "TYPE"),
        getopts::optflag("v", "output-duplicates", "do not use * to mark line suppression"),
        getopts::optopt("w", "width",
                        ("output BYTES bytes per output line. 32 is implied when BYTES is not \
                          specified."),
                        "BYTES"),
        getopts::optflag("h", "help", "display this help and exit."),
        getopts::optflag("", "version", "output version information and exit."),
    ];

    let matches = match getopts::getopts(args.tail(), &opts) {
        Ok(m) => m,
        Err(f) => panic!("Invalid options\n{}", f)
    };

    let input_offset_base = match parse_radix(matches.opt_str("A")) {
        Ok(r) => r,
        Err(f) => { panic!("Invalid -A/--address-radix\n{}", f) }
    };

    let mut fname;
    match args.last() {
        Some(n) => fname = n,
        None    => { panic!("Need fname for now") ; }
    };

    main(input_offset_base, fname.clone());

    0
}

fn main(input_offset_base: Radix, fname: String) {
    let mut f = match File::open(&Path::new(fname)) {
        Ok(f) => f,
        Err(e) => panic!("file error: {}", e)
    };

    let mut addr = 0;
    let bytes = &mut [b'\x00'; 16];
    loop {
        match f.read(bytes) {
            Ok(n) => {
                print!("{:07o}", addr);
                match input_offset_base {
                    Radix::Decimal => {},
                    Radix::Octal => {
                        for b in range(0, n / std::u16::BYTES) {
                            let bs = &bytes[(2 * b) .. (2 * b + 2)];
                            let p: u16 = (bs[1] as u16) << 8 | bs[0] as u16;
                            print!(" {:06o}", p);
                        }
                        if n % std::u16::BYTES == 1 {
                            print!(" {:06o}", bytes[n - 1]);
                        }
                    }
                    _ => { }
                };
                print!("\n");
                addr += n;
            },
            Err(_) => { println!("{:07o}", addr); break; }
        };
    };
}

fn parse_radix(radix_str: Option<String>) -> Result<Radix, &'static str> {
    match radix_str {
        None => Ok(Radix::Octal),
        Some(s) => {
            let st = s.into_bytes();
            if st.len() != 1 {
                Err("Radix must be one of [d, o, b, x]\n")
            } else {
                let radix: char = *(st.get(0)
                                      .expect("byte string of length 1 lacks a 0th elem")) as char;
                match radix {
                    'd' => Ok(Radix::Decimal),
                    'x' => Ok(Radix::Hexadecimal),
                    'o' => Ok(Radix::Octal),
                    'b' => Ok(Radix::Binary),
                    _ => Err("Radix must be one of [d, o, b, x]\n")
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Allocating pages of memory from the OS and carving them into individual
//! allocations.

use bit_vec::BitVec;
use heap::{GcHeap, HeapSessionId};
use marking::MarkingTracer;
use ptr::{Pointer, UntypedPointer};
use std::{cmp, mem, ptr};
use std::marker::PhantomData;
use std::ops::{Deref, DerefMut};
use traits::{IntoHeapAllocation, Tracer};

/// Non-inlined function that serves as an entry point to marking. This is used
/// for marking root set entries.
unsafe fn mark_entry_point<'h, T>(addr: UntypedPointer, tracer: &mut MarkingTracer)
where
    T: IntoHeapAllocation<'h>,
{
    let addr = addr.as_typed_ptr::<T::In>();

    if GcHeap::get_mark_bit::<T>(addr) {
        // If the mark bit is set, then this object is gray in the classic
        // tri-color sense: seen but we just popped it off the mark stack and
        // have't finished enumerating its outgoing edges.
        T::trace(addr.as_ref(), tracer);
    } else {
        // If the mark bit is not set, then this object is white in the classic
        // tri-color sense: freshly discovered to be live, and we now need to
        // set its mark bit and then process its edges or push it onto the mark
        // stack for later edge processing.
        tracer.visit::<T>(addr);
    }
}

/// A unique id for each type that implements `IntoHeapAllocation`.
///
/// Implementation note: function types don't support Eq, so we cast to a
/// meaningless pointer type.
#[derive(Debug, Hash, PartialEq, Eq)]
pub struct TypeId(*const ());

pub fn heap_type_id<'h, T: IntoHeapAllocation<'h>>() -> TypeId {
    TypeId(mark_entry_point::<T> as *const ())
}

pub(crate) const PAGE_SIZE: usize = 0x1000;

/// We rely on all bits to the right of this bit being 0 in addresses of
/// TypedPage instances.
pub(crate) const PAGE_ALIGN: usize = 0x1000;

fn is_aligned(ptr: *const ()) -> bool {
    ptr as usize & (PAGE_ALIGN - 1) == 0
}

pub struct PageHeader {
    pub heap: *mut GcHeap,
    next_page: *mut PageHeader,
    mark_bits: BitVec,
    allocated_bits: BitVec,
    mark_fn: unsafe fn(UntypedPointer, &mut MarkingTracer),
    freelist: *mut (),
}

impl PageHeader {
    pub fn find(ptr: UntypedPointer) -> *mut PageHeader {
        let header_addr = ptr.as_usize() & !(PAGE_ALIGN - 1);
        header_addr as *mut PageHeader
    }

    pub fn clear_mark_bits(&mut self) {
        self.mark_bits.clear();
    }

    /// True if nothing on this page is allocated.
    pub fn is_empty(&self) -> bool {
        self.allocated_bits.none()
    }

    pub unsafe fn mark(&self, ptr: UntypedPointer, tracer: &mut MarkingTracer) {
        (self.mark_fn)(ptr, tracer);
    }

    pub fn type_id(&self) -> TypeId {
        TypeId(self.mark_fn as *const ())
    }

    pub fn downcast_mut<'h, T>(&mut self) -> Option<&mut TypedPage<T::In>>
    where
        T: IntoHeapAllocation<'h>,
    {
        if heap_type_id::<T>() == self.type_id() {
            let ptr = self as *mut PageHeader as *mut TypedPage<T::In>;
            Some(unsafe { &mut *ptr })
        } else {
            None
        }
    }
}

pub struct TypedPage<U> {
    pub header: PageHeader,
    pub allocations: PhantomData<U>,
}

impl<U> Deref for TypedPage<U> {
    type Target = PageHeader;

    fn deref(&self) -> &PageHeader {
        &self.header
    }
}

impl<U> DerefMut for TypedPage<U> {
    fn deref_mut(&mut self) -> &mut PageHeader {
        &mut self.header
    }
}

/// Returns the smallest multiple of `k` that is at least `n`.
///
/// Panics if the answer is too big to fit in a `usize`.
fn round_up(n: usize, k: usize) -> usize {
    let a = n / k * k;
    if a == n {
        n
    } else {
        n + k
    }
}

impl<U> TypedPage<U> {
    /// The actual size of an allocation can't be smaller than the size of a
    /// pointer, due to the way we store the freelist by stealing a pointer
    /// from the allocation itself.
    fn allocation_size() -> usize {
        cmp::max(mem::size_of::<U>(), mem::size_of::<*mut U>())
    }

    fn allocation_align() -> usize {
        cmp::max(mem::align_of::<U>(), mem::align_of::<*mut U>())
    }

    /// Offset, in bytes, of the first allocation from the start of the page.
    pub(crate) fn first_allocation_offset() -> usize {
        round_up(mem::size_of::<PageHeader>(), Self::allocation_align())
    }

    /// Number of allocations that fit in a page.
    pub fn capacity() -> usize {
        (PAGE_SIZE - Self::first_allocation_offset()) / Self::allocation_size()
    }

    fn first_object_addr(&self) -> usize {
        (self as *const Self as usize) + Self::first_allocation_offset()
    }

    unsafe fn init_freelist(&mut self) {
        let mut addr = self.first_object_addr();
        for _ in 0..Self::capacity() {
            self.add_to_free_list(addr as *mut U);

            // This can't use `ptr = ptr.offset(1)` because if U is smaller
            // than a pointer, allocations are padded to pointer size.
            // `.offset(1)` doesn't know about the padding and therefore
            // wouldn't advance to the next allocation.
            addr += Self::allocation_size();
        }
    }

    /// Return the page containing the object `ptr` points to.
    pub fn find(ptr: Pointer<U>) -> *mut TypedPage<U> {
        let ptr: usize = ptr.into();
        let page_addr = ptr & !(PAGE_ALIGN - 1);
        page_addr as *mut TypedPage<U>
    }

    unsafe fn add_to_free_list(&mut self, p: *mut U) {
        let listp = p as *mut *mut ();
        *listp = self.header.freelist;
        assert_eq!(*listp, self.header.freelist);
        self.header.freelist = p as *mut ();
    }

    unsafe fn allocation_index(&self, ptr: Pointer<U>) -> usize {
        let base = self.first_object_addr();

        // Check that ptr is in range.
        assert!(ptr.as_void() >= base as _);
        assert!(ptr.as_void() < (base + (Self::capacity() * Self::allocation_size())) as _);

        let index = (ptr.as_usize() - base as usize) / Self::allocation_size();
        assert_eq!(
            (base + index * Self::allocation_size()) as *const (),
            ptr.as_void()
        );
        index
    }

    pub unsafe fn get_mark_bit(&mut self, ptr: Pointer<U>) -> bool {
        let index = self.allocation_index(ptr);
        self.header.mark_bits[index]
    }

    pub unsafe fn set_mark_bit(&mut self, ptr: Pointer<U>) {
        let index = self.allocation_index(ptr);
        self.header.mark_bits.set(index, true);
    }

    /// Allocate a `U`-sized-and-aligned region of uninitialized memory
    /// from this page.
    ///
    /// # Safety
    ///
    /// This is safe unless GC is happening.
    pub unsafe fn try_alloc(&mut self) -> Option<Pointer<U>> {
        let p = self.header.freelist;
        if p.is_null() {
            None
        } else {
            let listp = p as *mut *mut ();
            self.header.freelist = *listp;
            let ap = Pointer::new(p as *mut U);
            let index = self.allocation_index(ap);
            assert!(!self.header.allocated_bits[index]);
            self.header.allocated_bits.set(index, true);
            Some(ap)
        }
    }

    unsafe fn sweep(&mut self) {
        let mut addr = self.first_object_addr();
        for i in 0..Self::capacity() {
            if self.header.allocated_bits[i] && !self.header.mark_bits[i] {
                ptr::drop_in_place(addr as *mut U);
                self.header.allocated_bits.set(i, false);
                self.add_to_free_list(addr as *mut U);
            }
            addr += Self::allocation_size();
        }
    }
}

/// Sweep a page.
///
/// # Safety
///
/// This must be called only after a full mark phase, to avoid sweeping objects
/// that are still reachable.
unsafe fn sweep_entry_point<'h, T: IntoHeapAllocation<'h>>(header: &mut PageHeader) {
    header.downcast_mut::<T>().unwrap().sweep();
}

/// An unordered collection of memory pages that all share an allocation type.
pub struct PageSet {
    heap: *mut GcHeap,

    sweep_fn: unsafe fn(&mut PageHeader),

    /// The collection of pages.
    ///
    /// All pages in this collection have matching `.heap`, `.mark_fn`, and
    /// `.sweep_fn` fields.
    pages: *mut PageHeader,

    /// The maximum number of pages, or None for no limit.
    limit: Option<usize>,
}

impl Drop for PageSet {
    fn drop(&mut self) {
        unsafe {
            let mut p = self.pages;
            while !p.is_null() {
                let next = (*p).next_page;
                ptr::drop_in_place(p); // drop the header
                Vec::from_raw_parts(p as *mut u8, 0, PAGE_SIZE); // free the page
                p = next;
            }
        }
    }
}

impl PageSet {
    /// Create a new PageSet.
    ///
    /// # Safety
    ///
    /// Safe as long as `heap` is a valid pointer.
    pub unsafe fn new<'h, T: IntoHeapAllocation<'h>>(heap: *mut GcHeap) -> PageSet {
        PageSet {
            heap,
            sweep_fn: sweep_entry_point::<T>,
            pages: ptr::null_mut(),
            limit: None,
        }
    }

    /// Downcast to a typed PageSetRef.
    ///
    /// # Panics
    ///
    /// If T is not the actual allocation type for this page set.
    pub fn downcast_mut<'a, 'h, T>(&'a mut self) -> PageSetRef<'a, 'h, T>
    where
        T: IntoHeapAllocation<'h> + 'a,
    {
        assert_eq!(
            self.sweep_fn as *const (),
            sweep_entry_point::<T> as *const ()
        );

        PageSetRef {
            page_set: self,
            id: PhantomData,
            also: PhantomData,
        }
    }

    /// Clear mark bits from each page in this set.
    ///
    /// # Safety
    ///
    /// This must be called only at the beginning of a GC cycle.
    pub unsafe fn clear_mark_bits(&mut self) {
        let mut page = self.pages;
        while !page.is_null() {
            (*page).clear_mark_bits();
            page = (*page).next_page;
        }
    }

    /// Sweep all unmarked objects from all pages.
    ///
    /// # Safety
    ///
    /// Safe to call only as the final part of GC.
    pub unsafe fn sweep(&mut self) {
        let mut page = self.pages;
        while !page.is_null() {
            (self.sweep_fn)(&mut *page);
            page = (*page).next_page;
        }
    }

    /// Assert that nothing is allocated in this set of pages.
    pub fn all_pages_are_empty(&self) -> bool {
        unsafe {
            let mut page = self.pages;
            while !page.is_null() {
                if !(*page).is_empty() {
                    return false;
                }
                page = (*page).next_page;
            }
        }
        true
    }
}

pub struct PageSetRef<'a, 'h, T: IntoHeapAllocation<'h> + 'a> {
    page_set: &'a mut PageSet,
    id: HeapSessionId<'h>,
    also: PhantomData<&'a mut T>,
}

impl<'a, 'h, T: IntoHeapAllocation<'h> + 'a> PageSetRef<'a, 'h, T> {
    pub fn set_page_limit(&mut self, limit: Option<usize>) {
        self.page_set.limit = limit;
    }

    /// Allocate memory for a value of type `T::In`.
    ///
    /// # Safety
    ///
    /// Safe to call as long as GC is not happening.
    pub unsafe fn try_alloc(&mut self) -> Option<Pointer<T::In>> {
        // First, try to allocate from each existing page. (Obviously slow.)
        let mut page = self.page_set.pages;
        let mut page_count = 0;
        while !page.is_null() {
            let t_page = (*page).downcast_mut::<T>().unwrap();
            if let Some(ptr) = t_page.try_alloc() {
                return Some(ptr);
            }

            page_count += 1;
            page = (*page).next_page;
        }

        // If there is a limit and we already have at least that many pages, fail.
        match self.page_set.limit {
            Some(limit) if page_count >= limit => None,
            _ => self.new_page().try_alloc(),
        }
    }

    /// Allocate a page from the operating system.
    ///
    /// Initialize its header and freelist and link it into this page set's
    /// linked list of pages.
    fn new_page(&mut self) -> &mut TypedPage<T::In> {
        assert!({
            let size_of_page = mem::size_of::<TypedPage<T::In>>();
            let alloc_offset = TypedPage::<T::In>::first_allocation_offset();
            size_of_page <= alloc_offset
        });
        assert!({
            let alloc_offset = TypedPage::<T::In>::first_allocation_offset();
            let alloc_size = TypedPage::<T::In>::allocation_size();
            let capacity = TypedPage::<T::In>::capacity();
            alloc_offset + capacity * alloc_size <= PAGE_SIZE
        });

        let mut vec: Vec<u8> = Vec::with_capacity(PAGE_SIZE);
        let raw_page = vec.as_mut_ptr() as *mut ();

        // Rust makes no guarantee whatsoever that this will work.
        // If it doesn't, panic.
        assert!(is_aligned(raw_page));

        let page_ptr: *mut TypedPage<T::In> = raw_page as *mut TypedPage<T::In>;
        let capacity = TypedPage::<T::In>::capacity();
        unsafe {
            ptr::write(
                page_ptr,
                TypedPage {
                    header: PageHeader {
                        heap: self.page_set.heap,
                        next_page: self.page_set.pages,
                        mark_bits: BitVec::from_elem(capacity, false),
                        allocated_bits: BitVec::from_elem(capacity, false),
                        mark_fn: mark_entry_point::<T>,
                        freelist: ptr::null_mut(),
                    },
                    allocations: PhantomData,
                },
            );

            let page = &mut *page_ptr;
            page.init_freelist();

            // Remove the memory from the vector and link it into
            // the PageSet's linked list.
            mem::forget(vec);
            self.page_set.pages = &mut page.header;

            page
        }
    }
}
<commit_msg>Improved assertions. Poison memory on sweeping in test/debug builds.<commit_after>//! Allocating pages of memory from the OS and carving them into individual
//! allocations.

use bit_vec::BitVec;
use heap::{GcHeap, HeapSessionId};
use marking::MarkingTracer;
use ptr::{Pointer, UntypedPointer};
use std::{cmp, mem, ptr};
use std::marker::PhantomData;
use std::ops::{Deref, DerefMut};
use traits::{IntoHeapAllocation, Tracer};

/// Non-inlined function that serves as an entry point to marking. This is used
/// for marking root set entries.
unsafe fn mark_entry_point<'h, T>(addr: UntypedPointer, tracer: &mut MarkingTracer)
where
    T: IntoHeapAllocation<'h>,
{
    let addr = addr.as_typed_ptr::<T::In>();

    if GcHeap::get_mark_bit::<T>(addr) {
        // If the mark bit is set, then this object is gray in the classic
        // tri-color sense: seen but we just popped it off the mark stack and
        // have't finished enumerating its outgoing edges.
        T::trace(addr.as_ref(), tracer);
    } else {
        // If the mark bit is not set, then this object is white in the classic
        // tri-color sense: freshly discovered to be live, and we now need to
        // set its mark bit and then process its edges or push it onto the mark
        // stack for later edge processing.
        tracer.visit::<T>(addr);
    }
}

/// A unique id for each type that implements `IntoHeapAllocation`.
///
/// Implementation note: function types don't support Eq, so we cast to a
/// meaningless pointer type.
#[derive(Debug, Hash, PartialEq, Eq)]
pub struct TypeId(*const ());

pub fn heap_type_id<'h, T: IntoHeapAllocation<'h>>() -> TypeId {
    TypeId(mark_entry_point::<T> as *const ())
}

pub(crate) const PAGE_SIZE: usize = 0x1000;

/// We rely on all bits to the right of this bit being 0 in addresses of
/// TypedPage instances.
pub(crate) const PAGE_ALIGN: usize = 0x1000;

fn is_aligned(ptr: *const ()) -> bool {
    ptr as usize & (PAGE_ALIGN - 1) == 0
}

pub struct PageHeader {
    pub heap: *mut GcHeap,
    next_page: *mut PageHeader,
    mark_bits: BitVec,
    allocated_bits: BitVec,
    mark_fn: unsafe fn(UntypedPointer, &mut MarkingTracer),
    freelist: *mut (),
}

impl PageHeader {
    pub fn find(ptr: UntypedPointer) -> *mut PageHeader {
        let header_addr = ptr.as_usize() & !(PAGE_ALIGN - 1);
        assert!(header_addr != 0);
        header_addr as *mut PageHeader
    }

    pub fn clear_mark_bits(&mut self) {
        self.mark_bits.clear();
    }

    /// True if nothing on this page is allocated.
    pub fn is_empty(&self) -> bool {
        self.allocated_bits.none()
    }

    pub unsafe fn mark(&self, ptr: UntypedPointer, tracer: &mut MarkingTracer) {
        (self.mark_fn)(ptr, tracer);
    }

    pub fn type_id(&self) -> TypeId {
        TypeId(self.mark_fn as *const ())
    }

    pub fn downcast_mut<'h, T>(&mut self) -> Option<&mut TypedPage<T::In>>
    where
        T: IntoHeapAllocation<'h>,
    {
        if heap_type_id::<T>() == self.type_id() {
            let ptr = self as *mut PageHeader as *mut TypedPage<T::In>;
            Some(unsafe { &mut *ptr })
        } else {
            None
        }
    }
}

pub struct TypedPage<U> {
    pub header: PageHeader,
    pub allocations: PhantomData<U>,
}

impl<U> Deref for TypedPage<U> {
    type Target = PageHeader;

    fn deref(&self) -> &PageHeader {
        &self.header
    }
}

impl<U> DerefMut for TypedPage<U> {
    fn deref_mut(&mut self) -> &mut PageHeader {
        &mut self.header
    }
}

/// Returns the smallest multiple of `k` that is at least `n`.
///
/// Panics if the answer is too big to fit in a `usize`.
fn round_up(n: usize, k: usize) -> usize {
    let a = n / k * k;
    if a == n {
        n
    } else {
        n + k
    }
}

impl<U> TypedPage<U> {
    /// The actual size of an allocation can't be smaller than the size of a
    /// pointer, due to the way we store the freelist by stealing a pointer
    /// from the allocation itself.
    fn allocation_size() -> usize {
        cmp::max(mem::size_of::<U>(), mem::size_of::<*mut U>())
    }

    fn allocation_align() -> usize {
        cmp::max(mem::align_of::<U>(), mem::align_of::<*mut U>())
    }

    /// Offset, in bytes, of the first allocation from the start of the page.
    pub(crate) fn first_allocation_offset() -> usize {
        round_up(mem::size_of::<PageHeader>(), Self::allocation_align())
    }

    /// Number of allocations that fit in a page.
    pub fn capacity() -> usize {
        (PAGE_SIZE - Self::first_allocation_offset()) / Self::allocation_size()
    }

    fn first_object_addr(&self) -> usize {
        (self as *const Self as usize) + Self::first_allocation_offset()
    }

    unsafe fn init_freelist(&mut self) {
        let mut addr = self.first_object_addr();
        for _ in 0..Self::capacity() {
            self.add_to_free_list(addr as *mut U);

            // This can't use `ptr = ptr.offset(1)` because if U is smaller
            // than a pointer, allocations are padded to pointer size.
            // `.offset(1)` doesn't know about the padding and therefore
            // wouldn't advance to the next allocation.
            addr += Self::allocation_size();
        }
    }

    /// Return the page containing the object `ptr` points to.
    pub fn find(ptr: Pointer<U>) -> *mut TypedPage<U> {
        let ptr: usize = ptr.into();
        let page_addr = ptr & !(PAGE_ALIGN - 1);
        page_addr as *mut TypedPage<U>
    }

    unsafe fn add_to_free_list(&mut self, p: *mut U) {
        let listp = p as *mut *mut ();
        *listp = self.header.freelist;
        assert_eq!(*listp, self.header.freelist);
        self.header.freelist = p as *mut ();
    }

    unsafe fn allocation_index(&self, ptr: Pointer<U>) -> usize {
        let base = self.first_object_addr();

        // Check that ptr is in range.
        assert!(ptr.as_void() >= base as _);
        assert!(ptr.as_void() < (base + (Self::capacity() * Self::allocation_size())) as _);

        let index = (ptr.as_usize() - base as usize) / Self::allocation_size();
        assert_eq!(
            (base + index * Self::allocation_size()) as *const (),
            ptr.as_void()
        );
        index
    }

    pub unsafe fn get_mark_bit(&mut self, ptr: Pointer<U>) -> bool {
        let index = self.allocation_index(ptr);
        self.header.mark_bits[index]
    }

    pub unsafe fn set_mark_bit(&mut self, ptr: Pointer<U>) {
        let index = self.allocation_index(ptr);
        assert!(self.header.allocated_bits.get(index).unwrap(),
                "marking memory that isn't allocated (dangling pointer?)");
        self.header.mark_bits.set(index, true);
    }

    /// Allocate a `U`-sized-and-aligned region of uninitialized memory
    /// from this page.
    ///
    /// # Safety
    ///
    /// This is safe unless GC is happening.
    pub unsafe fn try_alloc(&mut self) -> Option<Pointer<U>> {
        let p = self.header.freelist;
        if p.is_null() {
            None
        } else {
            let listp = p as *mut *mut ();
            self.header.freelist = *listp;
            let ap = Pointer::new(p as *mut U);
            let index = self.allocation_index(ap);
            assert!(!self.header.allocated_bits[index]);
            self.header.allocated_bits.set(index, true);
            Some(ap)
        }
    }

    unsafe fn sweep(&mut self) {
        let mut addr = self.first_object_addr();
        for i in 0..Self::capacity() {
            if self.header.allocated_bits[i] && !self.header.mark_bits[i] {
                ptr::drop_in_place(addr as *mut U);
                if cfg!(debug_assertions) || cfg!(test) {
                    // Paint the unused memory with a known-bad value.
                    const SWEPT_BYTE: u8 = 0xf4;
                    ptr::write_bytes(addr as *mut U, SWEPT_BYTE, 1);
                }
                self.header.allocated_bits.set(i, false);
                self.add_to_free_list(addr as *mut U);
            }
            addr += Self::allocation_size();
        }
    }
}

/// Sweep a page.
///
/// # Safety
///
/// This must be called only after a full mark phase, to avoid sweeping objects
/// that are still reachable.
unsafe fn sweep_entry_point<'h, T: IntoHeapAllocation<'h>>(header: &mut PageHeader) {
    header.downcast_mut::<T>().expect("page header corrupted").sweep();
}

/// An unordered collection of memory pages that all share an allocation type.
pub struct PageSet {
    heap: *mut GcHeap,

    sweep_fn: unsafe fn(&mut PageHeader),

    /// The collection of pages.
    ///
    /// All pages in this collection have matching `.heap`, `.mark_fn`, and
    /// `.sweep_fn` fields.
    pages: *mut PageHeader,

    /// The maximum number of pages, or None for no limit.
    limit: Option<usize>,
}

impl Drop for PageSet {
    fn drop(&mut self) {
        unsafe {
            let mut p = self.pages;
            while !p.is_null() {
                let next = (*p).next_page;
                ptr::drop_in_place(p); // drop the header
                Vec::from_raw_parts(p as *mut u8, 0, PAGE_SIZE); // free the page
                p = next;
            }
        }
    }
}

impl PageSet {
    /// Create a new PageSet.
    ///
    /// # Safety
    ///
    /// Safe as long as `heap` is a valid pointer.
    pub unsafe fn new<'h, T: IntoHeapAllocation<'h>>(heap: *mut GcHeap) -> PageSet {
        PageSet {
            heap,
            sweep_fn: sweep_entry_point::<T>,
            pages: ptr::null_mut(),
            limit: None,
        }
    }

    /// Downcast to a typed PageSetRef.
    ///
    /// # Panics
    ///
    /// If T is not the actual allocation type for this page set.
    pub fn downcast_mut<'a, 'h, T>(&'a mut self) -> PageSetRef<'a, 'h, T>
    where
        T: IntoHeapAllocation<'h> + 'a,
    {
        assert_eq!(
            self.sweep_fn as *const (),
            sweep_entry_point::<T> as *const ()
        );

        PageSetRef {
            page_set: self,
            id: PhantomData,
            also: PhantomData,
        }
    }

    /// Clear mark bits from each page in this set.
    ///
    /// # Safety
    ///
    /// This must be called only at the beginning of a GC cycle.
    pub unsafe fn clear_mark_bits(&mut self) {
        let mut page = self.pages;
        while !page.is_null() {
            (*page).clear_mark_bits();
            page = (*page).next_page;
        }
    }

    /// Sweep all unmarked objects from all pages.
    ///
    /// # Safety
    ///
    /// Safe to call only as the final part of GC.
    pub unsafe fn sweep(&mut self) {
        let mut page = self.pages;
        while !page.is_null() {
            (self.sweep_fn)(&mut *page);
            page = (*page).next_page;
        }
    }

    /// Assert that nothing is allocated in this set of pages.
    pub fn all_pages_are_empty(&self) -> bool {
        unsafe {
            let mut page = self.pages;
            while !page.is_null() {
                if !(*page).is_empty() {
                    return false;
                }
                page = (*page).next_page;
            }
        }
        true
    }
}

pub struct PageSetRef<'a, 'h, T: IntoHeapAllocation<'h> + 'a> {
    page_set: &'a mut PageSet,
    id: HeapSessionId<'h>,
    also: PhantomData<&'a mut T>,
}

impl<'a, 'h, T: IntoHeapAllocation<'h> + 'a> PageSetRef<'a, 'h, T> {
    pub fn set_page_limit(&mut self, limit: Option<usize>) {
        self.page_set.limit = limit;
    }

    /// Allocate memory for a value of type `T::In`.
    ///
    /// # Safety
    ///
    /// Safe to call as long as GC is not happening.
    pub unsafe fn try_alloc(&mut self) -> Option<Pointer<T::In>> {
        // First, try to allocate from each existing page. (Obviously slow.)
        let mut page = self.page_set.pages;
        let mut page_count = 0;
        while !page.is_null() {
            let t_page = (*page).downcast_mut::<T>().unwrap();
            if let Some(ptr) = t_page.try_alloc() {
                return Some(ptr);
            }

            page_count += 1;
            page = (*page).next_page;
        }

        // If there is a limit and we already have at least that many pages, fail.
        match self.page_set.limit {
            Some(limit) if page_count >= limit => None,
            _ => self.new_page().try_alloc(),
        }
    }

    /// Allocate a page from the operating system.
    ///
    /// Initialize its header and freelist and link it into this page set's
    /// linked list of pages.
    fn new_page(&mut self) -> &mut TypedPage<T::In> {
        assert!({
            let size_of_page = mem::size_of::<TypedPage<T::In>>();
            let alloc_offset = TypedPage::<T::In>::first_allocation_offset();
            size_of_page <= alloc_offset
        });
        assert!({
            let alloc_offset = TypedPage::<T::In>::first_allocation_offset();
            let alloc_size = TypedPage::<T::In>::allocation_size();
            let capacity = TypedPage::<T::In>::capacity();
            alloc_offset + capacity * alloc_size <= PAGE_SIZE
        });

        let mut vec: Vec<u8> = Vec::with_capacity(PAGE_SIZE);
        let raw_page = vec.as_mut_ptr() as *mut ();

        // Rust makes no guarantee whatsoever that this will work.
        // If it doesn't, panic.
        assert!(is_aligned(raw_page));

        let page_ptr: *mut TypedPage<T::In> = raw_page as *mut TypedPage<T::In>;
        let capacity = TypedPage::<T::In>::capacity();
        unsafe {
            ptr::write(
                page_ptr,
                TypedPage {
                    header: PageHeader {
                        heap: self.page_set.heap,
                        next_page: self.page_set.pages,
                        mark_bits: BitVec::from_elem(capacity, false),
                        allocated_bits: BitVec::from_elem(capacity, false),
                        mark_fn: mark_entry_point::<T>,
                        freelist: ptr::null_mut(),
                    },
                    allocations: PhantomData,
                },
            );

            let page = &mut *page_ptr;
            page.init_freelist();

            // Remove the memory from the vector and link it into
            // the PageSet's linked list.
            mem::forget(vec);
            self.page_set.pages = &mut page.header;

            page
        }
    }
}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Panic handling and stack unwinding

use core::fmt::{Arguments, Write};
use super::io::term;
use vga::Color;

#[lang = "panic_fmt"]
#[no_mangle] #[inline(never)] #[cold]
pub extern "C" fn rust_begin_unwind( args: Arguments, file: &'static str
                               , line: usize )
                               -> ! {
    write!(term::CONSOLE.lock()
                        .set_colors(Color::White, Color::Red)
                        .clear()
          , "KERNEL PANIC in {} at line {}\
            \nSomething has gone horribly wrong: {}. \
            This is fine."
          , file, line, args
          );
    loop { }
}

#[lang = "stack_exhausted"]
#[no_mangle] #[inline(never)] #[cold]
pub extern "C" fn __morestack() -> ! {
    println!("stack exhausted");
    loop { }
}

#[allow(non_snake_case)]
#[no_mangle] #[inline(never)] #[cold]
pub extern "C" fn _Unwind_Resume() -> ! {
    println!("UNWIND!");
    loop { }
}

// #[lang = "begin_unwind"]
// #[no_mangle] #[inline(never)] #[cold]
// pub fn begin_unwind<M: Send>(msg: M
//                             , file_line: &(&'static str, usize))  -> !
// { loop { } }
<commit_msg>Kernel panics made slightly more readable<commit_after>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Panic handling and stack unwinding

use core::fmt::{Arguments, Write};
use super::io::term;
use vga::Color;

#[lang = "panic_fmt"]
#[no_mangle] #[inline(never)] #[cold]
pub extern "C" fn rust_begin_unwind( args: Arguments, file: &'static str
                               , line: usize )
                               -> ! {
    write!(term::CONSOLE.lock()
                        .set_colors(Color::White, Color::Red)
                        .clear()
          , "KERNEL PANIC in {} at line {}\
            \nSomething has gone horribly wrong:\n {} \
            This is fine."
          , file, line, args
          );
    loop { }
}

#[lang = "stack_exhausted"]
#[no_mangle] #[inline(never)] #[cold]
pub extern "C" fn __morestack() -> ! {
    println!("stack exhausted");
    loop { }
}

#[allow(non_snake_case)]
#[no_mangle] #[inline(never)] #[cold]
pub extern "C" fn _Unwind_Resume() -> ! {
    println!("UNWIND!");
    loop { }
}

// #[lang = "begin_unwind"]
// #[no_mangle] #[inline(never)] #[cold]
// pub fn begin_unwind<M: Send>(msg: M
//                             , file_line: &(&'static str, usize))  -> !
// { loop { } }
<|endoftext|>


<commit_before>use super::ByteSize;

impl std::str::FromStr for ByteSize {
    type Err = String;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value.parse::<u64>() {
            Ok(v) => return Ok(Self(v)), // simple case, bytes
            Err(_) => {}
        };
        let number: String = value
            .chars()
            .take_while(|c| c.is_digit(10) || c == &'.')
            .collect();
        match number.parse::<f64>() {
            Ok(v) => {
                let suffix: String = value
                    .chars()
                    .skip_while(|c| c.is_whitespace() || c.is_digit(10) || c == &'.')
                    .collect();
                match suffix.parse::<Unit>() {
                    Ok(u) => Ok(Self((v * u.factor() as f64) as u64)),
                    Err(error) => Err(format!(
                        "couldn't parse {:?} into a known SI unit, {}",
                        suffix, error
                    )),
                }
            }
            Err(error) => Err(format!(
                "couldn't parse {:?} into a ByteSize, {}",
                value, error
            )),
        }
    }
}

enum Unit {
    Byte,
    // power of tens
    KiloByte,
    MegaByte,
    GigaByte,
    TeraByte,
    PetaByte,
    // power of twos
    KibiByte,
    MebiByte,
    GibiByte,
    TebiByte,
    PebiByte,
}

impl Unit {
    const B: u64 = super::B;
    const KB: u64 = super::KB;
    const KIB: u64 = super::KIB;
    const MB: u64 = super::MB;
    const MIB: u64 = super::MIB;
    const GB: u64 = super::GB;
    const GIB: u64 = super::GIB;
    const TB: u64 = super::TB;
    const TIB: u64 = super::TIB;
    const PB: u64 = super::PB;
    const PIB: u64 = super::PIB;

    fn factor(&self) -> u64 {
        match self {
            Self::Byte => Self::B,
            // power of tens
            Self::KiloByte => Self::KB,
            Self::MegaByte => Self::MB,
            Self::GigaByte => Self::GB,
            Self::TeraByte => Self::TB,
            Self::PetaByte => Self::PB,
            // power of twos
            Self::KibiByte => Self::KIB,
            Self::MebiByte => Self::MIB,
            Self::GibiByte => Self::GIB,
            Self::TebiByte => Self::TIB,
            Self::PebiByte => Self::PIB,
        }
    }
}

impl std::str::FromStr for Unit {
    type Err = String;

    fn from_str(unit: &str) -> Result<Self, Self::Err> {
        match unit.to_lowercase().as_str() {
            "b" => Ok(Self::Byte),
            // power of tens
            "k" | "kb" => Ok(Self::KiloByte),
            "m" | "mb" => Ok(Self::MegaByte),
            "g" | "gb" => Ok(Self::GigaByte),
            "t" | "tb" => Ok(Self::TeraByte),
            "p" | "pb" => Ok(Self::PetaByte),
            // power of twos
            "ki" | "kib" => Ok(Self::KibiByte),
            "mi" | "mib" => Ok(Self::MebiByte),
            "gi" | "gib" => Ok(Self::GibiByte),
            "ti" | "tib" => Ok(Self::TebiByte),
            "pi" | "pib" => Ok(Self::PebiByte),
            _ => Err(format!("couldn't parse unit of {:?}", unit)),
        }
    }
}

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

    #[test]
    fn when_ok() {
        // shortcut for writing test cases
        fn parse(s: &str) -> u64 {
            s.parse::<ByteSize>().unwrap().0
        }

        assert_eq!("0".parse::<ByteSize>().unwrap().0, 0);
        assert_eq!(parse("0"), 0);
        assert_eq!(parse("500"), 500);
        assert_eq!(parse("1K"), 1 * Unit::KB);
        assert_eq!(parse("1Ki"), 1 * Unit::KIB);
        assert_eq!(parse("1.5Ki"), (1.5 * Unit::KIB as f64) as u64);
        assert_eq!(parse("1KiB"), 1 * Unit::KIB);
        assert_eq!(parse("1.5KiB"), (1.5 * Unit::KIB as f64) as u64);
        assert_eq!(parse("3 MB"), 3 * Unit::MB);
        assert_eq!(parse("4 MiB"), 4 * Unit::MIB);
        assert_eq!(parse("6 GB"), 6 * Unit::GB);
        assert_eq!(parse("4 GiB"), 4 * Unit::GIB);
        assert_eq!(parse("88TB"), 88 * Unit::TB);
        assert_eq!(parse("521TiB"), 521 * Unit::TIB);
        assert_eq!(parse("8 PB"), 8 * Unit::PB);
        assert_eq!(parse("8P"), 8 * Unit::PB);
        assert_eq!(parse("12 PiB"), 12 * Unit::PIB);
    }

    #[test]
    fn when_err() {
        // shortcut for writing test cases
        fn parse(s: &str) -> Result<ByteSize, String> {
            s.parse::<ByteSize>()
        }

        assert!(parse("").is_err());
        assert!(parse("a124GB").is_err());
    }

    #[test]
    fn to_and_from_str() {
        // shortcut for writing test cases
        fn parse(s: &str) -> u64 {
            s.parse::<ByteSize>().unwrap().0
        }

        assert_eq!(parse(&format!("{}", parse("128GB"))), 128 * Unit::GB);
        assert_eq!(
            parse(&super::super::to_string(parse("128.000 GiB"), true)),
            128 * Unit::GIB
        );
    }
}
<commit_msg>cosmetic<commit_after>use super::ByteSize;

impl std::str::FromStr for ByteSize {
    type Err = String;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value.parse::<u64>() {
            Ok(v) => return Ok(Self(v)), // simple case, bytes
            Err(_) => {}
        };
        let number: String = value
            .chars()
            .take_while(|c| c.is_digit(10) || c == &'.')
            .collect();
        match number.parse::<f64>() {
            Ok(v) => {
                let suffix: String = value
                    .chars()
                    .skip_while(|c| c.is_whitespace() || c.is_digit(10) || c == &'.')
                    .collect();
                match suffix.parse::<Unit>() {
                    Ok(u) => Ok(Self((v * u.factor() as f64) as u64)),
                    Err(error) => Err(format!(
                        "couldn't parse {:?} into a known SI unit, {}",
                        suffix, error
                    )),
                }
            }
            Err(error) => Err(format!(
                "couldn't parse {:?} into a ByteSize, {}",
                value, error
            )),
        }
    }
}

enum Unit {
    Byte,
    // power of tens
    KiloByte,
    MegaByte,
    GigaByte,
    TeraByte,
    PetaByte,
    // power of twos
    KibiByte,
    MebiByte,
    GibiByte,
    TebiByte,
    PebiByte,
}

impl Unit {
    const B: u64 = super::B;
    const KB: u64 = super::KB;
    const KIB: u64 = super::KIB;
    const MB: u64 = super::MB;
    const MIB: u64 = super::MIB;
    const GB: u64 = super::GB;
    const GIB: u64 = super::GIB;
    const TB: u64 = super::TB;
    const TIB: u64 = super::TIB;
    const PB: u64 = super::PB;
    const PIB: u64 = super::PIB;

    fn factor(&self) -> u64 {
        match self {
            Self::Byte => Self::B,
            // power of tens
            Self::KiloByte => Self::KB,
            Self::MegaByte => Self::MB,
            Self::GigaByte => Self::GB,
            Self::TeraByte => Self::TB,
            Self::PetaByte => Self::PB,
            // power of twos
            Self::KibiByte => Self::KIB,
            Self::MebiByte => Self::MIB,
            Self::GibiByte => Self::GIB,
            Self::TebiByte => Self::TIB,
            Self::PebiByte => Self::PIB,
        }
    }
}

impl std::str::FromStr for Unit {
    type Err = String;

    fn from_str(unit: &str) -> Result<Self, Self::Err> {
        match unit.to_lowercase().as_str() {
            "b" => Ok(Self::Byte),
            // power of tens
            "k" | "kb" => Ok(Self::KiloByte),
            "m" | "mb" => Ok(Self::MegaByte),
            "g" | "gb" => Ok(Self::GigaByte),
            "t" | "tb" => Ok(Self::TeraByte),
            "p" | "pb" => Ok(Self::PetaByte),
            // power of twos
            "ki" | "kib" => Ok(Self::KibiByte),
            "mi" | "mib" => Ok(Self::MebiByte),
            "gi" | "gib" => Ok(Self::GibiByte),
            "ti" | "tib" => Ok(Self::TebiByte),
            "pi" | "pib" => Ok(Self::PebiByte),
            _ => Err(format!("couldn't parse unit of {:?}", unit)),
        }
    }
}

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

    #[test]
    fn when_ok() {
        // shortcut for writing test cases
        fn parse(s: &str) -> u64 {
            s.parse::<ByteSize>().unwrap().0
        }

        assert_eq!("0".parse::<ByteSize>().unwrap().0, 0);
        assert_eq!(parse("0"), 0);
        assert_eq!(parse("500"), 500);
        assert_eq!(parse("1K"), 1 * Unit::KB);
        assert_eq!(parse("1Ki"), 1 * Unit::KIB);
        assert_eq!(parse("1.5Ki"), (1.5 * Unit::KIB as f64) as u64);
        assert_eq!(parse("1KiB"), 1 * Unit::KIB);
        assert_eq!(parse("1.5KiB"), (1.5 * Unit::KIB as f64) as u64);
        assert_eq!(parse("3 MB"), 3 * Unit::MB);
        assert_eq!(parse("4 MiB"), 4 * Unit::MIB);
        assert_eq!(parse("6 GB"), 6 * Unit::GB);
        assert_eq!(parse("4 GiB"), 4 * Unit::GIB);
        assert_eq!(parse("88TB"), 88 * Unit::TB);
        assert_eq!(parse("521TiB"), 521 * Unit::TIB);
        assert_eq!(parse("8 PB"), 8 * Unit::PB);
        assert_eq!(parse("8P"), 8 * Unit::PB);
        assert_eq!(parse("12 PiB"), 12 * Unit::PIB);
    }

    #[test]
    fn when_err() {
        // shortcut for writing test cases
        fn parse(s: &str) -> Result<ByteSize, String> {
            s.parse::<ByteSize>()
        }

        assert!(parse("").is_err());
        assert!(parse("a124GB").is_err());
    }

    #[test]
    fn to_and_from_str() {
        // shortcut for writing test cases
        fn parse(s: &str) -> u64 {
            s.parse::<ByteSize>().unwrap().0
        }

        assert_eq!(parse(&format!("{}", parse("128GB"))), 128 * Unit::GB);
        assert_eq!(
            parse(&crate::to_string(parse("128.000 GiB"), true)),
            128 * Unit::GIB
        );
    }
}
<|endoftext|>


<commit_before>//! The parse module is responsible for parsing events published by clients.

use regex::Regex;

use std::str::{self, FromStr};

use error::{Error, CapellaResult};

/// MetricType defines what kind of metric was parsed from a client.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum MetricType {
    Counter,
    Gauge,
    Histogram,
    Meter,
    Timer,
}

impl FromStr for MetricType {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "c" => Ok(MetricType::Counter),
            "g" => Ok(MetricType::Gauge),
            "h" => Ok(MetricType::Histogram),
            "m" => Ok(MetricType::Meter),
            "ms" => Ok(MetricType::Timer),
            _ => Err(Error::Parse),
        }
    }
}

/// A Metric defines a published client event.
#[derive(Debug, PartialEq)]
pub struct Metric {
    name: String,
    value: i64,
    metric_type: MetricType,
    sample_rate: Option<f32>,
}

impl Metric {
    pub fn new() -> Metric {
        Metric {
            name: String::new(),
            value: 0,
            metric_type: MetricType::Meter,
            sample_rate: None,
        }
    }
}

/// The parse_metric function trys to break down a single UDP packet into a single metric.
pub fn parse_metric(packet: &[u8]) -> CapellaResult<Metric> {
    lazy_static! {
        // TODO: Learn how to write this over multiple lines.
        static ref PATTERN: Regex = Regex::new(r"(?P<name>\w+):((?P<sign>\-|\+))?(?P<val>\d+)\|(?P<type>\w+)(\|@(?P<rate>\d+\.\d+))?").unwrap();
    }

    if let Ok(val) = str::from_utf8(packet) {
        if !PATTERN.is_match(val) {
            return Err(Error::Parse);
        }
    } else {
        return Err(Error::Parse);
    }

    // We know this matched before so it is safe to unwrap.
    let caps = PATTERN.captures(str::from_utf8(packet).unwrap()).unwrap();

    let mut metric = Metric::new();
    // These are required to match.
    let name = caps.name("name").unwrap().as_str();
    let value = caps.name("val").unwrap().as_str().parse::<i64>().map_err(Error::from)?;
    let metric_type = caps.name("type").unwrap().as_str();

    metric.name = String::from(name);
    metric.value = value;
    metric.metric_type = metric_type.parse::<MetricType>()?;

    // Now see if there were optional values added in.
    if let Some(sign) = caps.name("sign") {
        let s = sign.as_str();
        if s == "-" {
            metric.value *= -1;
        }
    }

    if let Some(rate) = caps.name("rate") {
        let r = rate.as_str().parse::<f32>().map_err(Error::from)?;
        metric.sample_rate = Some(r);
    }

    Ok(metric)
}

#[cfg(test)]
mod tests {
    use super::{Metric, MetricType, parse_metric};

    #[test]
    fn good_simple_meter() {
        let packet = b"test:1|m";
        let m1 = parse_metric(packet).unwrap();

        let mut m2 = Metric::new();
        m2.name = String::from("test");
        m2.value = 1;

        assert_eq!(m1, m2);
    }

    #[test]
    fn bad_parse_double_colon() {
        let packet = b"test::1|c";
        assert!(parse_metric(packet).is_err());
    }

    #[test]
    fn good_timing_with_rate() {
        let packet = b"test:1|ms|@0.1";
        let m1 = parse_metric(packet).unwrap();

        let mut m2 = Metric::new();
        m2.name = String::from("test");
        m2.value = 1;
        m2.metric_type = MetricType::Timer;
        m2.sample_rate = Some(0.1);

        assert_eq!(m1, m2);
    }

    #[test]
    fn bad_parse_float_value() {
        let packet = b"test:1.0|g";
        assert!(parse_metric(packet).is_err());
    }

    #[test]
    fn bad_parse_unknown_metric_type() {
        let packet = b"test:1|a";
        assert!(parse_metric(packet).is_err());
    }
}
<commit_msg>adjust pattern to allow for nested metric names<commit_after>//! The parse module is responsible for parsing events published by clients.

use regex::Regex;

use std::str::{self, FromStr};

use error::{Error, CapellaResult};

/// MetricType defines what kind of metric was parsed from a client.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum MetricType {
    Counter,
    Gauge,
    Histogram,
    Meter,
    Timer,
}

impl FromStr for MetricType {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "c" => Ok(MetricType::Counter),
            "g" => Ok(MetricType::Gauge),
            "h" => Ok(MetricType::Histogram),
            "m" => Ok(MetricType::Meter),
            "ms" => Ok(MetricType::Timer),
            _ => Err(Error::Parse),
        }
    }
}

/// A Metric defines a published client event.
#[derive(Debug, PartialEq)]
pub struct Metric {
    name: String,
    value: i64,
    metric_type: MetricType,
    sample_rate: Option<f32>,
}

impl Metric {
    pub fn new() -> Metric {
        Metric {
            name: String::new(),
            value: 0,
            metric_type: MetricType::Meter,
            sample_rate: None,
        }
    }
}

/// The parse_metric function trys to break down a single UDP packet into a single metric.
pub fn parse_metric(packet: &[u8]) -> CapellaResult<Metric> {
    lazy_static! {
        // TODO: Learn how to write this over multiple lines.
        static ref PATTERN: Regex = Regex::new(r"(?P<name>[\w\.]+):((?P<sign>\-|\+))?(?P<val>\d+)\|(?P<type>\w+)(\|@(?P<rate>\d+\.\d+))?").unwrap();
    }

    if let Ok(val) = str::from_utf8(packet) {
        if !PATTERN.is_match(val) {
            return Err(Error::Parse);
        }
    } else {
        return Err(Error::Parse);
    }

    // We know this matched before so it is safe to unwrap.
    let caps = PATTERN.captures(str::from_utf8(packet).unwrap()).unwrap();

    let mut metric = Metric::new();
    // These are required to match.
    let name = caps.name("name").unwrap().as_str();
    let value = caps.name("val").unwrap().as_str().parse::<i64>().map_err(Error::from)?;
    let metric_type = caps.name("type").unwrap().as_str();

    metric.name = String::from(name);
    metric.value = value;
    metric.metric_type = metric_type.parse::<MetricType>()?;

    // Now see if there were optional values added in.
    if let Some(sign) = caps.name("sign") {
        let s = sign.as_str();
        if s == "-" {
            metric.value *= -1;
        }
    }

    if let Some(rate) = caps.name("rate") {
        let r = rate.as_str().parse::<f32>().map_err(Error::from)?;
        metric.sample_rate = Some(r);
    }

    Ok(metric)
}

#[cfg(test)]
mod tests {
    use super::{Metric, MetricType, parse_metric};

    #[test]
    fn good_simple_meter() {
        let packet = b"test:1|m";
        let m1 = parse_metric(packet).unwrap();

        let mut m2 = Metric::new();
        m2.name = String::from("test");
        m2.value = 1;

        assert_eq!(m1, m2);
    }

    #[test]
    fn bad_parse_double_colon() {
        let packet = b"test::1|c";
        assert!(parse_metric(packet).is_err());
    }

    #[test]
    fn good_timing_with_rate() {
        let packet = b"test:1|ms|@0.1";
        let m1 = parse_metric(packet).unwrap();

        let mut m2 = Metric::new();
        m2.name = String::from("test");
        m2.value = 1;
        m2.metric_type = MetricType::Timer;
        m2.sample_rate = Some(0.1);

        assert_eq!(m1, m2);
    }

    #[test]
    fn bad_parse_float_value() {
        let packet = b"test:1.0|g";
        assert!(parse_metric(packet).is_err());
    }

    #[test]
    fn bad_parse_unknown_metric_type() {
        let packet = b"test:1|a";
        assert!(parse_metric(packet).is_err());
    }

    #[test]
    fn good_nested_metric_name() {
        let packet = b"test.nested.name:1|c";
        let m1 = parse_metric(packet).unwrap();

        let mut m2 = Metric::new();
        m2.name = String::from("test.nested.name");
        m2.value = 1;
        m2.metric_type = MetricType::Counter;

        assert_eq!(m1, m2);
    }
}
<|endoftext|>


<commit_before>use std::net::{SocketAddr, IpAddr};
use std::io;
use std::time::Duration;

use rand;
use mioco::udp::UdpSocket;
use mioco::tcp::{TcpStream};
use trust_dns::op::{Message, MessageType, ResponseCode, Query, OpCode, Edns};

use utils::{Result, Error, CloneExt, MessageExt, WithTimeout, Future, AsDisplay};
use utils::with_timeout::TcpStreamExt;
use transport::{DnsTransport, DnsMsgTransport};

pub const EDNS_VER: u8 = 0;
pub const EDNS_MAX_PAYLOAD: u16 = 1200;

// Idea: Use DNSSEC to check whether a domain is poisoned by GFW

fn query_core<T: DnsMsgTransport>(q: Query, mut transport: T, enable_edns: bool) -> Result<Message> {
    let mut msg : Message = Message::new();
    msg.message_type(MessageType::Query);
    msg.id(rand::random());
    msg.op_code(OpCode::Query);
    msg.recursion_desired(true);
    if enable_edns {
        let mut edns = Edns::new();
        edns.set_version(EDNS_VER);
        edns.set_dnssec_ok(false); // TODO
        edns.set_max_payload(EDNS_MAX_PAYLOAD);
        msg.set_edns(edns);
    }
    msg.add_query(q);
    debug!("[{}] {}", transport, msg.as_disp());
    try!(transport.send_msg(&msg));
    loop {
        let resp = try!(transport.recv_msg());
        if ! resp.is_resp_for(&msg) {
            warn!("[{}][{}] Invalid response for {}: {:?}",
                  transport, msg.get_id(), msg.as_disp(), resp);
            continue;
        }
        if resp.get_response_code() == ResponseCode::FormErr && enable_edns {
            // Maybe the server doesn't implement EDNS?
            return query_core(msg.get_queries()[0].clone(), transport, false);
        }
        debug!("[{}] {} -> {}", transport, resp.get_queries()[0].as_disp(), resp.as_disp());
        return Ok(resp);
    }
}
fn get_bind_addr(target: &IpAddr) -> SocketAddr {
    SocketAddr::new(match *target {
        IpAddr::V4(_) => "0.0.0.0",
        IpAddr::V6(_) => "::0",
    }.parse().unwrap(), 0)
}
pub fn query(q: Query, addr: IpAddr, timeout: Duration) -> Result<Message> {
    let target = SocketAddr::new(addr, 53);
    let mut transport = try!(UdpSocket::bound(&get_bind_addr(&addr))).with_timeout(timeout);
    match query_core(q, transport.bound(Some(&target)), true) {
        Ok(msg) => {
            if msg.is_truncated() {
                // Try TCP
                let tcp_result = TcpStream::connect_with_timeout(&target, timeout)
                .map_err(|e| e.into())
                .and_then(|stream| stream.set_nodelay(true)
                                         .or_else(|_| Ok(()))
                                         .map(|_| stream))
                .and_then(|stream| query_core(
                    msg.get_queries()[0].clone(),
                    stream.with_timeout(timeout).bound(Some(&target)),
                    true,
                ));
                Ok(tcp_result.unwrap_or(msg))
            } else {
                Ok(msg)
            }
        },
        Err(e) => Err(e),
    }
}
/* FIXME: Commented out to silence warning. This may be useful again in future.
pub fn query_multiple_build_futures(q: &Query, servers: &[IpAddr]) -> Vec<Future<Result<Message>>> {
    servers.iter()
    .cloned()
    .map(move |server| {
        let qc = q.clone();
        Future::from_fn(move || query(qc, server, true))
    })
    .collect()
}
*/
pub fn query_multiple_handle_futures(futures: &mut Vec<Future<Result<Message>>>) -> Result<Message> {
    if futures.is_empty() {
        return Err(io::ErrorKind::InvalidInput.into());
    }
    loop {
        let result_index = Future::wait_any(futures);
        let result = futures.remove(result_index).consume();
        let mut should_return = futures.is_empty();
        match result {
            Ok(ref msg) => {
                match msg.get_response_code() {
                    ResponseCode::ServFail |
                    ResponseCode::NotImp |
                    ResponseCode::Refused => {
                        // Fall below to wait for remaining servers, if any
                    },
                    _ => { should_return = true; },
                }
            },
            Err(Error::Io(ref err)) if err.kind() == io::ErrorKind::TimedOut => {
                debug!("Timed out");
            },
            Err(ref err) => {
                debug!("Error {:?}", err);
            },
        };
        if should_return {
            return result;
        }
    }
}
/* FIXME: Commented out to silence warning. This may be useful again in future.
pub fn query_multiple(q: &Query, servers: &[IpAddr]) -> Result<Message> {
    let mut futures = query_multiple_build_futures(q, servers);
    query_multiple_handle_futures(&mut futures)
}
*/

#[cfg(test)]
mod tests {
    use std::time::Duration;
    use trust_dns::op::*;
    use trust_dns::rr::*;

    use super::*;
    use ::mioco_config_start;

    #[test]
    fn simple_query() {
        mioco_config_start(|| {
            let mut q = Query::new();
            q.name(Name::parse("www.google.com", Some(&Name::root())).unwrap());
            let result = query(q, "8.8.8.8".parse().unwrap(), Duration::from_secs(5)).unwrap();
            assert!(result.get_answers().len() > 0);
        }).unwrap();
    }
}
<commit_msg>Cleanup<commit_after>use std::net::{SocketAddr, IpAddr};
use std::io;
use std::time::Duration;

use rand;
use mioco::udp::UdpSocket;
use mioco::tcp::{TcpStream};
use trust_dns::op::{Message, MessageType, ResponseCode, Query, OpCode, Edns};

use utils::{Result, Error, CloneExt, MessageExt, WithTimeout, Future, AsDisplay};
use utils::with_timeout::TcpStreamExt;
use transport::{DnsTransport, DnsMsgTransport};

pub const EDNS_VER: u8 = 0;
pub const EDNS_MAX_PAYLOAD: u16 = 1200;

// Idea: Use DNSSEC to check whether a domain is poisoned by GFW

fn query_core<T: DnsMsgTransport>(q: Query, mut transport: T, enable_edns: bool) -> Result<Message> {
    let mut msg : Message = Message::new();
    msg.message_type(MessageType::Query);
    msg.id(rand::random());
    msg.op_code(OpCode::Query);
    msg.recursion_desired(true);
    if enable_edns {
        let mut edns = Edns::new();
        edns.set_version(EDNS_VER);
        edns.set_dnssec_ok(false); // TODO
        edns.set_max_payload(EDNS_MAX_PAYLOAD);
        msg.set_edns(edns);
    }
    msg.add_query(q);
    debug!("[{}] {}", transport, msg.as_disp());
    try!(transport.send_msg(&msg));
    loop {
        let resp = try!(transport.recv_msg());
        if ! resp.is_resp_for(&msg) {
            warn!("[{}][{}] Invalid response for {}: {:?}",
                  transport, msg.get_id(), msg.as_disp(), resp);
            continue;
        }
        if resp.get_response_code() == ResponseCode::FormErr && enable_edns {
            // Maybe the server doesn't implement EDNS?
            return query_core(msg.get_queries()[0].clone(), transport, false);
        }
        debug!("[{}] {} -> {}", transport, resp.get_queries()[0].as_disp(), resp.as_disp());
        return Ok(resp);
    }
}
fn get_bind_addr(target: &IpAddr) -> SocketAddr {
    SocketAddr::new(match *target {
        IpAddr::V4(_) => "0.0.0.0",
        IpAddr::V6(_) => "::0",
    }.parse().unwrap(), 0)
}
pub fn query(q: Query, addr: IpAddr, timeout: Duration) -> Result<Message> {
    let target = SocketAddr::new(addr, 53);
    let mut transport = try!(UdpSocket::bound(&get_bind_addr(&addr))).with_timeout(timeout);
    match query_core(q, transport.bound(Some(&target)), true) {
        Ok(msg) => {
            if msg.is_truncated() {
                // Try TCP
                let tcp_result = TcpStream::connect_with_timeout(&target, timeout)
                .map_err(|e| e.into())
                .and_then(|stream| stream.set_nodelay(true)
                                         .or_else(|_| Ok(()))
                                         .map(|_| stream))
                .and_then(|stream| query_core(
                    msg.get_queries()[0].clone(),
                    stream.with_timeout(timeout).bound(Some(&target)),
                    true,
                ));
                Ok(tcp_result.unwrap_or(msg))
            } else {
                Ok(msg)
            }
        },
        Err(e) => Err(e),
    }
}
pub fn query_multiple_handle_futures(futures: &mut Vec<Future<Result<Message>>>) -> Result<Message> {
    if futures.is_empty() {
        return Err(io::ErrorKind::InvalidInput.into());
    }
    loop {
        let result_index = Future::wait_any(futures);
        let result = futures.remove(result_index).consume();
        let mut should_return = futures.is_empty();
        match result {
            Ok(ref msg) => {
                match msg.get_response_code() {
                    ResponseCode::ServFail |
                    ResponseCode::NotImp |
                    ResponseCode::Refused => {
                        // Fall below to wait for remaining servers, if any
                    },
                    _ => { should_return = true; },
                }
            },
            Err(Error::Io(ref err)) if err.kind() == io::ErrorKind::TimedOut => {
                debug!("Timed out");
            },
            Err(ref err) => {
                debug!("Error {:?}", err);
            },
        };
        if should_return {
            return result;
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;
    use trust_dns::op::*;
    use trust_dns::rr::*;

    use super::*;
    use ::mioco_config_start;

    #[test]
    fn simple_query() {
        mioco_config_start(|| {
            let mut q = Query::new();
            q.name(Name::parse("www.google.com", Some(&Name::root())).unwrap());
            let result = query(q, "8.8.8.8".parse().unwrap(), Duration::from_secs(5)).unwrap();
            assert!(result.get_answers().len() > 0);
        }).unwrap();
    }
}
<|endoftext|>


<commit_before>// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.

use crate::grpc_sys::{self, grpc_resource_quota};
use std::ffi::CString;
use std::ptr;

/// ResourceQuota represents a bound on memory and thread usage by the gRPC.
/// NOTE: The management of threads created in grpc-core don't use ResourceQuota.
/// TODO: Manage the poller threads created in grpc-rs with this ResourceQuota later.
pub struct ResourceQuota {
    raw: *mut grpc_resource_quota,
}

impl ResourceQuota {
    /// Create a control block for resource quota. If a name is
    /// not declared for this control block, a name is automatically
    /// generated in grpc core.
    pub fn new(name: Option<&str>) -> ResourceQuota {
        match name {
            Some(name_str) => {
                let name_cstr = CString::new(name_str).unwrap();
                ResourceQuota {
                    raw: unsafe { grpc_sys::grpc_resource_quota_create(name_cstr.as_ptr() as _) },
                }
            }
            None => ResourceQuota {
                raw: unsafe { grpc_sys::grpc_resource_quota_create(ptr::null()) },
            },
        }
    }

    /// Resize this ResourceQuota to a new memory size.
    pub fn resize_memory(self, new_size: usize) -> ResourceQuota {
        unsafe { grpc_sys::grpc_resource_quota_resize(self.raw, new_size) };
        self
    }

    pub(crate) fn get_ptr(&self) -> *mut grpc_resource_quota {
        self.raw
    }
}

impl Drop for ResourceQuota {
    fn drop(&mut self) {
        unsafe {
            grpc_sys::grpc_resource_quota_unref(self.raw);
        }
    }
}
<commit_msg>implement Clone for ResourceQuota (#568)<commit_after>// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.

use crate::grpc_sys::{self, grpc_resource_quota};
use std::ffi::CString;
use std::ptr;

/// ResourceQuota represents a bound on memory and thread usage by the gRPC.
/// NOTE: The management of threads created in grpc-core don't use ResourceQuota.
/// TODO: Manage the poller threads created in grpc-rs with this ResourceQuota later.
pub struct ResourceQuota {
    raw: *mut grpc_resource_quota,
}

impl ResourceQuota {
    /// Create a control block for resource quota. If a name is
    /// not declared for this control block, a name is automatically
    /// generated in grpc core.
    pub fn new(name: Option<&str>) -> ResourceQuota {
        match name {
            Some(name_str) => {
                let name_cstr = CString::new(name_str).unwrap();
                ResourceQuota {
                    raw: unsafe { grpc_sys::grpc_resource_quota_create(name_cstr.as_ptr() as _) },
                }
            }
            None => ResourceQuota {
                raw: unsafe { grpc_sys::grpc_resource_quota_create(ptr::null()) },
            },
        }
    }

    /// Resize this ResourceQuota to a new memory size.
    pub fn resize_memory(self, new_size: usize) -> ResourceQuota {
        unsafe { grpc_sys::grpc_resource_quota_resize(self.raw, new_size) };
        self
    }

    pub(crate) fn get_ptr(&self) -> *mut grpc_resource_quota {
        self.raw
    }
}

impl Clone for ResourceQuota {
    fn clone(&self) -> Self {
        unsafe {
            grpc_sys::grpc_resource_quota_ref(self.raw);
        }
        Self { raw: self.raw }
    }
}

impl Drop for ResourceQuota {
    fn drop(&mut self) {
        unsafe {
            grpc_sys::grpc_resource_quota_unref(self.raw);
        }
    }
}
<|endoftext|>