<commit_before>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

mod msg_on_socket;

use std::io::Result;
use std::marker::PhantomData;
use std::os::unix::io::{AsRawFd, RawFd};
use std::pin::Pin;
use std::task::{Context, Poll};

use futures::Stream;
use libc::{EWOULDBLOCK, O_NONBLOCK};

use cros_async::fd_executor::add_read_waker;
use sys_util::{
    add_fd_flags, clear_fd_flags, error, handle_eintr, net::UnixSeqpacket, Error as SysError,
    ScmSocket,
};

pub use crate::msg_on_socket::*;
pub use msg_on_socket_derive::*;

/// Create a pair of socket. Request is send in one direction while response is in the other
/// direction.
pub fn pair<Request: MsgOnSocket, Response: MsgOnSocket>(
) -> Result<(MsgSocket<Request, Response>, MsgSocket<Response, Request>)> {
    let (sock1, sock2) = UnixSeqpacket::pair()?;
    let requester = MsgSocket::new(sock1);
    let responder = MsgSocket::new(sock2);
    Ok((requester, responder))
}

/// Bidirection sock that support both send and recv.
pub struct MsgSocket<I: MsgOnSocket, O: MsgOnSocket> {
    sock: UnixSeqpacket,
    _i: PhantomData<I>,
    _o: PhantomData<O>,
}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSocket<I, O> {
    // Create a new MsgSocket.
    pub fn new(s: UnixSeqpacket) -> MsgSocket<I, O> {
        MsgSocket {
            sock: s,
            _i: PhantomData,
            _o: PhantomData,
        }
    }

    // Creates an async receiver that implements `futures::Stream`.
    pub fn async_receiver(&mut self) -> MsgResult<AsyncReceiver<I, O>> {
        AsyncReceiver::new(self)
    }
}

/// One direction socket that only supports sending.
pub struct Sender<M: MsgOnSocket> {
    sock: UnixSeqpacket,
    _m: PhantomData<M>,
}

impl<M: MsgOnSocket> Sender<M> {
    /// Create a new sender sock.
    pub fn new(s: UnixSeqpacket) -> Sender<M> {
        Sender {
            sock: s,
            _m: PhantomData,
        }
    }
}

/// One direction socket that only supports receiving.
pub struct Receiver<M: MsgOnSocket> {
    sock: UnixSeqpacket,
    _m: PhantomData<M>,
}

