<commit_before>use std::cmp;
use std::iter::IntoIterator;

#[doc(hidden)]
#[derive(Clone, Debug)]
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct ZipOption<A, B> {
  a: A,
  b: B,
  // index and len are only used by the specialized version of zip
  index: usize,
  len: usize,
}

/// Zips to iterators together to the longest length
/// via Option<(Option<A>, Option<B>)>
pub trait ZipOpt {
  /// Zip to iterators to longest length via Option<(Option<A>, Option<B>)> results.
  /// # Example
  /// ```
  /// use array_tool::iter::ZipOpt;
  ///
  /// let a = vec!["a","b","c", "d"];
  /// let b = vec!["c","d"];
  /// let mut x = a.iter().zip_option(b.iter());
  ///
  /// assert_eq!(x.next(), Some((Some(&"a"), Some(&"c"))));
  /// assert_eq!(x.next(), Some((Some(&"b"), Some(&"d"))));
  /// assert_eq!(x.next(), Some((Some(&"c"), None)));
  /// assert_eq!(x.next(), Some((Some(&"d"), None)));
  /// assert_eq!(x.next(), None);
  /// ```
  ///
  /// # Output
  /// ```text
  /// vec![ "a", "b", "c", "d" ]
  /// ```
  fn zip_option<U: Iterator>(self, other: U) -> ZipOption<Self, U>
    where Self: Sized, U: IntoIterator;
}

impl<I: Iterator> ZipOpt for I {
  fn zip_option<U: Iterator>(self, other: U) -> ZipOption<Self, U>
    where Self: Sized, U: IntoIterator {

    ZipOption::new(self, other)
  }
}

impl<A, B> Iterator for ZipOption<A, B> where A: Iterator, B: Iterator {
  type Item = (Option<A::Item>, Option<B::Item>);

  #[inline]
  fn next(&mut self) -> Option<Self::Item> {
    ZipImpl::next(self)
  }

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

  #[inline]
  fn nth(&mut self, n: usize) -> Option<Self::Item> {
    ZipImpl::nth(self, n)
  }
}

#[doc(hidden)]
impl<A, B> DoubleEndedIterator for ZipOption<A, B> where
A: DoubleEndedIterator + ExactSizeIterator,
B: DoubleEndedIterator + ExactSizeIterator,
{
  #[inline]
  fn next_back(&mut self) -> Option<(Option<A::Item>, Option<B::Item>)> {
    ZipImpl::next_back(self)
  }
}

#[doc(hidden)]
trait ZipImpl<A, B> {
  type Item;
  fn new(a: A, b: B) -> Self;
  fn next(&mut self) -> Option<Self::Item>;
  fn size_hint(&self) -> (usize, Option<usize>);
  fn nth(&mut self, n: usize) -> Option<Self::Item>;
  fn super_nth(&mut self, mut n: usize) -> Option<Self::Item> {
    while let Some(x) = self.next() {
      if n == 0 { return Some(x) }
      n -= 1;
    }
    None
  }
  fn next_back(&mut self) -> Option<Self::Item>
    where A: DoubleEndedIterator + ExactSizeIterator,
          B: DoubleEndedIterator + ExactSizeIterator;
}

#[doc(hidden)]
impl<A, B> ZipImpl<A, B> for ZipOption<A, B>
where A: Iterator, B: Iterator {
  type Item = (Option<A::Item>, Option<B::Item>);
  fn new(a: A, b: B) -> Self {
    ZipOption {
      a,
      b,
      index: 0, // unused
      len: 0, // unused
    }
  }

  #[inline]
  fn next(&mut self) -> Option<(Option<A::Item>, Option<B::Item>)> {
    let first = self.a.next();
    let second = self.b.next();

    if first.is_some() || second.is_some() {
      Some((first, second))
    } else {
      None
    }
  }

  #[inline]
  fn nth(&mut self, n: usize) -> Option<Self::Item> {
    self.super_nth(n)
  }

  #[inline]
  fn next_back(&mut self) -> Option<(Option<A::Item>, Option<B::Item>)>
    where A: DoubleEndedIterator + ExactSizeIterator,
          B: DoubleEndedIterator + ExactSizeIterator {
    let a_sz = self.a.len();
    let b_sz = self.b.len();
    if a_sz != b_sz {
      // Adjust a, b to equal length
      if a_sz > b_sz {
        for _ in 0..a_sz - b_sz { self.a.next_back(); }
      } else {
        for _ in 0..b_sz - a_sz { self.b.next_back(); }
      }
    }
    match (self.a.next_back(), self.b.next_back()) {
      (None, None) => None,
      (f,s) => Some((f, s)),
    }
  }

  #[inline]
  fn size_hint(&self) -> (usize, Option<usize>) {
    let (a_lower, a_upper) = self.a.size_hint();
    let (b_lower, b_upper) = self.b.size_hint();

    let lower = cmp::min(a_lower, b_lower);

    let upper = match (a_upper, b_upper) {
      (Some(x), Some(y)) => Some(cmp::max(x,y)),
      (Some(x), None) => Some(x),
      (None, Some(y)) => Some(y),
      (None, None) => None
    };

    (lower, upper)
  }
}
<commit_msg>white space<commit_after>use std::cmp;
use std::iter::IntoIterator;

#[doc(hidden)]
#[derive(Clone, Debug)]
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct ZipOption<A, B> {
  a: A,
  b: B,
  // index and len are only used by the specialized version of zip
  index: usize,
  len: usize,
}

/// Zips to iterators together to the longest length
/// via Option<(Option<A>, Option<B>)>
pub trait ZipOpt {
  /// Zip to iterators to longest length via Option<(Option<A>, Option<B>)> results.
  /// # Example
  /// ```
  /// use array_tool::iter::ZipOpt;
  ///
  /// let a = vec!["a","b","c", "d"];
  /// let b = vec!["c","d"];
  /// let mut x = a.iter().zip_option(b.iter());
  ///
  /// assert_eq!(x.next(), Some((Some(&"a"), Some(&"c"))));
  /// assert_eq!(x.next(), Some((Some(&"b"), Some(&"d"))));
  /// assert_eq!(x.next(), Some((Some(&"c"), None)));
  /// assert_eq!(x.next(), Some((Some(&"d"), None)));
  /// assert_eq!(x.next(), None);
  /// ```
  ///
  /// # Output
  /// ```text
  /// vec![ "a", "b", "c", "d" ]
  /// ```
  fn zip_option<U: Iterator>(self, other: U) -> ZipOption<Self, U>
    where Self: Sized, U: IntoIterator;
}

impl<I: Iterator> ZipOpt for I {
  fn zip_option<U: Iterator>(self, other: U) -> ZipOption<Self, U>
    where Self: Sized, U: IntoIterator {

    ZipOption::new(self, other)
  }
}

impl<A, B> Iterator for ZipOption<A, B> where A: Iterator, B: Iterator {
  type Item = (Option<A::Item>, Option<B::Item>);

  #[inline]
  fn next(&mut self) -> Option<Self::Item> {
    ZipImpl::next(self)
  }

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

  #[inline]
  fn nth(&mut self, n: usize) -> Option<Self::Item> {
    ZipImpl::nth(self, n)
  }
}

#[doc(hidden)]
impl<A, B> DoubleEndedIterator for ZipOption<A, B> where
A: DoubleEndedIterator + ExactSizeIterator,
B: DoubleEndedIterator + ExactSizeIterator,
{
  #[inline]
  fn next_back(&mut self) -> Option<(Option<A::Item>, Option<B::Item>)> {
    ZipImpl::next_back(self)
  }
}

#[doc(hidden)]
trait ZipImpl<A, B> {
  type Item;
  fn new(a: A, b: B) -> Self;
  fn next(&mut self) -> Option<Self::Item>;
  fn size_hint(&self) -> (usize, Option<usize>);
  fn nth(&mut self, n: usize) -> Option<Self::Item>;
  fn super_nth(&mut self, mut n: usize) -> Option<Self::Item> {
    while let Some(x) = self.next() {
      if n == 0 { return Some(x) }
      n -= 1;
    }
    None
  }
  fn next_back(&mut self) -> Option<Self::Item>
    where A: DoubleEndedIterator + ExactSizeIterator,
          B: DoubleEndedIterator + ExactSizeIterator;
}

#[doc(hidden)]
impl<A, B> ZipImpl<A, B> for ZipOption<A, B>
  where A: Iterator, B: Iterator {
  type Item = (Option<A::Item>, Option<B::Item>);
  fn new(a: A, b: B) -> Self {
    ZipOption {
      a,
      b,
      index: 0, // unused
      len: 0, // unused
    }
  }

  #[inline]
  fn next(&mut self) -> Option<(Option<A::Item>, Option<B::Item>)> {
    let first = self.a.next();
    let second = self.b.next();

    if first.is_some() || second.is_some() {
      Some((first, second))
    } else {
      None
    }
  }

  #[inline]
  fn nth(&mut self, n: usize) -> Option<Self::Item> {
    self.super_nth(n)
  }

  #[inline]
  fn next_back(&mut self) -> Option<(Option<A::Item>, Option<B::Item>)>
    where A: DoubleEndedIterator + ExactSizeIterator,
          B: DoubleEndedIterator + ExactSizeIterator {
    let a_sz = self.a.len();
    let b_sz = self.b.len();
    if a_sz != b_sz {
      // Adjust a, b to equal length
      if a_sz > b_sz {
        for _ in 0..a_sz - b_sz { self.a.next_back(); }
      } else {
        for _ in 0..b_sz - a_sz { self.b.next_back(); }
      }
    }
    match (self.a.next_back(), self.b.next_back()) {
      (None, None) => None,
      (f,s) => Some((f, s)),
    }
  }

  #[inline]
  fn size_hint(&self) -> (usize, Option<usize>) {
    let (a_lower, a_upper) = self.a.size_hint();
    let (b_lower, b_upper) = self.b.size_hint();

    let lower = cmp::min(a_lower, b_lower);

    let upper = match (a_upper, b_upper) {
      (Some(x), Some(y)) => Some(cmp::max(x,y)),
      (Some(x), None) => Some(x),
      (None, Some(y)) => Some(y),
      (None, None) => None
    };

    (lower, upper)
  }
}
<|endoftext|>


