<commit_before>extern crate yaml_rust;
extern crate regex;
extern crate rusqlite;

use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::net::TcpListener;
use std::time::Duration;
use std::sync::Arc;
use std::thread;
use std::process::exit;
use yaml_rust::YamlLoader;
use regex::bytes::RegexSetBuilder;
use rusqlite::Connection;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

#[derive(Copy, Clone)]
struct App {
  print_ascii: bool,
  print_binary: bool,
  sql_logging: bool
}

fn to_hex(bytes: &[u8]) -> String {
  let mut count = 0;
  let mut result = String::with_capacity(67);
  let mut dotline = String::with_capacity(18);

  for byte in bytes.iter() {
    if count != 0 && count%8 == 0 {
      result.push(' ');
      if count%16 == 0 {
        dotline.push('\n');
        result.push_str(&dotline);
        result.reserve(67);
        dotline.truncate(0);
      }
      else { dotline.push(' '); }
    }
    result.push_str(&format!("{:02X} ", byte));
    if *byte > 31 && *byte < 127 { dotline.push(*byte as char) }
    else if *byte == 0 { dotline.push('-'); }
    else { dotline.push('.'); }
    count += 1;
  }
  while count%16 != 0 {
    if count%8 == 0 { result.push(' '); }
    result.push_str("   ");
    count += 1;
  }
  if dotline.len() != 0 { result.push_str(&dotline); }
  result
}