impl<M: MsgOnSocket> Receiver<M> {
    /// Create a new receiver sock.
    pub fn new(s: UnixSeqpacket) -> Receiver<M> {
        Receiver {
            sock: s,
            _m: PhantomData,
        }
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRef<UnixSeqpacket> for MsgSocket<I, O> {
    fn as_ref(&self) -> &UnixSeqpacket {
        &self.sock
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRawFd for MsgSocket<I, O> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

impl<M: MsgOnSocket> AsRef<UnixSeqpacket> for Sender<M> {
    fn as_ref(&self) -> &UnixSeqpacket {
        &self.sock
    }
}

impl<M: MsgOnSocket> AsRawFd for Sender<M> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

impl<M: MsgOnSocket> AsRef<UnixSeqpacket> for Receiver<M> {
    fn as_ref(&self) -> &UnixSeqpacket {
        &self.sock
    }
}

impl<M: MsgOnSocket> AsRawFd for Receiver<M> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

/// Types that could send a message.
pub trait MsgSender: AsRef<UnixSeqpacket> {
    type M: MsgOnSocket;
    fn send(&self, msg: &Self::M) -> MsgResult<()> {
        let msg_size = msg.msg_size();
        let fd_size = msg.fd_count();
        let mut msg_buffer: Vec<u8> = vec![0; msg_size];
        let mut fd_buffer: Vec<RawFd> = vec![0; fd_size];

        let fd_size = msg.write_to_buffer(&mut msg_buffer, &mut fd_buffer)?;
        let sock: &UnixSeqpacket = self.as_ref();
        if fd_size == 0 {
            handle_eintr!(sock.send(&msg_buffer))
                .map_err(|e| MsgError::Send(SysError::new(e.raw_os_error().unwrap_or(0))))?;
        } else {
            sock.send_with_fds(&msg_buffer[..], &fd_buffer[0..fd_size])
                .map_err(MsgError::Send)?;
        }
        Ok(())
    }
}

/// Types that could receive a message.
pub trait MsgReceiver: AsRef<UnixSeqpacket> {
    type M: MsgOnSocket;
    fn recv(&self) -> MsgResult<Self::M> {
        let sock: &UnixSeqpacket = self.as_ref();

        let (msg_buffer, fd_buffer) = {
            if Self::M::uses_fd() {
                sock.recv_as_vec_with_fds()
                    .map_err(|e| MsgError::Recv(SysError::new(e.raw_os_error().unwrap_or(0))))?
            } else {
                (
                    sock.recv_as_vec().map_err(|e| {
                        MsgError::Recv(SysError::new(e.raw_os_error().unwrap_or(0)))
                    })?,
                    vec![],
                )
            }
        };

        if msg_buffer.len() == 0 && Self::M::fixed_size() != Some(0) {
            return Err(MsgError::RecvZero);
        }

        if let Some(fixed_size) = Self::M::fixed_size() {
            if fixed_size != msg_buffer.len() {
                return Err(MsgError::BadRecvSize {
                    expected: fixed_size,
                    actual: msg_buffer.len(),
                });
            }
        }

        // Safe because fd buffer is read from socket.
        let (v, read_fd_size) = unsafe { Self::M::read_from_buffer(&msg_buffer, &fd_buffer)? };
        if fd_buffer.len() != read_fd_size {
            return Err(MsgError::NotExpectFd);
        }
        Ok(v)
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSender for MsgSocket<I, O> {
    type M = I;
}
impl<I: MsgOnSocket, O: MsgOnSocket> MsgReceiver for MsgSocket<I, O> {
    type M = O;
}

impl<I: MsgOnSocket> MsgSender for Sender<I> {
    type M = I;
}
impl<O: MsgOnSocket> MsgReceiver for Receiver<O> {
    type M = O;
}

/// Asynchronous adaptor for `MsgSocket`.
pub struct AsyncReceiver<'a, I: MsgOnSocket, O: MsgOnSocket> {
    inner: &'a mut MsgSocket<I, O>,
    done: bool, // Have hit an error and the Stream should return null when polled.
}

impl<'a, I: MsgOnSocket, O: MsgOnSocket> AsyncReceiver<'a, I, O> {
    fn new(msg_socket: &mut MsgSocket<I, O>) -> MsgResult<AsyncReceiver<I, O>> {
        add_fd_flags(msg_socket.as_raw_fd(), O_NONBLOCK).map_err(MsgError::SettingFdFlags)?;
        Ok(AsyncReceiver {
            inner: msg_socket,
            done: false,
        })
    }
}

impl<'a, I: MsgOnSocket, O: MsgOnSocket> Drop for AsyncReceiver<'a, I, O> {
    fn drop(&mut self) {
        if let Err(e) = clear_fd_flags(self.inner.as_raw_fd(), O_NONBLOCK) {
            error!(
                "Failed to restore non-blocking behavior to message socket: {}",
                e
            );
        }
    }
}

impl<'a, I: MsgOnSocket, O: MsgOnSocket> Stream for AsyncReceiver<'a, I, O> {
    type Item = MsgResult<O>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        if self.done {
            return Poll::Ready(None);
        }

        let ret = match self.inner.recv() {
            Ok(msg) => Ok(Poll::Ready(Some(Ok(msg)))),
            Err(MsgError::Recv(e)) => {
                if e.errno() == EWOULDBLOCK {
                    add_read_waker(self.inner.as_raw_fd(), cx.waker().clone())
                        .map(|_| Poll::Pending)
                        .map_err(MsgError::AddingWaker)
                } else {
                    Err(MsgError::Recv(e))
                }
            }
            Err(e) => Err(e),
        };

        match ret {
            Ok(p) => p,
            Err(e) => {
                // Indicate something went wrong and no more events will be provided.
                self.done = true;
                Poll::Ready(Some(Err(e)))
            }
        }
    }
}
<commit_msg>msg_socket: AsyncReceiver doesn't need a mut ref<commit_after>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

mod msg_on_socket;

use std::io::Result;
use std::marker::PhantomData;
use std::os::unix::io::{AsRawFd, RawFd};
use std::pin::Pin;
use std::task::{Context, Poll};

use futures::Stream;
use libc::{EWOULDBLOCK, O_NONBLOCK};

use cros_async::fd_executor::add_read_waker;
use sys_util::{
    add_fd_flags, clear_fd_flags, error, handle_eintr, net::UnixSeqpacket, Error as SysError,
    ScmSocket,
};

pub use crate::msg_on_socket::*;
pub use msg_on_socket_derive::*;

/// Create a pair of socket. Request is send in one direction while response is in the other
/// direction.
pub fn pair<Request: MsgOnSocket, Response: MsgOnSocket>(
) -> Result<(MsgSocket<Request, Response>, MsgSocket<Response, Request>)> {
    let (sock1, sock2) = UnixSeqpacket::pair()?;
    let requester = MsgSocket::new(sock1);
    let responder = MsgSocket::new(sock2);
    Ok((requester, responder))
}

/// Bidirection sock that support both send and recv.
pub struct MsgSocket<I: MsgOnSocket, O: MsgOnSocket> {
    sock: UnixSeqpacket,
    _i: PhantomData<I>,
    _o: PhantomData<O>,
}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSocket<I, O> {
    // Create a new MsgSocket.
    pub fn new(s: UnixSeqpacket) -> MsgSocket<I, O> {
        MsgSocket {
            sock: s,
            _i: PhantomData,
            _o: PhantomData,
        }
    }

    // Creates an async receiver that implements `futures::Stream`.
    pub fn async_receiver(&mut self) -> MsgResult<AsyncReceiver<I, O>> {
        AsyncReceiver::new(self)
    }
}

/// One direction socket that only supports sending.
pub struct Sender<M: MsgOnSocket> {
    sock: UnixSeqpacket,
    _m: PhantomData<M>,
}

impl<M: MsgOnSocket> Sender<M> {
    /// Create a new sender sock.
    pub fn new(s: UnixSeqpacket) -> Sender<M> {
        Sender {
            sock: s,
            _m: PhantomData,
        }
    }
}

/// One direction socket that only supports receiving.
pub struct Receiver<M: MsgOnSocket> {
    sock: UnixSeqpacket,
    _m: PhantomData<M>,
}

impl<M: MsgOnSocket> Receiver<M> {
    /// Create a new receiver sock.
    pub fn new(s: UnixSeqpacket) -> Receiver<M> {
        Receiver {
            sock: s,
            _m: PhantomData,
        }
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRef<UnixSeqpacket> for MsgSocket<I, O> {
    fn as_ref(&self) -> &UnixSeqpacket {
        &self.sock
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRawFd for MsgSocket<I, O> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

impl<M: MsgOnSocket> AsRef<UnixSeqpacket> for Sender<M> {
    fn as_ref(&self) -> &UnixSeqpacket {
        &self.sock
    }
}

impl<M: MsgOnSocket> AsRawFd for Sender<M> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

impl<M: MsgOnSocket> AsRef<UnixSeqpacket> for Receiver<M> {
    fn as_ref(&self) -> &UnixSeqpacket {
        &self.sock
    }
}

impl<M: MsgOnSocket> AsRawFd for Receiver<M> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

/// Types that could send a message.
pub trait MsgSender: AsRef<UnixSeqpacket> {
    type M: MsgOnSocket;
    fn send(&self, msg: &Self::M) -> MsgResult<()> {
        let msg_size = msg.msg_size();
        let fd_size = msg.fd_count();
        let mut msg_buffer: Vec<u8> = vec![0; msg_size];
        let mut fd_buffer: Vec<RawFd> = vec![0; fd_size];

        let fd_size = msg.write_to_buffer(&mut msg_buffer, &mut fd_buffer)?;
        let sock: &UnixSeqpacket = self.as_ref();
        if fd_size == 0 {
            handle_eintr!(sock.send(&msg_buffer))
                .map_err(|e| MsgError::Send(SysError::new(e.raw_os_error().unwrap_or(0))))?;
        } else {
            sock.send_with_fds(&msg_buffer[..], &fd_buffer[0..fd_size])
                .map_err(MsgError::Send)?;
        }
        Ok(())
    }
}

/// Types that could receive a message.
pub trait MsgReceiver: AsRef<UnixSeqpacket> {
    type M: MsgOnSocket;
    fn recv(&self) -> MsgResult<Self::M> {
        let sock: &UnixSeqpacket = self.as_ref();

        let (msg_buffer, fd_buffer) = {
            if Self::M::uses_fd() {
                sock.recv_as_vec_with_fds()
                    .map_err(|e| MsgError::Recv(SysError::new(e.raw_os_error().unwrap_or(0))))?
            } else {
                (
                    sock.recv_as_vec().map_err(|e| {
                        MsgError::Recv(SysError::new(e.raw_os_error().unwrap_or(0)))
                    })?,
                    vec![],
                )
            }
        };

        if msg_buffer.len() == 0 && Self::M::fixed_size() != Some(0) {
            return Err(MsgError::RecvZero);
        }

        if let Some(fixed_size) = Self::M::fixed_size() {
            if fixed_size != msg_buffer.len() {
                return Err(MsgError::BadRecvSize {
                    expected: fixed_size,
                    actual: msg_buffer.len(),
                });
            }
        }

        // Safe because fd buffer is read from socket.
        let (v, read_fd_size) = unsafe { Self::M::read_from_buffer(&msg_buffer, &fd_buffer)? };
        if fd_buffer.len() != read_fd_size {
            return Err(MsgError::NotExpectFd);
        }
        Ok(v)
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSender for MsgSocket<I, O> {
    type M = I;
}
impl<I: MsgOnSocket, O: MsgOnSocket> MsgReceiver for MsgSocket<I, O> {
    type M = O;
}

impl<I: MsgOnSocket> MsgSender for Sender<I> {
    type M = I;
}
impl<O: MsgOnSocket> MsgReceiver for Receiver<O> {
    type M = O;
}

/// Asynchronous adaptor for `MsgSocket`.
pub struct AsyncReceiver<'a, I: MsgOnSocket, O: MsgOnSocket> {
    inner: &'a MsgSocket<I, O>,
    done: bool, // Have hit an error and the Stream should return null when polled.
}

impl<'a, I: MsgOnSocket, O: MsgOnSocket> AsyncReceiver<'a, I, O> {
    fn new(msg_socket: &MsgSocket<I, O>) -> MsgResult<AsyncReceiver<I, O>> {
        add_fd_flags(msg_socket.as_raw_fd(), O_NONBLOCK).map_err(MsgError::SettingFdFlags)?;
        Ok(AsyncReceiver {
            inner: msg_socket,
            done: false,
        })
    }
}

impl<'a, I: MsgOnSocket, O: MsgOnSocket> Drop for AsyncReceiver<'a, I, O> {
    fn drop(&mut self) {
        if let Err(e) = clear_fd_flags(self.inner.as_raw_fd(), O_NONBLOCK) {
            error!(
                "Failed to restore non-blocking behavior to message socket: {}",
                e
            );
        }
    }
}

impl<'a, I: MsgOnSocket, O: MsgOnSocket> Stream for AsyncReceiver<'a, I, O> {
    type Item = MsgResult<O>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        if self.done {
            return Poll::Ready(None);
        }

        let ret = match self.inner.recv() {
            Ok(msg) => Ok(Poll::Ready(Some(Ok(msg)))),
            Err(MsgError::Recv(e)) => {
                if e.errno() == EWOULDBLOCK {
                    add_read_waker(self.inner.as_raw_fd(), cx.waker().clone())
                        .map(|_| Poll::Pending)
                        .map_err(MsgError::AddingWaker)
                } else {
                    Err(MsgError::Recv(e))
                }
            }
            Err(e) => Err(e),
        };

        match ret {
            Ok(p) => p,
            Err(e) => {
                // Indicate something went wrong and no more events will be provided.
                self.done = true;
                Poll::Ready(Some(Err(e)))
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

mod msg_on_socket;

use std::io::Result;
use std::marker::PhantomData;
use std::ops::Deref;
use std::os::unix::io::{AsRawFd, RawFd};

use sys_util::{handle_eintr, net::UnixSeqpacket, Error as SysError, ScmSocket};

pub use crate::msg_on_socket::*;
pub use msg_on_socket_derive::*;

/// Create a pair of socket. Request is send in one direction while response is in the other
/// direction.
pub fn pair<Request: MsgOnSocket, Response: MsgOnSocket>(
) -> Result<(MsgSocket<Request, Response>, MsgSocket<Response, Request>)> {
    let (sock1, sock2) = UnixSeqpacket::pair()?;
    let requester = MsgSocket {
        sock: sock1,
        _i: PhantomData,
        _o: PhantomData,
    };
    let responder = MsgSocket {
        sock: sock2,
        _i: PhantomData,
        _o: PhantomData,
    };
    Ok((requester, responder))
}

/// Bidirection sock that support both send and recv.
pub struct MsgSocket<I: MsgOnSocket, O: MsgOnSocket> {
    sock: UnixSeqpacket,
    _i: PhantomData<I>,
    _o: PhantomData<O>,
}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSocket<I, O> {
    // Create a new MsgSocket.
    pub fn new(s: UnixSeqpacket) -> MsgSocket<I, O> {
        MsgSocket {
            sock: s,
            _i: PhantomData,
            _o: PhantomData,
        }
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> Deref for MsgSocket<I, O> {
    type Target = UnixSeqpacket;
    fn deref(&self) -> &Self::Target {
        &self.sock
    }
}

/// One direction socket that only supports sending.
pub struct Sender<M: MsgOnSocket> {
    sock: UnixSeqpacket,
    _m: PhantomData<M>,
}

impl<M: MsgOnSocket> Sender<M> {
    /// Create a new sender sock.
    pub fn new(s: UnixSeqpacket) -> Sender<M> {
        Sender {
            sock: s,
            _m: PhantomData,
        }
    }
}

/// One direction socket that only supports receiving.
pub struct Receiver<M: MsgOnSocket> {
    sock: UnixSeqpacket,
    _m: PhantomData<M>,
}

impl<M: MsgOnSocket> Receiver<M> {
    /// Create a new receiver sock.
    pub fn new(s: UnixSeqpacket) -> Receiver<M> {
        Receiver {
            sock: s,
            _m: PhantomData,
        }
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRef<UnixSeqpacket> for MsgSocket<I, O> {
    fn as_ref(&self) -> &UnixSeqpacket {
        &self.sock
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRawFd for MsgSocket<I, O> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

impl<M: MsgOnSocket> AsRef<UnixSeqpacket> for Sender<M> {
    fn as_ref(&self) -> &UnixSeqpacket {
        &self.sock
    }
}

impl<M: MsgOnSocket> AsRawFd for Sender<M> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

impl<M: MsgOnSocket> AsRef<UnixSeqpacket> for Receiver<M> {
    fn as_ref(&self) -> &UnixSeqpacket {
        &self.sock
    }
}

impl<M: MsgOnSocket> AsRawFd for Receiver<M> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

/// Types that could send a message.
pub trait MsgSender<M: MsgOnSocket>: AsRef<UnixSeqpacket> {
    fn send(&self, msg: &M) -> MsgResult<()> {
        let msg_size = M::msg_size();
        let fd_size = M::max_fd_count();
        let mut msg_buffer: Vec<u8> = vec![0; msg_size];
        let mut fd_buffer: Vec<RawFd> = vec![0; fd_size];

        let fd_size = msg.write_to_buffer(&mut msg_buffer, &mut fd_buffer)?;
        let sock: &UnixSeqpacket = self.as_ref();
        if fd_size == 0 {
            handle_eintr!(sock.send(&msg_buffer))
                .map_err(|e| MsgError::Send(SysError::new(e.raw_os_error().unwrap_or(0))))?;
        } else {
            sock.send_with_fds(&msg_buffer[..], &fd_buffer[0..fd_size])
                .map_err(MsgError::Send)?;
        }
        Ok(())
    }
}

/// Types that could receive a message.
pub trait MsgReceiver<M: MsgOnSocket>: AsRef<UnixSeqpacket> {
    fn recv(&self) -> MsgResult<M> {
        let msg_size = M::msg_size();
        let fd_size = M::max_fd_count();
        let mut msg_buffer: Vec<u8> = vec![0; msg_size];
        let mut fd_buffer: Vec<RawFd> = vec![0; fd_size];

        let sock: &UnixSeqpacket = self.as_ref();

        let (recv_msg_size, recv_fd_size) = {
            if fd_size == 0 {
                let size = sock
                    .recv(&mut msg_buffer)
                    .map_err(|e| MsgError::Recv(SysError::new(e.raw_os_error().unwrap_or(0))))?;
                (size, 0)
            } else {
                sock.recv_with_fds(&mut msg_buffer, &mut fd_buffer)
                    .map_err(MsgError::Recv)?
            }
        };
        if msg_size != recv_msg_size {
            return Err(MsgError::BadRecvSize {
                expected: msg_size,
                actual: recv_msg_size,
            });
        }
        // Safe because fd buffer is read from socket.
        let (v, read_fd_size) = unsafe {
            M::read_from_buffer(&msg_buffer[0..recv_msg_size], &fd_buffer[0..recv_fd_size])?
        };
        if recv_fd_size != read_fd_size {
            return Err(MsgError::NotExpectFd);
        }
        Ok(v)
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSender<I> for MsgSocket<I, O> {}
impl<I: MsgOnSocket, O: MsgOnSocket> MsgReceiver<O> for MsgSocket<I, O> {}

impl<M: MsgOnSocket> MsgSender<M> for Sender<M> {}
impl<M: MsgOnSocket> MsgReceiver<M> for Receiver<M> {}
<commit_msg>msg_socket: Change trait type parameters to associated types<commit_after>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

mod msg_on_socket;

use std::io::Result;
use std::marker::PhantomData;
use std::ops::Deref;
use std::os::unix::io::{AsRawFd, RawFd};

use sys_util::{handle_eintr, net::UnixSeqpacket, Error as SysError, ScmSocket};

pub use crate::msg_on_socket::*;
pub use msg_on_socket_derive::*;

/// Create a pair of socket. Request is send in one direction while response is in the other
/// direction.
pub fn pair<Request: MsgOnSocket, Response: MsgOnSocket>(
) -> Result<(MsgSocket<Request, Response>, MsgSocket<Response, Request>)> {
    let (sock1, sock2) = UnixSeqpacket::pair()?;
    let requester = MsgSocket {
        sock: sock1,
        _i: PhantomData,
        _o: PhantomData,
    };
    let responder = MsgSocket {
        sock: sock2,
        _i: PhantomData,
        _o: PhantomData,
    };
    Ok((requester, responder))
}

/// Bidirection sock that support both send and recv.
pub struct MsgSocket<I: MsgOnSocket, O: MsgOnSocket> {
    sock: UnixSeqpacket,
    _i: PhantomData<I>,
    _o: PhantomData<O>,
}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSocket<I, O> {
    // Create a new MsgSocket.
    pub fn new(s: UnixSeqpacket) -> MsgSocket<I, O> {
        MsgSocket {
            sock: s,
            _i: PhantomData,
            _o: PhantomData,
        }
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> Deref for MsgSocket<I, O> {
    type Target = UnixSeqpacket;
    fn deref(&self) -> &Self::Target {
        &self.sock
    }
}

/// One direction socket that only supports sending.
pub struct Sender<M: MsgOnSocket> {
    sock: UnixSeqpacket,
    _m: PhantomData<M>,
}

impl<M: MsgOnSocket> Sender<M> {
    /// Create a new sender sock.
    pub fn new(s: UnixSeqpacket) -> Sender<M> {
        Sender {
            sock: s,
            _m: PhantomData,
        }
    }
}

/// One direction socket that only supports receiving.
pub struct Receiver<M: MsgOnSocket> {
    sock: UnixSeqpacket,
    _m: PhantomData<M>,
}

impl<M: MsgOnSocket> Receiver<M> {
    /// Create a new receiver sock.
    pub fn new(s: UnixSeqpacket) -> Receiver<M> {
        Receiver {
            sock: s,
            _m: PhantomData,
        }
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRef<UnixSeqpacket> for MsgSocket<I, O> {
    fn as_ref(&self) -> &UnixSeqpacket {
        &self.sock
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRawFd for MsgSocket<I, O> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

impl<M: MsgOnSocket> AsRef<UnixSeqpacket> for Sender<M> {
    fn as_ref(&self) -> &UnixSeqpacket {
        &self.sock
    }
}

impl<M: MsgOnSocket> AsRawFd for Sender<M> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

impl<M: MsgOnSocket> AsRef<UnixSeqpacket> for Receiver<M> {
    fn as_ref(&self) -> &UnixSeqpacket {
        &self.sock
    }
}

impl<M: MsgOnSocket> AsRawFd for Receiver<M> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

/// Types that could send a message.
pub trait MsgSender: AsRef<UnixSeqpacket> {
    type M: MsgOnSocket;
    fn send(&self, msg: &Self::M) -> MsgResult<()> {
        let msg_size = Self::M::msg_size();
        let fd_size = Self::M::max_fd_count();
        let mut msg_buffer: Vec<u8> = vec![0; msg_size];
        let mut fd_buffer: Vec<RawFd> = vec![0; fd_size];

        let fd_size = msg.write_to_buffer(&mut msg_buffer, &mut fd_buffer)?;
        let sock: &UnixSeqpacket = self.as_ref();
        if fd_size == 0 {
            handle_eintr!(sock.send(&msg_buffer))
                .map_err(|e| MsgError::Send(SysError::new(e.raw_os_error().unwrap_or(0))))?;
        } else {
            sock.send_with_fds(&msg_buffer[..], &fd_buffer[0..fd_size])
                .map_err(MsgError::Send)?;
        }
        Ok(())
    }
}

/// Types that could receive a message.
pub trait MsgReceiver: AsRef<UnixSeqpacket> {
    type M: MsgOnSocket;
    fn recv(&self) -> MsgResult<Self::M> {
        let msg_size = Self::M::msg_size();
        let fd_size = Self::M::max_fd_count();
        let mut msg_buffer: Vec<u8> = vec![0; msg_size];
        let mut fd_buffer: Vec<RawFd> = vec![0; fd_size];

        let sock: &UnixSeqpacket = self.as_ref();

        let (recv_msg_size, recv_fd_size) = {
            if fd_size == 0 {
                let size = sock
                    .recv(&mut msg_buffer)
                    .map_err(|e| MsgError::Recv(SysError::new(e.raw_os_error().unwrap_or(0))))?;
                (size, 0)
            } else {
                sock.recv_with_fds(&mut msg_buffer, &mut fd_buffer)
                    .map_err(MsgError::Recv)?
            }
        };
        if msg_size != recv_msg_size {
            return Err(MsgError::BadRecvSize {
                expected: msg_size,
                actual: recv_msg_size,
            });
        }
        // Safe because fd buffer is read from socket.
        let (v, read_fd_size) = unsafe {
            Self::M::read_from_buffer(&msg_buffer[0..recv_msg_size], &fd_buffer[0..recv_fd_size])?
        };
        if recv_fd_size != read_fd_size {
            return Err(MsgError::NotExpectFd);
        }
        Ok(v)
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSender for MsgSocket<I, O> {
    type M = I;
}
impl<I: MsgOnSocket, O: MsgOnSocket> MsgReceiver for MsgSocket<I, O> {
    type M = O;
}

impl<I: MsgOnSocket> MsgSender for Sender<I> {
    type M = I;
}
impl<O: MsgOnSocket> MsgReceiver for Receiver<O> {
    type M = O;
}
<|endoftext|>


<commit_before>use core::fmt::Arguments;

#[cfg(not(test))]
#[lang = "stack_exhausted"]
pub fn stack_exhausted() -> ! {
    unimplemented!()
}

#[cfg(not(test))]
#[lang = "eh_personality"]
pub fn eh_personality() -> ! {
    unimplemented!()
}

#[cfg(not(test))]
#[lang = "panic_fmt"]
pub fn panic_fmt(_: Arguments, _: &(&'static str, u32)) -> ! {
    loop {}
}
<commit_msg>lang: Add implementation of panic_fmt<commit_after>use core::fmt;

use ::xen::sched::crash;

#[cfg(not(test))]
#[lang = "stack_exhausted"]
pub fn stack_exhausted() -> ! {
    unimplemented!()
}

#[cfg(not(test))]
#[lang = "eh_personality"]
pub fn eh_personality() -> ! {
    unimplemented!()
}

#[cfg(not(test))]
#[lang = "panic_fmt"]
pub extern fn panic_impl(msg: fmt::Arguments, file: &'static str, line: u32) {
    println!("Panic at '{}:{}' with message '{}'", file, line, msg);
    ::xen::console::console().flush();

    crash();

    loop {}
}
<|endoftext|>


<commit_before>// https://leetcode.com/problems/count-asterisks/
pub fn count_asterisks(s: String) -> i32 {
    todo!()
}

fn main() {
    println!("{}", count_asterisks(String::from("l|*e*et|c**o|*de|"))); // 2
    println!("{}", count_asterisks(String::from("iamprogrammer"))); // 0
    println!(
        "{}",
        count_asterisks(String::from("yo|uar|e**|b|e***au|tifu|l"))
    ); // 5
}
<commit_msg>Solve problem 2315<commit_after>// https://leetcode.com/problems/count-asterisks/
pub fn count_asterisks(s: String) -> i32 {
    let mut answer = 0;
    let mut j = 0;
    for i in s.split("|") {
        if j % 2 != 0 {
            j += 1;
            continue;
        }
        answer += (i.chars().filter(|a| *a == '*').count()) as i32;
        j += 1;
    }
    answer
}

fn main() {
    println!("{}", count_asterisks(String::from("l|*e*et|c**o|*de|"))); // 2
    println!("{}", count_asterisks(String::from("iamprogrammer"))); // 0
    println!(
        "{}",
        count_asterisks(String::from("yo|uar|e**|b|e***au|tifu|l"))
    ); // 5
}
<|endoftext|>


<commit_before>// https://leetcode.com/problems/time-needed-to-rearrange-a-binary-string/
pub fn seconds_to_remove_occurrences(s: String) -> i32 {
    let mut answer = 0;
    let mut chars: Vec<char> = s.chars().collect();
    let mut i = 0;
    while i < chars.len() - 1 {
        if chars[i] == '0' && chars[i + 1] == '1' {
            chars.swap(i, i + 1);
            answer += 1;
        }
        i += 1;
    }
    answer
}

fn main() {
    println!("{}", seconds_to_remove_occurrences("0110101".to_string())); // 4
    println!("{}", seconds_to_remove_occurrences("11100".to_string())); // 0
    println!(
        "{}",
        seconds_to_remove_occurrences("1000111001010010010101010101".to_string())
    ); // 15
    println!("{}", seconds_to_remove_occurrences("1100101".to_string())); // 3
}
<commit_msg>Update problem 2380<commit_after>// https://leetcode.com/problems/time-needed-to-rearrange-a-binary-string/
pub fn seconds_to_remove_occurrences(s: String) -> i32 {
    let mut answer = 0;
    let mut chars: Vec<char> = s.chars().collect();
    loop {
        let mut i = 0;
        let mut count = 0;
        while i < chars.len() - 1 {
            if chars[i] == '0' && chars[i + 1] == '1' {
                chars.swap(i, i + 1);
                count += 1;
            }
            i += 1;
        }
        if count == 0 {
            break;
        }
        answer += 1;
    }
    answer
}

fn main() {
    println!("{}", seconds_to_remove_occurrences("0110101".to_string())); // 4
    println!("{}", seconds_to_remove_occurrences("11100".to_string())); // 0
    println!(
        "{}",
        seconds_to_remove_occurrences("1000111001010010010101010101".to_string())
    ); // 15
    println!("{}", seconds_to_remove_occurrences("1100101".to_string())); // 3
}
<|endoftext|>


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

use std::mem::{self, ManuallyDrop};
use std::thread;

/// Dummy type that panics when its destructor is invoked.
#[derive(Debug)]
struct DropBomb();

impl Drop for DropBomb {
    fn drop(&mut self) {
        let message = "Trying to drop an AliasedCell, which may still have aliases outstanding.";
        if thread::panicking() {
            eprintln!("{}", message);
        } else {
            panic!(message);
        }
    }
}

/// A wrapper for unsafely aliased memory locations.
///
/// `AliasedCell' makes sure that its inner value
/// is completely inaccessible from safe code.
/// Once an `AliasedCell` has been constructed from some value,
/// until the value is moved out again with the (unsafe) `into_inner()` method,
/// the only way to access the wrapped value
/// is through the unsafe `alias_mut()` method.
///
/// This is useful for FFI calls that take raw pointers as input,
/// and hold on to them even after returning control to the caller.
/// Since Rust's type system is not aware of such aliases,
/// it cannot provide the usual guarantees about validity of pointers
/// and exclusiveness of mutable pointers.
/// This means that any code that has access to the memory in question
/// is inherently unsafe as long as such untracked aliases exist.
/// Putting the value in an `AliasedCell` before the FFI call,
/// and only taking it out again
/// once the caller has ensured that all aliases have been dropped
/// (most likely through another FFI call),
/// makes certain that any such unsafe access to the aliased value
/// can only happen from code marked as `unsafe`.
///
/// An `AliasedCell` should never be simply dropped in normal use.
/// Rather, it should always be freed explicitly with `into_inner()`,
/// signalling that there are no outstanding aliases anymore.
/// When an `AliasedCell` is dropped unexpectedly,
/// we have to assume that it likely still has untracked aliases,
/// and thus freeing the memory would probably be unsound.
/// Therefore, the `drop()` implementation of `AliasedCell`
/// leaks the inner value instead of freeing it;
/// and throws a panic.
/// Leaking the memory, while undesirable in general,
/// keeps the memory accessible to any outstanding aliases.
/// This is the only way to retain soundness during unwinding,
/// or when the panic gets caught.
///
/// Note that making FFI access through untracked aliases
/// requires the value to have a stable memory location --
/// typically by living on the heap rather than on the stack.
/// If the value isn't already in a heap-allocated container
/// such as `Box<>`, `Vec<>`, or `String`,
/// it is the caller's responsibility to wrap it in a `Box<>` explicitly.
/// `AliasedCell` itself cannot ensure that the address remains stable
/// when the `AliasedCell` gets moved around.
#[derive(Debug)]
pub struct AliasedCell<T> {
    value: ManuallyDrop<T>,
    drop_bomb: DropBomb,
}

impl<T> AliasedCell<T> {
    /// Wrap the provided value in an `AliasedCell`, making it inaccessible from safe code.
    pub fn new(value: T) -> AliasedCell<T> {
        AliasedCell {
            value: ManuallyDrop::new(value),
            drop_bomb: DropBomb(),
        }
    }

    /// Get a pointer to the inner value.
    ///
    /// Note that this yields a regular reference.
    /// To actually get an untracked alias,
    /// it needs to be cast or coerced into a raw pointer.
    /// This usually happens implicitly though
    /// when calling an FFI function (or any other function)
    /// taking a raw pointer as argument.
    ///
    /// `alias_mut()` can be called any number of times:
    /// the wrapper doesn't keep track of the number of outstanding aliases --
    /// the caller is responsible for making sure that no aliases are left
    /// before invoking `into_inner()`.
    /// If you need to track the number of aliases,
    /// wrap the inner value in an `Rc<>` or `Arc` --
    /// this way, the reference count will also be inaccessible from safe code.
    pub unsafe fn alias_mut(&mut self) -> &mut T {
        &mut self.value
    }

    /// Get a shared (immutable) pointer to the inner value.
    ///
    /// With this method it's possible to get an alias
    /// while only holding a shared reference to the `AliasedCell`.
    ///
    /// Since all the unsafe aliases are untracked,
    /// it's up to the callers to make sure no shared aliases are used
    /// while the data might actually be mutated elsewhere
    /// through some outstanding mutable aliases.
    pub unsafe fn alias(&self) -> &T {
        &self.value // orig &self.inner
    }

    /// Move out the wrapped value, making it accessible from safe code again.
    pub unsafe fn into_inner(self) -> T {
        mem::forget(self.drop_bomb);
        ManuallyDrop::into_inner(self.value)
    }
}

/// Some basic tests.
///
/// Note: These mostly just check that various expected usage scenarios
/// can be compiled and basically work.
/// We can't verify though that the usage is actually sound;
/// nor do we check whether invalid usage is indeed prevented by the compiler...
///
/// (The latter could probably be remedied though
/// with some sort of compile-fail tests.)
#[cfg(test)]
mod tests {
    use super::AliasedCell;

    unsafe fn mutate_value(addr: *mut [i32; 4]) {
        let value = addr.as_mut().unwrap();
        value[1] += value[3];
    }

    struct Mutator {
        addr: *mut [i32; 4],
        ascended: bool,
    }

    impl Mutator {
        unsafe fn new(addr: *mut [i32; 4]) -> Mutator {
            Mutator {
                addr: addr,
                ascended: false,
            }
        }

        fn ascend(&mut self) {
            self.ascended = true;
        }

        unsafe fn mutate(&mut self) {
            let value = self.addr.as_mut().unwrap();
            if self.ascended {
                value[3] += value[2];
            } else {
                value[1] += value[3];
            }
        }
    }

    #[test]
    fn noop_roundtrip() {
        let value = [1, 3, 3, 7];
        let cell = AliasedCell::new(Box::new(value));
        let new_value = unsafe {
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 3, 3, 7]);
    }

    #[test]
    fn unused_alias() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        unsafe {
            cell.alias_mut().as_mut();
        }
        let new_value = unsafe {
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 3, 3, 7]);
    }

    #[test]
    fn mutate() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        unsafe {
            mutate_value(cell.alias_mut().as_mut());
        }
        let new_value = unsafe {
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 10, 3, 7]);
    }

    /// Verify that we can take multiple aliases.
    #[test]
    fn mutate_twice() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        unsafe {
            mutate_value(cell.alias_mut().as_mut());
        }
        unsafe {
            mutate_value(cell.alias_mut().as_mut());
        }
        let new_value = unsafe {
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 17, 3, 7]);
    }

    /// Verify that we can do basic safe manipulations between unsafe blocks.
    #[test]
    fn moves() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        unsafe {
            mutate_value(cell.alias_mut().as_mut());
        }
        let mut cell2 = cell;
        unsafe {
            mutate_value(cell2.alias_mut().as_mut());
        }
        let cell3 = cell2;
        let new_value = unsafe {
            *cell3.into_inner()
        };
        assert_eq!(new_value, [1, 17, 3, 7]);
    }

    /// Verify that alias can be used at a later point.
    #[test]
    fn mutate_deferred() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        let mut mutator = unsafe {
            Mutator::new(cell.alias_mut().as_mut())
        };
        unsafe {
            mutator.mutate();
        }
        let new_value = unsafe {
            drop(mutator);
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 10, 3, 7]);
    }

    #[test]
    fn mutate_deferred_twice() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        let mut mutator = unsafe {
            Mutator::new(cell.alias_mut().as_mut())
        };
        unsafe {
            mutator.mutate();
        }
        unsafe {
            mutator.mutate();
        }
        let new_value = unsafe {
            drop(mutator);
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 17, 3, 7]);
    }

    /// Further safe manipulations.
    #[test]
    fn deferred_moves() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        let mutator = unsafe {
            Mutator::new(cell.alias_mut().as_mut())
        };
        let cell2 = cell;
        let mut mutator2 = mutator;
        unsafe {
            mutator2.mutate();
        }
        let cell3 = cell2;
        let mutator3 = mutator2;
        let new_value = unsafe {
            drop(mutator3);
            *cell3.into_inner()
        };
        assert_eq!(new_value, [1, 10, 3, 7]);
    }

    /// Non-trivial safe manipulation.
    #[test]
    fn safe_frobbing() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        let mut mutator = unsafe {
            Mutator::new(cell.alias_mut().as_mut())
        };
        unsafe {
            mutator.mutate();
        }
        mutator.ascend();
        unsafe {
            mutator.mutate();
        }
        let new_value = unsafe {
            drop(mutator);
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 10, 3, 10]);
    }

    /// Verify that two aliases can exist simultaneously.
    #[test]
    fn two_mutators() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        let mut mutator1 = unsafe {
            Mutator::new(cell.alias_mut().as_mut())
        };
        unsafe {
            mutator1.mutate();
        }
        let mut mutator2 = unsafe {
            Mutator::new(cell.alias_mut().as_mut())
        };
        unsafe {
            mutator2.mutate();
        }
        let new_value = unsafe {
            drop(mutator1);
            drop(mutator2);
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 17, 3, 7]);
    }

    #[test]
    #[should_panic(expected = "Trying to drop an AliasedCell, which may still have aliases outstanding.")]
    fn invalid_drop() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        unsafe {
            let mut mutator = Mutator::new(cell.alias_mut().as_mut());
            mutator.mutate();
            drop(cell);
        }
    }

    /// Verify that we skip the panic-on-drop while unwinding from another panic.
    #[test]
    #[should_panic(expected = "bye!")]
    fn panic() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        unsafe {
            let mut mutator = Mutator::new(cell.alias_mut().as_mut());
            mutator.mutate();
            panic!("bye!");
        }
    }
}
<commit_msg>chore: use `panic!("{}", ...)` in prep. for Rust 2021 edition<commit_after>// Copyright 2018 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::mem::{self, ManuallyDrop};
use std::thread;

/// Dummy type that panics when its destructor is invoked.
#[derive(Debug)]
struct DropBomb();

impl Drop for DropBomb {
    fn drop(&mut self) {
        let message = "Trying to drop an AliasedCell, which may still have aliases outstanding.";
        if thread::panicking() {
            eprintln!("{}", message);
        } else {
            panic!("{}", message);
        }
    }
}

/// A wrapper for unsafely aliased memory locations.
///
/// `AliasedCell' makes sure that its inner value
/// is completely inaccessible from safe code.
/// Once an `AliasedCell` has been constructed from some value,
/// until the value is moved out again with the (unsafe) `into_inner()` method,
/// the only way to access the wrapped value
/// is through the unsafe `alias_mut()` method.
///
/// This is useful for FFI calls that take raw pointers as input,
/// and hold on to them even after returning control to the caller.
/// Since Rust's type system is not aware of such aliases,
/// it cannot provide the usual guarantees about validity of pointers
/// and exclusiveness of mutable pointers.
/// This means that any code that has access to the memory in question
/// is inherently unsafe as long as such untracked aliases exist.
/// Putting the value in an `AliasedCell` before the FFI call,
/// and only taking it out again
/// once the caller has ensured that all aliases have been dropped
/// (most likely through another FFI call),
/// makes certain that any such unsafe access to the aliased value
/// can only happen from code marked as `unsafe`.
///
/// An `AliasedCell` should never be simply dropped in normal use.
/// Rather, it should always be freed explicitly with `into_inner()`,
/// signalling that there are no outstanding aliases anymore.
/// When an `AliasedCell` is dropped unexpectedly,
/// we have to assume that it likely still has untracked aliases,
/// and thus freeing the memory would probably be unsound.
/// Therefore, the `drop()` implementation of `AliasedCell`
/// leaks the inner value instead of freeing it;
/// and throws a panic.
/// Leaking the memory, while undesirable in general,
/// keeps the memory accessible to any outstanding aliases.
/// This is the only way to retain soundness during unwinding,
/// or when the panic gets caught.
///
/// Note that making FFI access through untracked aliases
/// requires the value to have a stable memory location --
/// typically by living on the heap rather than on the stack.
/// If the value isn't already in a heap-allocated container
/// such as `Box<>`, `Vec<>`, or `String`,
/// it is the caller's responsibility to wrap it in a `Box<>` explicitly.
/// `AliasedCell` itself cannot ensure that the address remains stable
/// when the `AliasedCell` gets moved around.
#[derive(Debug)]
pub struct AliasedCell<T> {
    value: ManuallyDrop<T>,
    drop_bomb: DropBomb,
}

impl<T> AliasedCell<T> {
    /// Wrap the provided value in an `AliasedCell`, making it inaccessible from safe code.
    pub fn new(value: T) -> AliasedCell<T> {
        AliasedCell {
            value: ManuallyDrop::new(value),
            drop_bomb: DropBomb(),
        }
    }

    /// Get a pointer to the inner value.
    ///
    /// Note that this yields a regular reference.
    /// To actually get an untracked alias,
    /// it needs to be cast or coerced into a raw pointer.
    /// This usually happens implicitly though
    /// when calling an FFI function (or any other function)
    /// taking a raw pointer as argument.
    ///
    /// `alias_mut()` can be called any number of times:
    /// the wrapper doesn't keep track of the number of outstanding aliases --
    /// the caller is responsible for making sure that no aliases are left
    /// before invoking `into_inner()`.
    /// If you need to track the number of aliases,
    /// wrap the inner value in an `Rc<>` or `Arc` --
    /// this way, the reference count will also be inaccessible from safe code.
    pub unsafe fn alias_mut(&mut self) -> &mut T {
        &mut self.value
    }

    /// Get a shared (immutable) pointer to the inner value.
    ///
    /// With this method it's possible to get an alias
    /// while only holding a shared reference to the `AliasedCell`.
    ///
    /// Since all the unsafe aliases are untracked,
    /// it's up to the callers to make sure no shared aliases are used
    /// while the data might actually be mutated elsewhere
    /// through some outstanding mutable aliases.
    pub unsafe fn alias(&self) -> &T {
        &self.value // orig &self.inner
    }

    /// Move out the wrapped value, making it accessible from safe code again.
    pub unsafe fn into_inner(self) -> T {
        mem::forget(self.drop_bomb);
        ManuallyDrop::into_inner(self.value)
    }
}

/// Some basic tests.
///
/// Note: These mostly just check that various expected usage scenarios
/// can be compiled and basically work.
/// We can't verify though that the usage is actually sound;
/// nor do we check whether invalid usage is indeed prevented by the compiler...
///
/// (The latter could probably be remedied though
/// with some sort of compile-fail tests.)
#[cfg(test)]
mod tests {
    use super::AliasedCell;

    unsafe fn mutate_value(addr: *mut [i32; 4]) {
        let value = addr.as_mut().unwrap();
        value[1] += value[3];
    }

    struct Mutator {
        addr: *mut [i32; 4],
        ascended: bool,
    }

    impl Mutator {
        unsafe fn new(addr: *mut [i32; 4]) -> Mutator {
            Mutator {
                addr: addr,
                ascended: false,
            }
        }

        fn ascend(&mut self) {
            self.ascended = true;
        }

        unsafe fn mutate(&mut self) {
            let value = self.addr.as_mut().unwrap();
            if self.ascended {
                value[3] += value[2];
            } else {
                value[1] += value[3];
            }
        }
    }

    #[test]
    fn noop_roundtrip() {
        let value = [1, 3, 3, 7];
        let cell = AliasedCell::new(Box::new(value));
        let new_value = unsafe {
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 3, 3, 7]);
    }

    #[test]
    fn unused_alias() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        unsafe {
            cell.alias_mut().as_mut();
        }
        let new_value = unsafe {
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 3, 3, 7]);
    }

    #[test]
    fn mutate() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        unsafe {
            mutate_value(cell.alias_mut().as_mut());
        }
        let new_value = unsafe {
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 10, 3, 7]);
    }

    /// Verify that we can take multiple aliases.
    #[test]
    fn mutate_twice() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        unsafe {
            mutate_value(cell.alias_mut().as_mut());
        }
        unsafe {
            mutate_value(cell.alias_mut().as_mut());
        }
        let new_value = unsafe {
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 17, 3, 7]);
    }

    /// Verify that we can do basic safe manipulations between unsafe blocks.
    #[test]
    fn moves() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        unsafe {
            mutate_value(cell.alias_mut().as_mut());
        }
        let mut cell2 = cell;
        unsafe {
            mutate_value(cell2.alias_mut().as_mut());
        }
        let cell3 = cell2;
        let new_value = unsafe {
            *cell3.into_inner()
        };
        assert_eq!(new_value, [1, 17, 3, 7]);
    }

    /// Verify that alias can be used at a later point.
    #[test]
    fn mutate_deferred() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        let mut mutator = unsafe {
            Mutator::new(cell.alias_mut().as_mut())
        };
        unsafe {
            mutator.mutate();
        }
        let new_value = unsafe {
            drop(mutator);
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 10, 3, 7]);
    }