<commit_before>//! Contains information for keybindings.

use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::hash::{Hash, Hasher};

use rustwlc::xkb::{Keysym, NameFlags};
use rustwlc::types::*; // Need * for bitflags...

use super::commands;
use registry::get_command;
use super::layout::tree;
use super::layout::container::{Layout, ContainerType};
use super::lua;

/// Register default keypresses to a map
macro_rules! register_defaults {
    ( $map:expr; $($func:expr, $press:expr);+ ) => {
        $(
            let _ = $map.insert($press, Arc::new($func));
            )+
    }
}

/// Generate switch_workspace methods and register them in $map
macro_rules! gen_switch_workspace {
    ($map:expr; $($b:ident, $n:expr);+) => {
        $(fn $b() {
            trace!("Switching to workspace {}", $n);
            if let Ok(mut tree)  = tree::try_lock_tree() {
                tree.switch_to_workspace(&$n);
            }
        }
          register_defaults!( $map; $b, keypress!("Alt", $n) );
          )+
    };
}

/// Generate move_container_to methods and register them in $map
macro_rules! gen_move_to_workspace {
    ($map:expr; $($b:ident, $n:expr);+) => {
        $(fn $b() {
            trace!("Moving active container to {}", $n);
            if let Ok(mut tree) = tree::try_lock_tree() {
                tree.send_active_to_workspace(&$n);
            }
        }
          register_defaults!( $map; $b, KeyPress::from_key_names(vec!["Alt", "Shift"],
                                                                vec![$n]).unwrap());
        )+
    };
}

lazy_static! {
    static ref BINDINGS: RwLock<HashMap<KeyPress, KeyEvent>> = {
        let mut map = HashMap::<KeyPress, KeyEvent>::new();

        gen_move_to_workspace!(map;
                               move_to_workspace_1, "1";
                               move_to_workspace_2, "2";
                               move_to_workspace_3, "3";
                               move_to_workspace_4, "4";
                               move_to_workspace_5, "5";
                               move_to_workspace_6, "6";
                               move_to_workspace_7, "7";
                               move_to_workspace_8, "8";
                               move_to_workspace_9, "9";
                               move_to_workspace_0, "0");

        gen_switch_workspace!(map; switch_workspace_1, "1";
                              switch_workspace_2, "2";
                              switch_workspace_3, "3";
                              switch_workspace_4, "4";
                              switch_workspace_5, "5";
                              switch_workspace_6, "6";
                              switch_workspace_7, "7";
                              switch_workspace_8, "8";
                              switch_workspace_9, "9";
                              switch_workspace_0, "0");

        register_defaults! {
            map;
            quit_fn, keypress!("Alt", "Escape");
            terminal_fn, keypress!("Alt", "Return");
            dmenu_fn, keypress!("Alt", "d");
            pointer_fn, keypress!("Alt", "p");
            horizontal_vertical_switch_fn, keypress!("Alt", "e");
            split_vertical, keypress!("Alt", "v");
            split_horizontal, keypress!("Alt", "h");
            move_left, keypress!("Alt", "left");
            move_right, keypress!("Alt", "right");
            move_up, keypress!("Alt", "up");
            move_down, keypress!("Alt", "down")
        }

        RwLock::new(map)
    };
}

fn move_up() {
    use layout::tree::Direction;
    if let Ok(mut tree) = tree::try_lock_tree() {
        tree.move_focus(Direction::Up)
    }
}

fn move_down() {
    use layout::tree::Direction;
    if let Ok(mut tree) = tree::try_lock_tree() {
        tree.move_focus(Direction::Down)
    }
}

fn move_left() {
    use layout::tree::Direction;
    if let Ok(mut tree) = tree::try_lock_tree() {
        tree.move_focus(Direction::Left)
    }
}

fn move_right() {
    use layout::tree::Direction;
    if let Ok(mut tree) = tree::try_lock_tree() {
        tree.move_focus(Direction::Right)
    }
}

fn split_vertical() {
    trace!("Splitting vertically");
    if let Ok(mut tree) = tree::try_lock_tree() {
        tree.toggle_active_layout(Layout::Vertical);
    }
}

fn split_horizontal() {
    trace!("Splitting horizontally");
    if let Ok(mut tree) = tree::try_lock_tree() {
        tree.toggle_active_layout(Layout::Horizontal);
    }
}

fn horizontal_vertical_switch_fn() {
    if let Ok(mut tree) = tree::try_lock_tree() {
        tree.toggle_active_horizontal();
        tree.layout_active_of(ContainerType::Root);
    } else {
        error!("Could not grab tree")
    }
}

fn terminal_fn() {
    use std::process::Command;
    use std::env;

    let term = env::var("WAYLAND_TERMINAL")
        .unwrap_or("weston-terminal".to_string());

    Command::new("sh")
        .arg("-c")
        .arg(term)
        .spawn().expect("Error launching terminal");
}

fn dmenu_fn() {
    use std::process::Command;
    Command::new("sh")
        .arg("-c")
        .arg("dmenu_run")
        .spawn().expect("Error launching terminal");
}

fn pointer_fn() {
    use lua::LuaQuery;
    let code = "if wm == nil then print('wm table does not exist')\n\
                elseif wm.pointer == nil then print('wm.pointer table does not exist')\n\
                else\n\
                print('get_position is a func, preparing execution')
                local x, y = wm.pointer.get_position()\n\
                print('The cursor is at ' .. x .. ', ' .. y)\n\
                end".to_string();
    lua::send(LuaQuery::Execute(code))
        .expect("Error telling Lua to get pointer coords");
}

fn quit_fn() {
    info!("handler: Esc keypress!");
    ::rustwlc::terminate();
}

#[derive(Eq, PartialEq, Clone, Debug)]
pub struct KeyPress {
    modifiers: KeyMod,
    keys: Vec<Keysym>
}

pub fn init() {
    macro_rules! insert_all {
        ( $($press:expr => $name:expr);+ ) => {
            register(vec![
                $( ($press, get_command(&$name.to_string())
                  .expect("Unable to register default command!")) ),+
                ]);
        }
    }

    insert_all! {
        keypress!("Alt", "Escape") => "quit";
        keypress!("Alt", "Return") => "launch_terminal";
        keypress!("Alt", "d") => "launch_dmenu";
        //keypress!("Alt", "l") => "dmenu_eval";
        /*KeyPress::from_key_names(vec!["Alt", "Shift"], vec!["l"])
            .expect("Unable to create default keypress")
            => "dmenu_lua_dofile";*/

        keypress!("Alt", "1") => "switch_workspace_1";
        keypress!("Alt", "2") => "switch_workspace_2";
        keypress!("Alt", "3") => "switch_workspace_3";
        keypress!("Alt", "4") => "switch_workspace_4";
        keypress!("Alt", "5") => "switch_workspace_5";
        keypress!("Alt", "6") => "switch_workspace_6";
        keypress!("Alt", "7") => "switch_workspace_7";
        keypress!("Alt", "8") => "switch_workspace_8";
        keypress!("Alt", "9") => "switch_workspace_9";
        keypress!("Alt", "0") => "switch_workspace_0"
    }
}

/// Parses a KeyMod from key names.
pub fn keymod_from_names(keys: Vec<&str>) -> Result<KeyMod, String> {
    let mut result = KeyMod::empty();
    for key in keys {
        match key.to_lowercase().as_str() {
            "shift"            => result = result | MOD_SHIFT,
            "control" | "ctrl" => result = result | MOD_CTRL,
            "alt"              => result = result | MOD_ALT,
            "mod2"             => result = result | MOD_MOD2,
            "mod3"             => result = result | MOD_MOD3,
            "mod4" | "super" | "logo" => result = result | MOD_MOD4,
            "mod5" | "5mod5me" => result = result | MOD_MOD5,
            err => return Err(format!("Invalid modifier: {}", err))
        }
    }
    return Ok(result);
}

impl KeyPress {
    /// Creates a new KeyPress struct from a list of modifier and key names
    pub fn from_key_names(mods: Vec<&str>, keys: Vec<&str>) -> Result<KeyPress, String> {
        keymod_from_names(mods).and_then(|mods| {
            let mut syms: Vec<Keysym> = Vec::with_capacity(keys.len());
            for name in keys {
                // Parse a keysym for each given key
                if let Some(sym) = Keysym::from_name(name.to_string(),
                                                     NameFlags::None) {
                    syms.push(sym);
                }
                // If lowercase cannot be parsed, try case insensitive
                else if let Some(sym) = Keysym::from_name(name.to_string(),
                                                          NameFlags::CaseInsensitive) {
                    syms.push(sym);
                }
                else {
                    return Err(format!("Invalid key: {}", name));
                }
            }
            // Sort and dedup to make sure hashes are the same
            syms.sort_by_key(|s| s.get_code());
            syms.dedup();
            return Ok(KeyPress { modifiers: mods, keys: syms });
        })
    }

    /// Creates a KeyPress from keys that are pressed at the moment
    pub fn new(mods: KeyMod, mut keys: Vec<Keysym>) -> KeyPress {
        // Sort and dedup to make sure hashes are the same
        keys.sort_by_key(|k| k.get_code());
        keys.dedup();

        KeyPress { modifiers: mods, keys: keys }
    }
}

impl Hash for KeyPress {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        hasher.write_u32(self.modifiers.bits());
        for ref sym in &self.keys {
            hasher.write_u32(sym.get_code());
        }
    }
}

/// The type of function a key press handler is.
pub type KeyEvent = Arc<Fn() + Send + Sync>;

/// Get a key mapping from the list.
pub fn get(key: &KeyPress) -> Option<KeyEvent> {
    let bindings = BINDINGS.read()
        .expect("Keybindings/get: unable to lock keybindings");
    match bindings.get(key) {
        None => None,
        Some(val) => Some(val.clone())
    }
}

