<commit_before>mod crc;
#[macro_use]
mod macros;
mod types;

pub use self::crc::{crc8, crc16};
pub use self::types::{ErrorKind, ByteStream, ReadStream};

use nom::IResult;
use metadata::{Metadata, metadata_parser};

/// An interface for parsing through some type of producer to a byte stream.
///
/// External parsers get passed in and consumes the bytes held internally
/// and outputs the `Result` of that parser.
pub trait StreamProducer {
  fn parse<F, T>(&mut self, f: F) -> Result<T, ErrorKind>
   where F: FnOnce(&[u8]) -> IResult<&[u8], T>;
}

// Convert one to four byte slices into an unsigned 32-bit number.
//
// NOTE: This assumes big-endian since most numbers in the FLAC binary are
// that endianness.
#[inline]
pub fn to_u32(bytes: &[u8]) -> u32 {
  let length = bytes.len();

  debug_assert!(length <= 4);

  (0..length).fold(0, |result, i|
    result + ((bytes[i] as u32) << ((length - 1 - i) * 8))
  )
}

// Extends a signed value of a specific bit size to a larger bit size.
//
// NOTE: This assumes that the larger bit size will be 32 bit since that is
// the largest sample size supported in FLAC.
pub fn extend_sign(value: u32, bit_count: usize) -> i32 {
  if bit_count >= 32 || value < (1 << (bit_count - 1)) {
    value as i32
  } else {
    (value as i32).wrapping_sub(1 << bit_count)
  }
}

#[derive(PartialEq, Eq)]
enum ParserState {
  Header,
  StreamInfo,
  Metadata
}

fn parser<'a>(input: &'a [u8], state: &mut ParserState)
              -> IResult<&'a [u8], Metadata> {
  let mut slice = input;
  let error     = nom::Err::Code(nom::ErrorKind::Custom(10));

  if *state == ParserState::Header {
    let (i, _) = try_parse!(slice, tag!("fLaC"));

    slice  = i;
    *state = ParserState::StreamInfo;
  }

  match *state {
    ParserState::StreamInfo => {
      let (i, block) = try_parse!(slice, metadata_parser);

      if let metadata::Data::StreamInfo(_) = block.data {
        *state = ParserState::Metadata;

        IResult::Done(i, block)
      } else {
        IResult::Error(error)
      }
    }
    ParserState::Metadata   => metadata_parser(slice),
    _                       => IResult::Error(error),
  }
}

pub fn many_metadata<S, F>(stream: &mut S, mut f: F) -> bool
 where S: StreamProducer,
       F: FnMut(Metadata) {
  let mut state    = ParserState::Header;
  let mut is_error = false;

  loop {
    match stream.parse(|i| parser(i, &mut state)) {
      Ok(block)                => {
        let is_last = block.is_last;

        f(block);

        if is_last {
          break;
        }
      }
      Err(ErrorKind::Continue) => continue,
      Err(_)                   => {
        is_error = true;

        break;
      }
    }
  }

  is_error
}

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

  #[test]
  #[should_panic]
  fn test_panic_to_u32() {
    to_u32(&[0x00, 0x00, 0x00, 0x00, 0x00]);
  }

  #[test]
  fn test_to_u32() {
    let bytes = [0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef];

    assert_eq!(to_u32(&bytes[0..1]), 0x00000001);
    assert_eq!(to_u32(&bytes[3..5]), 0x00006789);
    assert_eq!(to_u32(&bytes[1..4]), 0x00234567);
    assert_eq!(to_u32(&bytes[4..]), 0x89abcdef);
  }

  #[test]
  fn test_extend_sign() {
    assert_eq!(extend_sign(32, 6), -32);
    assert_eq!(extend_sign(31, 6), 31);
    assert_eq!(extend_sign(128, 8), -128);
    assert_eq!(extend_sign(127, 8), 127);

    assert_eq!(extend_sign(2048, 12), -2048);
    assert_eq!(extend_sign(2047, 12), 2047);
    assert_eq!(extend_sign(32768, 16), -32768);
    assert_eq!(extend_sign(32767, 16), 32767);

    assert_eq!(extend_sign(8388608, 24), -8388608);
    assert_eq!(extend_sign(8388607, 24), 8388607);
    assert_eq!(extend_sign(2147483648, 32), -2147483648);
    assert_eq!(extend_sign(2147483647, 32), 2147483647);
  }
}
<commit_msg>Add missing imports<commit_after>mod crc;
#[macro_use]
mod macros;
mod types;

pub use self::crc::{crc8, crc16};
pub use self::types::{ErrorKind, ByteStream, ReadStream};

use nom::{self, IResult};
use metadata::{self, Metadata, metadata_parser};

/// An interface for parsing through some type of producer to a byte stream.
///
/// External parsers get passed in and consumes the bytes held internally
/// and outputs the `Result` of that parser.
pub trait StreamProducer {
  fn parse<F, T>(&mut self, f: F) -> Result<T, ErrorKind>
   where F: FnOnce(&[u8]) -> IResult<&[u8], T>;
}

// Convert one to four byte slices into an unsigned 32-bit number.
//
// NOTE: This assumes big-endian since most numbers in the FLAC binary are
// that endianness.
#[inline]
pub fn to_u32(bytes: &[u8]) -> u32 {
  let length = bytes.len();

  debug_assert!(length <= 4);

  (0..length).fold(0, |result, i|
    result + ((bytes[i] as u32) << ((length - 1 - i) * 8))
  )
}

// Extends a signed value of a specific bit size to a larger bit size.
//
// NOTE: This assumes that the larger bit size will be 32 bit since that is
// the largest sample size supported in FLAC.
pub fn extend_sign(value: u32, bit_count: usize) -> i32 {
  if bit_count >= 32 || value < (1 << (bit_count - 1)) {
    value as i32
  } else {
    (value as i32).wrapping_sub(1 << bit_count)
  }
}

#[derive(PartialEq, Eq)]
enum ParserState {
  Header,
  StreamInfo,
  Metadata
}

fn parser<'a>(input: &'a [u8], state: &mut ParserState)
              -> IResult<&'a [u8], Metadata> {
  let mut slice = input;
  let error     = nom::Err::Code(nom::ErrorKind::Custom(10));

  if *state == ParserState::Header {
    let (i, _) = try_parse!(slice, tag!("fLaC"));

    slice  = i;
    *state = ParserState::StreamInfo;
  }

  match *state {
    ParserState::StreamInfo => {
      let (i, block) = try_parse!(slice, metadata_parser);

      if let metadata::Data::StreamInfo(_) = block.data {
        *state = ParserState::Metadata;

        IResult::Done(i, block)
      } else {
        IResult::Error(error)
      }
    }
    ParserState::Metadata   => metadata_parser(slice),
    _                       => IResult::Error(error),
  }
}

pub fn many_metadata<S, F>(stream: &mut S, mut f: F) -> bool
 where S: StreamProducer,
       F: FnMut(Metadata) {
  let mut state    = ParserState::Header;
  let mut is_error = false;

  loop {
    match stream.parse(|i| parser(i, &mut state)) {
      Ok(block)                => {
        let is_last = block.is_last;

        f(block);

        if is_last {
          break;
        }
      }
      Err(ErrorKind::Continue) => continue,
      Err(_)                   => {
        is_error = true;

        break;
      }
    }
  }

  is_error
}

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

  #[test]
  #[should_panic]
  fn test_panic_to_u32() {
    to_u32(&[0x00, 0x00, 0x00, 0x00, 0x00]);
  }

  #[test]
  fn test_to_u32() {
    let bytes = [0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef];

    assert_eq!(to_u32(&bytes[0..1]), 0x00000001);
    assert_eq!(to_u32(&bytes[3..5]), 0x00006789);
    assert_eq!(to_u32(&bytes[1..4]), 0x00234567);
    assert_eq!(to_u32(&bytes[4..]), 0x89abcdef);
  }

  #[test]
  fn test_extend_sign() {
    assert_eq!(extend_sign(32, 6), -32);
    assert_eq!(extend_sign(31, 6), 31);
    assert_eq!(extend_sign(128, 8), -128);
    assert_eq!(extend_sign(127, 8), 127);

    assert_eq!(extend_sign(2048, 12), -2048);
    assert_eq!(extend_sign(2047, 12), 2047);
    assert_eq!(extend_sign(32768, 16), -32768);
    assert_eq!(extend_sign(32767, 16), 32767);

    assert_eq!(extend_sign(8388608, 24), -8388608);
    assert_eq!(extend_sign(8388607, 24), 8388607);
    assert_eq!(extend_sign(2147483648, 32), -2147483648);
    assert_eq!(extend_sign(2147483647, 32), 2147483647);
  }
}
<|endoftext|>


<commit_before>// Copyright 2015-2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

// Inspired by Paul Woolcock's cargo-fmt (https://github.com/pwoolcoc/cargo-fmt/)

#![cfg(not(test))]
#![deny(warnings)]

extern crate getopts;
extern crate rustc_serialize;

use std::env;
use std::io::Write;
use std::path::PathBuf;
use std::process::{Command, ExitStatus};
use std::str;

use getopts::Options;
use rustc_serialize::json::Json;

fn main() {
    let exit_status = execute();
    std::io::stdout().flush().unwrap();
    std::process::exit(exit_status);
}

fn execute() -> i32 {
    let success = 0;
    let failure = 1;

    let mut opts = getopts::Options::new();
    opts.optflag("h", "help", "show this message");
    opts.optflag("q", "quiet", "no output printed to stdout");
    opts.optflag("v", "verbose", "use verbose output");

    let matches = match opts.parse(env::args().skip(1).take_while(|a| a != "--")) {
        Ok(m) => m,
        Err(e) => {
            print_usage(&opts, &e.to_string());
            return failure;
        }
    };

    let verbosity = match (matches.opt_present("v"), matches.opt_present("q")) {
        (false, false) => Verbosity::Normal,
        (false, true) => Verbosity::Quiet,
        (true, false) => Verbosity::Verbose,
        (true, true) => {
            print_usage(&opts, "quiet mode and verbose mode are not compatible");
            return failure;
        }
    };

    if matches.opt_present("h") {
        print_usage(&opts, "");
        return success;
    }

    match format_crate(verbosity) {
        Err(e) => {
            print_usage(&opts, &e.to_string());
            failure
        }
        Ok(status) => {
            if status.success() {
                success
            } else {
                status.code().unwrap_or(failure)
            }
        }
    }
}

