<commit_before>pub type blkcnt_t = i64;
pub type blksize_t = i64;
pub type c_char = i8;
pub type c_long = i64;
pub type c_ulong = u64;
pub type ino_t = u64;
pub type nlink_t = u64;
pub type off_t = i64;
pub type suseconds_t = i64;
pub type time_t = i64;
pub type wchar_t = i32;
pub type clock_t = i64;

s! {
    pub struct aiocb {
        pub aio_fildes: ::c_int,
        pub aio_lio_opcode: ::c_int,
        pub aio_reqprio: ::c_int,
        pub aio_buf: *mut ::c_void,
        pub aio_nbytes: ::size_t,
        pub aio_sigevent: ::sigevent,
        __next_prio: *mut aiocb,
        __abs_prio: ::c_int,
        __policy: ::c_int,
        __error_code: ::c_int,
        __return_value: ::ssize_t,
        pub aio_offset: off_t,
        __glibc_reserved: [::c_char; 32]
    }

    pub struct stat {
        pub st_dev: ::c_ulong,
        st_pad1: [::c_long; 2],
        pub st_ino: ::ino_t,
        pub st_mode: ::mode_t,
        pub st_nlink: ::nlink_t,
        pub st_uid: ::uid_t,
        pub st_gid: ::gid_t,
        pub st_rdev: ::c_ulong,
        st_pad2: [::c_ulong; 1],
        pub st_size: ::off_t,
        st_pad3: ::c_long,
        pub st_atime: ::time_t,
        pub st_atime_nsec: ::c_long,
        pub st_mtime: ::time_t,
        pub st_mtime_nsec: ::c_long,
        pub st_ctime: ::time_t,
        pub st_ctime_nsec: ::c_long,
        pub st_blksize: ::blksize_t,
        st_pad4: ::c_long,
        pub st_blocks: ::blkcnt_t,
        st_pad5: [::c_long; 7],
    }

    pub struct stat64 {
        pub st_dev: ::c_ulong,
        st_pad1: [::c_long; 2],
        pub st_ino: ::ino64_t,
        pub st_mode: ::mode_t,
        pub st_nlink: ::nlink_t,
        pub st_uid: ::uid_t,
        pub st_gid: ::gid_t,
        pub st_rdev: ::c_ulong,
        st_pad2: [::c_long; 2],
        pub st_size: ::off64_t,
        pub st_atime: ::time_t,
        pub st_atime_nsec: ::c_long,
        pub st_mtime: ::time_t,
        pub st_mtime_nsec: ::c_long,
        pub st_ctime: ::time_t,
        pub st_ctime_nsec: ::c_long,
        pub st_blksize: ::blksize_t,
        st_pad3: ::c_long,
        pub st_blocks: ::blkcnt64_t,
        st_pad5: [::c_long; 7],
    }

    pub struct pthread_attr_t {
        __size: [::c_ulong; 7]
    }

    pub struct sigaction {
        pub sa_flags: ::c_int,
        pub sa_sigaction: ::sighandler_t,
        pub sa_mask: sigset_t,
        pub sa_restorer: ::dox::Option<extern fn()>,
    }

    pub struct stack_t {
        pub ss_sp: *mut ::c_void,
        pub ss_size: ::size_t,
        pub ss_flags: ::c_int,
    }

    pub struct sigset_t {
        __size: [::c_ulong; 16],
    }

    pub struct siginfo_t {
        pub si_signo: ::c_int,
        pub si_code: ::c_int,
        pub si_errno: ::c_int,
        _pad: ::c_int,
        _pad2: [::c_long; 14],
    }

    pub struct ipc_perm {
        pub __key: ::key_t,
        pub uid: ::uid_t,
        pub gid: ::gid_t,
        pub cuid: ::uid_t,
        pub cgid: ::gid_t,
        pub mode: ::c_uint,
        pub __seq: ::c_ushort,
        __pad1: ::c_ushort,
        __unused1: ::c_ulong,
        __unused2: ::c_ulong
    }

    pub struct shmid_ds {
        pub shm_perm: ::ipc_perm,
        pub shm_segsz: ::size_t,
        pub shm_atime: ::time_t,
        pub shm_dtime: ::time_t,
        pub shm_ctime: ::time_t,
        pub shm_cpid: ::pid_t,
        pub shm_lpid: ::pid_t,
        pub shm_nattch: ::shmatt_t,
        __unused4: ::c_ulong,
        __unused5: ::c_ulong
    }

    pub struct msqid_ds {
        pub msg_perm: ::ipc_perm,
        pub msg_stime: ::time_t,
        pub msg_rtime: ::time_t,
        pub msg_ctime: ::time_t,
        __msg_cbytes: ::c_ulong,
        pub msg_qnum: ::msgqnum_t,
        pub msg_qbytes: ::msglen_t,
        pub msg_lspid: ::pid_t,
        pub msg_lrpid: ::pid_t,
        __glibc_reserved4: ::c_ulong,
        __glibc_reserved5: ::c_ulong,
    }

    pub struct statfs {
        pub f_type: ::c_long,
        pub f_bsize: ::c_long,
        pub f_frsize: ::c_long,
        pub f_blocks: ::fsblkcnt_t,
        pub f_bfree: ::fsblkcnt_t,
        pub f_files: ::fsblkcnt_t,
        pub f_ffree: ::fsblkcnt_t,
        pub f_bavail: ::fsblkcnt_t,
        pub f_fsid: ::fsid_t,

        pub f_namelen: ::c_long,
        f_spare: [::c_long; 6],
    }

    pub struct msghdr {
        pub msg_name: *mut ::c_void,
        pub msg_namelen: ::socklen_t,
        pub msg_iov: *mut ::iovec,
        pub msg_iovlen: ::size_t,
        pub msg_control: *mut ::c_void,
        pub msg_controllen: ::size_t,
        pub msg_flags: ::c_int,
    }

    pub struct cmsghdr {
        pub cmsg_len: ::size_t,
        pub cmsg_level: ::c_int,
        pub cmsg_type: ::c_int,
    }

    pub struct termios {
        pub c_iflag: ::tcflag_t,
        pub c_oflag: ::tcflag_t,
        pub c_cflag: ::tcflag_t,
        pub c_lflag: ::tcflag_t,
        pub c_line: ::cc_t,
        pub c_cc: [::cc_t; ::NCCS],
    }

    pub struct sysinfo {
        pub uptime: ::c_long,
        pub loads: [::c_ulong; 3],
        pub totalram: ::c_ulong,
        pub freeram: ::c_ulong,
        pub sharedram: ::c_ulong,
        pub bufferram: ::c_ulong,
        pub totalswap: ::c_ulong,
        pub freeswap: ::c_ulong,
        pub procs: ::c_ushort,
        pub pad: ::c_ushort,
        pub totalhigh: ::c_ulong,
        pub freehigh: ::c_ulong,
        pub mem_unit: ::c_uint,
        pub _f: [::c_char; 0],
    }
}

pub const __SIZEOF_PTHREAD_CONDATTR_T: usize = 4;
pub const __SIZEOF_PTHREAD_MUTEXATTR_T: usize = 4;
pub const __SIZEOF_PTHREAD_MUTEX_T: usize = 40;
pub const __SIZEOF_PTHREAD_RWLOCK_T: usize = 56;

pub const RLIM_INFINITY: ::rlim_t = 0xffff_ffff_ffff_ffff;

pub const SYS_gettid: ::c_long = 5178;   // Valid for n64
<commit_msg>Add missing struct flock to mips64<commit_after>pub type blkcnt_t = i64;
pub type blksize_t = i64;
pub type c_char = i8;
pub type c_long = i64;
pub type c_ulong = u64;
pub type ino_t = u64;
pub type nlink_t = u64;
pub type off_t = i64;
pub type suseconds_t = i64;
pub type time_t = i64;
pub type wchar_t = i32;
pub type clock_t = i64;

