<commit_before>// Copyright 2013-2019 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Windows console handling

// FIXME (#13400): this is only a tiny fraction of the Windows console api

use crate::color;
use crate::Attr;
use crate::Error;
use crate::Result;
use crate::Terminal;
use std::io;
use std::io::prelude::*;
use std::ops::Deref;
use std::ptr;

use winapi::shared::minwindef::{DWORD, WORD};
use winapi::um::consoleapi::{GetConsoleMode, SetConsoleMode};
use winapi::um::fileapi::{CreateFileA, OPEN_EXISTING};
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::wincon::FillConsoleOutputAttribute;
use winapi::um::wincon::{
    FillConsoleOutputCharacterW, GetConsoleScreenBufferInfo, CONSOLE_SCREEN_BUFFER_INFO, COORD,
};
use winapi::um::wincon::{SetConsoleCursorPosition, SetConsoleTextAttribute};
use winapi::um::wincon::{BACKGROUND_INTENSITY, ENABLE_VIRTUAL_TERMINAL_PROCESSING};
use winapi::um::winnt::{FILE_SHARE_WRITE, GENERIC_READ, GENERIC_WRITE, HANDLE};

/// Console info which can be used by a Terminal implementation
/// which uses the Win32 Console API.
pub struct WinConsoleInfo {
    def_foreground: color::Color,
    def_background: color::Color,
    foreground: color::Color,
    background: color::Color,
    reverse: bool,
    secure: bool,
    standout: bool,
}

/// A Terminal implementation which uses the Win32 Console API.
pub struct WinConsole<T> {
    buf: T,
    info: WinConsoleInfo,
}

fn color_to_bits(color: color::Color) -> u16 {
    // magic numbers from mingw-w64's wincon.h

    let bits = match color % 8 {
        color::BLACK => 0,
        color::BLUE => 0x1,
        color::GREEN => 0x2,
        color::RED => 0x4,
        color::YELLOW => 0x2 | 0x4,
        color::MAGENTA => 0x1 | 0x4,
        color::CYAN => 0x1 | 0x2,
        color::WHITE => 0x1 | 0x2 | 0x4,
        _ => unreachable!(),
    };

    if color >= 8 {
        bits | 0x8
    } else {
        bits
    }
}

fn bits_to_color(bits: u16) -> color::Color {
    let color = match bits & 0x7 {
        0 => color::BLACK,
        0x1 => color::BLUE,
        0x2 => color::GREEN,
        0x4 => color::RED,
        0x6 => color::YELLOW,
        0x5 => color::MAGENTA,
        0x3 => color::CYAN,
        0x7 => color::WHITE,
        _ => unreachable!(),
    };

    color | (bits as u32 & 0x8) // copy the hi-intensity bit
}

struct HandleWrapper {
    inner: HANDLE,
}

impl HandleWrapper {
    fn new(h: HANDLE) -> HandleWrapper {
        HandleWrapper { inner: h }
    }
}

impl Drop for HandleWrapper {
    fn drop(&mut self) {
        if self.inner != INVALID_HANDLE_VALUE {
            unsafe {
                CloseHandle(self.inner);
            }
        }
    }
}

impl Deref for HandleWrapper {
    type Target = HANDLE;
    fn deref(&self) -> &HANDLE {
        &self.inner
    }
}

/// Just get a handle to the current console buffer whatever it is
fn conout() -> io::Result<HandleWrapper> {
    let name = b"CONOUT$\0";
    let handle = unsafe {
        CreateFileA(
            name.as_ptr() as *const i8,
            GENERIC_READ | GENERIC_WRITE,
            FILE_SHARE_WRITE,
            ptr::null_mut(),
            OPEN_EXISTING,
            0,
            ptr::null_mut(),
        )
    };
    if handle == INVALID_HANDLE_VALUE {
        Err(io::Error::last_os_error())
    } else {
        Ok(HandleWrapper::new(handle))
    }
}

unsafe fn set_flag(handle: HANDLE, flag: DWORD) -> io::Result<()> {
    let mut curr_mode: DWORD = 0;
    if GetConsoleMode(handle, &mut curr_mode) == 0 {
        return Err(io::Error::last_os_error());
    }

    if SetConsoleMode(handle, curr_mode | flag) == 0 {
        return Err(io::Error::last_os_error());
    }
    return Ok(());
}

/// Check if console supports ansi codes (should succeed on Windows 10)
pub fn supports_ansi() -> bool {
    conout()
        .and_then(|handle| unsafe { set_flag(*handle, ENABLE_VIRTUAL_TERMINAL_PROCESSING) })
        .is_ok()
}

// This test will only pass if it is running in an actual console, probably
#[test]
fn test_conout() {
    assert!(conout().is_ok())
}

#[rustversion::before(1.36)]
unsafe fn get_console_screen_buffer_info(handle: HANDLE) -> io::Result<CONSOLE_SCREEN_BUFFER_INFO> {
    let mut buffer_info = ::std::mem::uninitialized();
    if GetConsoleScreenBufferInfo(handle, &mut buffer_info) == 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(buffer_info)
    }
}
#[rustversion::since(1.36)]
unsafe fn get_console_screen_buffer_info(handle: HANDLE) -> io::Result<CONSOLE_SCREEN_BUFFER_INFO> {
    let mut buffer_info = ::std::mem::MaybeUninit::uninit().assume_init();
    if GetConsoleScreenBufferInfo(handle, &mut buffer_info) == 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(buffer_info)
    }
}

impl WinConsoleInfo {
    /// Returns `Err` whenever console info cannot be retrieved for some
    /// reason.
    pub fn from_env() -> io::Result<WinConsoleInfo> {
        let fg;
        let bg;
        let handle = conout()?;
        unsafe {
            let buffer_info = get_console_screen_buffer_info(*handle)?;
            fg = bits_to_color(buffer_info.wAttributes);
            bg = bits_to_color(buffer_info.wAttributes >> 4);
        }
        Ok(WinConsoleInfo {
            def_foreground: fg,
            def_background: bg,
            foreground: fg,
            background: bg,
            reverse: false,
            secure: false,
            standout: false,
        })
    }
}

impl<T: Write + Send> WinConsole<T> {
    fn apply(&mut self) -> io::Result<()> {
        let out = conout()?;
        let _unused = self.buf.flush();

        let (mut fg, bg) = if self.info.reverse {
            (self.info.background, self.info.foreground)
        } else {
            (self.info.foreground, self.info.background)
        };

        if self.info.secure {
            fg = bg;
        }

        let mut accum: WORD = 0;

        accum |= color_to_bits(fg);
        accum |= color_to_bits(bg) << 4;

        if self.info.standout {
            accum |= BACKGROUND_INTENSITY;
        } else {
            accum &= BACKGROUND_INTENSITY ^ 0xFF;
        }

        unsafe {
            SetConsoleTextAttribute(*out, accum);
        }
        Ok(())
    }

    /// Create a new WinConsole with the given WinConsoleInfo and out
    pub fn new_with_consoleinfo(out: T, info: WinConsoleInfo) -> WinConsole<T> {
        WinConsole { buf: out, info }
    }

    /// Returns `Err` whenever the terminal cannot be created for some
    /// reason.
    pub fn new(out: T) -> io::Result<WinConsole<T>> {
        let info = WinConsoleInfo::from_env()?;
        Ok(Self::new_with_consoleinfo(out, info))
    }
}

impl<T: Write> Write for WinConsole<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.buf.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.buf.flush()
    }
}

impl<T: Write + Send> Terminal for WinConsole<T> {
    type Output = T;

    fn fg(&mut self, color: color::Color) -> Result<()> {
        self.info.foreground = color;
        self.apply()?;

        Ok(())
    }

    fn bg(&mut self, color: color::Color) -> Result<()> {
        self.info.background = color;
        self.apply()?;

        Ok(())
    }

    fn attr(&mut self, attr: Attr) -> Result<()> {
        match attr {
            Attr::ForegroundColor(f) => {
                self.info.foreground = f;
                self.apply()?;
                Ok(())
            }
            Attr::BackgroundColor(b) => {
                self.info.background = b;
                self.apply()?;
                Ok(())
            }
            Attr::Reverse => {
                self.info.reverse = true;
                self.apply()?;
                Ok(())
            }
            Attr::Secure => {
                self.info.secure = true;
                self.apply()?;
                Ok(())
            }
            Attr::Standout(v) => {
                self.info.standout = v;
                self.apply()?;
                Ok(())
            }
            _ => Err(Error::NotSupported),
        }
    }

    fn supports_attr(&self, attr: Attr) -> bool {
        match attr {
            Attr::ForegroundColor(_)
            | Attr::BackgroundColor(_)
            | Attr::Standout(_)
            | Attr::Reverse
            | Attr::Secure => true,
            _ => false,
        }
    }

    fn reset(&mut self) -> Result<()> {
        self.info.foreground = self.info.def_foreground;
        self.info.background = self.info.def_background;
        self.info.reverse = false;
        self.info.secure = false;
        self.info.standout = false;
        self.apply()?;

        Ok(())
    }

    fn supports_reset(&self) -> bool {
        true
    }

    fn supports_color(&self) -> bool {
        true
    }

    fn cursor_up(&mut self) -> Result<()> {
        let _unused = self.buf.flush();
        let handle = conout()?;
        unsafe {
            let buffer_info = get_console_screen_buffer_info(*handle)?;
            let (x, y) = (
                buffer_info.dwCursorPosition.X,
                buffer_info.dwCursorPosition.Y,
            );
            if y == 0 {
                // Even though this might want to be a CursorPositionInvalid, on Unix there
                // is no checking to see if the cursor is already on the first line.
                // I'm not sure what the ideal behavior is, but I think it'd be silly to have
                // cursor_up fail in this case.
                Ok(())
            } else {
                let pos = COORD { X: x, Y: y - 1 };
                if SetConsoleCursorPosition(*handle, pos) != 0 {
                    Ok(())
                } else {
                    Err(io::Error::last_os_error().into())
                }
            }
        }
    }

    fn delete_line(&mut self) -> Result<()> {
        let _unused = self.buf.flush();
        let handle = conout()?;
        unsafe {
            let buffer_info = get_console_screen_buffer_info(*handle)?;
            let pos = buffer_info.dwCursorPosition;
            let size = buffer_info.dwSize;
            let num = (size.X - pos.X) as DWORD;
            let mut written = 0;
            // 0x0020u16 is ' ' (space) in UTF-16 (same as ascii)
            if FillConsoleOutputCharacterW(*handle, 0x0020, num, pos, &mut written) == 0 {
                return Err(io::Error::last_os_error().into());
            }
            if FillConsoleOutputAttribute(*handle, 0, num, pos, &mut written) == 0 {
                return Err(io::Error::last_os_error().into());
            }
            // Similar reasoning for not failing as in cursor_up -- it doesn't even make
            // sense to
            // me that these APIs could have written 0, unless the terminal is width zero.
            Ok(())
        }
    }

    fn carriage_return(&mut self) -> Result<()> {
        let _unused = self.buf.flush();
        let handle = conout()?;
        unsafe {
            let buffer_info = get_console_screen_buffer_info(*handle)?;
            let COORD { X: x, Y: y } = buffer_info.dwCursorPosition;
            if x == 0 {
                Err(Error::CursorDestinationInvalid)
            } else {
                let pos = COORD { X: 0, Y: y };
                if SetConsoleCursorPosition(*handle, pos) != 0 {
                    Ok(())
                } else {
                    Err(io::Error::last_os_error().into())
                }
            }
        }
    }

