<commit_before>extern crate tokio_core;
extern crate futures;
extern crate tk_http;
extern crate tk_listen;
extern crate wonky_tanks_server;

use std::cell::{RefCell, Cell};
use std::collections::hash_map::HashMap;
use std::fmt::Write;
use std::rc::Rc;

use tokio_core::reactor::Core;
use tokio_core::net::TcpListener;

use futures::{Stream, Future};
use futures::future::{FutureResult, ok};
use futures::sync::mpsc::{unbounded, UnboundedSender};

use tk_http::{Status};
use tk_http::server::buffered::{Request, BufferedDispatcher};
use tk_http::server::{Encoder, EncoderDone, Config as ServerConfig, Proto, Error as ServerError};
use tk_http::websocket::{Packet, Config as WebsocketConfig, Dispatcher, Frame, Loop, Error as WebsocketError};

use tk_listen::ListenExt;

use wonky_tanks_server::*;

const LOBBY: &'static str = "/";

type SharedData = Rc<RefCell<HashMap<CountedString, ChannelState>>>;
type TransientData = (SharedData, String);

#[derive(Clone, Default)]
pub struct ChannelState {
    capacity: Option<u8>,
    participants: Vec<Rc<UnboundedSender<Packet>>>,
}

impl std::fmt::Debug for ChannelState {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        f
            .debug_struct("ChannelState")
            .field("capacity", &self.capacity)
            .field(
                "participants",
                &format_args!(
                    "{:x}",
                    wonky_tanks_server::IterContentsDebug(||
                        self
                            .participants
                            .iter()
                            .map(|i| (&i) as &UnboundedSender<_> as *const _ as usize)
                    )
                )
            )
            .finish()
    }
}

fn service<S>(directory: &Directory, data: &mut TransientData, req: Request, mut e: Encoder<S>) -> FutureResult<EncoderDone<S>, ServerError> {
    if let Some(ws) = req.websocket_handshake() {
        data.1.push_str(req.path());
        e.status(Status::SwitchingProtocol);
        e.add_header("Connection", "upgrade").unwrap();
        e.add_header("Upgrade", "websocket").unwrap();
        e.format_header("Sec-Websocket-Accept", &ws.accept).unwrap();
        e.done_headers().unwrap();
        return ok(e.done())
    }

    if let Some((ctype, data)) = directory.get(req.path()) {
        e.status(Status::Ok);
        e.add_length(data.len() as u64).unwrap();
        e.add_header("Content-Type", ctype).unwrap();
        if e.done_headers().unwrap() {
            e.write_body(&data);
        }
        return ok(e.done())
    }

    let not_found = "No data found".as_bytes();
    e.status(Status::NotFound);
    e.add_length(not_found.len() as u64).unwrap();
    e.add_header("Content-Type", "text/plain").unwrap();
    if e.done_headers().unwrap() {
        e.write_body(not_found);
    }
    ok(e.done())
}

struct WebsocketConnection {
    shared_data: SharedData,
    sender: Rc<UnboundedSender<Packet>>,
    token: CountedString,
}

impl WebsocketConnection {
    pub fn new(data: TransientData, sender: Rc<UnboundedSender<Packet>>) -> Self {
        let (shared_data, token) = data;
        let token: CountedString = token.into();
        let mut data: String = Default::default();
        match match {
            let shared_data = &mut *shared_data.borrow_mut();
            let shared_data = shared_data.entry(token.clone());
            #[cfg(debug_assertions)]
            println!("Previous data: {:?}", &shared_data);
            let &mut ChannelState { ref mut participants, capacity } = shared_data.or_insert_with(Default::default);
            participants.push(sender.clone());
            data.clear();
            write!(data, "count: {}", participants.len()).unwrap();
            let (last, participants) = participants.split_last().unwrap();
            for participant in participants {
                drop(participant.unbounded_send(Packet::Text(data.clone())));
            }
            drop(last.unbounded_send(Packet::Text(data.clone())));
            (capacity, participants.len())
        } {
            (Some(capacity), 1) if 1 < capacity => Some(true),
            (None, 1) => Some(true),
            (Some(capacity), count) if count == capacity as usize && count != 1 => Some(false),
            _ => None,
        } {
            Some(toggle) => {
                if let Some(lobby) = shared_data.try_borrow().unwrap().get(LOBBY) {
                    data.clear();
                    write!(data, "{}: {}", if toggle { "+" } else { "-" }, token).unwrap();
                    let (last, participants) = lobby.participants.split_last().unwrap();
                    for participant in participants {
                        drop(participant.unbounded_send(Packet::Text(data.clone())));
                    }
                    drop(last.unbounded_send(Packet::Text(data)));
                }
            },
            None => {},
        }
        #[cfg(debug_assertions)]
        println!("Created connection: {:x}/{}", (&sender) as &UnboundedSender<_> as *const _ as usize, token);
        WebsocketConnection { shared_data, sender, token }
    }

    fn is_lobby(&self) -> bool {
        LOBBY == &*self.token.0
    }
}

impl Drop for WebsocketConnection {
    fn drop(&mut self) {
        let &mut WebsocketConnection {
            ref mut shared_data,
            ref sender,
            ref token
        } = self;
        use std::collections::hash_map::Entry::Occupied;
        if Some(1) != {
            let mut map = shared_data.borrow_mut();
            let mut entry =
                if let Occupied(entry) = map.entry(token.clone()) {
                    entry
                } else {
                    panic!("Bad state; expected {} to be occupied: {:?}", token, shared_data)
                };
            if let (capacity, 0) = {
                let &mut ChannelState { ref mut participants, capacity } = entry.get_mut();
                participants.retain(|target| !Rc::ptr_eq(target, sender));
                (capacity, participants.len())
            } {
                entry.remove_entry();
                capacity
            } else {
                Some(1)
            }
        } {
            if let Some(lobby) = shared_data.try_borrow().unwrap().get(LOBBY) {
                let data = format!("-: {}", token);
                let (last, participants) = lobby.participants.split_last().unwrap();
                for participant in participants {
                    drop(participant.unbounded_send(Packet::Text(data.clone())));
                }
                drop(last.unbounded_send(Packet::Text(data)));
            }
        }
        #[cfg(debug_assertions)]
        println!("Dropped connection: {:x}/{}", (&self.sender) as &UnboundedSender<_> as *const _ as usize, token);
    }
}

impl Dispatcher for WebsocketConnection {
    type Future = FutureResult<(), WebsocketError>;

    fn frame(&mut self, frame: &Frame) -> FutureResult<(), WebsocketError> {
        if self.is_lobby() { return ok(()) }
        let &mut WebsocketConnection { ref shared_data, ref sender, ref token } = self;

        match frame {
            &Frame::Text(t) => {
                // FIXME - parse state changes
            },
            &Frame::Binary(b) => {
                let &ChannelState { ref participants, .. } = &(&*shared_data.borrow())[token];
                if participants.len() == 1 { return ok(()) }
                let mut participants = participants.iter().filter(|&participant| !Rc::ptr_eq(participant, sender));
                let last = participants.next_back().unwrap();
                let to_send = Packet::Binary(Vec::from(b));
                for participant in participants {
                    drop(participant.unbounded_send(to_send.clone()));
                }
                drop(last.unbounded_send(to_send));
            },
            _ => {},
        }

        ok(())
    }
}

fn main() {
    let directory = &Directory::new(std::env::current_dir().unwrap().as_path());
    let shared_data: SharedData = Default::default();
    let mut lp = Core::new().unwrap();
    let h1 = lp.handle();

    let addr = "0.0.0.0:5000".parse().unwrap();
    let listener = TcpListener::bind(&addr, &lp.handle()).unwrap();
    let cfg = ServerConfig::new().done();
    let wcfg = WebsocketConfig::new().done();

    let pipeline = listener.incoming();
    let pipeline =
        pipeline.map(
        move |(socket, addr)| {
            let wcfg = wcfg.clone();
            let h2 = h1.clone();

            let data: TransientData = (shared_data.clone(), Default::default());
            let data = Rc::new(RefCell::new(data));
            let http_data = data.clone();
            let websocket_data = Cell::new(Some(data));

            Proto::new(
                socket,
                &cfg,
                BufferedDispatcher::new_with_websockets(
                    addr,
                    &h1,
                    move |p1, p2|
                        service(directory, &mut *http_data.borrow_mut(),p1, p2),
                    move |outp, inp| {
                        let (sender, receiver) = unbounded();
                        let sender = Rc::new(sender);
                        Loop::server(
                            outp,
                            inp,
                            receiver.map_err(move |e| format!("{:?} : stream closed", e)),
                            WebsocketConnection::new(
                                Rc::try_unwrap(websocket_data.replace(None).unwrap()).unwrap().into_inner(),
                                sender
                            ),
                            &wcfg,
                            &h2
                        )
                            .map_err(|e| println!("Websocket err: {:?}", e))
                    }
                ),
                &h1
            )
                .map_err(|e| {println!("Connection error: {:?}", e);})
                .then(|_| Ok(())) // don't fail, please
        }
    );
    let pipeline = pipeline.listen(1000);

    lp.run(pipeline).unwrap();
}<commit_msg>Remove extra qualifier for IterContentsDebug<commit_after>extern crate tokio_core;
extern crate futures;
extern crate tk_http;
extern crate tk_listen;
extern crate wonky_tanks_server;

use std::cell::{RefCell, Cell};
use std::collections::hash_map::HashMap;
use std::fmt::Write;
use std::rc::Rc;

use tokio_core::reactor::Core;
use tokio_core::net::TcpListener;

use futures::{Stream, Future};
use futures::future::{FutureResult, ok};
use futures::sync::mpsc::{unbounded, UnboundedSender};

use tk_http::{Status};
use tk_http::server::buffered::{Request, BufferedDispatcher};
use tk_http::server::{Encoder, EncoderDone, Config as ServerConfig, Proto, Error as ServerError};
use tk_http::websocket::{Packet, Config as WebsocketConfig, Dispatcher, Frame, Loop, Error as WebsocketError};

use tk_listen::ListenExt;

use wonky_tanks_server::*;

const LOBBY: &'static str = "/";

type SharedData = Rc<RefCell<HashMap<CountedString, ChannelState>>>;
type TransientData = (SharedData, String);

