<commit_before>use crate::hotkey::{Hook, KeyCode};
use crate::{HotkeyConfig, SharedTimer};

pub use crate::hotkey::{Error, Result};

// This enum might be better situated in hotkey_config, but the last method should stay in this file
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum Hotkey {
    Split,
    /// The key to use for resetting the current attempt.
    Reset,
    /// The key to use for undoing the last split.
    Undo,
    /// The key to use for skipping the current split.
    Skip,
    /// The key to use for pausing the current attempt and starting a new
    /// attempt.
    Pause,
    /// The key to use for removing all the pause times from the current time.
    UndoAllPauses,
    /// The key to use for switching to the previous comparison.
    PreviousComparison,
    /// The key to use for switching to the next comparison.
    NextComparison,
    /// The key to use for toggling between the `Real Time` and `Game Time`
    /// timing methods.
    ToggleTimingMethod,
}

impl Hotkey {
    fn set_keycode(self, config: &mut HotkeyConfig, keycode: Option<KeyCode>) {
        match self {
            Hotkey::Split => config.split = keycode,
            Hotkey::Reset => config.reset = keycode,
            Hotkey::Undo => config.undo = keycode,
            Hotkey::Skip => config.skip = keycode,
            Hotkey::Pause => config.pause = keycode,
            Hotkey::UndoAllPauses => config.undo_all_pauses = keycode,
            Hotkey::PreviousComparison => config.previous_comparison = keycode,
            Hotkey::NextComparison => config.next_comparison = keycode,
            Hotkey::ToggleTimingMethod => config.toggle_timing_method = keycode,
        }
    }
    fn get_keycode(self, config: &HotkeyConfig) -> Option<KeyCode> {
        match self {
            Hotkey::Split => config.split,
            Hotkey::Reset => config.reset,
            Hotkey::Undo => config.undo,
            Hotkey::Skip => config.skip,
            Hotkey::Pause => config.pause,
            Hotkey::UndoAllPauses => config.undo_all_pauses,
            Hotkey::PreviousComparison => config.previous_comparison,
            Hotkey::NextComparison => config.next_comparison,
            Hotkey::ToggleTimingMethod => config.toggle_timing_method,
        }
    }
    fn callback(self, timer: SharedTimer) -> Box<dyn FnMut() + Send + 'static>{
        match self {
            Hotkey::Split => Box::new(move || timer.write().split_or_start()),
            Hotkey::Reset => Box::new(move || timer.write().reset(true)),
            Hotkey::Undo => Box::new(move || timer.write().undo_split()),
            Hotkey::Skip => Box::new(move || timer.write().skip_split()),
            Hotkey::Pause => Box::new(move || timer.write().toggle_pause_or_start()),
            Hotkey::UndoAllPauses => Box::new(move || timer.write().undo_all_pauses()),
            Hotkey::PreviousComparison => Box::new(move || timer.write().switch_to_previous_comparison()),
            Hotkey::NextComparison => Box::new(move || timer.write().switch_to_next_comparison()),
            Hotkey::ToggleTimingMethod => Box::new(move || timer.write().toggle_timing_method()),
        }
    }
}

/// With a Hotkey System the runner can use hotkeys on their keyboard to control
/// the Timer. The hotkeys are global, so the application doesn't need to be in
/// focus. The behavior of the hotkeys depends on the platform and is stubbed
/// out on platforms that don't support hotkeys. You can turn off a Hotkey
/// System temporarily. By default the Hotkey System is activated.
pub struct HotkeySystem {
    config: HotkeyConfig,
    hook: Hook,
    timer: SharedTimer,
    is_active: bool,
}

impl HotkeySystem {
    /// Creates a new Hotkey System for a Timer with the default hotkeys.
    pub fn new(timer: SharedTimer) -> Result<Self> {
        Self::with_config(timer, Default::default())
    }

    /// Creates a new Hotkey System for a Timer with a custom configuration for
    /// the hotkeys.
    pub fn with_config(timer: SharedTimer, config: HotkeyConfig) -> Result<Self> {
        let mut hotkey_system = Self {
            config : config,
            hook : Hook::new()?,
            timer : timer,
            is_active : false,
        };
        hotkey_system.activate()?;
        Ok(hotkey_system)
    }
    // This method should never be public, because it might mess up the internal state and we might
    // leak a registered hotkey
    unsafe fn register_raw(&mut self, hotkey: Hotkey) -> Result<()> {
        let inner = self.timer.clone();
        if let Some(keycode) = hotkey.get_keycode(&self.config) {
            println!("Registering Keycode {:?} as {:?}", keycode, hotkey);
            self.hook.register(keycode, hotkey.callback(inner))?;
        }
        Ok(())
    }
    fn register(&mut self, hotkey: Hotkey, keycode: Option<KeyCode>) -> Result<()> {
        hotkey.set_keycode(&mut self.config, keycode);
        unsafe { self.register_raw(hotkey) }
    }
    // This method should never be public, because it might mess up the internal state and we might
    // leak a registered hotkey
    unsafe fn unregister_raw(&mut self, hotkey: Hotkey) -> Result<()> {
        if let Some(keycode) = hotkey.get_keycode(&self.config) {
            println!("Unregistering Keycode {:?} as {:?}", keycode, hotkey);
            self.hook.unregister(keycode)?;
        }
        Ok(())
    }
    fn unregister(&mut self, hotkey: Hotkey) -> Result<()> {
        hotkey.set_keycode(&mut self.config, None);
        unsafe { self.unregister_raw(hotkey) }
    }
    fn set_hotkey(&mut self, hotkey: Hotkey, keycode: Option<KeyCode>) -> Result<()> {
        // FixMe: We do not check whether the keycode is already in use
        if hotkey.get_keycode(&self.config) == keycode {
            return Ok(());
        }
        if self.is_active {
            self.unregister(hotkey)?;
            self.register(hotkey, keycode)?;
        }
        Ok(())
    }

