<commit_before>use consts::{OpCode, ZkError, ZkState};
use proto::{ByteBuf, ConnectRequest, ConnectResponse, ReadFrom, ReplyHeader, RequestHeader, WriteTo};

use byteorder::{BigEndian, ByteOrder};
use bytes::{Buf, RingBuf};
use mio::{EventLoop, EventSet, Handler, PollOpt, Sender, Timeout, Token, TryRead, TryWrite};
use mio::tcp::TcpStream;
use std::collections::VecDeque;
use std::io;
use std::io::Cursor;
use std::mem;
use std::net::SocketAddr;
use std::sync::mpsc::SyncSender;
use std::time::Duration;

const ZK: Token = Token(1);

lazy_static! {
    static ref PING: ByteBuf =
    RequestHeader{xid: -2, opcode: OpCode::Ping}.to_len_prefixed_buf().unwrap();
}

pub struct RawRequest {
    pub opcode: OpCode,
    pub data: ByteBuf,
    pub listener: Option<SyncSender<RawResponse>>
}

pub struct RawResponse {
    pub header: ReplyHeader,
    pub data: ByteBuf
}

struct Hosts {
    addrs: Vec<SocketAddr>,
    index: usize
}

impl Hosts {
    fn new(addrs: Vec<SocketAddr>) -> Hosts {
        Hosts{addrs: addrs, index: 0}
    }

    fn next(&mut self) -> &SocketAddr {
        let addr = &self.addrs[self.index];
        if self.addrs.len() == self.index+1 {
            self.index = 0;
        } else {
            self.index += 1;
        }
        addr
    }
}

struct ZkHandler {
    sock: TcpStream,
    state: ZkState,
    hosts: Hosts,
    buffer: VecDeque<RawRequest>,
    inflight: VecDeque<RawRequest>,
    response: RingBuf,
    timeout: Option<Timeout>,
    timeout_ms: u64,
    event_sender: Sender<RawResponse>,
}

impl ZkHandler {
    fn new(addrs: Vec<SocketAddr>, timeout_ms: u64, event_sender: Sender<RawResponse>) -> ZkHandler {
        ZkHandler{
            sock: unsafe{mem::dropped()},
            state: ZkState::NotConnected,
            hosts: Hosts::new(addrs),
            buffer: VecDeque::new(),
            inflight: VecDeque::new(),
            response: RingBuf::new(1024 * 1024), // TODO server reads max up to 1MB, otherwise drops the connection, size should be 1MB + tcp rcvBufsize
            timeout: None,
            timeout_ms: timeout_ms,
            event_sender: event_sender
        }
    }

    fn register(&mut self, event_loop: &mut EventLoop<Self>, events: EventSet) {
        event_loop.register_opt(&self.sock, ZK, events, PollOpt::edge() | PollOpt::oneshot()).unwrap();
    }

    fn handle_response(&mut self, event_loop: &mut EventLoop<Self>) {
        loop {
            if self.response.remaining() <= 4 {
                return
            }
            let len = BigEndian::read_i32(&self.response.bytes()[..4]) as usize;

            debug!("response chunk len = {} buf len is {}", len, self.response.bytes().len());

            if self.response.remaining() - 4 < len {
                return
            } else {

                self.response.advance(4);

                {
                    let mut data = Cursor::new(&self.response.bytes()[..len]);

                    debug!("handle_response in {:?} state [{}]", self.state, data.bytes().len());

                    // Could be moved to a handle_data() function
                    // or this whole stuff could be move to an Iterator
                    if self.state != ZkState::NotConnected {
                        let header = match ReplyHeader::read_from(&mut data) {
                            Ok(header) => header,
                            Err(e) => panic!("Failed to parse ReplyHeader {:?}", e) // TODO skip this chunk
                        };
                        let response = RawResponse{header: header, data: Cursor::new(data.bytes().to_vec())}; // TODO COPY!
                        match response.header.xid {
                            -1 => {
                                debug!("handle_response Got a watch event!");
                                self.event_sender.send(response).unwrap();
                            },
                            -2 => {
                                debug!("handle_response Got a ping response!");
                            },
                            _ => match self.inflight.pop_front() {
                                Some(ref request) => {
                                    Self::send_response(request, response);
                                    if request.opcode == OpCode::CloseSession {
                                        self.state = ZkState::Closed;
                                        event_loop.shutdown();
                                    }
                                },
                                None => panic!("Shouldn't happen, no inflight request")
                            }
                        }
                    } else {
                        self.inflight.pop_front(); // drop the connect request
                        let conn_resp = match ConnectResponse::read_from(&mut data) {
                            Ok(conn_resp) => conn_resp,
                            Err(e) => panic!("Failed to parse ConnectResponse {:?}", e) // TODO skip this chunk
                        };
                        info!("Connected: {:?}", conn_resp);
                        self.timeout_ms = conn_resp.timeout;
                        self.state = ZkState::Connected;
                    }
                }

                self.response.advance(len);
            }
        }
    }

    fn send_response(request: &RawRequest, response: RawResponse) {
        match request.listener {
            Some(ref listener) => {
                debug!("handle_response Opcode is {:?}", request.opcode);
                listener.send(response).unwrap();
            },
            None => debug!("Nobody is interested in respone {:?}", request.opcode)
        }
    }

    fn clear_ping_timeout(&mut self, event_loop: &mut EventLoop<Self>) {
        if let Some(timeout) = self.timeout {
            event_loop.clear_timeout(timeout);
        }
    }

    fn reconnect(&mut self, event_loop: &mut EventLoop<Self>) {
        info!("Connecting to new server");

        self.buffer.clear(); // drop all requests

        self.response.mark(); self.response.reset(); // TODO drop all read bytes once RingBuf.clear() is merged

        //self.state = ZkState::Connecting;

        self.clear_ping_timeout(event_loop);

        // TODO reconnect loop
        let host = self.hosts.next().clone();
        self.sock = match TcpStream::connect(&host) {
            Ok(sock) => sock,
            Err(e) => panic!("Failed to open connection {:?}: {:?}", host, e)
        };

        let request = self.connect_request();

        self.buffer.push_back(request);

        self.register(event_loop, EventSet::all());
    }

    fn connect_request(&self) -> RawRequest {
        let conn_req = ConnectRequest::from(ConnectResponse::initial(self.timeout_ms), 0);
        let buf = conn_req.to_len_prefixed_buf().unwrap();
        RawRequest{opcode: OpCode::Auth, data: buf, listener: None}
    }
}

impl Handler for ZkHandler {
    type Message = RawRequest;
    type Timeout = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Self>, token: Token, events: EventSet) {
        // Not sure that we need to write, but we always need to read, because of watches
        let mut next_events = EventSet::all();
        next_events.remove(EventSet::writable());

        debug!("ready {:?} {:?}", token, events);
        if events.is_writable() {
            loop {
                match self.buffer.pop_front() {
                    Some(mut request) => match self.sock.try_write_buf(&mut request.data) {
                        Ok(Some(written)) if written > 0 => {
                            debug!("Written {:?} bytes", written);
                            if request.data.has_remaining() {
                                self.buffer.push_front(request);
                                break
                            } else {
                                self.inflight.push_back(request);

                                // Sent a full message, clear the ping timeout
                                self.clear_ping_timeout(event_loop);
                            }
                        },
                        Ok(None) => trace!("Spurious write"),
                        Ok(Some(_)) => warn!("Connection closed: write"),
                        Err(e) => error!("Failed to write socket: {:?}", e)
                    },
                    None => break
                }
            }
            if !self.buffer.is_empty() {
                // Output buffer still has content, so we need to write again!
                next_events.insert(EventSet::writable());
            }
            self.timeout = Some(event_loop.timeout_ms((), self.timeout_ms).unwrap());
        }
        if events.is_readable() {
            match self.sock.try_read_buf(&mut self.response) {
                Ok(Some(read)) if read > 0 => {
                    debug!("Read {:?} bytes", read);
                    self.handle_response(event_loop);
                },
                Ok(None) => trace!("Spurious read"),
                Ok(Some(_)) => warn!("Connection closed: read"),
                Err(e) => error!("Failed to read socket: {:?}", e)
            }
        }
        if events.is_hup() {
            if self.state != ZkState::Closed {
                self.state = ZkState::NotConnected;
                self.reconnect(event_loop);
            }
        }

        // This tick is done, subscribe to a forthcoming one
        self.register(event_loop, next_events);
    }

    fn notify(&mut self, event_loop: &mut EventLoop<Self>, request: Self::Message) {
        debug!("notify {:?}", request.opcode);
        if self.state != ZkState::Closed {
            self.buffer.push_back(request);
            self.register(event_loop, EventSet::all());
        } else {
            let header = ReplyHeader{xid: 0, zxid: 0, err: ZkError::ConnectionLoss as i32};
            let response = RawResponse{header: header, data: ByteBuf::new(vec![])};
            Self::send_response(&request, response);
        }
    }

    fn timeout(&mut self, event_loop: &mut EventLoop<Self>, _: Self::Timeout) {
        debug!("Pinging {:?}", self.sock.peer_addr().unwrap());
        let ping = RawRequest{
            opcode: OpCode::Ping,
            data: PING.clone(),
            listener: None
        };
        event_loop.channel().send(ping).unwrap();
    }
}

pub struct ZkIo {
    event_loop: EventLoop<ZkHandler>,
    handler: ZkHandler,
}

impl ZkIo {
    pub fn new(addrs: Vec<SocketAddr>, timeout: Duration, event_sender: Sender<RawResponse>) -> ZkIo {
        let event_loop = EventLoop::new().unwrap();
        let timeout_ms = timeout.as_secs() * 1000 + timeout.subsec_nanos() as u64 / 1000000;
        let handler = ZkHandler::new(addrs, timeout_ms, event_sender);
        ZkIo{event_loop: event_loop, handler: handler}
    }

    pub fn sender(&self) -> Sender<RawRequest> {
        self.event_loop.channel()
    }

    pub fn run(mut self) -> io::Result<()> {
        self.handler.reconnect(&mut self.event_loop);
        self.event_loop.run(&mut self.handler)
    }
}
<commit_msg>Result::unwrap to ok().expect()<commit_after>use consts::{OpCode, ZkError, ZkState};
use proto::{ByteBuf, ConnectRequest, ConnectResponse, ReadFrom, ReplyHeader, RequestHeader, WriteTo};

use byteorder::{BigEndian, ByteOrder};
use bytes::{Buf, RingBuf};
use mio::{EventLoop, EventSet, Handler, PollOpt, Sender, Timeout, Token, TryRead, TryWrite};
use mio::tcp::TcpStream;
use std::collections::VecDeque;
use std::io;
use std::io::Cursor;
use std::mem;
use std::net::SocketAddr;
use std::sync::mpsc::SyncSender;
use std::time::Duration;

const ZK: Token = Token(1);

lazy_static! {
    static ref PING: ByteBuf =
    RequestHeader{xid: -2, opcode: OpCode::Ping}.to_len_prefixed_buf().unwrap();
}

pub struct RawRequest {
    pub opcode: OpCode,
    pub data: ByteBuf,
    pub listener: Option<SyncSender<RawResponse>>
}

pub struct RawResponse {
    pub header: ReplyHeader,
    pub data: ByteBuf
}

struct Hosts {
    addrs: Vec<SocketAddr>,
    index: usize
}

impl Hosts {
    fn new(addrs: Vec<SocketAddr>) -> Hosts {
        Hosts{addrs: addrs, index: 0}
    }

