<commit_before>// Copyright 2016 Itoa Developers
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::{io, mem, ptr, slice};

#[inline]
pub fn write<W: io::Write + ?Sized, V: Integer>(wr: &mut W, value: V) -> io::Result<()> {
    value.write(wr)
}

pub trait Integer {
    fn write<W: io::Write + ?Sized>(self, &mut W) -> io::Result<()>;
}

const DEC_DIGITS_LUT: &'static[u8] =
    b"0001020304050607080910111213141516171819\
      2021222324252627282930313233343536373839\
      4041424344454647484950515253545556575859\
      6061626364656667686970717273747576777879\
      8081828384858687888990919293949596979899";

// Adaptation of the original implementation at
// https://github.com/rust-lang/rust/blob/b8214dc6c6fc20d0a660fb5700dca9ebf51ebe89/src/libcore/fmt/num.rs#L188-L266
macro_rules! impl_Integer {
    ($($t:ident),* as $conv_fn:ident) => ($(
    impl Integer for $t {
        #[allow(unused_comparisons)]
        fn write<W: io::Write + ?Sized>(self, wr: &mut W) -> io::Result<()> {
            let is_nonnegative = self >= 0;
            let mut n = if is_nonnegative {
                self as $conv_fn
            } else {
                try!(wr.write_all(b"-"));
                // convert the negative num to positive by summing 1 to it's 2 complement
                (!(self as $conv_fn)).wrapping_add(1)
            };
            let mut buf: [u8; 20] = unsafe { mem::uninitialized() };
            let mut curr = buf.len() as isize;
            let buf_ptr = buf.as_mut_ptr();
            let lut_ptr = DEC_DIGITS_LUT.as_ptr();

            unsafe {
                // eagerly decode 4 characters at a time
                if <$t>::max_value() as u64 >= 10000 {
                    while n >= 10000 {
                        let rem = (n % 10000) as isize;
                        n /= 10000;

                        let d1 = (rem / 100) << 1;
                        let d2 = (rem % 100) << 1;
                        curr -= 4;
                        ptr::copy_nonoverlapping(lut_ptr.offset(d1), buf_ptr.offset(curr), 2);
                        ptr::copy_nonoverlapping(lut_ptr.offset(d2), buf_ptr.offset(curr + 2), 2);
                    }
                }

                // if we reach here numbers are <= 9999, so at most 4 chars long
                let mut n = n as isize; // possibly reduce 64bit math

                // decode 2 more chars, if > 2 chars
                if n >= 100 {
                    let d1 = (n % 100) << 1;
                    n /= 100;
                    curr -= 2;
                    ptr::copy_nonoverlapping(lut_ptr.offset(d1), buf_ptr.offset(curr), 2);
                }

                // decode last 1 or 2 chars
                if n < 10 {
                    curr -= 1;
                    *buf_ptr.offset(curr) = (n as u8) + 48;
                } else {
                    let d1 = n << 1;
                    curr -= 2;
                    ptr::copy_nonoverlapping(lut_ptr.offset(d1), buf_ptr.offset(curr), 2);
                }
            }

            wr.write_all(unsafe {
                slice::from_raw_parts(buf_ptr.offset(curr), buf.len() - curr as usize)
            })
        }
    })*);
}

impl_Integer!(i8, u8, i16, u16, i32, u32 as u32);
impl_Integer!(i64, u64 as u64);
#[cfg(target_pointer_width = "16")]
impl_Integer!(isize, usize as u16);
#[cfg(target_pointer_width = "32")]
impl_Integer!(isize, usize as u32);
#[cfg(target_pointer_width = "64")]
impl_Integer!(isize, usize as u64);
<commit_msg>Make itoa::write take a W instead of a &mut W<commit_after>// Copyright 2016 Itoa Developers
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::{io, mem, ptr, slice};

#[inline]
pub fn write<W: io::Write, V: Integer>(wr: W, value: V) -> io::Result<()> {
    value.write(wr)
}

pub trait Integer {
    fn write<W: io::Write>(self, W) -> io::Result<()>;
}

const DEC_DIGITS_LUT: &'static[u8] =
    b"0001020304050607080910111213141516171819\
      2021222324252627282930313233343536373839\
      4041424344454647484950515253545556575859\
      6061626364656667686970717273747576777879\
      8081828384858687888990919293949596979899";

// Adaptation of the original implementation at
// https://github.com/rust-lang/rust/blob/b8214dc6c6fc20d0a660fb5700dca9ebf51ebe89/src/libcore/fmt/num.rs#L188-L266
macro_rules! impl_Integer {
    ($($t:ident),* as $conv_fn:ident) => ($(
    impl Integer for $t {
        #[allow(unused_comparisons)]
        fn write<W: io::Write>(self, mut wr: W) -> io::Result<()> {
            let is_nonnegative = self >= 0;
            let mut n = if is_nonnegative {
                self as $conv_fn
            } else {
                try!(wr.write_all(b"-"));
                // convert the negative num to positive by summing 1 to it's 2 complement
                (!(self as $conv_fn)).wrapping_add(1)
            };
            let mut buf: [u8; 20] = unsafe { mem::uninitialized() };
            let mut curr = buf.len() as isize;
            let buf_ptr = buf.as_mut_ptr();
            let lut_ptr = DEC_DIGITS_LUT.as_ptr();

            unsafe {
                // eagerly decode 4 characters at a time
                if <$t>::max_value() as u64 >= 10000 {
                    while n >= 10000 {
                        let rem = (n % 10000) as isize;
                        n /= 10000;

                        let d1 = (rem / 100) << 1;
                        let d2 = (rem % 100) << 1;
                        curr -= 4;
                        ptr::copy_nonoverlapping(lut_ptr.offset(d1), buf_ptr.offset(curr), 2);
                        ptr::copy_nonoverlapping(lut_ptr.offset(d2), buf_ptr.offset(curr + 2), 2);
                    }
                }

                // if we reach here numbers are <= 9999, so at most 4 chars long
                let mut n = n as isize; // possibly reduce 64bit math

                // decode 2 more chars, if > 2 chars
                if n >= 100 {
                    let d1 = (n % 100) << 1;
                    n /= 100;
                    curr -= 2;
                    ptr::copy_nonoverlapping(lut_ptr.offset(d1), buf_ptr.offset(curr), 2);
                }

                // decode last 1 or 2 chars
                if n < 10 {
                    curr -= 1;
                    *buf_ptr.offset(curr) = (n as u8) + 48;
                } else {
                    let d1 = n << 1;
                    curr -= 2;
                    ptr::copy_nonoverlapping(lut_ptr.offset(d1), buf_ptr.offset(curr), 2);
                }
            }

            wr.write_all(unsafe {
                slice::from_raw_parts(buf_ptr.offset(curr), buf.len() - curr as usize)
            })
        }
    })*);
}

impl_Integer!(i8, u8, i16, u16, i32, u32 as u32);
impl_Integer!(i64, u64 as u64);
#[cfg(target_pointer_width = "16")]
impl_Integer!(isize, usize as u16);
#[cfg(target_pointer_width = "32")]
impl_Integer!(isize, usize as u32);
#[cfg(target_pointer_width = "64")]
impl_Integer!(isize, usize as u64);
<|endoftext|>


<commit_before>#![allow(
    // We support older compilers
    clippy::match_like_matches_macro,
)]

mod cmd;
mod config;
mod edit;
mod error;
mod fmt;
mod opts;

use std::env;
use std::ffi::OsString;
use std::fs;
use std::io::{self, BufRead, Write};
use std::path::{Path, PathBuf};
use std::process::{self, Command, Stdio};
use std::str::FromStr;

use atty::Stream::{Stderr, Stdout};
use bat::assets::HighlightingAssets;
use bat::PrettyPrinter;
use quote::quote;
use structopt::StructOpt;
use termcolor::{Color::Green, ColorChoice, ColorSpec, StandardStream, WriteColor};

use crate::cmd::Line;
use crate::config::Config;
use crate::error::Result;
use crate::opts::Coloring::*;
use crate::opts::{Args, Coloring, Opts};

fn main() {
    let result = cargo_expand_or_run_nightly();
    process::exit(match result {
        Ok(code) => code,
        Err(err) => {
            let _ = writeln!(io::stderr(), "{}", err);
            1
        }
    });
}

fn cargo_expand_or_run_nightly() -> Result<i32> {
    const NO_RUN_NIGHTLY: &str = "CARGO_EXPAND_NO_RUN_NIGHTLY";

    let maybe_nightly = !definitely_not_nightly();
    if maybe_nightly || env::var_os(NO_RUN_NIGHTLY).is_some() {
        return cargo_expand();
    }

    let mut nightly = Command::new("cargo");
    nightly.arg("+nightly");
    nightly.arg("expand");

    let mut args = env::args_os().peekable();
    args.next().unwrap(); // cargo
    if args.peek().map_or(false, |arg| arg == "expand") {
        args.next().unwrap(); // expand
    }
    nightly.args(args);

    // Hopefully prevent infinite re-run loop.
    nightly.env(NO_RUN_NIGHTLY, "");

    let status = nightly.status()?;

    Ok(match status.code() {
        Some(code) => code,
        None => {
            if status.success() {
                0
            } else {
                1
            }
        }
    })
}

fn definitely_not_nightly() -> bool {
    let mut cmd = Command::new(cargo_binary());
    cmd.arg("--version");

    let output = match cmd.output() {
        Ok(output) => output,
        Err(_) => return false,
    };

    let version = match String::from_utf8(output.stdout) {
        Ok(version) => version,
        Err(_) => return false,
    };

    version.starts_with("cargo 1") && !version.contains("nightly")
}

fn cargo_binary() -> OsString {
    env::var_os("CARGO").unwrap_or_else(|| "cargo".to_owned().into())
}