    /// Sets the key to use for splitting and starting a new attempt.
    pub fn set_split(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Split, hotkey)
    }

    /// Sets the key to use for resetting the current attempt.
    pub fn set_reset(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Reset, hotkey)
    }

    /// Sets the key to use for pausing the current attempt and starting a new
    /// attempt.
    pub fn set_pause(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Pause, hotkey)
    }

    /// Sets the key to use for skipping the current split.
    pub fn set_skip(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Skip, hotkey)
    }

    /// Sets the key to use for undoing the last split.
    pub fn set_undo(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Undo, hotkey)
    }

    /// Sets the key to use for switching to the previous comparison.
    pub fn set_previous_comparison(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::PreviousComparison, hotkey)
    }

    /// Sets the key to use for switching to the next comparison.
    pub fn set_next_comparison(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::NextComparison, hotkey)
    }

    /// Sets the key to use for removing all the pause times from the current
    /// time.
    pub fn set_undo_all_pauses(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::UndoAllPauses, hotkey)
    }

    /// Sets the key to use for toggling between the `Real Time` and `Game Time`
    /// timing methods.
    pub fn set_toggle_timing_method(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::ToggleTimingMethod, hotkey)
    }

    /// Deactivates the Hotkey System. No hotkeys will go through until it gets
    /// activated again. If it's already deactivated, nothing happens.
    pub fn deactivate(&mut self) -> Result<()> {
        if self.is_active {
            unsafe {
                self.unregister_raw(Hotkey::Split)?;
                self.unregister_raw(Hotkey::Reset)?;
                self.unregister_raw(Hotkey::Undo)?;
                self.unregister_raw(Hotkey::Skip)?;
                self.unregister_raw(Hotkey::Pause)?;
                self.unregister_raw(Hotkey::UndoAllPauses)?;
                self.unregister_raw(Hotkey::PreviousComparison)?;
                self.unregister_raw(Hotkey::NextComparison)?;
                self.unregister_raw(Hotkey::ToggleTimingMethod)?;
            }
        }
        self.is_active = false;
        Ok(())
    }

    /// Activates a previously deactivated Hotkey System. If it's already
    /// active, nothing happens.
    pub fn activate(&mut self) -> Result<()> {
        if !self.is_active {
            unsafe {
                self.register_raw(Hotkey::Split)?;
                self.register_raw(Hotkey::Reset)?;
                self.register_raw(Hotkey::Undo)?;
                self.register_raw(Hotkey::Skip)?;
                self.register_raw(Hotkey::Pause)?;
                self.register_raw(Hotkey::UndoAllPauses)?;
                self.register_raw(Hotkey::PreviousComparison)?;
                self.register_raw(Hotkey::NextComparison)?;
                self.register_raw(Hotkey::ToggleTimingMethod)?;
            }
        }
        self.is_active = true;
        Ok(())
    }

    /// Returns true if the Hotkey System is active, false otherwise.
    pub fn is_active(&self) -> bool {
        self.is_active
    }

    /// Returns the hotkey configuration currently in use by the Hotkey System.
    pub fn config(&self) -> HotkeyConfig {
        self.config
    }

    /// Applies a new hotkey configuration to the Hotkey System. Each hotkey is
    /// changed to the one specified in the configuration. This operation may
    /// fail if you provide a hotkey configuration where a hotkey is used for
    /// multiple operations.
    pub fn set_config(&mut self, config: HotkeyConfig) -> Result<()> {
        self.set_split(config.split)?;
        self.set_reset(config.reset)?;
        self.set_undo(config.undo)?;
        self.set_skip(config.skip)?;
        self.set_pause(config.pause)?;
        self.set_previous_comparison(config.previous_comparison)?;
        self.set_next_comparison(config.next_comparison)?;
        self.set_undo_all_pauses(config.undo_all_pauses)?;
        self.set_toggle_timing_method(config.toggle_timing_method)?;

        Ok(())
    }
}
<commit_msg>cargo fmt<commit_after>use crate::hotkey::{Hook, KeyCode};
use crate::{HotkeyConfig, SharedTimer};

pub use crate::hotkey::{Error, Result};

// This enum might be better situated in hotkey_config, but the last method should stay in this file
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum Hotkey {
    Split,
    /// The key to use for resetting the current attempt.
    Reset,
    /// The key to use for undoing the last split.
    Undo,
    /// The key to use for skipping the current split.
    Skip,
    /// The key to use for pausing the current attempt and starting a new
    /// attempt.
    Pause,
    /// The key to use for removing all the pause times from the current time.
    UndoAllPauses,
    /// The key to use for switching to the previous comparison.
    PreviousComparison,
    /// The key to use for switching to the next comparison.
    NextComparison,
    /// The key to use for toggling between the `Real Time` and `Game Time`
    /// timing methods.
    ToggleTimingMethod,
}

impl Hotkey {
    fn set_keycode(self, config: &mut HotkeyConfig, keycode: Option<KeyCode>) {
        match self {
            Hotkey::Split => config.split = keycode,
            Hotkey::Reset => config.reset = keycode,
            Hotkey::Undo => config.undo = keycode,
            Hotkey::Skip => config.skip = keycode,
            Hotkey::Pause => config.pause = keycode,
            Hotkey::UndoAllPauses => config.undo_all_pauses = keycode,
            Hotkey::PreviousComparison => config.previous_comparison = keycode,
            Hotkey::NextComparison => config.next_comparison = keycode,
            Hotkey::ToggleTimingMethod => config.toggle_timing_method = keycode,
        }
    }
    fn get_keycode(self, config: &HotkeyConfig) -> Option<KeyCode> {
        match self {
            Hotkey::Split => config.split,
            Hotkey::Reset => config.reset,
            Hotkey::Undo => config.undo,
            Hotkey::Skip => config.skip,
            Hotkey::Pause => config.pause,
            Hotkey::UndoAllPauses => config.undo_all_pauses,
            Hotkey::PreviousComparison => config.previous_comparison,
            Hotkey::NextComparison => config.next_comparison,
            Hotkey::ToggleTimingMethod => config.toggle_timing_method,
        }
    }
    fn callback(self, timer: SharedTimer) -> Box<dyn FnMut() + Send + 'static> {
        match self {
            Hotkey::Split => Box::new(move || timer.write().split_or_start()),
            Hotkey::Reset => Box::new(move || timer.write().reset(true)),
            Hotkey::Undo => Box::new(move || timer.write().undo_split()),
            Hotkey::Skip => Box::new(move || timer.write().skip_split()),
            Hotkey::Pause => Box::new(move || timer.write().toggle_pause_or_start()),
            Hotkey::UndoAllPauses => Box::new(move || timer.write().undo_all_pauses()),
            Hotkey::PreviousComparison => {
                Box::new(move || timer.write().switch_to_previous_comparison())
            }
            Hotkey::NextComparison => Box::new(move || timer.write().switch_to_next_comparison()),
            Hotkey::ToggleTimingMethod => Box::new(move || timer.write().toggle_timing_method()),
        }
    }
}

/// With a Hotkey System the runner can use hotkeys on their keyboard to control
/// the Timer. The hotkeys are global, so the application doesn't need to be in
/// focus. The behavior of the hotkeys depends on the platform and is stubbed
/// out on platforms that don't support hotkeys. You can turn off a Hotkey
/// System temporarily. By default the Hotkey System is activated.
pub struct HotkeySystem {
    config: HotkeyConfig,
    hook: Hook,
    timer: SharedTimer,
    is_active: bool,
}

impl HotkeySystem {
    /// Creates a new Hotkey System for a Timer with the default hotkeys.
    pub fn new(timer: SharedTimer) -> Result<Self> {
        Self::with_config(timer, Default::default())
    }

    /// Creates a new Hotkey System for a Timer with a custom configuration for
    /// the hotkeys.
    pub fn with_config(timer: SharedTimer, config: HotkeyConfig) -> Result<Self> {
        let mut hotkey_system = Self {
            config: config,
            hook: Hook::new()?,
            timer: timer,
            is_active: false,
        };
        hotkey_system.activate()?;
        Ok(hotkey_system)
    }
    // This method should never be public, because it might mess up the internal state and we might
    // leak a registered hotkey
    unsafe fn register_raw(&mut self, hotkey: Hotkey) -> Result<()> {
        let inner = self.timer.clone();
        if let Some(keycode) = hotkey.get_keycode(&self.config) {
            println!("Registering Keycode {:?} as {:?}", keycode, hotkey);
            self.hook.register(keycode, hotkey.callback(inner))?;
        }
        Ok(())
    }
    fn register(&mut self, hotkey: Hotkey, keycode: Option<KeyCode>) -> Result<()> {
        hotkey.set_keycode(&mut self.config, keycode);
        unsafe { self.register_raw(hotkey) }
    }
    // This method should never be public, because it might mess up the internal state and we might
    // leak a registered hotkey
    unsafe fn unregister_raw(&mut self, hotkey: Hotkey) -> Result<()> {
        if let Some(keycode) = hotkey.get_keycode(&self.config) {
            println!("Unregistering Keycode {:?} as {:?}", keycode, hotkey);
            self.hook.unregister(keycode)?;
        }
        Ok(())
    }
    fn unregister(&mut self, hotkey: Hotkey) -> Result<()> {
        hotkey.set_keycode(&mut self.config, None);
        unsafe { self.unregister_raw(hotkey) }
    }
    fn set_hotkey(&mut self, hotkey: Hotkey, keycode: Option<KeyCode>) -> Result<()> {
        // FixMe: We do not check whether the keycode is already in use
        if hotkey.get_keycode(&self.config) == keycode {
            return Ok(());
        }
        if self.is_active {
            self.unregister(hotkey)?;
            self.register(hotkey, keycode)?;
        }
        Ok(())
    }

