<commit_before>// Copyright 2015-2016 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

//! The host wrapper for communicating with a remote host.

use czmq::{ZCert, ZMsg, ZSock, ZSockType};
use error::Error;
use Result;
use std::path::Path;
use super::*;
use zfilexfer;

impl Host {
    /// Create a new Host to represent your managed host.
    pub fn new() -> Host {
        Host {
            hostname: None,
            api_sock: None,
            file_sock: None,
        }
    }

    #[cfg(test)]
    pub fn test_new(hostname: Option<String>, api_sock: Option<ZSock>, file_sock: Option<ZSock>) -> Host {
        let host = Host {
            hostname: hostname,
            api_sock: api_sock,
            file_sock: file_sock,
        };

        host
    }

    pub fn connect(&mut self, hostname: &str, api_port: u32, file_port: u32, auth_server: &str) -> Result<()> {
        self.hostname = Some(hostname.to_string());

        let user_cert = try!(ZCert::load("user.crt"));
        let server_cert = try!(self.lookup_server_cert(hostname, auth_server, &user_cert));

        let api_sock = ZSock::new(ZSockType::REQ);
        user_cert.apply(&api_sock);
        api_sock.set_curve_serverkey(server_cert.public_txt());
        api_sock.set_linger(5000);
        try!(api_sock.connect(&format!("tcp://{}:{}", hostname, api_port)));
        self.api_sock = Some(api_sock);

        let file_sock = ZSock::new(ZSockType::DEALER);
        user_cert.apply(&file_sock);
        file_sock.set_curve_serverkey(server_cert.public_txt());
        file_sock.set_linger(5000);
        try!(file_sock.connect(&format!("tcp://{}:{}", hostname, file_port)));
        self.file_sock = Some(file_sock);

        Ok(())
    }

    fn lookup_server_cert(&self, hostname: &str, auth_server: &str, user_cert: &ZCert) -> Result<ZCert> {
        let auth_cert = try!(ZCert::load("auth.crt"));

        let auth_sock = ZSock::new(ZSockType::REQ);
        user_cert.apply(&auth_sock);
        auth_sock.set_curve_serverkey(auth_cert.public_txt());
        auth_sock.set_sndtimeo(Some(5000));
        auth_sock.set_rcvtimeo(Some(5000));
        try!(auth_sock.connect(&format!("tcp://{}", auth_server)));

        // Get server cert from Auth server
        let msg = ZMsg::new();
        try!(msg.addstr("cert::lookup"));
        try!(msg.addstr(hostname));
        try!(msg.send(&auth_sock));

        let reply = try!(ZMsg::recv(&auth_sock));

        if reply.size() != 2 {
            return Err(Error::HostResponse);
        }

        match try!(reply.popstr().unwrap().or(Err(Error::HostResponse))).as_ref() {
            "Ok" => {
                let pk = try!(reply.popstr().unwrap().or(Err(Error::HostResponse)));
                Ok(ZCert::from_txt(&pk, "0000000000000000000000000000000000000000"))
            },
            "Err" => Err(Error::Agent(try!(reply.popstr().unwrap().or(Err(Error::HostResponse))))),
            _ => unreachable!(),
        }
    }

    pub fn close(&mut self) -> Result<()> {
        if self.api_sock.is_some() {
            self.api_sock.take().unwrap();
        }

        if self.file_sock.is_some() {
            self.file_sock.take().unwrap();
        }

        Ok(())
    }

    #[doc(hidden)]
    pub fn send(&mut self, msg: ZMsg) -> Result<()> {
        if self.api_sock.is_none() {
            return Err(Error::HostDisconnected);
        }

        try!(msg.send(self.api_sock.as_mut().unwrap()));
        Ok(())
    }

    #[doc(hidden)]
    pub fn send_file<P: AsRef<Path>>(&mut self, file: &zfilexfer::File, local_path: P) -> Result<()> {
        if self.file_sock.is_none() {
            return Err(Error::HostDisconnected);
        }

        try!(file.send(self.file_sock.as_ref().unwrap(), local_path));
        Ok(())
    }

    #[doc(hidden)]
    pub fn recv(&self, min: usize, max: Option<usize>) -> Result<ZMsg> {
        let msg = try!(ZMsg::recv(self.api_sock.as_ref().unwrap()));
        try!(Self::extract_header(&msg));

        // Check msg size
        if msg.size() < min || (max.is_some() && msg.size() > max.unwrap()) {
            Err(Error::HostResponse)
        } else {
            Ok(msg)
        }
    }

    fn extract_header(msg: &ZMsg) -> Result<()> {
        if msg.size() == 0 {
            return Err(Error::HostResponse);
        }

        match try!(msg.popstr().unwrap().or(Err(Error::HostResponse))).as_ref() {
            "Ok" => Ok(()),
            "Err" => {
                if msg.size() == 0 {
                    Err(Error::HostResponse)
                } else {
                    Err(Error::Agent(try!(msg.popstr().unwrap().or(Err(Error::HostResponse)))))
                }
            },
            _ => Err(Error::HostResponse),
        }
    }
}

#[cfg(test)]
mod tests {
    use {create_project_fs, Host, mock_auth_server};
    use czmq::{ZMsg, ZSys};
    use std::fs;
    use std::thread::spawn;
    use tempdir::TempDir;
    use zfilexfer::File;

    #[test]
    fn test_connect_close() {
        create_project_fs();
        let (handle, auth_server) = mock_auth_server();

        let mut host = Host::new();
        assert!(host.connect("localhost", 7101, 7102, &auth_server).is_ok());
        assert!(host.close().is_ok());

        handle.join().unwrap();
    }

    #[test]
    fn test_send_recv() {
        ZSys::init();

        let (client, server) = ZSys::create_pipe().unwrap();

        let mut host1 = Host::test_new(None, Some(client), None);
        let mut host2 = Host::test_new(None, Some(server), None);

        let msg = ZMsg::new();
        msg.addstr("Ok").unwrap();
        msg.addstr("moo").unwrap();
        msg.addstr("cow").unwrap();
        host1.send(msg).unwrap();

        let reply = host2.recv(2, Some(2)).unwrap();
        assert_eq!(reply.popstr().unwrap().unwrap(), "moo");
        assert_eq!(reply.popstr().unwrap().unwrap(), "cow");

        let msg = ZMsg::new();
        msg.addstr("No header").unwrap();
        host2.send(msg).unwrap();

        assert!(host1.recv(0, None).is_err());
        let msg = ZMsg::new();
        msg.addstr("Err").unwrap();
        host1.send(msg).unwrap();

        assert!(host2.recv(0, None).is_err());
    }

    #[test]
    fn test_send_file() {
        ZSys::init();

        let tempdir = TempDir::new("host_test_send_file").unwrap();
        let path = format!("{}/file.txt", tempdir.path().to_str().unwrap());
        fs::File::create(&path).unwrap();
        let file = File::open(&path, None).unwrap();

        let (client, server) = ZSys::create_pipe().unwrap();
        client.set_rcvtimeo(Some(500));
        server.set_rcvtimeo(Some(500));

        let handle = spawn(move|| {
            let msg = ZMsg::recv(&server).unwrap();
            assert_eq!(msg.popstr().unwrap().unwrap(), "NEW");

            server.flush();
            server.send_str("Ok").unwrap();
        });

        let mut host = Host::test_new(None, None, Some(client));
        assert!(host.send_file(&file, &path).is_ok());

        handle.join().unwrap();
    }
}
<commit_msg>Added timeouts to sockets<commit_after>// Copyright 2015-2016 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

//! The host wrapper for communicating with a remote host.

use czmq::{ZCert, ZMsg, ZSock, ZSockType};
use error::Error;
use Result;
use std::path::Path;
use super::*;
use zfilexfer;

impl Host {
    /// Create a new Host to represent your managed host.
    pub fn new() -> Host {
        Host {
            hostname: None,
            api_sock: None,
            file_sock: None,
        }
    }

    #[cfg(test)]
    pub fn test_new(hostname: Option<String>, api_sock: Option<ZSock>, file_sock: Option<ZSock>) -> Host {
        let host = Host {
            hostname: hostname,
            api_sock: api_sock,
            file_sock: file_sock,
        };

        host
    }

    pub fn connect(&mut self, hostname: &str, api_port: u32, file_port: u32, auth_server: &str) -> Result<()> {
        self.hostname = Some(hostname.to_string());

        let user_cert = try!(ZCert::load("user.crt"));
        let server_cert = try!(self.lookup_server_cert(hostname, auth_server, &user_cert));

        let api_sock = ZSock::new(ZSockType::REQ);
        user_cert.apply(&api_sock);
        api_sock.set_curve_serverkey(server_cert.public_txt());
        api_sock.set_sndtimeo(Some(5000));
        api_sock.set_rcvtimeo(Some(5000));
        try!(api_sock.connect(&format!("tcp://{}:{}", hostname, api_port)));
        self.api_sock = Some(api_sock);

        let file_sock = ZSock::new(ZSockType::DEALER);
        user_cert.apply(&file_sock);
        file_sock.set_curve_serverkey(server_cert.public_txt());
        file_sock.set_sndtimeo(Some(5000));
        file_sock.set_rcvtimeo(Some(5000));
        try!(file_sock.connect(&format!("tcp://{}:{}", hostname, file_port)));
        self.file_sock = Some(file_sock);

        Ok(())
    }

    fn lookup_server_cert(&self, hostname: &str, auth_server: &str, user_cert: &ZCert) -> Result<ZCert> {
        let auth_cert = try!(ZCert::load("auth.crt"));

        let auth_sock = ZSock::new(ZSockType::REQ);
        user_cert.apply(&auth_sock);
        auth_sock.set_curve_serverkey(auth_cert.public_txt());
        auth_sock.set_sndtimeo(Some(5000));
        auth_sock.set_rcvtimeo(Some(5000));
        try!(auth_sock.connect(&format!("tcp://{}", auth_server)));

        // Get server cert from Auth server
        let msg = ZMsg::new();
        try!(msg.addstr("cert::lookup"));
        try!(msg.addstr(hostname));
        try!(msg.send(&auth_sock));

        let reply = try!(ZMsg::recv(&auth_sock));

        if reply.size() != 2 {
            return Err(Error::HostResponse);
        }

        match try!(reply.popstr().unwrap().or(Err(Error::HostResponse))).as_ref() {
            "Ok" => {
                let pk = try!(reply.popstr().unwrap().or(Err(Error::HostResponse)));
                Ok(ZCert::from_txt(&pk, "0000000000000000000000000000000000000000"))
            },
            "Err" => Err(Error::Agent(try!(reply.popstr().unwrap().or(Err(Error::HostResponse))))),
            _ => unreachable!(),
        }
    }

    pub fn close(&mut self) -> Result<()> {
        if self.api_sock.is_some() {
            self.api_sock.take().unwrap();
        }

        if self.file_sock.is_some() {
            self.file_sock.take().unwrap();
        }

        Ok(())
    }

    #[doc(hidden)]
    pub fn send(&mut self, msg: ZMsg) -> Result<()> {
        if self.api_sock.is_none() {
            return Err(Error::HostDisconnected);
        }

        try!(msg.send(self.api_sock.as_mut().unwrap()));
        Ok(())
    }