fn cargo_expand() -> Result<i32> {
    let Opts::Expand(args) = Opts::from_args();
    let config = config::deserialize();

    if args.themes {
        for theme in HighlightingAssets::from_binary().themes() {
            let _ = writeln!(io::stdout(), "{}", theme);
        }
        return Ok(0);
    }

    let rustfmt;
    match (&args.item, args.ugly) {
        (Some(item), true) => {
            let _ = writeln!(
                io::stderr(),
                "ERROR: cannot expand single item ({}) in ugly mode.",
                item,
            );
            return Ok(1);
        }
        (Some(item), false) => {
            rustfmt = which_rustfmt();
            if rustfmt.is_none() {
                let _ = writeln!(
                    io::stderr(),
                    "ERROR: cannot expand single item ({}) without rustfmt.",
                    item,
                );
                let _ = writeln!(
                    io::stderr(),
                    "Install rustfmt by running `rustup component add rustfmt --toolchain nightly`.",
                );
                return Ok(1);
            }
        }
        (None, true) => rustfmt = None,
        (None, false) => rustfmt = which_rustfmt(),
    }

    let mut builder = tempfile::Builder::new();
    builder.prefix("cargo-expand");
    let outdir = builder.tempdir().expect("failed to create tmp file");
    let outfile_path = outdir.path().join("expanded");
    let color = get_color(&args, &config);

    // Run cargo
    let mut cmd = Command::new(cargo_binary());
    apply_args(&mut cmd, &args, &color, &outfile_path);
    let code = filter_err(&mut cmd, ignore_cargo_err)?;

    if !outfile_path.exists() {
        return Ok(1);
    }

    let mut content = fs::read_to_string(&outfile_path)?;
    if content.is_empty() {
        let _ = writeln!(io::stderr(), "ERROR: rustc produced no expanded output");
        return Ok(if code == 0 { 1 } else { code });
    }

    // Run rustfmt
    if let Some(rustfmt) = rustfmt {
        // Work around rustfmt not being able to parse paths containing $crate.
        // This placeholder should be the same width as $crate to preserve
        // alignments.
        const DOLLAR_CRATE_PLACEHOLDER: &str = "Ξcrate";
        content = content.replace("$crate", DOLLAR_CRATE_PLACEHOLDER);

        // Discard comments, which are misplaced by the compiler
        if let Ok(mut syntax_tree) = syn::parse_file(&content) {
            edit::sanitize(&mut syntax_tree);
            if let Some(filter) = args.item {
                syntax_tree.shebang = None;
                syntax_tree.attrs.clear();
                syntax_tree.items = filter.apply_to(&syntax_tree);
                if syntax_tree.items.is_empty() {
                    let _ = writeln!(io::stderr(), "WARNING: no such item: {}", filter);
                    return Ok(1);
                }
            }
            content = quote!(#syntax_tree).to_string();
        }
        fs::write(&outfile_path, content)?;

        fmt::write_rustfmt_config(&outdir)?;

        let output = Command::new(&rustfmt)
            .arg("--edition=2018")
            .arg(&outfile_path)
            .stderr(Stdio::null())
            .output();
        if let Ok(output) = output {
            if !output.status.success() {
                // Probably was the wrong edition.
                let _status = Command::new(&rustfmt)
                    .arg("--edition=2015")
                    .arg(&outfile_path)
                    .stderr(Stdio::null())
                    .status();
            }
        }

        content = fs::read_to_string(&outfile_path)?;
        content = content.replace(DOLLAR_CRATE_PLACEHOLDER, "$crate");
    }

    // Run pretty printer
    let theme = args.theme.or(config.theme);
    let none_theme = theme.as_deref() == Some("none");
    let do_color = match color {
        Always => true,
        Never => false,
        Auto => !none_theme && atty::is(Stdout),
    };
    let _ = writeln!(io::stderr());
    if do_color {
        let mut pretty_printer = PrettyPrinter::new();
        pretty_printer
            .input_from_bytes(content.as_bytes())
            .language("rust")
            .tab_width(Some(4))
            .true_color(false)
            .header(false)
            .line_numbers(false)
            .grid(false);
        if let Some(theme) = theme {
            pretty_printer.theme(theme);
        }

        // Ignore any errors.
        let _ = pretty_printer.print();
    } else {
        let _ = write!(io::stdout(), "{}", content);
    }

    Ok(0)
}

fn which_rustfmt() -> Option<PathBuf> {
    match env::var_os("RUSTFMT") {
        Some(which) => {
            if which.is_empty() {
                None
            } else {
                Some(PathBuf::from(which))
            }
        }
        None => toolchain_find::find_installed_component("rustfmt"),
    }
}

// Based on https://github.com/rsolomo/cargo-check
fn apply_args(cmd: &mut Command, args: &Args, color: &Coloring, outfile: &Path) {
    let mut line = Line::new("cargo");

    line.arg("rustc");

    if args.tests && args.test.is_none() {
        line.arg("--profile=test");
    } else {
        line.arg("--profile=check");
    }

    if args.release {
        line.arg("--release");
    }

    if let Some(features) = &args.features {
        line.arg("--features");
        line.arg(features);
    }

    if args.all_features {
        line.arg("--all-features");
    }

    if args.no_default_features {
        line.arg("--no-default-features");
    }

    if args.lib {
        line.arg("--lib");
    }

    if let Some(bin) = &args.bin {
        line.arg("--bin");
        line.arg(bin);
    }

    if let Some(example) = &args.example {
        line.arg("--example");
        line.arg(example);
    }

    if let Some(test) = &args.test {
        line.arg("--test");
        line.arg(test);
    }

    if let Some(bench) = &args.bench {
        line.arg("--bench");
        line.arg(bench);
    }

    if let Some(target) = &args.target {
        line.arg("--target");
        line.arg(target);
    }

    if let Some(target_dir) = &args.target_dir {
        line.arg("--target-dir");
        line.arg(target_dir);
    }

    if let Some(manifest_path) = &args.manifest_path {
        line.arg("--manifest-path");
        line.arg(manifest_path);
    }

    if let Some(package) = &args.package {
        line.arg("--package");
        line.arg(package);
    }

    if let Some(jobs) = args.jobs {
        line.arg("--jobs");
        line.arg(jobs.to_string());
    }

    if args.verbose {
        line.arg("--verbose");
    }

    line.arg("--color");
    match color {
        Coloring::Auto => line.arg(if atty::is(Stderr) { "always" } else { "never" }),
        color => line.arg(color.to_string()),
    }

    if args.frozen {
        line.arg("--frozen");
    }

    if args.locked {
        line.arg("--locked");
    }

    for unstable_flag in &args.unstable_flags {
        line.arg("-Z");
        line.arg(unstable_flag);
    }

    line.arg("--");

    line.arg("-o");
    line.arg(outfile);
    line.arg("-Zunstable-options");
    line.arg("--pretty=expanded");

    if args.verbose {
        let mut display = line.clone();
        display.insert(0, "+nightly");
        print_command(display, color);
    }

    cmd.args(line);
}

fn print_command(line: Line, color: &Coloring) {
    let color_choice = match color {
        Coloring::Auto => ColorChoice::Auto,
        Coloring::Always => ColorChoice::Always,
        Coloring::Never => ColorChoice::Never,
    };

    let mut stream = StandardStream::stderr(color_choice);
    let _ = stream.set_color(ColorSpec::new().set_bold(true).set_fg(Some(Green)));
    let _ = write!(stream, "{:>12}", "Running");
    let _ = stream.reset();
    let _ = writeln!(stream, " `{}`", line);
}

fn filter_err(cmd: &mut Command, ignore: fn(&str) -> bool) -> io::Result<i32> {
    let mut child = cmd.stderr(Stdio::piped()).spawn()?;
    let mut stderr = io::BufReader::new(child.stderr.take().unwrap());
    let mut line = String::new();
    while let Ok(n) = stderr.read_line(&mut line) {
        if n == 0 {
            break;
        }
        if !ignore(&line) {
            let _ = write!(io::stderr(), "{}", line);
        }
        line.clear();
    }
    let code = child.wait()?.code().unwrap_or(1);
    Ok(code)
}

fn ignore_cargo_err(line: &str) -> bool {
    if line.trim().is_empty() {
        return true;
    }

    let discarded_lines = [
        "ignoring specified output filename because multiple outputs were \
         requested",
        "ignoring specified output filename for 'link' output because multiple \
         outputs were requested",
        "ignoring --out-dir flag due to -o flag",
        "ignoring -C extra-filename flag due to -o flag",
        "due to multiple output types requested, the explicitly specified \
         output file name will be adapted for each output type",
        "warning emitted",
        "warnings emitted",
    ];
    for s in &discarded_lines {
        if line.contains(s) {
            return true;
        }
    }

    false
}

fn get_color(args: &Args, config: &Config) -> Coloring {
    if let Some(value) = args.color {
        return value;
    }

    if let Some(value) = config.color.as_ref() {
        match Coloring::from_str(value.as_str()) {
            Ok(color) => return color,
            Err(err) => {
                let _ = writeln!(
                    io::stderr(),
                    "WARNING: invalid color in cargo config: {}",
                    err
                );
            }
        }
    }

    Coloring::Auto // default
}
<commit_msg>Support cargo-expand built with panic=abort<commit_after>#![allow(
    // We support older compilers
    clippy::match_like_matches_macro,
)]

mod cmd;
mod config;
mod edit;
mod error;
mod fmt;
mod opts;

use std::env;
use std::ffi::OsString;
use std::fs;
use std::io::{self, BufRead, Write};
use std::path::{Path, PathBuf};
use std::process::{self, Command, Stdio};
use std::str::FromStr;

use atty::Stream::{Stderr, Stdout};
use bat::assets::HighlightingAssets;
use bat::PrettyPrinter;
use quote::quote;
use structopt::StructOpt;
use termcolor::{Color::Green, ColorChoice, ColorSpec, StandardStream, WriteColor};

use crate::cmd::Line;
use crate::config::Config;
use crate::error::Result;
use crate::opts::Coloring::*;
use crate::opts::{Args, Coloring, Opts};

fn main() {
    let result = cargo_expand_or_run_nightly();
    process::exit(match result {
        Ok(code) => code,
        Err(err) => {
            let _ = writeln!(io::stderr(), "{}", err);
            1
        }
    });
}

fn cargo_expand_or_run_nightly() -> Result<i32> {
    const NO_RUN_NIGHTLY: &str = "CARGO_EXPAND_NO_RUN_NIGHTLY";

    let maybe_nightly = !definitely_not_nightly();
    if maybe_nightly || env::var_os(NO_RUN_NIGHTLY).is_some() {
        return cargo_expand();
    }

    let mut nightly = Command::new("cargo");
    nightly.arg("+nightly");
    nightly.arg("expand");

    let mut args = env::args_os().peekable();
    args.next().unwrap(); // cargo
    if args.peek().map_or(false, |arg| arg == "expand") {
        args.next().unwrap(); // expand
    }
    nightly.args(args);

    // Hopefully prevent infinite re-run loop.
    nightly.env(NO_RUN_NIGHTLY, "");

    let status = nightly.status()?;

    Ok(match status.code() {
        Some(code) => code,
        None => {
            if status.success() {
                0
            } else {
                1
            }
        }
    })
}

fn definitely_not_nightly() -> bool {
    let mut cmd = Command::new(cargo_binary());
    cmd.arg("--version");

    let output = match cmd.output() {
        Ok(output) => output,
        Err(_) => return false,
    };

    let version = match String::from_utf8(output.stdout) {
        Ok(version) => version,
        Err(_) => return false,
    };

    version.starts_with("cargo 1") && !version.contains("nightly")
}

fn cargo_binary() -> OsString {
    env::var_os("CARGO").unwrap_or_else(|| "cargo".to_owned().into())
}

