<commit_before>use std::env;
use std::fs::read_link;
use std::io::{stdout, stderr};
use std::path::Path;
use std::str::FromStr;

use libc::pid_t;
use argparse::{ArgumentParser};
use unshare::{Command};

use config::command::CommandInfo;
use config::command::WriteMode;

use super::setup;
use super::Wrapper;
use super::util::find_cmd;
use process_util::{run_and_wait, convert_status};
use process_util::{set_fake_uidmap};

pub fn commandline_cmd(command: &CommandInfo,
    wrapper: &Wrapper, mut cmdline: Vec<String>)
    -> Result<i32, String>
{
    if command.run.len() == 0 {
        return Err(format!(
            r#"Command has empty "run" parameter. Nothing to run."#));
    }
    // TODO(tailhook) detect other shells too
    let has_args = command.accepts_arguments
            .unwrap_or(&command.run[0][..] != "/bin/sh");
    let mut args = Vec::new();
    if !has_args {
        let mut ap = ArgumentParser::new();
        ap.set_description(command.description.as_ref()
            .map(|x| &x[..]).unwrap_or(""));
        ap.stop_on_first_argument(true);
        match ap.parse(cmdline, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(122);
            }
        }
    } else {
        cmdline.remove(0);
        args.extend(cmdline.into_iter());
    }
    let mut cmdline = command.run.clone();
    cmdline.extend(args.into_iter());

    let pid: pid_t = try!(read_link(&Path::new("/proc/self"))
        .map_err(|e| format!("Can't read /proc/self: {}", e))
        .and_then(|v| v.to_str().and_then(|x| FromStr::from_str(x).ok())
            .ok_or(format!("Can't parse pid: {:?}", v))));
    try!(setup::setup_base_filesystem(
        wrapper.project_root, wrapper.ext_settings));

    let cconfig = try!(wrapper.config.containers.get(&command.container)
        .ok_or(format!("Container {} not found", command.container)));

    let write_mode = match command.write_mode {
        WriteMode::read_only => setup::WriteMode::ReadOnly,
        WriteMode::transient_hard_link_copy
        => setup::WriteMode::TransientHardlinkCopy(pid),
    };
    let cont_ver = wrapper.root.as_ref().unwrap();
    let mut setup_info = setup::SetupInfo::from_container(&cconfig);
    setup_info
        .volumes(&command.volumes)
        .write_mode(write_mode);
    try!(setup::setup_filesystem(&setup_info, &cont_ver));

    let mut env = try!(setup::get_environment(cconfig, &wrapper.settings));
    for (k, v) in command.environ.iter() {
        env.insert(k.clone(), v.clone());
    }
    let cpath = try!(find_cmd(&cmdline.remove(0), &env));

    let mut cmd = Command::new(&cpath);
    cmd.args(&cmdline);
    if let Some(euid) = command.external_user_id {
        try!(set_fake_uidmap(&mut cmd, command.user_id, euid));
    }
    cmd.uid(command.user_id);
    cmd.gid(command.group_id);
    cmd.groups(command.supplementary_gids.clone());
    if let Some(ref wd) = command.work_dir {
        cmd.current_dir(Path::new("/work").join(&wd));
    } else {
        cmd.current_dir(env::var("_VAGGA_WORKDIR")
                        .unwrap_or("/work".to_string()));
    }
    for (ref k, ref v) in env.iter() {
        cmd.env(k, v);
    }

    run_and_wait(&mut cmd)
    .map(convert_status)
}
<commit_msg>Correct propagation of environment vars in non-supervise command<commit_after>use std::env;
use std::fs::read_link;
use std::io::{stdout, stderr};
use std::path::Path;
use std::str::FromStr;

use libc::pid_t;
use argparse::{ArgumentParser};
use unshare::{Command};

use config::command::CommandInfo;
use config::command::WriteMode;

use super::setup;
use super::Wrapper;
use super::util::find_cmd;
use process_util::{run_and_wait, convert_status, copy_env_vars};
use process_util::{set_fake_uidmap};

pub fn commandline_cmd(command: &CommandInfo,
    wrapper: &Wrapper, mut cmdline: Vec<String>)
    -> Result<i32, String>
{
    if command.run.len() == 0 {
        return Err(format!(
            r#"Command has empty "run" parameter. Nothing to run."#));
    }
    // TODO(tailhook) detect other shells too
    let has_args = command.accepts_arguments
            .unwrap_or(&command.run[0][..] != "/bin/sh");
    let mut args = Vec::new();
    if !has_args {
        let mut ap = ArgumentParser::new();
        ap.set_description(command.description.as_ref()
            .map(|x| &x[..]).unwrap_or(""));
        ap.stop_on_first_argument(true);
        match ap.parse(cmdline, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(122);
            }
        }
    } else {
        cmdline.remove(0);
        args.extend(cmdline.into_iter());
    }
    let mut cmdline = command.run.clone();
    cmdline.extend(args.into_iter());

    let pid: pid_t = try!(read_link(&Path::new("/proc/self"))
        .map_err(|e| format!("Can't read /proc/self: {}", e))
        .and_then(|v| v.to_str().and_then(|x| FromStr::from_str(x).ok())
            .ok_or(format!("Can't parse pid: {:?}", v))));
    try!(setup::setup_base_filesystem(
        wrapper.project_root, wrapper.ext_settings));

    let cconfig = try!(wrapper.config.containers.get(&command.container)
        .ok_or(format!("Container {} not found", command.container)));

    let write_mode = match command.write_mode {
        WriteMode::read_only => setup::WriteMode::ReadOnly,
        WriteMode::transient_hard_link_copy
        => setup::WriteMode::TransientHardlinkCopy(pid),
    };
    let cont_ver = wrapper.root.as_ref().unwrap();
    let mut setup_info = setup::SetupInfo::from_container(&cconfig);
    setup_info
        .volumes(&command.volumes)
        .write_mode(write_mode);
    try!(setup::setup_filesystem(&setup_info, &cont_ver));

    let mut env = try!(setup::get_environment(cconfig, &wrapper.settings));
    for (k, v) in command.environ.iter() {
        env.insert(k.clone(), v.clone());
    }
    let cpath = try!(find_cmd(&cmdline.remove(0), &env));

    let mut cmd = Command::new(&cpath);
    cmd.args(&cmdline);
    cmd.env_clear();
    copy_env_vars(&mut cmd, &wrapper.settings);
    if let Some(euid) = command.external_user_id {
        try!(set_fake_uidmap(&mut cmd, command.user_id, euid));
    }
    cmd.uid(command.user_id);
    cmd.gid(command.group_id);
    cmd.groups(command.supplementary_gids.clone());
    if let Some(ref wd) = command.work_dir {
        cmd.current_dir(Path::new("/work").join(&wd));
    } else {
        cmd.current_dir(env::var("_VAGGA_WORKDIR")
                        .unwrap_or("/work".to_string()));
    }
    for (ref k, ref v) in env.iter() {
        cmd.env(k, v);
    }

    run_and_wait(&mut cmd)
    .map(convert_status)
}
<|endoftext|>


<commit_before>
use std::str::CharIndices;
use std::iter::Peekable;

#[derive(PartialEq, Debug)]
pub enum Token<'a> {
  Eof(usize, usize),
  Bang(usize, usize),
  Dollar(usize, usize),
  LeftParen(usize, usize),
  RightParen(usize, usize),
  Spread(usize, usize),
  Colon(usize, usize),
  Equals(usize, usize),
  At(usize, usize),
  LeftBracket(usize, usize),
  RightBracket(usize, usize),
  LeftBrace(usize, usize),
  RightBrace(usize, usize),
  Pipe(usize, usize),
  Name(usize, usize, &'a str),
  Int(usize, usize, &'a str),
  Float(usize, usize, &'a str),
  String(usize, usize, &'a str),
}

macro_rules! take {
  ($lexer: ident, $($p: pat)|*) => ({
    let mut is_taken = false;
    if let Some(&(p, c)) = $lexer.iter.peek() {
      match c {
        $(
          $p => {
            $lexer.iter.next();
            $lexer.hi = p;
            is_taken = true;
          }
        )*
        _ => {}
      }
    }
    is_taken
  });
}

macro_rules! take_while {
  ($lexer: ident, $($p: pat)|*) => ({
    let mut is_taken = false;
    while let Some(&(p, c)) = $lexer.iter.peek() {
      match c {
        $(
          $p => {
            $lexer.iter.next();
            is_taken = true;
          }
        )*
        _ => {
          $lexer.hi = p;
          break;
        }
      }
    }
    is_taken
  });
}

macro_rules! take_while_not {
  ($lexer: ident, $($p: pat)|*) => ({
    let mut is_taken = false;
    while let Some(&(p, c)) = $lexer.iter.peek() {
      match c {
        $(
          $p => {
            $lexer.hi = p;
            break;
          }
        )*
        _ => {
            $lexer.iter.next();
            is_taken = true;
        }
      }
    }
    is_taken
  });
}

macro_rules! peek {
  ($lexer: ident, $($p: pat)|*) => ({
    let mut is_match = false;
    if let Some(&(_, c)) = $lexer.iter.peek() {
      match c {
        $(
          $p => {
            is_match = true;
          }
        )*
        _ => {}
      }
    }
    is_match
  });
}

macro_rules! take_eof {
  ($lexer: ident) => ({
    if $lexer.iter.peek().is_none() {
      $lexer.hi = $lexer.input_len;
      true
    } else {
      false
    }
  });
}

pub struct Lexer<'a> {
  input: &'a str,
  input_len: usize,
  iter: Peekable<CharIndices<'a>>,
  is_eof: bool,
  lo: usize,
  hi: usize,
}

impl<'a> Lexer<'a> {
  pub fn new(input: &'a str) -> Lexer<'a> {
    Lexer {
      input: input,
      input_len: input.len(),
      iter: input.char_indices().peekable(),
      is_eof: false,
      lo: 0,
      hi: 0,
    }
  }
}

fn skip_ignored_token(lexer: &mut Lexer) {
  loop {
    if take!(lexer, '#') {
      take_while_not!(lexer, '\r' | '\n');
    }
    if !take_while!(lexer, '\u{feff}' | ' ' | '\t' | '\r' | '\n' | ',') {
      break;
    }
  }
}

fn scan_spread(lexer: &mut Lexer) -> bool {
  if take!(lexer, '.') {
    if take!(lexer, '.') {
      if take!(lexer, '.') {
        return true;
      }
    }
  }
  return false;
}

fn scan_name(lexer: &mut Lexer) {
  take_while!(lexer, '_' | 'a'...'z' | 'A'...'Z' | '0'...'9');
  take_eof!(lexer);
}

// Return the raw string without unescape character
fn scan_string(lexer: &mut Lexer) -> (bool, bool) {
  let mut is_bad_escape = false;
  let mut is_terminated = false;
  if take!(lexer, '"') {
    lexer.lo = lexer.lo + 1;
  }
  loop {
    if take!(lexer, '\\') {
      if !take!(lexer, '\\' | '"' | 'b' | 'f' | 'n' | 'r' | 't' | '/' | 'u') {
        is_bad_escape = true;
        break;
      }
    }
    take_while_not!(lexer, '"' | '\\');
    if take!(lexer, '"') {
      is_terminated = true;
      break;
    }
  }
  (is_terminated, is_bad_escape)
}

// Int:   (+|-)?(0|[1-9][0-9]*)
// Float: (+|-)?(0|[1-9][0-9]*)(\.[0-9]+)?((E|e)(+|-)?[0-9]+)?
fn scan_number(lexer: &mut Lexer) -> (bool, bool) {
  let mut is_float = false;
  take!(lexer, '+' | '-');
  // Integer part
  if take!(lexer, '1'...'9') {
    take_while!(lexer, '0'...'9');
    if take!(lexer, ' ' | '\t' | '\r' | '\n' | ',') || take_eof!(lexer) {
      return (true, is_float);
    }
  } else if take!(lexer, '0') {
    if take!(lexer, ' ' | '\t' | '\r' | '\n' | ',') || take_eof!(lexer) {
      return (true, is_float);
    }
  }
  // Fractional part
  if take!(lexer, '.') {
    if !take_while!(lexer, '0'...'9') {
      return (false, is_float);
    }
    is_float = true;
  }
  // Exponent part
  if !is_float && peek!(lexer, 'E' | 'e') {
    return (false, is_float);
  }
  if take!(lexer, 'E' | 'e') {
    if take!(lexer, '+' | '-') {
      if !take_while!(lexer, '0'...'9') {
        return (false, is_float);
      }
    } else {
      take_while!(lexer, '0'...'9');
    }
  }

  if take!(lexer, ' ' | '\t' | '\r' | '\n' | ',') || take_eof!(lexer) {
    return (true, is_float);
  }
  return (false, is_float);
}

impl<'a> Iterator for Lexer<'a> {
  type Item = Result<Token<'a>, &'a str>;
  fn next(&mut self) -> Option<Result<Token<'a>, &'a str>> {
    if !self.is_eof && take_eof!(self) {
      self.is_eof = true;
      return Some(Ok(Token::Eof(self.lo, self.hi)));
    }
    if self.is_eof {
      return None;
    }
    skip_ignored_token(self);

    let mut item = None;
    if let Some(&(p, c)) = self.iter.peek() {
      self.lo = p;
      self.hi = p;
      item = match c {
        '!' => {
          self.iter.next();
          Some(Ok(Token::Bang(self.lo, self.hi)))
        }
        '$' => {
          self.iter.next();
          Some(Ok(Token::Dollar(self.lo, self.hi)))
        }
        '(' => {
          self.iter.next();
          Some(Ok(Token::LeftParen(self.lo, self.hi)))
        }
        ')' => {
          self.iter.next();
          Some(Ok(Token::RightParen(self.lo, self.hi)))
        }
        ':' => {
          self.iter.next();
          Some(Ok(Token::Colon(self.lo, self.hi)))
        }
        '=' => {
          self.iter.next();
          Some(Ok(Token::Equals(self.lo, self.hi)))
        }
        '@' => {
          self.iter.next();
          Some(Ok(Token::At(self.lo, self.hi)))
        }
        '[' => {
          self.iter.next();
          Some(Ok(Token::LeftBracket(self.lo, self.hi)))
        }
        ']' => {
          self.iter.next();
          Some(Ok(Token::RightBracket(self.lo, self.hi)))
        }
        '{' => {
          self.iter.next();
          Some(Ok(Token::LeftBrace(self.lo, self.hi)))
        }
        '}' => {
          self.iter.next();
          Some(Ok(Token::RightBrace(self.lo, self.hi)))
        }
        '|' => {
          self.iter.next();
          Some(Ok(Token::Pipe(self.lo, self.hi)))
        }
        '.' => {
          if scan_spread(self) {
            Some(Ok(Token::Spread(self.lo, self.hi)))
          } else {
            Some(Err("Unexpected character."))
          }
        }
        '_' | 'a'...'z' | 'A'...'Z' => {
          scan_name(self);
          Some(Ok(Token::Name(self.lo, self.hi, &self.input[self.lo..self.hi])))
        }
        '-' | '0'...'9' => {
          match scan_number(self) {
            (false, false) => Some(Err("Invalid integer number.")),
            (false, true) => Some(Err("Invalid float number.")),
            (true, false) => Some(Ok(Token::Int(self.lo, self.hi, &self.input[self.lo..self.hi]))),
            (true, true) => Some(Ok(Token::Float(self.lo, self.hi, &self.input[self.lo..self.hi]))),
          }
        }
        '"' => {
          match scan_string(self) {
            (false, _) => Some(Err("Unterminated string.")),
            (_, true) => Some(Err("Bad character escape sequence.")),
            _ => Some(Ok(Token::String(self.lo, self.hi, &self.input[self.lo..self.hi]))),
          }
        }
        _ => Some(Err("Unexpected character.")),
      }
    }
    item
  }
}

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

  fn test_next_token(lexer: &mut Lexer, expected: Token) {
    let next = lexer.next();
    assert_eq!(Some(Ok(expected)), next);
  }

  fn test_with_lex(input: &str, expeced: Token) {
    let mut lexer = Lexer::new(&input);
    test_next_token(&mut lexer, expeced);
  }

  fn test_str(input: &str, expeced: &str) {
    let len = expeced.len();
    test_with_lex(input, Token::String(1, len + 1, expeced))
  }

  #[test]
  fn test_string() {
    test_str(r#""simple""#, r#"simple"#);
    test_str(r#"" white space ""#, r#" white space "#);
    test_str(r#""quote \"""#, r#"quote \""#);
    test_str(r#""slashes \\ \/""#, r#"slashes \\ \/"#);
    test_str(r#""unicode \u1234\u5678\u90AB\uCDEF""#,
             r#"unicode \u1234\u5678\u90AB\uCDEF"#);
  }

  #[test]
  fn test_name() {
    let mut lexer = Lexer::new("     name1    name2   \"simple\" \"  white space   \" other_name");
    test_next_token(&mut lexer, Token::Name(5, 10,"name1"));
    test_next_token(&mut lexer, Token::Name(14,19,"name2"));
    test_next_token(&mut lexer, Token::String(23,29,"simple"));
    test_next_token(&mut lexer, Token::String(32,48,"  white space   "));
    test_next_token(&mut lexer, Token::Name(50,60,"other_name"));
    test_next_token(&mut lexer, Token::Eof(50,60));
    assert_eq!(None, lexer.next());
  }

  #[test]
  fn test_number() {
    let mut lexer = Lexer::new("-9 0 9 -0 -9 1234 -1234 0.0 -0.0 1.0 -1.012 -0.101 1.0e2 -1.0E2 \
                                0.00e-002 -0.00E+002 9.1E+0");
    test_next_token(&mut lexer, Token::Int(0,2,    "-9"));
    test_next_token(&mut lexer, Token::Int(3,4,    "0"));
    test_next_token(&mut lexer, Token::Int(5,6,    "9"));
    test_next_token(&mut lexer, Token::Int(7,9,    "-0"));
    test_next_token(&mut lexer, Token::Int(10,12,  "-9"));
    test_next_token(&mut lexer, Token::Int(13,17,  "1234"));
    test_next_token(&mut lexer, Token::Int(18,23,  "-1234"));
    test_next_token(&mut lexer, Token::Float(24,27,"0.0"));
    test_next_token(&mut lexer, Token::Float(28,32,"-0.0"));
    test_next_token(&mut lexer, Token::Float(33,36,"1.0"));
    test_next_token(&mut lexer, Token::Float(37,43,"-1.012"));
    test_next_token(&mut lexer, Token::Float(44,50,"-0.101"));
    test_next_token(&mut lexer, Token::Float(51,56,"1.0e2"));
    test_next_token(&mut lexer, Token::Float(57,63,"-1.0E2"));
    test_next_token(&mut lexer, Token::Float(64,73,"0.00e-002"));
    test_next_token(&mut lexer, Token::Float(74,84,"-0.00E+002"));
    test_next_token(&mut lexer, Token::Float(85,91,"9.1E+0"));
    test_next_token(&mut lexer, Token::Eof(85,91));
    assert_eq!(None, lexer.next());
  }

  #[test]
  fn test_symbol() {
    let mut lexer = Lexer::new("@ !  :$ =    [  {   (|] } )      ...");
    test_next_token(&mut lexer, Token::At(0,0));
    test_next_token(&mut lexer, Token::Bang(2,2));
    test_next_token(&mut lexer, Token::Colon(5,5));
    test_next_token(&mut lexer, Token::Dollar(6,6));
    test_next_token(&mut lexer, Token::Equals(8,8));
    test_next_token(&mut lexer, Token::LeftBracket(13,13));
    test_next_token(&mut lexer, Token::LeftBrace(16,16));
    test_next_token(&mut lexer, Token::LeftParen(20,20));
    test_next_token(&mut lexer, Token::Pipe(21,21));
    test_next_token(&mut lexer, Token::RightBracket(22,22));
    test_next_token(&mut lexer, Token::RightBrace(24,24));
    test_next_token(&mut lexer, Token::RightParen(26,26));
    test_next_token(&mut lexer, Token::Spread(33,35));
    test_next_token(&mut lexer, Token::Eof(33,36));
    assert_eq!(None, lexer.next());
  }
}
<commit_msg>feat: add Punctuator<commit_after>
use std::str::CharIndices;
use std::iter::Peekable;

#[derive(PartialEq, Debug)]
pub enum Token<'a> {
  Eof(usize, usize),
  Punctuator(Punctuator, usize),
  Name(&'a str, usize, usize),
  IntValue(&'a str, usize, usize),
  FloatValue(&'a str, usize, usize),
  StringValue(&'a str, usize, usize),
}

#[derive(PartialEq, Debug)]
pub enum Punctuator {
  Bang,
  Dollar,
  LeftParen,
  RightParen,
  Spread(usize),
  Colon,
  Equals,
  At,
  LeftBracket,
  RightBracket,
  LeftBrace,
  RightBrace,
  Pipe,
}

macro_rules! take {
  ($lexer: ident, $($p: pat)|*) => ({
    let mut is_taken = false;
    if let Some(&(p, c)) = $lexer.iter.peek() {
      match c {
        $(
          $p => {
            $lexer.iter.next();
            $lexer.hi = p;
            is_taken = true;
          }
        )*
        _ => {}
      }
    }
    is_taken
  });
}

macro_rules! take_while {
  ($lexer: ident, $($p: pat)|*) => ({
    let mut is_taken = false;
    while let Some(&(p, c)) = $lexer.iter.peek() {
      match c {
        $(
          $p => {
            $lexer.iter.next();
            is_taken = true;
          }
        )*
        _ => {
          $lexer.hi = p;
          break;
        }
      }
    }
    is_taken
  });
}

macro_rules! take_while_not {
  ($lexer: ident, $($p: pat)|*) => ({
    let mut is_taken = false;
    while let Some(&(p, c)) = $lexer.iter.peek() {
      match c {
        $(
          $p => {
            $lexer.hi = p;
            break;
          }
        )*
        _ => {
            $lexer.iter.next();
            is_taken = true;
        }
      }
    }
    is_taken
  });
}

macro_rules! peek {
  ($lexer: ident, $($p: pat)|*) => ({
    let mut is_match = false;
    if let Some(&(_, c)) = $lexer.iter.peek() {
      match c {
        $(
          $p => {
            is_match = true;
          }
        )*
        _ => {}
      }
    }
    is_match
  });
}

macro_rules! take_eof {
  ($lexer: ident) => ({
    if $lexer.iter.peek().is_none() {
      $lexer.hi = $lexer.input_len;
      true
    } else {
      false
    }
  });
}

pub struct Lexer<'a> {
  input: &'a str,
  input_len: usize,
  iter: Peekable<CharIndices<'a>>,
  is_eof: bool,
  lo: usize,
  hi: usize,
}

impl<'a> Lexer<'a> {
  pub fn new(input: &'a str) -> Lexer<'a> {
    Lexer {
      input: input,
      input_len: input.len(),
      iter: input.char_indices().peekable(),
      is_eof: false,
      lo: 0,
      hi: 0,
    }
  }
}

fn skip_ignored_token(lexer: &mut Lexer) {
  loop {
    if take!(lexer, '#') {
      take_while_not!(lexer, '\r' | '\n');
    }
    if !take_while!(lexer, '\u{feff}' | ' ' | '\t' | '\r' | '\n' | ',') {
      break;
    }
  }
}

fn scan_spread(lexer: &mut Lexer) -> bool {
  if take!(lexer, '.') {
    if take!(lexer, '.') {
      if take!(lexer, '.') {
        return true;
      }
    }
  }
  return false;
}

fn scan_name(lexer: &mut Lexer) {
  take_while!(lexer, '_' | 'a'...'z' | 'A'...'Z' | '0'...'9');
  take_eof!(lexer);
}

// Return the raw string without unescape character
fn scan_string(lexer: &mut Lexer) -> (bool, bool) {
  let mut is_bad_escape = false;
  let mut is_terminated = false;
  if take!(lexer, '"') {
    lexer.lo = lexer.lo + 1;
  }
  loop {
    if take!(lexer, '\\') {
      if !take!(lexer, '\\' | '"' | 'b' | 'f' | 'n' | 'r' | 't' | '/' | 'u') {
        is_bad_escape = true;
        break;
      }
    }
    take_while_not!(lexer, '"' | '\\');
    if take!(lexer, '"') {
      is_terminated = true;
      break;
    }
  }
  (is_terminated, is_bad_escape)
}

// Int:   (+|-)?(0|[1-9][0-9]*)
// Float: (+|-)?(0|[1-9][0-9]*)(\.[0-9]+)?((E|e)(+|-)?[0-9]+)?
fn scan_number(lexer: &mut Lexer) -> (bool, bool) {
  let mut is_float = false;
  take!(lexer, '+' | '-');
  // Integer part
  if take!(lexer, '1'...'9') {
    take_while!(lexer, '0'...'9');
    if take!(lexer, ' ' | '\t' | '\r' | '\n' | ',') || take_eof!(lexer) {
      return (true, is_float);
    }
  } else if take!(lexer, '0') {
    if take!(lexer, ' ' | '\t' | '\r' | '\n' | ',') || take_eof!(lexer) {
      return (true, is_float);
    }
  }
  // Fractional part
  if take!(lexer, '.') {
    if !take_while!(lexer, '0'...'9') {
      return (false, is_float);
    }
    is_float = true;
  }
  // Exponent part
  if !is_float && peek!(lexer, 'E' | 'e') {
    return (false, is_float);
  }
  if take!(lexer, 'E' | 'e') {
    if take!(lexer, '+' | '-') {
      if !take_while!(lexer, '0'...'9') {
        return (false, is_float);
      }
    } else {
      take_while!(lexer, '0'...'9');
    }
  }

  if take!(lexer, ' ' | '\t' | '\r' | '\n' | ',') || take_eof!(lexer) {
    return (true, is_float);
  }
  return (false, is_float);
}

impl<'a> Iterator for Lexer<'a> {
  type Item = Result<Token<'a>, &'a str>;
  fn next(&mut self) -> Option<Result<Token<'a>, &'a str>> {
    if !self.is_eof && take_eof!(self) {
      self.is_eof = true;
      return Some(Ok(Token::Eof(self.lo, self.hi)));
    }
    if self.is_eof {
      return None;
    }
    skip_ignored_token(self);

    let mut item = None;
    if let Some(&(p, c)) = self.iter.peek() {
      self.lo = p;
      self.hi = p;
      item = match c {
        '!' => {
          self.iter.next();
          Some(Ok(Token::Bang(self.lo, self.hi)))
        }
        '$' => {
          self.iter.next();
          Some(Ok(Token::Dollar(self.lo, self.hi)))
        }
        '(' => {
          self.iter.next();
          Some(Ok(Token::LeftParen(self.lo, self.hi)))
        }
        ')' => {
          self.iter.next();
          Some(Ok(Token::RightParen(self.lo, self.hi)))
        }
        ':' => {
          self.iter.next();
          Some(Ok(Token::Colon(self.lo, self.hi)))
        }
        '=' => {
          self.iter.next();
          Some(Ok(Token::Equals(self.lo, self.hi)))
        }
        '@' => {
          self.iter.next();
          Some(Ok(Token::At(self.lo, self.hi)))
        }
        '[' => {
          self.iter.next();
          Some(Ok(Token::LeftBracket(self.lo, self.hi)))
        }
        ']' => {
          self.iter.next();
          Some(Ok(Token::RightBracket(self.lo, self.hi)))
        }
        '{' => {
          self.iter.next();
          Some(Ok(Token::LeftBrace(self.lo, self.hi)))
        }
        '}' => {
          self.iter.next();
          Some(Ok(Token::RightBrace(self.lo, self.hi)))
        }
        '|' => {
          self.iter.next();
          Some(Ok(Token::Pipe(self.lo, self.hi)))
        }
        '.' => {
          if scan_spread(self) {
            Some(Ok(Token::Spread(self.lo, self.hi)))
          } else {
            Some(Err("Unexpected character."))
          }
        }
        '_' | 'a'...'z' | 'A'...'Z' => {
          scan_name(self);
          Some(Ok(Token::Name(self.lo, self.hi, &self.input[self.lo..self.hi])))
        }
        '-' | '0'...'9' => {
          match scan_number(self) {
            (false, false) => Some(Err("Invalid integer number.")),
            (false, true) => Some(Err("Invalid float number.")),
            (true, false) => Some(Ok(Token::Int(self.lo, self.hi, &self.input[self.lo..self.hi]))),
            (true, true) => Some(Ok(Token::Float(self.lo, self.hi, &self.input[self.lo..self.hi]))),
          }
        }
        '"' => {
          match scan_string(self) {
            (false, _) => Some(Err("Unterminated string.")),
            (_, true) => Some(Err("Bad character escape sequence.")),
            _ => Some(Ok(Token::String(self.lo, self.hi, &self.input[self.lo..self.hi]))),
          }
        }
        _ => Some(Err("Unexpected character.")),
      }
    }
    item
  }
}

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

  fn test_next_token(lexer: &mut Lexer, expected: Token) {
    let next = lexer.next();
    assert_eq!(Some(Ok(expected)), next);
  }

  fn test_with_lex(input: &str, expeced: Token) {
    let mut lexer = Lexer::new(&input);
    test_next_token(&mut lexer, expeced);
  }

  fn test_str(input: &str, expeced: &str) {
    let len = expeced.len();
    test_with_lex(input, Token::String(1, len + 1, expeced))
  }

  #[test]
  fn test_string() {
    test_str(r#""simple""#, r#"simple"#);
    test_str(r#"" white space ""#, r#" white space "#);
    test_str(r#""quote \"""#, r#"quote \""#);
    test_str(r#""slashes \\ \/""#, r#"slashes \\ \/"#);
    test_str(r#""unicode \u1234\u5678\u90AB\uCDEF""#,
             r#"unicode \u1234\u5678\u90AB\uCDEF"#);
  }

  #[test]
  fn test_name() {
    let mut lexer = Lexer::new("     name1    name2   \"simple\" \"  white space   \" other_name");
    test_next_token(&mut lexer, Token::Name(5, 10,"name1"));
    test_next_token(&mut lexer, Token::Name(14,19,"name2"));
    test_next_token(&mut lexer, Token::String(23,29,"simple"));
    test_next_token(&mut lexer, Token::String(32,48,"  white space   "));
    test_next_token(&mut lexer, Token::Name(50,60,"other_name"));
    test_next_token(&mut lexer, Token::Eof(50,60));
    assert_eq!(None, lexer.next());
  }

  #[test]
  fn test_number() {
    let mut lexer = Lexer::new("-9 0 9 -0 -9 1234 -1234 0.0 -0.0 1.0 -1.012 -0.101 1.0e2 -1.0E2 \
                                0.00e-002 -0.00E+002 9.1E+0");
    test_next_token(&mut lexer, Token::Int(0,2,    "-9"));
    test_next_token(&mut lexer, Token::Int(3,4,    "0"));
    test_next_token(&mut lexer, Token::Int(5,6,    "9"));
    test_next_token(&mut lexer, Token::Int(7,9,    "-0"));
    test_next_token(&mut lexer, Token::Int(10,12,  "-9"));
    test_next_token(&mut lexer, Token::Int(13,17,  "1234"));
    test_next_token(&mut lexer, Token::Int(18,23,  "-1234"));
    test_next_token(&mut lexer, Token::Float(24,27,"0.0"));
    test_next_token(&mut lexer, Token::Float(28,32,"-0.0"));
    test_next_token(&mut lexer, Token::Float(33,36,"1.0"));
    test_next_token(&mut lexer, Token::Float(37,43,"-1.012"));
    test_next_token(&mut lexer, Token::Float(44,50,"-0.101"));
    test_next_token(&mut lexer, Token::Float(51,56,"1.0e2"));
    test_next_token(&mut lexer, Token::Float(57,63,"-1.0E2"));
    test_next_token(&mut lexer, Token::Float(64,73,"0.00e-002"));
    test_next_token(&mut lexer, Token::Float(74,84,"-0.00E+002"));
    test_next_token(&mut lexer, Token::Float(85,91,"9.1E+0"));
    test_next_token(&mut lexer, Token::Eof(85,91));
    assert_eq!(None, lexer.next());
  }

  #[test]
  fn test_symbol() {
    let mut lexer = Lexer::new("@ !  :$ =    [  {   (|] } )      ...");
    test_next_token(&mut lexer, Token::At(0,0));
    test_next_token(&mut lexer, Token::Bang(2,2));
    test_next_token(&mut lexer, Token::Colon(5,5));
    test_next_token(&mut lexer, Token::Dollar(6,6));
    test_next_token(&mut lexer, Token::Equals(8,8));
    test_next_token(&mut lexer, Token::LeftBracket(13,13));
    test_next_token(&mut lexer, Token::LeftBrace(16,16));
    test_next_token(&mut lexer, Token::LeftParen(20,20));
    test_next_token(&mut lexer, Token::Pipe(21,21));
    test_next_token(&mut lexer, Token::RightBracket(22,22));
    test_next_token(&mut lexer, Token::RightBrace(24,24));
    test_next_token(&mut lexer, Token::RightParen(26,26));
    test_next_token(&mut lexer, Token::Spread(33,35));
    test_next_token(&mut lexer, Token::Eof(33,36));
    assert_eq!(None, lexer.next());
  }
}
<|endoftext|>


<commit_before>use std::slice;

use goblin::elf::section_header::SHT_GNU_HASH;
use goblin::elf::sym::Sym;
use goblin::elf::Elf;
use goblin::elf32::gnu_hash::GnuHash as GnuHash32;
use goblin::elf64::gnu_hash::GnuHash as GnuHash64;

#[repr(C)]
#[repr(align(64))] // Align to cache lines
pub struct AlignedData<T: ?Sized>(T);

fn parse_gnu_hash_section(base: &[u8], symbol_name: &str) -> Result<Sym, &'static str> {
    let obj = Elf::parse(base).map_err(|_| "cannot parse ELF file")?;
    let hash_section = obj
        .section_headers
        .iter()
        .find(|s| s.sh_type == SHT_GNU_HASH)
        .ok_or("object does not contain .gnu.hash section")?;
    let hashtab: &[u8] = unsafe {
        let addr = base.as_ptr().add(hash_section.sh_offset as usize);
        let size = hash_section.sh_size as usize;
        slice::from_raw_parts(addr, size)
    };
    let dynsyms = obj.dynsyms.to_vec();
    let section = unsafe {
        if obj.is_64 {
            GnuHash64::from_raw_table(hashtab, &dynsyms)?.find(symbol_name, &obj.dynstrtab)
        } else {
            GnuHash32::from_raw_table(hashtab, &dynsyms)?.find(symbol_name, &obj.dynstrtab)
        }
    };
    section.copied().ok_or("cannot find symbol")
}

// Use lazy_parse and assembles the Elf with only parts we care
fn parse_text_section_size_lazy(base: &[u8]) -> Result<u64, &'static str> {
    let header = Elf::parse_header(base).map_err(|_| "parse elf header error")?;
    // dummy Elf with only header
    let mut obj = Elf::lazy_parse(header).map_err(|_| "cannot parse ELF file")?;

    use goblin::container::{Container, Ctx};
    use goblin::elf::SectionHeader;
    use goblin::strtab::Strtab;

    let ctx = Ctx {
        le: scroll::Endian::Little,
        container: Container::Big,
    };

    obj.section_headers =
        SectionHeader::parse(base, header.e_shoff as usize, header.e_shnum as usize, ctx)
            .map_err(|_| "parse section headers error")?;

    let strtab_idx = header.e_shstrndx as usize;
    let strtab_shdr = &obj.section_headers[strtab_idx];
    let strtab = Strtab::parse(
        base,
        strtab_shdr.sh_offset as usize,
        strtab_shdr.sh_size as usize,
        0x0,
    )
    .map_err(|_| "parse string table error")?;
    for (_, section) in obj.section_headers.iter().enumerate() {
        let section_name = strtab.get(section.sh_name).unwrap().unwrap();
        if section_name == ".text" {
            return Ok(section.sh_size);
        }
    }

    Err("Didn't find text section")
}

#[test]
fn test_parse_gnu_hash_section_64bit() {
    static ALIGNED_DATA: &AlignedData<[u8]> =
        &AlignedData(*include_bytes!("bins/elf/gnu_hash/hello.so"));

    assert_eq!(
        parse_gnu_hash_section(&ALIGNED_DATA.0, "helloWorld"),
        Ok(Sym {
            st_name: 97,
            st_info: 0x12,
            st_other: 0,
            st_shndx: 12,
            st_value: 0x65a,
            st_size: 33,
        })
    );
    assert_eq!(
        parse_gnu_hash_section(&ALIGNED_DATA.0, "_edata"),
        Ok(Sym {
            st_name: 130,
            st_info: 0x10,
            st_other: 0,
            st_shndx: 22,
            st_value: 0x0020_1030,
            st_size: 0,
        })
    );
    assert_eq!(
        parse_gnu_hash_section(&ALIGNED_DATA.0, "__gmon_start__"),
        Err("cannot find symbol"),
    );
}

#[test]
fn test_parse_gnu_hash_section_32bit() {
    static ALIGNED_DATA: &AlignedData<[u8]> =
        &AlignedData(*include_bytes!("bins/elf/gnu_hash/hello32.so"));

    assert_eq!(
        parse_gnu_hash_section(&ALIGNED_DATA.0, "helloWorld"),
        Ok(Sym {
            st_name: 97,
            st_info: 0x12,
            st_other: 0,
            st_shndx: 12,
            st_value: 0x4ed,
            st_size: 49,
        })
    );
    assert_eq!(
        parse_gnu_hash_section(&ALIGNED_DATA.0, "_edata"),
        Ok(Sym {
            st_name: 130,
            st_info: 0x10,
            st_other: 0,
            st_shndx: 22,
            st_value: 0x2018,
            st_size: 0,
        })
    );
    assert_eq!(
        parse_gnu_hash_section(&ALIGNED_DATA.0, "__gmon_start__"),
        Err("cannot find symbol"),
    );
}

#[test]
fn test_parse_text_section_size_lazy() {
    static ALIGNED_DATA: &AlignedData<[u8]> =
        &AlignedData(*include_bytes!("bins/elf/gnu_hash/hello.so"));
    assert_eq!(parse_text_section_size_lazy(&ALIGNED_DATA.0), Ok(0x126));
}

#[test]
fn test_oom() {
    use goblin::container::{Container, Ctx};
    use scroll::Pwrite;

    fn test_oom(data: &mut [u8]) {
        let mut modified_data = data.to_vec();
        let mut elf = Elf::parse(&data).unwrap();
        let endianness = elf.header.endianness().unwrap();
        let ctx = Ctx::new(
            if elf.is_64 {
                Container::Big
            } else {
                Container::Little
            },
            endianness,
        );
        let original_e_phnum = elf.header.e_phnum;
        let original_e_shnum = elf.header.e_shnum;

        // Way too many program headers
        elf.header.e_phnum = 1000;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        assert!(Elf::parse(&modified_data).is_err());

        // Possible overflow of program headers
        elf.header.e_phnum = std::u16::MAX;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        assert!(Elf::parse(&modified_data).is_err());

        // Back to original
        elf.header.e_phnum = original_e_phnum;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        assert!(Elf::parse(&modified_data).is_ok());

        // Way too many section headers
        elf.header.e_shnum = 1000;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        assert!(Elf::parse(&modified_data).is_err());

        // Fallback to large empty section header's sh_size
        elf.header.e_shnum = 0;
        elf.section_headers[0].sh_size = std::u64::MAX;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        modified_data
            .pwrite_with(
                elf.section_headers[0].clone(),
                elf.header.e_shoff as usize,
                ctx,
            )
            .unwrap();
        assert!(Elf::parse(&modified_data).is_err());

        // Possible overflow of section headers
        elf.header.e_shnum = std::u16::MAX;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        assert!(Elf::parse(&modified_data).is_err());

        // Back to original
        elf.header.e_shnum = original_e_shnum;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        assert!(Elf::parse(&modified_data).is_ok());
    }

    let aligned_data: &mut AlignedData<[u8]> =
        &mut AlignedData(*include_bytes!("bins/elf/gnu_hash/hello32.so"));
    test_oom(&mut aligned_data.0);

    let aligned_data: &mut AlignedData<[u8]> =
        &mut AlignedData(*include_bytes!("bins/elf/gnu_hash/hello.so"));
    test_oom(&mut aligned_data.0);
}
<commit_msg>tests.elf: Replace unchecked unsafe slicing with regular slicing<commit_after>use goblin::elf::section_header::SHT_GNU_HASH;
use goblin::elf::sym::Sym;
use goblin::elf::Elf;
use goblin::elf32::gnu_hash::GnuHash as GnuHash32;
use goblin::elf64::gnu_hash::GnuHash as GnuHash64;

#[repr(C)]
#[repr(align(64))] // Align to cache lines
pub struct AlignedData<T: ?Sized>(T);

fn parse_gnu_hash_section(base: &[u8], symbol_name: &str) -> Result<Sym, &'static str> {
    let obj = Elf::parse(base).map_err(|_| "cannot parse ELF file")?;
    let hash_section = obj
        .section_headers
        .iter()
        .find(|s| s.sh_type == SHT_GNU_HASH)
        .ok_or("object does not contain .gnu.hash section")?;
    let (start, length) = (
        hash_section.sh_offset as usize,
        hash_section.sh_size as usize,
    );
    let hashtab: &[u8] = &base[start..(start + length)];
    let dynsyms = obj.dynsyms.to_vec();
    let section = unsafe {
        if obj.is_64 {
            GnuHash64::from_raw_table(hashtab, &dynsyms)?.find(symbol_name, &obj.dynstrtab)
        } else {
            GnuHash32::from_raw_table(hashtab, &dynsyms)?.find(symbol_name, &obj.dynstrtab)
        }
    };
    section.copied().ok_or("cannot find symbol")
}

// Use lazy_parse and assembles the Elf with only parts we care
fn parse_text_section_size_lazy(base: &[u8]) -> Result<u64, &'static str> {
    let header = Elf::parse_header(base).map_err(|_| "parse elf header error")?;
    // dummy Elf with only header
    let mut obj = Elf::lazy_parse(header).map_err(|_| "cannot parse ELF file")?;

    use goblin::container::{Container, Ctx};
    use goblin::elf::SectionHeader;
    use goblin::strtab::Strtab;

    let ctx = Ctx {
        le: scroll::Endian::Little,
        container: Container::Big,
    };

    obj.section_headers =
        SectionHeader::parse(base, header.e_shoff as usize, header.e_shnum as usize, ctx)
            .map_err(|_| "parse section headers error")?;

    let strtab_idx = header.e_shstrndx as usize;
    let strtab_shdr = &obj.section_headers[strtab_idx];
    let strtab = Strtab::parse(
        base,
        strtab_shdr.sh_offset as usize,
        strtab_shdr.sh_size as usize,
        0x0,
    )
    .map_err(|_| "parse string table error")?;
    for (_, section) in obj.section_headers.iter().enumerate() {
        let section_name = strtab.get(section.sh_name).unwrap().unwrap();
        if section_name == ".text" {
            return Ok(section.sh_size);
        }
    }

    Err("Didn't find text section")
}

#[test]
fn test_parse_gnu_hash_section_64bit() {
    static ALIGNED_DATA: &AlignedData<[u8]> =
        &AlignedData(*include_bytes!("bins/elf/gnu_hash/hello.so"));

    assert_eq!(
        parse_gnu_hash_section(&ALIGNED_DATA.0, "helloWorld"),
        Ok(Sym {
            st_name: 97,
            st_info: 0x12,
            st_other: 0,
            st_shndx: 12,
            st_value: 0x65a,
            st_size: 33,
        })
    );
    assert_eq!(
        parse_gnu_hash_section(&ALIGNED_DATA.0, "_edata"),
        Ok(Sym {
            st_name: 130,
            st_info: 0x10,
            st_other: 0,
            st_shndx: 22,
            st_value: 0x0020_1030,
            st_size: 0,
        })
    );
    assert_eq!(
        parse_gnu_hash_section(&ALIGNED_DATA.0, "__gmon_start__"),
        Err("cannot find symbol"),
    );
}

#[test]
fn test_parse_gnu_hash_section_32bit() {
    static ALIGNED_DATA: &AlignedData<[u8]> =
        &AlignedData(*include_bytes!("bins/elf/gnu_hash/hello32.so"));

    assert_eq!(
        parse_gnu_hash_section(&ALIGNED_DATA.0, "helloWorld"),
        Ok(Sym {
            st_name: 97,
            st_info: 0x12,
            st_other: 0,
            st_shndx: 12,
            st_value: 0x4ed,
            st_size: 49,
        })
    );
    assert_eq!(
        parse_gnu_hash_section(&ALIGNED_DATA.0, "_edata"),
        Ok(Sym {
            st_name: 130,
            st_info: 0x10,
            st_other: 0,
            st_shndx: 22,
            st_value: 0x2018,
            st_size: 0,
        })
    );
    assert_eq!(
        parse_gnu_hash_section(&ALIGNED_DATA.0, "__gmon_start__"),
        Err("cannot find symbol"),
    );
}

#[test]
fn test_parse_text_section_size_lazy() {
    static ALIGNED_DATA: &AlignedData<[u8]> =
        &AlignedData(*include_bytes!("bins/elf/gnu_hash/hello.so"));
    assert_eq!(parse_text_section_size_lazy(&ALIGNED_DATA.0), Ok(0x126));
}

#[test]
fn test_oom() {
    use goblin::container::{Container, Ctx};
    use scroll::Pwrite;

    fn test_oom(data: &mut [u8]) {
        let mut modified_data = data.to_vec();
        let mut elf = Elf::parse(&data).unwrap();
        let endianness = elf.header.endianness().unwrap();
        let ctx = Ctx::new(
            if elf.is_64 {
                Container::Big
            } else {
                Container::Little
            },
            endianness,
        );
        let original_e_phnum = elf.header.e_phnum;
        let original_e_shnum = elf.header.e_shnum;

        // Way too many program headers
        elf.header.e_phnum = 1000;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        assert!(Elf::parse(&modified_data).is_err());

        // Possible overflow of program headers
        elf.header.e_phnum = std::u16::MAX;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        assert!(Elf::parse(&modified_data).is_err());

        // Back to original
        elf.header.e_phnum = original_e_phnum;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        assert!(Elf::parse(&modified_data).is_ok());

        // Way too many section headers
        elf.header.e_shnum = 1000;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        assert!(Elf::parse(&modified_data).is_err());

        // Fallback to large empty section header's sh_size
        elf.header.e_shnum = 0;
        elf.section_headers[0].sh_size = std::u64::MAX;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        modified_data
            .pwrite_with(
                elf.section_headers[0].clone(),
                elf.header.e_shoff as usize,
                ctx,
            )
            .unwrap();
        assert!(Elf::parse(&modified_data).is_err());

        // Possible overflow of section headers
        elf.header.e_shnum = std::u16::MAX;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        assert!(Elf::parse(&modified_data).is_err());

        // Back to original
        elf.header.e_shnum = original_e_shnum;
        modified_data
            .pwrite_with(elf.header, 0, endianness)
            .unwrap();
        assert!(Elf::parse(&modified_data).is_ok());
    }

    let aligned_data: &mut AlignedData<[u8]> =
        &mut AlignedData(*include_bytes!("bins/elf/gnu_hash/hello32.so"));
    test_oom(&mut aligned_data.0);

    let aligned_data: &mut AlignedData<[u8]> =
        &mut AlignedData(*include_bytes!("bins/elf/gnu_hash/hello.so"));
    test_oom(&mut aligned_data.0);
}
<|endoftext|>


<commit_before>extern crate natural_constants;

#[test]
fn test_math() {
    use natural_constants::math::*;

    println!("{}", golden_ratio)
}

#[test]
fn test_conversion() {
    use natural_constants::conversion::*;

    let mass_in_g = 50_000.0; // 50 kg
    assert_eq!(to_kilo(mass_in_g), 50.0);
    assert_eq!(from_kilo(50.0), mass_in_g);

    let energy_in_j = 125_000_000.0; // 125 MJoule
    assert_eq!(to_mega(energy_in_j), 125.0);
    assert_eq!(from_mega(125.0), energy_in_j);

    let file_size_in_bytes1 = 73_000_000_000.0; // 74 gigabytes (not gibibytes)
    assert_eq!(to_giga(file_size_in_bytes1), 73.0);
    assert_eq!(from_giga(73.0), file_size_in_bytes1);

    let file_size_in_bytes2 = 312_000_000_000_000.0; // 312 terabytes (not tebibytes)
    assert_eq!(to_tera(file_size_in_bytes2), 312.0);
    assert_eq!(from_tera(312.0), file_size_in_bytes2);

    let force_in_kn = 780.23; // kilo newton
    assert_eq!(kilo_to_mega(force_in_kn), 0.78023);
    assert_eq!(mega_to_kilo(0.78023), force_in_kn);
}

#[test]
fn limits() {
    use natural_constants::math::*;
    assert_eq!(golden_ratio, (1.0 + 5.0_f64.sqrt())/2.0);
}
<commit_msg>Add tests for conversion kelvin to deg and vice versa<commit_after>extern crate natural_constants;

#[test]
fn test_math() {
    use natural_constants::math::*;

    println!("{}", golden_ratio)
}

#[test]
fn test_conversion() {
    use natural_constants::conversion::*;

    let mass_in_g = 50_000.0; // 50 kg
    assert_eq!(to_kilo(mass_in_g), 50.0);
    assert_eq!(from_kilo(50.0), mass_in_g);

    let energy_in_j = 125_000_000.0; // 125 MJoule
    assert_eq!(to_mega(energy_in_j), 125.0);
    assert_eq!(from_mega(125.0), energy_in_j);

    let file_size_in_bytes1 = 73_000_000_000.0; // 74 gigabytes (not gibibytes)
    assert_eq!(to_giga(file_size_in_bytes1), 73.0);
    assert_eq!(from_giga(73.0), file_size_in_bytes1);

    let file_size_in_bytes2 = 312_000_000_000_000.0; // 312 terabytes (not tebibytes)
    assert_eq!(to_tera(file_size_in_bytes2), 312.0);
    assert_eq!(from_tera(312.0), file_size_in_bytes2);

    let force_in_kn = 780.23; // kilo newton
    assert_eq!(kilo_to_mega(force_in_kn), 0.78023);
    assert_eq!(mega_to_kilo(0.78023), force_in_kn);

    let temp_in_kelvin = 1.0;
    assert_eq!(kelvin_to_deg_cel(temp_in_kelvin), -272.15);
    assert_eq!(deg_cel_to_kelvin(-272.15), temp_in_kelvin);
}

#[test]
fn limits() {
    use natural_constants::math::*;
    assert_eq!(golden_ratio, (1.0 + 5.0_f64.sqrt())/2.0);
}
<|endoftext|>


<commit_before>extern crate rand;
extern crate spectra;

use rand::{Rng, thread_rng};
use spectra::anim::spline::*;

#[test]
fn hold() {
  let spline = Spline::from_keys(vec![
    Key::new(0., 10., Interpolation::Step(1.)),
    Key::new(24., 100., Interpolation::Step(1.)),
    Key::new(45., -3.34, Interpolation::Step(1.))
  ]);

  assert_eq!(spline.sample(0.), Some(10.));
  assert_eq!(spline.sample(2.), Some(10.));
  assert_eq!(spline.sample(23.), Some(10.));
  assert_eq!(spline.sample(24.), Some(100.));
  assert_eq!(spline.sample(44.), Some(100.));
  assert_eq!(spline.sample(44.), Some(100.));
  assert_eq!(spline.sample(45.), None);
  assert_eq!(spline.sample(45347.,), None);
  assert_eq!(spline.sample(45347.,), None);
}

#[test]
fn linear() {
  let spline = Spline::from_keys(vec![
    Key::new(0., 10., Interpolation::Linear),
    Key::new(10., 20., Interpolation::Linear)
  ]);

  assert_eq!(spline.sample(0.), Some(10.));
  assert_eq!(spline.sample(10.), None);
  assert_eq!(spline.sample(5.), Some(15.));
}

#[test]
fn keys_sorted() {
  let nb = 10000;
  let mut rng = thread_rng();
  let mut keys = Vec::with_capacity(nb);

  for _ in 0..nb {
    let t = rng.gen::<f32>().abs();
    let v: f32 = rng.gen();
    let key = Key::new(t, v, Interpolation::Step(1.));

    keys.push(key);
  }

  let anim_param = Spline::from_keys(keys);

  let mut t = 0.;
  for key in anim_param.into_iter() {
    assert!(t <= key.t, "t: {}, key.t: {}", t, key.t);
    t = key.t;
  }
}
<commit_msg>Remove tests. We’ll re-introduce new ones later.<commit_after><|endoftext|>


<commit_before>use std::env;
use std::process::Command;

#[derive(Debug)]
pub struct NiceOutput {
    pub status: i32,
    pub stdout: String,
    pub stderr: String,
}

fn run(args: Vec<&str>) -> NiceOutput {
    let mut directory = env::current_exe().unwrap();
            directory.pop(); // chop off exe name but leave "debug"
    let name = format!("to{}", env::consts::EXE_SUFFIX);
    let binary = directory.join(&name);

    let mut command = Command::new(binary);
            command.args(&args);
    let output = command.output().unwrap();
    let stdout = trim(output.stdout);
    let stderr = trim(output.stderr);

    return NiceOutput {
        status: output.status.code().unwrap(),
        stdout: stdout,
        stderr: stderr,
    };
}

fn trim(output: Vec<u8>) -> String {
    let string = String::from_utf8(output).unwrap();
    return String::from(string.trim());
}

#[test]
fn help_long() {
    let result = run(vec!["--help"]);
    assert_eq!(result.status, 0);
}

#[test]
fn help_short() {
    let result = run(vec!["-h"]);
    assert_eq!(result.status, 0);
}

#[test]
fn version_long() {
    let result = run(vec!["--version"]);
    let version = env!("CARGO_PKG_VERSION");
    let expected = format!("to {}", version);

    assert_eq!(result.stdout, expected);
    assert_eq!(result.status, 0);
}

#[test]
fn version_short() {
    let result = run(vec!["-V"]);
    let version = env!("CARGO_PKG_VERSION");
    let expected = format!("to {}", version);

    assert_eq!(result.stdout, expected);
    assert_eq!(result.status, 0);
}

#[test]
fn save_long() {
    let result = run(vec!["--save"]);
    assert_eq!(result.status, 0);
}

#[test]
fn save_short() {
    let result = run(vec!["-s"]);
    assert_eq!(result.status, 0);
}
<commit_msg>Cleans up integration tests.<commit_after>use std::env;
use std::process::Command;

#[derive(Debug)]
pub struct NiceOutput {
    pub status: i32,
    pub stdout: String,
    pub stderr: String,
}

fn run(args: Vec<&str>) -> NiceOutput {
    let mut directory = env::current_exe().unwrap();
            directory.pop(); // chop off exe name but leave "debug"
    let name = format!("to{}", env::consts::EXE_SUFFIX);
    let binary = directory.join(&name);

    let mut command = Command::new(binary);
            command.args(&args);
    let output = command.output().unwrap();
    let stdout = trim(output.stdout);
    let stderr = trim(output.stderr);

    return NiceOutput {
        status: output.status.code().unwrap(),
        stdout: stdout,
        stderr: stderr,
    };
}

fn trim(output: Vec<u8>) -> String {
    let string = String::from_utf8(output).unwrap();
    return String::from(string.trim());
}

#[test]
fn help() {
    let result = run(vec!["--help"]);
    assert_eq!(result.status, 0);

    let result = run(vec!["-h"]);
    assert_eq!(result.status, 0);
}

#[test]
fn version() {
    let version = env!("CARGO_PKG_VERSION");
    let expected = format!("to {}", version);

    let result = run(vec!["--version"]);
    assert_eq!(result.stdout, expected);
    assert_eq!(result.status, 0);

    let result = run(vec!["-V"]);
    assert_eq!(result.stdout, expected);
    assert_eq!(result.status, 0);
}

#[test]
fn save() {
    let result = run(vec!["--save"]);
    assert_eq!(result.status, 0);

    let result = run(vec!["-s"]);
    assert_eq!(result.status, 0);
}

#[test]
fn info() {
    let result = run(vec!["--info"]);
    assert_eq!(result.status, 0);

    let result = run(vec!["-i"]);
    assert_eq!(result.status, 0);
}
<|endoftext|>


<commit_before>use cap::Cap;
use crates_index::*;
use rayon::iter::ParallelIterator;
use std::alloc;

#[global_allocator]
static ALLOCATOR: Cap<alloc::System> = Cap::new(alloc::System, usize::max_value());

#[test]
fn mem_usage() {
    let index = Index::new_cargo_default().unwrap();

    let before = ALLOCATOR.allocated();
    // let all_crates: Vec<_> = index.crates().collect();
    let all_crates: Vec<_> = index.crates_parallel().map(|c| c.unwrap()).collect();
    let after = ALLOCATOR.allocated();
    let used = after - before;
    assert!(all_crates.len() > 76200);
    eprintln!(
        "used mem: {}B for {} crates, {}B per crate",
        used,
        all_crates.len(),
        used / all_crates.len()
    );
    assert!(used / all_crates.len() < 5000);
}
<commit_msg>The index keeps growing<commit_after>use cap::Cap;
use crates_index::*;
use rayon::iter::ParallelIterator;
use std::alloc;

#[global_allocator]
static ALLOCATOR: Cap<alloc::System> = Cap::new(alloc::System, usize::max_value());

#[test]
fn mem_usage() {
    let index = Index::new_cargo_default().unwrap();

    let before = ALLOCATOR.allocated();
    // let all_crates: Vec<_> = index.crates().collect();
    let all_crates: Vec<_> = index.crates_parallel().map(|c| c.unwrap()).collect();
    let after = ALLOCATOR.allocated();
    let used = after - before;
    assert!(all_crates.len() > 89000);
    eprintln!(
        "used mem: {}B for {} crates, {}B per crate",
        used,
        all_crates.len(),
        used / all_crates.len()
    );
    assert!(used / all_crates.len() < 5100);
}
<|endoftext|>


<commit_before>
use rocket::Route;
use rocket::State;

use comm;

use utils::IdRange;

pub enum OutgoingMessage {
    GetPressure(u32),
    GetTemperature(u32),
    GetHumidity(u32),
    GetLightLevel(u32),
}

enum IncomingMessage {
    Pressure(u32, f32),
    Temperature(u32, f32),
    Humidity(u32, f32),
    RetVal(i32),
}

#[get("/<id_range>/pressure")]
fn query_pressure(id_range: IdRange, comm_state: State<comm::CommState>) -> String {
    format!("{:?}", id_range)
}

#[get("/<id_range>/temperature")]
fn query_temperature(id_range: IdRange, comm_state: State<comm::CommState>) -> String {
    format!("{:?}", id_range)
}

#[get("/<id_range>/humidity")]
fn query_humidity(id_range: IdRange, comm_state: State<comm::CommState>) -> String {
    format!("{:?}", id_range)
}

#[get("/<id_range>/light_level")]
fn query_light_level(id_range: IdRange, comm_state: State<comm::CommState>) -> String {
    format!("{:?}", id_range)
}

pub fn get_routes() -> Vec<Route> {
    routes![query_pressure,
            query_temperature,
            query_humidity,
            query_light_level]
}
<commit_msg>Implemented meteo message ser/des<commit_after>
use rocket::Route;
use rocket::State;
use rocket_contrib::Json;

use comm;

use utils::IdRange;

use std::str::FromStr;
use std::collections::HashMap;
use std::time::Duration;

#[derive(Debug)]
struct MeteoError;

pub enum OutgoingMessage {
    GetPressure(u32),
    GetTemperature(u32),
    GetHumidity(u32),
    GetLightLevel(u32),
}

impl From<OutgoingMessage> for String {
    fn from(msg: OutgoingMessage) -> String {
        match msg {
            OutgoingMessage::GetPressure(ch) => format!("GET_PRESSURE,{}", ch),
            OutgoingMessage::GetTemperature(ch) => format!("GET_TEMPERATURE,{}", ch),
            OutgoingMessage::GetHumidity(ch) => format!("GET_HUMIDITY,{}", ch),
            OutgoingMessage::GetLightLevel(ch) => format!("GET_LIGHT_LEVEL,{}", ch),
        }
    }
}

enum IncomingMessage {
    Pressure(u32, f32),
    Temperature(u32, f32),
    Humidity(u32, f32),
    LightLevel(u32, f32),
    RetVal(i32),
}

impl FromStr for IncomingMessage {
    type Err = MeteoError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut tokens = s.split(',');

        if let Some(msg_type) = tokens.next() {
            match msg_type {
                "PRESSURE_REPLY" => {
                    let ch_num = tokens.next().ok_or(MeteoError)?.parse().map_err(
                        |_| MeteoError,
                    )?;
                    let val = tokens.next().ok_or(MeteoError)?.parse().map_err(
                        |_| MeteoError,
                    )?;

                    Ok(IncomingMessage::Pressure(ch_num, val))
                }
                "TEMPERATURE_REPLY" => {
                    let ch_num = tokens.next().ok_or(MeteoError)?.parse().map_err(
                        |_| MeteoError,
                    )?;
                    let val = tokens.next().ok_or(MeteoError)?.parse().map_err(
                        |_| MeteoError,
                    )?;

                    Ok(IncomingMessage::Temperature(ch_num, val))
                }
                "HUMIDITY_REPLY" => {
                    let ch_num = tokens.next().ok_or(MeteoError)?.parse().map_err(
                        |_| MeteoError,
                    )?;
                    let val = tokens.next().ok_or(MeteoError)?.parse().map_err(
                        |_| MeteoError,
                    )?;

                    Ok(IncomingMessage::Humidity(ch_num, val))
                }
                "LIGHT_LEVEL_REPLY" => {
                    let ch_num = tokens.next().ok_or(MeteoError)?.parse().map_err(
                        |_| MeteoError,
                    )?;
                    let val = tokens.next().ok_or(MeteoError)?.parse().map_err(
                        |_| MeteoError,
                    )?;

                    Ok(IncomingMessage::LightLevel(ch_num, val))
                }
                "RET_VAL" => {
                    let ret_val = tokens.next().ok_or(MeteoError)?.parse().map_err(
                        |_| MeteoError,
                    )?;

                    Ok(IncomingMessage::RetVal(ret_val))
                }
                _ => Err(MeteoError),
            }
        } else {
            Err(MeteoError)
        }
    }
}

fn transfer(
    comm_ch: &comm::CommChannelTx,
    msg: OutgoingMessage,
) -> Result<IncomingMessage, MeteoError> {
    let msg_str = msg.into();
    debug!("Sending: {}", msg_str);

    let response_channel = comm_ch.send(msg_str).map_err(|e| {
        warn!("{}", e);
        MeteoError
    })?;

    let raw_response_msg = response_channel
        .recv_timeout(Duration::from_secs(3))
        .map_err(|e| {
            warn!("{}", e);
            MeteoError
        })?;

    debug!("Response: {}", raw_response_msg);

    raw_response_msg.parse()
}

#[get("/<id_range>/pressure")]
fn query_pressure(
    id_range: IdRange,
    comm_state: State<comm::CommState>,
) -> Result<Json<HashMap<u32, f32>>, MeteoError> {

    let comm = comm_state.get_comm_channel();

    let mut press_map = HashMap::new();

    for id in id_range {
        let resp = transfer(&comm, OutgoingMessage::GetPressure(id))?;

        match resp {
            IncomingMessage::Pressure(ch, val) if ch == id => {
                press_map.insert(ch, val);
            }
            IncomingMessage::RetVal(_val) => {}
            _ => return Err(MeteoError),
        }
    }

    Ok(Json(press_map))
}

#[get("/<id_range>/temperature")]
fn query_temperature(
    id_range: IdRange,
    comm_state: State<comm::CommState>,
) -> Result<Json<HashMap<u32, f32>>, MeteoError> {

    let comm = comm_state.get_comm_channel();

    let mut press_map = HashMap::new();

    for id in id_range {
        let resp = transfer(&comm, OutgoingMessage::GetTemperature(id))?;

        match resp {
            IncomingMessage::Temperature(ch, val) if ch == id => {
                press_map.insert(ch, val);
            }
            IncomingMessage::RetVal(_val) => {}
            _ => return Err(MeteoError),
        }
    }

    Ok(Json(press_map))
}

#[get("/<id_range>/humidity")]
fn query_humidity(
    id_range: IdRange,
    comm_state: State<comm::CommState>,
) -> Result<Json<HashMap<u32, f32>>, MeteoError> {

    let comm = comm_state.get_comm_channel();

    let mut press_map = HashMap::new();

    for id in id_range {
        let resp = transfer(&comm, OutgoingMessage::GetHumidity(id))?;

        match resp {
            IncomingMessage::Humidity(ch, val) if ch == id => {
                press_map.insert(ch, val);
            }
            IncomingMessage::RetVal(_val) => {}
            _ => return Err(MeteoError),
        }
    }

    Ok(Json(press_map))
}

#[get("/<id_range>/light_level")]
fn query_light_level(
    id_range: IdRange,
    comm_state: State<comm::CommState>,
) -> Result<Json<HashMap<u32, f32>>, MeteoError> {

    let comm = comm_state.get_comm_channel();

    let mut press_map = HashMap::new();

    for id in id_range {
        let resp = transfer(&comm, OutgoingMessage::GetLightLevel(id))?;

        match resp {
            IncomingMessage::LightLevel(ch, val) if ch == id => {
                press_map.insert(ch, val);
            }
            IncomingMessage::RetVal(_val) => {}
            _ => return Err(MeteoError),
        }
    }

    Ok(Json(press_map))
}

pub fn get_routes() -> Vec<Route> {
    routes![
        query_pressure,
        query_temperature,
        query_humidity,
        query_light_level,
    ]
}
<|endoftext|>


<commit_before>use ewma::EWMA;

struct Meter {
    value: i64,
    ewma1: EWMA
}

impl Meter {
}

fn new() -> Meter {
    let i = -5.0f64/60.0f64/1f64;

    return Meter {
        value: 0,
        ewma1: EWMA::new(1f64 - i)
    }
}<commit_msg>More complete meter<commit_after>extern crate time;

use self::time::get_time;
use self::time::Timespec;

use std::sync::Mutex;

use ewma::EWMA;

pub struct MeterSnapshot{
    count: i64,
    rate1: f64,
    rate5: f64,
    rate15: f64,
    rate_mean: f64
}

pub struct Meter {
    data: Mutex<MeterSnapshot>,
    a1: EWMA,
    a5: EWMA,
    a15: EWMA,
    start: Timespec
}

impl Meter {
    pub fn mark(&self, n: i64) {
        let mut s = self.data.lock();
        s.count += n;
        self.a1.update(n as uint);
        self.a5.update(n as uint);
        self.a15.update(n as uint);
    }

    pub fn snapshot(&self) -> MeterSnapshot{
        let s = self.data.lock();
        MeterSnapshot {
            count: s.count,
            rate1: s.rate1,
            rate5: s.rate5,
            rate15: s.rate15,
            rate_mean: s.rate_mean
        }
    }

    fn update_snapshot(&self) {
        let mut s = self.data.lock();
        s.rate1 = self.a1.rate();
        s.rate5 = self.a5.rate();
        s.rate15 = self.a15.rate();

        let diff = get_time() - self.start;
        s.rate_mean = s.count as f64 / diff.num_seconds() as f64;
    }

    pub fn tick(&mut self) {
        self.data.lock();
        self.a1.tick();
        self.a5.tick();
        self.a15.tick();
        self.update_snapshot()
    }

    pub fn new() -> Meter {
        let i = -5.0f64/60.0f64/1f64;

        let data = MeterSnapshot{
            count: 0i64,
            rate1: 0f64,
            rate5: 0f64,
            rate15: 0f64,
            rate_mean: 0f64
        };

        Meter {
            data: Mutex::new(data),
            a1: EWMA::new(1f64 - i),
            a5: EWMA::new(5f64 - i),
            a15: EWMA::new(15f64 - i),
            start: get_time()
        }
    }
}

#[cfg(test)]
mod test {
    use meter::{Meter,MeterSnapshot};

    #[test]
    fn zero() {
        let m: Meter = Meter::new();
        let s: MeterSnapshot = m.snapshot();
        assert_eq!(s.count, 0);
    }

    #[test]
    fn non_zero() {
        let m: Meter = Meter::new();
        m.mark(3);
        let s: MeterSnapshot = m.snapshot();
        assert_eq!(s.count, 3);
    }

    #[test]
    fn snapshot() {
        let m: Meter = Meter::new();
        m.mark(1);
        m.mark(1);

        let s = m.snapshot();

        m.mark(1);

        assert_eq!(s.count, 2);
        assert_eq!(m.snapshot().count, 3);
    }
}<|endoftext|>


<commit_before>/*
 * Copyright (C) 2017 Ryan Huang
 *
 * 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 <https://www.gnu.org/licenses/>.
 */

pub struct Model {
    pub board: Board,
    pub turn: Turn,
    pub white_pieces: u32,
    pub white_hexes: u32,
    pub black_pieces: u32,
    pub black_hexes: u32,
    pub selected_piece: Option<FieldCoord>,
}

impl Model {
    pub fn new() -> Model {
        Model {
            board: Board::new(),
            turn: Turn::White,
            white_pieces: 18,
            white_hexes: 0,
            black_pieces: 18,
            black_hexes: 0,
            selected_piece: None,
        }
    }
}

pub enum Turn {
    White,
    Black,
}

impl Turn {
    pub fn switch_turns(&mut self) {
        *self = match *self {
            Turn::White => Turn::Black,
            Turn::Black => Turn::White,
        }
    }
}

pub struct Board {
    /*
                ____
               /    \
          ____/  +y  \____
         /    \      /    \
        /      \____/  +x  \
        \      /    \      /
         \____/      \____/
         /    \      /    \
        /  -x  \____/      \
        \      /    \      /
         \____/  -y  \____/
              \      /
               \____/

        The hex board uses an axial coordinate system with (0, 0) at the center. The x-axis slopes
        up to the right, and the y-axis goes up and down. The board is stored as a dense 2D array,
        because a ragged array won't work (since each row would have to be a different type).

        See http://www.redblobgames.com/grids/hexagons/#coordinates-axial for more info.
    */
    board: [[Option<Hex>; 5]; 5],
}

// Fields are numbered clockwise from the top. Even indicies are black, odd indicies are white.
type Hex = [Field; 6];

// There's no need to store the color of the piece on this field, since only white pieces can be on
// white fields, and vice versa.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Field {
    Piece,
    Empty,
}

impl Board {
    #[cfg_attr(rustfmt, rustfmt_skip)]
    /// Create a new board with the "Laurentius" starting position.
    pub fn new() -> Board {
        let mut board = Board {
            board: [[None; 5]; 5]
        };

        // (0, 0) is the only empty hex.
        board.set_hex(&HexCoord::new(0, 0), Some([Field::Empty; 6]));

        // All other hexes have exactly two pieces on them in the starting position.
        let piece_locations = [
            (-2,  2, 0, 4),
            (-2,  1, 0, 3),
            (-2,  0, 3, 5),
            (-1,  2, 1, 4),
            (-1,  1, 0, 4),
            (-1,  0, 3, 5),
            (-1, -1, 2, 5),
            ( 0,  2, 1, 5),
            ( 0,  1, 1, 5),
            ( 0, -1, 2, 4),
            ( 0, -2, 2, 4),
            ( 1,  1, 2, 5),
            ( 1,  0, 0, 2),
            ( 1, -1, 1, 3),
            ( 1, -2, 1, 4),
            ( 2,  0, 0, 2),
            ( 2, -1, 0, 3),
            ( 2, -2, 1, 3),
        ];

        for &(x, y, f1, f2) in &piece_locations {
            let mut hex = [Field::Empty; 6];
            hex[f1] = Field::Piece;
            hex[f2] = Field::Piece;
            board.set_hex(&HexCoord::new(x, y), Some(hex));
        }

        board
    }

    pub fn get_field(&self, coord: &FieldCoord) -> &Field {
        match *self.get_hex(&coord.to_hex()) {
            Some(ref hex) => &hex[coord.f as usize],
            None => panic!(
                "Cannot get field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    fn set_field(&mut self, coord: &FieldCoord, field: Field) {
        match *self.get_hex(&coord.to_hex()) {
            Some(mut hex) => hex[coord.f as usize] = field,
            None => panic!(
                "Cannot set field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    /// Return fields that share an edge with the given field. These fields are always the opposite
    /// color of the given field. If all of a piece's edge neighbors are occupied, that piece might
    /// be capturable.
    pub fn get_field_edge_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        let mut neighbors = vec![
            // There are always two edge neighbors on the same hex as the given field
            FieldCoord::new(coord.x, coord.y, (coord.f + 1) % 6),
            FieldCoord::new(coord.x, coord.y, (coord.f + 5) % 6),
        ];

        for (i, neighbor) in self.get_hex_neighbors(&coord.to_hex()) {
            if coord.f == i {
                let f = (i + 3) % 6;
                neighbors.push(neighbor.to_field(f));
                break;
            }
        }

        neighbors
    }
    /// Return fields that share a vertex with the given field and have the same color as the given
    /// field. Pieces can move to fields that are vertex neighbors of the field they are on.
    pub fn get_field_vertex_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        // A field's vertex neighbors can be defined as the edge neighbors of its edge neighbors
        let mut neighbors = vec![];
        for neighbor in self.get_field_edge_neighbors(coord) {
            neighbors.append(&mut self.get_field_edge_neighbors(&neighbor));
        }
        // A field is not its own neighbor
        neighbors.into_iter().filter(|n| n != coord).collect()
    }
    pub fn move_piece(&mut self, from: &FieldCoord, to: &FieldCoord) {
        assert_eq!(
            self.get_field(from),
            &Field::Piece,
            "There is no piece at {:?} to move",
            from
        );
        assert!(
            self.get_field_vertex_neighbors(from).contains(to),
            "Cannot move piece at {:?} to non-(vertex neighbor) {:?}",
            from,
            to
        );
        assert_eq!(
            self.get_field(to),
            &Field::Empty,
            "Cannot move piece at {:?} to occupied neighbor {:?}",
            from,
            to
        );

        self.set_field(from, Field::Empty);
        self.set_field(to, Field::Piece);
    }
    pub fn remove_piece(&mut self, coord: &FieldCoord) {
        assert_eq!(
            self.get_field(coord),
            &Field::Piece,
            "There is no piece at {:?} to remove",
            coord
        );
        self.set_field(coord, Field::Empty);
    }

    fn get_hex(&self, coord: &HexCoord) -> &Option<Hex> {
        &self.board[(coord.x + 2) as usize][(coord.y + 2) as usize]
    }
    fn set_hex(&mut self, coord: &HexCoord, hex: Option<Hex>) {
        self.board[(coord.x + 2) as usize][(coord.y + 2) as usize] = hex;
    }
    fn is_hex_extant(&self, coord: &HexCoord) -> bool {
        self.get_hex(coord).is_some()
    }
    // We return a Vec of tuples so that get_hex_field_neighbors and is_hex_removable know which
    // neighbors are on which side of the hex. They need to know this for different reasons:
    //   * get_hex_field_neighbors: the index of each neighboring field depends on which hex
    //                              neighbor that field neighbor is on
    //   * is_hex_removable: a hex is removable only if it is "attached to the board by 3 or less
    //                       adjacent sides"
    fn get_hex_neighbors(&self, coord: &HexCoord) -> Vec<(u32, HexCoord)> {
        let neighbors = [
            (coord.x, coord.y + 1),
            (coord.x + 1, coord.y),
            (coord.x + 1, coord.y - 1),
            (coord.x, coord.y - 1),
            (coord.x - 1, coord.y),
            (coord.x - 1, coord.y + 1),
        ];

        neighbors
            .iter()
            .enumerate()
            .filter(|&(_, &(x, y))| HexCoord::is_valid_coord(x, y))
            .filter_map(|(i, &(x, y))| {
                let coord = HexCoord::new(x, y);
                if self.is_hex_extant(&coord) {
                    Some((i as u32, coord))
                } else {
                    None
                }
            })
            .collect()
    }
    /// Return fields that share an edge with the given hex and are outside of the given hex. If a
    /// hex is removed from the board, pieces occupying that hex's field neighbors might be
    /// capturable.
    pub fn get_hex_field_neighbors(&self, coord: &HexCoord) -> Vec<FieldCoord> {
        let mut neighbors = vec![];

        for (i, neighbor) in self.get_hex_neighbors(coord) {
            let f = (i + 3) % 6;
            neighbors.push(neighbor.to_field(f));
        }
        neighbors
    }
    /// A hex is removable (and must be removed) if it is empty and is "attached to the board by 3
    /// or less adjacent sides."
    pub fn is_hex_removable(&self, coord: &HexCoord) -> bool {
        match *self.get_hex(coord) {
            Some(hex) => if hex.iter().any(|&f| f == Field::Piece) {
                return false;
            },
            None => panic!("The hex at {:?} has already been removed", coord),
        }

        let neighbor_idxs: Vec<u32> = self.get_hex_neighbors(coord).iter().map(|&(i, _)| i).collect();
        let neighbor_idxs_slice = neighbor_idxs.as_slice();

        match neighbor_idxs_slice.len() {
            0 => panic!("A hex at {:?} is empty and has no neighbors", coord),
            1 => true,
            2 => {
                let valid_idx_combos = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [0, 5]];
                valid_idx_combos.iter().any(|c| c == neighbor_idxs_slice)
            }
            3 => {
                let valid_idx_combos = [
                    [0, 1, 2],
                    [1, 2, 3],
                    [2, 3, 4],
                    [3, 4, 5],
                    [0, 4, 5],
                    [0, 1, 5],
                ];
                valid_idx_combos.iter().any(|c| c == neighbor_idxs_slice)
            }
            _ => false,
        }
    }
    pub fn remove_hex(&mut self, coord: &HexCoord) {
        assert!(
            self.is_hex_removable(coord),
            "Cannot remove the hex at {:?} which is {:?} and has {} neighbors",
            coord,
            self.get_hex(coord),
            self.get_hex_neighbors(coord).len(),
        );
        self.set_hex(coord, None);
    }
    /// > extant (adj.): Still in existence; not destroyed, lost, or extinct (The Free Dictionary)
    ///
    /// Return the coordinates of the hexes that have not been removed yet.
    pub fn extant_hexes(&self) -> Vec<HexCoord> {
        let mut coords = vec![];

        for x in -2..3 {
            for y in -2..3 {
                if HexCoord::is_valid_coord(x, y) {
                    let coord = HexCoord::new(x, y);
                    if self.is_hex_extant(&coord) {
                        coords.push(coord);
                    }
                }
            }
        }
        coords
    }
}

#[derive(Debug, PartialEq)]
pub struct FieldCoord {
    x: i32,
    y: i32,
    f: u32,
}

#[derive(Clone, Copy, Debug)]
pub struct HexCoord {
    x: i32,
    y: i32,
}

impl FieldCoord {
    pub fn new(x: i32, y: i32, f: u32) -> FieldCoord {
        assert!(Self::is_valid_coord(x, y, f));
        FieldCoord { x, y, f }
    }
    pub fn x(&self) -> i32 {
        self.x
    }
    pub fn y(&self) -> i32 {
        self.y
    }
    pub fn f(&self) -> u32 {
        self.f
    }
    pub fn to_hex(&self) -> HexCoord {
        HexCoord::new(self.x, self.y)
    }
    pub fn is_valid_coord(x: i32, y: i32, f: u32) -> bool {
        (x + y).abs() <= 2 && x.abs() <= 2 && y.abs() <= 2 && f < 6
    }
}

impl HexCoord {
    pub fn new(x: i32, y: i32) -> HexCoord {
        assert!(Self::is_valid_coord(x, y));
        HexCoord { x, y }
    }
    pub fn x(&self) -> i32 {
        self.x
    }
    pub fn y(&self) -> i32 {
        self.y
    }
    pub fn to_field(&self, f: u32) -> FieldCoord {
        FieldCoord::new(self.x, self.y, f)
    }
    pub fn is_valid_coord(x: i32, y: i32) -> bool {
        (x + y).abs() <= 2 && x.abs() <= 2 && y.abs() <= 2
    }
}
<commit_msg>Fix set_field<commit_after>/*
 * Copyright (C) 2017 Ryan Huang
 *
 * 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 <https://www.gnu.org/licenses/>.
 */

pub struct Model {
    pub board: Board,
    pub turn: Turn,
    pub white_pieces: u32,
    pub white_hexes: u32,
    pub black_pieces: u32,
    pub black_hexes: u32,
    pub selected_piece: Option<FieldCoord>,
}

impl Model {
    pub fn new() -> Model {
        Model {
            board: Board::new(),
            turn: Turn::White,
            white_pieces: 18,
            white_hexes: 0,
            black_pieces: 18,
            black_hexes: 0,
            selected_piece: None,
        }
    }
}

pub enum Turn {
    White,
    Black,
}

impl Turn {
    pub fn switch_turns(&mut self) {
        *self = match *self {
            Turn::White => Turn::Black,
            Turn::Black => Turn::White,
        }
    }
}

pub struct Board {
    /*
                ____
               /    \
          ____/  +y  \____
         /    \      /    \
        /      \____/  +x  \
        \      /    \      /
         \____/      \____/
         /    \      /    \
        /  -x  \____/      \
        \      /    \      /
         \____/  -y  \____/
              \      /
               \____/

        The hex board uses an axial coordinate system with (0, 0) at the center. The x-axis slopes
        up to the right, and the y-axis goes up and down. The board is stored as a dense 2D array,
        because a ragged array won't work (since each row would have to be a different type).

        See http://www.redblobgames.com/grids/hexagons/#coordinates-axial for more info.
    */
    board: [[Option<Hex>; 5]; 5],
}

// Fields are numbered clockwise from the top. Even indicies are black, odd indicies are white.
type Hex = [Field; 6];

// There's no need to store the color of the piece on this field, since only white pieces can be on
// white fields, and vice versa.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Field {
    Piece,
    Empty,
}

impl Board {
    #[cfg_attr(rustfmt, rustfmt_skip)]
    /// Create a new board with the "Laurentius" starting position.
    pub fn new() -> Board {
        let mut board = Board {
            board: [[None; 5]; 5]
        };

        // (0, 0) is the only empty hex.
        board.set_hex(&HexCoord::new(0, 0), Some([Field::Empty; 6]));

        // All other hexes have exactly two pieces on them in the starting position.
        let piece_locations = [
            (-2,  2, 0, 4),
            (-2,  1, 0, 3),
            (-2,  0, 3, 5),
            (-1,  2, 1, 4),
            (-1,  1, 0, 4),
            (-1,  0, 3, 5),
            (-1, -1, 2, 5),
            ( 0,  2, 1, 5),
            ( 0,  1, 1, 5),
            ( 0, -1, 2, 4),
            ( 0, -2, 2, 4),
            ( 1,  1, 2, 5),
            ( 1,  0, 0, 2),
            ( 1, -1, 1, 3),
            ( 1, -2, 1, 4),
            ( 2,  0, 0, 2),
            ( 2, -1, 0, 3),
            ( 2, -2, 1, 3),
        ];

        for &(x, y, f1, f2) in &piece_locations {
            let mut hex = [Field::Empty; 6];
            hex[f1] = Field::Piece;
            hex[f2] = Field::Piece;
            board.set_hex(&HexCoord::new(x, y), Some(hex));
        }

        board
    }

    pub fn get_field(&self, coord: &FieldCoord) -> &Field {
        match *self.get_hex(&coord.to_hex()) {
            Some(ref hex) => &hex[coord.f as usize],
            None => panic!(
                "Cannot get field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    fn set_field(&mut self, coord: &FieldCoord, field: Field) {
        match *self.get_hex(&coord.to_hex()) {
            Some(mut hex) => {
                hex[coord.f as usize] = field;
                self.set_hex(&coord.to_hex(), Some(hex));
            }
            None => panic!(
                "Cannot set field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    /// Return fields that share an edge with the given field. These fields are always the opposite
    /// color of the given field. If all of a piece's edge neighbors are occupied, that piece might
    /// be capturable.
    pub fn get_field_edge_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        let mut neighbors = vec![
            // There are always two edge neighbors on the same hex as the given field
            FieldCoord::new(coord.x, coord.y, (coord.f + 1) % 6),
            FieldCoord::new(coord.x, coord.y, (coord.f + 5) % 6),
        ];

        for (i, neighbor) in self.get_hex_neighbors(&coord.to_hex()) {
            if coord.f == i {
                let f = (i + 3) % 6;
                neighbors.push(neighbor.to_field(f));
                break;
            }
        }

        neighbors
    }
    /// Return fields that share a vertex with the given field and have the same color as the given
    /// field. Pieces can move to fields that are vertex neighbors of the field they are on.
    pub fn get_field_vertex_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        // A field's vertex neighbors can be defined as the edge neighbors of its edge neighbors
        let mut neighbors = vec![];
        for neighbor in self.get_field_edge_neighbors(coord) {
            neighbors.append(&mut self.get_field_edge_neighbors(&neighbor));
        }
        // A field is not its own neighbor
        neighbors.into_iter().filter(|n| n != coord).collect()
    }
    pub fn move_piece(&mut self, from: &FieldCoord, to: &FieldCoord) {
        assert_eq!(
            self.get_field(from),
            &Field::Piece,
            "There is no piece at {:?} to move",
            from
        );
        assert!(
            self.get_field_vertex_neighbors(from).contains(to),
            "Cannot move piece at {:?} to non-(vertex neighbor) {:?}",
            from,
            to
        );
        assert_eq!(
            self.get_field(to),
            &Field::Empty,
            "Cannot move piece at {:?} to occupied neighbor {:?}",
            from,
            to
        );

        self.set_field(from, Field::Empty);
        self.set_field(to, Field::Piece);
    }
    pub fn remove_piece(&mut self, coord: &FieldCoord) {
        assert_eq!(
            self.get_field(coord),
            &Field::Piece,
            "There is no piece at {:?} to remove",
            coord
        );
        self.set_field(coord, Field::Empty);
    }

    fn get_hex(&self, coord: &HexCoord) -> &Option<Hex> {
        &self.board[(coord.x + 2) as usize][(coord.y + 2) as usize]
    }
    fn set_hex(&mut self, coord: &HexCoord, hex: Option<Hex>) {
        self.board[(coord.x + 2) as usize][(coord.y + 2) as usize] = hex;
    }
    fn is_hex_extant(&self, coord: &HexCoord) -> bool {
        self.get_hex(coord).is_some()
    }
    // We return a Vec of tuples so that get_hex_field_neighbors and is_hex_removable know which
    // neighbors are on which side of the hex. They need to know this for different reasons:
    //   * get_hex_field_neighbors: the index of each neighboring field depends on which hex
    //                              neighbor that field neighbor is on
    //   * is_hex_removable: a hex is removable only if it is "attached to the board by 3 or less
    //                       adjacent sides"
    fn get_hex_neighbors(&self, coord: &HexCoord) -> Vec<(u32, HexCoord)> {
        let neighbors = [
            (coord.x, coord.y + 1),
            (coord.x + 1, coord.y),
            (coord.x + 1, coord.y - 1),
            (coord.x, coord.y - 1),
            (coord.x - 1, coord.y),
            (coord.x - 1, coord.y + 1),
        ];

        neighbors
            .iter()
            .enumerate()
            .filter(|&(_, &(x, y))| HexCoord::is_valid_coord(x, y))
            .filter_map(|(i, &(x, y))| {
                let coord = HexCoord::new(x, y);
                if self.is_hex_extant(&coord) {
                    Some((i as u32, coord))
                } else {
                    None
                }
            })
            .collect()
    }
    /// Return fields that share an edge with the given hex and are outside of the given hex. If a
    /// hex is removed from the board, pieces occupying that hex's field neighbors might be
    /// capturable.
    pub fn get_hex_field_neighbors(&self, coord: &HexCoord) -> Vec<FieldCoord> {
        let mut neighbors = vec![];

        for (i, neighbor) in self.get_hex_neighbors(coord) {
            let f = (i + 3) % 6;
            neighbors.push(neighbor.to_field(f));
        }
        neighbors
    }
    /// A hex is removable (and must be removed) if it is empty and is "attached to the board by 3
    /// or less adjacent sides."
    pub fn is_hex_removable(&self, coord: &HexCoord) -> bool {
        match *self.get_hex(coord) {
            Some(hex) => if hex.iter().any(|&f| f == Field::Piece) {
                return false;
            },
            None => panic!("The hex at {:?} has already been removed", coord),
        }

        let neighbor_idxs: Vec<u32> = self.get_hex_neighbors(coord).iter().map(|&(i, _)| i).collect();
        let neighbor_idxs_slice = neighbor_idxs.as_slice();

        match neighbor_idxs_slice.len() {
            0 => panic!("A hex at {:?} is empty and has no neighbors", coord),
            1 => true,
            2 => {
                let valid_idx_combos = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [0, 5]];
                valid_idx_combos.iter().any(|c| c == neighbor_idxs_slice)
            }
            3 => {
                let valid_idx_combos = [
                    [0, 1, 2],
                    [1, 2, 3],
                    [2, 3, 4],
                    [3, 4, 5],
                    [0, 4, 5],
                    [0, 1, 5],
                ];
                valid_idx_combos.iter().any(|c| c == neighbor_idxs_slice)
            }
            _ => false,
        }
    }
    pub fn remove_hex(&mut self, coord: &HexCoord) {
        assert!(
            self.is_hex_removable(coord),
            "Cannot remove the hex at {:?} which is {:?} and has {} neighbors",
            coord,
            self.get_hex(coord),
            self.get_hex_neighbors(coord).len(),
        );
        self.set_hex(coord, None);
    }
    /// > extant (adj.): Still in existence; not destroyed, lost, or extinct (The Free Dictionary)
    ///
    /// Return the coordinates of the hexes that have not been removed yet.
    pub fn extant_hexes(&self) -> Vec<HexCoord> {
        let mut coords = vec![];

        for x in -2..3 {
            for y in -2..3 {
                if HexCoord::is_valid_coord(x, y) {
                    let coord = HexCoord::new(x, y);
                    if self.is_hex_extant(&coord) {
                        coords.push(coord);
                    }
                }
            }
        }
        coords
    }
}

#[derive(Debug, PartialEq)]
pub struct FieldCoord {
    x: i32,
    y: i32,
    f: u32,
}

#[derive(Clone, Copy, Debug)]
pub struct HexCoord {
    x: i32,
    y: i32,
}

impl FieldCoord {
    pub fn new(x: i32, y: i32, f: u32) -> FieldCoord {
        assert!(Self::is_valid_coord(x, y, f));
        FieldCoord { x, y, f }
    }
    pub fn x(&self) -> i32 {
        self.x
    }
    pub fn y(&self) -> i32 {
        self.y
    }
    pub fn f(&self) -> u32 {
        self.f
    }
    pub fn to_hex(&self) -> HexCoord {
        HexCoord::new(self.x, self.y)
    }
    pub fn is_valid_coord(x: i32, y: i32, f: u32) -> bool {
        (x + y).abs() <= 2 && x.abs() <= 2 && y.abs() <= 2 && f < 6
    }
}

impl HexCoord {
    pub fn new(x: i32, y: i32) -> HexCoord {
        assert!(Self::is_valid_coord(x, y));
        HexCoord { x, y }
    }
    pub fn x(&self) -> i32 {
        self.x
    }
    pub fn y(&self) -> i32 {
        self.y
    }
    pub fn to_field(&self, f: u32) -> FieldCoord {
        FieldCoord::new(self.x, self.y, f)
    }
    pub fn is_valid_coord(x: i32, y: i32) -> bool {
        (x + y).abs() <= 2 && x.abs() <= 2 && y.abs() <= 2
    }
}
<|endoftext|>