<commit_before>//! Shadowsocks manager connecting interface

#[cfg(unix)]
use std::io::ErrorKind;
use std::{fmt, io, net::SocketAddr};

use tokio::net::UdpSocket;
#[cfg(unix)]
use tokio::net::{unix::SocketAddr as UnixSocketAddr, UnixDatagram};

use crate::{
    config::ManagerAddr,
    context::Context,
    net::{ConnectOpts, UdpSocket as ShadowUdpSocket},
};

/// Address accepted from Manager
#[derive(Debug)]
pub enum ManagerSocketAddr {
    SocketAddr(SocketAddr),
    #[cfg(unix)]
    UnixSocketAddr(UnixSocketAddr),
}

impl ManagerSocketAddr {
    /// Check if it is unnamed (not binded to any valid address), only valid for `UnixSocketAddr`
    pub fn is_unnamed(&self) -> bool {
        match *self {
            ManagerSocketAddr::SocketAddr(..) => false,
            #[cfg(unix)]
            ManagerSocketAddr::UnixSocketAddr(ref s) => s.is_unnamed(),
        }
    }
}

impl fmt::Display for ManagerSocketAddr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            ManagerSocketAddr::SocketAddr(ref saddr) => fmt::Display::fmt(saddr, f),
            #[cfg(unix)]
            ManagerSocketAddr::UnixSocketAddr(ref saddr) => fmt::Debug::fmt(saddr, f),
        }
    }
}

/// Datagram socket for manager
///
/// For *nix system, this is a wrapper for both UDP socket and Unix socket
pub enum ManagerDatagram {
    UdpDatagram(UdpSocket),
    #[cfg(unix)]
    UnixDatagram(UnixDatagram),
}

impl ManagerDatagram {
    /// Create a `ManagerDatagram` binding to requested `bind_addr`
    pub async fn bind(context: &Context, bind_addr: &ManagerAddr) -> io::Result<ManagerDatagram> {
        match *bind_addr {
            ManagerAddr::SocketAddr(ref saddr) => Ok(ManagerDatagram::UdpDatagram(
                ShadowUdpSocket::listen(saddr).await?.into(),
            )),
            ManagerAddr::DomainName(ref dname, port) => {
                let (_, socket) =
                    lookup_then!(context, dname, port, |saddr| { ShadowUdpSocket::listen(&saddr).await })?;

                Ok(ManagerDatagram::UdpDatagram(socket.into()))
            }
            #[cfg(unix)]
            ManagerAddr::UnixSocketAddr(ref path) => {
                use std::fs;

                // Remove it first incase it is already exists
                let _ = fs::remove_file(path);

                Ok(ManagerDatagram::UnixDatagram(UnixDatagram::bind(path)?))
            }
        }
    }

    /// Create a `ManagerDatagram` for sending data to manager
    pub async fn connect(
        context: &Context,
        bind_addr: &ManagerAddr,
        connect_opts: &ConnectOpts,
    ) -> io::Result<ManagerDatagram> {
        match *bind_addr {
            ManagerAddr::SocketAddr(sa) => ManagerDatagram::connect_socket_addr(sa, connect_opts).await,

            ManagerAddr::DomainName(ref dname, port) => {
                // Try connect to all socket addresses
                lookup_then!(context, dname, port, |addr| {
                    ManagerDatagram::connect_socket_addr(addr, connect_opts).await
                })
                .map(|(_, d)| d)
            }

            #[cfg(unix)]
            // For unix socket, it doesn't need to bind to any valid address
            // Because manager won't response to you
            ManagerAddr::UnixSocketAddr(..) => Ok(ManagerDatagram::UnixDatagram(UnixDatagram::unbound()?)),
        }
    }

    async fn connect_socket_addr(sa: SocketAddr, connect_opts: &ConnectOpts) -> io::Result<ManagerDatagram> {
        let socket = ShadowUdpSocket::connect_with_opts(&sa, connect_opts).await?;
        Ok(ManagerDatagram::UdpDatagram(socket.into()))
    }

    /// Receives data from the socket.
    pub async fn recv(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match *self {
            ManagerDatagram::UdpDatagram(ref mut udp) => udp.recv(buf).await,
            #[cfg(unix)]
            ManagerDatagram::UnixDatagram(ref mut unix) => unix.recv(buf).await,
        }
    }

    /// Receives data from the socket.
    pub async fn recv_from(&mut self, buf: &mut [u8]) -> io::Result<(usize, ManagerSocketAddr)> {
        match *self {
            ManagerDatagram::UdpDatagram(ref mut udp) => {
                let (s, addr) = udp.recv_from(buf).await?;
                Ok((s, ManagerSocketAddr::SocketAddr(addr)))
            }
            #[cfg(unix)]
            ManagerDatagram::UnixDatagram(ref mut unix) => {
                let (s, addr) = unix.recv_from(buf).await?;
                Ok((s, ManagerSocketAddr::UnixSocketAddr(addr)))
            }
        }
    }

    /// Sends data to the socket
    pub async fn send(&mut self, buf: &[u8]) -> io::Result<usize> {
        match *self {
            ManagerDatagram::UdpDatagram(ref mut udp) => udp.send(buf).await,
            #[cfg(unix)]
            ManagerDatagram::UnixDatagram(ref mut unix) => unix.send(buf).await,
        }
    }

    /// Sends data to the socket to the specified address.
    pub async fn send_to(&mut self, buf: &[u8], target: &ManagerSocketAddr) -> io::Result<usize> {
        match *self {
            ManagerDatagram::UdpDatagram(ref mut udp) => match *target {
                ManagerSocketAddr::SocketAddr(ref saddr) => udp.send_to(buf, saddr).await,
                #[cfg(unix)]
                ManagerSocketAddr::UnixSocketAddr(..) => {
                    let err = io::Error::new(ErrorKind::InvalidInput, "udp datagram requires IP address target");
                    Err(err)
                }
            },
            #[cfg(unix)]
            ManagerDatagram::UnixDatagram(ref mut unix) => match *target {
                ManagerSocketAddr::UnixSocketAddr(ref saddr) => match saddr.as_pathname() {
                    Some(paddr) => unix.send_to(buf, paddr).await,
                    None => {
                        let err = io::Error::new(ErrorKind::InvalidInput, "target address must not be unnamed");
                        Err(err)
                    }
                },
                ManagerSocketAddr::SocketAddr(..) => {
                    let err = io::Error::new(ErrorKind::InvalidInput, "unix datagram requires path address target");
                    Err(err)
                }
            },
        }
    }

    /// Sends data on the socket to the specified manager address
    pub async fn send_to_manager(&mut self, buf: &[u8], context: &Context, target: &ManagerAddr) -> io::Result<usize> {
        match *self {
            ManagerDatagram::UdpDatagram(ref mut udp) => match *target {
                ManagerAddr::SocketAddr(ref saddr) => udp.send_to(buf, saddr).await,
                ManagerAddr::DomainName(ref dname, port) => {
                    let (_, n) = lookup_then!(context, dname, port, |saddr| { udp.send_to(buf, saddr).await })?;
                    Ok(n)
                }
                #[cfg(unix)]
                ManagerAddr::UnixSocketAddr(..) => {
                    let err = io::Error::new(ErrorKind::InvalidInput, "udp datagram requires IP address target");
                    Err(err)
                }
            },
            #[cfg(unix)]
            ManagerDatagram::UnixDatagram(ref mut unix) => match *target {
                ManagerAddr::UnixSocketAddr(ref paddr) => unix.send_to(buf, paddr).await,
                ManagerAddr::SocketAddr(..) | ManagerAddr::DomainName(..) => {
                    let err = io::Error::new(ErrorKind::InvalidInput, "unix datagram requires path address target");
                    Err(err)
                }
            },
        }
    }

    /// Returns the local address that this socket is bound to.
    pub fn local_addr(&self) -> io::Result<ManagerSocketAddr> {
        match *self {
            ManagerDatagram::UdpDatagram(ref socket) => socket.local_addr().map(ManagerSocketAddr::SocketAddr),
            #[cfg(unix)]
            ManagerDatagram::UnixDatagram(ref dgram) => dgram.local_addr().map(ManagerSocketAddr::UnixSocketAddr),
        }
    }
}
<commit_msg>fixed ssserver reporting to manager listening on UDS<commit_after>//! Shadowsocks manager connecting interface

#[cfg(unix)]
use std::io::ErrorKind;
use std::{fmt, io, net::SocketAddr};

use tokio::net::UdpSocket;
#[cfg(unix)]
use tokio::net::{unix::SocketAddr as UnixSocketAddr, UnixDatagram};

use crate::{
    config::ManagerAddr,
    context::Context,
    net::{ConnectOpts, UdpSocket as ShadowUdpSocket},
};

/// Address accepted from Manager
#[derive(Debug)]
pub enum ManagerSocketAddr {
    SocketAddr(SocketAddr),
    #[cfg(unix)]
    UnixSocketAddr(UnixSocketAddr),
}

impl ManagerSocketAddr {
    /// Check if it is unnamed (not binded to any valid address), only valid for `UnixSocketAddr`
    pub fn is_unnamed(&self) -> bool {
        match *self {
            ManagerSocketAddr::SocketAddr(..) => false,
            #[cfg(unix)]
            ManagerSocketAddr::UnixSocketAddr(ref s) => s.is_unnamed(),
        }
    }
}

impl fmt::Display for ManagerSocketAddr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            ManagerSocketAddr::SocketAddr(ref saddr) => fmt::Display::fmt(saddr, f),
            #[cfg(unix)]
            ManagerSocketAddr::UnixSocketAddr(ref saddr) => fmt::Debug::fmt(saddr, f),
        }
    }
}

/// Datagram socket for manager
///
/// For *nix system, this is a wrapper for both UDP socket and Unix socket
pub enum ManagerDatagram {
    UdpDatagram(UdpSocket),
    #[cfg(unix)]
    UnixDatagram(UnixDatagram),
}

impl ManagerDatagram {
    /// Create a `ManagerDatagram` binding to requested `bind_addr`
    pub async fn bind(context: &Context, bind_addr: &ManagerAddr) -> io::Result<ManagerDatagram> {
        match *bind_addr {
            ManagerAddr::SocketAddr(ref saddr) => Ok(ManagerDatagram::UdpDatagram(
                ShadowUdpSocket::listen(saddr).await?.into(),
            )),
            ManagerAddr::DomainName(ref dname, port) => {
                let (_, socket) =
                    lookup_then!(context, dname, port, |saddr| { ShadowUdpSocket::listen(&saddr).await })?;

                Ok(ManagerDatagram::UdpDatagram(socket.into()))
            }
            #[cfg(unix)]
            ManagerAddr::UnixSocketAddr(ref path) => {
                use std::fs;

                // Remove it first incase it is already exists
                let _ = fs::remove_file(path);

                Ok(ManagerDatagram::UnixDatagram(UnixDatagram::bind(path)?))
            }
        }
    }

    /// Create a `ManagerDatagram` for sending data to manager
    pub async fn connect(
        context: &Context,
        bind_addr: &ManagerAddr,
        connect_opts: &ConnectOpts,
    ) -> io::Result<ManagerDatagram> {
        match *bind_addr {
            ManagerAddr::SocketAddr(sa) => ManagerDatagram::connect_socket_addr(sa, connect_opts).await,

            ManagerAddr::DomainName(ref dname, port) => {
                // Try connect to all socket addresses
                lookup_then!(context, dname, port, |addr| {
                    ManagerDatagram::connect_socket_addr(addr, connect_opts).await
                })
                .map(|(_, d)| d)
            }

            #[cfg(unix)]
            // For unix socket, it doesn't need to bind to any valid address
            // Because manager won't response to you
            ManagerAddr::UnixSocketAddr(ref path) => {
                let dgram = UnixDatagram::unbound()?;
                dgram.connect(path)?;
                Ok(ManagerDatagram::UnixDatagram(dgram))
            }
        }
    }

    async fn connect_socket_addr(sa: SocketAddr, connect_opts: &ConnectOpts) -> io::Result<ManagerDatagram> {
        let socket = ShadowUdpSocket::connect_with_opts(&sa, connect_opts).await?;
        Ok(ManagerDatagram::UdpDatagram(socket.into()))
    }

    /// Receives data from the socket.
    pub async fn recv(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match *self {
            ManagerDatagram::UdpDatagram(ref mut udp) => udp.recv(buf).await,
            #[cfg(unix)]
            ManagerDatagram::UnixDatagram(ref mut unix) => unix.recv(buf).await,
        }
    }