fn cargo_expand() -> Result<i32> {
    let Opts::Expand(args) = Opts::from_args();
    let config = config::deserialize();

    if args.themes {
        for theme in HighlightingAssets::from_binary().themes() {
            let _ = writeln!(io::stdout(), "{}", theme);
        }
        return Ok(0);
    }

    let rustfmt;
    match (&args.item, args.ugly) {
        (Some(item), true) => {
            let _ = writeln!(
                io::stderr(),
                "ERROR: cannot expand single item ({}) in ugly mode.",
                item,
            );
            return Ok(1);
        }
        (Some(item), false) => {
            rustfmt = which_rustfmt();
            if rustfmt.is_none() {
                let _ = writeln!(
                    io::stderr(),
                    "ERROR: cannot expand single item ({}) without rustfmt.",
                    item,
                );
                let _ = writeln!(
                    io::stderr(),
                    "Install rustfmt by running `rustup component add rustfmt --toolchain nightly`.",
                );
                return Ok(1);
            }
        }
        (None, true) => rustfmt = None,
        (None, false) => rustfmt = which_rustfmt(),
    }

    let mut builder = tempfile::Builder::new();
    builder.prefix("cargo-expand");
    let outdir = builder.tempdir().expect("failed to create tmp file");
    let outfile_path = outdir.path().join("expanded");
    let color = get_color(&args, &config);

    // Run cargo
    let mut cmd = Command::new(cargo_binary());
    apply_args(&mut cmd, &args, &color, &outfile_path);
    let code = filter_err(&mut cmd, ignore_cargo_err)?;

    if !outfile_path.exists() {
        return Ok(1);
    }

    let mut content = fs::read_to_string(&outfile_path)?;
    if content.is_empty() {
        let _ = writeln!(io::stderr(), "ERROR: rustc produced no expanded output");
        return Ok(if code == 0 { 1 } else { code });
    }

    // Run rustfmt
    if let Some(rustfmt) = rustfmt {
        // Work around rustfmt not being able to parse paths containing $crate.
        // This placeholder should be the same width as $crate to preserve
        // alignments.
        const DOLLAR_CRATE_PLACEHOLDER: &str = "Ξcrate";
        content = content.replace("$crate", DOLLAR_CRATE_PLACEHOLDER);

        // Support cargo-expand built with panic=abort, as otherwise proc-macro2
        // ends up using a catch_unwind.
        proc_macro2::fallback::force();

        // Discard comments, which are misplaced by the compiler
        if let Ok(mut syntax_tree) = syn::parse_file(&content) {
            edit::sanitize(&mut syntax_tree);
            if let Some(filter) = args.item {
                syntax_tree.shebang = None;
                syntax_tree.attrs.clear();
                syntax_tree.items = filter.apply_to(&syntax_tree);
                if syntax_tree.items.is_empty() {
                    let _ = writeln!(io::stderr(), "WARNING: no such item: {}", filter);
                    return Ok(1);
                }
            }
            content = quote!(#syntax_tree).to_string();
        }
        fs::write(&outfile_path, content)?;

        fmt::write_rustfmt_config(&outdir)?;

        let output = Command::new(&rustfmt)
            .arg("--edition=2018")
            .arg(&outfile_path)
            .stderr(Stdio::null())
            .output();
        if let Ok(output) = output {
            if !output.status.success() {
                // Probably was the wrong edition.
                let _status = Command::new(&rustfmt)
                    .arg("--edition=2015")
                    .arg(&outfile_path)
                    .stderr(Stdio::null())
                    .status();
            }
        }

        content = fs::read_to_string(&outfile_path)?;
        content = content.replace(DOLLAR_CRATE_PLACEHOLDER, "$crate");
    }

    // Run pretty printer
    let theme = args.theme.or(config.theme);
    let none_theme = theme.as_deref() == Some("none");
    let do_color = match color {
        Always => true,
        Never => false,
        Auto => !none_theme && atty::is(Stdout),
    };
    let _ = writeln!(io::stderr());
    if do_color {
        let mut pretty_printer = PrettyPrinter::new();
        pretty_printer
            .input_from_bytes(content.as_bytes())
            .language("rust")
            .tab_width(Some(4))
            .true_color(false)
            .header(false)
            .line_numbers(false)
            .grid(false);
        if let Some(theme) = theme {
            pretty_printer.theme(theme);
        }

        // Ignore any errors.
        let _ = pretty_printer.print();
    } else {
        let _ = write!(io::stdout(), "{}", content);
    }

    Ok(0)
}

fn which_rustfmt() -> Option<PathBuf> {
    match env::var_os("RUSTFMT") {
        Some(which) => {
            if which.is_empty() {
                None
            } else {
                Some(PathBuf::from(which))
            }
        }
        None => toolchain_find::find_installed_component("rustfmt"),
    }
}

// Based on https://github.com/rsolomo/cargo-check
fn apply_args(cmd: &mut Command, args: &Args, color: &Coloring, outfile: &Path) {
    let mut line = Line::new("cargo");

    line.arg("rustc");

    if args.tests && args.test.is_none() {
        line.arg("--profile=test");
    } else {
        line.arg("--profile=check");
    }

    if args.release {
        line.arg("--release");
    }

    if let Some(features) = &args.features {
        line.arg("--features");
        line.arg(features);
    }

    if args.all_features {
        line.arg("--all-features");
    }

    if args.no_default_features {
        line.arg("--no-default-features");
    }

    if args.lib {
        line.arg("--lib");
    }

    if let Some(bin) = &args.bin {
        line.arg("--bin");
        line.arg(bin);
    }

    if let Some(example) = &args.example {
        line.arg("--example");
        line.arg(example);
    }

    if let Some(test) = &args.test {
        line.arg("--test");
        line.arg(test);
    }

    if let Some(bench) = &args.bench {
        line.arg("--bench");
        line.arg(bench);
    }

    if let Some(target) = &args.target {
        line.arg("--target");
        line.arg(target);
    }

    if let Some(target_dir) = &args.target_dir {
        line.arg("--target-dir");
        line.arg(target_dir);
    }

    if let Some(manifest_path) = &args.manifest_path {
        line.arg("--manifest-path");
        line.arg(manifest_path);
    }

    if let Some(package) = &args.package {
        line.arg("--package");
        line.arg(package);
    }

    if let Some(jobs) = args.jobs {
        line.arg("--jobs");
        line.arg(jobs.to_string());
    }

    if args.verbose {
        line.arg("--verbose");
    }

    line.arg("--color");
    match color {
        Coloring::Auto => line.arg(if atty::is(Stderr) { "always" } else { "never" }),
        color => line.arg(color.to_string()),
    }

    if args.frozen {
        line.arg("--frozen");
    }

    if args.locked {
        line.arg("--locked");
    }

    for unstable_flag in &args.unstable_flags {
        line.arg("-Z");
        line.arg(unstable_flag);
    }

    line.arg("--");

    line.arg("-o");
    line.arg(outfile);
    line.arg("-Zunstable-options");
    line.arg("--pretty=expanded");

    if args.verbose {
        let mut display = line.clone();
        display.insert(0, "+nightly");
        print_command(display, color);
    }

    cmd.args(line);
}

fn print_command(line: Line, color: &Coloring) {
    let color_choice = match color {
        Coloring::Auto => ColorChoice::Auto,
        Coloring::Always => ColorChoice::Always,
        Coloring::Never => ColorChoice::Never,
    };

    let mut stream = StandardStream::stderr(color_choice);
    let _ = stream.set_color(ColorSpec::new().set_bold(true).set_fg(Some(Green)));
    let _ = write!(stream, "{:>12}", "Running");
    let _ = stream.reset();
    let _ = writeln!(stream, " `{}`", line);
}

fn filter_err(cmd: &mut Command, ignore: fn(&str) -> bool) -> io::Result<i32> {
    let mut child = cmd.stderr(Stdio::piped()).spawn()?;
    let mut stderr = io::BufReader::new(child.stderr.take().unwrap());
    let mut line = String::new();
    while let Ok(n) = stderr.read_line(&mut line) {
        if n == 0 {
            break;
        }
        if !ignore(&line) {
            let _ = write!(io::stderr(), "{}", line);
        }
        line.clear();
    }
    let code = child.wait()?.code().unwrap_or(1);
    Ok(code)
}

fn ignore_cargo_err(line: &str) -> bool {
    if line.trim().is_empty() {
        return true;
    }

    let discarded_lines = [
        "ignoring specified output filename because multiple outputs were \
         requested",
        "ignoring specified output filename for 'link' output because multiple \
         outputs were requested",
        "ignoring --out-dir flag due to -o flag",
        "ignoring -C extra-filename flag due to -o flag",
        "due to multiple output types requested, the explicitly specified \
         output file name will be adapted for each output type",
        "warning emitted",
        "warnings emitted",
    ];
    for s in &discarded_lines {
        if line.contains(s) {
            return true;
        }
    }

    false
}

fn get_color(args: &Args, config: &Config) -> Coloring {
    if let Some(value) = args.color {
        return value;
    }

    if let Some(value) = config.color.as_ref() {
        match Coloring::from_str(value.as_str()) {
            Ok(color) => return color,
            Err(err) => {
                let _ = writeln!(
                    io::stderr(),
                    "WARNING: invalid color in cargo config: {}",
                    err
                );
            }
        }
    }

    Coloring::Auto // default
}
<|endoftext|>


<commit_before>#![feature(std_misc)]
#![cfg_attr(test, feature(test))]

#[cfg(test)] extern crate test;

use std::ops::Index;
use std::default::Default;
use std::hash::{Hasher, Hash, SipHasher};
use std::iter::{FromIterator, IntoIterator};
use std::borrow::Borrow;
use std::collections::HashMap;
use std::collections::hash_map::{Keys, IntoIter, Iter};
use std::collections::hash_state::{HashState, DefaultState};

/// Distribution doesn't require a cryptographically secure hash, and by
/// default will not use one.
pub trait Distribution<K, H = SipHasher> {
  type Quantity;

  fn len(&self) -> usize;
  fn get<Q: ?Sized>(&self, k: &Q) -> Self::Quantity
    where K: Borrow<Q>, Q: Hash + Eq;
  fn clear(&mut self);
  fn insert(&mut self, k: K);
  fn remove<Q: ?Sized>(&mut self, k: &Q)
    where K: Borrow<Q>, Q: Hash + Eq;
}

/// Implementation of a Frequency Distribution in Rust. Keeps track of how many
/// times an object appears in a larger context (for example, how many times a
/// token appears in a piece of text). The underlying data structure of the
/// Frequency Distribution is a HashMap, so the object that is being counted
/// must be hashable.
///
/// # Example
///
/// ```rust
/// use freqdist::{Distribution, FrequencyDistribution};
///
/// let mut fdist = FrequencyDistribution::new();
///
/// fdist.insert("hello");
/// fdist.insert("hello");
/// fdist.insert("goodbye");
///
/// assert_eq!(fdist.get(&"hello"), 2);
/// ```
pub struct FrequencyDistribution<K, S = DefaultState<SipHasher>> {
  hashmap: HashMap<K, usize, S>,
  sum_counts: usize
}

/// Used as a return value if a key is not found.
static ZERO: &'static usize = &0;

impl<K> FrequencyDistribution<K>
  where K: Eq + Hash
{
  /// Creates a new FrequencyDistribution where the size of the
  /// HashMap is unknown.
  #[inline]
  pub fn new() -> FrequencyDistribution<K> {
    FrequencyDistribution::with_hash_state(Default::default())
  }

  /// Creates a new FrequencyDistribution where the size of the HashMap
  /// is known, or a estimate can be made.
  #[inline]
  pub fn with_capacity(
    size: usize
  ) -> FrequencyDistribution<K> {
    FrequencyDistribution::with_capacity_and_hash_state(size, Default::default())
  }
}

impl<K, S, H> FrequencyDistribution<K, S>
  where K: Eq + Hash,
        S: HashState<Hasher=H>,
        H: Hasher
{
  /// Creates a new FrequencyDistrbution with a hash state and size, where
  /// the size is known or can be estimated.
  #[inline]
  pub fn with_capacity_and_hash_state(
    size: usize,
    state: S
  ) -> FrequencyDistribution<K, S> {
    FrequencyDistribution {
      hashmap: HashMap::with_capacity_and_hash_state(size, state),
      sum_counts: 0
    }
  }

  /// Creates a new FrequencyDistribution with a hash state and default size.
  #[inline]
  pub fn with_hash_state(state: S) -> FrequencyDistribution<K, S> {
    FrequencyDistribution {
      hashmap: HashMap::with_hash_state(state),
      sum_counts: 0
    }
  }
}

impl<K, S, H> FrequencyDistribution<K, S>
  where K: Eq + Hash,
        S: HashState<Hasher=H>,
        H: Hasher
{
  /// Iterator over the keys.
  #[inline]
  #[stable]
  pub fn keys(&self) -> Keys<K, usize> {
    self.hashmap.keys()
  }

  /// Iterator over the key, frequency pairs.
  #[inline]
  #[stable]
  pub fn iter(&self) -> Iter<K, usize> {
    self.hashmap.iter()
  }

  /// Iterator over the non-zero frequency keys.
  ///
  /// # Example
  ///
  /// ```rust
  /// use std::iter::{FromIterator, IntoIterator};
  /// use freqdist::FrequencyDistribution;
  ///
  /// let existing = vec![
  ///   ("shoes", 1),
  ///   ("scarf", 0),
  ///   ("shirt", 13),
  ///   ("pants", 4)
  /// ];
  ///
  /// let fdist: FrequencyDistribution<&str> =
  ///   FromIterator::from_iter(existing.into_iter());
  /// let mut iter = fdist.iter_non_zero();
  ///
  /// assert_eq!(*iter.next().unwrap(), "shirt");
  /// assert_eq!(*iter.next().unwrap(), "shoes");
  /// assert_eq!(*iter.next().unwrap(), "pants");
  /// assert!(iter.next().is_none());
  /// ```
  #[inline]
  #[unstable]
  pub fn iter_non_zero(&self) -> NonZeroKeysIter<K> {
    NonZeroKeysIter { iter: self.iter() }
  }

  /// Sum of the total number of items counted thus far.
  #[inline]
  #[stable]
  pub fn sum_counts(&self) -> usize {
    self.sum_counts
  }

  /// Inserts a value sizeo the hashmap if it does not exist with a new quantity
  /// specified by the increment. If the value already exists, increments by
  /// the specified amount.
  #[inline]
  fn insert_or_incr_by(&mut self, k: K, incr: usize) {
    if !self.hashmap.contains_key(&k) {
      self.hashmap.insert(k, incr);
    } else {
      *self.hashmap.get_mut(&k).unwrap() += incr;
    }

    self.sum_counts += incr;
  }
}

impl<K> Default for FrequencyDistribution<K>
  where K: Eq + Hash
{
  /// Creates a default FrequencyDistribution with an XXHasher.
  #[inline]
  #[stable]
  fn default() -> FrequencyDistribution<K> {
    FrequencyDistribution::new()
  }
}

impl<K, S, H> FromIterator<(K, usize)> for FrequencyDistribution<K, S>
  where K: Eq + Hash,
        S: HashState<Hasher = H> + Default,
        H: Hasher
{
  /// Iterates through an iterator, and creates a new FrequencyDistribution from
  /// it. The iterator should be an iterator over keys and frequencies. If a
  /// upper bounded `size_hsize` is available, then it is used, otherwise the lower
  /// bounded `size_hsize` is used.
  ///
  /// # Example
  ///
  /// ```rust
  /// use std::iter::FromIterator;
  /// use freqdist::{Distribution, FrequencyDistribution};
  ///
  /// let existing = vec![
  ///   ("apples", 3),
  ///   ("oranges", 4),
  ///   ("bannana", 7)
  /// ];
  ///
  /// let fdist: FrequencyDistribution<&str> =
  ///   FromIterator::from_iter(existing.into_iter());
  ///
  /// assert_eq!(fdist.get(&"apples"), 3);
  /// assert_eq!(fdist.get(&"oranges"), 4);
  /// assert_eq!(fdist.get(&"bannana"), 7);
  /// ```
  fn from_iter<T>(iter: T) -> FrequencyDistribution<K, S>
    where T: IntoIterator<Item = (K, usize)>
  {
    let iterator = iter.into_iter();
    let mut fdist = if iterator.size_hint().1.is_some() {
      FrequencyDistribution::with_capacity_and_hash_state(
        iterator.size_hint().1.unwrap(),
        Default::default())
    } else {
      FrequencyDistribution::with_capacity_and_hash_state(
        iterator.size_hint().0,
        Default::default())
    };

    for (k, freq) in iterator { fdist.insert_or_incr_by(k, freq); }

    fdist
  }
}

impl<K, S, H> Extend<(K, usize)> for FrequencyDistribution<K, S>
  where K: Eq + Hash,
        S: HashState<Hasher=H>,
        H: Hasher
{
  /// Extends the hashmap by adding the keys or updating the frequencies of the keys.
  #[inline]
  fn extend<T>(&mut self, iter: T)
    where T: IntoIterator<Item = (K, usize)>
  {
    for (k, freq) in iter.into_iter() { self.insert_or_incr_by(k, freq); }
  }
}

impl<K, S, H> Distribution<K, H> for FrequencyDistribution<K, S>
  where K: Eq + Hash,
        S: HashState<Hasher=H>,
        H: Hasher
{
  type Quantity = usize;

  /// Returns the number of entries in the distribution
  #[inline]
  #[stable]
  fn len(&self) -> usize {
    self.hashmap.len()
  }

  /// Gets the frequency in which the key occurs.
  #[inline]
  #[stable]
  fn get<Q: ?Sized>(&self, k: &Q) -> usize
    where K: Borrow<Q>, Q: Hash + Eq
  {
    self[k]
  }

  /// Clears the counts of all keys and clears all keys from
  /// the distribution.
  #[inline]
  #[stable]
  fn clear(&mut self) {
    self.hashmap.clear()
  }

  /// Updates the frequency of the value found with the key if it
  /// already exists. Otherwise, inserts the key sizeo the hashmap,
  /// and sets its frequency to 1.
  #[inline]
  #[stable]
  fn insert(&mut self, k: K) {
    self.insert_or_incr_by(k, 1);
  }

  /// Removes a Key and its associated value from the Distrbution.
  #[inline]
  #[stable]
  fn remove<Q: ?Sized>(&mut self, k: &Q)
    where K: Borrow<Q>, Q: Hash + Eq
  {
    match self.hashmap.remove(k) {
      Some(count) => self.sum_counts -= count,
      None        => ()
    }
  }
}

impl<K, S, H> IntoIterator for FrequencyDistribution<K, S>
  where K: Eq + Hash,
        S: HashState<Hasher=H>,
        H: Hasher
{
  type Item = (K, usize);
  type IntoIter = IntoIter<K, usize>;

  /// Consumes the distribution, and creates an iterator over the
  /// (Key, Quantity: usize) pairs.
  #[inline]
  #[stable]
  fn into_iter(self) -> IntoIter<K, usize> {
    self.hashmap.into_iter()
  }
}

impl<'a, K, S, H, Q: ?Sized> Index<&'a Q> for FrequencyDistribution<K, S>
  where K: Eq + Hash + Borrow<Q>,
        S: HashState<Hasher=H>,
        H: Hasher,
        Q: Eq + Hash
{
  type Output = usize;

  fn index<'b>(&'b self, index: &Q) -> &'b usize {
    self.hashmap.get(index).unwrap_or(ZERO)
  }
}