fn print_usage(opts: &Options, reason: &str) {
    let msg = format!("{}\nusage: cargo fmt [options]", reason);
    println!("{}\nThis utility formats all bin and lib files of the current crate using rustfmt. \
              Arguments after `--` are passed to rustfmt.",
             opts.usage(&msg));
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Verbosity {
    Verbose,
    Normal,
    Quiet,
}

fn format_crate(verbosity: Verbosity) -> Result<ExitStatus, std::io::Error> {
    let targets = try!(get_targets());

    // Currently only bin and lib files get formatted
    let files: Vec<_> = targets
        .into_iter()
        .filter(|t| t.kind.should_format())
        .inspect(|t| if verbosity == Verbosity::Verbose {
                     println!("[{:?}] {:?}", t.kind, t.path)
                 })
        .map(|t| t.path)
        .collect();

    format_files(&files, &get_fmt_args(), verbosity)
}

fn get_fmt_args() -> Vec<String> {
    // All arguments after -- are passed to rustfmt
    env::args().skip_while(|a| a != "--").skip(1).collect()
}

#[derive(Debug)]
enum TargetKind {
    Lib, // dylib, staticlib, lib
    Bin, // bin
    Example, // example file
    Test, // test file
    Bench, // bench file
    CustomBuild, // build script
    Other, // plugin,...
}

impl TargetKind {
    fn should_format(&self) -> bool {
        match *self {
            TargetKind::Lib | TargetKind::Bin | TargetKind::Example | TargetKind::Test |
            TargetKind::Bench | TargetKind::CustomBuild => true,
            _ => false,
        }
    }
}

#[derive(Debug)]
pub struct Target {
    path: PathBuf,
    kind: TargetKind,
}

// Returns a vector of all compile targets of a crate
fn get_targets() -> Result<Vec<Target>, std::io::Error> {
    let mut targets: Vec<Target> = vec![];
    let output = try!(Command::new("cargo").arg("read-manifest").output());
    if output.status.success() {
        // None of the unwraps should fail if output of `cargo read-manifest` is correct
        let data = &String::from_utf8(output.stdout).unwrap();
        let json = Json::from_str(data).unwrap();
        let jtargets = json.find("targets").unwrap().as_array().unwrap();
        for jtarget in jtargets {
            targets.push(target_from_json(jtarget));
        }

        Ok(targets)
    } else {
        // This happens when cargo-fmt is not used inside a crate
        Err(std::io::Error::new(std::io::ErrorKind::NotFound,
                                str::from_utf8(&output.stderr).unwrap()))
    }
}

fn target_from_json(jtarget: &Json) -> Target {
    let jtarget = jtarget.as_object().unwrap();
    let path = PathBuf::from(jtarget.get("src_path").unwrap().as_string().unwrap());
    let kinds = jtarget.get("kind").unwrap().as_array().unwrap();
    let kind = match kinds[0].as_string().unwrap() {
        "bin" => TargetKind::Bin,
        "lib" | "dylib" | "staticlib" | "cdylib" | "rlib" => TargetKind::Lib,
        "test" => TargetKind::Test,
        "example" => TargetKind::Example,
        "bench" => TargetKind::Bench,
        "custom-build" => TargetKind::CustomBuild,
        _ => TargetKind::Other,
    };

    Target {
        path: path,
        kind: kind,
    }
}

fn format_files(files: &[PathBuf],
                fmt_args: &[String],
                verbosity: Verbosity)
                -> Result<ExitStatus, std::io::Error> {
    let stdout = if verbosity == Verbosity::Quiet {
        std::process::Stdio::null()
    } else {
        std::process::Stdio::inherit()
    };
    if verbosity == Verbosity::Verbose {
        print!("rustfmt");
        for a in fmt_args.iter() {
            print!(" {}", a);
        }
        for f in files.iter() {
            print!(" {}", f.display());
        }
        println!("");
    }
    let mut command = try!(Command::new("rustfmt")
        .stdout(stdout)
        .args(files)
        .args(fmt_args)
        .spawn()
        .map_err(|e| match e.kind() {
            std::io::ErrorKind::NotFound => {
                std::io::Error::new(std::io::ErrorKind::Other,
                                    "Could not run rustfmt, please make sure it is in your PATH.")
            }
            _ => e,
        }));
    command.wait()
}
<commit_msg>Implemented support for workspaces<commit_after>// Copyright 2015-2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

// Inspired by Paul Woolcock's cargo-fmt (https://github.com/pwoolcoc/cargo-fmt/)

#![cfg(not(test))]
#![deny(warnings)]

extern crate getopts;
extern crate rustc_serialize;

use std::env;
use std::io::Write;
use std::path::PathBuf;
use std::process::{Command, ExitStatus};
use std::str;
use std::collections::HashSet;
use std::iter::FromIterator;

use getopts::{Options, Matches};
use rustc_serialize::json::Json;

fn main() {
    let exit_status = execute();
    std::io::stdout().flush().unwrap();
    std::process::exit(exit_status);
}

fn execute() -> i32 {
    let success = 0;
    let failure = 1;

    let mut opts = getopts::Options::new();
    opts.optflag("h", "help", "show this message");
    opts.optflag("q", "quiet", "no output printed to stdout");
    opts.optflag("v", "verbose", "use verbose output");
    opts.optmulti("p",
                  "package",
                  "specify package to format (only usable in workspaces)",
                  "<package>");
    opts.optflag("", "all", "format all packages (only usable in workspaces)");

    let matches = match opts.parse(env::args().skip(1).take_while(|a| a != "--")) {
        Ok(m) => m,
        Err(e) => {
            print_usage(&opts, &e.to_string());
            return failure;
        }
    };

    let verbosity = match (matches.opt_present("v"), matches.opt_present("q")) {
        (false, false) => Verbosity::Normal,
        (false, true) => Verbosity::Quiet,
        (true, false) => Verbosity::Verbose,
        (true, true) => {
            print_usage(&opts, "quiet mode and verbose mode are not compatible");
            return failure;
        }
    };

    if matches.opt_present("h") {
        print_usage(&opts, "");
        return success;
    }

    let workspace_hitlist = WorkspaceHitlist::from_matches(&matches);

    match format_crate(verbosity, workspace_hitlist) {
        Err(e) => {
            print_usage(&opts, &e.to_string());
            failure
        }
        Ok(status) => {
            if status.success() {
                success
            } else {
                status.code().unwrap_or(failure)
            }
        }
    }
}

fn print_usage(opts: &Options, reason: &str) {
    let msg = format!("{}\nusage: cargo fmt [options]", reason);
    println!("{}\nThis utility formats all bin and lib files of the current crate using rustfmt. \
              Arguments after `--` are passed to rustfmt.",
             opts.usage(&msg));
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Verbosity {
    Verbose,
    Normal,
    Quiet,
}

fn format_crate(verbosity: Verbosity,
                workspace_hitlist: WorkspaceHitlist)
                -> Result<ExitStatus, std::io::Error> {
    let targets = try!(get_targets(workspace_hitlist));

    // Currently only bin and lib files get formatted
    let files: Vec<_> = targets
        .into_iter()
        .filter(|t| t.kind.should_format())
        .inspect(|t| if verbosity == Verbosity::Verbose {
                     println!("[{:?}] {:?}", t.kind, t.path)
                 })
        .map(|t| t.path)
        .collect();

    format_files(&files, &get_fmt_args(), verbosity)
}

fn get_fmt_args() -> Vec<String> {
    // All arguments after -- are passed to rustfmt
    env::args().skip_while(|a| a != "--").skip(1).collect()
}

#[derive(Debug)]
enum TargetKind {
    Lib, // dylib, staticlib, lib
    Bin, // bin
    Example, // example file
    Test, // test file
    Bench, // bench file
    CustomBuild, // build script
    Other, // plugin,...
}

impl TargetKind {
    fn should_format(&self) -> bool {
        match *self {
            TargetKind::Lib | TargetKind::Bin | TargetKind::Example | TargetKind::Test |
            TargetKind::Bench | TargetKind::CustomBuild => true,
            _ => false,
        }
    }
}

#[derive(Debug)]
pub struct Target {
    path: PathBuf,
    kind: TargetKind,
}

#[derive(Debug, PartialEq, Eq)]
pub enum WorkspaceHitlist {
    All,
    Some(Vec<String>),
    None,
}

impl WorkspaceHitlist {
    pub fn get_some<'a>(&'a self) -> Option<&'a [String]> {
        if let &WorkspaceHitlist::Some(ref hitlist) = self {
            Some(&hitlist)
        } else {
            None
        }
    }

    pub fn from_matches(matches: &Matches) -> WorkspaceHitlist {
        match (matches.opt_present("all"), matches.opt_present("p")) {
            (false, false) => WorkspaceHitlist::None,
            (true, _) => WorkspaceHitlist::All,
            (false, true) => WorkspaceHitlist::Some(matches.opt_strs("p")),
        }
    }
}

// Returns a vector of all compile targets of a crate
fn get_targets(workspace_hitlist: WorkspaceHitlist) -> Result<Vec<Target>, std::io::Error> {
    let mut targets: Vec<Target> = vec![];
    if workspace_hitlist == WorkspaceHitlist::None {
        let output = try!(Command::new("cargo").arg("read-manifest").output());
        if output.status.success() {
            // None of the unwraps should fail if output of `cargo read-manifest` is correct
            let data = &String::from_utf8(output.stdout).unwrap();
            let json = Json::from_str(data).unwrap();
            let jtargets = json.find("targets").unwrap().as_array().unwrap();
            for jtarget in jtargets {
                targets.push(target_from_json(jtarget));
            }

            return Ok(targets);
        }
        return Err(std::io::Error::new(std::io::ErrorKind::NotFound,
                                       str::from_utf8(&output.stderr).unwrap()));
    }
    // This happens when cargo-fmt is not used inside a crate or
    // is used inside a workspace.
    // To ensure backward compatability, we only use `cargo metadata` for workspaces.
    // TODO: Is it possible only use metadata or read-manifest
    let output = Command::new("cargo").arg("metadata")
        .arg("--no-deps")
        .output()?;
    if output.status.success() {
        let data = &String::from_utf8(output.stdout).unwrap();
        let json = Json::from_str(data).unwrap();
        let mut hitlist: HashSet<&String> = if workspace_hitlist != WorkspaceHitlist::All {
            HashSet::from_iter(workspace_hitlist.get_some().unwrap())
        } else {
            HashSet::new() // Unused
        };
        let members: Vec<&Json> = json.find("packages")
            .unwrap()
            .as_array()
            .unwrap()
            .into_iter()
            .filter(|member| if workspace_hitlist == WorkspaceHitlist::All {
                        true
                    } else {
                        let member_name = member.find("name").unwrap().as_string().unwrap();
                        hitlist.take(&member_name.to_string()).is_some()
                    })
            .collect();
        if hitlist.len() != 0 {
            // Mimick cargo of only outputting one <package> spec.
            return Err(std::io::Error::new(std::io::ErrorKind::InvalidInput,
                                           format!("package `{}` is not a member of the workspace",
                                                   hitlist.iter().next().unwrap())));
        }
        for member in members {
            let jtargets = member.find("targets").unwrap().as_array().unwrap();
            for jtarget in jtargets {
                targets.push(target_from_json(jtarget));
            }
        }
        return Ok(targets);
    }
    Err(std::io::Error::new(std::io::ErrorKind::NotFound,
                            str::from_utf8(&output.stderr).unwrap()))

}

fn target_from_json(jtarget: &Json) -> Target {
    let jtarget = jtarget.as_object().unwrap();
    let path = PathBuf::from(jtarget.get("src_path").unwrap().as_string().unwrap());
    let kinds = jtarget.get("kind").unwrap().as_array().unwrap();
    let kind = match kinds[0].as_string().unwrap() {
        "bin" => TargetKind::Bin,
        "lib" | "dylib" | "staticlib" | "cdylib" | "rlib" => TargetKind::Lib,
        "test" => TargetKind::Test,
        "example" => TargetKind::Example,
        "bench" => TargetKind::Bench,
        "custom-build" => TargetKind::CustomBuild,
        _ => TargetKind::Other,
    };

    Target {
        path: path,
        kind: kind,
    }
}