    #[test]
    fn mutate_deferred_twice() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        let mut mutator = unsafe {
            Mutator::new(cell.alias_mut().as_mut())
        };
        unsafe {
            mutator.mutate();
        }
        unsafe {
            mutator.mutate();
        }
        let new_value = unsafe {
            drop(mutator);
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 17, 3, 7]);
    }

    /// Further safe manipulations.
    #[test]
    fn deferred_moves() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        let mutator = unsafe {
            Mutator::new(cell.alias_mut().as_mut())
        };
        let cell2 = cell;
        let mut mutator2 = mutator;
        unsafe {
            mutator2.mutate();
        }
        let cell3 = cell2;
        let mutator3 = mutator2;
        let new_value = unsafe {
            drop(mutator3);
            *cell3.into_inner()
        };
        assert_eq!(new_value, [1, 10, 3, 7]);
    }

    /// Non-trivial safe manipulation.
    #[test]
    fn safe_frobbing() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        let mut mutator = unsafe {
            Mutator::new(cell.alias_mut().as_mut())
        };
        unsafe {
            mutator.mutate();
        }
        mutator.ascend();
        unsafe {
            mutator.mutate();
        }
        let new_value = unsafe {
            drop(mutator);
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 10, 3, 10]);
    }

    /// Verify that two aliases can exist simultaneously.
    #[test]
    fn two_mutators() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        let mut mutator1 = unsafe {
            Mutator::new(cell.alias_mut().as_mut())
        };
        unsafe {
            mutator1.mutate();
        }
        let mut mutator2 = unsafe {
            Mutator::new(cell.alias_mut().as_mut())
        };
        unsafe {
            mutator2.mutate();
        }
        let new_value = unsafe {
            drop(mutator1);
            drop(mutator2);
            *cell.into_inner()
        };
        assert_eq!(new_value, [1, 17, 3, 7]);
    }

    #[test]
    #[should_panic(expected = "Trying to drop an AliasedCell, which may still have aliases outstanding.")]
    fn invalid_drop() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        unsafe {
            let mut mutator = Mutator::new(cell.alias_mut().as_mut());
            mutator.mutate();
            drop(cell);
        }
    }

    /// Verify that we skip the panic-on-drop while unwinding from another panic.
    #[test]
    #[should_panic(expected = "bye!")]
    fn panic() {
        let value = [1, 3, 3, 7];
        let mut cell = AliasedCell::new(Box::new(value));
        unsafe {
            let mut mutator = Mutator::new(cell.alias_mut().as_mut());
            mutator.mutate();
            panic!("bye!");
        }
    }
}
<|endoftext|>