    fn get_ref<'a>(&'a self) -> &'a T {
        &self.buf
    }

    fn get_mut<'a>(&'a mut self) -> &'a mut T {
        &mut self.buf
    }

    fn into_inner(self) -> T
    where
        Self: Sized,
    {
        self.buf
    }
}
<commit_msg>Call .assume_init() after CONSOLE_SCREEN_BUFFER_INFO is initialized<commit_after>// Copyright 2013-2019 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Windows console handling

// FIXME (#13400): this is only a tiny fraction of the Windows console api

use crate::color;
use crate::Attr;
use crate::Error;
use crate::Result;
use crate::Terminal;
use std::io;
use std::io::prelude::*;
use std::ops::Deref;
use std::ptr;

use winapi::shared::minwindef::{DWORD, WORD};
use winapi::um::consoleapi::{GetConsoleMode, SetConsoleMode};
use winapi::um::fileapi::{CreateFileA, OPEN_EXISTING};
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::wincon::FillConsoleOutputAttribute;
use winapi::um::wincon::{
    FillConsoleOutputCharacterW, GetConsoleScreenBufferInfo, CONSOLE_SCREEN_BUFFER_INFO, COORD,
};
use winapi::um::wincon::{SetConsoleCursorPosition, SetConsoleTextAttribute};
use winapi::um::wincon::{BACKGROUND_INTENSITY, ENABLE_VIRTUAL_TERMINAL_PROCESSING};
use winapi::um::winnt::{FILE_SHARE_WRITE, GENERIC_READ, GENERIC_WRITE, HANDLE};

/// Console info which can be used by a Terminal implementation
/// which uses the Win32 Console API.
pub struct WinConsoleInfo {
    def_foreground: color::Color,
    def_background: color::Color,
    foreground: color::Color,
    background: color::Color,
    reverse: bool,
    secure: bool,
    standout: bool,
}

/// A Terminal implementation which uses the Win32 Console API.
pub struct WinConsole<T> {
    buf: T,
    info: WinConsoleInfo,
}

fn color_to_bits(color: color::Color) -> u16 {
    // magic numbers from mingw-w64's wincon.h

    let bits = match color % 8 {
        color::BLACK => 0,
        color::BLUE => 0x1,
        color::GREEN => 0x2,
        color::RED => 0x4,
        color::YELLOW => 0x2 | 0x4,
        color::MAGENTA => 0x1 | 0x4,
        color::CYAN => 0x1 | 0x2,
        color::WHITE => 0x1 | 0x2 | 0x4,
        _ => unreachable!(),
    };

    if color >= 8 {
        bits | 0x8
    } else {
        bits
    }
}

fn bits_to_color(bits: u16) -> color::Color {
    let color = match bits & 0x7 {
        0 => color::BLACK,
        0x1 => color::BLUE,
        0x2 => color::GREEN,
        0x4 => color::RED,
        0x6 => color::YELLOW,
        0x5 => color::MAGENTA,
        0x3 => color::CYAN,
        0x7 => color::WHITE,
        _ => unreachable!(),
    };

    color | (bits as u32 & 0x8) // copy the hi-intensity bit
}

struct HandleWrapper {
    inner: HANDLE,
}

impl HandleWrapper {
    fn new(h: HANDLE) -> HandleWrapper {
        HandleWrapper { inner: h }
    }
}

impl Drop for HandleWrapper {
    fn drop(&mut self) {
        if self.inner != INVALID_HANDLE_VALUE {
            unsafe {
                CloseHandle(self.inner);
            }
        }
    }
}

impl Deref for HandleWrapper {
    type Target = HANDLE;
    fn deref(&self) -> &HANDLE {
        &self.inner
    }
}

/// Just get a handle to the current console buffer whatever it is
fn conout() -> io::Result<HandleWrapper> {
    let name = b"CONOUT$\0";
    let handle = unsafe {
        CreateFileA(
            name.as_ptr() as *const i8,
            GENERIC_READ | GENERIC_WRITE,
            FILE_SHARE_WRITE,
            ptr::null_mut(),
            OPEN_EXISTING,
            0,
            ptr::null_mut(),
        )
    };
    if handle == INVALID_HANDLE_VALUE {
        Err(io::Error::last_os_error())
    } else {
        Ok(HandleWrapper::new(handle))
    }
}

unsafe fn set_flag(handle: HANDLE, flag: DWORD) -> io::Result<()> {
    let mut curr_mode: DWORD = 0;
    if GetConsoleMode(handle, &mut curr_mode) == 0 {
        return Err(io::Error::last_os_error());
    }

    if SetConsoleMode(handle, curr_mode | flag) == 0 {
        return Err(io::Error::last_os_error());
    }
    return Ok(());
}

/// Check if console supports ansi codes (should succeed on Windows 10)
pub fn supports_ansi() -> bool {
    conout()
        .and_then(|handle| unsafe { set_flag(*handle, ENABLE_VIRTUAL_TERMINAL_PROCESSING) })
        .is_ok()
}

// This test will only pass if it is running in an actual console, probably
#[test]
fn test_conout() {
    assert!(conout().is_ok())
}

#[rustversion::before(1.36)]
unsafe fn get_console_screen_buffer_info(handle: HANDLE) -> io::Result<CONSOLE_SCREEN_BUFFER_INFO> {
    let mut buffer_info = ::std::mem::uninitialized();
    if GetConsoleScreenBufferInfo(handle, &mut buffer_info) == 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(buffer_info)
    }
}
#[rustversion::since(1.36)]
unsafe fn get_console_screen_buffer_info(handle: HANDLE) -> io::Result<CONSOLE_SCREEN_BUFFER_INFO> {
    let mut buffer_info = ::std::mem::MaybeUninit::uninit();
    if GetConsoleScreenBufferInfo(handle, buffer_info.as_mut_ptr()) == 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(buffer_info.assume_init())
    }
}

impl WinConsoleInfo {
    /// Returns `Err` whenever console info cannot be retrieved for some
    /// reason.
    pub fn from_env() -> io::Result<WinConsoleInfo> {
        let fg;
        let bg;
        let handle = conout()?;
        unsafe {
            let buffer_info = get_console_screen_buffer_info(*handle)?;
            fg = bits_to_color(buffer_info.wAttributes);
            bg = bits_to_color(buffer_info.wAttributes >> 4);
        }
        Ok(WinConsoleInfo {
            def_foreground: fg,
            def_background: bg,
            foreground: fg,
            background: bg,
            reverse: false,
            secure: false,
            standout: false,
        })
    }
}

impl<T: Write + Send> WinConsole<T> {
    fn apply(&mut self) -> io::Result<()> {
        let out = conout()?;
        let _unused = self.buf.flush();

        let (mut fg, bg) = if self.info.reverse {
            (self.info.background, self.info.foreground)
        } else {
            (self.info.foreground, self.info.background)
        };

        if self.info.secure {
            fg = bg;
        }

        let mut accum: WORD = 0;

        accum |= color_to_bits(fg);
        accum |= color_to_bits(bg) << 4;

        if self.info.standout {
            accum |= BACKGROUND_INTENSITY;
        } else {
            accum &= BACKGROUND_INTENSITY ^ 0xFF;
        }

        unsafe {
            SetConsoleTextAttribute(*out, accum);
        }
        Ok(())
    }

    /// Create a new WinConsole with the given WinConsoleInfo and out
    pub fn new_with_consoleinfo(out: T, info: WinConsoleInfo) -> WinConsole<T> {
        WinConsole { buf: out, info }
    }

    /// Returns `Err` whenever the terminal cannot be created for some
    /// reason.
    pub fn new(out: T) -> io::Result<WinConsole<T>> {
        let info = WinConsoleInfo::from_env()?;
        Ok(Self::new_with_consoleinfo(out, info))
    }
}

impl<T: Write> Write for WinConsole<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.buf.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.buf.flush()
    }
}

impl<T: Write + Send> Terminal for WinConsole<T> {
    type Output = T;

    fn fg(&mut self, color: color::Color) -> Result<()> {
        self.info.foreground = color;
        self.apply()?;

        Ok(())
    }

    fn bg(&mut self, color: color::Color) -> Result<()> {
        self.info.background = color;
        self.apply()?;

        Ok(())
    }

    fn attr(&mut self, attr: Attr) -> Result<()> {
        match attr {
            Attr::ForegroundColor(f) => {
                self.info.foreground = f;
                self.apply()?;
                Ok(())
            }
            Attr::BackgroundColor(b) => {
                self.info.background = b;
                self.apply()?;
                Ok(())
            }
            Attr::Reverse => {
                self.info.reverse = true;
                self.apply()?;
                Ok(())
            }
            Attr::Secure => {
                self.info.secure = true;
                self.apply()?;
                Ok(())
            }
            Attr::Standout(v) => {
                self.info.standout = v;
                self.apply()?;
                Ok(())
            }
            _ => Err(Error::NotSupported),
        }
    }

    fn supports_attr(&self, attr: Attr) -> bool {
        match attr {
            Attr::ForegroundColor(_)
            | Attr::BackgroundColor(_)
            | Attr::Standout(_)
            | Attr::Reverse
            | Attr::Secure => true,
            _ => false,
        }
    }

    fn reset(&mut self) -> Result<()> {
        self.info.foreground = self.info.def_foreground;
        self.info.background = self.info.def_background;
        self.info.reverse = false;
        self.info.secure = false;
        self.info.standout = false;
        self.apply()?;

        Ok(())
    }

    fn supports_reset(&self) -> bool {
        true
    }

    fn supports_color(&self) -> bool {
        true
    }

    fn cursor_up(&mut self) -> Result<()> {
        let _unused = self.buf.flush();
        let handle = conout()?;
        unsafe {
            let buffer_info = get_console_screen_buffer_info(*handle)?;
            let (x, y) = (
                buffer_info.dwCursorPosition.X,
                buffer_info.dwCursorPosition.Y,
            );
            if y == 0 {
                // Even though this might want to be a CursorPositionInvalid, on Unix there
                // is no checking to see if the cursor is already on the first line.
                // I'm not sure what the ideal behavior is, but I think it'd be silly to have
                // cursor_up fail in this case.
                Ok(())
            } else {
                let pos = COORD { X: x, Y: y - 1 };
                if SetConsoleCursorPosition(*handle, pos) != 0 {
                    Ok(())
                } else {
                    Err(io::Error::last_os_error().into())
                }
            }
        }
    }

    fn delete_line(&mut self) -> Result<()> {
        let _unused = self.buf.flush();
        let handle = conout()?;
        unsafe {
            let buffer_info = get_console_screen_buffer_info(*handle)?;
            let pos = buffer_info.dwCursorPosition;
            let size = buffer_info.dwSize;
            let num = (size.X - pos.X) as DWORD;
            let mut written = 0;
            // 0x0020u16 is ' ' (space) in UTF-16 (same as ascii)
            if FillConsoleOutputCharacterW(*handle, 0x0020, num, pos, &mut written) == 0 {
                return Err(io::Error::last_os_error().into());
            }
            if FillConsoleOutputAttribute(*handle, 0, num, pos, &mut written) == 0 {
                return Err(io::Error::last_os_error().into());
            }
            // Similar reasoning for not failing as in cursor_up -- it doesn't even make
            // sense to
            // me that these APIs could have written 0, unless the terminal is width zero.
            Ok(())
        }
    }