    fn next(&mut self) -> &SocketAddr {
        let addr = &self.addrs[self.index];
        if self.addrs.len() == self.index+1 {
            self.index = 0;
        } else {
            self.index += 1;
        }
        addr
    }
}

struct ZkHandler {
    sock: TcpStream,
    state: ZkState,
    hosts: Hosts,
    buffer: VecDeque<RawRequest>,
    inflight: VecDeque<RawRequest>,
    response: RingBuf,
    timeout: Option<Timeout>,
    timeout_ms: u64,
    event_sender: Sender<RawResponse>,
}

impl ZkHandler {
    fn new(addrs: Vec<SocketAddr>, timeout_ms: u64, event_sender: Sender<RawResponse>) -> ZkHandler {
        ZkHandler{
            sock: unsafe{mem::dropped()},
            state: ZkState::NotConnected,
            hosts: Hosts::new(addrs),
            buffer: VecDeque::new(),
            inflight: VecDeque::new(),
            response: RingBuf::new(1024 * 1024), // TODO server reads max up to 1MB, otherwise drops the connection, size should be 1MB + tcp rcvBufsize
            timeout: None,
            timeout_ms: timeout_ms,
            event_sender: event_sender
        }
    }

    fn register(&mut self, event_loop: &mut EventLoop<Self>, events: EventSet) {
        event_loop.register_opt(&self.sock, ZK, events, PollOpt::edge() | PollOpt::oneshot())
        .ok().expect("Failed to register ZK handle");
    }

    fn handle_response(&mut self, event_loop: &mut EventLoop<Self>) {
        loop {
            if self.response.remaining() <= 4 {
                return
            }
            let len = BigEndian::read_i32(&self.response.bytes()[..4]) as usize;

            debug!("response chunk len = {} buf len is {}", len, self.response.bytes().len());

            if self.response.remaining() - 4 < len {
                return
            } else {

                self.response.advance(4);

                {
                    let mut data = Cursor::new(&self.response.bytes()[..len]);

                    debug!("handle_response in {:?} state [{}]", self.state, data.bytes().len());

                    // Could be moved to a handle_data() function
                    // or this whole stuff could be move to an Iterator
                    if self.state != ZkState::NotConnected {
                        let header = match ReplyHeader::read_from(&mut data) {
                            Ok(header) => header,
                            Err(e) => panic!("Failed to parse ReplyHeader {:?}", e) // TODO skip this chunk
                        };
                        let response = RawResponse{header: header, data: Cursor::new(data.bytes().to_vec())}; // TODO COPY!
                        match response.header.xid {
                            -1 => {
                                debug!("handle_response Got a watch event!");
                                self.event_sender.send(response).unwrap();
                            },
                            -2 => {
                                debug!("handle_response Got a ping response!");
                            },
                            _ => match self.inflight.pop_front() {
                                Some(ref request) => {
                                    Self::send_response(request, response);
                                    if request.opcode == OpCode::CloseSession {
                                        self.state = ZkState::Closed;
                                        event_loop.shutdown();
                                    }
                                },
                                None => panic!("Shouldn't happen, no inflight request")
                            }
                        }
                    } else {
                        self.inflight.pop_front(); // drop the connect request
                        let conn_resp = match ConnectResponse::read_from(&mut data) {
                            Ok(conn_resp) => conn_resp,
                            Err(e) => panic!("Failed to parse ConnectResponse {:?}", e) // TODO skip this chunk
                        };
                        info!("Connected: {:?}", conn_resp);
                        self.timeout_ms = conn_resp.timeout;
                        self.state = ZkState::Connected;
                    }
                }

                self.response.advance(len);
            }
        }
    }

    fn send_response(request: &RawRequest, response: RawResponse) {
        match request.listener {
            Some(ref listener) => {
                debug!("handle_response Opcode is {:?}", request.opcode);
                listener.send(response).unwrap();
            },
            None => debug!("Nobody is interested in respone {:?}", request.opcode)
        }
    }

    fn clear_ping_timeout(&mut self, event_loop: &mut EventLoop<Self>) {
        if let Some(timeout) = self.timeout {
            event_loop.clear_timeout(timeout);
        }
    }

    fn reconnect(&mut self, event_loop: &mut EventLoop<Self>) {
        info!("Connecting to new server");

        self.buffer.clear(); // drop all requests

        self.response.mark(); self.response.reset(); // TODO drop all read bytes once RingBuf.clear() is merged

        //self.state = ZkState::Connecting;

        self.clear_ping_timeout(event_loop);

        // TODO reconnect loop
        let host = self.hosts.next().clone();
        self.sock = match TcpStream::connect(&host) {
            Ok(sock) => sock,
            Err(e) => panic!("Failed to open connection {:?}: {:?}", host, e)
        };

        let request = self.connect_request();

        self.buffer.push_back(request);

        self.register(event_loop, EventSet::all());
    }

    fn connect_request(&self) -> RawRequest {
        let conn_req = ConnectRequest::from(ConnectResponse::initial(self.timeout_ms), 0);
        let buf = conn_req.to_len_prefixed_buf().unwrap();
        RawRequest{opcode: OpCode::Auth, data: buf, listener: None}
    }
}

impl Handler for ZkHandler {
    type Message = RawRequest;
    type Timeout = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Self>, token: Token, events: EventSet) {
        // Not sure that we need to write, but we always need to read, because of watches
        let mut next_events = EventSet::all();
        next_events.remove(EventSet::writable());

        debug!("ready {:?} {:?}", token, events);
        if events.is_writable() {
            loop {
                match self.buffer.pop_front() {
                    Some(mut request) => match self.sock.try_write_buf(&mut request.data) {
                        Ok(Some(written)) if written > 0 => {
                            debug!("Written {:?} bytes", written);
                            if request.data.has_remaining() {
                                self.buffer.push_front(request);
                                break
                            } else {
                                self.inflight.push_back(request);

                                // Sent a full message, clear the ping timeout
                                self.clear_ping_timeout(event_loop);
                            }
                        },
                        Ok(None) => trace!("Spurious write"),
                        Ok(Some(_)) => warn!("Connection closed: write"),
                        Err(e) => error!("Failed to write socket: {:?}", e)
                    },
                    None => break
                }
            }
            if !self.buffer.is_empty() {
                // Output buffer still has content, so we need to write again!
                next_events.insert(EventSet::writable());
            }
            self.timeout = Some(event_loop.timeout_ms((), self.timeout_ms).unwrap());
        }
        if events.is_readable() {
            match self.sock.try_read_buf(&mut self.response) {
                Ok(Some(read)) if read > 0 => {
                    debug!("Read {:?} bytes", read);
                    self.handle_response(event_loop);
                },
                Ok(None) => trace!("Spurious read"),
                Ok(Some(_)) => warn!("Connection closed: read"),
                Err(e) => error!("Failed to read socket: {:?}", e)
            }
        }
        if events.is_hup() {
            if self.state != ZkState::Closed {
                self.state = ZkState::NotConnected;
                self.reconnect(event_loop);
            }
        }

        // This tick is done, subscribe to a forthcoming one
        self.register(event_loop, next_events);
    }

    fn notify(&mut self, event_loop: &mut EventLoop<Self>, request: Self::Message) {
        debug!("notify {:?}", request.opcode);
        if self.state != ZkState::Closed {
            self.buffer.push_back(request);
            self.register(event_loop, EventSet::all());
        } else {
            let header = ReplyHeader{xid: 0, zxid: 0, err: ZkError::ConnectionLoss as i32};
            let response = RawResponse{header: header, data: ByteBuf::new(vec![])};
            Self::send_response(&request, response);
        }
    }

    fn timeout(&mut self, event_loop: &mut EventLoop<Self>, _: Self::Timeout) {
        debug!("Pinging {:?}", self.sock.peer_addr().unwrap());
        let ping = RawRequest{
            opcode: OpCode::Ping,
            data: PING.clone(),
            listener: None
        };
        event_loop.channel().send(ping).unwrap();
    }
}

pub struct ZkIo {
    event_loop: EventLoop<ZkHandler>,
    handler: ZkHandler,
}

impl ZkIo {
    pub fn new(addrs: Vec<SocketAddr>, timeout: Duration, event_sender: Sender<RawResponse>) -> ZkIo {
        let event_loop = EventLoop::new().unwrap();
        let timeout_ms = timeout.as_secs() * 1000 + timeout.subsec_nanos() as u64 / 1000000;
        let handler = ZkHandler::new(addrs, timeout_ms, event_sender);
        ZkIo{event_loop: event_loop, handler: handler}
    }

    pub fn sender(&self) -> Sender<RawRequest> {
        self.event_loop.channel()
    }

