<commit_before>
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)
}

// IntValue:   (+|-)?(0|[1-9][0-9]*)
// FloatValue: (+|-)?(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, '+' | '-');
  // IntValueeger 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>> {
    use self::Token::*;
    use self::Punctuator::*;

    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(Punctuator(Bang, self.lo)))
        }
        '$' => {
          self.iter.next();
          Some(Ok(Punctuator(Dollar, self.lo)))
        }
        '(' => {
          self.iter.next();
          Some(Ok(Punctuator(LeftParen, self.lo)))
        }
        ')' => {
          self.iter.next();
          Some(Ok(Punctuator(RightParen, self.lo)))
        }
        ':' => {
          self.iter.next();
          Some(Ok(Punctuator(Colon, self.lo)))
        }
        '=' => {
          self.iter.next();
          Some(Ok(Punctuator(Equals, self.lo)))
        }
        '@' => {
          self.iter.next();
          Some(Ok(Punctuator(At, self.lo)))
        }
        '[' => {
          self.iter.next();
          Some(Ok(Punctuator(LeftBracket, self.lo)))
        }
        ']' => {
          self.iter.next();
          Some(Ok(Punctuator(RightBracket, self.lo)))
        }
        '{' => {
          self.iter.next();
          Some(Ok(Punctuator(LeftBrace, self.lo)))
        }
        '}' => {
          self.iter.next();
          Some(Ok(Punctuator(RightBrace, self.lo)))
        }
        '|' => {
          self.iter.next();
          Some(Ok(Punctuator(Pipe, self.lo)))
        }
        '.' => {
          if scan_spread(self) {
            Some(Ok(Punctuator(Spread(self.lo), self.hi)))
          } else {
            Some(Err("Unexpected character."))
          }
        }
        '_' | 'a'...'z' | 'A'...'Z' => {
          scan_name(self);
          Some(Ok(Name(&self.input[self.lo..self.hi], 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(IntValue(&self.input[self.lo..self.hi], self.lo, self.hi))),
            (true, true) => Some(Ok(FloatValue(&self.input[self.lo..self.hi], self.lo, self.hi))),
          }
        }
        '"' => {
          match scan_string(self) {
            (false, _) => Some(Err("Unterminated string.")),
            (_, true) => Some(Err("Bad character escape sequence.")),
            _ => Some(Ok(StringValue(&self.input[self.lo..self.hi], self.lo, self.hi))),
          }
        }
        _ => Some(Err("Unexpected character.")),
      }
    }
    item
  }
}

#[cfg(test)]
mod tests {
  use super::Punctuator::*;
  use super::Token::*;
  use super::{Lexer, Token};

  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, StringValue(expeced, 1, len + 1))
  }

  #[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, Name("name1", 5, 10));
    test_next_token(&mut lexer, Name("name2", 14, 19));
    test_next_token(&mut lexer, StringValue("simple", 23, 29));
    test_next_token(&mut lexer, StringValue("  white space   ", 32, 48));
    test_next_token(&mut lexer, Name("other_name", 50, 60));
    test_next_token(&mut lexer, 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, IntValue("-9", 0, 2));
    test_next_token(&mut lexer, IntValue("0", 3, 4));
    test_next_token(&mut lexer, IntValue("9", 5, 6));
    test_next_token(&mut lexer, IntValue("-0", 7, 9));
    test_next_token(&mut lexer, IntValue("-9", 10, 12));
    test_next_token(&mut lexer, IntValue("1234", 13, 17));
    test_next_token(&mut lexer, IntValue("-1234", 18, 23));
    test_next_token(&mut lexer, FloatValue("0.0", 24, 27));
    test_next_token(&mut lexer, FloatValue("-0.0", 28, 32));
    test_next_token(&mut lexer, FloatValue("1.0", 33, 36));
    test_next_token(&mut lexer, FloatValue("-1.012", 37, 43));
    test_next_token(&mut lexer, FloatValue("-0.101", 44, 50));
    test_next_token(&mut lexer, FloatValue("1.0e2", 51, 56));
    test_next_token(&mut lexer, FloatValue("-1.0E2", 57, 63));
    test_next_token(&mut lexer, FloatValue("0.00e-002", 64, 73));
    test_next_token(&mut lexer, FloatValue("-0.00E+002", 74, 84));
    test_next_token(&mut lexer, FloatValue("9.1E+0", 85, 91));
    test_next_token(&mut lexer, Eof(85,91));
    assert_eq!(None, lexer.next());
  }

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

// IntValue:   (+|-)?(0|[1-9][0-9]*)
// FloatValue: (+|-)?(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>> {
    use self::Token::*;
    use self::Punctuator::*;

    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(Punctuator(Bang, self.lo)))
        }
        '$' => {
          self.iter.next();
          Some(Ok(Punctuator(Dollar, self.lo)))
        }
        '(' => {
          self.iter.next();
          Some(Ok(Punctuator(LeftParen, self.lo)))
        }
        ')' => {
          self.iter.next();
          Some(Ok(Punctuator(RightParen, self.lo)))
        }
        ':' => {
          self.iter.next();
          Some(Ok(Punctuator(Colon, self.lo)))
        }
        '=' => {
          self.iter.next();
          Some(Ok(Punctuator(Equals, self.lo)))
        }
        '@' => {
          self.iter.next();
          Some(Ok(Punctuator(At, self.lo)))
        }
        '[' => {
          self.iter.next();
          Some(Ok(Punctuator(LeftBracket, self.lo)))
        }
        ']' => {
          self.iter.next();
          Some(Ok(Punctuator(RightBracket, self.lo)))
        }
        '{' => {
          self.iter.next();
          Some(Ok(Punctuator(LeftBrace, self.lo)))
        }
        '}' => {
          self.iter.next();
          Some(Ok(Punctuator(RightBrace, self.lo)))
        }
        '|' => {
          self.iter.next();
          Some(Ok(Punctuator(Pipe, self.lo)))
        }
        '.' => {
          if scan_spread(self) {
            Some(Ok(Punctuator(Spread(self.lo), self.hi)))
          } else {
            Some(Err("Unexpected character."))
          }
        }
        '_' | 'a'...'z' | 'A'...'Z' => {
          scan_name(self);
          Some(Ok(Name(&self.input[self.lo..self.hi], 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(IntValue(&self.input[self.lo..self.hi], self.lo, self.hi))),
            (true, true) => Some(Ok(FloatValue(&self.input[self.lo..self.hi], self.lo, self.hi))),
          }
        }
        '"' => {
          match scan_string(self) {
            (false, _) => Some(Err("Unterminated string.")),
            (_, true) => Some(Err("Bad character escape sequence.")),
            _ => Some(Ok(StringValue(&self.input[self.lo..self.hi], self.lo, self.hi))),
          }
        }
        _ => Some(Err("Unexpected character.")),
      }
    }
    item
  }
}

#[cfg(test)]
mod tests {
  use super::Punctuator::*;
  use super::Token::*;
  use super::{Lexer, Token};

  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, StringValue(expeced, 1, len + 1))
  }

  #[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, Name("name1", 5, 10));
    test_next_token(&mut lexer, Name("name2", 14, 19));
    test_next_token(&mut lexer, StringValue("simple", 23, 29));
    test_next_token(&mut lexer, StringValue("  white space   ", 32, 48));
    test_next_token(&mut lexer, Name("other_name", 50, 60));
    test_next_token(&mut lexer, 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, IntValue("-9", 0, 2));
    test_next_token(&mut lexer, IntValue("0", 3, 4));
    test_next_token(&mut lexer, IntValue("9", 5, 6));
    test_next_token(&mut lexer, IntValue("-0", 7, 9));
    test_next_token(&mut lexer, IntValue("-9", 10, 12));
    test_next_token(&mut lexer, IntValue("1234", 13, 17));
    test_next_token(&mut lexer, IntValue("-1234", 18, 23));
    test_next_token(&mut lexer, FloatValue("0.0", 24, 27));
    test_next_token(&mut lexer, FloatValue("-0.0", 28, 32));
    test_next_token(&mut lexer, FloatValue("1.0", 33, 36));
    test_next_token(&mut lexer, FloatValue("-1.012", 37, 43));
    test_next_token(&mut lexer, FloatValue("-0.101", 44, 50));
    test_next_token(&mut lexer, FloatValue("1.0e2", 51, 56));
    test_next_token(&mut lexer, FloatValue("-1.0E2", 57, 63));
    test_next_token(&mut lexer, FloatValue("0.00e-002", 64, 73));
    test_next_token(&mut lexer, FloatValue("-0.00E+002", 74, 84));
    test_next_token(&mut lexer, FloatValue("9.1E+0", 85, 91));
    test_next_token(&mut lexer, Eof(85,91));
    assert_eq!(None, lexer.next());
  }

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


<commit_before>use crate::prolog::machine::machine_indices::*;
use crate::prolog::machine::stack::*;
use crate::prolog::machine::streams::*;

use std::mem;
use std::ops::IndexMut;

type Trail = Vec<(Ref, HeapCellValue)>;

#[derive(Clone, Copy)]
pub enum AttrVarPolicy {
    DeepCopy,
    StripAttributes
}

pub(crate)
trait CopierTarget: IndexMut<usize, Output = HeapCellValue> {
    fn deref(&self, val: Addr) -> Addr;
    fn push(&mut self, val: HeapCellValue);
    fn stack(&mut self) -> &mut Stack;
    fn store(&self, val: Addr) -> Addr;
    fn threshold(&self) -> usize;
}

pub(crate)
fn copy_term<T: CopierTarget>(target: T, addr: Addr, attr_var_policy: AttrVarPolicy) {
    let mut copy_term_state = CopyTermState::new(target, attr_var_policy);
    copy_term_state.copy_term_impl(addr);
}

struct CopyTermState<T: CopierTarget> {
    trail: Trail,
    scan: usize,
    old_h: usize,
    target: T,
    attr_var_policy: AttrVarPolicy,
}

impl<T: CopierTarget> CopyTermState<T> {
    fn new(target: T, attr_var_policy: AttrVarPolicy) -> Self {
        CopyTermState {
            trail: Trail::new(),
            scan: 0,
            old_h: target.threshold(),
            target,
            attr_var_policy
        }
    }

    #[inline]
    fn value_at_scan(&mut self) -> &mut HeapCellValue {
        let scan = self.scan;
        &mut self.target[scan]
    }