    /// Sets the key to use for splitting and starting a new attempt.
    pub fn set_split(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Split, hotkey)
    }

    /// Sets the key to use for resetting the current attempt.
    pub fn set_reset(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Reset, hotkey)
    }

    /// Sets the key to use for pausing the current attempt and starting a new
    /// attempt.
    pub fn set_pause(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Pause, hotkey)
    }

    /// Sets the key to use for skipping the current split.
    pub fn set_skip(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Skip, hotkey)
    }

    /// Sets the key to use for undoing the last split.
    pub fn set_undo(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::Undo, hotkey)
    }

    /// Sets the key to use for switching to the previous comparison.
    pub fn set_previous_comparison(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::PreviousComparison, hotkey)
    }

    /// Sets the key to use for switching to the next comparison.
    pub fn set_next_comparison(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::NextComparison, hotkey)
    }

    /// Sets the key to use for removing all the pause times from the current
    /// time.
    pub fn set_undo_all_pauses(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::UndoAllPauses, hotkey)
    }

    /// Sets the key to use for toggling between the `Real Time` and `Game Time`
    /// timing methods.
    pub fn set_toggle_timing_method(&mut self, hotkey: Option<KeyCode>) -> Result<()> {
        self.set_hotkey(Hotkey::ToggleTimingMethod, hotkey)
    }

    /// Deactivates the Hotkey System. No hotkeys will go through until it gets
    /// activated again. If it's already deactivated, nothing happens.
    pub fn deactivate(&mut self) -> Result<()> {
        if self.is_active {
            unsafe {
                self.unregister_raw(Hotkey::Split)?;
                self.unregister_raw(Hotkey::Reset)?;
                self.unregister_raw(Hotkey::Undo)?;
                self.unregister_raw(Hotkey::Skip)?;
                self.unregister_raw(Hotkey::Pause)?;
                self.unregister_raw(Hotkey::UndoAllPauses)?;
                self.unregister_raw(Hotkey::PreviousComparison)?;
                self.unregister_raw(Hotkey::NextComparison)?;
                self.unregister_raw(Hotkey::ToggleTimingMethod)?;
            }
        }
        self.is_active = false;
        Ok(())
    }

    /// Activates a previously deactivated Hotkey System. If it's already
    /// active, nothing happens.
    pub fn activate(&mut self) -> Result<()> {
        if !self.is_active {
            unsafe {
                self.register_raw(Hotkey::Split)?;
                self.register_raw(Hotkey::Reset)?;
                self.register_raw(Hotkey::Undo)?;
                self.register_raw(Hotkey::Skip)?;
                self.register_raw(Hotkey::Pause)?;
                self.register_raw(Hotkey::UndoAllPauses)?;
                self.register_raw(Hotkey::PreviousComparison)?;
                self.register_raw(Hotkey::NextComparison)?;
                self.register_raw(Hotkey::ToggleTimingMethod)?;
            }
        }
        self.is_active = true;
        Ok(())
    }

    /// Returns true if the Hotkey System is active, false otherwise.
    pub fn is_active(&self) -> bool {
        self.is_active
    }

    /// Returns the hotkey configuration currently in use by the Hotkey System.
    pub fn config(&self) -> HotkeyConfig {
        self.config
    }

    /// Applies a new hotkey configuration to the Hotkey System. Each hotkey is
    /// changed to the one specified in the configuration. This operation may
    /// fail if you provide a hotkey configuration where a hotkey is used for
    /// multiple operations.
    pub fn set_config(&mut self, config: HotkeyConfig) -> Result<()> {
        self.set_split(config.split)?;
        self.set_reset(config.reset)?;
        self.set_undo(config.undo)?;
        self.set_skip(config.skip)?;
        self.set_pause(config.pause)?;
        self.set_previous_comparison(config.previous_comparison)?;
        self.set_next_comparison(config.next_comparison)?;
        self.set_undo_all_pauses(config.undo_all_pauses)?;
        self.set_toggle_timing_method(config.toggle_timing_method)?;

        Ok(())
    }
}
<|endoftext|>


<commit_before>use libc;
use std::ptr;
use std::ffi::CStr;
use constants::*;
use csh;
use ffi;
use instruction::{Insn, Instructions};

/// An instance of the capstone disassembler
pub struct Capstone {
    csh: csh, // Opaque handle to cs_engine
}

pub type CsResult<T> = Result<T, CsErr>;

impl Capstone {
    /// Create a new instance of the decompiler
    ///
    /// ```
    /// use capstone::Capstone;
    /// use capstone::constants::*;
    /// let cs = Capstone::new(CsArch::ARCH_X86, CsMode::MODE_64);
    /// assert!(cs.is_ok());
    /// ```
    pub fn new(arch: CsArch, mode: CsMode) -> CsResult<Capstone> {
        let mut handle = 0;
        let err = unsafe { ffi::cs_open(arch, mode, &mut handle) };

        if CsErr::CS_ERR_OK == err {
            Ok(Capstone { csh: handle })
        } else {
            Err(err)
        }
    }

    /// Disassemble a &[u8] full of instructions
    ///
    /// Pass count = 0 to disassemble all instructions in the buffer
    pub fn disasm(&self, code: &[u8], addr: u64, count: isize) -> CsResult<Instructions> {
        let mut ptr: *const Insn = ptr::null();
        let insn_count = unsafe {
            ffi::cs_disasm(self.csh,
                      code.as_ptr(),
                      code.len() as libc::size_t,
                      addr,
                      count as libc::size_t,
                      &mut ptr)
        };
        if insn_count == 0 {
            return self.get_error_result();
        }

        Ok(unsafe { Instructions::from_raw_parts(ptr, insn_count as isize) })
    }

    /// Get error CsResult based on current errno
    fn get_error_result(&self) ->CsResult<Instructions> {
        Err(unsafe { ffi::cs_errno(self.csh) })
    }

    /// Convert a reg_id to a String naming the register
    pub fn reg_name(&self, reg_id: u64) -> Option<String> {
        let reg_name = unsafe {
            let _reg_name = ffi::cs_reg_name(self.csh, reg_id as libc::c_uint);
            if _reg_name == ptr::null() {
                return None
            }

            CStr::from_ptr(_reg_name).to_string_lossy().into_owned()
        };

        Some(reg_name)
    }

    /// Convert an instruction_id to a String naming the instruction
    pub fn insn_name(&self, insn_id: u64) -> Option<String> {
        let insn_name = unsafe {
            let _insn_name = ffi::cs_insn_name(self.csh, insn_id as libc::c_uint);
            if _insn_name == ptr::null() {
                return None
            }

            CStr::from_ptr(_insn_name).to_string_lossy().into_owned()
        };

        Some(insn_name)
    }
}

impl Drop for Capstone {
    fn drop(&mut self) {
        unsafe { ffi::cs_close(&mut self.csh) };
    }
}

/// Return tuple (major, minor) indicating the version of the capstone C library
pub fn lib_version() -> (u32, u32) {
    let mut major: libc::c_int = 0;
    let mut minor: libc::c_int = 0;
    let major_ptr: *mut libc::c_int = &mut major;
    let minor_ptr: *mut libc::c_int = &mut minor;

    let _ = unsafe { ffi::cs_version(major_ptr, minor_ptr) };

    (major as u32, minor as u32)
}
<commit_msg>Fix comment<commit_after>use libc;
use std::ptr;
use std::ffi::CStr;
use constants::*;
use csh;
use ffi;
use instruction::{Insn, Instructions};