fn format_files(files: &[PathBuf],
                fmt_args: &[String],
                verbosity: Verbosity)
                -> Result<ExitStatus, std::io::Error> {
    let stdout = if verbosity == Verbosity::Quiet {
        std::process::Stdio::null()
    } else {
        std::process::Stdio::inherit()
    };
    if verbosity == Verbosity::Verbose {
        print!("rustfmt");
        for a in fmt_args.iter() {
            print!(" {}", a);
        }
        for f in files.iter() {
            print!(" {}", f.display());
        }
        println!("");
    }
    let mut command = try!(Command::new("rustfmt")
        .stdout(stdout)
        .args(files)
        .args(fmt_args)
        .spawn()
        .map_err(|e| match e.kind() {
            std::io::ErrorKind::NotFound => {
                std::io::Error::new(std::io::ErrorKind::Other,
                                    "Could not run rustfmt, please make sure it is in your PATH.")
            }
            _ => e,
        }));
    command.wait()
}
<|endoftext|>


<commit_before>// Implements http://rosettacode.org/wiki/Atomic_updates
//
// This is mostly a straight port of the D version.  Originally, the "non-locking" Go solution was
// tried, because it was supposed to be faster than the version with Mutexes, but my experience was
// that this was not the case.  Perhaps it is true with green threads.  D's version was much faster
// and this version seems to achieve parity with the benchmarks on the Rosetta Code site (at least
// on my machine).  I am pretty sure it could be made faster, though--for example, the Mutex type
// we're using here was the fourth type I tried but the first to produce acceptable performance
// (previously I tried, in order, std::sync::RwLock, std::sync::Mutex, and std::sync::Semaphore)
// and this type still appears to have quite a bit of overhead.
#![feature(tuple_indexing)]

extern crate sync;

use std::io::timer::Timer;
use std::iter::AdditiveIterator;
use std::rand::{Rng, weak_rng};
use std::rand::distributions::{IndependentSample, Range};
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomics;
use std::time::duration::Duration;

// The reason I used a module here is simply to keep it clearer who can access what.  Rust
// protects against data races just fine, but it's not as good at protecting against deadlocks or
// other types of race conditions.
mod buckets {
    use std::sync::atomic::AtomicUint;
    use std::sync::atomics;
    use sync::mutex::Mutex;

    // We hardcode the number of buckets mostly for convenience.  Now that Rust has dynamically
    // sized types, this is possibly no longer a problem.

    // To expand: in Rust, there are two special kinds, Sync and Send, used for concurrency.
    //
    // If T is Sync, and you take an immutable reference of type &T, then it's safe to share between
    // threads.  Most types are Sync unless they contain non-threadsafe interior mutability: for
    // example, Cell and RefCell are not Sync, because they can be modified through a & reference
    // in a non-threadsafe way.  On the other hand, atomic types are Sync even though they can be
    // modified this way.  So are types accessible through a Mutex.
    //
    // The other type is Send.  Send's semantic meanining is "a type that can be sent between
    // tasks."  For most practical purposes, this means "has no non-'static references" (with a very
    // few exceptions, like the Rc type, which are non-Send anyway).  The idea is that it is not
    // safe to send a type between tasks if it has any non-static references, because one doesn't
    // know when the data it's referencing will be deallocated if it's on some other task's stack.
    //
    // Usually, that's a reasonable assumption.  But it means that if we want to put a slice in a
    // structure and share it between tasks, it can't be a &[T]--it has to be a bare [T].  Arc has
    // Send + Sync bounds for this reason.  We could also use a Vec, but this would introduce
    // unnecessary indirection.  With dynamically sized types, this problem may be solved--we can
    // have a dynamically sized struct that doesn't contain explicit references.  But until that is
    // fully baked, this seems like the sanest solution.
    //
    // (Another way to solve this would be associated constants, something Rust does not have yet).
    pub static N_BUCKETS: uint = 20;

    // We don't really have to hardcode the workers.  This is left over from the Go implementation.
    // All the counting statistics could be moved outside of buckets and probably should, since
    // they have no influence on the correctness of the algorithm.
    pub static N_WORKERS: uint = 2;

    struct Bucket {
        data: AtomicUint, // The actual data.  It is atomic because it is read (not written)
                          // outside the Mutex, unless a consistent snapshot is required.
        mutex: Mutex,     // The mutex used to synchronize writes and snapshot reads of the bucket.
                          // As the D solution says, using a per-bucket Mutex dramatically improves
                          // scalability compared to the alternatives.
    }

    pub struct Buckets {
        buckets: [Bucket, .. N_BUCKETS ], // Buckets containing values to be transferred.
        transfers: [AtomicUint, .. N_WORKERS ], // Statistics about total transfers this go-around.
    }

    impl Buckets {
        // Create a new Buckets instance.
        pub fn new(buckets: [uint, .. N_BUCKETS]) -> Buckets {
            // The unsafe initialization here is required because Bucket is not Clone (it can't be,
            // since neither AtomicUint nor Mutex are) and we would otherwise have to literally
            // write out N_BUCKETS different values, which would be painful.  As a result, we have
            // to be careful not to allow any failure here, or we'll segfault (by Drop-ing empty
            // buckets).
            let mut buckets_: [Bucket, .. N_BUCKETS] = unsafe { ::std::mem::uninitialized() };
            let mut transfers: [AtomicUint, .. N_WORKERS] = unsafe { ::std::mem::uninitialized() };
            for (dest, &src) in buckets_.iter_mut().zip(buckets.iter()){
                let bucket = Bucket { data: AtomicUint::new(src), mutex: Mutex::new() };
                // If we don't use an unsafe write(), the uninitialized mutex in the bucket will be
                // dropped.
                unsafe { ::std::ptr::write(dest as *mut _, bucket) }
            }
            for t in transfers.as_mut_slice().iter_mut() {
                *t = AtomicUint::new(0);
            }
            Buckets { buckets: buckets_, transfers: transfers }
        }

        // Get the value of the bucket at index i, or None if out of bounds.
        pub fn get(&self, i: uint) -> Option<uint> {
            // This is used as an estimate, and is used without the mutex lock, so there's no
            // compelling reason to demand consistency here.
            self.buckets.get(i).map( |b| b.data.load(atomics::Relaxed) )
        }

        // Transfer at most `amount` from the bucket at index `from` to that at index `to`, and
        // increment the transfer count for worker `worker` (like I said, that last part can likely
        // be done elsewhere).
        pub fn transfer(&self, from: uint, to: uint, amount: uint, worker: uint) {
            // The from == to check is important so we don't deadlock, since Rust mutexes are
            // nonreentrant.
            if from == to || N_BUCKETS <= from || N_BUCKETS <= to || N_WORKERS <= worker { return }
            // We know this won't fail, and the compiler seems to know as well.  However, if it
            // *did* fail, it wouldn't fail while we were holding mutexes (which can cause
            // problems since they may need to poison other tasks).
            let ref b1 = self.buckets[from];
            let ref b2 = self.buckets[to];
            // It's very important to lock our Mutexes in the same order everywhere to avoid
            // deadlock.  We arbitrarily choose the convention that we lock in ascending index
            // order.
            let (low, high) = if from < to { (b1, b2) } else { (b2, b1) };
            {
                // The reason we introduce a new scope here is that we want to make it clear how
                // long we're locking for.  Locks should be held as briefly as possible and
                // anything that happens here should really *require* the locks.
                let _s1 = low.mutex.lock();
                let _s2 = high.mutex.lock();
                // It is possible that SeqCst is too strong for this section, but it is hard to
                // test on x86 because it has unusually strong consistency by default.
                let v1 = b1.data.load(atomics::SeqCst);
                let real_amount = ::std::cmp::min(v1, amount);
                b1.data.store(v1 - real_amount, atomics::SeqCst);
                b2.data.fetch_add(real_amount, atomics::SeqCst);
            }
            // Doing this outside the critical section increases throughput substantially.  Since
            // this is just a summary statistic, it's okay for it to be a few off.  That's also why
            // we use Acquire semantics rather than AcqRel or SeqCst here--we only really care that
            // we synchronize when the transfer count is set to 0.
            self.transfers[worker].fetch_add(1, atomics::Acquire);
        }

        // Acquire a consistent snapshot of the state of the bucket list.  This should maintain the
        // invariant that total buckets are conserved.  Also returns the list of transfer counts.
        pub fn snapshot(&self) -> ([uint, .. N_BUCKETS], [uint, .. N_WORKERS ]) {
            // Since this method is called relatively rarely, we aren't too concerned about
            // performance here.
            let mut buckets = [0, .. N_BUCKETS];
            let mut transfers = [0, .. N_WORKERS];
            // We collect all the locks in order, being careful not to drop any until we're done
            // (so as to preserve consistency of the snapshot).
            let locks = buckets.iter_mut().zip(self.buckets.iter()).map( |(dest, src)| {
                let lock = src.mutex.lock();
                // Is SeqCst necessary here?  Maybe, maybe not, but when in doubt go with SeqCst.
                *dest = src.data.load(atomics::SeqCst);
                lock
            }).collect::<Vec<_>>();
            for (dest, src) in transfers.iter_mut().zip(self.transfers.iter()) {
                // We synchronize with the Acquire in transfer, making sure that our zeroing out
                // gets noticed.
                *dest = src.swap(0, atomics::Release);
            }
            // We can drop the locks before we return.  This probably gets optimized out, but it's
            // rarely a bad idea to drop locks explicitly.
            drop(locks);
            (buckets, transfers)
        }
    }
}

// Convenience method to create a distribution of buckets summing to initial_sum.
fn make_buckets(initial_sum: uint) -> buckets::Buckets {
    let mut buckets = [0, .. buckets::N_BUCKETS];
    let mut dist = initial_sum;
    for (i, b) in buckets.as_mut_slice().iter_mut().enumerate() {
        let v = dist / (buckets::N_BUCKETS - i);
        *b = v;
        dist -= v;
    }
    buckets::Buckets::new(buckets)
}

// The equalize task--it chooses two random buckets and tries to make their values the same.
fn equalize(bl: &buckets::Buckets, running: &AtomicBool, worker: uint) {
    // We preallocate the Range for improved performance.
    let between = Range::new(0, buckets::N_BUCKETS);
    // We use the weak random number generator for improved performance.
    let ref mut r = weak_rng();
    // Running is read Relaxed because it's not important that the task stop right away as long as
    // it happens eventually.
    while running.load(atomics::Relaxed) {
        let b1 = between.ind_sample(r);
        let b2 = between.ind_sample(r);
        let v1 = bl.get(b1).unwrap();
        let v2 = bl.get(b2).unwrap();
        if v1 > v2 {
            bl.transfer(b1, b2, (v1 - v2) / 2, worker)
        } else {
            bl.transfer(b2, b1, (v2 - v1) / 2, worker)
        }
    }
}

// The randomize task--it chooses two random buckets and randomly redistributes their values.
fn randomize(bl: &buckets::Buckets, running: &AtomicBool, worker: uint) {
    // We preallocate the Range for improved performance.
    let between = Range::new(0, buckets::N_BUCKETS);
    // We use the weak random number generator for improved performance.
    let ref mut r = weak_rng();
    // Running is read Relaxed because it's not important that the task stop right away as long as
    // it happens eventually.
    while running.load(atomics::Relaxed) {
        let b1 = between.ind_sample(r);
        let b2 = between.ind_sample(r);
        bl.transfer(b1, b2, r.gen_range(0, bl.get(b1).unwrap() + 1), worker);
    }
}