/// Register a new set of key mappings
#[allow(dead_code)]
pub fn register(values: Vec<(KeyPress, KeyEvent)>) {
    let mut bindings = BINDINGS.write()
        .expect("Keybindings/register: unable to lock keybindings");
    for value in values {
        bindings.insert(value.0, value.1);
    }
}

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

    fn test_cmd() {
        assert!(true);
    }

    fn keypress() -> KeyPress {
        keypress!("Ctrl", "t")
    }

    #[test]
    fn add_key() {
        require_rustwlc!();
        register(vec![(keypress(), Arc::new(test_cmd))]);
        assert!(get(&keypress()).is_some(), "Key not registered");
    }
}
<commit_msg>keys.rs fix<commit_after>//! Contains information for keybindings.

use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::hash::{Hash, Hasher};

use rustwlc::xkb::{Keysym, NameFlags};
use rustwlc::types::*; // Need * for bitflags...

use super::commands;

lazy_static! {
    static ref BINDINGS: RwLock<HashMap<KeyPress, KeyEvent>> =
        RwLock::new(HashMap::new());
}

#[derive(Eq, PartialEq, Clone, Debug)]
pub struct KeyPress {
    modifiers: KeyMod,
    keys: Vec<Keysym>
}

pub fn init() {
    macro_rules! insert_all {
        ( $($press:expr => $name:expr);+ ) => {
            register(vec![
                $( ($press, commands::get(&$name.to_string())
                  .expect("Unable to register default command!")) ),+
                ]);
        }
    }

    insert_all! {
        keypress!("Alt", "Escape") => "quit";
        keypress!("Alt", "Return") => "launch_terminal";
        keypress!("Alt", "d") => "launch_dmenu";
        keypress!("Alt", "l") => "dmenu_eval";
        KeyPress::from_key_names(vec!["Alt", "Shift"], vec!["l"])
            .expect("Unable to create default keypress")
            => "dmenu_lua_dofile";

        keypress!("Alt", "1") => "switch_workspace_1";
        keypress!("Alt", "2") => "switch_workspace_2";
        keypress!("Alt", "3") => "switch_workspace_3";
        keypress!("Alt", "4") => "switch_workspace_4";
        keypress!("Alt", "5") => "switch_workspace_5";
        keypress!("Alt", "6") => "switch_workspace_6";
        keypress!("Alt", "7") => "switch_workspace_7";
        keypress!("Alt", "8") => "switch_workspace_8";
        keypress!("Alt", "9") => "switch_workspace_9";
        keypress!("Alt", "0") => "switch_workspace_0"
    }
}

/// Parses a KeyMod from key names.
pub fn keymod_from_names(keys: Vec<&str>) -> Result<KeyMod, String> {
    let mut result = KeyMod::empty();
    for key in keys {
        match key.to_lowercase().as_str() {
            "shift"            => result = result | MOD_SHIFT,
            "control" | "ctrl" => result = result | MOD_CTRL,
            "alt"              => result = result | MOD_ALT,
            "mod2"             => result = result | MOD_MOD2,
            "mod3"             => result = result | MOD_MOD3,
            "mod4" | "super" | "logo" => result = result | MOD_MOD4,
            "mod5" | "5mod5me" => result = result | MOD_MOD5,
            err => return Err(format!("Invalid modifier: {}", err))
        }
    }
    return Ok(result);
}

impl KeyPress {
    /// Creates a new KeyPress struct from a list of modifier and key names
    pub fn from_key_names(mods: Vec<&str>, keys: Vec<&str>) -> Result<KeyPress, String> {
        keymod_from_names(mods).and_then(|mods| {
            let mut syms: Vec<Keysym> = Vec::with_capacity(keys.len());
            for name in keys {
                // Parse a keysym for each given key
                if let Some(sym) = Keysym::from_name(name.to_string(),
                                                     NameFlags::None) {
                    syms.push(sym);
                }
                // If lowercase cannot be parsed, try case insensitive
                else if let Some(sym) = Keysym::from_name(name.to_string(),
                                                          NameFlags::CaseInsensitive) {
                    syms.push(sym);
                }
                else {
                    return Err(format!("Invalid key: {}", name));
                }
            }
            // Sort and dedup to make sure hashes are the same
            syms.sort_by_key(|s| s.get_code());
            syms.dedup();
            return Ok(KeyPress { modifiers: mods, keys: syms });
        })
    }

    /// Creates a KeyPress from keys that are pressed at the moment
    pub fn new(mods: KeyMod, mut keys: Vec<Keysym>) -> KeyPress {
        // Sort and dedup to make sure hashes are the same
        keys.sort_by_key(|k| k.get_code());
        keys.dedup();

        KeyPress { modifiers: mods, keys: keys }
    }
}

impl Hash for KeyPress {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        hasher.write_u32(self.modifiers.bits());
        for ref sym in &self.keys {
            hasher.write_u32(sym.get_code());
        }
    }
}

/// The type of function a key press handler is.
pub type KeyEvent = Arc<Fn() + Send + Sync>;

/// Get a key mapping from the list.
pub fn get(key: &KeyPress) -> Option<KeyEvent> {
    let bindings = BINDINGS.read()
        .expect("Keybindings/get: unable to lock keybindings");
    match bindings.get(key) {
        None => None,
        Some(val) => Some(val.clone())
    }
}

/// Register a new set of key mappings
#[allow(dead_code)]
pub fn register(values: Vec<(KeyPress, KeyEvent)>) {
    let mut bindings = BINDINGS.write()
        .expect("Keybindings/register: unable to lock keybindings");
    for value in values {
        bindings.insert(value.0, value.1);
    }
}

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

    fn test_cmd() {
        assert!(true);
    }

    fn keypress() -> KeyPress {
        keypress!("Ctrl", "t")
    }

    #[test]
    fn add_key() {
        require_rustwlc!();
        register(vec![(keypress(), Arc::new(test_cmd))]);
        assert!(get(&keypress()).is_some(), "Key not registered");
    }
}
<|endoftext|>


<commit_before>extern crate piston;
extern crate graphics;
extern crate glutin_window;
extern crate opengl_graphics;

use piston::window::WindowSettings;
use piston::event_loop::*;
use piston::input::*;
use glutin_window::GlutinWindow as Window;
use opengl_graphics::{GlGraphics, OpenGL};
use graphics::*;

const OPENGL: OpenGL = OpenGL::V3_2;
const WIDTH: i64 = 800;
const HEIGHT: i64 = 640;

pub struct Creature {
    x: f64,
    y: f64,
    rot: f64,
    rect: [f64; 4],
}

pub struct App {
    gl: GlGraphics, // OpenGL drawing backend.
    up_d: bool,
    down_d: bool,
    left_d: bool,
    right_d: bool,
    player: Creature,
}

impl Creature {
    fn new() -> Self {
        Creature {
            x: (WIDTH / 2) as f64,
            y: (HEIGHT / 2) as f64,
            rot: 0.0,
            rect: rectangle::square(0.0, 0.0, 50.0),
        }
    }
    fn moves(&mut self, dx: f64, dy: f64) {
        self.x = border_add(self.x, dx, true);
        self.y = border_add(self.y, dy, false);
    }
}

impl App {
    fn new() -> Self {
        App {
            gl: GlGraphics::new(OPENGL), // OpenGL drawing backend.
            up_d: false,
            down_d: false,
            left_d: false,
            right_d: false,
            player: Creature::new(),
        }
    }

    fn render(&mut self, args: &RenderArgs) {

        const GREEN: [f32; 4] = [0.0, 1.0, 0.0, 1.0];
        const RED: [f32; 4] = [1.0, 0.0, 0.0, 1.0];
        const BLACK: [f32; 4] = [0.0, 0.0, 0.0, 1.0];
        const WHITE: [f32; 4] = [1.0, 1.0, 1.0, 1.0];

        let square_length = 20;
        let field_width = 640 / square_length - 2;
        let field_heigth = 800 / square_length - 2;
        let mut vec: Vec<[f64; 4]> = Vec::new();

        for i in 0..field_heigth {
            for j in 0..field_width {
                vec.push(rectangle::square((20 + square_length * i) as f64,
                                           (20 + square_length * j) as f64,
                                           square_length as f64));
            }
        }
        let (x, y) = (self.player.x as f64, self.player.y as f64);
        let player = self.player.rect;
                let rotation = self.player.rot;
        self.gl.draw(args.viewport(), |c, gl| {
            // Clear the screen.
            clear(GREEN, gl);

            let transform = c.transform
                .trans(x, y)
                .rot_rad(rotation)
                .trans(-25.0, -25.0);

            // Draw a box rotating around the middle of the screen.
            for i in 0..(field_heigth) {
                for j in 0..field_width {
                    rectangle(if i % 2 == j % 2 { BLACK } else { WHITE },
                              vec[i * field_width + j],
                              c.transform,
                              gl);
                }
            }

            rectangle(RED, player, transform, gl);
        });
    }

    fn update(&mut self, args: &UpdateArgs) {
        // Rotate 2 radians per second.
        self.player.rot += 2.0 * args.dt;
        if self.up_d {
            self.player.moves(0.0, -1.0);
        }
        if self.down_d {
            self.player.moves(0.0, 1.0);
        }
        if self.right_d {
            self.player.moves(1.0, 0.0);
        }
        if self.left_d {
            self.player.moves(-1.0, 0.0);
        }
    }

    fn on_input(&mut self, inp: Button, pressed: bool) {

        match inp {
            Button::Keyboard(Key::Up) => {
                self.up_d = pressed;
            }
            Button::Keyboard(Key::Down) => {
                self.down_d = pressed;
            }
            Button::Keyboard(Key::Left) => {
                self.left_d = pressed;
            }
            Button::Keyboard(Key::Right) => {
                self.right_d = pressed;
            }
            _ => {}

        }
    }
}

fn main() {
    // Change this to OpenGL::V2_1 if not working.

    // Create an Glutin window.
    let mut window: Window = WindowSettings::new("spinning-square", [800, 640])
        .opengl(OPENGL)
        .exit_on_esc(true)
        .build()
        .unwrap();

    // Create a new game and run it.
    let mut app = App::new();

    let mut events = window.events();

    while let Some(e) = events.next(&mut window) {
        if let Some(r) = e.render_args() {
            app.render(&r);
        }

        if let Some(u) = e.update_args() {
            app.update(&u);
        }

        if let Some(i) = e.press_args() {
            app.on_input(i, true);
        }
        if let Some(i) = e.release_args() {
            app.on_input(i, false);
        }
    }
}