/// An instance of the capstone disassembler
pub struct Capstone {
    csh: csh, // Opaque handle to cs_engine
}

pub type CsResult<T> = Result<T, CsErr>;

impl Capstone {
    /// Create a new instance of the disassembler
    ///
    /// ```
    /// use capstone::Capstone;
    /// use capstone::constants::*;
    /// let cs = Capstone::new(CsArch::ARCH_X86, CsMode::MODE_64);
    /// assert!(cs.is_ok());
    /// ```
    pub fn new(arch: CsArch, mode: CsMode) -> CsResult<Capstone> {
        let mut handle = 0;
        let err = unsafe { ffi::cs_open(arch, mode, &mut handle) };

        if CsErr::CS_ERR_OK == err {
            Ok(Capstone { csh: handle })
        } else {
            Err(err)
        }
    }

    /// Disassemble a &[u8] full of instructions
    ///
    /// Pass count = 0 to disassemble all instructions in the buffer
    pub fn disasm(&self, code: &[u8], addr: u64, count: isize) -> CsResult<Instructions> {
        let mut ptr: *const Insn = ptr::null();
        let insn_count = unsafe {
            ffi::cs_disasm(self.csh,
                      code.as_ptr(),
                      code.len() as libc::size_t,
                      addr,
                      count as libc::size_t,
                      &mut ptr)
        };
        if insn_count == 0 {
            return self.get_error_result();
        }

        Ok(unsafe { Instructions::from_raw_parts(ptr, insn_count as isize) })
    }

    /// Get error CsResult based on current errno
    fn get_error_result(&self) ->CsResult<Instructions> {
        Err(unsafe { ffi::cs_errno(self.csh) })
    }

    /// Convert a reg_id to a String naming the register
    pub fn reg_name(&self, reg_id: u64) -> Option<String> {
        let reg_name = unsafe {
            let _reg_name = ffi::cs_reg_name(self.csh, reg_id as libc::c_uint);
            if _reg_name == ptr::null() {
                return None
            }

            CStr::from_ptr(_reg_name).to_string_lossy().into_owned()
        };

        Some(reg_name)
    }

    /// Convert an instruction_id to a String naming the instruction
    pub fn insn_name(&self, insn_id: u64) -> Option<String> {
        let insn_name = unsafe {
            let _insn_name = ffi::cs_insn_name(self.csh, insn_id as libc::c_uint);
            if _insn_name == ptr::null() {
                return None
            }

            CStr::from_ptr(_insn_name).to_string_lossy().into_owned()
        };

        Some(insn_name)
    }
}

impl Drop for Capstone {
    fn drop(&mut self) {
        unsafe { ffi::cs_close(&mut self.csh) };
    }
}

/// Return tuple (major, minor) indicating the version of the capstone C library
pub fn lib_version() -> (u32, u32) {
    let mut major: libc::c_int = 0;
    let mut minor: libc::c_int = 0;
    let major_ptr: *mut libc::c_int = &mut major;
    let minor_ptr: *mut libc::c_int = &mut minor;

    let _ = unsafe { ffi::cs_version(major_ptr, minor_ptr) };

    (major as u32, minor as u32)
}
<|endoftext|>


<commit_before>use context::timeout;
use device::pic;
use device::serial::{COM1, COM2};
use time;

unsafe fn trigger(irq: u8) {
    extern {
        fn irq_trigger(irq: u8);
    }

    if irq < 16 {
        if irq >= 8 {
            pic::SLAVE.mask_set(irq - 8);
            pic::MASTER.ack();
            pic::SLAVE.ack();
        } else {
            pic::MASTER.mask_set(irq);
            pic::MASTER.ack();
        }
    }

    irq_trigger(irq);
}

pub unsafe fn acknowledge(irq: usize) {
    if irq < 16 {
        if irq >= 8 {
            pic::SLAVE.mask_clear(irq as u8 - 8);
        } else {
            pic::MASTER.mask_clear(irq as u8);
        }
    }
}

interrupt!(pit, {
    // Saves CPU time by not sending IRQ event irq_trigger(0);

    const PIT_RATE: u64 = 2250286;

    {
        let mut offset = time::OFFSET.lock();
        let sum = offset.1 + PIT_RATE;
        offset.1 = sum % 1000000000;
        offset.0 += sum / 1000000000;
    }

    pic::MASTER.ack();

    // Any better way of doing this?
    timeout::trigger();
});

interrupt!(keyboard, {
    trigger(1);
});

interrupt!(cascade, {
    // No need to do any operations on cascade
    pic::MASTER.ack();
});

interrupt!(com2, {
    COM2.lock().receive();
    pic::MASTER.ack();
});

interrupt!(com1, {
    COM1.lock().receive();
    pic::MASTER.ack();
});

interrupt!(lpt2, {
    trigger(5);
});

interrupt!(floppy, {
    trigger(6);
});

interrupt!(lpt1, {
    trigger(7);
});

interrupt!(rtc, {
    trigger(8);
});

interrupt!(pci1, {
    trigger(9);
});

interrupt!(pci2, {
    trigger(10);
});

interrupt!(pci3, {
    trigger(11);
});

interrupt!(mouse, {
    trigger(12);
});

interrupt!(fpu, {
    trigger(13);
});

interrupt!(ata1, {
    trigger(14);
});

interrupt!(ata2, {
    trigger(15);
});
<commit_msg>Fix PIT Interrupt Not Context Switching [irq]<commit_after>use context::timeout;
use device::pic;
use device::serial::{COM1, COM2};
use core::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};

use time;
use context;

static PIT_TICKS: AtomicUsize = ATOMIC_USIZE_INIT;

unsafe fn trigger(irq: u8) {
    extern {
        fn irq_trigger(irq: u8);
    }

    if irq < 16 {
        if irq >= 8 {
            pic::SLAVE.mask_set(irq - 8);
            pic::MASTER.ack();
            pic::SLAVE.ack();
        } else {
            pic::MASTER.mask_set(irq);
            pic::MASTER.ack();
        }
    }

    irq_trigger(irq);
}

pub unsafe fn acknowledge(irq: usize) {
    if irq < 16 {
        if irq >= 8 {
            pic::SLAVE.mask_clear(irq as u8 - 8);
        } else {
            pic::MASTER.mask_clear(irq as u8);
        }
    }
}

interrupt!(pit, {
    // Saves CPU time by not sending IRQ event irq_trigger(0);

    const PIT_RATE: u64 = 2250286;

    {
        let mut offset = time::OFFSET.lock();
        let sum = offset.1 + PIT_RATE;
        offset.1 = sum % 1000000000;
        offset.0 += sum / 1000000000;
    }

    pic::MASTER.ack();

    if PIT_TICKS.fetch_add(1, Ordering::SeqCst) % 10 == 0 {
        context::switch();
    }

    // Any better way of doing this?
    timeout::trigger();
});

interrupt!(keyboard, {
    trigger(1);
});

interrupt!(cascade, {
    // No need to do any operations on cascade
    pic::MASTER.ack();
});

interrupt!(com2, {
    COM2.lock().receive();
    pic::MASTER.ack();
});

interrupt!(com1, {
    COM1.lock().receive();
    pic::MASTER.ack();
});

interrupt!(lpt2, {
    trigger(5);
});

interrupt!(floppy, {
    trigger(6);
});

interrupt!(lpt1, {
    trigger(7);
});

interrupt!(rtc, {
    trigger(8);
});

interrupt!(pci1, {
    trigger(9);
});

interrupt!(pci2, {
    trigger(10);
});

interrupt!(pci3, {
    trigger(11);
});

interrupt!(mouse, {
    trigger(12);
});

interrupt!(fpu, {
    trigger(13);
});