    pub fn run(mut self) -> io::Result<()> {
        self.handler.reconnect(&mut self.event_loop);
        self.event_loop.run(&mut self.handler)
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use std::ops::{DerefMut, Deref};
use mio::util::Slab;
use mio::Socket;
use {EventPort, Promise, PromiseFulfiller, Result, new_promise_and_fulfiller};
use private::{with_current_event_loop};

pub trait AsyncRead {
    fn read<T>(self: Self, buf: T, min_bytes: usize, max_bytes: usize) -> Promise<(Self, T, usize)>
        where T: DerefMut<Target=[u8]>;
}

pub trait AsyncWrite {
    fn write<T>(self: Self, buf: T) -> Promise<(Self, T)>
        where T: Deref<Target=[u8]>;
}

#[derive(Copy, Clone)]
pub struct NetworkAddress {
    address: ::std::net::SocketAddr,
}

impl NetworkAddress {
    pub fn new(address: ::std::net::SocketAddr) -> NetworkAddress {
        NetworkAddress { address: address }
    }

    pub fn listen(self) -> Result<ConnectionReceiver> {
        let socket = try!(::mio::tcp::TcpSocket::v4());
        try!(socket.set_reuseaddr(true));
        try!(socket.bind(&self.address));
        let token = FdObserver::new();
        Ok(ConnectionReceiver { listener: try!(socket.listen(256)),
                                token: token })
    }

    pub fn connect(self) -> Promise<AsyncIoStream> {
        let socket = ::mio::tcp::TcpSocket::v4().unwrap();
        let token = FdObserver::new();
        let (stream, _) = socket.connect(&self.address).unwrap();
        with_current_event_loop(move |event_loop| {
            event_loop.event_port.borrow_mut().reactor.register_opt(&stream, token,
                                                                    ::mio::Interest::writable(),
                                                                    ::mio::PollOpt::edge()|
                                                                    ::mio::PollOpt::oneshot()).unwrap();
            let promise =
                event_loop.event_port.borrow_mut().handler.observers[token].when_becomes_writable();
            return promise.map(move |()| {
                // TODO check for error.

                return with_current_event_loop(move |event_loop| {
                    event_loop.event_port.borrow_mut().reactor.reregister(
                        &stream, token,
                        ::mio::Interest::writable() | ::mio::Interest::readable(),
                        ::mio::PollOpt::edge()).unwrap();
                    return Ok(AsyncIoStream { stream: stream, token: token });
                });

            });
        })
    }
}

pub struct ConnectionReceiver {
    listener: ::mio::tcp::TcpListener,
    token: ::mio::Token,
}

impl Drop for ConnectionReceiver {
    fn drop(&mut self) {
        // deregister the token
    }
}

impl ConnectionReceiver {
    pub fn accept(self) -> Promise<(ConnectionReceiver, AsyncIoStream)> {
        with_current_event_loop(move |event_loop| {
            let promise =
                event_loop.event_port.borrow_mut().handler.observers[self.token].when_becomes_readable();
            event_loop.event_port.borrow_mut().reactor.register_opt(&self.listener, self.token,
                                                                    ::mio::Interest::readable(),
                                                                    ::mio::PollOpt::edge()|
                                                                    ::mio::PollOpt::oneshot()).unwrap();
            return promise.map(move |()| {
                let stream = try!(self.listener.accept()).unwrap();
                let token = FdObserver::new();
                return Ok((self, AsyncIoStream::new(stream, token)));
            });
        })
    }
}

#[allow(dead_code)]
pub struct AsyncIoStream {
    stream: ::mio::tcp::TcpStream,
    token: ::mio::Token,
}

impl AsyncIoStream {
    fn new(stream: ::mio::tcp::TcpStream, token: ::mio::Token) -> AsyncIoStream {
        with_current_event_loop(move |event_loop| {
            event_loop.event_port.borrow_mut().reactor.register_opt(
                &stream, token,
                ::mio::Interest::readable() | ::mio::Interest::writable(),
                ::mio::PollOpt::level()).unwrap();
            AsyncIoStream { stream: stream, token: token }
        })
    }

    fn read_internal<T>(mut self,
                        mut buf: T,
                        start: usize,
                        min_bytes: usize,
                        max_bytes: usize) -> Promise<(Self, T, usize)> where T: DerefMut<Target=[u8]> {
        if start >= min_bytes {
            return Promise::fulfilled((self, buf, start));
        } else {
            with_current_event_loop(move |event_loop| {
                let promise =
                    event_loop.event_port.borrow_mut().handler.observers[self.token].when_becomes_readable();
                return promise.then(move |()| {
                    use mio::TryRead;
                    let n = try!(self.stream.read_slice(&mut buf[start..])).unwrap();
                    return Ok(self.read_internal(buf, start + n, min_bytes, max_bytes));
                });
            })
        }
    }

    fn write_internal<T>(mut self,
                         buf: T,
                         cursor: usize) -> Promise<(Self, T)> where T: Deref<Target=[u8]> {
        if cursor == buf.len() {
            return Promise::fulfilled((self, buf));
        } else {
            with_current_event_loop(move |event_loop| {
                let promise =
                    event_loop.event_port.borrow_mut().handler.observers[self.token].when_becomes_writable();
                return promise.then(move |()| {
                    use mio::TryWrite;
                    let n = try!(self.stream.write_slice(&buf[cursor..])).unwrap();
                    return Ok(self.write_internal(buf, cursor + n));
                });
            })
        }
    }
}

impl AsyncRead for AsyncIoStream {
    fn read<T>(self, buf: T,
               min_bytes: usize,
               max_bytes: usize) -> Promise<(Self, T, usize)> where T: DerefMut<Target=[u8]> {
        self.read_internal(buf, 0, min_bytes, max_bytes)
    }
}

impl AsyncWrite for AsyncIoStream {
    fn write<T>(self, buf: T) -> Promise<(Self, T)> where T: Deref<Target=[u8]> {
        self.write_internal(buf, 0)
    }
}

pub struct FdObserver {
    read_fulfiller: Option<Box<PromiseFulfiller<()>>>,
    write_fulfiller: Option<Box<PromiseFulfiller<()>>>,
}

impl FdObserver {
    pub fn new() -> ::mio::Token {
        with_current_event_loop(move |event_loop| {

            let observer = FdObserver { read_fulfiller: None, write_fulfiller: None };
            let event_port = &mut *event_loop.event_port.borrow_mut();
            let token = match event_port.handler.observers.insert(observer) {
                Ok(tok) => tok,
                Err(_obs) => panic!("could not insert observer."),
            };
            return token;
        })
    }

    pub fn when_becomes_readable(&mut self) -> Promise<()> {
        let (promise, fulfiller) = new_promise_and_fulfiller();
        self.read_fulfiller = Some(fulfiller);
        return promise;
    }

    pub fn when_becomes_writable(&mut self) -> Promise<()> {
        let (promise, fulfiller) = new_promise_and_fulfiller();
        self.write_fulfiller = Some(fulfiller);
        return promise;
    }
}

pub struct MioEventPort {
    handler: Handler,
    reactor: ::mio::EventLoop<Handler>,
}

pub struct Handler {
    observers: Slab<FdObserver>,
}

impl MioEventPort {
    pub fn new() -> Result<MioEventPort> {
        Ok(MioEventPort {
            handler: Handler { observers: Slab::new(100) },
            reactor: try!(::mio::EventLoop::new()),
        })
    }
}

impl ::mio::Handler for Handler {
    type Timeout = ();
    type Message = ();
    fn readable(&mut self, _event_loop: &mut ::mio::EventLoop<Handler>,
                token: ::mio::Token, _hint: ::mio::ReadHint) {
        match ::std::mem::replace(&mut self.observers[token].read_fulfiller, None) {
            Some(fulfiller) => {
                fulfiller.fulfill(())
            }
            None => {
                ()
            }
        }
    }
    fn writable(&mut self, _event_loop: &mut ::mio::EventLoop<Handler>, token: ::mio::Token) {
        match ::std::mem::replace(&mut self.observers[token].write_fulfiller, None) {
            Some(fulfiller) => fulfiller.fulfill(()),
            None => (),
        }
    }
}

impl EventPort for MioEventPort {
    fn wait(&mut self) -> bool {
        self.reactor.run_once(&mut self.handler).unwrap();
        return false;
    }

    fn poll(&mut self) -> bool {
        self.reactor.run_once(&mut self.handler).unwrap();
        return false;
    }
}

// alternate approach below =========

pub enum StepResult<S, T> {
   TheresMore(S),
   Done(T)
}

/// Intermediate state for a reading a T.
pub trait AsyncReadState<T> {

   /// Reads as much as possible without blocking. If done, returns the final T value. Otherwise
   /// returns the new intermediate state T.
   fn read_step(self, r: &mut ::std::io::Read) -> ::std::io::Result<StepResult<Self, T>>;
}

/// Gives back `r` once the T has been completely read.
pub fn read_async<R, S, T>(_r: R, _state: S) -> Promise<(R, T)>
  where R: ::std::io::Read + 'static,
        S: AsyncReadState<T> + 'static,
        T: 'static {
            unimplemented!();
}
<commit_msg>Always use edge-triggered mode. Delete unused things.<commit_after>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use std::ops::{DerefMut, Deref};
use mio::util::Slab;
use mio::Socket;
use {EventPort, Promise, PromiseFulfiller, Result, new_promise_and_fulfiller};
use private::{with_current_event_loop};

pub trait AsyncRead {
    fn read<T>(self: Self, buf: T, min_bytes: usize, max_bytes: usize) -> Promise<(Self, T, usize)>
        where T: DerefMut<Target=[u8]>;
}

pub trait AsyncWrite {
    fn write<T>(self: Self, buf: T) -> Promise<(Self, T)>
        where T: Deref<Target=[u8]>;
}

#[derive(Copy, Clone)]
pub struct NetworkAddress {
    address: ::std::net::SocketAddr,
}

impl NetworkAddress {
    pub fn new(address: ::std::net::SocketAddr) -> NetworkAddress {
        NetworkAddress { address: address }
    }

    pub fn listen(self) -> Result<ConnectionReceiver> {
        let socket = try!(::mio::tcp::TcpSocket::v4());
        try!(socket.set_reuseaddr(true));
        try!(socket.bind(&self.address));
        let token = FdObserver::new();
        Ok(ConnectionReceiver { listener: try!(socket.listen(256)),
                                token: token })
    }

    pub fn connect(self) -> Promise<AsyncIoStream> {
        let socket = ::mio::tcp::TcpSocket::v4().unwrap();
        let token = FdObserver::new();
        let (stream, _) = socket.connect(&self.address).unwrap();
        with_current_event_loop(move |event_loop| {
            event_loop.event_port.borrow_mut().reactor.register_opt(&stream, token,
                                                                    ::mio::Interest::writable(),
                                                                    ::mio::PollOpt::edge()|
                                                                    ::mio::PollOpt::oneshot()).unwrap();
            let promise =
                event_loop.event_port.borrow_mut().handler.observers[token].when_becomes_writable();
            return promise.map(move |()| {
                // TODO check for error.

                return with_current_event_loop(move |event_loop| {
                    event_loop.event_port.borrow_mut().reactor.reregister(
                        &stream, token,
                        ::mio::Interest::writable() | ::mio::Interest::readable(),
                        ::mio::PollOpt::edge()).unwrap();
                    return Ok(AsyncIoStream { stream: stream, token: token });
                });

            });
        })
    }
}

pub struct ConnectionReceiver {
    listener: ::mio::tcp::TcpListener,
    token: ::mio::Token,
}

impl Drop for ConnectionReceiver {
    fn drop(&mut self) {
        // deregister the token
    }
}

impl ConnectionReceiver {
    pub fn accept(self) -> Promise<(ConnectionReceiver, AsyncIoStream)> {
        with_current_event_loop(move |event_loop| {
            let promise =
                event_loop.event_port.borrow_mut().handler.observers[self.token].when_becomes_readable();
            event_loop.event_port.borrow_mut().reactor.register_opt(&self.listener, self.token,
                                                                    ::mio::Interest::readable(),
                                                                    ::mio::PollOpt::edge() |
                                                                    ::mio::PollOpt::oneshot()).unwrap();
            return promise.map(move |()| {
                let stream = try!(self.listener.accept()).unwrap();
                let token = FdObserver::new();
                return Ok((self, AsyncIoStream::new(stream, token)));
            });
        })
    }
}

#[allow(dead_code)]
pub struct AsyncIoStream {
    stream: ::mio::tcp::TcpStream,
    token: ::mio::Token,
}

impl AsyncIoStream {
    fn new(stream: ::mio::tcp::TcpStream, token: ::mio::Token) -> AsyncIoStream {
        with_current_event_loop(move |event_loop| {
            event_loop.event_port.borrow_mut().reactor.register_opt(
                &stream, token,
                ::mio::Interest::readable() | ::mio::Interest::writable(),
                ::mio::PollOpt::edge()).unwrap();
            AsyncIoStream { stream: stream, token: token }
        })
    }

    fn read_internal<T>(mut self,
                        mut buf: T,
                        start: usize,
                        min_bytes: usize,
                        max_bytes: usize) -> Promise<(Self, T, usize)> where T: DerefMut<Target=[u8]> {
        if start >= min_bytes {
            return Promise::fulfilled((self, buf, start));
        } else {
            with_current_event_loop(move |event_loop| {
                let promise =
                    event_loop.event_port.borrow_mut().handler.observers[self.token].when_becomes_readable();
                return promise.then(move |()| {
                    use mio::TryRead;
                    let n = try!(self.stream.read_slice(&mut buf[start..])).unwrap();
                    return Ok(self.read_internal(buf, start + n, min_bytes, max_bytes));
                });
            })
        }
    }