    fn carriage_return(&mut self) -> Result<()> {
        let _unused = self.buf.flush();
        let handle = conout()?;
        unsafe {
            let buffer_info = get_console_screen_buffer_info(*handle)?;
            let COORD { X: x, Y: y } = buffer_info.dwCursorPosition;
            if x == 0 {
                Err(Error::CursorDestinationInvalid)
            } else {
                let pos = COORD { X: 0, Y: y };
                if SetConsoleCursorPosition(*handle, pos) != 0 {
                    Ok(())
                } else {
                    Err(io::Error::last_os_error().into())
                }
            }
        }
    }

    fn get_ref<'a>(&'a self) -> &'a T {
        &self.buf
    }

    fn get_mut<'a>(&'a mut self) -> &'a mut T {
        &mut self.buf
    }

    fn into_inner(self) -> T
    where
        Self: Sized,
    {
        self.buf
    }
}
<|endoftext|>


<commit_before>use crate::model::{read_uint, write_uint};
use serde::Serialize;
use shakmaty::{Color, Outcome};
use std::{
    io::{self, Read, Write},
    ops::AddAssign,
};

#[derive(Debug, Default, Clone, Eq, PartialEq, Serialize)]
pub struct Stats {
    pub white: u64,
    pub draws: u64,
    pub black: u64,
}

impl From<Outcome> for Stats {
    fn from(outcome: Outcome) -> Stats {
        Stats {
            white: if outcome.winner() == Some(Color::White) {
                1
            } else {
                0
            },
            black: if outcome.winner() == Some(Color::Black) {
                1
            } else {
                0
            },
            draws: if outcome.winner().is_none() { 1 } else { 0 },
        }
    }
}

impl AddAssign for Stats {
    fn add_assign(&mut self, rhs: Stats) {
        self.white += rhs.white;
        self.draws += rhs.draws;
        self.black += rhs.black;
    }
}

impl Stats {
    pub fn total(&self) -> u64 {
        self.white + self.draws + self.black
    }

    pub fn is_empty(&self) -> bool {
        self.total() == 0
    }

    pub fn read<R: Read>(reader: &mut R) -> io::Result<Stats> {
        Ok(match read_uint(reader)? {
            0 => Stats { white: 1, draws: 0, black: 0 },
            1 => Stats { white: 0, draws: 0, black: 1 },
            2 => Stats { white: 0, draws: 1, black: 0 },
            white_plus_three => Stats {
                white: white_plus_three - 3,
                draws: read_uint(reader)?,
                black: read_uint(reader)?,
            },
        })
    }