interrupt!(ata1, {
    trigger(14);
});

interrupt!(ata2, {
    trigger(15);
});
<|endoftext|>


<commit_before>use std::cell::UnsafeCell;
use std::io;
use std::rc::Rc;

use io::{AsyncRead, AsyncWrite};
use nio;

pub struct ReadHalf<R> {
    inner: Rc<UnsafeCell<R>>,
}

pub struct WriteHalf<W> {
    inner: Rc<UnsafeCell<W>>,
}

#[inline]
pub fn split<T>(t: T) -> (ReadHalf<T>, WriteHalf<T>)
where
    T: AsyncRead + AsyncWrite,
{
    let io = Rc::new(UnsafeCell::new(t));
    (ReadHalf { inner: io.clone() }, WriteHalf { inner: io })
}

impl<R> ReadHalf<R> {
    #[inline]
    fn get_mut(&self) -> &mut R {
        unsafe { &mut *(&self.inner).get() }
    }
}

impl<W> WriteHalf<W> {
    #[inline]
    fn get_mut(&self) -> &mut W {
        unsafe { &mut *(&self.inner).get() }
    }
}

impl<R: AsyncRead> io::Read for ReadHalf<R> {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.get_mut().read(buf)
    }
}

impl<R: AsyncRead> nio::ReadV for ReadHalf<R> {
    #[inline]
    fn readv(&mut self, iovs: &[nio::IoVec]) -> io::Result<usize> {
        self.get_mut().readv(iovs)
    }
}

impl<R: AsyncRead> AsyncRead for ReadHalf<R> {
    #[inline]
    fn need_read(&mut self) -> io::Result<()> {
        self.get_mut().need_read()
    }

    #[inline]
    fn no_need_read(&mut self) -> io::Result<()> {
        self.get_mut().no_need_read()
    }

    #[inline]
    fn is_readable(&self) -> bool {
        self.get_mut().is_readable()
    }
}

impl<W: AsyncWrite> io::Write for WriteHalf<W> {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.get_mut().write(buf)
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        self.get_mut().flush()
    }
}

impl<W: AsyncWrite> nio::WriteV for WriteHalf<W> {
    #[inline]
    fn writev(&mut self, iovs: &[nio::IoVec]) -> io::Result<usize> {
        self.get_mut().writev(iovs)
    }
}

impl<W: AsyncWrite> AsyncWrite for WriteHalf<W> {
    #[inline]
    fn need_write(&mut self) -> io::Result<()> {
        self.get_mut().need_write()
    }

    #[inline]
    fn no_need_write(&mut self) -> io::Result<()> {
        self.get_mut().no_need_write()
    }

    #[inline]
    fn is_writable(&self) -> bool {
        self.get_mut().is_writable()
    }
}
<commit_msg>Clone for ReadHalf and WriteHalf<commit_after>use std::cell::UnsafeCell;
use std::io;
use std::rc::Rc;

use io::{AsyncRead, AsyncWrite};
use nio;

#[derive(Clone)]
pub struct ReadHalf<R> {
    inner: Rc<UnsafeCell<R>>,
}

#[derive(Clone)]
pub struct WriteHalf<W> {
    inner: Rc<UnsafeCell<W>>,
}

#[inline]
pub fn split<T>(t: T) -> (ReadHalf<T>, WriteHalf<T>)
where
    T: AsyncRead + AsyncWrite,
{
    let io = Rc::new(UnsafeCell::new(t));
    (ReadHalf { inner: io.clone() }, WriteHalf { inner: io })
}

impl<R> ReadHalf<R> {
    #[inline]
    fn get_mut(&self) -> &mut R {
        unsafe { &mut *(&self.inner).get() }
    }
}

impl<W> WriteHalf<W> {
    #[inline]
    fn get_mut(&self) -> &mut W {
        unsafe { &mut *(&self.inner).get() }
    }
}

impl<R: AsyncRead> io::Read for ReadHalf<R> {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.get_mut().read(buf)
    }
}

impl<R: AsyncRead> nio::ReadV for ReadHalf<R> {
    #[inline]
    fn readv(&mut self, iovs: &[nio::IoVec]) -> io::Result<usize> {
        self.get_mut().readv(iovs)
    }
}

impl<R: AsyncRead> AsyncRead for ReadHalf<R> {
    #[inline]
    fn need_read(&mut self) -> io::Result<()> {
        self.get_mut().need_read()
    }

    #[inline]
    fn no_need_read(&mut self) -> io::Result<()> {
        self.get_mut().no_need_read()
    }

    #[inline]
    fn is_readable(&self) -> bool {
        self.get_mut().is_readable()
    }
}

impl<W: AsyncWrite> io::Write for WriteHalf<W> {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.get_mut().write(buf)
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        self.get_mut().flush()
    }
}

impl<W: AsyncWrite> nio::WriteV for WriteHalf<W> {
    #[inline]
    fn writev(&mut self, iovs: &[nio::IoVec]) -> io::Result<usize> {
        self.get_mut().writev(iovs)
    }
}

impl<W: AsyncWrite> AsyncWrite for WriteHalf<W> {
    #[inline]
    fn need_write(&mut self) -> io::Result<()> {
        self.get_mut().need_write()
    }

    #[inline]
    fn no_need_write(&mut self) -> io::Result<()> {
        self.get_mut().no_need_write()
    }

    #[inline]
    fn is_writable(&self) -> bool {
        self.get_mut().is_writable()
    }
}
<|endoftext|>


<commit_before>use std::num::FromPrimitive;
use std::char;

pub enum Key {
    Unknown      = 0,
    Tab          = 9,
    Enter        = 13,
    CtrlQ        = 17,
    CtrlR        = 18,
    CtrlS        = 19,
    Esc          = 27,
    Space        = 32,
    Exclaim      = 33,
    Hash         = 35,
    Dollar       = 36,
    Percent      = 37,
    Ampersand    = 38,
    Quote        = 39,
    LeftParen    = 40,
    RightParen   = 41,
    Asterisk     = 42,
    Plus         = 43,
    Comma        = 44,
    Minus        = 45,
    Period       = 46,
    Slash        = 47,
    D0           = 48,
    D1           = 49,
    D2           = 50,
    D3           = 51,
    D4           = 52,
    D5           = 53,
    D6           = 54,
    D7           = 55,
    D8           = 56,
    D9           = 57,
    Colon        = 58,
    Semicolon    = 59,
    Less         = 60,
    Equals       = 61,
    Greater      = 62,
    Question     = 63,
    At           = 64,
    LeftBracket  = 91,
    Backslash    = 92,
    RightBracket = 93,
    Caret        = 94,
    Underscore   = 95,
    Backquote    = 96,
    A            = 97,
    B            = 98,
    C            = 99,
    D            = 100,
    E            = 101,
    F            = 102,
    G            = 103,
    H            = 104,
    I            = 105,
    J            = 106,
    K            = 107,
    L            = 108,
    M            = 109,
    N            = 110,
    O            = 111,
    P            = 112,
    Q            = 113,
    R            = 114,
    S            = 115,
    T            = 116,
    U            = 117,
    V            = 118,
    W            = 119,
    X            = 120,
    Y            = 121,
    Z            = 122,
    LeftBrace    = 123,
    Pipe         = 124,
    RightBrace   = 125,
    Tilde        = 126,
    Backspace    = 127,
    Right        = 65514,
    Left         = 65515,
    Down         = 65516,
    Up           = 65517,
    Delete       = 65522,
}

impl Key {
    #[inline(always)]
    pub fn code(&self) -> u32 {
        *self as u32
    }

    pub fn get_char(&self) -> Option<char> {
        char::from_u32(self.code())
    }
}