    fn write_internal<T>(mut self,
                         buf: T,
                         cursor: usize) -> Promise<(Self, T)> where T: Deref<Target=[u8]> {
        if cursor == buf.len() {
            return Promise::fulfilled((self, buf));
        } else {
            with_current_event_loop(move |event_loop| {
                let promise =
                    event_loop.event_port.borrow_mut().handler.observers[self.token].when_becomes_writable();
                return promise.then(move |()| {
                    use mio::TryWrite;
                    let n = try!(self.stream.write_slice(&buf[cursor..])).unwrap();
                    return Ok(self.write_internal(buf, cursor + n));
                });
            })
        }
    }
}

impl AsyncRead for AsyncIoStream {
    fn read<T>(self, buf: T,
               min_bytes: usize,
               max_bytes: usize) -> Promise<(Self, T, usize)> where T: DerefMut<Target=[u8]> {
        self.read_internal(buf, 0, min_bytes, max_bytes)
    }
}

impl AsyncWrite for AsyncIoStream {
    fn write<T>(self, buf: T) -> Promise<(Self, T)> where T: Deref<Target=[u8]> {
        self.write_internal(buf, 0)
    }
}

pub struct FdObserver {
    read_fulfiller: Option<Box<PromiseFulfiller<()>>>,
    write_fulfiller: Option<Box<PromiseFulfiller<()>>>,
}

impl FdObserver {
    pub fn new() -> ::mio::Token {
        with_current_event_loop(move |event_loop| {

            let observer = FdObserver { read_fulfiller: None, write_fulfiller: None };
            let event_port = &mut *event_loop.event_port.borrow_mut();
            let token = match event_port.handler.observers.insert(observer) {
                Ok(tok) => tok,
                Err(_obs) => panic!("could not insert observer."),
            };
            return token;
        })
    }

    pub fn when_becomes_readable(&mut self) -> Promise<()> {
        let (promise, fulfiller) = new_promise_and_fulfiller();
        self.read_fulfiller = Some(fulfiller);
        return promise;
    }

    pub fn when_becomes_writable(&mut self) -> Promise<()> {
        let (promise, fulfiller) = new_promise_and_fulfiller();
        self.write_fulfiller = Some(fulfiller);
        return promise;
    }
}

pub struct MioEventPort {
    handler: Handler,
    reactor: ::mio::EventLoop<Handler>,
}

pub struct Handler {
    observers: Slab<FdObserver>,
}

impl MioEventPort {
    pub fn new() -> Result<MioEventPort> {
        Ok(MioEventPort {
            handler: Handler { observers: Slab::new(100) },
            reactor: try!(::mio::EventLoop::new()),
        })
    }
}

impl ::mio::Handler for Handler {
    type Timeout = ();
    type Message = ();
    fn readable(&mut self, _event_loop: &mut ::mio::EventLoop<Handler>,
                token: ::mio::Token, _hint: ::mio::ReadHint) {
        match ::std::mem::replace(&mut self.observers[token].read_fulfiller, None) {
            Some(fulfiller) => {
                fulfiller.fulfill(())
            }
            None => {
                ()
            }
        }
    }
    fn writable(&mut self, _event_loop: &mut ::mio::EventLoop<Handler>, token: ::mio::Token) {
        match ::std::mem::replace(&mut self.observers[token].write_fulfiller, None) {
            Some(fulfiller) => fulfiller.fulfill(()),
            None => (),
        }
    }
}

impl EventPort for MioEventPort {
    fn wait(&mut self) -> bool {
        self.reactor.run_once(&mut self.handler).unwrap();
        return false;
    }

    fn poll(&mut self) -> bool {
        self.reactor.run_once(&mut self.handler).unwrap();
        return false;
    }
}
<|endoftext|>


<commit_before>use std::sync::Arc;

use wayland_backend::server::ObjectData;

use crate::{Client, DisplayHandle, Resource};

pub trait Dispatch<I: Resource>: Sized {
    type UserData: DestructionNotify + Send + Sync + 'static;

    fn request(
        &mut self,
        client: &Client,
        resource: &I,
        request: I::Request,
        data: &Self::UserData,
        dhandle: &mut DisplayHandle<'_, Self>,
        data_init: &mut DataInit<'_, Self>,
    );
}

pub trait DestructionNotify {
    #[cfg(not(tarpaulin_include))]
    fn object_destroyed(&self) {}
}

impl DestructionNotify for () {}

#[derive(Debug)]
pub struct ResourceData<I, U> {
    marker: std::marker::PhantomData<fn(I)>,
    pub udata: U,
}

#[derive(Debug)]
pub struct New<I> {
    id: I,
}

impl<I> New<I> {
    pub fn wrap(id: I) -> New<I> {
        New { id }
    }
}

#[derive(Debug)]
pub struct DataInit<'a, D> {
    store: &'a mut Option<Arc<dyn ObjectData<D>>>,
}

impl<'a, D> DataInit<'a, D> {
    pub fn init<I: Resource + 'static>(
        &mut self,
        resource: New<I>,
        data: <D as Dispatch<I>>::UserData,
    ) -> I
    where
        D: Dispatch<I>,
    {
        *self.store = Some(Arc::new(ResourceData::<I, _>::new(data)));
        resource.id
    }
}

impl<I, U> ResourceData<I, U> {
    pub(crate) fn new(udata: U) -> Self {
        ResourceData { marker: std::marker::PhantomData, udata }
    }
}

impl<
        I: Resource + 'static,
        U: DestructionNotify + Send + Sync + 'static,
        D: Dispatch<I, UserData = U>,
    > ObjectData<D> for ResourceData<I, U>
{
    fn request(
        self: Arc<Self>,
        handle: &mut wayland_backend::server::Handle<D>,
        data: &mut D,
        client_id: wayland_backend::server::ClientId,
        msg: wayland_backend::protocol::Message<wayland_backend::server::ObjectId>,
    ) -> Option<Arc<dyn ObjectData<D>>> {
        let mut dhandle = DisplayHandle::from_handle(handle);
        let client = match Client::from_id(&mut dhandle, client_id) {
            Ok(v) => v,
            Err(_) => {
                log::error!("Receiving a request from a dead client ?!");
                return None;
            }
        };

        let (resource, request) = match I::parse_request(&mut dhandle, msg) {
            Ok(v) => v,
            Err(e) => {
                log::warn!("Dispatching error encountered: {:?}, killing client.", e);
                // TODO: Kill client
                return None;
            }
        };
        let udata = resource.data::<U>().expect("Wrong user_data value for object");

        let mut new_data = None;

        data.request(
            &client,
            &resource,
            request,
            udata,
            &mut dhandle,
            &mut DataInit { store: &mut new_data },
        );

        new_data
    }

    fn destroyed(
        &self,
        _: wayland_backend::server::ClientId,
        _: wayland_backend::server::ObjectId,
    ) {
        self.udata.object_destroyed()
    }
}
<commit_msg>server: add dispatch delegation traits and docs<commit_after>use std::sync::Arc;

use wayland_backend::server::ObjectData;

use crate::{Client, DisplayHandle, Resource};

/// A trait which provides an implementation for handling a client's requests from a resource with some type
/// of associated user data.
pub trait Dispatch<I: Resource>: Sized {
    /// The user data associated with the type of resource.
    type UserData: DestructionNotify + Send + Sync + 'static;

    /// Called when a request from a client is processed.
    ///
    /// The implementation of this function will vary depending on what protocol is being implemented. Typically
    /// the server may respond to clients by sending events to the resource, or some other resource stored in
    /// the user data.
    fn request(
        &mut self,
        client: &Client,
        resource: &I,
        request: I::Request,
        data: &Self::UserData,
        dhandle: &mut DisplayHandle<'_, Self>,
        data_init: &mut DataInit<'_, Self>,
    );
}

/// A trait to be implemented on user data, which indicates when an object has been destroyed by a client.
pub trait DestructionNotify {
    /// Called when the object this user data is associated with has been destroyed.
    ///
    /// Note this type only provides an immutable reference, you will need to use interior mutability to change
    /// the inside of the object.
    ///
    /// Typically a [`Mutex`](std::sync::Mutex) would be used to have interior mutability.
    #[cfg(not(tarpaulin_include))]
    fn object_destroyed(&self) {}
}

impl DestructionNotify for () {}

#[derive(Debug)]
pub struct ResourceData<I, U> {
    marker: std::marker::PhantomData<fn(I)>,
    pub udata: U,
}

#[derive(Debug)]
pub struct New<I> {
    id: I,
}

impl<I> New<I> {
    pub fn wrap(id: I) -> New<I> {
        New { id }
    }
}

#[derive(Debug)]
pub struct DataInit<'a, D> {
    store: &'a mut Option<Arc<dyn ObjectData<D>>>,
}

impl<'a, D> DataInit<'a, D> {
    pub fn init<I: Resource + 'static>(
        &mut self,
        resource: New<I>,
        data: <D as Dispatch<I>>::UserData,
    ) -> I
    where
        D: Dispatch<I>,
    {
        *self.store = Some(Arc::new(ResourceData::<I, _>::new(data)));
        resource.id
    }
}

/*
 * Dispatch delegation helpers.
 */

/// The base trait used to define a delegate type to hand some type of resource.
pub trait DelegateDispatchBase<I: Resource> {
    /// The type of user data the delegate holds.
    type UserData: DestructionNotify + Send + Sync + 'static;
}

/// A trait which defines a delegate to handle some type of resource.
///
/// This trait is useful for building modular handlers of resources.
pub trait DelegateDispatch<
    I: Resource,
    D: Dispatch<I, UserData = <Self as DelegateDispatchBase<I>>::UserData>,
>: Sized + DelegateDispatchBase<I>
{
    /// Called when a request from a client is processed.
    ///
    /// The implementation of this function will vary depending on what protocol is being implemented. Typically
    /// the server may respond to clients by sending events to the resource, or some other resource stored in
    /// the user data.
    fn request(
        &mut self,
        client: &Client,
        resource: &I,
        request: I::Request,
        data: &Self::UserData,
        dhandle: &mut DisplayHandle<D>,
        data_init: &mut DataInit<'_, D>,
    );
}

impl<I, U> ResourceData<I, U> {
    pub(crate) fn new(udata: U) -> Self {
        ResourceData { marker: std::marker::PhantomData, udata }
    }
}