    /// Receives data from the socket.
    pub async fn recv_from(&mut self, buf: &mut [u8]) -> io::Result<(usize, ManagerSocketAddr)> {
        match *self {
            ManagerDatagram::UdpDatagram(ref mut udp) => {
                let (s, addr) = udp.recv_from(buf).await?;
                Ok((s, ManagerSocketAddr::SocketAddr(addr)))
            }
            #[cfg(unix)]
            ManagerDatagram::UnixDatagram(ref mut unix) => {
                let (s, addr) = unix.recv_from(buf).await?;
                Ok((s, ManagerSocketAddr::UnixSocketAddr(addr)))
            }
        }
    }

    /// Sends data to the socket
    pub async fn send(&mut self, buf: &[u8]) -> io::Result<usize> {
        match *self {
            ManagerDatagram::UdpDatagram(ref mut udp) => udp.send(buf).await,
            #[cfg(unix)]
            ManagerDatagram::UnixDatagram(ref mut unix) => unix.send(buf).await,
        }
    }

    /// Sends data to the socket to the specified address.
    pub async fn send_to(&mut self, buf: &[u8], target: &ManagerSocketAddr) -> io::Result<usize> {
        match *self {
            ManagerDatagram::UdpDatagram(ref mut udp) => match *target {
                ManagerSocketAddr::SocketAddr(ref saddr) => udp.send_to(buf, saddr).await,
                #[cfg(unix)]
                ManagerSocketAddr::UnixSocketAddr(..) => {
                    let err = io::Error::new(ErrorKind::InvalidInput, "udp datagram requires IP address target");
                    Err(err)
                }
            },
            #[cfg(unix)]
            ManagerDatagram::UnixDatagram(ref mut unix) => match *target {
                ManagerSocketAddr::UnixSocketAddr(ref saddr) => match saddr.as_pathname() {
                    Some(paddr) => unix.send_to(buf, paddr).await,
                    None => {
                        let err = io::Error::new(ErrorKind::InvalidInput, "target address must not be unnamed");
                        Err(err)
                    }
                },
                ManagerSocketAddr::SocketAddr(..) => {
                    let err = io::Error::new(ErrorKind::InvalidInput, "unix datagram requires path address target");
                    Err(err)
                }
            },
        }
    }

    /// Sends data on the socket to the specified manager address
    pub async fn send_to_manager(&mut self, buf: &[u8], context: &Context, target: &ManagerAddr) -> io::Result<usize> {
        match *self {
            ManagerDatagram::UdpDatagram(ref mut udp) => match *target {
                ManagerAddr::SocketAddr(ref saddr) => udp.send_to(buf, saddr).await,
                ManagerAddr::DomainName(ref dname, port) => {
                    let (_, n) = lookup_then!(context, dname, port, |saddr| { udp.send_to(buf, saddr).await })?;
                    Ok(n)
                }
                #[cfg(unix)]
                ManagerAddr::UnixSocketAddr(..) => {
                    let err = io::Error::new(ErrorKind::InvalidInput, "udp datagram requires IP address target");
                    Err(err)
                }
            },
            #[cfg(unix)]
            ManagerDatagram::UnixDatagram(ref mut unix) => match *target {
                ManagerAddr::UnixSocketAddr(ref paddr) => unix.send_to(buf, paddr).await,
                ManagerAddr::SocketAddr(..) | ManagerAddr::DomainName(..) => {
                    let err = io::Error::new(ErrorKind::InvalidInput, "unix datagram requires path address target");
                    Err(err)
                }
            },
        }
    }

    /// Returns the local address that this socket is bound to.
    pub fn local_addr(&self) -> io::Result<ManagerSocketAddr> {
        match *self {
            ManagerDatagram::UdpDatagram(ref socket) => socket.local_addr().map(ManagerSocketAddr::SocketAddr),
            #[cfg(unix)]
            ManagerDatagram::UnixDatagram(ref dgram) => dgram.local_addr().map(ManagerSocketAddr::UnixSocketAddr),
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate lazy_static;
extern crate quickcheck;
extern crate fail;
extern crate rand;
extern crate sled;
extern crate pagecache;

use std::collections::{BTreeMap, HashSet};
use std::sync::Mutex;

use quickcheck::{Arbitrary, Gen, QuickCheck, StdGen};

use sled::*;

#[derive(Debug, Clone)]
enum Op {
    Set(u8, u8),
    Del(u8),
    Restart,
    FailPoint(String),
}

impl Arbitrary for Op {
    fn arbitrary<G: Gen>(g: &mut G) -> Op {
        let fail_points = vec![
            "initial allocation",
            "initial allocation post",
            "buffer write",
            "buffer write post",
            "trailer write",
            "trailer write post",
            "zero garbage segment",
            "zero garbage segment post",
            "zero segment",
            "zero segment post",
            "snap write",
            "snap write len",
            "snap write crc",
            "snap write post",
            "snap write mv",
            "snap write mv post",
            "snap write rm old",
            "write_config bytes",
            "write_config crc",
            "write_config post",
        ];

        if g.gen_weighted_bool(10) {
            return Op::FailPoint((*g.choose(&fail_points).unwrap()).to_owned());
        }

        if g.gen_weighted_bool(10) {
            return Op::Restart;
        }

        let choice = g.gen_range(0, 2);

        match choice {
            0 => Op::Set(g.gen::<u8>(), g.gen::<u8>()),
            1 => Op::Del(g.gen::<u8>()),
            _ => panic!("impossible choice"),
        }
    }
}

#[derive(Debug, Clone)]
struct OpVec {
    ops: Vec<Op>,
}

impl Arbitrary for OpVec {
    fn arbitrary<G: Gen>(g: &mut G) -> OpVec {
        let mut ops = vec![];
        for _ in 0..g.gen_range(1, 100) {
            let op = Op::arbitrary(g);
            ops.push(op);

        }
        OpVec {
            ops: ops,
        }
    }

    fn shrink(&self) -> Box<Iterator<Item = OpVec>> {
        let mut smaller = vec![];
        for i in 0..self.ops.len() {
            let mut clone = self.clone();
            clone.ops.remove(i);
            smaller.push(clone);
        }

        Box::new(smaller.into_iter())
    }
}

fn prop_tree_crashes_nicely(
    ops: OpVec,
    blink_fanout: u8,
    snapshot_after: u8,
) -> bool {
    lazy_static! {
        // forces quickcheck to run one thread at a time
        static ref M: Mutex<()> = Mutex::new(());
    }

    let _lock = M.lock().unwrap();

    use self::Op::*;

    fail::setup();

    let config = ConfigBuilder::new()
        .temporary(true)
        .snapshot_after_ops(snapshot_after as usize + 1)
        .flush_every_ms(Some(1))
        .io_buf_size(10000)
        .blink_fanout(blink_fanout as usize + 2)
        .cache_capacity(40)
        .build();

    let mut tree = sled::Tree::start(config.clone()).unwrap();
    let mut reference = BTreeMap::new();
    let mut fail_points = HashSet::new();

    macro_rules! restart {
        () => {
            println!("restarting");
            drop(tree);
            tree = sled::Tree::start(config.clone()).unwrap();

            let tree_iter = tree.iter().map(|res| {
                let (ref tk, ref tv) = res.unwrap();
                (tk[0], tv[0])
            });
            let ref_iter =
                reference.iter().map(|(ref rk, ref rv)| (**rk, **rv));
            for (t, r) in tree_iter.zip(ref_iter) {
                assert_eq!(t, r);
            }
        }
    }

    macro_rules! clear_fps {
        () => {
            for fp in fail_points.drain() {
                fail::remove(fp);
            }
        }
    }

    macro_rules! fp_crash {
        ($e:expr) => {
            match $e {
                Ok(thing) => thing,
                Err(Error::FailPoint) => {
                    println!("got a crash");
                    clear_fps!();
                    restart!();
                    continue;
                }
                _ => {
                    println!("about to unwrap a weirdo {:?}", $e);
                    $e.unwrap()
                },
            }
        }
    }

    for op in ops.ops.into_iter() {
        match op {
            Set(k, v) => {
                fp_crash!(tree.set(vec![k], vec![v]));
                reference.insert(k, v);
            }
            Del(k) => {
                fp_crash!(tree.del(&*vec![k]));
                reference.remove(&k);
            }
            Restart => {
                restart!();
            }
            FailPoint(fp) => {
                fail_points.insert(fp.clone());
                println!("configuring fp at {}", fp);
                fail::cfg(&*fp, "return").unwrap();
            }
        }
    }

    fail::teardown();

    true
}

#[test]
fn quickcheck_tree_with_failpoints() {
    // use fewer tests for travis OSX builds that stall out all the time
    #[cfg(target_os = "macos")]
    let n_tests = 100;

    #[cfg(not(target_os = "macos"))]
    let n_tests = 500;

    QuickCheck::new()
        .gen(StdGen::new(rand::thread_rng(), 1))
        .tests(n_tests)
        .max_tests(10000)
        .quickcheck(prop_tree_crashes_nicely as fn(OpVec, u8, u8) -> bool);
}
<commit_msg>capture more errors in a clean way from the property function<commit_after>#[macro_use]
extern crate lazy_static;
extern crate quickcheck;
extern crate fail;
extern crate rand;
extern crate sled;
extern crate pagecache;

use std::collections::{BTreeMap, HashSet};
use std::sync::Mutex;

use quickcheck::{Arbitrary, Gen, QuickCheck, StdGen};

use sled::*;

#[derive(Debug, Clone)]
enum Op {
    Set(u8, u8),
    Del(u8),
    Restart,
    FailPoint(String),
}

impl Arbitrary for Op {
    fn arbitrary<G: Gen>(g: &mut G) -> Op {
        let fail_points = vec![
            "initial allocation",
            "initial allocation post",
            "buffer write",
            "buffer write post",
            "trailer write",
            "trailer write post",
            "zero garbage segment",
            "zero garbage segment post",
            "zero segment",
            "zero segment post",
            "snap write",
            "snap write len",
            "snap write crc",
            "snap write post",
            "snap write mv",
            "snap write mv post",
            "snap write rm old",
            "write_config bytes",
            "write_config crc",
            "write_config post",
        ];

        if g.gen_weighted_bool(10) {
            return Op::FailPoint((*g.choose(&fail_points).unwrap()).to_owned());
        }

        if g.gen_weighted_bool(10) {
            return Op::Restart;
        }

        let choice = g.gen_range(0, 2);

        match choice {
            0 => Op::Set(g.gen::<u8>(), g.gen::<u8>()),
            1 => Op::Del(g.gen::<u8>()),
            _ => panic!("impossible choice"),
        }
    }
}

#[derive(Debug, Clone)]
struct OpVec {
    ops: Vec<Op>,
}

impl Arbitrary for OpVec {
    fn arbitrary<G: Gen>(g: &mut G) -> OpVec {
        let mut ops = vec![];
        for _ in 0..g.gen_range(1, 100) {
            let op = Op::arbitrary(g);
            ops.push(op);

        }
        OpVec {
            ops: ops,
        }
    }

    fn shrink(&self) -> Box<Iterator<Item = OpVec>> {
        let mut smaller = vec![];
        for i in 0..self.ops.len() {
            let mut clone = self.clone();
            clone.ops.remove(i);
            smaller.push(clone);
        }

        Box::new(smaller.into_iter())
    }
}

fn prop_tree_crashes_nicely(ops: OpVec) -> bool {
    lazy_static! {
        // forces quickcheck to run one thread at a time
        static ref M: Mutex<()> = Mutex::new(());
    }

    let _lock = M.lock().unwrap();

    use self::Op::*;

    fail::setup();

    let config = ConfigBuilder::new()
        .temporary(true)
        .snapshot_after_ops(1)
        .flush_every_ms(Some(1))
        .io_buf_size(10000)
        .cache_capacity(40)
        .build();

    let mut tree = sled::Tree::start(config.clone()).unwrap();
    let mut reference = BTreeMap::new();
    let mut fail_points = HashSet::new();

    macro_rules! restart {
        () => {
            println!("restarting");
            drop(tree);
            let tree_res = sled::Tree::start(config.clone());
            if tree_res.is_err() {
                println!("failed to open tree: {:?}", tree_res.unwrap_err());
                clear_fps!();
                fail::teardown();
                return false;
            }

            tree = tree_res.unwrap();

            let tree_iter = tree.iter().map(|res| {
                let (ref tk, ref tv) = res.unwrap();
                (tk[0], tv[0])
            });
            let ref_iter =
                reference.iter().map(|(ref rk, ref rv)| (**rk, **rv));
            for (t, r) in tree_iter.zip(ref_iter) {
                if t != r {
                    println!("tree value {:?} failed to match expected reference {:?}", t, r);
                    clear_fps!();
                    fail::teardown();
                    return false;
                }
            }
        }
    }

    macro_rules! clear_fps {
        () => {
            for fp in fail_points.drain() {
                fail::remove(fp);
            }
        }
    }

    macro_rules! fp_crash {
        ($e:expr) => {
            match $e {
                Ok(thing) => thing,
                Err(Error::FailPoint) => {
                    println!("got a crash");
                    clear_fps!();
                    restart!();
                    continue;
                }
                _ => {
                    println!("about to unwrap a weirdo {:?}", $e);
                    clear_fps!();
                    fail::teardown();
                    return false;
                },
            }
        }
    }

    for op in ops.ops.into_iter() {
        match op {
            Set(k, v) => {
                fp_crash!(tree.set(vec![k], vec![v]));
                reference.insert(k, v);
            }
            Del(k) => {
                fp_crash!(tree.del(&*vec![k]));
                reference.remove(&k);
            }
            Restart => {
                restart!();
            }
            FailPoint(fp) => {
                fail_points.insert(fp.clone());
                println!("configuring fp at {}", fp);
                fail::cfg(&*fp, "return").unwrap();
            }
        }
    }

    fail::teardown();

    true
}

#[test]
fn quickcheck_tree_with_failpoints() {
    // use fewer tests for travis OSX builds that stall out all the time
    #[cfg(target_os = "macos")]
    let n_tests = 100;

    #[cfg(not(target_os = "macos"))]
    let n_tests = 500;

    QuickCheck::new()
        .gen(StdGen::new(rand::thread_rng(), 1))
        .tests(n_tests)
        .max_tests(10000)
        .quickcheck(prop_tree_crashes_nicely as fn(OpVec) -> bool);
}
<|endoftext|>


<commit_before>use rustc_hir::def_id::DefId;
use rustc_hir::intravisit::FnKind;
use rustc_index::bit_set::BitSet;
use rustc_index::vec::IndexVec;
use rustc_middle::hir::map::blocks::FnLikeNode;
use rustc_middle::mir::{BasicBlock, Body, ReadOnlyBodyAndCache, TerminatorKind, START_BLOCK};
use rustc_middle::ty::subst::InternalSubsts;
use rustc_middle::ty::{self, AssocItem, AssocItemContainer, Instance, TyCtxt};
use rustc_session::lint::builtin::UNCONDITIONAL_RECURSION;
use std::collections::VecDeque;

crate fn check<'tcx>(tcx: TyCtxt<'tcx>, body: &ReadOnlyBodyAndCache<'_, 'tcx>, def_id: DefId) {
    let hir_id = tcx.hir().as_local_hir_id(def_id).unwrap();

    if let Some(fn_like_node) = FnLikeNode::from_node(tcx.hir().get(hir_id)) {
        if let FnKind::Closure(_) = fn_like_node.kind() {
            // closures can't recur, so they don't matter.
            return;
        }

        check_fn_for_unconditional_recursion(tcx, body, def_id);
    }
}

fn check_fn_for_unconditional_recursion<'tcx>(
    tcx: TyCtxt<'tcx>,
    body: &ReadOnlyBodyAndCache<'_, 'tcx>,
    def_id: DefId,
) {
    let self_calls = find_blocks_calling_self(tcx, &body, def_id);

    // Stores a list of `Span`s for every basic block. Those are the spans of `Call` terminators
    // where we know that one of them will definitely be reached.
    let mut results = IndexVec::from_elem_n(vec![], body.basic_blocks().len());

    // We start the analysis at the self calls and work backwards.
    let mut queue: VecDeque<_> = self_calls.iter().collect();

    while let Some(bb) = queue.pop_front() {
        if !results[bb].is_empty() {
            // Already propagated.
            continue;
        }

        let locations = if self_calls.contains(bb) {
            // `bb` *is* a self-call.
            // We don't look at successors here because they are irrelevant here and we don't want
            // to lint them (eg. `f(); f()` should only lint the first call).
            vec![bb]
        } else {
            // If *all* successors of `bb` lead to a self-call, emit notes at all of their
            // locations.

            // Converging successors without unwind paths.
            let terminator = body[bb].terminator();
            let relevant_successors = match &terminator.kind {
                TerminatorKind::Call { destination: None, .. }
                | TerminatorKind::Yield { .. }
                | TerminatorKind::GeneratorDrop => None.into_iter().chain(&[]),
                TerminatorKind::SwitchInt { targets, .. } => None.into_iter().chain(targets),
                TerminatorKind::Goto { target }
                | TerminatorKind::Drop { target, .. }
                | TerminatorKind::DropAndReplace { target, .. }
                | TerminatorKind::Assert { target, .. }
                | TerminatorKind::FalseEdges { real_target: target, .. }
                | TerminatorKind::FalseUnwind { real_target: target, .. }
                | TerminatorKind::Call { destination: Some((_, target)), .. } => {
                    Some(target).into_iter().chain(&[])
                }
                TerminatorKind::Resume
                | TerminatorKind::Abort
                | TerminatorKind::Return
                | TerminatorKind::Unreachable => {
                    // We propagate backwards, so these should never be encountered here.
                    unreachable!("unexpected terminator {:?}", terminator.kind)
                }
            };

            // If all our successors are known to lead to self-calls, then we do too.
            let all_are_self_calls =
                relevant_successors.clone().all(|&succ| !results[succ].is_empty());

            if all_are_self_calls {
                // We'll definitely lead to a self-call. Merge all call locations of the successors
                // for linting them later.
                relevant_successors.flat_map(|&succ| results[succ].iter().copied()).collect()
            } else {
                // At least 1 successor does not always lead to a self-call, so we also don't.
                vec![]
            }
        };

        if !locations.is_empty() {
            // This is a newly confirmed-to-always-reach-self-call block.
            results[bb] = locations;

            // Propagate backwards through the CFG.
            debug!("propagate loc={:?} in {:?} -> {:?}", results[bb], bb, body.predecessors()[bb]);
            queue.extend(body.predecessors()[bb].iter().copied());
        }
    }

    debug!("unconditional recursion results: {:?}", results);

    let self_call_locations = &mut results[START_BLOCK];
    self_call_locations.sort();
    self_call_locations.dedup();

    if !self_call_locations.is_empty() {
        let hir_id = tcx.hir().as_local_hir_id(def_id).unwrap();
        let sp = tcx.sess.source_map().guess_head_span(tcx.hir().span(hir_id));
        tcx.struct_span_lint_hir(UNCONDITIONAL_RECURSION, hir_id, sp, |lint| {
            let mut db = lint.build("function cannot return without recursing");
            db.span_label(sp, "cannot return without recursing");
            // offer some help to the programmer.
            for bb in self_call_locations {
                let span = body.basic_blocks()[*bb].terminator().source_info.span;
                db.span_label(span, "recursive call site");
            }
            db.help("a `loop` may express intention better if this is on purpose");
            db.emit();
        });
    }
}

/// Finds blocks with `Call` terminators that would end up calling back into the same method.
fn find_blocks_calling_self<'tcx>(
    tcx: TyCtxt<'tcx>,
    body: &Body<'tcx>,
    def_id: DefId,
) -> BitSet<BasicBlock> {
    let param_env = tcx.param_env(def_id);

    // If this is trait/impl method, extract the trait's substs.
    let trait_substs_count = match tcx.opt_associated_item(def_id) {
        Some(AssocItem { container: AssocItemContainer::TraitContainer(trait_def_id), .. }) => {
            tcx.generics_of(trait_def_id).count()
        }
        _ => 0,
    };
    let trait_substs = &InternalSubsts::identity_for_item(tcx, def_id)[..trait_substs_count];

    let mut self_calls = BitSet::new_empty(body.basic_blocks().len());

    for (bb, data) in body.basic_blocks().iter_enumerated() {
        if let TerminatorKind::Call { func, .. } = &data.terminator().kind {
            let func_ty = func.ty(body, tcx);

            if let ty::FnDef(fn_def_id, substs) = func_ty.kind {
                let (call_fn_id, call_substs) =
                    if let Some(instance) = Instance::resolve(tcx, param_env, fn_def_id, substs) {
                        (instance.def_id(), instance.substs)
                    } else {
                        (fn_def_id, substs)
                    };

                // FIXME(#57965): Make this work across function boundaries

                // If this is a trait fn, the substs on the trait have to match, or we might be
                // calling into an entirely different method (for example, a call from the default
                // method in the trait to `<A as Trait<B>>::method`, where `A` and/or `B` are
                // specific types).
                let is_self_call =
                    call_fn_id == def_id && &call_substs[..trait_substs.len()] == trait_substs;

                if is_self_call {
                    self_calls.insert(bb);
                }
            }
        }
    }

    self_calls
}
<commit_msg>Further improve comments<commit_after>use rustc_hir::def_id::DefId;
use rustc_hir::intravisit::FnKind;
use rustc_index::bit_set::BitSet;
use rustc_index::vec::IndexVec;
use rustc_middle::hir::map::blocks::FnLikeNode;
use rustc_middle::mir::{BasicBlock, Body, ReadOnlyBodyAndCache, TerminatorKind, START_BLOCK};
use rustc_middle::ty::subst::InternalSubsts;
use rustc_middle::ty::{self, AssocItem, AssocItemContainer, Instance, TyCtxt};
use rustc_session::lint::builtin::UNCONDITIONAL_RECURSION;
use std::collections::VecDeque;