#[derive(Clone, Default)]
pub struct ChannelState {
    capacity: Option<u8>,
    participants: Vec<Rc<UnboundedSender<Packet>>>,
}

impl std::fmt::Debug for ChannelState {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        f
            .debug_struct("ChannelState")
            .field("capacity", &self.capacity)
            .field(
                "participants",
                &format_args!(
                    "{:x}",
                    IterContentsDebug(||
                        self
                            .participants
                            .iter()
                            .map(|i| (&i) as &UnboundedSender<_> as *const _ as usize)
                    )
                )
            )
            .finish()
    }
}

fn service<S>(directory: &Directory, data: &mut TransientData, req: Request, mut e: Encoder<S>) -> FutureResult<EncoderDone<S>, ServerError> {
    if let Some(ws) = req.websocket_handshake() {
        data.1.push_str(req.path());
        e.status(Status::SwitchingProtocol);
        e.add_header("Connection", "upgrade").unwrap();
        e.add_header("Upgrade", "websocket").unwrap();
        e.format_header("Sec-Websocket-Accept", &ws.accept).unwrap();
        e.done_headers().unwrap();
        return ok(e.done())
    }

    if let Some((ctype, data)) = directory.get(req.path()) {
        e.status(Status::Ok);
        e.add_length(data.len() as u64).unwrap();
        e.add_header("Content-Type", ctype).unwrap();
        if e.done_headers().unwrap() {
            e.write_body(&data);
        }
        return ok(e.done())
    }

    let not_found = "No data found".as_bytes();
    e.status(Status::NotFound);
    e.add_length(not_found.len() as u64).unwrap();
    e.add_header("Content-Type", "text/plain").unwrap();
    if e.done_headers().unwrap() {
        e.write_body(not_found);
    }
    ok(e.done())
}

struct WebsocketConnection {
    shared_data: SharedData,
    sender: Rc<UnboundedSender<Packet>>,
    token: CountedString,
}

impl WebsocketConnection {
    pub fn new(data: TransientData, sender: Rc<UnboundedSender<Packet>>) -> Self {
        let (shared_data, token) = data;
        let token: CountedString = token.into();
        let mut data: String = Default::default();
        match match {
            let shared_data = &mut *shared_data.borrow_mut();
            let shared_data = shared_data.entry(token.clone());
            #[cfg(debug_assertions)]
            println!("Previous data: {:?}", &shared_data);
            let &mut ChannelState { ref mut participants, capacity } = shared_data.or_insert_with(Default::default);
            participants.push(sender.clone());
            data.clear();
            write!(data, "count: {}", participants.len()).unwrap();
            let (last, participants) = participants.split_last().unwrap();
            for participant in participants {
                drop(participant.unbounded_send(Packet::Text(data.clone())));
            }
            drop(last.unbounded_send(Packet::Text(data.clone())));
            (capacity, participants.len())
        } {
            (Some(capacity), 1) if 1 < capacity => Some(true),
            (None, 1) => Some(true),
            (Some(capacity), count) if count == capacity as usize && count != 1 => Some(false),
            _ => None,
        } {
            Some(toggle) => {
                if let Some(lobby) = shared_data.try_borrow().unwrap().get(LOBBY) {
                    data.clear();
                    write!(data, "{}: {}", if toggle { "+" } else { "-" }, token).unwrap();
                    let (last, participants) = lobby.participants.split_last().unwrap();
                    for participant in participants {
                        drop(participant.unbounded_send(Packet::Text(data.clone())));
                    }
                    drop(last.unbounded_send(Packet::Text(data)));
                }
            },
            None => {},
        }
        #[cfg(debug_assertions)]
        println!("Created connection: {:x}/{}", (&sender) as &UnboundedSender<_> as *const _ as usize, token);
        WebsocketConnection { shared_data, sender, token }
    }

    fn is_lobby(&self) -> bool {
        LOBBY == &*self.token.0
    }
}

impl Drop for WebsocketConnection {
    fn drop(&mut self) {
        let &mut WebsocketConnection {
            ref mut shared_data,
            ref sender,
            ref token
        } = self;
        use std::collections::hash_map::Entry::Occupied;
        if Some(1) != {
            let mut map = shared_data.borrow_mut();
            let mut entry =
                if let Occupied(entry) = map.entry(token.clone()) {
                    entry
                } else {
                    panic!("Bad state; expected {} to be occupied: {:?}", token, shared_data)
                };
            if let (capacity, 0) = {
                let &mut ChannelState { ref mut participants, capacity } = entry.get_mut();
                participants.retain(|target| !Rc::ptr_eq(target, sender));
                (capacity, participants.len())
            } {
                entry.remove_entry();
                capacity
            } else {
                Some(1)
            }
        } {
            if let Some(lobby) = shared_data.try_borrow().unwrap().get(LOBBY) {
                let data = format!("-: {}", token);
                let (last, participants) = lobby.participants.split_last().unwrap();
                for participant in participants {
                    drop(participant.unbounded_send(Packet::Text(data.clone())));
                }
                drop(last.unbounded_send(Packet::Text(data)));
            }
        }
        #[cfg(debug_assertions)]
        println!("Dropped connection: {:x}/{}", (&self.sender) as &UnboundedSender<_> as *const _ as usize, token);
    }
}

impl Dispatcher for WebsocketConnection {
    type Future = FutureResult<(), WebsocketError>;

    fn frame(&mut self, frame: &Frame) -> FutureResult<(), WebsocketError> {
        if self.is_lobby() { return ok(()) }
        let &mut WebsocketConnection { ref shared_data, ref sender, ref token } = self;

        match frame {
            &Frame::Text(t) => {
                // FIXME - parse state changes
            },
            &Frame::Binary(b) => {
                let &ChannelState { ref participants, .. } = &(&*shared_data.borrow())[token];
                if participants.len() == 1 { return ok(()) }
                let mut participants = participants.iter().filter(|&participant| !Rc::ptr_eq(participant, sender));
                let last = participants.next_back().unwrap();
                let to_send = Packet::Binary(Vec::from(b));
                for participant in participants {
                    drop(participant.unbounded_send(to_send.clone()));
                }
                drop(last.unbounded_send(to_send));
            },
            _ => {},
        }

        ok(())
    }
}

fn main() {
    let directory = &Directory::new(std::env::current_dir().unwrap().as_path());
    let shared_data: SharedData = Default::default();
    let mut lp = Core::new().unwrap();
    let h1 = lp.handle();

    let addr = "0.0.0.0:5000".parse().unwrap();
    let listener = TcpListener::bind(&addr, &lp.handle()).unwrap();
    let cfg = ServerConfig::new().done();
    let wcfg = WebsocketConfig::new().done();

    let pipeline = listener.incoming();
    let pipeline =
        pipeline.map(
        move |(socket, addr)| {
            let wcfg = wcfg.clone();
            let h2 = h1.clone();

            let data: TransientData = (shared_data.clone(), Default::default());
            let data = Rc::new(RefCell::new(data));
            let http_data = data.clone();
            let websocket_data = Cell::new(Some(data));

            Proto::new(
                socket,
                &cfg,
                BufferedDispatcher::new_with_websockets(
                    addr,
                    &h1,
                    move |p1, p2|
                        service(directory, &mut *http_data.borrow_mut(),p1, p2),
                    move |outp, inp| {
                        let (sender, receiver) = unbounded();
                        let sender = Rc::new(sender);
                        Loop::server(
                            outp,
                            inp,
                            receiver.map_err(move |e| format!("{:?} : stream closed", e)),
                            WebsocketConnection::new(
                                Rc::try_unwrap(websocket_data.replace(None).unwrap()).unwrap().into_inner(),
                                sender
                            ),
                            &wcfg,
                            &h2
                        )
                            .map_err(|e| println!("Websocket err: {:?}", e))
                    }
                ),
                &h1
            )
                .map_err(|e| {println!("Connection error: {:?}", e);})
                .then(|_| Ok(())) // don't fail, please
        }
    );
    let pipeline = pipeline.listen(1000);

    lp.run(pipeline).unwrap();
}<|endoftext|>


<commit_before>//! Hostcalls that implement
//! [WASI](https://github.com/CraneStation/wasmtime-wasi/blob/wasi/docs/WASI-overview.md).
//!
//! This code borrows heavily from [wasmtime-wasi](https://github.com/CraneStation/wasmtime-wasi),
//! which in turn borrows from cloudabi-utils. See `LICENSE.wasmtime-wasi` for license information.
//!
//! This is currently a very incomplete prototype, only supporting the hostcalls required to run
//! `/examples/hello.c`, and using a bare-bones translation of the capabilities system rather than
//! something nice.

#![allow(non_camel_case_types)]
use crate::memory::*;
use crate::{host, wasm32};
use lucet_runtime::vmctx::{lucet_vmctx, Vmctx};
use std::collections::HashMap;
use std::fs::File;
use std::io::{stderr, stdin, stdout};
use std::os::unix::prelude::{AsRawFd, FileTypeExt, FromRawFd, RawFd};

#[no_mangle]
pub extern "C" fn __wasi_proc_exit(_vmctx: *mut lucet_vmctx, rval: wasm32::__wasi_exitcode_t) -> ! {
    std::process::exit(rval as i32)
}

pub struct WasiCtx {
    fds: HashMap<host::__wasi_fd_t, FdEntry>,
}

impl WasiCtx {
    pub fn new() -> WasiCtx {
        use nix::unistd::dup;

        let mut ctx = WasiCtx {
            fds: HashMap::new(),
        };
        ctx.insert_existing_fd(0, dup(stdin().as_raw_fd()).unwrap());
        ctx.insert_existing_fd(1, dup(stdout().as_raw_fd()).unwrap());
        ctx.insert_existing_fd(2, dup(stderr().as_raw_fd()).unwrap());
        ctx
    }

    fn get_fd_entry(
        &self,
        fd: host::__wasi_fd_t,
        rights_base: host::__wasi_rights_t,
        rights_inheriting: host::__wasi_rights_t,
    ) -> Result<&FdEntry, host::__wasi_errno_t> {
        if let Some(ref fe) = self.fds.get(&fd) {
            // validate rights
            if !fe.rights_base & rights_base != 0 || !fe.rights_inheriting & rights_inheriting != 0
            {
                Err(host::__WASI_ENOTCAPABLE as host::__wasi_errno_t)
            } else {
                Ok(fe)
            }
        } else {
            Err(host::__WASI_EBADF as host::__wasi_errno_t)
        }
    }