impl FromPrimitive for Key {
    fn from_u64(n: u64) -> Option<Key> {
        match n {
            0     => Some(Key::Unknown),
            9     => Some(Key::Tab),
            13    => Some(Key::Enter),
            17    => Some(Key::CtrlQ),
            18    => Some(Key::CtrlR),
            19    => Some(Key::CtrlS),
            27    => Some(Key::Esc),
            32    => Some(Key::Space),
            33    => Some(Key::Exclaim),
            35    => Some(Key::Hash),
            36    => Some(Key::Dollar),
            37    => Some(Key::Percent),
            38    => Some(Key::Ampersand),
            39    => Some(Key::Quote),
            40    => Some(Key::LeftParen),
            41    => Some(Key::RightParen),
            42    => Some(Key::Asterisk),
            43    => Some(Key::Plus),
            44    => Some(Key::Comma),
            45    => Some(Key::Minus),
            46    => Some(Key::Period),
            47    => Some(Key::Slash),
            48    => Some(Key::D0),
            49    => Some(Key::D1),
            50    => Some(Key::D2),
            51    => Some(Key::D3),
            52    => Some(Key::D4),
            53    => Some(Key::D5),
            54    => Some(Key::D6),
            55    => Some(Key::D7),
            56    => Some(Key::D8),
            57    => Some(Key::D9),
            58    => Some(Key::Colon),
            59    => Some(Key::Semicolon),
            60    => Some(Key::Less),
            61    => Some(Key::Equals),
            62    => Some(Key::Greater),
            63    => Some(Key::Question),
            64    => Some(Key::At),
            91    => Some(Key::LeftBracket),
            92    => Some(Key::Backslash),
            93    => Some(Key::RightBracket),
            94    => Some(Key::Caret),
            95    => Some(Key::Underscore),
            96    => Some(Key::Backquote),
            97    => Some(Key::A),
            98    => Some(Key::B),
            99    => Some(Key::C),
            100   => Some(Key::D),
            101   => Some(Key::E),
            102   => Some(Key::F),
            103   => Some(Key::G),
            104   => Some(Key::H),
            105   => Some(Key::I),
            106   => Some(Key::J),
            107   => Some(Key::K),
            108   => Some(Key::L),
            109   => Some(Key::M),
            110   => Some(Key::N),
            111   => Some(Key::O),
            112   => Some(Key::P),
            113   => Some(Key::Q),
            114   => Some(Key::R),
            115   => Some(Key::S),
            116   => Some(Key::T),
            117   => Some(Key::U),
            118   => Some(Key::V),
            119   => Some(Key::W),
            120   => Some(Key::X),
            121   => Some(Key::Y),
            122   => Some(Key::Z),
            123   => Some(Key::LeftBrace),
            124   => Some(Key::Pipe),
            125   => Some(Key::RightBrace),
            126   => Some(Key::Tilde),
            127   => Some(Key::Backspace),
            65514 => Some(Key::Right),
            65515 => Some(Key::Left),
            65516 => Some(Key::Down),
            65517 => Some(Key::Up),
            65522 => Some(Key::Delete),
            _     => None
        }
    }

    #[inline(always)]
    fn from_i64(n: i64) -> Option<Key> {
        FromPrimitive::from_u64(n as u64)
    }

    #[inline(always)]
    fn from_int(n: int) -> Option<Key> {
        FromPrimitive::from_u64(n as u64)
    }

    #[inline(always)]
    fn from_u16(n: u16) -> Option<Key> {
        FromPrimitive::from_u64(n as u64)
    }

    #[inline(always)]
    fn from_u32(n: u32) -> Option<Key> {
        FromPrimitive::from_u64(n as u64)
    }
}
<commit_msg>Added copy trait to Key - fixes #24<commit_after>use std::num::FromPrimitive;
use std::char;

#[deriving(Copy)]
pub enum Key {
    Unknown      = 0,
    Tab          = 9,
    Enter        = 13,
    CtrlQ        = 17,
    CtrlR        = 18,
    CtrlS        = 19,
    Esc          = 27,
    Space        = 32,
    Exclaim      = 33,
    Hash         = 35,
    Dollar       = 36,
    Percent      = 37,
    Ampersand    = 38,
    Quote        = 39,
    LeftParen    = 40,
    RightParen   = 41,
    Asterisk     = 42,
    Plus         = 43,
    Comma        = 44,
    Minus        = 45,
    Period       = 46,
    Slash        = 47,
    D0           = 48,
    D1           = 49,
    D2           = 50,
    D3           = 51,
    D4           = 52,
    D5           = 53,
    D6           = 54,
    D7           = 55,
    D8           = 56,
    D9           = 57,
    Colon        = 58,
    Semicolon    = 59,
    Less         = 60,
    Equals       = 61,
    Greater      = 62,
    Question     = 63,
    At           = 64,
    LeftBracket  = 91,
    Backslash    = 92,
    RightBracket = 93,
    Caret        = 94,
    Underscore   = 95,
    Backquote    = 96,
    A            = 97,
    B            = 98,
    C            = 99,
    D            = 100,
    E            = 101,
    F            = 102,
    G            = 103,
    H            = 104,
    I            = 105,
    J            = 106,
    K            = 107,
    L            = 108,
    M            = 109,
    N            = 110,
    O            = 111,
    P            = 112,
    Q            = 113,
    R            = 114,
    S            = 115,
    T            = 116,
    U            = 117,
    V            = 118,
    W            = 119,
    X            = 120,
    Y            = 121,
    Z            = 122,
    LeftBrace    = 123,
    Pipe         = 124,
    RightBrace   = 125,
    Tilde        = 126,
    Backspace    = 127,
    Right        = 65514,
    Left         = 65515,
    Down         = 65516,
    Up           = 65517,
    Delete       = 65522,
}

impl Key {
    #[inline(always)]
    pub fn code(&self) -> u32 {
        *self as u32
    }

    pub fn get_char(&self) -> Option<char> {
        char::from_u32(self.code())
    }
}

impl FromPrimitive for Key {
    fn from_u64(n: u64) -> Option<Key> {
        match n {
            0     => Some(Key::Unknown),
            9     => Some(Key::Tab),
            13    => Some(Key::Enter),
            17    => Some(Key::CtrlQ),
            18    => Some(Key::CtrlR),
            19    => Some(Key::CtrlS),
            27    => Some(Key::Esc),
            32    => Some(Key::Space),
            33    => Some(Key::Exclaim),
            35    => Some(Key::Hash),
            36    => Some(Key::Dollar),
            37    => Some(Key::Percent),
            38    => Some(Key::Ampersand),
            39    => Some(Key::Quote),
            40    => Some(Key::LeftParen),
            41    => Some(Key::RightParen),
            42    => Some(Key::Asterisk),
            43    => Some(Key::Plus),
            44    => Some(Key::Comma),
            45    => Some(Key::Minus),
            46    => Some(Key::Period),
            47    => Some(Key::Slash),
            48    => Some(Key::D0),
            49    => Some(Key::D1),
            50    => Some(Key::D2),
            51    => Some(Key::D3),
            52    => Some(Key::D4),
            53    => Some(Key::D5),
            54    => Some(Key::D6),
            55    => Some(Key::D7),
            56    => Some(Key::D8),
            57    => Some(Key::D9),
            58    => Some(Key::Colon),
            59    => Some(Key::Semicolon),
            60    => Some(Key::Less),
            61    => Some(Key::Equals),
            62    => Some(Key::Greater),
            63    => Some(Key::Question),
            64    => Some(Key::At),
            91    => Some(Key::LeftBracket),
            92    => Some(Key::Backslash),
            93    => Some(Key::RightBracket),
            94    => Some(Key::Caret),
            95    => Some(Key::Underscore),
            96    => Some(Key::Backquote),
            97    => Some(Key::A),
            98    => Some(Key::B),
            99    => Some(Key::C),
            100   => Some(Key::D),
            101   => Some(Key::E),
            102   => Some(Key::F),
            103   => Some(Key::G),
            104   => Some(Key::H),
            105   => Some(Key::I),
            106   => Some(Key::J),
            107   => Some(Key::K),
            108   => Some(Key::L),
            109   => Some(Key::M),
            110   => Some(Key::N),
            111   => Some(Key::O),
            112   => Some(Key::P),
            113   => Some(Key::Q),
            114   => Some(Key::R),
            115   => Some(Key::S),
            116   => Some(Key::T),
            117   => Some(Key::U),
            118   => Some(Key::V),
            119   => Some(Key::W),
            120   => Some(Key::X),
            121   => Some(Key::Y),
            122   => Some(Key::Z),
            123   => Some(Key::LeftBrace),
            124   => Some(Key::Pipe),
            125   => Some(Key::RightBrace),
            126   => Some(Key::Tilde),
            127   => Some(Key::Backspace),
            65514 => Some(Key::Right),
            65515 => Some(Key::Left),
            65516 => Some(Key::Down),
            65517 => Some(Key::Up),
            65522 => Some(Key::Delete),
            _     => None
        }
    }