crate fn check<'tcx>(tcx: TyCtxt<'tcx>, body: &ReadOnlyBodyAndCache<'_, 'tcx>, def_id: DefId) {
    let hir_id = tcx.hir().as_local_hir_id(def_id).unwrap();

    if let Some(fn_like_node) = FnLikeNode::from_node(tcx.hir().get(hir_id)) {
        if let FnKind::Closure(_) = fn_like_node.kind() {
            // closures can't recur, so they don't matter.
            return;
        }

        check_fn_for_unconditional_recursion(tcx, body, def_id);
    }
}

fn check_fn_for_unconditional_recursion<'tcx>(
    tcx: TyCtxt<'tcx>,
    body: &ReadOnlyBodyAndCache<'_, 'tcx>,
    def_id: DefId,
) {
    let self_calls = find_blocks_calling_self(tcx, &body, def_id);

    // Stores a list of `Span`s for every basic block. Those are the spans of self-calls where we
    // know that one of them will definitely be reached. If the list is empty, the block either
    // wasn't processed yet or will not always go to a self-call.
    let mut results = IndexVec::from_elem_n(vec![], body.basic_blocks().len());

    // We start the analysis at the self calls and work backwards.
    let mut queue: VecDeque<_> = self_calls.iter().collect();

    while let Some(bb) = queue.pop_front() {
        if !results[bb].is_empty() {
            // Already propagated.
            continue;
        }

        let locations = if self_calls.contains(bb) {
            // `bb` *is* a self-call.
            // We don't look at successors here because they are irrelevant here and we don't want
            // to lint them (eg. `f(); f()` should only lint the first call).
            vec![bb]
        } else {
            // If *all* successors of `bb` lead to a self-call, emit notes at all of their
            // locations.

            // Determine all "relevant" successors. We ignore successors only reached via unwinding.
            let terminator = body[bb].terminator();
            let relevant_successors = match &terminator.kind {
                TerminatorKind::Call { destination: None, .. }
                | TerminatorKind::Yield { .. }
                | TerminatorKind::GeneratorDrop => None.into_iter().chain(&[]),
                TerminatorKind::SwitchInt { targets, .. } => None.into_iter().chain(targets),
                TerminatorKind::Goto { target }
                | TerminatorKind::Drop { target, .. }
                | TerminatorKind::DropAndReplace { target, .. }
                | TerminatorKind::Assert { target, .. }
                | TerminatorKind::FalseEdges { real_target: target, .. }
                | TerminatorKind::FalseUnwind { real_target: target, .. }
                | TerminatorKind::Call { destination: Some((_, target)), .. } => {
                    Some(target).into_iter().chain(&[])
                }
                TerminatorKind::Resume
                | TerminatorKind::Abort
                | TerminatorKind::Return
                | TerminatorKind::Unreachable => {
                    // We propagate backwards, so these should never be encountered here.
                    unreachable!("unexpected terminator {:?}", terminator.kind)
                }
            };

            // If all our successors are known to lead to self-calls, then we do too.
            let all_are_self_calls =
                relevant_successors.clone().all(|&succ| !results[succ].is_empty());

            if all_are_self_calls {
                // We'll definitely lead to a self-call. Merge all call locations of the successors
                // for linting them later.
                relevant_successors.flat_map(|&succ| results[succ].iter().copied()).collect()
            } else {
                // At least 1 successor does not always lead to a self-call, so we also don't.
                vec![]
            }
        };

        if !locations.is_empty() {
            // This is a newly confirmed-to-always-reach-self-call block.
            results[bb] = locations;

            // Propagate backwards through the CFG.
            debug!("propagate loc={:?} in {:?} -> {:?}", results[bb], bb, body.predecessors()[bb]);
            queue.extend(body.predecessors()[bb].iter().copied());
        }
    }

    debug!("unconditional recursion results: {:?}", results);

    let self_call_locations = &mut results[START_BLOCK];
    self_call_locations.sort();
    self_call_locations.dedup();

    if !self_call_locations.is_empty() {
        let hir_id = tcx.hir().as_local_hir_id(def_id).unwrap();
        let sp = tcx.sess.source_map().guess_head_span(tcx.hir().span(hir_id));
        tcx.struct_span_lint_hir(UNCONDITIONAL_RECURSION, hir_id, sp, |lint| {
            let mut db = lint.build("function cannot return without recursing");
            db.span_label(sp, "cannot return without recursing");
            // offer some help to the programmer.
            for bb in self_call_locations {
                let span = body.basic_blocks()[*bb].terminator().source_info.span;
                db.span_label(span, "recursive call site");
            }
            db.help("a `loop` may express intention better if this is on purpose");
            db.emit();
        });
    }
}