    pub fn insert_existing_fd(&mut self, fd: host::__wasi_fd_t, rawfd: RawFd) {
        assert!(!self.fds.contains_key(&fd));
        self.fds.insert(fd, unsafe { FdEntry::from_raw_fd(rawfd) });
    }
}

struct FdEntry {
    fd_object: FdObject,
    rights_base: host::__wasi_rights_t,
    rights_inheriting: host::__wasi_rights_t,
}

impl FromRawFd for FdEntry {
    unsafe fn from_raw_fd(rawfd: RawFd) -> FdEntry {
        let (ty, mut rights_base, rights_inheriting) = {
            let file = File::from_raw_fd(rawfd);
            let ft = file.metadata().unwrap().file_type();
            // we just make a `File` here for convenience; we don't want it to close when it drops
            std::mem::forget(file);
            if ft.is_block_device() {
                (
                    host::__WASI_FILETYPE_BLOCK_DEVICE,
                    host::RIGHTS_BLOCK_DEVICE_BASE,
                    host::RIGHTS_BLOCK_DEVICE_INHERITING,
                )
            } else if ft.is_char_device() {
                if nix::unistd::isatty(rawfd).unwrap() {
                    (
                        host::__WASI_FILETYPE_CHARACTER_DEVICE,
                        host::RIGHTS_TTY_BASE,
                        host::RIGHTS_TTY_BASE,
                    )
                } else {
                    (
                        host::__WASI_FILETYPE_CHARACTER_DEVICE,
                        host::RIGHTS_CHARACTER_DEVICE_BASE,
                        host::RIGHTS_CHARACTER_DEVICE_INHERITING,
                    )
                }
            } else if ft.is_dir() {
                (
                    host::__WASI_FILETYPE_DIRECTORY,
                    host::RIGHTS_DIRECTORY_BASE,
                    host::RIGHTS_DIRECTORY_INHERITING,
                )
            } else if ft.is_file() {
                (
                    host::__WASI_FILETYPE_REGULAR_FILE,
                    host::RIGHTS_REGULAR_FILE_BASE,
                    host::RIGHTS_REGULAR_FILE_INHERITING,
                )
            } else if ft.is_socket() {
                use nix::sys::socket;
                match socket::getsockopt(rawfd, socket::sockopt::SockType).unwrap() {
                    socket::SockType::Datagram => (
                        host::__WASI_FILETYPE_SOCKET_DGRAM,
                        host::RIGHTS_SOCKET_BASE,
                        host::RIGHTS_SOCKET_INHERITING,
                    ),
                    socket::SockType::Stream => (
                        host::__WASI_FILETYPE_SOCKET_STREAM,
                        host::RIGHTS_SOCKET_BASE,
                        host::RIGHTS_SOCKET_INHERITING,
                    ),
                    s => panic!("unsupported socket type: {:?}", s),
                }
            } else if ft.is_fifo() {
                (
                    host::__WASI_FILETYPE_SOCKET_STREAM,
                    host::RIGHTS_SOCKET_BASE,
                    host::RIGHTS_SOCKET_INHERITING,
                )
            } else {
                panic!("unsupported file type: {:?}", ft);
            }
        };

        use nix::fcntl::{fcntl, OFlag, F_GETFL};
        let flags_bits = fcntl(rawfd, F_GETFL).expect("fcntl succeeds");
        let flags = OFlag::from_bits_truncate(flags_bits);
        let accmode = flags & OFlag::O_ACCMODE;
        if accmode == OFlag::O_RDONLY {
            rights_base &= !host::__WASI_RIGHT_FD_WRITE as host::__wasi_rights_t;
        } else if accmode == OFlag::O_WRONLY {
            rights_base &= !host::__WASI_RIGHT_FD_READ as host::__wasi_rights_t;
        }

        FdEntry {
            fd_object: FdObject {
                ty: ty as u8,
                rawfd,
            },
            rights_base,
            rights_inheriting,
        }
    }
}

struct FdObject {
    ty: host::__wasi_filetype_t,
    rawfd: RawFd,
    // TODO: directories
}

#[no_mangle]
pub extern "C" fn __wasi_fd_close(
    vmctx: *mut lucet_vmctx,
    fd: wasm32::__wasi_fd_t,
) -> wasm32::__wasi_errno_t {
    let vmctx = unsafe { Vmctx::from_raw(vmctx) };
    let mut ctx = unsafe { Box::from_raw(vmctx.embed_ctx() as *mut WasiCtx) };
    let fd = dec_fd(fd);
    if let Some(fdent) = ctx.fds.remove(&fd) {
        match nix::unistd::close(fdent.fd_object.rawfd) {
            Ok(_) => wasm32::__WASI_ESUCCESS,
            Err(e) => wasm32::errno_from_nix(e.as_errno().unwrap()),
        }
    } else {
        wasm32::__WASI_EBADF
    }
}

#[no_mangle]
pub extern "C" fn __wasi_fd_stat_get(
    vmctx: *mut lucet_vmctx,
    fd: wasm32::__wasi_fd_t,
    fdstat_ptr: wasm32::uintptr_t, // *mut wasm32::__wasi_fdstat_t
) -> wasm32::__wasi_errno_t {
    let mut vmctx = unsafe { Vmctx::from_raw(vmctx) };
    let ctx = unsafe { Box::from_raw(vmctx.embed_ctx() as *mut WasiCtx) };

    let host_fd = dec_fd(fd);
    let mut host_fdstat = match unsafe { dec_fdstat_byref(&mut vmctx, fdstat_ptr) } {
        Ok(host_fdstat) => host_fdstat,
        Err(e) => return enc_errno(e),
    };

    let errno = if let Some(ref fe) = ctx.fds.get(&host_fd) {
        host_fdstat.fs_filetype = fe.fd_object.ty;
        host_fdstat.fs_rights_base = fe.rights_base;
        host_fdstat.fs_rights_inheriting = fe.rights_inheriting;
        use nix::fcntl::{fcntl, OFlag, F_GETFL};
        match fcntl(fe.fd_object.rawfd, F_GETFL).map(OFlag::from_bits) {
            Ok(Some(flags)) => {
                if flags.contains(OFlag::O_APPEND) {
                    host_fdstat.fs_flags |= wasm32::__WASI_FDFLAG_APPEND;
                }
                if flags.contains(OFlag::O_DSYNC) {
                    host_fdstat.fs_flags |= wasm32::__WASI_FDFLAG_DSYNC;
                }
                if flags.contains(OFlag::O_NONBLOCK) {
                    host_fdstat.fs_flags |= wasm32::__WASI_FDFLAG_NONBLOCK;
                }
                if flags.contains(OFlag::O_RSYNC) {
                    host_fdstat.fs_flags |= wasm32::__WASI_FDFLAG_RSYNC;
                }
                if flags.contains(OFlag::O_SYNC) {
                    host_fdstat.fs_flags |= wasm32::__WASI_FDFLAG_SYNC;
                }
                wasm32::__WASI_ESUCCESS
            }
            Ok(None) => wasm32::__WASI_ENOSYS,
            Err(e) => wasm32::errno_from_nix(e.as_errno().unwrap()),
        }
    } else {
        return wasm32::__WASI_EBADF;
    };

    unsafe {
        enc_fdstat_byref(&mut vmctx, fdstat_ptr, host_fdstat)
            .expect("can write back into the pointer we read from");
    }

    errno
}

#[no_mangle]
pub extern "C" fn __wasi_fd_seek(
    vmctx: *mut lucet_vmctx,
    fd: wasm32::__wasi_fd_t,
    offset: wasm32::__wasi_filedelta_t,
    whence: wasm32::__wasi_whence_t,
    newoffset: wasm32::uintptr_t,
) -> wasm32::__wasi_errno_t {
    let mut vmctx = unsafe { Vmctx::from_raw(vmctx) };
    let ctx = unsafe { Box::from_raw(vmctx.embed_ctx() as *mut WasiCtx) };
    let fd = dec_fd(fd);
    let offset = dec_filedelta(offset);
    let whence = dec_whence(whence);

    let host_newoffset = {
        use nix::unistd::{lseek, Whence};
        let nwhence = match whence as u32 {
            host::__WASI_WHENCE_CUR => Whence::SeekCur,
            host::__WASI_WHENCE_END => Whence::SeekEnd,
            host::__WASI_WHENCE_SET => Whence::SeekSet,
            _ => return wasm32::__WASI_EINVAL,
        };

        let rights = if offset == 0 && whence as u32 == host::__WASI_WHENCE_CUR {
            host::__WASI_RIGHT_FD_TELL
        } else {
            host::__WASI_RIGHT_FD_SEEK | host::__WASI_RIGHT_FD_TELL
        };
        match ctx.get_fd_entry(fd, rights.into(), 0) {
            Ok(fe) => match lseek(fe.fd_object.rawfd, offset, nwhence) {
                Ok(newoffset) => newoffset,
                Err(e) => return wasm32::errno_from_nix(e.as_errno().unwrap()),
            },
            Err(e) => return enc_errno(e),
        }
    };

    unsafe {
        enc_filesize_byref(&mut vmctx, newoffset, host_newoffset as u64)
            .map(|_| wasm32::__WASI_ESUCCESS)
            .unwrap_or_else(|e| e)
    }
}