fn main() {
  let authorstring: String = str::replace(env!("CARGO_PKG_AUTHORS"), ":", "\n");
  println!("Portlurker v{}", VERSION);
  println!("{}", authorstring);
  println!("-----------------------------------------");

  let mut app = App { print_ascii: false, print_binary: false, sql_logging: false };
  let binary_matches = [
    ("SSL3.0 Record Protocol", r"^\x16\x03\x00..\x01"),
    ("TLS1.0 Record Protocol", r"^\x16\x03\x01..\x01"),
    ("TLS1.1 Record Protocol", r"^\x16\x03\x02..\x01"),
    ("TLS1.2 Record Protocol", r"^\x16\x03\x03..\x01"),
    ("SSL3.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x00"),
    ("TLS1.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x01"),
    ("TLS1.1 CLIENT_HELLO", r"^\x16....\x01...\x03\x02"),
    ("TLS1.2 CLIENT_HELLO", r"^\x16....\x01...\x03\x03"),
    ("SMB1 COMMAND NEGOTIATE", r"^....\xffSMB\x72"),
    ("SMB1 NT_STATUS Success", r"^....\xffSMB.[\x00-\x0f]"),
    ("SMB1 NT_STATUS Information", r"^....\xffSMB.[\x40-\x4f]"),
    ("SMB1 NT_STATUS Warning", r"^....\xffSMB.[\x80-\x8f]"),
    ("SMB1 NT_STATUS Error", r"^....\xffSMB.[\xc0-\xcf]"),
    ("SMB2 COMMAND NEGOTIATE", r"^\x00...\xfeSMB........\x00\x00"),
    ("SMB2 NT_STATUS Success", r"^\x00...\xfeSMB....[\x00-\x0f]"),
    ("SMB2 NT_STATUS Information", r"^\x00...\xfeSMB....[\x40-\x4f]"),
    ("SMB2 NT_STATUS Warning", r"^\x00...\xfeSMB....[\x80-\x8f]"),
    ("SMB2 NT_STATUS Error", r"^\x00...\xfeSMB....[\xc0-\xcf]"),
    ("MS-TDS PRELOGIN Request", r"^\x12\x01\x00.\x00\x00"),
    ("MS-TDS LOGIN Request", r"^\x10\x01\x00.\x00\x00"),
    ("SOCKS4 NOAUTH Request", r"^\x04\x01\x00\x50"),
    ("SOCKS5 NOAUTH Request", r"^\x05\x01\x00$"), // Tested ok-ish
    ("SOCKS5 USER/PASS Request", r"^\x05\x02\x00\x02$") // possibly broken
  ];
  let io_timeout = Duration::new(300, 0); // 5 minutes

  let mut config_str = String::new();
  match File::open("config.yml") {
      Ok(mut file)   => { file.read_to_string(&mut config_str).unwrap(); },
      Err(err)  => panic!("Unable to open configuration file: {}", err),
  }

  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  //println!("{:?}", config);
  if config["general"].is_badvalue() {
    println!("No 'general' section found in configuration file");
    exit(-1);
  }
  if config["ports"].is_badvalue() {
    println!("No 'ports' section found in configuration file");
    exit(-1);
  }
  if !config["general"]["print_ascii"].is_badvalue() {
    if config["general"]["print_ascii"].as_bool().unwrap() {
      app.print_ascii = true;
      println!("Printing ASCII");
    }
  }
  if !config["general"]["print_binary"].is_badvalue() {
    if config["general"]["print_binary"].as_bool().unwrap() {
      app.print_binary = true;
      println!("Printing binary in hexadecimal");
    }
  }
  if !config["general"]["sql_logging"].is_badvalue() {
    if config["general"]["sql_logging"].as_bool().unwrap() {
      app.sql_logging = true;
      println!("Logging events to portlurker.sqlite");
      let conn = Connection::open("portlurker.sqlite").expect("Failed to open or create database! Check your permissions or disk space.");
      conn.execute("CREATE TABLE IF NOT EXISTS connections (
        id         INTEGER PRIMARY KEY,
        time       INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
        remoteip   TEXT NOT NULL,
        remoteport INTEGER NOT NULL,
        localport  INTEGER NOT NULL
      )", &[]).expect("Failed to create table inside database!");
    }
  }
  let app = app; // Revert to immutable

  let mut patterns = Vec::with_capacity(binary_matches.len());
  for &(_, pattern) in binary_matches.into_iter() {
    patterns.push(pattern);
  }
  let regexset = RegexSetBuilder::new(patterns)
    .unicode(false)
    .dot_matches_new_line(false)
    .build().unwrap();

  println!("\nStarting listeners on the following ports:");

  for port in config["ports"].as_vec().unwrap() {
    if !port["tcp"].is_badvalue() {
      let portno = port["tcp"].as_i64().unwrap();
      println!("TCP port {}", portno);
      let mut banner = Arc::new(String::new());
      if let Some(x) = port["banner"].as_str() {
        Arc::get_mut(& mut banner).unwrap().push_str(x);
        println!("  with banner: {}", *banner);
      }
      let regexset = regexset.clone();
      thread::spawn(move || {
        let server = TcpListener::bind(("0.0.0.0", portno as u16)).unwrap();
        for res in server.incoming() {
          let mut stream = match res {
            Ok(stream) => stream,
            Err(e) => { println!("ACCEPT ERROR TCP {}: {}", portno, e.to_string()); continue; }
          };
          stream.set_read_timeout(Some(io_timeout)).expect("Failed to set read timeout on TcpStream");
          stream.set_write_timeout(Some(io_timeout)).expect("Failed to set write timeout on TcpStream");
          let addr = stream.peer_addr().unwrap();
          println!("CONNECT TCP {} from {}", portno, addr);

          if app.sql_logging {
            #[derive(Debug)]
            struct LoggedConnection {
              remoteip: String,
              remoteport: u16,
              localport: i64
            }

            let newdbentry = LoggedConnection {
              remoteip: addr.ip().to_string(),
              remoteport: addr.port(),
              localport: portno
            };

            let conn = Connection::open("./portlurker.sqlite").unwrap();
            conn.execute("INSERT INTO connections (remoteip, remoteport, localport) VALUES (?1, ?2, ?3)", &[&newdbentry.remoteip, &newdbentry.remoteport, &newdbentry.localport]).expect("Can't write new row into table!");
          }

          let regexset = regexset.clone();
          let banner = banner.clone();
          thread::spawn(move || {
            if banner.len() > 0 {
              match stream.write((*banner).as_bytes()) {
                Ok(_) => {
                  for line in (*banner).lines() {
                    println!("> {}", line);
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("WRITE TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("WRITE ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  return;
                }
              }
            }
            let mut buf: [u8; 2048] = [0; 2048];
            loop {
              match stream.read(&mut buf) {
                Ok(c) => {
                  if c == 0 {
                    println!("CLOSE TCP {} from {}", portno, addr);
                    break;
                  }
                  let mut printables = Vec::new();
                  let mut found = false;
                  let mut mbfound = false;
                  let mut mbstring = String::new();
                  let mut start = 0;
                  for i in 0..c {
                    if (buf[i] > 31 && buf[i] < 127) || buf[i] == 10 || buf[i] == 13 {
                      if !found {
                        start = i;
                        found = true;
                      }
                    }
                    else {
                      if found {
                        if i-start == 1 {
                          if (start > 0) && (buf[start-1] == 0) {
                            mbstring.push(buf[i-1] as char);
                            if !mbfound { mbfound = true; }
                          }
                        }
                        else { printables.push(&buf[start..i]); }
                        found = false;
                      }
                      else if mbfound {
                        mbstring.push('\n');
                        mbfound = false;
                      }
                    }
                  }
                  if found { printables.push(&buf[start..c]); }
                  if printables.len() == 1 && printables[0].len() == c {
                    if app.print_ascii {
                      let data = String::from_utf8_lossy(printables[0]);
                      for line in data.lines() {
                        println!("| {}", line);
                      }
                    }
                    else { println!("! Read {} bytes of printable ASCII", c); }
                  }
                  else {
                    println!("! Read {} bytes of binary", c);
                    for id in regexset.matches(&buf[..c]).into_iter() {
                      println!("^ Matches pattern {}", binary_matches[id].0);
                    }
                    for printable in printables {
                      if printable.len() > 3 {
                        let data = String::from_utf8_lossy(printable);
                        for line in data.lines() {
                          println!("$ {}", line);
                        }
                      }
                    }
                    for line in mbstring.lines() {
                      if line.len() > 3 { println!("% {}", line); }
                    }
                    if app.print_binary {
                      let hex = to_hex(&buf[..c]);
                      for line in hex.lines() { println!(". {}", line); }
                    }
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("READ TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("READ ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  break;
                }
              }
              match stream.take_error() {
                Ok(opt) => {
                  if opt.is_some() {
                    println!("ERROR TCP {} from {}: {}", portno, addr, opt.unwrap().to_string());
                    break;
                  }
                }
                Err(_) => {
                  println!("This shouldn't happen...");
                  break;
                }
              }
            }
          });
        }
      });
    }
    else if !port["udp"].is_badvalue() {
      println!("UDP port {}", port["udp"].as_i64().unwrap());
    }
    else {
      println!("Invalid port specification in configuration file");
    }
  }

  loop { thread::sleep(Duration::new(60, 0)); }
}
<commit_msg>Config file error handling, replaced panic for message + exit<commit_after>extern crate yaml_rust;
extern crate regex;
extern crate rusqlite;

use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::net::TcpListener;
use std::time::Duration;
use std::sync::Arc;
use std::thread;
use std::process::exit;
use yaml_rust::YamlLoader;
use regex::bytes::RegexSetBuilder;
use rusqlite::Connection;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

#[derive(Copy, Clone)]
struct App {
  print_ascii: bool,
  print_binary: bool,
  sql_logging: bool
}

fn to_hex(bytes: &[u8]) -> String {
  let mut count = 0;
  let mut result = String::with_capacity(67);
  let mut dotline = String::with_capacity(18);

  for byte in bytes.iter() {
    if count != 0 && count%8 == 0 {
      result.push(' ');
      if count%16 == 0 {
        dotline.push('\n');
        result.push_str(&dotline);
        result.reserve(67);
        dotline.truncate(0);
      }
      else { dotline.push(' '); }
    }
    result.push_str(&format!("{:02X} ", byte));
    if *byte > 31 && *byte < 127 { dotline.push(*byte as char) }
    else if *byte == 0 { dotline.push('-'); }
    else { dotline.push('.'); }
    count += 1;
  }
  while count%16 != 0 {
    if count%8 == 0 { result.push(' '); }
    result.push_str("   ");
    count += 1;
  }
  if dotline.len() != 0 { result.push_str(&dotline); }
  result
}

fn main() {
  let authorstring: String = str::replace(env!("CARGO_PKG_AUTHORS"), ":", "\n");
  println!("Portlurker v{}", VERSION);
  println!("{}", authorstring);
  println!("-----------------------------------------");

  let mut app = App { print_ascii: false, print_binary: false, sql_logging: false };
  let binary_matches = [
    ("SSL3.0 Record Protocol", r"^\x16\x03\x00..\x01"),
    ("TLS1.0 Record Protocol", r"^\x16\x03\x01..\x01"),
    ("TLS1.1 Record Protocol", r"^\x16\x03\x02..\x01"),
    ("TLS1.2 Record Protocol", r"^\x16\x03\x03..\x01"),
    ("SSL3.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x00"),
    ("TLS1.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x01"),
    ("TLS1.1 CLIENT_HELLO", r"^\x16....\x01...\x03\x02"),
    ("TLS1.2 CLIENT_HELLO", r"^\x16....\x01...\x03\x03"),
    ("SMB1 COMMAND NEGOTIATE", r"^....\xffSMB\x72"),
    ("SMB1 NT_STATUS Success", r"^....\xffSMB.[\x00-\x0f]"),
    ("SMB1 NT_STATUS Information", r"^....\xffSMB.[\x40-\x4f]"),
    ("SMB1 NT_STATUS Warning", r"^....\xffSMB.[\x80-\x8f]"),
    ("SMB1 NT_STATUS Error", r"^....\xffSMB.[\xc0-\xcf]"),
    ("SMB2 COMMAND NEGOTIATE", r"^\x00...\xfeSMB........\x00\x00"),
    ("SMB2 NT_STATUS Success", r"^\x00...\xfeSMB....[\x00-\x0f]"),
    ("SMB2 NT_STATUS Information", r"^\x00...\xfeSMB....[\x40-\x4f]"),
    ("SMB2 NT_STATUS Warning", r"^\x00...\xfeSMB....[\x80-\x8f]"),
    ("SMB2 NT_STATUS Error", r"^\x00...\xfeSMB....[\xc0-\xcf]"),
    ("MS-TDS PRELOGIN Request", r"^\x12\x01\x00.\x00\x00"),
    ("MS-TDS LOGIN Request", r"^\x10\x01\x00.\x00\x00"),
    ("SOCKS4 NOAUTH Request", r"^\x04\x01\x00\x50"),
    ("SOCKS5 NOAUTH Request", r"^\x05\x01\x00$"), // Tested ok-ish
    ("SOCKS5 USER/PASS Request", r"^\x05\x02\x00\x02$") // possibly broken
  ];
  let io_timeout = Duration::new(300, 0); // 5 minutes

  let mut config_str = String::new();
  match File::open("config.yml") {
      Ok(mut file)   => { file.read_to_string(&mut config_str).unwrap(); },
      Err(err)  => { println!("Unable to open configuration file: {}", err); exit(-1); },
  }

  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  //println!("{:?}", config);
  if config["general"].is_badvalue() {
    println!("No 'general' section found in configuration file");
    exit(-1);
  }
  if config["ports"].is_badvalue() {
    println!("No 'ports' section found in configuration file");
    exit(-1);
  }
  if !config["general"]["print_ascii"].is_badvalue() {
    if config["general"]["print_ascii"].as_bool().unwrap() {
      app.print_ascii = true;
      println!("Printing ASCII");
    }
  }
  if !config["general"]["print_binary"].is_badvalue() {
    if config["general"]["print_binary"].as_bool().unwrap() {
      app.print_binary = true;
      println!("Printing binary in hexadecimal");
    }
  }
  if !config["general"]["sql_logging"].is_badvalue() {
    if config["general"]["sql_logging"].as_bool().unwrap() {
      app.sql_logging = true;
      println!("Logging events to portlurker.sqlite");
      let conn = Connection::open("portlurker.sqlite").expect("Failed to open or create database! Check your permissions or disk space.");
      conn.execute("CREATE TABLE IF NOT EXISTS connections (
        id         INTEGER PRIMARY KEY,
        time       INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
        remoteip   TEXT NOT NULL,
        remoteport INTEGER NOT NULL,
        localport  INTEGER NOT NULL
      )", &[]).expect("Failed to create table inside database!");
    }
  }
  let app = app; // Revert to immutable

  let mut patterns = Vec::with_capacity(binary_matches.len());
  for &(_, pattern) in binary_matches.into_iter() {
    patterns.push(pattern);
  }
  let regexset = RegexSetBuilder::new(patterns)
    .unicode(false)
    .dot_matches_new_line(false)
    .build().unwrap();

  println!("\nStarting listeners on the following ports:");

  for port in config["ports"].as_vec().unwrap() {
    if !port["tcp"].is_badvalue() {
      let portno = port["tcp"].as_i64().unwrap();
      println!("TCP port {}", portno);
      let mut banner = Arc::new(String::new());
      if let Some(x) = port["banner"].as_str() {
        Arc::get_mut(& mut banner).unwrap().push_str(x);
        println!("  with banner: {}", *banner);
      }
      let regexset = regexset.clone();
      thread::spawn(move || {
        let server = TcpListener::bind(("0.0.0.0", portno as u16)).unwrap();
        for res in server.incoming() {
          let mut stream = match res {
            Ok(stream) => stream,
            Err(e) => { println!("ACCEPT ERROR TCP {}: {}", portno, e.to_string()); continue; }
          };
          stream.set_read_timeout(Some(io_timeout)).expect("Failed to set read timeout on TcpStream");
          stream.set_write_timeout(Some(io_timeout)).expect("Failed to set write timeout on TcpStream");
          let addr = stream.peer_addr().unwrap();
          println!("CONNECT TCP {} from {}", portno, addr);

          if app.sql_logging {
            #[derive(Debug)]
            struct LoggedConnection {
              remoteip: String,
              remoteport: u16,
              localport: i64
            }

            let newdbentry = LoggedConnection {
              remoteip: addr.ip().to_string(),
              remoteport: addr.port(),
              localport: portno
            };

            let conn = Connection::open("./portlurker.sqlite").unwrap();
            conn.execute("INSERT INTO connections (remoteip, remoteport, localport) VALUES (?1, ?2, ?3)", &[&newdbentry.remoteip, &newdbentry.remoteport, &newdbentry.localport]).expect("Can't write new row into table!");
          }

          let regexset = regexset.clone();
          let banner = banner.clone();
          thread::spawn(move || {
            if banner.len() > 0 {
              match stream.write((*banner).as_bytes()) {
                Ok(_) => {
                  for line in (*banner).lines() {
                    println!("> {}", line);
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("WRITE TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("WRITE ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  return;
                }
              }
            }
            let mut buf: [u8; 2048] = [0; 2048];
            loop {
              match stream.read(&mut buf) {
                Ok(c) => {
                  if c == 0 {
                    println!("CLOSE TCP {} from {}", portno, addr);
                    break;
                  }
                  let mut printables = Vec::new();
                  let mut found = false;
                  let mut mbfound = false;
                  let mut mbstring = String::new();
                  let mut start = 0;
                  for i in 0..c {
                    if (buf[i] > 31 && buf[i] < 127) || buf[i] == 10 || buf[i] == 13 {
                      if !found {
                        start = i;
                        found = true;
                      }
                    }
                    else {
                      if found {
                        if i-start == 1 {
                          if (start > 0) && (buf[start-1] == 0) {
                            mbstring.push(buf[i-1] as char);
                            if !mbfound { mbfound = true; }
                          }
                        }
                        else { printables.push(&buf[start..i]); }
                        found = false;
                      }
                      else if mbfound {
                        mbstring.push('\n');
                        mbfound = false;
                      }
                    }
                  }
                  if found { printables.push(&buf[start..c]); }
                  if printables.len() == 1 && printables[0].len() == c {
                    if app.print_ascii {
                      let data = String::from_utf8_lossy(printables[0]);
                      for line in data.lines() {
                        println!("| {}", line);
                      }
                    }
                    else { println!("! Read {} bytes of printable ASCII", c); }
                  }
                  else {
                    println!("! Read {} bytes of binary", c);
                    for id in regexset.matches(&buf[..c]).into_iter() {
                      println!("^ Matches pattern {}", binary_matches[id].0);
                    }
                    for printable in printables {
                      if printable.len() > 3 {
                        let data = String::from_utf8_lossy(printable);
                        for line in data.lines() {
                          println!("$ {}", line);
                        }
                      }
                    }
                    for line in mbstring.lines() {
                      if line.len() > 3 { println!("% {}", line); }
                    }
                    if app.print_binary {
                      let hex = to_hex(&buf[..c]);
                      for line in hex.lines() { println!(". {}", line); }
                    }
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("READ TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("READ ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  break;
                }
              }
              match stream.take_error() {
                Ok(opt) => {
                  if opt.is_some() {
                    println!("ERROR TCP {} from {}: {}", portno, addr, opt.unwrap().to_string());
                    break;
                  }
                }
                Err(_) => {
                  println!("This shouldn't happen...");
                  break;
                }
              }
            }
          });
        }
      });
    }
    else if !port["udp"].is_badvalue() {
      println!("UDP port {}", port["udp"].as_i64().unwrap());
    }
    else {
      println!("Invalid port specification in configuration file");
    }
  }

  loop { thread::sleep(Duration::new(60, 0)); }
}
<|endoftext|>


<commit_before>// Copyright (c) 2019, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::encoder::FrameInvariants;
use crate::quantize::ac_q;
use crate::quantize::select_ac_qi;
use crate::quantize::select_dc_qi;

// The number of frame sub-types for which we track distinct parameters.
pub const FRAME_NSUBTYPES: usize = 4;

pub const FRAME_SUBTYPE_I: usize = 0;
pub const FRAME_SUBTYPE_P: usize = 1;
pub const FRAME_SUBTYPE_B0: usize = 2;
pub const FRAME_SUBTYPE_B1: usize = 3;

// The scale of AV1 quantizer tables (relative to the pixel domain), i.e., Q3.
const QSCALE: i32 = 3;

// Integer binary logarithm of a 64-bit value.
// v: A 64-bit value.
// Returns floor(log2(v)) + 1, or 0 if v == 0.
// This is the number of bits that would be required to represent v in two's
//  complement notation with all of the leading zeros stripped.
// TODO: Mark const once leading_zeros() as a constant function stabilizes.
fn ilog64(v: i64) -> i32 {
  64 - (v.leading_zeros() as i32)
}

// Convert an integer into a Q57 fixed-point fraction.
// The integer must be in the range -64 to 63, inclusive.
const fn q57(v: i32) -> i64 {
  // TODO: Add assert if it ever becomes possible to do in a const function.
  (v as i64) << 57
}

#[cfg_attr(rustfmt, rustfmt_skip)]
const ATANH_LOG2: &[i64; 32] = &[
  0x32B803473F7AD0F4, 0x2F2A71BD4E25E916, 0x2E68B244BB93BA06,
  0x2E39FB9198CE62E4, 0x2E2E683F68565C8F, 0x2E2B850BE2077FC1,
  0x2E2ACC58FE7B78DB, 0x2E2A9E2DE52FD5F2, 0x2E2A92A338D53EEC,
  0x2E2A8FC08F5E19B6, 0x2E2A8F07E51A485E, 0x2E2A8ED9BA8AF388,
  0x2E2A8ECE2FE7384A, 0x2E2A8ECB4D3E4B1A, 0x2E2A8ECA94940FE8,
  0x2E2A8ECA6669811D, 0x2E2A8ECA5ADEDD6A, 0x2E2A8ECA57FC347E,
  0x2E2A8ECA57438A43, 0x2E2A8ECA57155FB4, 0x2E2A8ECA5709D510,
  0x2E2A8ECA5706F267, 0x2E2A8ECA570639BD, 0x2E2A8ECA57060B92,
  0x2E2A8ECA57060008, 0x2E2A8ECA5705FD25, 0x2E2A8ECA5705FC6C,
  0x2E2A8ECA5705FC3E, 0x2E2A8ECA5705FC33, 0x2E2A8ECA5705FC30,
  0x2E2A8ECA5705FC2F, 0x2E2A8ECA5705FC2F
];

// Computes the binary exponential of logq57.
// input: a log base 2 in Q57 format.
// output: a 64 bit integer in Q0 (no fraction).
// TODO: Mark const once we can use local variables in a const function.
fn bexp64(logq57: i64) -> i64 {
  let ipart = (logq57 >> 57) as i32;
  if ipart < 0 {
    return 0;
  }
  if ipart >= 63 {
    return 0x7FFFFFFFFFFFFFFF;
  }
  // z is the fractional part of the log in Q62 format.
  // We need 1 bit of headroom since the magnitude can get larger than 1
  //  during the iteration, and a sign bit.
  let mut z = logq57 - q57(ipart);
  let mut w: i64;
  if z != 0 {
    // Rust has 128 bit multiplies, so it should be possible to do this
    //  faster without losing accuracy.
    z <<= 5;
    // w is the exponential in Q61 format (since it also needs headroom and can
    //  get as large as 2.0); we could get another bit if we dropped the sign,
    //  but we'll recover that bit later anyway.
    // Ideally this should start out as
    //   \lim_{n->\infty} 2^{61}/\product_{i=1}^n \sqrt{1-2^{-2i}}
    //  but in order to guarantee convergence we have to repeat iterations 4,
    //  13 (=3*4+1), and 40 (=3*13+1, etc.), so it winds up somewhat larger.
    w = 0x26A3D0E401DD846D;
    let mut i: i64 = 0;
    loop {
      let mask = -((z < 0) as i64);
      w += ((w >> (i + 1)) + mask) ^ mask;
      z -= (ATANH_LOG2[i as usize] + mask) ^ mask;
      // Repeat iteration 4.
      if i >= 3 {
        break;
      }
      z *= 2;
      i += 1;
    }
    loop {
      let mask = -((z < 0) as i64);
      w += ((w >> (i + 1)) + mask) ^ mask;
      z -= (ATANH_LOG2[i as usize] + mask) ^ mask;
      // Repeat iteration 13.
      if i >= 12 {
        break;
      }
      z *= 2;
      i += 1;
    }
    while i < 32 {
      let mask = -((z < 0) as i64);
      w += ((w >> (i + 1)) + mask) ^ mask;
      z = (z - ((ATANH_LOG2[i as usize] + mask) ^ mask)) * 2;
      i += 1;
    }
    // Skip the remaining iterations unless we really require that much
    //  precision.
    // We could have bailed out earlier for smaller iparts, but that would
    //  require initializing w from a table, as the limit doesn't converge to
    //  61-bit precision until n=30.
    let mut wlo: i32 = 0;
    if ipart > 30 {
      // For these iterations, we just update the low bits, as the high bits
      //  can't possibly be affected.
      // OD_ATANH_LOG2 has also converged (it actually did so one iteration
      //  earlier, but that's no reason for an extra special case).
      loop {
        let mask = -((z < 0) as i64);
        wlo += (((w >> i) + mask) ^ mask) as i32;
        z -= (ATANH_LOG2[31] + mask) ^ mask;
        // Repeat iteration 40.
        if i >= 39 {
          break;
        }
        z *= 2;
        i += 1;
      }
      while i < 61 {
        let mask = -((z < 0) as i64);
        wlo += (((w >> i) + mask) ^ mask) as i32;
        z = (z - ((ATANH_LOG2[31] + mask) ^ mask)) * 2;
        i += 1;
      }
    }
    w = (w << 1) + (wlo as i64);
  } else {
    w = 1i64 << 62;
  }
  if ipart < 62 {
    w = ((w >> (61 - ipart)) + 1) >> 1;
  }
  w
}

// Computes the binary log of w.
// input: a 64-bit integer in Q0 (no fraction).
// output: a 64-bit log in Q57.
// TODO: Mark const once we can use local variables in a const function.
fn blog64(w: i64) -> i64 {
  let mut w = w;
  if w <= 0 {
    return -1;
  }
  let ipart = ilog64(w) - 1;
  if ipart > 61 {
    w >>= ipart - 61;
  } else {
    w <<= 61 - ipart;
  }
  // z is the fractional part of the log in Q61 format.
  let mut z: i64 = 0;
  if (w & (w - 1)) != 0 {
    // Rust has 128 bit multiplies, so it should be possible to do this
    //  faster without losing accuracy.
    // x and y are the cosh() and sinh(), respectively, in Q61 format.
    // We are computing z = 2*atanh(y/x) = 2*atanh((w - 1)/(w + 1)).
    let mut x = w + (1i64 << 61);
    let mut y = w - (1i64 << 61);
    for i in 0..4 {
      let mask = -((y < 0) as i64);
      z += ((ATANH_LOG2[i as usize] >> i) + mask) ^ mask;
      let u = x >> (i + 1);
      x -= ((y >> (i + 1)) + mask) ^ mask;
      y -= (u + mask) ^ mask;
    }
    // Repeat iteration 4.
    for i in 3..13 {
      let mask = -((y < 0) as i64);
      z += ((ATANH_LOG2[i as usize] >> i) + mask) ^ mask;
      let u = x >> (i + 1);
      x -= ((y >> (i + 1)) + mask) ^ mask;
      y -= (u + mask) ^ mask;
    }
    // Repeat iteration 13.
    for i in 12..32 {
      let mask = -((y < 0) as i64);
      z += ((ATANH_LOG2[i as usize] >> i) + mask) ^ mask;
      let u = x >> (i + 1);
      x -= ((y >> (i + 1)) + mask) ^ mask;
      y -= (u + mask) ^ mask;
    }
    // OD_ATANH_LOG2 has converged.
    for i in 32..40 {
      let mask = -((y < 0) as i64);
      z += ((ATANH_LOG2[31] >> i) + mask) ^ mask;
      let u = x >> (i + 1);
      x -= ((y >> (i + 1)) + mask) ^ mask;
      y -= (u + mask) ^ mask;
    }
    // Repeat iteration 40.
    for i in 39..62 {
      let mask = -((y < 0) as i64);
      z += ((ATANH_LOG2[31] >> i) + mask) ^ mask;
      let u = x >> (i + 1);
      x -= ((y >> (i + 1)) + mask) ^ mask;
      y -= (u + mask) ^ mask;
    }
    z = (z + 8) >> 4;
  }
  q57(ipart) + z
}

pub struct RCState {
  // TODO: Add state needed for actual bitrate targeting.
}

// TODO: Separate qi values for each color plane.
pub struct QuantizerParameters {
  // The full-precision log quantizer modulated by the current frame type upon
  //  which our quantizer indices are based (including any adjustments to
  //  prevent buffer overrun or underrun).
  // This is used when estimating the scale parameter once we know the actual
  //  bit usage of a frame.
  pub log_target_q: i64,
  pub dc_qi: u8,
  pub ac_qi: u8,
  pub lambda: f64
}

const Q57_SQUARE_EXP_SCALE: f64 =
  (2.0 * ::std::f64::consts::LN_2) / ((1i64 << 57) as f64);

impl QuantizerParameters {
  fn new_from_log_q(log_target_q: i64, bit_depth: i32) -> QuantizerParameters {
    let quantizer = bexp64(log_target_q + q57(QSCALE + bit_depth - 8));
    QuantizerParameters {
      log_target_q,
      dc_qi: select_dc_qi(quantizer, bit_depth as usize),
      ac_qi: select_ac_qi(quantizer, bit_depth as usize),
      lambda: (::std::f64::consts::LN_2 / 6.0)
        * ((log_target_q as f64) * Q57_SQUARE_EXP_SCALE).exp()
    }
  }
}

impl RCState {
  pub fn new() -> RCState {
    RCState {}
  }

  // TODO: Separate quantizers for Cb and Cr.
  pub fn select_qi(
    &self, fi: &FrameInvariants, fti: usize
  ) -> QuantizerParameters {
    // Rate control is not active.
    // Derive quantizer directly from frame type.
    // TODO: Rename "quantizer" something that indicates it is a quantizer
    //  index, and move it somewhere more sensible (or choose a better way to
    //  parameterize a "quality" configuration parameter).
    let base_qi = fi.config.quantizer;
    let bit_depth = fi.sequence.bit_depth as i32;
    // Adjust base_qi for the frame type.
    // TODO: Adjust the quantizer (not the index) instead to avoid issues with
    //  the non-linearity of the AV1 quantizer tables.
    let qi = if fti == FRAME_SUBTYPE_I {
      let q_boost = 15;
      base_qi.max(1 + q_boost).min(255 + q_boost) - q_boost
    } else {
      let q_drop = 15 * (fti - FRAME_SUBTYPE_P);
      base_qi.min(255 - q_drop) + q_drop
    };
    // We use the AC quantizer as the source quantizer since its quantizer
    //  tables have unique entries, while the DC tables do not.
    let quantizer = ac_q(qi as u8, 0, bit_depth as usize) as i64;
    // Get the log quantizer as Q57.
    let log_q = blog64(quantizer) - q57(QSCALE + bit_depth - 8);
    QuantizerParameters::new_from_log_q(log_q, bit_depth)
  }
}
<commit_msg>Pick a quantizer between similar DC and AC entries<commit_after>// Copyright (c) 2019, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::encoder::FrameInvariants;
use crate::quantize::ac_q;
use crate::quantize::dc_q;
use crate::quantize::select_ac_qi;
use crate::quantize::select_dc_qi;

// The number of frame sub-types for which we track distinct parameters.
pub const FRAME_NSUBTYPES: usize = 4;

pub const FRAME_SUBTYPE_I: usize = 0;
pub const FRAME_SUBTYPE_P: usize = 1;
pub const FRAME_SUBTYPE_B0: usize = 2;
pub const FRAME_SUBTYPE_B1: usize = 3;

// The scale of AV1 quantizer tables (relative to the pixel domain), i.e., Q3.
const QSCALE: i32 = 3;

// Integer binary logarithm of a 64-bit value.
// v: A 64-bit value.
// Returns floor(log2(v)) + 1, or 0 if v == 0.
// This is the number of bits that would be required to represent v in two's
//  complement notation with all of the leading zeros stripped.
// TODO: Mark const once leading_zeros() as a constant function stabilizes.
fn ilog64(v: i64) -> i32 {
  64 - (v.leading_zeros() as i32)
}

// Convert an integer into a Q57 fixed-point fraction.
// The integer must be in the range -64 to 63, inclusive.
const fn q57(v: i32) -> i64 {
  // TODO: Add assert if it ever becomes possible to do in a const function.
  (v as i64) << 57
}

#[cfg_attr(rustfmt, rustfmt_skip)]
const ATANH_LOG2: &[i64; 32] = &[
  0x32B803473F7AD0F4, 0x2F2A71BD4E25E916, 0x2E68B244BB93BA06,
  0x2E39FB9198CE62E4, 0x2E2E683F68565C8F, 0x2E2B850BE2077FC1,
  0x2E2ACC58FE7B78DB, 0x2E2A9E2DE52FD5F2, 0x2E2A92A338D53EEC,
  0x2E2A8FC08F5E19B6, 0x2E2A8F07E51A485E, 0x2E2A8ED9BA8AF388,
  0x2E2A8ECE2FE7384A, 0x2E2A8ECB4D3E4B1A, 0x2E2A8ECA94940FE8,
  0x2E2A8ECA6669811D, 0x2E2A8ECA5ADEDD6A, 0x2E2A8ECA57FC347E,
  0x2E2A8ECA57438A43, 0x2E2A8ECA57155FB4, 0x2E2A8ECA5709D510,
  0x2E2A8ECA5706F267, 0x2E2A8ECA570639BD, 0x2E2A8ECA57060B92,
  0x2E2A8ECA57060008, 0x2E2A8ECA5705FD25, 0x2E2A8ECA5705FC6C,
  0x2E2A8ECA5705FC3E, 0x2E2A8ECA5705FC33, 0x2E2A8ECA5705FC30,
  0x2E2A8ECA5705FC2F, 0x2E2A8ECA5705FC2F
];

// Computes the binary exponential of logq57.
// input: a log base 2 in Q57 format.
// output: a 64 bit integer in Q0 (no fraction).
// TODO: Mark const once we can use local variables in a const function.
fn bexp64(logq57: i64) -> i64 {
  let ipart = (logq57 >> 57) as i32;
  if ipart < 0 {
    return 0;
  }
  if ipart >= 63 {
    return 0x7FFFFFFFFFFFFFFF;
  }
  // z is the fractional part of the log in Q62 format.
  // We need 1 bit of headroom since the magnitude can get larger than 1
  //  during the iteration, and a sign bit.
  let mut z = logq57 - q57(ipart);
  let mut w: i64;
  if z != 0 {
    // Rust has 128 bit multiplies, so it should be possible to do this
    //  faster without losing accuracy.
    z <<= 5;
    // w is the exponential in Q61 format (since it also needs headroom and can
    //  get as large as 2.0); we could get another bit if we dropped the sign,
    //  but we'll recover that bit later anyway.
    // Ideally this should start out as
    //   \lim_{n->\infty} 2^{61}/\product_{i=1}^n \sqrt{1-2^{-2i}}
    //  but in order to guarantee convergence we have to repeat iterations 4,
    //  13 (=3*4+1), and 40 (=3*13+1, etc.), so it winds up somewhat larger.
    w = 0x26A3D0E401DD846D;
    let mut i: i64 = 0;
    loop {
      let mask = -((z < 0) as i64);
      w += ((w >> (i + 1)) + mask) ^ mask;
      z -= (ATANH_LOG2[i as usize] + mask) ^ mask;
      // Repeat iteration 4.
      if i >= 3 {
        break;
      }
      z *= 2;
      i += 1;
    }
    loop {
      let mask = -((z < 0) as i64);
      w += ((w >> (i + 1)) + mask) ^ mask;
      z -= (ATANH_LOG2[i as usize] + mask) ^ mask;
      // Repeat iteration 13.
      if i >= 12 {
        break;
      }
      z *= 2;
      i += 1;
    }
    while i < 32 {
      let mask = -((z < 0) as i64);
      w += ((w >> (i + 1)) + mask) ^ mask;
      z = (z - ((ATANH_LOG2[i as usize] + mask) ^ mask)) * 2;
      i += 1;
    }
    // Skip the remaining iterations unless we really require that much
    //  precision.
    // We could have bailed out earlier for smaller iparts, but that would
    //  require initializing w from a table, as the limit doesn't converge to
    //  61-bit precision until n=30.
    let mut wlo: i32 = 0;
    if ipart > 30 {
      // For these iterations, we just update the low bits, as the high bits
      //  can't possibly be affected.
      // OD_ATANH_LOG2 has also converged (it actually did so one iteration
      //  earlier, but that's no reason for an extra special case).
      loop {
        let mask = -((z < 0) as i64);
        wlo += (((w >> i) + mask) ^ mask) as i32;
        z -= (ATANH_LOG2[31] + mask) ^ mask;
        // Repeat iteration 40.
        if i >= 39 {
          break;
        }
        z *= 2;
        i += 1;
      }
      while i < 61 {
        let mask = -((z < 0) as i64);
        wlo += (((w >> i) + mask) ^ mask) as i32;
        z = (z - ((ATANH_LOG2[31] + mask) ^ mask)) * 2;
        i += 1;
      }
    }
    w = (w << 1) + (wlo as i64);
  } else {
    w = 1i64 << 62;
  }
  if ipart < 62 {
    w = ((w >> (61 - ipart)) + 1) >> 1;
  }
  w
}

// Computes the binary log of w.
// input: a 64-bit integer in Q0 (no fraction).
// output: a 64-bit log in Q57.
// TODO: Mark const once we can use local variables in a const function.
fn blog64(w: i64) -> i64 {
  let mut w = w;
  if w <= 0 {
    return -1;
  }
  let ipart = ilog64(w) - 1;
  if ipart > 61 {
    w >>= ipart - 61;
  } else {
    w <<= 61 - ipart;
  }
  // z is the fractional part of the log in Q61 format.
  let mut z: i64 = 0;
  if (w & (w - 1)) != 0 {
    // Rust has 128 bit multiplies, so it should be possible to do this
    //  faster without losing accuracy.
    // x and y are the cosh() and sinh(), respectively, in Q61 format.
    // We are computing z = 2*atanh(y/x) = 2*atanh((w - 1)/(w + 1)).
    let mut x = w + (1i64 << 61);
    let mut y = w - (1i64 << 61);
    for i in 0..4 {
      let mask = -((y < 0) as i64);
      z += ((ATANH_LOG2[i as usize] >> i) + mask) ^ mask;
      let u = x >> (i + 1);
      x -= ((y >> (i + 1)) + mask) ^ mask;
      y -= (u + mask) ^ mask;
    }
    // Repeat iteration 4.
    for i in 3..13 {
      let mask = -((y < 0) as i64);
      z += ((ATANH_LOG2[i as usize] >> i) + mask) ^ mask;
      let u = x >> (i + 1);
      x -= ((y >> (i + 1)) + mask) ^ mask;
      y -= (u + mask) ^ mask;
    }
    // Repeat iteration 13.
    for i in 12..32 {
      let mask = -((y < 0) as i64);
      z += ((ATANH_LOG2[i as usize] >> i) + mask) ^ mask;
      let u = x >> (i + 1);
      x -= ((y >> (i + 1)) + mask) ^ mask;
      y -= (u + mask) ^ mask;
    }
    // OD_ATANH_LOG2 has converged.
    for i in 32..40 {
      let mask = -((y < 0) as i64);
      z += ((ATANH_LOG2[31] >> i) + mask) ^ mask;
      let u = x >> (i + 1);
      x -= ((y >> (i + 1)) + mask) ^ mask;
      y -= (u + mask) ^ mask;
    }
    // Repeat iteration 40.
    for i in 39..62 {
      let mask = -((y < 0) as i64);
      z += ((ATANH_LOG2[31] >> i) + mask) ^ mask;
      let u = x >> (i + 1);
      x -= ((y >> (i + 1)) + mask) ^ mask;
      y -= (u + mask) ^ mask;
    }
    z = (z + 8) >> 4;
  }
  q57(ipart) + z
}

pub struct RCState {
  // TODO: Add state needed for actual bitrate targeting.
}

// TODO: Separate qi values for each color plane.
pub struct QuantizerParameters {
  // The full-precision log quantizer modulated by the current frame type upon
  //  which our quantizer indices are based (including any adjustments to
  //  prevent buffer overrun or underrun).
  // This is used when estimating the scale parameter once we know the actual
  //  bit usage of a frame.
  pub log_target_q: i64,
  pub dc_qi: u8,
  pub ac_qi: u8,
  pub lambda: f64
}

const Q57_SQUARE_EXP_SCALE: f64 =
  (2.0 * ::std::f64::consts::LN_2) / ((1i64 << 57) as f64);

impl QuantizerParameters {
  fn new_from_log_q(log_target_q: i64, bit_depth: i32) -> QuantizerParameters {
    let quantizer = bexp64(log_target_q + q57(QSCALE + bit_depth - 8));
    QuantizerParameters {
      log_target_q,
      dc_qi: select_dc_qi(quantizer, bit_depth as usize),
      ac_qi: select_ac_qi(quantizer, bit_depth as usize),
      lambda: (::std::f64::consts::LN_2 / 6.0)
        * ((log_target_q as f64) * Q57_SQUARE_EXP_SCALE).exp()
    }
  }
}

impl RCState {
  pub fn new() -> RCState {
    RCState {}
  }

  // TODO: Separate quantizers for Cb and Cr.
  pub fn select_qi(
    &self, fi: &FrameInvariants, fti: usize
  ) -> QuantizerParameters {
    // Rate control is not active.
    // Derive quantizer directly from frame type.
    // TODO: Rename "quantizer" something that indicates it is a quantizer
    //  index, and move it somewhere more sensible (or choose a better way to
    //  parameterize a "quality" configuration parameter).
    let base_qi = fi.config.quantizer;
    let bit_depth = fi.sequence.bit_depth as i32;
    // Adjust base_qi for the frame type.
    // TODO: Adjust the quantizer (not the index) instead to avoid issues with
    //  the non-linearity of the AV1 quantizer tables.
    let qi = if fti == FRAME_SUBTYPE_I {
      let q_boost = 15;
      base_qi.max(1 + q_boost).min(255 + q_boost) - q_boost
    } else {
      let q_drop = 15 * (fti - FRAME_SUBTYPE_P);
      base_qi.min(255 - q_drop) + q_drop
    };
    // We use the AC quantizer as the source quantizer since its quantizer
    //  tables have unique entries, while the DC tables do not.
    let quantizer = ac_q(qi as u8, 0, bit_depth as usize) as i64;
    // Pick the nearest DC entry since an exact match may be unavailable.
    let dc_qi = select_dc_qi(quantizer, bit_depth as usize);
    let dc_quantizer = dc_q(dc_qi as u8, 0, bit_depth as usize) as i64;
    // Get the log quantizer as Q57.
    let log_q = blog64(quantizer) - q57(QSCALE + bit_depth - 8);
    let log_dc_q = blog64(dc_quantizer) - q57(QSCALE + bit_depth - 8);
    // Target the midpoint of the chosen entries.
    QuantizerParameters::new_from_log_q((log_q + log_dc_q + 1) / 2, bit_depth)
  }
}
<|endoftext|>


<commit_before><commit_msg>signal: Actually make unix::bsd public<commit_after><|endoftext|>


<commit_before>extern crate env_logger;
extern crate log_panics;
extern crate log;
#[cfg(target_os = "android")]
extern crate android_logger;

use self::env_logger::Builder;
use self::log::{Record, LevelFilter, Level};
use std::env;
use std::io::Write;
#[cfg(target_os = "android")]
use self::android_logger::Filter;

pub struct LoggerUtils {}

impl LoggerUtils {
    pub fn init() {
        log_panics::init(); //Logging of panics is essential for android. As android does not log to stdout for native code
        if cfg!(target_os = "android") {
            //Set logging to off when deploying production android app.
            #[cfg(target_os = "android")]
                android_logger::init_once(
                Filter::default().with_min_level(log::Level::Trace)
            );
            info!("Logging for Android");
        } else {
            Builder::new()
                .format(|buf, record| writeln!(buf, "{:>5}|{:<30}|{:>35}:{:<4}| {}", record.level(), record.target(), record.file().get_or_insert(""), record.line().get_or_insert(0), record.args()))
                .filter(None, LevelFilter::Off)
                .parse(env::var("RUST_LOG").as_ref().map(String::as_str).unwrap_or(""))
                .try_init()
                .ok();
        }
    }
}

#[macro_export]
macro_rules! try_log {
    ($expr:expr) => (match $expr {
        Ok(val) => val,
        Err(err) => {
            error!("try_log! | {}", err);
            return Err(From::from(err))
        }
    })
}

macro_rules! _map_err {
    ($lvl:expr, $expr:expr) => (
        |err| {
            log!($lvl, "{} - {}", $expr, err);
            err
        }
    );
    ($lvl:expr) => (
        |err| {
            log!($lvl, "{}", err);
            err
        }
    )
}

#[macro_export]
macro_rules! map_err_err {
    () => ( _map_err!(::log::Level::Error) );
    ($($arg:tt)*) => ( _map_err!(::log::Level::Error, $($arg)*) )
}

#[macro_export]
macro_rules! map_err_trace {
    () => ( _map_err!(::log::Level::Trace) );
    ($($arg:tt)*) => ( _map_err!(::log::Level::Trace, $($arg)*) )
}

#[macro_export]
macro_rules! map_err_info {
    () => ( _map_err!(::log::Level::Info) );
    ($($arg:tt)*) => ( _map_err!(::log::Level::Info, $($arg)*) )
}
<commit_msg>Logger fix<commit_after>extern crate env_logger;
extern crate log_panics;
extern crate log;
#[cfg(target_os = "android")]
extern crate android_logger;

use self::env_logger::Builder;
use self::log::{Record, LevelFilter, Level};
use std::env;
use std::io::Write;
#[cfg(target_os = "android")]
use self::android_logger::Filter;

pub struct LoggerUtils {}

impl LoggerUtils {
    pub fn init() {
        log_panics::init(); //Logging of panics is essential for android. As android does not log to stdout for native code

        if cfg!(target_os = "android") {
            let log_filter =match env::var("SODIUM_LIB_DIR") {
                Ok(val) => Filter::parse(env::var("RUST_LOG").as_ref().map(String::as_str).unwrap_or("")).try_init().ok() ,
                Err(..) => Filter::default().with_min_level(log::Level::Off);
            };

            //Set logging to off when deploying production android app.
            #[cfg(target_os = "android")]
                android_logger::init_once(log_filter);
            info!("Logging for Android");
        } else {
            Builder::new()
                .format(|buf, record| writeln!(buf, "{:>5}|{:<30}|{:>35}:{:<4}| {}", record.level(), record.target(), record.file().get_or_insert(""), record.line().get_or_insert(0), record.args()))
                .filter(None, LevelFilter::Off)
                .parse(env::var("RUST_LOG").as_ref().map(String::as_str).unwrap_or(""))
                .try_init()
                .ok();
        }
    }
}

#[macro_export]
macro_rules! try_log {
    ($expr:expr) => (match $expr {
        Ok(val) => val,
        Err(err) => {
            error!("try_log! | {}", err);
            return Err(From::from(err))
        }
    })
}

macro_rules! _map_err {
    ($lvl:expr, $expr:expr) => (
        |err| {
            log!($lvl, "{} - {}", $expr, err);
            err
        }
    );
    ($lvl:expr) => (
        |err| {
            log!($lvl, "{}", err);
            err
        }
    )
}

#[macro_export]
macro_rules! map_err_err {
    () => ( _map_err!(::log::Level::Error) );
    ($($arg:tt)*) => ( _map_err!(::log::Level::Error, $($arg)*) )
}

#[macro_export]
macro_rules! map_err_trace {
    () => ( _map_err!(::log::Level::Trace) );
    ($($arg:tt)*) => ( _map_err!(::log::Level::Trace, $($arg)*) )
}

#[macro_export]
macro_rules! map_err_info {
    () => ( _map_err!(::log::Level::Info) );
    ($($arg:tt)*) => ( _map_err!(::log::Level::Info, $($arg)*) )
}
<|endoftext|>


<commit_before>use std::fmt;
use hex_slice::AsHex;
use rand::{random, Rng, XorShiftRng};
use std::rc::Rc;
use std::cell::RefCell;

/// Something that can extract information from an `InfoSource`.
pub trait InfoSink {
    /// The output data.
    type Out;
    /// Called by [InfoSource::draw](trait.InfoSource.html#tymethod.draw)
    fn sink<I: InfoSource>(&mut self, i: &mut I) -> Self::Out;
}

/// Something that an act as a source of test data.
pub trait InfoSource {
    /// Take a single byte from the source.
    fn draw_u8(&mut self) -> u8;

    /// Call F with access to the data source.
    fn draw<S: InfoSink>(&mut self, sink: S) -> S::Out
    where
        Self: Sized;
}

/// Generates data from an underlying Rng instance.
#[derive(Debug)]
pub struct RngSource<R> {
    rng: R,
}

/// An adapter that can record the data drawn from an underlying source.
#[derive(Debug)]
pub struct InfoRecorder<I> {
    inner: Rc<RefCell<I>>,
    pub(crate) data: Vec<u8>,
}

impl<'a, I: InfoSource + ?Sized> InfoSource for &'a mut I {
    fn draw_u8(&mut self) -> u8 {
        (**self).draw_u8()
    }
    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        sink.sink(self)
    }
}

impl<I> InfoRecorder<I> {
    /// Creates a recording InfoSource.
    pub fn new(inner: I) -> Self {
        InfoRecorder {
            inner: Rc::new(RefCell::new(inner)),
            data: Vec::new(),
        }
    }

    /// Extracts the data recorded.
    pub fn into_pool(self) -> InfoPool {
        InfoPool::of_vec(self.data)
    }
}

impl<I: InfoSource> InfoSource for InfoRecorder<I> {
    fn draw_u8(&mut self) -> u8 {
        let byte = self.inner.borrow_mut().draw_u8();
        self.data.push(byte);
        byte
    }

    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        let start = self.data.len();
        debug!("-> InfoRecorder::draw @{}", start);
        let res = sink.sink(self);
        let end = self.data.len();
        debug!("<- InfoRecorder::draw @{}", end);
        debug!("Span: {:?}", (start, end));
        res
    }
}

impl RngSource<XorShiftRng> {
    /// Creates a RngSource with a randomly seeded XorShift generator.
    pub fn new() -> Self {
        let rng = random::<XorShiftRng>();
        RngSource { rng }
    }
}

impl<R: Rng> InfoSource for RngSource<R> {
    fn draw_u8(&mut self) -> u8 {
        self.rng.gen::<u8>()
    }
    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        sink.sink(self)
    }
}

/// A pool of data that we can draw upon to generate other types of data.
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct InfoPool {
    pub(in data) data: Vec<u8>,
}

/// A handle to an info Pool that we can draw replayed bytes from, and zero after.
#[derive(Clone)]
pub struct InfoReplay<'a> {
    data: &'a [u8],
    off: usize,
}

impl fmt::Debug for InfoPool {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("InfoPool")
            .field("data", &format_args!("{:x}", self.data.as_hex()))
            .finish()
    }
}

/// The reasons why drawing data from a pool can fail.
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum DataError {
    /// Not enough data to generate a value
    PoolExhausted,
    /// One of our combinators said that we should not test this value.
    SkipItem,
}

impl InfoPool {
    /// Create an `InfoPool` with a given vector of bytes. (Mostly used for
    /// testing).
    pub fn of_vec(data: Vec<u8>) -> Self {
        InfoPool { data: data }
    }

    /// Create an `InfoPool` with a `size` length vector of random bytes
    /// using the generator `rng`. (Mostly used for testing).
    pub fn new() -> Self {
        Self { data: Vec::new() }
    }

    /// Allows access to the underlying buffer.
    pub fn buffer(&self) -> &[u8] {
        &*self.data
    }

    /// Creates a tap that allows drawing new information from this pool.
    pub fn replay(&self) -> InfoReplay {
        InfoReplay {
            data: &*self.data,
            off: 0,
        }
    }
}

impl<'a> InfoReplay<'a> {
    /// Consumes the next byte from this tap. Returns `Ok(x)` if successful,
    /// or `Err(DataError::PoolExhausted)` if we have reached the end.
    pub fn next_byte(&mut self) -> u8 {
        if let Some(res) = self.data.get(self.off).cloned() {
            self.off += 1;
            res
        } else {
            0
        }
    }
}

impl<'a> InfoSource for InfoReplay<'a> {
    fn draw_u8(&mut self) -> u8 {
        self.next_byte()
    }

    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        sink.sink(self)
    }
}

#[cfg(test)]
mod tests {
    extern crate env_logger;
    use super::*;
    impl<R: Rng> RngSource<R> {
        pub(crate) fn of(rng: R) -> Self {
            RngSource { rng }
        }
    }

    #[test]
    fn should_take_each_item_in_pool() {
        let p = InfoPool::of_vec(vec![0, 1, 2, 3]);
        let mut t = p.replay();
        assert_eq!(t.next_byte(), 0);
        assert_eq!(t.next_byte(), 1);
        assert_eq!(t.next_byte(), 2);
        assert_eq!(t.next_byte(), 3);
        assert_eq!(t.next_byte(), 0);
    }

    #[test]
    fn should_generate_random_data() {
        let trials = 1024usize;
        let mut vals = 0;
        let mut p = RngSource::of(XorShiftRng::new_unseeded());
        let error = 8;
        for _ in 0..trials {
            vals += p.draw_u8() as usize;
        }
        let mean = vals / trials;
        let expected = 128;
        assert!(
            (expected - error) < mean && (expected + error) > mean,
            "Expected {} trials to be ({}+/-{}); got {}",
            trials,
            expected,
            error,
            mean
        )
    }

    #[test]
    fn should_allow_restarting_read() {
        let mut p = InfoRecorder::new(RngSource::new());
        let mut v0 = Vec::new();
        {
            for _ in 0..4 {
                v0.push(p.draw_u8())
            }
        }

        let p = p.into_pool();
        let mut t = p.replay();
        let mut v1 = Vec::new();
        for _ in 0..4 {
            v1.push(t.draw_u8())
        }

        assert_eq!(v0, v1)
    }

    #[test]
    fn should_allow_restarting_child_reads() {
        struct FnSink<F>(F);
        impl<F: FnMut(&mut InfoSource) -> R, R> InfoSink for FnSink<F> {
            type Out = R;
            fn sink<I: InfoSource>(&mut self, k: &mut I) -> R {
                (self.0)(k as &mut InfoSource)
            }
        }

        let mut p = InfoRecorder::new(RngSource::new());
        let mut v0 = Vec::new();
        p.draw(FnSink(|src: &mut InfoSource| {
            for _ in 0..4 {
                let x: u8 = src.draw_u8();
                v0.push(x);
            }
        }));

        let p = p.into_pool();
        let mut t = p.replay();
        let mut v1 = Vec::new();
        for _ in 0..4 {
            v1.push(t.draw_u8())
        }

        assert_eq!(v0, v1)
    }

    #[test]
    fn should_allow_restarting_mixed_child_reads() {
        struct FnSink<F>(F);
        impl<F: FnMut(&mut InfoSource) -> R, R> InfoSink for FnSink<F> {
            type Out = R;
            fn sink<I: InfoSource>(&mut self, k: &mut I) -> R {
                (self.0)(k as &mut InfoSource)
            }
        }

        let mut p = InfoRecorder::new(RngSource::new());
        let mut v0 = Vec::new();

        for _ in 0..2 {
            let x: u8 = p.draw_u8();
            v0.push(x);
        }

        p.draw(FnSink(|src: &mut InfoSource| {
            for _ in 0..4 {
                let x: u8 = src.draw_u8();
                v0.push(x);
            }
        }));

        for _ in 0..2 {
            let x: u8 = p.draw_u8();
            v0.push(x);
        }

        let p = p.into_pool();
        let mut t = p.replay();
        let mut v1 = Vec::new();
        for _ in 0..8 {
            v1.push(t.draw_u8())
        }

        assert_eq!(v0, v1)
    }

    #[test]
    fn should_allow_borrowing_buffer() {
        let p = InfoPool::of_vec(vec![1]);
        assert_eq!(p.buffer(), &[1]);
    }

    #[test]
    fn tap_can_act_as_source() {
        let buf = vec![4, 3, 2, 1];
        let p = InfoPool::of_vec(buf.clone());
        let _: &InfoSource = &p.replay();
        let mut res = Vec::new();
        let mut it = p.replay();
        for _ in 0..4 {
            res.push(it.draw_u8())
        }
        assert_eq!(res, buf)
    }

    #[test]
    fn replay_can_act_as_source() {
        let buf = vec![4, 3, 2, 1];
        let p = InfoPool::of_vec(buf.clone());

        let mut res = Vec::new();
        let mut it = p.replay();
        for _ in 0..4 {
            res.push(it.draw_u8())
        }
        assert_eq!(res, buf)
    }

}
<commit_msg>Do not need Rc<RefCell<_>> here.<commit_after>use std::fmt;
use hex_slice::AsHex;
use rand::{random, Rng, XorShiftRng};
use std::rc::Rc;
use std::cell::RefCell;

/// Something that can extract information from an `InfoSource`.
pub trait InfoSink {
    /// The output data.
    type Out;
    /// Called by [InfoSource::draw](trait.InfoSource.html#tymethod.draw)
    fn sink<I: InfoSource>(&mut self, i: &mut I) -> Self::Out;
}

/// Something that an act as a source of test data.
pub trait InfoSource {
    /// Take a single byte from the source.
    fn draw_u8(&mut self) -> u8;

    /// Call F with access to the data source.
    fn draw<S: InfoSink>(&mut self, sink: S) -> S::Out
    where
        Self: Sized;
}

/// Generates data from an underlying Rng instance.
#[derive(Debug)]
pub struct RngSource<R> {
    rng: R,
}

/// An adapter that can record the data drawn from an underlying source.
#[derive(Debug)]
pub struct InfoRecorder<I> {
    inner: I,
    pub(crate) data: Vec<u8>,
}

impl<'a, I: InfoSource + ?Sized> InfoSource for &'a mut I {
    fn draw_u8(&mut self) -> u8 {
        (**self).draw_u8()
    }
    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        sink.sink(self)
    }
}

impl<I> InfoRecorder<I> {
    /// Creates a recording InfoSource.
    pub fn new(inner: I) -> Self {
        InfoRecorder {
            inner: inner,
            data: Vec::new(),
        }
    }

    /// Extracts the data recorded.
    pub fn into_pool(self) -> InfoPool {
        InfoPool::of_vec(self.data)
    }
}

impl<I: InfoSource> InfoSource for InfoRecorder<I> {
    fn draw_u8(&mut self) -> u8 {
        let byte = self.inner.draw_u8();
        self.data.push(byte);
        byte
    }

    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        let start = self.data.len();
        debug!("-> InfoRecorder::draw @{}", start);
        let res = sink.sink(self);
        let end = self.data.len();
        debug!("<- InfoRecorder::draw @{}", end);
        debug!("Span: {:?}", (start, end));
        res
    }
}