/// Iterator over entries with non-zero quantities.
#[stable]
pub struct NonZeroKeysIter<'a, K: 'a> {
  iter: Iter<'a, K, usize>
}

impl<'a, K: 'a> Iterator for NonZeroKeysIter<'a, K> {
  type Item = &'a K;

  #[inline]
  fn next(&mut self) -> Option<&'a K> {
    loop {
      match self.iter.next() {
        Some((k, c)) if *c > 0 => return Some(k),
        None                   => return None,
        _                      => ()
      }
    }
  }

  #[inline]
  fn size_hint(&self) -> (usize, Option<usize>) {
    self.iter.size_hint()
  }
}

#[test]
fn smoke_test_frequency_distribution_insert() {
  let words = vec!("alpha", "beta");
  let mut dist: FrequencyDistribution<&str> = FrequencyDistribution::new();

  dist.insert(words[0]);

  assert_eq!(dist.get(&words[0]), 1);

  dist.insert(words[1]);

  assert_eq!(dist.get(&words[1]), 1);

  for _ in (0..7us) { dist.insert(words[0]); }

  assert_eq!(dist.get(&words[0]), 8);
}

#[test]
fn smoke_test_frequency_distribution_iter() {
  let words = vec!(("a", 50us), ("b", 100us), ("c", 75us), ("d", 0us));
  let dist: FrequencyDistribution<&str> = FromIterator::from_iter(words.into_iter());

  assert_eq!(dist.get(&"a"), 50);
  assert_eq!(dist.get(&"b"), 100);
  assert_eq!(dist.get(&"c"), 75);

  let mut iter = dist.iter_non_zero();

  assert!(iter.next().is_some());
  assert!(iter.next().is_some());
  assert!(iter.next().is_some());
  assert!(iter.next().is_none());

  assert_eq!(dist.sum_counts(), 225);
}