//UTIL//////////////////////////////

fn border_add(add1: f64, add2: f64, width: bool) -> f64 {
    let border = (if width { WIDTH } else { HEIGHT }) as f64;
    let sum = add1 + add2;
    if sum < 0.0 {
        0.0
    } else if sum > border {
        border
    } else {
        sum
    }
}
<commit_msg>add dynamic in chosen size<commit_after>extern crate piston;
extern crate graphics;
extern crate glutin_window;
extern crate opengl_graphics;

use piston::window::WindowSettings;
use piston::event_loop::*;
use piston::input::*;
use glutin_window::GlutinWindow as Window;
use opengl_graphics::{GlGraphics, OpenGL};
use graphics::*;

const OPENGL: OpenGL = OpenGL::V3_2;
const WIDTH: i64 = 400;
const HEIGHT: i64 = 400;

pub struct Creature {
    x: f64,
    y: f64,
    rot: f64,
    rect: [f64; 4],
}

pub struct App {
    gl: GlGraphics, // OpenGL drawing backend.
    up_d: bool,
    down_d: bool,
    left_d: bool,
    right_d: bool,
    player: Creature,
}

impl Creature {
    fn new() -> Self {
        Creature {
            x: (WIDTH / 2) as f64,
            y: (HEIGHT / 2) as f64,
            rot: 0.0,
            rect: rectangle::square(0.0, 0.0, 50.0),
        }
    }
    fn moves(&mut self, dx: f64, dy: f64) {
        self.x = border_add(self.x, dx, true);
        self.y = border_add(self.y, dy, false);
    }
}

impl App {
    fn new() -> Self {
        App {
            gl: GlGraphics::new(OPENGL), // OpenGL drawing backend.
            up_d: false,
            down_d: false,
            left_d: false,
            right_d: false,
            player: Creature::new(),
        }
    }

    fn render(&mut self, args: &RenderArgs) {

        const GREEN: [f32; 4] = [0.0, 1.0, 0.0, 1.0];
        const RED: [f32; 4] = [1.0, 0.0, 0.0, 1.0];
        const BLACK: [f32; 4] = [0.0, 0.0, 0.0, 1.0];
        const WHITE: [f32; 4] = [1.0, 1.0, 1.0, 1.0];

        let square_length = 20;
        let field_width = WIDTH / square_length - 2;
        let field_heigth = HEIGHT / square_length - 2;
        let mut vec: Vec<[f64; 4]> = Vec::new();

        for i in 0..field_heigth {
            for j in 0..field_width {
                vec.push(rectangle::square((20 + square_length * j) as f64,
                                           (20 + square_length * i) as f64,
                                           square_length as f64));
            }
        }
        let (x, y) = (self.player.x as f64, self.player.y as f64);
        let player = self.player.rect;
                let rotation = self.player.rot;
        self.gl.draw(args.viewport(), |c, gl| {
            // Clear the screen.
            clear(GREEN, gl);

            let transform = c.transform
                .trans(x, y)
                .rot_rad(rotation)
                .trans(-25.0, -25.0);

            // Draw a box rotating around the middle of the screen.
            for i in 0..field_heigth {
                for j in 0..field_width {
                    rectangle(if i % 2 == j % 2 { BLACK } else { WHITE },
                              vec[(i * field_width + j) as usize],
                              c.transform,
                              gl);
                }
            }

            rectangle(RED, player, transform, gl);
        });
    }

    fn update(&mut self, args: &UpdateArgs) {
        // Rotate 2 radians per second.
        self.player.rot += 2.0 * args.dt;
        if self.up_d {
            self.player.moves(0.0, -1.0);
        }
        if self.down_d {
            self.player.moves(0.0, 1.0);
        }
        if self.right_d {
            self.player.moves(1.0, 0.0);
        }
        if self.left_d {
            self.player.moves(-1.0, 0.0);
        }
    }

    fn on_input(&mut self, inp: Button, pressed: bool) {

        match inp {
            Button::Keyboard(Key::Up) => {
                self.up_d = pressed;
            }
            Button::Keyboard(Key::Down) => {
                self.down_d = pressed;
            }
            Button::Keyboard(Key::Left) => {
                self.left_d = pressed;
            }
            Button::Keyboard(Key::Right) => {
                self.right_d = pressed;
            }
            _ => {}

        }
    }
}

fn main() {
    // Create an Glutin window.
    let mut window: Window = WindowSettings::new("spinning-square", [WIDTH as u32 , HEIGHT as u32])
        .opengl(OPENGL)
        .exit_on_esc(true)
        .build()
        .unwrap();

    // Create a new game and run it.
    let mut app = App::new();

    let mut events = window.events();

    while let Some(e) = events.next(&mut window) {
        if let Some(r) = e.render_args() {
            app.render(&r);
        }

        if let Some(u) = e.update_args() {
            app.update(&u);
        }

        if let Some(i) = e.press_args() {
            app.on_input(i, true);
        }
        if let Some(i) = e.release_args() {
            app.on_input(i, false);
        }
    }
}

//UTIL//////////////////////////////

fn border_add(add1: f64, add2: f64, width: bool) -> f64 {
    let border = (if width { WIDTH } else { HEIGHT }) as f64;
    let sum = add1 + add2;
    if sum < 0.0 {
        0.0
    } else if sum > border {
        border
    } else {
        sum
    }
}
<|endoftext|>


<commit_before>use csv;
use indicatif;
use lazy_static;
use prettytable;
use regex;
use rusqlite;
use rustyline;
use std::env;
use std::error::Error;
use std::fs::File;