s! {
    pub struct aiocb {
        pub aio_fildes: ::c_int,
        pub aio_lio_opcode: ::c_int,
        pub aio_reqprio: ::c_int,
        pub aio_buf: *mut ::c_void,
        pub aio_nbytes: ::size_t,
        pub aio_sigevent: ::sigevent,
        __next_prio: *mut aiocb,
        __abs_prio: ::c_int,
        __policy: ::c_int,
        __error_code: ::c_int,
        __return_value: ::ssize_t,
        pub aio_offset: off_t,
        __glibc_reserved: [::c_char; 32]
    }

    pub struct stat {
        pub st_dev: ::c_ulong,
        st_pad1: [::c_long; 2],
        pub st_ino: ::ino_t,
        pub st_mode: ::mode_t,
        pub st_nlink: ::nlink_t,
        pub st_uid: ::uid_t,
        pub st_gid: ::gid_t,
        pub st_rdev: ::c_ulong,
        st_pad2: [::c_ulong; 1],
        pub st_size: ::off_t,
        st_pad3: ::c_long,
        pub st_atime: ::time_t,
        pub st_atime_nsec: ::c_long,
        pub st_mtime: ::time_t,
        pub st_mtime_nsec: ::c_long,
        pub st_ctime: ::time_t,
        pub st_ctime_nsec: ::c_long,
        pub st_blksize: ::blksize_t,
        st_pad4: ::c_long,
        pub st_blocks: ::blkcnt_t,
        st_pad5: [::c_long; 7],
    }

    pub struct stat64 {
        pub st_dev: ::c_ulong,
        st_pad1: [::c_long; 2],
        pub st_ino: ::ino64_t,
        pub st_mode: ::mode_t,
        pub st_nlink: ::nlink_t,
        pub st_uid: ::uid_t,
        pub st_gid: ::gid_t,
        pub st_rdev: ::c_ulong,
        st_pad2: [::c_long; 2],
        pub st_size: ::off64_t,
        pub st_atime: ::time_t,
        pub st_atime_nsec: ::c_long,
        pub st_mtime: ::time_t,
        pub st_mtime_nsec: ::c_long,
        pub st_ctime: ::time_t,
        pub st_ctime_nsec: ::c_long,
        pub st_blksize: ::blksize_t,
        st_pad3: ::c_long,
        pub st_blocks: ::blkcnt64_t,
        st_pad5: [::c_long; 7],
    }

    pub struct pthread_attr_t {
        __size: [::c_ulong; 7]
    }

    pub struct sigaction {
        pub sa_flags: ::c_int,
        pub sa_sigaction: ::sighandler_t,
        pub sa_mask: sigset_t,
        pub sa_restorer: ::dox::Option<extern fn()>,
    }

    pub struct stack_t {
        pub ss_sp: *mut ::c_void,
        pub ss_size: ::size_t,
        pub ss_flags: ::c_int,
    }

    pub struct sigset_t {
        __size: [::c_ulong; 16],
    }

    pub struct siginfo_t {
        pub si_signo: ::c_int,
        pub si_code: ::c_int,
        pub si_errno: ::c_int,
        _pad: ::c_int,
        _pad2: [::c_long; 14],
    }

    pub struct ipc_perm {
        pub __key: ::key_t,
        pub uid: ::uid_t,
        pub gid: ::gid_t,
        pub cuid: ::uid_t,
        pub cgid: ::gid_t,
        pub mode: ::c_uint,
        pub __seq: ::c_ushort,
        __pad1: ::c_ushort,
        __unused1: ::c_ulong,
        __unused2: ::c_ulong
    }

    pub struct shmid_ds {
        pub shm_perm: ::ipc_perm,
        pub shm_segsz: ::size_t,
        pub shm_atime: ::time_t,
        pub shm_dtime: ::time_t,
        pub shm_ctime: ::time_t,
        pub shm_cpid: ::pid_t,
        pub shm_lpid: ::pid_t,
        pub shm_nattch: ::shmatt_t,
        __unused4: ::c_ulong,
        __unused5: ::c_ulong
    }

    pub struct msqid_ds {
        pub msg_perm: ::ipc_perm,
        pub msg_stime: ::time_t,
        pub msg_rtime: ::time_t,
        pub msg_ctime: ::time_t,
        __msg_cbytes: ::c_ulong,
        pub msg_qnum: ::msgqnum_t,
        pub msg_qbytes: ::msglen_t,
        pub msg_lspid: ::pid_t,
        pub msg_lrpid: ::pid_t,
        __glibc_reserved4: ::c_ulong,
        __glibc_reserved5: ::c_ulong,
    }

    pub struct statfs {
        pub f_type: ::c_long,
        pub f_bsize: ::c_long,
        pub f_frsize: ::c_long,
        pub f_blocks: ::fsblkcnt_t,
        pub f_bfree: ::fsblkcnt_t,
        pub f_files: ::fsblkcnt_t,
        pub f_ffree: ::fsblkcnt_t,
        pub f_bavail: ::fsblkcnt_t,
        pub f_fsid: ::fsid_t,

        pub f_namelen: ::c_long,
        f_spare: [::c_long; 6],
    }

    pub struct msghdr {
        pub msg_name: *mut ::c_void,
        pub msg_namelen: ::socklen_t,
        pub msg_iov: *mut ::iovec,
        pub msg_iovlen: ::size_t,
        pub msg_control: *mut ::c_void,
        pub msg_controllen: ::size_t,
        pub msg_flags: ::c_int,
    }

    pub struct cmsghdr {
        pub cmsg_len: ::size_t,
        pub cmsg_level: ::c_int,
        pub cmsg_type: ::c_int,
    }

    pub struct termios {
        pub c_iflag: ::tcflag_t,
        pub c_oflag: ::tcflag_t,
        pub c_cflag: ::tcflag_t,
        pub c_lflag: ::tcflag_t,
        pub c_line: ::cc_t,
        pub c_cc: [::cc_t; ::NCCS],
    }

    pub struct flock {
        pub l_type: ::c_short,
        pub l_whence: ::c_short,
        pub l_start: ::off_t,
        pub l_len: ::off_t,
        pub l_pid: ::pid_t,
    }

    pub struct sysinfo {
        pub uptime: ::c_long,
        pub loads: [::c_ulong; 3],
        pub totalram: ::c_ulong,
        pub freeram: ::c_ulong,
        pub sharedram: ::c_ulong,
        pub bufferram: ::c_ulong,
        pub totalswap: ::c_ulong,
        pub freeswap: ::c_ulong,
        pub procs: ::c_ushort,
        pub pad: ::c_ushort,
        pub totalhigh: ::c_ulong,
        pub freehigh: ::c_ulong,
        pub mem_unit: ::c_uint,
        pub _f: [::c_char; 0],
    }
}

pub const __SIZEOF_PTHREAD_CONDATTR_T: usize = 4;
pub const __SIZEOF_PTHREAD_MUTEXATTR_T: usize = 4;
pub const __SIZEOF_PTHREAD_MUTEX_T: usize = 40;
pub const __SIZEOF_PTHREAD_RWLOCK_T: usize = 56;

pub const RLIM_INFINITY: ::rlim_t = 0xffff_ffff_ffff_ffff;

pub const SYS_gettid: ::c_long = 5178;   // Valid for n64
<|endoftext|>


<commit_before>// Copyright (C) 2019 Alibaba Cloud Computing. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Traits and Structs to handle vhost-user requests from the slave to the master.

use libc;
use std::mem;
use std::os::unix::io::{AsRawFd, RawFd};
use std::os::unix::net::UnixStream;
use std::sync::{Arc, Mutex};

use super::connection::Endpoint;
use super::message::*;
use super::{Error, HandlerResult, Result};

/// Trait to handle vhost-user requests from the slave to the master.
pub trait VhostUserMasterReqHandler {
    // fn handle_iotlb_msg(&mut self, iotlb: VhostUserIotlb);
    // fn handle_vring_host_notifier(&mut self, area: VhostUserVringArea, fd: RawFd);