    #[doc(hidden)]
    pub fn send_file<P: AsRef<Path>>(&mut self, file: &zfilexfer::File, local_path: P) -> Result<()> {
        if self.file_sock.is_none() {
            return Err(Error::HostDisconnected);
        }

        try!(file.send(self.file_sock.as_ref().unwrap(), local_path));
        Ok(())
    }

    #[doc(hidden)]
    pub fn recv(&self, min: usize, max: Option<usize>) -> Result<ZMsg> {
        let msg = try!(ZMsg::recv(self.api_sock.as_ref().unwrap()));
        try!(Self::extract_header(&msg));

        // Check msg size
        if msg.size() < min || (max.is_some() && msg.size() > max.unwrap()) {
            Err(Error::HostResponse)
        } else {
            Ok(msg)
        }
    }

    fn extract_header(msg: &ZMsg) -> Result<()> {
        if msg.size() == 0 {
            return Err(Error::HostResponse);
        }

        match try!(msg.popstr().unwrap().or(Err(Error::HostResponse))).as_ref() {
            "Ok" => Ok(()),
            "Err" => {
                if msg.size() == 0 {
                    Err(Error::HostResponse)
                } else {
                    Err(Error::Agent(try!(msg.popstr().unwrap().or(Err(Error::HostResponse)))))
                }
            },
            _ => Err(Error::HostResponse),
        }
    }
}

#[cfg(test)]
mod tests {
    use {create_project_fs, Host, mock_auth_server};
    use czmq::{ZMsg, ZSys};
    use std::fs;
    use std::thread::spawn;
    use tempdir::TempDir;
    use zfilexfer::File;

    #[test]
    fn test_connect_close() {
        create_project_fs();
        let (handle, auth_server) = mock_auth_server();

        let mut host = Host::new();
        assert!(host.connect("localhost", 7101, 7102, &auth_server).is_ok());
        assert!(host.close().is_ok());

        handle.join().unwrap();
    }

    #[test]
    fn test_send_recv() {
        ZSys::init();

        let (client, server) = ZSys::create_pipe().unwrap();

        let mut host1 = Host::test_new(None, Some(client), None);
        let mut host2 = Host::test_new(None, Some(server), None);

        let msg = ZMsg::new();
        msg.addstr("Ok").unwrap();
        msg.addstr("moo").unwrap();
        msg.addstr("cow").unwrap();
        host1.send(msg).unwrap();

        let reply = host2.recv(2, Some(2)).unwrap();
        assert_eq!(reply.popstr().unwrap().unwrap(), "moo");
        assert_eq!(reply.popstr().unwrap().unwrap(), "cow");

        let msg = ZMsg::new();
        msg.addstr("No header").unwrap();
        host2.send(msg).unwrap();

        assert!(host1.recv(0, None).is_err());
        let msg = ZMsg::new();
        msg.addstr("Err").unwrap();
        host1.send(msg).unwrap();

        assert!(host2.recv(0, None).is_err());
    }

    #[test]
    fn test_send_file() {
        ZSys::init();

        let tempdir = TempDir::new("host_test_send_file").unwrap();
        let path = format!("{}/file.txt", tempdir.path().to_str().unwrap());
        fs::File::create(&path).unwrap();
        let file = File::open(&path, None).unwrap();

        let (client, server) = ZSys::create_pipe().unwrap();
        client.set_rcvtimeo(Some(500));
        server.set_rcvtimeo(Some(500));

        let handle = spawn(move|| {
            let msg = ZMsg::recv(&server).unwrap();
            assert_eq!(msg.popstr().unwrap().unwrap(), "NEW");

            server.flush();
            server.send_str("Ok").unwrap();
        });

        let mut host = Host::test_new(None, None, Some(client));
        assert!(host.send_file(&file, &path).is_ok());

        handle.join().unwrap();
    }
}
<|endoftext|>


<commit_before>use std::path::Path;
use std::io::prelude::*;
use std::fs::File;
#[cfg(not(test))] use std::env::home_dir;
use yaml_rust::{YamlLoader, YamlEmitter, Yaml};

mod parser;

static MUXED_FOLDER: &'static str = "muxed";

pub fn open(project_name: String) -> Vec<Command> {
    let yaml = YamlLoader::load_from_str(&read(path_string(project_name))).unwrap();
    parser::main(&yaml)
}

pub fn path_string(project_name: String) -> String {
    format!("{}/.{}/{}", homedir_string(), &MUXED_FOLDER.to_string(), project_name)
}

fn read(config_str: String) -> String {
    let path = Path::new(&config_str);
    let mut s = String::new();
    File::open(path).expect("Config Read error").read_to_string(&mut s);

    return s
}

fn parse_config() {
}

#[cfg(not(test))] fn homedir_string() -> String {
    match home_dir() {
        Some(dir) => format!("{}", dir.display()),
        None      => panic!("Impossible to get your home dir!")
    }
}

#[cfg(test)] fn homedir_string() -> String {
  String::from("/tmp")
}

#[test]
pub fn path_string_returns_muxed_inside_homedir() {
    let path = format!("{}", path_string("".to_string()));
    let new  = format!("{}", Path::new("/tmp/.muxed/").display());
    assert_eq!(path, new)
}
<commit_msg>Wrong file extension.<commit_after>use std::path::Path;
use std::io::prelude::*;
use std::fs::File;
#[cfg(not(test))] use std::env::home_dir;
use yaml_rust::{YamlLoader, YamlEmitter, Yaml};

mod parser;

static MUXED_FOLDER: &'static str = "muxed";

pub fn open(project_name: String) -> Vec<Command> {
    let yaml = YamlLoader::load_from_str(&read(path_string(project_name))).unwrap();
    parser::main(&yaml)
}

pub fn path_string(project_name: String) -> String {
    format!("{}/.{}/{}.yml", homedir_string(), &MUXED_FOLDER.to_string(), project_name)
}

fn read(config_str: String) -> String {
    let path = Path::new(&config_str);
    let mut s = String::new();
    File::open(path).expect("Config Read error").read_to_string(&mut s);

    return s
}

fn parse_config() {
}

#[cfg(not(test))] fn homedir_string() -> String {
    match home_dir() {
        Some(dir) => format!("{}", dir.display()),
        None      => panic!("Impossible to get your home dir!")
    }
}

#[cfg(test)] fn homedir_string() -> String {
  String::from("/tmp")
}

#[test]
pub fn path_string_returns_muxed_inside_homedir() {
    let path = format!("{}", path_string("".to_string()));
    let new  = format!("{}", Path::new("/tmp/.muxed/").display());
    assert_eq!(path, new)
}
<|endoftext|>


<commit_before>/*!
 * Poorly categorised functions for reading things used in multiple places.
 *
 * (That is, typically, more than one of request-or-response reading-or-writing.)
 *
 * TODO: refactor all this to store things in more usefully categorised places.
 */
use std::num::{Unsigned, NumCast, Int, cast};
use std::io::{IoError, IoResult, OtherIoError};
#[cfg(test)]
use std::io::MemReader;

// XXX: IoError ain't a good representation of this.
fn bad_input() -> IoError {
    IoError {
        kind: OtherIoError,
        desc: "invalid number",
        detail: None,
    }
}

const ASCII_ZERO: u8 = b'0';
const ASCII_NINE: u8 = b'9';
const ASCII_LOWER_A: u8 = b'a';
const ASCII_LOWER_F: u8 = b'f';
const ASCII_UPPER_A: u8 = b'A';
const ASCII_UPPER_F: u8 = b'F';

/**
 * Read a positive decimal integer from the given reader.
 *
 * # Arguments
 *
 * - `reader` - the reader to read the decimal digits from (and whatever byte comes next)
 * - `expected_end` - this function is called with the byte that follows the last decimal digit
 *   and should return `true` if that byte is what is expected to follow the number.
 *
 * # Return value
 *
 * - `None`, if the number overflows;
 * - `None`, if there is a leading zero;
 * - `None`, if all the digits are read and the `expected_end` function is false
 *
 * Should everything work as designed (i.e. none of these conditions occur) a `Some` is returned.
 */
pub fn read_decimal<R: Reader, N: Unsigned + NumCast + Int>
                   (reader: &mut R, expected_end: |u8| -> bool)
                   -> IoResult<N> {
    // Here and in `read_hexadecimal` there is the possibility of infinite sequence of zeroes. The
    // spec allows this, but it may not be a good thing to allow. It's not a particularly good
    // attack surface, though, because of the low return.
    let mut n: N = Int::zero();
    let mut got_content = false;
    let ten: N = cast(10u32).unwrap();
    loop {
        n = match reader.read_byte() {
            Ok(b@ASCII_ZERO...ASCII_NINE) => {
                // Written sanely, this is: n * 10 + (b - '0'), but we avoid
                // (semantically unsound) overflow by using checked operations.
                // There is no need in the b - '0' part as it is safe.
                match n.checked_mul(ten).and_then(
                        |n| n.checked_add(cast(b - ASCII_ZERO).unwrap())) {
                    Some(new_n) => new_n,
                    None => return Err(bad_input()),  // overflow
                }
            },
            Ok(b) if got_content && expected_end(b) => return Ok(n),
            Ok(_) => return Err(bad_input()),  // not a valid number
            Err(err) => return Err(err),  // I/O error
        };
        got_content = true;
    }
}

/**
 * Read a positive hexadecimal integer from the given reader.
 *
 * # Arguments
 *
 * - `reader` - the reader to read the hexadecimal digits from (and whatever byte comes next)
 * - `expected_end` - this function is called with the byte that follows the last hexadecimal digit
 *   and should return `true` if that byte is what is expected to follow the number.
 *
 * # Return value
 *
 * - `None`, if the number overflows;
 * - `None`, if there is a leading zero;
 * - `None`, if all the digits are read and the `expected_end` function is false
 *
 * Should everything work as designed (i.e. none of these conditions occur) a `Some` is returned.
 */
pub fn read_hexadecimal<R: Reader, N: Unsigned + NumCast + Int>
                       (reader: &mut R, expected_end: |u8| -> bool)
                       -> IoResult<N> {
    let mut n: N = Int::zero();
    let mut got_content = false;
    let sixteen: N = cast(16u32).unwrap();
    loop {
        n = match reader.read_byte() {
            Ok(b@ASCII_ZERO...ASCII_NINE) => {
                match n.checked_mul(sixteen).and_then(
                        |n| n.checked_add(cast(b - ASCII_ZERO).unwrap())) {
                    Some(new_n) => new_n,
                    None => return Err(bad_input()),  // overflow
                }
            },
            Ok(b@ASCII_LOWER_A...ASCII_LOWER_F) => {
                match n.checked_mul(sixteen).and_then(
                        |n| n.checked_add(cast(b - ASCII_LOWER_A + 10).unwrap())) {
                    Some(new_n) => new_n,
                    None => return Err(bad_input()),  // overflow
                }
            },
            Ok(b@ASCII_UPPER_A...ASCII_UPPER_F) => {
                match n.checked_mul(sixteen).and_then(
                        |n| n.checked_add(cast(b - ASCII_UPPER_A + 10).unwrap())) {
                    Some(new_n) => new_n,
                    None => return Err(bad_input()),  // overflow
                }
            },
            Ok(b) if got_content && expected_end(b) => return Ok(n),
            Ok(_) => return Err(bad_input()),  // not a valid number
            Err(err) => return Err(err),  // I/O error
        };
        got_content = true;
    }
}