/// Finds blocks with `Call` terminators that would end up calling back into the same method.
fn find_blocks_calling_self<'tcx>(
    tcx: TyCtxt<'tcx>,
    body: &Body<'tcx>,
    def_id: DefId,
) -> BitSet<BasicBlock> {
    let param_env = tcx.param_env(def_id);

    // If this is trait/impl method, extract the trait's substs.
    let trait_substs_count = match tcx.opt_associated_item(def_id) {
        Some(AssocItem { container: AssocItemContainer::TraitContainer(trait_def_id), .. }) => {
            tcx.generics_of(trait_def_id).count()
        }
        _ => 0,
    };
    let trait_substs = &InternalSubsts::identity_for_item(tcx, def_id)[..trait_substs_count];

    let mut self_calls = BitSet::new_empty(body.basic_blocks().len());

    for (bb, data) in body.basic_blocks().iter_enumerated() {
        if let TerminatorKind::Call { func, .. } = &data.terminator().kind {
            let func_ty = func.ty(body, tcx);

            if let ty::FnDef(fn_def_id, substs) = func_ty.kind {
                let (call_fn_id, call_substs) =
                    if let Some(instance) = Instance::resolve(tcx, param_env, fn_def_id, substs) {
                        (instance.def_id(), instance.substs)
                    } else {
                        (fn_def_id, substs)
                    };

                // FIXME(#57965): Make this work across function boundaries

                // If this is a trait fn, the substs on the trait have to match, or we might be
                // calling into an entirely different method (for example, a call from the default
                // method in the trait to `<A as Trait<B>>::method`, where `A` and/or `B` are
                // specific types).
                let is_self_call =
                    call_fn_id == def_id && &call_substs[..trait_substs.len()] == trait_substs;

                if is_self_call {
                    self_calls.insert(bb);
                }
            }
        }
    }

    self_calls
}
<|endoftext|>


<commit_before>/// Memory buffers for the benefit of `std::rt::io::net` which has slow read/write.

use std::rt::io::{Reader, Writer};
use std::cast::transmute_mut;
use std::cmp::min;
use std::ptr;

// 64KB chunks (moderately arbitrary)
static READ_BUF_SIZE: uint = 0x10000;
static WRITE_BUF_SIZE: uint = 0x10000;
// TODO: consider removing constants and giving a buffer size in the constructor

struct BufferedReader<'self, T> {
    wrapped: &'self mut T,
    buffer: [u8, ..READ_BUF_SIZE],
    // The current position in the buffer
    pos: uint,
    // The last valid position in the reader
    max: uint,
}

impl<'self, T: Reader> BufferedReader<'self, T> {
    pub fn new<'a>(reader: &'a mut T/*, buffer_size: uint*/) -> BufferedReader<'a, T> {
        BufferedReader {
            wrapped: reader,
            buffer: [0u8, ..READ_BUF_SIZE], //[0u8, ..buffer_size],
            pos: 0u,
            max: 0u,
        }
    }

    /// Poke a single byte back so it will be read next. For this to make sense, you must have just
    /// read that byte. If `self.pos` is 0 and `self.max` is not 0 (i.e. if the buffer is just
    /// filled
    /// Very great caution must be used in calling this as it will fail if `self.pos` is 0.
    pub fn poke_byte(&mut self, byte: u8) {
        match (self.pos, self.max) {
            (0, 0) => self.max = 1,
            (0, _) => fail!("poke called when buffer is full"),
            (_, _) => self.pos -= 1,
        }
        self.buffer[self.pos] = byte;
    }
}

impl<'self, T: Reader> Reader for ~BufferedReader<'self, T> {
    /// Read at most N bytes into `buf`, where N is the minimum of `buf.len()` and the buffer size.
    ///
    /// At present, this makes no attempt to fill its buffer proactively, instead waiting until you
    /// ask.
    fn read(&mut self, buf: &mut [u8]) -> Option<uint> {
        if self.pos == self.max {
            // Run out of buffered content, read some more
            match self.wrapped.read(self.buffer) {
                None => {
                    self.pos = 0;
                    self.max = 0;
                    return None
                },
                Some(i) => {
                    self.pos = 0;
                    self.max = i;
                },
            }
        }
        let size = min(self.max - self.pos, buf.len());
        unsafe {
            do buf.as_mut_buf |p_dst, _len_dst| {
                do self.buffer.as_imm_buf |p_src, _len_src| {
                    // Note that copy_memory works on bytes; good, u8 is byte-sized
                    ptr::copy_memory(p_dst, ptr::offset(p_src, self.pos), size)
                }
            }
        }
        self.pos += size;
        Some(size)
    }

    /// Return whether the Reader has reached the end of the stream AND exhausted its buffer.
    fn eof(&mut self) -> bool {
        self.pos == self.max && self.wrapped.eof()
    }
}

struct BufferedWriter<'self, T> {
    wrapped: &'self mut T,
    buffer: [u8, ..WRITE_BUF_SIZE],
    buflen: uint,

    /// Some things being written may not like flush() being called yet (e.g. explicitly fail!())
    /// The BufferedReader may need to be flushed for good control, but let it provide for such
    /// cases by not calling the wrapped object's flush method in turn.
    call_wrapped_flush: bool,
}

impl<'self, T: Writer> BufferedWriter<'self, T> {
    pub fn new<'a>(writer: &'a mut T, call_wrapped_flush: bool/*, buffer_size: uint*/)
    -> BufferedWriter<'a, T> {
        BufferedWriter {
            wrapped: writer,
            buffer: [0u8, ..WRITE_BUF_SIZE], //[0u8, ..buffer_size],
            buflen: 0u,
            call_wrapped_flush: call_wrapped_flush,
        }
    }
}

#[unsafe_destructor]
impl<'self, T: Writer> Drop for BufferedWriter<'self, T> {
    fn drop(&self) {
        // Clearly wouldn't be a good idea to finish without flushing!

        // TODO: blocked on https://github.com/mozilla/rust/issues/4252
        // Also compare usage of response.flush() in server.rs
        //unsafe { transmute_mut(self) }.flush();
    }
}

impl<'self, T: Writer> Writer for BufferedWriter<'self, T> {
    fn write(&mut self, buf: &[u8]) {
        if buf.len() + self.buflen > self.buffer.len() {
            // This is the lazy approach which may involve two writes where it's really not
            // warranted. Maybe deal with that later.
            if self.buflen > 0 {
                self.wrapped.write(self.buffer.slice_to(self.buflen));
                self.buflen = 0;
            }
            self.wrapped.write(buf);
            self.buflen = 0;
        } else {
            // Safely copy buf onto the "end" of self.buffer
            unsafe {
                do buf.as_imm_buf |p_src, len_src| {
                    do self.buffer.as_mut_buf |p_dst, _len_dst| {
                        // Note that copy_memory works on bytes; good, u8 is byte-sized
                        ptr::copy_memory(ptr::mut_offset(p_dst, self.buflen), p_src, len_src)
                    }
                }
            }
            self.buflen += buf.len();
            if self.buflen == self.buffer.len() {
                self.wrapped.write(self.buffer);
                self.buflen = 0;
            }
        }
    }

    fn flush(&mut self) {
        if self.buflen > 0 {
            self.wrapped.write(self.buffer.slice_to(self.buflen));
            self.buflen = 0;
        }
        if self.call_wrapped_flush {
            self.wrapped.flush();
        }
    }
}
<commit_msg>Make a more efficient BufferedReader.read_byte().<commit_after>/// Memory buffers for the benefit of `std::rt::io::net` which has slow read/write.

use std::rt::io::{Reader, Writer};
use std::cast::transmute_mut;
use std::cmp::min;
use std::ptr;

// 64KB chunks (moderately arbitrary)
static READ_BUF_SIZE: uint = 0x10000;
static WRITE_BUF_SIZE: uint = 0x10000;
// TODO: consider removing constants and giving a buffer size in the constructor

struct BufferedReader<'self, T> {
    wrapped: &'self mut T,
    buffer: [u8, ..READ_BUF_SIZE],
    // The current position in the buffer
    pos: uint,
    // The last valid position in the reader
    max: uint,
}

impl<'self, T: Reader> BufferedReader<'self, T> {
    pub fn new<'a>(reader: &'a mut T/*, buffer_size: uint*/) -> BufferedReader<'a, T> {
        BufferedReader {
            wrapped: reader,
            buffer: [0u8, ..READ_BUF_SIZE], //[0u8, ..buffer_size],
            pos: 0u,
            max: 0u,
        }
    }

    /// Poke a single byte back so it will be read next. For this to make sense, you must have just
    /// read that byte. If `self.pos` is 0 and `self.max` is not 0 (i.e. if the buffer is just
    /// filled
    /// Very great caution must be used in calling this as it will fail if `self.pos` is 0.
    pub fn poke_byte(&mut self, byte: u8) {
        match (self.pos, self.max) {
            (0, 0) => self.max = 1,
            (0, _) => fail!("poke called when buffer is full"),
            (_, _) => self.pos -= 1,
        }
        self.buffer[self.pos] = byte;
    }

    #[inline]
    fn fill_buffer(&mut self) -> bool {
        assert_eq!(self.pos, self.max);
        match self.wrapped.read(self.buffer) {
            None => {
                self.pos = 0;
                self.max = 0;
                false
            },
            Some(i) => {
                self.pos = 0;
                self.max = i;
                true
            },
        }
    }

    /// Slightly faster implementation of read_byte than that which is provided by ReaderUtil
    /// (which just uses `read()`)
    #[inline]
    pub fn read_byte(&mut self) -> Option<u8> {
        if self.pos == self.max && !self.fill_buffer() {
            // Run out of buffered content, no more to come
            return None;
        }
        self.pos += 1;
        Some(self.buffer[self.pos - 1])
    }
}

impl<'self, T: Reader> Reader for ~BufferedReader<'self, T> {
    /// Read at most N bytes into `buf`, where N is the minimum of `buf.len()` and the buffer size.
    ///
    /// At present, this makes no attempt to fill its buffer proactively, instead waiting until you
    /// ask.
    fn read(&mut self, buf: &mut [u8]) -> Option<uint> {
        if self.pos == self.max && !self.fill_buffer() {
            // Run out of buffered content, no more to come
            return None;
        }
        let size = min(self.max - self.pos, buf.len());
        unsafe {
            do buf.as_mut_buf |p_dst, _len_dst| {
                do self.buffer.as_imm_buf |p_src, _len_src| {
                    // Note that copy_memory works on bytes; good, u8 is byte-sized
                    ptr::copy_memory(p_dst, ptr::offset(p_src, self.pos), size)
                }
            }
        }
        self.pos += size;
        Some(size)
    }

    /// Return whether the Reader has reached the end of the stream AND exhausted its buffer.
    fn eof(&mut self) -> bool {
        self.pos == self.max && self.wrapped.eof()
    }
}

struct BufferedWriter<'self, T> {
    wrapped: &'self mut T,
    buffer: [u8, ..WRITE_BUF_SIZE],
    buflen: uint,

    /// Some things being written may not like flush() being called yet (e.g. explicitly fail!())
    /// The BufferedReader may need to be flushed for good control, but let it provide for such
    /// cases by not calling the wrapped object's flush method in turn.
    call_wrapped_flush: bool,
}

impl<'self, T: Writer> BufferedWriter<'self, T> {
    pub fn new<'a>(writer: &'a mut T, call_wrapped_flush: bool/*, buffer_size: uint*/)
    -> BufferedWriter<'a, T> {
        BufferedWriter {
            wrapped: writer,
            buffer: [0u8, ..WRITE_BUF_SIZE], //[0u8, ..buffer_size],
            buflen: 0u,
            call_wrapped_flush: call_wrapped_flush,
        }
    }
}