<commit_before>extern crate core;

use core::num::FromPrimitive;
use core::fmt;

use std::rand::random;

use Direction::{SpinUp, SpinDown, SpinSuper};
use Plan::{Trivial, OddBall};

#[derive(Copy)]
#[derive(Show)]
#[derive(PartialEq)]
enum Direction {
    SpinUp,
    SpinDown,
    SpinSuper,
}

impl fmt::String for Direction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let r = match self {
            &SpinUp => "SpinUp",
            &SpinDown => "SpinDown",
            &SpinSuper => "SpinSuper"
        };
        write!(f, "{}", r)
    }
}

impl fmt::String for Detector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use Detector::{D12, D3, D9};
        let r = match self {
            &D12 => "12 oclock",
            &D3 => "3 oclock",
            &D9 => "9 oclock"
        };
        write!(f, "{}", r)
    }
}

struct Pair<T> {
    lhs: T,
    rhs: T
}

struct Particle {
    spin: Direction
}

enum Plan {
    Trivial, // up-up-up -> down-down-down
    OddBall  // up-down-up -> down-up-down
}

#[derive(Show)]
enum Detector {
    D12, // 12 o'clock
    D3, // 3 o'clock
    D9, // 9 o' clock
}

impl Particle {
    fn new_pair () -> Pair<Particle> {
        let d1 = SpinSuper;
        let d2 = SpinSuper;
        let p1 = Particle{spin: d1};
        let p2 = Particle{spin: d2};

        return Pair{lhs: p1, rhs: p2};
    }

