<commit_before>///! Encapsulates the implementation of the syscall collectors.

use libc;
use num_cpus;
use chrono::{NaiveDateTime, DateTime, UTC, TimeZone};
use nix::Error as NixError;
use nix::errno;
use nix::fcntl::{open, O_RDWR, O_SYNC};
use nix::unistd::close;
use nix::sys::stat::S_IRUSR;
use nix::sys::mman::{mmap, munmap};
use nix::sys::mman::{MAP_SHARED, PROT_READ, PROT_WRITE};
use std::ptr;
use std::mem::size_of;

use syscall::{SyscallHdr, Syscall};
use syscall::syscall_table::SyscallTable;
use error::{Error, Result};

const RING_BUF_SIZE: usize = 8 * 1024 * 1024;
const BUFFER_EMPTY_WAIT_TIME_MS: u32 = 30;
const MAX_N_CONSECUTIVE_WAITS: usize = 4;

const PPM_IOCTL_MAGIC: u8 = 's' as u8;
const PPM_IOCTL_DISABLE_CAPTURE: u8 = 1;
const PPM_IOCTL_ENABLE_CAPTURE: u8 = 1;

ioctl!(none ioctl_start with PPM_IOCTL_MAGIC, PPM_IOCTL_ENABLE_CAPTURE);
ioctl!(none ioctl_stop with PPM_IOCTL_MAGIC, PPM_IOCTL_DISABLE_CAPTURE);

#[repr(C)]
struct RingBufferInfo {
    head: u32,
    tail: u32,
    num_syscalls: u64,
    num_drops_buffer: u64,
    num_drops_pf: u64,
    num_preemptions: u64,
    num_context_switches: u64
}

struct RingBufferDev {
    fd: i32,
    buffer: *mut libc::c_char,
    buffer_info: *mut RingBufferInfo,
    last_readsize: u32,
    next_syscall_hdr: *mut libc::c_char,
    len: u32
}

pub trait Collector {
    fn start(&mut self) -> Result<usize>;

    fn stop(&mut self) -> Result<()>;

    fn next(&mut self) -> Option<Syscall>;
}

pub struct RingBufferCollector {
    devs: Vec<RingBufferDev>,
    consecutive_waits: usize,
    syscall_table: SyscallTable
}

/// The implementation of the syscall collector based on the
/// Sysdig's rock solid kernel driver. The driver uses circular / ring
/// buffer data structure to publish the intercepted system calls from the
/// kernel space. This aims to provide extremely high throughput and efficiency.
impl Collector for RingBufferCollector {
    /// Opens available ring buffer devices and maps the corresponding
    /// buffer content and the buffer information to a memory region
    /// of the calling thread.
    ///
    /// If the previous operations are done successfully, it sends the
    /// IO code to the underlying device to start the capture from the
    /// kernel driver.
    ///
    /// Returns `Result::Ok(i)` where `i` holds the number of devices (CPUs) detected on the
    /// system or `Result::Err(e)` where `e` is the payload with a string containing the detailed
    /// error message.
    ///
    /// # Example
    ///
    /// ```
    /// let collector = RingBufferCollector::new();
    /// match collector.start() {
    ///     Ok(num_devs) => {
    ///
    ///     },
    ///     Err(e) => {
    ///     }
    /// }
    /// ```
    ///
    fn start(&mut self) -> Result<usize> {
        let num_devs = num_cpus::get();
        let len = RING_BUF_SIZE * 2;

        // open and map each of the char devices
        // which represent the ring buffers
        for i in 0..num_devs {
            match open(format!("/dev/sysdig{}", i).as_str(), O_RDWR | O_SYNC, S_IRUSR) {
                Ok(fd) => {
                    let buffer = mmap(ptr::null_mut(),
                                      len,
                                      PROT_READ,
                                      MAP_SHARED, fd, 0);
                    let buffer_info = mmap(ptr::null_mut(), size_of::<RingBufferInfo>(),
                                           PROT_READ | PROT_WRITE,
                                           MAP_SHARED, fd, 0);
                    if buffer.is_err() || buffer_info.is_err() {
                        munmap(buffer.unwrap() as *mut libc::c_void, len);
                        close(fd);
                        return Err(Error::RingBufferMapping);
                    }
                    let dev = RingBufferDev {
                        fd: fd,
                        buffer: buffer.unwrap() as *mut libc::c_char,
                        buffer_info: buffer_info.unwrap() as *mut RingBufferInfo,
                        last_readsize: 0,
                        next_syscall_hdr: ptr::null_mut(),
                        len: 0
                    };

                    self.devs.push(dev);
                    // send the ioctl code to start the capture
                    unsafe { ioctl_start(fd); }

                },
                Err(e) => {
                    match e {
                        NixError::Sys(err) => {
                            if err == errno::ENODEV {
                                continue;
                            } else if err == errno::EBUSY {
                                return Err(Error::TooManyCollectors);
                            } else {
                                return Err(Error::DeviceError);
                            }
                        },
                        NixError::InvalidPath => {}
                    }
                }
            }
        }
        Ok(num_devs)
    }

    fn stop(&mut self) -> Result<()> {
        let len = RING_BUF_SIZE * 2;
        for dev in &self.devs {
            unsafe { ioctl_stop(dev.fd); }
            munmap(dev.buffer as *mut libc::c_void, len);
            munmap(dev.buffer_info as *mut libc::c_void, size_of::<RingBufferInfo>());
            close(dev.fd);
        }
        Ok(())
    }

    /// Consumes the next available syscall event from the ring buffer. If the buffers
    /// are empty, the head and tail are updated accordingly to pick up the
    /// new generated syscall events. When `consecutive_waits` is above the defined threshold, the
    /// main thread is suspended for `BUFFER_EMPTY_WAIT_TIME_MS` milliseconds. This can give a
    /// chance for buffers to refill.
    ///
    /// Returns `Some(syscall)` where `syscall` stores a plethora of information about the
    /// intercepted syscall event. On error it returns `None`.
    ///
    /// # Example
    ///
    /// ```
    /// let collector = RingBufferCollector::new();
    /// loop {
    ///     match collector.next() {
    ///         Some(syscall) => {
    ///             println!("syscall name {}", syscall.name);
    ///         },
    ///         None => {}
    ///     }
    /// }
    /// ```
    ///
    fn next(&mut self) -> Option<Syscall> {
        let mut cpu = None;
        let mut syscall_hdr = ptr::null_mut();

        for (j, dev) in self.devs.iter().enumerate() {
            if dev.len == 0 {
                continue;
            } else {
                syscall_hdr = dev.next_syscall_hdr as *mut SyscallHdr;
                cpu = Some(j);
            }
        }
        if cpu != None {
            let cpuid = cpu.unwrap();
            unsafe {
                assert!(self.devs[cpuid].len >= (*syscall_hdr).len);
                let len = (*syscall_hdr).len as isize;
                self.devs[cpuid].len -= (*syscall_hdr).len;
                self.devs[cpuid].next_syscall_hdr = self.devs[cpuid].next_syscall_hdr.offset(len);
            };
        } else {
            // in case buffers are empty
            if self.check_next_wait() {
                unsafe { libc::usleep(BUFFER_EMPTY_WAIT_TIME_MS * 1000); }
                self.consecutive_waits += 1;
            }

            // update ring buffer's tail/head
            for mut dev in &mut self.devs {
                let buffer_info = dev.buffer_info;
                let ttail: usize = unsafe { ((*buffer_info).tail + dev.last_readsize) as usize };
                if ttail < RING_BUF_SIZE {
                    unsafe { (*dev.buffer_info).tail = ttail as u32 }
                } else {
                    unsafe { (*dev.buffer_info).tail = (ttail - RING_BUF_SIZE) as u32 }
                }

                let read_size = RingBufferCollector::get_buffer_readsize(buffer_info);
                dev.last_readsize = read_size;
                dev.len = read_size;

                unsafe { dev.next_syscall_hdr = dev.buffer.offset(ttail as isize); }
            }
        }

        // get syscall metadata and create an instance of
        // the `Syscall` structure which contains required syscall info
        if syscall_hdr.is_null() {
            None
        } else {
            let id = unsafe { (*syscall_hdr).id };
            let ts = unsafe { (*syscall_hdr).ts };

            match self.syscall_table.get_syscall_meta(id as usize) {
                Some(meta) => {
                    let timestamp = NaiveDateTime::from_timestamp((ts / 1000000000) as i64, 0);
                    let syscall = Syscall {
                        ts: DateTime::<UTC>::from_utc(timestamp, UTC),
                        name: meta.name.to_string(),
                        params: meta.build_params(syscall_hdr),
                    };
                    return Some(syscall);

                },
                None => { return None; }
            }
        }
    }
}

impl RingBufferCollector {

    pub fn new() -> RingBufferCollector {
        RingBufferCollector {
            devs: Vec::<RingBufferDev>::new(),
            consecutive_waits: 0,
            syscall_table: SyscallTable::default()
        }
    }

    /// Determines if we should wait while the ring buffers
    /// are refilled with new system call events. Returns `true` if the
    /// condition is met, or `false` otherwise.
    fn check_next_wait(&mut self) -> bool {
        let mut res = true;
        for dev in &self.devs {
            let read_size = RingBufferCollector::get_buffer_readsize(dev.buffer_info);

            if read_size > 20000 {
                self.consecutive_waits = 0;
                res = false;
            }
        }
        if res == false {
            return false
        }

        if self.consecutive_waits >= MAX_N_CONSECUTIVE_WAITS {
            self.consecutive_waits = 0;
            return false;
        } else {
            return true;
        }
    }

    /// Calculates the new size of the buffer to be read.
    fn get_buffer_readsize(buffer_info: *mut RingBufferInfo) -> u32 {

        let head: usize = unsafe { (*buffer_info).head as usize };
        let tail: usize = unsafe { (*buffer_info).tail as usize };

        (if tail > head { RING_BUF_SIZE - tail + head } else { head - tail }) as u32

    }
}

<commit_msg>refactored<commit_after>///! Encapsulates the implementation of the syscall collectors.

use libc;
use num_cpus;
use chrono::{NaiveDateTime, DateTime, UTC, TimeZone};
use nix::Error as NixError;
use nix::errno;
use nix::fcntl::{open, O_RDWR, O_SYNC};
use nix::unistd::close;
use nix::sys::stat::S_IRUSR;
use nix::sys::mman::{mmap, munmap};
use nix::sys::mman::{MAP_SHARED, PROT_READ, PROT_WRITE};
use std::ptr;
use std::mem::size_of;

use syscall::{Syscall, SyscallInfo};
use syscall::syscall_table::SyscallTable;
use error::{Error, Result};

const RING_BUF_SIZE: usize = 8 * 1024 * 1024;
const BUFFER_EMPTY_WAIT_TIME_MS: u32 = 30;
const MAX_N_CONSECUTIVE_WAITS: usize = 4;