// The display task--for a total time of `duration`, it displays information about the update
// process and checks to make sure that the invariant (that the total remains constant) is
// preserved.  It prints an update `nticks` times, evenly spaced.
fn display(bl: &buckets::Buckets, running: &AtomicBool, original_total: uint, duration: Duration, nticks: i32) {
    println!("transfers, N. transfers, buckets, buckets sum:");

    let mut timer = Timer::new().unwrap();
    let duration = duration / nticks;
    for _ in range(0, nticks) {
        // Get a consistent snapshot
        let (s, tc) = bl.snapshot();
        // Sum up the buckets
        let sum = s.iter().map( |&i| i ).sum();
        // Sum up the transfers.
        let n_transfers = tc.iter().map( |&i| i ).sum();
        // Print the relevant information.
        println!("{}, {}, {}, {}", tc[], n_transfers, s[], sum);
        // Check the invariant, failing if necessary.
        assert_eq!(sum, original_total);
        // Sleep before printing again.
        timer.sleep(duration);
    }
    // We're done--cleanly exit the other update tasks.
    running.store(false, atomics::Relaxed);
}

// Putting together all three tasks.
fn perform_atomic_updates(duration: Duration, original_total: uint, num_ticks: i32)
{
    // Worker IDs for the two updater tasks.
    static ID_EQUALIZE: uint = 0;
    static ID_RANDOMIZE: uint = 1;

    // `running` is an atomic boolean that we use to signal when to stop to the updater tasks.
    let running = AtomicBool::new(true);
    // We use an Arc here in order to share a reference to the buckets between threads.  Since the
    // Buckets are already Sync, and we don't need &mut references to them, there's no need to
    // create a Mutex here.
    let arc = Arc::new((make_buckets(original_total), running));
    // Cloning the arc bumps the reference count.
    let arc_ = arc.clone();
    // Start off the equalize task
    spawn(proc() { equalize(&arc_.0, &arc_.1, ID_EQUALIZE) });
    let arc_ = arc.clone();
    // Start off the randomize task
    spawn(proc() { randomize(&arc_.0, &arc_.1, ID_RANDOMIZE) });
    let (ref bl, ref running) = *arc;
    // Run the display task in the current thread, so failure propagates to the user.
    display(bl, running, original_total, duration, num_ticks);
}

static ORIGINAL_TOTAL: uint = 1000;
static NUM_TICKS: i32 = 10;

#[cfg(not(test))]
fn main() {
    // Run for 10 seconds
    perform_atomic_updates(Duration::seconds(10), ORIGINAL_TOTAL, NUM_TICKS);
}

#[test]
fn test_atomic_updates() {
    // Run for 1 second
    perform_atomic_updates(Duration::seconds(1), ORIGINAL_TOTAL, NUM_TICKS);
}
<commit_msg>Only test atomic updates for 1/10th of a second with test configuration.<commit_after>// Implements http://rosettacode.org/wiki/Atomic_updates
//
// This is mostly a straight port of the D version.  Originally, the "non-locking" Go solution was
// tried, because it was supposed to be faster than the version with Mutexes, but my experience was
// that this was not the case.  Perhaps it is true with green threads.  D's version was much faster
// and this version seems to achieve parity with the benchmarks on the Rosetta Code site (at least
// on my machine).  I am pretty sure it could be made faster, though--for example, the Mutex type
// we're using here was the fourth type I tried but the first to produce acceptable performance
// (previously I tried, in order, std::sync::RwLock, std::sync::Mutex, and std::sync::Semaphore)
// and this type still appears to have quite a bit of overhead.
#![feature(tuple_indexing)]

extern crate sync;

use std::io::timer::Timer;
use std::iter::AdditiveIterator;
use std::rand::{Rng, weak_rng};
use std::rand::distributions::{IndependentSample, Range};
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomics;
use std::time::duration::Duration;

// The reason I used a module here is simply to keep it clearer who can access what.  Rust
// protects against data races just fine, but it's not as good at protecting against deadlocks or
// other types of race conditions.
mod buckets {
    use std::sync::atomic::AtomicUint;
    use std::sync::atomics;
    use sync::mutex::Mutex;

    // We hardcode the number of buckets mostly for convenience.  Now that Rust has dynamically
    // sized types, this is possibly no longer a problem.

    // To expand: in Rust, there are two special kinds, Sync and Send, used for concurrency.
    //
    // If T is Sync, and you take an immutable reference of type &T, then it's safe to share between
    // threads.  Most types are Sync unless they contain non-threadsafe interior mutability: for
    // example, Cell and RefCell are not Sync, because they can be modified through a & reference
    // in a non-threadsafe way.  On the other hand, atomic types are Sync even though they can be
    // modified this way.  So are types accessible through a Mutex.
    //
    // The other type is Send.  Send's semantic meanining is "a type that can be sent between
    // tasks."  For most practical purposes, this means "has no non-'static references" (with a very
    // few exceptions, like the Rc type, which are non-Send anyway).  The idea is that it is not
    // safe to send a type between tasks if it has any non-static references, because one doesn't
    // know when the data it's referencing will be deallocated if it's on some other task's stack.
    //
    // Usually, that's a reasonable assumption.  But it means that if we want to put a slice in a
    // structure and share it between tasks, it can't be a &[T]--it has to be a bare [T].  Arc has
    // Send + Sync bounds for this reason.  We could also use a Vec, but this would introduce
    // unnecessary indirection.  With dynamically sized types, this problem may be solved--we can
    // have a dynamically sized struct that doesn't contain explicit references.  But until that is
    // fully baked, this seems like the sanest solution.
    //
    // (Another way to solve this would be associated constants, something Rust does not have yet).
    pub static N_BUCKETS: uint = 20;

    // We don't really have to hardcode the workers.  This is left over from the Go implementation.
    // All the counting statistics could be moved outside of buckets and probably should, since
    // they have no influence on the correctness of the algorithm.
    pub static N_WORKERS: uint = 2;

    struct Bucket {
        data: AtomicUint, // The actual data.  It is atomic because it is read (not written)
                          // outside the Mutex, unless a consistent snapshot is required.
        mutex: Mutex,     // The mutex used to synchronize writes and snapshot reads of the bucket.
                          // As the D solution says, using a per-bucket Mutex dramatically improves
                          // scalability compared to the alternatives.
    }

    pub struct Buckets {
        buckets: [Bucket, .. N_BUCKETS ], // Buckets containing values to be transferred.
        transfers: [AtomicUint, .. N_WORKERS ], // Statistics about total transfers this go-around.
    }

    impl Buckets {
        // Create a new Buckets instance.
        pub fn new(buckets: [uint, .. N_BUCKETS]) -> Buckets {
            // The unsafe initialization here is required because Bucket is not Clone (it can't be,
            // since neither AtomicUint nor Mutex are) and we would otherwise have to literally
            // write out N_BUCKETS different values, which would be painful.  As a result, we have
            // to be careful not to allow any failure here, or we'll segfault (by Drop-ing empty
            // buckets).
            let mut buckets_: [Bucket, .. N_BUCKETS] = unsafe { ::std::mem::uninitialized() };
            let mut transfers: [AtomicUint, .. N_WORKERS] = unsafe { ::std::mem::uninitialized() };
            for (dest, &src) in buckets_.iter_mut().zip(buckets.iter()){
                let bucket = Bucket { data: AtomicUint::new(src), mutex: Mutex::new() };
                // If we don't use an unsafe write(), the uninitialized mutex in the bucket will be
                // dropped.
                unsafe { ::std::ptr::write(dest as *mut _, bucket) }
            }
            for t in transfers.as_mut_slice().iter_mut() {
                *t = AtomicUint::new(0);
            }
            Buckets { buckets: buckets_, transfers: transfers }
        }

        // Get the value of the bucket at index i, or None if out of bounds.
        pub fn get(&self, i: uint) -> Option<uint> {
            // This is used as an estimate, and is used without the mutex lock, so there's no
            // compelling reason to demand consistency here.
            self.buckets.get(i).map( |b| b.data.load(atomics::Relaxed) )
        }

        // Transfer at most `amount` from the bucket at index `from` to that at index `to`, and
        // increment the transfer count for worker `worker` (like I said, that last part can likely
        // be done elsewhere).
        pub fn transfer(&self, from: uint, to: uint, amount: uint, worker: uint) {
            // The from == to check is important so we don't deadlock, since Rust mutexes are
            // nonreentrant.
            if from == to || N_BUCKETS <= from || N_BUCKETS <= to || N_WORKERS <= worker { return }
            // We know this won't fail, and the compiler seems to know as well.  However, if it
            // *did* fail, it wouldn't fail while we were holding mutexes (which can cause
            // problems since they may need to poison other tasks).
            let ref b1 = self.buckets[from];
            let ref b2 = self.buckets[to];
            // It's very important to lock our Mutexes in the same order everywhere to avoid
            // deadlock.  We arbitrarily choose the convention that we lock in ascending index
            // order.
            let (low, high) = if from < to { (b1, b2) } else { (b2, b1) };
            {
                // The reason we introduce a new scope here is that we want to make it clear how
                // long we're locking for.  Locks should be held as briefly as possible and
                // anything that happens here should really *require* the locks.
                let _s1 = low.mutex.lock();
                let _s2 = high.mutex.lock();
                // It is possible that SeqCst is too strong for this section, but it is hard to
                // test on x86 because it has unusually strong consistency by default.
                let v1 = b1.data.load(atomics::SeqCst);
                let real_amount = ::std::cmp::min(v1, amount);
                b1.data.store(v1 - real_amount, atomics::SeqCst);
                b2.data.fetch_add(real_amount, atomics::SeqCst);
            }
            // Doing this outside the critical section increases throughput substantially.  Since
            // this is just a summary statistic, it's okay for it to be a few off.  That's also why
            // we use Acquire semantics rather than AcqRel or SeqCst here--we only really care that
            // we synchronize when the transfer count is set to 0.
            self.transfers[worker].fetch_add(1, atomics::Acquire);
        }

        // Acquire a consistent snapshot of the state of the bucket list.  This should maintain the
        // invariant that total buckets are conserved.  Also returns the list of transfer counts.
        pub fn snapshot(&self) -> ([uint, .. N_BUCKETS], [uint, .. N_WORKERS ]) {
            // Since this method is called relatively rarely, we aren't too concerned about
            // performance here.
            let mut buckets = [0, .. N_BUCKETS];
            let mut transfers = [0, .. N_WORKERS];
            // We collect all the locks in order, being careful not to drop any until we're done
            // (so as to preserve consistency of the snapshot).
            let locks = buckets.iter_mut().zip(self.buckets.iter()).map( |(dest, src)| {
                let lock = src.mutex.lock();
                // Is SeqCst necessary here?  Maybe, maybe not, but when in doubt go with SeqCst.
                *dest = src.data.load(atomics::SeqCst);
                lock
            }).collect::<Vec<_>>();
            for (dest, src) in transfers.iter_mut().zip(self.transfers.iter()) {
                // We synchronize with the Acquire in transfer, making sure that our zeroing out
                // gets noticed.
                *dest = src.swap(0, atomics::Release);
            }
            // We can drop the locks before we return.  This probably gets optimized out, but it's
            // rarely a bad idea to drop locks explicitly.
            drop(locks);
            (buckets, transfers)
        }
    }
}

// Convenience method to create a distribution of buckets summing to initial_sum.
fn make_buckets(initial_sum: uint) -> buckets::Buckets {
    let mut buckets = [0, .. buckets::N_BUCKETS];
    let mut dist = initial_sum;
    for (i, b) in buckets.as_mut_slice().iter_mut().enumerate() {
        let v = dist / (buckets::N_BUCKETS - i);
        *b = v;
        dist -= v;
    }
    buckets::Buckets::new(buckets)
}

