<commit_before>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2016-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Initial 32-bit bootloader for x86_64
#![crate_name = "boot"]
// #![crate_type = "staticlib"]
#![feature(asm)]
#![feature(lang_items)]
#![feature(naked_functions)]
#![feature(static_recursion)]
#![feature(linkage)]
#![feature(struct_field_attributes)]
#![feature(stmt_expr_attributes)]
#![no_std]
// #![no_main]

const TABLE_LENGTH: usize = 512;

type Table = [u64; TABLE_LENGTH];

use core::ptr;

// #[lang = "panic_fmt"]
// #[cold]
// unsafe extern fn panic_fmt() -> ! {
//     let vga_buf: *mut u16 = 0xb8000 as *mut u16;
//     for (i, c) in b"boot panic occurred!".iter().enumerate() {
//         ptr::write_volatile( vga_buf.offset(i as isize)
//                            , 0x0200 + *c as u16);
//     }
//     loop { asm!("hlt") };
// }

#[repr(C, packed)]
pub struct Gdt { _null: u64
               , code: u64
               }

#[repr(C, packed)]
pub struct GdtPointer { /// the length of the GDT
                        pub limit: u16
                      , /// pointer to the GDT
                        pub base: &'static Gdt
                      }

impl GdtPointer {
    #[cold]
    #[inline(always)]
    #[naked]
    unsafe fn load (&self) {
        asm!("lgdt ($0)" :: "r"(self) : "memory");
    }
}

#[link_name = ".gdt64"]
#[link_section = ".gdt"]
#[no_mangle]
pub static GDT: Gdt
    = Gdt { _null: 0
          , code: (1<<44) | (1<<47) | (1<<41) | (1<<43) | (1<<53)
        //   , data: (1 << 44) | (1 << 47) | (1 << 41)
          };

// #[no_mangle]
// static mut GDT_PTR: GdtPointer
//     = GdtPointer { limit: 15
//                  , base: &GDT
//                  };

#[cfg(feature = "log")]
#[inline(always)]
#[naked]
fn boot_write(s: &[u8]) {
	// use core::ptr;
    let mut offset = 0;
    unsafe {
        let vga_buf: *mut u16 = 0xb8000 as *mut u16;
    	for c in s {
    		ptr::write_volatile( vga_buf.offset(offset)
                               , 0x0200 + *c as u16);
    		offset += 1;
    	}

    }

}

#[cfg(not(feature = "log"))]
fn boot_write(_s: &[u8]) { }

extern "C" {
    static mut pml4_table: Table;
    static mut pdp_table: Table;
    static mut pd_table: Table;    // static mut page_table: Table;
    // #[link_name = "__vga_buffer"]
    // static mut vga_buf: *mut u16;
}

// #[naked]
// #[inline(always)]
// unsafe fn set_cs(cs: u16) {
//     asm!("ljmpl $0, $$fake_label
//           fake_label: \n\t"
//         :
//         : "i"(cs)
//         : "ax"
//         : "volatile"
//         );
// }
#[cold]
#[inline(always)]
#[naked]
unsafe fn create_page_tables() {
    // 3. if everything is okay, create the page tables and start long mode
    const HUGE_PAGE_SIZE: u64 = 2 * 1024 * 1024; // 2 MiB

    //-- map the PML4 and PDP tables -----------------------------------------
    // recursive map last PML4 entry
    pml4_table[511] = (&pml4_table as *const _ as u64) | 0b11;
    // map first PML4 entry to PDP table
    pml4_table[0] = (&pdp_table as *const _ as u64) | 0b11;
    // map first PDPT entry to PD table
    pdp_table[0] = (&pd_table as *const _ as u64) | 0b11;

    boot_write(b"3.1");

    //-- map the PD table ----------------------------------------------------
    for (number, entry) in pd_table.iter_mut().enumerate() {
        // set each PD entry equal to the start address of the page (the page
        // number times the page's size)
        let addr = number as u64 * HUGE_PAGE_SIZE;
        // with the appropriate flags (present + writable + huge)
        // TODO: do we want to do this using bitflags, or is that too
        //       heavyweight for the boot module?
        //          - eliza, 1/23/2017
        *entry = addr | 0b10000011;
    }

    boot_write(b"3.2");
}
#[cold]
#[inline(always)]
#[naked]
unsafe fn set_long_mode() {

    // load PML4 addr to cr3
    asm!( "mov   cr3, $0" :: "r"(&pml4_table) :: "intel");
    boot_write(b"3.3");

    // enable PAE flag in cr4
    let mut cr4: usize;
    asm!("mov $0, cr4" : "=r"(cr4) ::: "intel");
    cr4 |= 1 << 5;
    asm!("mov cr4, $0" :: "r"(cr4) :: "intel");
    boot_write(b"3.4");

    // set the long mode bit in EFER MSR (model specific register)
    asm!( "mov   ecx, 0xC0000080
           rdmsr
           or    eax, 1 << 8
           wrmsr"
        :::: "intel");
    boot_write(b"3.5");

    // enable paging in cr0
    // let mut cr0: usize;
    // asm!( "mov $0, cr0" : "=r"(cr0) ::: "intel");
    // cr0 |= 1 << 31;
    // cr0 |= 1 << 16;
    // asm!( "mov cr0, $0" :: "r"(cr0) :: "intel");
    // boot_write(b"3.6");
    asm!( "mov eax, cr0
           or  eax, 1 << 31
           or  eax, 1 << 16
           mov cr0, eax"
        :::: "intel");
}

#[cold]
#[no_mangle]
#[naked]
pub unsafe extern "C" fn _start() {
    use core::mem;
    boot_write(b"0");
    asm!("cli");

    // 1. Move Multiboot info pointer to edi
    // let multiboot: usize;
    asm!("mov edi, ebx"
         : //"=r"(multiboot)
         ::: "intel"
    );
    boot_write(b"1");

    // 2. make sure the system supports SOS
    // TODO: port this from boot.asm
    boot_write(b"2");

    create_page_tables();
    set_long_mode();

    // 4. load the 64-bit GDT
    let gdt_ptr = GdtPointer { limit: mem::size_of_val(&GDT) as u16 - 1
                             , base: &GDT };

    // asm!("lgdt ($0)" :: "r"(&gdt_ptr) : "memory" );
    gdt_ptr.load();
    boot_write(b"4");

    // // 5. update selectors
    // asm!("mov ax, 0x10" :::: "intel");
    // boot_write(b"5.1");
    // // stack selector
    // asm!("mov ss, ax" :::: "intel");
    // boot_write(b"5.2");
    // // data selector
    // asm!("mov ds, ax" :::: "intel");
    // boot_write(b"5.3");
    // // extra selector
    // asm!("mov es, ax" :::: "intel");
    // boot_write(b"5.4");
    // loop { }
    // 6. jump to the 64-bit boot subroutine.
    // 0. Move the stack pointer to the top of the stack.
    asm!("ljmpl $$8, $$long_mode_init");

}
<commit_msg>[boot] make GDT loading more idiomatic<commit_after>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2016-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Initial 32-bit bootloader for x86_64
#![crate_name = "boot"]
// #![crate_type = "staticlib"]
#![feature(asm)]
#![feature(lang_items)]
#![feature(naked_functions)]
#![feature(static_recursion)]
#![feature(linkage)]
#![feature(struct_field_attributes)]
#![feature(stmt_expr_attributes)]
#![no_std]
// #![no_main]

const TABLE_LENGTH: usize = 512;

type Table = [u64; TABLE_LENGTH];

use core::ptr;
use core::convert;

#[repr(C, packed)]
pub struct Gdt { _null: u64
               , code: u64
               }

#[repr(C, packed)]
pub struct GdtPointer { /// the length of the GDT
                        pub limit: u16
                      , /// pointer to the GDT
                        pub base: &'static Gdt
                      }

impl GdtPointer {
    #[cold] #[inline(always)]
    unsafe fn load (&self) {
        asm!("lgdt ($0)" :: "r"(self) : "memory");
    }
}

impl convert::From<&'static Gdt> for GdtPointer {
    #[cold] #[inline(always)]
    fn from(gdt: &'static Gdt) -> Self {
        use core::mem::size_of_val;
        GdtPointer { limit: size_of_val(gdt) as u16 - 1
                   , base: gdt }
    }
}

#[link_name = ".gdt64"]
#[link_section = ".gdt"]
#[no_mangle]
pub static GDT: Gdt
    = Gdt { _null: 0
          , code: (1<<44) | (1<<47) | (1<<41) | (1<<43) | (1<<53)
        //   , data: (1 << 44) | (1 << 47) | (1 << 41)
          };

#[cfg(feature = "log")]
#[inline(always)]
#[naked]
fn boot_write(s: &[u8]) {
	// use core::ptr;
    let mut offset = 0;
    unsafe {
        let vga_buf: *mut u16 = 0xb8000 as *mut u16;
    	for c in s {
    		ptr::write_volatile( vga_buf.offset(offset)
                               , 0x0200 + *c as u16);
    		offset += 1;
    	}

    }

}

#[cfg(not(feature = "log"))]
fn boot_write(_s: &[u8]) { }

extern "C" {
    static mut pml4_table: Table;
    static mut pdp_table: Table;
    static mut pd_table: Table;    // static mut page_table: Table;
    // #[link_name = "__vga_buffer"]
    // static mut vga_buf: *mut u16;
}