    pub fn write<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        match *self {
            Stats { white: 1, draws: 0, black: 0 } => write_uint(writer, 0),
            Stats { white: 0, draws: 0, black: 1 } => write_uint(writer, 1),
            Stats { white: 0, draws: 1, black: 0 } => write_uint(writer, 2),
            Stats { white, draws, black } => {
                write_uint(writer, white + 3)?;
                write_uint(writer, draws)?;
                write_uint(writer, black)
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use quickcheck::{quickcheck, Arbitrary, Gen};
    use std::io::Cursor;

    impl Arbitrary for Stats {
        fn arbitrary(g: &mut Gen) -> Self {
            Stats {
                white: u64::from(u32::arbitrary(g)),
                draws: u64::from(u32::arbitrary(g)),
                black: u64::from(u32::arbitrary(g)),
            }
        }
    }

    quickcheck! {
        fn test_stats_roundtrip(stats: Stats) -> bool {
            let mut cursor = Cursor::new(Vec::new());
            stats.write(&mut cursor).unwrap();

            let mut cursor = Cursor::new(cursor.into_inner());
            Stats::read(&mut cursor).unwrap() == stats
        }
    }
}
<commit_msg>cargo fmt<commit_after>use crate::model::{read_uint, write_uint};
use serde::Serialize;
use shakmaty::{Color, Outcome};
use std::{
    io::{self, Read, Write},
    ops::AddAssign,
};

#[derive(Debug, Default, Clone, Eq, PartialEq, Serialize)]
pub struct Stats {
    pub white: u64,
    pub draws: u64,
    pub black: u64,
}

impl From<Outcome> for Stats {
    fn from(outcome: Outcome) -> Stats {
        Stats {
            white: if outcome.winner() == Some(Color::White) {
                1
            } else {
                0
            },
            black: if outcome.winner() == Some(Color::Black) {
                1
            } else {
                0
            },
            draws: if outcome.winner().is_none() { 1 } else { 0 },
        }
    }
}

impl AddAssign for Stats {
    fn add_assign(&mut self, rhs: Stats) {
        self.white += rhs.white;
        self.draws += rhs.draws;
        self.black += rhs.black;
    }
}

impl Stats {
    pub fn total(&self) -> u64 {
        self.white + self.draws + self.black
    }

    pub fn is_empty(&self) -> bool {
        self.total() == 0
    }

    pub fn read<R: Read>(reader: &mut R) -> io::Result<Stats> {
        Ok(match read_uint(reader)? {
            0 => Stats {
                white: 1,
                draws: 0,
                black: 0,
            },
            1 => Stats {
                white: 0,
                draws: 0,
                black: 1,
            },
            2 => Stats {
                white: 0,
                draws: 1,
                black: 0,
            },
            white_plus_three => Stats {
                white: white_plus_three - 3,
                draws: read_uint(reader)?,
                black: read_uint(reader)?,
            },
        })
    }

    pub fn write<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        match *self {
            Stats {
                white: 1,
                draws: 0,
                black: 0,
            } => write_uint(writer, 0),
            Stats {
                white: 0,
                draws: 0,
                black: 1,
            } => write_uint(writer, 1),
            Stats {
                white: 0,
                draws: 1,
                black: 0,
            } => write_uint(writer, 2),
            Stats {
                white,
                draws,
                black,
            } => {
                write_uint(writer, white + 3)?;
                write_uint(writer, draws)?;
                write_uint(writer, black)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use quickcheck::{quickcheck, Arbitrary, Gen};
    use std::io::Cursor;

    impl Arbitrary for Stats {
        fn arbitrary(g: &mut Gen) -> Self {
            Stats {
                white: u64::from(u32::arbitrary(g)),
                draws: u64::from(u32::arbitrary(g)),
                black: u64::from(u32::arbitrary(g)),
            }
        }
    }

    quickcheck! {
        fn test_stats_roundtrip(stats: Stats) -> bool {
            let mut cursor = Cursor::new(Vec::new());
            stats.write(&mut cursor).unwrap();

            let mut cursor = Cursor::new(cursor.into_inner());
            Stats::read(&mut cursor).unwrap() == stats
        }
    }
}
<|endoftext|>


<commit_before>use diesel::prelude::*;

use crate::app::App;
use crate::github::{github_api, team_url};
use crate::util::errors::{cargo_err, AppResult, NotFound};

use oauth2::AccessToken;

use crate::models::{Crate, CrateOwner, Owner, OwnerKind, User};
use crate::schema::{crate_owners, teams};
use crate::views::EncodableTeam;

/// For now, just a Github Team. Can be upgraded to other teams
/// later if desirable.
#[derive(Queryable, Identifiable, Serialize, Deserialize, Debug)]
pub struct Team {
    /// Unique table id
    pub id: i32,
    /// "github:org:team"
    /// An opaque unique ID, that was at one point parsed out to query Github.
    /// We only query membership with github using the github_id, though.
    /// This is the only name we should ever talk to Cargo about.
    pub login: String,
    /// The GitHub API works on team ID numbers. This can change, if a team
    /// is deleted and then recreated with the same name!!!
    pub github_id: i32,
    /// Sugary goodness
    pub name: Option<String>,
    pub avatar: Option<String>,
}

#[derive(Insertable, AsChangeset, Debug)]
#[table_name = "teams"]
pub struct NewTeam<'a> {
    pub login: &'a str,
    pub github_id: i32,
    pub name: Option<String>,
    pub avatar: Option<String>,
}

impl<'a> NewTeam<'a> {
    pub fn new(
        login: &'a str,
        github_id: i32,
        name: Option<String>,
        avatar: Option<String>,
    ) -> Self {
        NewTeam {
            login,
            github_id,
            name,
            avatar,
        }
    }

    pub fn create_or_update(&self, conn: &PgConnection) -> QueryResult<Team> {
        use crate::schema::teams::dsl::*;
        use diesel::insert_into;

        insert_into(teams)
            .values(self)
            .on_conflict(login)
            .do_update()
            .set(self)
            .get_result(conn)
    }
}

impl Team {
    /// Tries to create the Team in the DB (assumes a `:` has already been found).
    ///
    /// # Panics
    ///
    /// This function will panic if login contains less than 2 `:` characters.
    pub fn create_or_update(
        app: &App,
        conn: &PgConnection,
        login: &str,
        req_user: &User,
    ) -> AppResult<Self> {
        // must look like system:xxxxxxx
        let mut chunks = login.split(':');
        // unwrap is okay, split on an empty string still has 1 chunk
        match chunks.next().unwrap() {
            // github:rust-lang:owners
            "github" => {
                // unwrap is documented above as part of the calling contract
                let org = chunks.next().unwrap();
                let team = chunks.next().ok_or_else(|| {
                    cargo_err(
                        "missing github team argument; \
                         format is github:org:team",
                    )
                })?;
                Team::create_or_update_github_team(
                    app,
                    conn,
                    &login.to_lowercase(),
                    org,
                    team,
                    req_user,
                )
            }
            _ => Err(cargo_err(
                "unknown organization handler, \
                 only 'github:org:team' is supported",
            )),
        }
    }

    /// Tries to create or update a Github Team. Assumes `org` and `team` are
    /// correctly parsed out of the full `name`. `name` is passed as a
    /// convenience to avoid rebuilding it.
    fn create_or_update_github_team(
        app: &App,
        conn: &PgConnection,
        login: &str,
        org_name: &str,
        team_name: &str,
        req_user: &User,
    ) -> AppResult<Self> {
        // GET orgs/:org/teams
        // check that `team` is the `slug` in results, and grab its data

        // "sanitization"
        fn whitelist(c: char) -> bool {
            match c {
                'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_' => false,
                _ => true,
            }
        }

        if let Some(c) = org_name.chars().find(|c| whitelist(*c)) {
            return Err(cargo_err(&format_args!(
                "organization cannot contain special \
                 characters like {}",
                c
            )));
        }

        #[derive(Deserialize)]
        struct GithubTeam {
            id: i32,              // unique GH id (needed for membership queries)
            name: Option<String>, // Pretty name
        }

        let url = format!("/orgs/{}/teams/{}", org_name, team_name);
        let token = AccessToken::new(req_user.gh_access_token.clone());
        let team = github_api::<GithubTeam>(app, &url, &token).map_err(|_| {
            cargo_err(&format_args!(
                "could not find the github team {}/{}",
                org_name, team_name
            ))
        })?;

        if !team_with_gh_id_contains_user(app, team.id, req_user)? {
            return Err(cargo_err("only members of a team can add it as an owner"));
        }

        #[derive(Deserialize)]
        struct Org {
            avatar_url: Option<String>,
        }

        let url = format!("/orgs/{}", org_name);
        let org = github_api::<Org>(app, &url, &token)?;

        NewTeam::new(&login.to_lowercase(), team.id, team.name, org.avatar_url)
            .create_or_update(conn)
            .map_err(Into::into)
    }

    /// Phones home to Github to ask if this User is a member of the given team.
    /// Note that we're assuming that the given user is the one interested in
    /// the answer. If this is not the case, then we could accidentally leak
    /// private membership information here.
    pub fn contains_user(&self, app: &App, user: &User) -> AppResult<bool> {
        team_with_gh_id_contains_user(app, self.github_id, user)
    }

    pub fn owning(krate: &Crate, conn: &PgConnection) -> QueryResult<Vec<Owner>> {
        let base_query = CrateOwner::belonging_to(krate).filter(crate_owners::deleted.eq(false));
        let teams = base_query
            .inner_join(teams::table)
            .select(teams::all_columns)
            .filter(crate_owners::owner_kind.eq(OwnerKind::Team as i32))
            .load(conn)?
            .into_iter()
            .map(Owner::Team);

        Ok(teams.collect())
    }

    pub fn encodable(self) -> EncodableTeam {
        let Team {
            id,
            name,
            login,
            avatar,
            ..
        } = self;
        let url = team_url(&login);

        EncodableTeam {
            id,
            login,
            name,
            avatar,
            url: Some(url),
        }
    }
}

fn team_with_gh_id_contains_user(app: &App, github_id: i32, user: &User) -> AppResult<bool> {
    // GET teams/:team_id/memberships/:user_name
    // check that "state": "active"

    #[derive(Deserialize)]
    struct Membership {
        state: String,
    }

    let url = format!("/teams/{}/memberships/{}", &github_id, &user.gh_login);
    let token = AccessToken::new(user.gh_access_token.clone());
    let membership = match github_api::<Membership>(app, &url, &token) {
        // Officially how `false` is returned
        Err(ref e) if e.is::<NotFound>() => return Ok(false),
        x => x?,
    };

    // There is also `state: pending` for which we could possibly give
    // some feedback, but it's not obvious how that should work.
    Ok(membership.state == "active")
}
<commit_msg>Use `matches!` macro and rename sanitize function<commit_after>use diesel::prelude::*;

use crate::app::App;
use crate::github::{github_api, team_url};
use crate::util::errors::{cargo_err, AppResult, NotFound};

use oauth2::AccessToken;

use crate::models::{Crate, CrateOwner, Owner, OwnerKind, User};
use crate::schema::{crate_owners, teams};
use crate::views::EncodableTeam;

/// For now, just a Github Team. Can be upgraded to other teams
/// later if desirable.
#[derive(Queryable, Identifiable, Serialize, Deserialize, Debug)]
pub struct Team {
    /// Unique table id
    pub id: i32,
    /// "github:org:team"
    /// An opaque unique ID, that was at one point parsed out to query Github.
    /// We only query membership with github using the github_id, though.
    /// This is the only name we should ever talk to Cargo about.
    pub login: String,
    /// The GitHub API works on team ID numbers. This can change, if a team
    /// is deleted and then recreated with the same name!!!
    pub github_id: i32,
    /// Sugary goodness
    pub name: Option<String>,
    pub avatar: Option<String>,
}

#[derive(Insertable, AsChangeset, Debug)]
#[table_name = "teams"]
pub struct NewTeam<'a> {
    pub login: &'a str,
    pub github_id: i32,
    pub name: Option<String>,
    pub avatar: Option<String>,
}

impl<'a> NewTeam<'a> {
    pub fn new(
        login: &'a str,
        github_id: i32,
        name: Option<String>,
        avatar: Option<String>,
    ) -> Self {
        NewTeam {
            login,
            github_id,
            name,
            avatar,
        }
    }

    pub fn create_or_update(&self, conn: &PgConnection) -> QueryResult<Team> {
        use crate::schema::teams::dsl::*;
        use diesel::insert_into;

        insert_into(teams)
            .values(self)
            .on_conflict(login)
            .do_update()
            .set(self)
            .get_result(conn)
    }
}

impl Team {
    /// Tries to create the Team in the DB (assumes a `:` has already been found).
    ///
    /// # Panics
    ///
    /// This function will panic if login contains less than 2 `:` characters.
    pub fn create_or_update(
        app: &App,
        conn: &PgConnection,
        login: &str,
        req_user: &User,
    ) -> AppResult<Self> {
        // must look like system:xxxxxxx
        let mut chunks = login.split(':');
        // unwrap is okay, split on an empty string still has 1 chunk
        match chunks.next().unwrap() {
            // github:rust-lang:owners
            "github" => {
                // unwrap is documented above as part of the calling contract
                let org = chunks.next().unwrap();
                let team = chunks.next().ok_or_else(|| {
                    cargo_err(
                        "missing github team argument; \
                         format is github:org:team",
                    )
                })?;
                Team::create_or_update_github_team(
                    app,
                    conn,
                    &login.to_lowercase(),
                    org,
                    team,
                    req_user,
                )
            }
            _ => Err(cargo_err(
                "unknown organization handler, \
                 only 'github:org:team' is supported",
            )),
        }
    }

    /// Tries to create or update a Github Team. Assumes `org` and `team` are
    /// correctly parsed out of the full `name`. `name` is passed as a
    /// convenience to avoid rebuilding it.
    fn create_or_update_github_team(
        app: &App,
        conn: &PgConnection,
        login: &str,
        org_name: &str,
        team_name: &str,
        req_user: &User,
    ) -> AppResult<Self> {
        // GET orgs/:org/teams
        // check that `team` is the `slug` in results, and grab its data

        // "sanitization"
        fn is_allowed_char(c: char) -> bool {
            matches!(c, 'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_')
        }

        if let Some(c) = org_name.chars().find(|c| !is_allowed_char(*c)) {
            return Err(cargo_err(&format_args!(
                "organization cannot contain special \
                 characters like {}",
                c
            )));
        }

        #[derive(Deserialize)]
        struct GithubTeam {
            id: i32,              // unique GH id (needed for membership queries)
            name: Option<String>, // Pretty name
        }

        let url = format!("/orgs/{}/teams/{}", org_name, team_name);
        let token = AccessToken::new(req_user.gh_access_token.clone());
        let team = github_api::<GithubTeam>(app, &url, &token).map_err(|_| {
            cargo_err(&format_args!(
                "could not find the github team {}/{}",
                org_name, team_name
            ))
        })?;

        if !team_with_gh_id_contains_user(app, team.id, req_user)? {
            return Err(cargo_err("only members of a team can add it as an owner"));
        }

        #[derive(Deserialize)]
        struct Org {
            avatar_url: Option<String>,
        }

        let url = format!("/orgs/{}", org_name);
        let org = github_api::<Org>(app, &url, &token)?;

        NewTeam::new(&login.to_lowercase(), team.id, team.name, org.avatar_url)
            .create_or_update(conn)
            .map_err(Into::into)
    }

    /// Phones home to Github to ask if this User is a member of the given team.
    /// Note that we're assuming that the given user is the one interested in
    /// the answer. If this is not the case, then we could accidentally leak
    /// private membership information here.
    pub fn contains_user(&self, app: &App, user: &User) -> AppResult<bool> {
        team_with_gh_id_contains_user(app, self.github_id, user)
    }

    pub fn owning(krate: &Crate, conn: &PgConnection) -> QueryResult<Vec<Owner>> {
        let base_query = CrateOwner::belonging_to(krate).filter(crate_owners::deleted.eq(false));
        let teams = base_query
            .inner_join(teams::table)
            .select(teams::all_columns)
            .filter(crate_owners::owner_kind.eq(OwnerKind::Team as i32))
            .load(conn)?
            .into_iter()
            .map(Owner::Team);

        Ok(teams.collect())
    }

    pub fn encodable(self) -> EncodableTeam {
        let Team {
            id,
            name,
            login,
            avatar,
            ..
        } = self;
        let url = team_url(&login);

        EncodableTeam {
            id,
            login,
            name,
            avatar,
            url: Some(url),
        }
    }
}

fn team_with_gh_id_contains_user(app: &App, github_id: i32, user: &User) -> AppResult<bool> {
    // GET teams/:team_id/memberships/:user_name
    // check that "state": "active"

    #[derive(Deserialize)]
    struct Membership {
        state: String,
    }

    let url = format!("/teams/{}/memberships/{}", &github_id, &user.gh_login);
    let token = AccessToken::new(user.gh_access_token.clone());
    let membership = match github_api::<Membership>(app, &url, &token) {
        // Officially how `false` is returned
        Err(ref e) if e.is::<NotFound>() => return Ok(false),
        x => x?,
    };

    // There is also `state: pending` for which we could possibly give
    // some feedback, but it's not obvious how that should work.
    Ok(membership.state == "active")
}
<|endoftext|>


<commit_before>use std::collections::{HashMap, HashSet};
use std::f64;
use regex::Regex;
use rustc_serialize::json;

static DEFAULT_SMOOTHING: f64 = 1.0f64;

#[derive(Debug, Clone, RustcDecodable, RustcEncodable)]
/// Naive Bayes classifier
pub struct Classifier {
    vocab: HashSet<String>,
    num_examples: u32,
    smoothing: f64,
    classifications: HashMap<String, Classification>
}

#[derive(Debug, Clone, RustcDecodable, RustcEncodable)]
struct Classification {
    label: String,
    num_examples: u32,
    num_words: u32,
    probability: f64,
    default_word_probability: f64,
    words: HashMap<String, (u32, f64)>,
}

impl Classifier {

    /// Creates a new classifier
    pub fn new() -> Classifier {
        Classifier {
            vocab: HashSet::new(),
            num_examples: 0u32,
            smoothing: DEFAULT_SMOOTHING,
            classifications: HashMap::new(),
        }
    }

    /// Takes a document that has been tokenized into a vector of strings
    /// and a label and adds the document to the list of documents that the
    /// classifier is aware of and will train on next time the `train()` method is called
    pub fn add_document_tokenized(&mut self, document: &Vec<String>, label: &String) {
        if document.len() == 0 { return; }

        self.num_examples += 1;

        // make sure the classification already exists
        if !self.classifications.contains_key(label) {
            let c = Classification::new(label);
            self.classifications.insert(label.clone(), c);
        }

        let mut classification = self.classifications.get_mut(label).unwrap();

        for word in document.iter() {
            // add word to classification
            classification.add_word(word);

            // add word to vocab
            self.vocab.insert(word.to_string());
            // if self.vocab.contains_key(word) {
            //     *self.vocab.get_mut(word).unwrap() += 1;
            // } else {
            //     self.vocab.insert(word.to_string(), 1);
            // }
        }

        classification.num_examples += 1;
    }

    /// Takes a document and a label and tokenizes the document by
    /// breaking on whitespace characters. The document is added to the list
    /// of documents that the classifier is aware of and will train on next time the `train()` method is called
    pub fn add_document(&mut self, document: &String, label: &String) {
        self.add_document_tokenized(&split_document(document), label);
    }

    /// Adds a list of (document, label) tuples to the classifier
    pub fn add_documents(&mut self, examples: &Vec<(String, String)>) {
        for &(ref document, ref label) in examples.iter() {
            self.add_document(document, label);
        }
    }

    /// Adds a list of (tokenized document, label) tuples to the classifier
    pub fn add_documents_tokenized(&mut self, examples: &Vec<(Vec<String>, String)>) {
        for &(ref document, ref label) in examples.iter() {
            self.add_document_tokenized(document, label);
        }
    }

    /// Gets a vector of all of the labels that the classifier has seen so far
    pub fn get_labels(&self) -> Vec<String> {
        let labels: Vec<String> =
            self.classifications.values().map(|c| c.label.clone()).collect();
        labels
    }

    /// Sets the [smoothing](http://en.wikipedia.org/wiki/Additive_smoothing) value (must be greater than 0.0)
    pub fn set_smoothing(&mut self, smoothing: f64) {
        if smoothing <= 0.0 {
            panic!("smoothing value must be a positive number");
        }
        self.smoothing = smoothing;
    }

    /// Trains the classifier on the documents that have been observed so far
    pub fn train(&mut self) {
        for (_, classification) in self.classifications.iter_mut() {
            classification.train(&self.vocab, self.num_examples, self.smoothing);
        }
    }

    /// Takes an unlabeled document that has been tokenized into a vector of strings
    /// and then computes a classifying label for the document
    pub fn classify_tokenized(&self, document: &Vec<String>) -> String {
        let mut max_score = f64::NEG_INFINITY;
        let mut max_classification = None;

        for classification in self.classifications.values() {
            let score = classification.score_document(document, &self.vocab);
            if score > max_score {
                max_classification = Some(classification);
                max_score = score;
            }
        }

        max_classification.expect("no classification found").label.clone()
    }

    /// Takes an unlabeled document and tokenizes it by breaking on spaces and
    /// then computes a classifying label for the document
    pub fn classify(&self, document: &String) -> String {
        self.classify_tokenized(&split_document(document))
    }

    /// Similar to classify but instead of returning a single label, returns all
    /// labels and the probabilities of each one given the document
    pub fn get_document_probabilities_tokenized(&self, document: &Vec<String>) -> Vec<(String, f64)> {
        self.classifications.values().map(|classification| {
            let score = classification.score_document(document, &self.vocab);
            (classification.label.clone(), score)
        }).collect()
    }

    /// Similar to classify but instead of returning a single label, returns all
    /// labels and the probabilities of each one given the document
    pub fn get_document_probabilities(&self, document: &String) -> Vec<(String, f64)> {
        self.get_document_probabilities_tokenized(&split_document(document))
    }

    /// Encodes the classifier as a JSON string.
    pub fn to_json(&self) -> String {
        json::encode(self).ok().expect("encoding JSON failed")
    }

    /// Builds a new classifier from a JSON string
    pub fn from_json(encoded: &str) -> Classifier {
        let classifier: Classifier = json::decode(encoded).ok().expect("decoding JSON failed");
        classifier
    }

}

impl Classification {

    fn new(label: &String) -> Classification {
        Classification {
            label: label.clone(),
            num_examples: 0u32,
            num_words: 0u32,
            probability: 0.0f64,
            default_word_probability: 0.0f64,
            words: HashMap::new(),
        }
    }

    fn add_word(&mut self, word: &String) {
        self.num_words += 1;
        if self.words.contains_key(word) {
            self.words.get_mut(word).unwrap().0 += 1;
        } else {
            self.words.insert(word.clone(), (1, 0.0f64));
        }
    }

    fn train(&mut self, vocab: &HashSet<String>, total_examples: u32, smoothing: f64) {
        self.probability = self.num_examples as f64 / total_examples as f64;
        self.default_word_probability = smoothing / (self.num_words as f64 + smoothing * vocab.len() as f64);

        for word in vocab.iter() {
            if self.words.contains_key(word) {
                let word_entry = self.words.get_mut(word).unwrap();
                let word_count = word_entry.0;
                let p_word_given_label = (word_count as f64 + smoothing) / (self.num_words as f64 + smoothing * vocab.len() as f64);
                word_entry.1 = p_word_given_label;
            }
        }
    }

    fn score_document(&self, document: &Vec<String>, vocab: &HashSet<String>) -> f64 {
        let mut total = 0.0f64;
        for word in document.iter() {
            if vocab.contains(word) {
                let word_probability = match self.words.get(word) {
                    Some( &(_, p) ) => p,
                    None => self.default_word_probability,
                };
                total += word_probability.ln();
            }
        }
        self.probability.ln() + total
    }
}

fn split_document(document: &String) -> Vec<String> {
    let re = Regex::new(r"(\s)").unwrap();
    re.split(document).map(|s| s.to_string()).collect()
}<commit_msg>cleaned up code<commit_after>use std::collections::{HashMap, HashSet};
use std::f64;
use regex::Regex;
use rustc_serialize::json;

static DEFAULT_SMOOTHING: f64 = 1.0f64;

#[derive(Debug, Clone, RustcDecodable, RustcEncodable)]
/// Naive Bayes classifier
pub struct Classifier {
    vocab: HashSet<String>,
    num_examples: u32,
    smoothing: f64,
    classifications: HashMap<String, Classification>
}

#[derive(Debug, Clone, RustcDecodable, RustcEncodable)]
struct Classification {
    label: String,
    num_examples: u32,
    num_words: u32,
    probability: f64,
    default_word_probability: f64,
    words: HashMap<String, (u32, f64)>,
}

impl Classifier {

    /// Creates a new classifier
    pub fn new() -> Classifier {
        Classifier {
            vocab: HashSet::new(),
            num_examples: 0u32,
            smoothing: DEFAULT_SMOOTHING,
            classifications: HashMap::new(),
        }
    }

    /// Takes a document that has been tokenized into a vector of strings
    /// and a label and adds the document to the list of documents that the
    /// classifier is aware of and will train on next time the `train()` method is called
    pub fn add_document_tokenized(&mut self, document: &Vec<String>, label: &String) {
        if document.len() == 0 { return; }

        // make sure the classification already exists
        if !self.classifications.contains_key(label) {
            let c = Classification::new(label);
            self.classifications.insert(label.clone(), c);
        }

        let mut classification = self.classifications.get_mut(label).unwrap();

        for word in document.iter() {
            classification.add_word(word);
            self.vocab.insert(word.to_string());
        }

        self.num_examples += 1;
        classification.num_examples += 1;
    }

    /// Takes a document and a label and tokenizes the document by
    /// breaking on whitespace characters. The document is added to the list
    /// of documents that the classifier is aware of and will train on next time the `train()` method is called
    pub fn add_document(&mut self, document: &String, label: &String) {
        self.add_document_tokenized(&split_document(document), label);
    }

    /// Adds a list of (document, label) tuples to the classifier
    pub fn add_documents(&mut self, examples: &Vec<(String, String)>) {
        for &(ref document, ref label) in examples.iter() {
            self.add_document(document, label);
        }
    }

    /// Adds a list of (tokenized document, label) tuples to the classifier
    pub fn add_documents_tokenized(&mut self, examples: &Vec<(Vec<String>, String)>) {
        for &(ref document, ref label) in examples.iter() {
            self.add_document_tokenized(document, label);
        }
    }

    /// Gets a vector of all of the labels that the classifier has seen so far
    pub fn get_labels(&self) -> Vec<String> {
        let labels: Vec<String> =
            self.classifications.values().map(|c| c.label.clone()).collect();
        labels
    }

    /// Sets the [smoothing](http://en.wikipedia.org/wiki/Additive_smoothing) value (must be greater than 0.0)
    pub fn set_smoothing(&mut self, smoothing: f64) {
        if smoothing <= 0.0 {
            panic!("smoothing value must be a positive number");
        }
        self.smoothing = smoothing;
    }

    /// Trains the classifier on the documents that have been observed so far
    pub fn train(&mut self) {
        for (_, classification) in self.classifications.iter_mut() {
            classification.train(&self.vocab, self.num_examples, self.smoothing);
        }
    }

    /// Takes an unlabeled document that has been tokenized into a vector of strings
    /// and then computes a classifying label for the document
    pub fn classify_tokenized(&self, document: &Vec<String>) -> String {
        let mut max_score = f64::NEG_INFINITY;
        let mut max_classification = None;

        for classification in self.classifications.values() {
            let score = classification.score_document(document, &self.vocab);
            if score > max_score {
                max_classification = Some(classification);
                max_score = score;
            }
        }

        max_classification.expect("no classification found").label.clone()
    }

    /// Takes an unlabeled document and tokenizes it by breaking on spaces and
    /// then computes a classifying label for the document
    pub fn classify(&self, document: &String) -> String {
        self.classify_tokenized(&split_document(document))
    }

    /// Similar to classify but instead of returning a single label, returns all
    /// labels and the probabilities of each one given the document
    pub fn get_document_probabilities_tokenized(&self, document: &Vec<String>) -> Vec<(String, f64)> {
        self.classifications.values().map(|classification| {
            let score = classification.score_document(document, &self.vocab);
            (classification.label.clone(), score)
        }).collect()
    }

    /// Similar to classify but instead of returning a single label, returns all
    /// labels and the probabilities of each one given the document
    pub fn get_document_probabilities(&self, document: &String) -> Vec<(String, f64)> {
        self.get_document_probabilities_tokenized(&split_document(document))
    }

    /// Encodes the classifier as a JSON string.
    pub fn to_json(&self) -> String {
        json::encode(self).ok().expect("encoding JSON failed")
    }

    /// Builds a new classifier from a JSON string
    pub fn from_json(encoded: &str) -> Classifier {
        let classifier: Classifier = json::decode(encoded).ok().expect("decoding JSON failed");
        classifier
    }

}

impl Classification {

    fn new(label: &String) -> Classification {
        Classification {
            label: label.clone(),
            num_examples: 0u32,
            num_words: 0u32,
            probability: 0.0f64,
            default_word_probability: 0.0f64,
            words: HashMap::new(),
        }
    }

    fn add_word(&mut self, word: &String) {
        self.num_words += 1;
        if self.words.contains_key(word) {
            self.words.get_mut(word).unwrap().0 += 1;
        } else {
            self.words.insert(word.clone(), (1, 0.0f64));
        }
    }

    // trains this classification
    fn train(&mut self, vocab: &HashSet<String>, total_examples: u32, smoothing: f64) {
        // the probability of this classification
        self.probability = self.num_examples as f64 / total_examples as f64;
        // the probability of any word that has not been seen in a document
        // labeled with this classification's label
        self.default_word_probability = smoothing / (self.num_words as f64 + smoothing * vocab.len() as f64);

        for word in vocab.iter() {
            if self.words.contains_key(word) {
                let word_entry = self.words.get_mut(word).unwrap();
                let word_count = word_entry.0;
                let p_word_given_label = (word_count as f64 + smoothing) / (self.num_words as f64 + smoothing * vocab.len() as f64);
                word_entry.1 = p_word_given_label;
            }
        }
    }

    // retrieves the probability of the document given this classification
    // times the probability of this classification
    fn score_document(&self, document: &Vec<String>, vocab: &HashSet<String>) -> f64 {
        let mut total = 0.0f64;
        for word in document.iter() {
            if vocab.contains(word) {
                let word_probability = match self.words.get(word) {
                    Some( &(_, p) ) => p,
                    None => self.default_word_probability,
                };
                total += word_probability.ln();
            }
        }
        self.probability.ln() + total
    }
}

// splits a String on whitespaces
fn split_document(document: &String) -> Vec<String> {
    let re = Regex::new(r"(\s)").unwrap();
    re.split(document).map(|s| s.to_string()).collect()
}<|endoftext|>


<commit_before>extern crate conv;
extern crate rand;

use self::conv::prelude::*;
use neat::gene::Gene as Gene;
use neat::mutation::Mutation as Mutation;

#[derive(Debug, Clone)]
pub struct Genome{
    genes: Vec<Gene>,
    last_neuron_id: usize
}

const COMPATIBILITY_THRESHOLD: f64 = 1f64;
//const MUTATE_ADD_NEURON_PROBABILITY: f64 = 0.33f64;
const MUTATE_CONNECTION_WEIGHT: f64 = 0.33f64;
const MUTATE_ADD_CONNECTION: f64 = 0.33f64;

impl Genome{

    pub fn new() -> Genome {
        Genome {
            genes: vec![],
            last_neuron_id: 0 //we need 1 neuron to start to mutate
        }
    }

    pub fn mutate(&mut self) {
        let random = rand::random::<f64>();
        if random > MUTATE_ADD_CONNECTION + MUTATE_CONNECTION_WEIGHT && self.genes.len() > 0 {
            self.mutate_add_neuron();
        } else if random > MUTATE_ADD_CONNECTION && self.genes.len() > 0 {
            self.mutate_connection_weight();
        } else {
            self.mutate_add_connection();
        }
    }

    pub fn mate(&self, other: &Genome) -> Genome{
        if self.genes.len() > other.genes.len() {
            self.mate_genes(other)
        }else{
            other.mate_genes(self)
        }
    }

    fn mate_genes(&self, other: &Genome) -> Genome{
        let mut genome = Genome::new();
        for gene in &self.genes {
            genome.add_gene({
                if rand::random::<f64>() > 0.5f64 {
                    match other.genes.binary_search(&gene) {
                        Ok(position) => other.genes[position].clone(),
                        Err(_) => gene.clone()
                    }
                } else {
                    gene.clone()
                }
            });
        }
        genome
    }

    pub fn get_genes<'a>(&'a self)-> &'a Vec<Gene>{
        &self.genes
    }