impl<
        I: Resource + 'static,
        U: DestructionNotify + Send + Sync + 'static,
        D: Dispatch<I, UserData = U>,
    > ObjectData<D> for ResourceData<I, U>
{
    fn request(
        self: Arc<Self>,
        handle: &mut wayland_backend::server::Handle<D>,
        data: &mut D,
        client_id: wayland_backend::server::ClientId,
        msg: wayland_backend::protocol::Message<wayland_backend::server::ObjectId>,
    ) -> Option<Arc<dyn ObjectData<D>>> {
        let mut dhandle = DisplayHandle::from_handle(handle);
        let client = match Client::from_id(&mut dhandle, client_id) {
            Ok(v) => v,
            Err(_) => {
                log::error!("Receiving a request from a dead client ?!");
                return None;
            }
        };

        let (resource, request) = match I::parse_request(&mut dhandle, msg) {
            Ok(v) => v,
            Err(e) => {
                log::warn!("Dispatching error encountered: {:?}, killing client.", e);
                // TODO: Kill client
                return None;
            }
        };
        let udata = resource.data::<U>().expect("Wrong user_data value for object");

        let mut new_data = None;

        data.request(
            &client,
            &resource,
            request,
            udata,
            &mut dhandle,
            &mut DataInit { store: &mut new_data },
        );

        new_data
    }

    fn destroyed(
        &self,
        _: wayland_backend::server::ClientId,
        _: wayland_backend::server::ObjectId,
    ) {
        self.udata.object_destroyed()
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
// Copyright 2017 Paul Dicker.
// Copyright 2014-2017 Melissa O'Neill and PCG Project contributors
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! PCG random number generators

// This is the default multiplier used by PCG for 64-bit state.
const MULTIPLIER: u128 = 0x2360_ED05_1FC6_5DA4_4385_DF64_9FCC_F645;

use core::fmt;
use core::mem::transmute;
use rand_core::{RngCore, SeedableRng, Error, le};
#[cfg(feature="serde1")] use serde::{Serialize, Deserialize};

/// A PCG random number generator (XSL RR 128/64 (LCG) variant).
///
/// Permuted Congruential Generator with 128-bit state, internal Linear
/// Congruential Generator, and 64-bit output via "xorshift low (bits),
/// random rotation" output function.
///
/// This is a 128-bit LCG with explicitly chosen stream with the PCG-XSL-RR
/// output function. This combination is the standard `pcg64`.
///
/// Despite the name, this implementation uses 32 bytes (256 bit) space
/// comprising 128 bits of state and 128 bits stream selector. These are both
/// set by `SeedableRng`, using a 256-bit seed.
#[derive(Clone)]
#[cfg_attr(feature="serde1", derive(Serialize,Deserialize))]
pub struct Lcg128Xsl64 {
    state: u128,
    increment: u128,
}

/// `Lcg128Xsl64` is also officially known as `pcg64`.
pub type Pcg64 = Lcg128Xsl64;

impl Lcg128Xsl64 {
    /// Construct an instance compatible with PCG seed and stream.
    ///
    /// Note that PCG specifies default values for both parameters:
    ///
    /// - `state = 0xcafef00dd15ea5e5`
    /// - `stream = 0xa02bdbf7bb3c0a7ac28fa16a64abf96`
    pub fn new(state: u128, stream: u128) -> Self {
        // The increment must be odd, hence we discard one bit:
        let increment = (stream << 1) | 1;
        Lcg128Xsl64::from_state_incr(state, increment)
    }

    #[inline]
    fn from_state_incr(state: u128, increment: u128) -> Self {
        let mut pcg = Lcg128Xsl64 { state, increment };
        // Move away from inital value:
        pcg.state = pcg.state.wrapping_add(pcg.increment);
        pcg.step();
        pcg
    }

    #[inline]
    fn step(&mut self) {
        // prepare the LCG for the next round
        self.state = self.state
            .wrapping_mul(MULTIPLIER)
            .wrapping_add(self.increment);
    }
}

// Custom Debug implementation that does not expose the internal state
impl fmt::Debug for Lcg128Xsl64 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Lcg128Xsl64 {{}}")
    }
}

/// We use a single 255-bit seed to initialise the state and select a stream.
/// One `seed` bit (lowest bit of `seed[8]`) is ignored.
impl SeedableRng for Lcg128Xsl64 {
    type Seed = [u8; 32];

    fn from_seed(seed: Self::Seed) -> Self {
        let mut seed_u64 = [0u64; 4];
        le::read_u64_into(&seed, &mut seed_u64);
        let state = u128::from(seed_u64[0]) | (u128::from(seed_u64[1]) << 64);
        let incr = u128::from(seed_u64[2]) | (u128::from(seed_u64[3]) << 64);

        // The increment must be odd, hence we discard one bit:
        Lcg128Xsl64::from_state_incr(state, incr | 1)
    }
}

impl RngCore for Lcg128Xsl64 {
    #[inline]
    fn next_u32(&mut self) -> u32 {
        self.next_u64() as u32
    }

    #[inline]
    fn next_u64(&mut self) -> u64 {
        self.step();
        output_xsl_rr(self.state)
    }

    #[inline]
    fn fill_bytes(&mut self, dest: &mut [u8]) {
        fill_bytes_impl(self, dest)
    }

    #[inline]
    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.fill_bytes(dest);
        Ok(())
    }
}

/// A PCG random number generator (XSL 128/64 (MCG) variant).
///
/// Permuted Congruential Generator with 128-bit state, internal Multiplicative
/// Congruential Generator, and 64-bit output via "xorshift low (bits),
/// random rotation" output function.
///
/// This is a 128-bit MCG with the PCG-XSL-RR output function, also known as
/// `pcg64_fast`.
/// Note that compared to the standard `pcg64` (128-bit LCG with PCG-XSL-RR
/// output function), this RNG is faster, also has a long cycle, and still has
/// good performance on statistical tests.
#[derive(Clone)]
#[cfg_attr(feature="serde1", derive(Serialize,Deserialize))]
pub struct Mcg128Xsl64 {
    state: u128,
}

/// A friendly name for `Mcg128Xsl64` (also known as `pcg64_fast`).
pub type Pcg64Mcg = Mcg128Xsl64;

impl Mcg128Xsl64 {
    /// Construct an instance compatible with PCG seed.
    ///
    /// Note that PCG specifies a default value for the parameter:
    ///
    /// - `state = 0xcafef00dd15ea5e5`
    pub fn new(state: u128) -> Self {
        // Force low bit to 1, as in C version (C++ uses `state | 3` instead).
        Mcg128Xsl64 { state: state | 1 }
    }
}

// Custom Debug implementation that does not expose the internal state
impl fmt::Debug for Mcg128Xsl64 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Mcg128Xsl64 {{}}")
    }
}

/// We use a single 126-bit seed to initialise the state and select a stream.
/// Two `seed` bits (lowest order of last byte) are ignored.
impl SeedableRng for Mcg128Xsl64 {
    type Seed = [u8; 16];

    fn from_seed(seed: Self::Seed) -> Self {
        // Read as if a little-endian u128 value:
        let mut seed_u64 = [0u64; 2];
        le::read_u64_into(&seed, &mut seed_u64);
        let state = u128::from(seed_u64[0])  |
                    u128::from(seed_u64[1]) << 64;
        Mcg128Xsl64::new(state)
    }
}

impl RngCore for Mcg128Xsl64 {
    #[inline]
    fn next_u32(&mut self) -> u32 {
        self.next_u64() as u32
    }

    #[inline]
    fn next_u64(&mut self) -> u64 {
        self.state = self.state.wrapping_mul(MULTIPLIER);
        output_xsl_rr(self.state)
    }

    #[inline]
    fn fill_bytes(&mut self, dest: &mut [u8]) {
        fill_bytes_impl(self, dest)
    }

    #[inline]
    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.fill_bytes(dest);
        Ok(())
    }
}

#[inline(always)]
fn output_xsl_rr(state: u128) -> u64 {
    // Output function XSL RR ("xorshift low (bits), random rotation")
    // Constants are for 128-bit state, 64-bit output
    const XSHIFT: u32 = 64;     // (128 - 64 + 64) / 2
    const ROTATE: u32 = 122;    // 128 - 6

    let rot = (state >> ROTATE) as u32;
    let xsl = ((state >> XSHIFT) as u64) ^ (state as u64);
    xsl.rotate_right(rot)
}

#[inline(always)]
fn fill_bytes_impl<R: RngCore + ?Sized>(rng: &mut R, dest: &mut [u8]) {
    let mut left = dest;
    while left.len() >= 8 {
        let (l, r) = {left}.split_at_mut(8);
        left = r;
        let chunk: [u8; 8] = unsafe {
            transmute(rng.next_u64().to_le())
        };
        l.copy_from_slice(&chunk);
    }
    let n = left.len();
    if n > 0 {
        let chunk: [u8; 8] = unsafe {
            transmute(rng.next_u64().to_le())
        };
        left.copy_from_slice(&chunk[..n]);
    }
}
<commit_msg>rand_pcg: Replace transmutes with safe code<commit_after>// Copyright 2018 Developers of the Rand project.
// Copyright 2017 Paul Dicker.
// Copyright 2014-2017 Melissa O'Neill and PCG Project contributors
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! PCG random number generators

// This is the default multiplier used by PCG for 64-bit state.
const MULTIPLIER: u128 = 0x2360_ED05_1FC6_5DA4_4385_DF64_9FCC_F645;

use core::fmt;
use rand_core::{RngCore, SeedableRng, Error, le};
#[cfg(feature="serde1")] use serde::{Serialize, Deserialize};

/// A PCG random number generator (XSL RR 128/64 (LCG) variant).
///
/// Permuted Congruential Generator with 128-bit state, internal Linear
/// Congruential Generator, and 64-bit output via "xorshift low (bits),
/// random rotation" output function.
///
/// This is a 128-bit LCG with explicitly chosen stream with the PCG-XSL-RR
/// output function. This combination is the standard `pcg64`.
///
/// Despite the name, this implementation uses 32 bytes (256 bit) space
/// comprising 128 bits of state and 128 bits stream selector. These are both
/// set by `SeedableRng`, using a 256-bit seed.
#[derive(Clone)]
#[cfg_attr(feature="serde1", derive(Serialize,Deserialize))]
pub struct Lcg128Xsl64 {
    state: u128,
    increment: u128,
}

/// `Lcg128Xsl64` is also officially known as `pcg64`.
pub type Pcg64 = Lcg128Xsl64;

impl Lcg128Xsl64 {
    /// Construct an instance compatible with PCG seed and stream.
    ///
    /// Note that PCG specifies default values for both parameters:
    ///
    /// - `state = 0xcafef00dd15ea5e5`
    /// - `stream = 0xa02bdbf7bb3c0a7ac28fa16a64abf96`
    pub fn new(state: u128, stream: u128) -> Self {
        // The increment must be odd, hence we discard one bit:
        let increment = (stream << 1) | 1;
        Lcg128Xsl64::from_state_incr(state, increment)
    }

    #[inline]
    fn from_state_incr(state: u128, increment: u128) -> Self {
        let mut pcg = Lcg128Xsl64 { state, increment };
        // Move away from inital value:
        pcg.state = pcg.state.wrapping_add(pcg.increment);
        pcg.step();
        pcg
    }

    #[inline]
    fn step(&mut self) {
        // prepare the LCG for the next round
        self.state = self.state
            .wrapping_mul(MULTIPLIER)
            .wrapping_add(self.increment);
    }
}

// Custom Debug implementation that does not expose the internal state
impl fmt::Debug for Lcg128Xsl64 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Lcg128Xsl64 {{}}")
    }
}

/// We use a single 255-bit seed to initialise the state and select a stream.
/// One `seed` bit (lowest bit of `seed[8]`) is ignored.
impl SeedableRng for Lcg128Xsl64 {
    type Seed = [u8; 32];

    fn from_seed(seed: Self::Seed) -> Self {
        let mut seed_u64 = [0u64; 4];
        le::read_u64_into(&seed, &mut seed_u64);
        let state = u128::from(seed_u64[0]) | (u128::from(seed_u64[1]) << 64);
        let incr = u128::from(seed_u64[2]) | (u128::from(seed_u64[3]) << 64);

        // The increment must be odd, hence we discard one bit:
        Lcg128Xsl64::from_state_incr(state, incr | 1)
    }
}

impl RngCore for Lcg128Xsl64 {
    #[inline]
    fn next_u32(&mut self) -> u32 {
        self.next_u64() as u32
    }

    #[inline]
    fn next_u64(&mut self) -> u64 {
        self.step();
        output_xsl_rr(self.state)
    }

    #[inline]
    fn fill_bytes(&mut self, dest: &mut [u8]) {
        fill_bytes_impl(self, dest)
    }

    #[inline]
    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.fill_bytes(dest);
        Ok(())
    }
}

/// A PCG random number generator (XSL 128/64 (MCG) variant).
///
/// Permuted Congruential Generator with 128-bit state, internal Multiplicative
/// Congruential Generator, and 64-bit output via "xorshift low (bits),
/// random rotation" output function.
///
/// This is a 128-bit MCG with the PCG-XSL-RR output function, also known as
/// `pcg64_fast`.
/// Note that compared to the standard `pcg64` (128-bit LCG with PCG-XSL-RR
/// output function), this RNG is faster, also has a long cycle, and still has
/// good performance on statistical tests.
#[derive(Clone)]
#[cfg_attr(feature="serde1", derive(Serialize,Deserialize))]
pub struct Mcg128Xsl64 {
    state: u128,
}