#[test]
fn smoke_test_frequency_distribution_remove() {
  let words = vec!(("a", 50us), ("b", 100us), ("c", 25us));
  let mut dist: FrequencyDistribution<&str> = FromIterator::from_iter(words.into_iter());

  assert_eq!(dist.get(&"a"), 50);

  dist.remove(&"a");

  assert_eq!(dist.get(&"a"), 0);
  assert_eq!(dist.sum_counts(), 125);
}

#[test]
fn smoke_test_frequency_sum_counts() {
  let words = vec!(("a", 7us), ("b", 5us), ("c", 8us), ("d", 3us));
  let mut dist: FrequencyDistribution<&str> = FromIterator::from_iter(words.into_iter());

  assert_eq!(dist.sum_counts(), 23);

  dist.insert("e");

  assert_eq!(dist.sum_counts(), 24);
}
<commit_msg>replace 'us' with 'usize'<commit_after>#![feature(std_misc)]
#![cfg_attr(test, feature(test))]

#[cfg(test)] extern crate test;

use std::ops::Index;
use std::default::Default;
use std::hash::{Hasher, Hash, SipHasher};
use std::iter::{FromIterator, IntoIterator};
use std::borrow::Borrow;
use std::collections::HashMap;
use std::collections::hash_map::{Keys, IntoIter, Iter};
use std::collections::hash_state::{HashState, DefaultState};

/// Distribution doesn't require a cryptographically secure hash, and by
/// default will not use one.
pub trait Distribution<K, H = SipHasher> {
  type Quantity;

  fn len(&self) -> usize;
  fn get<Q: ?Sized>(&self, k: &Q) -> Self::Quantity
    where K: Borrow<Q>, Q: Hash + Eq;
  fn clear(&mut self);
  fn insert(&mut self, k: K);
  fn remove<Q: ?Sized>(&mut self, k: &Q)
    where K: Borrow<Q>, Q: Hash + Eq;
}

/// Implementation of a Frequency Distribution in Rust. Keeps track of how many
/// times an object appears in a larger context (for example, how many times a
/// token appears in a piece of text). The underlying data structure of the
/// Frequency Distribution is a HashMap, so the object that is being counted
/// must be hashable.
///
/// # Example
///
/// ```rust
/// use freqdist::{Distribution, FrequencyDistribution};
///
/// let mut fdist = FrequencyDistribution::new();
///
/// fdist.insert("hello");
/// fdist.insert("hello");
/// fdist.insert("goodbye");
///
/// assert_eq!(fdist.get(&"hello"), 2);
/// ```
pub struct FrequencyDistribution<K, S = DefaultState<SipHasher>> {
  hashmap: HashMap<K, usize, S>,
  sum_counts: usize
}

/// Used as a return value if a key is not found.
static ZERO: &'static usize = &0;

impl<K> FrequencyDistribution<K>
  where K: Eq + Hash
{
  /// Creates a new FrequencyDistribution where the size of the
  /// HashMap is unknown.
  #[inline]
  pub fn new() -> FrequencyDistribution<K> {
    FrequencyDistribution::with_hash_state(Default::default())
  }

  /// Creates a new FrequencyDistribution where the size of the HashMap
  /// is known, or a estimate can be made.
  #[inline]
  pub fn with_capacity(
    size: usize
  ) -> FrequencyDistribution<K> {
    FrequencyDistribution::with_capacity_and_hash_state(size, Default::default())
  }
}

impl<K, S, H> FrequencyDistribution<K, S>
  where K: Eq + Hash,
        S: HashState<Hasher=H>,
        H: Hasher
{
  /// Creates a new FrequencyDistrbution with a hash state and size, where
  /// the size is known or can be estimated.
  #[inline]
  pub fn with_capacity_and_hash_state(
    size: usize,
    state: S
  ) -> FrequencyDistribution<K, S> {
    FrequencyDistribution {
      hashmap: HashMap::with_capacity_and_hash_state(size, state),
      sum_counts: 0
    }
  }

  /// Creates a new FrequencyDistribution with a hash state and default size.
  #[inline]
  pub fn with_hash_state(state: S) -> FrequencyDistribution<K, S> {
    FrequencyDistribution {
      hashmap: HashMap::with_hash_state(state),
      sum_counts: 0
    }
  }
}

impl<K, S, H> FrequencyDistribution<K, S>
  where K: Eq + Hash,
        S: HashState<Hasher=H>,
        H: Hasher
{
  /// Iterator over the keys.
  #[inline]
  #[stable]
  pub fn keys(&self) -> Keys<K, usize> {
    self.hashmap.keys()
  }

  /// Iterator over the key, frequency pairs.
  #[inline]
  #[stable]
  pub fn iter(&self) -> Iter<K, usize> {
    self.hashmap.iter()
  }

  /// Iterator over the non-zero frequency keys.
  ///
  /// # Example
  ///
  /// ```rust
  /// use std::iter::{FromIterator, IntoIterator};
  /// use freqdist::FrequencyDistribution;
  ///
  /// let existing = vec![
  ///   ("shoes", 1),
  ///   ("scarf", 0),
  ///   ("shirt", 13),
  ///   ("pants", 4)
  /// ];
  ///
  /// let fdist: FrequencyDistribution<&str> =
  ///   FromIterator::from_iter(existing.into_iter());
  /// let mut iter = fdist.iter_non_zero();
  ///
  /// assert_eq!(*iter.next().unwrap(), "shirt");
  /// assert_eq!(*iter.next().unwrap(), "shoes");
  /// assert_eq!(*iter.next().unwrap(), "pants");
  /// assert!(iter.next().is_none());
  /// ```
  #[inline]
  #[unstable]
  pub fn iter_non_zero(&self) -> NonZeroKeysIter<K> {
    NonZeroKeysIter { iter: self.iter() }
  }

  /// Sum of the total number of items counted thus far.
  #[inline]
  #[stable]
  pub fn sum_counts(&self) -> usize {
    self.sum_counts
  }

  /// Inserts a value sizeo the hashmap if it does not exist with a new quantity
  /// specified by the increment. If the value already exists, increments by
  /// the specified amount.
  #[inline]
  fn insert_or_incr_by(&mut self, k: K, incr: usize) {
    if !self.hashmap.contains_key(&k) {
      self.hashmap.insert(k, incr);
    } else {
      *self.hashmap.get_mut(&k).unwrap() += incr;
    }

    self.sum_counts += incr;
  }
}

impl<K> Default for FrequencyDistribution<K>
  where K: Eq + Hash
{
  /// Creates a default FrequencyDistribution with an XXHasher.
  #[inline]
  #[stable]
  fn default() -> FrequencyDistribution<K> {
    FrequencyDistribution::new()
  }
}

impl<K, S, H> FromIterator<(K, usize)> for FrequencyDistribution<K, S>
  where K: Eq + Hash,
        S: HashState<Hasher = H> + Default,
        H: Hasher
{
  /// Iterates through an iterator, and creates a new FrequencyDistribution from
  /// it. The iterator should be an iterator over keys and frequencies. If a
  /// upper bounded `size_hsize` is available, then it is used, otherwise the lower
  /// bounded `size_hsize` is used.
  ///
  /// # Example
  ///
  /// ```rust
  /// use std::iter::FromIterator;
  /// use freqdist::{Distribution, FrequencyDistribution};
  ///
  /// let existing = vec![
  ///   ("apples", 3),
  ///   ("oranges", 4),
  ///   ("bannana", 7)
  /// ];
  ///
  /// let fdist: FrequencyDistribution<&str> =
  ///   FromIterator::from_iter(existing.into_iter());
  ///
  /// assert_eq!(fdist.get(&"apples"), 3);
  /// assert_eq!(fdist.get(&"oranges"), 4);
  /// assert_eq!(fdist.get(&"bannana"), 7);
  /// ```
  fn from_iter<T>(iter: T) -> FrequencyDistribution<K, S>
    where T: IntoIterator<Item = (K, usize)>
  {
    let iterator = iter.into_iter();
    let mut fdist = if iterator.size_hint().1.is_some() {
      FrequencyDistribution::with_capacity_and_hash_state(
        iterator.size_hint().1.unwrap(),
        Default::default())
    } else {
      FrequencyDistribution::with_capacity_and_hash_state(
        iterator.size_hint().0,
        Default::default())
    };

    for (k, freq) in iterator { fdist.insert_or_incr_by(k, freq); }

    fdist
  }
}

impl<K, S, H> Extend<(K, usize)> for FrequencyDistribution<K, S>
  where K: Eq + Hash,
        S: HashState<Hasher=H>,
        H: Hasher
{
  /// Extends the hashmap by adding the keys or updating the frequencies of the keys.
  #[inline]
  fn extend<T>(&mut self, iter: T)
    where T: IntoIterator<Item = (K, usize)>
  {
    for (k, freq) in iter.into_iter() { self.insert_or_incr_by(k, freq); }
  }
}

impl<K, S, H> Distribution<K, H> for FrequencyDistribution<K, S>
  where K: Eq + Hash,
        S: HashState<Hasher=H>,
        H: Hasher
{
  type Quantity = usize;

  /// Returns the number of entries in the distribution
  #[inline]
  #[stable]
  fn len(&self) -> usize {
    self.hashmap.len()
  }

  /// Gets the frequency in which the key occurs.
  #[inline]
  #[stable]
  fn get<Q: ?Sized>(&self, k: &Q) -> usize
    where K: Borrow<Q>, Q: Hash + Eq
  {
    self[k]
  }

  /// Clears the counts of all keys and clears all keys from
  /// the distribution.
  #[inline]
  #[stable]
  fn clear(&mut self) {
    self.hashmap.clear()
  }

  /// Updates the frequency of the value found with the key if it
  /// already exists. Otherwise, inserts the key sizeo the hashmap,
  /// and sets its frequency to 1.
  #[inline]
  #[stable]
  fn insert(&mut self, k: K) {
    self.insert_or_incr_by(k, 1);
  }

  /// Removes a Key and its associated value from the Distrbution.
  #[inline]
  #[stable]
  fn remove<Q: ?Sized>(&mut self, k: &Q)
    where K: Borrow<Q>, Q: Hash + Eq
  {
    match self.hashmap.remove(k) {
      Some(count) => self.sum_counts -= count,
      None        => ()
    }
  }
}

impl<K, S, H> IntoIterator for FrequencyDistribution<K, S>
  where K: Eq + Hash,
        S: HashState<Hasher=H>,
        H: Hasher
{
  type Item = (K, usize);
  type IntoIter = IntoIter<K, usize>;

  /// Consumes the distribution, and creates an iterator over the
  /// (Key, Quantity: usize) pairs.
  #[inline]
  #[stable]
  fn into_iter(self) -> IntoIter<K, usize> {
    self.hashmap.into_iter()
  }
}