const PPM_IOCTL_MAGIC: u8 = 's' as u8;
const PPM_IOCTL_DISABLE_CAPTURE: u8 = 1;
const PPM_IOCTL_ENABLE_CAPTURE: u8 = 1;

ioctl!(none ioctl_start with PPM_IOCTL_MAGIC, PPM_IOCTL_ENABLE_CAPTURE);
ioctl!(none ioctl_stop with PPM_IOCTL_MAGIC, PPM_IOCTL_DISABLE_CAPTURE);

#[repr(C)]
struct RingBufferInfo {
    head: u32,
    tail: u32,
    num_syscalls: u64,
    num_drops_buffer: u64,
    num_drops_pf: u64,
    num_preemptions: u64,
    num_context_switches: u64
}

struct RingBufferDev {
    fd: i32,
    buffer: *mut libc::c_char,
    buffer_info: *mut RingBufferInfo,
    last_readsize: u32,
    next_syscall: *mut libc::c_char,
    len: u32
}

pub trait Collector {
    fn start(&mut self) -> Result<usize>;

    fn stop(&mut self) -> Result<()>;

    fn next(&mut self) -> Option<SyscallInfo>;
}

pub struct RingBufferCollector {
    devs: Vec<RingBufferDev>,
    consecutive_waits: usize,
    syscall_table: SyscallTable
}

/// The implementation of the syscall collector based on the
/// Sysdig's rock solid kernel driver. The driver uses circular / ring
/// buffer data structure to publish the intercepted system calls from the
/// kernel space. This aims to provide extremely high throughput and efficiency.
impl Collector for RingBufferCollector {
    /// Opens the available ring buffer devices and maps the corresponding
    /// buffer content and the buffer information to a memory region
    /// of the calling thread.
    ///
    /// If the previous operations are done successfully, it sends the
    /// IO code to the underlying device to start the capture from the
    /// kernel driver.
    ///
    /// Returns `Result::Ok(i)` where `i` is the number of devices (CPUs) detected on the
    /// system or `Result::Err(e)` where `e` is the payload with a string containing the detailed
    /// error message.
    ///
    /// # Example
    ///
    /// ```
    /// let collector = RingBufferCollector::new();
    /// match collector.start() {
    ///     Ok(num_devs) => {
    ///
    ///     },
    ///     Err(e) => {
    ///     }
    /// }
    /// ```
    ///
    fn start(&mut self) -> Result<usize> {
        let num_devs = num_cpus::get();
        let len = RING_BUF_SIZE * 2;

        // open and map each of the char devices
        // which represent the ring buffers
        for i in 0..num_devs {
            match open(format!("/dev/sysdig{}", i).as_str(), O_RDWR | O_SYNC, S_IRUSR) {
                Ok(fd) => {
                    let buffer = mmap(ptr::null_mut(),
                                      len,
                                      PROT_READ,
                                      MAP_SHARED, fd, 0);
                    let buffer_info = mmap(ptr::null_mut(), size_of::<RingBufferInfo>(),
                                           PROT_READ | PROT_WRITE,
                                           MAP_SHARED, fd, 0);
                    if buffer.is_err() || buffer_info.is_err() {
                        munmap(buffer.unwrap() as *mut libc::c_void, len);
                        close(fd);
                        return Err(Error::RingBufferMapping);
                    }
                    let dev = RingBufferDev {
                        fd: fd,
                        buffer: buffer.unwrap() as *mut libc::c_char,
                        buffer_info: buffer_info.unwrap() as *mut RingBufferInfo,
                        last_readsize: 0,
                        next_syscall: ptr::null_mut(),
                        len: 0
                    };

                    self.devs.push(dev);
                    // ~ send the ioctl code to start the capture
                    unsafe { ioctl_start(fd); }

                },
                Err(e) => {
                    match e {
                        NixError::Sys(err) => {
                            if err == errno::ENODEV {
                                continue;
                            } else if err == errno::EBUSY {
                                return Err(Error::TooManyCollectors);
                            } else {
                                return Err(Error::DeviceError);
                            }
                        },
                        NixError::InvalidPath => {}
                    }
                }
            }
        }
        Ok(num_devs)
    }

    fn stop(&mut self) -> Result<()> {
        let len = RING_BUF_SIZE * 2;
        for dev in &self.devs {
            unsafe { ioctl_stop(dev.fd); }
            munmap(dev.buffer as *mut libc::c_void, len);
            munmap(dev.buffer_info as *mut libc::c_void, size_of::<RingBufferInfo>());
            close(dev.fd);
        }
        Ok(())
    }

    /// Consumes the next available syscall event from the ring buffer. If the buffers
    /// are empty, the head and tail are updated accordingly to pick up the
    /// new generated syscall events. When `consecutive_waits` is above the defined threshold, the
    /// main thread is suspended for `BUFFER_EMPTY_WAIT_TIME_MS` milliseconds. This can give a
    /// chance for buffers to refill.
    ///
    /// Returns `Some(syscall)` where `syscall` stores a plethora of information about the
    /// intercepted syscall event. On error it returns `None`.
    ///
    /// # Example
    ///
    /// ```
    /// let collector = RingBufferCollector::new();
    /// loop {
    ///     match collector.next() {
    ///         Some(syscall_info) => {
    ///             println!("syscall name {}", syscall_info.name);
    ///         },
    ///         None => {}
    ///     }
    /// }
    /// ```
    ///
    fn next(&mut self) -> Option<SyscallInfo> {
        let mut cpu = None;
        let mut syscall = ptr::null_mut();

        for (j, dev) in self.devs.iter().enumerate() {
            if dev.len == 0 {
                continue;
            } else {
                syscall = dev.next_syscall as *mut Syscall;
                cpu = Some(j);
            }
        }
        if cpu != None {
            let cpuid = cpu.unwrap();
            unsafe {
                assert!(self.devs[cpuid].len >= (*syscall).len);
                let len = (*syscall).len as isize;
                self.devs[cpuid].len -= (*syscall).len;
                self.devs[cpuid].next_syscall = self.devs[cpuid].next_syscall.offset(len);
            };
        } else {
            // ~ in case buffers are empty
            if self.check_next_wait() {
                unsafe { libc::usleep(BUFFER_EMPTY_WAIT_TIME_MS * 1000); }
                self.consecutive_waits += 1;
            }

            for mut dev in &mut self.devs {
                let buffer_info = dev.buffer_info;
                let ttail: usize = unsafe { ((*buffer_info).tail + dev.last_readsize) as usize };
                if ttail < RING_BUF_SIZE {
                    unsafe { (*dev.buffer_info).tail = ttail as u32 }
                } else {
                    unsafe { (*dev.buffer_info).tail = (ttail - RING_BUF_SIZE) as u32 }
                }

                let read_size = Self::get_buffer_readsize(buffer_info);
                dev.last_readsize = read_size;
                dev.len = read_size;

                unsafe { dev.next_syscall = dev.buffer.offset(ttail as isize); }
            }
        }

        if syscall.is_null() {
            None
        } else {
            let id = unsafe { (*syscall).id };
            let ts = unsafe { (*syscall).ts };

            match self.syscall_table.get_syscall_meta(id as usize) {
                Some(meta) => {
                    let timestamp = NaiveDateTime::from_timestamp((ts / 1000000000) as i64, 0);
                    let syscall_info = SyscallInfo {
                        ts: DateTime::<UTC>::from_utc(timestamp, UTC),
                        name: meta.name.to_string(),
                        params: meta.build_params(syscall),
                    };
                    return Some(syscall_info);

                },
                None => { return None; }
            }
        }
    }
}

impl RingBufferCollector {

    pub fn new() -> RingBufferCollector {
        RingBufferCollector {
            devs: Vec::<RingBufferDev>::new(),
            consecutive_waits: 0,
            syscall_table: SyscallTable::default()
        }
    }

    /// Determines if we should wait while the ring buffers
    /// are refilled with new system call events. Returns `true` if the
    /// condition is met, or `false` otherwise.
    fn check_next_wait(&mut self) -> bool {
        let mut res = true;
        for dev in &self.devs {
            let read_size = Self::get_buffer_readsize(dev.buffer_info);

            if read_size > 20000 {
                self.consecutive_waits = 0;
                res = false;
            }
        }
        if res == false {
            return false
        }

        if self.consecutive_waits >= MAX_N_CONSECUTIVE_WAITS {
            self.consecutive_waits = 0;
            return false;
        } else {
            return true;
        }
    }

    /// Calculates the new size of the buffer to be read.
    fn get_buffer_readsize(buffer_info: *mut RingBufferInfo) -> u32 {
        let head: usize = unsafe { (*buffer_info).head as usize };
        let tail: usize = unsafe { (*buffer_info).tail as usize };

        (if tail > head { RING_BUF_SIZE - tail + head } else { head - tail }) as u32

    }
}
<|endoftext|>


<commit_before>// Copyright 2016 GFX developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

#[repr(u64)]
#[allow(non_camel_case_types)]
#[derive(Debug, Copy, Clone)]
pub enum MTLPixelFormat {
   Invalid                = 0,
   A8Unorm                = 1,
   R8Unorm                = 10,
   R8Snorm                = 12,
   R8Uint                 = 13,
   R8Sint                 = 14,
   R16Unorm               = 20,
   R16Snorm               = 22,
   R16Uint                = 23,
   R16Sint                = 24,
   R16Float               = 25,
   RG8Unorm               = 30,
   RG8Snorm               = 32,
   RG8Uint                = 33,
   RG8Sint                = 34,
   R32Uint                = 53,
   R32Sint                = 54,
   R32Float               = 55,
   RG16Unorm              = 60,
   RG16Snorm              = 62,
   RG16Uint               = 63,
   RG16Sint               = 64,
   RG16Float              = 65,
   RGBA8Unorm             = 70,
   RGBA8Unorm_sRGB        = 71,
   RGBA8Snorm             = 72,
   RGBA8Uint              = 73,
   RGBA8Sint              = 74,
   BGRA8Unorm             = 80,
   BGRA8Unorm_sRGB        = 81,
   RGB10A2Unorm           = 90,
   RGB10A2Uint            = 91,
   RG11B10Float           = 92,
   RGB9E5Float            = 93,
   RG32Uint               = 103,
   RG32Sint               = 104,
   RG32Float              = 105,
   RGBA16Unorm            = 110,
   RGBA16Snorm            = 112,
   RGBA16Uint             = 113,
   RGBA16Sint             = 114,
   RGBA16Float            = 115,
   RGBA32Uint             = 123,
   RGBA32Sint             = 124,
   RGBA32Float            = 125,
   BC1_RGBA               = 130,
   BC1_RGBA_sRGB          = 131,
   BC2_RGBA               = 132,
   BC2_RGBA_sRGB          = 133,
   BC3_RGBA               = 134,
   BC3_RGBA_sRGB          = 135,
   BC4_RUnorm             = 140,
   BC4_RSnorm             = 141,
   BC5_RGUnorm            = 142,
   BC5_RGSnorm            = 143,
   BC6H_RGBFloat          = 150,
   BC6H_RGBUfloat         = 151,
   BC7_RGBAUnorm          = 152,
   BC7_RGBAUnorm_sRGB     = 153,
   GBGR422                = 240,
   BGRG422                = 241,
   Depth32Float           = 252,
   Stencil8               = 253,
   Depth24Unorm_Stencil8  = 255,
   Depth32Float_Stencil8  = 260,
}
<commit_msg>Add missing constants from MTLPixelFormat and make it derive Eq and PartialEq<commit_after>// Copyright 2016 GFX developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