/**
 * Read an HTTP-Version (e.g. "HTTP/1.1") from a stream.
 *
 * # Arguments
 *
 * - `reader` - the reader to read the HTTP-Version from (and whatever byte comes next)
 * - `expected_end` - this function is called with the byte that follows the HTTP-Version
 *   and should return `true` if that byte is what is expected to follow the HTTP-Version.
 *
 * # Return value
 *
 * - `None`, if the HTTP-Version is malformed in any way;
 * - `None`, if the `expected_end` function returns false;
 * - A `Some`, if all goes well.
 */
#[inline]
pub fn read_http_version<R: Reader>
                        (reader: &mut R, expected_end: |u8| -> bool)
                        -> IoResult<(uint, uint)> {
    // I'd read into a [0u8, ..5], but that buffer is not guaranteed to be
    // filled, so I must read it byte by byte to guarantee correctness.
    // (Sure, no sane person/library would send the first packet with "HTT"
    // and leave the "P/1.1" to the next packet, but it's *possible*.)
    let b0 = try!(reader.read_byte());
    let b1 = try!(reader.read_byte());
    let b2 = try!(reader.read_byte());
    let b3 = try!(reader.read_byte());
    let b4 = try!(reader.read_byte());
    if (b0 != b'h' && b0 != b'H') ||
       (b1 != b't' && b1 != b'T') ||
       (b2 != b't' && b2 != b'T') ||
       (b3 != b'p' && b3 != b'P') ||
       b4 != b'/' {
        return Err(bad_input());
    }

    let major = try!(read_decimal(reader, |b| b == b'.'));
    let minor = try!(read_decimal(reader, expected_end));
    Ok((major, minor))
}

// I couldn't think what to call it. Ah well. It's just trivial syntax sugar, anyway.
macro_rules! test_reads {
    ($func:ident $($value:expr => $expected:expr),*) => {{
        $(
            assert_eq!(
                concat_idents!(read_, $func)(&mut MemReader::new($value.bytes().collect::<Vec<_>>()),
                                             |b| b == 0).ok(),
                $expected);
        )*
    }}
}

#[test]
fn test_read_http_version() {
    test_reads!(http_version
                "HTTP/25.17\0" => Some((25, 17)),
                "http/1.0\0" => Some((1, 0)),
                "http 1.0\0" => None,
                "HTTP/1.0.\0" => None,
                "HTTP/1.0.\0" => None
    );
}

#[test]
fn test_read_decimal() {
    test_reads!(decimal
                "0\0" => Some(0u),
                "0\0" => Some(0u8),
                "0\0" => Some(0u64),
                "9\0" => Some(9u8),
                "42\0" => Some(42u64),
                "0123456789\0" => Some(123456789u64),
                "00000000000000000000000031\0" => Some(31u64),

                // End of stream in middle of number
                "0" => None::<u64>,

                // No number
                "\0" => None::<u64>,

                // Invalid character
                "0a\0" => None::<u64>,

                // At overflow bounds
                "255\0" => Some(255u8),
                "256\0" => None::<u8>,
                "256\0" => Some(256u16)
    );
}

#[test]
fn test_read_hexadecimal() {
    test_reads!(hexadecimal
                "0\0" => Some(0u),
                "0\0" => Some(0u8),
                "0\0" => Some(0u64),
                "f\0" => Some(0xfu8),
                "42\0" => Some(0x42u64),
                "012345\0" => Some(0x12345u64),
                "00000000000000000000000031\0" => Some(0x31u64),
                "0123456789AbCdEf\0" => Some(0x123456789abcdefu64),

                // End of stream in middle of number
                "0" => None::<u64>,

                // No number
                "\0" => None::<u64>,

                // Invalid character
                "fg\0" => None::<u64>,

                // At overflow bounds
                "ff\0" => Some(0xffu8),
                "100\0" => None::<u8>,
                "100\0" => Some(0x100u16)
    );
}
<commit_msg>Fix for rust nightly<commit_after>/*!
 * Poorly categorised functions for reading things used in multiple places.
 *
 * (That is, typically, more than one of request-or-response reading-or-writing.)
 *
 * TODO: refactor all this to store things in more usefully categorised places.
 */
use std::num::{UnsignedInt, NumCast, Int, cast};
use std::io::{IoError, IoResult, OtherIoError};
#[cfg(test)]
use std::io::MemReader;

// XXX: IoError ain't a good representation of this.
fn bad_input() -> IoError {
    IoError {
        kind: OtherIoError,
        desc: "invalid number",
        detail: None,
    }
}

const ASCII_ZERO: u8 = b'0';
const ASCII_NINE: u8 = b'9';
const ASCII_LOWER_A: u8 = b'a';
const ASCII_LOWER_F: u8 = b'f';
const ASCII_UPPER_A: u8 = b'A';
const ASCII_UPPER_F: u8 = b'F';

/**
 * Read a positive decimal integer from the given reader.
 *
 * # Arguments
 *
 * - `reader` - the reader to read the decimal digits from (and whatever byte comes next)
 * - `expected_end` - this function is called with the byte that follows the last decimal digit
 *   and should return `true` if that byte is what is expected to follow the number.
 *
 * # Return value
 *
 * - `None`, if the number overflows;
 * - `None`, if there is a leading zero;
 * - `None`, if all the digits are read and the `expected_end` function is false
 *
 * Should everything work as designed (i.e. none of these conditions occur) a `Some` is returned.
 */
pub fn read_decimal<R: Reader, N: UnsignedInt + NumCast + Int>
                   (reader: &mut R, expected_end: |u8| -> bool)
                   -> IoResult<N> {
    // Here and in `read_hexadecimal` there is the possibility of infinite sequence of zeroes. The
    // spec allows this, but it may not be a good thing to allow. It's not a particularly good
    // attack surface, though, because of the low return.
    let mut n: N = Int::zero();
    let mut got_content = false;
    let ten: N = cast(10u32).unwrap();
    loop {
        n = match reader.read_byte() {
            Ok(b@ASCII_ZERO...ASCII_NINE) => {
                // Written sanely, this is: n * 10 + (b - '0'), but we avoid
                // (semantically unsound) overflow by using checked operations.
                // There is no need in the b - '0' part as it is safe.
                match n.checked_mul(ten).and_then(
                        |n| n.checked_add(cast(b - ASCII_ZERO).unwrap())) {
                    Some(new_n) => new_n,
                    None => return Err(bad_input()),  // overflow
                }
            },
            Ok(b) if got_content && expected_end(b) => return Ok(n),
            Ok(_) => return Err(bad_input()),  // not a valid number
            Err(err) => return Err(err),  // I/O error
        };
        got_content = true;
    }
}

/**
 * Read a positive hexadecimal integer from the given reader.
 *
 * # Arguments
 *
 * - `reader` - the reader to read the hexadecimal digits from (and whatever byte comes next)
 * - `expected_end` - this function is called with the byte that follows the last hexadecimal digit
 *   and should return `true` if that byte is what is expected to follow the number.
 *
 * # Return value
 *
 * - `None`, if the number overflows;
 * - `None`, if there is a leading zero;
 * - `None`, if all the digits are read and the `expected_end` function is false
 *
 * Should everything work as designed (i.e. none of these conditions occur) a `Some` is returned.
 */
pub fn read_hexadecimal<R: Reader, N: UnsignedInt + NumCast + Int>
                       (reader: &mut R, expected_end: |u8| -> bool)
                       -> IoResult<N> {
    let mut n: N = Int::zero();
    let mut got_content = false;
    let sixteen: N = cast(16u32).unwrap();
    loop {
        n = match reader.read_byte() {
            Ok(b@ASCII_ZERO...ASCII_NINE) => {
                match n.checked_mul(sixteen).and_then(
                        |n| n.checked_add(cast(b - ASCII_ZERO).unwrap())) {
                    Some(new_n) => new_n,
                    None => return Err(bad_input()),  // overflow
                }
            },
            Ok(b@ASCII_LOWER_A...ASCII_LOWER_F) => {
                match n.checked_mul(sixteen).and_then(
                        |n| n.checked_add(cast(b - ASCII_LOWER_A + 10).unwrap())) {
                    Some(new_n) => new_n,
                    None => return Err(bad_input()),  // overflow
                }
            },
            Ok(b@ASCII_UPPER_A...ASCII_UPPER_F) => {
                match n.checked_mul(sixteen).and_then(
                        |n| n.checked_add(cast(b - ASCII_UPPER_A + 10).unwrap())) {
                    Some(new_n) => new_n,
                    None => return Err(bad_input()),  // overflow
                }
            },
            Ok(b) if got_content && expected_end(b) => return Ok(n),
            Ok(_) => return Err(bad_input()),  // not a valid number
            Err(err) => return Err(err),  // I/O error
        };
        got_content = true;
    }
}

/**
 * Read an HTTP-Version (e.g. "HTTP/1.1") from a stream.
 *
 * # Arguments
 *
 * - `reader` - the reader to read the HTTP-Version from (and whatever byte comes next)
 * - `expected_end` - this function is called with the byte that follows the HTTP-Version
 *   and should return `true` if that byte is what is expected to follow the HTTP-Version.
 *
 * # Return value
 *
 * - `None`, if the HTTP-Version is malformed in any way;
 * - `None`, if the `expected_end` function returns false;
 * - A `Some`, if all goes well.
 */
#[inline]
pub fn read_http_version<R: Reader>
                        (reader: &mut R, expected_end: |u8| -> bool)
                        -> IoResult<(uint, uint)> {
    // I'd read into a [0u8, ..5], but that buffer is not guaranteed to be
    // filled, so I must read it byte by byte to guarantee correctness.
    // (Sure, no sane person/library would send the first packet with "HTT"
    // and leave the "P/1.1" to the next packet, but it's *possible*.)
    let b0 = try!(reader.read_byte());
    let b1 = try!(reader.read_byte());
    let b2 = try!(reader.read_byte());
    let b3 = try!(reader.read_byte());
    let b4 = try!(reader.read_byte());
    if (b0 != b'h' && b0 != b'H') ||
       (b1 != b't' && b1 != b'T') ||
       (b2 != b't' && b2 != b'T') ||
       (b3 != b'p' && b3 != b'P') ||
       b4 != b'/' {
        return Err(bad_input());
    }

    let major = try!(read_decimal(reader, |b| b == b'.'));
    let minor = try!(read_decimal(reader, expected_end));
    Ok((major, minor))
}

// I couldn't think what to call it. Ah well. It's just trivial syntax sugar, anyway.
macro_rules! test_reads {
    ($func:ident $($value:expr => $expected:expr),*) => {{
        $(
            assert_eq!(
                concat_idents!(read_, $func)(&mut MemReader::new($value.bytes().collect::<Vec<_>>()),
                                             |b| b == 0).ok(),
                $expected);
        )*
    }}
}

#[test]
fn test_read_http_version() {
    test_reads!(http_version
                "HTTP/25.17\0" => Some((25, 17)),
                "http/1.0\0" => Some((1, 0)),
                "http 1.0\0" => None,
                "HTTP/1.0.\0" => None,
                "HTTP/1.0.\0" => None
    );
}

#[test]
fn test_read_decimal() {
    test_reads!(decimal
                "0\0" => Some(0u),
                "0\0" => Some(0u8),
                "0\0" => Some(0u64),
                "9\0" => Some(9u8),
                "42\0" => Some(42u64),
                "0123456789\0" => Some(123456789u64),
                "00000000000000000000000031\0" => Some(31u64),

                // End of stream in middle of number
                "0" => None::<u64>,

                // No number
                "\0" => None::<u64>,

                // Invalid character
                "0a\0" => None::<u64>,

                // At overflow bounds
                "255\0" => Some(255u8),
                "256\0" => None::<u8>,
                "256\0" => Some(256u16)
    );
}