#[cold]
#[inline(always)]
#[naked]
unsafe fn create_page_tables() {
    // 3. if everything is okay, create the page tables and start long mode
    const HUGE_PAGE_SIZE: u64 = 2 * 1024 * 1024; // 2 MiB

    //-- map the PML4 and PDP tables -----------------------------------------
    // recursive map last PML4 entry
    pml4_table[511] = (&pml4_table as *const _ as u64) | 0b11;
    // map first PML4 entry to PDP table
    pml4_table[0] = (&pdp_table as *const _ as u64) | 0b11;
    // map first PDPT entry to PD table
    pdp_table[0] = (&pd_table as *const _ as u64) | 0b11;

    boot_write(b"3.1");

    //-- map the PD table ----------------------------------------------------
    for (number, entry) in pd_table.iter_mut().enumerate() {
        // set each PD entry equal to the start address of the page (the page
        // number times the page's size)
        let addr = number as u64 * HUGE_PAGE_SIZE;
        // with the appropriate flags (present + writable + huge)
        // TODO: do we want to do this using bitflags, or is that too
        //       heavyweight for the boot module?
        //          - eliza, 1/23/2017
        *entry = addr | 0b10000011;
    }

    boot_write(b"3.2");
}
#[cold]
#[inline(always)]
#[naked]
unsafe fn set_long_mode() {

    // load PML4 addr to cr3
    asm!( "mov   cr3, $0" :: "r"(&pml4_table) :: "intel");
    boot_write(b"3.3");

    // enable PAE flag in cr4
    let mut cr4: usize;
    asm!("mov $0, cr4" : "=r"(cr4) ::: "intel");
    cr4 |= 1 << 5;
    asm!("mov cr4, $0" :: "r"(cr4) :: "intel");
    boot_write(b"3.4");

    // set the long mode bit in EFER MSR (model specific register)
    asm!( "mov   ecx, 0xC0000080
           rdmsr
           or    eax, 1 << 8
           wrmsr"
        :::: "intel");
    boot_write(b"3.5");

    // enable paging in cr0
    // let mut cr0: usize;
    // asm!( "mov $0, cr0" : "=r"(cr0) ::: "intel");
    // cr0 |= 1 << 31;
    // cr0 |= 1 << 16;
    // asm!( "mov cr0, $0" :: "r"(cr0) :: "intel");
    // boot_write(b"3.6");
    asm!( "mov eax, cr0
           or  eax, 1 << 31
           or  eax, 1 << 16
           mov cr0, eax"
        :::: "intel");
}

#[cold]
#[no_mangle]
#[naked]
pub unsafe extern "C" fn _start() {
    boot_write(b"0");
    asm!("cli");

    // 1. Move Multiboot info pointer to edi
    // let multiboot: usize;
    asm!("mov edi, ebx"
         : //"=r"(multiboot)
         ::: "intel"
    );
    boot_write(b"1");

    // 2. make sure the system supports SOS
    // TODO: port this from boot.asm
    boot_write(b"2");

    create_page_tables();
    set_long_mode();

    // 4. load the 64-bit GDT
    GdtPointer::from(&GDT).load();

    boot_write(b"4");

    // // 5. update selectors
    // asm!("mov ax, 0x10" :::: "intel");
    // boot_write(b"5.1");
    // // stack selector
    // asm!("mov ss, ax" :::: "intel");
    // boot_write(b"5.2");
    // // data selector
    // asm!("mov ds, ax" :::: "intel");
    // boot_write(b"5.3");
    // // extra selector
    // asm!("mov es, ax" :::: "intel");
    // boot_write(b"5.4");
    // loop { }
    // 6. jump to the 64-bit boot subroutine.
    // 0. Move the stack pointer to the top of the stack.
    asm!("ljmpl $$8, $$long_mode_init");

}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Kernel IO.
//!
//! This module should eventually abstract over architecture-specific
//! implementation.
pub mod term;
pub mod keyboard;
use core::ops;
use core::fmt;

macro_rules! println {
    ($fmt:expr) => (print!(concat!($fmt, "\n")));
    ($fmt:expr, $($arg:tt)*) => (print!(concat!($fmt, "\n"), $($arg)*));
}

macro_rules! print {
    ($($arg:tt)*) => ({
            use core::fmt::Write;
            $crate::io::term::CONSOLE.lock()
                                     .write_fmt(format_args!($($arg)*))
                                     .unwrap();
    });
}

/// The `Read` trait allows for reading bytes from a source.
///
/// Implementors of the `Read` trait are sometimes called 'readers'.
///
/// Readers are defined by one required method, `read()`. Each call to `read`
/// will attempt to pull bytes from this source into a provided buffer. A
/// number of other methods are implemented in terms of `read()`, giving
/// implementors a number of ways to read bytes while only needing to implement
/// a single method.
///
/// This is basically a brain-dead reimplementation of the standard
/// library's `Read` trait. Most of the methods available on the
/// standard lib's `Read` are not yet implemented.
pub trait Read {
    type Error;
    /// Pull some bytes from this source into the specified buffer, returning
    /// how many bytes were read.
    ///
    /// This function does not provide any guarantees about whether it blocks
    /// waiting for data, but if an object needs to block for a read but cannot
    /// it will typically signal this via an `Err` return value.
    ///
    /// If the return value of this method is `Ok(n)`, then it must be
    /// guaranteed that `0 <= n <= buf.len()`. A nonzero `n` value indicates
    /// that the buffer `buf` has been filled in with `n` bytes of data from
    /// this source. If `n` is `0`, then it can indicate one of two scenarios:
    ///
    /// 1. This reader has reached its "end of file" and will likely no longer
    ///    be able to produce bytes. Note that this does not mean that the
    ///    reader will *always* no longer be able to produce bytes.
    /// 2. The buffer specified was 0 bytes in length.
    ///
    /// No guarantees are provided about the contents of `buf` when this
    /// function is called, implementations cannot rely on any property of the
    /// contents of `buf` being true. It is recommended that implementations
    /// only write data to `buf` instead of reading its contents.
    ///
    /// # Errors
    ///
    /// If this function encounters any form of I/O or other error, an error
    /// variant will be returned. If an error is returned then it must be
    /// guaranteed that no bytes were read.
    fn read(&mut self, buf: &mut [u8])     -> Result<usize, Self::Error>;

    /// Read all bytes until EOF in this source, placing them into `buf`.
    ///
    /// All bytes read from this source will be appended to the specified buffer
    /// `buf`. This function will continuously call `read` to append more data
    /// to `buf` until `read` returns either `Ok(0)` or an error.
    ///
    /// If successful, this function will return the total number of bytes read.
    fn read_all(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error>;
}

/// A trait for objects which are byte-oriented sinks.
///
/// Implementors of the `Write` trait are sometimes called 'writers'.
///
/// Writers are defined by two required methods, `write()` and `flush()`:
///
/// * The `write()` method will attempt to write some data into the object,
///   returning how many bytes were successfully written.
///
/// * The `flush()` method is useful for adaptors and explicit buffers
///   themselves for ensuring that all buffered data has been pushed out to the
///   'true sink'.
///
/// This is basically a brain-dead reimplementation of the standard
/// library's `Write` trait. Most of the methods available on the
/// standard lib's `Write` are not yet implemented.
pub trait Write {
    type Error;

    /// Write a buffer into this object, returning how many bytes were written.
    ///
    /// This function will attempt to write the entire contents of `buf`, but
    /// the entire write may not succeed, or the write may also generate an
    /// error. A call to `write` represents *at most one* attempt to write to
    /// any wrapped object.
    ///
    /// Calls to `write` are not guaranteed to block waiting for data to be
    /// written, and a write which would otherwise block can be indicated through
    /// an `Err` variant.
    ///
    /// If the return value is `Ok(n)` then it must be guaranteed that
    /// `0 <= n <= buf.len()`. A return value of `0` typically means that the
    /// underlying object is no longer able to accept bytes and will likely not
    /// be able to in the future as well, or that the buffer provided is empty.
    ///
    /// # Errors
    ///
    /// Each call to `write` may generate an I/O error indicating that the
    /// operation could not be completed. If an error is returned then no bytes
    /// in the buffer were written to this writer.
    ///
    /// It is **not** considered an error if the entire buffer could not be
    /// written to this writer.
    fn write(&mut self, buf: &[u8])       -> Result<usize, Self::Error>;
}

impl<'a, 'b, E> ops::Shl<&'a [u8]> for &'b mut Write<Error=E>
where E: fmt::Debug {
    type Output = Self;

    /// Fakes the C++ `<<` operator for IOStreams on Write.
    fn shl(self, rhs: &'a [u8]) -> Self::Output {
        self.write(rhs)
            .unwrap();
        self
    }
}
<commit_msg>[kernel] add a quick safe wrapper for CPU ports<commit_after>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Kernel IO.
//!
//! This module should eventually abstract over architecture-specific
//! implementation.
pub mod term;
pub mod keyboard;
pub mod serial;

use arch::cpu;

use core::{ ops, fmt };
use core::marker::PhantomData;

macro_rules! println {
    ($fmt:expr) => (print!(concat!($fmt, "\n")));
    ($fmt:expr, $($arg:tt)*) => (print!(concat!($fmt, "\n"), $($arg)*));
}

macro_rules! print {
    ($($arg:tt)*) => ({
            use core::fmt::Write;
            $crate::io::term::CONSOLE.lock()
                                     .write_fmt(format_args!($($arg)*))
                                     .unwrap();
    });
}

/// The `Read` trait allows for reading bytes from a source.
///
/// Implementors of the `Read` trait are sometimes called 'readers'.
///
/// Readers are defined by one required method, `read()`. Each call to `read`
/// will attempt to pull bytes from this source into a provided buffer. A
/// number of other methods are implemented in terms of `read()`, giving
/// implementors a number of ways to read bytes while only needing to implement
/// a single method.
///
/// This is basically a brain-dead reimplementation of the standard
/// library's `Read` trait. Most of the methods available on the
/// standard lib's `Read` are not yet implemented.
pub trait Read {
    type Error;
    /// Pull some bytes from this source into the specified buffer, returning
    /// how many bytes were read.
    ///
    /// This function does not provide any guarantees about whether it blocks
    /// waiting for data, but if an object needs to block for a read but cannot
    /// it will typically signal this via an `Err` return value.
    ///
    /// If the return value of this method is `Ok(n)`, then it must be
    /// guaranteed that `0 <= n <= buf.len()`. A nonzero `n` value indicates
    /// that the buffer `buf` has been filled in with `n` bytes of data from
    /// this source. If `n` is `0`, then it can indicate one of two scenarios:
    ///
    /// 1. This reader has reached its "end of file" and will likely no longer
    ///    be able to produce bytes. Note that this does not mean that the
    ///    reader will *always* no longer be able to produce bytes.
    /// 2. The buffer specified was 0 bytes in length.
    ///
    /// No guarantees are provided about the contents of `buf` when this
    /// function is called, implementations cannot rely on any property of the
    /// contents of `buf` being true. It is recommended that implementations
    /// only write data to `buf` instead of reading its contents.
    ///
    /// # Errors
    ///
    /// If this function encounters any form of I/O or other error, an error
    /// variant will be returned. If an error is returned then it must be
    /// guaranteed that no bytes were read.
    fn read(&mut self, buf: &mut [u8])     -> Result<usize, Self::Error>;