    #[inline(always)]
    fn from_i64(n: i64) -> Option<Key> {
        FromPrimitive::from_u64(n as u64)
    }

    #[inline(always)]
    fn from_int(n: int) -> Option<Key> {
        FromPrimitive::from_u64(n as u64)
    }

    #[inline(always)]
    fn from_u16(n: u16) -> Option<Key> {
        FromPrimitive::from_u64(n as u64)
    }

    #[inline(always)]
    fn from_u32(n: u32) -> Option<Key> {
        FromPrimitive::from_u64(n as u64)
    }
}
<|endoftext|>


<commit_before>//! The Actual RethinkDB Commands

use errors::*;
use conn::{Connection, ConnectOpts};
use serde_json;
use ql2::proto;
use std::io::Write;
use byteorder::{WriteBytesExt, LittleEndian};
use std::io::Read;
use std::str;
use protobuf::ProtobufEnum;
use byteorder::ReadBytesExt;
use super::session::Client;
use super::{Result, r};
use std::fmt::Debug;

pub struct RootCommand(Result<String>);
pub struct Command;
pub struct Query;

pub trait IntoCommandArg {
    fn to_arg(&self) -> Result<String>;
}

impl<'a> IntoCommandArg for &'a str {
    fn to_arg(&self) -> Result<String> {
        Ok(self.to_string())
    }
}

impl<'a> IntoCommandArg for &'a String {
    fn to_arg(&self) -> Result<String> {
        Ok(format!("{}", **self))
    }
}

impl IntoCommandArg for serde_json::Value {
    fn to_arg(&self) -> Result<String> {
        serde_json::to_string(self).map_err(|e| From::from(DriverError::Json(e)))
    }
}

impl Client {
    pub fn connection(&self) -> ConnectOpts {
        Self::config().read().clone()
    }

    pub fn db_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DB_CREATE,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     None)))
    }

    pub fn db_drop<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DB_DROP,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     None)))
    }

    pub fn db<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DB,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     None)))
    }

    pub fn table_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let config = Client::config().read();
        r.db(config.db).table_create(name)
    }

    pub fn table<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let config = Client::config().read();
        r.db(config.db).table(name)
    }

    pub fn object(&self) -> serde_json::builder::ObjectBuilder {
        serde_json::builder::ObjectBuilder::new()
    }

    pub fn array(&self) -> serde_json::builder::ArrayBuilder {
        serde_json::builder::ArrayBuilder::new()
    }
}

impl RootCommand {
    pub fn table_create<T>(self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::TABLE_CREATE,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     Some(&commands))))
    }

    pub fn table<T>(self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::TABLE,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     Some(&commands))))
    }

    pub fn insert<T>(self, data: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::INSERT,
                                     Some(data),
                                     None,
                                     Some(&commands))))
    }

    pub fn delete(self) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DELETE, None as Option<&str>, None as Option<&str>, Some(&commands))))
    }

    pub fn filter<T>(self, filter: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::FILTER,
                                     Some(filter),
                                     None,
                                     Some(&commands))))
    }

    pub fn run(self) -> Result<String> {
        let logger = Client::logger().read();
        trace!(logger, "Calling r.run()");
        let commands = try!(self.0);
        let cfg = Client::config().read();
        let query = Query::wrap(proto::Query_QueryType::START, Some(&commands), None);
        debug!(logger, "{}", query);
        let mut conn = try!(Client::conn());
        // Try sending the query
        {
            let mut i = 0;
            let mut write = true;
            let mut connect = false;
            while i < cfg.retries {
                debug!(logger, "Getting connection...");
                if connect {
                    drop(&mut conn);
                    conn = match Client::conn() {
                        Ok(c) => c,
                        Err(error) => {
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(From::from(error));
                            } else {
                                debug!(logger, "Failed getting a connection. Retrying...");
                                i += 1;
                                continue;
                            }
                        }
                    };
                }
                debug!(logger, "Connection aquired.");
                conn.token += 1;
                if write {
                    debug!(logger, "Writing query...");
                    if let Err(error) = Query::write(&query, &mut conn) {
                        connect = true;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to write query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                    debug!(logger, "Query written...");
                    connect = false;
                }
                debug!(logger, "Reading query...");
                match Query::read(&mut conn) {
                    Ok(resp) => {
                        let result = try!(str::from_utf8(&resp));
                        debug!(logger, "{}", result);
                        // If the write operation failed, retry it
                        // {"t":18,"e":4100000,"r":["Cannot perform write: primary replica for
                        // shard [\"\", +inf) not available"],"b":[]}
                        let msg = r#"{"t":18,"e":4100000,"r":["Cannot perform write: primary replica for shard"#;
                        if result.starts_with(msg) {
                            write = true;
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(
                                From::from(
                                    AvailabilityError::OpFailed(
                                        String::from("Not available")
                                        )));
                            } else {
                                debug!(logger, "Write operation failed. Retrying...");
                                i += 1;
                                continue;
                            }
                        } else {
                            debug!(logger, "Query successfully read.");
                            // This is a successful operation
                            break;
                        }
                    }
                    Err(error) => {
                        write = false;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to read query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                }
            }
        }
        Ok(String::new())
    }
}

impl Command {
    pub fn wrap<T>(command: proto::Term_TermType,
                arguments: Option<T>,
                options: Option<T>,
                commands: Option<&str>)
                -> String
        where T: IntoCommandArg
    {
        let mut cmds = format!("[{},", command.value());
        let mut args = String::new();
        if let Some(commands) = commands {
            args.push_str(commands);
            if arguments.is_some() {
                args.push_str(",");
            }
        }
        if let Some(arguments) = arguments {
            args.push_str(&arguments.to_arg().unwrap());
        }
        cmds.push_str(&format!("[{}]", args));
        if let Some(options) = options {
            cmds.push_str(&format!(",{{{}}}", options.to_arg().unwrap()));
        }
        cmds.push(']');
        cmds
    }
}

impl Query {
    pub fn wrap(query_type: proto::Query_QueryType,
                query: Option<&str>,
                options: Option<&str>)
                -> String {
        let mut qry = format!("[{}", query_type.value());
        if let Some(query) = query {
            qry.push_str(&format!(",{}", query));
        }
        if let Some(options) = options {
            qry.push_str(&format!(",{}", options));
        }
        qry.push_str("]");
        qry
    }