impl<'a, K, S, H, Q: ?Sized> Index<&'a Q> for FrequencyDistribution<K, S>
  where K: Eq + Hash + Borrow<Q>,
        S: HashState<Hasher=H>,
        H: Hasher,
        Q: Eq + Hash
{
  type Output = usize;

  fn index<'b>(&'b self, index: &Q) -> &'b usize {
    self.hashmap.get(index).unwrap_or(ZERO)
  }
}

/// Iterator over entries with non-zero quantities.
#[stable]
pub struct NonZeroKeysIter<'a, K: 'a> {
  iter: Iter<'a, K, usize>
}

impl<'a, K: 'a> Iterator for NonZeroKeysIter<'a, K> {
  type Item = &'a K;

  #[inline]
  fn next(&mut self) -> Option<&'a K> {
    loop {
      match self.iter.next() {
        Some((k, c)) if *c > 0 => return Some(k),
        None                   => return None,
        _                      => ()
      }
    }
  }

  #[inline]
  fn size_hint(&self) -> (usize, Option<usize>) {
    self.iter.size_hint()
  }
}

#[test]
fn smoke_test_frequency_distribution_insert() {
  let words = vec!("alpha", "beta");
  let mut dist: FrequencyDistribution<&str> = FrequencyDistribution::new();

  dist.insert(words[0]);

  assert_eq!(dist.get(&words[0]), 1);

  dist.insert(words[1]);

  assert_eq!(dist.get(&words[1]), 1);

  for _ in (0..7u32) { dist.insert(words[0]); }

  assert_eq!(dist.get(&words[0]), 8);
}

#[test]
fn smoke_test_frequency_distribution_iter() {
  let words = vec!(("a", 50usize), ("b", 100usize), ("c", 75usize), ("d", 0usize));
  let dist: FrequencyDistribution<&str> = FromIterator::from_iter(words.into_iter());

  assert_eq!(dist.get(&"a"), 50);
  assert_eq!(dist.get(&"b"), 100);
  assert_eq!(dist.get(&"c"), 75);

  let mut iter = dist.iter_non_zero();

  assert!(iter.next().is_some());
  assert!(iter.next().is_some());
  assert!(iter.next().is_some());
  assert!(iter.next().is_none());

  assert_eq!(dist.sum_counts(), 225);
}

#[test]
fn smoke_test_frequency_distribution_remove() {
  let words = vec!(("a", 50usize), ("b", 100usize), ("c", 25usize));
  let mut dist: FrequencyDistribution<&str> = FromIterator::from_iter(words.into_iter());

  assert_eq!(dist.get(&"a"), 50);

  dist.remove(&"a");

  assert_eq!(dist.get(&"a"), 0);
  assert_eq!(dist.sum_counts(), 125);
}

#[test]
fn smoke_test_frequency_sum_counts() {
  let words = vec!(("a", 7usize), ("b", 5usize), ("c", 8usize), ("d", 3usize));
  let mut dist: FrequencyDistribution<&str> = FromIterator::from_iter(words.into_iter());

  assert_eq!(dist.sum_counts(), 23);

  dist.insert("e");

  assert_eq!(dist.sum_counts(), 24);
}
<|endoftext|>


<commit_before>//! Muxed. A tmux project manager with no runtime dependencies.
extern crate getopts;
extern crate libc;
extern crate yaml_rust;
extern crate rand;

use std::env;

mod tmux;
mod command;
mod project;

use project::parser;
use project::processor;

/// The main execution method.
/// Verify all the arguments and options passed are valid for the application.
fn main() {
    let args: Vec<String> = env::args().collect();
    //let program = args[0].clone();
    let input = args[1].clone();

    let yaml     = project::main(&input);
    let commands = parser::main(&yaml, &input);
    processor::main(&commands)
}
<commit_msg>Begin writing first round of code docs.<commit_after>//! Muxed. A tmux project manager with no runtime dependencies.
extern crate getopts;
extern crate libc;
extern crate yaml_rust;
extern crate rand;

use std::env;

mod tmux;
mod command;
mod project;

use project::parser;
use project::processor;

/// The main execution method.
/// Currently accepts a single option. The option represents a configuration
/// file in the same naming format. Given a project file name `projectName.yml`
/// in the `~/.muxed/` directory.
///
/// # Examples
///
/// ~/.muxed/projectName.yml
///
/// ```
/// root: ~/projects/muxed/
/// windows:
///     - cargo: "cargo build"
///     - vim: "vim ."
///     - git: ""
/// ```
///
/// You can run the command:
///
/// ```
/// $ ./muxed projectName
/// ```
pub fn main() {
    let args: Vec<String> = env::args().collect();
    //let program = args[0].clone();
    let input = args[1].clone();

    let yaml     = project::main(&input);
    let commands = parser::main(&yaml, &input);
    processor::main(&commands)
}
<|endoftext|>


<commit_before>#![crate_name = "raft"]
#![crate_type="lib"]

#![feature(core)]
#![feature(io)]
#![feature(std_misc)]
#![feature(rand)]
extern crate "rustc-serialize" as rustc_serialize;
pub mod interchange;

use std::old_io::net::ip::SocketAddr;
use std::old_io::net::udp::UdpSocket;
use std::old_io::timer::Timer;
use std::time::Duration;
use std::thread::Thread;
use std::rand::{thread_rng, Rng};
use std::sync::mpsc::{channel, Sender, Receiver};
use std::str;
use std::collections::HashMap;
use rustc_serialize::{json, Encodable, Decodable};

// Enums and variants.
use interchange::{ClientRequest, RemoteProcedureCall, RemoteProcedureResponse};
// Data structures.
use interchange::{AppendEntries, RequestVote};
use interchange::{AppendRequest, IndexRange};
use NodeState::{Leader, Follower, Candidate};

// The maximum size of the read buffer.
const BUFFER_SIZE: usize = 4096;
const HEARTBEAT_MIN: i64 = 150;
const HEARTBEAT_MAX: i64 = 300;

/// The Raft Distributed Consensus Algorithm requires two RPC calls to be available:
///
///   * `append_entries` which is used as both a heartbeat (with no payload) and the primary
///     interface for requests.
///   * `request_vote` which is used by candidates during campaigns to obtain a vote.
///
pub trait Raft<T: Encodable + Decodable + Send + Clone> {
    /// Returns (term, success)
    fn append_entries(term: u64, leader_id: u64, prev_log_index: u64,
                      prev_log_term: u64, entries: Vec<T>,
                      leader_commit: u64) -> (u64, bool);

    /// Returns (term, voteGranted)
    fn request_vote(term: u64, candidate_id: u64, last_log_index: u64,
                    last_log_term: u64) -> (u64, bool);
}

/// A `RaftNode` acts as a replicated state machine. The server's role in the cluster depends on it's
/// own status. It will maintain both volatile state (which can be safely lost) and persistent
/// state (which must be carefully stored and kept safe).
pub struct RaftNode<T: Encodable + Decodable + Send + Clone> {
    // Raft related.
    state: NodeState,
    persistent_state: PersistentState<T>,
    volatile_state: VolatileState,
    // Auxilary Data.
    // TODO: This should probably be split off.
    // All nodes need to know this otherwise they can't effectively lead or hold elections.
    leader_id: Option<u64>,
    own_id: u64,
    nodes: HashMap<u64, SocketAddr>,
    heartbeat: Receiver<()>,
    socket: UdpSocket,
    req_recv: Receiver<ClientRequest<T>>,
    res_send: Sender<Result<Vec<T>, String>>,
}