fn _normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)").unwrap();
    }
    RE.replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: String,
) -> Vec<String> {
    let mut num_rows = 0;
    let f = File::open(path).unwrap();
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::Reader::from_reader(f);

    let normalized_cols =
        reader
            .headers()
            .unwrap()
            .iter()
            .map(_normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).unwrap();
        while let Some(row) = records.next() {
            stmt.execute(&row.unwrap()).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    normalized_cols
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = conn
        .prepare(&line)
        .map_err(|e| e.description().to_owned())?;

    let mut table = prettytable::Table::new();
    let mut title_row = prettytable::Row::new(vec![]);
    for col in stmt.column_names() {
        title_row.add_cell(prettytable::Cell::new(col));
    }
    table.set_titles(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    table.set_format(*prettytable::format::consts::FORMAT_NO_LINESEP_WITH_TITLE);
    while let Ok(Some(r)) = results.next() {
        let mut row = prettytable::Row::new(vec![]);
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(prettytable::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    table.printstd();
    return Ok(());
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = conn
        .prepare(&query)
        .map_err(|e| e.description().to_owned())?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    return Ok(());
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

fn main() {
    let mut paths = env::args().skip(1);

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
        ".export",
    ]
    .iter()
    .map(|s| s.to_string())
    .collect::<Vec<String>>();

    if paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", paths.next().unwrap());
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in paths.enumerate() {
            let mut col_names = _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path);
            base_words.append(&mut col_names);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }
}
<commit_msg>clippy<commit_after>use csv;
use indicatif;
use lazy_static;
use prettytable;
use regex;
use rusqlite;
use rustyline;
use std::env;
use std::error::Error;
use std::fs::File;

fn _normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)").unwrap();
    }
    RE.replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: String,
) -> Vec<String> {
    let mut num_rows = 0;
    let f = File::open(path).unwrap();
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::Reader::from_reader(f);

    let normalized_cols =
        reader
            .headers()
            .unwrap()
            .iter()
            .map(_normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).unwrap();
        while let Some(row) = records.next() {
            stmt.execute(&row.unwrap()).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    normalized_cols
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = conn
        .prepare(&line)
        .map_err(|e| e.description().to_owned())?;

    let mut table = prettytable::Table::new();
    let mut title_row = prettytable::Row::new(vec![]);
    for col in stmt.column_names() {
        title_row.add_cell(prettytable::Cell::new(col));
    }
    table.set_titles(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    table.set_format(*prettytable::format::consts::FORMAT_NO_LINESEP_WITH_TITLE);
    while let Ok(Some(r)) = results.next() {
        let mut row = prettytable::Row::new(vec![]);
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(prettytable::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    table.printstd();
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = conn
        .prepare(&query)
        .map_err(|e| e.description().to_owned())?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

fn main() {
    let mut paths = env::args().skip(1);

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
        ".export",
    ]
    .iter()
    .map(|s| s.to_string())
    .collect::<Vec<String>>();

    if paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", paths.next().unwrap());
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in paths.enumerate() {
            let mut col_names = _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path);
            base_words.append(&mut col_names);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }
}
<|endoftext|>


<commit_before>

#![feature(path_ext)]
extern crate mysql;
extern crate toml;
extern crate rustc_serialize;
#[macro_use]
extern crate lazy_static;

mod lib;

use mysql::conn::MyOpts;
use std::default::Default;
use mysql::conn::pool;
use mysql::error::MyError;
use mysql::conn::pool::MyPooledConn;
use mysql::value::from_value;

use std::error::Error;

use std::ascii::AsciiExt;

use toml::Table;
use toml::Value;

use lib::config;
use lib::downloader::DownloadDB;
use lib::downloader::Downloader;
use lib::downloader::DownloadError;
use lib::converter::Converter;

use std::fs::rename;

static VERSION : &'static str = "0.1"; // String not valid
static SLEEP_MS: u32 = 5000;

lazy_static! {
    static ref CONFIG: config::Config = {
        println!("Starting yayd-backend v{}",&VERSION);
        config::init_config() //return
    };

}

//#[repr(i8)] broken, enum not usable as of #10292
enum QueryCodes {
    Waiting = 0,
    InProgress = 1,
    Finished = 2,
    Failed = 3,
}

fn main() {
    let opts = mysql_options();
    let pool = match pool::MyPool::new(opts) {
        Ok(conn) => { println!("Connected successfully."); conn},
        Err(err) => panic!("Unable to establish a connection!\n{}",err),
    };
    let converter = Converter::new(&CONFIG.general.ffmpeg_bin, &CONFIG.codecs.audio, pool.clone());
    loop {
        if let Some(result) = request_entry(& pool) {
            if result.playlist {
                println!("Playlist not supported atm!");
                //TODO: set playlist entry to errg
            }
            let qid = result.qid.clone();                 //&QueryCodes::InProgress as i32
            set_query_code(&mut pool.get_conn().unwrap(), &1, &result.qid).ok().expect("Failed to set query code!");

            let code = if handle_download(result, None, &converter) {
                2//QueryCodes::Finished as i32
            } else {
                3//QueryCodes::Failed as i32
            };
            set_query_code(&mut pool.get_conn().unwrap(), &code,&qid).ok().expect("Failed to set query code!");
        } else {
            println!("Pausing..");
            std::thread::sleep_ms(SLEEP_MS);
        }
    }

    // let downloader = downloader::Downloader::new(download_db);
    // downloader.download_video();

    println!("EOL!");
}

///Download handler
///Used by the playlist/file handler to download one file
///Based on the quality it needs to download audio & video splitted & convert them together
///In case of a playlist download it depends on the target download folder & if should bezipped
///In case of a DMCA we need to download the file via the socket connector,
///which will output a mp3, or if requested, a hard quality depending on the max-available of
///the video url. Thus in case of a DMCA we can't pick a quality anymore.
///Also the filename depends on the socket output then.
///
///If it's a non-zipped single file, the file is moved after a success download,convert etc to the
///main folder from which it should be downloadable.
///The original non-ascii & url_encode name of the file is stored in the DB
fn handle_download(downl_db: DownloadDB, folder: Option<String>, converter: &Converter) -> bool{
    let dbcopy = downl_db.clone(); //copy, all implement copy & no &'s in use
    let download = Downloader::new(downl_db, &CONFIG.general);
    //get filename, check for DMCA
    let mut dmca = false; // "succ." dmca -> file already downloaded
    let name = match download.get_file_name() { // get filename
        Ok(v) => v,
        Err(DownloadError::DMCAError) => { //now request via lib.. // k if( k == Err(DownloadError::DMCAError) )
            println!("DMCA error!");
            match download.lib_request_video() {
                Err(err) => { println!("Offliberty-call error {:?}", err); return false; },
                Ok(v) => { dmca = true; v },
            }
        },
        Err(e) => { // unknown error / video private etc.. abort
            println!("Unknown error: {:?}", e);
            //TODO: add error descr (change enum etc)
            set_query_state(&dbcopy.pool.clone(),&dbcopy.qid, "Error!");
            return false;
        },
    };

    println!("Filename: {}", name);

    if(dmca){
        //TODO: insert title name for file,
        //copy file to download folder
        return false;
    }

    let file_path = format_file_patH(&dbcopy.qid, folder.clone());
    let save_path = &format_save_path(folder.clone(),&name, &download);
    //download video, which is video/audio(m4a)
    download.download_video(&file_path, false);

    if is_split_container(&dbcopy.quality) { // download audio file & convert together
        let audio_path = format_audio_path(&dbcopy.qid, folder.clone());

        println!("Downloading audio.. {}", audio_path);
        download.download_video(&audio_path, true);
        match converter.merge_files(&dbcopy.qid,&file_path, &audio_path,&save_path) {
            Err(e) => {println!("merge error: {:?}",e); return false;},
            Ok(()) => {},
        }
    }else{
        if download.is_audio(){ // if audio-> convert m4a to mp3, which converts directly to downl. dir
            //TODO: convert -> saves already ?
        }else{
            //converter.
            //match rename(format!("{}/{}", folder));
            //TODO: move file
        }

    }

    //TODO: download file, convert if audio ?!
    true
}

fn get_file_ext<'a>(download: &Downloader) -> &'a str {
    match download.is_audio() {
        true => "mp3",
        false => "mp4",
    }
}

fn format_file_patH(qid: &i64, folder: Option<String>) -> String {
    match folder {
        Some(v) => format!("{}/{}/{}", &CONFIG.general.save_dir, v, qid),
        None => format!("{}/{}", &CONFIG.general.save_dir, qid),
    }
}

///Formats the audio path, based on the qid & optional folders
fn format_audio_path(qid: &i64, folder: Option<String>) -> String {
    match folder {
        Some(v) => format!("{}/{}/{}a", &CONFIG.general.save_dir, v, qid),
        None => format!("{}/{}a", &CONFIG.general.save_dir, qid),
    }
}

///Format save path, which stays inside the optional folder, if set for zipping later
fn format_save_path<'a>(folder: Option<String>, name: &str, download: &'a Downloader) -> String {
    match folder {
        Some(v) => format!("{}/{}/{}.{}", &CONFIG.general.save_dir, v, url_encode(name),get_file_ext(download)),
        None => format!("{}/{}.{}", &CONFIG.general.download_dir, url_encode(name),get_file_ext(download)),
    }
}

///Set the state of the current query, also in dependence of the code, see QueryCodes
fn set_query_state(pool: & pool::MyPool,qid: &i64 , state: &str){ // same here
    let mut conn = pool.get_conn().unwrap();
    let mut stmt = conn.prepare("UPDATE querydetails SET status = ? WHERE qid = ?").unwrap();
    let result = stmt.execute(&[&state,qid]); // why is this var needed ?!
    match result {
        Ok(_) => (),
        Err(why) => println!("Error setting query state: {}",why),
    }
}

///Return whether the quality is a split container: video only
///as specified in the docs
fn is_split_container(quality: &i16) -> bool {
    match *quality {
        141 | 83 | 82 | 84 | 85 => false,
        _ => true,
    }
}

///Return an url-conform String
fn url_encode(input: &str) -> String {
    // iterator over input, apply function to each element(function
    input.chars().map(|char| {
        match char {
            ' ' | '?' | '!' | '\\' | '/' | '.' | '(' | ')' | '[' | ']' => '_',
            '&' => '-',
            c if c.is_ascii() => c,
            _ => '_'
        }
    }).collect()
    // match for each char, then do collect: loop through the iterator, collect all elements
    // into container FromIterator
}

///Request an entry from the DB that should be handled
fn request_entry(pool: & pool::MyPool) -> Option<DownloadDB> {
    let mut conn = pool.get_conn().unwrap();
    let mut stmt = conn.prepare("SELECT queries.qid,url,type,quality FROM querydetails \
                    INNER JOIN queries \
                    ON querydetails.qid = queries.qid \
                    WHERE querydetails.code = 0 \
                    ORDER BY queries.created \
                    LIMIT 1").unwrap();
    let mut result = stmt.execute(&[]).unwrap();
    let result = match result.next() {
        Some(val) => val.unwrap(),
        None => {return None; },
    };
    println!("Result: {:?}", result[0]);
    println!("result str: {}", result[1].into_str());
    let download_db = DownloadDB { url: from_value::<String>(&result[1]),
                                                quality: from_value::<i16>(&result[3]),
                                                qid: from_value::<i64>(&result[0]),
                                                audioquality: CONFIG.codecs.audio,
                                                folder: CONFIG.general.save_dir.clone(),
                                                pool: pool.clone(),
                                                playlist: false, //TEMP
                                                compress: false };
    Some(download_db)
}

fn set_query_code(conn: & mut MyPooledConn, code: &i8, qid: &i64) -> Result<(), DownloadError> { // same here
    let mut stmt = conn.prepare("UPDATE querydetails SET code = ? WHERE qid = ?").unwrap();
    let result = stmt.execute(&[code,qid]); // why is this var needed ?!
    match result {
        Ok(_) => Ok(()),
        Err(why) => Err(DownloadError::DBError(why.description().into())),
    }
}

///Set options for the connection
fn mysql_options() -> MyOpts {
    //let dbconfig = CONFIG.get("db").unwrap().clone();
    //let dbconfig = dbconfig.as_table().unwrap(); // shadow binding to workaround borrow / lifetime problems

    MyOpts {
        tcp_addr: Some(CONFIG.db.ip.clone()),
        tcp_port: CONFIG.db.port,
        user: Some(CONFIG.db.user.clone()),
        pass: Some(CONFIG.db.password.clone()),
        db_name: Some(CONFIG.db.db.clone()),
        ..Default::default() // set other to default
    }
}

///Converts a toml::Value to a Option<String> for mysql::MyOpts
#[allow(deprecated)]
fn get_option_string(table: & Table,key: & str) -> Option<String> {
    let val: Value = table.get(key).unwrap().clone();
    if let toml::Value::String(s) = val {
        Some(s)
    } else { unreachable!() }
}<commit_msg>+set query sates +add file to table<commit_after>

#![feature(path_ext)]
extern crate mysql;
extern crate toml;
extern crate rustc_serialize;
#[macro_use]
extern crate lazy_static;

mod lib;

use mysql::conn::MyOpts;
use std::default::Default;
use mysql::conn::pool;
use mysql::error::MyError;
use mysql::conn::pool::MyPooledConn;
use mysql::value::from_value;

use std::error::Error;

use std::ascii::AsciiExt;

use toml::Table;
use toml::Value;

use lib::config;
use lib::downloader::DownloadDB;
use lib::downloader::Downloader;
use lib::downloader::DownloadError;
use lib::converter::Converter;

use std::fs::rename;

static VERSION : &'static str = "0.1"; // String not valid
static SLEEP_MS: u32 = 5000;

lazy_static! {
    static ref CONFIG: config::Config = {
        println!("Starting yayd-backend v{}",&VERSION);
        config::init_config() //return
    };

}

//#[repr(i8)] broken, enum not usable as of #10292
enum QueryCodes {
    Waiting = 0,
    InProgress = 1,
    Finished = 2,
    Failed = 3,
}

fn main() {
    let opts = mysql_options();
    let pool = match pool::MyPool::new(opts) {
        Ok(conn) => { println!("Connected successfully."); conn},
        Err(err) => panic!("Unable to establish a connection!\n{}",err),
    };
    let converter = Converter::new(&CONFIG.general.ffmpeg_bin, &CONFIG.codecs.audio, pool.clone());
    loop {
        if let Some(result) = request_entry(& pool) {
            if result.playlist {
                println!("Playlist not supported atm!");
                //TODO: set playlist entry to errg
            }
            let qid = result.qid.clone();                 //&QueryCodes::InProgress as i32
            set_query_code(&mut pool.get_conn().unwrap(), &1, &result.qid).ok().expect("Failed to set query code!");
            set_query_state(&pool.clone(),&result.qid, "started");
            let code = if handle_download(result, None, &converter) {
                2//QueryCodes::Finished as i32
            } else {
                3//QueryCodes::Failed as i32
            };
            set_query_code(&mut pool.get_conn().unwrap(), &code,&qid).ok().expect("Failed to set query code!");
            let state = if code = 2 {
                "finished"
            } else {
                "failed"
            };
            set_query_state(&pool.clone(),&result.qid, state);
        } else {
            println!("Pausing..");
            std::thread::sleep_ms(SLEEP_MS);
        }
    }

    // let downloader = downloader::Downloader::new(download_db);
    // downloader.download_video();

    println!("EOL!");
}

///Download handler
///Used by the playlist/file handler to download one file
///Based on the quality it needs to download audio & video splitted & convert them together
///In case of a playlist download it depends on the target download folder & if should bezipped
///In case of a DMCA we need to download the file via the socket connector,
///which will output a mp3, or if requested, a hard quality depending on the max-available of
///the video url. Thus in case of a DMCA we can't pick a quality anymore.
///Also the filename depends on the socket output then.
///
///If it's a non-zipped single file, the file is moved after a success download,convert etc to the
///main folder from which it should be downloadable.
///The original non-ascii & url_encode name of the file is stored in the DB
fn handle_download(downl_db: DownloadDB, folder: Option<String>, converter: &Converter) -> bool{
    let is_zipped = match folder {
        Some(v) => true,
        None => false,
    }

    let dbcopy = downl_db.clone(); //copy, all implement copy & no &'s in use
    let download = Downloader::new(downl_db, &CONFIG.general);
    //get filename, check for DMCA
    let mut dmca = false; // "succ." dmca -> file already downloaded
    let name = match download.get_file_name() { // get filename
        Ok(v) => v,
        Err(DownloadError::DMCAError) => { //now request via lib.. // k if( k == Err(DownloadError::DMCAError) )
            println!("DMCA error!");
            match download.lib_request_video() {
                Err(err) => { println!("Offliberty-call error {:?}", err); return false; },
                Ok(v) => { dmca = true; v },
            }
        },
        Err(e) => { // unknown error / video private etc.. abort
            println!("Unknown error: {:?}", e);
            //TODO: add error descr (change enum etc)
            set_query_state(&dbcopy.pool.clone(),&dbcopy.qid, "Error!");
            return false;
        },
    };

    let name_http_valid = url_encode(name);

    println!("Filename: {}", name);

    if(dmca){
        //TODO: insert title name for file,
        //copy file to download folder
        return true;
    }

    let file_path = format_file_patH(&dbcopy.qid, folder.clone());
    let save_path = &format_save_path(folder.clone(),&name, &download);

    let is_splitted_format = is_split_container(&dbcopy.quality);
    let total_steps = if is_splitted_format {
        4
    } else {
        2
    };
    //update progress
    update_steps(&pool.clone(),&result.qid, 1, total_steps));

    //download video, which is video/audio(m4a)
    download.download_video(&file_path, false);

    update_steps(&pool.clone(),&result.qid, 2, total_steps));

    if is_splitted_format { // download audio file & convert together
        let audio_path = format_audio_path(&dbcopy.qid, folder.clone());

        println!("Downloading audio.. {}", audio_path);
        download.download_video(&audio_path, true);

        update_steps(&pool.clone(),&result.qid, 3, total_steps));

        match converter.merge_files(&dbcopy.qid,&file_path, &audio_path,&save_path) {
            Err(e) => {println!("merge error: {:?}",e); return false;},
            Ok(()) => {},
        }

        update_steps(&pool.clone(),&result.qid, 4, total_steps));

    }else{
        if download.is_audio(){ // if audio-> convert m4a to mp3, which converts directly to downl. dir
            //TODO: convert -> saves already ?
        }else{
            //converter.
            //match rename(format!("{}/{}", folder));
            //TODO: move file
        }

    }

    if is_zipped { // add file to list, except it's for zip-compression later (=folder set)
        add_file_entry(&dbcopy.pool.clone(),&name_http_valid, &name);
    }

    //TODO: download file, convert if audio ?!
    true
}