    // only allow connected nodes
    pub fn inject_gene(&mut self, in_neuron_id: usize, out_neuron_id: usize, weight: f64) {
        let max_neuron_id = self.last_neuron_id + 1;

        if in_neuron_id == out_neuron_id && in_neuron_id > max_neuron_id {
            panic!("Try to create a gene neuron unconnected, max neuron id {}, {} -> {}", max_neuron_id, in_neuron_id, out_neuron_id);
        }

        assert!(in_neuron_id <= max_neuron_id, format!("in_neuron_id {} is greater than max allowed id {}", in_neuron_id, max_neuron_id));
        assert!(out_neuron_id <= max_neuron_id, format!("out_neuron_id {} is greater than max allowed id {}", out_neuron_id, max_neuron_id));

        if in_neuron_id > self.last_neuron_id {
            self.last_neuron_id = in_neuron_id;
        }
        if out_neuron_id > self.last_neuron_id {
            self.last_neuron_id = out_neuron_id;
        }
        self.create_gene(in_neuron_id, out_neuron_id, weight)
    }

    pub fn len(&self) -> usize{
        self.last_neuron_id + 1 //first neuron id is 0
    }

    fn create_gene(&mut self, in_neuron_id: usize, out_neuron_id: usize, weight: f64) {
        let gene = Gene {
            in_neuron_id: in_neuron_id,
            out_neuron_id: out_neuron_id,
            weight: weight,
            ..Default::default()
        };
        self.add_gene(gene);
    }