// The equalize task--it chooses two random buckets and tries to make their values the same.
fn equalize(bl: &buckets::Buckets, running: &AtomicBool, worker: uint) {
    // We preallocate the Range for improved performance.
    let between = Range::new(0, buckets::N_BUCKETS);
    // We use the weak random number generator for improved performance.
    let ref mut r = weak_rng();
    // Running is read Relaxed because it's not important that the task stop right away as long as
    // it happens eventually.
    while running.load(atomics::Relaxed) {
        let b1 = between.ind_sample(r);
        let b2 = between.ind_sample(r);
        let v1 = bl.get(b1).unwrap();
        let v2 = bl.get(b2).unwrap();
        if v1 > v2 {
            bl.transfer(b1, b2, (v1 - v2) / 2, worker)
        } else {
            bl.transfer(b2, b1, (v2 - v1) / 2, worker)
        }
    }
}

// The randomize task--it chooses two random buckets and randomly redistributes their values.
fn randomize(bl: &buckets::Buckets, running: &AtomicBool, worker: uint) {
    // We preallocate the Range for improved performance.
    let between = Range::new(0, buckets::N_BUCKETS);
    // We use the weak random number generator for improved performance.
    let ref mut r = weak_rng();
    // Running is read Relaxed because it's not important that the task stop right away as long as
    // it happens eventually.
    while running.load(atomics::Relaxed) {
        let b1 = between.ind_sample(r);
        let b2 = between.ind_sample(r);
        bl.transfer(b1, b2, r.gen_range(0, bl.get(b1).unwrap() + 1), worker);
    }
}

// The display task--for a total time of `duration`, it displays information about the update
// process and checks to make sure that the invariant (that the total remains constant) is
// preserved.  It prints an update `nticks` times, evenly spaced.
fn display(bl: &buckets::Buckets, running: &AtomicBool, original_total: uint, duration: Duration, nticks: i32) {
    println!("transfers, N. transfers, buckets, buckets sum:");

    let mut timer = Timer::new().unwrap();
    let duration = duration / nticks;
    for _ in range(0, nticks) {
        // Get a consistent snapshot
        let (s, tc) = bl.snapshot();
        // Sum up the buckets
        let sum = s.iter().map( |&i| i ).sum();
        // Sum up the transfers.
        let n_transfers = tc.iter().map( |&i| i ).sum();
        // Print the relevant information.
        println!("{}, {}, {}, {}", tc[], n_transfers, s[], sum);
        // Check the invariant, failing if necessary.
        assert_eq!(sum, original_total);
        // Sleep before printing again.
        timer.sleep(duration);
    }
    // We're done--cleanly exit the other update tasks.
    running.store(false, atomics::Relaxed);
}

// Putting together all three tasks.
fn perform_atomic_updates(duration: Duration, original_total: uint, num_ticks: i32)
{
    // Worker IDs for the two updater tasks.
    static ID_EQUALIZE: uint = 0;
    static ID_RANDOMIZE: uint = 1;

    // `running` is an atomic boolean that we use to signal when to stop to the updater tasks.
    let running = AtomicBool::new(true);
    // We use an Arc here in order to share a reference to the buckets between threads.  Since the
    // Buckets are already Sync, and we don't need &mut references to them, there's no need to
    // create a Mutex here.
    let arc = Arc::new((make_buckets(original_total), running));
    // Cloning the arc bumps the reference count.
    let arc_ = arc.clone();
    // Start off the equalize task
    spawn(proc() { equalize(&arc_.0, &arc_.1, ID_EQUALIZE) });
    let arc_ = arc.clone();
    // Start off the randomize task
    spawn(proc() { randomize(&arc_.0, &arc_.1, ID_RANDOMIZE) });
    let (ref bl, ref running) = *arc;
    // Run the display task in the current thread, so failure propagates to the user.
    display(bl, running, original_total, duration, num_ticks);
}

static ORIGINAL_TOTAL: uint = 1000;
static NUM_TICKS: i32 = 10;

#[cfg(not(test))]
fn main() {
    // Run for 10 seconds
    perform_atomic_updates(Duration::seconds(10), ORIGINAL_TOTAL, NUM_TICKS);
}

#[test]
fn test_atomic_updates() {
    // Run for 1/10th of a second
    perform_atomic_updates(Duration::seconds(1) / 10, ORIGINAL_TOTAL, NUM_TICKS);
}
<|endoftext|>


<commit_before>use super::attestation_parent_hashes::{attestation_parent_hashes, ParentHashesError};
use super::db::stores::{BeaconBlockAtSlotError, BeaconBlockStore, ValidatorStore};
use super::db::{ClientDB, DBError};
use super::message_generation::generate_signed_message;
use super::signature_verification::{
    verify_aggregate_signature_for_indices, SignatureVerificationError,
};
use super::types::Hash256;
use super::types::{AttestationRecord, AttesterMap};
use std::collections::HashSet;
use std::sync::Arc;

#[derive(Debug, PartialEq)]
pub enum AttestationValidationError {
    ParentSlotTooHigh,
    ParentSlotTooLow,
    BlockSlotTooHigh,
    BlockSlotTooLow,
    JustifiedSlotIncorrect,
    InvalidJustifiedBlockHash,
    TooManyObliqueHashes,
    BadCurrentHashes,
    BadObliqueHashes,
    BadAttesterMap,
    IntWrapping,
    PublicKeyCorrupt,
    NoPublicKeyForValidator,
    BadBitfieldLength,
    InvalidBitfield,
    InvalidBitfieldEndBits,
    NoSignatures,
    NonZeroTrailingBits,
    BadAggregateSignature,
    DBError(String),
    OutOfBoundsBitfieldIndex,
}

/// The context against which some attestation should be validated.
pub struct AttestationValidationContext<T>
where
    T: ClientDB + Sized,
{
    /// The slot as determined by the system time.
    pub block_slot: u64,
    /// The slot of the parent of the block that contained this attestation.
    pub parent_block_slot: u64,
    /// The cycle_length as determined by the chain configuration.
    pub cycle_length: u8,
    /// The last justified slot as per the client's view of the canonical chain.
    pub last_justified_slot: u64,
    /// A vec of the hashes of the blocks preceeding the present slot.
    pub recent_block_hashes: Arc<Vec<Hash256>>,
    /// The store containing block information.
    pub block_store: Arc<BeaconBlockStore<T>>,
    /// The store containing validator information.
    pub validator_store: Arc<ValidatorStore<T>>,
    /// A map of (slot, shard_id) to the attestation set of validation indices.
    pub attester_map: Arc<AttesterMap>,
}

impl<T> AttestationValidationContext<T>
where
    T: ClientDB,
{
    /// Validate a (fully deserialized) AttestationRecord against this context.
    ///
    /// The function will return a HashSet of validator indices (canonical validator indices not
    /// attestation indices) if the validation passed successfully, or an error otherwise.
    ///
    /// The attestation's aggregate signature will be verified, therefore the function must able to
    /// access all required validation public keys via the `validator_store`.
    pub fn validate_attestation(
        &self,
        a: &AttestationRecord,
    ) -> Result<HashSet<usize>, AttestationValidationError> {
        /*
         * The attesation slot must be less than or equal to the parent of the slot of the block
         * that contained the attestation.
         */
        if a.slot > self.parent_block_slot {
            return Err(AttestationValidationError::ParentSlotTooHigh);
        }

        /*
         * The slot of this attestation must not be more than cycle_length + 1 distance
         * from the parent_slot of block that contained it.
         */
        if a.slot < self
            .parent_block_slot
            .saturating_sub(u64::from(self.cycle_length).saturating_add(1))
        {
            return Err(AttestationValidationError::ParentSlotTooLow);
        }

        /*
         * The attestation justified slot must not be higher than the last_justified_slot of the
         * context.
         */
        if a.justified_slot > self.last_justified_slot {
            return Err(AttestationValidationError::JustifiedSlotIncorrect);
        }

        /*
         * There is no need to include more oblique parents hashes than there are blocks
         * in a cycle.
         */
        if a.oblique_parent_hashes.len() > usize::from(self.cycle_length) {
            return Err(AttestationValidationError::TooManyObliqueHashes);
        }

        /*
         * Retrieve the set of attestation indices for this slot and shard id.
         *
         * This is an array mapping the order that validators will appear in the bitfield to the
         * canonincal index of a validator.
         */
        let attestation_indices = self
            .attester_map
            .get(&(a.slot, a.shard_id))
            .ok_or(AttestationValidationError::BadAttesterMap)?;

        /*
         * The bitfield must be no longer than the minimum required to represent each validator in the
         * attestation indices for this slot and shard id.
         */
        if a.attester_bitfield.num_bytes() != bytes_for_bits(attestation_indices.len()) {
            return Err(AttestationValidationError::BadBitfieldLength);
        }

        /*
         * If there are excess bits in the bitfield because the number of a validators in not a
         * multiple of 8, reject this attestation record.
         *
         * Allow extra set bits would permit mutliple different byte layouts (and therefore hashes) to
         * refer to the same AttesationRecord.
         */
        if a.attester_bitfield.len() > attestation_indices.len() {
            return Err(AttestationValidationError::InvalidBitfieldEndBits);
        }

        /*
         * Generate the parent hashes for this attestation
         */
        let parent_hashes = attestation_parent_hashes(
            self.cycle_length,
            self.block_slot,
            a.slot,
            &self.recent_block_hashes,
            &a.oblique_parent_hashes,
        )?;

        /*
         * The specified justified block hash supplied in the attestation must be in the chain at
         * the given slot number.
         *
         * First, we find the latest parent hash from the parent_hashes array. Then, using the
         * block store (database) we iterate back through the blocks until we find (or fail to
         * find) the justified block hash referenced in the attestation record.
         */
        let latest_parent_hash = parent_hashes
            .last()
            .ok_or(AttestationValidationError::BadCurrentHashes)?;
        match self
            .block_store
            .block_at_slot(&latest_parent_hash, a.justified_slot)?
        {
            Some((ref hash, _)) if *hash == a.justified_block_hash.to_vec() => (),
            _ => return Err(AttestationValidationError::InvalidJustifiedBlockHash),
        };

        /*
         * Generate the message that this attestation aggregate signature must sign across.
         */
        let signed_message = {
            generate_signed_message(
                a.slot,
                &parent_hashes,
                a.shard_id,
                &a.shard_block_hash,
                a.justified_slot,
            )
        };

        let voted_hashset = verify_aggregate_signature_for_indices(
            &signed_message,
            &a.aggregate_sig,
            &attestation_indices,
            &a.attester_bitfield,
            &self.validator_store,
        )?;

        /*
         * If the hashset of voters is None, the signature verification failed.
         */
        match voted_hashset {
            None => Err(AttestationValidationError::BadAggregateSignature),
            Some(hashset) => Ok(hashset),
        }
    }
}

fn bytes_for_bits(bits: usize) -> usize {
    (bits.saturating_sub(1) / 8) + 1
}

impl From<ParentHashesError> for AttestationValidationError {
    fn from(e: ParentHashesError) -> Self {
        match e {
            ParentHashesError::BadCurrentHashes => AttestationValidationError::BadCurrentHashes,
            ParentHashesError::BadObliqueHashes => AttestationValidationError::BadObliqueHashes,
            ParentHashesError::SlotTooLow => AttestationValidationError::BlockSlotTooLow,
            ParentHashesError::SlotTooHigh => AttestationValidationError::BlockSlotTooHigh,
            ParentHashesError::IntWrapping => AttestationValidationError::IntWrapping,
        }
    }
}