fn update_steps(pool: & pool::MyPool, qid: &i64, state: &str, step: i32, max_steps: i32){
    set_query_state(&pool,&result.qid, format!("{}|{}", step, max_steps));
}

fn get_file_ext<'a>(download: &Downloader) -> &'a str {
    match download.is_audio() {
        true => "mp3",
        false => "mp4",
    }
}

fn format_file_patH(qid: &i64, folder: Option<String>) -> String {
    match folder {
        Some(v) => format!("{}/{}/{}", &CONFIG.general.save_dir, v, qid),
        None => format!("{}/{}", &CONFIG.general.save_dir, qid),
    }
}

///Formats the audio path, based on the qid & optional folders
fn format_audio_path(qid: &i64, folder: Option<String>) -> String {
    match folder {
        Some(v) => format!("{}/{}/{}a", &CONFIG.general.save_dir, v, qid),
        None => format!("{}/{}a", &CONFIG.general.save_dir, qid),
    }
}

///Format save path, which stays inside the optional folder, if set for zipping later
fn format_save_path<'a>(folder: Option<String>, name: &str, download: &'a Downloader) -> String {
    match folder {
        Some(v) => format!("{}/{}/{}.{}", &CONFIG.general.save_dir, v, url_encode(name),get_file_ext(download)),
        None => format!("{}/{}.{}", &CONFIG.general.download_dir, url_encode(name),get_file_ext(download)),
    }
}

///Set the state of the current query, also in dependence of the code, see QueryCodes
fn set_query_state(pool: & pool::MyPool,qid: &i64 , state: &str){ // same here
    let mut conn = pool.get_conn().unwrap();
    let mut stmt = conn.prepare("UPDATE querydetails SET status = ? WHERE qid = ?").unwrap();
    let result = stmt.execute(&[&state,qid]); // why is this var needed ?!
    match result {
        Ok(_) => (),
        Err(why) => println!("Error setting query state: {}",why),
    }
}

fn add_file_entry(pool: & pool::MyPool, fid: &i64, name: &str, &real_name: &str){
    let mut conn = pool.get_conn().unwrap();
    let mut stmt = conn.prepare("INSERT INTO files (fid,name,rname,valid) VALUES (?,?,?,?)").unwrap();
    let result = stmt.execute(&[fid,real_name,name,&1]); // why is this var needed ?!
    match result {
        Ok(_) => (),
        Err(why) => println!("Error adding file: {}",why),
    }
}

///Return whether the quality is a split container: video only
///as specified in the docs
fn is_split_container(quality: &i16) -> bool {
    match *quality {
        141 | 83 | 82 | 84 | 85 => false,
        _ => true,
    }
}

///Return an url-conform String
fn url_encode(input: &str) -> String {
    // iterator over input, apply function to each element(function
    input.chars().map(|char| {
        match char {
            ' ' | '?' | '!' | '\\' | '/' | '.' | '(' | ')' | '[' | ']' => '_',
            '&' => '-',
            c if c.is_ascii() => c,
            _ => '_'
        }
    }).collect()
    // match for each char, then do collect: loop through the iterator, collect all elements
    // into container FromIterator
}

///Request an entry from the DB that should be handled
fn request_entry(pool: & pool::MyPool) -> Option<DownloadDB> {
    let mut conn = pool.get_conn().unwrap();
    let mut stmt = conn.prepare("SELECT queries.qid,url,type,quality FROM querydetails \
                    INNER JOIN queries \
                    ON querydetails.qid = queries.qid \
                    WHERE querydetails.code = 0 \
                    ORDER BY queries.created \
                    LIMIT 1").unwrap();
    let mut result = stmt.execute(&[]).unwrap();
    let result = match result.next() {
        Some(val) => val.unwrap(),
        None => {return None; },
    };
    println!("Result: {:?}", result[0]);
    println!("result str: {}", result[1].into_str());
    let download_db = DownloadDB { url: from_value::<String>(&result[1]),
                                                quality: from_value::<i16>(&result[3]),
                                                qid: from_value::<i64>(&result[0]),
                                                audioquality: CONFIG.codecs.audio,
                                                folder: CONFIG.general.save_dir.clone(),
                                                pool: pool.clone(),
                                                playlist: false, //TEMP
                                                compress: false };
    Some(download_db)
}

fn set_query_code(conn: & mut MyPooledConn, code: &i8, qid: &i64) -> Result<(), DownloadError> { // same here
    let mut stmt = conn.prepare("UPDATE querydetails SET code = ? WHERE qid = ?").unwrap();
    let result = stmt.execute(&[code,qid]); // why is this var needed ?!
    match result {
        Ok(_) => Ok(()),
        Err(why) => Err(DownloadError::DBError(why.description().into())),
    }
}

///Set options for the connection
fn mysql_options() -> MyOpts {
    //let dbconfig = CONFIG.get("db").unwrap().clone();
    //let dbconfig = dbconfig.as_table().unwrap(); // shadow binding to workaround borrow / lifetime problems

    MyOpts {
        tcp_addr: Some(CONFIG.db.ip.clone()),
        tcp_port: CONFIG.db.port,
        user: Some(CONFIG.db.user.clone()),
        pass: Some(CONFIG.db.password.clone()),
        db_name: Some(CONFIG.db.db.clone()),
        ..Default::default() // set other to default
    }
}