#[repr(u64)]
#[allow(non_camel_case_types)]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum MTLPixelFormat {
   Invalid                = 0,
   A8Unorm                = 1,
   R8Unorm                = 10,
   R8Snorm                = 12,
   R8Uint                 = 13,
   R8Sint                 = 14,
   R16Unorm               = 20,
   R16Snorm               = 22,
   R16Uint                = 23,
   R16Sint                = 24,
   R16Float               = 25,
   RG8Unorm               = 30,
   RG8Snorm               = 32,
   RG8Uint                = 33,
   RG8Sint                = 34,
   R32Uint                = 53,
   R32Sint                = 54,
   R32Float               = 55,
   RG16Unorm              = 60,
   RG16Snorm              = 62,
   RG16Uint               = 63,
   RG16Sint               = 64,
   RG16Float              = 65,
   RGBA8Unorm             = 70,
   RGBA8Unorm_sRGB        = 71,
   RGBA8Snorm             = 72,
   RGBA8Uint              = 73,
   RGBA8Sint              = 74,
   BGRA8Unorm             = 80,
   BGRA8Unorm_sRGB        = 81,
   RGB10A2Unorm           = 90,
   RGB10A2Uint            = 91,
   RG11B10Float           = 92,
   RGB9E5Float            = 93,
   BGR10A2Unorm           = 94,
   RG32Uint               = 103,
   RG32Sint               = 104,
   RG32Float              = 105,
   RGBA16Unorm            = 110,
   RGBA16Snorm            = 112,
   RGBA16Uint             = 113,
   RGBA16Sint             = 114,
   RGBA16Float            = 115,
   RGBA32Uint             = 123,
   RGBA32Sint             = 124,
   RGBA32Float            = 125,
   BC1_RGBA               = 130,
   BC1_RGBA_sRGB          = 131,
   BC2_RGBA               = 132,
   BC2_RGBA_sRGB          = 133,
   BC3_RGBA               = 134,
   BC3_RGBA_sRGB          = 135,
   BC4_RUnorm             = 140,
   BC4_RSnorm             = 141,
   BC5_RGUnorm            = 142,
   BC5_RGSnorm            = 143,
   BC6H_RGBFloat          = 150,
   BC6H_RGBUfloat         = 151,
   BC7_RGBAUnorm          = 152,
   BC7_RGBAUnorm_sRGB     = 153,
   GBGR422                = 240,
   BGRG422                = 241,
   Depth16Unorm           = 250,
   Depth32Float           = 252,
   Stencil8               = 253,
   Depth24Unorm_Stencil8  = 255,
   Depth32Float_Stencil8  = 260,
   X32_Stencil8           = 261,
   X24_Stencil8           = 262
}
<|endoftext|>


<commit_before>use std::io;
use std::fs;
use std::io::Write;
use std::io::BufReader;
use std::io::BufRead;

pub fn store_event(name: String, details: Vec<String>) -> io::Result<()> {
    let mut f = try!(fs::File::create("/tmp/data.csv"));
    let mut row_data = details.to_vec();
    row_data.insert(0, name);
    let row_string = row_data.join(",");
    try!(f.write_all(row_string.as_bytes()));
    try!(f.write_all("\n".as_bytes()));
    Ok(())
}

pub fn current_task() -> Result<String, String>{
    let f = fs::File::open("/tmp/data.csv").expect("read failed");
    let f = BufReader::new(f);
    let mut event_line : String = "".to_string();
    for line in f.lines() {
        let line_copy = line.unwrap();
        if line_copy[..5] == "start".to_string() {
            event_line = line_copy.to_string().split(",").nth(0).expect("not found").to_string();
        }
    }
    Ok(event_line)
}
<commit_msg>Append to data file instead of truncating<commit_after>use std::io;
use std::fs;
use std::fs::OpenOptions;
use std::io::Write;
use std::io::BufReader;
use std::io::BufRead;

pub fn store_event(name: String, details: Vec<String>) -> io::Result<()> {
    let mut f = OpenOptions::new()
                .read(true)
                .write(true)
                .create(true)
                .append(true)
                .open("/tmp/data.csv")
                .expect("couldn't open data file");
    let mut row_data = details.to_vec();
    row_data.insert(0, name);
    let row_string = row_data.join(",");
    try!(f.write_all(row_string.as_bytes()));
    try!(f.write_all("\n".as_bytes()));
    Ok(())
}

pub fn current_task() -> Result<String, String>{
    let f = fs::File::open("/tmp/data.csv").expect("read failed");
    let f = BufReader::new(f);
    let mut event_line : String = "".to_string();
    for line in f.lines() {
        let line_copy = line.unwrap();
        if line_copy[..5] == "start".to_string() {
            event_line = line_copy.to_string().split(",").nth(1).expect("not found").to_string();
        }
    }
    Ok(event_line)
}
<|endoftext|>


<commit_before>//! The datum module provides a data structure, Datum, that implements the Unify trait.
//! Datum aims to be a drop-in for any algorithm in ai_kit that operates on the Unify trait.

#[cfg(feature = "with-constraint")]
use constraints::ConstraintValue;
use core;
use std;
use std::collections::{BTreeMap, HashMap};
use std::str;
use utils;

#[derive(Clone, Debug, Serialize, Deserialize, PartialOrd)]
pub enum Datum {
    #[serde(rename = "null")] Nil,
    #[serde(rename = "bool")] Bool(bool),
    #[serde(rename = "str")] String(String),
    #[serde(rename = "int")] Int(i64),
    #[serde(rename = "float")] Float(f64),
    #[serde(rename = "var")] Variable(String),
    #[serde(rename = "vec")] Vector(Vec<Datum>),
    #[serde(rename = "map")] Map(BTreeMap<String, Datum>),
    #[serde(rename = "fn")] Function { head: Box<Datum>, args: Vec<Datum> },
}

impl Datum {
    pub fn to_string(&self) -> Option<String> {
        match *self {
            Datum::String(ref value) => Some(value.clone()),
            _ => None,
        }
    }

    pub fn to_int(&self) -> Option<i64> {
        match *self {
            Datum::Int(ref value) => Some(value.clone()),
            _ => None,
        }
    }

    pub fn to_float(&self) -> Option<f64> {
        match *self {
            Datum::Float(ref f_value) => Some(f_value.clone()),
            Datum::Int(ref i_value) => Some(i_value.clone() as f64),
            _ => None,
        }
    }

    pub fn to_variable(&self) -> Option<String> {
        match *self {
            Datum::Variable(ref value) => Some(value.clone()),
            _ => None,
        }
    }

    pub fn pprint(&self) -> String {
        match *self {
            Datum::Nil => format!("nil"),
            Datum::Bool(ref b) => format!("{}", b),
            Datum::String(ref s) => format!("{}", s),
            Datum::Int(ref i) => format!("{}", i),
            Datum::Float(ref f) => format!("{}", f),
            Datum::Variable(ref v) => format!("{}", v),
            Datum::Vector(ref args) => {
                let elements: Vec<String> = args.iter().map(|e| e.pprint()).collect();
                format!("({})", elements.join(","))
            }
            Datum::Map(ref args) => {
                let elements: Vec<String> = args.iter()
                    .map(|(k, v)| format!("{} => {}", k, v.pprint()))
                    .collect();
                format!("({})", elements.join(","))
            }
            Datum::Function { ref head, ref args } => {
                let elements: Vec<String> = args.iter().map(|e| e.pprint()).collect();
                format!("({} ({}))", head, elements.join(","))
            }
        }
    }
}

impl Default for Datum {
    fn default() -> Self {
        Datum::Nil
    }
}

impl PartialEq for Datum {
    fn eq(&self, other: &Datum) -> bool {
        match *self {
            Datum::Nil => match *other {
                Datum::Nil => true,
                _ => false,
            },
            Datum::Bool(b) => match *other {
                Datum::Bool(b2) => b == b2,
                _ => false,
            },
            Datum::String(ref s) => match *other {
                Datum::String(ref s2) => s == s2,
                _ => false,
            },
            Datum::Int(ref i) => match *other {
                Datum::Int(ref i2) => i == i2,
                _ => false,
            },
            Datum::Float(ref f) => match *other {
                Datum::Float(ref f2) => f == f2,
                _ => false,
            },
            Datum::Variable(ref v) => match *other {
                Datum::Variable(ref v2) => v == v2,
                _ => false,
            },
            Datum::Vector(ref args) => match *other {
                Datum::Vector(ref args2) => args == args2,
                _ => false,
            },
            Datum::Map(ref args) => match *other {
                Datum::Map(ref args2) => args == args2,
                _ => false,
            },
            Datum::Function { ref head, ref args } => match *other {
                Datum::Function {
                    head: ref head2,
                    args: ref args2,
                } if args.len() == args2.len() =>
                {
                    head == head2 && args == args2
                }
                _ => false,
            },
        }
    }
}

impl Eq for Datum {}
impl Ord for Datum {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.partial_cmp(other).unwrap()
    }
}

impl core::BindingsValue for Datum {
    fn variable(s: &String) -> Option<Self> {
        Some(Datum::Variable(s.clone()))
    }
    fn to_variable(&self) -> Option<String> {
        self.to_variable()
    }
}

#[cfg(feature = "with-constraint")]
impl ConstraintValue for Datum {
    fn to_float(&self) -> Option<f64> {
        self.to_float()
    }
    fn float(c: f64) -> Self {
        Datum::Float(c)
    }
}

impl core::Unify<Datum> for Datum {
    fn unify(&self, other: &Datum, bindings: &core::Bindings<Self>) -> Option<core::Bindings<Self>> {
        fn unify_args(args: &Vec<Datum>, args2: &Vec<Datum>, bindings: &core::Bindings<Datum>) -> Option<core::Bindings<Datum>> {
            utils::fold_while_some(
                bindings.clone(),
                &mut args.iter().zip(args2.iter()),
                &|bindings: core::Bindings<Datum>, (ref a, ref b): (&Datum, &Datum)| a.unify(&b, &bindings),
            )
        }

        fn unify_maps(
            args: &BTreeMap<String, Datum>,
            args2: &BTreeMap<String, Datum>,
            bindings: &core::Bindings<Datum>,
        ) -> Option<core::Bindings<Datum>> {
            if args.len() != args2.len() {
                None
            } else {
                utils::fold_while_some(
                    bindings.clone(),
                    &mut args.iter(),
                    &|bindings: core::Bindings<Datum>, (ref k, ref v): (&String, &Datum)| args2.get(*k).and_then(|v2| v.unify(v2, &bindings)),
                )
            }
        }
        match *self {
            Datum::Variable(ref var_name) => bindings.update_bindings(var_name, other),
            Datum::Vector(ref args) => match *other {
                Datum::Vector(ref args2) if args.len() == args2.len() => unify_args(args, args2, &bindings),
                _ => None,
            },
            Datum::Map(ref args) => match *other {
                Datum::Map(ref args2) if args.len() == args2.len() => unify_maps(args, args2, &bindings),
                _ => None,
            },
            _ => match *other {
                Datum::Variable(ref var_name) => bindings.update_bindings(var_name, self),
                _ => {
                    if self == other {
                        Some(bindings.clone())
                    } else {
                        None
                    }
                }
            },
        }
    }