/// A friendly name for `Mcg128Xsl64` (also known as `pcg64_fast`).
pub type Pcg64Mcg = Mcg128Xsl64;

impl Mcg128Xsl64 {
    /// Construct an instance compatible with PCG seed.
    ///
    /// Note that PCG specifies a default value for the parameter:
    ///
    /// - `state = 0xcafef00dd15ea5e5`
    pub fn new(state: u128) -> Self {
        // Force low bit to 1, as in C version (C++ uses `state | 3` instead).
        Mcg128Xsl64 { state: state | 1 }
    }
}

// Custom Debug implementation that does not expose the internal state
impl fmt::Debug for Mcg128Xsl64 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Mcg128Xsl64 {{}}")
    }
}

/// We use a single 126-bit seed to initialise the state and select a stream.
/// Two `seed` bits (lowest order of last byte) are ignored.
impl SeedableRng for Mcg128Xsl64 {
    type Seed = [u8; 16];

    fn from_seed(seed: Self::Seed) -> Self {
        // Read as if a little-endian u128 value:
        let mut seed_u64 = [0u64; 2];
        le::read_u64_into(&seed, &mut seed_u64);
        let state = u128::from(seed_u64[0])  |
                    u128::from(seed_u64[1]) << 64;
        Mcg128Xsl64::new(state)
    }
}

impl RngCore for Mcg128Xsl64 {
    #[inline]
    fn next_u32(&mut self) -> u32 {
        self.next_u64() as u32
    }

    #[inline]
    fn next_u64(&mut self) -> u64 {
        self.state = self.state.wrapping_mul(MULTIPLIER);
        output_xsl_rr(self.state)
    }

    #[inline]
    fn fill_bytes(&mut self, dest: &mut [u8]) {
        fill_bytes_impl(self, dest)
    }

    #[inline]
    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.fill_bytes(dest);
        Ok(())
    }
}

#[inline(always)]
fn output_xsl_rr(state: u128) -> u64 {
    // Output function XSL RR ("xorshift low (bits), random rotation")
    // Constants are for 128-bit state, 64-bit output
    const XSHIFT: u32 = 64;     // (128 - 64 + 64) / 2
    const ROTATE: u32 = 122;    // 128 - 6

    let rot = (state >> ROTATE) as u32;
    let xsl = ((state >> XSHIFT) as u64) ^ (state as u64);
    xsl.rotate_right(rot)
}

#[inline(always)]
fn fill_bytes_impl<R: RngCore + ?Sized>(rng: &mut R, dest: &mut [u8]) {
    let mut left = dest;
    while left.len() >= 8 {
        let (l, r) = {left}.split_at_mut(8);
        left = r;
        let chunk: [u8; 8] = rng.next_u64().to_le_bytes();
        l.copy_from_slice(&chunk);
    }
    let n = left.len();
    if n > 0 {
        let chunk: [u8; 8] = rng.next_u64().to_le_bytes();
        left.copy_from_slice(&chunk[..n]);
    }
}
<|endoftext|>


<commit_before>//! Additional functionality for numerics.
//!
//! This module provides some extra types that are useful when doing numerical
//! work. See the individual documentation for each piece for more information.

#![stable(feature = "rust1", since = "1.0.0")]
#![allow(missing_docs)]

#[cfg(test)]
mod tests;

#[cfg(test)]
mod benches;

#[stable(feature = "rust1", since = "1.0.0")]
pub use core::num::Wrapping;
#[stable(feature = "rust1", since = "1.0.0")]
pub use core::num::{FpCategory, ParseFloatError, ParseIntError, TryFromIntError};

#[stable(feature = "signed_nonzero", since = "1.34.0")]
pub use core::num::{NonZeroI128, NonZeroI16, NonZeroI32, NonZeroI64, NonZeroI8, NonZeroIsize};
#[stable(feature = "nonzero", since = "1.28.0")]
pub use core::num::{NonZeroU128, NonZeroU16, NonZeroU32, NonZeroU64, NonZeroU8, NonZeroUsize};

 #[stable(feature = "int_error_matching", since = "1.54.0")]
pub use core::num::IntErrorKind;

#[cfg(test)]
use crate::fmt;
#[cfg(test)]
use crate::ops::{Add, Div, Mul, Rem, Sub};

/// Helper function for testing numeric operations
#[cfg(test)]
pub fn test_num<T>(ten: T, two: T)
where
    T: PartialEq
        + Add<Output = T>
        + Sub<Output = T>
        + Mul<Output = T>
        + Div<Output = T>
        + Rem<Output = T>
        + fmt::Debug
        + Copy,
{
    assert_eq!(ten.add(two), ten + two);
    assert_eq!(ten.sub(two), ten - two);
    assert_eq!(ten.mul(two), ten * two);
    assert_eq!(ten.div(two), ten / two);
    assert_eq!(ten.rem(two), ten % two);
}
<commit_msg>run tidy<commit_after>//! Additional functionality for numerics.
//!
//! This module provides some extra types that are useful when doing numerical
//! work. See the individual documentation for each piece for more information.

#![stable(feature = "rust1", since = "1.0.0")]
#![allow(missing_docs)]

#[cfg(test)]
mod tests;

#[cfg(test)]
mod benches;

#[stable(feature = "rust1", since = "1.0.0")]
pub use core::num::Wrapping;
#[stable(feature = "rust1", since = "1.0.0")]
pub use core::num::{FpCategory, ParseFloatError, ParseIntError, TryFromIntError};

#[stable(feature = "signed_nonzero", since = "1.34.0")]
pub use core::num::{NonZeroI128, NonZeroI16, NonZeroI32, NonZeroI64, NonZeroI8, NonZeroIsize};
#[stable(feature = "nonzero", since = "1.28.0")]
pub use core::num::{NonZeroU128, NonZeroU16, NonZeroU32, NonZeroU64, NonZeroU8, NonZeroUsize};

#[stable(feature = "int_error_matching", since = "1.54.0")]
pub use core::num::IntErrorKind;

#[cfg(test)]
use crate::fmt;
#[cfg(test)]
use crate::ops::{Add, Div, Mul, Rem, Sub};

/// Helper function for testing numeric operations
#[cfg(test)]
pub fn test_num<T>(ten: T, two: T)
where
    T: PartialEq
        + Add<Output = T>
        + Sub<Output = T>
        + Mul<Output = T>
        + Div<Output = T>
        + Rem<Output = T>
        + fmt::Debug
        + Copy,
{
    assert_eq!(ten.add(two), ten + two);
    assert_eq!(ten.sub(two), ten - two);
    assert_eq!(ten.mul(two), ten * two);
    assert_eq!(ten.div(two), ten / two);
    assert_eq!(ten.rem(two), ten % two);
}
<|endoftext|>


<commit_before>//! This module implements an efficient transposition table.

use std;
use std::cell::UnsafeCell;
use std::mem::transmute;

#[derive(Copy, Clone)]
pub struct EntryData {
    move16: u16,
    value: i16,
    eval_value: i16,
    gen_bound: u8,
    depth: u8,
}

impl EntryData {
    #[inline(always)]
    pub fn move16(&self) -> u16 {
        self.move16
    }

    #[inline(always)]
    pub fn value(&self) -> i16 {
        self.value
    }

    #[inline(always)]
    pub fn eval_value(&self) -> i16 {
        self.eval_value
    }

    #[inline(always)]
    pub fn depth(&self) -> u8 {
        self.depth
    }

    #[inline(always)]
    pub fn bound(&self) -> u8 {
        self.gen_bound & 0b11
    }
}

// Represents transposition table entry.
//
// It is 16 bytes, defined as below:
//
// * key        64 bit
// * move16     16 bit
// * value      16 bit
// * eval value 16 bit
// * generation  6 bit
// * bound type  2 bit
// * depth       8 bit
#[derive(Copy, Clone)]
struct Entry {
    key: u64,
    data: EntryData,
}

impl Default for Entry {
    fn default() -> Entry {
        Entry {
            key: 0,
            data: unsafe { transmute(0u64) },
        }
    }
}

impl Entry {
    fn new(key: u64,
           value: i16,
           bound: u8,
           depth: u8,
           move16: u16,
           generation: u8,
           eval_value: i16)
           -> Entry {
        Entry {
            key: key,
            data: EntryData {
                move16: move16,
                value: value,
                eval_value: eval_value,
                gen_bound: generation | bound,
                depth: depth, // TODO: DEPTH_NONE?
            },
        }
    }

    #[inline(always)]
    fn key(&self) -> u64 {
        self.key
    }

    #[inline(always)]
    fn data(&self) -> u64 {
        unsafe { transmute(self.data) }
    }

    #[inline(always)]
    fn generation(&self) -> u8 {
        self.data.gen_bound & 0b11111100
    }
}

#[derive(Default, Clone)]
struct Cluster {
    entries: [Entry; 4],
}

pub struct TranspositionTable {
    generation: u8,
    cluster_count: usize,
    table: UnsafeCell<Vec<Cluster>>,
}

impl TranspositionTable {
    pub fn new() -> TranspositionTable {
        TranspositionTable {
            generation: 0,
            cluster_count: 1,
            table: UnsafeCell::new(vec![Default::default()]),
        }
    }

    pub fn new_search(&mut self) {
        self.generation += 0b100;  // Lower 2 bits are used by bound type
    }

    pub fn resize(&mut self, size_mb: usize) {
        // The cluster count should be in the form "2**n", so the best
        // we can do is to ensure that the new cluster count will be
        // as close as possible but no greater than the requested one.
        let requested_cluster_count = (size_mb * 1024 * 1024) / std::mem::size_of::<Cluster>();

        // First we make sure `requested_cluster_count` is exceeded.
        let mut new_cluster_count = 1;
        while requested_cluster_count >= new_cluster_count {
            new_cluster_count <<= 1;
        }

        // Then we make one step back, making sure that the new
        // cluster count is not zero.
        if new_cluster_count > 1 {
            new_cluster_count >>= 1;
        } else {
            new_cluster_count = 1;
        }
        if new_cluster_count != self.cluster_count {
            self.cluster_count = new_cluster_count;
            self.table = UnsafeCell::new(vec![Default::default(); new_cluster_count]);
        }
    }

    pub fn clear(&mut self) {
        self.table = UnsafeCell::new(vec![Default::default(); self.cluster_count]);
    }

    pub fn probe(&self, key: u64) -> Option<EntryData> {
        let cluster = self.cluster_mut(key);
        for entry in cluster.entries.iter_mut() {
            if entry.key() ^ entry.data() == key {
                // Transposition table hit. Refresh entry's generation
                // and return the entry.
                entry.data.gen_bound = self.generation | entry.data.bound();
                return Some(entry.data);
            }
        }
        // Transposition table miss.
        None
    }

    pub fn store(&self, key: u64, mut data: EntryData) {
        let mut entries = self.cluster_mut(key).entries;
        let mut replace_index = 0;
        for i in 0..4 {
            let e = entries[i];
            if e.key() == 0 || e.key() ^ e.data() == key {
                // This is either an empty entry, or we overwrite an
                // old entry for the same key.
                if data.move16 == 0 {
                    // Preserve any existing move.
                    data.move16 = e.data.move16;
                }
                replace_index = i;
                break;
            }
            // Implement replace strategy
            // if (e.generation() == self.generation || e.data.bound() == Bound::BOUND_EXACT) -
            //    (entries[replace_index].generation() == self.generation) -
            //    (e.data.depth < entries[replace_index].data.depth) < 0 {
            //     replace_index = i;

            // }
        }
        unsafe {
            entries.get_unchecked_mut(replace_index).key = key ^ transmute::<EntryData, u64>(data);
            entries.get_unchecked_mut(replace_index).data = data;
        }
    }