impl From<BeaconBlockAtSlotError> for AttestationValidationError {
    fn from(e: BeaconBlockAtSlotError) -> Self {
        match e {
            BeaconBlockAtSlotError::DBError(s) => AttestationValidationError::DBError(s),
            _ => AttestationValidationError::InvalidJustifiedBlockHash,
        }
    }
}

impl From<DBError> for AttestationValidationError {
    fn from(e: DBError) -> Self {
        AttestationValidationError::DBError(e.message)
    }
}

impl From<SignatureVerificationError> for AttestationValidationError {
    fn from(e: SignatureVerificationError) -> Self {
        match e {
            SignatureVerificationError::BadValidatorIndex => {
                AttestationValidationError::BadAttesterMap
            }
            SignatureVerificationError::PublicKeyCorrupt => {
                AttestationValidationError::PublicKeyCorrupt
            }
            SignatureVerificationError::NoPublicKeyForValidator => {
                AttestationValidationError::NoPublicKeyForValidator
            }
            SignatureVerificationError::DBError(s) => AttestationValidationError::DBError(s),
            SignatureVerificationError::OutOfBoundsBitfieldIndex => {
                AttestationValidationError::OutOfBoundsBitfieldIndex
            }
        }
    }
}
<commit_msg>rustfmt edits<commit_after>use super::attestation_parent_hashes::{attestation_parent_hashes, ParentHashesError};
use super::db::stores::{BeaconBlockAtSlotError, BeaconBlockStore, ValidatorStore};
use super::db::{ClientDB, DBError};
use super::message_generation::generate_signed_message;
use super::signature_verification::{
    verify_aggregate_signature_for_indices, SignatureVerificationError,
};
use super::types::Hash256;
use super::types::{AttestationRecord, AttesterMap};
use std::collections::HashSet;
use std::sync::Arc;

#[derive(Debug, PartialEq)]
pub enum AttestationValidationError {
    ParentSlotTooHigh,
    ParentSlotTooLow,
    BlockSlotTooHigh,
    BlockSlotTooLow,
    JustifiedSlotIncorrect,
    InvalidJustifiedBlockHash,
    TooManyObliqueHashes,
    BadCurrentHashes,
    BadObliqueHashes,
    BadAttesterMap,
    IntWrapping,
    PublicKeyCorrupt,
    NoPublicKeyForValidator,
    BadBitfieldLength,
    InvalidBitfield,
    InvalidBitfieldEndBits,
    NoSignatures,
    NonZeroTrailingBits,
    BadAggregateSignature,
    DBError(String),
    OutOfBoundsBitfieldIndex,
}

/// The context against which some attestation should be validated.
pub struct AttestationValidationContext<T>
where
    T: ClientDB + Sized,
{
    /// The slot as determined by the system time.
    pub block_slot: u64,
    /// The slot of the parent of the block that contained this attestation.
    pub parent_block_slot: u64,
    /// The cycle_length as determined by the chain configuration.
    pub cycle_length: u8,
    /// The last justified slot as per the client's view of the canonical chain.
    pub last_justified_slot: u64,
    /// A vec of the hashes of the blocks preceeding the present slot.
    pub recent_block_hashes: Arc<Vec<Hash256>>,
    /// The store containing block information.
    pub block_store: Arc<BeaconBlockStore<T>>,
    /// The store containing validator information.
    pub validator_store: Arc<ValidatorStore<T>>,
    /// A map of (slot, shard_id) to the attestation set of validation indices.
    pub attester_map: Arc<AttesterMap>,
}

impl<T> AttestationValidationContext<T>
where
    T: ClientDB,
{
    /// Validate a (fully deserialized) AttestationRecord against this context.
    ///
    /// The function will return a HashSet of validator indices (canonical validator indices not
    /// attestation indices) if the validation passed successfully, or an error otherwise.
    ///
    /// The attestation's aggregate signature will be verified, therefore the function must able to
    /// access all required validation public keys via the `validator_store`.
    pub fn validate_attestation(
        &self,
        a: &AttestationRecord,
    ) -> Result<HashSet<usize>, AttestationValidationError> {
        /*
         * The attesation slot must be less than or equal to the parent of the slot of the block
         * that contained the attestation.
         */
        if a.slot > self.parent_block_slot {
            return Err(AttestationValidationError::ParentSlotTooHigh);
        }

        /*
         * The slot of this attestation must not be more than cycle_length + 1 distance
         * from the parent_slot of block that contained it.
         */
        if a.slot
            < self
                .parent_block_slot
                .saturating_sub(u64::from(self.cycle_length).saturating_add(1))
        {
            return Err(AttestationValidationError::ParentSlotTooLow);
        }

        /*
         * The attestation justified slot must not be higher than the last_justified_slot of the
         * context.
         */
        if a.justified_slot > self.last_justified_slot {
            return Err(AttestationValidationError::JustifiedSlotIncorrect);
        }

        /*
         * There is no need to include more oblique parents hashes than there are blocks
         * in a cycle.
         */
        if a.oblique_parent_hashes.len() > usize::from(self.cycle_length) {
            return Err(AttestationValidationError::TooManyObliqueHashes);
        }

        /*
         * Retrieve the set of attestation indices for this slot and shard id.
         *
         * This is an array mapping the order that validators will appear in the bitfield to the
         * canonincal index of a validator.
         */
        let attestation_indices = self
            .attester_map
            .get(&(a.slot, a.shard_id))
            .ok_or(AttestationValidationError::BadAttesterMap)?;

        /*
         * The bitfield must be no longer than the minimum required to represent each validator in the
         * attestation indices for this slot and shard id.
         */
        if a.attester_bitfield.num_bytes() != bytes_for_bits(attestation_indices.len()) {
            return Err(AttestationValidationError::BadBitfieldLength);
        }

        /*
         * If there are excess bits in the bitfield because the number of a validators in not a
         * multiple of 8, reject this attestation record.
         *
         * Allow extra set bits would permit mutliple different byte layouts (and therefore hashes) to
         * refer to the same AttesationRecord.
         */
        if a.attester_bitfield.len() > attestation_indices.len() {
            return Err(AttestationValidationError::InvalidBitfieldEndBits);
        }

        /*
         * Generate the parent hashes for this attestation
         */
        let parent_hashes = attestation_parent_hashes(
            self.cycle_length,
            self.block_slot,
            a.slot,
            &self.recent_block_hashes,
            &a.oblique_parent_hashes,
        )?;

        /*
         * The specified justified block hash supplied in the attestation must be in the chain at
         * the given slot number.
         *
         * First, we find the latest parent hash from the parent_hashes array. Then, using the
         * block store (database) we iterate back through the blocks until we find (or fail to
         * find) the justified block hash referenced in the attestation record.
         */
        let latest_parent_hash = parent_hashes
            .last()
            .ok_or(AttestationValidationError::BadCurrentHashes)?;
        match self
            .block_store
            .block_at_slot(&latest_parent_hash, a.justified_slot)?
        {
            Some((ref hash, _)) if *hash == a.justified_block_hash.to_vec() => (),
            _ => return Err(AttestationValidationError::InvalidJustifiedBlockHash),
        };

        /*
         * Generate the message that this attestation aggregate signature must sign across.
         */
        let signed_message = {
            generate_signed_message(
                a.slot,
                &parent_hashes,
                a.shard_id,
                &a.shard_block_hash,
                a.justified_slot,
            )
        };

        let voted_hashset = verify_aggregate_signature_for_indices(
            &signed_message,
            &a.aggregate_sig,
            &attestation_indices,
            &a.attester_bitfield,
            &self.validator_store,
        )?;

        /*
         * If the hashset of voters is None, the signature verification failed.
         */
        match voted_hashset {
            None => Err(AttestationValidationError::BadAggregateSignature),
            Some(hashset) => Ok(hashset),
        }
    }
}

fn bytes_for_bits(bits: usize) -> usize {
    (bits.saturating_sub(1) / 8) + 1
}

impl From<ParentHashesError> for AttestationValidationError {
    fn from(e: ParentHashesError) -> Self {
        match e {
            ParentHashesError::BadCurrentHashes => AttestationValidationError::BadCurrentHashes,
            ParentHashesError::BadObliqueHashes => AttestationValidationError::BadObliqueHashes,
            ParentHashesError::SlotTooLow => AttestationValidationError::BlockSlotTooLow,
            ParentHashesError::SlotTooHigh => AttestationValidationError::BlockSlotTooHigh,
            ParentHashesError::IntWrapping => AttestationValidationError::IntWrapping,
        }
    }
}

impl From<BeaconBlockAtSlotError> for AttestationValidationError {
    fn from(e: BeaconBlockAtSlotError) -> Self {
        match e {
            BeaconBlockAtSlotError::DBError(s) => AttestationValidationError::DBError(s),
            _ => AttestationValidationError::InvalidJustifiedBlockHash,
        }
    }
}

impl From<DBError> for AttestationValidationError {
    fn from(e: DBError) -> Self {
        AttestationValidationError::DBError(e.message)
    }
}