#[unsafe_destructor]
impl<'self, T: Writer> Drop for BufferedWriter<'self, T> {
    fn drop(&self) {
        // Clearly wouldn't be a good idea to finish without flushing!

        // TODO: blocked on https://github.com/mozilla/rust/issues/4252
        // Also compare usage of response.flush() in server.rs
        //unsafe { transmute_mut(self) }.flush();
    }
}

impl<'self, T: Writer> Writer for BufferedWriter<'self, T> {
    fn write(&mut self, buf: &[u8]) {
        if buf.len() + self.buflen > self.buffer.len() {
            // This is the lazy approach which may involve two writes where it's really not
            // warranted. Maybe deal with that later.
            if self.buflen > 0 {
                self.wrapped.write(self.buffer.slice_to(self.buflen));
                self.buflen = 0;
            }
            self.wrapped.write(buf);
            self.buflen = 0;
        } else {
            // Safely copy buf onto the "end" of self.buffer
            unsafe {
                do buf.as_imm_buf |p_src, len_src| {
                    do self.buffer.as_mut_buf |p_dst, _len_dst| {
                        // Note that copy_memory works on bytes; good, u8 is byte-sized
                        ptr::copy_memory(ptr::mut_offset(p_dst, self.buflen), p_src, len_src)
                    }
                }
            }
            self.buflen += buf.len();
            if self.buflen == self.buffer.len() {
                self.wrapped.write(self.buffer);
                self.buflen = 0;
            }
        }
    }

    fn flush(&mut self) {
        if self.buflen > 0 {
            self.wrapped.write(self.buffer.slice_to(self.buflen));
            self.buflen = 0;
        }
        if self.call_wrapped_flush {
            self.wrapped.flush();
        }
    }
}
<|endoftext|>


<commit_before>extern crate simple_excel_writer;
use std::io::{Cursor, Read};
use excel::*;
use simple_excel_writer as excel;

fn creates_and_saves_an_excel_sheet_driver(filename: Option<&str>) -> Option<Vec<u8>> {
    let mut wb = if let Some(name) = filename {
        excel::Workbook::create(name)
    } else {
        excel::Workbook::create_in_memory()
    };

    let mut ws = wb.create_sheet("test_sheet");
    ws.add_column(Column { width: 60.0 });
    ws.add_column(Column { width: 30.0 });
    ws.add_column(Column { width: 10.0 });
    ws.add_column(Column { width: 60.0 });
    wb.write_sheet(&mut ws, |sw| {
        sw.append_row(row!["Name", "Title", "Success", "Remark"])
            .unwrap();
        sw.append_row(row![
            "Mary",
            "Acountant",
            false,
            r#"<xml><tag>"" & 'World'</tag></xml>"#
        ])
        .unwrap();
        sw.append_row(row![
            "Mary",
            "Programmer",
            true,
            "<xml><tag>\"Hello\" & 'World'</tag></xml>"
        ])
        .unwrap();
        sw.append_row(row!["Marly", "Mary", "Success", "Success", true, 500.])
    })
    .expect("Write excel error!");

    let mut ws = wb.create_sheet("test_sheet<2");

    wb.write_sheet(&mut ws, |sw| {
        sw.append_row(row!["Name", "Title", "Success"]).unwrap();
        sw.append_row(row!["Mary", "This", true]).unwrap();

        #[cfg(feature = "chrono")]
        sw.append_row(row![
            chrono::NaiveDate::from_ymd(2020, 10, 15).and_hms(18, 27, 11),
            chrono::NaiveDate::from_ymd(2020, 10, 16)
        ])
        .unwrap();
        Ok(())
    })
    .expect("Write excel error!");

    let mut ws = wb.create_sheet("test_sheet3 is very long and breaks the limit of 31 charcters");

    wb.write_sheet(&mut ws, |sw| {
        sw.append_row(row!["Name", "Title", "Success"]).unwrap();
        sw.append_row(row!["Mary", "Sgt Monkey", true])
    })
    .expect("Write excel error!");

    wb.close().expect("Close excel error!")
}

fn get_file_as_str_from_zip(mem_file: &Vec<u8>, file_name: &str) -> String {
    let mut archive = zip::read::ZipArchive::new(Cursor::new(mem_file)).unwrap();
    let mut styles_file = archive.by_name(file_name).unwrap();
    let mut temp_buf = vec![];
    let _ = styles_file.read_to_end(&mut temp_buf).unwrap();
    std::str::from_utf8(&temp_buf[..]).unwrap().to_string()
}

#[test]
fn creates_and_saves_an_excel_sheet() {
    let file_test = creates_and_saves_an_excel_sheet_driver(Some("test.xlsx"));
    assert!(file_test.is_none());

    let in_memory_test = creates_and_saves_an_excel_sheet_driver(None);
    assert!(in_memory_test.is_some());
}

#[test]
fn creates_file_and_checks_default_style() {
    let expected_xml = r#"<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
"#.to_string();
    let mem_file = creates_and_saves_an_excel_sheet_driver(None).unwrap();
    let result = get_file_as_str_from_zip(&mem_file, "xl/styles.xml");
    assert_eq!(expected_xml, result, "The style sheet should match!");
}

#[test]
fn creates_file_with_custom_number_format_and_checks_style() {
    let expected_xml = r##"<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
"##.to_string();
    let mut wb = excel::Workbook::create_in_memory();
    let dollar_id = wb.add_number_format("\"$\"#,##0.00".to_string());
    let weight_id = wb.add_number_format("#,##0.00\" KG\"".to_string());
    let diamond_id = wb.add_number_format("#,##0.0\"<>\"".to_string());

    let mut ws = wb.create_sheet("test_sheet");
    ws.add_column(Column { width: 20.0 });
    ws.add_column(Column { width: 20.0 });
    ws.add_column(Column { width: 20.0 });
    wb.write_sheet(&mut ws, |sw| {
        sw.append_row(row!["Cost", "Weight", "Symbol"]).unwrap();
        sw.append_row(row![(20.1, Some(dollar_id)), (50.12,  Some(weight_id)), (700.0,  Some(diamond_id))])
    })
    .expect("Write excel error!");

    let mem_file = wb.close().unwrap().unwrap();
    let result = get_file_as_str_from_zip(&mem_file, "xl/styles.xml");
    assert_eq!(expected_xml, result, "The style sheet should match!");

    let sheet = get_file_as_str_from_zip(&mem_file, "xl/worksheets/sheet1.xml");
    assert!(sheet.contains("<c r=\"A2\" s=\"3\"><v>20.1</v></c>"), "First cell should reference the 3rd index of the cellXfs list");
    assert!(sheet.contains("<c r=\"B2\" s=\"4\"><v>50.12</v></c>"), "First cell should reference the 4th index of the cellXfs list");
    assert!(sheet.contains("<c r=\"C2\" s=\"5\"><v>700</v></c>"), "First cell should reference the 5th (last) index of the cellXfs list");
}<commit_msg>Fixing test function variable name<commit_after>extern crate simple_excel_writer;
use std::io::{Cursor, Read};
use excel::*;
use simple_excel_writer as excel;

fn creates_and_saves_an_excel_sheet_driver(filename: Option<&str>) -> Option<Vec<u8>> {
    let mut wb = if let Some(name) = filename {
        excel::Workbook::create(name)
    } else {
        excel::Workbook::create_in_memory()
    };

    let mut ws = wb.create_sheet("test_sheet");
    ws.add_column(Column { width: 60.0 });
    ws.add_column(Column { width: 30.0 });
    ws.add_column(Column { width: 10.0 });
    ws.add_column(Column { width: 60.0 });
    wb.write_sheet(&mut ws, |sw| {
        sw.append_row(row!["Name", "Title", "Success", "Remark"])
            .unwrap();
        sw.append_row(row![
            "Mary",
            "Acountant",
            false,
            r#"<xml><tag>"" & 'World'</tag></xml>"#
        ])
        .unwrap();
        sw.append_row(row![
            "Mary",
            "Programmer",
            true,
            "<xml><tag>\"Hello\" & 'World'</tag></xml>"
        ])
        .unwrap();
        sw.append_row(row!["Marly", "Mary", "Success", "Success", true, 500.])
    })
    .expect("Write excel error!");

    let mut ws = wb.create_sheet("test_sheet<2");

    wb.write_sheet(&mut ws, |sw| {
        sw.append_row(row!["Name", "Title", "Success"]).unwrap();
        sw.append_row(row!["Mary", "This", true]).unwrap();

        #[cfg(feature = "chrono")]
        sw.append_row(row![
            chrono::NaiveDate::from_ymd(2020, 10, 15).and_hms(18, 27, 11),
            chrono::NaiveDate::from_ymd(2020, 10, 16)
        ])
        .unwrap();
        Ok(())
    })
    .expect("Write excel error!");

    let mut ws = wb.create_sheet("test_sheet3 is very long and breaks the limit of 31 charcters");

    wb.write_sheet(&mut ws, |sw| {
        sw.append_row(row!["Name", "Title", "Success"]).unwrap();
        sw.append_row(row!["Mary", "Sgt Monkey", true])
    })
    .expect("Write excel error!");

    wb.close().expect("Close excel error!")
}

fn get_file_as_str_from_zip(mem_file: &Vec<u8>, file_name: &str) -> String {
    let mut archive = zip::read::ZipArchive::new(Cursor::new(mem_file)).unwrap();
    let mut zip_file = archive.by_name(file_name).unwrap();
    let mut temp_buf = vec![];
    let _ = zip_file.read_to_end(&mut temp_buf).unwrap();
    std::str::from_utf8(&temp_buf[..]).unwrap().to_string()
}

#[test]
fn creates_and_saves_an_excel_sheet() {
    let file_test = creates_and_saves_an_excel_sheet_driver(Some("test.xlsx"));
    assert!(file_test.is_none());

    let in_memory_test = creates_and_saves_an_excel_sheet_driver(None);
    assert!(in_memory_test.is_some());
}

#[test]
fn creates_file_and_checks_default_style() {
    let expected_xml = r#"<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
"#.to_string();
    let mem_file = creates_and_saves_an_excel_sheet_driver(None).unwrap();
    let result = get_file_as_str_from_zip(&mem_file, "xl/styles.xml");
    assert_eq!(expected_xml, result, "The style sheet should match!");
}

#[test]
fn creates_file_with_custom_number_format_and_checks_style() {
    let expected_xml = r##"<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
"##.to_string();
    let mut wb = excel::Workbook::create_in_memory();
    let dollar_id = wb.add_number_format("\"$\"#,##0.00".to_string());
    let weight_id = wb.add_number_format("#,##0.00\" KG\"".to_string());
    let diamond_id = wb.add_number_format("#,##0.0\"<>\"".to_string());

    let mut ws = wb.create_sheet("test_sheet");
    ws.add_column(Column { width: 20.0 });
    ws.add_column(Column { width: 20.0 });
    ws.add_column(Column { width: 20.0 });
    wb.write_sheet(&mut ws, |sw| {
        sw.append_row(row!["Cost", "Weight", "Symbol"]).unwrap();
        sw.append_row(row![(20.1, Some(dollar_id)), (50.12,  Some(weight_id)), (700.0,  Some(diamond_id))])
    })
    .expect("Write excel error!");

    let mem_file = wb.close().unwrap().unwrap();
    let result = get_file_as_str_from_zip(&mem_file, "xl/styles.xml");
    assert_eq!(expected_xml, result, "The style sheet should match!");

    let sheet = get_file_as_str_from_zip(&mem_file, "xl/worksheets/sheet1.xml");
    assert!(sheet.contains("<c r=\"A2\" s=\"3\"><v>20.1</v></c>"), "First cell should reference the 3rd index of the cellXfs list");
    assert!(sheet.contains("<c r=\"B2\" s=\"4\"><v>50.12</v></c>"), "First cell should reference the 4th index of the cellXfs list");
    assert!(sheet.contains("<c r=\"C2\" s=\"5\"><v>700</v></c>"), "First cell should reference the 5th (last) index of the cellXfs list");
}<|endoftext|>


<commit_before>extern crate codespan;
extern crate codespan_reporting;
#[macro_use]
extern crate structopt;

use structopt::StructOpt;

use codespan::{CodeMap, Span};
use codespan_reporting::termcolor::StandardStream;
use codespan_reporting::{emit, ColorArg, Diagnostic, Label, Severity};

#[derive(Debug, StructOpt)]
#[structopt(name = "emit")]
pub struct Opts {
    /// Configure coloring of output
    #[structopt(
        long = "color",
        parse(try_from_str),
        default_value = "auto",
        raw(possible_values = "ColorArg::VARIANTS", case_insensitive = "true")
    )]
    pub color: ColorArg,
}