    fn apply_bindings(&self, bindings: &core::Bindings<Self>) -> Option<Self> {
        fn apply_bindings_to_args(args: &Vec<Datum>, bindings: &core::Bindings<Datum>) -> Option<Vec<Datum>> {
            utils::map_while_some(&mut args.iter(), &|arg| {
                arg.apply_bindings(bindings)
            })
        }
        fn apply_bindings_to_map(args: &BTreeMap<String, Datum>, bindings: &core::Bindings<Datum>) -> Option<BTreeMap<String, Datum>> {
            utils::map_while_some(&mut args.iter(), &|(k, v)| {
                v.apply_bindings(bindings)
                    .and_then(|v| Some((k.clone(), v)))
            }).and_then(|tuple_vec| Some(tuple_vec.into_iter().collect::<BTreeMap<String, Datum>>()))
        }
        match *self {
            Datum::Variable(ref var_name) => bindings
                .get_binding(var_name)
                .or_else(|| Some(self.clone())),
            Datum::Vector(ref args) => apply_bindings_to_args(args, bindings).and_then(|args| Some(Datum::Vector(args))),
            Datum::Map(ref args) => apply_bindings_to_map(args, bindings).and_then(|args| Some(Datum::Map(args))),
            _ => Some(self.clone()),
        }
    }

    fn variables(&self) -> Vec<String> {
        match *self {
            Datum::Variable(ref v) => vec![v.clone()],
            Datum::Vector(ref args) => {
                let mut variables = Vec::new();
                for arg in args.iter() {
                    variables.extend(arg.variables().into_iter());
                }
                variables
            }
            Datum::Map(ref args) => {
                let mut variables = Vec::new();
                for v in args.values() {
                    variables.extend(v.variables().into_iter());
                }
                variables
            }
            _ => Vec::new(),
        }
    }

    fn rename_variables(&self, renamed_variables: &HashMap<String, String>) -> Self {
        match *self {
            Datum::Variable(ref v) => renamed_variables
                .get(v)
                .and_then(|new_v| Some(Datum::Variable(new_v.clone())))
                .or_else(|| Some(self.clone()))
                .unwrap(),
            Datum::Vector(ref args) => Datum::Vector(
                args.iter()
                    .map(|arg| arg.rename_variables(renamed_variables))
                    .collect(),
            ),
            Datum::Map(ref args) => Datum::Map(
                args.iter()
                    .map(|(k, v)| (k.clone(), v.rename_variables(renamed_variables)))
                    .collect(),
            ),
            _ => self.clone(),
        }
    }

    fn nil() -> Self {
        Datum::Nil
    }
}

impl std::fmt::Display for Datum {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", self.pprint())
    }
}

#[cfg(test)]
mod tests;
<commit_msg>Additional helper functions for the `Datum` struct<commit_after>//! The datum module provides a data structure, Datum, that implements the Unify trait.
//! Datum aims to be a drop-in for any algorithm in ai_kit that operates on the Unify trait.

#[cfg(feature = "with-constraint")]
use constraints::ConstraintValue;
use core;
use std;
use std::collections::{BTreeMap, HashMap};
use std::str;
use utils;

#[derive(Clone, Debug, Serialize, Deserialize, PartialOrd)]
pub enum Datum {
    #[serde(rename = "null")] Nil,
    #[serde(rename = "bool")] Bool(bool),
    #[serde(rename = "str")] String(String),
    #[serde(rename = "int")] Int(i64),
    #[serde(rename = "float")] Float(f64),
    #[serde(rename = "var")] Variable(String),
    #[serde(rename = "vec")] Vector(Vec<Datum>),
    #[serde(rename = "map")] Map(BTreeMap<String, Datum>),
    #[serde(rename = "fn")] Function { head: Box<Datum>, args: Vec<Datum> },
}

impl Datum {
    pub fn to_bool(&self) -> Option<bool> {
        match *self {
            Datum::Bool(ref value) => Some(value.clone()),
            _ => None,
        }
    }

    pub fn to_string(&self) -> Option<String> {
        match *self {
            Datum::String(ref value) => Some(value.clone()),
            _ => None,
        }
    }

    pub fn to_int(&self) -> Option<i64> {
        match *self {
            Datum::Int(ref value) => Some(value.clone()),
            _ => None,
        }
    }

    pub fn to_float(&self) -> Option<f64> {
        match *self {
            Datum::Float(ref f_value) => Some(f_value.clone()),
            Datum::Int(ref i_value) => Some(i_value.clone() as f64),
            _ => None,
        }
    }

    pub fn to_variable(&self) -> Option<String> {
        match *self {
            Datum::Variable(ref value) => Some(value.clone()),
            _ => None,
        }
    }

    pub fn function_head<'a>(&'a self) -> Option<&'a Box<Self>> {
        match *self {
            Datum::Function { ref head, .. } => Some(head),
            _ => None,
        }
    }

    pub fn function_args<'a>(&'a self) -> Option<&'a Vec<Self>> {
        match *self {
            Datum::Function { ref args, .. } => Some(args),
            _ => None,
        }
    }

    pub fn pprint(&self) -> String {
        match *self {
            Datum::Nil => format!("nil"),
            Datum::Bool(ref b) => format!("{}", b),
            Datum::String(ref s) => format!("{}", s),
            Datum::Int(ref i) => format!("{}", i),
            Datum::Float(ref f) => format!("{}", f),
            Datum::Variable(ref v) => format!("{}", v),
            Datum::Vector(ref args) => {
                let elements: Vec<String> = args.iter().map(|e| e.pprint()).collect();
                format!("({})", elements.join(","))
            }
            Datum::Map(ref args) => {
                let elements: Vec<String> = args.iter()
                    .map(|(k, v)| format!("{} => {}", k, v.pprint()))
                    .collect();
                format!("({})", elements.join(","))
            }
            Datum::Function { ref head, ref args } => {
                let elements: Vec<String> = args.iter().map(|e| e.pprint()).collect();
                format!("({} ({}))", head, elements.join(","))
            }
        }
    }

    pub fn is_nil(&self) -> bool {
        match *self {
            Datum::Nil => true,
            _ => false,
        }
    }

    pub fn is_bool(&self) -> bool {
        match *self {
            Datum::Bool(_) => true,
            _ => false,
        }
    }

    pub fn is_int(&self) -> bool {
        match *self {
            Datum::Int(_) => true,
            _ => false,
        }
    }

    pub fn is_float(&self) -> bool {
        match *self {
            Datum::Float(_) => true,
            _ => false,
        }
    }

    pub fn is_string(&self) -> bool {
        match *self {
            Datum::String(_) => true,
            _ => false,
        }
    }

    pub fn is_variable(&self) -> bool {
        match *self {
            Datum::Variable(_) => true,
            _ => false,
        }
    }

    pub fn is_vector(&self) -> bool {
        match *self {
            Datum::Vector(_) => true,
            _ => false,
        }
    }

    pub fn is_map(&self) -> bool {
        match *self {
            Datum::Map { .. } => true,
            _ => false,
        }
    }

    pub fn is_function(&self) -> bool {
        match *self {
            Datum::Function { .. } => true,
            _ => false,
        }
    }
}

impl Default for Datum {
    fn default() -> Self {
        Datum::Nil
    }
}

impl PartialEq for Datum {
    fn eq(&self, other: &Datum) -> bool {
        match *self {
            Datum::Nil => match *other {
                Datum::Nil => true,
                _ => false,
            },
            Datum::Bool(b) => match *other {
                Datum::Bool(b2) => b == b2,
                _ => false,
            },
            Datum::String(ref s) => match *other {
                Datum::String(ref s2) => s == s2,
                _ => false,
            },
            Datum::Int(ref i) => match *other {
                Datum::Int(ref i2) => i == i2,
                _ => false,
            },
            Datum::Float(ref f) => match *other {
                Datum::Float(ref f2) => f == f2,
                _ => false,
            },
            Datum::Variable(ref v) => match *other {
                Datum::Variable(ref v2) => v == v2,
                _ => false,
            },
            Datum::Vector(ref args) => match *other {
                Datum::Vector(ref args2) => args == args2,
                _ => false,
            },
            Datum::Map(ref args) => match *other {
                Datum::Map(ref args2) => args == args2,
                _ => false,
            },
            Datum::Function { ref head, ref args } => match *other {
                Datum::Function {
                    head: ref head2,
                    args: ref args2,
                } if args.len() == args2.len() =>
                {
                    head == head2 && args == args2
                }
                _ => false,
            },
        }
    }
}

impl Eq for Datum {}
impl Ord for Datum {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.partial_cmp(other).unwrap()
    }
}

impl core::BindingsValue for Datum {
    fn variable(s: &String) -> Option<Self> {
        Some(Datum::Variable(s.clone()))
    }
    fn to_variable(&self) -> Option<String> {
        self.to_variable()
    }
}

#[cfg(feature = "with-constraint")]
impl ConstraintValue for Datum {
    fn to_float(&self) -> Option<f64> {
        self.to_float()
    }
    fn float(c: f64) -> Self {
        Datum::Float(c)
    }
}

impl core::Unify<Datum> for Datum {
    fn unify(&self, other: &Datum, bindings: &core::Bindings<Self>) -> Option<core::Bindings<Self>> {
        fn unify_args(args: &Vec<Datum>, args2: &Vec<Datum>, bindings: &core::Bindings<Datum>) -> Option<core::Bindings<Datum>> {
            utils::fold_while_some(
                bindings.clone(),
                &mut args.iter().zip(args2.iter()),
                &|bindings: core::Bindings<Datum>, (ref a, ref b): (&Datum, &Datum)| a.unify(&b, &bindings),
            )
        }

        fn unify_maps(
            args: &BTreeMap<String, Datum>,
            args2: &BTreeMap<String, Datum>,
            bindings: &core::Bindings<Datum>,
        ) -> Option<core::Bindings<Datum>> {
            if args.len() != args2.len() {
                None
            } else {
                utils::fold_while_some(
                    bindings.clone(),
                    &mut args.iter(),
                    &|bindings: core::Bindings<Datum>, (ref k, ref v): (&String, &Datum)| args2.get(*k).and_then(|v2| v.unify(v2, &bindings)),
                )
            }
        }
        match *self {
            Datum::Variable(ref var_name) => bindings.update_bindings(var_name, other),
            Datum::Vector(ref args) => match *other {
                Datum::Vector(ref args2) if args.len() == args2.len() => unify_args(args, args2, &bindings),
                _ => None,
            },
            Datum::Map(ref args) => match *other {
                Datum::Map(ref args2) if args.len() == args2.len() => unify_maps(args, args2, &bindings),
                _ => None,
            },
            _ => match *other {
                Datum::Variable(ref var_name) => bindings.update_bindings(var_name, self),
                _ => {
                    if self == other {
                        Some(bindings.clone())
                    } else {
                        None
                    }
                }
            },
        }
    }