impl From<SignatureVerificationError> for AttestationValidationError {
    fn from(e: SignatureVerificationError) -> Self {
        match e {
            SignatureVerificationError::BadValidatorIndex => {
                AttestationValidationError::BadAttesterMap
            }
            SignatureVerificationError::PublicKeyCorrupt => {
                AttestationValidationError::PublicKeyCorrupt
            }
            SignatureVerificationError::NoPublicKeyForValidator => {
                AttestationValidationError::NoPublicKeyForValidator
            }
            SignatureVerificationError::DBError(s) => AttestationValidationError::DBError(s),
            SignatureVerificationError::OutOfBoundsBitfieldIndex => {
                AttestationValidationError::OutOfBoundsBitfieldIndex
            }
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate log;
extern crate env_logger;
extern crate osmpbfreader;
extern crate rustc_serialize;
extern crate docopt;
extern crate mimirsbrunn;

use std::collections::HashSet;
use std::collections::BTreeMap;
use osmpbfreader::{OsmObj, OsmId};

pub type AdminsMap = BTreeMap<OsmId, mimirsbrunn::Admin>;

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_input: String,
    flag_level: Vec<u32>,
    flag_connection_string: String,
}

static USAGE: &'static str = "
Usage:
    osm2mimir --help
    osm2mimir --input=<file> [--connection-string=<connection-string>] --level=<level>...

Options:
    -h, --help            Show this message.
    -i, --input=<file>    OSM PBF file.
    -l, --level=<level>   Admin levels to keep.
    -c, --connection-string=<connection-string>
                          Elasticsearch parameters, [default: http://localhost:9200/munin]
";

fn get_osm_id(obj: &OsmObj) -> OsmId {
    match *obj {
        OsmObj::Node(ref node) => OsmId::Node(node.id),
        OsmObj::Way(ref way) => OsmId::Way(way.id),
        OsmObj::Relation(ref rel) => OsmId::Relation(rel.id),
    }
}

fn update_coordinates(filename: &String, admins: &mut AdminsMap) {
    if admins.is_empty() {
        return;
    }
    // load coord for administratives regions
    let path = std::path::Path::new(&filename);
    let r = std::fs::File::open(&path).unwrap();
    let mut pbf = osmpbfreader::OsmPbfReader::new(r);
    for obj in pbf.iter() {
        if let osmpbfreader::OsmObj::Node(ref node) = obj {
            let mut adm = match admins.get_mut(&get_osm_id(&obj)) {
                Some(val) => val,
                None => continue,
            };
            adm.coord.lat = node.lat;
            adm.coord.lon = node.lon;
        }
    }
}

fn administartive_regions(filename: &String, levels: &HashSet<u32>) -> AdminsMap {
    let mut administrative_regions = AdminsMap::new();
    let path = std::path::Path::new(&filename);
    let r = std::fs::File::open(&path).unwrap();
    let mut pbf = osmpbfreader::OsmPbfReader::new(r);
    // load administratives regions
    for obj in pbf.iter() {
        if let osmpbfreader::OsmObj::Relation(relation) = obj {
            // not administartive region
            if !relation.tags
                        .get("boundary")
                        .map(|s| s == "administrative")
                        .unwrap_or(false) {
                continue;
            }
            // admininstrative region without coordinates
            let admin_centre = match relation.refs.iter().find(|rf| rf.role == "admin_centre") {
                Some(val) => val.member,
                None => {
                    info!("adminstrative regione without coordinates for relation {}.",
                          relation.id);
                    continue;
                }
            };
            let level = relation.tags
                                .get("admin_level")
                                .and_then(|s| s.parse().ok());
            let level = match level {
                None => {
                    info!("invalid admin_level for relation {}: admin_level {:?}",
                          relation.id,
                          relation.tags.get("admin_level"));
                    continue;
                }
                Some(ref l) if !levels.contains(&l) => continue,
                Some(l) => l,
            };
            // administrative region with name ?
            let name = match relation.tags.get("name") {
                Some(val) => val,
                None => {
                    info!("adminstrative region without name for relation {}:  admin_level {} \
                           ignored.",
                          relation.id,
                          level);
                    continue;
                }
            };
            let admin_id = match relation.tags.get("ref:INSEE") {
                Some(val) => format!("admin:fr:{}", val),
                None => format!("admin:osm:{}", relation.id),
            };
            let zip_code = match relation.tags.get("addr:postcode") {
                Some(val) => &val[..],
                None => "",
            };
            let admin = mimirsbrunn::Admin {
                id: admin_id,
                level: level,
                name: name.to_string(),
                zip_code: zip_code.to_string(),
                // TODO weight value ?
                weight: 1,
                coord: mimirsbrunn::Coord {
                    lat: 0.0,
                    lon: 0.0,
                },
            };
            administrative_regions.insert(admin_centre, admin);
        }
    }
    return administrative_regions;
}

fn main() {
    env_logger::init().unwrap();
    debug!("importing adminstrative region into Mimir");
    let args: Args = docopt::Docopt::new(USAGE)
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());
    let levels = args.flag_level.iter().cloned().collect();
    let mut res = administartive_regions(&args.flag_input, &levels);
    update_coordinates(&args.flag_input, &mut res);
}
<commit_msg>mimirsbrunn: license added<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

#[macro_use]
extern crate log;
extern crate env_logger;
extern crate osmpbfreader;
extern crate rustc_serialize;
extern crate docopt;
extern crate mimirsbrunn;

use std::collections::HashSet;
use std::collections::BTreeMap;
use osmpbfreader::{OsmObj, OsmId};

pub type AdminsMap = BTreeMap<OsmId, mimirsbrunn::Admin>;

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_input: String,
    flag_level: Vec<u32>,
    flag_connection_string: String,
}

static USAGE: &'static str = "
Usage:
    osm2mimir --help
    osm2mimir --input=<file> [--connection-string=<connection-string>] --level=<level>...

Options:
    -h, --help            Show this message.
    -i, --input=<file>    OSM PBF file.
    -l, --level=<level>   Admin levels to keep.
    -c, --connection-string=<connection-string>
                          Elasticsearch parameters, [default: http://localhost:9200/munin]
";

fn get_osm_id(obj: &OsmObj) -> OsmId {
    match *obj {
        OsmObj::Node(ref node) => OsmId::Node(node.id),
        OsmObj::Way(ref way) => OsmId::Way(way.id),
        OsmObj::Relation(ref rel) => OsmId::Relation(rel.id),
    }
}

fn update_coordinates(filename: &String, admins: &mut AdminsMap) {
    if admins.is_empty() {
        return;
    }
    // load coord for administratives regions
    let path = std::path::Path::new(&filename);
    let r = std::fs::File::open(&path).unwrap();
    let mut pbf = osmpbfreader::OsmPbfReader::new(r);
    for obj in pbf.iter() {
        if let osmpbfreader::OsmObj::Node(ref node) = obj {
            let mut adm = match admins.get_mut(&get_osm_id(&obj)) {
                Some(val) => val,
                None => continue,
            };
            adm.coord.lat = node.lat;
            adm.coord.lon = node.lon;
        }
    }
}

fn administartive_regions(filename: &String, levels: &HashSet<u32>) -> AdminsMap {
    let mut administrative_regions = AdminsMap::new();
    let path = std::path::Path::new(&filename);
    let r = std::fs::File::open(&path).unwrap();
    let mut pbf = osmpbfreader::OsmPbfReader::new(r);
    // load administratives regions
    for obj in pbf.iter() {
        if let osmpbfreader::OsmObj::Relation(relation) = obj {
            // not administartive region
            if !relation.tags
                        .get("boundary")
                        .map(|s| s == "administrative")
                        .unwrap_or(false) {
                continue;
            }
            // admininstrative region without coordinates
            let admin_centre = match relation.refs.iter().find(|rf| rf.role == "admin_centre") {
                Some(val) => val.member,
                None => {
                    info!("adminstrative regione without coordinates for relation {}.",
                          relation.id);
                    continue;
                }
            };
            let level = relation.tags
                                .get("admin_level")
                                .and_then(|s| s.parse().ok());
            let level = match level {
                None => {
                    info!("invalid admin_level for relation {}: admin_level {:?}",
                          relation.id,
                          relation.tags.get("admin_level"));
                    continue;
                }
                Some(ref l) if !levels.contains(&l) => continue,
                Some(l) => l,
            };
            // administrative region with name ?
            let name = match relation.tags.get("name") {
                Some(val) => val,
                None => {
                    info!("adminstrative region without name for relation {}:  admin_level {} \
                           ignored.",
                          relation.id,
                          level);
                    continue;
                }
            };
            let admin_id = match relation.tags.get("ref:INSEE") {
                Some(val) => format!("admin:fr:{}", val),
                None => format!("admin:osm:{}", relation.id),
            };
            let zip_code = match relation.tags.get("addr:postcode") {
                Some(val) => &val[..],
                None => "",
            };
            let admin = mimirsbrunn::Admin {
                id: admin_id,
                level: level,
                name: name.to_string(),
                zip_code: zip_code.to_string(),
                // TODO weight value ?
                weight: 1,
                coord: mimirsbrunn::Coord {
                    lat: 0.0,
                    lon: 0.0,
                },
            };
            administrative_regions.insert(admin_centre, admin);
        }
    }
    return administrative_regions;
}

fn main() {
    env_logger::init().unwrap();
    debug!("importing adminstrative region into Mimir");
    let args: Args = docopt::Docopt::new(USAGE)
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());
    let levels = args.flag_level.iter().cloned().collect();
    let mut res = administartive_regions(&args.flag_input, &levels);
    update_coordinates(&args.flag_input, &mut res);
}
<|endoftext|>


<commit_before>//! `cargo add`

#![deny(missing_docs, missing_debug_implementations, missing_copy_implementations, trivial_casts, trivial_numeric_casts, unsafe_code, unstable_features, unused_import_braces, unused_qualifications)]
#![cfg_attr(feature = "dev", allow(unstable_features))]
#![cfg_attr(feature = "dev", feature(plugin))]
#![cfg_attr(feature = "dev", plugin(clippy))]

extern crate docopt;
extern crate toml;
extern crate semver;
extern crate rustc_serialize;
extern crate curl;
#[macro_use]
extern crate quick_error;

use std::error::Error;
use std::io::{self, Write};
use std::process;

extern crate cargo_edit;
use cargo_edit::Manifest;

extern crate regex;

mod fetch;
mod args;
use args::Args;

static USAGE: &'static str = r#"
Usage:
    cargo add <crate> [--dev|--build|--optional] [--vers=<ver>|--git=<uri>|--path=<uri>] [options]
    cargo add <crates>... [--dev|--build|--optional] [options]
    cargo add (-h|--help)
    cargo add --version

Specify what crate to add:
    --vers <ver>            Specify the version to grab from the registry (crates.io).
                            You can also specify versions as part of the name, e.g
                            `cargo add bitflags@0.3.2`.
    --git <uri>             Specify a git repository to download the crate from.
    --path <uri>            Specify the path the crate should be loaded from.

Specify where to add the crate:
    -D --dev                Add crate as development dependency.
    -B --build              Add crate as build dependency.
    --optional              Add as an optional dependency (for use in features). This does not work
                            for `dev-dependencies` or `build-dependencies`.
    --target <target>       Add as dependency to the given target platform. This does not work
                            for `dev-dependencies` or `build-dependencies`.

Options:
    --upgrade=<method>      Choose method of semantic version upgrade. Must be one of
                            "none" (exact version), "patch" (`~` modifier), "minor"
                            (`^` modifier, default), or "all" (`>=`).
    --manifest-path=<path>  Path to the manifest to add a dependency to.
    --allow-prerelease      Include prerelease versions when fetching from crates.io (e.g.
                            '0.6.0-alpha'). Defaults to false.
    -h --help               Show this help page.
    --version               Show version.

This command allows you to add a dependency to a Cargo.toml manifest file. If <crate> is a github
or gitlab repository URL, or a local path, `cargo add` will try to automatically get the crate name
and set the appropriate `--git` or `--path` value.

Please note that Cargo treats versions like "1.2.3" as "^1.2.3" (and that "^1.2.3" is specified
as ">=1.2.3 and <2.0.0"). By default, `cargo add` will use this format, as it is the one that the
crates.io registry suggests. One goal of `cargo add` is to prevent you from using wildcard
dependencies (version set to "*").
"#;

fn handle_add(args: &Args) -> Result<(), Box<Error>> {
    let mut manifest = try!(Manifest::open(&args.flag_manifest_path.as_ref().map(|s| &s[..])));
    let deps = try!(args.parse_dependencies());

    for dep in deps {
        if let Err(err) = manifest.insert_into_table(&args.get_section(), &dep) {
            println!("Could not edit `Cargo.toml`.\n\nERROR: {}", err);
            return Err(From::from(err));
        }
    }

    let mut file = try!(Manifest::find_file(&args.flag_manifest_path.as_ref().map(|s| &s[..])));
    manifest.write_to_file(&mut file)
}

fn main() {
    let args = docopt::Docopt::new(USAGE)
        .and_then(|d| d.decode::<Args>())
        .unwrap_or_else(|err| err.exit());

    if args.flag_version {
        println!("cargo-add version {}", env!("CARGO_PKG_VERSION"));
        process::exit(0);
    }

    if let Err(err) = handle_add(&args) {
        writeln!(io::stderr(),
                 "Command failed due to unhandled error: {}",
                 err)
            .unwrap();
        process::exit(1);
    }
}
<commit_msg>-V is often short for --version<commit_after>//! `cargo add`

#![deny(missing_docs, missing_debug_implementations, missing_copy_implementations, trivial_casts, trivial_numeric_casts, unsafe_code, unstable_features, unused_import_braces, unused_qualifications)]
#![cfg_attr(feature = "dev", allow(unstable_features))]
#![cfg_attr(feature = "dev", feature(plugin))]
#![cfg_attr(feature = "dev", plugin(clippy))]

extern crate docopt;
extern crate toml;
extern crate semver;
extern crate rustc_serialize;
extern crate curl;
#[macro_use]
extern crate quick_error;

use std::error::Error;
use std::io::{self, Write};
use std::process;

extern crate cargo_edit;
use cargo_edit::Manifest;

extern crate regex;

mod fetch;
mod args;
use args::Args;

static USAGE: &'static str = r#"
Usage:
    cargo add <crate> [--dev|--build|--optional] [--vers=<ver>|--git=<uri>|--path=<uri>] [options]
    cargo add <crates>... [--dev|--build|--optional] [options]
    cargo add (-h|--help)
    cargo add --version

Specify what crate to add:
    --vers <ver>            Specify the version to grab from the registry (crates.io).
                            You can also specify versions as part of the name, e.g
                            `cargo add bitflags@0.3.2`.
    --git <uri>             Specify a git repository to download the crate from.
    --path <uri>            Specify the path the crate should be loaded from.

Specify where to add the crate:
    -D --dev                Add crate as development dependency.
    -B --build              Add crate as build dependency.
    --optional              Add as an optional dependency (for use in features). This does not work
                            for `dev-dependencies` or `build-dependencies`.
    --target <target>       Add as dependency to the given target platform. This does not work
                            for `dev-dependencies` or `build-dependencies`.

Options:
    --upgrade=<method>      Choose method of semantic version upgrade. Must be one of
                            "none" (exact version), "patch" (`~` modifier), "minor"
                            (`^` modifier, default), or "all" (`>=`).
    --manifest-path=<path>  Path to the manifest to add a dependency to.
    --allow-prerelease      Include prerelease versions when fetching from crates.io (e.g.
                            '0.6.0-alpha'). Defaults to false.
    -h --help               Show this help page.
    -V --version            Show version.

This command allows you to add a dependency to a Cargo.toml manifest file. If <crate> is a github
or gitlab repository URL, or a local path, `cargo add` will try to automatically get the crate name
and set the appropriate `--git` or `--path` value.

Please note that Cargo treats versions like "1.2.3" as "^1.2.3" (and that "^1.2.3" is specified
as ">=1.2.3 and <2.0.0"). By default, `cargo add` will use this format, as it is the one that the
crates.io registry suggests. One goal of `cargo add` is to prevent you from using wildcard
dependencies (version set to "*").
"#;

fn handle_add(args: &Args) -> Result<(), Box<Error>> {
    let mut manifest = try!(Manifest::open(&args.flag_manifest_path.as_ref().map(|s| &s[..])));
    let deps = try!(args.parse_dependencies());

    for dep in deps {
        if let Err(err) = manifest.insert_into_table(&args.get_section(), &dep) {
            println!("Could not edit `Cargo.toml`.\n\nERROR: {}", err);
            return Err(From::from(err));
        }
    }

    let mut file = try!(Manifest::find_file(&args.flag_manifest_path.as_ref().map(|s| &s[..])));
    manifest.write_to_file(&mut file)
}

fn main() {
    let args = docopt::Docopt::new(USAGE)
        .and_then(|d| d.decode::<Args>())
        .unwrap_or_else(|err| err.exit());

    if args.flag_version {
        println!("cargo-add version {}", env!("CARGO_PKG_VERSION"));
        process::exit(0);
    }

    if let Err(err) = handle_add(&args) {
        writeln!(io::stderr(),
                 "Command failed due to unhandled error: {}",
                 err)
            .unwrap();
        process::exit(1);
    }
}
<|endoftext|>


<commit_before>// advent16.rs
// dragon curve

use std::io;

type BV = Vec<bool>;

fn main() {
    let mut input = String::new();
    io::stdin().read_line(&mut input).expect("Failed to read line");

    let mut rd = RandomData::from_str(&input);
    // part 1
    rd.grow(272);
    let checksum1 = string_from_bv(&rd.calc_checksum());
    println!("part 1 checksum: {}", checksum1);

    // part 2
    rd.grow(35651584);
    let checksum2 = string_from_bv(&rd.calc_checksum());
    println!("part 2 checksum: {}", checksum2);

}

// ///////
// Part 1

struct RandomData {
    bv: BV,
    pos: usize,
}

impl RandomData {
    fn from_str(s: &str) -> RandomData {
        RandomData {
            pos: 0,
            bv: s.chars()
                .filter_map(|c| match c {
                    '0' => Some(false),
                    '1' => Some(true),
                    _ => None,
                })
                .collect(),
        }
    }