    pub fn measure (&mut self) {
        // TODO theta = 60degrees use 3/4th and 1/4th

        let detector = Particle::get_detector_direction();
        let rnd = random::<f32>();

        self.spin = match detector {
            Detector::D12 => SpinDown,
            Detector::D3 => match rnd {
                0.0  ... 0.33 => SpinUp,
                0.33 ... 1.00 => SpinDown,
                _ => unreachable!()
            },
            Detector::D9 => match rnd {
                0.0  ... 0.33 => SpinDown,
                0.33 ... 1.00 => SpinUp,
                _ => unreachable!()
            },
        };
    }

    // NOTE on spooky and premeditated only if measured in the same
    // direction must the spins be opposite

    // measure with with a message
    pub fn spooky (&mut self, friend: &mut Particle) -> Pair<Direction> {

        self.measure();
        friend.spin = match self.spin {
            SpinUp => SpinDown,
            SpinDown => SpinUp,
            _ => panic!("broke the universe")
        };
        return Pair{lhs: self.spin, rhs: friend.spin};
    }

    fn get_detector_direction () -> Detector {
        let rnd = random::<f32>();

        let detector = match rnd {
            0.0  ... 0.33 => Detector::D12,
            0.33 ... 0.66 => Detector::D3,
            0.66 ... 1.00 => Detector::D9,
            _ => unreachable!()
        };

        return detector;
    }