#[test]
fn test_read_hexadecimal() {
    test_reads!(hexadecimal
                "0\0" => Some(0u),
                "0\0" => Some(0u8),
                "0\0" => Some(0u64),
                "f\0" => Some(0xfu8),
                "42\0" => Some(0x42u64),
                "012345\0" => Some(0x12345u64),
                "00000000000000000000000031\0" => Some(0x31u64),
                "0123456789AbCdEf\0" => Some(0x123456789abcdefu64),

                // End of stream in middle of number
                "0" => None::<u64>,

                // No number
                "\0" => None::<u64>,

                // Invalid character
                "fg\0" => None::<u64>,

                // At overflow bounds
                "ff\0" => Some(0xffu8),
                "100\0" => None::<u8>,
                "100\0" => Some(0x100u16)
    );
}
<|endoftext|>


<commit_before>//! HyperLogLog implementation.
use bytecount;
use std::cmp;
use std::collections::hash_map::DefaultHasher;
use std::fmt;
use std::hash::{BuildHasher, Hash, Hasher};

use hash_utils::MyBuildHasherDefault;

/// A HyperLogLog is a data structure to count unique elements on a data stream.
///
/// # Examples
/// ```
/// use pdatastructs::hyperloglog::HyperLogLog;
///
/// // set up filter
/// let address_bits = 4;  // so we store 2^4 = 16 registers in total
/// let mut hll = HyperLogLog::new(address_bits);
///
/// // add some data
/// hll.add(&"my super long string");
/// hll.add(&"another super long string");
/// hll.add(&"my super long string");  // again
///
/// // later
/// assert_eq!(hll.count(), 2);
/// ```
///
/// # Applications
/// - an approximative `COUNT(DISTINCT x)` in SQL
/// - count distinct elements in a data stream
///
/// # How It Works
/// The HyperLogLog consists of `2^b` 8bit counters. Each counter is initialized to 0.
///
/// During insertion, a hash `h(x)` is calculated. The first `b` bits of the hash function are used
/// to address a register, the other bits are used to create a number `p` which essentially counts
/// the number of leading 0-bits (or in other words: the leftmost 1-bit). The addressed register is
/// then updated to the maximum of its current value and `p`.
///
/// The calculation of the count is based on `1 / Sum_0^{2^b} (2^-register_i)` with a bunch of
/// factors a corrections applied (see paper or source code).
///
/// # Implementation
/// In contrast to the official HyperLogLog, a 64 bit instead of a 32 bit hash function is used. The
/// registers always allocate 8 bits and are not compressed.
///
/// # See Also
/// - `std::collections::HashSet`: can be used to get the exact count but requires you to store
///   each and every element
/// - `pdatastructs::bloomfilter::BloomFilter` and `pdatastructs::cuckoofilter::CuckooFilter`: when
///   you also want to check if a single element is in the observed set
///
/// # References
/// - ["HyperLogLog: the analysis of a near-optimal cardinality estimation algorithm", Philippe
///   Flajolet, Éric Fusy, Olivier Gandouet, Frédéric Meunier, 2007](http://dmod.eu/deca/ft_gateway.cfm.pdf)
/// - [Wikipedia: HyperLogLog](https://en.wikipedia.org/wiki/HyperLogLog)
#[derive(Clone)]
pub struct HyperLogLog<B = MyBuildHasherDefault<DefaultHasher>>
where
    B: BuildHasher + Clone + Eq,
{
    registers: Vec<u8>,
    b: usize,
    buildhasher: B,
}

impl HyperLogLog {
    /// Creates a new, empty HyperLogLog.
    ///
    /// - `b` number of bits used for register selection, number of registers within the
    ///   HyperLogLog will be `2^b`. `b` must be in `[4, 16]`
    ///
    /// Panics when `b` is out of bounds.
    pub fn new(b: usize) -> Self {
        let bh = MyBuildHasherDefault::<DefaultHasher>::default();
        Self::with_hash(b, bh)
    }
}

impl<B> HyperLogLog<B>
where
    B: BuildHasher + Clone + Eq,
{
    /// Same as `new` but with a specific `BuildHasher`.
    pub fn with_hash(b: usize, buildhasher: B) -> Self {
        assert!(
            (b >= 4) & (b <= 16),
            "b ({}) must be larger or equal than 4 and smaller or equal than 16",
            b
        );

        let m = (1 as usize) << b;
        let registers = vec![0; m];
        Self {
            registers,
            b,
            buildhasher,
        }
    }

    /// Get number of bits used for register selection.
    pub fn b(&self) -> usize {
        self.b
    }

    /// Get number of registers.
    pub fn m(&self) -> usize {
        self.registers.len()
    }

    /// Get `BuildHasher`.
    pub fn buildhasher(&self) -> &B {
        &self.buildhasher
    }

    /// Adds an element to the HyperLogLog.
    pub fn add<T>(&mut self, obj: &T)
    where
        T: Hash,
    {
        let mut hasher = self.buildhasher.build_hasher();
        obj.hash(&mut hasher);
        let h: u64 = hasher.finish();

        // split h into:
        //  - w = 64 - b upper bits
        //  - j = b lower bits
        let w = h >> self.b;
        let j = h - (w << self.b); // no 1 as in the paper since register indices are 0-based

        // p = leftmost bit (1-based count)
        let p = w.leading_zeros() + 1 - (self.b as u32);

        let m_old = self.registers[j as usize];
        self.registers[j as usize] = cmp::max(m_old, p as u8);
    }

    /// Guess the number of unique elements seen by the HyperLogLog.
    pub fn count(&self) -> usize {
        let m = self.registers.len() as f64;

        let z = 1f64 / self
            .registers
            .iter()
            .map(|&x| 2f64.powi(-(i32::from(x))))
            .sum::<f64>();

        let am = if m >= 128. {
            0.7213 / (1. + 1.079 / m)
        } else if m >= 64. {
            0.709
        } else if m >= 32. {
            0.697
        } else {
            0.673
        };

        let e = am * m * m * z;

        let e_star = if e <= 5. / 2. * m {
            // small range correction
            let v = bytecount::count(&self.registers, 0);
            if v != 0 {
                m * (m / (v as f64)).ln()
            } else {
                e
            }
        } else if e <= 1. / 30. * 2f64.powi(32) {
            // intermediate range => no correction
            e
        } else {
            // large range correction
            -(2f64.powi(32)) * (1. - e / 2f64.powi(32)).ln()
        };

        e_star as usize
    }

    /// Merge w/ another HyperLogLog.
    ///
    /// This HyperLogLog will then have the same state as if all elements seen by `other` where
    /// directly added to `self`.
    ///
    /// Panics when `b` or `buildhasher` parameter of `self` and `other` do not match.
    pub fn merge(&mut self, other: &Self) {
        assert_eq!(
            self.b, other.b,
            "b must be equal (left={}, right={})",
            self.b, other.b
        );
        assert!(
            self.buildhasher == other.buildhasher,
            "buildhasher must be equal"
        );

        self.registers = self
            .registers
            .iter()
            .zip(other.registers.iter())
            .map(|x| cmp::max(x.0, x.1))
            .cloned()
            .collect();
    }

    /// Empties the HyperLogLog.
    pub fn clear(&mut self) {
        self.registers = vec![0; self.registers.len()];
    }

    /// Checks whether the HyperLogLog has never seen an element.
    pub fn is_empty(&self) -> bool {
        self.registers.iter().all(|&x| x == 0)
    }
}

impl fmt::Debug for HyperLogLog {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "HyperLogLog {{ b: {} }}", self.b)
    }
}

impl<T> Extend<T> for HyperLogLog
where
    T: Hash,
{
    fn extend<S: IntoIterator<Item = T>>(&mut self, iter: S) {
        for elem in iter {
            self.add(&elem);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::HyperLogLog;
    use hash_utils::BuildHasherSeeded;

    #[test]
    #[should_panic(expected = "b (3) must be larger or equal than 4 and smaller or equal than 16")]
    fn new_panics_b3() {
        HyperLogLog::new(3);
    }

    #[test]
    #[should_panic(expected = "b (17) must be larger or equal than 4 and smaller or equal than 16")]
    fn new_panics_b17() {
        HyperLogLog::new(17);
    }

    #[test]
    fn getter() {
        let hll = HyperLogLog::new(8);
        assert_eq!(hll.b(), 8);
        assert_eq!(hll.m(), 1 << 8);
        hll.buildhasher();
    }

    #[test]
    fn empty() {
        let hll = HyperLogLog::new(8);
        assert_eq!(hll.count(), 0);
        assert!(hll.is_empty());
    }

    #[test]
    fn add_b4_n1k() {
        let mut hll = HyperLogLog::new(4);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 571);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b8_n1k() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 966);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b12_n1k() {
        let mut hll = HyperLogLog::new(12);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 984);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n1k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 998);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b8_n10k() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10196);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b12_n10k() {
        let mut hll = HyperLogLog::new(12);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10303);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n10k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10055);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n100k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..100000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 100551);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n1m() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..1000000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 1000226);
        assert!(!hll.is_empty());
    }

    #[test]
    fn clear() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..1000 {
            hll.add(&i);
        }
        hll.clear();
        assert_eq!(hll.count(), 0);
        assert!(hll.is_empty());
    }

    #[test]
    fn clone() {
        let mut hll1 = HyperLogLog::new(12);
        for i in 0..500 {
            hll1.add(&i);
        }
        let c1a = hll1.count();

        let hll2 = hll1.clone();
        assert_eq!(hll2.count(), c1a);

        for i in 501..1000 {
            hll1.add(&i);
        }
        let c1b = hll1.count();
        assert_ne!(c1b, c1a);
        assert_eq!(hll2.count(), c1a);
    }

    #[test]
    fn merge() {
        let mut hll1 = HyperLogLog::new(12);
        let mut hll2 = HyperLogLog::new(12);
        let mut hll = HyperLogLog::new(12);
        for i in 0..500 {
            hll.add(&i);
            hll1.add(&i);
        }
        for i in 501..1000 {
            hll.add(&i);
            hll2.add(&i);
        }
        assert_ne!(hll.count(), hll1.count());
        assert_ne!(hll.count(), hll2.count());

        hll1.merge(&hll2);
        assert_eq!(hll.count(), hll1.count());
    }

    #[test]
    #[should_panic(expected = "b must be equal (left=5, right=12)")]
    fn merge_panics_p() {
        let mut hll1 = HyperLogLog::new(5);
        let hll2 = HyperLogLog::new(12);
        hll1.merge(&hll2);
    }

    #[test]
    #[should_panic(expected = "buildhasher must be equal")]
    fn merge_panics_buildhasher() {
        let mut hll1 = HyperLogLog::with_hash(12, BuildHasherSeeded::new(0));
        let hll2 = HyperLogLog::with_hash(12, BuildHasherSeeded::new(1));
        hll1.merge(&hll2);
    }

    #[test]
    fn debug() {
        let hll = HyperLogLog::new(12);
        assert_eq!(format!("{:?}", hll), "HyperLogLog { b: 12 }");
    }

    #[test]
    fn extend() {
        let mut hll = HyperLogLog::new(4);
        hll.extend(0..1000);
        assert_eq!(hll.count(), 571);
        assert!(!hll.is_empty());
    }
}
<commit_msg>hyperloglog: extent references<commit_after>//! HyperLogLog implementation.
use bytecount;
use std::cmp;
use std::collections::hash_map::DefaultHasher;
use std::fmt;
use std::hash::{BuildHasher, Hash, Hasher};