    fn mutate_add_connection(&mut self) {
        let mut rng = rand::thread_rng();
        let neuron_ids_to_connect = {
            if self.last_neuron_id == 0 {
                vec![0, 0]
            } else {
                rand::sample(&mut rng, 0..self.last_neuron_id + 1, 2)
            }
        };
        self.add_connection(neuron_ids_to_connect[0], neuron_ids_to_connect[1]);
    }

    fn mutate_connection_weight(&mut self) {
        let mut rng = rand::thread_rng();
        let selected_gene = rand::sample(&mut rng, 0..self.genes.len(), 1)[0];
        let gene = &mut self.genes[selected_gene];
        Mutation::connection_weight(gene);
    }

    fn mutate_add_neuron(&mut self) {
        let (gene1, gene2) = {
            let mut rng = rand::thread_rng();
            let selected_gene = rand::sample(&mut rng, 0..self.genes.len(), 1)[0];
            let gene = &mut self.genes[selected_gene];
            self.last_neuron_id += 1;
            Mutation::add_neuron(gene, self.last_neuron_id)
        };
        self.add_gene(gene1);
        self.add_gene(gene2);
    }

    fn add_connection(&mut self, in_neuron_id: usize, out_neuron_id: usize) {
        let gene = Mutation::add_connection(in_neuron_id, out_neuron_id);
        self.add_gene(gene);
    }

    fn add_gene(&mut self, gene: Gene){
        if !self.genes.contains(&gene) {
            self.genes.push(gene);
        }
    }

    pub fn is_same_specie(&self, other: &Genome) -> bool{
        self.compatibility_distance(other) < COMPATIBILITY_THRESHOLD
    }

    pub fn total_weights(&self) -> f64{
        let mut total = 0f64;
        for gene in &self.genes {
            total += gene.weight;
        }
        total
    }

    pub fn total_genes(&self) -> usize{
        self.genes.len()
    }

    //http://nn.cs.utexas.edu/downloads/papers/stanley.ec02.pdf - Pag. 110
    //I have consider disjoint and excess genes as the same
    fn compatibility_distance(&self, other: &Genome) -> f64 {
        //TODO: optimize this method
        let c2 = 0.5f64;
        let c3 = 0.5f64;

        //Number of excess
        let n1 = self.genes.len().value_as::<f64>().unwrap();
        let n2 = other.genes.len().value_as::<f64>().unwrap();
        let n = n1.max(n2);

        if n == 0f64 {
            return 0f64; //no genes in any genome, the genomes are equal
        }

        let matching_genes  = self.genes.iter().filter(|i1_gene| other.genes.contains(i1_gene)).collect::<Vec<&Gene>>();

        let n3 = matching_genes.len().value_as::<f64>().unwrap();

        //Disjoint genes
        let d = n1 + n2 - (2f64 * n3);

        //average weight differences of matching genes
        let w1 = matching_genes.iter().fold(0f64, |acc, &m_gene| acc + (m_gene.weight + other.genes.get(other.genes.binary_search(m_gene).unwrap()).unwrap().weight)).abs();

        let w = w1 / n3;

        //compatibility distance
        let delta = (c2 * d / n) + c3 * w;
        delta
    }
}

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

    #[test]
    fn mutation_connection_weight(){
        let mut genome = Genome::new();
        genome.create_gene(1, 1, 1f64);
        let orig_gene = genome.genes[0].clone();
        genome.mutate_connection_weight();

        assert!(genome.genes[0].weight != orig_gene.weight);
    }

    #[test]
    fn mutation_add_connection(){
        let mut genome = Genome::new();
        genome.add_connection(1, 2);

        assert!(genome.genes[0].in_neuron_id == 1);
        assert!(genome.genes[0].out_neuron_id == 2);
    }

    #[test]
    fn mutation_add_neuron(){
        let mut genome = Genome::new();
        genome.mutate_add_connection();
        genome.mutate_add_neuron();
        assert!(!genome.genes[0].enabled);
        assert!(genome.genes[1].in_neuron_id == genome.genes[0].in_neuron_id);
        assert!(genome.genes[1].out_neuron_id == 1);
        assert!(genome.genes[2].in_neuron_id == 1);
        assert!(genome.genes[2].out_neuron_id == genome.genes[0].out_neuron_id);
    }

    #[test]
    #[should_panic(expected = "Try to create a gene neuron unconnected, max neuron id 1, 2 -> 2")]
    fn try_to_inject_a_unconnected_neuron_gene_should_panic(){
        let mut genome1 = Genome::new();
        genome1.inject_gene(2, 2, 0.5f64);
    }

    #[test]
    fn two_genomes_without_differences_should_be_in_same_specie(){
        let mut genome1 = Genome::new();
        genome1.create_gene(1, 1, 1f64);
        genome1.create_gene(1, 2, 1f64);
        let mut genome2 = Genome::new();
        genome2.create_gene(1, 1, 0f64);
        genome2.create_gene(1, 2, 0f64);
        genome2.create_gene(1, 3, 0f64);
        assert!(genome1.is_same_specie(&genome2));
    }

    #[test]
    fn two_genomes_with_enought_difference_should_be_in_different_species(){
        let mut genome1 = Genome::new();
        genome1.create_gene(1, 1, 1f64);
        genome1.create_gene(1, 2, 1f64);
        let mut genome2 = Genome::new();
        genome2.create_gene(1, 3, 1f64);
        genome2.create_gene(1, 4, 1f64);
        assert!(!genome1.is_same_specie(&genome2));
    }

    #[test]
    fn already_existing_gene_should_be_not_duplicated(){
        let mut genome1 = Genome::new();
        genome1.create_gene(0, 0, 1f64);
        genome1.add_connection(0, 0);
        assert!(genome1.genes.len() == 1);
        assert!(genome1.get_genes()[0].weight == 1f64);
    }
}
<commit_msg>Fix genome length calculation<commit_after>extern crate conv;
extern crate rand;

use self::conv::prelude::*;
use neat::gene::Gene as Gene;
use neat::mutation::Mutation as Mutation;

#[derive(Debug, Clone)]
pub struct Genome{
    genes: Vec<Gene>,
    last_neuron_id: usize
}

const COMPATIBILITY_THRESHOLD: f64 = 1f64;
const MUTATE_CONNECTION_WEIGHT: f64 = 0.33f64;
const MUTATE_ADD_CONNECTION: f64 = 0.33f64;

impl Genome{

    pub fn new() -> Genome {
        Genome {
            genes: vec![],
            last_neuron_id: 0 //we need 1 neuron to start to mutate
        }
    }

    pub fn mutate(&mut self) {
        let random = rand::random::<f64>();
        if random > MUTATE_ADD_CONNECTION + MUTATE_CONNECTION_WEIGHT && self.genes.len() > 0 {
            self.mutate_add_neuron();
        } else if random > MUTATE_ADD_CONNECTION && self.genes.len() > 0 {
            self.mutate_connection_weight();
        } else {
            self.mutate_add_connection();
        }
    }

    pub fn mate(&self, other: &Genome) -> Genome{
        if self.genes.len() > other.genes.len() {
            self.mate_genes(other)
        }else{
            other.mate_genes(self)
        }
    }

    fn mate_genes(&self, other: &Genome) -> Genome{
        let mut genome = Genome::new();
        for gene in &self.genes {
            genome.add_gene({
                if rand::random::<f64>() > 0.5f64 {
                    match other.genes.binary_search(&gene) {
                        Ok(position) => other.genes[position].clone(),
                        Err(_) => gene.clone()
                    }
                } else {
                    gene.clone()
                }
            });
        }
        genome
    }