    pub fn hidden_information (&mut self, friend: &mut Particle, plan: Plan) -> Pair<Direction> {

        let rnd = random::<f32>();

        let detector = Particle::get_detector_direction();
        let spin = match plan {
            Trivial => SpinUp,
            OddBall => match detector {
                Detector::D3 => SpinDown,
                _ => SpinUp
            }
        };
        friend.spin = spin;

        let detector = Particle::get_detector_direction();
        let spin = match plan {
            Trivial => SpinDown,
            OddBall => match detector {
                Detector::D3 => SpinUp,
                _ => SpinDown
            }
        };
        self.spin = spin;

        return Pair{lhs: self.spin, rhs: friend.spin};
    }
}

// hidden_information should give +55.6% difference, spooky would give 50%
fn main () {

    let mut particle = Particle{spin: SpinSuper};
    particle.measure();
    println!("particle {}", particle.spin);

    let mut trials: f64 = 1000f64;
    let mut num_different: f64 = 0f64;

    let particles = Particle::new_pair();
    let mut lhs = particles.lhs;
    let mut rhs = particles.rhs;

    for _ in range(0, trials as usize) {
        let particles = Particle::new_pair();
        lhs = particles.lhs;
        rhs = particles.rhs;

        lhs.spooky(&mut rhs);

        if lhs.spin != rhs.spin {
            num_different += 1.0;
        }
    }

    println!("lhs.spin {}, rhs.spin {} num_different {}", lhs.spin, rhs.spin, num_different);

    println!("num_different {}%", 100f64 * (trials - num_different) / trials as f64);

    trials = 1000f64;
    num_different = 0f64;

    for _ in range(0, trials as usize) {
        let particles = Particle::new_pair();
        lhs = particles.lhs;
        rhs = particles.rhs;
        lhs.hidden_information(&mut rhs, OddBall);

        //println!("lhs.spin {}, rhs.spin {}", lhs.spin, rhs.spin);
        if lhs.spin != rhs.spin {
            num_different += 1.0;
        }
    }

    println!("Percent the same {}%", 100f64 * (trials - num_different) / trials as f64);

}
<commit_msg>print should be<commit_after>extern crate core;

use core::num::FromPrimitive;
use core::fmt;

use std::rand::random;

use Direction::{SpinUp, SpinDown, SpinSuper};
use Plan::{Trivial, OddBall};

#[derive(Copy)]
#[derive(Show)]
#[derive(PartialEq)]
enum Direction {
    SpinUp,
    SpinDown,
    SpinSuper,
}

impl fmt::String for Direction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let r = match self {
            &SpinUp => "SpinUp",
            &SpinDown => "SpinDown",
            &SpinSuper => "SpinSuper"
        };
        write!(f, "{}", r)
    }
}

impl fmt::String for Detector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use Detector::{D12, D3, D9};
        let r = match self {
            &D12 => "12 oclock",
            &D3 => "3 oclock",
            &D9 => "9 oclock"
        };
        write!(f, "{}", r)
    }
}

struct Pair<T> {
    lhs: T,
    rhs: T
}

struct Particle {
    spin: Direction
}

enum Plan {
    Trivial, // up-up-up -> down-down-down
    OddBall  // up-down-up -> down-up-down
}

#[derive(Show)]
enum Detector {
    D12, // 12 o'clock
    D3, // 3 o'clock
    D9, // 9 o' clock
}

impl Particle {
    fn new_pair () -> Pair<Particle> {
        let d1 = SpinSuper;
        let d2 = SpinSuper;
        let p1 = Particle{spin: d1};
        let p2 = Particle{spin: d2};

        return Pair{lhs: p1, rhs: p2};
    }

    pub fn measure (&mut self) {
        // TODO theta = 60degrees use 3/4th and 1/4th

        let detector = Particle::get_detector_direction();
        let rnd = random::<f32>();

        self.spin = match detector {
            Detector::D12 => SpinDown,
            Detector::D3 => match rnd {
                0.0  ... 0.33 => SpinUp,
                0.33 ... 1.00 => SpinDown,
                _ => unreachable!()
            },
            Detector::D9 => match rnd {
                0.0  ... 0.33 => SpinDown,
                0.33 ... 1.00 => SpinUp,
                _ => unreachable!()
            },
        };
    }

    // NOTE on spooky and premeditated only if measured in the same
    // direction must the spins be opposite

    // measure with with a message
    pub fn spooky (&mut self, friend: &mut Particle) -> Pair<Direction> {

        self.measure();
        friend.spin = match self.spin {
            SpinUp => SpinDown,
            SpinDown => SpinUp,
            _ => panic!("broke the universe")
        };
        return Pair{lhs: self.spin, rhs: friend.spin};
    }

    fn get_detector_direction () -> Detector {
        let rnd = random::<f32>();

        let detector = match rnd {
            0.0  ... 0.33 => Detector::D12,
            0.33 ... 0.66 => Detector::D3,
            0.66 ... 1.00 => Detector::D9,
            _ => unreachable!()
        };

        return detector;
    }

    pub fn hidden_information (&mut self, friend: &mut Particle, plan: Plan) -> Pair<Direction> {

        let rnd = random::<f32>();

        let detector = Particle::get_detector_direction();
        let spin = match plan {
            Trivial => SpinUp,
            OddBall => match detector {
                Detector::D3 => SpinDown,
                _ => SpinUp
            }
        };
        friend.spin = spin;

        let detector = Particle::get_detector_direction();
        let spin = match plan {
            Trivial => SpinDown,
            OddBall => match detector {
                Detector::D3 => SpinUp,
                _ => SpinDown
            }
        };
        self.spin = spin;

        return Pair{lhs: self.spin, rhs: friend.spin};
    }
}

// hidden_information should give +55.6% difference, spooky would give 50%
fn main () {

    let mut particle = Particle{spin: SpinSuper};
    particle.measure();
    println!("particle {}", particle.spin);

    let mut trials: f64 = 1000f64;
    let mut num_different: f64 = 0f64;

    let particles = Particle::new_pair();
    let mut lhs = particles.lhs;
    let mut rhs = particles.rhs;

    for _ in range(0, trials as usize) {
        let particles = Particle::new_pair();
        lhs = particles.lhs;
        rhs = particles.rhs;

        lhs.spooky(&mut rhs);

        if lhs.spin != rhs.spin {
            num_different += 1.0;
        }
    }

    println!("lhs.spin {}, rhs.spin {} num_different {}", lhs.spin, rhs.spin, num_different);

    println!("num_different {}%", 100f64 * (trials - num_different) / trials as f64);

    trials = 1000f64;
    num_different = 0f64;

    for _ in range(0, trials as usize) {
        let particles = Particle::new_pair();
        lhs = particles.lhs;
        rhs = particles.rhs;
        lhs.hidden_information(&mut rhs, OddBall);

        //println!("lhs.spin {}, rhs.spin {}", lhs.spin, rhs.spin);
        if lhs.spin != rhs.spin {
            num_different += 1.0;
        }
    }

    println!("Percent the same {}%", 100f64 * (num_different) / trials as f64);
    println!("Should be about 5/9th or {}%", 100.0 * 5.0/9.0);

}
<|endoftext|>


<commit_before>#![allow(unused_imports)]

extern crate ansi_term;
extern crate colored;

use ansi_term::*;
use colored::*;

macro_rules! test_simple_color {
    ($string:expr, $colored_name:ident, $ansi_term_name:ident) => {
        #[test]
        fn $colored_name() {
            let s = format!("{} {}", $string, stringify!($colored_name));
            assert_eq!(
                s.$colored_name().to_string(),
                Colour::$ansi_term_name.paint(s).to_string()
            )
        }
    };
}

mod compat_colors {
    use super::ansi_term::*;
    use super::colored::*;

    test_simple_color!("test string", black, Black);
    test_simple_color!("test string", red, Red);
    test_simple_color!("test string", green, Green);
    test_simple_color!("test string", yellow, Yellow);
    test_simple_color!("test string", blue, Blue);
    test_simple_color!("test string", magenta, Purple);
    test_simple_color!("test string", cyan, Cyan);
    test_simple_color!("test string", white, White);
}

macro_rules! test_simple_style {
    ($string:expr, $style:ident) => {
        #[test]
        fn $style() {
            let s = format!("{} {}", $string, stringify!($style));
            assert_eq!(
                s.$style().to_string(),
                ansi_term::Style::new().$style().paint(s).to_string()
            )
        }
    };
}

mod compat_styles {
    use super::ansi_term;
    use super::ansi_term::*;
    use super::colored;
    use super::colored::*;