    /// Read all bytes until EOF in this source, placing them into `buf`.
    ///
    /// All bytes read from this source will be appended to the specified buffer
    /// `buf`. This function will continuously call `read` to append more data
    /// to `buf` until `read` returns either `Ok(0)` or an error.
    ///
    /// If successful, this function will return the total number of bytes read.
    fn read_all(&mut self, buf: &mut [u8]) -> Result<usize, Self::Error>;
}

/// A trait for objects which are byte-oriented sinks.
///
/// Implementors of the `Write` trait are sometimes called 'writers'.
///
/// Writers are defined by two required methods, `write()` and `flush()`:
///
/// * The `write()` method will attempt to write some data into the object,
///   returning how many bytes were successfully written.
///
/// * The `flush()` method is useful for adaptors and explicit buffers
///   themselves for ensuring that all buffered data has been pushed out to the
///   'true sink'.
///
/// This is basically a brain-dead reimplementation of the standard
/// library's `Write` trait. Most of the methods available on the
/// standard lib's `Write` are not yet implemented.
pub trait Write {
    type Error;

    /// Write a buffer into this object, returning how many bytes were written.
    ///
    /// This function will attempt to write the entire contents of `buf`, but
    /// the entire write may not succeed, or the write may also generate an
    /// error. A call to `write` represents *at most one* attempt to write to
    /// any wrapped object.
    ///
    /// Calls to `write` are not guaranteed to block waiting for data to be
    /// written, and a write which would otherwise block can be indicated through
    /// an `Err` variant.
    ///
    /// If the return value is `Ok(n)` then it must be guaranteed that
    /// `0 <= n <= buf.len()`. A return value of `0` typically means that the
    /// underlying object is no longer able to accept bytes and will likely not
    /// be able to in the future as well, or that the buffer provided is empty.
    ///
    /// # Errors
    ///
    /// Each call to `write` may generate an I/O error indicating that the
    /// operation could not be completed. If an error is returned then no bytes
    /// in the buffer were written to this writer.
    ///
    /// It is **not** considered an error if the entire buffer could not be
    /// written to this writer.
    fn write(&mut self, buf: &[u8])       -> Result<usize, Self::Error>;
}

impl<'a, 'b, E> ops::Shl<&'a [u8]> for &'b mut Write<Error=E>
where E: fmt::Debug {
    type Output = Self;

    /// Fakes the C++ `<<` operator for IOStreams on Write.
    fn shl(self, rhs: &'a [u8]) -> Self::Output {
        self.write(rhs)
            .unwrap();
        self
    }
}

/// Safe typed port wrapper
pub struct Port<T> { raw_port: cpu::Port
                   , typ: PhantomData<T>
                   }

macro_rules! make_ports {
    ( $( $t:ty, $read:ident, $out:ident ),+ ) => {
        $(
            impl Port<$t> {
                #[inline]
                pub const fn new(number: u16) -> Self {
                    // TODO: can we check if the port number is valid
                    unsafe {
                        Port { raw_port: cpu::Port::new(number)
                             , typ: PhantomData::<$t>
                             }
                    }
                }

                #[inline]
                pub fn read(&self) -> $t {
                    unsafe { self.raw_port.$read() }
                }

                #[inline]
                pub fn write(&self, data: $t) {
                    unsafe { self.raw_port.$out(data) }
                }
            }
        )+
    }
}

make_ports! { u8, in8, out8
            , u16, in16, out16
            , u32, in32, out32
            }

#[cfg(arch="x86_64")]
make_ports! { u64, in64, out64 }
<|endoftext|>


<commit_before>// Copyright (c) 2016 Anatoly Ikorsky
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use BoxFuture;
use errors::*;
use std::fmt;
use std::fs::File;
use lib_futures::Async::NotReady;
use lib_futures::Async::Ready;
#[cfg(feature = "ssl")]
use lib_futures::{Future, IntoFuture};
use lib_futures::Poll;
use lib_futures::stream;
use io::futures::ConnectingStream;
use io::futures::new_connecting_stream;
use io::futures::new_write_packet;
use io::futures::WritePacket;
#[cfg(feature = "ssl")]
use native_tls::{Certificate, Pkcs12, TlsConnector};
use myc::packets::{PacketParser, ParseResult, RawPacket};
use opts::SslOpts;
use std::cmp;
use std::io;
use std::io::Read;
use std::net::ToSocketAddrs;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::reactor::Handle;
use tokio_io::AsyncRead;
use tokio_io::AsyncWrite;
#[cfg(feature = "ssl")]
use tokio_tls::{TlsConnectorExt, TlsStream};

pub mod futures;

#[derive(Debug)]
enum Endpoint {
    Plain(TcpStream),
    #[cfg(feature = "ssl")]
    Secure(TlsStream<TcpStream>),
}

impl Endpoint {
    pub fn set_keepalive_ms(&self, ms: Option<u32>) -> Result<()> {
        let ms = ms.map(|val| Duration::from_millis(val as u64));
        match *self {
            Endpoint::Plain(ref stream) => stream.set_keepalive(ms)?,
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref stream) => stream.get_ref().get_ref().set_keepalive(ms)?,
        }
        Ok(())
    }

    #[cfg(feature = "ssl")]
    pub fn make_secure(self, domain: String, ssl_opts: SslOpts) -> BoxFuture<Self> {
        let fut = ::std::fs::File::open(ssl_opts.pkcs12_path())
            .map_err(Error::from)
            .and_then(|mut file| {
                let mut der = vec![];
                file.read_to_end(&mut der)?;
                let identity = Pkcs12::from_der(&*der, ssl_opts.password().unwrap_or(""))
                    .chain_err(|| "Can't parse der")?;
                let mut builder = TlsConnector::builder().chain_err(
                    || "Can't create TlsConnectorBuilder)",
                )?;
                match ssl_opts.root_cert_path() {
                    Some(root_cert_path) => {
                        let mut root_cert_der = vec![];
                        let mut root_cert_file = File::open(root_cert_path)?;
                        root_cert_file.read_to_end(&mut root_cert_der)?;
                        let root_cert = Certificate::from_der(&*root_cert_der)
                            .chain_err(|| "Can't parse root certificate")?;
                        builder.add_root_certificate(root_cert)
                            .chain_err(|| "Can't add root certificate")?;
                    },
                    None => (),
                }
                builder.identity(identity).chain_err(
                    || "Can't set identity for TlsConnectorBuilder",
                )?;
                builder.build().chain_err(
                    || "Can't build TlsConnectorBuilder",
                )
            })
            .into_future()
            .and_then(move |tls_connector| match self {
                Endpoint::Plain(stream) => {
                    let fut = if ssl_opts.skip_domain_validation() {
                        TlsConnectorExt::danger_connect_async_without_providing_domain_for_certificate_verification_and_server_name_indication(&tls_connector, stream)
                    } else {
                        TlsConnectorExt::connect_async(&tls_connector, &*domain, stream)
                    };
                    fut.then(|result| result.chain_err(|| "Can't connect TlsConnector"))
                }
                Endpoint::Secure(_) => unreachable!(),
            })
            .map(|tls_stream| Endpoint::Secure(tls_stream));
        Box::new(fut)
    }
}

impl From<TcpStream> for Endpoint {
    fn from(stream: TcpStream) -> Self {
        Endpoint::Plain(stream)
    }
}

#[cfg(feature = "ssl")]
impl From<TlsStream<TcpStream>> for Endpoint {
    fn from(stream: TlsStream<TcpStream>) -> Self {
        Endpoint::Secure(stream)
    }
}

impl io::Read for Endpoint {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.read(buf),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.read(buf),
        }
    }
}

impl io::Write for Endpoint {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.write(buf),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.write(buf),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.flush(),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.flush(),
        }
    }
}

impl AsyncRead for Endpoint {
    unsafe fn prepare_uninitialized_buffer(&self, _buf: &mut [u8]) -> bool {
        match *self {
            Endpoint::Plain(ref stream) => stream.prepare_uninitialized_buffer(_buf),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref stream) => stream.prepare_uninitialized_buffer(_buf),
        }
    }
}

impl AsyncWrite for Endpoint {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.shutdown(),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.shutdown(),
        }
    }
}

/// Stream connected to MySql server.
pub struct Stream {
    endpoint: Option<Endpoint>,
    closed: bool,
    next_packet: Option<ParseResult>,
    buf: Option<Vec<u8>>,
}