    pub fn get_genes<'a>(&'a self)-> &'a Vec<Gene>{
        &self.genes
    }

    // only allow connected nodes
    pub fn inject_gene(&mut self, in_neuron_id: usize, out_neuron_id: usize, weight: f64) {
        let max_neuron_id = self.last_neuron_id + 1;

        if in_neuron_id == out_neuron_id && in_neuron_id > max_neuron_id {
            panic!("Try to create a gene neuron unconnected, max neuron id {}, {} -> {}", max_neuron_id, in_neuron_id, out_neuron_id);
        }

        assert!(in_neuron_id <= max_neuron_id, format!("in_neuron_id {} is greater than max allowed id {}", in_neuron_id, max_neuron_id));
        assert!(out_neuron_id <= max_neuron_id, format!("out_neuron_id {} is greater than max allowed id {}", out_neuron_id, max_neuron_id));

        self.create_gene(in_neuron_id, out_neuron_id, weight)
    }

    pub fn len(&self) -> usize{
        self.last_neuron_id + 1 //first neuron id is 0
    }

    fn create_gene(&mut self, in_neuron_id: usize, out_neuron_id: usize, weight: f64) {
        let gene = Gene {
            in_neuron_id: in_neuron_id,
            out_neuron_id: out_neuron_id,
            weight: weight,
            ..Default::default()
        };
        self.add_gene(gene);
    }

    fn mutate_add_connection(&mut self) {
        let mut rng = rand::thread_rng();
        let neuron_ids_to_connect = {
            if self.last_neuron_id == 0 {
                vec![0, 0]
            } else {
                rand::sample(&mut rng, 0..self.last_neuron_id + 1, 2)
            }
        };
        self.add_connection(neuron_ids_to_connect[0], neuron_ids_to_connect[1]);
    }

    fn mutate_connection_weight(&mut self) {
        let mut rng = rand::thread_rng();
        let selected_gene = rand::sample(&mut rng, 0..self.genes.len(), 1)[0];
        let gene = &mut self.genes[selected_gene];
        Mutation::connection_weight(gene);
    }

    fn mutate_add_neuron(&mut self) {
        let (gene1, gene2) = {
            let mut rng = rand::thread_rng();
            let selected_gene = rand::sample(&mut rng, 0..self.genes.len(), 1)[0];
            let gene = &mut self.genes[selected_gene];
            self.last_neuron_id += 1;
            Mutation::add_neuron(gene, self.last_neuron_id)
        };
        self.add_gene(gene1);
        self.add_gene(gene2);
    }

    fn add_connection(&mut self, in_neuron_id: usize, out_neuron_id: usize) {
        let gene = Mutation::add_connection(in_neuron_id, out_neuron_id);
        self.add_gene(gene);
    }

    fn add_gene(&mut self, gene: Gene){
        if gene.in_neuron_id > self.last_neuron_id {
            self.last_neuron_id = gene.in_neuron_id;
        }
        if gene.out_neuron_id > self.last_neuron_id {
            self.last_neuron_id = gene.out_neuron_id;
        }
        match self.genes.binary_search(&gene) {
            Ok(pos) => self.genes[pos].enabled = true,
            Err(_) => self.genes.push(gene)
        }
    }

    pub fn is_same_specie(&self, other: &Genome) -> bool{
        self.compatibility_distance(other) < COMPATIBILITY_THRESHOLD
    }

    pub fn total_weights(&self) -> f64{
        let mut total = 0f64;
        for gene in &self.genes {
            total += gene.weight;
        }
        total
    }

    pub fn total_genes(&self) -> usize{
        self.genes.len()
    }

    //http://nn.cs.utexas.edu/downloads/papers/stanley.ec02.pdf - Pag. 110
    //I have consider disjoint and excess genes as the same
    fn compatibility_distance(&self, other: &Genome) -> f64 {
        //TODO: optimize this method
        let c2 = 0.5f64;
        let c3 = 0.5f64;

        //Number of excess
        let n1 = self.genes.len().value_as::<f64>().unwrap();
        let n2 = other.genes.len().value_as::<f64>().unwrap();
        let n = n1.max(n2);

        if n == 0f64 {
            return 0f64; //no genes in any genome, the genomes are equal
        }

        let matching_genes  = self.genes.iter().filter(|i1_gene| other.genes.contains(i1_gene)).collect::<Vec<&Gene>>();

        let n3 = matching_genes.len().value_as::<f64>().unwrap();

        //Disjoint genes
        let d = n1 + n2 - (2f64 * n3);

        //average weight differences of matching genes
        let w1 = matching_genes.iter().fold(0f64, |acc, &m_gene| acc + (m_gene.weight + other.genes.get(other.genes.binary_search(m_gene).unwrap()).unwrap().weight)).abs();

        let w = w1 / n3;

        //compatibility distance
        let delta = (c2 * d / n) + c3 * w;
        delta
    }
}

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

    #[test]
    fn mutation_connection_weight(){
        let mut genome = Genome::new();
        genome.inject_gene(0, 0, 1f64);
        let orig_gene = genome.genes[0].clone();
        genome.mutate_connection_weight();

        assert!(genome.genes[0].weight != orig_gene.weight);
    }

    #[test]
    fn mutation_add_connection(){
        let mut genome = Genome::new();
        genome.add_connection(1, 2);

        assert!(genome.genes[0].in_neuron_id == 1);
        assert!(genome.genes[0].out_neuron_id == 2);
    }

    #[test]
    fn mutation_add_neuron(){
        let mut genome = Genome::new();
        genome.mutate_add_connection();
        genome.mutate_add_neuron();
        assert!(!genome.genes[0].enabled);
        assert!(genome.genes[1].in_neuron_id == genome.genes[0].in_neuron_id);
        assert!(genome.genes[1].out_neuron_id == 1);
        assert!(genome.genes[2].in_neuron_id == 1);
        assert!(genome.genes[2].out_neuron_id == genome.genes[0].out_neuron_id);
    }

    #[test]
    #[should_panic(expected = "Try to create a gene neuron unconnected, max neuron id 1, 2 -> 2")]
    fn try_to_inject_a_unconnected_neuron_gene_should_panic(){
        let mut genome1 = Genome::new();
        genome1.inject_gene(2, 2, 0.5f64);
    }

    #[test]
    fn two_genomes_without_differences_should_be_in_same_specie(){
        let mut genome1 = Genome::new();
        genome1.inject_gene(0, 0, 1f64);
        genome1.inject_gene(0, 1, 1f64);
        let mut genome2 = Genome::new();
        genome2.inject_gene(0, 0, 0f64);
        genome2.inject_gene(0, 1, 0f64);
        genome2.inject_gene(0, 2, 0f64);
        assert!(genome1.is_same_specie(&genome2));
    }

    #[test]
    fn two_genomes_with_enought_difference_should_be_in_different_species(){
        let mut genome1 = Genome::new();
        genome1.inject_gene(0, 0, 1f64);
        genome1.inject_gene(0, 1, 1f64);
        let mut genome2 = Genome::new();
        genome2.inject_gene(0, 0, 1f64);
        genome2.inject_gene(0, 1, 1f64);
        genome2.inject_gene(0, 2, 1f64);
        genome2.inject_gene(0, 3, 1f64);
        assert!(!genome1.is_same_specie(&genome2));
    }

    #[test]
    fn already_existing_gene_should_be_not_duplicated(){
        let mut genome1 = Genome::new();
        genome1.inject_gene(0, 0, 1f64);
        genome1.add_connection(0, 0);
        assert!(genome1.genes.len() == 1);
        assert!(genome1.get_genes()[0].weight == 1f64);
    }

    #[test]
    fn adding_an_existing_gene_disabled_should_enable_original(){
        let mut genome1 = Genome::new();
        genome1.inject_gene(0, 1, 0f64);
        genome1.mutate_add_neuron();
        assert!(!genome1.genes[0].enabled);
        assert!(genome1.genes.len() == 3);
        genome1.add_connection(0, 1);
        assert!(genome1.genes[0].enabled);
        assert!(genome1.genes[0].weight == 0f64);
        assert!(genome1.genes.len() == 3);
    }
}
<|endoftext|>


<commit_before>use rustc::lint::*;
use rustc_front::hir::*;
use utils::span_lint;
use rustc::middle::ty::{TypeAndMut, TypeVariants, MethodCall};

declare_lint! {
    pub UNNECESSARY_MUT_PASSED,
    Warn,
    "an argument is passed as a mutable reference although the function/method only demands an \
     immutable reference"
}

#[derive(Copy,Clone)]
pub struct UnnecessaryMutPassed;

impl LintPass for UnnecessaryMutPassed {
    fn get_lints(&self) -> LintArray {
        lint_array!(UNNECESSARY_MUT_PASSED)
    }
}

impl LateLintPass for UnnecessaryMutPassed {
    fn check_expr(&mut self, cx: &LateContext, e: &Expr) {
        match e.node {
            ExprCall(ref fn_expr, ref arguments) => {
                let borrowed_table = cx.tcx.tables.borrow();
                let funtion_type = match borrowed_table.node_types.get(&fn_expr.id) {
                    Some(funtion_type) => funtion_type,
                    None => unreachable!(), // A function with unknown type is called.
                                            // If this happened the compiler would have aborted the
                                            // compilation long ago.
                };
                if let TypeVariants::TyBareFn(_, ref b) = funtion_type.sty {
                    let parameters = b.sig.skip_binder().inputs.clone();
                    for (argument, parameter) in arguments.iter().zip(parameters.iter()) {
                        match parameter.sty {
                            TypeVariants::TyRef(_, TypeAndMut {ty: _, mutbl: MutImmutable}) |
                            TypeVariants::TyRawPtr(TypeAndMut {ty: _, mutbl: MutImmutable}) => {
                                if let Expr_::ExprAddrOf(MutMutable, _) = argument.node {
                                    if let ExprPath(_, path) = fn_expr.node.clone() {
                                        span_lint(cx, UNNECESSARY_MUT_PASSED,
                                                  argument.span, &format!("This argument of the \
                                                  function \"{}\" doesn't need to be mutable", path));
                                    }
                                }
                            },
                            _ => {}
                        }
                    }
                }
            },
            ExprMethodCall(ref name, _, ref arguments) => {
                let method_call = MethodCall::expr(e.id);
                let borrowed_table = cx.tcx.tables.borrow();
                let method_type = match borrowed_table.method_map.get(&method_call) {
                    Some(method_type) => method_type,
                    None => unreachable!(), // Just like above, this should never happen.
                };
                if let TypeVariants::TyBareFn(_, ref b) = method_type.ty.sty {
                    let parameters = b.sig.skip_binder().inputs.iter().clone();
                    for (argument, parameter) in arguments.iter().zip(parameters).skip(1) {
                        // Skip the first argument and the first parameter because it is the
                        // struct the function is called on.
                        match parameter.sty {
                            TypeVariants::TyRef(_, TypeAndMut {ty: _, mutbl: MutImmutable}) |
                            TypeVariants::TyRawPtr(TypeAndMut {ty: _, mutbl: MutImmutable}) => {
                                if let Expr_::ExprAddrOf(MutMutable, _) = argument.node {
                                    span_lint(cx, UNNECESSARY_MUT_PASSED,
                                              argument.span, &format!("This argument of the \
                                              method \"{}\" doesn't need to be mutable",
                                              name.node.as_str()));
                                }
                            },
                            _ => {}
                        }
                    }
                }
            },
            _ => {}
        }
    }
}
<commit_msg>Remove unnecessary clones and add helper function<commit_after>use rustc::lint::*;
use rustc_front::hir::*;
use utils::span_lint;
use rustc::middle::ty::{TypeAndMut, TypeVariants, MethodCall, TyS};
use syntax::ptr::P;