    test_simple_style!("test string", bold);
    test_simple_style!("test string", dimmed);
    test_simple_style!("test string", italic);
    test_simple_style!("test string", underline);
    test_simple_style!("test string", blink);
    test_simple_style!("test string", reverse);
    test_simple_style!("test string", hidden);
}

macro_rules! test_simple_bgcolor {
    ($string:expr, $colored_name:ident, $ansi_term_name:ident) => {
        #[test]
        fn $colored_name() {
            let s = format!("{} {}", $string, stringify!($colored_name));
            assert_eq!(
                s.$colored_name().to_string(),
                ansi_term::Style::default()
                    .on(ansi_term::Colour::$ansi_term_name)
                    .paint(s)
                    .to_string()
            )
        }
    };
}

mod compat_bgcolors {
    use super::ansi_term;
    use super::ansi_term::*;
    use super::colored;
    use super::colored::*;

    test_simple_bgcolor!("test string", on_black, Black);
    test_simple_bgcolor!("test string", on_red, Red);
    test_simple_bgcolor!("test string", on_green, Green);
    test_simple_bgcolor!("test string", on_yellow, Yellow);
    test_simple_bgcolor!("test string", on_blue, Blue);
    test_simple_bgcolor!("test string", on_magenta, Purple);
    test_simple_bgcolor!("test string", on_cyan, Cyan);
    test_simple_bgcolor!("test string", on_white, White);
}

mod compat_complex {
    use super::ansi_term;
    use super::ansi_term::*;
    use super::colored;
    use super::colored::*;

    #[test]
    fn complex1() {
        let s = "test string";
        let ansi = Colour::Red.on(Colour::Black).bold().italic().paint(s);
        assert_eq!(
            ansi.to_string(),
            s.red().bold().italic().on_black().to_string()
        );
    }

    #[test]
    fn complex2() {
        let s = "test string";
        let ansi = Colour::Green.on(Colour::Yellow).underline().paint(s);
        assert_eq!(
            ansi.to_string(),
            s.green().on_yellow().underline().to_string()
        );
    }
}
<commit_msg>test: Add testing for overriding an old style<commit_after>#![allow(unused_imports)]

extern crate ansi_term;
extern crate colored;

use ansi_term::*;
use colored::*;

macro_rules! test_simple_color {
    ($string:expr, $colored_name:ident, $ansi_term_name:ident) => {
        #[test]
        fn $colored_name() {
            let s = format!("{} {}", $string, stringify!($colored_name));
            assert_eq!(
                s.$colored_name().to_string(),
                Colour::$ansi_term_name.paint(s).to_string()
            )
        }
    };
}

mod compat_colors {
    use super::ansi_term::*;
    use super::colored::*;

    test_simple_color!("test string", black, Black);
    test_simple_color!("test string", red, Red);
    test_simple_color!("test string", green, Green);
    test_simple_color!("test string", yellow, Yellow);
    test_simple_color!("test string", blue, Blue);
    test_simple_color!("test string", magenta, Purple);
    test_simple_color!("test string", cyan, Cyan);
    test_simple_color!("test string", white, White);
}

macro_rules! test_simple_style {
    ($string:expr, $style:ident) => {
        #[test]
        fn $style() {
            let s = format!("{} {}", $string, stringify!($style));
            assert_eq!(
                s.$style().to_string(),
                ansi_term::Style::new().$style().paint(s).to_string()
            )
        }
    };
}

mod compat_styles {
    use super::ansi_term;
    use super::ansi_term::*;
    use super::colored;
    use super::colored::*;

    test_simple_style!("test string", bold);
    test_simple_style!("test string", dimmed);
    test_simple_style!("test string", italic);
    test_simple_style!("test string", underline);
    test_simple_style!("test string", blink);
    test_simple_style!("test string", reverse);
    test_simple_style!("test string", hidden);
}

macro_rules! test_simple_bgcolor {
    ($string:expr, $colored_name:ident, $ansi_term_name:ident) => {
        #[test]
        fn $colored_name() {
            let s = format!("{} {}", $string, stringify!($colored_name));
            assert_eq!(
                s.$colored_name().to_string(),
                ansi_term::Style::default()
                    .on(ansi_term::Colour::$ansi_term_name)
                    .paint(s)
                    .to_string()
            )
        }
    };
}

mod compat_bgcolors {
    use super::ansi_term;
    use super::ansi_term::*;
    use super::colored;
    use super::colored::*;

    test_simple_bgcolor!("test string", on_black, Black);
    test_simple_bgcolor!("test string", on_red, Red);
    test_simple_bgcolor!("test string", on_green, Green);
    test_simple_bgcolor!("test string", on_yellow, Yellow);
    test_simple_bgcolor!("test string", on_blue, Blue);
    test_simple_bgcolor!("test string", on_magenta, Purple);
    test_simple_bgcolor!("test string", on_cyan, Cyan);
    test_simple_bgcolor!("test string", on_white, White);
}

mod compat_complex {
    use super::ansi_term;
    use super::ansi_term::*;
    use super::colored;
    use super::colored::*;

    #[test]
    fn complex1() {
        let s = "test string";
        let ansi = Colour::Red.on(Colour::Black).bold().italic().paint(s);
        assert_eq!(
            ansi.to_string(),
            s.red().bold().italic().on_black().to_string()
        );
    }

    #[test]
    fn complex2() {
        let s = "test string";
        let ansi = Colour::Green.on(Colour::Yellow).underline().paint(s);
        assert_eq!(
            ansi.to_string(),
            s.green().on_yellow().underline().to_string()
        );
    }
}

mod compat_overrides {
    use super::ansi_term;
    use super::ansi_term::*;
    use super::colored;
    use super::colored::*;

    #[test]
    fn overrides1() {
        let s = "test string";
        let ansi = Colour::Red.on(Colour::Black).on(Colour::Blue).paint(s);
        assert_eq!(
            ansi.to_string(),
            s.red().on_blue().to_string()
        );
    }

    #[test]
    fn overrides2() {
        let s = "test string";
        let ansi = Colour::Green.on(Colour::Yellow).paint(s);
        assert_eq!(
            ansi.to_string(),
            s.green().on_yellow().green().on_yellow().to_string()
        );
    }
}
<|endoftext|>


<commit_before>use self::block_producer_service::{BeaconBlockGrpcClient, BlockProducerService};
use self::duties::{DutiesManager, DutiesManagerService, EpochDutiesMap};
use crate::attester_service::{AttestationGrpcClient, AttesterService};
use crate::config::ClientConfig;
use attester::test_utils::EpochMap;
use attester::{test_utils::LocalSigner as AttesterLocalSigner, Attester};
use block_proposer::{test_utils::LocalSigner as BlockProposerLocalSigner, BlockProducer};
use bls::Keypair;
use clap::{App, Arg};
use grpcio::{ChannelBuilder, EnvBuilder};
use protos::services_grpc::{
    AttestationServiceClient, BeaconBlockServiceClient, ValidatorServiceClient,
};
use slog::{error, info, o, Drain};
use slot_clock::SystemTimeSlotClock;
use std::path::PathBuf;
use std::sync::Arc;
use std::thread;
use types::ChainSpec;

mod attester_service;
mod block_producer_service;
mod config;
mod duties;