impl fmt::Debug for Stream {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Stream (endpoint={:?})", self.endpoint)
    }
}

impl Stream {
    pub fn connect<S>(addr: S, handle: &Handle) -> ConnectingStream
    where
        S: ToSocketAddrs,
    {
        new_connecting_stream(addr, handle)
    }

    pub fn write_packet(self, data: Vec<u8>, seq_id: u8) -> WritePacket {
        new_write_packet(self, data, seq_id)
    }

    pub fn set_keepalive_ms(&self, ms: Option<u32>) -> Result<()> {
        match self.endpoint {
            Some(ref endpoint) => endpoint.set_keepalive_ms(ms),
            None => unreachable!(),
        }
    }

    #[cfg(not(feature = "ssl"))]
    #[allow(unused)]
    pub fn make_secure(self, domain: String, ssl_opts: SslOpts) -> BoxFuture<Self> {
        panic!("Ssl connection requires `ssl` feature");
    }

    #[cfg(feature = "ssl")]
    pub fn make_secure(mut self, domain: String, ssl_opts: SslOpts) -> BoxFuture<Self> {
        match self.endpoint.take() {
            Some(endpoint) => {
                let fut = endpoint.make_secure(domain, ssl_opts).map(|endpoint| {
                    self.endpoint = Some(endpoint);
                    self
                });
                Box::new(fut)
            }
            None => unreachable!(),
        }
    }
}

impl io::Read for Stream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.read(buf),
            None => unreachable!(),
        }
    }
}

impl io::Write for Stream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.write(buf),
            None => unreachable!(),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.flush(),
            None => unreachable!(),
        }
    }
}

impl AsyncRead for Stream {
    unsafe fn prepare_uninitialized_buffer(&self, _buf: &mut [u8]) -> bool {
        match self.endpoint {
            Some(ref endpoint) => endpoint.prepare_uninitialized_buffer(_buf),
            None => unreachable!(),
        }
    }
}

impl AsyncWrite for Stream {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.shutdown(),
            None => unreachable!(),
        }
    }
}

impl stream::Stream for Stream {
    type Item = (RawPacket, u8);
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<(RawPacket, u8)>, Error> {
        // should read everything from self.endpoint
        let mut would_block = false;
        if !self.closed {
            let mut buf = [0u8; 4096];
            loop {
                match self.read(&mut buf[..]) {
                    Ok(0) => {
                        break;
                    }
                    Ok(size) => {
                        let buf_handle = self.buf.as_mut().unwrap();
                        buf_handle.extend_from_slice(&buf[..size]);
                    }
                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                        would_block = true;
                        break;
                    }
                    Err(error) => {
                        self.closed = true;
                        return Err(Error::from(error));
                    }
                };
            }
        } else {
            return Ok(Ready(None));
        }

        // need to call again if there is a data in self.buf
        // or data was written to packet parser
        let mut should_poll = false;

        let next_packet = self.next_packet.take().expect(
            "Stream.next_packet should not be None",
        );
        let next_packet = match next_packet {
            ParseResult::Done(packet, seq_id) => {
                self.next_packet = Some(PacketParser::empty().parse());
                return Ok(Ready(Some((packet, seq_id))));
            }
            ParseResult::Incomplete(mut new_packet, needed) => {
                let buf_handle = self.buf.as_mut().unwrap();
                let buf_len = buf_handle.len();

                let to = cmp::min(needed, buf_len);
                new_packet.extend_from_slice(&buf_handle[..to]);
                unsafe {
                    let src = buf_handle.as_ptr().offset(to as isize);
                    let dst = buf_handle.as_mut_ptr();
                    let len = buf_handle.len() - to;
                    ::std::ptr::copy(src, dst, len);
                    buf_handle.set_len(len);
                }

                if buf_len != 0 {
                    should_poll = true;
                }

                new_packet
            }
        };

        self.next_packet = Some(next_packet.parse());

        if should_poll || !would_block {
            self.poll()
        } else {
            Ok(NotReady)
        }
    }
}
<commit_msg>Turn recursive io future polling into a loop<commit_after>// Copyright (c) 2016 Anatoly Ikorsky
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use BoxFuture;
use errors::*;
use std::fmt;
use std::fs::File;
use lib_futures::Async::NotReady;
use lib_futures::Async::Ready;
#[cfg(feature = "ssl")]
use lib_futures::{Future, IntoFuture};
use lib_futures::Poll;
use lib_futures::stream;
use io::futures::ConnectingStream;
use io::futures::new_connecting_stream;
use io::futures::new_write_packet;
use io::futures::WritePacket;
#[cfg(feature = "ssl")]
use native_tls::{Certificate, Pkcs12, TlsConnector};
use myc::packets::{PacketParser, ParseResult, RawPacket};
use opts::SslOpts;
use std::cmp;
use std::io;
use std::io::Read;
use std::net::ToSocketAddrs;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::reactor::Handle;
use tokio_io::AsyncRead;
use tokio_io::AsyncWrite;
#[cfg(feature = "ssl")]
use tokio_tls::{TlsConnectorExt, TlsStream};

pub mod futures;

#[derive(Debug)]
enum Endpoint {
    Plain(TcpStream),
    #[cfg(feature = "ssl")]
    Secure(TlsStream<TcpStream>),
}

impl Endpoint {
    pub fn set_keepalive_ms(&self, ms: Option<u32>) -> Result<()> {
        let ms = ms.map(|val| Duration::from_millis(val as u64));
        match *self {
            Endpoint::Plain(ref stream) => stream.set_keepalive(ms)?,
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref stream) => stream.get_ref().get_ref().set_keepalive(ms)?,
        }
        Ok(())
    }

    #[cfg(feature = "ssl")]
    pub fn make_secure(self, domain: String, ssl_opts: SslOpts) -> BoxFuture<Self> {
        let fut = ::std::fs::File::open(ssl_opts.pkcs12_path())
            .map_err(Error::from)
            .and_then(|mut file| {
                let mut der = vec![];
                file.read_to_end(&mut der)?;
                let identity = Pkcs12::from_der(&*der, ssl_opts.password().unwrap_or(""))
                    .chain_err(|| "Can't parse der")?;
                let mut builder = TlsConnector::builder().chain_err(
                    || "Can't create TlsConnectorBuilder)",
                )?;
                match ssl_opts.root_cert_path() {
                    Some(root_cert_path) => {
                        let mut root_cert_der = vec![];
                        let mut root_cert_file = File::open(root_cert_path)?;
                        root_cert_file.read_to_end(&mut root_cert_der)?;
                        let root_cert = Certificate::from_der(&*root_cert_der)
                            .chain_err(|| "Can't parse root certificate")?;
                        builder.add_root_certificate(root_cert)
                            .chain_err(|| "Can't add root certificate")?;
                    },
                    None => (),
                }
                builder.identity(identity).chain_err(
                    || "Can't set identity for TlsConnectorBuilder",
                )?;
                builder.build().chain_err(
                    || "Can't build TlsConnectorBuilder",
                )
            })
            .into_future()
            .and_then(move |tls_connector| match self {
                Endpoint::Plain(stream) => {
                    let fut = if ssl_opts.skip_domain_validation() {
                        TlsConnectorExt::danger_connect_async_without_providing_domain_for_certificate_verification_and_server_name_indication(&tls_connector, stream)
                    } else {
                        TlsConnectorExt::connect_async(&tls_connector, &*domain, stream)
                    };
                    fut.then(|result| result.chain_err(|| "Can't connect TlsConnector"))
                }
                Endpoint::Secure(_) => unreachable!(),
            })
            .map(|tls_stream| Endpoint::Secure(tls_stream));
        Box::new(fut)
    }
}

impl From<TcpStream> for Endpoint {
    fn from(stream: TcpStream) -> Self {
        Endpoint::Plain(stream)
    }
}

#[cfg(feature = "ssl")]
impl From<TlsStream<TcpStream>> for Endpoint {
    fn from(stream: TlsStream<TcpStream>) -> Self {
        Endpoint::Secure(stream)
    }
}

impl io::Read for Endpoint {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.read(buf),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.read(buf),
        }
    }
}

impl io::Write for Endpoint {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.write(buf),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.write(buf),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.flush(),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.flush(),
        }
    }
}

impl AsyncRead for Endpoint {
    unsafe fn prepare_uninitialized_buffer(&self, _buf: &mut [u8]) -> bool {
        match *self {
            Endpoint::Plain(ref stream) => stream.prepare_uninitialized_buffer(_buf),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref stream) => stream.prepare_uninitialized_buffer(_buf),
        }
    }
}

impl AsyncWrite for Endpoint {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.shutdown(),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.shutdown(),
        }
    }
}

/// Stream connected to MySql server.
pub struct Stream {
    endpoint: Option<Endpoint>,
    closed: bool,
    next_packet: Option<ParseResult>,
    buf: Option<Vec<u8>>,
}

impl fmt::Debug for Stream {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Stream (endpoint={:?})", self.endpoint)
    }
}

impl Stream {
    pub fn connect<S>(addr: S, handle: &Handle) -> ConnectingStream
    where
        S: ToSocketAddrs,
    {
        new_connecting_stream(addr, handle)
    }

    pub fn write_packet(self, data: Vec<u8>, seq_id: u8) -> WritePacket {
        new_write_packet(self, data, seq_id)
    }

    pub fn set_keepalive_ms(&self, ms: Option<u32>) -> Result<()> {
        match self.endpoint {
            Some(ref endpoint) => endpoint.set_keepalive_ms(ms),
            None => unreachable!(),
        }
    }

    #[cfg(not(feature = "ssl"))]
    #[allow(unused)]
    pub fn make_secure(self, domain: String, ssl_opts: SslOpts) -> BoxFuture<Self> {
        panic!("Ssl connection requires `ssl` feature");
    }