#[no_mangle]
pub extern "C" fn __wasi_fd_write(
    vmctx: *mut lucet_vmctx,
    fd: wasm32::__wasi_fd_t,
    iovs_ptr: wasm32::uintptr_t,
    iovs_len: wasm32::size_t,
    nwritten: wasm32::uintptr_t,
) -> wasm32::__wasi_errno_t {
    use nix::sys::uio::{writev, IoVec};

    let mut vmctx = unsafe { Vmctx::from_raw(vmctx) };
    let ctx = unsafe { Box::from_raw(vmctx.embed_ctx() as *mut WasiCtx) };
    let fd = dec_fd(fd);
    let iovs = match unsafe { dec_ciovec_slice(&mut vmctx, iovs_ptr, iovs_len) } {
        Ok(iovs) => iovs,
        Err(e) => return enc_errno(e),
    };

    let fe = match ctx.get_fd_entry(fd, host::__WASI_RIGHT_FD_WRITE.into(), 0) {
        Ok(fe) => fe,
        Err(e) => return enc_errno(e),
    };

    let iovs: Vec<IoVec<&[u8]>> = iovs
        .iter()
        .map(|iov| unsafe { host::ciovec_to_nix(iov) })
        .collect();

    let host_nwritten = match writev(fe.fd_object.rawfd, &iovs) {
        Ok(len) => len,
        Err(e) => return wasm32::errno_from_nix(e.as_errno().unwrap()),
    };

    unsafe {
        enc_usize_byref(&mut vmctx, nwritten, host_nwritten)
            .map(|_| wasm32::__WASI_ESUCCESS)
            .unwrap_or_else(|e| e)
    }
}
<commit_msg>hostcalls: use new embed ctx methods<commit_after>//! Hostcalls that implement
//! [WASI](https://github.com/CraneStation/wasmtime-wasi/blob/wasi/docs/WASI-overview.md).
//!
//! This code borrows heavily from [wasmtime-wasi](https://github.com/CraneStation/wasmtime-wasi),
//! which in turn borrows from cloudabi-utils. See `LICENSE.wasmtime-wasi` for license information.
//!
//! This is currently a very incomplete prototype, only supporting the hostcalls required to run
//! `/examples/hello.c`, and using a bare-bones translation of the capabilities system rather than
//! something nice.

#![allow(non_camel_case_types)]
use crate::memory::*;
use crate::{host, wasm32};
use lucet_runtime::vmctx::{lucet_vmctx, Vmctx};
use std::collections::HashMap;
use std::fs::File;
use std::io::{stderr, stdin, stdout};
use std::os::unix::prelude::{AsRawFd, FileTypeExt, FromRawFd, RawFd};

#[no_mangle]
pub extern "C" fn __wasi_proc_exit(_vmctx: *mut lucet_vmctx, rval: wasm32::__wasi_exitcode_t) -> ! {
    std::process::exit(rval as i32)
}

pub struct WasiCtx {
    fds: HashMap<host::__wasi_fd_t, FdEntry>,
}

impl WasiCtx {
    pub fn new() -> WasiCtx {
        use nix::unistd::dup;

        let mut ctx = WasiCtx {
            fds: HashMap::new(),
        };
        ctx.insert_existing_fd(0, dup(stdin().as_raw_fd()).unwrap());
        ctx.insert_existing_fd(1, dup(stdout().as_raw_fd()).unwrap());
        ctx.insert_existing_fd(2, dup(stderr().as_raw_fd()).unwrap());
        ctx
    }

    fn get_fd_entry(
        &self,
        fd: host::__wasi_fd_t,
        rights_base: host::__wasi_rights_t,
        rights_inheriting: host::__wasi_rights_t,
    ) -> Result<&FdEntry, host::__wasi_errno_t> {
        if let Some(ref fe) = self.fds.get(&fd) {
            // validate rights
            if !fe.rights_base & rights_base != 0 || !fe.rights_inheriting & rights_inheriting != 0
            {
                Err(host::__WASI_ENOTCAPABLE as host::__wasi_errno_t)
            } else {
                Ok(fe)
            }
        } else {
            Err(host::__WASI_EBADF as host::__wasi_errno_t)
        }
    }

    pub fn insert_existing_fd(&mut self, fd: host::__wasi_fd_t, rawfd: RawFd) {
        assert!(!self.fds.contains_key(&fd));
        self.fds.insert(fd, unsafe { FdEntry::from_raw_fd(rawfd) });
    }
}

struct FdEntry {
    fd_object: FdObject,
    rights_base: host::__wasi_rights_t,
    rights_inheriting: host::__wasi_rights_t,
}

impl FromRawFd for FdEntry {
    unsafe fn from_raw_fd(rawfd: RawFd) -> FdEntry {
        let (ty, mut rights_base, rights_inheriting) = {
            let file = File::from_raw_fd(rawfd);
            let ft = file.metadata().unwrap().file_type();
            // we just make a `File` here for convenience; we don't want it to close when it drops
            std::mem::forget(file);
            if ft.is_block_device() {
                (
                    host::__WASI_FILETYPE_BLOCK_DEVICE,
                    host::RIGHTS_BLOCK_DEVICE_BASE,
                    host::RIGHTS_BLOCK_DEVICE_INHERITING,
                )
            } else if ft.is_char_device() {
                if nix::unistd::isatty(rawfd).unwrap() {
                    (
                        host::__WASI_FILETYPE_CHARACTER_DEVICE,
                        host::RIGHTS_TTY_BASE,
                        host::RIGHTS_TTY_BASE,
                    )
                } else {
                    (
                        host::__WASI_FILETYPE_CHARACTER_DEVICE,
                        host::RIGHTS_CHARACTER_DEVICE_BASE,
                        host::RIGHTS_CHARACTER_DEVICE_INHERITING,
                    )
                }
            } else if ft.is_dir() {
                (
                    host::__WASI_FILETYPE_DIRECTORY,
                    host::RIGHTS_DIRECTORY_BASE,
                    host::RIGHTS_DIRECTORY_INHERITING,
                )
            } else if ft.is_file() {
                (
                    host::__WASI_FILETYPE_REGULAR_FILE,
                    host::RIGHTS_REGULAR_FILE_BASE,
                    host::RIGHTS_REGULAR_FILE_INHERITING,
                )
            } else if ft.is_socket() {
                use nix::sys::socket;
                match socket::getsockopt(rawfd, socket::sockopt::SockType).unwrap() {
                    socket::SockType::Datagram => (
                        host::__WASI_FILETYPE_SOCKET_DGRAM,
                        host::RIGHTS_SOCKET_BASE,
                        host::RIGHTS_SOCKET_INHERITING,
                    ),
                    socket::SockType::Stream => (
                        host::__WASI_FILETYPE_SOCKET_STREAM,
                        host::RIGHTS_SOCKET_BASE,
                        host::RIGHTS_SOCKET_INHERITING,
                    ),
                    s => panic!("unsupported socket type: {:?}", s),
                }
            } else if ft.is_fifo() {
                (
                    host::__WASI_FILETYPE_SOCKET_STREAM,
                    host::RIGHTS_SOCKET_BASE,
                    host::RIGHTS_SOCKET_INHERITING,
                )
            } else {
                panic!("unsupported file type: {:?}", ft);
            }
        };

        use nix::fcntl::{fcntl, OFlag, F_GETFL};
        let flags_bits = fcntl(rawfd, F_GETFL).expect("fcntl succeeds");
        let flags = OFlag::from_bits_truncate(flags_bits);
        let accmode = flags & OFlag::O_ACCMODE;
        if accmode == OFlag::O_RDONLY {
            rights_base &= !host::__WASI_RIGHT_FD_WRITE as host::__wasi_rights_t;
        } else if accmode == OFlag::O_WRONLY {
            rights_base &= !host::__WASI_RIGHT_FD_READ as host::__wasi_rights_t;
        }

        FdEntry {
            fd_object: FdObject {
                ty: ty as u8,
                rawfd,
            },
            rights_base,
            rights_inheriting,
        }
    }
}

struct FdObject {
    ty: host::__wasi_filetype_t,
    rawfd: RawFd,
    // TODO: directories
}

#[no_mangle]
pub extern "C" fn __wasi_fd_close(
    vmctx: *mut lucet_vmctx,
    fd: wasm32::__wasi_fd_t,
) -> wasm32::__wasi_errno_t {
    let mut vmctx = unsafe { Vmctx::from_raw(vmctx) };
    let ctx: &mut WasiCtx = vmctx.get_embed_ctx_mut();
    let fd = dec_fd(fd);
    if let Some(fdent) = ctx.fds.remove(&fd) {
        match nix::unistd::close(fdent.fd_object.rawfd) {
            Ok(_) => wasm32::__WASI_ESUCCESS,
            Err(e) => wasm32::errno_from_nix(e.as_errno().unwrap()),
        }
    } else {
        wasm32::__WASI_EBADF
    }
}

#[no_mangle]
pub extern "C" fn __wasi_fd_stat_get(
    vmctx: *mut lucet_vmctx,
    fd: wasm32::__wasi_fd_t,
    fdstat_ptr: wasm32::uintptr_t, // *mut wasm32::__wasi_fdstat_t
) -> wasm32::__wasi_errno_t {
    let mut vmctx = unsafe { Vmctx::from_raw(vmctx) };
    let ctx: &mut WasiCtx = vmctx.get_embed_ctx_mut();

    let host_fd = dec_fd(fd);
    let mut host_fdstat = match unsafe { dec_fdstat_byref(&mut vmctx, fdstat_ptr) } {
        Ok(host_fdstat) => host_fdstat,
        Err(e) => return enc_errno(e),
    };

    let errno = if let Some(ref fe) = ctx.fds.get(&host_fd) {
        host_fdstat.fs_filetype = fe.fd_object.ty;
        host_fdstat.fs_rights_base = fe.rights_base;
        host_fdstat.fs_rights_inheriting = fe.rights_inheriting;
        use nix::fcntl::{fcntl, OFlag, F_GETFL};
        match fcntl(fe.fd_object.rawfd, F_GETFL).map(OFlag::from_bits) {
            Ok(Some(flags)) => {
                if flags.contains(OFlag::O_APPEND) {
                    host_fdstat.fs_flags |= wasm32::__WASI_FDFLAG_APPEND;
                }
                if flags.contains(OFlag::O_DSYNC) {
                    host_fdstat.fs_flags |= wasm32::__WASI_FDFLAG_DSYNC;
                }
                if flags.contains(OFlag::O_NONBLOCK) {
                    host_fdstat.fs_flags |= wasm32::__WASI_FDFLAG_NONBLOCK;
                }
                if flags.contains(OFlag::O_RSYNC) {
                    host_fdstat.fs_flags |= wasm32::__WASI_FDFLAG_RSYNC;
                }
                if flags.contains(OFlag::O_SYNC) {
                    host_fdstat.fs_flags |= wasm32::__WASI_FDFLAG_SYNC;
                }
                wasm32::__WASI_ESUCCESS
            }
            Ok(None) => wasm32::__WASI_ENOSYS,
            Err(e) => wasm32::errno_from_nix(e.as_errno().unwrap()),
        }
    } else {
        return wasm32::__WASI_EBADF;
    };

    unsafe {
        enc_fdstat_byref(&mut vmctx, fdstat_ptr, host_fdstat)
            .expect("can write back into the pointer we read from");
    }

    errno
}