declare_lint! {
    pub UNNECESSARY_MUT_PASSED,
    Warn,
    "an argument is passed as a mutable reference although the function/method only demands an \
     immutable reference"
}

#[derive(Copy,Clone)]
pub struct UnnecessaryMutPassed;

impl LintPass for UnnecessaryMutPassed {
    fn get_lints(&self) -> LintArray {
        lint_array!(UNNECESSARY_MUT_PASSED)
    }
}

impl LateLintPass for UnnecessaryMutPassed {
    fn check_expr(&mut self, cx: &LateContext, e: &Expr) {
        let borrowed_table = cx.tcx.tables.borrow();
        match e.node {
            ExprCall(ref fn_expr, ref arguments) => {
                match borrowed_table.node_types.get(&fn_expr.id) {
                    Some(function_type) => {
                        if let ExprPath(_, ref path) = fn_expr.node {
                            check_arguments(cx, &arguments, function_type,
                                            &format!("{}", path));
                        }
                    },
                    None => unreachable!(), // A function with unknown type is called.
                                            // If this happened the compiler would have aborted the
                                            // compilation long ago.
                };

            },
            ExprMethodCall(ref name, _, ref arguments) => {
                let method_call = MethodCall::expr(e.id);
                match borrowed_table.method_map.get(&method_call) {
                    Some(method_type) => check_arguments(cx, &arguments, method_type.ty,
                                                         &format!("{}", name.node.as_str())),
                    None => unreachable!(), // Just like above, this should never happen.
                };
            },
            _ => {}
        }
    }
}

fn check_arguments(cx: &LateContext, arguments: &[P<Expr>], type_definition: &TyS, name: &str) {
    if let TypeVariants::TyBareFn(_, ref fn_type) = type_definition.sty {
        let parameters = &fn_type.sig.skip_binder().inputs;
        for (argument, parameter) in arguments.iter().zip(parameters.iter()) {
            match parameter.sty {
                TypeVariants::TyRef(_, TypeAndMut {ty: _, mutbl: MutImmutable}) |
                TypeVariants::TyRawPtr(TypeAndMut {ty: _, mutbl: MutImmutable}) => {
                    if let Expr_::ExprAddrOf(MutMutable, _) = argument.node {
                        span_lint(cx, UNNECESSARY_MUT_PASSED,
                                  argument.span, &format!("The function/method \"{}\" \
                                  doesn't need a mutable reference",
                                  name));
                    }
                },
                _ => {}
            }
        }
    }
}
<|endoftext|>


<commit_before>/// Quasi-quotation macro that accepts input like the [`quote!`] macro but uses
/// type inference to figure out a return type for those tokens.
///
/// [`quote!`]: https://docs.rs/quote/0.6/quote/index.html
///
/// The return type can be any syntax tree node that implements the [`Parse`]
/// trait.
///
/// [`Parse`]: parse::Parse
///
/// ```edition2018
/// use quote::quote;
/// use syn::{parse_quote, Stmt};
///
/// fn main() {
///     let name = quote!(v);
///     let ty = quote!(u8);
///
///     let stmt: Stmt = parse_quote! {
///         let #name: #ty = Default::default();
///     };
///
///     println!("{:#?}", stmt);
/// }
/// ```
///
/// *This macro is available if Syn is built with the `"parsing"` feature,
/// although interpolation of syntax tree nodes into the quoted tokens is only
/// supported if Syn is built with the `"printing"` feature as well.*
///
/// # Example
///
/// The following helper function adds a bound `T: HeapSize` to every type
/// parameter `T` in the input generics.
///
/// ```edition2018
/// use syn::{parse_quote, Generics, GenericParam};
///
/// // Add a bound `T: HeapSize` to every type parameter T.
/// fn add_trait_bounds(mut generics: Generics) -> Generics {
///     for param in &mut generics.params {
///         if let GenericParam::Type(type_param) = param {
///             type_param.bounds.push(parse_quote!(HeapSize));
///         }
///     }
///     generics
/// }
/// ```
///
/// # Special cases
///
/// This macro can parse the following additional types as a special case even
/// though they do not implement the `Parse` trait.
///
/// - [`Attribute`] — parses one attribute, allowing either outer like `#[...]`
///   or inner like `#![...]`
/// - [`Punctuated<T, P>`] — parses zero or more `T` separated by punctuation
///   `P` with optional trailing punctuation
///
/// [`Punctuated<T, P>`]: punctuated::Punctuated
///
/// # Panics
///
/// Panics if the tokens fail to parse as the expected syntax tree type. The
/// caller is responsible for ensuring that the input tokens are syntactically
/// valid.
//
// TODO: allow Punctuated to be inferred as intra doc link, currently blocked on
// https://github.com/rust-lang/rust/issues/62834
#[macro_export(local_inner_macros)]
macro_rules! parse_quote {
    ($($tt:tt)*) => {
        $crate::parse_quote::parse($crate::export::From::from(quote_impl!($($tt)*)))
    };
}

#[doc(hidden)]
#[macro_export]
macro_rules! quote_impl {
    ($($tt:tt)*) => {
        $crate::export::quote::quote!($($tt)*)
    };
}

////////////////////////////////////////////////////////////////////////////////
// Can parse any type that implements Parse.

use crate::parse::{Parse, ParseStream, Parser, Result};
use proc_macro2::TokenStream;

// Not public API.
#[doc(hidden)]
pub fn parse<T: ParseQuote>(token_stream: TokenStream) -> T {
    let parser = T::parse;
    match parser.parse2(token_stream) {
        Ok(t) => t,
        Err(err) => panic!("{}", err),
    }
}

// Not public API.
#[doc(hidden)]
pub trait ParseQuote: Sized {
    fn parse(input: ParseStream) -> Result<Self>;
}

impl<T: Parse> ParseQuote for T {
    fn parse(input: ParseStream) -> Result<Self> {
        <T as Parse>::parse(input)
    }
}

////////////////////////////////////////////////////////////////////////////////
// Any other types that we want `parse_quote!` to be able to parse.

use crate::punctuated::Punctuated;
#[cfg(any(feature = "full", feature = "derive"))]
use crate::{attr, Attribute};

#[cfg(any(feature = "full", feature = "derive"))]
impl ParseQuote for Attribute {
    fn parse(input: ParseStream) -> Result<Self> {
        if input.peek(Token![#]) && input.peek2(Token![!]) {
            attr::parsing::single_parse_inner(input)
        } else {
            attr::parsing::single_parse_outer(input)
        }
    }
}

impl<T: Parse, P: Parse> ParseQuote for Punctuated<T, P> {
    fn parse(input: ParseStream) -> Result<Self> {
        Self::parse_terminated(input)
    }
}
<commit_msg>Collapse quote_impl macro<commit_after>/// Quasi-quotation macro that accepts input like the [`quote!`] macro but uses
/// type inference to figure out a return type for those tokens.
///
/// [`quote!`]: https://docs.rs/quote/0.6/quote/index.html
///
/// The return type can be any syntax tree node that implements the [`Parse`]
/// trait.
///
/// [`Parse`]: parse::Parse
///
/// ```edition2018
/// use quote::quote;
/// use syn::{parse_quote, Stmt};
///
/// fn main() {
///     let name = quote!(v);
///     let ty = quote!(u8);
///
///     let stmt: Stmt = parse_quote! {
///         let #name: #ty = Default::default();
///     };
///
///     println!("{:#?}", stmt);
/// }
/// ```
///
/// *This macro is available if Syn is built with the `"parsing"` feature,
/// although interpolation of syntax tree nodes into the quoted tokens is only
/// supported if Syn is built with the `"printing"` feature as well.*
///
/// # Example
///
/// The following helper function adds a bound `T: HeapSize` to every type
/// parameter `T` in the input generics.
///
/// ```edition2018
/// use syn::{parse_quote, Generics, GenericParam};
///
/// // Add a bound `T: HeapSize` to every type parameter T.
/// fn add_trait_bounds(mut generics: Generics) -> Generics {
///     for param in &mut generics.params {
///         if let GenericParam::Type(type_param) = param {
///             type_param.bounds.push(parse_quote!(HeapSize));
///         }
///     }
///     generics
/// }
/// ```
///
/// # Special cases
///
/// This macro can parse the following additional types as a special case even
/// though they do not implement the `Parse` trait.
///
/// - [`Attribute`] — parses one attribute, allowing either outer like `#[...]`
///   or inner like `#![...]`
/// - [`Punctuated<T, P>`] — parses zero or more `T` separated by punctuation
///   `P` with optional trailing punctuation
///
/// [`Punctuated<T, P>`]: punctuated::Punctuated
///
/// # Panics
///
/// Panics if the tokens fail to parse as the expected syntax tree type. The
/// caller is responsible for ensuring that the input tokens are syntactically
/// valid.
//
// TODO: allow Punctuated to be inferred as intra doc link, currently blocked on
// https://github.com/rust-lang/rust/issues/62834
#[macro_export(local_inner_macros)]
macro_rules! parse_quote {
    ($($tt:tt)*) => {
        $crate::parse_quote::parse(
            $crate::export::From::from(
                $crate::export::quote::quote!($($tt)*)
            )
        )
    };
}

////////////////////////////////////////////////////////////////////////////////
// Can parse any type that implements Parse.

use crate::parse::{Parse, ParseStream, Parser, Result};
use proc_macro2::TokenStream;

// Not public API.
#[doc(hidden)]
pub fn parse<T: ParseQuote>(token_stream: TokenStream) -> T {
    let parser = T::parse;
    match parser.parse2(token_stream) {
        Ok(t) => t,
        Err(err) => panic!("{}", err),
    }
}

// Not public API.
#[doc(hidden)]
pub trait ParseQuote: Sized {
    fn parse(input: ParseStream) -> Result<Self>;
}

impl<T: Parse> ParseQuote for T {
    fn parse(input: ParseStream) -> Result<Self> {
        <T as Parse>::parse(input)
    }
}

////////////////////////////////////////////////////////////////////////////////
// Any other types that we want `parse_quote!` to be able to parse.

use crate::punctuated::Punctuated;
#[cfg(any(feature = "full", feature = "derive"))]
use crate::{attr, Attribute};

#[cfg(any(feature = "full", feature = "derive"))]
impl ParseQuote for Attribute {
    fn parse(input: ParseStream) -> Result<Self> {
        if input.peek(Token![#]) && input.peek2(Token![!]) {
            attr::parsing::single_parse_inner(input)
        } else {
            attr::parsing::single_parse_outer(input)
        }
    }
}

impl<T: Parse, P: Parse> ParseQuote for Punctuated<T, P> {
    fn parse(input: ParseStream) -> Result<Self> {
        Self::parse_terminated(input)
    }
}
<|endoftext|>


<commit_before><commit_msg>Added mapping on `EnvMBE`s.<commit_after><|endoftext|>