fn main() {
    // Logging
    let decorator = slog_term::TermDecorator::new().build();
    let drain = slog_term::CompactFormat::new(decorator).build().fuse();
    let drain = slog_async::Async::new(drain).build().fuse();
    let log = slog::Logger::root(drain, o!());

    // CLI
    let matches = App::new("Lighthouse Validator Client")
        .version("0.0.1")
        .author("Sigma Prime <contact@sigmaprime.io>")
        .about("Eth 2.0 Validator Client")
        .arg(
            Arg::with_name("datadir")
                .long("datadir")
                .value_name("DIR")
                .help("Data directory for keys and databases.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("server")
                .long("server")
                .value_name("server")
                .help("Address to connect to BeaconNode.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("spec")
                .long("spec")
                .value_name("spec")
                .short("s")
                .help("Configuration of Beacon Chain")
                .takes_value(true)
                .possible_values(&["foundation", "few_validators"])
                .default_value("foundation"),
        )
        .get_matches();

    let mut config = ClientConfig::default();

    // Custom datadir
    if let Some(dir) = matches.value_of("datadir") {
        config.data_dir = PathBuf::from(dir.to_string());
    }

    // Custom server port
    if let Some(server_str) = matches.value_of("server") {
        if let Ok(addr) = server_str.parse::<u16>() {
            config.server = addr.to_string();
        } else {
            error!(log, "Invalid address"; "server" => server_str);
            return;
        }
    }

    // TODO: Permit loading a custom spec from file.
    // Custom spec
    if let Some(spec_str) = matches.value_of("spec") {
        match spec_str {
            "foundation" => config.spec = ChainSpec::foundation(),
            "few_validators" => config.spec = ChainSpec::few_validators(),
            // Should be impossible due to clap's `possible_values(..)` function.
            _ => unreachable!(),
        };
    }

    // Log configuration
    info!(log, "";
          "data_dir" => &config.data_dir.to_str(),
          "server" => &config.server);

    // Beacon node gRPC beacon block endpoints.
    let beacon_block_grpc_client = {
        let env = Arc::new(EnvBuilder::new().build());
        let ch = ChannelBuilder::new(env).connect(&config.server);
        Arc::new(BeaconBlockServiceClient::new(ch))
    };

    // Beacon node gRPC validator endpoints.
    let validator_grpc_client = {
        let env = Arc::new(EnvBuilder::new().build());
        let ch = ChannelBuilder::new(env).connect(&config.server);
        Arc::new(ValidatorServiceClient::new(ch))
    };

    //Beacon node gRPC attester endpoints.
    let attester_grpc_client = {
        let env = Arc::new(EnvBuilder::new().build());
        let ch = ChannelBuilder::new(env).connect(&config.server);
        Arc::new(AttestationServiceClient::new(ch))
    };

    // Spec
    let spec = Arc::new(config.spec.clone());

    // Clock for determining the present slot.
    // TODO: this shouldn't be a static time, instead it should be pulled from the beacon node.
    // https://github.com/sigp/lighthouse/issues/160
    let genesis_time = 1_549_935_547;
    let slot_clock = {
        info!(log, "Genesis time"; "unix_epoch_seconds" => genesis_time);
        let clock = SystemTimeSlotClock::new(genesis_time, spec.seconds_per_slot)
            .expect("Unable to instantiate SystemTimeSlotClock.");
        Arc::new(clock)
    };

    let poll_interval_millis = spec.seconds_per_slot * 1000 / 10; // 10% epoch time precision.
    info!(log, "Starting block producer service"; "polls_per_epoch" => spec.seconds_per_slot * 1000 / poll_interval_millis);

    /*
     * Start threads.
     */
    let mut threads = vec![];
    // TODO: keypairs are randomly generated; they should be loaded from a file or generated.
    // https://github.com/sigp/lighthouse/issues/160
    let keypairs = vec![Keypair::random()];

    for keypair in keypairs {
        info!(log, "Starting validator services"; "validator" => keypair.pk.concatenated_hex_id());
        let duties_map = Arc::new(EpochDutiesMap::new(spec.slots_per_epoch));
        let epoch_map_for_attester = Arc::new(EpochMap::new(spec.slots_per_epoch));

        // Spawn a new thread to maintain the validator's `EpochDuties`.
        let duties_manager_thread = {
            let spec = spec.clone();
            let duties_map = duties_map.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let beacon_node = validator_grpc_client.clone();
            let pubkey = keypair.pk.clone();
            thread::spawn(move || {
                let manager = DutiesManager {
                    duties_map,
                    pubkey,
                    spec,
                    slot_clock,
                    beacon_node,
                };
                let mut duties_manager_service = DutiesManagerService {
                    manager,
                    poll_interval_millis,
                    log,
                };

                duties_manager_service.run();
            })
        };

        // Spawn a new thread to perform block production for the validator.
        let producer_thread = {
            let spec = spec.clone();
            let signer = Arc::new(BlockProposerLocalSigner::new(keypair.clone()));
            let duties_map = duties_map.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let client = Arc::new(BeaconBlockGrpcClient::new(beacon_block_grpc_client.clone()));
            thread::spawn(move || {
                let block_producer =
                    BlockProducer::new(spec, duties_map, slot_clock, client, signer);
                let mut block_producer_service = BlockProducerService {
                    block_producer,
                    poll_interval_millis,
                    log,
                };

                block_producer_service.run();
            })
        };

        // Spawn a new thread for attestation for the validator.
        let attester_thread = {
            let signer = Arc::new(AttesterLocalSigner::new(keypair.clone()));
            let epoch_map = epoch_map_for_attester.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let client = Arc::new(AttestationGrpcClient::new(attester_grpc_client.clone()));
            thread::spawn(move || {
                let attester = Attester::new(epoch_map, slot_clock, client, signer);
                let mut attester_service = AttesterService {
                    attester,
                    poll_interval_millis,
                    log,
                };

                attester_service.run();
            })
        };

        threads.push((duties_manager_thread, producer_thread, attester_thread));
    }

    // Naively wait for all the threads to complete.
    for tuple in threads {
        let (manager, producer, attester) = tuple;
        let _ = producer.join();
        let _ = manager.join();
        let _ = attester.join();
    }
}
<commit_msg>Add BeaconNodeInfo RPC to validator client<commit_after>use crate::attester_service::{AttestationGrpcClient, AttesterService};
use crate::block_producer_service::{BeaconBlockGrpcClient, BlockProducerService};
use crate::config::ClientConfig;
use crate::duties::{DutiesManager, DutiesManagerService, EpochDutiesMap};
use attester::test_utils::EpochMap;
use attester::{test_utils::LocalSigner as AttesterLocalSigner, Attester};
use block_proposer::{test_utils::LocalSigner as BlockProposerLocalSigner, BlockProducer};
use bls::Keypair;
use clap::{App, Arg};
use grpcio::{ChannelBuilder, EnvBuilder};
use protos::services::{Empty, NodeInfo};
use protos::services_grpc::{
    AttestationServiceClient, BeaconBlockServiceClient, BeaconNodeServiceClient,
    ValidatorServiceClient,
};
use slog::{error, info, o, Drain};
use slot_clock::SystemTimeSlotClock;
use std::path::PathBuf;
use std::sync::Arc;
use std::thread;
use types::ChainSpec;

mod attester_service;
mod block_producer_service;
mod config;
mod duties;

fn main() {
    // Logging
    let decorator = slog_term::TermDecorator::new().build();
    let drain = slog_term::CompactFormat::new(decorator).build().fuse();
    let drain = slog_async::Async::new(drain).build().fuse();
    let log = slog::Logger::root(drain, o!());

    // CLI
    let matches = App::new("Lighthouse Validator Client")
        .version("0.0.1")
        .author("Sigma Prime <contact@sigmaprime.io>")
        .about("Eth 2.0 Validator Client")
        .arg(
            Arg::with_name("datadir")
                .long("datadir")
                .value_name("DIR")
                .help("Data directory for keys and databases.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("server")
                .long("server")
                .value_name("server")
                .help("Address to connect to BeaconNode.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("spec")
                .long("spec")
                .value_name("spec")
                .short("s")
                .help("Configuration of Beacon Chain")
                .takes_value(true)
                .possible_values(&["foundation", "few_validators"])
                .default_value("foundation"),
        )
        .get_matches();

    let mut config = ClientConfig::default();

    // Custom datadir
    if let Some(dir) = matches.value_of("datadir") {
        config.data_dir = PathBuf::from(dir.to_string());
    }

    // Custom server port
    if let Some(server_str) = matches.value_of("server") {
        if let Ok(addr) = server_str.parse::<u16>() {
            config.server = addr.to_string();
        } else {
            error!(log, "Invalid address"; "server" => server_str);
            return;
        }
    }

    // TODO: Permit loading a custom spec from file.
    // Custom spec
    if let Some(spec_str) = matches.value_of("spec") {
        match spec_str {
            "foundation" => config.spec = ChainSpec::foundation(),
            "few_validators" => config.spec = ChainSpec::few_validators(),
            // Should be impossible due to clap's `possible_values(..)` function.
            _ => unreachable!(),
        };
    }

    // Log configuration
    info!(log, "";
          "data_dir" => &config.data_dir.to_str(),
          "server" => &config.server);

    // Beacon node gRPC beacon node endpoints.
    let beacon_node_grpc_client = {
        let env = Arc::new(EnvBuilder::new().build());
        let ch = ChannelBuilder::new(env).connect(&config.server);
        Arc::new(BeaconNodeServiceClient::new(ch))
    };

    // Beacon node gRPC beacon block endpoints.
    let beacon_block_grpc_client = {
        let env = Arc::new(EnvBuilder::new().build());
        let ch = ChannelBuilder::new(env).connect(&config.server);
        Arc::new(BeaconBlockServiceClient::new(ch))
    };

    // Beacon node gRPC validator endpoints.
    let validator_grpc_client = {
        let env = Arc::new(EnvBuilder::new().build());
        let ch = ChannelBuilder::new(env).connect(&config.server);
        Arc::new(ValidatorServiceClient::new(ch))
    };

    //Beacon node gRPC attester endpoints.
    let attester_grpc_client = {
        let env = Arc::new(EnvBuilder::new().build());
        let ch = ChannelBuilder::new(env).connect(&config.server);
        Arc::new(AttestationServiceClient::new(ch))
    };

    // retrieve node information
    let node_info = beacon_node_grpc_client.info(&Empty::new());

    info!(log, "Beacon node info: {:?}", node_info);

    // Spec
    let spec = Arc::new(config.spec.clone());

    let genesis_time = 1_549_935_547;
    let slot_clock = {
        info!(log, "Genesis time"; "unix_epoch_seconds" => genesis_time);
        let clock = SystemTimeSlotClock::new(genesis_time, spec.seconds_per_slot)
            .expect("Unable to instantiate SystemTimeSlotClock.");
        Arc::new(clock)
    };

    let poll_interval_millis = spec.seconds_per_slot * 1000 / 10; // 10% epoch time precision.
    info!(log, "Starting block producer service"; "polls_per_epoch" => spec.seconds_per_slot * 1000 / poll_interval_millis);

    /*
     * Start threads.
     */
    let mut threads = vec![];
    // TODO: keypairs are randomly generated; they should be loaded from a file or generated.
    // https://github.com/sigp/lighthouse/issues/160
    let keypairs = vec![Keypair::random()];

    for keypair in keypairs {
        info!(log, "Starting validator services"; "validator" => keypair.pk.concatenated_hex_id());
        let duties_map = Arc::new(EpochDutiesMap::new(spec.slots_per_epoch));
        let epoch_map_for_attester = Arc::new(EpochMap::new(spec.slots_per_epoch));

        // Spawn a new thread to maintain the validator's `EpochDuties`.
        let duties_manager_thread = {
            let spec = spec.clone();
            let duties_map = duties_map.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let beacon_node = validator_grpc_client.clone();
            let pubkey = keypair.pk.clone();
            thread::spawn(move || {
                let manager = DutiesManager {
                    duties_map,
                    pubkey,
                    spec,
                    slot_clock,
                    beacon_node,
                };
                let mut duties_manager_service = DutiesManagerService {
                    manager,
                    poll_interval_millis,
                    log,
                };

                duties_manager_service.run();
            })
        };

        // Spawn a new thread to perform block production for the validator.
        let producer_thread = {
            let spec = spec.clone();
            let signer = Arc::new(BlockProposerLocalSigner::new(keypair.clone()));
            let duties_map = duties_map.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let client = Arc::new(BeaconBlockGrpcClient::new(beacon_block_grpc_client.clone()));
            thread::spawn(move || {
                let block_producer =
                    BlockProducer::new(spec, duties_map, slot_clock, client, signer);
                let mut block_producer_service = BlockProducerService {
                    block_producer,
                    poll_interval_millis,
                    log,
                };

                block_producer_service.run();
            })
        };

        // Spawn a new thread for attestation for the validator.
        let attester_thread = {
            let signer = Arc::new(AttesterLocalSigner::new(keypair.clone()));
            let epoch_map = epoch_map_for_attester.clone();
            let slot_clock = slot_clock.clone();
            let log = log.clone();
            let client = Arc::new(AttestationGrpcClient::new(attester_grpc_client.clone()));
            thread::spawn(move || {
                let attester = Attester::new(epoch_map, slot_clock, client, signer);
                let mut attester_service = AttesterService {
                    attester,
                    poll_interval_millis,
                    log,
                };

                attester_service.run();
            })
        };

        threads.push((duties_manager_thread, producer_thread, attester_thread));
    }

    // Naively wait for all the threads to complete.
    for tuple in threads {
        let (manager, producer, attester) = tuple;
        let _ = producer.join();
        let _ = manager.join();
        let _ = attester.join();
    }
}
<|endoftext|>