    #[cfg(feature = "ssl")]
    pub fn make_secure(mut self, domain: String, ssl_opts: SslOpts) -> BoxFuture<Self> {
        match self.endpoint.take() {
            Some(endpoint) => {
                let fut = endpoint.make_secure(domain, ssl_opts).map(|endpoint| {
                    self.endpoint = Some(endpoint);
                    self
                });
                Box::new(fut)
            }
            None => unreachable!(),
        }
    }
}

impl io::Read for Stream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.read(buf),
            None => unreachable!(),
        }
    }
}

impl io::Write for Stream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.write(buf),
            None => unreachable!(),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.flush(),
            None => unreachable!(),
        }
    }
}

impl AsyncRead for Stream {
    unsafe fn prepare_uninitialized_buffer(&self, _buf: &mut [u8]) -> bool {
        match self.endpoint {
            Some(ref endpoint) => endpoint.prepare_uninitialized_buffer(_buf),
            None => unreachable!(),
        }
    }
}

impl AsyncWrite for Stream {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.shutdown(),
            None => unreachable!(),
        }
    }
}

impl stream::Stream for Stream {
    type Item = (RawPacket, u8);
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<(RawPacket, u8)>, Error> {
        loop {
            // should read everything from self.endpoint
            let mut would_block = false;
            if !self.closed {
                let mut buf = [0u8; 4096];
                loop {
                    match self.read(&mut buf[..]) {
                        Ok(0) => {
                            break;
                        }
                        Ok(size) => {
                            let buf_handle = self.buf.as_mut().unwrap();
                            buf_handle.extend_from_slice(&buf[..size]);
                        }
                        Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                            would_block = true;
                            break;
                        }
                        Err(error) => {
                            self.closed = true;
                            return Err(Error::from(error));
                        }
                    };
                }
            } else {
                return Ok(Ready(None));
            }

            // need to call again if there is a data in self.buf
            // or data was written to packet parser
            let mut should_poll = false;

            let next_packet = self.next_packet.take().expect(
                "Stream.next_packet should not be None",
                );
            let next_packet = match next_packet {
                ParseResult::Done(packet, seq_id) => {
                    self.next_packet = Some(PacketParser::empty().parse());
                    return Ok(Ready(Some((packet, seq_id))));
                }
                ParseResult::Incomplete(mut new_packet, needed) => {
                    let buf_handle = self.buf.as_mut().unwrap();
                    let buf_len = buf_handle.len();

                    let to = cmp::min(needed, buf_len);
                    new_packet.extend_from_slice(&buf_handle[..to]);
                    unsafe {
                        let src = buf_handle.as_ptr().offset(to as isize);
                        let dst = buf_handle.as_mut_ptr();
                        let len = buf_handle.len() - to;
                        ::std::ptr::copy(src, dst, len);
                        buf_handle.set_len(len);
                    }

                    if buf_len != 0 {
                        should_poll = true;
                    }

                    new_packet
                }
            };

            self.next_packet = Some(next_packet.parse());

            if !should_poll && would_block {
                return Ok(NotReady)
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::sync::Arc;
use std::{error, fmt, io};

/// The most generic error message in ReQL
#[derive(Debug, Clone)]
pub enum Error {
    Compile(String),
    Runtime(Runtime),
    Driver(Driver),
}

impl error::Error for Error {}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::Compile(msg) => write!(f, "compile error; {}", msg),
            Self::Runtime(msg) => write!(f, "runtime error; {}", msg),
            Self::Driver(msg) => write!(f, "client error; {}", msg),
        }
    }
}

/// The parent class of all runtime errors
///
/// All errors on the server unrelated to compilation. Programs may use this to catch any runtime
/// error, but the server will always return a more specific error class.
#[derive(Debug, Clone)]
pub enum Runtime {
    /// The query contains a logical impossibility, such as adding a number to a string.
    QueryLogic(String),
    NonExistence(String),
    ResourceLimit(String),
    User(String),
    Internal(String),
    Availability(Availability),
    Permission(String),
}

impl From<Runtime> for Error {
    fn from(err: Runtime) -> Error {
        Error::Runtime(err)
    }
}

impl fmt::Display for Runtime {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::QueryLogic(msg) => write!(f, "query logic; {}", msg),
            Self::NonExistence(msg) => write!(f, "non-existance error; {}", msg),
            Self::ResourceLimit(msg) => write!(f, "resource limit error; {}", msg),
            Self::User(msg) => write!(f, "user error; {}", msg),
            Self::Internal(msg) => write!(f, "internal error; {}", msg),
            Self::Availability(msg) => write!(f, "availability error; {}", msg),
            Self::Permission(msg) => write!(f, "permission error; {}", msg),
        }
    }
}

/// A server in the cluster is unavailable
///
/// The parent class of `OpFailedError` and `OpIndeterminateError`. Programs may use this
/// to catch any availability error, but the server will always return one of this class’s
/// children.
#[derive(Debug, Clone)]
pub enum Availability {
    OpFailed(String),
    OpIndeterminate(String),
}

impl From<Availability> for Error {
    fn from(err: Availability) -> Error {
        Runtime::Availability(err).into()
    }
}

impl fmt::Display for Availability {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::OpFailed(msg) => write!(f, "operation failed; {}", msg),
            Self::OpIndeterminate(msg) => write!(f, "operation indeterminate; {}", msg),
        }
    }
}

/// An error has occurred within the driver
///
/// This may be a driver bug, or it may be an unfulfillable command, such as an unserializable
/// query.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum Driver {
    Auth(String),
    ConnectionBroken,
    ConnectionLocked,
    Io(io::ErrorKind, String),
    Json(Arc<serde_json::Error>),
    Other(String),
}

impl From<Driver> for Error {
    fn from(err: Driver) -> Error {
        Error::Driver(err)
    }
}

impl fmt::Display for Driver {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::Auth(msg) => write!(f, "auth error; {}", msg),
            Self::ConnectionBroken => write!(f, "connection broken"),
            Self::ConnectionLocked => write!(
                f,
                "another query is running a changefeed on this connection"
            ),
            Self::Io(_, error) => write!(f, "{}", error),
            Self::Json(error) => write!(f, "{}", error),
            Self::Other(msg) => write!(f, "{}", msg),
        }
    }
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Driver::Io(err.kind(), err.to_string()).into()
    }
}

impl From<serde_json::Error> for Error {
    fn from(err: serde_json::Error) -> Error {
        Driver::Json(Arc::new(err)).into()
    }
}
<commit_msg>fix typo<commit_after>use std::sync::Arc;
use std::{error, fmt, io};

/// The most generic error message in ReQL
#[derive(Debug, Clone)]
pub enum Error {
    Compile(String),
    Runtime(Runtime),
    Driver(Driver),
}

impl error::Error for Error {}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::Compile(msg) => write!(f, "compile error; {}", msg),
            Self::Runtime(msg) => write!(f, "runtime error; {}", msg),
            Self::Driver(msg) => write!(f, "client error; {}", msg),
        }
    }
}

/// The parent class of all runtime errors
///
/// All errors on the server unrelated to compilation. Programs may use this to catch any runtime
/// error, but the server will always return a more specific error class.
#[derive(Debug, Clone)]
pub enum Runtime {
    /// The query contains a logical impossibility, such as adding a number to a string.
    QueryLogic(String),
    NonExistence(String),
    ResourceLimit(String),
    User(String),
    Internal(String),
    Availability(Availability),
    Permission(String),
}

impl From<Runtime> for Error {
    fn from(err: Runtime) -> Error {
        Error::Runtime(err)
    }
}

impl fmt::Display for Runtime {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::QueryLogic(msg) => write!(f, "query logic; {}", msg),
            Self::NonExistence(msg) => write!(f, "non-existence error; {}", msg),
            Self::ResourceLimit(msg) => write!(f, "resource limit error; {}", msg),
            Self::User(msg) => write!(f, "user error; {}", msg),
            Self::Internal(msg) => write!(f, "internal error; {}", msg),
            Self::Availability(msg) => write!(f, "availability error; {}", msg),
            Self::Permission(msg) => write!(f, "permission error; {}", msg),
        }
    }
}

/// A server in the cluster is unavailable
///
/// The parent class of `OpFailedError` and `OpIndeterminateError`. Programs may use this
/// to catch any availability error, but the server will always return one of this class’s
/// children.
#[derive(Debug, Clone)]
pub enum Availability {
    OpFailed(String),
    OpIndeterminate(String),
}

impl From<Availability> for Error {
    fn from(err: Availability) -> Error {
        Runtime::Availability(err).into()
    }
}

impl fmt::Display for Availability {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::OpFailed(msg) => write!(f, "operation failed; {}", msg),
            Self::OpIndeterminate(msg) => write!(f, "operation indeterminate; {}", msg),
        }
    }
}

/// An error has occurred within the driver
///
/// This may be a driver bug, or it may be an unfulfillable command, such as an unserializable
/// query.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum Driver {
    Auth(String),
    ConnectionBroken,
    ConnectionLocked,
    Io(io::ErrorKind, String),
    Json(Arc<serde_json::Error>),
    Other(String),
}

impl From<Driver> for Error {
    fn from(err: Driver) -> Error {
        Error::Driver(err)
    }
}

impl fmt::Display for Driver {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::Auth(msg) => write!(f, "auth error; {}", msg),
            Self::ConnectionBroken => write!(f, "connection broken"),
            Self::ConnectionLocked => write!(
                f,
                "another query is running a changefeed on this connection"
            ),
            Self::Io(_, error) => write!(f, "{}", error),
            Self::Json(error) => write!(f, "{}", error),
            Self::Other(msg) => write!(f, "{}", msg),
        }
    }
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Driver::Io(err.kind(), err.to_string()).into()
    }
}