    fn apply_bindings(&self, bindings: &core::Bindings<Self>) -> Option<Self> {
        fn apply_bindings_to_args(args: &Vec<Datum>, bindings: &core::Bindings<Datum>) -> Option<Vec<Datum>> {
            utils::map_while_some(&mut args.iter(), &|arg| {
                arg.apply_bindings(bindings)
            })
        }
        fn apply_bindings_to_map(args: &BTreeMap<String, Datum>, bindings: &core::Bindings<Datum>) -> Option<BTreeMap<String, Datum>> {
            utils::map_while_some(&mut args.iter(), &|(k, v)| {
                v.apply_bindings(bindings)
                    .and_then(|v| Some((k.clone(), v)))
            }).and_then(|tuple_vec| Some(tuple_vec.into_iter().collect::<BTreeMap<String, Datum>>()))
        }
        match *self {
            Datum::Variable(ref var_name) => bindings
                .get_binding(var_name)
                .or_else(|| Some(self.clone())),
            Datum::Vector(ref args) => apply_bindings_to_args(args, bindings).and_then(|args| Some(Datum::Vector(args))),
            Datum::Map(ref args) => apply_bindings_to_map(args, bindings).and_then(|args| Some(Datum::Map(args))),
            _ => Some(self.clone()),
        }
    }

    fn variables(&self) -> Vec<String> {
        match *self {
            Datum::Variable(ref v) => vec![v.clone()],
            Datum::Vector(ref args) => {
                let mut variables = Vec::new();
                for arg in args.iter() {
                    variables.extend(arg.variables().into_iter());
                }
                variables
            }
            Datum::Map(ref args) => {
                let mut variables = Vec::new();
                for v in args.values() {
                    variables.extend(v.variables().into_iter());
                }
                variables
            }
            _ => Vec::new(),
        }
    }

    fn rename_variables(&self, renamed_variables: &HashMap<String, String>) -> Self {
        match *self {
            Datum::Variable(ref v) => renamed_variables
                .get(v)
                .and_then(|new_v| Some(Datum::Variable(new_v.clone())))
                .or_else(|| Some(self.clone()))
                .unwrap(),
            Datum::Vector(ref args) => Datum::Vector(
                args.iter()
                    .map(|arg| arg.rename_variables(renamed_variables))
                    .collect(),
            ),
            Datum::Map(ref args) => Datum::Map(
                args.iter()
                    .map(|(k, v)| (k.clone(), v.rename_variables(renamed_variables)))
                    .collect(),
            ),
            _ => self.clone(),
        }
    }

    fn nil() -> Self {
        Datum::Nil
    }
}

impl std::fmt::Display for Datum {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", self.pprint())
    }
}

#[cfg(test)]
mod tests;
<|endoftext|>


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

use num::BigUint;

use Construct;
use Data;
use Count;
use Of;
use ToIndex;
use ToPos;
use Zero;

/// Dimension is natural number, position is the same as index.
pub struct Dimension<T = Data>(PhantomData<T>);

impl<T> Construct for Dimension<T> {
    fn new() -> Dimension<T> { Dimension(PhantomData) }
}

impl Count<usize> for Dimension<Data> {
    type N = usize;
    fn count(&self, dim: &usize) -> usize { *dim }
}

impl Count<BigUint> for Dimension<Data> {
    type N = BigUint;
    fn count(&self, dim: &BigUint) -> BigUint { dim.clone() }
}

impl<T, U> Count<U> for Dimension<Of<T>>
    where T: Construct + Count<U>
{
    type N = <T as Count<U>>::N;
    fn count(&self, dim: &U) -> Self::N {
        let of: T = Construct::new();
        of.count(dim)
    }
}

impl Zero<usize, usize> for Dimension<Data> {
    fn zero(&self, _dim: &usize) -> usize { 0 }
}

impl<T, U, V>
Zero<U, V> for Dimension<Of<T>>
    where T: Construct + Zero<U, V>
{
    fn zero(&self, dim: &U) -> V {
        let of: T = Construct::new();
        of.zero(dim)
    }
}

impl ToIndex<usize, usize> for Dimension<Data> {
    fn to_index(&self, _dim: &usize, pos: &usize) -> usize { *pos }
}

impl<T, U, V> ToIndex<U, V> for Dimension<Of<T>>
    where
        T: Construct + ToIndex<U, V>
{
    fn to_index(&self, dim: &U, pos: &V) -> usize {
        let of: T = Construct::new();
        of.to_index(dim, pos)
    }
}

impl ToPos<usize, usize> for Dimension<Data> {
    fn to_pos(&self, _dim: &usize, index: usize, pos: &mut usize) {
        *pos = index;
    }
}

impl<T, U, V> ToPos<U, V> for Dimension<Of<T>>
    where
        T: Construct + ToPos<U, V>
{
    fn to_pos(&self, dim: &U, index: usize, pos: &mut V) {
        let of: T = Construct::new();
        of.to_pos(dim, index, pos);
    }
}

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

    #[test]
    fn features() {
        is_complete::<Dimension, usize, usize>();
        is_complete::<Dimension<Of<Pair>>, usize, (usize, usize)>();
    }
}
<commit_msg>Impl `Zero<BigUint, BigUint>` for `Dimension`<commit_after>use std::marker::PhantomData;

use num::BigUint;

use Construct;
use Data;
use Count;
use Of;
use ToIndex;
use ToPos;
use Zero;

/// Dimension is natural number, position is the same as index.
pub struct Dimension<T = Data>(PhantomData<T>);

impl<T> Construct for Dimension<T> {
    fn new() -> Dimension<T> { Dimension(PhantomData) }
}

impl Count<usize> for Dimension<Data> {
    type N = usize;
    fn count(&self, dim: &usize) -> usize { *dim }
}

impl Count<BigUint> for Dimension<Data> {
    type N = BigUint;
    fn count(&self, dim: &BigUint) -> BigUint { dim.clone() }
}

impl<T, U> Count<U> for Dimension<Of<T>>
    where T: Construct + Count<U>
{
    type N = <T as Count<U>>::N;
    fn count(&self, dim: &U) -> Self::N {
        let of: T = Construct::new();
        of.count(dim)
    }
}

impl Zero<usize, usize> for Dimension<Data> {
    fn zero(&self, _dim: &usize) -> usize { 0 }
}

impl Zero<BigUint, BigUint> for Dimension<Data> {
    fn zero(&self, _dim: &BigUint) -> BigUint { 0usize.into() }
}

impl<T, U, V>
Zero<U, V> for Dimension<Of<T>>
    where T: Construct + Zero<U, V>
{
    fn zero(&self, dim: &U) -> V {
        let of: T = Construct::new();
        of.zero(dim)
    }
}

impl ToIndex<usize, usize> for Dimension<Data> {
    fn to_index(&self, _dim: &usize, pos: &usize) -> usize { *pos }
}

impl<T, U, V> ToIndex<U, V> for Dimension<Of<T>>
    where
        T: Construct + ToIndex<U, V>
{
    fn to_index(&self, dim: &U, pos: &V) -> usize {
        let of: T = Construct::new();
        of.to_index(dim, pos)
    }
}

impl ToPos<usize, usize> for Dimension<Data> {
    fn to_pos(&self, _dim: &usize, index: usize, pos: &mut usize) {
        *pos = index;
    }
}

impl<T, U, V> ToPos<U, V> for Dimension<Of<T>>
    where
        T: Construct + ToPos<U, V>
{
    fn to_pos(&self, dim: &U, index: usize, pos: &mut V) {
        let of: T = Construct::new();
        of.to_pos(dim, index, pos);
    }
}

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

    #[test]
    fn features() {
        is_complete::<Dimension, usize, usize>();
        is_complete::<Dimension<Of<Pair>>, usize, (usize, usize)>();
    }
}
<|endoftext|>


<commit_before>use rand::{Rng, thread_rng};

#[derive(Clone, Copy, PartialEq)]
pub enum Direction {
    North,
    East,
    South,
    West,
}

impl Direction {
    pub fn opposite(&self) -> Direction {
        match *self {
            Direction::North => Direction::South,
            Direction::East => Direction::West,
            Direction::South => Direction::North,
            Direction::West => Direction::East,
        }
    }

    pub fn all() -> [Direction; 4] {
        [
            Direction::North,
            Direction::East,
            Direction::South,
            Direction::West,
        ]
    }

    pub fn random() -> Direction {
        match thread_rng().gen_range(0, 4) {
            0 => Direction::North,
            1 => Direction::East,
            2 => Direction::South,
            3 => Direction::West,
            _ => unreachable!()
        }
    }
}

<commit_msg>Add Direction::from_positions()<commit_after>use rand::{Rng, thread_rng};

#[derive(Clone, Copy, PartialEq)]
pub enum Direction {
    North,
    East,
    South,
    West,
}

impl Direction {
    pub fn opposite(&self) -> Direction {
        match *self {
            Direction::North => Direction::South,
            Direction::East => Direction::West,
            Direction::South => Direction::North,
            Direction::West => Direction::East,
        }
    }

    pub fn all() -> [Direction; 4] {
        [
            Direction::North,
            Direction::East,
            Direction::South,
            Direction::West,
        ]
    }

    pub fn random() -> Direction {
        match thread_rng().gen_range(0, 4) {
            0 => Direction::North,
            1 => Direction::East,
            2 => Direction::South,
            3 => Direction::West,
            _ => unreachable!()
        }
    }

    /// Returns the direction travelled from `pos1` to `pos2` if the positions
    /// represent travel in exactly north, south, east or west direction, or
    /// returns an error.
    pub fn from_positions(pos1: &[u8; 2], pos2: &[u8; 2]) -> Result<Direction, &'static str> {
        let x_diff = pos1[0] as i16 - pos2[0] as i16;
        let y_diff = pos1[1] as i16 - pos2[1] as i16;

        if x_diff == 0 && y_diff > 0 {
            Ok(Direction::North)
        } else if x_diff == 0 && y_diff < 0 {
            Ok(Direction::South)
        } else if x_diff > 0 && y_diff == 0 {
            Ok(Direction::West)
        } else if x_diff < 0 && y_diff == 0 {
            Ok(Direction::East)
        } else {
            Err("positions do not represent a supported direction")
        }
    }
}