    fn grow(&mut self, new_len: usize) {
        let old_len = self.bv.len();
        if new_len > old_len {
            self.bv.reserve(new_len - old_len);
            for _ in old_len..new_len {
                if self.pos == 0 {
                    self.pos = self.bv.len();
                    self.bv.push(false);
                } else {
                    self.pos -= 1;
                    let new_val = !self.bv[self.pos];
                    self.bv.push(new_val);
                }
            }
        }
    }

    fn calc_checksum(&self) -> BV {
        let disk_len = self.bv.len();
        let trailing = count_trailing_zeros(disk_len);

        if trailing == 0 {
            self.bv.clone()
        } else {
            let chunk_size = 1 << trailing;

            // checksum of each chunk is true iff the number of true elements is even
            self.bv
                .chunks(chunk_size)
                .map(|chunk| chunk.iter().filter(|&&x| x).count() % 2 == 0)
                .collect()
        }
    }
}

fn count_trailing_zeros(x: usize) -> u32 {
    // Convert trailing 0's to 1's and set everything else to 0. Then just count 1's
    // http://graphics.stanford.edu/~seander/bithacks.html#ZerosOnRightLinear
    ((x ^ (x - 1)) >> 1).count_ones()
}

#[allow(ptr_arg)]
fn string_from_bv(bv: &BV) -> String {
    bv.iter().map(|&x| if x { '1' } else { '0' }).collect()
}

// //////
// Tests
#[cfg(test)]
mod tests {
    use super::RandomData;
    use super::count_trailing_zeros;
    use super::string_from_bv;

    #[test]
    fn test_rd_from_str() {
        let bv = RandomData::from_str("100").bv;
        assert_eq!(3, bv.len());
        assert_eq!(true, bv[0]);
        assert_eq!(false, bv[1]);
        assert_eq!(false, bv[2]);
    }

    #[test]
    fn test_rd_grow() {
        does_rd_grow_match("1", "100");
        does_rd_grow_match("0", "001");
        does_rd_grow_match("11111", "11111000000");
        does_rd_grow_match("111100001010", "1111000010100101011110000");
        // check a few more iterations too
        does_rd_grow_match("0", "001001100011011");
    }

    fn does_rd_grow_match(input: &str, expect: &str) {
        let len = expect.len();
        let mut rd_in = RandomData::from_str(input);
        rd_in.grow(len);
        assert_eq!(&string_from_bv(&rd_in.bv), expect);
    }

    #[test]
    fn test_count_trailing_zeros() {
        assert_eq!(0, count_trailing_zeros(0b1));
        assert_eq!(1, count_trailing_zeros(0b1110));
        assert_eq!(3, count_trailing_zeros(0b10101000));
    }

    #[test]
    fn test_rd_checksum() {
        does_rd_checksum_match("110010110100", 12, "100");
        does_rd_checksum_match("10000", 20, "01100");
    }

    fn does_rd_checksum_match(input: &str, len: usize, expect: &str) {
        let mut rd_in = RandomData::from_str(input);
        rd_in.grow(len);
        assert_eq!(&string_from_bv(&rd_in.calc_checksum()), expect);
    }
}
<commit_msg>Suppress warning about unknown lint<commit_after>// advent16.rs
// dragon curve

use std::io;

type BV = Vec<bool>;

fn main() {
    let mut input = String::new();
    io::stdin().read_line(&mut input).expect("Failed to read line");

    let mut rd = RandomData::from_str(&input);
    // part 1
    rd.grow(272);
    let checksum1 = string_from_bv(&rd.calc_checksum());
    println!("part 1 checksum: {}", checksum1);

    // part 2
    rd.grow(35651584);
    let checksum2 = string_from_bv(&rd.calc_checksum());
    println!("part 2 checksum: {}", checksum2);

}

// ///////
// Part 1

struct RandomData {
    bv: BV,
    pos: usize,
}

impl RandomData {
    fn from_str(s: &str) -> RandomData {
        RandomData {
            pos: 0,
            bv: s.chars()
                .filter_map(|c| match c {
                    '0' => Some(false),
                    '1' => Some(true),
                    _ => None,
                })
                .collect(),
        }
    }

    fn grow(&mut self, new_len: usize) {
        let old_len = self.bv.len();
        if new_len > old_len {
            self.bv.reserve(new_len - old_len);
            for _ in old_len..new_len {
                if self.pos == 0 {
                    self.pos = self.bv.len();
                    self.bv.push(false);
                } else {
                    self.pos -= 1;
                    let new_val = !self.bv[self.pos];
                    self.bv.push(new_val);
                }
            }
        }
    }

    fn calc_checksum(&self) -> BV {
        let disk_len = self.bv.len();
        let trailing = count_trailing_zeros(disk_len);

        if trailing == 0 {
            self.bv.clone()
        } else {
            let chunk_size = 1 << trailing;

            // checksum of each chunk is true iff the number of true elements is even
            self.bv
                .chunks(chunk_size)
                .map(|chunk| chunk.iter().filter(|&&x| x).count() % 2 == 0)
                .collect()
        }
    }
}

fn count_trailing_zeros(x: usize) -> u32 {
    // Convert trailing 0's to 1's and set everything else to 0. Then just count 1's
    // http://graphics.stanford.edu/~seander/bithacks.html#ZerosOnRightLinear
    ((x ^ (x - 1)) >> 1).count_ones()
}

#[allow(unknown_lints)]
#[allow(ptr_arg)]
fn string_from_bv(bv: &BV) -> String {
    bv.iter().map(|&x| if x { '1' } else { '0' }).collect()
}

// //////
// Tests
#[cfg(test)]
mod tests {
    use super::RandomData;
    use super::count_trailing_zeros;
    use super::string_from_bv;

    #[test]
    fn test_rd_from_str() {
        let bv = RandomData::from_str("100").bv;
        assert_eq!(3, bv.len());
        assert_eq!(true, bv[0]);
        assert_eq!(false, bv[1]);
        assert_eq!(false, bv[2]);
    }

    #[test]
    fn test_rd_grow() {
        does_rd_grow_match("1", "100");
        does_rd_grow_match("0", "001");
        does_rd_grow_match("11111", "11111000000");
        does_rd_grow_match("111100001010", "1111000010100101011110000");
        // check a few more iterations too
        does_rd_grow_match("0", "001001100011011");
    }

    fn does_rd_grow_match(input: &str, expect: &str) {
        let len = expect.len();
        let mut rd_in = RandomData::from_str(input);
        rd_in.grow(len);
        assert_eq!(&string_from_bv(&rd_in.bv), expect);
    }

    #[test]
    fn test_count_trailing_zeros() {
        assert_eq!(0, count_trailing_zeros(0b1));
        assert_eq!(1, count_trailing_zeros(0b1110));
        assert_eq!(3, count_trailing_zeros(0b10101000));
    }

    #[test]
    fn test_rd_checksum() {
        does_rd_checksum_match("110010110100", 12, "100");
        does_rd_checksum_match("10000", 20, "01100");
    }

    fn does_rd_checksum_match(input: &str, len: usize, expect: &str) {
        let mut rd_in = RandomData::from_str(input);
        rd_in.grow(len);
        assert_eq!(&string_from_bv(&rd_in.calc_checksum()), expect);
    }
}
<|endoftext|>