impl From<serde_json::Error> for Error {
    fn from(err: serde_json::Error) -> Error {
        Driver::Json(Arc::new(err)).into()
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Eli Friedman.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![allow(unused_variables)]

use conversions::Convert128;
use simd::i8x16;
use __m128;
use __m128i;
use __m128d;

// Declarations copied from the llvmint crate.
#[allow(improper_ctypes)]
extern {
    #[link_name = "llvm.x86.sse3.ldu.dq"]
    pub fn sse3_ldu_dq(a: *mut i8) -> i8x16;
}

/// addsubpd
#[inline]
pub fn _mm_addsub_pd (a: __m128d, b: __m128d) -> __m128d {
    unimplemented!()
}
/// addsubps
#[inline]
pub fn _mm_addsub_ps (a: __m128, b: __m128) -> __m128 {
    unimplemented!()
}
/// haddpd
#[inline]
pub fn _mm_hadd_pd (a: __m128d, b: __m128d) -> __m128d {
    unimplemented!()
}
/// haddps
#[inline]
pub fn _mm_hadd_ps (a: __m128, b: __m128) -> __m128 {
    unimplemented!()
}
/// hsubpd
#[inline]
pub fn _mm_hsub_pd (a: __m128d, b: __m128d) -> __m128d {
    unimplemented!()
}
/// hsubps
#[inline]
pub fn _mm_hsub_ps (a: __m128, b: __m128) -> __m128 {
    unimplemented!()
}
/// lddqu
#[inline]
pub unsafe fn _mm_lddqu_si128(mem_addr: *const __m128i) -> __m128i {
    sse3_ldu_dq(mem_addr as *mut i8).as_i64x2()
}
/// movddup
#[inline]
pub unsafe fn _mm_loaddup_pd (mem_addr: *const f64) ->  __m128d {
    unimplemented!()
}
/// movddup
#[inline]
pub fn _mm_movedup_pd (a: __m128d) ->  __m128d {
    __m128d::new(a.extract(0), a.extract(0))
}
/// movshdup
#[inline]
pub fn _mm_movehdup_ps (a: __m128) -> __m128 {
    unimplemented!()
}
/// movsldup
#[inline]
pub fn _mm_moveldup_ps (a: __m128) -> __m128 {
    unimplemented!()
}
<commit_msg>SSE3 implementations.<commit_after>// Copyright 2016 Eli Friedman.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![allow(unused_variables)]

use conversions::Convert128;
use simd::i8x16;
use simd::x86::sse3::{Sse3F64x2, Sse3F32x4};
use __m128;
use __m128i;
use __m128d;
use std::ptr::copy_nonoverlapping;

// Declarations copied from the llvmint crate.
#[allow(improper_ctypes)]
extern {
    #[link_name = "llvm.x86.sse3.ldu.dq"]
    pub fn sse3_ldu_dq(a: *mut i8) -> i8x16;
}

/// addsubpd
#[inline]
pub fn _mm_addsub_pd(a: __m128d, b: __m128d) -> __m128d {
    a.addsub(b)
}
/// addsubps
#[inline]
pub fn _mm_addsub_ps(a: __m128, b: __m128) -> __m128 {
    a.addsub(b)
}
/// haddpd
#[inline]
pub fn _mm_hadd_pd(a: __m128d, b: __m128d) -> __m128d {
    a.hadd(b)
}
/// haddps
#[inline]
pub fn _mm_hadd_ps(a: __m128, b: __m128) -> __m128 {
    a.hadd(b)
}
/// hsubpd
#[inline]
pub fn _mm_hsub_pd(a: __m128d, b: __m128d) -> __m128d {
    a.hsub(b)
}
/// hsubps
#[inline]
pub fn _mm_hsub_ps(a: __m128, b: __m128) -> __m128 {
    a.hsub(b)
}
/// lddqu
#[inline]
pub unsafe fn _mm_lddqu_si128(mem_addr: *const __m128i) -> __m128i {
    sse3_ldu_dq(mem_addr as *mut i8).as_i64x2()
}
/// movddup
#[inline]
pub unsafe fn _mm_loaddup_pd(mem_addr: *const f64) ->  __m128d {
    let mut f: f64 = 0.;
    let mem_addr = mem_addr as *const u8;
    let pf = &mut f as *mut f64 as *mut u8;
    copy_nonoverlapping(mem_addr, pf, 8);
    __m128d::new(f, f)
}
/// movddup
#[inline]
pub fn _mm_movedup_pd(a: __m128d) ->  __m128d {
    __m128d::new(a.extract(0), a.extract(0))
}
/// movshdup
#[inline]
pub fn _mm_movehdup_ps(a: __m128) -> __m128 {
    a.replace(0, a.extract(1)).replace(2, a.extract(3))
}
/// movsldup
#[inline]
pub fn _mm_moveldup_ps(a: __m128) -> __m128 {
    a.replace(1, a.extract(0)).replace(3, a.extract(2))
}
<|endoftext|>


<commit_before>use std::marker::PhantomData;
use std::iter::FromIterator;

use rand;
use num::traits::FromPrimitive;

pub struct GenerateCtx<'a, R: ?Sized + 'a> {
    pub rng: &'a mut R,
    pub size: usize
}

impl <'a, R: ?Sized + 'a> GenerateCtx<'a, R> {
    fn new(rng: &'a mut R, size: usize) -> Self {
        GenerateCtx { rng: rng, size: size }
    }

    #[inline]
    fn chop<'b>(&'b mut self) -> GenerateCtx<'b, R>
        where 'a: 'b
    {
        Self::new(self.rng, self.size/2)
    }

    fn gen_size(&mut self) -> usize
        where R: rand::Rng + Sized
    {
        match self.size {
            0 => 0,
            size @ _ if size == <usize>::max_value() => self.rng.gen(),
            size @ _ => self.rng.gen_range(0, size + 1)
        }
    }
}

pub trait Generator {
    type Output;

    fn generate<R: rand::Rng>(&self, &mut GenerateCtx<R>) -> <Self as Generator>::Output;
}

impl <'a, G: Generator> Generator for &'a G {
    type Output = G::Output;

    #[inline]
    fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> Self::Output {
        (*self).generate(ctx)
    }
}

#[derive(Copy, Clone)]
pub struct Constant<T>(pub T);

impl <T: Clone> Generator for Constant<T> {
    type Output = T;

    #[inline]
    fn generate<R: rand::Rng>(&self, _: &mut GenerateCtx<R>) -> T {
        self.0.clone()
    }
}

macro_rules! tuple_impls {
    ($($name:ident),*) => {
        impl <$($name: Generator),*> Generator for ($($name,)*) {
            type Output = ($($name::Output,)*);

            #[inline]
            #[allow(unused_variables, non_snake_case)]
            fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> Self::Output {
                let ( $(ref $name,)* ) = *self;
                ($($name.generate(ctx),)*)
            }
        }
    }
}

macro_tuples_impl!{tuple_impls}

pub struct FnGenerator<T> {
    f: Box<Fn(&mut GenerateCtx<rand::Rng>) -> T>
}

impl <T> Generator for FnGenerator<T> {
    type Output = T;

    fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> T {
        let mut boxed_ctx = GenerateCtx { rng: ctx.rng as &mut rand::Rng, size: ctx.size };
        (self.f)(&mut boxed_ctx)
    }
}

impl <T> FnGenerator<T> {
    fn new<F: Fn(&mut GenerateCtx<rand::Rng>) -> T + 'static>(f: F) -> FnGenerator<T> {
        FnGenerator { f: Box::new(f) }
    }
}

fn testy() -> FnGenerator<u32> {
    FnGenerator::new(|ctx| ctx.rng.next_u32())
}

pub struct IntegerGenerator<X>(PhantomData<fn() -> X>);

macro_rules! int_impls {
    ($($ty:ty),*) => {
        $(
            impl Generator for IntegerGenerator<$ty>
            {
                type Output = $ty;

                fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> $ty {
                    if ctx.size == 0 { return 0; }
                    let cast_size = <$ty>::from_usize(ctx.size);
                    let upper_bound = cast_size.and_then(|s| s.checked_add(1));
                    let lower_bound = cast_size.and_then(|s| s.checked_mul(-1));
                    match (lower_bound, upper_bound) {
                        (Some(lower), Some(upper)) => ctx.rng.gen_range(lower, upper),
                        _ => ctx.rng.gen()
                    }
                }
            }
        )*
    }
}

int_impls! { i8, i16, i32, i64, isize }

pub struct UnsignedIntegerGenerator<X>(PhantomData<fn() -> X>);

macro_rules! uint_impls {
    ($($ty:ty),*) => {
        $(
            impl Generator for UnsignedIntegerGenerator<$ty>
            {
                type Output = $ty;

                fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> $ty {
                    if ctx.size == 0 { return 0; }
                    let upper_bound = <$ty>::from_usize(ctx.size).and_then(|s| s.checked_add(1));
                    match upper_bound {
                        Some(upper) => ctx.rng.gen_range(0, upper),
                        _ => ctx.rng.gen()
                    }
                }
            }
        )*
    }
}

uint_impls! { u8, u16, u32, u64, usize, i8, i16, i32, i64, isize }

pub struct FromIteratorGenerator<C, G> {
    generator: G,
    _marker: PhantomData<fn() -> C>
}

impl <C, G> FromIteratorGenerator<C, G>
    where FromIteratorGenerator<C, G>: Generator
{
    fn new(generator: G) -> Self {
        FromIteratorGenerator { generator: generator, _marker: PhantomData }
    }
}