impl RngSource<XorShiftRng> {
    /// Creates a RngSource with a randomly seeded XorShift generator.
    pub fn new() -> Self {
        let rng = random::<XorShiftRng>();
        RngSource { rng }
    }
}

impl<R: Rng> InfoSource for RngSource<R> {
    fn draw_u8(&mut self) -> u8 {
        self.rng.gen::<u8>()
    }
    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        sink.sink(self)
    }
}

/// A pool of data that we can draw upon to generate other types of data.
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct InfoPool {
    pub(in data) data: Vec<u8>,
}

/// A handle to an info Pool that we can draw replayed bytes from, and zero after.
#[derive(Clone)]
pub struct InfoReplay<'a> {
    data: &'a [u8],
    off: usize,
}

impl fmt::Debug for InfoPool {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("InfoPool")
            .field("data", &format_args!("{:x}", self.data.as_hex()))
            .finish()
    }
}

/// The reasons why drawing data from a pool can fail.
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum DataError {
    /// Not enough data to generate a value
    PoolExhausted,
    /// One of our combinators said that we should not test this value.
    SkipItem,
}

impl InfoPool {
    /// Create an `InfoPool` with a given vector of bytes. (Mostly used for
    /// testing).
    pub fn of_vec(data: Vec<u8>) -> Self {
        InfoPool { data: data }
    }