fn main() {
    let opts = Opts::from_args();
    let mut code_map = CodeMap::new();

    let source = r##"
(define test 123)
(+ test "")
()
"##;
    let file_map = code_map.add_filemap("test".into(), source.to_string());

    let str_start = file_map.byte_index(2.into(), 8.into()).unwrap();
    let error = Diagnostic::new(Severity::Error, "Unexpected type in `+` application")
        .with_label(
            Label::new_primary(Span::from_offset(str_start, 2.into()))
                .with_message("Expected integer but got string"),
        )
        .with_label(
            Label::new_secondary(Span::from_offset(str_start, 2.into()))
                .with_message("Expected integer but got string"),
        );

    let line_start = file_map.byte_index(2.into(), 0.into()).unwrap();
    let warning = Diagnostic::new(
        Severity::Warning,
        "`+` function has no effect unless its result is used",
    ).with_label(Label::new_primary(Span::from_offset(line_start, 11.into())));

    let diagnostics = [error, warning];

    let writer = StandardStream::stderr(opts.color.into());
    for diagnostic in &diagnostics {
        emit(&mut writer.lock(), &code_map, &diagnostic).unwrap();
        println!();
    }
}
<commit_msg>Update emitter example to use a diagnostic code<commit_after>extern crate codespan;
extern crate codespan_reporting;
#[macro_use]
extern crate structopt;

use structopt::StructOpt;

use codespan::{CodeMap, Span};
use codespan_reporting::termcolor::StandardStream;
use codespan_reporting::{emit, ColorArg, Diagnostic, Label, Severity};

#[derive(Debug, StructOpt)]
#[structopt(name = "emit")]
pub struct Opts {
    /// Configure coloring of output
    #[structopt(
        long = "color",
        parse(try_from_str),
        default_value = "auto",
        raw(possible_values = "ColorArg::VARIANTS", case_insensitive = "true")
    )]
    pub color: ColorArg,
}