impl <C, G> Generator for FromIteratorGenerator<C, G>
    where G: Generator,
          C: FromIterator<G::Output>
{
    type Output = C;

    fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> Self::Output {
        let size = ctx.gen_size();
        let mut chopped_ctx = ctx.chop();
        (0..size).map(|_| self.generator.generate(&mut chopped_ctx)).collect()
    }
}

pub struct OptionGenerator<C, G> {
    generator: G,
    _marker: PhantomData<fn() -> C>
}

impl <C, G> OptionGenerator<C, G>
    where OptionGenerator<C, G>: Generator
{
    fn new(generator: G) -> Self { OptionGenerator { generator: generator, _marker: PhantomData } }
}

impl <G: Generator> Generator for OptionGenerator<Option<G::Output>, G> {
    type Output = Option<G::Output>;

    fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> Self::Output {
        match ctx.rng.gen() {
            true => Some(self.generator.generate(ctx)),
            false => None
        }
    }
}

pub struct ResultGenerator<GOk, GErr> {
    g_ok: GOk,
    g_err: GErr
}

impl <GOk: Generator, GErr: Generator> Generator for ResultGenerator<GOk, GErr> {
    type Output = Result<GOk::Output, GErr::Output>;

    fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> Self::Output {
        match ctx.rng.gen() {
            true => Ok(self.g_ok.generate(ctx)),
            false => Err(self.g_err.generate(ctx))
        }
    }
}

pub struct SimpleGenerator<T>(PhantomData<fn() -> T>);
impl Generator for SimpleGenerator<bool> {
    type Output = bool;

    fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> Self::Output {
        ctx.rng.gen()
    }
}
<commit_msg>Adding a couple of tests.<commit_after>use std::marker::PhantomData;
use std::iter::FromIterator;

use rand;
use num::traits::FromPrimitive;

pub struct GenerateCtx<'a, R: ?Sized + 'a> {
    pub rng: &'a mut R,
    pub size: usize
}

impl <'a, R: ?Sized + 'a> GenerateCtx<'a, R> {
    fn new(rng: &'a mut R, size: usize) -> Self {
        GenerateCtx { rng: rng, size: size }
    }

    #[inline]
    fn chop<'b>(&'b mut self) -> GenerateCtx<'b, R>
        where 'a: 'b
    {
        Self::new(self.rng, self.size/2)
    }

    fn gen_size(&mut self) -> usize
        where R: rand::Rng + Sized
    {
        match self.size {
            0 => 0,
            size @ _ if size == <usize>::max_value() => self.rng.gen(),
            size @ _ => self.rng.gen_range(0, size + 1)
        }
    }
}

pub trait Generator {
    type Output;

    fn generate<R: rand::Rng>(&self, &mut GenerateCtx<R>) -> <Self as Generator>::Output;
}

impl <'a, G: Generator> Generator for &'a G {
    type Output = G::Output;

    #[inline]
    fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> Self::Output {
        (*self).generate(ctx)
    }
}

#[derive(Copy, Clone)]
pub struct Constant<T>(pub T);

impl <T: Clone> Generator for Constant<T> {
    type Output = T;

    #[inline]
    fn generate<R: rand::Rng>(&self, _: &mut GenerateCtx<R>) -> T {
        self.0.clone()
    }
}

macro_rules! tuple_impls {
    ($($name:ident),*) => {
        impl <$($name: Generator),*> Generator for ($($name,)*) {
            type Output = ($($name::Output,)*);

            #[inline]
            #[allow(unused_variables, non_snake_case)]
            fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> Self::Output {
                let ( $(ref $name,)* ) = *self;
                ($($name.generate(ctx),)*)
            }
        }
    }
}

macro_tuples_impl!{tuple_impls}

pub struct FnGenerator<T> {
    f: Box<Fn(&mut GenerateCtx<rand::Rng>) -> T>
}

impl <T> Generator for FnGenerator<T> {
    type Output = T;

    fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> T {
        let mut boxed_ctx = GenerateCtx { rng: ctx.rng as &mut rand::Rng, size: ctx.size };
        (self.f)(&mut boxed_ctx)
    }
}

impl <T> FnGenerator<T> {
    fn new<F: Fn(&mut GenerateCtx<rand::Rng>) -> T + 'static>(f: F) -> FnGenerator<T> {
        FnGenerator { f: Box::new(f) }
    }
}

pub struct IntegerGenerator<X>(PhantomData<fn() -> X>);

impl <X> IntegerGenerator<X> where IntegerGenerator<X>: Generator
{
    fn new() -> Self { IntegerGenerator(PhantomData) }
}

macro_rules! int_impls {
    ($($ty:ty),*) => {
        $(
            impl Generator for IntegerGenerator<$ty>
            {
                type Output = $ty;

                fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> $ty {
                    if ctx.size == 0 { return 0; }
                    let cast_size = <$ty>::from_usize(ctx.size);
                    let upper_bound = cast_size.and_then(|s| s.checked_add(1));
                    let lower_bound = cast_size.and_then(|s| s.checked_mul(-1));
                    match (lower_bound, upper_bound) {
                        (Some(lower), Some(upper)) => ctx.rng.gen_range(lower, upper),
                        _ => ctx.rng.gen()
                    }
                }
            }
        )*
    }
}

int_impls! { i8, i16, i32, i64, isize }

pub struct UnsignedIntegerGenerator<X>(PhantomData<fn() -> X>);

impl <X> UnsignedIntegerGenerator<X> where UnsignedIntegerGenerator<X>: Generator
{
    fn new() -> Self { UnsignedIntegerGenerator(PhantomData) }
}

macro_rules! uint_impls {
    ($($ty:ty),*) => {
        $(
            impl Generator for UnsignedIntegerGenerator<$ty>
            {
                type Output = $ty;

                fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> $ty {
                    if ctx.size == 0 { return 0; }
                    let upper_bound = <$ty>::from_usize(ctx.size).and_then(|s| s.checked_add(1));
                    match upper_bound {
                        Some(upper) => ctx.rng.gen_range(0, upper),
                        _ => ctx.rng.gen()
                    }
                }
            }
        )*
    }
}

uint_impls! { u8, u16, u32, u64, usize, i8, i16, i32, i64, isize }

pub struct FromIteratorGenerator<C, G> {
    generator: G,
    _marker: PhantomData<fn() -> C>
}

impl <C, G> FromIteratorGenerator<C, G>
    where FromIteratorGenerator<C, G>: Generator
{
    fn new(generator: G) -> Self {
        FromIteratorGenerator { generator: generator, _marker: PhantomData }
    }
}

impl <C, G> Generator for FromIteratorGenerator<C, G>
    where G: Generator,
          C: FromIterator<G::Output>
{
    type Output = C;

    fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> Self::Output {
        let size = ctx.gen_size();
        let mut chopped_ctx = ctx.chop();
        (0..size).map(|_| self.generator.generate(&mut chopped_ctx)).collect()
    }
}

pub struct OptionGenerator<C, G> {
    generator: G,
    _marker: PhantomData<fn() -> C>
}

impl <C, G> OptionGenerator<C, G>
    where OptionGenerator<C, G>: Generator
{
    fn new(generator: G) -> Self { OptionGenerator { generator: generator, _marker: PhantomData } }
}

impl <G: Generator> Generator for OptionGenerator<Option<G::Output>, G> {
    type Output = Option<G::Output>;

    fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> Self::Output {
        match ctx.rng.gen() {
            true => Some(self.generator.generate(ctx)),
            false => None
        }
    }
}

pub struct ResultGenerator<GOk, GErr> {
    g_ok: GOk,
    g_err: GErr
}

impl <GOk: Generator, GErr: Generator> Generator for ResultGenerator<GOk, GErr> {
    type Output = Result<GOk::Output, GErr::Output>;

    fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> Self::Output {
        match ctx.rng.gen() {
            true => Ok(self.g_ok.generate(ctx)),
            false => Err(self.g_err.generate(ctx))
        }
    }
}

pub struct SimpleGenerator<T>(PhantomData<fn() -> T>);
impl Generator for SimpleGenerator<bool> {
    type Output = bool;

    fn generate<R: rand::Rng>(&self, ctx: &mut GenerateCtx<R>) -> Self::Output {
        ctx.rng.gen()
    }
}

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

    use rand;

    #[test]
    fn gen_unit() {
        let mut rng = rand::thread_rng(); let mut ctx = GenerateCtx::new(&mut rng, 5);
        assert_eq!(().generate(&mut ctx), ());
    }

    #[test]
    fn gen_u8() {
        let mut rng = rand::thread_rng(); let mut ctx = GenerateCtx::new(&mut rng, 5);
        let gen = UnsignedIntegerGenerator::<u8>::new();
        rep(&mut || { let n = gen.generate(&mut ctx); assert!(n <= 5); });
    }

    #[test]
    fn gen_i8() {
        let mut rng = rand::thread_rng(); let mut ctx = GenerateCtx::new(&mut rng, 5);
        let gen = UnsignedIntegerGenerator::<i8>::new();
        rep(&mut || { let n = gen.generate(&mut ctx); assert!((n >= -5) && (n <= 5)); });
    }

    fn rep<F>(f: &mut F) where F: FnMut() -> () {
        for _ in 0..100 {
            f()
        }
    }
}
<|endoftext|>


<commit_before>use socket::Socket;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Source {
    DISABLED,
    NOW,
    CH1_PE,
    CH1_NE,
    CH2_PE,
    CH2_NE,
    EXT_PE,
    EXT_NE,
    AWG_PE,
    AWG_NE,
}