#[no_mangle]
pub extern "C" fn __wasi_fd_seek(
    vmctx: *mut lucet_vmctx,
    fd: wasm32::__wasi_fd_t,
    offset: wasm32::__wasi_filedelta_t,
    whence: wasm32::__wasi_whence_t,
    newoffset: wasm32::uintptr_t,
) -> wasm32::__wasi_errno_t {
    let mut vmctx = unsafe { Vmctx::from_raw(vmctx) };
    let ctx: &mut WasiCtx = vmctx.get_embed_ctx_mut();
    let fd = dec_fd(fd);
    let offset = dec_filedelta(offset);
    let whence = dec_whence(whence);

    let host_newoffset = {
        use nix::unistd::{lseek, Whence};
        let nwhence = match whence as u32 {
            host::__WASI_WHENCE_CUR => Whence::SeekCur,
            host::__WASI_WHENCE_END => Whence::SeekEnd,
            host::__WASI_WHENCE_SET => Whence::SeekSet,
            _ => return wasm32::__WASI_EINVAL,
        };

        let rights = if offset == 0 && whence as u32 == host::__WASI_WHENCE_CUR {
            host::__WASI_RIGHT_FD_TELL
        } else {
            host::__WASI_RIGHT_FD_SEEK | host::__WASI_RIGHT_FD_TELL
        };
        match ctx.get_fd_entry(fd, rights.into(), 0) {
            Ok(fe) => match lseek(fe.fd_object.rawfd, offset, nwhence) {
                Ok(newoffset) => newoffset,
                Err(e) => return wasm32::errno_from_nix(e.as_errno().unwrap()),
            },
            Err(e) => return enc_errno(e),
        }
    };

    unsafe {
        enc_filesize_byref(&mut vmctx, newoffset, host_newoffset as u64)
            .map(|_| wasm32::__WASI_ESUCCESS)
            .unwrap_or_else(|e| e)
    }
}