use hash_utils::MyBuildHasherDefault;

/// A HyperLogLog is a data structure to count unique elements on a data stream.
///
/// # Examples
/// ```
/// use pdatastructs::hyperloglog::HyperLogLog;
///
/// // set up filter
/// let address_bits = 4;  // so we store 2^4 = 16 registers in total
/// let mut hll = HyperLogLog::new(address_bits);
///
/// // add some data
/// hll.add(&"my super long string");
/// hll.add(&"another super long string");
/// hll.add(&"my super long string");  // again
///
/// // later
/// assert_eq!(hll.count(), 2);
/// ```
///
/// # Applications
/// - an approximative `COUNT(DISTINCT x)` in SQL
/// - count distinct elements in a data stream
///
/// # How It Works
/// The HyperLogLog consists of `2^b` 8bit counters. Each counter is initialized to 0.
///
/// During insertion, a hash `h(x)` is calculated. The first `b` bits of the hash function are used
/// to address a register, the other bits are used to create a number `p` which essentially counts
/// the number of leading 0-bits (or in other words: the leftmost 1-bit). The addressed register is
/// then updated to the maximum of its current value and `p`.
///
/// The calculation of the count is based on `1 / Sum_0^{2^b} (2^-register_i)` with a bunch of
/// factors a corrections applied (see paper or source code).
///
/// # Implementation
/// In contrast to the official HyperLogLog, a 64 bit instead of a 32 bit hash function is used. The
/// registers always allocate 8 bits and are not compressed.
///
/// # See Also
/// - `std::collections::HashSet`: can be used to get the exact count but requires you to store
///   each and every element
/// - `pdatastructs::bloomfilter::BloomFilter` and `pdatastructs::cuckoofilter::CuckooFilter`: when
///   you also want to check if a single element is in the observed set
///
/// # References
/// - ["HyperLogLog: the analysis of a near-optimal cardinality estimation algorithm", Philippe
///   Flajolet, Éric Fusy, Olivier Gandouet, Frédéric Meunier, 2007](http://dmod.eu/deca/ft_gateway.cfm.pdf)
/// - ["HyperLogLog in Practice: Algorithmic Engineering of a State of The Art Cardinality Estimation Algorithm", Stefan
///   Heule, Marc Nunkesser, Alexander Hall, 2013](https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/40671.pdf)
/// - [Wikipedia: HyperLogLog](https://en.wikipedia.org/wiki/HyperLogLog)
#[derive(Clone)]
pub struct HyperLogLog<B = MyBuildHasherDefault<DefaultHasher>>
where
    B: BuildHasher + Clone + Eq,
{
    registers: Vec<u8>,
    b: usize,
    buildhasher: B,
}

impl HyperLogLog {
    /// Creates a new, empty HyperLogLog.
    ///
    /// - `b` number of bits used for register selection, number of registers within the
    ///   HyperLogLog will be `2^b`. `b` must be in `[4, 16]`
    ///
    /// Panics when `b` is out of bounds.
    pub fn new(b: usize) -> Self {
        let bh = MyBuildHasherDefault::<DefaultHasher>::default();
        Self::with_hash(b, bh)
    }
}

impl<B> HyperLogLog<B>
where
    B: BuildHasher + Clone + Eq,
{
    /// Same as `new` but with a specific `BuildHasher`.
    pub fn with_hash(b: usize, buildhasher: B) -> Self {
        assert!(
            (b >= 4) & (b <= 16),
            "b ({}) must be larger or equal than 4 and smaller or equal than 16",
            b
        );

        let m = (1 as usize) << b;
        let registers = vec![0; m];
        Self {
            registers,
            b,
            buildhasher,
        }
    }

    /// Get number of bits used for register selection.
    pub fn b(&self) -> usize {
        self.b
    }

    /// Get number of registers.
    pub fn m(&self) -> usize {
        self.registers.len()
    }

    /// Get `BuildHasher`.
    pub fn buildhasher(&self) -> &B {
        &self.buildhasher
    }

    /// Adds an element to the HyperLogLog.
    pub fn add<T>(&mut self, obj: &T)
    where
        T: Hash,
    {
        let mut hasher = self.buildhasher.build_hasher();
        obj.hash(&mut hasher);
        let h: u64 = hasher.finish();

        // split h into:
        //  - w = 64 - b upper bits
        //  - j = b lower bits
        let w = h >> self.b;
        let j = h - (w << self.b); // no 1 as in the paper since register indices are 0-based

        // p = leftmost bit (1-based count)
        let p = w.leading_zeros() + 1 - (self.b as u32);

        let m_old = self.registers[j as usize];
        self.registers[j as usize] = cmp::max(m_old, p as u8);
    }

    /// Guess the number of unique elements seen by the HyperLogLog.
    pub fn count(&self) -> usize {
        let m = self.registers.len() as f64;

        let z = 1f64 / self
            .registers
            .iter()
            .map(|&x| 2f64.powi(-(i32::from(x))))
            .sum::<f64>();

        let am = if m >= 128. {
            0.7213 / (1. + 1.079 / m)
        } else if m >= 64. {
            0.709
        } else if m >= 32. {
            0.697
        } else {
            0.673
        };

        let e = am * m * m * z;

        let e_star = if e <= 5. / 2. * m {
            // small range correction
            let v = bytecount::count(&self.registers, 0);
            if v != 0 {
                m * (m / (v as f64)).ln()
            } else {
                e
            }
        } else if e <= 1. / 30. * 2f64.powi(32) {
            // intermediate range => no correction
            e
        } else {
            // large range correction
            -(2f64.powi(32)) * (1. - e / 2f64.powi(32)).ln()
        };

        e_star as usize
    }

    /// Merge w/ another HyperLogLog.
    ///
    /// This HyperLogLog will then have the same state as if all elements seen by `other` where
    /// directly added to `self`.
    ///
    /// Panics when `b` or `buildhasher` parameter of `self` and `other` do not match.
    pub fn merge(&mut self, other: &Self) {
        assert_eq!(
            self.b, other.b,
            "b must be equal (left={}, right={})",
            self.b, other.b
        );
        assert!(
            self.buildhasher == other.buildhasher,
            "buildhasher must be equal"
        );

        self.registers = self
            .registers
            .iter()
            .zip(other.registers.iter())
            .map(|x| cmp::max(x.0, x.1))
            .cloned()
            .collect();
    }

    /// Empties the HyperLogLog.
    pub fn clear(&mut self) {
        self.registers = vec![0; self.registers.len()];
    }

    /// Checks whether the HyperLogLog has never seen an element.
    pub fn is_empty(&self) -> bool {
        self.registers.iter().all(|&x| x == 0)
    }
}

impl fmt::Debug for HyperLogLog {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "HyperLogLog {{ b: {} }}", self.b)
    }
}

impl<T> Extend<T> for HyperLogLog
where
    T: Hash,
{
    fn extend<S: IntoIterator<Item = T>>(&mut self, iter: S) {
        for elem in iter {
            self.add(&elem);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::HyperLogLog;
    use hash_utils::BuildHasherSeeded;

    #[test]
    #[should_panic(expected = "b (3) must be larger or equal than 4 and smaller or equal than 16")]
    fn new_panics_b3() {
        HyperLogLog::new(3);
    }

    #[test]
    #[should_panic(expected = "b (17) must be larger or equal than 4 and smaller or equal than 16")]
    fn new_panics_b17() {
        HyperLogLog::new(17);
    }

    #[test]
    fn getter() {
        let hll = HyperLogLog::new(8);
        assert_eq!(hll.b(), 8);
        assert_eq!(hll.m(), 1 << 8);
        hll.buildhasher();
    }

    #[test]
    fn empty() {
        let hll = HyperLogLog::new(8);
        assert_eq!(hll.count(), 0);
        assert!(hll.is_empty());
    }

    #[test]
    fn add_b4_n1k() {
        let mut hll = HyperLogLog::new(4);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 571);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b8_n1k() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 966);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b12_n1k() {
        let mut hll = HyperLogLog::new(12);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 984);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n1k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 998);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b8_n10k() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10196);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b12_n10k() {
        let mut hll = HyperLogLog::new(12);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10303);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n10k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10055);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n100k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..100000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 100551);
        assert!(!hll.is_empty());
    }

    #[test]
    fn add_b16_n1m() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..1000000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 1000226);
        assert!(!hll.is_empty());
    }

    #[test]
    fn clear() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..1000 {
            hll.add(&i);
        }
        hll.clear();
        assert_eq!(hll.count(), 0);
        assert!(hll.is_empty());
    }

    #[test]
    fn clone() {
        let mut hll1 = HyperLogLog::new(12);
        for i in 0..500 {
            hll1.add(&i);
        }
        let c1a = hll1.count();

        let hll2 = hll1.clone();
        assert_eq!(hll2.count(), c1a);

        for i in 501..1000 {
            hll1.add(&i);
        }
        let c1b = hll1.count();
        assert_ne!(c1b, c1a);
        assert_eq!(hll2.count(), c1a);
    }

    #[test]
    fn merge() {
        let mut hll1 = HyperLogLog::new(12);
        let mut hll2 = HyperLogLog::new(12);
        let mut hll = HyperLogLog::new(12);
        for i in 0..500 {
            hll.add(&i);
            hll1.add(&i);
        }
        for i in 501..1000 {
            hll.add(&i);
            hll2.add(&i);
        }
        assert_ne!(hll.count(), hll1.count());
        assert_ne!(hll.count(), hll2.count());

        hll1.merge(&hll2);
        assert_eq!(hll.count(), hll1.count());
    }

    #[test]
    #[should_panic(expected = "b must be equal (left=5, right=12)")]
    fn merge_panics_p() {
        let mut hll1 = HyperLogLog::new(5);
        let hll2 = HyperLogLog::new(12);
        hll1.merge(&hll2);
    }

    #[test]
    #[should_panic(expected = "buildhasher must be equal")]
    fn merge_panics_buildhasher() {
        let mut hll1 = HyperLogLog::with_hash(12, BuildHasherSeeded::new(0));
        let hll2 = HyperLogLog::with_hash(12, BuildHasherSeeded::new(1));
        hll1.merge(&hll2);
    }

    #[test]
    fn debug() {
        let hll = HyperLogLog::new(12);
        assert_eq!(format!("{:?}", hll), "HyperLogLog { b: 12 }");
    }

    #[test]
    fn extend() {
        let mut hll = HyperLogLog::new(4);
        hll.extend(0..1000);
        assert_eq!(hll.count(), 571);
        assert!(!hll.is_empty());
    }
}
<|endoftext|>