    pub fn write(query: &str, conn: &mut Connection) -> Result<()> {
        let query = query.as_bytes();
        let token = conn.token;
        if let Err(error) = conn.stream.write_u64::<LittleEndian>(token) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_u32::<LittleEndian>(query.len() as u32) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_all(query) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.flush() {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(())
    }

    pub fn read(conn: &mut Connection) -> Result<Vec<u8>> {
        // @TODO use response_token to implement parallel reads and writes?
        // let response_token = try!(conn.stream.read_u64::<LittleEndian>());
        let _ = match conn.stream.read_u64::<LittleEndian>() {
            Ok(token) => token,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let len = match conn.stream.read_u32::<LittleEndian>() {
            Ok(len) => len,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let mut resp = vec![0u8; len as usize];
        if let Err(error) = conn.stream.read_exact(&mut resp) {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(resp)
    }
}
<commit_msg>Use String instead of &String<commit_after>//! The Actual RethinkDB Commands

use errors::*;
use conn::{Connection, ConnectOpts};
use serde_json;
use ql2::proto;
use std::io::Write;
use byteorder::{WriteBytesExt, LittleEndian};
use std::io::Read;
use std::str;
use protobuf::ProtobufEnum;
use byteorder::ReadBytesExt;
use super::session::Client;
use super::{Result, r};
use std::fmt::Debug;

pub struct RootCommand(Result<String>);
pub struct Command;
pub struct Query;

pub trait IntoCommandArg {
    fn to_arg(&self) -> Result<String>;
}

impl<'a> IntoCommandArg for &'a str {
    fn to_arg(&self) -> Result<String> {
        Ok(self.to_string())
    }
}

impl IntoCommandArg for String {
    fn to_arg(&self) -> Result<String> {
        Ok(format!("{}", self))
    }
}

impl IntoCommandArg for serde_json::Value {
    fn to_arg(&self) -> Result<String> {
        serde_json::to_string(self).map_err(|e| From::from(DriverError::Json(e)))
    }
}

impl Client {
    pub fn connection(&self) -> ConnectOpts {
        Self::config().read().clone()
    }

    pub fn db_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DB_CREATE,
                                     Some(format!("{:?}", name)),
                                     None,
                                     None)))
    }

    pub fn db_drop<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DB_DROP,
                                     Some(format!("{:?}", name)),
                                     None,
                                     None)))
    }

    pub fn db<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DB,
                                     Some(format!("{:?}", name)),
                                     None,
                                     None)))
    }

    pub fn table_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let config = Client::config().read();
        r.db(config.db).table_create(name)
    }

    pub fn table<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let config = Client::config().read();
        r.db(config.db).table(name)
    }

    pub fn object(&self) -> serde_json::builder::ObjectBuilder {
        serde_json::builder::ObjectBuilder::new()
    }

    pub fn array(&self) -> serde_json::builder::ArrayBuilder {
        serde_json::builder::ArrayBuilder::new()
    }
}

impl RootCommand {
    pub fn table_create<T>(self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::TABLE_CREATE,
                                     Some(format!("{:?}", name)),
                                     None,
                                     Some(&commands))))
    }

    pub fn table<T>(self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::TABLE,
                                     Some(format!("{:?}", name)),
                                     None,
                                     Some(&commands))))
    }

    pub fn insert<T>(self, data: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::INSERT,
                                     Some(data),
                                     None,
                                     Some(&commands))))
    }

    pub fn delete(self) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DELETE, None as Option<&str>, None as Option<&str>, Some(&commands))))
    }

    pub fn filter<T>(self, filter: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::FILTER,
                                     Some(filter),
                                     None,
                                     Some(&commands))))
    }

    pub fn run(self) -> Result<String> {
        let logger = Client::logger().read();
        trace!(logger, "Calling r.run()");
        let commands = try!(self.0);
        let cfg = Client::config().read();
        let query = Query::wrap(proto::Query_QueryType::START, Some(&commands), None);
        debug!(logger, "{}", query);
        let mut conn = try!(Client::conn());
        // Try sending the query
        {
            let mut i = 0;
            let mut write = true;
            let mut connect = false;
            while i < cfg.retries {
                debug!(logger, "Getting connection...");
                if connect {
                    drop(&mut conn);
                    conn = match Client::conn() {
                        Ok(c) => c,
                        Err(error) => {
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(From::from(error));
                            } else {
                                debug!(logger, "Failed getting a connection. Retrying...");
                                i += 1;
                                continue;
                            }
                        }
                    };
                }
                debug!(logger, "Connection aquired.");
                conn.token += 1;
                if write {
                    debug!(logger, "Writing query...");
                    if let Err(error) = Query::write(&query, &mut conn) {
                        connect = true;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to write query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                    debug!(logger, "Query written...");
                    connect = false;
                }
                debug!(logger, "Reading query...");
                match Query::read(&mut conn) {
                    Ok(resp) => {
                        let result = try!(str::from_utf8(&resp));
                        debug!(logger, "{}", result);
                        // If the write operation failed, retry it
                        // {"t":18,"e":4100000,"r":["Cannot perform write: primary replica for
                        // shard [\"\", +inf) not available"],"b":[]}
                        let msg = r#"{"t":18,"e":4100000,"r":["Cannot perform write: primary replica for shard"#;
                        if result.starts_with(msg) {
                            write = true;
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(
                                From::from(
                                    AvailabilityError::OpFailed(
                                        String::from("Not available")
                                        )));
                            } else {
                                debug!(logger, "Write operation failed. Retrying...");
                                i += 1;
                                continue;
                            }
                        } else {
                            debug!(logger, "Query successfully read.");
                            // This is a successful operation
                            break;
                        }
                    }
                    Err(error) => {
                        write = false;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to read query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                }
            }
        }
        Ok(String::new())
    }
}

impl Command {
    pub fn wrap<T>(command: proto::Term_TermType,
                arguments: Option<T>,
                options: Option<T>,
                commands: Option<&str>)
                -> String
        where T: IntoCommandArg
    {
        let mut cmds = format!("[{},", command.value());
        let mut args = String::new();
        if let Some(commands) = commands {
            args.push_str(commands);
            if arguments.is_some() {
                args.push_str(",");
            }
        }
        if let Some(arguments) = arguments {
            args.push_str(&arguments.to_arg().unwrap());
        }
        cmds.push_str(&format!("[{}]", args));
        if let Some(options) = options {
            cmds.push_str(&format!(",{{{}}}", options.to_arg().unwrap()));
        }
        cmds.push(']');
        cmds
    }
}

impl Query {
    pub fn wrap(query_type: proto::Query_QueryType,
                query: Option<&str>,
                options: Option<&str>)
                -> String {
        let mut qry = format!("[{}", query_type.value());
        if let Some(query) = query {
            qry.push_str(&format!(",{}", query));
        }
        if let Some(options) = options {
            qry.push_str(&format!(",{}", options));
        }
        qry.push_str("]");
        qry
    }

    pub fn write(query: &str, conn: &mut Connection) -> Result<()> {
        let query = query.as_bytes();
        let token = conn.token;
        if let Err(error) = conn.stream.write_u64::<LittleEndian>(token) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_u32::<LittleEndian>(query.len() as u32) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_all(query) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.flush() {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(())
    }

    pub fn read(conn: &mut Connection) -> Result<Vec<u8>> {
        // @TODO use response_token to implement parallel reads and writes?
        // let response_token = try!(conn.stream.read_u64::<LittleEndian>());
        let _ = match conn.stream.read_u64::<LittleEndian>() {
            Ok(token) => token,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let len = match conn.stream.read_u32::<LittleEndian>() {
            Ok(len) => len,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let mut resp = vec![0u8; len as usize];
        if let Err(error) = conn.stream.read_exact(&mut resp) {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(resp)
    }
}
<|endoftext|>