    #[inline(always)]
    fn cluster_mut(&self, key: u64) -> &mut Cluster {
        let cluster_index = (key & (self.cluster_count - 1) as u64) as usize;
        unsafe { (&mut *self.table.get()).get_unchecked_mut(cluster_index) }
    }

    // void TranspositionTable::store(const Key key, Value v, Bound b, Depth d, Move m, Value statV) {

    //   TTEntry *tte, *replace;
    //   uint16_t key16 = key >> 48; // Use the high 16 bits as key inside the cluster

    // tte = replace = first_entry(key);

    //   for (unsigned i = 0; i < TTClusterSize; ++i, ++tte)
    //   {
    //       if (!tte->key16 || tte->key16 == key16) // Empty or overwrite old
    //       {
    //           if (!m)
    //               m = tte->move(); // Preserve any existing ttMove

    //           replace = tte;
    //           break;
    //       }

    //       // Implement replace strategy
    //       if (  ((    tte->genBound8 & 0xFC) == generation || tte->bound() == BOUND_EXACT)
    //           - ((replace->genBound8 & 0xFC) == generation)
    //           - (tte->depth8 < replace->depth8) < 0)
    //           replace = tte;
    //   }

    //   replace->save(key16, v, b, d, m, generation, statV);
    // }

    // const TTEntry* TranspositionTable::probe(const Key key) const {

    //   TTEntry* tte = first_entry(key);
    //   uint16_t key16 = key >> 48;

    //   for (unsigned i = 0; i < TTClusterSize; ++i, ++tte)
    //       if (tte->key16 == key16)
    //       {
    //           tte->genBound8 = generation | tte->bound(); // Refresh
    //           return tte;
    //       }

    //   return NULL;
    // }
}

pub enum Bound {
    BOUND_NONE = 0,
    BOUND_UPPER = 1,
    BOUND_LOWER = 2,
    BOUND_EXACT = 3,
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::{Cluster, Entry};
    use std;
    use basetypes::*;

    #[test]
    fn test_cluster_size() {
        assert_eq!(std::mem::size_of::<Cluster>(), 64);
        assert_eq!(std::mem::size_of::<Entry>(), 16);
    }

    #[test]
    fn test_tt_resize() {
        let mut tt = TranspositionTable::new();
        assert_eq!(unsafe { &*tt.table.get() }.capacity(), 1);
        tt.resize(1);
        assert_eq!(unsafe { &*tt.table.get() }.capacity(),
                   1024 * 1024 / std::mem::size_of::<Cluster>());
        tt.clear();
    }
}
<commit_msg>more work on tt<commit_after>//! This module implements a database that stores results of
//! previously performed searches.
//!
//! Using such a database (called "transposition table") is a way to
//! greatly reduce the search space of a chess tree with little
//! negative impact. Chess programs, during their brute-force search,
//! encounter the same positions again and again, but from different
//! sequences of moves, which is called a transposition.
//!
//! When the search encounters a transposition, it is beneficial to
//! "remember" what was determined last time the position was
//! examined, rather than redoing the entire search again. For this
//! reason, chess programs have a transposition table, which is a
//! large hash table storing information about positions previously
//! searched, how deeply they were searched, and what we concluded
//! about them.

use std;
use std::cell::UnsafeCell;
use std::mem::transmute;

/// Stores information about a particular position.
#[derive(Copy, Clone)]
pub struct EntryData {
    move16: u16,
    value: i16,
    eval_value: i16,
    gen_bound: u8,
    depth: u8,
}

impl EntryData {
    /// Creates a new instance.
    ///
    /// The following information can be stored:
    ///
    /// * `value` -- the value assigned to the position,
    ///
    /// * `bound` -- the meaning of the assigned value ("Exact",
    ///    "Upper Bound", or "Lower Bound"),
    ///
    /// * `depth` -- the depth of search,
    ///
    /// * `move16` -- best or refutation move,
    ///
    /// * `eval_value` -- the calculated static evaluation for the
    ///    position.
    pub fn new(value: i16, bound: u8, depth: u8, move16: u16, eval_value: i16) -> EntryData {
        assert!(bound <= 0b11);
        assert!(depth < 128);
        EntryData {
            move16: move16,
            value: value,
            eval_value: eval_value,
            gen_bound: bound, // Stores the entry's generation and the bound.
            depth: depth,
        }
    }

    #[inline(always)]
    pub fn move16(&self) -> u16 {
        self.move16
    }

    #[inline(always)]
    pub fn value(&self) -> i16 {
        self.value
    }

    #[inline(always)]
    pub fn eval_value(&self) -> i16 {
        self.eval_value
    }

    #[inline(always)]
    pub fn bound(&self) -> u8 {
        self.gen_bound & 0b11
    }

    #[inline(always)]
    pub fn depth(&self) -> u8 {
        self.depth
    }
}

// Represents transposition table entry.
//
// It is 16 bytes, defined as below:
//
// * key        64 bit
// * move16     16 bit
// * value      16 bit
// * eval value 16 bit
// * generation  6 bit
// * bound type  2 bit
// * depth       8 bit
#[derive(Copy, Clone)]
struct Entry {
    key: u64,
    data: EntryData,
}

impl Default for Entry {
    fn default() -> Entry {
        Entry {
            key: 0,
            data: unsafe { transmute(0u64) },
        }
    }
}

impl Entry {
    // Returns the whole entry data as `u64`.
    //
    // This is needed for implementing the lock-less probing and
    // storing.
    #[inline(always)]
    fn data_u64(&self) -> u64 {
        unsafe { transmute(self.data) }
    }

    #[inline(always)]
    fn generation(&self) -> u8 {
        self.data.gen_bound & 0b11111100
    }

    #[inline(always)]
    fn set_generation(&mut self, generation: u8) {
        assert_eq!(generation & 0b11, 0);
        self.data.gen_bound = generation | self.data.bound()
    }
}

/// A large hash-table that stores results of previously performed
/// searches.
pub struct TranspositionTable {
    generation: u8,
    cluster_count: usize,
    table: UnsafeCell<Vec<[Entry; 4]>>,
}

impl TranspositionTable {
    /// Creates a new transposition table.
    ///
    /// The newly created table has the minimum possible size. Before
    /// using the new table for anything, `resize()` should be called
    /// on it, specifying the desired size.
    pub fn new() -> TranspositionTable {
        TranspositionTable {
            generation: 0,
            cluster_count: 1,
            table: UnsafeCell::new(vec![Default::default()]),
        }
    }

    /// Signals that a new search is about to begin.
    ///
    /// Internally, the transposition table maintains a generation
    /// counter that is used to implement an efficient replacement
    /// strategy. This method increases this counter.
    pub fn new_search(&mut self) {
        self.generation += 0b100;  // Lower 2 bits are used by bound type
        assert_eq!(self.generation & 0b11, 0);
    }

    /// Resizes the transpositon table.
    ///
    /// `size_mb` is the desired new size in Mbytes. All entries in
    /// the table will be lost.
    pub fn resize(&mut self, size_mb: usize) {
        // The cluster count should be in the form "2**n", so the best
        // we can do is to ensure that the new cluster count will be
        // as close as possible but no greater than the requested one.
        let requested_cluster_count = (size_mb * 1024 * 1024) / std::mem::size_of::<[Entry; 4]>();

        // First, make sure `requested_cluster_count` is exceeded.
        let mut new_cluster_count = 1;
        while requested_cluster_count >= new_cluster_count {
            new_cluster_count <<= 1;
        }

        // Then make one step back, being careful so that the new
        // cluster count is not zero.
        if new_cluster_count > 1 {
            new_cluster_count >>= 1;
        } else {
            new_cluster_count = 1;
        }

        // Finally, reallocate the vector of clusters.
        if new_cluster_count != self.cluster_count {
            self.cluster_count = new_cluster_count;
            self.table = UnsafeCell::new(vec![Default::default(); new_cluster_count]);
        }
    }

    /// Clears the transpositon table.
    ///
    /// All remaining entries in the table will be removed.
    pub fn clear(&mut self) {
        self.table = UnsafeCell::new(vec![Default::default(); self.cluster_count]);
    }

    /// Probes for data by a specific key.
    #[inline]
    pub fn probe(&self, key: u64) -> Option<EntryData> {
        // `probe` and `store` implement a clever lock-less hashing
        // method. Rather than to store two disjoint items, the key is
        // stored xored with data, while data is stored additionally
        // as usual.
        let cluster = unsafe { self.cluster_mut(key) };
        for entry in cluster.iter_mut() {
            if entry.key ^ entry.data_u64() == key {
                // If `key` and `data` were written simultaneously by
                // different search instances with different keys, the
                // error will yield in a mismatch of the comparison
                // (except the rare and inherent key collisions).
                entry.set_generation(self.generation);  // Update the generation.!!! bug!!!!
                return Some(entry.data);
            }
        }
        None
    }

    /// Stores data by a specific key.
    ///
    /// After being stored, the data might be retrieved by
    /// `probe(key)`. This is not guaranteed though, because in the
    /// meantime it might have been overwritten.
    #[inline]
    pub fn store(&self, key: u64, mut data: EntryData) {
        data.gen_bound |= self.generation;  // Set the generation.
        let mut cluster = unsafe { self.cluster_mut(key) };
        let mut replace_index = 0;
        let mut replace_score = 0xff;
        for (i, entry) in cluster.iter_mut().enumerate() {
            // Check if this is an empty entry, or an old entry for
            // the same key.
            if entry.key == 0 || entry.key ^ entry.data_u64() == key {
                if data.move16 == 0 {
                    data.move16 = entry.data.move16;  // Preserve any existing move.
                }
                replace_index = i;
                break;
            }
            // The replaced entry will be the entry with the lowest
            // score.
            let entry_score = self.calc_score(entry);
            if entry_score < replace_score {
                replace_index = i;
                replace_score = entry_score;
            }
        }
        unsafe {
            cluster.get_unchecked_mut(replace_index).key = key ^ transmute::<EntryData, u64>(data);
            cluster.get_unchecked_mut(replace_index).data = data;
        }
    }

    // A helper method for `store`.
    //
    // Implements our replacement strategy. Should return higher
    // values for the entries that are move likely to save CPU work.
    #[inline]
    fn calc_score(&self, entry: &Entry) -> u8 {
        (if entry.generation() == self.generation {
            128
        } else {
            0
        }) + entry.data.depth
    }