<commit_before>use crate::{
    db::Database,
    model::{
        AnnoLichess, GameInfo, GameInfoPlayer, Mode, PersonalEntry, PersonalKeyBuilder,
        PersonalStatus, UserId, UserName,
    },
    util::NevermindExt as _,
};
use axum::http::StatusCode;
use clap::Clap;
use futures_util::StreamExt;
use rustc_hash::FxHashMap;
use shakmaty::{
    uci::Uci, variant::VariantPosition, zobrist::Zobrist, ByColor, CastlingMode, Color, Outcome,
    Position,
};
use std::{
    sync::Arc,
    hash::{Hash, Hasher, BuildHasher},
    time::{Duration, SystemTime},
    collections::hash_map::{RandomState},
};
use tokio::{
    sync::{
        mpsc::{self, error::SendTimeoutError},
        oneshot,
    },
    task::JoinHandle,
};

mod lila;

use lila::{Game, Lila};

#[derive(Clap, Clone)]
pub struct IndexerOpt {
    #[clap(long = "lila", default_value = "https://lichess.org")]
    lila: String,
    #[clap(long = "bearer")]
    bearer: Option<String>,
    #[clap(long = "indexers", default_value = "16")]
    indexers: usize,
}

#[derive(Clone)]
pub struct IndexerStub {
    random_state: RandomState,
    txs: Vec<mpsc::Sender<IndexerMessage>>,
}

impl IndexerStub {
    pub fn spawn(db: Arc<Database>, opt: IndexerOpt) -> (IndexerStub, Vec<JoinHandle<()>>) {
        let mut txs = Vec::with_capacity(opt.indexers);
        let mut join_handles = Vec::with_capacity(opt.indexers);
        for idx in 0..opt.indexers {
            let (tx, rx) = mpsc::channel(500);
            txs.push(tx);
            join_handles.push(tokio::spawn(
                    IndexerActor {
                        idx,
                        rx,
                        db: db.clone(),
                        lila: Lila::new(opt.clone()),
                    }
                    .run(),
                ));
        }
        let random_state = RandomState::new();
        (IndexerStub { txs, random_state }, join_handles)
    }

    pub async fn index_player(&self, player: UserName) -> oneshot::Receiver<()> {
        let (req, res) = oneshot::channel();

        let mut responsible_indexer = self.random_state.build_hasher();
        UserId::from(player.clone()).hash(&mut responsible_indexer);

        match self
            .txs[responsible_indexer.finish() as usize % self.txs.len()]
            .send_timeout(
                IndexerMessage::IndexPlayer {
                    player,
                    callback: req,
                },
                Duration::from_secs(10),
            )
            .await
        {
            Ok(_) => (),
            Err(SendTimeoutError::Timeout(_)) => {
                log::error!("indexer queue full for more than 10 seconds")
            }
            Err(SendTimeoutError::Closed(_)) => panic!("indexer died"),
        };

        res
    }
}

struct IndexerActor {
    idx: usize,
    rx: mpsc::Receiver<IndexerMessage>,
    db: Arc<Database>,
    lila: Lila,
}

impl IndexerActor {
    async fn run(mut self) {
        while let Some(msg) = self.rx.recv().await {
            match msg {
                IndexerMessage::IndexPlayer { callback, player } => {
                    self.index_player(player).await;
                    callback.send(()).nevermind("requester gone away");
                }
            }
        }
    }

    async fn index_player(&self, player: UserName) {
        let player_id = UserId::from(player.clone());
        let mut status = self
            .db
            .queryable()
            .get_player_status(&player_id)
            .expect("get player status")
            .unwrap_or_default();

        let since_created_at = match status
            .maybe_revisit_ongoing()
            .or_else(|| status.maybe_index())
        {
            Some(since) => since,
            None => {
                log::debug!("not reindexing {} so soon", player);
                return;
            }
        };

        log::info!(
            "indexer {} starting {} (created_at >= {})",
            self.idx,
            player,
            since_created_at
        );
        let mut games = match self.lila.user_games(&player, since_created_at).await {
            Ok(games) => games,
            Err(err) if err.status() == Some(StatusCode::NOT_FOUND) => {
                log::warn!("indexer did not find player {}", player);
                return;
            }
            Err(err) => {
                log::error!("indexer {}: request failed: {}", self.idx, err);
                return;
            }
        };

        let hash = ByColor::new_with(|color| PersonalKeyBuilder::with_user_pov(&player_id, color));
        let mut num_games = 0;
        while let Some(game) = games.next().await {
            let game = match game {
                Ok(game) => game,
                Err(err) => {
                    log::error!("{}", err);
                    continue;
                }
            };

            self.index_game(&player, &hash, game, &mut status);

            num_games += 1;
            if num_games % 1024 == 0 {
                log::info!("indexer {}: indexed {} games for {}", self.idx, num_games, player);
            }
        }

        status.indexed_at = SystemTime::now();
        self.db
            .queryable()
            .put_player_status(&player_id, status)
            .expect("put player status");
        log::info!("indexer {}: finished indexing {} games for {}", self.idx, num_games, player);
    }

    fn index_game(
        &self,
        player: &UserName,
        hash: &ByColor<PersonalKeyBuilder>,
        game: Game,
        status: &mut PersonalStatus,
    ) {
        status.latest_created_at = game.created_at;

        if game.status.is_ongoing() {
            if status.revisit_ongoing_created_at.is_none() {
                log::debug!("will revisit ongoing game {} eventually", game.id);
                status.revisit_ongoing_created_at = Some(game.created_at);
            }
            return;
        }

        if game.status.is_unindexable() {
            return;
        }

        let color = if Some(player) == game.user_name(Color::White) {
            Color::White
        } else if Some(player) == game.user_name(Color::Black) {
            Color::Black
        } else {
            return;
        };
        let year = AnnoLichess::from_time(game.last_move_at);
        let outcome = Outcome::from_winner(game.winner);

        let queryable = self.db.queryable();
        if queryable
            .get_game_info(game.id)
            .expect("get game info")
            .map_or(false, |info| info.indexed.into_color(color))
        {
            log::debug!(
                "{}/{} already indexed",
                game.id,
                color.fold("white", "black")
            );
            return;
        }
        queryable
            .merge_game_info(
                game.id,
                GameInfo {
                    winner: outcome.winner(),
                    speed: game.speed,
                    rated: game.rated,
                    year: year.year(),
                    players: game.players.map(|p| GameInfoPlayer {
                        name: p.user.map(|p| p.name.to_string()),
                        rating: p.rating,
                    }),
                    indexed: ByColor::new_with(|c| color == c),
                },
            )
            .expect("put game info");

        let variant = game.variant.into();
        let pos = match game.initial_fen {
            Some(fen) => VariantPosition::from_setup(variant, &fen, CastlingMode::Chess960),
            None => Ok(VariantPosition::new(variant)),
        };

        let mut pos: Zobrist<_, u128> = match pos {
            Ok(pos) => Zobrist::new(pos),
            Err(err) => {
                log::warn!("not indexing {}: {}", game.id, err);
                return;
            }
        };

        // Build an intermediate table to remove loops (due to repetitions).
        let mut table: FxHashMap<u128, Uci> =
            FxHashMap::with_capacity_and_hasher(game.moves.len(), Default::default());

        for (ply, san) in game.moves.into_iter().enumerate() {
            let m = match san.to_move(&pos) {
                Ok(m) => m,
                Err(err) => {
                    log::warn!("cutting off {} at ply {}: {}: {}", game.id, ply, err, san);
                    break;
                }
            };

            let uci = m.to_uci(CastlingMode::Chess960);
            table.insert(pos.zobrist_hash(), uci);

            pos.play_unchecked(&m);
        }

        for (zobrist, uci) in table {
            queryable
                .merge_personal(
                    hash.by_color(color)
                        .with_zobrist(variant, zobrist)
                        .with_year(year),
                    PersonalEntry::new_single(
                        uci.clone(),
                        game.speed,
                        Mode::from_rated(game.rated),
                        game.id,
                        outcome,
                    ),
                )
                .expect("merge personal");
        }
    }
}

enum IndexerMessage {
    IndexPlayer {
        player: UserName,
        callback: oneshot::Sender<()>,
    },
}

pub enum IndexerStatus {
    Ongoing,
    Completed,
    Failed,
}
<commit_msg>cargo fmt<commit_after>use crate::{
    db::Database,
    model::{
        AnnoLichess, GameInfo, GameInfoPlayer, Mode, PersonalEntry, PersonalKeyBuilder,
        PersonalStatus, UserId, UserName,
    },
    util::NevermindExt as _,
};
use axum::http::StatusCode;
use clap::Clap;
use futures_util::StreamExt;
use rustc_hash::FxHashMap;
use shakmaty::{
    uci::Uci, variant::VariantPosition, zobrist::Zobrist, ByColor, CastlingMode, Color, Outcome,
    Position,
};
use std::{
    collections::hash_map::RandomState,
    hash::{BuildHasher, Hash, Hasher},
    sync::Arc,
    time::{Duration, SystemTime},
};
use tokio::{
    sync::{
        mpsc::{self, error::SendTimeoutError},
        oneshot,
    },
    task::JoinHandle,
};

mod lila;

use lila::{Game, Lila};

#[derive(Clap, Clone)]
pub struct IndexerOpt {
    #[clap(long = "lila", default_value = "https://lichess.org")]
    lila: String,
    #[clap(long = "bearer")]
    bearer: Option<String>,
    #[clap(long = "indexers", default_value = "16")]
    indexers: usize,
}

#[derive(Clone)]
pub struct IndexerStub {
    random_state: RandomState,
    txs: Vec<mpsc::Sender<IndexerMessage>>,
}

impl IndexerStub {
    pub fn spawn(db: Arc<Database>, opt: IndexerOpt) -> (IndexerStub, Vec<JoinHandle<()>>) {
        let mut txs = Vec::with_capacity(opt.indexers);
        let mut join_handles = Vec::with_capacity(opt.indexers);
        for idx in 0..opt.indexers {
            let (tx, rx) = mpsc::channel(500);
            txs.push(tx);
            join_handles.push(tokio::spawn(
                IndexerActor {
                    idx,
                    rx,
                    db: db.clone(),
                    lila: Lila::new(opt.clone()),
                }
                .run(),
            ));
        }
        let random_state = RandomState::new();
        (IndexerStub { txs, random_state }, join_handles)
    }

    pub async fn index_player(&self, player: UserName) -> oneshot::Receiver<()> {
        let (req, res) = oneshot::channel();

        let mut responsible_indexer = self.random_state.build_hasher();
        UserId::from(player.clone()).hash(&mut responsible_indexer);

        match self.txs[responsible_indexer.finish() as usize % self.txs.len()]
            .send_timeout(
                IndexerMessage::IndexPlayer {
                    player,
                    callback: req,
                },
                Duration::from_secs(10),
            )
            .await
        {
            Ok(_) => (),
            Err(SendTimeoutError::Timeout(_)) => {
                log::error!("indexer queue full for more than 10 seconds")
            }
            Err(SendTimeoutError::Closed(_)) => panic!("indexer died"),
        };

        res
    }
}

struct IndexerActor {
    idx: usize,
    rx: mpsc::Receiver<IndexerMessage>,
    db: Arc<Database>,
    lila: Lila,
}

impl IndexerActor {
    async fn run(mut self) {
        while let Some(msg) = self.rx.recv().await {
            match msg {
                IndexerMessage::IndexPlayer { callback, player } => {
                    self.index_player(player).await;
                    callback.send(()).nevermind("requester gone away");
                }
            }
        }
    }