    /// Handle device configuration change notifications from the slave.
    fn handle_config_change(&mut self) -> HandlerResult<u64> {
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    /// Handle virtio-fs map file requests from the slave.
    fn fs_slave_map(&mut self, _fs: &VhostUserFSSlaveMsg, fd: RawFd) -> HandlerResult<u64> {
        // Safe because we have just received the rawfd from kernel.
        unsafe { libc::close(fd) };
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    /// Handle virtio-fs unmap file requests from the slave.
    fn fs_slave_unmap(&mut self, _fs: &VhostUserFSSlaveMsg) -> HandlerResult<u64> {
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    /// Handle virtio-fs sync file requests from the slave.
    fn fs_slave_sync(&mut self, _fs: &VhostUserFSSlaveMsg) -> HandlerResult<u64> {
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    /// Handle virtio-fs file IO requests from the slave.
    fn fs_slave_io(&mut self, _fs: &VhostUserFSSlaveMsg, fd: RawFd) -> HandlerResult<u64> {
        // Safe because we have just received the rawfd from kernel.
        unsafe { libc::close(fd) };
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }
}

/// A vhost-user master request endpoint which relays all received requests from the slave to the
/// provided request handler.
pub struct MasterReqHandler<S: VhostUserMasterReqHandler> {
    // underlying Unix domain socket for communication
    sub_sock: Endpoint<SlaveReq>,
    tx_sock: UnixStream,
    // the VirtIO backend device object
    backend: Arc<Mutex<S>>,
    // whether the endpoint has encountered any failure
    error: Option<i32>,
}

impl<S: VhostUserMasterReqHandler> MasterReqHandler<S> {
    /// Create a vhost-user slave request handler.
    /// This opens a pair of connected anonymous sockets.
    /// Returns Self and the socket that must be sent to the slave via SET_SLAVE_REQ_FD.
    pub fn new(backend: Arc<Mutex<S>>) -> Result<Self> {
        let (tx, rx) = UnixStream::pair().map_err(Error::SocketError)?;

        Ok(MasterReqHandler {
            sub_sock: Endpoint::<SlaveReq>::from_stream(rx),
            tx_sock: tx,
            backend,
            error: None,
        })
    }

    /// Get the raw fd to send to the slave as slave communication channel.
    pub fn get_tx_raw_fd(&self) -> RawFd {
        self.tx_sock.as_raw_fd()
    }

    /// Mark endpoint as failed or normal state.
    pub fn set_failed(&mut self, error: i32) {
        self.error = Some(error);
    }

    /// Receive and handle one incoming request message from the slave.
    /// The caller needs to:
    /// . serialize calls to this function
    /// . decide what to do when errer happens
    /// . optional recover from failure
    pub fn handle_request(&mut self) -> Result<u64> {
        // Return error if the endpoint is already in failed state.
        self.check_state()?;

        // The underlying communication channel is a Unix domain socket in
        // stream mode, and recvmsg() is a little tricky here. To successfully
        // receive attached file descriptors, we need to receive messages and
        // corresponding attached file descriptors in this way:
        // . recv messsage header and optional attached file
        // . validate message header
        // . recv optional message body and payload according size field in
        //   message header
        // . validate message body and optional payload
        let (hdr, rfds) = self.sub_sock.recv_header()?;
        let rfds = self.check_attached_rfds(&hdr, rfds)?;
        let (size, buf) = match hdr.get_size() {
            0 => (0, vec![0u8; 0]),
            len => {
                let (size2, rbuf) = self.sub_sock.recv_data(len as usize)?;
                if size2 != len as usize {
                    return Err(Error::InvalidMessage);
                }
                (size2, rbuf)
            }
        };

        let res = match hdr.get_code() {
            SlaveReq::CONFIG_CHANGE_MSG => {
                self.check_msg_size(&hdr, size, 0)?;
                self.backend
                    .lock()
                    .unwrap()
                    .handle_config_change()
                    .map_err(Error::ReqHandlerError)
            }
            SlaveReq::FS_MAP => {
                let msg = self.extract_msg_body::<VhostUserFSSlaveMsg>(&hdr, size, &buf)?;
                self.backend
                    .lock()
                    .unwrap()
                    .fs_slave_map(msg, rfds.unwrap()[0])
                    .map_err(Error::ReqHandlerError)
            }
            SlaveReq::FS_UNMAP => {
                let msg = self.extract_msg_body::<VhostUserFSSlaveMsg>(&hdr, size, &buf)?;
                self.backend
                    .lock()
                    .unwrap()
                    .fs_slave_unmap(msg)
                    .map_err(Error::ReqHandlerError)
            }
            SlaveReq::FS_SYNC => {
                let msg = self.extract_msg_body::<VhostUserFSSlaveMsg>(&hdr, size, &buf)?;
                self.backend
                    .lock()
                    .unwrap()
                    .fs_slave_sync(msg)
                    .map_err(Error::ReqHandlerError)
            }
            SlaveReq::FS_IO => {
                let msg = self.extract_msg_body::<VhostUserFSSlaveMsg>(&hdr, size, &buf)?;
                self.backend
                    .lock()
                    .unwrap()
                    .fs_slave_io(msg, rfds.unwrap()[0])
                    .map_err(Error::ReqHandlerError)
            }
            _ => Err(Error::InvalidMessage),
        };

        self.send_ack_message(&hdr, &res)?;

        res
    }

    fn check_state(&self) -> Result<()> {
        match self.error {
            Some(e) => Err(Error::SocketBroken(std::io::Error::from_raw_os_error(e))),
            None => Ok(()),
        }
    }

    fn check_msg_size(
        &self,
        hdr: &VhostUserMsgHeader<SlaveReq>,
        size: usize,
        expected: usize,
    ) -> Result<()> {
        if hdr.get_size() as usize != expected
            || hdr.is_reply()
            || hdr.get_version() != 0x1
            || size != expected
        {
            return Err(Error::InvalidMessage);
        }
        Ok(())
    }

    fn check_attached_rfds(
        &self,
        hdr: &VhostUserMsgHeader<SlaveReq>,
        rfds: Option<Vec<RawFd>>,
    ) -> Result<Option<Vec<RawFd>>> {
        match hdr.get_code() {
            SlaveReq::FS_MAP | SlaveReq::FS_IO => {
                // Expect an fd set with a single fd.
                match rfds {
                    None => Err(Error::InvalidMessage),
                    Some(fds) => {
                        if fds.len() != 1 {
                            Endpoint::<SlaveReq>::close_rfds(Some(fds));
                            Err(Error::InvalidMessage)
                        } else {
                            Ok(Some(fds))
                        }
                    }
                }
            }
            _ => {
                if rfds.is_some() {
                    Endpoint::<SlaveReq>::close_rfds(rfds);
                    return Err(Error::InvalidMessage);
                } else {
                    Ok(rfds)
                }
            }
        }
    }

    fn extract_msg_body<'a, T: Sized + VhostUserMsgValidator>(
        &self,
        hdr: &VhostUserMsgHeader<SlaveReq>,
        size: usize,
        buf: &'a [u8],
    ) -> Result<&'a T> {
        self.check_msg_size(hdr, size, mem::size_of::<T>())?;
        let msg = unsafe { &*(buf.as_ptr() as *const T) };
        if !msg.is_valid() {
            return Err(Error::InvalidMessage);
        }
        Ok(msg)
    }

    fn new_reply_header<T: Sized>(
        &self,
        req: &VhostUserMsgHeader<SlaveReq>,
    ) -> Result<VhostUserMsgHeader<SlaveReq>> {
        if mem::size_of::<T>() > MAX_MSG_SIZE {
            return Err(Error::InvalidParam);
        }
        self.check_state()?;
        Ok(VhostUserMsgHeader::new(
            req.get_code(),
            VhostUserHeaderFlag::REPLY.bits(),
            mem::size_of::<T>() as u32,
        ))
    }

    fn send_ack_message(
        &mut self,
        req: &VhostUserMsgHeader<SlaveReq>,
        res: &Result<u64>,
    ) -> Result<()> {
        if req.is_need_reply() {
            let hdr = self.new_reply_header::<VhostUserU64>(req)?;
            let def_err = libc::EINVAL;
            let val = match res {
                Ok(n) => *n,
                Err(e) => match &*e {
                    Error::ReqHandlerError(ioerr) => match ioerr.raw_os_error() {
                        Some(rawerr) => -rawerr as u64,
                        None => -def_err as u64,
                    },
                    _ => -def_err as u64,
                },
            };
            let msg = VhostUserU64::new(val);
            self.sub_sock.send_message(&hdr, &msg, None)?;
        }
        Ok(())
    }
}

impl<S: VhostUserMasterReqHandler> AsRawFd for MasterReqHandler<S> {
    fn as_raw_fd(&self) -> RawFd {
        self.sub_sock.as_raw_fd()
    }
}
<commit_msg>vhost-user: refine the MasterReqHandler trait<commit_after>// Copyright (C) 2019-2021 Alibaba Cloud. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use libc;
use std::mem;
use std::os::unix::io::{AsRawFd, RawFd};
use std::os::unix::net::UnixStream;
use std::sync::{Arc, Mutex};

use super::connection::Endpoint;
use super::message::*;
use super::{Error, HandlerResult, Result};

/// Define services provided by masters for the slave communication channel.
///
/// The vhost-user specification defines a slave communication channel, by which slaves could
/// request services from masters. The [VhostUserMasterReqHandler] trait defines services provided
/// by masters, and it's used both on the master side and slave side.
/// - on the slave side, a stub forwarder implementing [VhostUserMasterReqHandler] will proxy
///   service requests to masters. The [SlaveFsCacheReq] is an example stub forwarder.
/// - on the master side, the [MasterReqHandler] will forward service requests to a handler
///   implementing [VhostUserMasterReqHandler].
///
/// The [VhostUserMasterReqHandler] trait is design with interior mutability to improve performance
/// for multi-threading.
///
/// [VhostUserMasterReqHandler]: trait.VhostUserMasterReqHandler.html
/// [MasterReqHandler]: struct.MasterReqHandler.html
/// [SlaveFsCacheReq]: struct.SlaveFsCacheReq.html
pub trait VhostUserMasterReqHandler {
    /// Handle device configuration change notifications.
    fn handle_config_change(&self) -> HandlerResult<u64> {
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    /// Handle virtio-fs map file requests.
    fn fs_slave_map(&self, _fs: &VhostUserFSSlaveMsg, fd: RawFd) -> HandlerResult<u64> {
        // Safe because we have just received the rawfd from kernel.
        unsafe { libc::close(fd) };
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    /// Handle virtio-fs unmap file requests.
    fn fs_slave_unmap(&self, _fs: &VhostUserFSSlaveMsg) -> HandlerResult<u64> {
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    /// Handle virtio-fs sync file requests.
    fn fs_slave_sync(&self, _fs: &VhostUserFSSlaveMsg) -> HandlerResult<u64> {
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    /// Handle virtio-fs file IO requests.
    fn fs_slave_io(&self, _fs: &VhostUserFSSlaveMsg, fd: RawFd) -> HandlerResult<u64> {
        // Safe because we have just received the rawfd from kernel.
        unsafe { libc::close(fd) };
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    // fn handle_iotlb_msg(&mut self, iotlb: VhostUserIotlb);
    // fn handle_vring_host_notifier(&mut self, area: VhostUserVringArea, fd: RawFd);
}

/// A helper trait mirroring [VhostUserMasterReqHandler] but without interior mutability.
///
/// [VhostUserMasterReqHandler]: trait.VhostUserMasterReqHandler.html
pub trait VhostUserMasterReqHandlerMut {
    /// Handle device configuration change notifications.
    fn handle_config_change(&mut self) -> HandlerResult<u64> {
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    /// Handle virtio-fs map file requests.
    fn fs_slave_map(&mut self, _fs: &VhostUserFSSlaveMsg, fd: RawFd) -> HandlerResult<u64> {
        // Safe because we have just received the rawfd from kernel.
        unsafe { libc::close(fd) };
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    /// Handle virtio-fs unmap file requests.
    fn fs_slave_unmap(&mut self, _fs: &VhostUserFSSlaveMsg) -> HandlerResult<u64> {
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    /// Handle virtio-fs sync file requests.
    fn fs_slave_sync(&mut self, _fs: &VhostUserFSSlaveMsg) -> HandlerResult<u64> {
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    /// Handle virtio-fs file IO requests.
    fn fs_slave_io(&mut self, _fs: &VhostUserFSSlaveMsg, fd: RawFd) -> HandlerResult<u64> {
        // Safe because we have just received the rawfd from kernel.
        unsafe { libc::close(fd) };
        Err(std::io::Error::from_raw_os_error(libc::ENOSYS))
    }

    // fn handle_iotlb_msg(&mut self, iotlb: VhostUserIotlb);
    // fn handle_vring_host_notifier(&mut self, area: VhostUserVringArea, fd: RawFd);
}

impl<S: VhostUserMasterReqHandlerMut> VhostUserMasterReqHandler for Mutex<S> {
    fn handle_config_change(&self) -> HandlerResult<u64> {
        self.lock().unwrap().handle_config_change()
    }

    fn fs_slave_map(&self, fs: &VhostUserFSSlaveMsg, fd: RawFd) -> HandlerResult<u64> {
        self.lock().unwrap().fs_slave_map(fs, fd)
    }

    fn fs_slave_unmap(&self, fs: &VhostUserFSSlaveMsg) -> HandlerResult<u64> {
        self.lock().unwrap().fs_slave_unmap(fs)
    }

    fn fs_slave_sync(&self, fs: &VhostUserFSSlaveMsg) -> HandlerResult<u64> {
        self.lock().unwrap().fs_slave_sync(fs)
    }

    fn fs_slave_io(&self, fs: &VhostUserFSSlaveMsg, fd: RawFd) -> HandlerResult<u64> {
        self.lock().unwrap().fs_slave_io(fs, fd)
    }
}

/// Server to handle service requests from slaves from the slave communication channel.
///
/// The [MasterReqHandler] acts as a server on the master side, to handle service requests from
/// slaves on the slave communication channel. It's actually a proxy invoking the registered
/// handler implementing [VhostUserMasterReqHandler] to do the real work.
///
/// [MasterReqHandler]: struct.MasterReqHandler.html
/// [VhostUserMasterReqHandler]: trait.VhostUserMasterReqHandler.html
pub struct MasterReqHandler<S: VhostUserMasterReqHandler> {
    // underlying Unix domain socket for communication
    sub_sock: Endpoint<SlaveReq>,
    tx_sock: UnixStream,
    // Protocol feature VHOST_USER_PROTOCOL_F_REPLY_ACK has been negotiated.
    reply_ack_negotiated: bool,
    // the VirtIO backend device object
    backend: Arc<S>,
    // whether the endpoint has encountered any failure
    error: Option<i32>,
}

impl<S: VhostUserMasterReqHandler> MasterReqHandler<S> {
    /// Create a server to handle service requests from slaves on the slave communication channel.
    ///
    /// This opens a pair of connected anonymous sockets to form the slave communication channel.
    /// The socket fd returned by [Self::get_tx_raw_fd()] should be sent to the slave by
    /// [VhostUserMaster::set_slave_request_fd()].
    ///
    /// [Self::get_tx_raw_fd()]: struct.MasterReqHandler.html#method.get_tx_raw_fd
    /// [VhostUserMaster::set_slave_request_fd()]: trait.VhostUserMaster.html#tymethod.set_slave_request_fd
    pub fn new(backend: Arc<S>) -> Result<Self> {
        let (tx, rx) = UnixStream::pair().map_err(Error::SocketError)?;

        Ok(MasterReqHandler {
            sub_sock: Endpoint::<SlaveReq>::from_stream(rx),
            tx_sock: tx,
            reply_ack_negotiated: false,
            backend,
            error: None,
        })
    }

    /// Get the socket fd for the slave to communication with the master.
    ///
    /// The returned fd should be sent to the slave by [VhostUserMaster::set_slave_request_fd()].
    ///
    /// [VhostUserMaster::set_slave_request_fd()]: trait.VhostUserMaster.html#tymethod.set_slave_request_fd
    pub fn get_tx_raw_fd(&self) -> RawFd {
        self.tx_sock.as_raw_fd()
    }

    /// Set the negotiation state of the `VHOST_USER_PROTOCOL_F_REPLY_ACK` protocol feature.
    ///
    /// When the `VHOST_USER_PROTOCOL_F_REPLY_ACK` protocol feature has been negotiated,
    /// the "REPLY_ACK" flag will be set in the message header for every slave to master request
    /// message.
    pub fn set_reply_ack_flag(&mut self, enable: bool) {
        self.reply_ack_negotiated = enable;
    }

    /// Mark endpoint as failed or in normal state.
    pub fn set_failed(&mut self, error: i32) {
        if error == 0 {
            self.error = None;
        } else {
            self.error = Some(error);
        }
    }

    /// Main entrance to server slave request from the slave communication channel.
    ///
    /// The caller needs to:
    /// - serialize calls to this function
    /// - decide what to do when errer happens
    /// - optional recover from failure
    pub fn handle_request(&mut self) -> Result<u64> {
        // Return error if the endpoint is already in failed state.
        self.check_state()?;

        // The underlying communication channel is a Unix domain socket in
        // stream mode, and recvmsg() is a little tricky here. To successfully
        // receive attached file descriptors, we need to receive messages and
        // corresponding attached file descriptors in this way:
        // . recv messsage header and optional attached file
        // . validate message header
        // . recv optional message body and payload according size field in
        //   message header
        // . validate message body and optional payload
        let (hdr, rfds) = self.sub_sock.recv_header()?;
        let rfds = self.check_attached_rfds(&hdr, rfds)?;
        let (size, buf) = match hdr.get_size() {
            0 => (0, vec![0u8; 0]),
            len => {
                if len as usize > MAX_MSG_SIZE {
                    return Err(Error::InvalidMessage);
                }
                let (size2, rbuf) = self.sub_sock.recv_data(len as usize)?;
                if size2 != len as usize {
                    return Err(Error::InvalidMessage);
                }
                (size2, rbuf)
            }
        };

        let res = match hdr.get_code() {
            SlaveReq::CONFIG_CHANGE_MSG => {
                self.check_msg_size(&hdr, size, 0)?;
                self.backend
                    .handle_config_change()
                    .map_err(Error::ReqHandlerError)
            }
            SlaveReq::FS_MAP => {
                let msg = self.extract_msg_body::<VhostUserFSSlaveMsg>(&hdr, size, &buf)?;
                // check_attached_rfds() has validated rfds
                self.backend
                    .fs_slave_map(&msg, rfds.unwrap()[0])
                    .map_err(Error::ReqHandlerError)
            }
            SlaveReq::FS_UNMAP => {
                let msg = self.extract_msg_body::<VhostUserFSSlaveMsg>(&hdr, size, &buf)?;
                self.backend
                    .fs_slave_unmap(&msg)
                    .map_err(Error::ReqHandlerError)
            }
            SlaveReq::FS_SYNC => {
                let msg = self.extract_msg_body::<VhostUserFSSlaveMsg>(&hdr, size, &buf)?;
                self.backend
                    .fs_slave_sync(&msg)
                    .map_err(Error::ReqHandlerError)
            }
            SlaveReq::FS_IO => {
                let msg = self.extract_msg_body::<VhostUserFSSlaveMsg>(&hdr, size, &buf)?;
                // check_attached_rfds() has validated rfds
                self.backend
                    .fs_slave_io(&msg, rfds.unwrap()[0])
                    .map_err(Error::ReqHandlerError)
            }
            _ => Err(Error::InvalidMessage),
        };

        self.send_ack_message(&hdr, &res)?;

        res
    }

    fn check_state(&self) -> Result<()> {
        match self.error {
            Some(e) => Err(Error::SocketBroken(std::io::Error::from_raw_os_error(e))),
            None => Ok(()),
        }
    }

    fn check_msg_size(
        &self,
        hdr: &VhostUserMsgHeader<SlaveReq>,
        size: usize,
        expected: usize,
    ) -> Result<()> {
        if hdr.get_size() as usize != expected
            || hdr.is_reply()
            || hdr.get_version() != 0x1
            || size != expected
        {
            return Err(Error::InvalidMessage);
        }
        Ok(())
    }

    fn check_attached_rfds(
        &self,
        hdr: &VhostUserMsgHeader<SlaveReq>,
        rfds: Option<Vec<RawFd>>,
    ) -> Result<Option<Vec<RawFd>>> {
        match hdr.get_code() {
            SlaveReq::FS_MAP | SlaveReq::FS_IO => {
                // Expect an fd set with a single fd.
                match rfds {
                    None => Err(Error::InvalidMessage),
                    Some(fds) => {
                        if fds.len() != 1 {
                            Endpoint::<SlaveReq>::close_rfds(Some(fds));
                            Err(Error::InvalidMessage)
                        } else {
                            Ok(Some(fds))
                        }
                    }
                }
            }
            _ => {
                if rfds.is_some() {
                    Endpoint::<SlaveReq>::close_rfds(rfds);
                    return Err(Error::InvalidMessage);
                } else {
                    Ok(rfds)
                }
            }
        }
    }

    fn extract_msg_body<T: Sized + VhostUserMsgValidator>(
        &self,
        hdr: &VhostUserMsgHeader<SlaveReq>,
        size: usize,
        buf: &[u8],
    ) -> Result<T> {
        self.check_msg_size(hdr, size, mem::size_of::<T>())?;
        let msg = unsafe { std::ptr::read_unaligned(buf.as_ptr() as *const T) };
        if !msg.is_valid() {
            return Err(Error::InvalidMessage);
        }
        Ok(msg)
    }

    fn new_reply_header<T: Sized>(
        &self,
        req: &VhostUserMsgHeader<SlaveReq>,
    ) -> Result<VhostUserMsgHeader<SlaveReq>> {
        if mem::size_of::<T>() > MAX_MSG_SIZE {
            return Err(Error::InvalidParam);
        }
        self.check_state()?;
        Ok(VhostUserMsgHeader::new(
            req.get_code(),
            VhostUserHeaderFlag::REPLY.bits(),
            mem::size_of::<T>() as u32,
        ))
    }

    fn send_ack_message(
        &mut self,
        req: &VhostUserMsgHeader<SlaveReq>,
        res: &Result<u64>,
    ) -> Result<()> {
        if self.reply_ack_negotiated && req.is_need_reply() {
            let hdr = self.new_reply_header::<VhostUserU64>(req)?;
            let def_err = libc::EINVAL;
            let val = match res {
                Ok(n) => *n,
                Err(e) => match &*e {
                    Error::ReqHandlerError(ioerr) => match ioerr.raw_os_error() {
                        Some(rawerr) => -rawerr as u64,
                        None => -def_err as u64,
                    },
                    _ => -def_err as u64,
                },
            };
            let msg = VhostUserU64::new(val);
            self.sub_sock.send_message(&hdr, &msg, None)?;
        }
        Ok(())
    }
}

impl<S: VhostUserMasterReqHandler> AsRawFd for MasterReqHandler<S> {
    fn as_raw_fd(&self) -> RawFd {
        self.sub_sock.as_raw_fd()
    }
}
<|endoftext|>


<commit_before>// https://leetcode.com/problems/number-of-submatrices-that-sum-to-target/
pub fn num_submatrix_sum_target(matrix: Vec<Vec<i32>>, target: i32) -> i32 {
    let mut answer = 0;
    let max_rows = matrix.len();
    let max_cols = if max_rows > 0 { matrix[0].len() } else { 0 };
    let mut prefix_sums = vec![vec![0; max_cols]; max_rows];
    for row in 0..max_rows {
        for col in 0..max_cols {
            prefix_sums[row][col] = if col == 0 {
                matrix[row][col]
            } else {
                prefix_sums[row][col - 1] + matrix[row][col]
            };
        }
    }
    for row1 in 0..max_rows {
        for col1 in 0..max_cols {
            for row2 in row1..max_rows {
                for col2 in col1..max_cols {
                    let mut sum = 0;
                    for i in row1..=row2 {
                        sum += if col1 == 0 {
                            prefix_sums[i][col2]
                        } else {
                            prefix_sums[i][col2] - prefix_sums[i][col1 - 1]
                        };
                    }
                    if sum == target {
                        answer += 1;
                    }
                }
            }
        }
    }
    answer
}

fn main() {
    println!(
        "{}",
        num_submatrix_sum_target(vec![vec![0, 0, 0], vec![1, -1, 1], vec![0, 1, 0]], 0)
    ); // 14
    println!(
        "{}",
        num_submatrix_sum_target(vec![vec![0, 1, 0], vec![1, 1, 1], vec![0, 1, 0]], 0)
    ); // 4
    println!(
        "{}",
        num_submatrix_sum_target(vec![vec![1, -1], vec![-1, 1]], 0)
    ); // 5
    println!("{}", num_submatrix_sum_target(vec![vec![904]], 0)); // 0
    println!(
        "{}",
        num_submatrix_sum_target(
            vec![
                vec![0, 0, 0, 1, 1],
                vec![1, 1, 1, 0, 1],
                vec![1, 1, 1, 1, 0],
                vec![0, 0, 0, 1, 0],
                vec![0, 0, 0, 1, 1]
            ],
            0
        )
    ); // 28
}
<commit_msg>Solve problem 1074<commit_after>// https://leetcode.com/problems/number-of-submatrices-that-sum-to-target/
pub fn num_submatrix_sum_target(matrix: Vec<Vec<i32>>, target: i32) -> i32 {
    let mut answer = 0;
    let max_rows = matrix.len();
    let max_cols = if max_rows > 0 { matrix[0].len() } else { 0 };
    let mut prefix_sums = vec![vec![0; max_cols]; max_rows];
    for row in 0..max_rows {
        let mut prefix_sum = 0;
        for col in 0..max_cols {
            prefix_sum += matrix[row][col];
            prefix_sums[row][col] = prefix_sum
                + if row > 0 {
                    prefix_sums[row - 1][col]
                } else {
                    0
                };
        }
    }
    for row1 in 0..max_rows {
        for col1 in 0..max_cols {
            for row2 in row1..max_rows {
                for col2 in col1..max_cols {
                    let a = if row1 > 0 {
                        prefix_sums[row1 - 1][col2]
                    } else {
                        0
                    };
                    let b = if col1 > 0 {
                        prefix_sums[row2][col1 - 1]
                            - if row1 > 0 {
                                prefix_sums[row1 - 1][col1 - 1]
                            } else {
                                0
                            }
                    } else {
                        0
                    };
                    let sum = prefix_sums[row2][col2] - a - b;
                    if sum == target {
                        answer += 1;
                    }
                }
            }
        }
    }
    answer
}

fn main() {
    println!(
        "{}",
        num_submatrix_sum_target(vec![vec![0, 0, 0], vec![1, -1, 1], vec![0, 1, 0]], 0)
    ); // 14
    println!(
        "{}",
        num_submatrix_sum_target(vec![vec![0, 1, 0], vec![1, 1, 1], vec![0, 1, 0]], 0)
    ); // 4
    println!(
        "{}",
        num_submatrix_sum_target(vec![vec![1, -1], vec![-1, 1]], 0)
    ); // 5
    println!("{}", num_submatrix_sum_target(vec![vec![904]], 0)); // 0
    println!(
        "{}",
        num_submatrix_sum_target(
            vec![
                vec![0, 0, 0, 1, 1],
                vec![1, 1, 1, 0, 1],
                vec![1, 1, 1, 1, 0],
                vec![0, 0, 0, 1, 0],
                vec![0, 0, 0, 1, 1]
            ],
            0
        )
    ); // 28
}
<|endoftext|>


<commit_before>//! Resource manager to allocate and switch between framebuffers.

use context::{Context, Framebuffer, Texture, Renderbuffer};
use either::Either;

#[path = "../error.rs"]
mod error;

/// The target to every rendering call.
pub enum RenderTarget {
    /// The screen (main framebuffer).
    Screen,
    /// An off-screen buffer.
    Offscreen(OffscreenBuffers),
}

/// OpenGL identifiers to an off-screen buffer.
pub struct OffscreenBuffers {
    texture: Texture,
    depth: Either<Texture, Renderbuffer>,
}

impl RenderTarget {
    /// Returns an opengl handle to the off-screen texture buffer.
    ///
    /// Returns `None` if the texture is off-screen.
    pub fn texture_id(&self) -> Option<&Texture> {
        match *self {
            RenderTarget::Screen => None,
            RenderTarget::Offscreen(ref o) => Some(&o.texture),
        }
    }

    /// Returns an opengl handle to the off-screen depth buffer.
    ///
    /// Returns `None` if the texture is off-screen.
    #[cfg(not(any(target_arch = "wasm32", target_arch = "asmjs")))]
    pub fn depth_id(&self) -> Option<&Either<Texture, Renderbuffer>> {
        match *self {
            RenderTarget::Screen => None,
            RenderTarget::Offscreen(ref o) => Some(&o.depth),
        }
    }

    /// Resizes this render target.
    pub fn resize(&mut self, w: f32, h: f32) {
        let ctxt = Context::get();

        match *self {
            RenderTarget::Screen => {
                verify!(ctxt.viewport(0, 0, w as i32, h as i32));
            }
            RenderTarget::Offscreen(ref o) => {
                // Update the fbo
                verify!(ctxt.bind_texture(Context::TEXTURE_2D, Some(&o.texture)));
                verify!(ctxt.tex_image2d(
                    Context::TEXTURE_2D,
                    0,
                    Context::RGBA as i32,
                    w as i32,
                    h as i32,
                    0,
                    Context::RGBA,
                    None
                ));
                verify!(ctxt.bind_texture(Context::TEXTURE_2D, None));

                match &o.depth {
                    Either::Left(texture) => {
                        verify!(ctxt.bind_texture(Context::TEXTURE_2D, Some(texture)));
                        verify!(ctxt.tex_image2d(
                            Context::TEXTURE_2D,
                            0,
                            Context::DEPTH_COMPONENT as i32,
                            w as i32,
                            h as i32,
                            0,
                            Context::DEPTH_COMPONENT,
                            None
                        ));
                        verify!(ctxt.bind_texture(Context::TEXTURE_2D, None));
                    }
                    Either::Right(renderbuffer) => {
                        verify!(ctxt.bind_renderbuffer(Some(renderbuffer)));
                        verify!(ctxt.renderbuffer_storage(Context::DEPTH_COMPONENT, w as i32, h as i32));
                    }
                }
            }
        }
    }
}

/// A framebuffer manager. It is a simple to to switch between an off-screen framebuffer and the
/// default (window) framebuffer.
pub struct FramebufferManager {
    fbo_onscreen: bool,
    fbo: Framebuffer,
}

impl FramebufferManager {
    /// Creates a new framebuffer manager.
    pub fn new() -> FramebufferManager {
        let ctxt = Context::get();

        // create an off-screen framebuffer
        let fbo = ctxt
            .create_framebuffer()
            .expect("Framebuffer creation failed.");

        // ensure that the current framebuffer is the screen
        verify!(ctxt.bind_framebuffer(Context::FRAMEBUFFER, None));

        FramebufferManager {
            fbo_onscreen: true,
            fbo: fbo,
        }
    }

    /// Creates a new render target. A render target is the combination of a color buffer and a
    /// depth buffer.
    pub fn new_render_target(width: usize, height: usize, create_depth_texture: bool) -> RenderTarget {
        let ctxt = Context::get();

        /* Texture */
        verify!(ctxt.active_texture(Context::TEXTURE0));
        let fbo_texture = verify!(
            ctxt.create_texture()
                .expect("Failde to create framebuffer object texture.")
        );
        verify!(ctxt.bind_texture(Context::TEXTURE_2D, Some(&fbo_texture)));
        verify!(ctxt.tex_parameteri(
            Context::TEXTURE_2D,
            Context::TEXTURE_MAG_FILTER,
            Context::LINEAR as i32
        ));
        verify!(ctxt.tex_parameteri(
            Context::TEXTURE_2D,
            Context::TEXTURE_MIN_FILTER,
            Context::LINEAR as i32
        ));
        verify!(ctxt.tex_parameteri(
            Context::TEXTURE_2D,
            Context::TEXTURE_WRAP_S,
            Context::CLAMP_TO_EDGE as i32
        ));
        verify!(ctxt.tex_parameteri(
            Context::TEXTURE_2D,
            Context::TEXTURE_WRAP_T,
            Context::CLAMP_TO_EDGE as i32
        ));
        verify!(ctxt.tex_image2d(
            Context::TEXTURE_2D,
            0,
            Context::RGBA as i32,
            width as i32,
            height as i32,
            0,
            Context::RGBA,
            None
        ));
        verify!(ctxt.bind_texture(Context::TEXTURE_2D, None));

        /* Depth buffer */
        if create_depth_texture && cfg!(not(any(target_arch = "wasm32", target_arch = "asmjs"))) {
            verify!(ctxt.active_texture(Context::TEXTURE1));
            let fbo_depth = verify!(ctxt.create_texture().expect("Failed to create a texture."));
            verify!(ctxt.bind_texture(Context::TEXTURE_2D, Some(&fbo_depth)));
            verify!(ctxt.tex_parameteri(
                Context::TEXTURE_2D,
                Context::TEXTURE_MAG_FILTER,
                Context::LINEAR as i32
            ));
            verify!(ctxt.tex_parameteri(
                Context::TEXTURE_2D,
                Context::TEXTURE_MIN_FILTER,
                Context::LINEAR as i32
            ));
            verify!(ctxt.tex_parameteri(
                Context::TEXTURE_2D,
                Context::TEXTURE_WRAP_S,
                Context::CLAMP_TO_EDGE as i32
            ));
            verify!(ctxt.tex_parameteri(
                Context::TEXTURE_2D,
                Context::TEXTURE_WRAP_T,
                Context::CLAMP_TO_EDGE as i32
            ));
            verify!(ctxt.tex_image2di(
                Context::TEXTURE_2D,
                0,
                Context::DEPTH_COMPONENT as i32,
                width as i32,
                height as i32,
                0,
                Context::DEPTH_COMPONENT,
                None
            ));
            verify!(ctxt.bind_texture(Context::TEXTURE_2D, None));

            RenderTarget::Offscreen(OffscreenBuffers {
                texture: fbo_texture,
                depth: Either::Left(fbo_depth),
            })
        } else {
            // Create a renderbuffer instead of the texture for the depth.
            let renderbuffer = verify!(ctxt.create_renderbuffer()).expect("Failed to create a renderbuffer.");
            verify!(ctxt.bind_renderbuffer(Some(&renderbuffer)));
            verify!(ctxt.renderbuffer_storage(Context::DEPTH_COMPONENT16, width as i32, height as i32));
            RenderTarget::Offscreen(OffscreenBuffers {
                texture: fbo_texture,
                depth: Either::Right(renderbuffer),
            })
        }
    }

    /// Returns the render target associated with the screen.
    pub fn screen() -> RenderTarget {
        RenderTarget::Screen
    }

    /// Selects a specific render target
    pub fn select(&mut self, target: &RenderTarget) {
        match *target {
            RenderTarget::Screen => {
                self.select_onscreen();
            }
            RenderTarget::Offscreen(ref o) => {
                let ctxt = Context::get();
                self.select_fbo();

                // FIXME: don't switch if the current texture is
                // already o.texture ?
                verify!(ctxt.framebuffer_texture2d(
                    Context::FRAMEBUFFER,
                    Context::COLOR_ATTACHMENT0,
                    Context::TEXTURE_2D,
                    Some(&o.texture),
                    0
                ));

                match &o.depth {
                    Either::Left(texture) => {
                        verify!(ctxt.framebuffer_texture2d(
                            Context::FRAMEBUFFER,
                            Context::DEPTH_ATTACHMENT,
                            Context::TEXTURE_2D,
                            Some(texture),
                            0
                        ));
                    }
                    Either::Right(renderbuffer) => {
                        verify!(ctxt.framebuffer_renderbuffer(
                            Context::DEPTH_ATTACHMENT,
                            Some(renderbuffer),
                        ))
                    }
                }
            }
        }
    }

    fn select_onscreen(&mut self) {
        if !self.fbo_onscreen {
            verify!(Context::get().bind_framebuffer(Context::FRAMEBUFFER, None));
            self.fbo_onscreen = true;
        }
    }

    fn select_fbo(&mut self) {
        if self.fbo_onscreen {
            verify!(Context::get().bind_framebuffer(Context::FRAMEBUFFER, Some(&self.fbo)));
            self.fbo_onscreen = false;
        }
    }
}

impl Drop for FramebufferManager {
    fn drop(&mut self) {
        let ctxt = Context::get();
        if ctxt.is_framebuffer(Some(&self.fbo)) {
            verify!(ctxt.bind_framebuffer(Context::FRAMEBUFFER, None));
            verify!(ctxt.delete_framebuffer(Some(&self.fbo)));
        }
    }
}

impl Drop for OffscreenBuffers {
    fn drop(&mut self) {
        let ctxt = Context::get();
        if ctxt.is_texture(Some(&self.texture)) {
            verify!(ctxt.delete_texture(Some(&self.texture)));
        }

        match &self.depth {
            Either::Left(texture) => {
                if ctxt.is_texture(Some(texture)) {
                    verify!(ctxt.delete_texture(Some(texture)));
                }
            }
            Either::Right(renderbuffer) => {
                if ctxt.is_renderbuffer(Some(renderbuffer)) {
                    verify!(ctxt.delete_renderbuffer(Some(renderbuffer)));
                }
            }
        }

    }
}
<commit_msg>Correctly initialize the renderbuffer storage.<commit_after>//! Resource manager to allocate and switch between framebuffers.

use context::{Context, Framebuffer, Renderbuffer, Texture};
use either::Either;

#[path = "../error.rs"]
mod error;

/// The target to every rendering call.
pub enum RenderTarget {
    /// The screen (main framebuffer).
    Screen,
    /// An off-screen buffer.
    Offscreen(OffscreenBuffers),
}

/// OpenGL identifiers to an off-screen buffer.
pub struct OffscreenBuffers {
    texture: Texture,
    depth: Either<Texture, Renderbuffer>,
}

impl RenderTarget {
    /// Returns an opengl handle to the off-screen texture buffer.
    ///
    /// Returns `None` if the texture is off-screen.
    pub fn texture_id(&self) -> Option<&Texture> {
        match *self {
            RenderTarget::Screen => None,
            RenderTarget::Offscreen(ref o) => Some(&o.texture),
        }
    }

    /// Returns an opengl handle to the off-screen depth buffer.
    ///
    /// Returns `None` if the texture is off-screen.
    #[cfg(not(any(target_arch = "wasm32", target_arch = "asmjs")))]
    pub fn depth_id(&self) -> Option<&Either<Texture, Renderbuffer>> {
        match *self {
            RenderTarget::Screen => None,
            RenderTarget::Offscreen(ref o) => Some(&o.depth),
        }
    }

    /// Resizes this render target.
    pub fn resize(&mut self, w: f32, h: f32) {
        let ctxt = Context::get();

        match *self {
            RenderTarget::Screen => {
                verify!(ctxt.viewport(0, 0, w as i32, h as i32));
            }
            RenderTarget::Offscreen(ref o) => {
                // Update the fbo
                verify!(ctxt.bind_texture(Context::TEXTURE_2D, Some(&o.texture)));
                verify!(ctxt.tex_image2d(
                    Context::TEXTURE_2D,
                    0,
                    Context::RGBA as i32,
                    w as i32,
                    h as i32,
                    0,
                    Context::RGBA,
                    None
                ));
                verify!(ctxt.bind_texture(Context::TEXTURE_2D, None));

                match &o.depth {
                    Either::Left(texture) => {
                        verify!(ctxt.bind_texture(Context::TEXTURE_2D, Some(texture)));
                        verify!(ctxt.tex_image2d(
                            Context::TEXTURE_2D,
                            0,
                            Context::DEPTH_COMPONENT as i32,
                            w as i32,
                            h as i32,
                            0,
                            Context::DEPTH_COMPONENT,
                            None
                        ));
                        verify!(ctxt.bind_texture(Context::TEXTURE_2D, None));
                    }
                    Either::Right(renderbuffer) => {
                        verify!(ctxt.bind_renderbuffer(Some(renderbuffer)));
                        verify!(ctxt.renderbuffer_storage(
                            Context::DEPTH_COMPONENT16,
                            w as i32,
                            h as i32
                        ));
                        verify!(ctxt.bind_renderbuffer(None));
                    }
                }
            }
        }
    }
}

/// A framebuffer manager. It is a simple to to switch between an off-screen framebuffer and the
/// default (window) framebuffer.
pub struct FramebufferManager {
    fbo_onscreen: bool,
    fbo: Framebuffer,
}

impl FramebufferManager {
    /// Creates a new framebuffer manager.
    pub fn new() -> FramebufferManager {
        let ctxt = Context::get();

        // create an off-screen framebuffer
        let fbo = ctxt
            .create_framebuffer()
            .expect("Framebuffer creation failed.");

        // ensure that the current framebuffer is the screen
        verify!(ctxt.bind_framebuffer(Context::FRAMEBUFFER, None));

        FramebufferManager {
            fbo_onscreen: true,
            fbo: fbo,
        }
    }

    /// Creates a new render target. A render target is the combination of a color buffer and a
    /// depth buffer.
    pub fn new_render_target(
        width: usize,
        height: usize,
        create_depth_texture: bool,
    ) -> RenderTarget {
        let ctxt = Context::get();

        /* Texture */
        verify!(ctxt.active_texture(Context::TEXTURE0));
        let fbo_texture = verify!(ctxt
            .create_texture()
            .expect("Failde to create framebuffer object texture."));
        verify!(ctxt.bind_texture(Context::TEXTURE_2D, Some(&fbo_texture)));
        verify!(ctxt.tex_parameteri(
            Context::TEXTURE_2D,
            Context::TEXTURE_MAG_FILTER,
            Context::LINEAR as i32
        ));
        verify!(ctxt.tex_parameteri(
            Context::TEXTURE_2D,
            Context::TEXTURE_MIN_FILTER,
            Context::LINEAR as i32
        ));
        verify!(ctxt.tex_parameteri(
            Context::TEXTURE_2D,
            Context::TEXTURE_WRAP_S,
            Context::CLAMP_TO_EDGE as i32
        ));
        verify!(ctxt.tex_parameteri(
            Context::TEXTURE_2D,
            Context::TEXTURE_WRAP_T,
            Context::CLAMP_TO_EDGE as i32
        ));
        verify!(ctxt.tex_image2d(
            Context::TEXTURE_2D,
            0,
            Context::RGBA as i32,
            width as i32,
            height as i32,
            0,
            Context::RGBA,
            None
        ));
        verify!(ctxt.bind_texture(Context::TEXTURE_2D, None));

        /* Depth buffer */
        if create_depth_texture && cfg!(not(any(target_arch = "wasm32", target_arch = "asmjs"))) {
            verify!(ctxt.active_texture(Context::TEXTURE1));
            let fbo_depth = verify!(ctxt.create_texture().expect("Failed to create a texture."));
            verify!(ctxt.bind_texture(Context::TEXTURE_2D, Some(&fbo_depth)));
            verify!(ctxt.tex_parameteri(
                Context::TEXTURE_2D,
                Context::TEXTURE_MAG_FILTER,
                Context::LINEAR as i32
            ));
            verify!(ctxt.tex_parameteri(
                Context::TEXTURE_2D,
                Context::TEXTURE_MIN_FILTER,
                Context::LINEAR as i32
            ));
            verify!(ctxt.tex_parameteri(
                Context::TEXTURE_2D,
                Context::TEXTURE_WRAP_S,
                Context::CLAMP_TO_EDGE as i32
            ));
            verify!(ctxt.tex_parameteri(
                Context::TEXTURE_2D,
                Context::TEXTURE_WRAP_T,
                Context::CLAMP_TO_EDGE as i32
            ));
            verify!(ctxt.tex_image2di(
                Context::TEXTURE_2D,
                0,
                Context::DEPTH_COMPONENT as i32,
                width as i32,
                height as i32,
                0,
                Context::DEPTH_COMPONENT,
                None
            ));
            verify!(ctxt.bind_texture(Context::TEXTURE_2D, None));

            RenderTarget::Offscreen(OffscreenBuffers {
                texture: fbo_texture,
                depth: Either::Left(fbo_depth),
            })
        } else {
            // Create a renderbuffer instead of the texture for the depth.
            let renderbuffer =
                verify!(ctxt.create_renderbuffer()).expect("Failed to create a renderbuffer.");
            verify!(ctxt.bind_renderbuffer(Some(&renderbuffer)));
            verify!(ctxt.renderbuffer_storage(
                Context::DEPTH_COMPONENT16,
                width as i32,
                height as i32
            ));
            verify!(ctxt.bind_renderbuffer(None));

            RenderTarget::Offscreen(OffscreenBuffers {
                texture: fbo_texture,
                depth: Either::Right(renderbuffer),
            })
        }
    }

    /// Returns the render target associated with the screen.
    pub fn screen() -> RenderTarget {
        RenderTarget::Screen
    }

    /// Selects a specific render target
    pub fn select(&mut self, target: &RenderTarget) {
        match *target {
            RenderTarget::Screen => {
                self.select_onscreen();
            }
            RenderTarget::Offscreen(ref o) => {
                let ctxt = Context::get();
                self.select_fbo();

                // FIXME: don't switch if the current texture is
                // already o.texture ?
                verify!(ctxt.framebuffer_texture2d(
                    Context::FRAMEBUFFER,
                    Context::COLOR_ATTACHMENT0,
                    Context::TEXTURE_2D,
                    Some(&o.texture),
                    0
                ));

                match &o.depth {
                    Either::Left(texture) => {
                        verify!(ctxt.framebuffer_texture2d(
                            Context::FRAMEBUFFER,
                            Context::DEPTH_ATTACHMENT,
                            Context::TEXTURE_2D,
                            Some(texture),
                            0
                        ));
                    }
                    Either::Right(renderbuffer) => verify!(ctxt
                        .framebuffer_renderbuffer(Context::DEPTH_ATTACHMENT, Some(renderbuffer))),
                }
            }
        }
    }

    fn select_onscreen(&mut self) {
        if !self.fbo_onscreen {
            verify!(Context::get().bind_framebuffer(Context::FRAMEBUFFER, None));
            self.fbo_onscreen = true;
        }
    }

    fn select_fbo(&mut self) {
        if self.fbo_onscreen {
            verify!(Context::get().bind_framebuffer(Context::FRAMEBUFFER, Some(&self.fbo)));
            self.fbo_onscreen = false;
        }
    }
}

impl Drop for FramebufferManager {
    fn drop(&mut self) {
        let ctxt = Context::get();
        if ctxt.is_framebuffer(Some(&self.fbo)) {
            verify!(ctxt.bind_framebuffer(Context::FRAMEBUFFER, None));
            verify!(ctxt.delete_framebuffer(Some(&self.fbo)));
        }
    }
}

impl Drop for OffscreenBuffers {
    fn drop(&mut self) {
        let ctxt = Context::get();
        if ctxt.is_texture(Some(&self.texture)) {
            verify!(ctxt.delete_texture(Some(&self.texture)));
        }

        match &self.depth {
            Either::Left(texture) => {
                if ctxt.is_texture(Some(texture)) {
                    verify!(ctxt.delete_texture(Some(texture)));
                }
            }
            Either::Right(renderbuffer) => {
                if ctxt.is_renderbuffer(Some(renderbuffer)) {
                    verify!(ctxt.delete_renderbuffer(Some(renderbuffer)));
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// Finds items that are externally reachable, to determine which items
// need to have their metadata (and possibly their AST) serialized.
// All items that can be referred to through an exported name are
// reachable, and when a reachable thing is inline or generic, it
// makes all other generics or inline functions that it references
// reachable as well.

import syntax::ast::*;
import syntax::{visit, ast_util, ast_map};
import syntax::ast_util::def_id_of_def;
import syntax::attr;
import std::map::hashmap;

export map, find_reachable;

type map = std::map::hashmap<node_id, ()>;

type ctx = {exp_map: resolve::exp_map,
            tcx: ty::ctxt,
            method_map: typeck::method_map,
            rmap: map};

fn find_reachable(crate_mod: _mod, exp_map: resolve::exp_map,
                  tcx: ty::ctxt, method_map: typeck::method_map) -> map {
    let rmap = std::map::int_hash();
    let cx = {exp_map: exp_map, tcx: tcx, method_map: method_map, rmap: rmap};
    traverse_public_mod(cx, crate_mod);
    traverse_all_resources(cx, crate_mod);
    rmap
}

fn traverse_exports(cx: ctx, vis: [@view_item]) -> bool {
    let mut found_export = false;
    for vec::each(vis) {|vi|
        alt vi.node {
          view_item_export(vps) {
            found_export = true;
            for vec::each(vps) {|vp|
                alt vp.node {
                  view_path_simple(_, _, id) | view_path_glob(_, id) |
                  view_path_list(_, _, id) {
                    traverse_export(cx, id);
                  }
                }
            }
          }
          _ {}
        }
    }
    found_export
}

fn traverse_export(cx: ctx, exp_id: node_id) {
    option::iter(cx.exp_map.find(exp_id)) {|defs|
        for vec::each(defs) {|def| traverse_def_id(cx, def.id); }
    }
}

fn traverse_def_id(cx: ctx, did: def_id) {
    if did.crate != local_crate { ret; }
    alt cx.tcx.items.get(did.node) {
      ast_map::node_item(item, _) { traverse_public_item(cx, item); }
      ast_map::node_method(_, impl_id, _) { traverse_def_id(cx, impl_id); }
      ast_map::node_native_item(item, _, _) { cx.rmap.insert(item.id, ()); }
      ast_map::node_variant(v, _, _) { cx.rmap.insert(v.node.id, ()); }
      _ {}
    }
}

fn traverse_public_mod(cx: ctx, m: _mod) {
    if !traverse_exports(cx, m.view_items) {
        // No exports, so every local item is exported
        for vec::each(m.items) {|item| traverse_public_item(cx, item); }
    }
}

fn traverse_public_item(cx: ctx, item: @item) {
    if cx.rmap.contains_key(item.id) { ret; }
    cx.rmap.insert(item.id, ());
    alt item.node {
      item_mod(m) { traverse_public_mod(cx, m); }
      item_native_mod(nm) {
          if !traverse_exports(cx, nm.view_items) {
              for vec::each(nm.items) {|item| cx.rmap.insert(item.id, ()); }
          }
      }
      item_res(_, tps, blk, _, _, _) {
        // resources seem to be unconditionally inlined
        traverse_inline_body(cx, blk);
      }
      item_fn(_, tps, blk) {
        if tps.len() > 0u ||
           attr::find_inline_attr(item.attrs) != attr::ia_none {
            traverse_inline_body(cx, blk);
        }
      }
      item_impl(tps, _, _, _, ms) {
        for vec::each(ms) {|m|
            if tps.len() > 0u || m.tps.len() > 0u ||
               attr::find_inline_attr(m.attrs) != attr::ia_none {
                cx.rmap.insert(m.id, ());
                traverse_inline_body(cx, m.body);
            }
        }
      }
      item_class(tps, _ifaces, items, ctor, m_dtor, _) {
        cx.rmap.insert(ctor.node.id, ());
        option::iter(m_dtor) {|dtor|
            cx.rmap.insert(dtor.node.id, ());
        }
        for vec::each(items) {|item|
            alt item.node {
              class_method(m) {
                cx.rmap.insert(m.id, ());
                if tps.len() > 0u ||
                   attr::find_inline_attr(m.attrs) != attr::ia_none {
                    traverse_inline_body(cx, m.body);
                }
              }
              _ {}
            }
        }
      }
      item_const(*) | item_ty(*) |
      item_enum(*) | item_iface(*) {}
    }
}

fn traverse_inline_body(cx: ctx, body: blk) {
    fn traverse_expr(e: @expr, cx: ctx, v: visit::vt<ctx>) {
        alt e.node {
          expr_path(_) {
            traverse_def_id(cx, def_id_of_def(cx.tcx.def_map.get(e.id)));
          }
          expr_field(_, _, _) {
            alt cx.method_map.find(e.id) {
              some(typeck::method_static(did)) { traverse_def_id(cx, did); }
              _ {}
            }
          }
          _ {}
        }
        visit::visit_expr(e, cx, v);
    }
    // Don't ignore nested items: for example if a generic fn contains a
    // generic impl (as in deque::create), we need to monomorphize the
    // impl as well
    fn traverse_item(i: @item, cx: ctx, _v: visit::vt<ctx>) {
      traverse_public_item(cx, i);
    }
     visit::visit_block(body, cx, visit::mk_vt(@{
        visit_expr: traverse_expr,
        visit_item: traverse_item
        with *visit::default_visitor()
    }));
}

fn traverse_all_resources(cx: ctx, crate_mod: _mod) {
    visit::visit_mod(crate_mod, ast_util::dummy_sp(), 0, cx, visit::mk_vt(@{
        visit_expr: {|_e, _cx, _v|},
        visit_item: {|i, cx, v|
            visit::visit_item(i, cx, v);
            alt i.node {
              item_res(_, _, _, _, _, _) {
                traverse_public_item(cx, i);
              }
              _ {}
            }
        }
        with *visit::default_visitor()
    }));
}
<commit_msg>In reachability, visit class ctors<commit_after>// Finds items that are externally reachable, to determine which items
// need to have their metadata (and possibly their AST) serialized.
// All items that can be referred to through an exported name are
// reachable, and when a reachable thing is inline or generic, it
// makes all other generics or inline functions that it references
// reachable as well.

import syntax::ast::*;
import syntax::{visit, ast_util, ast_map};
import syntax::ast_util::def_id_of_def;
import syntax::attr;
import std::map::hashmap;

export map, find_reachable;

type map = std::map::hashmap<node_id, ()>;

type ctx = {exp_map: resolve::exp_map,
            tcx: ty::ctxt,
            method_map: typeck::method_map,
            rmap: map};

fn find_reachable(crate_mod: _mod, exp_map: resolve::exp_map,
                  tcx: ty::ctxt, method_map: typeck::method_map) -> map {
    let rmap = std::map::int_hash();
    let cx = {exp_map: exp_map, tcx: tcx, method_map: method_map, rmap: rmap};
    traverse_public_mod(cx, crate_mod);
    traverse_all_resources(cx, crate_mod);
    rmap
}

fn traverse_exports(cx: ctx, vis: [@view_item]) -> bool {
    let mut found_export = false;
    for vec::each(vis) {|vi|
        alt vi.node {
          view_item_export(vps) {
            found_export = true;
            for vec::each(vps) {|vp|
                alt vp.node {
                  view_path_simple(_, _, id) | view_path_glob(_, id) |
                  view_path_list(_, _, id) {
                    traverse_export(cx, id);
                  }
                }
            }
          }
          _ {}
        }
    }
    found_export
}

fn traverse_export(cx: ctx, exp_id: node_id) {
    option::iter(cx.exp_map.find(exp_id)) {|defs|
        for vec::each(defs) {|def| traverse_def_id(cx, def.id); }
    }
}

fn traverse_def_id(cx: ctx, did: def_id) {
    if did.crate != local_crate { ret; }
    alt cx.tcx.items.get(did.node) {
      ast_map::node_item(item, _) { traverse_public_item(cx, item); }
      ast_map::node_method(_, impl_id, _) { traverse_def_id(cx, impl_id); }
      ast_map::node_native_item(item, _, _) { cx.rmap.insert(item.id, ()); }
      ast_map::node_variant(v, _, _) { cx.rmap.insert(v.node.id, ()); }
      // If it's a class ctor, consider the parent reachable
      ast_map::node_ctor(_, _, ast_map::class_ctor(_, parent_id), _) {
        traverse_def_id(cx, parent_id);
      }
      _ {}
    }
}

fn traverse_public_mod(cx: ctx, m: _mod) {
    if !traverse_exports(cx, m.view_items) {
        // No exports, so every local item is exported
        for vec::each(m.items) {|item| traverse_public_item(cx, item); }
    }
}

fn traverse_public_item(cx: ctx, item: @item) {
    if cx.rmap.contains_key(item.id) { ret; }
    cx.rmap.insert(item.id, ());
    alt item.node {
      item_mod(m) { traverse_public_mod(cx, m); }
      item_native_mod(nm) {
          if !traverse_exports(cx, nm.view_items) {
              for vec::each(nm.items) {|item| cx.rmap.insert(item.id, ()); }
          }
      }
      item_res(_, tps, blk, _, _, _) {
        // resources seem to be unconditionally inlined
        traverse_inline_body(cx, blk);
      }
      item_fn(_, tps, blk) {
        if tps.len() > 0u ||
           attr::find_inline_attr(item.attrs) != attr::ia_none {
            traverse_inline_body(cx, blk);
        }
      }
      item_impl(tps, _, _, _, ms) {
        for vec::each(ms) {|m|
            if tps.len() > 0u || m.tps.len() > 0u ||
               attr::find_inline_attr(m.attrs) != attr::ia_none {
                cx.rmap.insert(m.id, ());
                traverse_inline_body(cx, m.body);
            }
        }
      }
      item_class(tps, _ifaces, items, ctor, m_dtor, _) {
        cx.rmap.insert(ctor.node.id, ());
        option::iter(m_dtor) {|dtor|
            cx.rmap.insert(dtor.node.id, ());
        }
        for vec::each(items) {|item|
            alt item.node {
              class_method(m) {
                cx.rmap.insert(m.id, ());
                if tps.len() > 0u ||
                   attr::find_inline_attr(m.attrs) != attr::ia_none {
                    traverse_inline_body(cx, m.body);
                }
              }
              _ {}
            }
        }
      }
      item_const(*) | item_ty(*) |
      item_enum(*) | item_iface(*) {}
    }
}

fn traverse_inline_body(cx: ctx, body: blk) {
    fn traverse_expr(e: @expr, cx: ctx, v: visit::vt<ctx>) {
        alt e.node {
          expr_path(_) {
            traverse_def_id(cx, def_id_of_def(cx.tcx.def_map.get(e.id)));
          }
          expr_field(_, _, _) {
            alt cx.method_map.find(e.id) {
              some(typeck::method_static(did)) { traverse_def_id(cx, did); }
              _ {}
            }
          }
          _ {}
        }
        visit::visit_expr(e, cx, v);
    }
    // Don't ignore nested items: for example if a generic fn contains a
    // generic impl (as in deque::create), we need to monomorphize the
    // impl as well
    fn traverse_item(i: @item, cx: ctx, _v: visit::vt<ctx>) {
      traverse_public_item(cx, i);
    }
     visit::visit_block(body, cx, visit::mk_vt(@{
        visit_expr: traverse_expr,
        visit_item: traverse_item
        with *visit::default_visitor()
    }));
}

fn traverse_all_resources(cx: ctx, crate_mod: _mod) {
    visit::visit_mod(crate_mod, ast_util::dummy_sp(), 0, cx, visit::mk_vt(@{
        visit_expr: {|_e, _cx, _v|},
        visit_item: {|i, cx, v|
            visit::visit_item(i, cx, v);
            alt i.node {
              item_res(_, _, _, _, _, _) {
                traverse_public_item(cx, i);
              }
              _ {}
            }
        }
        with *visit::default_visitor()
    }));
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

use futures::{future, prelude::*};
use publisher::Publisher as _;
use std::{
    collections::HashMap,
    io,
    net::SocketAddr,
    sync::{Arc, Mutex},
    time::Duration,
};
use subscriber::Subscriber as _;
use tarpc::{
    client, context,
    server::{self, Handler},
};
use tokio_serde::formats::Json;

pub mod subscriber {
    #[tarpc::service]
    pub trait Subscriber {
        async fn receive(message: String);
    }
}

pub mod publisher {
    use std::net::SocketAddr;

    #[tarpc::service]
    pub trait Publisher {
        async fn broadcast(message: String);
        async fn subscribe(id: u32, address: SocketAddr) -> Result<(), String>;
        async fn unsubscribe(id: u32);
    }
}

#[derive(Clone, Debug)]
struct Subscriber {
    id: u32,
}

#[tarpc::server]
impl subscriber::Subscriber for Subscriber {
    async fn receive(self, _: context::Context, message: String) {
        eprintln!("{} received message: {}", self.id, message);
    }
}

impl Subscriber {
    async fn listen(id: u32, config: server::Config) -> io::Result<SocketAddr> {
        let incoming = tarpc::serde_transport::tcp::listen("localhost:0", Json::default)
            .await?
            .filter_map(|r| future::ready(r.ok()));
        let addr = incoming.get_ref().local_addr();
        tokio::spawn(
            server::new(config)
                .incoming(incoming)
                .take(1)
                .respond_with(Subscriber { id }.serve()),
        );
        Ok(addr)
    }
}

#[derive(Clone, Debug)]
struct Publisher {
    clients: Arc<Mutex<HashMap<u32, subscriber::SubscriberClient>>>,
}

impl Publisher {
    fn new() -> Publisher {
        Publisher {
            clients: Arc::new(Mutex::new(HashMap::new())),
        }
    }
}

#[tarpc::server]
impl publisher::Publisher for Publisher {
    async fn broadcast(self, _: context::Context, message: String) {
        let mut clients = self.clients.lock().unwrap().clone();
        for client in clients.values_mut() {
            // Ignore failing subscribers. In a real pubsub,
            // you'd want to continually retry until subscribers
            // ack.
            let _ = client.receive(context::current(), message.clone()).await;
        }
    }

    async fn subscribe(self, _: context::Context, id: u32, addr: SocketAddr) -> Result<(), String> {
        let conn = tarpc::serde_transport::tcp::connect(addr, Json::default())
            .await
            .map_err(|e| e.to_string())?;
        let subscriber = subscriber::SubscriberClient::new(client::Config::default(), conn)
            .spawn()
            .map_err(|e| e.to_string())?;
        eprintln!("Subscribing {}.", id);
        self.clients.lock().unwrap().insert(id, subscriber);
        Ok(())
    }

    async fn unsubscribe(self, _: context::Context, id: u32) {
        eprintln!("Unsubscribing {}", id);
        let mut clients = self.clients.lock().unwrap();
        if clients.remove(&id).is_none() {
            eprintln!(
                "Client {} not found. Existings clients: {:?}",
                id, &*clients
            );
        }
    }
}

#[tokio::main]
async fn main() -> io::Result<()> {
    env_logger::init();

    let transport = tarpc::serde_transport::tcp::listen("localhost:0", Json::default)
        .await?
        .filter_map(|r| future::ready(r.ok()));
    let publisher_addr = transport.get_ref().local_addr();
    tokio::spawn(
        transport
            .take(1)
            .map(server::BaseChannel::with_defaults)
            .respond_with(Publisher::new().serve()),
    );

    let subscriber1 = Subscriber::listen(0, server::Config::default()).await?;
    let subscriber2 = Subscriber::listen(1, server::Config::default()).await?;

    let publisher_conn = tarpc::serde_transport::tcp::connect(publisher_addr, Json::default());
    let publisher_conn = publisher_conn.await?;
    let mut publisher =
        publisher::PublisherClient::new(client::Config::default(), publisher_conn).spawn()?;

    if let Err(e) = publisher
        .subscribe(context::current(), 0, subscriber1)
        .await?
    {
        eprintln!("Couldn't subscribe subscriber 0: {}", e);
    }
    if let Err(e) = publisher
        .subscribe(context::current(), 1, subscriber2)
        .await?
    {
        eprintln!("Couldn't subscribe subscriber 1: {}", e);
    }

    println!("Broadcasting...");
    publisher
        .broadcast(context::current(), "hello to all".to_string())
        .await?;
    publisher.unsubscribe(context::current(), 1).await?;
    publisher
        .broadcast(context::current(), "hi again".to_string())
        .await?;
    drop(publisher);

    tokio::time::delay_for(Duration::from_millis(100)).await;
    println!("Done.");

    Ok(())
}
<commit_msg>Rewrite pubsub example to have the subscriber connect to the publisher.<commit_after>// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

use futures::{
    future::{self, AbortHandle},
    prelude::*,
};
use log::info;
use publisher::Publisher as _;
use std::{
    collections::HashMap,
    io,
    net::SocketAddr,
    sync::{Arc, Mutex},
};
use subscriber::Subscriber as _;
use tarpc::{
    client, context,
    serde_transport::tcp,
    server::{self, Channel},
};
use tokio::net::ToSocketAddrs;
use tokio_serde::formats::Json;

pub mod subscriber {
    #[tarpc::service]
    pub trait Subscriber {
        async fn receive(message: String);
    }
}

pub mod publisher {
    #[tarpc::service]
    pub trait Publisher {
        async fn broadcast(message: String);
    }
}

#[derive(Clone, Debug)]
struct Subscriber {
    local_addr: SocketAddr,
}

#[tarpc::server]
impl subscriber::Subscriber for Subscriber {
    async fn receive(self, _: context::Context, message: String) {
        info!("{} received message: {}", self.local_addr, message);
    }
}

struct SubscriberHandle(AbortHandle);

impl Drop for SubscriberHandle {
    fn drop(&mut self) {
        self.0.abort();
    }
}

impl Subscriber {
    async fn connect(publisher_addr: impl ToSocketAddrs) -> io::Result<SubscriberHandle> {
        let publisher = tcp::connect(publisher_addr, Json::default()).await?;
        let local_addr = publisher.local_addr()?;
        let (handler, abort_handle) = future::abortable(
            server::BaseChannel::with_defaults(publisher)
                .respond_with(Subscriber { local_addr }.serve())
                .execute(),
        );
        tokio::spawn(handler);
        Ok(SubscriberHandle(abort_handle))
    }
}

#[derive(Clone, Debug)]
struct Publisher {
    clients: Arc<Mutex<HashMap<SocketAddr, subscriber::SubscriberClient>>>,
}

struct PublisherAddrs {
    publisher: SocketAddr,
    subscriptions: SocketAddr,
}

impl Publisher {
    async fn start(self) -> io::Result<PublisherAddrs> {
        let mut connecting_publishers = tcp::listen("localhost:0", Json::default).await?;

        let publisher_addrs = PublisherAddrs {
            publisher: connecting_publishers.local_addr(),
            subscriptions: Self::start_subscription_manager(self.clients.clone()).await?,
        };

        tokio::spawn(async move {
            // Because this is just an example, we know there will only be one publisher. In more
            // realistic code, this would be a loop to continually accept new publisher
            // connections.
            let publisher = connecting_publishers.next().await.unwrap().unwrap();
            info!("[{}] publisher connected.", publisher.peer_addr().unwrap());

            server::BaseChannel::with_defaults(publisher)
                .respond_with(self.serve())
                .execute()
                .await
        });

        Ok(publisher_addrs)
    }

    async fn start_subscription_manager(
        clients: Arc<Mutex<HashMap<SocketAddr, subscriber::SubscriberClient>>>,
    ) -> io::Result<SocketAddr> {
        let mut connecting_subscribers = tcp::listen("localhost:0", Json::default)
            .await?
            .filter_map(|r| future::ready(r.ok()));
        let new_subscriber_addr = connecting_subscribers.get_ref().local_addr();

        tokio::spawn(async move {
            while let Some(conn) = connecting_subscribers.next().await {
                let subscriber_addr = conn.peer_addr().unwrap();
                info!("[{}] subscriber connected.", subscriber_addr);

                let tarpc::client::NewClient {
                    client: subscriber,
                    dispatch,
                } = subscriber::SubscriberClient::new(client::Config::default(), conn);
                clients.lock().unwrap().insert(subscriber_addr, subscriber);

                let dropped_clients = clients.clone();
                tokio::spawn(async move {
                    match dispatch.await {
                        Ok(()) => info!("[{:?}] subscriber connection closed", subscriber_addr),
                        Err(e) => info!(
                            "[{:?}] subscriber connection broken: {}",
                            subscriber_addr, e
                        ),
                    }
                    dropped_clients.lock().unwrap().remove(&subscriber_addr);
                });
            }
        });

        Ok(new_subscriber_addr)
    }
}

#[tarpc::server]
impl publisher::Publisher for Publisher {
    async fn broadcast(self, _: context::Context, message: String) {
        info!("received message to broadcast.");
        let mut clients = self.clients.lock().unwrap().clone();
        let mut publications = Vec::new();
        for client in clients.values_mut() {
            publications.push(client.receive(context::current(), message.clone()));
        }
        // Ignore failing subscribers. In a real pubsub, you'd want to continually retry until
        // subscribers ack. Of course, a lot would be different in a real pubsub :)
        future::join_all(publications).await;
    }
}

#[tokio::main]
async fn main() -> io::Result<()> {
    env_logger::init();

    let clients = Arc::new(Mutex::new(HashMap::new()));
    let addrs = Publisher { clients }.start().await?;

    let mut publisher = publisher::PublisherClient::new(
        client::Config::default(),
        tcp::connect(addrs.publisher, Json::default()).await?,
    )
    .spawn()?;

    let _subscriber0 = Subscriber::connect(addrs.subscriptions).await?;
    publisher
        .broadcast(context::current(), "hello to one".to_string())
        .await?;

    let _subscriber1 = Subscriber::connect(addrs.subscriptions).await?;
    publisher
        .broadcast(context::current(), "hello to all".to_string())
        .await?;

    info!("done.");

    Ok(())
}
<|endoftext|>