<|endoftext|>


<commit_before>//! The driver trait which all drivers will implement.

use error::{Error, ErrorCode};
use query::{Condition, SortRule, Range, Query};
use schema::Type;
use value::{Value, ValueIter};

pub trait Driver {
  /// Connects to a driver and returns a driver instance. After calling this
  /// the driver is ready to roll!
  fn connect(uri: &str) -> Result<Self, Error> where Self: Sized;

  /// Lazily read some values from the driver.
  ///
  /// Designed against a couple of database specifications. Including the
  /// following:
  ///
  /// - [SQL `SELECT` statement][1].
  /// - [MongoDB `find` command][2].
  ///
  /// [1]: http://www.postgresql.org/docs/current/static/sql-select.html
  /// [2]: https://docs.mongodb.org/manual/reference/command/find/
  fn read(
    &self,
    type_: &Type,
    condition: Condition,
    sort: Vec<SortRule>,
    range: Range,
    query: Query
  ) -> Result<ValueIter, Error>;

  /// Read a single value from the driver. The default implementation uses the
  /// driver read method with no sort and a limit of one.
  fn read_one(
    &self,
    type_: &Type,
    condition: Condition,
    query: Query
  ) -> Result<Value, Error> {
    let mut values: Vec<_> = try!(self.read(
      type_,
      condition,
      Default::default(),
      Range::new(None, Some(1)),
      query
    )).collect();

    if values.len() > 1 {
      Err(Error::internal("Read with a limit of one returned more than one value."))
    } else if let Some(value) = values.pop() {
      Ok(value)
    } else {
      Err(Error::new(ErrorCode::NotFound, "No value was found for the condition.", None))
    }
  }
}
<commit_msg>doc: driver trait<commit_after>use error::{Error, ErrorCode};
use query::{Condition, SortRule, Range, Query};
use schema::Type;
use value::{Value, ValueIter};

/// The driver trait which all drivers will implement. Designed to be
/// interoperable with any data source, however the driver also assumes a
/// collection based data model.
///
/// ## Collection Based Data Model
/// Any value the driver can access is assumed to have a specific “type”
/// associated with it. This way the driver can map to structures like
/// collections in MongoDB, tables or views in a SQL database like PostgreSQL,
/// or labels in Neo4j.
///
/// In addition, operations on the driver are done in a CRUD fashion for a
/// specific type. Looking at the specific methods, like `read`, heavy
/// inspiration is taken from SQL databases and MongoDB which do all their CRUD
/// on collections. The unit for CRUD in an Ardite driver is a type.
///
/// Originally the driver was designed to work strictly like Falcor or GraphQL
/// by assuming a “graph” like structure for *all* data. This turns out to be a
/// nice abstraction, however it is inneficient for creating performant and
/// flexible systems. Abandoning the graph structure of drivers also helps the
/// driver implementors who might find it incredibly difficult and repetitive
/// to copy the same graph interface on their relational database.
///
/// ## Relationship Between the Schema and the Driver
/// The driver and the schema should never interact with each other. The schema
/// should be managed completely by Ardite and used to validate queries and
/// values before being sent to the driver. The driver should only be a
/// low-level consistent interface to a data source and should not make any
/// intelligent decisions about how an Ardite program should run.
///
/// Some unique driver-specific features are encouraged to be included in the
/// driver, such as PostgreSQL computed columns. However, these features should
/// *not* be derived from the schema, these features *should* be consistent,
/// and the driver must consistently perform all other expected functionality
/// *before* implementing new features.
///
/// This sharp seperation is most primarily for a seperation of concerns. It
/// allows developers working with the driver higher up in Ardite to make basic
/// assumptions about the driver. The seperation also prevents code reuse. The
/// `validate_query` or `validate_value` methods on `Schema` may be tempting to
/// use in the driver. There should never be any confusion over where in the
/// algorithm these methods should be called.
pub trait Driver {
  /// Connects to a driver and returns a driver instance. After calling this
  /// the driver is ready to roll!
  ///
  /// No schema definition is provided to the driver in its construction step.
  fn connect(uri: &str) -> Result<Self, Error> where Self: Sized;

  /// Lazily read some values from the driver.
  ///
  /// Designed against a couple of database specifications. Including the
  /// following:
  ///
  /// - [SQL `SELECT` statement][1].
  /// - [MongoDB `find` command][2].
  ///
  /// [1]: http://www.postgresql.org/docs/current/static/sql-select.html
  /// [2]: https://docs.mongodb.org/manual/reference/command/find/
  fn read(
    &self,
    type_: &Type,
    condition: Condition,
    sort: Vec<SortRule>,
    range: Range,
    query: Query
  ) -> Result<ValueIter, Error>;

  /// Read a single value from the driver. The default implementation uses the
  /// driver read method with a range of one.
  ///
  /// If a condition matches more than one value (while not recommended for
  /// this method) the first of these values, using the default sorting
  /// algorithm of the database, is returned.
  ///
  /// This method may be optionally optimized by the driver.
  fn read_one(
    &self,
    type_: &Type,
    condition: Condition,
    query: Query
  ) -> Result<Value, Error> {
    let mut values: Vec<_> = try!(self.read(
      type_,
      condition,
      Default::default(),
      Range::new(None, Some(1)),
      query
    )).collect();

    if values.len() > 1 {
      Err(Error::internal("Read with a limit of one returned more than one value."))
    } else if let Some(value) = values.pop() {
      Ok(value)
    } else {
      Err(Error::new(ErrorCode::NotFound, "No value was found for the condition.", None))
    }
  }
}
<|endoftext|>


<commit_before>use std::thread;
use std::cell::UnsafeCell;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver, RecvError};
use basetypes::*;
use chess_move::*;
use tt::*;
use position::Position;

pub enum Command {
    Search {
        search_id: usize,
        position: Position,
        depth: u8,
        lower_bound: Value,
        upper_bound: Value,
    },
    Stop,
    Exit,
}

pub enum Report {
    Progress {
        search_id: usize,
        searched_nodes: NodeCount,
        depth: u8,
    },
    Done {
        search_id: usize,
        searched_nodes: NodeCount,
        value: Option<Value>,
    },
}

pub fn run_deepening(tt: Arc<TranspositionTable>,
                     commands: Receiver<Command>,
                     reports: Sender<Report>) {
    let (slave_commands_tx, slave_commands_rx) = channel();
    let (slave_reports_tx, slave_reports_rx) = channel();
    let slave = thread::spawn(move || {
        run(tt, slave_commands_rx, slave_reports_tx);
    });
    let mut pending_command = None;
    loop {
        // If there is a pending command, we take it, otherwise we
        // block and wait to receive a new one.
        let command = match pending_command.take() {
            Some(cmd) => cmd,
            None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
        };

        match command {
            Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                let mut searched_nodes_final = 0;
                let mut value_final = None;
                'depthloop: for n in 1..(depth + 1) {
                    slave_commands_tx.send(Command::Search {
                                         search_id: n as usize,
                                         position: position.clone(),
                                         depth: n,
                                         lower_bound: lower_bound,
                                         upper_bound: upper_bound,
                                     })
                                     .unwrap();
                    loop {
                        match slave_reports_rx.recv().unwrap() {
                            Report::Progress { searched_nodes, .. } => {
                                reports.send(Report::Progress {
                                           search_id: search_id,
                                           searched_nodes: searched_nodes_final + searched_nodes,
                                           depth: n,
                                       })
                                       .ok();
                                if pending_command.is_none() {
                                    pending_command = match commands.try_recv() {
                                        Ok(cmd) => {
                                            slave_commands_tx.send(Command::Stop).unwrap();
                                            Some(cmd)
                                        }
                                        _ => None,
                                    }
                                }
                            }
                            Report::Done { searched_nodes, value, .. } => {
                                searched_nodes_final += searched_nodes;
                                if n == depth {
                                    value_final = value;
                                }
                                if pending_command.is_some() {
                                    break 'depthloop;
                                }
                                break;
                            }
                        }
                    }
                }
                reports.send(Report::Done {
                           search_id: search_id,
                           searched_nodes: searched_nodes_final,
                           value: value_final,
                       })
                       .ok();
            }
            Command::Stop => {
                slave_commands_tx.send(Command::Stop).unwrap();
                continue;
            }
            Command::Exit => {
                slave_commands_tx.send(Command::Exit).unwrap();
                break;
            }
        }
    }
    slave.join().unwrap();
}

pub fn run(tt: Arc<TranspositionTable>, commands: Receiver<Command>, reports: Sender<Report>) {
    thread_local!(
        static MOVE_STACK: UnsafeCell<MoveStack> = UnsafeCell::new(MoveStack::new())
    );
    MOVE_STACK.with(|s| {
        let mut move_stack = unsafe { &mut *s.get() };
        let mut pending_command = None;
        loop {
            // If there is a pending command, we take it, otherwise we
            // block and wait to receive a new one.
            let command = match pending_command.take() {
                Some(cmd) => cmd,
                None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
            };

            match command {
                Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                    let mut report = |n| {
                        reports.send(Report::Progress {
                                   search_id: search_id,
                                   searched_nodes: n,
                                   depth: depth,
                               })
                               .ok();
                        match commands.try_recv() {
                            Ok(x) => {
                                pending_command = Some(x);
                                Err(TerminatedSearch)
                            }
                            _ => Ok(()),
                        }
                    };
                    let mut search = Search::new(position, &tt, move_stack, &mut report);
                    let value = search.run(lower_bound, upper_bound, depth).ok();
                    reports.send(Report::Progress {
                               search_id: search_id,
                               searched_nodes: search.node_count(),
                               depth: depth,
                           })
                           .ok();
                    reports.send(Report::Done {
                               search_id: search_id,
                               searched_nodes: search.node_count(),
                               value: value,
                           })
                           .ok();
                    search.reset();
                }
                Command::Stop => continue,
                Command::Exit => break,
            }
        }
    })
}

/// Represents a terminated search condition.
pub struct TerminatedSearch;

enum NodePhase {
    Pristine,
    TriedHashMove,
    GeneratedMoves,
}

struct NodeState {
    phase: NodePhase,
    entry: EntryData,
}

const NODE_COUNT_REPORT_INTERVAL: NodeCount = 10000;

pub struct Search<'a> {
    tt: &'a TranspositionTable,
    position: Position,
    moves: &'a mut MoveStack,
    moves_starting_ply: usize,
    state_stack: Vec<NodeState>,
    reported_nodes: NodeCount,
    unreported_nodes: NodeCount,
    report_function: &'a mut FnMut(NodeCount) -> Result<(), TerminatedSearch>,
}