    fn copy_list(&mut self, addr: usize) {
        if let Addr::Lis(h) = self.target[addr + 1].as_addr(addr + 1) {
            if h >= self.old_h {
                *self.value_at_scan() = HeapCellValue::Addr(Addr::Lis(h));
                self.scan += 1;

                return;
            }
        }

        let threshold = self.target.threshold();

        *self.value_at_scan() = HeapCellValue::Addr(Addr::Lis(threshold));

        for i in 0 .. 2 {
            let hcv = self.target[addr + i].context_free_clone();
            self.target.push(hcv);
        }

        let cdr = self.target.store(self.target.deref(Addr::HeapCell(addr + 1)));

        if let Addr::Lis(_) = cdr {
            let tail_addr = self.target[addr + 1].as_addr(addr + 1);

            self.trail.push((
                Ref::HeapCell(addr + 1),
                HeapCellValue::Addr(tail_addr),
            ));

            self.target[addr + 1] = HeapCellValue::Addr(Addr::Lis(threshold));
        }

        self.scan += 1;
    }

    fn copy_partial_string(&mut self, addr: usize, n: usize) {
        if let &HeapCellValue::Addr(Addr::PStrLocation(h, _)) = &self.target[addr] {
            if h >= self.old_h {
                *self.value_at_scan() = HeapCellValue::Addr(Addr::PStrLocation(h, n));
                self.scan += 1;

                return;
            }
        }

        let threshold = self.target.threshold();

        *self.value_at_scan() =
            HeapCellValue::Addr(Addr::PStrLocation(threshold, 0));

        self.scan += 1;

        let (pstr, has_tail) =
            match &self.target[addr] {
                &HeapCellValue::PartialString(ref pstr, has_tail) => {
                    (pstr.clone_from_offset(n), has_tail)
                }
                _ => {
                    unreachable!()
                }
            };

        self.target.push(HeapCellValue::PartialString(pstr, has_tail));

        let replacement = HeapCellValue::Addr(Addr::PStrLocation(threshold, 0));

        let trail_item = mem::replace(
            &mut self.target[addr],
            replacement,
        );

        self.trail.push((
            Ref::HeapCell(addr),
            trail_item,
        ));

        if has_tail {
            let tail_addr = self.target[addr + 1].as_addr(addr + 1);
            self.target.push(HeapCellValue::Addr(tail_addr));
        }
    }

    fn reinstantiate_var(&mut self, addr: Addr, frontier: usize) {
        match addr {
            Addr::HeapCell(h) => {
                self.target[frontier] = HeapCellValue::Addr(Addr::HeapCell(frontier));
                self.target[h] = HeapCellValue::Addr(Addr::HeapCell(frontier));

                self.trail.push((
                    Ref::HeapCell(h),
                    HeapCellValue::Addr(Addr::HeapCell(h)),
                ));
            }
            Addr::StackCell(fr, sc) => {
                self.target[frontier] = HeapCellValue::Addr(Addr::HeapCell(frontier));
                self.target.stack().index_and_frame_mut(fr)[sc] = Addr::HeapCell(frontier);

                self.trail.push((
                    Ref::StackCell(fr, sc),
                    HeapCellValue::Addr(Addr::StackCell(fr, sc)),
                ));
            }
            Addr::AttrVar(h) => {
                let threshold = if let AttrVarPolicy::DeepCopy = self.attr_var_policy {
                    self.target.threshold()
                } else {
                    frontier
                };

                self.target[frontier] = HeapCellValue::Addr(Addr::HeapCell(threshold));
                self.target[h] = HeapCellValue::Addr(Addr::HeapCell(frontier));

                self.trail.push((
                    Ref::AttrVar(h),
                    HeapCellValue::Addr(Addr::AttrVar(h)),
                ));

                if let AttrVarPolicy::DeepCopy = self.attr_var_policy {
                    self.target.push(HeapCellValue::Addr(Addr::AttrVar(threshold)));

                    let list_val = self.target[h + 1].context_free_clone();
                    self.target.push(list_val);
                }
            }
            _ => {
                unreachable!()
            }
        }
    }

    fn copy_var(&mut self, addr: Addr) {
        let rd = self.target.store(self.target.deref(addr));

        match rd {
            Addr::AttrVar(h) | Addr::HeapCell(h) if h >= self.old_h => {
                *self.value_at_scan() = HeapCellValue::Addr(rd);
                self.scan += 1;
            }
            _ if addr == rd => {
                self.reinstantiate_var(addr, self.scan);
                self.scan += 1;
            }
            _ => {
                *self.value_at_scan() = HeapCellValue::Addr(rd);
            }
        }
    }

    fn copy_stream(&mut self, addr: usize) {
        let threshold = self.target.threshold();

        let trail_item = mem::replace(
            &mut self.target[addr],
            HeapCellValue::Addr(Addr::Stream(threshold)),
        );

        self.trail.push((
            Ref::HeapCell(addr),
            trail_item,
        ));

        self.target.push(HeapCellValue::Stream(Stream::null_stream()));

        self.scan += 1;
    }

    fn copy_structure(&mut self, addr: usize) {
        match self.target[addr].context_free_clone() {
            HeapCellValue::NamedStr(arity, name, fixity) => {
                let threshold = self.target.threshold();

                *self.value_at_scan() = HeapCellValue::Addr(Addr::Str(threshold));

                let trail_item = mem::replace(
                    &mut self.target[addr],
                    HeapCellValue::Addr(Addr::Str(threshold)),
                );

                self.trail.push((
                    Ref::HeapCell(addr),
                    trail_item,
                ));

                self.target.push(HeapCellValue::NamedStr(arity, name, fixity));

                for i in 0..arity {
                    let hcv = self.target[addr + 1 + i].context_free_clone();
                    self.target.push(hcv);
                }
            }
            HeapCellValue::Addr(Addr::Str(addr)) => {
                *self.value_at_scan() = HeapCellValue::Addr(Addr::Str(addr))
            }
            _ => {
                unreachable!()
            }
        }

        self.scan += 1;
    }

    fn copy_term_impl(&mut self, addr: Addr) {
        self.scan = self.target.threshold();
        self.target.push(HeapCellValue::Addr(addr));

        while self.scan < self.target.threshold() {
            match self.value_at_scan() {
                &mut HeapCellValue::Addr(addr) => {
                    match addr {
                        Addr::Con(h) => {
                            let addr = self.target[h].as_addr(h);

                            if addr == Addr::Con(h) {
                                *self.value_at_scan() = self.target[h].context_free_clone();
                            } else {
                                *self.value_at_scan() = HeapCellValue::Addr(addr);
                            }
                        }
                        Addr::Lis(h) => {
                            self.copy_list(h);
                        }
                        addr @ Addr::AttrVar(_) |
                        addr @ Addr::HeapCell(_) |
                        addr @ Addr::StackCell(..) => {
                            self.copy_var(addr);
                        }
                        Addr::Str(addr) => {
                            self.copy_structure(addr);
                        }
                        Addr::PStrLocation(addr, n) => {
                            self.copy_partial_string(addr, n);
                        }
                        Addr::Stream(h) => {
                            self.copy_stream(h);
                        }
                        _ => {
                            self.scan += 1;
                        }
                    }
                }
                _ => {
                    self.scan += 1;
                }
            }
        }

        self.unwind_trail();
    }