    /// Create an `InfoPool` with a `size` length vector of random bytes
    /// using the generator `rng`. (Mostly used for testing).
    pub fn new() -> Self {
        Self { data: Vec::new() }
    }

    /// Allows access to the underlying buffer.
    pub fn buffer(&self) -> &[u8] {
        &*self.data
    }

    /// Creates a tap that allows drawing new information from this pool.
    pub fn replay(&self) -> InfoReplay {
        InfoReplay {
            data: &*self.data,
            off: 0,
        }
    }
}

impl<'a> InfoReplay<'a> {
    /// Consumes the next byte from this tap. Returns `Ok(x)` if successful,
    /// or `Err(DataError::PoolExhausted)` if we have reached the end.
    pub fn next_byte(&mut self) -> u8 {
        if let Some(res) = self.data.get(self.off).cloned() {
            self.off += 1;
            res
        } else {
            0
        }
    }
}

impl<'a> InfoSource for InfoReplay<'a> {
    fn draw_u8(&mut self) -> u8 {
        self.next_byte()
    }

    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        sink.sink(self)
    }
}

#[cfg(test)]
mod tests {
    extern crate env_logger;
    use super::*;
    impl<R: Rng> RngSource<R> {
        pub(crate) fn of(rng: R) -> Self {
            RngSource { rng }
        }
    }