impl<'a> Search<'a> {
    /// Creates a new instance.
    pub fn new(p: Position,
               tt: &'a TranspositionTable,
               moves: &'a mut MoveStack,
               report_function: &'a mut FnMut(NodeCount) -> Result<(), TerminatedSearch>)
               -> Search<'a> {
        let moves_starting_ply = moves.ply();
        Search {
            tt: tt,
            position: p,
            moves: moves,
            moves_starting_ply: moves_starting_ply,
            state_stack: Vec::with_capacity(32),
            reported_nodes: 0,
            unreported_nodes: 0,
            report_function: report_function,
        }
    }

    /// Performs a principal variation search and returns a result.
    ///
    /// **Important note**: This method may leave un-restored move
    /// lists in the move stack. Call `reset` if you want the move
    /// stack to be restored to the state it had when the search
    /// instance was created.
    pub fn run(&mut self,
               mut alpha: Value, // lower bound
               beta: Value, // upper bound
               depth: u8)
               -> Result<Value, TerminatedSearch> {
        assert!(alpha < beta);

        let entry = self.node_begin();

        // Check if the TT entry gives the result.
        if entry.depth() >= depth {
            let value = entry.value();
            let bound = entry.bound();
            if (value >= beta && bound == BOUND_LOWER) ||
               (value <= alpha && bound == BOUND_UPPER) || (bound == BOUND_EXACT) {
                self.node_end();
                return Ok(value);
            }
        }

        // Initial guests for the final result.
        let mut bound = BOUND_UPPER;
        let mut best_move = Move::invalid();

        if depth == 0 {
            // On leaf nodes, do quiescence search.
            let (value, nodes) = self.position
                                     .evaluate_quiescence(alpha, beta, Some(entry.eval_value()));
            try!(self.report_progress(nodes));

            // See how good this position is.
            if value >= beta {
                alpha = beta;
                bound = BOUND_LOWER;
            } else if value > alpha {
                alpha = value;
                bound = BOUND_EXACT;
            }

        } else {
            // On non-leaf nodes, try moves.
            let mut no_moves_yet = true;
            while let Some(m) = self.do_move() {
                try!(self.report_progress(1));

                // Make a recursive call.
                let value = if no_moves_yet {
                    // The first move we analyze with a fully open window
                    // (alpha, beta). If this happens to be a good move,
                    // it will probably raise `alpha`.
                    no_moves_yet = false;
                    -try!(self.run(-beta, -alpha, depth - 1))
                } else {
                    // For the next moves we first try to prove that they
                    // are not better than our current best move. For this
                    // purpose we analyze them with a null window (alpha,
                    // alpha + 1). This is faster than a full window
                    // search. Only if we are certain that the move is
                    // better than our current best move, we do a
                    // full-window search.
                    match -try!(self.run(-alpha - 1, -alpha, depth - 1)) {
                        x if x <= alpha => x,
                        _ => -try!(self.run(-beta, -alpha, depth - 1)),
                    }
                };
                self.undo_move();

                // See how good this move was.
                if value >= beta {
                    // This move is so good, that the opponent will
                    // probably not allow this line of play to
                    // happen. Therefore we should not lose any more time
                    // on this position.
                    alpha = beta;
                    bound = BOUND_LOWER;
                    best_move = m;
                    break;
                }
                if value > alpha {
                    // We found a new best move.
                    alpha = value;
                    bound = BOUND_EXACT;
                    best_move = m;
                }
            }

            // Check if we are in a final position (no legal moves).
            if no_moves_yet {
                alpha = self.position.evaluate_final();
                bound = BOUND_EXACT;
            }
        }

        self.store(alpha, bound, depth, best_move);
        self.node_end();
        Ok(alpha)
    }

    /// Returns the number of searched positions.
    #[inline(always)]
    pub fn node_count(&self) -> NodeCount {
        self.reported_nodes + self.unreported_nodes
    }

    /// Resets the instance to the state it had when it was created.
    #[inline]
    pub fn reset(&mut self) {
        while self.moves.ply() > self.moves_starting_ply {
            self.moves.restore();
        }
        self.state_stack.clear();
        self.reported_nodes = 0;
        self.unreported_nodes = 0;
    }

    // Declares that we are starting to process a new node.
    //
    // Each recursive call to `run` begins with a call to
    // `node_begin`. The returned value is a TT entry telling
    // everything we know about the current position.
    #[inline]
    fn node_begin(&mut self) -> EntryData {
        // Consult the transposition table.
        let entry = if let Some(e) = self.tt.probe(self.position.hash()) {
            e
        } else {
            EntryData::new(0, BOUND_NONE, 0, 0, self.position.evaluate_static())
        };
        self.state_stack.push(NodeState {
            phase: NodePhase::Pristine,
            entry: entry,
        });
        entry
    }

    // Declares that we are done processing the current node.
    //
    // Each recursive call to `run` ends with a call to `node_end`.
    #[inline]
    fn node_end(&mut self) {
        if let NodePhase::Pristine = self.state_stack.last().unwrap().phase {
            // For pristine nodes we have not saved the move list
            // yet, so we should not restore it.
        } else {
            self.moves.restore();
        }
        self.state_stack.pop();
    }

    // Plays the next legal move in the current position and returns
    // it.
    //
    // Each call to `do_move` for the same position will play and
    // return a different move. When all legal moves has been played,
    // `None` will be returned. `do_move` will do whatever it can to
    // play the best moves first, and the worst last. It will also try
    // to be efficient, for example it will generate the list of all
    // pseudo-legal moves at the last possible moment.
    #[inline]
    fn do_move(&mut self) -> Option<Move> {
        let state = self.state_stack.last_mut().unwrap();

        if let NodePhase::Pristine = state.phase {
            // We save the move list at the last possible moment,
            // because most of the nodes are leafs.
            self.moves.save();

            // We always try the hash move first.
            state.phase = NodePhase::TriedHashMove;
            if state.entry.move16() != 0 {
                if let Some(m) = self.position.try_move_digest(state.entry.move16()) {
                    if self.position.do_move(m) {
                        return Some(m);
                    }
                }
            }
        }

        if let NodePhase::TriedHashMove = state.phase {
            // After the hash move, we generate all pseudo-legal
            // moves. But we should not forget to remove the already
            // tried hash move from the list.
            self.position.generate_moves(self.moves);
            if state.entry.move16() != 0 {
                self.moves.remove_move(state.entry.move16());
            }
            state.phase = NodePhase::GeneratedMoves;
        }

        // For the last, we spit the generated moves out.
        while let Some(m) = self.moves.remove_best_move() {
            if self.position.do_move(m) {
                return Some(m);
            }
        }
        None
    }

    // Takes the last played move back.
    #[inline]
    fn undo_move(&mut self) {
        self.position.undo_move();
    }

    // Stores updated node information in the transposition table.
    #[inline]
    fn store(&mut self, value: Value, bound: BoundType, depth: u8, best_move: Move) {
        let entry = &self.state_stack.last().unwrap().entry;
        let move16 = match best_move.digest() {
            0 => entry.move16(),
            x => x,
        };
        self.tt.store(self.position.hash(),
                      EntryData::new(value, bound, depth, move16, entry.eval_value()));
    }

    // Reports search progress.
    //
    // From time to time, we should report how many nodes had been
    // searched since the beginning of the search. This also gives an
    // opportunity for the search to be terminated.
    #[inline]
    fn report_progress(&mut self, new_nodes: NodeCount) -> Result<(), TerminatedSearch> {
        self.unreported_nodes += new_nodes;
        if self.unreported_nodes > NODE_COUNT_REPORT_INTERVAL {
            self.reported_nodes += self.unreported_nodes;
            self.unreported_nodes = 0;
            try!((*self.report_function)(self.reported_nodes));
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::Search;
    use chess_move::*;
    use tt::*;
    use position::Position;

    #[test]
    fn test_search() {
        let p = Position::from_fen("8/8/8/8/3q3k/7n/6PP/2Q2R1K b - - 0 1").ok().unwrap();
        let tt = TranspositionTable::new();
        let mut moves = MoveStack::new();
        let mut report = |_| Ok(());
        let mut search = Search::new(p, &tt, &mut moves, &mut report);
        let value = search.run(-30000, 30000, 2)
                          .ok()
                          .unwrap();
        assert!(value < -300);
        search.reset();
        let value = search.run(-30000, 30000, 4)
                          .ok()
                          .unwrap();
        assert!(value >= 20000);
    }
}
<commit_msg>comment<commit_after>use std::thread;
use std::cell::UnsafeCell;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver, RecvError};
use basetypes::*;
use chess_move::*;
use tt::*;
use position::Position;

pub enum Command {
    Search {
        search_id: usize,
        position: Position,
        depth: u8,
        lower_bound: Value,
        upper_bound: Value,
    },
    Stop,
    Exit,
}

pub enum Report {
    Progress {
        search_id: usize,
        searched_nodes: NodeCount,
        depth: u8,
    },
    Done {
        search_id: usize,
        searched_nodes: NodeCount,
        value: Option<Value>,
    },
}

pub fn run_deepening(tt: Arc<TranspositionTable>,
                     commands: Receiver<Command>,
                     reports: Sender<Report>) {
    let (slave_commands_tx, slave_commands_rx) = channel();
    let (slave_reports_tx, slave_reports_rx) = channel();
    let slave = thread::spawn(move || {
        run(tt, slave_commands_rx, slave_reports_tx);
    });
    let mut pending_command = None;
    loop {
        // If there is a pending command, we take it, otherwise we
        // block and wait to receive a new one.
        let command = match pending_command.take() {
            Some(cmd) => cmd,
            None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
        };

        match command {
            Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                let mut searched_nodes_final = 0;
                let mut value_final = None;
                'depthloop: for n in 1..(depth + 1) {
                    slave_commands_tx.send(Command::Search {
                                         search_id: n as usize,
                                         position: position.clone(),
                                         depth: n,
                                         lower_bound: lower_bound,
                                         upper_bound: upper_bound,
                                     })
                                     .unwrap();
                    loop {
                        match slave_reports_rx.recv().unwrap() {
                            Report::Progress { searched_nodes, .. } => {
                                reports.send(Report::Progress {
                                           search_id: search_id,
                                           searched_nodes: searched_nodes_final + searched_nodes,
                                           depth: n,
                                       })
                                       .ok();
                                if pending_command.is_none() {
                                    pending_command = match commands.try_recv() {
                                        Ok(cmd) => {
                                            slave_commands_tx.send(Command::Stop).unwrap();
                                            Some(cmd)
                                        }
                                        _ => None,
                                    }
                                }
                            }
                            Report::Done { searched_nodes, value, .. } => {
                                searched_nodes_final += searched_nodes;
                                if n == depth {
                                    value_final = value;
                                }
                                if pending_command.is_some() {
                                    break 'depthloop;
                                }
                                break;
                            }
                        }
                    }
                }
                reports.send(Report::Done {
                           search_id: search_id,
                           searched_nodes: searched_nodes_final,
                           value: value_final,
                       })
                       .ok();
            }
            Command::Stop => {
                slave_commands_tx.send(Command::Stop).unwrap();
                continue;
            }
            Command::Exit => {
                slave_commands_tx.send(Command::Exit).unwrap();
                break;
            }
        }
    }
    slave.join().unwrap();
}