///Converts a toml::Value to a Option<String> for mysql::MyOpts
#[allow(deprecated)]
fn get_option_string(table: & Table,key: & str) -> Option<String> {
    let val: Value = table.get(key).unwrap().clone();
    if let toml::Value::String(s) = val {
        Some(s)
    } else { unreachable!() }
}<|endoftext|>


<commit_before>extern crate base64;
extern crate clap;
extern crate hyper;
extern crate hyper_rustls;
extern crate pem;
#[macro_use]
extern crate prettytable;
extern crate serde_json;

extern crate ct_tools;

use std::fs::File;
use std::io::Read;

use ct_tools::{censys, crtsh};
use ct_tools::ct::submit_cert_to_logs;
use ct_tools::google::fetch_trusted_ct_logs;

fn pems_to_chain(data: &str) -> Vec<Vec<u8>> {
    return pem::parse_many(data)
               .into_iter()
               .filter(|p| p.tag == "CERTIFICATE")
               .map(|p| p.contents)
               .collect();
}

fn submit(paths: clap::Values) {
    let http_client = hyper::Client::with_connector(
        hyper::net::HttpsConnector::new(hyper_rustls::TlsClient::new())
    );
    // TODO: timeout on the http_client, but for submit_cert_to_logs only, not build_chain
    let logs = fetch_trusted_ct_logs(&http_client);

    for path in paths {
        println!("Submitting {} ...", path);

        let mut contents = String::new();
        File::open(path)
            .unwrap()
            .read_to_string(&mut contents)
            .unwrap();

        let mut chain = pems_to_chain(&contents);
        if chain.len() == 1 {
            // TODO: There's got to be some way to do this ourselves, instead of using crt.sh as a
            // glorified AIA chaser.
            println!("Only one certificate in chain, using crt.sh to build a full chain ...");
            chain = crtsh::build_chain_for_cert(&http_client, &chain[0]);
        }
        let scts = submit_cert_to_logs(&http_client, &logs, &chain);

        println!("Find the cert on crt.sh: {}",
                 crtsh::url_for_cert(&chain[0]));
        let mut table = prettytable::Table::new();
        table.add_row(row!["Log", "SCT"]);
        for (log, sct) in scts {
            table.add_row(row![log.description, base64::encode(&sct.to_raw_bytes())]);
        }
        table.printstd();
        println!();
        println!();
    }
}

fn check(paths: clap::Values) {
    let http_client = hyper::Client::with_connector(
        hyper::net::HttpsConnector::new(hyper_rustls::TlsClient::new())
    );

    for path in paths {
        let mut contents = String::new();
        File::open(path)
            .unwrap()
            .read_to_string(&mut contents)
            .unwrap();

        let chain = pems_to_chain(&contents);

        let is_logged = censys::is_cert_logged(&http_client, &chain[0]);
        if is_logged {
            println!("{} was already logged", path);
        } else {
            println!("{} has not been logged", path);
        }
    }
}

fn main() {
    let matches = clap::App::new("ct-submitter")
        .subcommand(clap::SubCommand::with_name("submit")
                        .about("Directly submits certificates to CT logs")
                        .arg(clap::Arg::with_name("path")
                                 .multiple(true)
                                 .required(true)
                                 .help("Path to certificate or chain")))
        .subcommand(clap::SubCommand::with_name("check")
                        .about("Checks whether a certificate exists in CT logs")
                        .arg(clap::Arg::with_name("path")
                                 .multiple(true)
                                 .required(true)
                                 .help("Path to certificate or chain")))
        .get_matches();

    if let Some(matches) = matches.subcommand_matches("submit") {
        submit(matches.values_of("path").unwrap());
    } else if let Some(matches) = matches.subcommand_matches("check") {
        check(matches.values_of("path").unwrap());
    }
}
<commit_msg>changed name here too<commit_after>extern crate base64;
extern crate clap;
extern crate hyper;
extern crate hyper_rustls;
extern crate pem;
#[macro_use]
extern crate prettytable;
extern crate serde_json;

extern crate ct_tools;

use std::fs::File;
use std::io::Read;

use ct_tools::{censys, crtsh};
use ct_tools::ct::submit_cert_to_logs;
use ct_tools::google::fetch_trusted_ct_logs;

fn pems_to_chain(data: &str) -> Vec<Vec<u8>> {
    return pem::parse_many(data)
               .into_iter()
               .filter(|p| p.tag == "CERTIFICATE")
               .map(|p| p.contents)
               .collect();
}

fn submit(paths: clap::Values) {
    let http_client = hyper::Client::with_connector(
        hyper::net::HttpsConnector::new(hyper_rustls::TlsClient::new())
    );
    // TODO: timeout on the http_client, but for submit_cert_to_logs only, not build_chain
    let logs = fetch_trusted_ct_logs(&http_client);

    for path in paths {
        println!("Submitting {} ...", path);

        let mut contents = String::new();
        File::open(path)
            .unwrap()
            .read_to_string(&mut contents)
            .unwrap();

        let mut chain = pems_to_chain(&contents);
        if chain.len() == 1 {
            // TODO: There's got to be some way to do this ourselves, instead of using crt.sh as a
            // glorified AIA chaser.
            println!("Only one certificate in chain, using crt.sh to build a full chain ...");
            chain = crtsh::build_chain_for_cert(&http_client, &chain[0]);
        }
        let scts = submit_cert_to_logs(&http_client, &logs, &chain);

        println!("Find the cert on crt.sh: {}",
                 crtsh::url_for_cert(&chain[0]));
        let mut table = prettytable::Table::new();
        table.add_row(row!["Log", "SCT"]);
        for (log, sct) in scts {
            table.add_row(row![log.description, base64::encode(&sct.to_raw_bytes())]);
        }
        table.printstd();
        println!();
        println!();
    }
}

fn check(paths: clap::Values) {
    let http_client = hyper::Client::with_connector(
        hyper::net::HttpsConnector::new(hyper_rustls::TlsClient::new())
    );

    for path in paths {
        let mut contents = String::new();
        File::open(path)
            .unwrap()
            .read_to_string(&mut contents)
            .unwrap();

        let chain = pems_to_chain(&contents);

        let is_logged = censys::is_cert_logged(&http_client, &chain[0]);
        if is_logged {
            println!("{} was already logged", path);
        } else {
            println!("{} has not been logged", path);
        }
    }
}

fn main() {
    let matches = clap::App::new("ct-tools")
        .subcommand(clap::SubCommand::with_name("submit")
                        .about("Directly submits certificates to CT logs")
                        .arg(clap::Arg::with_name("path")
                                 .multiple(true)
                                 .required(true)
                                 .help("Path to certificate or chain")))
        .subcommand(clap::SubCommand::with_name("check")
                        .about("Checks whether a certificate exists in CT logs")
                        .arg(clap::Arg::with_name("path")
                                 .multiple(true)
                                 .required(true)
                                 .help("Path to certificate or chain")))
        .get_matches();

    if let Some(matches) = matches.subcommand_matches("submit") {
        submit(matches.values_of("path").unwrap());
    } else if let Some(matches) = matches.subcommand_matches("check") {
        check(matches.values_of("path").unwrap());
    }
}
<|endoftext|>


<commit_before>extern crate image;
extern crate rand;
extern crate docopt;
extern crate rustc_serialize;
extern crate num_cpus;
extern crate scoped_threadpool;
extern crate clock_ticks;
extern crate tray_rust;

use std::vec::Vec;
use std::iter;
use std::path::PathBuf;
use std::io::ErrorKind;

use rand::StdRng;
use docopt::Docopt;

use tray_rust::film::{self, ImageSample, Colorf};
use tray_rust::geometry::{Geometry, Instance, Emitter};
use tray_rust::sampler::{self, Sampler};
use tray_rust::scene;
use tray_rust::integrator::Integrator;

static USAGE: &'static str = "
Usage: tray_rust <scenefile> [options]

Options:
  -o <file>     Specify the output file or directory to save the image or frames. Supported formats are
                PNG, JPG and PPM. Default is 'frame<#>.png'.
  -n <number>   Specify the number of threads to use for rendering. Defaults to the number of cores
                on the system.
  -h, --help    Show this message.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_scenefile: String,
    flag_o: Option<String>,
    flag_n: Option<u32>,
}

/// Threads are each sent a sender end of the channel that is
/// read from by the render target thread which then saves the
/// values recieved to the render target
fn thread_work(spp: usize, queue: &sampler::BlockQueue, scene: &scene::Scene,
               target: &film::RenderTarget, light_list: &Vec<&Emitter>) {
    let mut sampler = sampler::LowDiscrepancy::new(queue.block_dim(), spp);
    let mut sample_pos = Vec::with_capacity(sampler.max_spp());
    let mut time_samples: Vec<_> = iter::repeat(0.0).take(sampler.max_spp()).collect();
    let block_dim = queue.block_dim();
    let mut block_samples = Vec::with_capacity(sampler.max_spp() * (block_dim.0 * block_dim.1) as usize);
    let mut rng = match StdRng::new() {
        Ok(r) => r,
        Err(e) => { println!("Failed to get StdRng, {}", e); return }
    };
    // Grab a block from the queue and start working on it, submitting samples
    // to the render target thread after each pixel
    for b in queue.iter() {
        sampler.select_block(b);
        let mut pixel_samples = 0;
        while sampler.has_samples() {
            // Get samples for a pixel and render them
            sampler.get_samples(&mut sample_pos, &mut rng);
            sampler.get_samples_1d(&mut time_samples[..], &mut rng);
            for (s, t) in sample_pos.iter().zip(time_samples.iter()) {
                let mut ray = scene.camera.generate_ray(s, *t);
                if let Some(hit) = scene.intersect(&mut ray) {
                    let c = scene.integrator.illumination(scene, light_list, &ray,
                                                          &hit, &mut sampler, &mut rng).clamp();
                    block_samples.push(ImageSample::new(s.0, s.1, c));
                } else {
                    block_samples.push(ImageSample::new(s.0, s.1, Colorf::black()));
                }
            }
            // If the samples are ok the samples for the next pixel start at the end of the current
            // pixel's samples
            if sampler.report_results(&block_samples[pixel_samples..]) {
                pixel_samples = block_samples.len();
            }
        }
        target.write(&block_samples, sampler.get_region());
        block_samples.clear();
    }
}