    #[test]
    fn should_take_each_item_in_pool() {
        let p = InfoPool::of_vec(vec![0, 1, 2, 3]);
        let mut t = p.replay();
        assert_eq!(t.next_byte(), 0);
        assert_eq!(t.next_byte(), 1);
        assert_eq!(t.next_byte(), 2);
        assert_eq!(t.next_byte(), 3);
        assert_eq!(t.next_byte(), 0);
    }

    #[test]
    fn should_generate_random_data() {
        let trials = 1024usize;
        let mut vals = 0;
        let mut p = RngSource::of(XorShiftRng::new_unseeded());
        let error = 8;
        for _ in 0..trials {
            vals += p.draw_u8() as usize;
        }
        let mean = vals / trials;
        let expected = 128;
        assert!(
            (expected - error) < mean && (expected + error) > mean,
            "Expected {} trials to be ({}+/-{}); got {}",
            trials,
            expected,
            error,
            mean
        )
    }

    #[test]
    fn should_allow_restarting_read() {
        let mut p = InfoRecorder::new(RngSource::new());
        let mut v0 = Vec::new();
        {
            for _ in 0..4 {
                v0.push(p.draw_u8())
            }
        }

        let p = p.into_pool();
        let mut t = p.replay();
        let mut v1 = Vec::new();
        for _ in 0..4 {
            v1.push(t.draw_u8())
        }

        assert_eq!(v0, v1)
    }