fn main() {
    let opts = Opts::from_args();
    let mut code_map = CodeMap::new();

    let source = r##"
(define test 123)
(+ test "")
()
"##;
    let file_map = code_map.add_filemap("test".into(), source.to_string());

    let str_start = file_map.byte_index(2.into(), 8.into()).unwrap();
    let error = Diagnostic::new(Severity::Error, "Unexpected type in `+` application")
        .with_label(
            Label::new_primary(Span::from_offset(str_start, 2.into()))
                .with_message("Expected integer but got string"),
        )
        .with_label(
            Label::new_secondary(Span::from_offset(str_start, 2.into()))
                .with_message("Expected integer but got string"),
        )
        .with_code("E0001");

    let line_start = file_map.byte_index(2.into(), 0.into()).unwrap();
    let warning = Diagnostic::new(
        Severity::Warning,
        "`+` function has no effect unless its result is used",
    ).with_label(Label::new_primary(Span::from_offset(line_start, 11.into())));

    let diagnostics = [error, warning];

    let writer = StandardStream::stderr(opts.color.into());
    for diagnostic in &diagnostics {
        emit(&mut writer.lock(), &code_map, &diagnostic).unwrap();
        println!();
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use ipc_channel::ipc::IpcSharedMemory;
use piston_image::{self, DynamicImage, ImageFormat};
use webrender_traits;

#[derive(Clone, Copy, Deserialize, Eq, PartialEq, Serialize, HeapSizeOf)]
pub enum PixelFormat {
    K8,         // Luminance channel only
    KA8,        // Luminance + alpha
    RGB8,       // RGB, 8 bits per channel
    RGBA8,      // RGB + alpha, 8 bits per channel
}

#[derive(Clone, Deserialize, Serialize, HeapSizeOf)]
pub struct Image {
    pub width: u32,
    pub height: u32,
    pub format: PixelFormat,
    #[ignore_heap_size_of = "Defined in ipc-channel"]
    pub bytes: IpcSharedMemory,
    #[ignore_heap_size_of = "Defined in webrender_traits"]
    pub id: Option<webrender_traits::ImageKey>,
}

#[derive(Clone, Deserialize, Eq, PartialEq, Serialize, HeapSizeOf)]
pub struct ImageMetadata {
    pub width: u32,
    pub height: u32,
}

// FIXME: Images must not be copied every frame. Instead we should atomically
// reference count them.

// TODO(pcwalton): Speed up with SIMD, or better yet, find some way to not do this.
fn byte_swap_and_premultiply(data: &mut [u8]) {
    let length = data.len();

    for i in (0..length).step_by(4) {
        let r = data[i + 2];
        let g = data[i + 1];
        let b = data[i + 0];

        data[i + 0] = r;
        data[i + 1] = g;
        data[i + 2] = b;
    }
}

pub fn load_from_memory(buffer: &[u8]) -> Option<Image> {
    if buffer.is_empty() {
        return None;
    }

    let image_fmt_result = detect_image_format(buffer);
    match image_fmt_result {
        Err(msg) => {
            debug!("{}", msg);
            None
        }
        Ok(_) => {
            match piston_image::load_from_memory(buffer) {
                Ok(image) => {
                    let mut rgba = match image {
                        DynamicImage::ImageRgba8(rgba) => rgba,
                        image => image.to_rgba()
                    };
                    byte_swap_and_premultiply(&mut *rgba);
                    Some(Image {
                        width: rgba.width(),
                        height: rgba.height(),
                        format: PixelFormat::RGBA8,
                        bytes: IpcSharedMemory::from_bytes(&*rgba),
                        id: None,
                    })
                }
                Err(e) => {
                    debug!("Image decoding error: {:?}", e);
                    None
                }
            }
        }
    }
}

// https://developer.mozilla.org/en-US/docs/Web/HTML/Element/img
pub fn detect_image_format(buffer: &[u8]) -> Result<ImageFormat, &str> {
    if is_gif(buffer)         { Ok(ImageFormat::GIF)
    } else if is_jpeg(buffer) { Ok(ImageFormat::JPEG)
    } else if is_png(buffer)  { Ok(ImageFormat::PNG)
    } else if is_bmp(buffer)  { Ok(ImageFormat::BMP)
    } else if is_ico(buffer)  { Ok(ImageFormat::ICO)
    } else { Err("Image Format Not Supported") }
}

fn is_gif(buffer: &[u8]) -> bool {
    match buffer {
        &[b'G', b'I', b'F', b'8', n, b'a', ..] if n == b'7' || n == b'9' => true,
        _ => false
    }
}

fn is_jpeg(buffer: &[u8]) -> bool {
    buffer.starts_with(&[0xff, 0xd8, 0xff])
}

fn is_png(buffer: &[u8]) -> bool {
    buffer.starts_with(&[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A])
}

fn is_bmp(buffer: &[u8]) -> bool {
    buffer.starts_with(&[0x42, 0x4D])
}

fn is_ico(buffer: &[u8]) -> bool {
    buffer.starts_with(&[0x00, 0x00, 0x01, 0x00])
}
<commit_msg>Use documentation comments for PixelFormat.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use ipc_channel::ipc::IpcSharedMemory;
use piston_image::{self, DynamicImage, ImageFormat};
use webrender_traits;

#[derive(Clone, Copy, Deserialize, Eq, PartialEq, Serialize, HeapSizeOf)]
pub enum PixelFormat {
    /// Luminance channel only
    K8,
    /// Luminance + alpha
    KA8,
    /// RGB, 8 bits per channel
    RGB8,
    /// RGB + alpha, 8 bits per channel
    RGBA8,
}

#[derive(Clone, Deserialize, Serialize, HeapSizeOf)]
pub struct Image {
    pub width: u32,
    pub height: u32,
    pub format: PixelFormat,
    #[ignore_heap_size_of = "Defined in ipc-channel"]
    pub bytes: IpcSharedMemory,
    #[ignore_heap_size_of = "Defined in webrender_traits"]
    pub id: Option<webrender_traits::ImageKey>,
}

#[derive(Clone, Deserialize, Eq, PartialEq, Serialize, HeapSizeOf)]
pub struct ImageMetadata {
    pub width: u32,
    pub height: u32,
}

// FIXME: Images must not be copied every frame. Instead we should atomically
// reference count them.

// TODO(pcwalton): Speed up with SIMD, or better yet, find some way to not do this.
fn byte_swap_and_premultiply(data: &mut [u8]) {
    let length = data.len();

    for i in (0..length).step_by(4) {
        let r = data[i + 2];
        let g = data[i + 1];
        let b = data[i + 0];

        data[i + 0] = r;
        data[i + 1] = g;
        data[i + 2] = b;
    }
}

pub fn load_from_memory(buffer: &[u8]) -> Option<Image> {
    if buffer.is_empty() {
        return None;
    }

    let image_fmt_result = detect_image_format(buffer);
    match image_fmt_result {
        Err(msg) => {
            debug!("{}", msg);
            None
        }
        Ok(_) => {
            match piston_image::load_from_memory(buffer) {
                Ok(image) => {
                    let mut rgba = match image {
                        DynamicImage::ImageRgba8(rgba) => rgba,
                        image => image.to_rgba()
                    };
                    byte_swap_and_premultiply(&mut *rgba);
                    Some(Image {
                        width: rgba.width(),
                        height: rgba.height(),
                        format: PixelFormat::RGBA8,
                        bytes: IpcSharedMemory::from_bytes(&*rgba),
                        id: None,
                    })
                }
                Err(e) => {
                    debug!("Image decoding error: {:?}", e);
                    None
                }
            }
        }
    }
}

// https://developer.mozilla.org/en-US/docs/Web/HTML/Element/img
pub fn detect_image_format(buffer: &[u8]) -> Result<ImageFormat, &str> {
    if is_gif(buffer)         { Ok(ImageFormat::GIF)
    } else if is_jpeg(buffer) { Ok(ImageFormat::JPEG)
    } else if is_png(buffer)  { Ok(ImageFormat::PNG)
    } else if is_bmp(buffer)  { Ok(ImageFormat::BMP)
    } else if is_ico(buffer)  { Ok(ImageFormat::ICO)
    } else { Err("Image Format Not Supported") }
}

fn is_gif(buffer: &[u8]) -> bool {
    match buffer {
        &[b'G', b'I', b'F', b'8', n, b'a', ..] if n == b'7' || n == b'9' => true,
        _ => false
    }
}

fn is_jpeg(buffer: &[u8]) -> bool {
    buffer.starts_with(&[0xff, 0xd8, 0xff])
}

fn is_png(buffer: &[u8]) -> bool {
    buffer.starts_with(&[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A])
}

fn is_bmp(buffer: &[u8]) -> bool {
    buffer.starts_with(&[0x42, 0x4D])
}

fn is_ico(buffer: &[u8]) -> bool {
    buffer.starts_with(&[0x00, 0x00, 0x01, 0x00])
}
<|endoftext|>


<commit_before>// Copyright 2022 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! This module implements a lightweight and safe decoder interface over
//! `libva`. It is designed to concentrate all calls to unsafe methods in one
//! place

#![deny(missing_docs)]

mod bindings;
mod buffer;
mod buffer_type;
mod config;
mod context;
mod display;
mod generic_value;
mod image;
mod picture;
mod status;
mod surface;
mod usage_hint;

pub use bindings::constants;
pub use bindings::VAConfigAttrib;
pub use bindings::VAConfigAttribType;
pub use bindings::VAEntrypoint;
pub use bindings::VAImageFormat;
pub use bindings::VAProfile;
pub use bindings::VASurfaceAttribType;
pub use buffer::*;
pub use buffer_type::*;
pub use config::*;
pub use context::*;
pub use display::*;
pub use generic_value::*;
pub use image::*;
pub use picture::*;
pub use surface::*;
pub use usage_hint::*;

#[cfg(test)]
mod tests {
    use std::rc::Rc;

    use super::*;

    /// Returns a 32-bit CRC for the visible part of `image`, which must be in NV12 format.
    fn crc_nv12_image(image: &Image) -> u32 {
        let data = image.as_ref();
        let va_image = image.image();
        let offsets = &va_image.offsets;
        let pitches = &va_image.pitches;
        let width = va_image.width as usize;
        let height = va_image.height as usize;

        // We only support NV12 images
        assert_eq!(va_image.format.fourcc, u32::from_ne_bytes(*b"NV12"));
        // Consistency check
        assert_eq!(va_image.num_planes, 2);

        let mut hasher = crc32fast::Hasher::new();

        let offset = offsets[0] as usize;
        let pitch = pitches[0] as usize;
        let y_plane = data[offset..(offset + pitch * height)]
            .chunks(pitch)
            .map(|line| &line[0..width]);

        let offset = offsets[1] as usize;
        let pitch = pitches[1] as usize;
        let uv_plane = data[offset..(offset + pitch * ((height + 1) / 2))]
            .chunks(pitch)
            .map(|line| &line[0..width]);

        for line in y_plane.chain(uv_plane) {
            hasher.update(line);
        }

        hasher.finalize()
    }

    #[test]
    // Ignore this test by default as it requires libva-compatible hardware.
    #[ignore]
    fn libva_utils_mpeg2vldemo() {
        // Adapted from <https://github.com/intel/libva-utils/blob/master/decode/mpeg2vldemo.cpp>
        let display = Rc::new(Display::open().unwrap());

        assert!(!display.query_vendor_string().unwrap().is_empty());
        let profiles = display.query_config_profiles().unwrap();
        assert!(!profiles.is_empty());

        let profile = bindings::VAProfile::VAProfileMPEG2Main;
        let entrypoints = display.query_config_entrypoints(profile).unwrap();
        assert!(!entrypoints.is_empty());
        assert!(entrypoints
            .iter()
            .any(|e| *e == bindings::VAEntrypoint::VAEntrypointVLD));

        let format = bindings::constants::VA_RT_FORMAT_YUV420;
        let width = 16;
        let height = 16;

        let mut attrs = vec![bindings::VAConfigAttrib {
            type_: bindings::VAConfigAttribType::VAConfigAttribRTFormat,
            value: 0,
        }];

        let entrypoint = bindings::VAEntrypoint::VAEntrypointVLD;
        display
            .get_config_attributes(profile, entrypoint, &mut attrs)
            .unwrap();
        assert!(attrs[0].value != bindings::constants::VA_ATTRIB_NOT_SUPPORTED);
        assert!(attrs[0].value & bindings::constants::VA_RT_FORMAT_YUV420 != 0);

        let config = display
            .create_config(Some(attrs), profile, entrypoint)
            .unwrap();

        let mut surfaces = display
            .create_surfaces(
                format,
                None,
                width,
                height,
                Some(UsageHint::USAGE_HINT_DECODER),
                1,
            )
            .unwrap();
        let context = Rc::new(
            display
                .create_context(
                    &config,
                    width as i32,
                    (((height + 15) / 16) * 16) as i32,
                    Some(&surfaces),
                    true,
                )
                .unwrap(),
        );

        // The picture data is adapted from libva-utils at decode/mpeg2vldemo.cpp
        // Data dump of a 16x16 MPEG2 video clip,it has one I frame
        let mut mpeg2_clip: Vec<u8> = vec![
            0x00, 0x00, 0x01, 0xb3, 0x01, 0x00, 0x10, 0x13, 0xff, 0xff, 0xe0, 0x18, 0x00, 0x00,
            0x01, 0xb5, 0x14, 0x8a, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x01, 0xb8, 0x00, 0x08,
            0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x0f, 0xff, 0xf8, 0x00, 0x00, 0x01, 0xb5,
            0x8f, 0xff, 0xf3, 0x41, 0x80, 0x00, 0x00, 0x01, 0x01, 0x13, 0xe1, 0x00, 0x15, 0x81,
            0x54, 0xe0, 0x2a, 0x05, 0x43, 0x00, 0x2d, 0x60, 0x18, 0x01, 0x4e, 0x82, 0xb9, 0x58,
            0xb1, 0x83, 0x49, 0xa4, 0xa0, 0x2e, 0x05, 0x80, 0x4b, 0x7a, 0x00, 0x01, 0x38, 0x20,
            0x80, 0xe8, 0x05, 0xff, 0x60, 0x18, 0xe0, 0x1d, 0x80, 0x98, 0x01, 0xf8, 0x06, 0x00,
            0x54, 0x02, 0xc0, 0x18, 0x14, 0x03, 0xb2, 0x92, 0x80, 0xc0, 0x18, 0x94, 0x42, 0x2c,
            0xb2, 0x11, 0x64, 0xa0, 0x12, 0x5e, 0x78, 0x03, 0x3c, 0x01, 0x80, 0x0e, 0x80, 0x18,
            0x80, 0x6b, 0xca, 0x4e, 0x01, 0x0f, 0xe4, 0x32, 0xc9, 0xbf, 0x01, 0x42, 0x69, 0x43,
            0x50, 0x4b, 0x01, 0xc9, 0x45, 0x80, 0x50, 0x01, 0x38, 0x65, 0xe8, 0x01, 0x03, 0xf3,
            0xc0, 0x76, 0x00, 0xe0, 0x03, 0x20, 0x28, 0x18, 0x01, 0xa9, 0x34, 0x04, 0xc5, 0xe0,
            0x0b, 0x0b, 0x04, 0x20, 0x06, 0xc0, 0x89, 0xff, 0x60, 0x12, 0x12, 0x8a, 0x2c, 0x34,
            0x11, 0xff, 0xf6, 0xe2, 0x40, 0xc0, 0x30, 0x1b, 0x7a, 0x01, 0xa9, 0x0d, 0x00, 0xac,
            0x64,
        ];

        let picture_coding_extension =
            MPEG2PictureCodingExtension::new(0, 3, 0, 1, 0, 0, 0, 0, 0, 1, 1);
        let pic_param = PictureParameterBufferMPEG2::new(
            16,
            16,
            0xffffffff,
            0xffffffff,
            1,
            0xffff,
            &picture_coding_extension,
        );

        let pic_param = BufferType::PictureParameter(PictureParameter::MPEG2(pic_param));

        let iq_matrix = IQMatrixBufferMPEG2::new(
            1,
            1,
            0,
            0,
            [
                8, 16, 16, 19, 16, 19, 22, 22, 22, 22, 22, 22, 26, 24, 26, 27, 27, 27, 26, 26, 26,
                26, 27, 27, 27, 29, 29, 29, 34, 34, 34, 29, 29, 29, 27, 27, 29, 29, 32, 32, 34, 34,
                37, 38, 37, 35, 35, 34, 35, 38, 38, 40, 40, 40, 48, 48, 46, 46, 56, 56, 58, 69, 69,
                83,
            ],
            [
                16, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                0, 0, 0, 0, 0, 0, 0, 0, 0,
            ],
            [0; 64],
            [0; 64],
        );

        let iq_matrix = BufferType::IQMatrix(IQMatrix::MPEG2(iq_matrix));

        let slice_param = SliceParameterBufferMPEG2::new(150, 0, 0, 38, 0, 0, 2, 0);

        let slice_param = BufferType::SliceParameter(SliceParameter::MPEG2(slice_param));

        let test_data_offset = 47;
        let slice_data = BufferType::SliceData(mpeg2_clip.drain(test_data_offset..).collect());

        let buffers = vec![
            context.create_buffer(pic_param).unwrap(),
            context.create_buffer(slice_param).unwrap(),
            context.create_buffer(iq_matrix).unwrap(),
            context.create_buffer(slice_data).unwrap(),
        ];

        let mut picture = Picture::new(0, Rc::clone(&context), surfaces.remove(0));
        for buffer in buffers {
            picture.add_buffer(buffer);
        }

        // Actual client code can just chain the calls.
        let picture = picture.begin().unwrap();
        let picture = picture.render().unwrap();
        let picture = picture.end().unwrap();
        let mut picture = picture.sync().unwrap();

        // Test whether we can map the resulting surface to obtain the raw yuv
        // data
        let image_fmts = display.query_image_formats().unwrap();
        let image_fmt = image_fmts
            .into_iter()
            .find(|f| f.fourcc == bindings::constants::VA_FOURCC_NV12)
            .expect("No valid VAImageFormat found for NV12");

        let image = Image::new(&mut picture, image_fmt, width, height, false).unwrap();

        assert_eq!(crc_nv12_image(&image), 0xa5713e52);
    }
}
<commit_msg>media: libva: Also export VASurfaceID<commit_after>// Copyright 2022 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! This module implements a lightweight and safe decoder interface over
//! `libva`. It is designed to concentrate all calls to unsafe methods in one
//! place

#![deny(missing_docs)]

mod bindings;
mod buffer;
mod buffer_type;
mod config;
mod context;
mod display;
mod generic_value;
mod image;
mod picture;
mod status;
mod surface;
mod usage_hint;

pub use bindings::constants;
pub use bindings::VAConfigAttrib;
pub use bindings::VAConfigAttribType;
pub use bindings::VAEntrypoint;
pub use bindings::VAImageFormat;
pub use bindings::VAProfile;
pub use bindings::VASurfaceAttribType;
pub use bindings::VASurfaceID;
pub use buffer::*;
pub use buffer_type::*;
pub use config::*;
pub use context::*;
pub use display::*;
pub use generic_value::*;
pub use image::*;
pub use picture::*;
pub use surface::*;
pub use usage_hint::*;

#[cfg(test)]
mod tests {
    use std::rc::Rc;

    use super::*;

    /// Returns a 32-bit CRC for the visible part of `image`, which must be in NV12 format.
    fn crc_nv12_image(image: &Image) -> u32 {
        let data = image.as_ref();
        let va_image = image.image();
        let offsets = &va_image.offsets;
        let pitches = &va_image.pitches;
        let width = va_image.width as usize;
        let height = va_image.height as usize;

        // We only support NV12 images
        assert_eq!(va_image.format.fourcc, u32::from_ne_bytes(*b"NV12"));
        // Consistency check
        assert_eq!(va_image.num_planes, 2);

        let mut hasher = crc32fast::Hasher::new();

        let offset = offsets[0] as usize;
        let pitch = pitches[0] as usize;
        let y_plane = data[offset..(offset + pitch * height)]
            .chunks(pitch)
            .map(|line| &line[0..width]);

        let offset = offsets[1] as usize;
        let pitch = pitches[1] as usize;
        let uv_plane = data[offset..(offset + pitch * ((height + 1) / 2))]
            .chunks(pitch)
            .map(|line| &line[0..width]);

        for line in y_plane.chain(uv_plane) {
            hasher.update(line);
        }

        hasher.finalize()
    }

    #[test]
    // Ignore this test by default as it requires libva-compatible hardware.
    #[ignore]
    fn libva_utils_mpeg2vldemo() {
        // Adapted from <https://github.com/intel/libva-utils/blob/master/decode/mpeg2vldemo.cpp>
        let display = Rc::new(Display::open().unwrap());

        assert!(!display.query_vendor_string().unwrap().is_empty());
        let profiles = display.query_config_profiles().unwrap();
        assert!(!profiles.is_empty());

        let profile = bindings::VAProfile::VAProfileMPEG2Main;
        let entrypoints = display.query_config_entrypoints(profile).unwrap();
        assert!(!entrypoints.is_empty());
        assert!(entrypoints
            .iter()
            .any(|e| *e == bindings::VAEntrypoint::VAEntrypointVLD));

        let format = bindings::constants::VA_RT_FORMAT_YUV420;
        let width = 16;
        let height = 16;

        let mut attrs = vec![bindings::VAConfigAttrib {
            type_: bindings::VAConfigAttribType::VAConfigAttribRTFormat,
            value: 0,
        }];

        let entrypoint = bindings::VAEntrypoint::VAEntrypointVLD;
        display
            .get_config_attributes(profile, entrypoint, &mut attrs)
            .unwrap();
        assert!(attrs[0].value != bindings::constants::VA_ATTRIB_NOT_SUPPORTED);
        assert!(attrs[0].value & bindings::constants::VA_RT_FORMAT_YUV420 != 0);

        let config = display
            .create_config(Some(attrs), profile, entrypoint)
            .unwrap();

        let mut surfaces = display
            .create_surfaces(
                format,
                None,
                width,
                height,
                Some(UsageHint::USAGE_HINT_DECODER),
                1,
            )
            .unwrap();
        let context = Rc::new(
            display
                .create_context(
                    &config,
                    width as i32,
                    (((height + 15) / 16) * 16) as i32,
                    Some(&surfaces),
                    true,
                )
                .unwrap(),
        );

        // The picture data is adapted from libva-utils at decode/mpeg2vldemo.cpp
        // Data dump of a 16x16 MPEG2 video clip,it has one I frame
        let mut mpeg2_clip: Vec<u8> = vec![
            0x00, 0x00, 0x01, 0xb3, 0x01, 0x00, 0x10, 0x13, 0xff, 0xff, 0xe0, 0x18, 0x00, 0x00,
            0x01, 0xb5, 0x14, 0x8a, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x01, 0xb8, 0x00, 0x08,
            0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x0f, 0xff, 0xf8, 0x00, 0x00, 0x01, 0xb5,
            0x8f, 0xff, 0xf3, 0x41, 0x80, 0x00, 0x00, 0x01, 0x01, 0x13, 0xe1, 0x00, 0x15, 0x81,
            0x54, 0xe0, 0x2a, 0x05, 0x43, 0x00, 0x2d, 0x60, 0x18, 0x01, 0x4e, 0x82, 0xb9, 0x58,
            0xb1, 0x83, 0x49, 0xa4, 0xa0, 0x2e, 0x05, 0x80, 0x4b, 0x7a, 0x00, 0x01, 0x38, 0x20,
            0x80, 0xe8, 0x05, 0xff, 0x60, 0x18, 0xe0, 0x1d, 0x80, 0x98, 0x01, 0xf8, 0x06, 0x00,
            0x54, 0x02, 0xc0, 0x18, 0x14, 0x03, 0xb2, 0x92, 0x80, 0xc0, 0x18, 0x94, 0x42, 0x2c,
            0xb2, 0x11, 0x64, 0xa0, 0x12, 0x5e, 0x78, 0x03, 0x3c, 0x01, 0x80, 0x0e, 0x80, 0x18,
            0x80, 0x6b, 0xca, 0x4e, 0x01, 0x0f, 0xe4, 0x32, 0xc9, 0xbf, 0x01, 0x42, 0x69, 0x43,
            0x50, 0x4b, 0x01, 0xc9, 0x45, 0x80, 0x50, 0x01, 0x38, 0x65, 0xe8, 0x01, 0x03, 0xf3,
            0xc0, 0x76, 0x00, 0xe0, 0x03, 0x20, 0x28, 0x18, 0x01, 0xa9, 0x34, 0x04, 0xc5, 0xe0,
            0x0b, 0x0b, 0x04, 0x20, 0x06, 0xc0, 0x89, 0xff, 0x60, 0x12, 0x12, 0x8a, 0x2c, 0x34,
            0x11, 0xff, 0xf6, 0xe2, 0x40, 0xc0, 0x30, 0x1b, 0x7a, 0x01, 0xa9, 0x0d, 0x00, 0xac,
            0x64,
        ];

        let picture_coding_extension =
            MPEG2PictureCodingExtension::new(0, 3, 0, 1, 0, 0, 0, 0, 0, 1, 1);
        let pic_param = PictureParameterBufferMPEG2::new(
            16,
            16,
            0xffffffff,
            0xffffffff,
            1,
            0xffff,
            &picture_coding_extension,
        );

        let pic_param = BufferType::PictureParameter(PictureParameter::MPEG2(pic_param));

        let iq_matrix = IQMatrixBufferMPEG2::new(
            1,
            1,
            0,
            0,
            [
                8, 16, 16, 19, 16, 19, 22, 22, 22, 22, 22, 22, 26, 24, 26, 27, 27, 27, 26, 26, 26,
                26, 27, 27, 27, 29, 29, 29, 34, 34, 34, 29, 29, 29, 27, 27, 29, 29, 32, 32, 34, 34,
                37, 38, 37, 35, 35, 34, 35, 38, 38, 40, 40, 40, 48, 48, 46, 46, 56, 56, 58, 69, 69,
                83,
            ],
            [
                16, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                0, 0, 0, 0, 0, 0, 0, 0, 0,
            ],
            [0; 64],
            [0; 64],
        );

        let iq_matrix = BufferType::IQMatrix(IQMatrix::MPEG2(iq_matrix));

        let slice_param = SliceParameterBufferMPEG2::new(150, 0, 0, 38, 0, 0, 2, 0);

        let slice_param = BufferType::SliceParameter(SliceParameter::MPEG2(slice_param));

        let test_data_offset = 47;
        let slice_data = BufferType::SliceData(mpeg2_clip.drain(test_data_offset..).collect());

        let buffers = vec![
            context.create_buffer(pic_param).unwrap(),
            context.create_buffer(slice_param).unwrap(),
            context.create_buffer(iq_matrix).unwrap(),
            context.create_buffer(slice_data).unwrap(),
        ];

        let mut picture = Picture::new(0, Rc::clone(&context), surfaces.remove(0));
        for buffer in buffers {
            picture.add_buffer(buffer);
        }

        // Actual client code can just chain the calls.
        let picture = picture.begin().unwrap();
        let picture = picture.render().unwrap();
        let picture = picture.end().unwrap();
        let mut picture = picture.sync().unwrap();

        // Test whether we can map the resulting surface to obtain the raw yuv
        // data
        let image_fmts = display.query_image_formats().unwrap();
        let image_fmt = image_fmts
            .into_iter()
            .find(|f| f.fourcc == bindings::constants::VA_FOURCC_NV12)
            .expect("No valid VAImageFormat found for NV12");

        let image = Image::new(&mut picture, image_fmt, width, height, false).unwrap();

        assert_eq!(crc_nv12_image(&image), 0xa5713e52);
    }
}
<|endoftext|>


<commit_before>//  rpc-perf - RPC Performance Testing
//  Copyright 2015 Twitter, Inc
//
//  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 byteorder::{ByteOrder, BigEndian, WriteBytesExt};

use consts;

#[derive(Clone)]
pub struct Buffer {
    buffer: Vec<u8>,
}

impl Buffer {
    pub fn new() -> Buffer {
        let mut buffer = Vec::<u8>::new();
        buffer.resize(4, 0);

        Buffer { buffer: buffer }
    }

    /// returns the Vec<u8> from the `Buffer`
    pub fn into_vec(self) -> Vec<u8> {
        self.buffer
    }

    /// add protocol version to buffer
    pub fn protocol_header(&mut self) -> &Self {
        self.buffer.extend_from_slice(&[128, 1, 0, 1]);
        self
    }

    /// write the framed length to the buffer
    #[inline]
    pub fn frame(&mut self) -> &Self {
        let bytes = self.buffer.len() - 4;
        BigEndian::write_i32(&mut self.buffer[..4], bytes as i32);
        self
    }

    /// add method name to buffer
    #[inline]
    pub fn method_name(&mut self, method: &str) -> &Self {
        self.write_str(method)
    }

    /// add sequence id to buffer
    #[inline]
    pub fn sequence_id(&mut self, id: i32) -> &Self {
        self.write_i32(id as i32)
    }

    /// add stop sequence to buffer
    pub fn stop(&mut self) -> &Self {
        self.write_bytes(&[consts::STOP])
    }

    // write an i16 to the buffer
    #[inline]
    pub fn write_i16(&mut self, value: i16) -> &Self {
        self.buffer.write_i16::<BigEndian>(value).unwrap();
        self
    }

    // write an i32 to the buffer
    #[inline]
    pub fn write_i32(&mut self, value: i32) -> &Self {
        self.buffer.write_i32::<BigEndian>(value).unwrap();
        self
    }

    // write an i64 to the buffer
    #[inline]
    pub fn write_i64(&mut self, value: i64) -> &Self {
        self.buffer.write_i64::<BigEndian>(value).unwrap();
        self
    }

    #[inline]
    pub fn write_f64(&mut self, value: f64) -> &Self {
        self.buffer.write_f64::<BigEndian>(value).unwrap();
        self
    }

    // write a literal byte sequence to the buffer
    #[inline]
    pub fn write_bytes(&mut self, bytes: &[u8]) -> &Self {
        self.buffer.extend_from_slice(bytes);
        self
    }

    // write bool to the buffer
    #[inline]
    pub fn write_bool(&mut self, b: bool) -> &Self {
        self.buffer.extend_from_slice(&[(b as u8)]);
        self
    }

    #[inline]
    pub fn write_str(&mut self, string: &str) -> &Self {
        let string = string.as_bytes();
        self.write_i32(string.len() as i32);
        self.buffer.extend_from_slice(string);
        self
    }
}

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

    fn test_buff<F>(expected: &[u8], f: F)
        where F: FnOnce(&mut Buffer) -> ()
    {

        let mut b = Buffer::new();
        f(&mut b);
        assert_eq!(b.into_vec().as_slice(), expected);
    }

    #[test]
    fn into_vec() {
        test_buff(&[0, 0, 0, 0], |_| {});
    }

    #[test]
    fn test_protocol_header() {
        test_buff(&[0, 0, 0, 0, 128, 1, 0, 1], |b| {
            b.protocol_header();
        });
    }

    #[test]
    fn test_sequence_id() {
        test_buff(&[0, 0, 0, 4, 0, 0, 0, 0], |b| {
            b.sequence_id(0_i32);
            b.frame();
        });
    }

    #[test]
    fn test_method_name() {
        test_buff(&[0, 0, 0, 0, 0, 0, 0, 4, 112, 105, 110, 103], |b| {
            b.method_name("ping");
        });
    }

    #[test]
    fn test_stop() {
        test_buff(&vec![0, 0, 0, 0, 0], |b| {
            b.stop();
        });
    }
}
<commit_msg>impl Default for Buffer<commit_after>//  rpc-perf - RPC Performance Testing
//  Copyright 2015 Twitter, Inc
//
//  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 byteorder::{ByteOrder, BigEndian, WriteBytesExt};

use consts;

#[derive(Clone)]
pub struct Buffer {
    buffer: Vec<u8>,
}

impl Default for Buffer {
    fn default() -> Buffer {
        let mut buffer = Vec::<u8>::new();
        buffer.resize(4, 0);

        Buffer { buffer: buffer }
    }
}

impl Buffer {
    pub fn new() -> Buffer {
        Buffer::default()
    }

    /// returns the Vec<u8> from the `Buffer`
    pub fn into_vec(self) -> Vec<u8> {
        self.buffer
    }

    /// add protocol version to buffer
    pub fn protocol_header(&mut self) -> &Self {
        self.buffer.extend_from_slice(&[128, 1, 0, 1]);
        self
    }

    /// write the framed length to the buffer
    #[inline]
    pub fn frame(&mut self) -> &Self {
        let bytes = self.buffer.len() - 4;
        BigEndian::write_i32(&mut self.buffer[..4], bytes as i32);
        self
    }

    /// add method name to buffer
    #[inline]
    pub fn method_name(&mut self, method: &str) -> &Self {
        self.write_str(method)
    }

    /// add sequence id to buffer
    #[inline]
    pub fn sequence_id(&mut self, id: i32) -> &Self {
        self.write_i32(id as i32)
    }

    /// add stop sequence to buffer
    pub fn stop(&mut self) -> &Self {
        self.write_bytes(&[consts::STOP])
    }

    // write an i16 to the buffer
    #[inline]
    pub fn write_i16(&mut self, value: i16) -> &Self {
        self.buffer.write_i16::<BigEndian>(value).unwrap();
        self
    }

    // write an i32 to the buffer
    #[inline]
    pub fn write_i32(&mut self, value: i32) -> &Self {
        self.buffer.write_i32::<BigEndian>(value).unwrap();
        self
    }

    // write an i64 to the buffer
    #[inline]
    pub fn write_i64(&mut self, value: i64) -> &Self {
        self.buffer.write_i64::<BigEndian>(value).unwrap();
        self
    }

    #[inline]
    pub fn write_f64(&mut self, value: f64) -> &Self {
        self.buffer.write_f64::<BigEndian>(value).unwrap();
        self
    }

    // write a literal byte sequence to the buffer
    #[inline]
    pub fn write_bytes(&mut self, bytes: &[u8]) -> &Self {
        self.buffer.extend_from_slice(bytes);
        self
    }

    // write bool to the buffer
    #[inline]
    pub fn write_bool(&mut self, b: bool) -> &Self {
        self.buffer.extend_from_slice(&[(b as u8)]);
        self
    }

    #[inline]
    pub fn write_str(&mut self, string: &str) -> &Self {
        let string = string.as_bytes();
        self.write_i32(string.len() as i32);
        self.buffer.extend_from_slice(string);
        self
    }
}

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

    fn test_buff<F>(expected: &[u8], f: F)
        where F: FnOnce(&mut Buffer) -> ()
    {

        let mut b = Buffer::new();
        f(&mut b);
        assert_eq!(b.into_vec().as_slice(), expected);
    }

    #[test]
    fn into_vec() {
        test_buff(&[0, 0, 0, 0], |_| {});
    }

    #[test]
    fn test_protocol_header() {
        test_buff(&[0, 0, 0, 0, 128, 1, 0, 1], |b| {
            b.protocol_header();
        });
    }

    #[test]
    fn test_sequence_id() {
        test_buff(&[0, 0, 0, 4, 0, 0, 0, 0], |b| {
            b.sequence_id(0_i32);
            b.frame();
        });
    }

    #[test]
    fn test_method_name() {
        test_buff(&[0, 0, 0, 0, 0, 0, 0, 4, 112, 105, 110, 103], |b| {
            b.method_name("ping");
        });
    }

    #[test]
    fn test_stop() {
        test_buff(&vec![0, 0, 0, 0, 0], |b| {
            b.stop();
        });
    }
}
<|endoftext|>