pub fn run(tt: Arc<TranspositionTable>, commands: Receiver<Command>, reports: Sender<Report>) {
    thread_local!(
        static MOVE_STACK: UnsafeCell<MoveStack> = UnsafeCell::new(MoveStack::new())
    );
    MOVE_STACK.with(|s| {
        let mut move_stack = unsafe { &mut *s.get() };
        let mut pending_command = None;
        loop {
            // If there is a pending command, we take it, otherwise we
            // block and wait to receive a new one.
            let command = match pending_command.take() {
                Some(cmd) => cmd,
                None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
            };

            match command {
                Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                    let mut report = |n| {
                        reports.send(Report::Progress {
                                   search_id: search_id,
                                   searched_nodes: n,
                                   depth: depth,
                               })
                               .ok();
                        match commands.try_recv() {
                            Ok(x) => {
                                pending_command = Some(x);
                                Err(TerminatedSearch)
                            }
                            _ => Ok(()),
                        }
                    };
                    let mut search = Search::new(position, &tt, move_stack, &mut report);
                    let value = search.run(lower_bound, upper_bound, depth).ok();
                    reports.send(Report::Progress {
                               search_id: search_id,
                               searched_nodes: search.node_count(),
                               depth: depth,
                           })
                           .ok();
                    reports.send(Report::Done {
                               search_id: search_id,
                               searched_nodes: search.node_count(),
                               value: value,
                           })
                           .ok();
                    search.reset();
                }
                Command::Stop => continue,
                Command::Exit => break,
            }
        }
    })
}

/// Represents a terminated search condition.
pub struct TerminatedSearch;

enum NodePhase {
    Pristine,
    TriedHashMove,
    GeneratedMoves,
}

struct NodeState {
    phase: NodePhase,
    entry: EntryData,
}

const NODE_COUNT_REPORT_INTERVAL: NodeCount = 10000;

pub struct Search<'a> {
    tt: &'a TranspositionTable,
    position: Position,
    moves: &'a mut MoveStack,
    moves_starting_ply: usize,
    state_stack: Vec<NodeState>,
    reported_nodes: NodeCount,
    unreported_nodes: NodeCount,
    report_function: &'a mut FnMut(NodeCount) -> Result<(), TerminatedSearch>,
}

impl<'a> Search<'a> {
    /// Creates a new instance.
    ///
    /// `report_function` should be a function that registers the
    /// search progress. It will be called with the number of searched
    /// positions from the beginning of the search to this moment.
    pub fn new(p: Position,
               tt: &'a TranspositionTable,
               moves: &'a mut MoveStack,
               report_function: &'a mut FnMut(NodeCount) -> Result<(), TerminatedSearch>)
               -> Search<'a> {
        let moves_starting_ply = moves.ply();
        Search {
            tt: tt,
            position: p,
            moves: moves,
            moves_starting_ply: moves_starting_ply,
            state_stack: Vec::with_capacity(32),
            reported_nodes: 0,
            unreported_nodes: 0,
            report_function: report_function,
        }
    }

    /// Performs a principal variation search and returns a result.
    ///
    /// **Important note**: This method may leave un-restored move
    /// lists in the move stack. Call `reset` if you want the move
    /// stack to be restored to the state it had when the search
    /// instance was created.
    pub fn run(&mut self,
               mut alpha: Value, // lower bound
               beta: Value, // upper bound
               depth: u8)
               -> Result<Value, TerminatedSearch> {
        assert!(alpha < beta);

        let entry = self.node_begin();

        // Check if the TT entry gives the result.
        if entry.depth() >= depth {
            let value = entry.value();
            let bound = entry.bound();
            if (value >= beta && bound == BOUND_LOWER) ||
               (value <= alpha && bound == BOUND_UPPER) || (bound == BOUND_EXACT) {
                self.node_end();
                return Ok(value);
            }
        }

        // Initial guests for the final result.
        let mut bound = BOUND_UPPER;
        let mut best_move = Move::invalid();

        if depth == 0 {
            // On leaf nodes, do quiescence search.
            let (value, nodes) = self.position
                                     .evaluate_quiescence(alpha, beta, Some(entry.eval_value()));
            try!(self.report_progress(nodes));

            // See how good this position is.
            if value >= beta {
                alpha = beta;
                bound = BOUND_LOWER;
            } else if value > alpha {
                alpha = value;
                bound = BOUND_EXACT;
            }

        } else {
            // On non-leaf nodes, try moves.
            let mut no_moves_yet = true;
            while let Some(m) = self.do_move() {
                try!(self.report_progress(1));

                // Make a recursive call.
                let value = if no_moves_yet {
                    // The first move we analyze with a fully open window
                    // (alpha, beta). If this happens to be a good move,
                    // it will probably raise `alpha`.
                    no_moves_yet = false;
                    -try!(self.run(-beta, -alpha, depth - 1))
                } else {
                    // For the next moves we first try to prove that they
                    // are not better than our current best move. For this
                    // purpose we analyze them with a null window (alpha,
                    // alpha + 1). This is faster than a full window
                    // search. Only if we are certain that the move is
                    // better than our current best move, we do a
                    // full-window search.
                    match -try!(self.run(-alpha - 1, -alpha, depth - 1)) {
                        x if x <= alpha => x,
                        _ => -try!(self.run(-beta, -alpha, depth - 1)),
                    }
                };
                self.undo_move();

                // See how good this move was.
                if value >= beta {
                    // This move is so good, that the opponent will
                    // probably not allow this line of play to
                    // happen. Therefore we should not lose any more time
                    // on this position.
                    alpha = beta;
                    bound = BOUND_LOWER;
                    best_move = m;
                    break;
                }
                if value > alpha {
                    // We found a new best move.
                    alpha = value;
                    bound = BOUND_EXACT;
                    best_move = m;
                }
            }

            // Check if we are in a final position (no legal moves).
            if no_moves_yet {
                alpha = self.position.evaluate_final();
                bound = BOUND_EXACT;
            }
        }

        self.store(alpha, bound, depth, best_move);
        self.node_end();
        Ok(alpha)
    }

    /// Returns the number of searched positions.
    #[inline(always)]
    pub fn node_count(&self) -> NodeCount {
        self.reported_nodes + self.unreported_nodes
    }

    /// Resets the instance to the state it had when it was created.
    #[inline]
    pub fn reset(&mut self) {
        while self.moves.ply() > self.moves_starting_ply {
            self.moves.restore();
        }
        self.state_stack.clear();
        self.reported_nodes = 0;
        self.unreported_nodes = 0;
    }

    // Declares that we are starting to process a new node.
    //
    // Each recursive call to `run` begins with a call to
    // `node_begin`. The returned value is a TT entry telling
    // everything we know about the current position.
    #[inline]
    fn node_begin(&mut self) -> EntryData {
        // Consult the transposition table.
        let entry = if let Some(e) = self.tt.probe(self.position.hash()) {
            e
        } else {
            EntryData::new(0, BOUND_NONE, 0, 0, self.position.evaluate_static())
        };
        self.state_stack.push(NodeState {
            phase: NodePhase::Pristine,
            entry: entry,
        });
        entry
    }

    // Declares that we are done processing the current node.
    //
    // Each recursive call to `run` ends with a call to `node_end`.
    #[inline]
    fn node_end(&mut self) {
        if let NodePhase::Pristine = self.state_stack.last().unwrap().phase {
            // For pristine nodes we have not saved the move list
            // yet, so we should not restore it.
        } else {
            self.moves.restore();
        }
        self.state_stack.pop();
    }

    // Plays the next legal move in the current position and returns
    // it.
    //
    // Each call to `do_move` for the same position will play and
    // return a different move. When all legal moves has been played,
    // `None` will be returned. `do_move` will do whatever it can to
    // play the best moves first, and the worst last. It will also try
    // to be efficient, for example it will generate the list of all
    // pseudo-legal moves at the last possible moment.
    #[inline]
    fn do_move(&mut self) -> Option<Move> {
        let state = self.state_stack.last_mut().unwrap();

        if let NodePhase::Pristine = state.phase {
            // We save the move list at the last possible moment,
            // because most of the nodes are leafs.
            self.moves.save();

            // We always try the hash move first.
            state.phase = NodePhase::TriedHashMove;
            if state.entry.move16() != 0 {
                if let Some(m) = self.position.try_move_digest(state.entry.move16()) {
                    if self.position.do_move(m) {
                        return Some(m);
                    }
                }
            }
        }

        if let NodePhase::TriedHashMove = state.phase {
            // After the hash move, we generate all pseudo-legal
            // moves. But we should not forget to remove the already
            // tried hash move from the list.
            self.position.generate_moves(self.moves);
            if state.entry.move16() != 0 {
                self.moves.remove_move(state.entry.move16());
            }
            state.phase = NodePhase::GeneratedMoves;
        }

        // For the last, we spit the generated moves out.
        while let Some(m) = self.moves.remove_best_move() {
            if self.position.do_move(m) {
                return Some(m);
            }
        }
        None
    }

    // Takes the last played move back.
    #[inline]
    fn undo_move(&mut self) {
        self.position.undo_move();
    }

    // Stores updated node information in the transposition table.
    #[inline]
    fn store(&mut self, value: Value, bound: BoundType, depth: u8, best_move: Move) {
        let entry = &self.state_stack.last().unwrap().entry;
        let move16 = match best_move.digest() {
            0 => entry.move16(),
            x => x,
        };
        self.tt.store(self.position.hash(),
                      EntryData::new(value, bound, depth, move16, entry.eval_value()));
    }

    // Reports search progress.
    //
    // From time to time, we should report how many nodes had been
    // searched since the beginning of the search. This also gives an
    // opportunity for the search to be terminated.
    #[inline]
    fn report_progress(&mut self, new_nodes: NodeCount) -> Result<(), TerminatedSearch> {
        self.unreported_nodes += new_nodes;
        if self.unreported_nodes > NODE_COUNT_REPORT_INTERVAL {
            self.reported_nodes += self.unreported_nodes;
            self.unreported_nodes = 0;
            try!((*self.report_function)(self.reported_nodes));
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::Search;
    use chess_move::*;
    use tt::*;
    use position::Position;

    #[test]
    fn test_search() {
        let p = Position::from_fen("8/8/8/8/3q3k/7n/6PP/2Q2R1K b - - 0 1").ok().unwrap();
        let tt = TranspositionTable::new();
        let mut moves = MoveStack::new();
        let mut report = |_| Ok(());
        let mut search = Search::new(p, &tt, &mut moves, &mut report);
        let value = search.run(-30000, 30000, 2)
                          .ok()
                          .unwrap();
        assert!(value < -300);
        search.reset();
        let value = search.run(-30000, 30000, 4)
                          .ok()
                          .unwrap();
        assert!(value >= 20000);
    }
}
<|endoftext|>