    #[test]
    fn should_allow_restarting_child_reads() {
        struct FnSink<F>(F);
        impl<F: FnMut(&mut InfoSource) -> R, R> InfoSink for FnSink<F> {
            type Out = R;
            fn sink<I: InfoSource>(&mut self, k: &mut I) -> R {
                (self.0)(k as &mut InfoSource)
            }
        }

        let mut p = InfoRecorder::new(RngSource::new());
        let mut v0 = Vec::new();
        p.draw(FnSink(|src: &mut InfoSource| {
            for _ in 0..4 {
                let x: u8 = src.draw_u8();
                v0.push(x);
            }
        }));

        let p = p.into_pool();
        let mut t = p.replay();
        let mut v1 = Vec::new();
        for _ in 0..4 {
            v1.push(t.draw_u8())
        }

        assert_eq!(v0, v1)
    }

    #[test]
    fn should_allow_restarting_mixed_child_reads() {
        struct FnSink<F>(F);
        impl<F: FnMut(&mut InfoSource) -> R, R> InfoSink for FnSink<F> {
            type Out = R;
            fn sink<I: InfoSource>(&mut self, k: &mut I) -> R {
                (self.0)(k as &mut InfoSource)
            }
        }

        let mut p = InfoRecorder::new(RngSource::new());
        let mut v0 = Vec::new();

        for _ in 0..2 {
            let x: u8 = p.draw_u8();
            v0.push(x);
        }

        p.draw(FnSink(|src: &mut InfoSource| {
            for _ in 0..4 {
                let x: u8 = src.draw_u8();
                v0.push(x);
            }
        }));

        for _ in 0..2 {
            let x: u8 = p.draw_u8();
            v0.push(x);
        }

        let p = p.into_pool();
        let mut t = p.replay();
        let mut v1 = Vec::new();
        for _ in 0..8 {
            v1.push(t.draw_u8())
        }

        assert_eq!(v0, v1)
    }

    #[test]
    fn should_allow_borrowing_buffer() {
        let p = InfoPool::of_vec(vec![1]);
        assert_eq!(p.buffer(), &[1]);
    }

    #[test]
    fn tap_can_act_as_source() {
        let buf = vec![4, 3, 2, 1];
        let p = InfoPool::of_vec(buf.clone());
        let _: &InfoSource = &p.replay();
        let mut res = Vec::new();
        let mut it = p.replay();
        for _ in 0..4 {
            res.push(it.draw_u8())
        }
        assert_eq!(res, buf)
    }

    #[test]
    fn replay_can_act_as_source() {
        let buf = vec![4, 3, 2, 1];
        let p = InfoPool::of_vec(buf.clone());

        let mut res = Vec::new();
        let mut it = p.replay();
        for _ in 0..4 {
            res.push(it.draw_u8())
        }
        assert_eq!(res, buf)
    }

}
<|endoftext|>


<commit_before>#![deny(clippy::all, clippy::pedantic)]

use anyhow::anyhow;
use std::error::Error as _;
use std::io;
use thiserror::Error;

#[test]
fn test_transparent_struct() {
    #[derive(Error, Debug)]
    #[error(transparent)]
    struct Error(ErrorKind);

    #[derive(Error, Debug)]
    enum ErrorKind {
        #[error("E0")]
        E0,
        #[error("E1")]
        E1(#[from] io::Error),
    }

    let error = Error(ErrorKind::E0);
    assert_eq!("E0", error.to_string());
    assert!(error.source().is_none());

    let io = io::Error::new(io::ErrorKind::Other, "oh no!");
    let error = Error(ErrorKind::from(io));
    assert_eq!("E1", error.to_string());
    error.source().unwrap().downcast_ref::<io::Error>().unwrap();
}

#[test]
fn test_transparent_enum() {
    #[derive(Error, Debug)]
    enum Error {
        #[error("this failed")]
        This,
        #[error(transparent)]
        Other(anyhow::Error),
    }

    let error = Error::This;
    assert_eq!("this failed", error.to_string());

    let error = Error::Other(anyhow!("inner").context("outer"));
    assert_eq!("outer", error.to_string());
    assert_eq!("inner", error.source().unwrap().to_string());
}

#[test]
fn test_anyhow() {
    #[derive(Error, Debug)]
    #[error(transparent)]
    struct Any(#[from] anyhow::Error);

    let error = Any::from(anyhow!("inner").context("outer"));
    assert_eq!("outer", error.to_string());
    assert_eq!("inner", error.source().unwrap().to_string());
}
<commit_msg>Add regression test for issue 113<commit_after>#![deny(clippy::all, clippy::pedantic)]

use anyhow::anyhow;
use std::error::Error as _;
use std::io;
use thiserror::Error;

#[test]
fn test_transparent_struct() {
    #[derive(Error, Debug)]
    #[error(transparent)]
    struct Error(ErrorKind);

    #[derive(Error, Debug)]
    enum ErrorKind {
        #[error("E0")]
        E0,
        #[error("E1")]
        E1(#[from] io::Error),
    }

    let error = Error(ErrorKind::E0);
    assert_eq!("E0", error.to_string());
    assert!(error.source().is_none());

    let io = io::Error::new(io::ErrorKind::Other, "oh no!");
    let error = Error(ErrorKind::from(io));
    assert_eq!("E1", error.to_string());
    error.source().unwrap().downcast_ref::<io::Error>().unwrap();
}

#[test]
fn test_transparent_enum() {
    #[derive(Error, Debug)]
    enum Error {
        #[error("this failed")]
        This,
        #[error(transparent)]
        Other(anyhow::Error),
    }

    let error = Error::This;
    assert_eq!("this failed", error.to_string());

    let error = Error::Other(anyhow!("inner").context("outer"));
    assert_eq!("outer", error.to_string());
    assert_eq!("inner", error.source().unwrap().to_string());
}

#[test]
fn test_anyhow() {
    #[derive(Error, Debug)]
    #[error(transparent)]
    struct Any(#[from] anyhow::Error);

    let error = Any::from(anyhow!("inner").context("outer"));
    assert_eq!("outer", error.to_string());
    assert_eq!("inner", error.source().unwrap().to_string());
}

#[test]
fn test_non_static() {
    #[derive(Error, Debug)]
    #[error(transparent)]
    struct Error<'a> {
        inner: ErrorKind<'a>,
    }

    #[derive(Error, Debug)]
    enum ErrorKind<'a> {
        #[error("unexpected token: {:?}", token)]
        Unexpected { token: &'a str },
    }

    let error = Error {
        inner: ErrorKind::Unexpected { token: "error" },
    };
    assert_eq!("unexpected token: \"error\"", error.to_string());
    assert!(error.source().is_none());
}
<|endoftext|>


<commit_before>use std::sync::Arc;
use itertools::Itertools;

use fall_tree::{AstNode, AstClass, Text};
use fall_tree::search::{child_of_type, child_of_type_exn};
use fall_tree::visitor::{Visitor, NodeVisitor};

use super::{Analysis, DiagnosticSink};
use ::{Expr, CallExpr, SynRule, IDENT, SIMPLE_STRING, RefKind};

#[derive(Debug, Eq, PartialEq, Clone)]
pub enum CallKind<'f> {
    Eof,
    Any,
    Commit,

    Not(Expr<'f>),
    Layer(Expr<'f>, Expr<'f>),
    WithSkip(Expr<'f>, Expr<'f>),

    Enter(u32, Expr<'f>),
    Exit(u32, Expr<'f>),
    IsIn(u32),

    RuleCall(SynRule<'f>, Arc<Vec<(u32, Expr<'f>)>>),
    // TODO: u32 is unclear
    PrevIs(Arc<Vec<usize>>)
}

pub (super) fn resolve<'f>(a: &Analysis<'f>, d: &mut DiagnosticSink, call: CallExpr<'f>) -> Option<CallKind<'f>> {
    let n_args = call.args().count();
    let expect_args = |d: &mut DiagnosticSink, n_expected: usize| {
        if n_expected != n_args {
            d.error(
                call.node(),
                format!("Wrong number of arguments, expected {}, got {}", n_expected, n_args)
            )
        }
    };

    let zero = |d: &mut DiagnosticSink, kind: CallKind<'f>| {
        expect_args(d, 0);
        Some(kind)
    };
    let mut any = |d: &mut DiagnosticSink| zero(d, CallKind::Any);
    let mut commit = |d: &mut DiagnosticSink| zero(d, CallKind::Commit);
    let mut eof = |d: &mut DiagnosticSink| zero(d, CallKind::Eof);

    let one = |d: &mut DiagnosticSink, kind: &mut FnMut(&mut DiagnosticSink, Expr<'f>) -> Option<CallKind<'f>>| {
        expect_args(d, 1);
        call.args().next().and_then(|arg| kind(d, arg))
    };
    let mut not = |d: &mut DiagnosticSink| one(d, &mut |_, arg| Some(CallKind::Not(arg)));
    let mut is_in = |d: &mut DiagnosticSink| one(d, &mut |d, arg| resolve_context(a, d, arg).map(CallKind::IsIn));

    let two = |d: &mut DiagnosticSink, kind: &mut FnMut(&mut DiagnosticSink, Expr<'f>, Expr<'f>) -> Option<CallKind<'f>>| {
        expect_args(d, 2);
        call.args().next_tuple().and_then(|(arg1, arg2)| kind(d, arg1, arg2))
    };
    let mut layer = |d: &mut DiagnosticSink| two(d, &mut |_, arg1, arg2| {
        Some(CallKind::Layer(arg1, arg2))
    });
    let mut with_skip = |d: &mut DiagnosticSink| two(d, &mut |_, arg1, arg2| {
        Some(CallKind::WithSkip(arg1, arg2))
    });
    let mut enter = |d: &mut DiagnosticSink| two(d, &mut |d, arg1, arg2| {
        resolve_context(a, d, arg1).map(|ctx| CallKind::Enter(ctx, arg2))
    });
    let mut exit = |d: &mut DiagnosticSink| two(d, &mut |d, arg1, arg2| {
        resolve_context(a, d, arg1).map(|ctx| CallKind::Exit(ctx, arg2))
    });

    let mut prev_is = |d: &mut DiagnosticSink| {
        let mut args = Vec::new();
        for arg in call.args() {
            let mut ok = false;
            if let Expr::RefExpr(expr) = arg {
                if let Some(RefKind::RuleReference(rule)) = expr.resolve() {
                    if let Some(ty) = rule.resolve_ty() {
                        args.push(ty);
                        ok = true;
                    }
                }
            }
            if !ok {
                d.error(call.node(), "<prev_is> arguments must be public rules")
            }
        }
        Some(CallKind::PrevIs(Arc::new(args)))
    };

    let build_in: Vec<(&str, &mut FnMut(&mut DiagnosticSink) -> Option<CallKind<'f>>)> = vec![
        ("any", &mut any),
        ("commit", &mut commit),
        ("eof", &mut eof),
        ("not", &mut not),
        ("is_in", &mut is_in),
        ("layer", &mut layer),
        ("with_skip", &mut with_skip),
        ("enter", &mut enter),
        ("exit", &mut exit),
        ("prev_is", &mut prev_is)
    ];

    for (name, kind) in build_in.into_iter() {
        if call.fn_name() == name {
            return kind(d);
        }
    }

    if let Some(rule) = a.file.resolve_rule(call.fn_name()) {
        if let Some(parameters) = rule.parameters() {
            let n_expected = parameters.parameters().count();
            if n_expected != n_args {
                d.error(
                    call.node(),
                    format!("Expected {} arguments, got {}", n_expected, n_args)
                )
            }

            let params = parameters.parameters()
                .map(|p| p.idx())
                .zip(call.args())
                .collect();
            return Some(CallKind::RuleCall(rule, Arc::new(params)));
        }
    }

    d.error(
        child_of_type_exn(call.node(), IDENT),
        "Unresolved reference"
    );
    return None;
}

pub fn contexts<'f>(a: &Analysis<'f>) -> Vec<Text<'f>> {
    Visitor(Vec::new())
        .visit::<CallExpr, _>(|contexts, call| {
            if call.fn_name() == "is_in" || call.fn_name() == "enter" || call.fn_name() == "exit" {
                if let Some(ctx) = call.args().next().and_then(context_name) {
                    contexts.push(ctx);
                }
            }
        })
        .walk_recursively_children_first(a.file().node())
}

fn context_name<'f>(ctx: Expr<'f>) -> Option<Text<'f>> {
    child_of_type(ctx.node(), SIMPLE_STRING)
        .map(|n| ::ast_ext::lit_body(n.text()))
}