    async fn index_player(&self, player: UserName) {
        let player_id = UserId::from(player.clone());
        let mut status = self
            .db
            .queryable()
            .get_player_status(&player_id)
            .expect("get player status")
            .unwrap_or_default();

        let since_created_at = match status
            .maybe_revisit_ongoing()
            .or_else(|| status.maybe_index())
        {
            Some(since) => since,
            None => {
                log::debug!("not reindexing {} so soon", player);
                return;
            }
        };

        log::info!(
            "indexer {} starting {} (created_at >= {})",
            self.idx,
            player,
            since_created_at
        );
        let mut games = match self.lila.user_games(&player, since_created_at).await {
            Ok(games) => games,
            Err(err) if err.status() == Some(StatusCode::NOT_FOUND) => {
                log::warn!("indexer did not find player {}", player);
                return;
            }
            Err(err) => {
                log::error!("indexer {}: request failed: {}", self.idx, err);
                return;
            }
        };

        let hash = ByColor::new_with(|color| PersonalKeyBuilder::with_user_pov(&player_id, color));
        let mut num_games = 0;
        while let Some(game) = games.next().await {
            let game = match game {
                Ok(game) => game,
                Err(err) => {
                    log::error!("{}", err);
                    continue;
                }
            };

            self.index_game(&player, &hash, game, &mut status);

            num_games += 1;
            if num_games % 1024 == 0 {
                log::info!(
                    "indexer {}: indexed {} games for {}",
                    self.idx,
                    num_games,
                    player
                );
            }
        }

        status.indexed_at = SystemTime::now();
        self.db
            .queryable()
            .put_player_status(&player_id, status)
            .expect("put player status");
        log::info!(
            "indexer {}: finished indexing {} games for {}",
            self.idx,
            num_games,
            player
        );
    }

    fn index_game(
        &self,
        player: &UserName,
        hash: &ByColor<PersonalKeyBuilder>,
        game: Game,
        status: &mut PersonalStatus,
    ) {
        status.latest_created_at = game.created_at;

        if game.status.is_ongoing() {
            if status.revisit_ongoing_created_at.is_none() {
                log::debug!("will revisit ongoing game {} eventually", game.id);
                status.revisit_ongoing_created_at = Some(game.created_at);
            }
            return;
        }

        if game.status.is_unindexable() {
            return;
        }

        let color = if Some(player) == game.user_name(Color::White) {
            Color::White
        } else if Some(player) == game.user_name(Color::Black) {
            Color::Black
        } else {
            return;
        };
        let year = AnnoLichess::from_time(game.last_move_at);
        let outcome = Outcome::from_winner(game.winner);

        let queryable = self.db.queryable();
        if queryable
            .get_game_info(game.id)
            .expect("get game info")
            .map_or(false, |info| info.indexed.into_color(color))
        {
            log::debug!(
                "{}/{} already indexed",
                game.id,
                color.fold("white", "black")
            );
            return;
        }
        queryable
            .merge_game_info(
                game.id,
                GameInfo {
                    winner: outcome.winner(),
                    speed: game.speed,
                    rated: game.rated,
                    year: year.year(),
                    players: game.players.map(|p| GameInfoPlayer {
                        name: p.user.map(|p| p.name.to_string()),
                        rating: p.rating,
                    }),
                    indexed: ByColor::new_with(|c| color == c),
                },
            )
            .expect("put game info");

        let variant = game.variant.into();
        let pos = match game.initial_fen {
            Some(fen) => VariantPosition::from_setup(variant, &fen, CastlingMode::Chess960),
            None => Ok(VariantPosition::new(variant)),
        };

        let mut pos: Zobrist<_, u128> = match pos {
            Ok(pos) => Zobrist::new(pos),
            Err(err) => {
                log::warn!("not indexing {}: {}", game.id, err);
                return;
            }
        };

        // Build an intermediate table to remove loops (due to repetitions).
        let mut table: FxHashMap<u128, Uci> =
            FxHashMap::with_capacity_and_hasher(game.moves.len(), Default::default());

        for (ply, san) in game.moves.into_iter().enumerate() {
            let m = match san.to_move(&pos) {
                Ok(m) => m,
                Err(err) => {
                    log::warn!("cutting off {} at ply {}: {}: {}", game.id, ply, err, san);
                    break;
                }
            };

            let uci = m.to_uci(CastlingMode::Chess960);
            table.insert(pos.zobrist_hash(), uci);

            pos.play_unchecked(&m);
        }

        for (zobrist, uci) in table {
            queryable
                .merge_personal(
                    hash.by_color(color)
                        .with_zobrist(variant, zobrist)
                        .with_year(year),
                    PersonalEntry::new_single(
                        uci.clone(),
                        game.speed,
                        Mode::from_rated(game.rated),
                        game.id,
                        outcome,
                    ),
                )
                .expect("merge personal");
        }
    }
}

enum IndexerMessage {
    IndexPlayer {
        player: UserName,
        callback: oneshot::Sender<()>,
    },
}

pub enum IndexerStatus {
    Ongoing,
    Completed,
    Failed,
}
<|endoftext|>


<commit_before>//! Support for reading ELF files.
//!
//! Defines traits to abstract over the difference between PE32/PE32+,
//! and implements read functionality in terms of these traits.
//!
//! This module reuses some of the COFF functionality.
//!
//! Also provides `PeFile` and related types which implement the `Object` trait.

mod file;
pub use file::*;

mod section;
pub use section::*;

pub use super::coff::{SectionTable, SymbolTable};
<commit_msg>Update mod.rs<commit_after>//! Support for reading PE files.
//!
//! Defines traits to abstract over the difference between PE32/PE32+,
//! and implements read functionality in terms of these traits.
//!
//! This module reuses some of the COFF functionality.
//!
//! Also provides `PeFile` and related types which implement the `Object` trait.

mod file;
pub use file::*;

mod section;
pub use section::*;

pub use super::coff::{SectionTable, SymbolTable};
<|endoftext|>


<commit_before>pub trait ReplyRenderer {
    fn render(&self) -> String;
}

mod text;
mod image;
mod voice;
mod video;
mod music;
mod articles;
mod transfer_customer_service;

pub use self::text::TextReply;
pub use self::image::ImageReply;
pub use self::voice::VoiceReply;
pub use self::video::VideoReply;
pub use self::music::MusicReply;
pub use self::articles::ArticlesReply;
pub use self::transfer_customer_service::TransferCustomerServiceReply;

#[derive(Debug, Clone)]
pub enum Reply {
    TextReply(TextReply),
    ImageReply(ImageReply),
    VoiceReply(VoiceReply),
    VideoReply(VideoReply),
    MusicReply(MusicReply),
    ArticlesReply(ArticlesReply),
    TransferCustomerServiceReply(TransferCustomerServiceReply),
}
<commit_msg>Add Reply::render method<commit_after>pub trait ReplyRenderer {
    fn render(&self) -> String;
}

mod text;
mod image;
mod voice;
mod video;
mod music;
mod articles;
mod transfer_customer_service;

pub use self::text::TextReply;
pub use self::image::ImageReply;
pub use self::voice::VoiceReply;
pub use self::video::VideoReply;
pub use self::music::MusicReply;
pub use self::articles::ArticlesReply;
pub use self::transfer_customer_service::TransferCustomerServiceReply;

#[derive(Debug, Clone)]
pub enum Reply {
    TextReply(TextReply),
    ImageReply(ImageReply),
    VoiceReply(VoiceReply),
    VideoReply(VideoReply),
    MusicReply(MusicReply),
    ArticlesReply(ArticlesReply),
    TransferCustomerServiceReply(TransferCustomerServiceReply),
}

impl Reply {
    pub fn render(&self) -> String {
        let reply = match *self {
            Reply::TextReply(ref r) => r.render(),
            Reply::ImageReply(ref r) => r.render(),
            Reply::VoiceReply(ref r) => r.render(),
            Reply::VideoReply(ref r) => r.render(),
            Reply::MusicReply(ref r) => r.render(),
            Reply::ArticlesReply(ref r) => r.render(),
            Reply::TransferCustomerServiceReply(ref r) => r.render(),
        };
        reply
    }
}
<|endoftext|>


<commit_before>//! generate_series virtual table.
//! Port of C [generate_series "function"](http://www.sqlite.org/cgi/src/finfo?name=ext/misc/series.c).
use std::default::Default;
use libc;

use {Connection, Error, Result};
use ffi;
use vtab::{self, declare_vtab, Context, IndexInfo, Values, VTab, VTabCursor};

/// Register the "generate_series" module.
pub fn load_module(conn: &Connection) -> Result<()> {
    let aux: Option<()> = None;
    conn.create_module("generate_series", &SERIES_MODULE, aux)
}

init_module!(SERIES_MODULE,
             SeriesTab,
             SeriesTabCursor,
             None,
             series_connect,
             series_best_index,
             series_disconnect,
             None,
             series_open,
             series_close,
             series_filter,
             series_next,
             series_eof,
             series_column,
             series_rowid);

// Column numbers
// const SERIES_COLUMN_VALUE : libc::c_int = 0;
const SERIES_COLUMN_START: libc::c_int = 1;
const SERIES_COLUMN_STOP: libc::c_int = 2;
const SERIES_COLUMN_STEP: libc::c_int = 3;

bitflags! {
    #[repr(C)]
    flags QueryPlanFlags: ::libc::c_int {
        // start = $value  -- constraint exists
        const START = 1,
        // stop = $value   -- constraint exists
        const STOP  = 2,
        // step = $value   -- constraint exists
        const STEP  = 4,
        // output in descending order
        const DESC  = 8,
        // Both start and stop
        const BOTH  = START.bits | STOP.bits,
    }
}

/// An instance of the Series virtual table
#[repr(C)]
struct SeriesTab {
    /// Base class. Must be first
    base: ffi::sqlite3_vtab,
}

impl VTab<SeriesTabCursor> for SeriesTab {
    fn connect(db: *mut ffi::sqlite3,
               _aux: *mut libc::c_void,
               _args: &[&[u8]])
               -> Result<SeriesTab> {
        let vtab = SeriesTab { base: Default::default() };
        try!(declare_vtab(db,
                          "CREATE TABLE x(value,start hidden,stop hidden,step hidden)"));
        Ok(vtab)
    }

    fn best_index(&self, info: &mut IndexInfo) -> Result<()> {
        // The query plan bitmask
        let mut idx_num: QueryPlanFlags = QueryPlanFlags::empty();
        // Index of the start= constraint
        let mut start_idx = None;
        // Index of the stop= constraint
        let mut stop_idx = None;
        // Index of the step= constraint
        let mut step_idx = None;
        for i in 0..info.num_of_constraint() {
            if !info.is_constraint_usable(i) {
                continue;
            }
            if info.constraint_operator(i) != vtab::SQLITE_INDEX_CONSTRAINT_EQ {
                continue;
            }
            match info.constraint_column(i) {
                SERIES_COLUMN_START => {
                    start_idx = Some(i);
                    idx_num |= START;
                }
                SERIES_COLUMN_STOP => {
                    stop_idx = Some(i);
                    idx_num |= STOP;
                }
                SERIES_COLUMN_STEP => {
                    step_idx = Some(i);
                    idx_num |= STEP;
                }
                _ => {}
            };
        }

        let mut num_of_arg = 0;
        if let Some(start_idx) = start_idx {
            num_of_arg += 1;
            info.set_argv_index(start_idx, num_of_arg);
            info.set_omit(start_idx, true);
        }
        if let Some(stop_idx) = stop_idx {
            num_of_arg += 1;
            info.set_argv_index(stop_idx, num_of_arg);
            info.set_omit(stop_idx, true);
        }
        if let Some(step_idx) = step_idx {
            num_of_arg += 1;
            info.set_argv_index(step_idx, num_of_arg);
            info.set_omit(step_idx, true);
        }
        if idx_num.contains(BOTH) {
            // Both start= and stop= boundaries are available.
            info.set_estimated_cost((2 - if idx_num.contains(STEP) { 1 } else { 0 }) as f64);
            info.set_estimated_rows(1000);
            if info.num_of_order_by() == 1 {
                if info.is_order_by_desc(0) {
                    idx_num |= DESC;
                }
                info.set_order_by_consumed(true);
            }
        } else {
            info.set_estimated_cost(2147483647f64);
            info.set_estimated_rows(2147483647);
        }
        info.set_idx_num(idx_num.bits());
        Ok(())
    }