/// The implementation of the RaftNode. In most use cases, creating a `RaftNode` should just be
/// done via `::new()`.
///
/// ```
/// use raft::RaftNode;
/// use std::old_io::net::ip::SocketAddr;
/// use std::old_io::net::ip::IpAddr::Ipv4Addr;
/// use std::collections::HashMap;
///
/// let mut nodes = HashMap::new();
/// nodes.insert(1, SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11111 });
/// nodes.insert(2, SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11112 });
/// // Create the nodes.
/// let node = RaftNode::<String>::start(1, nodes.clone());
/// ```
impl<T: Encodable + Decodable + Send + Clone> RaftNode<T> {
    /// Creates a new RaftNode with the neighbors specified. `id` should be a valid index into
    /// `nodes`. The idea is that you can use the same `nodes` on all of the clients and only vary
    /// `id`.
    pub fn start(id: u64, nodes: HashMap<u64, SocketAddr>) -> (Sender<ClientRequest<T>>, Receiver<Result<Vec<T>, String>>) {
        // TODO: Check index.
        // Setup the socket, make it not block.
        let own_socket_addr = nodes.get(&id)
            .unwrap().clone(); // TODO: Can we do better?
        let mut socket = UdpSocket::bind(own_socket_addr)
            .unwrap(); // TODO: Can we do better?
        socket.set_read_timeout(Some(0));
        // Communication channels.
        let (req_send, req_recv) = channel::<ClientRequest<T>>();
        let (res_send, res_recv) = channel::<Result<Vec<T>, String>>();
        // Fire up the thread.
        Thread::spawn(move || {
            // Start up a RNG and Timer
            let mut rng = thread_rng();
            let mut timer = Timer::new().unwrap();
            // We need a read buffer.
            let mut read_buffer = [0; BUFFER_SIZE];
            // Create the struct.
            let mut raft_node = RaftNode {
                state: Follower,
                persistent_state: PersistentState {
                    current_term: 0, // TODO: Double Check.
                    voted_for: 0,    // TODO: Better type?
                    log: Vec::<T>::new(),
                },
                volatile_state: VolatileState {
                    commit_index: 0,
                    last_applied: 0,
                },
                leader_id: None,
                own_id: id,
                nodes: nodes,
                // Blank timer for now.
                heartbeat: timer.oneshot(Duration::milliseconds(rng.gen_range::<i64>(HEARTBEAT_MIN, HEARTBEAT_MAX))), // If this fails we're in trouble.
                socket: socket,
                req_recv: req_recv,
                res_send: res_send,
            };
            // This is the main, strongly typed state machine. It loops indefinitely for now. It
            // would be nice if this was event based.
            loop {
                raft_node.tick();
            }
        });
        (req_send, res_recv)
    }
    /// This is the main tick for a leader node.
    fn tick(&mut self) {
        // We need a read buffer.
        let mut read_buffer = [0; BUFFER_SIZE];
        // If socket has data.
        match self.socket.recv_from(&mut read_buffer) {
            Ok((num_read, source)) => { // Something on the socket.
                // TODO: Verify this is a legitimate request, just check if it's
                //       in the cluster for now?
                // This is possibly an RPC from another node. Try to parse it out
                // and determine what to do based on it's variant.
                let data = str::from_utf8(&mut read_buffer[.. num_read])
                    .unwrap();
                if let Ok(rpc) = json::decode::<RemoteProcedureCall<T>>(data) {
                    match rpc {
                        RemoteProcedureCall::RequestVote(call) =>
                            self.handle_request_vote(call, source),
                        RemoteProcedureCall::AppendEntries(call) =>
                            self.handle_append_entries(call, source),
                    }
                } else if let Ok(rpr) = json::decode::<RemoteProcedureResponse>(data) {
                    match rpr {
                        RemoteProcedureResponse::Accepted { .. } =>
                            self.handle_accepted(rpr, source),
                        RemoteProcedureResponse::Rejected { .. } =>
                            self.handle_rejected(rpr, source),
                    }
                }
            },
            Err(_) => (),                 // Nothing on the socket.
        }
        // If channel has data.
        match self.req_recv.try_recv() {
            Ok(request) => {              // Something in channel.
                match request {
                    ClientRequest::IndexRange(request) =>
                        self.handle_index_range(request),
                    ClientRequest::AppendRequest(request) =>
                        self.handle_append_request(request),
                }
            },
            Err(_) => (),               // Nothing in channel.
        }
        // If timer has fired.
        match self.heartbeat.try_recv() {
            Ok(_) => {                  // Timer has fired.
                // A heartbeat has fired.
                self.handle_timer()
            },
            Err(_) => (),               // Timer hasn't fired.
        }
    }
    /// A lookup for index -> SocketAddr
    fn lookup(&self, index: u64) -> Option<&SocketAddr> {
        self.nodes.get(&index)
    }
    /// When a `Follower`'s heartbeat times out it's time to start a campaign for election and
    /// become a `Candidate`. If successful, the `RaftNode` will transistion state into a `Leader`,
    /// otherwise it will become `Follower` again.
    /// This function accepts a `Follower` and transforms it into a `Candidate` then attempts to
    /// issue `RequestVote` remote procedure calls to other known nodes. If a majority come back
    /// accepted, it will become the leader.
    fn campaign(&mut self) {
        self.state = match self.state {
            Follower => Candidate,
            _ => panic!("Should not campaign while not a follower!")
        };
        // TODO: Issue `RequestVote` to known nodes.
        unimplemented!()
            // We rely on the loop to handle incoming responses regarding `RequestVote`, don't worry
            // about that here.
    }
    /// Handles a `RemoteProcedureCall::RequestVote` call.
    ///
    ///   * Reply false if term < currentTerm.
    ///   * If votedFor is null or candidateId, and candidate’s log is at least as up-to-date as
    ///     receiver’s log, grant vote.
    fn handle_request_vote(&mut self, call: RequestVote, source: SocketAddr) {
        match self.state {
            Leader(ref state) => {
                // Re-assert leadership.
                let rpr = RemoteProcedureResponse::Rejected {
                    term: self.persistent_state.current_term,
                    current_leader: self.leader_id.unwrap(), // Should be self.
                };
                let encoded = json::encode::<RemoteProcedureResponse>(&rpr)
                    .unwrap();
                self.socket.send_to(encoded.as_bytes(), source)
                    .unwrap(); // TODO: Can we do better?
            },
            Follower => {
                // Do checks and respond appropriately.
                unimplemented!();
            },
            Candidate => {
                // TODO ???
                unimplemented!();
            }
        }
    }
    /// Handles an `AppendEntries` request from a caller.
    ///
    fn handle_append_entries(&mut self, call: AppendEntries<T>, source: SocketAddr) {
        match self.state {
            Leader(ref state) => {
                unimplemented!();
            },
            Follower => {
                unimplemented!();
            },
            Candidate => {
                unimplemented!();
            },
        }
    }
    fn handle_append_request(&mut self, request: AppendRequest<T>) {
        match self.state {
            Leader(ref state) => {
                unimplemented!();
            },
            Follower => {
                unimplemented!();
            },
            Candidate => {
                unimplemented!();
            },
        }
    }
    fn handle_index_range(&mut self, request: IndexRange) {
        unimplemented!();
    }
    fn handle_accepted(&mut self, response: RemoteProcedureResponse, source: SocketAddr) {
        unimplemented!();
    }
    fn handle_rejected(&mut self, response: RemoteProcedureResponse, source: SocketAddr) {
        unimplemented!();
    }
    fn handle_timer(&mut self) {
        match self.state {
            Leader(ref state) => {
                // Send heartbeats.
                unimplemented!();
            },
            Follower => self.campaign(),
            Candidate => panic!("Candidate should not have their heartbeat fire."),
        }
    }

}

/// The RPC calls required by the Raft protocol.
impl<T: Encodable + Decodable + Send + Clone> Raft<T> for RaftNode<T> {
    /// Returns (term, success)
    fn append_entries(term: u64, leader_id: u64, prev_log_index: u64,
                      prev_log_term: u64, entries: Vec<T>,
                      leader_commit: u64) -> (u64, bool) {
        (0, false) // TODO: Implement
    }
    /// Returns (term, voteGranted)
    fn request_vote(term: u64, candidate_id: u64, last_log_index: u64,
                    last_log_term: u64) -> (u64, bool) {
        (0, false) // TODO: Implement
    }
}

/// Nodes can either be:
///
///   * A `Follower`, which replicates AppendEntries requests and votes for it's leader.
///   * A `Leader`, which leads the cluster by serving incoming requests, ensuring data is
///     replicated, and issuing heartbeats..
///   * A `Candidate`, which campaigns in an election and may become a `Leader` (if it gets enough
///     votes) or a `Follower`, if it hears from a `Leader`.
#[derive(PartialEq, Eq)]
pub enum NodeState {
    Follower,
    Leader(LeaderState),
    Candidate,
}

/// Persistent state
/// **Must be updated to stable storage before RPC response.**
pub struct PersistentState<T: Encodable + Decodable + Send + Clone> {
    current_term: u64,
    voted_for: u64, // Better way? Can we use a IpAddr?
    log: Vec<T>,
}

/// Volatile state
#[derive(Copy)]
pub struct VolatileState {
    commit_index: u64,
    last_applied: u64
}

/// Leader Only
/// **Reinitialized after election.**
#[derive(PartialEq, Eq)]
pub struct LeaderState {
    next_index: Vec<u64>,
    match_index: Vec<u64>
}
<commit_msg>More work on event loop<commit_after>#![crate_name = "raft"]
#![crate_type="lib"]

#![feature(core)]
#![feature(io)]
#![feature(std_misc)]
#![feature(rand)]
extern crate "rustc-serialize" as rustc_serialize;
pub mod interchange;

use std::old_io::net::ip::SocketAddr;
use std::old_io::net::udp::UdpSocket;
use std::old_io::timer::Timer;
use std::time::Duration;
use std::thread::Thread;
use std::rand::{thread_rng, Rng, ThreadRng};
use std::sync::mpsc::{channel, Sender, Receiver};
use std::str;
use std::collections::HashMap;
use rustc_serialize::{json, Encodable, Decodable};

// Enums and variants.
use interchange::{ClientRequest, RemoteProcedureCall, RemoteProcedureResponse};
// Data structures.
use interchange::{AppendEntries, RequestVote};
use interchange::{AppendRequest, IndexRange};
use NodeState::{Leader, Follower, Candidate};

// The maximum size of the read buffer.
const BUFFER_SIZE: usize = 4096;
const HEARTBEAT_MIN: i64 = 150;
const HEARTBEAT_MAX: i64 = 300;

/// The Raft Distributed Consensus Algorithm requires two RPC calls to be available:
///
///   * `append_entries` which is used as both a heartbeat (with no payload) and the primary
///     interface for requests.
///   * `request_vote` which is used by candidates during campaigns to obtain a vote.
///
pub trait Raft<T: Encodable + Decodable + Send + Clone> {
    /// Returns (term, success)
    fn append_entries(term: u64, leader_id: u64, prev_log_index: u64,
                      prev_log_term: u64, entries: Vec<T>,
                      leader_commit: u64) -> (u64, bool);

    /// Returns (term, voteGranted)
    fn request_vote(term: u64, candidate_id: u64, last_log_index: u64,
                    last_log_term: u64) -> (u64, bool);
}

/// A `RaftNode` acts as a replicated state machine. The server's role in the cluster depends on it's
/// own status. It will maintain both volatile state (which can be safely lost) and persistent
/// state (which must be carefully stored and kept safe).
pub struct RaftNode<T: Encodable + Decodable + Send + Clone> {
    // Raft related.
    state: NodeState,
    persistent_state: PersistentState<T>,
    volatile_state: VolatileState,
    // Auxilary Data.
    // TODO: This should probably be split off.
    // All nodes need to know this otherwise they can't effectively lead or hold elections.
    leader_id: Option<u64>,
    own_id: u64,
    nodes: HashMap<u64, SocketAddr>,
    heartbeat: Receiver<()>,
    rng: ThreadRng,
    timer: Timer,
    socket: UdpSocket,
    req_recv: Receiver<ClientRequest<T>>,
    res_send: Sender<Result<Vec<T>, String>>,
}