#[no_mangle]
pub extern "C" fn __wasi_fd_write(
    vmctx: *mut lucet_vmctx,
    fd: wasm32::__wasi_fd_t,
    iovs_ptr: wasm32::uintptr_t,
    iovs_len: wasm32::size_t,
    nwritten: wasm32::uintptr_t,
) -> wasm32::__wasi_errno_t {
    use nix::sys::uio::{writev, IoVec};

    let mut vmctx = unsafe { Vmctx::from_raw(vmctx) };
    let ctx: &mut WasiCtx = vmctx.get_embed_ctx_mut();
    let fd = dec_fd(fd);
    let iovs = match unsafe { dec_ciovec_slice(&mut vmctx, iovs_ptr, iovs_len) } {
        Ok(iovs) => iovs,
        Err(e) => return enc_errno(e),
    };

    let fe = match ctx.get_fd_entry(fd, host::__WASI_RIGHT_FD_WRITE.into(), 0) {
        Ok(fe) => fe,
        Err(e) => return enc_errno(e),
    };

    let iovs: Vec<IoVec<&[u8]>> = iovs
        .iter()
        .map(|iov| unsafe { host::ciovec_to_nix(iov) })
        .collect();

    let host_nwritten = match writev(fe.fd_object.rawfd, &iovs) {
        Ok(len) => len,
        Err(e) => return wasm32::errno_from_nix(e.as_errno().unwrap()),
    };

    unsafe {
        enc_usize_byref(&mut vmctx, nwritten, host_nwritten)
            .map(|_| wasm32::__WASI_ESUCCESS)
            .unwrap_or_else(|e| e)
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 The Gfx-rs Developers.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use syntax::{ast, ext};
use syntax::ext::build::AstBuilder;
use syntax::ext::deriving::generic;
use syntax::codemap;
use syntax::parse::token;
use syntax::ptr::P;
use syntax::ext::base::ItemDecorator;

#[derive(Copy, PartialEq, Debug)]
enum Param {
    Uniform,
    Block,
    Texture,
}

#[derive(Copy, Debug)]
enum ParamError {
    DeprecatedTexture,
}

/// Classify variable types (`i32`, `TextureParam`, etc) into the `Param`
fn classify(node: &ast::Ty_) -> Result<Param, ParamError> {
    match *node {
        ast::TyPath(_,ref path) => match path.segments.last() {
            Some(segment) => match segment.identifier.name.as_str() {
                "RawBufferHandle" => Ok(Param::Block),
                "TextureParam" => Ok(Param::Texture),
                "TextureHandle" => Err(ParamError::DeprecatedTexture),
                _ => Ok(Param::Uniform),
            },
            None => Ok(Param::Uniform),
        },
        _ => Ok(Param::Uniform),
    }
}

/// Generates the the method body for `gfx::shade::ParamValues::create_link`
fn method_create(cx: &mut ext::base::ExtCtxt, span: codemap::Span,
                 substr: &generic::Substructure,
                 link_name: &str,
                 path_root: ast::Ident) -> P<ast::Expr> {
    let link_ident = cx.ident_of(link_name);
    match *substr.fields {
        generic::StaticStruct(definition, generic::Named(ref fields)) => {
            let init_expr = cx.expr_struct_ident(
                span, link_ident,
                fields.iter().map(|&(fname, fspan)| {
                    cx.field_imm(fspan, fname, cx.expr_none(fspan))
                }).collect()
            );
            let class_info: Vec<(Param, P<ast::Expr>)> = definition.fields.iter()
                    .zip(fields.iter()).scan((), |_, (def, &(fname, fspan))|
                match classify(&def.node.ty.node) {
                    Ok(c) => {
                        let name = match super::find_name(cx, span, &def.node.attrs) {
                            Some(name) => name,
                            None => token::get_ident(fname),
                        };
                        Some((c, cx.expr_str(fspan, name)))
                    },
                    Err(e) => {
                        cx.span_err(fspan, &format!(
                            "Unrecognized parameter ({:?}) type {:?}",
                            fname.as_str(), e
                            )
                        );
                        None
                    },
                }
            ).collect();
            let gen_arms = |&: ptype: Param, var: ast::Ident| -> Vec<ast::Arm> {
                class_info.iter().zip(fields.iter())
                          .filter(|&(&(ref class, _), _)| *class == ptype)
                          .map(|(&(_, ref name_expr), &(fname, fspan))|
                    cx.arm(fspan,
                        vec![cx.pat_lit(fspan, name_expr.clone())],
                        quote_expr!(cx,
                            out.$fname = Some(i as $path_root::gfx::shade::$var)
                        )
                    )
                ).collect()
            };
            let uniform_arms = gen_arms(Param::Uniform, cx.ident_of("VarUniform"));
            let block_arms = gen_arms(Param::Block, cx.ident_of("VarBlock"));
            let texture_arms = gen_arms(Param::Texture, cx.ident_of("VarTexture"));
            let input = &substr.nonself_args[1];
            quote_expr!(cx, {
                let mut out = $init_expr;
                for (i, u) in $input.uniforms.iter().enumerate() {
                    let _ = i; // suppress warning about unused i
                    match &u.name[..] {
                        $uniform_arms
                        _ => return Err($path_root::gfx::shade::
                            ParameterError::MissingUniform(u.name.clone())),
                    }
                }
                for (i, b) in $input.blocks.iter().enumerate() {
                    let _ = i; // suppress warning about unused i
                    match &b.name[..] {
                        $block_arms
                        _ => return Err($path_root::gfx::shade::
                            ParameterError::MissingBlock(b.name.clone())),
                    }
                }
                for (i, t) in $input.textures.iter().enumerate() {
                    let _ = i; // suppress warning about unused i
                    match &t.name[..] {
                        $texture_arms
                        _ => return Err($path_root::gfx::shade::
                            ParameterError::MissingTexture(t.name.clone())),
                    }
                }
                Ok(out)
            })
        },
        _ => {
            cx.span_err(span, "Unable to implement `ShaderParam::create_link()` on a non-structure");
            cx.expr_tuple(span, vec![])
        },
    }
}

/// Generates the the method body for `gfx::shade::ParamValues::fill_params`
fn method_fill(cx: &mut ext::base::ExtCtxt, span: codemap::Span,
               substr: &generic::Substructure,
               definition: P<ast::StructDef>,
               path_root: ast::Ident)
               -> P<ast::Expr> {
    match *substr.fields {
        generic::Struct(ref fields) => {
            let out = &substr.nonself_args[1];
            let max_num = cx.expr_usize(span, fields.len());
            let mut calls = vec![
                cx.stmt_item(span, cx.item_use_simple(
                    span,
                    ast::Inherited,
                    cx.path(span, vec![
                        cx.ident_of("self"),
                        path_root,
                        cx.ident_of("gfx"),
                        cx.ident_of("shade"),
                        cx.ident_of("ToUniform"),
                    ])
                )),
                quote_stmt!(cx, $out.uniforms.reserve($max_num);),
                quote_stmt!(cx, $out.blocks.reserve($max_num);),
                quote_stmt!(cx, $out.textures.reserve($max_num);),
            ];
            calls.extend(definition.fields.iter().zip(fields.iter())
                                   .map(|(def, f)| {
                let value_id = &f.self_;
                let var_id = cx.expr_field_access(
                    span,
                    substr.nonself_args[0].clone(),
                    f.name.unwrap()
                    );
                match classify(&def.node.ty.node) {
                    Ok(Param::Uniform) => quote_stmt!(cx,
                        $var_id.map_or((), |id| {
                            if $out.uniforms.len() <= id as usize {
                                unsafe { $out.uniforms.set_len(id as usize + 1) }
                            }
                            *$out.uniforms.get_mut(id as usize).unwrap() = $value_id.to_uniform()
                        })
                    ),
                    Ok(Param::Block)   => quote_stmt!(cx,
                        $var_id.map_or((), |id| {
                            if $out.blocks.len() <= id as usize {
                                unsafe { $out.blocks.set_len(id as usize + 1) }
                            }
                            *$out.blocks.get_mut(id as usize).unwrap() = {$value_id}
                        })
                    ),
                    Ok(Param::Texture) => quote_stmt!(cx,
                        $var_id.map_or((), |id| {
                            if $out.textures.len() <= id as usize {
                                unsafe { $out.textures.set_len(id as usize + 1) }
                            }
                            *$out.textures.get_mut(id as usize).unwrap() = {$value_id}
                        })
                    ),
                    Err(_) => {
                        cx.span_err(span, &format!(
                            "Invalid uniform: {:?}",
                            f.name.unwrap().as_str(),
                            )
                        );
                        cx.stmt_expr(cx.expr_usize(span, 0))
                    },
                }
            }));
            cx.expr_block(cx.block_all(span, calls, None))
        },
        _ => {
            cx.span_err(span, "Unable to implement `ShaderParam::bind()` on a non-structure");
            cx.expr_tuple(span, vec![])
        }
    }
}

/// A helper function that translates variable type (`i32`, `TextureHandle`, etc)
/// into the corresponding shader var id type (`VarUniform`, `VarBlock`, or `VarTexture`)
fn node_to_var_type(cx: &mut ext::base::ExtCtxt,
                    span: codemap::Span, node: &ast::Ty_,
                    path_root: ast::Ident) -> P<ast::Ty> {
    let id = match classify(node) {
        Ok(Param::Uniform) => "VarUniform",
        Ok(Param::Block)   => "VarBlock",
        Ok(Param::Texture) => "VarTexture",
        Err(ParamError::DeprecatedTexture) => {
            cx.span_err(span, "Use gfx::shade::TextureParam for texture vars instead of gfx::shade::TextureHandle");
            ""
        },
    };
    cx.ty_option(cx.ty_path(
        cx.path(span, vec![
            path_root,
            cx.ident_of("gfx"),
            cx.ident_of("shade"),
            cx.ident_of(id),
        ]),
    ))
}

#[derive(Copy)]
pub struct ShaderParam;

impl ItemDecorator for ShaderParam {
    /// Decorator for `shader_param` attribute
    fn expand(&self, context: &mut ext::base::ExtCtxt, span: codemap::Span,
              meta_item: &ast::MetaItem, item: &ast::Item,
              push: &mut FnMut(P<ast::Item>)) {
        // Insert the `gfx` reexport module
        let path_root = super::extern_crate_hack(context, span, |i| (*push)(i));

        // constructing the Link struct
        let (base_def, link_def) = match item.node {
            ast::ItemStruct(ref definition, ref generics) => {
                if generics.lifetimes.len() > 0 {
                    context.bug("Lifetimes are not allowed in ShaderParam struct");
                }
                (definition, ast::StructDef {
                    fields: definition.fields.iter()
                        .map(|f| codemap::Spanned {
                            node: ast::StructField_ {
                                kind: f.node.kind,
                                id: f.node.id,
                                ty: node_to_var_type(context, f.span, &f.node.ty.node, path_root),
                                attrs: Vec::new(),
                            },
                            span: f.span,
                        }).collect(),
                    ctor_id: None,
                })
            },
            _ => {
                context.span_err(span, "Only free-standing named structs allowed to derive ShaderParam");
                return;
            }
        };
        let link_name = format!("_{}Link", item.ident.as_str());
        let link_ident = context.ident_of(&link_name);
        let link_ty = generic::ty::Literal(
            generic::ty::Path::new_local(&link_name)
        );
        let link_item = context.item_struct(span, link_ident, link_def)
                               .map(|mut item| {
            item.attrs.push(context.attribute(span,
                context.meta_list(span, token::InternedString::new("derive"), vec![
                        context.meta_word(span, token::InternedString::new("Copy")),
                        context.meta_word(span, token::InternedString::new("Debug")),
                ])
            ));
            item
        });
        (*push)(link_item);

        // process meta parameters (expecting none)
        let resources = match meta_item.node {
            ast::MetaList(_, ref list) if list.len() == 1 => {
                use std::old_path::BytesContainer;
                match list[0].node {
                    ast::MetaWord(ref r) => r.container_as_str().unwrap(),
                    _ => {
                        context.span_err(list[0].span, "Invalid parameter");
                        ""
                    }
                }
            },
            _ => {
                context.span_err(meta_item.span,
                    "Use as #[shader_param(<ResourcesType>)]");
                ""
            }
        };

        // #[derive ShaderParam]
        let trait_def = generic::TraitDef {
            span: span,
            attributes: Vec::new(),
            path: generic::ty::Path::new(
                vec![super::EXTERN_CRATE_HACK, "gfx", "shade", "ShaderParam"],
            ),
            additional_bounds: Vec::new(),
            generics: generic::ty::LifetimeBounds::empty(), //TODO
            methods: vec![
                generic::MethodDef {
                    name: "create_link",
                    generics: generic::ty::LifetimeBounds::empty(),
                    explicit_self: None,
                    args: vec![
                        generic::ty::Literal(generic::ty::Path {
                            path: vec!["Option"],
                            lifetime: None,
                            params: vec![
                                box generic::ty::Ptr(
                                    box generic::ty::Self_,
                                    generic::ty::Borrowed(None, ast::MutImmutable)
                                ),
                            ],
                            global: false,
                        }),
                        generic::ty::Ptr(
                            box generic::ty::Literal(generic::ty::Path::new(
                                vec![super::EXTERN_CRATE_HACK, "gfx", "ProgramInfo"])),
                            generic::ty::Borrowed(None, ast::MutImmutable)
                        ),
                    ],
                    ret_ty: generic::ty::Literal(
                        generic::ty::Path {
                            path: vec!["Result"],
                            lifetime: None,
                            params: vec![
                                box link_ty.clone(),
                                box generic::ty::Literal(generic::ty::Path::new(
                                    vec![super::EXTERN_CRATE_HACK, "gfx", "shade", "ParameterError"]
                                ))
                            ],
                            global: false,
                        },
                    ),
                    attributes: Vec::new(),
                    combine_substructure: generic::combine_substructure(box |cx, span, sub|
                        method_create(cx, span, sub, &link_name, path_root)
                    ),
                },
                generic::MethodDef {
                    name: "fill_params",
                    generics: generic::ty::LifetimeBounds::empty(),
                    explicit_self: Some(Some(generic::ty::Borrowed(
                        None, ast::MutImmutable
                    ))),
                    args: vec![
                        generic::ty::Ptr(
                            box link_ty.clone(),
                            generic::ty::Borrowed(None, ast::MutImmutable)
                        ),
                        generic::ty::Literal(
                            generic::ty::Path {
                                path: vec![super::EXTERN_CRATE_HACK, "gfx", "shade", "ParamValues"],
                                lifetime: None,
                                params: vec![box generic::ty::Literal(
                                    generic::ty::Path::new_local(resources)
                                )],
                                global: false,
                            },
                        ),
                    ],
                    ret_ty: generic::ty::Tuple(Vec::new()),
                    attributes: Vec::new(),
                    combine_substructure: generic::combine_substructure(box |cx, span, sub|
                        method_fill(cx, span, sub, base_def.clone(), path_root)
                    ),
                },
            ],
            associated_types: vec![
                (context.ident_of("Resources"), generic::ty::Literal(
                    generic::ty::Path::new_local(resources)
                )),
                (context.ident_of("Link"), link_ty),
            ],
        };
        let fixup = |item| {
            (*push)(super::fixup_extern_crate_paths(item, path_root))
        };
        trait_def.expand(context, meta_item, item, fixup);
    }
}
<commit_msg>Rewrote #[shader_param] expansion using pure AST<commit_after>// Copyright 2015 The Gfx-rs Developers.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use syntax::{abi, ast, codemap, ext};
use syntax::ext::base::ItemDecorator;
use syntax::ext::build::AstBuilder;
use syntax::owned_slice::OwnedSlice;
use syntax::parse::token;
use syntax::ptr::P;

#[derive(Copy, PartialEq, Debug)]
enum Param {
    Uniform,
    Block,
    Texture,
    Special,
}

#[derive(Copy, Debug)]
enum ParamError {
    DeprecatedTexture,
}

/// Classify variable types (`i32`, `TextureParam`, etc) into the `Param`
fn classify(node: &ast::Ty_) -> Result<Param, ParamError> {
    match *node {
        ast::TyPath(_,ref path) => match path.segments.last() {
            Some(segment) => match segment.identifier.name.as_str() {
                "RawBufferHandle" => Ok(Param::Block),
                "TextureParam" => Ok(Param::Texture),
                "TextureHandle" => Err(ParamError::DeprecatedTexture),
                "PhantomData"   => Ok(Param::Special),
                _ => Ok(Param::Uniform),
            },
            None => Ok(Param::Uniform),
        },
        _ => Ok(Param::Uniform),
    }
}

/// Generates the the method body for `gfx::shade::ParamValues::create_link`
fn method_create(cx: &mut ext::base::ExtCtxt,
                 span: codemap::Span,
                 definition: &ast::StructDef,
                 input: ast::Ident,
                 link_ident: ast::Ident,
                 path_root: ast::Ident)
                 -> P<ast::Block> {
    let init_expr = cx.expr_struct_ident(
        span, link_ident,
        definition.fields.iter().scan((), |_, field| {
            field.node.ident().map(|name|
                cx.field_imm(field.span, name, cx.expr_none(field.span))
            )
        }).collect()
    );
    let class_info: Vec<(Param, P<ast::Expr>)> = definition.fields.iter().scan((), |_, field|
        match (field.node.ident(), classify(&field.node.ty.node)) {
            (None, _) => {
                cx.span_err(field.span, "Named fields are required for `ShaderParam`");
                None
            },
            (Some(fname), Ok(c)) => {
                let name = match super::find_name(cx, field.span, &field.node.attrs) {
                    Some(name) => name,
                    None => token::get_ident(fname),
                };
                Some((c, cx.expr_str(field.span, name)))
            },
            (Some(fname), Err(e)) => {
                cx.span_err(field.span, &format!(
                    "Unrecognized parameter ({:?}) type {:?}",
                    fname.as_str(), e
                ));
                None
            },
        }
    ).collect();
    let gen_arms = |&: ptype: Param, var: ast::Ident| -> Vec<ast::Arm> {
        class_info.iter().zip(definition.fields.iter())
                  .filter(|&(&(ref class, _), _)| *class == ptype)
                  .scan((), |_, (&(_, ref name_expr), field)|
            field.node.ident().map(|name| quote_arm!(cx,
                $name_expr => {out.$name = Some(i as $path_root::gfx::shade::$var)}
            ))
        ).collect()
    };
    let uniform_arms = gen_arms(Param::Uniform, cx.ident_of("VarUniform"));
    let block_arms = gen_arms(Param::Block, cx.ident_of("VarBlock"));
    let texture_arms = gen_arms(Param::Texture, cx.ident_of("VarTexture"));
    let expr = quote_expr!(cx, {
        let mut out = $init_expr;
        for (i, u) in $input.uniforms.iter().enumerate() {
            let _ = i; // suppress warning about unused i
            match &u.name[..] {
                $uniform_arms
                _ => return Err($path_root::gfx::shade::
                    ParameterError::MissingUniform(u.name.clone())),
            }
        }
        for (i, b) in $input.blocks.iter().enumerate() {
            let _ = i; // suppress warning about unused i
            match &b.name[..] {
                $block_arms
                _ => return Err($path_root::gfx::shade::
                    ParameterError::MissingBlock(b.name.clone())),
            }
        }
        for (i, t) in $input.textures.iter().enumerate() {
            let _ = i; // suppress warning about unused i
            match &t.name[..] {
                $texture_arms
                _ => return Err($path_root::gfx::shade::
                    ParameterError::MissingTexture(t.name.clone())),
            }
        }
        Ok(out)
    });
    cx.block_expr(expr)
}

/// Generates the the method body for `gfx::shade::ParamValues::fill_params`
fn method_fill(cx: &mut ext::base::ExtCtxt,
               span: codemap::Span,
               definition: &ast::StructDef,
               path_root: ast::Ident)
               -> P<ast::Block> {
    let max_num = cx.expr_usize(span, definition.fields.len());
    let mut calls = vec![
        quote_stmt!(cx, use self::$path_root::gfx::shade::ToUniform;),
        quote_stmt!(cx, out.uniforms.reserve($max_num);),
        quote_stmt!(cx, out.blocks.reserve($max_num);),
        quote_stmt!(cx, out.textures.reserve($max_num);),
    ];
    calls.extend(definition.fields.iter().scan((), |_, field| {
        let name = match field.node.ident() {
            Some(n) => n,
            None => {
                cx.span_err(span, "Named fields are required for `ShaderParam`");
                return None
            }
        };
        classify(&field.node.ty.node).ok().map(|param| match param {
            Param::Uniform => quote_stmt!(cx,
                link.$name.map_or((), |id| {
                    if out.uniforms.len() <= id as usize {
                        unsafe { out.uniforms.set_len(id as usize + 1) }
                    }
                    *out.uniforms.get_mut(id as usize).unwrap() = self.$name.to_uniform()
                })
            ),
            Param::Block   => quote_stmt!(cx,
                link.$name.map_or((), |id| {
                    if out.blocks.len() <= id as usize {
                        unsafe { out.blocks.set_len(id as usize + 1) }
                    }
                    *out.blocks.get_mut(id as usize).unwrap() = {self.$name}
                })
            ),
            Param::Texture => quote_stmt!(cx,
                link.$name.map_or((), |id| {
                    if out.textures.len() <= id as usize {
                        unsafe { out.textures.set_len(id as usize + 1) }
                    }
                    *out.textures.get_mut(id as usize).unwrap() = {self.$name}
                })
            ),
            Param::Special => quote_stmt!(cx, ()),
        })
    }));
    cx.block_all(span, calls, None)
}

/// A helper function that translates variable type (`i32`, `TextureHandle`, etc)
/// into the corresponding shader var id type (`VarUniform`, `VarBlock`, or `VarTexture`)
fn node_to_var_type(cx: &mut ext::base::ExtCtxt,
                    span: codemap::Span, node: &ast::Ty_,
                    path_root: ast::Ident) -> P<ast::Ty> {
    let id = cx.ident_of(match classify(node) {
        Ok(Param::Uniform) => "VarUniform",
        Ok(Param::Block)   => "VarBlock",
        Ok(Param::Texture) => "VarTexture",
        Ok(Param::Special) => return quote_ty!(cx, Option<()>),
        Err(ParamError::DeprecatedTexture) => {
            cx.span_err(span, "Use gfx::shade::TextureParam for texture vars instead of gfx::shade::TextureHandle");
            ""
        },
    });
    quote_ty!(cx, Option<$path_root::gfx::shade::$id>)
}

fn impl_type(cx: &mut ext::base::ExtCtxt, span: codemap::Span, name: &str, type_ident: ast::Ident) -> ast::ImplItem {
    ast::TypeImplItem(P(ast::Typedef {
        id: ast::DUMMY_NODE_ID,
        span: span,
        ident: cx.ident_of(name),
        vis: ast::Visibility::Inherited,
        attrs: Vec::new(),
        typ: cx.ty_ident(span, type_ident),
    }))
}

fn impl_method(cx: &mut ext::base::ExtCtxt, span: codemap::Span, name: &str,
               with_self: bool, declaration: P<ast::FnDecl>, body: P<ast::Block>)
               -> ast::ImplItem {
    ast::MethodImplItem(P(ast::Method {
        attrs: Vec::new(),
        id: ast::DUMMY_NODE_ID,
        span: span,
        node: ast::MethDecl(
            cx.ident_of(name),
            ast::Generics {
                lifetimes: Vec::new(),
                ty_params: OwnedSlice::empty(),
                where_clause: ast::WhereClause {
                    id: ast::DUMMY_NODE_ID,
                    predicates: Vec::new(),
                },
            },
            abi::Abi::Rust,
            codemap::Spanned {
                node: if with_self {
                        ast::SelfRegion(None, ast::MutImmutable, cx.ident_of("self"))
                    }else {
                        ast::SelfStatic
                },
                span: span,
            },
            ast::Unsafety::Normal,
            declaration,
            body,
            ast::Visibility::Inherited
        )
    }))
}

#[derive(Copy)]
pub struct ShaderParam;
impl ItemDecorator for ShaderParam {
    /// Decorator for `shader_param` attribute
    fn expand(&self, context: &mut ext::base::ExtCtxt, span: codemap::Span,
              meta_item: &ast::MetaItem, item: &ast::Item,
              push: &mut FnMut(P<ast::Item>)) {
        // Insert the `gfx` reexport module
        let extern_hack = context.ident_of(super::EXTERN_CRATE_HACK);
        let path_root = super::extern_crate_hack(context, span, |i| (*push)(i));

        // constructing the Link struct
        let (base_def, generics, link_def) = match item.node {
            ast::ItemStruct(ref definition, ref generics) => {
                (definition, generics.clone(), ast::StructDef {
                    fields: definition.fields.iter()
                        .map(|f| codemap::Spanned {
                            node: ast::StructField_ {
                                kind: f.node.kind,
                                id: f.node.id,
                                ty: node_to_var_type(context, f.span, &f.node.ty.node, path_root),
                                attrs: Vec::new(),
                            },
                            span: f.span,
                        }).collect(),
                    ctor_id: None,
                })
            },
            _ => {
                context.span_err(span, "Only free-standing named structs allowed to derive ShaderParam");
                return;
            }
        };

        // derive and push
        let link_name = format!("_{}Link", item.ident.as_str());
        let link_ident = context.ident_of(&link_name);
        let link_item = context.item_struct(span, link_ident, link_def)
                               .map(|mut item| {
            item.attrs.push(context.attribute(span,
                context.meta_list(span, token::InternedString::new("derive"), vec![
                        context.meta_word(span, token::InternedString::new("Copy")),
                        context.meta_word(span, token::InternedString::new("Debug")),
                ])
            ));
            item
        });
        (*push)(link_item);

        // process meta parameters (expecting none)
        match meta_item.node {
            ast::MetaWord(_) => (), //expected
            _ => {
                context.span_err(meta_item.span, "#[shader_param] needs no param");
            }
        }

        // find the generic `Resources` bound
        let resource_ident = match generics.ty_params.iter().find(|typ|
            typ.bounds.iter().find(|b| match **b {
                ast::TraitTyParamBound(ref poly_trait, _) =>
                    poly_trait.trait_ref.path.segments.last().unwrap().identifier.as_str() == "Resources",
                ast::RegionTyParamBound(_) => false,
            }).is_some()
        ){
            Some(typ) => typ.ident,
            None => {
                context.span_err(meta_item.span, "#[shader_param] unable to find generic `gfx::Resources` bound");
                context.ident_of("R")
            }
        };

        // construct `create_link()`
        let lifetimes = generics.lifetimes.iter().map(|ld| ld.lifetime).collect();
        let generic_parameters = generics.ty_params.iter().map(|ty|
            context.ty_ident(span, ty.ident)
        ).collect();
        let struct_ty = context.ty_path(context.path_all(
            span, false,
            vec![item.ident],
            lifetimes,
            generic_parameters,
            Vec::new(),
        ));
        let create_param = context.ident_of("params");
        let body_create = method_create(context, span, base_def, create_param, link_ident, path_root);
        let decl_create = context.fn_decl(
            vec![
                ast::Arg {
                    ty: quote_ty!(context, Option<&$struct_ty>),
                    pat: context.pat_wild(span),
                    id: ast::DUMMY_NODE_ID,
                },
                ast::Arg {
                    ty: quote_ty!(context, &$extern_hack::gfx::ProgramInfo),
                    pat: context.pat_ident(span, create_param),
                    id: ast::DUMMY_NODE_ID,
                },
            ],
            context.ty_path(context.path_all(
                span, false, vec![context.ident_of("Result")],
                Vec::new(), vec![
                    context.ty_ident(span, link_ident),
                    quote_ty!(context, $extern_hack::gfx::shade::ParameterError),
                ], Vec::new()
            )),
        );

        // construct `fill_params()`
        let body_fill = method_fill(context, span, base_def, path_root);
        let decl_fill = context.fn_decl(
             vec![
                ast::Arg::new_self(span, ast::MutImmutable, context.ident_of("self")),
                ast::Arg {
                    ty: quote_ty!(context, &$link_ident),
                    pat: context.pat_ident(span, context.ident_of("link")),
                    id: ast::DUMMY_NODE_ID,
                },
                ast::Arg {
                    ty: quote_ty!(context, $extern_hack::gfx::shade::ParamValues<$resource_ident>),
                    pat: context.pat_ident(span, context.ident_of("out")),
                    id: ast::DUMMY_NODE_ID,
                },
            ],
            quote_ty!(context, ())
        );

        // construct implementations for types and methods
        let impls = vec![
            impl_type(context, span, "Resources", resource_ident),
            impl_type(context, span, "Link", link_ident),
            impl_method(context, span, "create_link", false, decl_create, body_create),
            impl_method(context, span, "fill_params", true, decl_fill, body_fill),
        ];

        // final implementation item
        let item = context.item(span, item.ident, Vec::new(), ast::Item_::ItemImpl(
                ast::Unsafety::Normal,
                ast::ImplPolarity::Positive,
                generics,
                Some(context.trait_ref(context.path(span, vec![
                    extern_hack,
                    context.ident_of("gfx"),
                    context.ident_of("shade"),
                    context.ident_of("ShaderParam"),
                ]))),
                struct_ty,
                impls
        ));
        (*push)(super::fixup_extern_crate_paths(item, path_root));
    }
}
<|endoftext|>


<commit_before>#![feature(pin, arbitrary_self_types, futures_api)]

use futures::channel::oneshot;
use futures::executor::{block_on, LocalPool};
use futures::future::{self, FutureExt, LocalFutureObj};
use futures::task::LocalSpawn;
use std::cell::{Cell, RefCell};
use std::rc::Rc;
use std::thread;

fn send_shared_oneshot_and_wait_on_multiple_threads(threads_number: u32) {
    let (tx, rx) = oneshot::channel::<i32>();
    let f = rx.shared();
    let join_handles = (0..threads_number)
        .map(|_| {
            let cloned_future = f.clone();
            thread::spawn(move || {
                assert_eq!(block_on(cloned_future).unwrap(), 6);
            })
        })
        .collect::<Vec<_>>();

    tx.send(6).unwrap();

    assert_eq!(block_on(f).unwrap(), 6);
    for join_handle in join_handles {
        join_handle.join().unwrap();
    }
}

#[test]
fn one_thread() {
    send_shared_oneshot_and_wait_on_multiple_threads(1);
}

#[test]
fn two_threads() {
    send_shared_oneshot_and_wait_on_multiple_threads(2);
}

#[test]
fn many_threads() {
    send_shared_oneshot_and_wait_on_multiple_threads(1000);
}

/* ToDo: This requires FutureExt::select to be implemented
#[test]
fn drop_on_one_task_ok() {
    let (tx, rx) = oneshot::channel::<u32>();
    let f1 = rx.shared();
    let f2 = f1.clone();

    let (tx2, rx2) = oneshot::channel::<u32>();

    let t1 = thread::spawn(|| {
        let f = f1.map_err(|_| ()).map(|x| *x).select(rx2.map_err(|_| ()));
        drop(block_on(f));
    });

    let (tx3, rx3) = oneshot::channel::<u32>();

    let t2 = thread::spawn(|| {
        let _ = block_on(f2.map(|x| tx3.send(*x).unwrap()).map_err(|_| ()));
    });

    tx2.send(11).unwrap(); // cancel `f1`
    t1.join().unwrap();

    tx.send(42).unwrap(); // Should cause `f2` and then `rx3` to get resolved.
    let result = block_on(rx3).unwrap();
    assert_eq!(result, 42);
    t2.join().unwrap();
}*/

#[test]
fn drop_in_poll() {
    let slot1 = Rc::new(RefCell::new(None));
    let slot2 = slot1.clone();

    let future1 = future::lazy(move |_| {
        slot2.replace(None); // Drop future
        1
    }).shared();

    let future2 = LocalFutureObj::new(Box::new(future1.clone()));
    slot1.replace(Some(future2));

    assert_eq!(block_on(future1), 1);
}

#[test]
fn peek() {
    let mut local_pool = LocalPool::new();
    let spawn = &mut local_pool.spawner();

    let (tx0, rx0) = oneshot::channel::<i32>();
    let f1 = rx0.shared();
    let f2 = f1.clone();

    // Repeated calls on the original or clone do not change the outcome.
    for _ in 0..2 {
        assert!(f1.peek().is_none());
        assert!(f2.peek().is_none());
    }

    // Completing the underlying future has no effect, because the value has not been `poll`ed in.
    tx0.send(42).unwrap();
    for _ in 0..2 {
        assert!(f1.peek().is_none());
        assert!(f2.peek().is_none());
    }

    // Once the Shared has been polled, the value is peekable on the clone.
    spawn.spawn_local_obj(LocalFutureObj::new(Box::new(f1.map(|_| ())))).unwrap();
    local_pool.run();
    for _ in 0..2 {
        assert_eq!(f2.peek().unwrap(), Ok(42));
    }
}

struct CountClone(Rc<Cell<i32>>);

impl Clone for CountClone {
    fn clone(&self) -> Self {
        self.0.set(self.0.get() + 1);
        CountClone(self.0.clone())
    }
}

#[test]
fn dont_clone_in_single_owner_shared_future() {
    let counter = CountClone(Rc::new(Cell::new(0)));
    let (tx, rx) = oneshot::channel();

    let rx = rx.shared();

    tx.send(counter).ok().unwrap();

    assert_eq!(block_on(rx).unwrap().0.get(), 0);
}

#[test]
fn dont_do_unnecessary_clones_on_output() {
    let counter = CountClone(Rc::new(Cell::new(0)));
    let (tx, rx) = oneshot::channel();

    let rx = rx.shared();

    tx.send(counter).ok().unwrap();

    assert_eq!(block_on(rx.clone()).unwrap().0.get(), 1);
    assert_eq!(block_on(rx.clone()).unwrap().0.get(), 2);
    assert_eq!(block_on(rx).unwrap().0.get(), 2);
}
<commit_msg>Fix peek test<commit_after>#![feature(pin, arbitrary_self_types, futures_api)]

use futures::channel::oneshot;
use futures::executor::{block_on, LocalPool};
use futures::future::{self, FutureExt, LocalFutureObj};
use futures::task::LocalSpawn;
use std::cell::{Cell, RefCell};
use std::rc::Rc;
use std::thread;

fn send_shared_oneshot_and_wait_on_multiple_threads(threads_number: u32) {
    let (tx, rx) = oneshot::channel::<i32>();
    let f = rx.shared();
    let join_handles = (0..threads_number)
        .map(|_| {
            let cloned_future = f.clone();
            thread::spawn(move || {
                assert_eq!(block_on(cloned_future).unwrap(), 6);
            })
        })
        .collect::<Vec<_>>();

    tx.send(6).unwrap();

    assert_eq!(block_on(f).unwrap(), 6);
    for join_handle in join_handles {
        join_handle.join().unwrap();
    }
}

#[test]
fn one_thread() {
    send_shared_oneshot_and_wait_on_multiple_threads(1);
}

#[test]
fn two_threads() {
    send_shared_oneshot_and_wait_on_multiple_threads(2);
}

#[test]
fn many_threads() {
    send_shared_oneshot_and_wait_on_multiple_threads(1000);
}

/* ToDo: This requires FutureExt::select to be implemented
#[test]
fn drop_on_one_task_ok() {
    let (tx, rx) = oneshot::channel::<u32>();
    let f1 = rx.shared();
    let f2 = f1.clone();

    let (tx2, rx2) = oneshot::channel::<u32>();

    let t1 = thread::spawn(|| {
        let f = f1.map_err(|_| ()).map(|x| *x).select(rx2.map_err(|_| ()));
        drop(block_on(f));
    });

    let (tx3, rx3) = oneshot::channel::<u32>();

    let t2 = thread::spawn(|| {
        let _ = block_on(f2.map(|x| tx3.send(*x).unwrap()).map_err(|_| ()));
    });

    tx2.send(11).unwrap(); // cancel `f1`
    t1.join().unwrap();

    tx.send(42).unwrap(); // Should cause `f2` and then `rx3` to get resolved.
    let result = block_on(rx3).unwrap();
    assert_eq!(result, 42);
    t2.join().unwrap();
}*/

#[test]
fn drop_in_poll() {
    let slot1 = Rc::new(RefCell::new(None));
    let slot2 = slot1.clone();

    let future1 = future::lazy(move |_| {
        slot2.replace(None); // Drop future
        1
    }).shared();

    let future2 = LocalFutureObj::new(Box::new(future1.clone()));
    slot1.replace(Some(future2));

    assert_eq!(block_on(future1), 1);
}

#[test]
fn peek() {
    let mut local_pool = LocalPool::new();
    let spawn = &mut local_pool.spawner();

    let (tx0, rx0) = oneshot::channel::<i32>();
    let f1 = rx0.shared();
    let f2 = f1.clone();

    // Repeated calls on the original or clone do not change the outcome.
    for _ in 0..2 {
        assert!(f1.peek().is_none());
        assert!(f2.peek().is_none());
    }

    // Completing the underlying future has no effect, because the value has not been `poll`ed in.
    tx0.send(42).unwrap();
    for _ in 0..2 {
        assert!(f1.peek().is_none());
        assert!(f2.peek().is_none());
    }

    // Once the Shared has been polled, the value is peekable on the clone.
    spawn.spawn_local_obj(LocalFutureObj::new(Box::new(f1.map(|_| ())))).unwrap();
    local_pool.run();
    for _ in 0..2 {
        assert_eq!(*f2.peek().unwrap(), Ok(42));
    }
}

struct CountClone(Rc<Cell<i32>>);

impl Clone for CountClone {
    fn clone(&self) -> Self {
        self.0.set(self.0.get() + 1);
        CountClone(self.0.clone())
    }
}

#[test]
fn dont_clone_in_single_owner_shared_future() {
    let counter = CountClone(Rc::new(Cell::new(0)));
    let (tx, rx) = oneshot::channel();

    let rx = rx.shared();

    tx.send(counter).ok().unwrap();

    assert_eq!(block_on(rx).unwrap().0.get(), 0);
}

#[test]
fn dont_do_unnecessary_clones_on_output() {
    let counter = CountClone(Rc::new(Cell::new(0)));
    let (tx, rx) = oneshot::channel();

    let rx = rx.shared();

    tx.send(counter).ok().unwrap();

    assert_eq!(block_on(rx.clone()).unwrap().0.get(), 1);
    assert_eq!(block_on(rx.clone()).unwrap().0.get(), 2);
    assert_eq!(block_on(rx).unwrap().0.get(), 2);
}
<|endoftext|>