fn resolve_context(a: &Analysis, d: &mut DiagnosticSink, ctx: Expr) -> Option<u32> {
    if let Some(name) = context_name(ctx) {
        a.contexts().iter()
            .position(|&c| c == name)
            .map(|usize_| usize_ as u32)
    } else {
        d.error(ctx.node(), "Context should be a single quoted string");
        None
    }
}

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

    use fall_tree::search;
    use fall_tree::search::ast;

    use ::test_util::parse_with_caret;

    #[test]
    fn test_resolve_no_args() {
        check_resolved("rule foo { <^eof> }", "Eof");
        check_resolved("rule foo { <^any> }", "Any");
        check_resolved("rule foo { <^commit> }", "Commit");
    }

    #[test]
    fn test_resolve_simple_args() {
        check_resolved("rule foo { <^not a> }", "Not(RefExpr@[16; 17))");
        check_resolved("rule foo { <^with_skip a b> }", "WithSkip(RefExpr@[22; 23), RefExpr@[24; 25))");
        check_resolved("rule foo { <^layer a b> }", "Layer(RefExpr@[18; 19), RefExpr@[20; 21))");
    }

    #[test]
    fn test_resolve_layer() {
        check_resolved("rule foo { <^enter 'ctx' a> }", "Enter(0, RefExpr@[24; 25))");
        check_resolved("rule foo { <^exit  'ctx' a> }", "Exit(0, RefExpr@[24; 25))");

        check_resolved("rule foo { <is_in 'ctx1'> <^enter 'ctx2' a> }", "Enter(1, RefExpr@[40; 41))");
        check_resolved("rule foo { <is_in 'ctx1'> <^exit  'ctx2' a> }", "Exit(1, RefExpr@[40; 41))");

        check_resolved(
            "rule foo { <enter 'ctx1' a> <enter 'ctx2' b> <^is_in 'ctx2'> }",
            "IsIn(1)"
        );

        check(
            r" rule foo { <^is_in foo>}",
            None,
            r#"[(foo, "Context should be a single quoted string")]"#,
        )
    }

    #[test]
    fn test_not() {
        resolve_call("rule foo { <^not bar>}", |c, _| {
            let c = c.unwrap();
            assert_eq!(format!("{:?}", c), "Not(RefExpr@[16; 19))");
        })
    }

    #[test]
    fn test_unresolved_call() {
        check(
            r"rule foo { <^abracadabra>}",
            None,
            r#"[(abracadabra, "Unresolved reference")]"#
        );
    }

    #[test]
    fn wrong_arity() {
        check(
            "rule foo { <^eof foo> }",
            Some(CallKind::Eof),
            r#"[(<eof foo>, "Wrong number of arguments, expected 0, got 1")]"#
        );

        check(
            "rule foo { <^any foo bar> }",
            Some(CallKind::Any),
            r#"[(<any foo bar>, "Wrong number of arguments, expected 0, got 2")]"#
        );

        check(
            "rule foo { <^commit foo bar baz> }",
            Some(CallKind::Commit),
            r#"[(<commit foo bar baz>, "Wrong number of arguments, expected 0, got 3")]"#
        );

        check(
            "rule foo { <^not> }",
            None,
            r#"[(<not>, "Wrong number of arguments, expected 1, got 0")]"#
        );
    }

    #[test]
    fn check_prev_is() {
        check_resolved(
            "pub rule foo {} pub rule bar {} rule baz { <^prev_is foo bar> }",
            "PrevIs([1, 2])"
        )
    }

    #[test]
    fn call_custom_rule() {
        check_resolved(
            "rule foo { <^bar a b>} rule bar(x, y) { }",
            "RuleCall(SynRule@[22; 40), [(0, RefExpr@[16; 17)), (1, RefExpr@[18; 19))])"
        );
        ::analysis::check_diagnostics(
            "rule foo { <bar <eof>>} rule bar(x, y) { }",
            "<bar <eof>>: E Expected 2 arguments, got 1"
        )
    }

    fn check(text: &str, kind: Option<CallKind>, diagnostics: &str) {
        resolve_call(text, |c, a| {
            assert_eq!(c, kind);
            assert_eq!(a.debug_diagnostics(), diagnostics);
        })
    }

    fn check_resolved(text: &str, kind: &str) {
        resolve_call(text, |c, a| {
            let c = c.unwrap();
            assert_eq!(format!("{:?}", c), kind);
            assert_eq!(a.debug_diagnostics(), "[]");
        })
    }

    fn resolve_call<F: FnOnce(Option<CallKind>, Analysis)>(text: &str, f: F) {
        let (file, caret) = parse_with_caret(&text);
        let call = {
            let leaf = search::find_leaf_at_offset(file.root(), caret)
                .left_biased().unwrap();
            ast::ancestor_exn(leaf)
        };
        let a = Analysis::new(::ast(&file));
        let c = a.resolve_call(call);
        f(c, a)
    }
}

<commit_msg>This isn't gonna be fixed soon :)<commit_after>use std::sync::Arc;
use itertools::Itertools;

use fall_tree::{AstNode, AstClass, Text};
use fall_tree::search::{child_of_type, child_of_type_exn};
use fall_tree::visitor::{Visitor, NodeVisitor};

use super::{Analysis, DiagnosticSink};
use ::{Expr, CallExpr, SynRule, IDENT, SIMPLE_STRING, RefKind};

#[derive(Debug, Eq, PartialEq, Clone)]
pub enum CallKind<'f> {
    Eof,
    Any,
    Commit,

    Not(Expr<'f>),
    Layer(Expr<'f>, Expr<'f>),
    WithSkip(Expr<'f>, Expr<'f>),

    Enter(u32, Expr<'f>),
    Exit(u32, Expr<'f>),
    IsIn(u32),

    RuleCall(SynRule<'f>, Arc<Vec<(u32, Expr<'f>)>>),
    PrevIs(Arc<Vec<usize>>)
}

pub (super) fn resolve<'f>(a: &Analysis<'f>, d: &mut DiagnosticSink, call: CallExpr<'f>) -> Option<CallKind<'f>> {
    let n_args = call.args().count();
    let expect_args = |d: &mut DiagnosticSink, n_expected: usize| {
        if n_expected != n_args {
            d.error(
                call.node(),
                format!("Wrong number of arguments, expected {}, got {}", n_expected, n_args)
            )
        }
    };

    let zero = |d: &mut DiagnosticSink, kind: CallKind<'f>| {
        expect_args(d, 0);
        Some(kind)
    };
    let mut any = |d: &mut DiagnosticSink| zero(d, CallKind::Any);
    let mut commit = |d: &mut DiagnosticSink| zero(d, CallKind::Commit);
    let mut eof = |d: &mut DiagnosticSink| zero(d, CallKind::Eof);

    let one = |d: &mut DiagnosticSink, kind: &mut FnMut(&mut DiagnosticSink, Expr<'f>) -> Option<CallKind<'f>>| {
        expect_args(d, 1);
        call.args().next().and_then(|arg| kind(d, arg))
    };
    let mut not = |d: &mut DiagnosticSink| one(d, &mut |_, arg| Some(CallKind::Not(arg)));
    let mut is_in = |d: &mut DiagnosticSink| one(d, &mut |d, arg| resolve_context(a, d, arg).map(CallKind::IsIn));

    let two = |d: &mut DiagnosticSink, kind: &mut FnMut(&mut DiagnosticSink, Expr<'f>, Expr<'f>) -> Option<CallKind<'f>>| {
        expect_args(d, 2);
        call.args().next_tuple().and_then(|(arg1, arg2)| kind(d, arg1, arg2))
    };
    let mut layer = |d: &mut DiagnosticSink| two(d, &mut |_, arg1, arg2| {
        Some(CallKind::Layer(arg1, arg2))
    });
    let mut with_skip = |d: &mut DiagnosticSink| two(d, &mut |_, arg1, arg2| {
        Some(CallKind::WithSkip(arg1, arg2))
    });
    let mut enter = |d: &mut DiagnosticSink| two(d, &mut |d, arg1, arg2| {
        resolve_context(a, d, arg1).map(|ctx| CallKind::Enter(ctx, arg2))
    });
    let mut exit = |d: &mut DiagnosticSink| two(d, &mut |d, arg1, arg2| {
        resolve_context(a, d, arg1).map(|ctx| CallKind::Exit(ctx, arg2))
    });

    let mut prev_is = |d: &mut DiagnosticSink| {
        let mut args = Vec::new();
        for arg in call.args() {
            let mut ok = false;
            if let Expr::RefExpr(expr) = arg {
                if let Some(RefKind::RuleReference(rule)) = expr.resolve() {
                    if let Some(ty) = rule.resolve_ty() {
                        args.push(ty);
                        ok = true;
                    }
                }
            }
            if !ok {
                d.error(call.node(), "<prev_is> arguments must be public rules")
            }
        }
        Some(CallKind::PrevIs(Arc::new(args)))
    };

    let build_in: Vec<(&str, &mut FnMut(&mut DiagnosticSink) -> Option<CallKind<'f>>)> = vec![
        ("any", &mut any),
        ("commit", &mut commit),
        ("eof", &mut eof),
        ("not", &mut not),
        ("is_in", &mut is_in),
        ("layer", &mut layer),
        ("with_skip", &mut with_skip),
        ("enter", &mut enter),
        ("exit", &mut exit),
        ("prev_is", &mut prev_is)
    ];

    for (name, kind) in build_in.into_iter() {
        if call.fn_name() == name {
            return kind(d);
        }
    }

    if let Some(rule) = a.file.resolve_rule(call.fn_name()) {
        if let Some(parameters) = rule.parameters() {
            let n_expected = parameters.parameters().count();
            if n_expected != n_args {
                d.error(
                    call.node(),
                    format!("Expected {} arguments, got {}", n_expected, n_args)
                )
            }

            let params = parameters.parameters()
                .map(|p| p.idx())
                .zip(call.args())
                .collect();
            return Some(CallKind::RuleCall(rule, Arc::new(params)));
        }
    }

    d.error(
        child_of_type_exn(call.node(), IDENT),
        "Unresolved reference"
    );
    return None;
}

pub fn contexts<'f>(a: &Analysis<'f>) -> Vec<Text<'f>> {
    Visitor(Vec::new())
        .visit::<CallExpr, _>(|contexts, call| {
            if call.fn_name() == "is_in" || call.fn_name() == "enter" || call.fn_name() == "exit" {
                if let Some(ctx) = call.args().next().and_then(context_name) {
                    contexts.push(ctx);
                }
            }
        })
        .walk_recursively_children_first(a.file().node())
}

fn context_name<'f>(ctx: Expr<'f>) -> Option<Text<'f>> {
    child_of_type(ctx.node(), SIMPLE_STRING)
        .map(|n| ::ast_ext::lit_body(n.text()))
}