/// Render the scene in parallel using `n` threads and write the result to the render target
fn render_parallel(rt: &mut film::RenderTarget, scene: &scene::Scene, pool: &mut scoped_threadpool::Pool, spp: usize){
    let dim = rt.dimensions();
    let block_queue = sampler::BlockQueue::new((dim.0 as u32, dim.1 as u32), (8, 8));
    let light_list: Vec<_> = scene.bvh.into_iter().filter_map(|x| {
        match x {
            &Instance::Emitter(ref e) => Some(e),
            _ => None,
        }
    }).collect();
    assert!(!light_list.is_empty(), "At least one light is required");
    let n = pool.thread_count();
    pool.scoped(|scope| {
        for _ in 0..n {
            let b = &block_queue;
            let ref r = rt;
            let l = &light_list;
            scope.execute(move || {
                thread_work(spp, b, scene, r, l);
            });
        }
    });
}

fn main() {
    let args: Args = Docopt::new(USAGE).and_then(|d| d.decode()).unwrap_or_else(|e| e.exit());
    let n = match args.flag_n {
        Some(n) => n,
        None => num_cpus::get() as u32,
    };
    let out_path = match &args.flag_o {
        &Some(ref f) => {
            let p = PathBuf::from(f);
            // If we're writing to a directory make sure it exists
            if p.extension() == None {
                match std::fs::create_dir(p.as_path()) {
                    Err(e) => {
                        if e.kind() != ErrorKind::AlreadyExists {
                            panic!("Failed to create output directory");
                        }
                    },
                    Ok(_) => {},
                }
            }
            p
        },
        &None => PathBuf::from("./"),
    };

    let (mut scene, mut rt, spp, frame_info) = scene::Scene::load_file(&args.arg_scenefile[..]);
    let image_dim = rt.dimensions();
    println!("Rendering using {} threads\n--------------------", n);
    let mut pool = scoped_threadpool::Pool::new(n);
    let scene_start = clock_ticks::precise_time_s();

    let time_step = frame_info.time / frame_info.frames as f32;
    for i in frame_info.start..frame_info.end + 1 {
        let frame_start = i as f32 * time_step;
        let frame_end = (i as f32 + 1.0) * time_step;
        scene.camera.update_shutter(frame_start, frame_end);
        // TODO: How often to re-build the BVH?
        println!("Frame {}: re-building bvh for {} to {}", i, frame_start, frame_end);
        scene.bvh.rebuild(frame_start, frame_end);
        println!("Frame {}: rendering for {} to {}", i, frame_start, frame_end);
        let start = clock_ticks::precise_time_s();
        render_parallel(&mut rt, &scene, &mut pool, spp);
        let time = clock_ticks::precise_time_s() - start;
        println!("Frame {}: Rendering took {}s", i, time);

        let img = rt.get_render();
        let out_file = match out_path.extension() {
            Some(_) => out_path.clone(),
            None => out_path.join(PathBuf::from(format!("frame{:05}.png", i))),
        };
        match image::save_buffer(&out_file.as_path(), &img[..], image_dim.0 as u32, image_dim.1 as u32, image::RGB(8)) {
            Ok(_) => {},
            Err(e) => println!("Error saving image, {}", e),
        };
        rt.clear();
        println!("Frame {}: Rendered to '{}'\n--------------------", i, out_file.display());
    }
    let time = clock_ticks::precise_time_s() - scene_start;
    println!("Rendering entire sequence took {}s", time);
}

<commit_msg>Consistent capitalization of rendering time info<commit_after>extern crate image;
extern crate rand;
extern crate docopt;
extern crate rustc_serialize;
extern crate num_cpus;
extern crate scoped_threadpool;
extern crate clock_ticks;
extern crate tray_rust;

use std::vec::Vec;
use std::iter;
use std::path::PathBuf;
use std::io::ErrorKind;

use rand::StdRng;
use docopt::Docopt;

use tray_rust::film::{self, ImageSample, Colorf};
use tray_rust::geometry::{Geometry, Instance, Emitter};
use tray_rust::sampler::{self, Sampler};
use tray_rust::scene;
use tray_rust::integrator::Integrator;

static USAGE: &'static str = "
Usage: tray_rust <scenefile> [options]

Options:
  -o <file>     Specify the output file or directory to save the image or frames. Supported formats are
                PNG, JPG and PPM. Default is 'frame<#>.png'.
  -n <number>   Specify the number of threads to use for rendering. Defaults to the number of cores
                on the system.
  -h, --help    Show this message.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_scenefile: String,
    flag_o: Option<String>,
    flag_n: Option<u32>,
}

/// Threads are each sent a sender end of the channel that is
/// read from by the render target thread which then saves the
/// values recieved to the render target
fn thread_work(spp: usize, queue: &sampler::BlockQueue, scene: &scene::Scene,
               target: &film::RenderTarget, light_list: &Vec<&Emitter>) {
    let mut sampler = sampler::LowDiscrepancy::new(queue.block_dim(), spp);
    let mut sample_pos = Vec::with_capacity(sampler.max_spp());
    let mut time_samples: Vec<_> = iter::repeat(0.0).take(sampler.max_spp()).collect();
    let block_dim = queue.block_dim();
    let mut block_samples = Vec::with_capacity(sampler.max_spp() * (block_dim.0 * block_dim.1) as usize);
    let mut rng = match StdRng::new() {
        Ok(r) => r,
        Err(e) => { println!("Failed to get StdRng, {}", e); return }
    };
    // Grab a block from the queue and start working on it, submitting samples
    // to the render target thread after each pixel
    for b in queue.iter() {
        sampler.select_block(b);
        let mut pixel_samples = 0;
        while sampler.has_samples() {
            // Get samples for a pixel and render them
            sampler.get_samples(&mut sample_pos, &mut rng);
            sampler.get_samples_1d(&mut time_samples[..], &mut rng);
            for (s, t) in sample_pos.iter().zip(time_samples.iter()) {
                let mut ray = scene.camera.generate_ray(s, *t);
                if let Some(hit) = scene.intersect(&mut ray) {
                    let c = scene.integrator.illumination(scene, light_list, &ray,
                                                          &hit, &mut sampler, &mut rng).clamp();
                    block_samples.push(ImageSample::new(s.0, s.1, c));
                } else {
                    block_samples.push(ImageSample::new(s.0, s.1, Colorf::black()));
                }
            }
            // If the samples are ok the samples for the next pixel start at the end of the current
            // pixel's samples
            if sampler.report_results(&block_samples[pixel_samples..]) {
                pixel_samples = block_samples.len();
            }
        }
        target.write(&block_samples, sampler.get_region());
        block_samples.clear();
    }
}

/// Render the scene in parallel using `n` threads and write the result to the render target
fn render_parallel(rt: &mut film::RenderTarget, scene: &scene::Scene, pool: &mut scoped_threadpool::Pool, spp: usize){
    let dim = rt.dimensions();
    let block_queue = sampler::BlockQueue::new((dim.0 as u32, dim.1 as u32), (8, 8));
    let light_list: Vec<_> = scene.bvh.into_iter().filter_map(|x| {
        match x {
            &Instance::Emitter(ref e) => Some(e),
            _ => None,
        }
    }).collect();
    assert!(!light_list.is_empty(), "At least one light is required");
    let n = pool.thread_count();
    pool.scoped(|scope| {
        for _ in 0..n {
            let b = &block_queue;
            let ref r = rt;
            let l = &light_list;
            scope.execute(move || {
                thread_work(spp, b, scene, r, l);
            });
        }
    });
}

fn main() {
    let args: Args = Docopt::new(USAGE).and_then(|d| d.decode()).unwrap_or_else(|e| e.exit());
    let n = match args.flag_n {
        Some(n) => n,
        None => num_cpus::get() as u32,
    };
    let out_path = match &args.flag_o {
        &Some(ref f) => {
            let p = PathBuf::from(f);
            // If we're writing to a directory make sure it exists
            if p.extension() == None {
                match std::fs::create_dir(p.as_path()) {
                    Err(e) => {
                        if e.kind() != ErrorKind::AlreadyExists {
                            panic!("Failed to create output directory");
                        }
                    },
                    Ok(_) => {},
                }
            }
            p
        },
        &None => PathBuf::from("./"),
    };

    let (mut scene, mut rt, spp, frame_info) = scene::Scene::load_file(&args.arg_scenefile[..]);
    let image_dim = rt.dimensions();
    println!("Rendering using {} threads\n--------------------", n);
    let mut pool = scoped_threadpool::Pool::new(n);
    let scene_start = clock_ticks::precise_time_s();

    let time_step = frame_info.time / frame_info.frames as f32;
    for i in frame_info.start..frame_info.end + 1 {
        let frame_start = i as f32 * time_step;
        let frame_end = (i as f32 + 1.0) * time_step;
        scene.camera.update_shutter(frame_start, frame_end);
        // TODO: How often to re-build the BVH?
        println!("Frame {}: re-building bvh for {} to {}", i, frame_start, frame_end);
        scene.bvh.rebuild(frame_start, frame_end);
        println!("Frame {}: rendering for {} to {}", i, frame_start, frame_end);
        let start = clock_ticks::precise_time_s();
        render_parallel(&mut rt, &scene, &mut pool, spp);
        let time = clock_ticks::precise_time_s() - start;
        println!("Frame {}: rendering took {}s", i, time);

        let img = rt.get_render();
        let out_file = match out_path.extension() {
            Some(_) => out_path.clone(),
            None => out_path.join(PathBuf::from(format!("frame{:05}.png", i))),
        };
        match image::save_buffer(&out_file.as_path(), &img[..], image_dim.0 as u32, image_dim.1 as u32, image::RGB(8)) {
            Ok(_) => {},
            Err(e) => println!("Error saving image, {}", e),
        };
        rt.clear();
        println!("Frame {}: rendered to '{}'\n--------------------", i, out_file.display());
    }
    let time = clock_ticks::precise_time_s() - scene_start;
    println!("Rendering entire sequence took {}s", time);
}

<|endoftext|>