/// The implementation of the RaftNode. In most use cases, creating a `RaftNode` should just be
/// done via `::new()`.
///
/// ```
/// use raft::RaftNode;
/// use std::old_io::net::ip::SocketAddr;
/// use std::old_io::net::ip::IpAddr::Ipv4Addr;
/// use std::collections::HashMap;
///
/// let mut nodes = HashMap::new();
/// nodes.insert(1, SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11111 });
/// nodes.insert(2, SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: 11112 });
/// // Create the nodes.
/// let node = RaftNode::<String>::start(1, nodes.clone());
/// ```
impl<T: Encodable + Decodable + Send + Clone> RaftNode<T> {
    /// Creates a new RaftNode with the neighbors specified. `id` should be a valid index into
    /// `nodes`. The idea is that you can use the same `nodes` on all of the clients and only vary
    /// `id`.
    pub fn start(id: u64, nodes: HashMap<u64, SocketAddr>) -> (Sender<ClientRequest<T>>, Receiver<Result<Vec<T>, String>>) {
        // TODO: Check index.
        // Setup the socket, make it not block.
        let own_socket_addr = nodes.get(&id)
            .unwrap().clone(); // TODO: Can we do better?
        let mut socket = UdpSocket::bind(own_socket_addr)
            .unwrap(); // TODO: Can we do better?
        socket.set_read_timeout(Some(0));
        // Communication channels.
        let (req_send, req_recv) = channel::<ClientRequest<T>>();
        let (res_send, res_recv) = channel::<Result<Vec<T>, String>>();
        // Fire up the thread.
        Thread::spawn(move || {
            // Start up a RNG and Timer
            let mut rng = thread_rng();
            let mut timer = Timer::new().unwrap();
            // We need a read buffer.
            let mut read_buffer = [0; BUFFER_SIZE];
            // Create the struct.
            let mut raft_node = RaftNode {
                state: Follower,
                persistent_state: PersistentState {
                    current_term: 0, // TODO: Double Check.
                    voted_for: None,
                    log: Vec::<(u64, T)>::new(),
                },
                volatile_state: VolatileState {
                    commit_index: 0,
                    last_applied: 0,
                },
                leader_id: None,
                own_id: id,
                nodes: nodes,
                // Blank timer for now.
                heartbeat: timer.oneshot(Duration::milliseconds(rng.gen_range::<i64>(HEARTBEAT_MIN, HEARTBEAT_MAX))), // If this fails we're in trouble.
                timer: timer,
                rng: rng,
                socket: socket,
                req_recv: req_recv,
                res_send: res_send,
            };
            // This is the main, strongly typed state machine. It loops indefinitely for now. It
            // would be nice if this was event based.
            loop {
                raft_node.tick();
            }
        });
        (req_send, res_recv)
    }
    /// This is the main tick for a leader node.
    fn tick(&mut self) {
        // We need a read buffer.
        let mut read_buffer = [0; BUFFER_SIZE];
        // If socket has data.
        match self.socket.recv_from(&mut read_buffer) {
            Ok((num_read, source)) => { // Something on the socket.
                // TODO: Verify this is a legitimate request, just check if it's
                //       in the cluster for now?
                // This is possibly an RPC from another node. Try to parse it out
                // and determine what to do based on it's variant.
                let data = str::from_utf8(&mut read_buffer[.. num_read])
                    .unwrap();
                if let Ok(rpc) = json::decode::<RemoteProcedureCall<T>>(data) {
                    match rpc {
                        RemoteProcedureCall::RequestVote(call) =>
                            self.handle_request_vote(call, source),
                        RemoteProcedureCall::AppendEntries(call) =>
                            self.handle_append_entries(call, source),
                    }
                } else if let Ok(rpr) = json::decode::<RemoteProcedureResponse>(data) {
                    match rpr {
                        RemoteProcedureResponse::Accepted { .. } =>
                            self.handle_accepted(rpr, source),
                        RemoteProcedureResponse::Rejected { .. } =>
                            self.handle_rejected(rpr, source),
                    }
                }
            },
            Err(_) => (),                 // Nothing on the socket.
        }
        // If channel has data.
        match self.req_recv.try_recv() {
            Ok(request) => {              // Something in channel.
                match request {
                    ClientRequest::IndexRange(request) =>
                        self.handle_index_range(request),
                    ClientRequest::AppendRequest(request) =>
                        self.handle_append_request(request),
                }
            },
            Err(_) => (),               // Nothing in channel.
        }
        // If timer has fired.
        match self.heartbeat.try_recv() {
            Ok(_) => {                  // Timer has fired.
                // A heartbeat has fired.
                self.handle_timer()
            },
            Err(_) => (),               // Timer hasn't fired.
        }
    }
    /// A lookup for index -> SocketAddr
    fn lookup(&self, index: u64) -> Option<&SocketAddr> {
        self.nodes.get(&index)
    }
    fn other_nodes(&self) -> Vec<SocketAddr> {
        self.nodes.iter().filter_map(|(&k, &v)| {
            if k == self.own_id {
                None
            } else {
                Some(v)
            }
        }).collect::<Vec<SocketAddr>>()
    }
    /// When a `Follower`'s heartbeat times out it's time to start a campaign for election and
    /// become a `Candidate`. If successful, the `RaftNode` will transistion state into a `Leader`,
    /// otherwise it will become `Follower` again.
    /// This function accepts a `Follower` and transforms it into a `Candidate` then attempts to
    /// issue `RequestVote` remote procedure calls to other known nodes. If a majority come back
    /// accepted, it will become the leader.
    fn campaign(&mut self) {
        // On conversion to Candidate:
        // * Increment current_term
        // * Vote for self
        // * Reset election timer
        // * Send RequestVote RPC to all other nodes.
        self.state = match self.state {
            Follower | Candidate => Candidate,
            _ => panic!("Should not campaign as a Leader!")
        };
        self.persistent_state.current_term += 1;
        self.persistent_state.voted_for = Some(self.own_id); // TODO: Is this correct?
        self.reset_timer();
        // Issue `RequestVote` to known nodes.
        let request = RemoteProcedureCall::request_vote(
            self.persistent_state.current_term,
            self.persistent_state.voted_for.unwrap(), // TODO: Safe because we just set it. But correct
            self.volatile_state.last_applied,
            0); // TODO: Get this.
        self.broadcast(request);
            // We rely on the loop to handle incoming responses regarding `RequestVote`, don't worry
            // about that here.
    }
    /// Handles a `RemoteProcedureCall::RequestVote` call.
    ///
    ///   * Reply false if term < currentTerm.
    ///   * If votedFor is null or candidateId, and candidate’s log is at least as up-to-date as
    ///     receiver’s log, grant vote.
    fn handle_request_vote(&mut self, call: RequestVote, source: SocketAddr) {
        match self.state {
            Leader(_) => {
                // Re-assert leadership.
                let rpr = RemoteProcedureResponse::Rejected {
                    term: self.persistent_state.current_term,
                    current_leader: self.leader_id.unwrap(), // Should be self.
                };
                let encoded = json::encode::<RemoteProcedureResponse>(&rpr)
                    .unwrap();
                self.send_to(encoded.as_bytes(), source)
                    .unwrap(); // TODO: Can we do better?
            },
            Follower => {
                let term_ok = self.persistent_state.current_term < call.term;
                let vote_ok = self.persistent_state.voted_for.is_none();
                let log_idx_ok = self.volatile_state.last_applied < call.last_log_index;
                let log_term_ok = true; // TODO.
                // If voted_for is None or the requester's ID, and their log is up to date, accept.
                unimplemented!();
            },
            Candidate => {
                // TODO ???
                unimplemented!();
            }
        }
    }
    /// Handles an `AppendEntries` request from a caller.
    ///
    fn handle_append_entries(&mut self, call: AppendEntries<T>, source: SocketAddr) {
        match self.state {
            Leader(ref state) => {
                unimplemented!();
            },
            Follower => {
                unimplemented!();
            },
            Candidate => {
                unimplemented!();
            },
        }
    }
    fn handle_append_request(&mut self, request: AppendRequest<T>) {
        match self.state {
            Leader(ref state) => {
                unimplemented!();
            },
            Follower => {
                unimplemented!();
            },
            Candidate => {
                unimplemented!();
            },
        }
    }
    fn handle_index_range(&mut self, request: IndexRange) {
        unimplemented!();
    }
    fn handle_accepted(&mut self, response: RemoteProcedureResponse, source: SocketAddr) {
        unimplemented!();
    }
    fn handle_rejected(&mut self, response: RemoteProcedureResponse, source: SocketAddr) {
        unimplemented!();
    }
    fn handle_timer(&mut self) {
        match self.state {
            Leader(ref state) => {
                // Send heartbeats.
                unimplemented!();
            },
            Follower => self.campaign(),
            Candidate => self.campaign(),
        }
    }
    fn reset_timer(&mut self) {
        self.heartbeat = self.timer.oneshot(Duration::milliseconds(self.rng.gen_range::<i64>(HEARTBEAT_MIN, HEARTBEAT_MAX)));
    }
    fn broadcast(&mut self, rpc: RemoteProcedureCall<T>) -> Result<(), std::old_io::IoError> {
        let encoded = json::encode::<RemoteProcedureCall<T>>(&rpc)
            .unwrap();
        for node in self.other_nodes() {
            try!(self.send_to(encoded.as_bytes(), node));
        }
        Ok(())
    }
    // TODO: Improve "message" to not be &[u8]
    fn send_to(&mut self, message: &[u8], node: SocketAddr) -> Result<(), std::old_io::IoError> {
        self.socket.send_to(message, node)
    }
}

/// The RPC calls required by the Raft protocol.
impl<T: Encodable + Decodable + Send + Clone> Raft<T> for RaftNode<T> {
    /// Returns (term, success)
    fn append_entries(term: u64, leader_id: u64, prev_log_index: u64,
                      prev_log_term: u64, entries: Vec<T>,
                      leader_commit: u64) -> (u64, bool) {
        (0, false) // TODO: Implement
    }
    /// Returns (term, voteGranted)
    fn request_vote(term: u64, candidate_id: u64, last_log_index: u64,
                    last_log_term: u64) -> (u64, bool) {
        (0, false) // TODO: Implement
    }
}

/// Nodes can either be:
///
///   * A `Follower`, which replicates AppendEntries requests and votes for it's leader.
///   * A `Leader`, which leads the cluster by serving incoming requests, ensuring data is
///     replicated, and issuing heartbeats..
///   * A `Candidate`, which campaigns in an election and may become a `Leader` (if it gets enough
///     votes) or a `Follower`, if it hears from a `Leader`.
#[derive(PartialEq, Eq)]
pub enum NodeState {
    Follower,
    Leader(LeaderState),
    Candidate,
}

/// Persistent state
/// **Must be updated to stable storage before RPC response.**
pub struct PersistentState<T: Encodable + Decodable + Send + Clone> {
    current_term: u64,
    voted_for: Option<u64>, // request_vote cares if this is `None`
    log: Vec<(u64, T)>,
}

/// Volatile state
#[derive(Copy)]
pub struct VolatileState {
    commit_index: u64,
    last_applied: u64
}

/// Leader Only
/// **Reinitialized after election.**
#[derive(PartialEq, Eq)]
pub struct LeaderState {
    next_index: Vec<u64>,
    match_index: Vec<u64>
}
<|endoftext|>


<commit_before>//! This is a very early work-in-progress binding to various Linux Kernel APIs.
//!
//! It is not yet ready for use in your projects. Once version 0.1 or higher
//! is released, you are welcome to start using it :)

#![cfg(target_os="linux")]

pub use std::os::*;
pub use std::os::raw::*;
<commit_msg>Add a unit test to ensure that ctypes are imported<commit_after>//! This is a very early work-in-progress binding to various Linux Kernel APIs.
//!
//! It is not yet ready for use in your projects. Once version 0.1 or higher
//! is released, you are welcome to start using it :)

#![cfg(target_os="linux")]

pub use std::os::*;
pub use std::os::raw::*;

#[cfg(test)]
mod tests {

    #[allow(unused_variables)]
    #[test]
    fn ensure_types_exist() {
        let schar: ::c_schar = -3;
        let uchar: ::c_uchar = 2;
        let achar: ::c_char = 62;

        let ashort: ::c_short = -5162;
        let ushort: ::c_ushort = 65000;

        let aint: ::c_int = 26327;
        let uint: ::c_uint = 20000026;

        let long: ::c_long = 75473765327;
        let ulong: ::c_ulong = 26294762868676748;

        let float: ::c_float = 2462347.426f32;
        let double: ::c_double = 2694237684327.4326237637f64;

        assert!(true);//Well, we haven't crashed! I guess it worked
    }

}
<|endoftext|>