    // A helper method for `probe` and `store`.
    #[inline]
    unsafe fn cluster_mut(&self, key: u64) -> &mut [Entry; 4] {
        let cluster_index = (key & (self.cluster_count - 1) as u64) as usize;
        self.table.get().as_mut().unwrap().get_unchecked_mut(cluster_index)
    }
}

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

    #[test]
    fn test_cluster_size() {
        assert_eq!(std::mem::size_of::<[Entry; 4]>(), 64);
        assert_eq!(std::mem::size_of::<Entry>(), 16);
    }

    #[test]
    fn test_tt_resize() {
        let mut tt = TranspositionTable::new();
        assert_eq!(unsafe { &*tt.table.get() }.capacity(), 1);
        tt.resize(1);
        assert_eq!(unsafe { &*tt.table.get() }.capacity(),
                   1024 * 1024 / std::mem::size_of::<[Entry; 4]>());
        tt.clear();
    }

    #[test]
    fn test_store_and_probe() {
        let mut tt = TranspositionTable::new();
        assert!(tt.probe(1).is_none());
        let data = EntryData::new(0, 0, 100, 666, 0);
        assert_eq!(data.depth(), 100);
        assert_eq!(data.move16(), 666);
        tt.store(1, data);
        assert_eq!(tt.probe(1).unwrap().depth(), 100);
        tt.store(1, EntryData::new(0, 0, 100, 666, 0));
        assert_eq!(tt.probe(1).unwrap().depth(), 100);
        assert_eq!(tt.probe(1).unwrap().move16(), 666);
        for i in 2..100 {
            tt.store(i, EntryData::new(i as i16, 0, i as u8, i as u16, i as i16));
        }
        assert_eq!(tt.probe(1).unwrap().depth(), 100);
        assert_eq!(tt.probe(99).unwrap().depth(), 99);
        assert_eq!(tt.probe(98).unwrap().depth(), 98);
        assert_eq!(tt.probe(97).unwrap().depth(), 97);
        tt.clear();
        assert!(tt.probe(1).is_none());
        tt.store(1, data);
        tt.new_search();
        tt.probe(1);
        assert!(tt.probe(1).is_some());
    }
}
<|endoftext|>


<commit_before>use crate::syntax::Atom::{self, *};
use proc_macro2::{Literal, Span, TokenStream};
use quote::ToTokens;
use std::collections::HashSet;
use std::fmt::{self, Display};
use std::str::FromStr;
use syn::{Error, Expr, Lit, Result, Token, UnOp};

pub struct DiscriminantSet {
    repr: Option<Atom>,
    values: HashSet<Discriminant>,
    previous: Option<Discriminant>,
}

#[derive(Copy, Clone, Hash, Eq, PartialEq)]
pub struct Discriminant {
    negative: bool,
    magnitude: u32,
}

impl DiscriminantSet {
    pub fn new() -> Self {
        DiscriminantSet {
            repr: None,
            values: HashSet::new(),
            previous: None,
        }
    }

    pub fn insert(&mut self, expr: &Expr) -> Result<Discriminant> {
        let (discriminant, repr) = expr_to_discriminant(expr)?;
        self.repr = self.repr.or(repr);
        insert(self, discriminant)
    }

    pub fn insert_next(&mut self) -> Result<Discriminant> {
        let discriminant = match self.previous {
            None => Discriminant::zero(),
            Some(mut discriminant) if discriminant.negative => {
                discriminant.magnitude -= 1;
                if discriminant.magnitude == 0 {
                    discriminant.negative = false;
                }
                discriminant
            }
            Some(mut discriminant) => {
                if discriminant.magnitude == u32::MAX {
                    let msg = format!("discriminant overflow on value after {}", u32::MAX);
                    return Err(Error::new(Span::call_site(), msg));
                }
                discriminant.magnitude += 1;
                discriminant
            }
        };
        insert(self, discriminant)
    }
}

fn expr_to_discriminant(expr: &Expr) -> Result<(Discriminant, Option<Atom>)> {
    match expr {
        Expr::Lit(expr) => {
            if let Lit::Int(lit) = &expr.lit {
                let discriminant = lit.base10_parse::<Discriminant>()?;
                let repr = parse_int_suffix(lit.suffix())?;
                return Ok((discriminant, repr));
            }
        }
        Expr::Unary(unary) => {
            if let UnOp::Neg(_) = unary.op {
                let (mut discriminant, repr) = expr_to_discriminant(&unary.expr)?;
                discriminant.negative ^= true;
                return Ok((discriminant, repr));
            }
        }
        _ => {}
    }
    Err(Error::new_spanned(
        expr,
        "enums with non-integer literal discriminants are not supported yet",
    ))
}

fn insert(set: &mut DiscriminantSet, discriminant: Discriminant) -> Result<Discriminant> {
    if set.values.insert(discriminant) {
        set.previous = Some(discriminant);
        Ok(discriminant)
    } else {
        let msg = format!("discriminant value `{}` already exists", discriminant);
        Err(Error::new(Span::call_site(), msg))
    }
}

impl Discriminant {
    fn zero() -> Self {
        Discriminant {
            negative: false,
            magnitude: 0,
        }
    }
}

impl Display for Discriminant {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.negative {
            f.write_str("-")?;
        }
        Display::fmt(&self.magnitude, f)
    }
}

impl ToTokens for Discriminant {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        if self.negative {
            Token![-](Span::call_site()).to_tokens(tokens);
        }
        Literal::u32_unsuffixed(self.magnitude).to_tokens(tokens);
    }
}

impl FromStr for Discriminant {
    type Err = Error;

    fn from_str(mut s: &str) -> Result<Self> {
        let negative = s.starts_with('-');
        if negative {
            s = &s[1..];
        }
        match s.parse::<u32>() {
            Ok(magnitude) => Ok(Discriminant {
                negative,
                magnitude,
            }),
            Err(_) => Err(Error::new(
                Span::call_site(),
                "discriminant value outside of supported range",
            )),
        }
    }
}

fn parse_int_suffix(suffix: &str) -> Result<Option<Atom>> {
    if suffix.is_empty() {
        return Ok(None);
    }
    if let Some(atom) = Atom::from_str(suffix) {
        match atom {
            U8 | U16 | U32 | U64 | Usize | I8 | I16 | I32 | I64 | Isize => return Ok(Some(atom)),
            _ => {}
        }
    }
    let msg = format!("unrecognized integer suffix: `{}`", suffix);
    Err(Error::new(Span::call_site(), msg))
}
<commit_msg>Add const bound data for various discriminant reprs<commit_after>use crate::syntax::Atom::{self, *};
use proc_macro2::{Literal, Span, TokenStream};
use quote::ToTokens;
use std::collections::HashSet;
use std::fmt::{self, Display};
use std::str::FromStr;
use syn::{Error, Expr, Lit, Result, Token, UnOp};

pub struct DiscriminantSet {
    repr: Option<Atom>,
    values: HashSet<Discriminant>,
    previous: Option<Discriminant>,
}

#[derive(Copy, Clone, Hash, Eq, PartialEq)]
pub struct Discriminant {
    negative: bool,
    magnitude: u32,
}

impl DiscriminantSet {
    pub fn new() -> Self {
        DiscriminantSet {
            repr: None,
            values: HashSet::new(),
            previous: None,
        }
    }

    pub fn insert(&mut self, expr: &Expr) -> Result<Discriminant> {
        let (discriminant, repr) = expr_to_discriminant(expr)?;
        self.repr = self.repr.or(repr);
        insert(self, discriminant)
    }

    pub fn insert_next(&mut self) -> Result<Discriminant> {
        let discriminant = match self.previous {
            None => Discriminant::zero(),
            Some(mut discriminant) if discriminant.negative => {
                discriminant.magnitude -= 1;
                if discriminant.magnitude == 0 {
                    discriminant.negative = false;
                }
                discriminant
            }
            Some(mut discriminant) => {
                if discriminant.magnitude == u32::MAX {
                    let msg = format!("discriminant overflow on value after {}", u32::MAX);
                    return Err(Error::new(Span::call_site(), msg));
                }
                discriminant.magnitude += 1;
                discriminant
            }
        };
        insert(self, discriminant)
    }
}

fn expr_to_discriminant(expr: &Expr) -> Result<(Discriminant, Option<Atom>)> {
    match expr {
        Expr::Lit(expr) => {
            if let Lit::Int(lit) = &expr.lit {
                let discriminant = lit.base10_parse::<Discriminant>()?;
                let repr = parse_int_suffix(lit.suffix())?;
                return Ok((discriminant, repr));
            }
        }
        Expr::Unary(unary) => {
            if let UnOp::Neg(_) = unary.op {
                let (mut discriminant, repr) = expr_to_discriminant(&unary.expr)?;
                discriminant.negative ^= true;
                return Ok((discriminant, repr));
            }
        }
        _ => {}
    }
    Err(Error::new_spanned(
        expr,
        "enums with non-integer literal discriminants are not supported yet",
    ))
}

fn insert(set: &mut DiscriminantSet, discriminant: Discriminant) -> Result<Discriminant> {
    if set.values.insert(discriminant) {
        set.previous = Some(discriminant);
        Ok(discriminant)
    } else {
        let msg = format!("discriminant value `{}` already exists", discriminant);
        Err(Error::new(Span::call_site(), msg))
    }
}

impl Discriminant {
    const fn zero() -> Self {
        Discriminant {
            negative: false,
            magnitude: 0,
        }
    }

    const fn pos(u: u32) -> Self {
        Discriminant {
            negative: false,
            magnitude: u,
        }
    }

    const fn neg(i: i32) -> Self {
        Discriminant {
            negative: i < 0,
            // This is `i.abs() as u32` but without overflow on MIN. Uses the
            // fact that MIN.wrapping_abs() wraps back to MIN whose binary
            // representation is 1<<31, and thus the `as u32` conversion
            // produces 1<<31 too which happens to be the correct unsigned
            // magnitude.
            magnitude: i.wrapping_abs() as u32,
        }
    }
}

impl Display for Discriminant {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.negative {
            f.write_str("-")?;
        }
        Display::fmt(&self.magnitude, f)
    }
}

impl ToTokens for Discriminant {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        if self.negative {
            Token![-](Span::call_site()).to_tokens(tokens);
        }
        Literal::u32_unsuffixed(self.magnitude).to_tokens(tokens);
    }
}

impl FromStr for Discriminant {
    type Err = Error;

    fn from_str(mut s: &str) -> Result<Self> {
        let negative = s.starts_with('-');
        if negative {
            s = &s[1..];
        }
        match s.parse::<u32>() {
            Ok(magnitude) => Ok(Discriminant {
                negative,
                magnitude,
            }),
            Err(_) => Err(Error::new(
                Span::call_site(),
                "discriminant value outside of supported range",
            )),
        }
    }
}

fn parse_int_suffix(suffix: &str) -> Result<Option<Atom>> {
    if suffix.is_empty() {
        return Ok(None);
    }
    if let Some(atom) = Atom::from_str(suffix) {
        match atom {
            U8 | U16 | U32 | U64 | Usize | I8 | I16 | I32 | I64 | Isize => return Ok(Some(atom)),
            _ => {}
        }
    }
    let msg = format!("unrecognized integer suffix: `{}`", suffix);
    Err(Error::new(Span::call_site(), msg))
}

struct Bounds {
    repr: Atom,
    min: Discriminant,
    max: Discriminant,
}

const BOUNDS: [Bounds; 6] = [
    Bounds {
        repr: U8,
        min: Discriminant::zero(),
        max: Discriminant::pos(u8::MAX as u32),
    },
    Bounds {
        repr: I8,
        min: Discriminant::neg(i8::MIN as i32),
        max: Discriminant::pos(i8::MAX as u32),
    },
    Bounds {
        repr: U16,
        min: Discriminant::zero(),
        max: Discriminant::pos(u16::MAX as u32),
    },
    Bounds {
        repr: I16,
        min: Discriminant::neg(i16::MIN as i32),
        max: Discriminant::pos(i16::MAX as u32),
    },
    Bounds {
        repr: U32,
        min: Discriminant::zero(),
        max: Discriminant::pos(u32::MAX),
    },
    Bounds {
        repr: I32,
        min: Discriminant::neg(i32::MIN),
        max: Discriminant::pos(i32::MAX as u32),
    },
];
<|endoftext|>