impl ::std::fmt::Display for Source {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &Source::DISABLED => "DISABLED",
            &Source::NOW => "NOW",
            &Source::CH1_PE => "CH1_PE",
            &Source::CH1_NE => "CH1_NE",
            &Source::CH2_PE => "CH2_PE",
            &Source::CH2_NE => "CH2_NE",
            &Source::EXT_PE => "EXT_PE",
            &Source::EXT_NE => "EXT_NE",
            &Source::AWG_PE => "AWG_PE",
            &Source::AWG_NE => "AWG_NE",
        };

        write!(f, "{}", display)
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum State {
    WAIT,
    TD,
    UNKNOW,
}

impl ::std::convert::From<String> for State {
    fn from(s: String) -> Self {
        match s.as_str() {
            "WAIT" => State::WAIT,
            "TD" => State::TD,
            _ => State::UNKNOW,
        }
    }
}

#[derive(Clone)]
pub struct Trigger {
    socket: ::std::cell::RefCell<Socket>,
}

impl Trigger {
    pub fn new(socket: Socket) -> Self {
        Trigger {
            socket: ::std::cell::RefCell::new(socket),
        }
    }

    /**
     * Trigger immediately or set trigger source & edge.
     */
    pub fn enable(&self, source: Source) {
        self.send(format!("ACQ:TRIG {}", source));
    }

    /**
     * Disable triggering.
     */
    pub fn disable(&self) {
        self.enable(Source::DISABLED);
    }

    /**
     *  Get trigger status.
     *
     *  If DISABLED -> TD else WAIT.
     */
    pub fn get_state(&self) -> State {
        self.send("ACQ:TRIG:STAT?");

        self.receive()
            .into()
    }

    /**
     * Set trigger delay in samples.
     */
    pub fn set_delay(&self, delay: u16) {
        self.send(format!("ACQ:TRIG:DLY {}", delay));
    }

    /**
     * Get trigger delay in samples.
     */
    pub fn get_delay(&self) -> u16 {
        self.send("ACQ:TRIG:DLY?");

        self.receive()
            .parse()
            .unwrap()
    }

    /**
     * Set trigger delay in ns.
     */
    pub fn set_delay_in_ns(&self, delay: u8) {
        self.send(format!("ACQ:TRIG:DLY:NS {}", delay));
    }

    /**
     * Get trigger delay in ns.
     */
    pub fn get_delay_in_ns(&self) -> u8 {
        self.send("ACQ:TRIG:DLY:NS?");

        self.receive()
            .replace("ns", "")
            .parse()
            .unwrap()
    }

    /**
     * Set trigger level in mV.
     */
    pub fn set_level(&self, level: u8) {
        self.send(format!("ACQ:TRIG:LEV {}", level));
    }

    /**
     * Get trigger level in mV.
     */
    pub fn get_level(&self) -> String {
        self.send("ACQ:TRIG:LEV?");

        self.receive()
    }

    fn send<D>(&self, message: D)
        where D: ::std::fmt::Display
    {
        let mut socket = self.socket.borrow_mut();

        socket.send(message);
    }

    fn receive(&self) -> String {
        let mut socket = self.socket.borrow_mut();

        socket.receive()
    }
}

#[cfg(test)]
mod test {
    macro_rules! trigger_assert {
        ($f:ident, $e:expr) => {
            let (rx, trigger) = create_trigger();

            trigger.$f();
            assert_eq!($e, rx.recv().unwrap());
        }
    }

    #[test]
    fn test_enable() {
        let (rx, trigger) = create_trigger();

        trigger.enable(::trigger::Source::NOW);
        assert_eq!("ACQ:TRIG NOW\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_disable() {
        trigger_assert!(disable, "ACQ:TRIG DISABLED\r\n");
    }

    #[test]
    fn test_get_state() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_state(), ::trigger::State::WAIT);
    }

    #[test]
    fn test_set_delay() {
        let (rx, trigger) = create_trigger();

        trigger.set_delay(0);
        assert_eq!("ACQ:TRIG:DLY 0\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_delay() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_delay(), 2314);
    }

    #[test]
    fn test_set_delay_in_ns() {
        let (rx, trigger) = create_trigger();

        trigger.set_delay_in_ns(0);
        assert_eq!("ACQ:TRIG:DLY:NS 0\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_delay_in_ns() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_delay_in_ns(), 128);
    }

    #[test]
    fn test_set_level() {
        let (rx, trigger) = create_trigger();

        trigger.set_level(0);
        assert_eq!("ACQ:TRIG:LEV 0\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_level() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_level().as_str(), "123mV");
    }

    fn create_trigger() -> (::std::sync::mpsc::Receiver<String>, ::trigger::Trigger) {
        let (addr, rx) = ::test::launch_server();
        let socket = ::socket::Socket::new(addr);

        (rx, ::trigger::Trigger::new(socket))
    }
}
<commit_msg>[trigger] trigger::get_level() returns u8<commit_after>use socket::Socket;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Source {
    DISABLED,
    NOW,
    CH1_PE,
    CH1_NE,
    CH2_PE,
    CH2_NE,
    EXT_PE,
    EXT_NE,
    AWG_PE,
    AWG_NE,
}

impl ::std::fmt::Display for Source {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &Source::DISABLED => "DISABLED",
            &Source::NOW => "NOW",
            &Source::CH1_PE => "CH1_PE",
            &Source::CH1_NE => "CH1_NE",
            &Source::CH2_PE => "CH2_PE",
            &Source::CH2_NE => "CH2_NE",
            &Source::EXT_PE => "EXT_PE",
            &Source::EXT_NE => "EXT_NE",
            &Source::AWG_PE => "AWG_PE",
            &Source::AWG_NE => "AWG_NE",
        };

        write!(f, "{}", display)
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum State {
    WAIT,
    TD,
    UNKNOW,
}

impl ::std::convert::From<String> for State {
    fn from(s: String) -> Self {
        match s.as_str() {
            "WAIT" => State::WAIT,
            "TD" => State::TD,
            _ => State::UNKNOW,
        }
    }
}

#[derive(Clone)]
pub struct Trigger {
    socket: ::std::cell::RefCell<Socket>,
}

impl Trigger {
    pub fn new(socket: Socket) -> Self {
        Trigger {
            socket: ::std::cell::RefCell::new(socket),
        }
    }

    /**
     * Trigger immediately or set trigger source & edge.
     */
    pub fn enable(&self, source: Source) {
        self.send(format!("ACQ:TRIG {}", source));
    }

    /**
     * Disable triggering.
     */
    pub fn disable(&self) {
        self.enable(Source::DISABLED);
    }

    /**
     *  Get trigger status.
     *
     *  If DISABLED -> TD else WAIT.
     */
    pub fn get_state(&self) -> State {
        self.send("ACQ:TRIG:STAT?");

        self.receive()
            .into()
    }

    /**
     * Set trigger delay in samples.
     */
    pub fn set_delay(&self, delay: u16) {
        self.send(format!("ACQ:TRIG:DLY {}", delay));
    }

    /**
     * Get trigger delay in samples.
     */
    pub fn get_delay(&self) -> u16 {
        self.send("ACQ:TRIG:DLY?");

        self.receive()
            .parse()
            .unwrap()
    }

    /**
     * Set trigger delay in ns.
     */
    pub fn set_delay_in_ns(&self, delay: u8) {
        self.send(format!("ACQ:TRIG:DLY:NS {}", delay));
    }

    /**
     * Get trigger delay in ns.
     */
    pub fn get_delay_in_ns(&self) -> u8 {
        self.send("ACQ:TRIG:DLY:NS?");

        self.receive()
            .replace("ns", "")
            .parse()
            .unwrap()
    }

    /**
     * Set trigger level in mV.
     */
    pub fn set_level(&self, level: u8) {
        self.send(format!("ACQ:TRIG:LEV {}", level));
    }

    /**
     * Get trigger level in mV.
     */
    pub fn get_level(&self) -> u8 {
        self.send("ACQ:TRIG:LEV?");

        self.receive()
            .replace("mV", "")
            .parse()
            .unwrap()
    }

    fn send<D>(&self, message: D)
        where D: ::std::fmt::Display
    {
        let mut socket = self.socket.borrow_mut();

        socket.send(message);
    }

    fn receive(&self) -> String {
        let mut socket = self.socket.borrow_mut();

        socket.receive()
    }
}

#[cfg(test)]
mod test {
    macro_rules! trigger_assert {
        ($f:ident, $e:expr) => {
            let (rx, trigger) = create_trigger();

            trigger.$f();
            assert_eq!($e, rx.recv().unwrap());
        }
    }

    #[test]
    fn test_enable() {
        let (rx, trigger) = create_trigger();

        trigger.enable(::trigger::Source::NOW);
        assert_eq!("ACQ:TRIG NOW\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_disable() {
        trigger_assert!(disable, "ACQ:TRIG DISABLED\r\n");
    }

    #[test]
    fn test_get_state() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_state(), ::trigger::State::WAIT);
    }

    #[test]
    fn test_set_delay() {
        let (rx, trigger) = create_trigger();

        trigger.set_delay(0);
        assert_eq!("ACQ:TRIG:DLY 0\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_delay() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_delay(), 2314);
    }

    #[test]
    fn test_set_delay_in_ns() {
        let (rx, trigger) = create_trigger();

        trigger.set_delay_in_ns(0);
        assert_eq!("ACQ:TRIG:DLY:NS 0\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_delay_in_ns() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_delay_in_ns(), 128);
    }

    #[test]
    fn test_set_level() {
        let (rx, trigger) = create_trigger();

        trigger.set_level(0);
        assert_eq!("ACQ:TRIG:LEV 0\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_level() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_level(), 123);
    }

    fn create_trigger() -> (::std::sync::mpsc::Receiver<String>, ::trigger::Trigger) {
        let (addr, rx) = ::test::launch_server();
        let socket = ::socket::Socket::new(addr);

        (rx, ::trigger::Trigger::new(socket))
    }
}
<|endoftext|>