fn resolve_context(a: &Analysis, d: &mut DiagnosticSink, ctx: Expr) -> Option<u32> {
    if let Some(name) = context_name(ctx) {
        a.contexts().iter()
            .position(|&c| c == name)
            .map(|usize_| usize_ as u32)
    } else {
        d.error(ctx.node(), "Context should be a single quoted string");
        None
    }
}

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

    use fall_tree::search;
    use fall_tree::search::ast;

    use ::test_util::parse_with_caret;

    #[test]
    fn test_resolve_no_args() {
        check_resolved("rule foo { <^eof> }", "Eof");
        check_resolved("rule foo { <^any> }", "Any");
        check_resolved("rule foo { <^commit> }", "Commit");
    }

    #[test]
    fn test_resolve_simple_args() {
        check_resolved("rule foo { <^not a> }", "Not(RefExpr@[16; 17))");
        check_resolved("rule foo { <^with_skip a b> }", "WithSkip(RefExpr@[22; 23), RefExpr@[24; 25))");
        check_resolved("rule foo { <^layer a b> }", "Layer(RefExpr@[18; 19), RefExpr@[20; 21))");
    }

    #[test]
    fn test_resolve_layer() {
        check_resolved("rule foo { <^enter 'ctx' a> }", "Enter(0, RefExpr@[24; 25))");
        check_resolved("rule foo { <^exit  'ctx' a> }", "Exit(0, RefExpr@[24; 25))");

        check_resolved("rule foo { <is_in 'ctx1'> <^enter 'ctx2' a> }", "Enter(1, RefExpr@[40; 41))");
        check_resolved("rule foo { <is_in 'ctx1'> <^exit  'ctx2' a> }", "Exit(1, RefExpr@[40; 41))");

        check_resolved(
            "rule foo { <enter 'ctx1' a> <enter 'ctx2' b> <^is_in 'ctx2'> }",
            "IsIn(1)"
        );

        check(
            r" rule foo { <^is_in foo>}",
            None,
            r#"[(foo, "Context should be a single quoted string")]"#,
        )
    }

    #[test]
    fn test_not() {
        resolve_call("rule foo { <^not bar>}", |c, _| {
            let c = c.unwrap();
            assert_eq!(format!("{:?}", c), "Not(RefExpr@[16; 19))");
        })
    }

    #[test]
    fn test_unresolved_call() {
        check(
            r"rule foo { <^abracadabra>}",
            None,
            r#"[(abracadabra, "Unresolved reference")]"#
        );
    }

    #[test]
    fn wrong_arity() {
        check(
            "rule foo { <^eof foo> }",
            Some(CallKind::Eof),
            r#"[(<eof foo>, "Wrong number of arguments, expected 0, got 1")]"#
        );

        check(
            "rule foo { <^any foo bar> }",
            Some(CallKind::Any),
            r#"[(<any foo bar>, "Wrong number of arguments, expected 0, got 2")]"#
        );

        check(
            "rule foo { <^commit foo bar baz> }",
            Some(CallKind::Commit),
            r#"[(<commit foo bar baz>, "Wrong number of arguments, expected 0, got 3")]"#
        );

        check(
            "rule foo { <^not> }",
            None,
            r#"[(<not>, "Wrong number of arguments, expected 1, got 0")]"#
        );
    }

    #[test]
    fn check_prev_is() {
        check_resolved(
            "pub rule foo {} pub rule bar {} rule baz { <^prev_is foo bar> }",
            "PrevIs([1, 2])"
        )
    }

    #[test]
    fn call_custom_rule() {
        check_resolved(
            "rule foo { <^bar a b>} rule bar(x, y) { }",
            "RuleCall(SynRule@[22; 40), [(0, RefExpr@[16; 17)), (1, RefExpr@[18; 19))])"
        );
        ::analysis::check_diagnostics(
            "rule foo { <bar <eof>>} rule bar(x, y) { }",
            "<bar <eof>>: E Expected 2 arguments, got 1"
        )
    }

    fn check(text: &str, kind: Option<CallKind>, diagnostics: &str) {
        resolve_call(text, |c, a| {
            assert_eq!(c, kind);
            assert_eq!(a.debug_diagnostics(), diagnostics);
        })
    }

    fn check_resolved(text: &str, kind: &str) {
        resolve_call(text, |c, a| {
            let c = c.unwrap();
            assert_eq!(format!("{:?}", c), kind);
            assert_eq!(a.debug_diagnostics(), "[]");
        })
    }

    fn resolve_call<F: FnOnce(Option<CallKind>, Analysis)>(text: &str, f: F) {
        let (file, caret) = parse_with_caret(&text);
        let call = {
            let leaf = search::find_leaf_at_offset(file.root(), caret)
                .left_biased().unwrap();
            ast::ancestor_exn(leaf)
        };
        let a = Analysis::new(::ast(&file));
        let c = a.resolve_call(call);
        f(c, a)
    }
}

<|endoftext|>


<commit_before>#![feature(time2)]
#[macro_use] extern crate log;
extern crate rand;
extern crate byteorder;
extern crate mqtt3;
extern crate netopt;

mod error;
mod sub;
mod client;

pub use sub::{
    ToSubTopics,
    ToUnSubTopics
};

pub use client::{
    Client,
    ClientOptions
};

use std::sync::Arc;
use std::ops;
use std::time::Duration;
use mqtt3::{QoS, ToTopicPath};
use error::Result;

const MAX_QOS: QoS = mqtt3::QoS::AtLeastOnce;

pub trait Mqttc {
    fn publish<T: ToTopicPath, P: ToPayload>(&mut self, topic: T, payload: P, pubopt: PubOpt) -> Result<()>;
    fn subscribe<S: ToSubTopics>(&mut self, subs: S) -> Result<()>;
    fn unsubscribe<U: ToUnSubTopics>(&mut self, unsubs: U) -> Result<()>;
    fn disconnect(self) -> Result<()>;
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ClientState {
    Handshake,
    Connected,
    Disconnected
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReconnectMethod {
    ForeverDisconnect,
    ReconnectAfter(Duration)
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PubOpt(u8);

impl PubOpt {
    #[inline]
    pub fn at_most_once() -> PubOpt {
        PubOpt(0x00)
    }

    #[inline]
    pub fn at_least_once() -> PubOpt {
        PubOpt(0x01)
    }

    #[inline]
    pub fn exactly_once() -> PubOpt {
        PubOpt(0x02)
    }

    #[inline]
    pub fn retain() -> PubOpt {
        PubOpt(0x04)
    }

    #[inline]
    pub fn bits(&self) -> u8 {
        self.0
    }

    pub fn qos(&self) -> QoS {
        if (self.0 & PubOpt::exactly_once().bits()) != 0 {
            return QoS::ExactlyOnce;
        }
        if (self.0 & PubOpt::at_least_once().bits()) != 0 {
            return QoS::AtLeastOnce;
        }
        QoS::AtMostOnce
    }

    pub fn is_retain(&self) -> bool {
        (self.0 & PubOpt::retain().bits()) != 0
    }
}

impl ops::BitOr for PubOpt {
    type Output = PubOpt;

    #[inline]
    fn bitor(self, other: PubOpt) -> PubOpt {
        PubOpt(self.bits() | other.bits())
    }
}

impl ops::BitXor for PubOpt {
    type Output = PubOpt;

    #[inline]
    fn bitxor(self, other: PubOpt) -> PubOpt {
        PubOpt(self.bits() ^ other.bits())
    }
}

impl ops::BitAnd for PubOpt {
    type Output = PubOpt;

    #[inline]
    fn bitand(self, other: PubOpt) -> PubOpt {
        PubOpt(self.bits() & other.bits())
    }
}

impl ops::Sub for PubOpt {
    type Output = PubOpt;

    #[inline]
    fn sub(self, other: PubOpt) -> PubOpt {
        PubOpt(self.bits() & !other.bits())
    }
}

impl ops::Not for PubOpt {
    type Output = PubOpt;

    #[inline]
    fn not(self) -> PubOpt {
        PubOpt(!self.bits() & 0b111)
    }
}

pub type Payload = Arc<Vec<u8>>;

pub trait ToPayload {
    fn to_payload(&self) -> Payload;
}

impl ToPayload for String {
    fn to_payload(&self) -> Payload {
        Arc::new(self.clone().into_bytes())
    }
}

impl<'a> ToPayload for &'a str {
    fn to_payload(&self) -> Payload {
        Arc::new(self.as_bytes().to_vec())
    }
}

impl ToPayload for Vec<u8> {
    fn to_payload(&self) -> Payload {
        Arc::new(self.clone())
    }
}

impl ToPayload for Arc<Vec<u8>> {
    fn to_payload(&self) -> Payload {
        self.clone()
    }
}

#[cfg(test)]
mod test {
    use super::PubOpt;
    use mqtt3::QoS;

    #[test]
    fn pubopt_test() {
        let pubopt = PubOpt::at_least_once();
        assert_eq!(pubopt.qos(), QoS::AtLeastOnce);
        assert!(!pubopt.is_retain());

        let pubopt = PubOpt::at_most_once();
        assert_eq!(pubopt.qos(), QoS::AtMostOnce);
        assert!(!pubopt.is_retain());

        let pubopt = PubOpt::exactly_once();
        assert_eq!(pubopt.qos(), QoS::ExactlyOnce);
        assert!(!pubopt.is_retain());

        let pubopt = PubOpt::at_least_once() | PubOpt::retain();
        assert_eq!(pubopt.qos(), QoS::AtLeastOnce);
        assert!(pubopt.is_retain());

        let pubopt = PubOpt::at_most_once() | PubOpt::retain();
        assert_eq!(pubopt.qos(), QoS::AtMostOnce);
        assert!(pubopt.is_retain());

        let pubopt = PubOpt::exactly_once() | PubOpt::retain();
        assert_eq!(pubopt.qos(), QoS::ExactlyOnce);
        assert!(pubopt.is_retain());
    }
}
<commit_msg>mqttc: Add PubOpt::new<commit_after>#![feature(time2)]
#[macro_use] extern crate log;
extern crate rand;
extern crate byteorder;
extern crate mqtt3;
extern crate netopt;

mod error;
mod sub;
mod client;

pub use sub::{
    ToSubTopics,
    ToUnSubTopics
};

pub use client::{
    Client,
    ClientOptions
};

use std::sync::Arc;
use std::ops;
use std::time::Duration;
use mqtt3::{QoS, ToTopicPath};
use error::Result;

const MAX_QOS: QoS = mqtt3::QoS::AtLeastOnce;

pub trait Mqttc {
    fn publish<T: ToTopicPath, P: ToPayload>(&mut self, topic: T, payload: P, pubopt: PubOpt) -> Result<()>;
    fn subscribe<S: ToSubTopics>(&mut self, subs: S) -> Result<()>;
    fn unsubscribe<U: ToUnSubTopics>(&mut self, unsubs: U) -> Result<()>;
    fn disconnect(self) -> Result<()>;
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ClientState {
    Handshake,
    Connected,
    Disconnected
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReconnectMethod {
    ForeverDisconnect,
    ReconnectAfter(Duration)
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PubOpt(u8);

impl PubOpt {
    pub fn new(qos: QoS, retain: bool) -> PubOpt {
        let mut opt = PubOpt(qos.to_u8());
        if retain {
            opt = opt & PubOpt::retain();
        }
        opt
    }

    #[inline]
    pub fn at_most_once() -> PubOpt {
        PubOpt(0x00)
    }

    #[inline]
    pub fn at_least_once() -> PubOpt {
        PubOpt(0x01)
    }

    #[inline]
    pub fn exactly_once() -> PubOpt {
        PubOpt(0x02)
    }

    #[inline]
    pub fn retain() -> PubOpt {
        PubOpt(0x04)
    }

    #[inline]
    pub fn bits(&self) -> u8 {
        self.0
    }

    pub fn qos(&self) -> QoS {
        if (self.0 & PubOpt::exactly_once().bits()) != 0 {
            return QoS::ExactlyOnce;
        }
        if (self.0 & PubOpt::at_least_once().bits()) != 0 {
            return QoS::AtLeastOnce;
        }
        QoS::AtMostOnce
    }

    pub fn is_retain(&self) -> bool {
        (self.0 & PubOpt::retain().bits()) != 0
    }
}

impl ops::BitOr for PubOpt {
    type Output = PubOpt;

    #[inline]
    fn bitor(self, other: PubOpt) -> PubOpt {
        PubOpt(self.bits() | other.bits())
    }
}

impl ops::BitXor for PubOpt {
    type Output = PubOpt;

    #[inline]
    fn bitxor(self, other: PubOpt) -> PubOpt {
        PubOpt(self.bits() ^ other.bits())
    }
}

impl ops::BitAnd for PubOpt {
    type Output = PubOpt;

    #[inline]
    fn bitand(self, other: PubOpt) -> PubOpt {
        PubOpt(self.bits() & other.bits())
    }
}

impl ops::Sub for PubOpt {
    type Output = PubOpt;

    #[inline]
    fn sub(self, other: PubOpt) -> PubOpt {
        PubOpt(self.bits() & !other.bits())
    }
}

impl ops::Not for PubOpt {
    type Output = PubOpt;

    #[inline]
    fn not(self) -> PubOpt {
        PubOpt(!self.bits() & 0b111)
    }
}

pub type Payload = Arc<Vec<u8>>;

pub trait ToPayload {
    fn to_payload(&self) -> Payload;
}

impl ToPayload for String {
    fn to_payload(&self) -> Payload {
        Arc::new(self.clone().into_bytes())
    }
}

impl<'a> ToPayload for &'a str {
    fn to_payload(&self) -> Payload {
        Arc::new(self.as_bytes().to_vec())
    }
}

impl ToPayload for Vec<u8> {
    fn to_payload(&self) -> Payload {
        Arc::new(self.clone())
    }
}

impl ToPayload for Arc<Vec<u8>> {
    fn to_payload(&self) -> Payload {
        self.clone()
    }
}

#[cfg(test)]
mod test {
    use super::PubOpt;
    use mqtt3::QoS;

    #[test]
    fn pubopt_test() {
        let pubopt = PubOpt::at_least_once();
        assert_eq!(pubopt.qos(), QoS::AtLeastOnce);
        assert!(!pubopt.is_retain());

        let pubopt = PubOpt::at_most_once();
        assert_eq!(pubopt.qos(), QoS::AtMostOnce);
        assert!(!pubopt.is_retain());

        let pubopt = PubOpt::exactly_once();
        assert_eq!(pubopt.qos(), QoS::ExactlyOnce);
        assert!(!pubopt.is_retain());

        let pubopt = PubOpt::at_least_once() | PubOpt::retain();
        assert_eq!(pubopt.qos(), QoS::AtLeastOnce);
        assert!(pubopt.is_retain());

        let pubopt = PubOpt::at_most_once() | PubOpt::retain();
        assert_eq!(pubopt.qos(), QoS::AtMostOnce);
        assert!(pubopt.is_retain());

        let pubopt = PubOpt::exactly_once() | PubOpt::retain();
        assert_eq!(pubopt.qos(), QoS::ExactlyOnce);
        assert!(pubopt.is_retain());
    }
}
<|endoftext|>