    fn open(&self) -> Result<SeriesTabCursor> {
        Ok(SeriesTabCursor::new())
    }
}

/// A cursor for the Series virtual table
#[derive(Default)]
#[repr(C)]
struct SeriesTabCursor {
    /// Base class. Must be first
    base: ffi::sqlite3_vtab_cursor,
    /// True to count down rather than up
    is_desc: bool,
    /// The rowid
    row_id: i64,
    /// Current value ("value")
    value: i64,
    /// Mimimum value ("start")
    min_value: i64,
    /// Maximum value ("stop")
    max_value: i64,
    /// Increment ("step")
    step: i64,
}

impl SeriesTabCursor {
    fn new() -> SeriesTabCursor {
        Default::default()
    }
}
impl VTabCursor<SeriesTab> for SeriesTabCursor {
    fn vtab(&self) -> &mut SeriesTab {
        unsafe { &mut *(self.base.pVtab as *mut SeriesTab) }
    }
    fn filter(&mut self,
              idx_num: libc::c_int,
              _idx_str: Option<&str>,
              args: &Values)
              -> Result<()> {
        let idx_num = QueryPlanFlags::from_bits_truncate(idx_num);
        let mut i = 0;
        if idx_num.contains(START) {
            self.min_value = try!(args.get(i));
            i += 1;
        } else {
            self.min_value = 0;
        }
        if idx_num.contains(STOP) {
            self.max_value = try!(args.get(i));
            i += 1;
        } else {
            self.max_value = 0xffffffff;
        }
        if idx_num.contains(STEP) {
            self.step = try!(args.get(i));
            if self.step < 1 {
                self.step = 1;
            }
        } else {
            self.step = 1;
        };
        self.is_desc = idx_num.contains(DESC);
        if self.is_desc {
            self.value = self.max_value;
            if self.step > 1 {
                self.value -= (self.max_value - self.min_value) % self.step;
            }
        } else {
            self.value = self.min_value;
        }
        self.row_id = 0;
        Ok(())
    }
    fn next(&mut self) -> Result<()> {
        if self.is_desc {
            self.value -= self.step;
        } else {
            self.value += self.step;
        }
        self.row_id += 1;
        Ok(())
    }
    fn eof(&self) -> bool {
        if self.is_desc {
            self.value < self.min_value
        } else {
            self.value > self.max_value
        }
    }
    fn column(&self, ctx: &mut Context, i: libc::c_int) -> Result<()> {
        let x = match i {
            SERIES_COLUMN_START => self.min_value,
            SERIES_COLUMN_STOP => self.max_value,
            SERIES_COLUMN_STEP => self.step,
            _ => self.value,
        };
        ctx.set_result(&x);
        Ok(())
    }
    fn rowid(&self) -> Result<i64> {
        Ok(self.row_id)
    }
}
<commit_msg>Test series virtual table<commit_after>//! generate_series virtual table.
//! Port of C [generate_series "function"](http://www.sqlite.org/cgi/src/finfo?name=ext/misc/series.c).
use std::default::Default;
use libc;

use {Connection, Error, Result};
use ffi;
use vtab::{self, declare_vtab, Context, IndexInfo, Values, VTab, VTabCursor};

/// Register the "generate_series" module.
pub fn load_module(conn: &Connection) -> Result<()> {
    let aux: Option<()> = None;
    conn.create_module("generate_series", &SERIES_MODULE, aux)
}

init_module!(SERIES_MODULE,
             SeriesTab,
             SeriesTabCursor,
             None,
             series_connect,
             series_best_index,
             series_disconnect,
             None,
             series_open,
             series_close,
             series_filter,
             series_next,
             series_eof,
             series_column,
             series_rowid);

// Column numbers
// const SERIES_COLUMN_VALUE : libc::c_int = 0;
const SERIES_COLUMN_START: libc::c_int = 1;
const SERIES_COLUMN_STOP: libc::c_int = 2;
const SERIES_COLUMN_STEP: libc::c_int = 3;

bitflags! {
    #[repr(C)]
    flags QueryPlanFlags: ::libc::c_int {
        // start = $value  -- constraint exists
        const START = 1,
        // stop = $value   -- constraint exists
        const STOP  = 2,
        // step = $value   -- constraint exists
        const STEP  = 4,
        // output in descending order
        const DESC  = 8,
        // Both start and stop
        const BOTH  = START.bits | STOP.bits,
    }
}

/// An instance of the Series virtual table
#[repr(C)]
struct SeriesTab {
    /// Base class. Must be first
    base: ffi::sqlite3_vtab,
}

impl VTab<SeriesTabCursor> for SeriesTab {
    fn connect(db: *mut ffi::sqlite3,
               _aux: *mut libc::c_void,
               _args: &[&[u8]])
               -> Result<SeriesTab> {
        let vtab = SeriesTab { base: Default::default() };
        try!(declare_vtab(db,
                          "CREATE TABLE x(value,start hidden,stop hidden,step hidden)"));
        Ok(vtab)
    }

    fn best_index(&self, info: &mut IndexInfo) -> Result<()> {
        // The query plan bitmask
        let mut idx_num: QueryPlanFlags = QueryPlanFlags::empty();
        // Index of the start= constraint
        let mut start_idx = None;
        // Index of the stop= constraint
        let mut stop_idx = None;
        // Index of the step= constraint
        let mut step_idx = None;
        for i in 0..info.num_of_constraint() {
            if !info.is_constraint_usable(i) {
                continue;
            }
            if info.constraint_operator(i) != vtab::SQLITE_INDEX_CONSTRAINT_EQ {
                continue;
            }
            match info.constraint_column(i) {
                SERIES_COLUMN_START => {
                    start_idx = Some(i);
                    idx_num |= START;
                }
                SERIES_COLUMN_STOP => {
                    stop_idx = Some(i);
                    idx_num |= STOP;
                }
                SERIES_COLUMN_STEP => {
                    step_idx = Some(i);
                    idx_num |= STEP;
                }
                _ => {}
            };
        }

        let mut num_of_arg = 0;
        if let Some(start_idx) = start_idx {
            num_of_arg += 1;
            info.set_argv_index(start_idx, num_of_arg);
            info.set_omit(start_idx, true);
        }
        if let Some(stop_idx) = stop_idx {
            num_of_arg += 1;
            info.set_argv_index(stop_idx, num_of_arg);
            info.set_omit(stop_idx, true);
        }
        if let Some(step_idx) = step_idx {
            num_of_arg += 1;
            info.set_argv_index(step_idx, num_of_arg);
            info.set_omit(step_idx, true);
        }
        if idx_num.contains(BOTH) {
            // Both start= and stop= boundaries are available.
            info.set_estimated_cost((2 - if idx_num.contains(STEP) { 1 } else { 0 }) as f64);
            info.set_estimated_rows(1000);
            if info.num_of_order_by() == 1 {
                if info.is_order_by_desc(0) {
                    idx_num |= DESC;
                }
                info.set_order_by_consumed(true);
            }
        } else {
            info.set_estimated_cost(2147483647f64);
            info.set_estimated_rows(2147483647);
        }
        info.set_idx_num(idx_num.bits());
        Ok(())
    }

    fn open(&self) -> Result<SeriesTabCursor> {
        Ok(SeriesTabCursor::new())
    }
}

/// A cursor for the Series virtual table
#[derive(Default)]
#[repr(C)]
struct SeriesTabCursor {
    /// Base class. Must be first
    base: ffi::sqlite3_vtab_cursor,
    /// True to count down rather than up
    is_desc: bool,
    /// The rowid
    row_id: i64,
    /// Current value ("value")
    value: i64,
    /// Mimimum value ("start")
    min_value: i64,
    /// Maximum value ("stop")
    max_value: i64,
    /// Increment ("step")
    step: i64,
}

impl SeriesTabCursor {
    fn new() -> SeriesTabCursor {
        Default::default()
    }
}
impl VTabCursor<SeriesTab> for SeriesTabCursor {
    fn vtab(&self) -> &mut SeriesTab {
        unsafe { &mut *(self.base.pVtab as *mut SeriesTab) }
    }
    fn filter(&mut self,
              idx_num: libc::c_int,
              _idx_str: Option<&str>,
              args: &Values)
              -> Result<()> {
        let idx_num = QueryPlanFlags::from_bits_truncate(idx_num);
        let mut i = 0;
        if idx_num.contains(START) {
            self.min_value = try!(args.get(i));
            i += 1;
        } else {
            self.min_value = 0;
        }
        if idx_num.contains(STOP) {
            self.max_value = try!(args.get(i));
            i += 1;
        } else {
            self.max_value = 0xffffffff;
        }
        if idx_num.contains(STEP) {
            self.step = try!(args.get(i));
            if self.step < 1 {
                self.step = 1;
            }
        } else {
            self.step = 1;
        };
        self.is_desc = idx_num.contains(DESC);
        if self.is_desc {
            self.value = self.max_value;
            if self.step > 1 {
                self.value -= (self.max_value - self.min_value) % self.step;
            }
        } else {
            self.value = self.min_value;
        }
        self.row_id = 0;
        Ok(())
    }
    fn next(&mut self) -> Result<()> {
        if self.is_desc {
            self.value -= self.step;
        } else {
            self.value += self.step;
        }
        self.row_id += 1;
        Ok(())
    }
    fn eof(&self) -> bool {
        if self.is_desc {
            self.value < self.min_value
        } else {
            self.value > self.max_value
        }
    }
    fn column(&self, ctx: &mut Context, i: libc::c_int) -> Result<()> {
        let x = match i {
            SERIES_COLUMN_START => self.min_value,
            SERIES_COLUMN_STOP => self.max_value,
            SERIES_COLUMN_STEP => self.step,
            _ => self.value,
        };
        ctx.set_result(&x);
        Ok(())
    }
    fn rowid(&self) -> Result<i64> {
        Ok(self.row_id)
    }
}

#[cfg(test)]
mod test {
    use Connection;
    use vtab::series;
    use ffi;

    #[test]
    fn test_series_module() {
        let version = unsafe { ffi::sqlite3_libversion_number() };
        if version < 3008012 {
            return;
        }

        let db = Connection::open_in_memory().unwrap();
        series::load_module(&db).unwrap();

        let mut s = db.prepare("SELECT * FROM generate_series(0,20,5)").unwrap();

        let series = s.query_map(&[], |row| row.get(0))
            .unwrap();

        let mut expected = 0;
        for value in series {
            assert_eq!(expected, value.unwrap());
            expected += 5;
        }
    }
}
<|endoftext|>