    fn unwind_trail(&mut self) {
        for (r, value) in self.trail.drain(0..) {
            match r {
                Ref::AttrVar(h) | Ref::HeapCell(h) =>
                    self.target[h] = value,
                Ref::StackCell(fr, sc) =>
                    self.target.stack().index_and_frame_mut(fr)[sc] = value.as_addr(0),
            }
        }
    }
}
<commit_msg>use copy_term/3 to print residual goals (#254)<commit_after>use crate::prolog::machine::machine_indices::*;
use crate::prolog::machine::stack::*;
use crate::prolog::machine::streams::*;

use std::mem;
use std::ops::IndexMut;

type Trail = Vec<(Ref, HeapCellValue)>;

#[derive(Clone, Copy)]
pub enum AttrVarPolicy {
    DeepCopy,
    StripAttributes
}

pub(crate)
trait CopierTarget: IndexMut<usize, Output = HeapCellValue> {
    fn deref(&self, val: Addr) -> Addr;
    fn push(&mut self, val: HeapCellValue);
    fn stack(&mut self) -> &mut Stack;
    fn store(&self, val: Addr) -> Addr;
    fn threshold(&self) -> usize;
}

pub(crate)
fn copy_term<T: CopierTarget>(target: T, addr: Addr, attr_var_policy: AttrVarPolicy) {
    let mut copy_term_state = CopyTermState::new(target, attr_var_policy);
    copy_term_state.copy_term_impl(addr);
}

struct CopyTermState<T: CopierTarget> {
    trail: Trail,
    scan: usize,
    old_h: usize,
    target: T,
    attr_var_policy: AttrVarPolicy,
}

impl<T: CopierTarget> CopyTermState<T> {
    fn new(target: T, attr_var_policy: AttrVarPolicy) -> Self {
        CopyTermState {
            trail: Trail::new(),
            scan: 0,
            old_h: target.threshold(),
            target,
            attr_var_policy
        }
    }

    #[inline]
    fn value_at_scan(&mut self) -> &mut HeapCellValue {
        let scan = self.scan;
        &mut self.target[scan]
    }

    fn trail_list_cell(&mut self, addr: usize, threshold: usize) {
        let trail_item = mem::replace(
            &mut self.target[addr],
            HeapCellValue::Addr(Addr::Lis(threshold)),
        );

        self.trail.push((
            Ref::HeapCell(addr),
            trail_item,
        ));
    }

    fn copy_list(&mut self, addr: usize) {
        for offset in 0 .. 2 {
            if let Addr::Lis(h) = self.target[addr + offset].as_addr(addr + offset) {
                if h >= self.old_h {
                    *self.value_at_scan() = HeapCellValue::Addr(Addr::Lis(h));
                    self.scan += 1;

                    return;
                }
            }
        }

        let threshold = self.target.threshold();

        *self.value_at_scan() = HeapCellValue::Addr(Addr::Lis(threshold));

        for i in 0 .. 2 {
            let hcv = self.target[addr + i].context_free_clone();
            self.target.push(hcv);
        }

        let cdr = self.target.store(self.target.deref(Addr::HeapCell(addr + 1)));

        if !cdr.is_ref() {
            self.trail_list_cell(addr + 1, threshold);
        } else {
            let car = self.target.store(self.target.deref(Addr::HeapCell(addr)));

            if !car.is_ref() {
                self.trail_list_cell(addr, threshold);
            }
        }

        self.scan += 1;
    }

    fn copy_partial_string(&mut self, addr: usize, n: usize) {
        if let &HeapCellValue::Addr(Addr::PStrLocation(h, _)) = &self.target[addr] {
            if h >= self.old_h {
                *self.value_at_scan() = HeapCellValue::Addr(Addr::PStrLocation(h, n));
                self.scan += 1;

                return;
            }
        }

        let threshold = self.target.threshold();

        *self.value_at_scan() =
            HeapCellValue::Addr(Addr::PStrLocation(threshold, 0));

        self.scan += 1;

        let (pstr, has_tail) =
            match &self.target[addr] {
                &HeapCellValue::PartialString(ref pstr, has_tail) => {
                    (pstr.clone_from_offset(n), has_tail)
                }
                _ => {
                    unreachable!()
                }
            };

        self.target.push(HeapCellValue::PartialString(pstr, has_tail));

        let replacement = HeapCellValue::Addr(Addr::PStrLocation(threshold, 0));

        let trail_item = mem::replace(
            &mut self.target[addr],
            replacement,
        );

        self.trail.push((
            Ref::HeapCell(addr),
            trail_item,
        ));

        if has_tail {
            let tail_addr = self.target[addr + 1].as_addr(addr + 1);
            self.target.push(HeapCellValue::Addr(tail_addr));
        }
    }

    fn reinstantiate_var(&mut self, addr: Addr, frontier: usize) {
        match addr {
            Addr::HeapCell(h) => {
                self.target[frontier] = HeapCellValue::Addr(Addr::HeapCell(frontier));
                self.target[h] = HeapCellValue::Addr(Addr::HeapCell(frontier));

                self.trail.push((
                    Ref::HeapCell(h),
                    HeapCellValue::Addr(Addr::HeapCell(h)),
                ));
            }
            Addr::StackCell(fr, sc) => {
                self.target[frontier] = HeapCellValue::Addr(Addr::HeapCell(frontier));
                self.target.stack().index_and_frame_mut(fr)[sc] = Addr::HeapCell(frontier);

                self.trail.push((
                    Ref::StackCell(fr, sc),
                    HeapCellValue::Addr(Addr::StackCell(fr, sc)),
                ));
            }
            Addr::AttrVar(h) => {
                let threshold = if let AttrVarPolicy::DeepCopy = self.attr_var_policy {
                    self.target.threshold()
                } else {
                    frontier
                };

                self.target[frontier] = HeapCellValue::Addr(Addr::HeapCell(threshold));
                self.target[h] = HeapCellValue::Addr(Addr::HeapCell(threshold));

                self.trail.push((
                    Ref::AttrVar(h),
                    HeapCellValue::Addr(Addr::AttrVar(h)),
                ));

                if let AttrVarPolicy::DeepCopy = self.attr_var_policy {
                    self.target.push(HeapCellValue::Addr(Addr::AttrVar(threshold)));

                    let list_val = self.target[h + 1].context_free_clone();
                    self.target.push(list_val);
                }
            }
            _ => {
                unreachable!()
            }
        }
    }

    fn copy_var(&mut self, addr: Addr) {
        let rd = self.target.store(self.target.deref(addr));

        match rd {
            Addr::AttrVar(h) | Addr::HeapCell(h) if h >= self.old_h => {
                *self.value_at_scan() = HeapCellValue::Addr(rd);
                self.scan += 1;
            }
            _ if addr == rd => {
                self.reinstantiate_var(addr, self.scan);
                self.scan += 1;
            }
            _ => {
                *self.value_at_scan() = HeapCellValue::Addr(rd);
            }
        }
    }

    fn copy_stream(&mut self, addr: usize) {
        let threshold = self.target.threshold();

        let trail_item = mem::replace(
            &mut self.target[addr],
            HeapCellValue::Addr(Addr::Stream(threshold)),
        );

        self.trail.push((
            Ref::HeapCell(addr),
            trail_item,
        ));

        self.target.push(HeapCellValue::Stream(Stream::null_stream()));

        self.scan += 1;
    }

    fn copy_structure(&mut self, addr: usize) {
        match self.target[addr].context_free_clone() {
            HeapCellValue::NamedStr(arity, name, fixity) => {
                let threshold = self.target.threshold();

                *self.value_at_scan() = HeapCellValue::Addr(Addr::Str(threshold));

                let trail_item = mem::replace(
                    &mut self.target[addr],
                    HeapCellValue::Addr(Addr::Str(threshold)),
                );

                self.trail.push((
                    Ref::HeapCell(addr),
                    trail_item,
                ));

                self.target.push(HeapCellValue::NamedStr(arity, name, fixity));

                for i in 0..arity {
                    let hcv = self.target[addr + 1 + i].context_free_clone();
                    self.target.push(hcv);
                }
            }
            HeapCellValue::Addr(Addr::Str(addr)) => {
                *self.value_at_scan() = HeapCellValue::Addr(Addr::Str(addr))
            }
            _ => {
                unreachable!()
            }
        }

        self.scan += 1;
    }

    fn copy_term_impl(&mut self, addr: Addr) {
        self.scan = self.target.threshold();
        self.target.push(HeapCellValue::Addr(addr));

        while self.scan < self.target.threshold() {
            match self.value_at_scan() {
                &mut HeapCellValue::Addr(addr) => {
                    match addr {
                        Addr::Con(h) => {
                            let addr = self.target[h].as_addr(h);

                            if addr == Addr::Con(h) {
                                *self.value_at_scan() = self.target[h].context_free_clone();
                            } else {
                                *self.value_at_scan() = HeapCellValue::Addr(addr);
                            }
                        }
                        Addr::Lis(h) if h >= self.old_h => {
                            self.scan += 1;
                        }
                        Addr::Lis(h) => {
                            self.copy_list(h);
                        }
                        addr @ Addr::AttrVar(_) |
                        addr @ Addr::HeapCell(_) |
                        addr @ Addr::StackCell(..) => {
                            self.copy_var(addr);
                        }
                        Addr::Str(addr) => {
                            self.copy_structure(addr);
                        }
                        Addr::PStrLocation(addr, n) => {
                            self.copy_partial_string(addr, n);
                        }
                        Addr::Stream(h) => {
                            self.copy_stream(h);
                        }
                        _ => {
                            self.scan += 1;
                        }
                    }
                }
                _ => {
                    self.scan += 1;
                }
            }
        }

        self.unwind_trail();
    }

    fn unwind_trail(&mut self) {
        for (r, value) in self.trail.drain(0..) {
            match r {
                Ref::AttrVar(h) | Ref::HeapCell(h) =>
                    self.target[h] = value,
                Ref::StackCell(fr, sc) =>
                    self.target.stack().index_and_frame_mut(fr)[sc] = value.as_addr(0),
            }
        }
    }
}
<|endoftext|>


<commit_before>use decoders::Image;
extern crate itertools;
use self::itertools::Itertools;

pub fn camera_to_lab(img: &Image, inb: &[f32]) -> Vec<f32> {
  let mut out: Vec<f32> = vec![0.0; (img.width*img.height*3) as usize];
  let cmatrix = cam_to_xyz_matrix(img);

  let mut opos = 0;
  for pos in (0..(img.height*img.width*4)).step(4) {
    let cr = inb[pos];
    let cg = inb[pos+1];
    let cb = inb[pos+2];
    let ce = inb[pos+3];

    let x = cr * cmatrix[0][0] + cg * cmatrix[0][1] + cb * cmatrix[0][2] + ce * cmatrix[0][3];
    let y = cr * cmatrix[1][0] + cg * cmatrix[1][1] + cb * cmatrix[1][2] + ce * cmatrix[1][3];
    let z = cr * cmatrix[2][0] + cg * cmatrix[2][1] + cb * cmatrix[2][2] + ce * cmatrix[2][3];

    let (l,a,b) = xyz_to_lab(x,y,z);

    out[opos+0] = l;
    out[opos+1] = a;
    out[opos+2] = b;

    opos += 3;
  }

  out
}

pub fn lab_to_rec709(img: &Image, buf: &mut Vec<f32>) {
  let cmatrix = xyz_to_rec709_matrix();

  for pos in (0..(img.height*img.width*3)).step(3) {
    let l = buf[pos];
    let a = buf[pos+1];
    let b = buf[pos+2];

    let (x,y,z) = lab_to_xyz(l,a,b);

    let r = x * cmatrix[0][0] + y * cmatrix[0][1] + z * cmatrix[0][2];
    let g = x * cmatrix[1][0] + y * cmatrix[1][1] + z * cmatrix[1][2];
    let b = x * cmatrix[2][0] + y * cmatrix[2][1] + z * cmatrix[2][2];

    buf[pos+0] = r;
    buf[pos+1] = g;
    buf[pos+2] = b;
  }
}

fn inverse(inm: [[f32;3];3]) -> [[f32;3];3] {
  let invdet = 1.0 / (
    inm[0][0] * (inm[1][1] * inm[2][2] - inm[2][1] * inm[1][2]) -
    inm[0][1] * (inm[1][0] * inm[2][2] - inm[1][2] * inm[2][0]) +
    inm[0][2] * (inm[1][0] * inm[2][1] - inm[1][1] * inm[2][0])
  );

  let mut out = [[0.0; 3];3];
  out[0][0] =  (inm[1][1]*inm[2][2] - inm[2][1]*inm[1][2]) * invdet;
  out[0][1] = -(inm[0][1]*inm[2][2] - inm[0][2]*inm[2][1]) * invdet;
  out[0][2] =  (inm[0][1]*inm[1][2] - inm[0][2]*inm[1][1]) * invdet;
  out[1][0] = -(inm[1][0]*inm[2][2] - inm[1][2]*inm[2][0]) * invdet;
  out[1][1] =  (inm[0][0]*inm[2][2] - inm[0][2]*inm[2][0]) * invdet;
  out[1][2] = -(inm[0][0]*inm[1][2] - inm[1][0]*inm[0][2]) * invdet;
  out[2][0] =  (inm[1][0]*inm[2][1] - inm[2][0]*inm[1][1]) * invdet;
  out[2][1] = -(inm[0][0]*inm[2][1] - inm[2][0]*inm[0][1]) * invdet;
  out[2][2] =  (inm[0][0]*inm[1][1] - inm[1][0]*inm[0][1]) * invdet;

  out
}

fn pseudoinverse(inm: [[f32;3];4]) -> [[f32;4];3] {
  let mut temp: [[f32;6];3] = [[0.0; 6];3];

  for i in 0..3 {
    for j in 0..6 {
      temp[i][j] = if j == i+3 { 1.0 } else { 0.0 };
    }
    for j in 0..3 {
      for k in 0..4 {
        temp[i][j] += inm[k][i] * inm[k][j];
      }
    }
  }

  for i in 0..3 {
    let mut num = temp[i][i];
    for j in 0..6 {
      temp[i][j] /= num;
    }
    for k in 0..3 {
      if k == i { continue }
      num = temp[k][i];
      for j in 0..6 {
        temp[k][j] -= temp[i][j] * num;
      }
    }
  }

  let mut out: [[f32;4];3] = [[0.0; 4];3];

  for i in 0..4 {
    for j in 0..3 {
      out[j][i] = 0.0;
      for k in 0..3 {
        out[j][i] += temp[j][k+3] * inm[i][k];
      }
    }
  }

  out
}

fn cam_to_xyz_matrix(img: &Image) -> [[f32;4];3] {
  let mut xyz_to_cam = [[0.0;3];4];
  for i in 0..12 {
    xyz_to_cam[i/3][i%3] = img.color_matrix[i] as f32;
  }

  let d65_white = [0.9547,1.0,1.08883];
  let rgb_to_xyz = [
  // sRGB D65
    [ 0.412453, 0.357580, 0.180423 ],
    [ 0.212671, 0.715160, 0.072169 ],
    [ 0.019334, 0.119193, 0.950227 ],
  ];

  // Multiply RGB matrix
  let mut rgb_to_cam = [[0.0;3];4];
  for i in 0..4 {
    for j in 0..3 {
      rgb_to_cam[i][j] = 0.0;
      for k in 0..3 {
        rgb_to_cam[i][j] += xyz_to_cam[i][k] * rgb_to_xyz[k][j];
      }
    }
  }

  // Normalize rgb_to_cam so that rgb_to_cam * (1,1,1) is (1,1,1,1)
  for i in 0..4 {
    let mut num = 0.0;
    for j in 0..3 {
      num += rgb_to_cam[i][j];
    }
    for j in 0..3 {
      rgb_to_cam[i][j] = if num == 0.0 {
        0.0
      }  else {
        rgb_to_cam[i][j] / num
      };
    }
  }

  let cam_to_rgb = pseudoinverse(rgb_to_cam);
  let mut cam_to_xyz = [[0.0;4];3];
  // Multiply RGB matrix and adjust white to get a cam_to_xyz
  for i in 0..3 {
    for j in 0..4 {
      cam_to_xyz[i][j] = 0.0;
      for k in 0..3 {
        cam_to_xyz[i][j] += cam_to_rgb[k][j] * rgb_to_xyz[i][k] / d65_white[i];
      }
    }
  }

  cam_to_xyz
}

fn xyz_to_rec709_matrix() -> [[f32;3];3] {
  let rgb_to_xyz = [
  // sRGB D65
    [ 0.412453, 0.357580, 0.180423 ],
    [ 0.212671, 0.715160, 0.072169 ],
    [ 0.019334, 0.119193, 0.950227 ],
  ];

  inverse(rgb_to_xyz)
}

fn xyz_to_lab(x: f32, y: f32, z: f32) -> (f32,f32,f32) {
  // D50 White
  let xw = 0.9642; let yw = 1.000; let zw = 0.8249;

  let l = 116.0 * labf(y/yw) - 16.0;
  let a = 500.0 * (labf(x/xw) - labf(y/yw));
  let b = 200.0 * (labf(y/yw) - labf(z/zw));

  (l/100.0,(a+128.0)/256.0,(b+128.0)/256.0)
}

static CBRT_MAXVALS: usize = 1 << 16; // 2^16 should be enough precision
lazy_static! {
  static ref CBRT_LOOKUP: Vec<f32> = {
    let mut lookup: Vec<f32> = vec![0.0; CBRT_MAXVALS+1];
    for i in 0..(CBRT_MAXVALS+1) {
      let v = (i as f32) / (CBRT_MAXVALS as f32);
      lookup[i] = v.cbrt();
    }
    lookup
  };
}

fn labf(val: f32) -> f32 {
  let cutoff = (6.0/29.0)*(6.0/29.0)*(6.0/29.0);
  let multiplier = (1.0/3.0) * (29.0/6.0) * (29.0/6.0);
  let constant = 4.0 / 29.0;

  if val > cutoff {
    if val > 0.0 && val < 1.0 { // use the lookup table
      CBRT_LOOKUP[(val*(CBRT_MAXVALS as f32)) as usize]
    } else {
      val.cbrt()
    }
  } else {
    val * multiplier + constant
  }
}

fn lab_to_xyz(l: f32, a: f32, b: f32) -> (f32,f32,f32) {
  // D50 White
  let xw = 0.9642; let yw = 1.000; let zw = 0.8249;

  let cl = l * 100.0;
  let ca = (a * 256.0) - 128.0;
  let cb = (b * 256.0) - 128.0;

  let x = xw * labinvf((1.0/116.0) * (cl+16.0) + (1.0/500.0) * ca);
  let y = yw * labinvf((1.0/116.0) * (cl+16.0));
  let z = zw * labinvf((1.0/116.0) * (cl+16.0) - (1.0/200.0) * cb);

  (x,y,z)
}

fn labinvf(val: f32) -> f32 {
  let cutoff = 6.0 / 29.0;
  let multiplier = 3.0 * (6.0/29.0) * (6.0/29.0);
  let constant = multiplier * (-4.0 / 29.0);

  if val > cutoff {
    val * val * val
  } else {
    val * multiplier + constant
  }
}
<commit_msg>colorspaces: cleanup code using chunks/chunks_mut<commit_after>use decoders::Image;

pub fn camera_to_lab(img: &Image, inb: &[f32]) -> Vec<f32> {
  let mut out: Vec<f32> = vec![0.0; (img.width*img.height*3) as usize];
  let cmatrix = cam_to_xyz_matrix(img);

  for (pixin, pixout) in inb.chunks(4).zip(out.chunks_mut(3)) {
    let r = pixin[0];
    let g = pixin[1];
    let b = pixin[2];
    let e = pixin[3];

    let x = r * cmatrix[0][0] + g * cmatrix[0][1] + b * cmatrix[0][2] + e * cmatrix[0][3];
    let y = r * cmatrix[1][0] + g * cmatrix[1][1] + b * cmatrix[1][2] + e * cmatrix[1][3];
    let z = r * cmatrix[2][0] + g * cmatrix[2][1] + b * cmatrix[2][2] + e * cmatrix[2][3];

    let (l,a,b) = xyz_to_lab(x,y,z);

    pixout[0] = l;
    pixout[1] = a;
    pixout[2] = b;
  }

  out
}

pub fn lab_to_rec709(_: &Image, buf: &mut Vec<f32>) {
  let cmatrix = xyz_to_rec709_matrix();

  for pix in buf.chunks_mut(3) {
    let l = pix[0];
    let a = pix[1];
    let b = pix[2];

    let (x,y,z) = lab_to_xyz(l,a,b);

    let r = x * cmatrix[0][0] + y * cmatrix[0][1] + z * cmatrix[0][2];
    let g = x * cmatrix[1][0] + y * cmatrix[1][1] + z * cmatrix[1][2];
    let b = x * cmatrix[2][0] + y * cmatrix[2][1] + z * cmatrix[2][2];

    pix[0] = r;
    pix[1] = g;
    pix[2] = b;
  }
}

fn inverse(inm: [[f32;3];3]) -> [[f32;3];3] {
  let invdet = 1.0 / (
    inm[0][0] * (inm[1][1] * inm[2][2] - inm[2][1] * inm[1][2]) -
    inm[0][1] * (inm[1][0] * inm[2][2] - inm[1][2] * inm[2][0]) +
    inm[0][2] * (inm[1][0] * inm[2][1] - inm[1][1] * inm[2][0])
  );

  let mut out = [[0.0; 3];3];
  out[0][0] =  (inm[1][1]*inm[2][2] - inm[2][1]*inm[1][2]) * invdet;
  out[0][1] = -(inm[0][1]*inm[2][2] - inm[0][2]*inm[2][1]) * invdet;
  out[0][2] =  (inm[0][1]*inm[1][2] - inm[0][2]*inm[1][1]) * invdet;
  out[1][0] = -(inm[1][0]*inm[2][2] - inm[1][2]*inm[2][0]) * invdet;
  out[1][1] =  (inm[0][0]*inm[2][2] - inm[0][2]*inm[2][0]) * invdet;
  out[1][2] = -(inm[0][0]*inm[1][2] - inm[1][0]*inm[0][2]) * invdet;
  out[2][0] =  (inm[1][0]*inm[2][1] - inm[2][0]*inm[1][1]) * invdet;
  out[2][1] = -(inm[0][0]*inm[2][1] - inm[2][0]*inm[0][1]) * invdet;
  out[2][2] =  (inm[0][0]*inm[1][1] - inm[1][0]*inm[0][1]) * invdet;

  out
}

fn pseudoinverse(inm: [[f32;3];4]) -> [[f32;4];3] {
  let mut temp: [[f32;6];3] = [[0.0; 6];3];

  for i in 0..3 {
    for j in 0..6 {
      temp[i][j] = if j == i+3 { 1.0 } else { 0.0 };
    }
    for j in 0..3 {
      for k in 0..4 {
        temp[i][j] += inm[k][i] * inm[k][j];
      }
    }
  }

  for i in 0..3 {
    let mut num = temp[i][i];
    for j in 0..6 {
      temp[i][j] /= num;
    }
    for k in 0..3 {
      if k == i { continue }
      num = temp[k][i];
      for j in 0..6 {
        temp[k][j] -= temp[i][j] * num;
      }
    }
  }

  let mut out: [[f32;4];3] = [[0.0; 4];3];

  for i in 0..4 {
    for j in 0..3 {
      out[j][i] = 0.0;
      for k in 0..3 {
        out[j][i] += temp[j][k+3] * inm[i][k];
      }
    }
  }

  out
}

fn cam_to_xyz_matrix(img: &Image) -> [[f32;4];3] {
  let mut xyz_to_cam = [[0.0;3];4];
  for i in 0..12 {
    xyz_to_cam[i/3][i%3] = img.color_matrix[i] as f32;
  }

  let d65_white = [0.9547,1.0,1.08883];
  let rgb_to_xyz = [
  // sRGB D65
    [ 0.412453, 0.357580, 0.180423 ],
    [ 0.212671, 0.715160, 0.072169 ],
    [ 0.019334, 0.119193, 0.950227 ],
  ];

  // Multiply RGB matrix
  let mut rgb_to_cam = [[0.0;3];4];
  for i in 0..4 {
    for j in 0..3 {
      rgb_to_cam[i][j] = 0.0;
      for k in 0..3 {
        rgb_to_cam[i][j] += xyz_to_cam[i][k] * rgb_to_xyz[k][j];
      }
    }
  }

  // Normalize rgb_to_cam so that rgb_to_cam * (1,1,1) is (1,1,1,1)
  for i in 0..4 {
    let mut num = 0.0;
    for j in 0..3 {
      num += rgb_to_cam[i][j];
    }
    for j in 0..3 {
      rgb_to_cam[i][j] = if num == 0.0 {
        0.0
      }  else {
        rgb_to_cam[i][j] / num
      };
    }
  }

  let cam_to_rgb = pseudoinverse(rgb_to_cam);
  let mut cam_to_xyz = [[0.0;4];3];
  // Multiply RGB matrix and adjust white to get a cam_to_xyz
  for i in 0..3 {
    for j in 0..4 {
      cam_to_xyz[i][j] = 0.0;
      for k in 0..3 {
        cam_to_xyz[i][j] += cam_to_rgb[k][j] * rgb_to_xyz[i][k] / d65_white[i];
      }
    }
  }

  cam_to_xyz
}

fn xyz_to_rec709_matrix() -> [[f32;3];3] {
  let rgb_to_xyz = [
  // sRGB D65
    [ 0.412453, 0.357580, 0.180423 ],
    [ 0.212671, 0.715160, 0.072169 ],
    [ 0.019334, 0.119193, 0.950227 ],
  ];

  inverse(rgb_to_xyz)
}

fn xyz_to_lab(x: f32, y: f32, z: f32) -> (f32,f32,f32) {
  // D50 White
  let xw = 0.9642; let yw = 1.000; let zw = 0.8249;

  let l = 116.0 * labf(y/yw) - 16.0;
  let a = 500.0 * (labf(x/xw) - labf(y/yw));
  let b = 200.0 * (labf(y/yw) - labf(z/zw));

  (l/100.0,(a+128.0)/256.0,(b+128.0)/256.0)
}

static CBRT_MAXVALS: usize = 1 << 16; // 2^16 should be enough precision
lazy_static! {
  static ref CBRT_LOOKUP: Vec<f32> = {
    let mut lookup: Vec<f32> = vec![0.0; CBRT_MAXVALS+1];
    for i in 0..(CBRT_MAXVALS+1) {
      let v = (i as f32) / (CBRT_MAXVALS as f32);
      lookup[i] = v.cbrt();
    }
    lookup
  };
}

fn labf(val: f32) -> f32 {
  let cutoff = (6.0/29.0)*(6.0/29.0)*(6.0/29.0);
  let multiplier = (1.0/3.0) * (29.0/6.0) * (29.0/6.0);
  let constant = 4.0 / 29.0;

  if val > cutoff {
    if val > 0.0 && val < 1.0 { // use the lookup table
      CBRT_LOOKUP[(val*(CBRT_MAXVALS as f32)) as usize]
    } else {
      val.cbrt()
    }
  } else {
    val * multiplier + constant
  }
}

fn lab_to_xyz(l: f32, a: f32, b: f32) -> (f32,f32,f32) {
  // D50 White
  let xw = 0.9642; let yw = 1.000; let zw = 0.8249;

  let cl = l * 100.0;
  let ca = (a * 256.0) - 128.0;
  let cb = (b * 256.0) - 128.0;

  let x = xw * labinvf((1.0/116.0) * (cl+16.0) + (1.0/500.0) * ca);
  let y = yw * labinvf((1.0/116.0) * (cl+16.0));
  let z = zw * labinvf((1.0/116.0) * (cl+16.0) - (1.0/200.0) * cb);

  (x,y,z)
}

fn labinvf(val: f32) -> f32 {
  let cutoff = 6.0 / 29.0;
  let multiplier = 3.0 * (6.0/29.0) * (6.0/29.0);
  let constant = multiplier * (-4.0 / 29.0);

  if val > cutoff {
    val * val * val
  } else {
    val * multiplier + constant
  }
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT

use std::collections::HashMap;
use std::convert::TryFrom;

use clap::{Arg, App};

use Instruction::*;

enum Instruction {
    SetMask(Mask),
    SetReg(u64, u64),
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Mask {
    and: u64,
    or:  u64,
    x: Vec<u8>,
}
impl Mask {
    pub fn new() -> Mask {
        Mask { and: 0, or: 0, x: vec![] }
    }

    #[inline]
    pub fn apply_v1(&self, v: u64) -> u64 { (v | self.or) & self.and }

// Average of 2206 iterations
// real    0m    9.106ms
// user    0m    7.293ms
// sys     0m    1.735ms
// cached page faults  1,179
    pub fn apply_v2_b(&self, v: u64) -> Vec<u64> {
        let mut rv = vec![0; 2_usize.pow(self.x.len() as u32)];
        rv[0] = v | self.or;
        for (i, idx) in self.x.iter().enumerate() {
            let m: u64 = 1 << *idx;
            let mut end = 2_usize.pow((i+1) as u32) - 1;
            let mut n = 2_usize.pow(i as u32);
            while n > 0 {
                n -= 1;
                rv[end] = rv[n]|m; end -= 1;
                rv[end] = rv[n]&(!m); end -= 1;
            }
        }
        return rv;
    }

// Average of 341 iterations
// real    0m    9.187ms
// user    0m    7.467ms
// sys     0m    1.641ms
// cached page faults  1,185
//
// Average of 2155 iterations
// real    0m    9.68ms
// user    0m    7.208ms
// sys     0m    1.779ms
    pub fn apply_v2(&self, v: u64) -> Vec<u64> {
        let mut rv = vec![ v | self.or ];
        for i in self.x.iter() {
            let mut tmp = Vec::with_capacity(rv.len() * 2);
            let m: u64 = 1 << *i;
            for v in rv {
                tmp.push(v|m);
                tmp.push(v&(!m));
            }
            rv = tmp;
        }
        return rv;
    }
}
impl std::convert::TryFrom<&str> for Mask {
    type Error = String;
    fn try_from(src: &str) -> Result<Self, Self::Error> {
        let (mut a, mut o) = (0, 0);
        let mut vec = Vec::new();
        for (i, ch) in src.chars().rev().enumerate() {
            match ch {
                '0' => { },
                '1' => { o |= 1 << i; a |= 1 << i; },
                'X' => { a |= 1 << i; vec.push(i as u8); },
                _   => { return Err(format!("Unexpected character '{}'", ch)) },
            }
        }
        return Ok(Mask { and: a, or: o, x: vec});
    }
}

fn records(fname: &str) -> Vec<Instruction> {
    let contents = std::fs::read_to_string(fname)
        .unwrap_or_else(|err| panic!("Error reading {}: {}", fname, err));

    return contents.lines().enumerate().map(
        |(lineno, chunk)| {
            if chunk.starts_with("mask = ") {
                SetMask(Mask::try_from(&chunk[7..]).unwrap_or_else(|err| panic!("Parse error at '{}' in {} in record {}: {}", chunk, fname, lineno+1, err)))
            } else {
                let v: Vec<&str> = chunk.split("] = ").collect();
                let reg = (v.get(0).unwrap_or_else(|| panic!("Parse error at '{}' in {} in record {}", chunk, fname, lineno+1)))[4..]
                            .parse().unwrap_or_else(|err| panic!("Parse error at '{}' in {} in record {}: {}", chunk, fname, lineno+1, err));
                let val = (v.get(1).unwrap_or_else(|| panic!("Parse error at '{}' in {} in record {}", chunk, fname, lineno+1)))
                            .parse().unwrap_or_else(|err| panic!("Parse error at '{}' in {} in record {}: {}", chunk, fname, lineno+1, err));
                SetReg(reg, val)
            }
        }
    ).collect();
}

fn main() {
    let matches = App::new("Advent of code 2020, Day 14 Solution")
        .arg(Arg::with_name("FILE").help("Input file to process").index(1))
        .get_matches();
    let fname = matches.value_of("FILE").unwrap_or("14.in");

    let setup = records(fname);
    let mut reg = HashMap::new();
    let mut mask = Mask::new();
    for cmd in setup.iter() {
        match cmd {
            SetMask(m)   => { mask = m.clone(); },
            SetReg(k, v) => { reg.insert(*k, mask.apply_v1(*v)); },
        }
    }
    println!("Part 1: Sum {}", reg.values().sum::<u64>());

    let mut reg = HashMap::new();
    let mut mask = Mask::new();
    for cmd in setup.iter() {
        match cmd {
            SetMask(m)   => { mask = m.clone(); },
            SetReg(k, v) => { for addr in mask.apply_v2(*k) { reg.insert(addr, *v); } },
        }
    }
    println!("Part 2: Sum {}", reg.values().sum::<u64>());
}
<commit_msg>bit shifts, not .pow()!<commit_after>// SPDX-License-Identifier: MIT

use std::collections::HashMap;
use std::convert::TryFrom;

use clap::{Arg, App};

use Instruction::*;

enum Instruction {
    SetMask(Mask),
    SetReg(u64, u64),
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Mask {
    and: u64,
    or:  u64,
    x: Vec<u8>,
}
impl Mask {
    pub fn new() -> Mask {
        Mask { and: 0, or: 0, x: vec![] }
    }

    #[inline]
    pub fn apply_v1(&self, v: u64) -> u64 { (v | self.or) & self.and }

// Average of 2105 iterations
// real    0m    8.900ms
// user    0m    7.171ms
// sys     0m    1.648ms
// cached page faults  1,177
    pub fn apply_v2(&self, v: u64) -> Vec<u64> {
        let mut rv = vec![0; 1<<(self.x.len())];
        rv[0] = v | self.or;
        for (i, idx) in self.x.iter().enumerate() {
            let m: u64 = 1 << *idx;
            let mut end = (1<<(i+1)) - 1;
            let mut n = 1<<i;
            while n > 0 { // about 75000 iterations total
                n -= 1;
                rv[end] = rv[n]|m;    end -= 1;
                rv[end] = rv[n]&(!m); end -= 1;
            }
        }
        return rv;
    }

// Average of 341 iterations
// real    0m    9.187ms
// user    0m    7.467ms
// sys     0m    1.641ms
// cached page faults  1,185
//
// Average of 2155 iterations
// real    0m    9.68ms
// user    0m    7.208ms
// sys     0m    1.779ms
    pub fn apply_v2_a(&self, v: u64) -> Vec<u64> {
        let mut rv = vec![ v | self.or ];
        for i in self.x.iter() {
            let mut tmp = Vec::with_capacity(rv.len() * 2);
            let m: u64 = 1 << *i;
            for v in rv {
                tmp.push(v|m);
                tmp.push(v&(!m));
            }
            rv = tmp;
        }
        return rv;
    }
}
impl std::convert::TryFrom<&str> for Mask {
    type Error = String;
    fn try_from(src: &str) -> Result<Self, Self::Error> {
        let (mut a, mut o) = (0, 0);
        let mut vec = Vec::new();
        for (i, ch) in src.chars().rev().enumerate() {
            match ch {
                '0' => { },
                '1' => { o |= 1 << i; a |= 1 << i; },
                'X' => { a |= 1 << i; vec.push(i as u8); },
                _   => { return Err(format!("Unexpected character '{}'", ch)) },
            }
        }
        return Ok(Mask { and: a, or: o, x: vec});
    }
}

fn records(fname: &str) -> Vec<Instruction> {
    let contents = std::fs::read_to_string(fname)
        .unwrap_or_else(|err| panic!("Error reading {}: {}", fname, err));

    return contents.lines().enumerate().map(
        |(lineno, chunk)| {
            if chunk.starts_with("mask = ") {
                SetMask(Mask::try_from(&chunk[7..]).unwrap_or_else(|err| panic!("Parse error at '{}' in {} in record {}: {}", chunk, fname, lineno+1, err)))
            } else {
                let v: Vec<&str> = chunk.split("] = ").collect();
                let reg = (v.get(0).unwrap_or_else(|| panic!("Parse error at '{}' in {} in record {}", chunk, fname, lineno+1)))[4..]
                            .parse().unwrap_or_else(|err| panic!("Parse error at '{}' in {} in record {}: {}", chunk, fname, lineno+1, err));
                let val = (v.get(1).unwrap_or_else(|| panic!("Parse error at '{}' in {} in record {}", chunk, fname, lineno+1)))
                            .parse().unwrap_or_else(|err| panic!("Parse error at '{}' in {} in record {}: {}", chunk, fname, lineno+1, err));
                SetReg(reg, val)
            }
        }
    ).collect();
}

fn main() {
    let matches = App::new("Advent of code 2020, Day 14 Solution")
        .arg(Arg::with_name("FILE").help("Input file to process").index(1))
        .get_matches();
    let fname = matches.value_of("FILE").unwrap_or("14.in");

    let setup = records(fname);
    let mut reg = HashMap::new();
    let mut mask = Mask::new();
    for cmd in setup.iter() {
        match cmd {
            SetMask(m)   => { mask = m.clone(); },
            SetReg(k, v) => { reg.insert(*k, mask.apply_v1(*v)); },
        }
    }
    println!("Part 1: Sum {}", reg.values().sum::<u64>());

    let mut reg = HashMap::new();
    let mut mask = Mask::new();
    for cmd in setup.iter() {
        match cmd {
            SetMask(m)   => { mask = m.clone(); },
            SetReg(k, v) => { for addr in mask.apply_v2(*k) { reg.insert(addr, *v); } },
        }
    }
    println!("Part 2: Sum {}", reg.values().sum::<u64>());
}
<|endoftext|>


<commit_before>use core::mem;
use core::ptr::{self, NonNull};
use core::slice;

#[repr(C)]
pub struct RustSlice {
    repr: NonNull<[()]>,
}

impl RustSlice {
    pub fn from_ref<T>(slice: &[T]) -> Self {
        let ptr = NonNull::from(slice).cast::<T>();
        let len = slice.len();
        Self::from_raw_parts(ptr, len)
    }

    pub fn from_mut<T>(slice: &mut [T]) -> Self {
        let ptr = NonNull::from(&mut *slice).cast::<T>();
        let len = slice.len();
        Self::from_raw_parts(ptr, len)
    }

    pub unsafe fn as_slice<'a, T>(self) -> &'a [T] {
        let ptr = self.as_ptr();
        let len = self.len();
        slice::from_raw_parts(ptr.as_ptr(), len)
    }

    pub unsafe fn as_mut_slice<'a, T>(self) -> &'a mut [T] {
        let ptr = self.as_ptr();
        let len = self.len();
        slice::from_raw_parts_mut(ptr.as_ptr(), len)
    }

    pub(crate) fn from_raw_parts<T>(ptr: NonNull<T>, len: usize) -> Self {
        // TODO: use NonNull::from_raw_parts(ptr.cast(), len) when stable.
        // https://doc.rust-lang.org/nightly/std/ptr/struct.NonNull.html#method.from_raw_parts
        // https://github.com/rust-lang/rust/issues/81513
        let ptr = ptr::slice_from_raw_parts_mut(ptr.as_ptr().cast(), len);
        RustSlice {
            repr: unsafe { NonNull::new_unchecked(ptr) },
        }
    }

    pub(crate) fn as_ptr<T>(&self) -> NonNull<T> {
        self.repr.cast()
    }

    pub(crate) fn len(&self) -> usize {
        // TODO: use self.repr.len() when stable.
        // https://doc.rust-lang.org/nightly/std/ptr/struct.NonNull.html#method.len
        // https://github.com/rust-lang/rust/issues/71146
        let slice = unsafe { self.repr.as_ref() };
        slice.len()
    }
}

const_assert_eq!(
    mem::size_of::<Option<RustSlice>>(),
    mem::size_of::<RustSlice>(),
);
<commit_msg>Resolve improper ctype lint affecting RustSlice<commit_after>use core::mem::{self, MaybeUninit};
use core::ptr::{self, NonNull};
use core::slice;

// ABI compatible with C++ rust::Slice<T> (not necessarily &[T]).
#[repr(C)]
pub struct RustSlice {
    repr: [MaybeUninit<usize>; mem::size_of::<NonNull<[()]>>() / mem::size_of::<usize>()],
}

impl RustSlice {
    pub fn from_ref<T>(slice: &[T]) -> Self {
        let ptr = NonNull::from(slice).cast::<T>();
        let len = slice.len();
        Self::from_raw_parts(ptr, len)
    }

    pub fn from_mut<T>(slice: &mut [T]) -> Self {
        let ptr = NonNull::from(&mut *slice).cast::<T>();
        let len = slice.len();
        Self::from_raw_parts(ptr, len)
    }

    pub unsafe fn as_slice<'a, T>(self) -> &'a [T] {
        let ptr = self.as_ptr();
        let len = self.len();
        slice::from_raw_parts(ptr.as_ptr(), len)
    }

    pub unsafe fn as_mut_slice<'a, T>(self) -> &'a mut [T] {
        let ptr = self.as_ptr();
        let len = self.len();
        slice::from_raw_parts_mut(ptr.as_ptr(), len)
    }

    pub(crate) fn from_raw_parts<T>(ptr: NonNull<T>, len: usize) -> Self {
        // TODO: use NonNull::from_raw_parts(ptr.cast(), len) when stable.
        // https://doc.rust-lang.org/nightly/std/ptr/struct.NonNull.html#method.from_raw_parts
        // https://github.com/rust-lang/rust/issues/81513
        let ptr = ptr::slice_from_raw_parts_mut(ptr.as_ptr().cast(), len);
        unsafe { mem::transmute::<NonNull<[()]>, RustSlice>(NonNull::new_unchecked(ptr)) }
    }

    pub(crate) fn as_ptr<T>(&self) -> NonNull<T> {
        let rust_slice = RustSlice { repr: self.repr };
        let repr = unsafe { mem::transmute::<RustSlice, NonNull<[()]>>(rust_slice) };
        repr.cast()
    }

    pub(crate) fn len(&self) -> usize {
        let rust_slice = RustSlice { repr: self.repr };
        let repr = unsafe { mem::transmute::<RustSlice, NonNull<[()]>>(rust_slice) };
        // TODO: use repr.len() when stable.
        // https://doc.rust-lang.org/nightly/std/ptr/struct.NonNull.html#method.len
        // https://github.com/rust-lang/rust/issues/71146
        unsafe { repr.as_ref() }.len()
    }
}

const_assert_eq!(mem::size_of::<NonNull<[()]>>(), mem::size_of::<RustSlice>());
const_assert_eq!(
    mem::align_of::<NonNull<[()]>>(),
    mem::align_of::<RustSlice>(),
);
<|endoftext|>


<commit_before>// Copyright 2020 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::chain;
use crate::conn::{Message, MessageHandler, Tracker};
use crate::core::core::{self, hash::Hash, hash::Hashed, CompactBlock};

use crate::msg::{
	BanReason, GetPeerAddrs, Headers, Locator, Msg, PeerAddrs, Ping, Pong, TxHashSetArchive,
	TxHashSetRequest, Type,
};
use crate::types::{Error, NetAdapter, PeerInfo};
use chrono::prelude::Utc;
use rand::{thread_rng, Rng};
use std::cmp;
use std::fs::{self, File, OpenOptions};
use std::io::BufWriter;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Instant;

pub struct Protocol {
	adapter: Arc<dyn NetAdapter>,
	peer_info: PeerInfo,
	state_sync_requested: Arc<AtomicBool>,
}

impl Protocol {
	pub fn new(
		adapter: Arc<dyn NetAdapter>,
		peer_info: PeerInfo,
		state_sync_requested: Arc<AtomicBool>,
	) -> Protocol {
		Protocol {
			adapter,
			peer_info,
			state_sync_requested,
		}
	}
}

impl MessageHandler for Protocol {
	fn consume(
		&self,
		mut msg: Message,
		stopped: Arc<AtomicBool>,
		tracker: Arc<Tracker>,
	) -> Result<Option<Msg>, Error> {
		let adapter = &self.adapter;

		// If we received a msg from a banned peer then log and drop it.
		// If we are getting a lot of these then maybe we are not cleaning
		// banned peers up correctly?
		if adapter.is_banned(self.peer_info.addr) {
			debug!(
				"handler: consume: peer {:?} banned, received: {:?}, dropping.",
				self.peer_info.addr, msg.header.msg_type,
			);
			return Ok(None);
		}

		match msg.header.msg_type {
			Type::Ping => {
				let ping: Ping = msg.body()?;
				adapter.peer_difficulty(self.peer_info.addr, ping.total_difficulty, ping.height);

				Ok(Some(Msg::new(
					Type::Pong,
					Pong {
						total_difficulty: adapter.total_difficulty()?,
						height: adapter.total_height()?,
					},
					self.peer_info.version,
				)?))
			}

			Type::Pong => {
				let pong: Pong = msg.body()?;
				adapter.peer_difficulty(self.peer_info.addr, pong.total_difficulty, pong.height);
				Ok(None)
			}

			Type::BanReason => {
				let ban_reason: BanReason = msg.body()?;
				error!("handle_payload: BanReason {:?}", ban_reason);
				Ok(None)
			}

			Type::TransactionKernel => {
				let h: Hash = msg.body()?;
				debug!(
					"handle_payload: received tx kernel: {}, msg_len: {}",
					h, msg.header.msg_len
				);
				adapter.tx_kernel_received(h, &self.peer_info)?;
				Ok(None)
			}

			Type::GetTransaction => {
				let h: Hash = msg.body()?;
				debug!(
					"handle_payload: GetTransaction: {}, msg_len: {}",
					h, msg.header.msg_len,
				);
				let tx = adapter.get_transaction(h);
				if let Some(tx) = tx {
					Ok(Some(Msg::new(
						Type::Transaction,
						tx,
						self.peer_info.version,
					)?))
				} else {
					Ok(None)
				}
			}

			Type::Transaction => {
				debug!(
					"handle_payload: received tx: msg_len: {}",
					msg.header.msg_len
				);
				let tx: core::Transaction = msg.body()?;
				adapter.transaction_received(tx, false)?;
				Ok(None)
			}

			Type::StemTransaction => {
				debug!(
					"handle_payload: received stem tx: msg_len: {}",
					msg.header.msg_len
				);
				let tx: core::Transaction = msg.body()?;
				adapter.transaction_received(tx, true)?;
				Ok(None)
			}

			Type::GetBlock => {
				let h: Hash = msg.body()?;
				trace!(
					"handle_payload: GetBlock: {}, msg_len: {}",
					h,
					msg.header.msg_len,
				);

				let bo = adapter.get_block(h);
				if let Some(b) = bo {
					return Ok(Some(Msg::new(Type::Block, b, self.peer_info.version)?));
				}
				Ok(None)
			}

			Type::Block => {
				debug!(
					"handle_payload: received block: msg_len: {}",
					msg.header.msg_len
				);
				let b: core::UntrustedBlock = msg.body()?;

				// We default to NONE opts here as we do not know know yet why this block was
				// received.
				// If we requested this block from a peer due to our node syncing then
				// the peer adapter will override opts to reflect this.
				adapter.block_received(b.into(), &self.peer_info, chain::Options::NONE)?;
				Ok(None)
			}

			Type::GetCompactBlock => {
				let h: Hash = msg.body()?;
				if let Some(b) = adapter.get_block(h) {
					let cb: CompactBlock = b.into();
					Ok(Some(Msg::new(
						Type::CompactBlock,
						cb,
						self.peer_info.version,
					)?))
				} else {
					Ok(None)
				}
			}

			Type::CompactBlock => {
				debug!(
					"handle_payload: received compact block: msg_len: {}",
					msg.header.msg_len
				);
				let b: core::UntrustedCompactBlock = msg.body()?;

				adapter.compact_block_received(b.into(), &self.peer_info)?;
				Ok(None)
			}

			Type::GetHeaders => {
				// load headers from the locator
				let loc: Locator = msg.body()?;
				let headers = adapter.locate_headers(&loc.hashes)?;

				// serialize and send all the headers over
				Ok(Some(Msg::new(
					Type::Headers,
					Headers { headers },
					self.peer_info.version,
				)?))
			}

			// "header first" block propagation - if we have not yet seen this block
			// we can go request it from some of our peers
			Type::Header => {
				let header: core::UntrustedBlockHeader = msg.body()?;
				adapter.header_received(header.into(), &self.peer_info)?;
				Ok(None)
			}

			Type::Headers => {
				let mut total_bytes_read = 0;

				// Read the count (u16) so we now how many headers to read.
				let (count, bytes_read): (u16, _) = msg.streaming_read()?;
				total_bytes_read += bytes_read;

				// Read chunks of headers off the stream and pass them off to the adapter.
				let chunk_size = 32;
				for chunk in (0..count).collect::<Vec<_>>().chunks(chunk_size) {
					let mut headers = vec![];
					for _ in chunk {
						let (header, bytes_read) =
							msg.streaming_read::<core::UntrustedBlockHeader>()?;
						headers.push(header.into());
						total_bytes_read += bytes_read;
					}
					adapter.headers_received(&headers, &self.peer_info)?;
				}

				// Now check we read the correct total number of bytes off the stream.
				if total_bytes_read != msg.header.msg_len {
					return Err(Error::MsgLen);
				}

				Ok(None)
			}

			Type::GetPeerAddrs => {
				let get_peers: GetPeerAddrs = msg.body()?;
				let peers = adapter.find_peer_addrs(get_peers.capabilities);
				Ok(Some(Msg::new(
					Type::PeerAddrs,
					PeerAddrs { peers },
					self.peer_info.version,
				)?))
			}

			Type::PeerAddrs => {
				let peer_addrs: PeerAddrs = msg.body()?;
				adapter.peer_addrs_received(peer_addrs.peers);
				Ok(None)
			}

			Type::TxHashSetRequest => {
				let sm_req: TxHashSetRequest = msg.body()?;
				debug!(
					"handle_payload: txhashset req for {} at {}",
					sm_req.hash, sm_req.height
				);

				let txhashset_header = self.adapter.txhashset_archive_header()?;
				let txhashset_header_hash = txhashset_header.hash();
				let txhashset = self.adapter.txhashset_read(txhashset_header_hash);

				if let Some(txhashset) = txhashset {
					let file_sz = txhashset.reader.metadata()?.len();
					let mut resp = Msg::new(
						Type::TxHashSetArchive,
						&TxHashSetArchive {
							height: txhashset_header.height as u64,
							hash: txhashset_header_hash,
							bytes: file_sz,
						},
						self.peer_info.version,
					)?;
					resp.add_attachment(txhashset.reader);
					Ok(Some(resp))
				} else {
					Ok(None)
				}
			}

			Type::TxHashSetArchive => {
				let sm_arch: TxHashSetArchive = msg.body()?;
				debug!(
					"handle_payload: txhashset archive for {} at {}. size={}",
					sm_arch.hash, sm_arch.height, sm_arch.bytes,
				);
				if !self.adapter.txhashset_receive_ready() {
					error!(
						"handle_payload: txhashset archive received but SyncStatus not on TxHashsetDownload",
					);
					return Err(Error::BadMessage);
				}
				if !self.state_sync_requested.load(Ordering::Relaxed) {
					error!("handle_payload: txhashset archive received but from the wrong peer",);
					return Err(Error::BadMessage);
				}
				// Update the sync state requested status
				self.state_sync_requested.store(false, Ordering::Relaxed);

				let download_start_time = Utc::now();
				self.adapter
					.txhashset_download_update(download_start_time, 0, sm_arch.bytes);

				let nonce: u32 = thread_rng().gen_range(0, 1_000_000);
				let tmp = self.adapter.get_tmpfile_pathname(format!(
					"txhashset-{}-{}.zip",
					download_start_time.timestamp(),
					nonce
				));
				let mut now = Instant::now();
				let mut save_txhashset_to_file = |file| -> Result<(), Error> {
					let mut tmp_zip =
						BufWriter::new(OpenOptions::new().write(true).create_new(true).open(file)?);
					let total_size = sm_arch.bytes as usize;
					let mut downloaded_size: usize = 0;
					let mut request_size = cmp::min(48_000, total_size);
					while request_size > 0 {
						let size = msg.copy_attachment(request_size, &mut tmp_zip)?;
						downloaded_size += size;
						request_size = cmp::min(48_000, total_size - downloaded_size);
						self.adapter.txhashset_download_update(
							download_start_time,
							downloaded_size as u64,
							total_size as u64,
						);
						if now.elapsed().as_secs() > 10 {
							now = Instant::now();
							debug!(
								"handle_payload: txhashset archive: {}/{}",
								downloaded_size, total_size
							);
						}
						// Increase received bytes quietly (without affecting the counters).
						// Otherwise we risk banning a peer as "abusive".
						tracker.inc_quiet_received(size as u64);

						// check the close channel
						if stopped.load(Ordering::Relaxed) {
							debug!("stopping txhashset download early");
							return Err(Error::ConnectionClose);
						}
					}
					debug!(
						"handle_payload: txhashset archive: {}/{} ... DONE",
						downloaded_size, total_size
					);
					tmp_zip
						.into_inner()
						.map_err(|_| Error::Internal)?
						.sync_all()?;
					Ok(())
				};

				if let Err(e) = save_txhashset_to_file(tmp.clone()) {
					error!(
						"handle_payload: txhashset archive save to file fail. err={:?}",
						e
					);
					return Err(e);
				}

				trace!(
					"handle_payload: txhashset archive save to file {:?} success",
					tmp,
				);

				let tmp_zip = File::open(tmp.clone())?;
				let res = self
					.adapter
					.txhashset_write(sm_arch.hash, tmp_zip, &self.peer_info)?;

				debug!(
					"handle_payload: txhashset archive for {} at {}, DONE. Data Ok: {}",
					sm_arch.hash, sm_arch.height, !res
				);

				if let Err(e) = fs::remove_file(tmp.clone()) {
					warn!("fail to remove tmp file: {:?}. err: {}", tmp, e);
				}

				Ok(None)
			}
			Type::Error | Type::Hand | Type::Shake => {
				debug!("Received an unexpected msg: {:?}", msg.header.msg_type);
				Ok(None)
			}
		}
	}
}
<commit_msg>Reduce number of allocations in Headers message read (#3301)<commit_after>// Copyright 2020 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::chain;
use crate::conn::{Message, MessageHandler, Tracker};
use crate::core::core::{self, hash::Hash, hash::Hashed, CompactBlock};

use crate::msg::{
	BanReason, GetPeerAddrs, Headers, Locator, Msg, PeerAddrs, Ping, Pong, TxHashSetArchive,
	TxHashSetRequest, Type,
};
use crate::types::{Error, NetAdapter, PeerInfo};
use chrono::prelude::Utc;
use rand::{thread_rng, Rng};
use std::cmp;
use std::fs::{self, File, OpenOptions};
use std::io::BufWriter;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Instant;

pub struct Protocol {
	adapter: Arc<dyn NetAdapter>,
	peer_info: PeerInfo,
	state_sync_requested: Arc<AtomicBool>,
}

impl Protocol {
	pub fn new(
		adapter: Arc<dyn NetAdapter>,
		peer_info: PeerInfo,
		state_sync_requested: Arc<AtomicBool>,
	) -> Protocol {
		Protocol {
			adapter,
			peer_info,
			state_sync_requested,
		}
	}
}

impl MessageHandler for Protocol {
	fn consume(
		&self,
		mut msg: Message,
		stopped: Arc<AtomicBool>,
		tracker: Arc<Tracker>,
	) -> Result<Option<Msg>, Error> {
		let adapter = &self.adapter;

		// If we received a msg from a banned peer then log and drop it.
		// If we are getting a lot of these then maybe we are not cleaning
		// banned peers up correctly?
		if adapter.is_banned(self.peer_info.addr) {
			debug!(
				"handler: consume: peer {:?} banned, received: {:?}, dropping.",
				self.peer_info.addr, msg.header.msg_type,
			);
			return Ok(None);
		}

		match msg.header.msg_type {
			Type::Ping => {
				let ping: Ping = msg.body()?;
				adapter.peer_difficulty(self.peer_info.addr, ping.total_difficulty, ping.height);

				Ok(Some(Msg::new(
					Type::Pong,
					Pong {
						total_difficulty: adapter.total_difficulty()?,
						height: adapter.total_height()?,
					},
					self.peer_info.version,
				)?))
			}

			Type::Pong => {
				let pong: Pong = msg.body()?;
				adapter.peer_difficulty(self.peer_info.addr, pong.total_difficulty, pong.height);
				Ok(None)
			}

			Type::BanReason => {
				let ban_reason: BanReason = msg.body()?;
				error!("handle_payload: BanReason {:?}", ban_reason);
				Ok(None)
			}

			Type::TransactionKernel => {
				let h: Hash = msg.body()?;
				debug!(
					"handle_payload: received tx kernel: {}, msg_len: {}",
					h, msg.header.msg_len
				);
				adapter.tx_kernel_received(h, &self.peer_info)?;
				Ok(None)
			}

			Type::GetTransaction => {
				let h: Hash = msg.body()?;
				debug!(
					"handle_payload: GetTransaction: {}, msg_len: {}",
					h, msg.header.msg_len,
				);
				let tx = adapter.get_transaction(h);
				if let Some(tx) = tx {
					Ok(Some(Msg::new(
						Type::Transaction,
						tx,
						self.peer_info.version,
					)?))
				} else {
					Ok(None)
				}
			}

			Type::Transaction => {
				debug!(
					"handle_payload: received tx: msg_len: {}",
					msg.header.msg_len
				);
				let tx: core::Transaction = msg.body()?;
				adapter.transaction_received(tx, false)?;
				Ok(None)
			}

			Type::StemTransaction => {
				debug!(
					"handle_payload: received stem tx: msg_len: {}",
					msg.header.msg_len
				);
				let tx: core::Transaction = msg.body()?;
				adapter.transaction_received(tx, true)?;
				Ok(None)
			}

			Type::GetBlock => {
				let h: Hash = msg.body()?;
				trace!(
					"handle_payload: GetBlock: {}, msg_len: {}",
					h,
					msg.header.msg_len,
				);

				let bo = adapter.get_block(h);
				if let Some(b) = bo {
					return Ok(Some(Msg::new(Type::Block, b, self.peer_info.version)?));
				}
				Ok(None)
			}

			Type::Block => {
				debug!(
					"handle_payload: received block: msg_len: {}",
					msg.header.msg_len
				);
				let b: core::UntrustedBlock = msg.body()?;

				// We default to NONE opts here as we do not know know yet why this block was
				// received.
				// If we requested this block from a peer due to our node syncing then
				// the peer adapter will override opts to reflect this.
				adapter.block_received(b.into(), &self.peer_info, chain::Options::NONE)?;
				Ok(None)
			}

			Type::GetCompactBlock => {
				let h: Hash = msg.body()?;
				if let Some(b) = adapter.get_block(h) {
					let cb: CompactBlock = b.into();
					Ok(Some(Msg::new(
						Type::CompactBlock,
						cb,
						self.peer_info.version,
					)?))
				} else {
					Ok(None)
				}
			}

			Type::CompactBlock => {
				debug!(
					"handle_payload: received compact block: msg_len: {}",
					msg.header.msg_len
				);
				let b: core::UntrustedCompactBlock = msg.body()?;

				adapter.compact_block_received(b.into(), &self.peer_info)?;
				Ok(None)
			}

			Type::GetHeaders => {
				// load headers from the locator
				let loc: Locator = msg.body()?;
				let headers = adapter.locate_headers(&loc.hashes)?;

				// serialize and send all the headers over
				Ok(Some(Msg::new(
					Type::Headers,
					Headers { headers },
					self.peer_info.version,
				)?))
			}

			// "header first" block propagation - if we have not yet seen this block
			// we can go request it from some of our peers
			Type::Header => {
				let header: core::UntrustedBlockHeader = msg.body()?;
				adapter.header_received(header.into(), &self.peer_info)?;
				Ok(None)
			}

			Type::Headers => {
				let mut total_bytes_read = 0;

				// Read the count (u16) so we now how many headers to read.
				let (count, bytes_read): (u16, _) = msg.streaming_read()?;
				total_bytes_read += bytes_read;

				// Read chunks of headers off the stream and pass them off to the adapter.
				let chunk_size = 32u16;
				let mut headers = Vec::with_capacity(chunk_size as usize);
				for i in 1..=count {
					let (header, bytes_read) =
						msg.streaming_read::<core::UntrustedBlockHeader>()?;
					headers.push(header.into());
					total_bytes_read += bytes_read;
					if i % chunk_size == 0 || i == count {
						adapter.headers_received(&headers, &self.peer_info)?;
						headers.clear();
					}
				}

				// Now check we read the correct total number of bytes off the stream.
				if total_bytes_read != msg.header.msg_len {
					return Err(Error::MsgLen);
				}

				Ok(None)
			}

			Type::GetPeerAddrs => {
				let get_peers: GetPeerAddrs = msg.body()?;
				let peers = adapter.find_peer_addrs(get_peers.capabilities);
				Ok(Some(Msg::new(
					Type::PeerAddrs,
					PeerAddrs { peers },
					self.peer_info.version,
				)?))
			}

			Type::PeerAddrs => {
				let peer_addrs: PeerAddrs = msg.body()?;
				adapter.peer_addrs_received(peer_addrs.peers);
				Ok(None)
			}

			Type::TxHashSetRequest => {
				let sm_req: TxHashSetRequest = msg.body()?;
				debug!(
					"handle_payload: txhashset req for {} at {}",
					sm_req.hash, sm_req.height
				);

				let txhashset_header = self.adapter.txhashset_archive_header()?;
				let txhashset_header_hash = txhashset_header.hash();
				let txhashset = self.adapter.txhashset_read(txhashset_header_hash);

				if let Some(txhashset) = txhashset {
					let file_sz = txhashset.reader.metadata()?.len();
					let mut resp = Msg::new(
						Type::TxHashSetArchive,
						&TxHashSetArchive {
							height: txhashset_header.height as u64,
							hash: txhashset_header_hash,
							bytes: file_sz,
						},
						self.peer_info.version,
					)?;
					resp.add_attachment(txhashset.reader);
					Ok(Some(resp))
				} else {
					Ok(None)
				}
			}

			Type::TxHashSetArchive => {
				let sm_arch: TxHashSetArchive = msg.body()?;
				debug!(
					"handle_payload: txhashset archive for {} at {}. size={}",
					sm_arch.hash, sm_arch.height, sm_arch.bytes,
				);
				if !self.adapter.txhashset_receive_ready() {
					error!(
						"handle_payload: txhashset archive received but SyncStatus not on TxHashsetDownload",
					);
					return Err(Error::BadMessage);
				}
				if !self.state_sync_requested.load(Ordering::Relaxed) {
					error!("handle_payload: txhashset archive received but from the wrong peer",);
					return Err(Error::BadMessage);
				}
				// Update the sync state requested status
				self.state_sync_requested.store(false, Ordering::Relaxed);

				let download_start_time = Utc::now();
				self.adapter
					.txhashset_download_update(download_start_time, 0, sm_arch.bytes);

				let nonce: u32 = thread_rng().gen_range(0, 1_000_000);
				let tmp = self.adapter.get_tmpfile_pathname(format!(
					"txhashset-{}-{}.zip",
					download_start_time.timestamp(),
					nonce
				));
				let mut now = Instant::now();
				let mut save_txhashset_to_file = |file| -> Result<(), Error> {
					let mut tmp_zip =
						BufWriter::new(OpenOptions::new().write(true).create_new(true).open(file)?);
					let total_size = sm_arch.bytes as usize;
					let mut downloaded_size: usize = 0;
					let mut request_size = cmp::min(48_000, total_size);
					while request_size > 0 {
						let size = msg.copy_attachment(request_size, &mut tmp_zip)?;
						downloaded_size += size;
						request_size = cmp::min(48_000, total_size - downloaded_size);
						self.adapter.txhashset_download_update(
							download_start_time,
							downloaded_size as u64,
							total_size as u64,
						);
						if now.elapsed().as_secs() > 10 {
							now = Instant::now();
							debug!(
								"handle_payload: txhashset archive: {}/{}",
								downloaded_size, total_size
							);
						}
						// Increase received bytes quietly (without affecting the counters).
						// Otherwise we risk banning a peer as "abusive".
						tracker.inc_quiet_received(size as u64);

						// check the close channel
						if stopped.load(Ordering::Relaxed) {
							debug!("stopping txhashset download early");
							return Err(Error::ConnectionClose);
						}
					}
					debug!(
						"handle_payload: txhashset archive: {}/{} ... DONE",
						downloaded_size, total_size
					);
					tmp_zip
						.into_inner()
						.map_err(|_| Error::Internal)?
						.sync_all()?;
					Ok(())
				};

				if let Err(e) = save_txhashset_to_file(tmp.clone()) {
					error!(
						"handle_payload: txhashset archive save to file fail. err={:?}",
						e
					);
					return Err(e);
				}

				trace!(
					"handle_payload: txhashset archive save to file {:?} success",
					tmp,
				);

				let tmp_zip = File::open(tmp.clone())?;
				let res = self
					.adapter
					.txhashset_write(sm_arch.hash, tmp_zip, &self.peer_info)?;

				debug!(
					"handle_payload: txhashset archive for {} at {}, DONE. Data Ok: {}",
					sm_arch.hash, sm_arch.height, !res
				);

				if let Err(e) = fs::remove_file(tmp.clone()) {
					warn!("fail to remove tmp file: {:?}. err: {}", tmp, e);
				}

				Ok(None)
			}
			Type::Error | Type::Hand | Type::Shake => {
				debug!("Received an unexpected msg: {:?}", msg.header.msg_type);
				Ok(None)
			}
		}
	}
}
<|endoftext|>


<commit_before>pub mod screen;
mod main_menu_screen;
mod battle_screen;

pub use self::screen::*;
pub use self::main_menu_screen::*;
pub use self::battle_screen::*;
<commit_msg>Import fix<commit_after>pub mod screen;
mod main_menu;
mod battle;

pub use self::screen::*;
pub use self::main_menu::*;
pub use self::battle::*;
<|endoftext|>