<commit_before>//! Safe wrappers around Erlang binaries.
//!
//! Rustler provides two binary types: [`Binary`] and [`OwnedBinary`].  Both
//! represent a contiguous region `u8`s, and they both use the Erlang allocator. The
//! primary difference between the two is their ownership semantics.
//!
//! The _owned_ in `OwnedBinary` refers to the fact that it owns the binary it
//! wraps. The _owner_ of an `OwnedBinary` is free to modify its contents. Ownership
//! lasts until it is dropped or consumed by converting it into a regular
//! `Binary`. An `OwnedBinary` cannot be copied or cloned and is thus always moved.
//!
//! The `Binary` type is an immutable shared-reference to a binary. `Binary`s are
//! cheap to copy: all copies of a `Binary` point to the original `Binary`'s
//! data. Additionally, a `Binary`'s lifetime is tied to that of the NIF's [`Env`],
//! preventing outstanding references to the data after a NIF returns.
//!
//! # Examples
//!
//! Constructing an `OwnedBinary`:
//!
//! ```no_run
//! # use rustler::OwnedBinary;
//! {
//!     let mut bin = OwnedBinary::new(5).expect("allocation failed");
//!     bin.as_mut_slice().copy_from_slice("hello".as_bytes());
//! } // <- `bin` is dropped here
//! ```
//!
//! The following NIF takes a binary as its only parameter and returns a new binary
//! where each element is exclusive-or'ed with a constant:
//!
//! ```no_run
//! # use rustler::{Env, OwnedBinary, Binary, NifResult, Error};
//! #[rustler::nif]
//! fn xor_example<'a>(env: Env<'a>, bin: Binary<'a>) -> NifResult<Binary<'a>> {
//!     let mut owned: OwnedBinary = bin.to_owned().ok_or(Error::Term(Box::new("no mem")))?;
//!     for byte in owned.as_mut_slice() {
//!         *byte ^= 0xAA;
//!     }
//!
//!     // Ownership of `owned`'s data is transferred to `env` on the
//!     // following line, so no additional heap allocations are incurred.
//!     Ok(Binary::from_owned(owned, env))
//! }
//! ```
//!
//! The contents of a newly-allocated `OwnedBinary` is not initialized to any
//! particular value. If your usage of the binary requires the it's data to be
//! zeroed, for example, then you must explicit zero it. In this example, we
//! manually zeroize the binary before passing it as slice to a third party
//! function.
//!
//! ```no_run
//! # fn some_third_party_api(buf: &mut [u8]) {
//! #     for elem in buf {
//! #         if *elem == 0 { *elem = 1 } else { panic!("Not a zero!") }
//! #     }
//! # }
//! # use rustler::{Env, OwnedBinary, Binary, NifResult, Error};
//! #[rustler::nif]
//! fn wrapper_for_some_<'a>(env: Env<'a>) -> NifResult<Binary<'a>> {
//!     let mut owned = OwnedBinary::new(100).ok_or(Error::Term(Box::new("no mem")))?;
//!     for byte in owned.as_mut_slice() {
//!         *byte = 0;
//!     }
//!
//!     // Some third party API which requires the slice to be all zeros on entry.
//!     some_third_party_api(owned.as_mut_slice());
//!
//!     // The imaginary API call presumedly filled in our binary with meaningful
//!     // data, so let's return it.
//!     Ok(Binary::from_owned(owned, env))
//! }
//!
//! ```
//!
//! [`Binary`]: struct.Binary.html
//! [`Env`]: ../../env/struct.Env.html
//! [`OwnedBinary`]: struct.OwnedBinary.html

use crate::{
    wrapper::binary::{alloc, realloc, ErlNifBinary},
    Decoder, Encoder, Env, Error, NifResult, Term,
};
use std::{
    borrow::{Borrow, BorrowMut},
    io::Write,
    mem::MaybeUninit,
    ops::{Deref, DerefMut},
};

/// An mutable smart-pointer to an Erlang binary.
///
/// See [module-level doc](index.html) for more information.
pub struct OwnedBinary(ErlNifBinary);

impl<'a> OwnedBinary {
    pub unsafe fn from_raw(inner: ErlNifBinary) -> OwnedBinary {
        OwnedBinary(inner)
    }

    /// Allocates a new `OwnedBinary` with size `size`.
    ///
    /// Memory is not initialized. If uninitialized memory is undesirable, set it
    /// manually.
    ///
    /// # Errors
    ///
    /// If allocation fails, `None` is returned.
    pub fn new(size: usize) -> Option<OwnedBinary> {
        unsafe { alloc(size) }.map(OwnedBinary)
    }

    /// Copies `src`'s data into a new `OwnedBinary`.
    ///
    /// # Errors
    ///
    /// If allocation fails, `None` is returned.
    pub fn from_unowned(src: &Binary) -> Option<OwnedBinary> {
        OwnedBinary::new(src.len()).map(|mut b| {
            b.as_mut_slice().copy_from_slice(src);
            b
        })
    }

    /// Attempts to reallocate `self` with the new size.
    ///
    /// Memory outside the range of the original binary will not be initialized. If
    /// uninitialized memory is undesirable, set it manually.
    ///
    /// # Errors
    ///
    /// If reallocation fails, `false` is returned. Data remains intact on error.
    #[must_use]
    pub fn realloc(&mut self, size: usize) -> bool {
        unsafe { realloc(&mut self.0, size) }
    }

    /// Attempts to reallocate `self` with the new size.
    ///
    /// If reallocation fails, it will perform a copy instead.
    ///
    /// Memory outside the range of the original binary will not be initialized. If
    /// uninitialized memory is undesirable, set it manually.
    pub fn realloc_or_copy(&mut self, size: usize) {
        if !self.realloc(size) {
            let mut new = OwnedBinary::new(size).unwrap();
            if let Ok(num_written) = new.as_mut_slice().write(self.as_slice()) {
                if !(num_written == self.len() || num_written == new.len()) {
                    panic!("Could not copy binary");
                }
                ::std::mem::swap(&mut self.0, &mut new.0);
            } else {
                panic!("Could not copy binary");
            }
        }
    }

    /// Extracts a slice containing the entire binary.
    pub fn as_slice(&self) -> &[u8] {
        unsafe { ::std::slice::from_raw_parts(self.0.data, self.0.size) }
    }

    /// Extracts a mutable slice of the entire binary.
    pub fn as_mut_slice(&mut self) -> &mut [u8] {
        unsafe { ::std::slice::from_raw_parts_mut(self.0.data, self.0.size) }
    }

    /// Consumes `self` and returns an immutable `Binary`.
    ///
    /// This method is the mirror of [`Binary::from_owned`], and they can be used
    /// interchangeably.
    ///
    /// [`Binary::from_owned`]: struct.Binary.html#method.from_owned
    pub fn release(self, env: Env) -> Binary {
        Binary::from_owned(self, env)
    }
}

impl Borrow<[u8]> for OwnedBinary {
    fn borrow(&self) -> &[u8] {
        self.as_slice()
    }
}
impl BorrowMut<[u8]> for OwnedBinary {
    fn borrow_mut(&mut self) -> &mut [u8] {
        self.as_mut_slice()
    }
}
impl Deref for OwnedBinary {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        self.as_slice()
    }
}
impl DerefMut for OwnedBinary {
    fn deref_mut(&mut self) -> &mut [u8] {
        self.as_mut_slice()
    }
}

impl Drop for OwnedBinary {
    fn drop(&mut self) {
        unsafe { rustler_sys::enif_release_binary(&mut self.0) };
    }
}

unsafe impl Send for OwnedBinary {}

/// An immutable smart-pointer to an Erlang binary.
///
/// See [module-level doc](index.html) for more information.
#[derive(Copy, Clone)]
pub struct Binary<'a> {
    inner: ErlNifBinary,
    term: Term<'a>,
}

impl<'a> Binary<'a> {
    /// Consumes `owned` and returns an immutable `Binary`.
    pub fn from_owned(owned: OwnedBinary, env: Env<'a>) -> Self {
        // We are transferring ownership of `owned`'s data to the
        // environment. Therefore, we need to prevent `owned`'s destructor being
        // called at the end of this scope. The least error-prone solution (compared
        // to `mem::forget()`) is to wrap `owned` in a `ManuallyDrop` and EXPLICITLY
        // NOT CALL `ManuallyDrop::drop()`.
        let mut owned = std::mem::ManuallyDrop::new(owned);
        let term = unsafe {
            Term::new(
                env,
                rustler_sys::enif_make_binary(env.as_c_arg(), &mut owned.0),
            )
        };
        Binary {
            inner: owned.0,
            term,
        }
    }

    /// Copies `self`'s data into a new `OwnedBinary`.
    ///
    /// # Errors
    ///
    /// If allocation fails, an error will be returned.
    #[allow(clippy::wrong_self_convention)]
    pub fn to_owned(&self) -> Option<OwnedBinary> {
        OwnedBinary::from_unowned(self)
    }

    /// Creates a `Binary` from `term`.
    ///
    /// # Errors
    ///
    /// If `term` is not a binary, an error will be returned.
    pub fn from_term(term: Term<'a>) -> Result<Self, Error> {
        let mut binary = MaybeUninit::uninit();
        if unsafe {
            rustler_sys::enif_inspect_binary(
                term.get_env().as_c_arg(),
                term.as_c_arg(),
                binary.as_mut_ptr(),
            )
        } == 0
        {
            return Err(Error::BadArg);
        }
        Ok(Binary {
            inner: unsafe { binary.assume_init() },
            term,
        })
    }

    /// Creates a `Binary` from `term`.
    ///
    /// # Errors
    ///
    /// If `term` is not an `iolist`, an error will be returned.
    pub fn from_iolist(term: Term<'a>) -> Result<Self, Error> {
        let mut binary = MaybeUninit::uninit();
        if unsafe {
            rustler_sys::enif_inspect_iolist_as_binary(
                term.get_env().as_c_arg(),
                term.as_c_arg(),
                binary.as_mut_ptr(),
            )
        } == 0
        {
            return Err(Error::BadArg);
        }
        Ok(Binary {
            inner: unsafe { binary.assume_init() },
            term,
        })
    }

    /// Returns an Erlang term representation of `self`.
    #[allow(clippy::wrong_self_convention)]
    pub fn to_term<'b>(&self, env: Env<'b>) -> Term<'b> {
        self.term.in_env(env)
    }

    /// Extracts a slice containing the entire binary.
    pub fn as_slice(&self) -> &'a [u8] {
        unsafe { ::std::slice::from_raw_parts(self.inner.data, self.inner.size) }
    }

    /// Returns a new view into the same binary.
    ///
    /// This method is analogous to subslicing (e.g. `some_data[offset..length]`) in
    /// that it does not copy nor allocate data.
    ///
    /// # Errors
    ///
    /// If `offset + length` is out of bounds, an error will be returned.
    pub fn make_subbinary(&self, offset: usize, length: usize) -> NifResult<Binary<'a>> {
        let min_len = length.checked_add(offset);
        if min_len.ok_or(Error::BadArg)? > self.inner.size {
            return Err(Error::BadArg);
        }

        let raw_term = unsafe {
            rustler_sys::enif_make_sub_binary(
                self.term.get_env().as_c_arg(),
                self.term.as_c_arg(),
                offset,
                length,
            )
        };
        let term = unsafe { Term::new(self.term.get_env(), raw_term) };
        // This should never fail, as we are always passing in a binary term.
        Ok(Binary::from_term(term).ok().unwrap())
    }
}

impl<'a> Borrow<[u8]> for Binary<'a> {
    fn borrow(&self) -> &[u8] {
        self.as_slice()
    }
}
impl<'a> Deref for Binary<'a> {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        self.as_slice()
    }
}

impl<'a> Decoder<'a> for Binary<'a> {
    fn decode(term: Term<'a>) -> Result<Self, Error> {
        Binary::from_term(term)
    }
}
impl<'a> Encoder for Binary<'a> {
    fn encode<'b>(&self, env: Env<'b>) -> Term<'b> {
        self.to_term(env)
    }
}

/// ## Binary terms
impl<'a> Term<'a> {
    pub fn into_binary(self) -> NifResult<Binary<'a>> {
        Binary::from_term(self)
    }
}
<commit_msg>Implement Hash and Equality for binary types<commit_after>//! Safe wrappers around Erlang binaries.
//!
//! Rustler provides two binary types: [`Binary`] and [`OwnedBinary`].  Both
//! represent a contiguous region `u8`s, and they both use the Erlang allocator. The
//! primary difference between the two is their ownership semantics.
//!
//! The _owned_ in `OwnedBinary` refers to the fact that it owns the binary it
//! wraps. The _owner_ of an `OwnedBinary` is free to modify its contents. Ownership
//! lasts until it is dropped or consumed by converting it into a regular
//! `Binary`. An `OwnedBinary` cannot be copied or cloned and is thus always moved.
//!
//! The `Binary` type is an immutable shared-reference to a binary. `Binary`s are
//! cheap to copy: all copies of a `Binary` point to the original `Binary`'s
//! data. Additionally, a `Binary`'s lifetime is tied to that of the NIF's [`Env`],
//! preventing outstanding references to the data after a NIF returns.
//!
//! # Examples
//!
//! Constructing an `OwnedBinary`:
//!
//! ```no_run
//! # use rustler::OwnedBinary;
//! {
//!     let mut bin = OwnedBinary::new(5).expect("allocation failed");
//!     bin.as_mut_slice().copy_from_slice("hello".as_bytes());
//! } // <- `bin` is dropped here
//! ```
//!
//! The following NIF takes a binary as its only parameter and returns a new binary
//! where each element is exclusive-or'ed with a constant:
//!
//! ```no_run
//! # use rustler::{Env, OwnedBinary, Binary, NifResult, Error};
//! #[rustler::nif]
//! fn xor_example<'a>(env: Env<'a>, bin: Binary<'a>) -> NifResult<Binary<'a>> {
//!     let mut owned: OwnedBinary = bin.to_owned().ok_or(Error::Term(Box::new("no mem")))?;
//!     for byte in owned.as_mut_slice() {
//!         *byte ^= 0xAA;
//!     }
//!
//!     // Ownership of `owned`'s data is transferred to `env` on the
//!     // following line, so no additional heap allocations are incurred.
//!     Ok(Binary::from_owned(owned, env))
//! }
//! ```
//!
//! The contents of a newly-allocated `OwnedBinary` is not initialized to any
//! particular value. If your usage of the binary requires the it's data to be
//! zeroed, for example, then you must explicit zero it. In this example, we
//! manually zeroize the binary before passing it as slice to a third party
//! function.
//!
//! ```no_run
//! # fn some_third_party_api(buf: &mut [u8]) {
//! #     for elem in buf {
//! #         if *elem == 0 { *elem = 1 } else { panic!("Not a zero!") }
//! #     }
//! # }
//! # use rustler::{Env, OwnedBinary, Binary, NifResult, Error};
//! #[rustler::nif]
//! fn wrapper_for_some_<'a>(env: Env<'a>) -> NifResult<Binary<'a>> {
//!     let mut owned = OwnedBinary::new(100).ok_or(Error::Term(Box::new("no mem")))?;
//!     for byte in owned.as_mut_slice() {
//!         *byte = 0;
//!     }
//!
//!     // Some third party API which requires the slice to be all zeros on entry.
//!     some_third_party_api(owned.as_mut_slice());
//!
//!     // The imaginary API call presumedly filled in our binary with meaningful
//!     // data, so let's return it.
//!     Ok(Binary::from_owned(owned, env))
//! }
//!
//! ```
//!
//! [`Binary`]: struct.Binary.html
//! [`Env`]: ../../env/struct.Env.html
//! [`OwnedBinary`]: struct.OwnedBinary.html

use crate::{
    wrapper::binary::{alloc, realloc, ErlNifBinary},
    Decoder, Encoder, Env, Error, NifResult, Term,
};
use std::{
    borrow::{Borrow, BorrowMut},
    hash::{Hash, Hasher},
    io::Write,
    mem::MaybeUninit,
    ops::{Deref, DerefMut},
};

/// An mutable smart-pointer to an Erlang binary.
///
/// See [module-level doc](index.html) for more information.
pub struct OwnedBinary(ErlNifBinary);

impl<'a> OwnedBinary {
    pub unsafe fn from_raw(inner: ErlNifBinary) -> OwnedBinary {
        OwnedBinary(inner)
    }

    /// Allocates a new `OwnedBinary` with size `size`.
    ///
    /// Memory is not initialized. If uninitialized memory is undesirable, set it
    /// manually.
    ///
    /// # Errors
    ///
    /// If allocation fails, `None` is returned.
    pub fn new(size: usize) -> Option<OwnedBinary> {
        unsafe { alloc(size) }.map(OwnedBinary)
    }

    /// Copies `src`'s data into a new `OwnedBinary`.
    ///
    /// # Errors
    ///
    /// If allocation fails, `None` is returned.
    pub fn from_unowned(src: &Binary) -> Option<OwnedBinary> {
        OwnedBinary::new(src.len()).map(|mut b| {
            b.as_mut_slice().copy_from_slice(src);
            b
        })
    }

    /// Attempts to reallocate `self` with the new size.
    ///
    /// Memory outside the range of the original binary will not be initialized. If
    /// uninitialized memory is undesirable, set it manually.
    ///
    /// # Errors
    ///
    /// If reallocation fails, `false` is returned. Data remains intact on error.
    #[must_use]
    pub fn realloc(&mut self, size: usize) -> bool {
        unsafe { realloc(&mut self.0, size) }
    }

    /// Attempts to reallocate `self` with the new size.
    ///
    /// If reallocation fails, it will perform a copy instead.
    ///
    /// Memory outside the range of the original binary will not be initialized. If
    /// uninitialized memory is undesirable, set it manually.
    pub fn realloc_or_copy(&mut self, size: usize) {
        if !self.realloc(size) {
            let mut new = OwnedBinary::new(size).unwrap();
            if let Ok(num_written) = new.as_mut_slice().write(self.as_slice()) {
                if !(num_written == self.len() || num_written == new.len()) {
                    panic!("Could not copy binary");
                }
                ::std::mem::swap(&mut self.0, &mut new.0);
            } else {
                panic!("Could not copy binary");
            }
        }
    }

    /// Extracts a slice containing the entire binary.
    pub fn as_slice(&self) -> &[u8] {
        unsafe { ::std::slice::from_raw_parts(self.0.data, self.0.size) }
    }

    /// Extracts a mutable slice of the entire binary.
    pub fn as_mut_slice(&mut self) -> &mut [u8] {
        unsafe { ::std::slice::from_raw_parts_mut(self.0.data, self.0.size) }
    }

    /// Consumes `self` and returns an immutable `Binary`.
    ///
    /// This method is the mirror of [`Binary::from_owned`], and they can be used
    /// interchangeably.
    ///
    /// [`Binary::from_owned`]: struct.Binary.html#method.from_owned
    pub fn release(self, env: Env) -> Binary {
        Binary::from_owned(self, env)
    }
}

impl Borrow<[u8]> for OwnedBinary {
    fn borrow(&self) -> &[u8] {
        self.as_slice()
    }
}
impl BorrowMut<[u8]> for OwnedBinary {
    fn borrow_mut(&mut self) -> &mut [u8] {
        self.as_mut_slice()
    }
}
impl Deref for OwnedBinary {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        self.as_slice()
    }
}
impl DerefMut for OwnedBinary {
    fn deref_mut(&mut self) -> &mut [u8] {
        self.as_mut_slice()
    }
}
impl Hash for OwnedBinary {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.as_slice().hash(state);
    }
}
impl PartialEq for OwnedBinary {
    fn eq(&self, other: &Self) -> bool {
        self.as_slice() == other.as_slice()
    }
}
impl Eq for OwnedBinary {}
impl PartialEq<Binary<'_>> for OwnedBinary {
    fn eq(&self, other: &Binary) -> bool {
        self.as_slice() == other.as_slice()
    }
}

impl Drop for OwnedBinary {
    fn drop(&mut self) {
        unsafe { rustler_sys::enif_release_binary(&mut self.0) };
    }
}

unsafe impl Send for OwnedBinary {}

/// An immutable smart-pointer to an Erlang binary.
///
/// See [module-level doc](index.html) for more information.
#[derive(Copy, Clone)]
pub struct Binary<'a> {
    inner: ErlNifBinary,
    term: Term<'a>,
}

impl<'a> Binary<'a> {
    /// Consumes `owned` and returns an immutable `Binary`.
    pub fn from_owned(owned: OwnedBinary, env: Env<'a>) -> Self {
        // We are transferring ownership of `owned`'s data to the
        // environment. Therefore, we need to prevent `owned`'s destructor being
        // called at the end of this scope. The least error-prone solution (compared
        // to `mem::forget()`) is to wrap `owned` in a `ManuallyDrop` and EXPLICITLY
        // NOT CALL `ManuallyDrop::drop()`.
        let mut owned = std::mem::ManuallyDrop::new(owned);
        let term = unsafe {
            Term::new(
                env,
                rustler_sys::enif_make_binary(env.as_c_arg(), &mut owned.0),
            )
        };
        Binary {
            inner: owned.0,
            term,
        }
    }

    /// Copies `self`'s data into a new `OwnedBinary`.
    ///
    /// # Errors
    ///
    /// If allocation fails, an error will be returned.
    #[allow(clippy::wrong_self_convention)]
    pub fn to_owned(&self) -> Option<OwnedBinary> {
        OwnedBinary::from_unowned(self)
    }

    /// Creates a `Binary` from `term`.
    ///
    /// # Errors
    ///
    /// If `term` is not a binary, an error will be returned.
    pub fn from_term(term: Term<'a>) -> Result<Self, Error> {
        let mut binary = MaybeUninit::uninit();
        if unsafe {
            rustler_sys::enif_inspect_binary(
                term.get_env().as_c_arg(),
                term.as_c_arg(),
                binary.as_mut_ptr(),
            )
        } == 0
        {
            return Err(Error::BadArg);
        }
        Ok(Binary {
            inner: unsafe { binary.assume_init() },
            term,
        })
    }

    /// Creates a `Binary` from `term`.
    ///
    /// # Errors
    ///
    /// If `term` is not an `iolist`, an error will be returned.
    pub fn from_iolist(term: Term<'a>) -> Result<Self, Error> {
        let mut binary = MaybeUninit::uninit();
        if unsafe {
            rustler_sys::enif_inspect_iolist_as_binary(
                term.get_env().as_c_arg(),
                term.as_c_arg(),
                binary.as_mut_ptr(),
            )
        } == 0
        {
            return Err(Error::BadArg);
        }
        Ok(Binary {
            inner: unsafe { binary.assume_init() },
            term,
        })
    }

    /// Returns an Erlang term representation of `self`.
    #[allow(clippy::wrong_self_convention)]
    pub fn to_term<'b>(&self, env: Env<'b>) -> Term<'b> {
        self.term.in_env(env)
    }

    /// Extracts a slice containing the entire binary.
    pub fn as_slice(&self) -> &'a [u8] {
        unsafe { ::std::slice::from_raw_parts(self.inner.data, self.inner.size) }
    }

    /// Returns a new view into the same binary.
    ///
    /// This method is analogous to subslicing (e.g. `some_data[offset..length]`) in
    /// that it does not copy nor allocate data.
    ///
    /// # Errors
    ///
    /// If `offset + length` is out of bounds, an error will be returned.
    pub fn make_subbinary(&self, offset: usize, length: usize) -> NifResult<Binary<'a>> {
        let min_len = length.checked_add(offset);
        if min_len.ok_or(Error::BadArg)? > self.inner.size {
            return Err(Error::BadArg);
        }

        let raw_term = unsafe {
            rustler_sys::enif_make_sub_binary(
                self.term.get_env().as_c_arg(),
                self.term.as_c_arg(),
                offset,
                length,
            )
        };
        let term = unsafe { Term::new(self.term.get_env(), raw_term) };
        // This should never fail, as we are always passing in a binary term.
        Ok(Binary::from_term(term).ok().unwrap())
    }
}

impl<'a> Borrow<[u8]> for Binary<'a> {
    fn borrow(&self) -> &[u8] {
        self.as_slice()
    }
}
impl<'a> Deref for Binary<'a> {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        self.as_slice()
    }
}

impl<'a> Decoder<'a> for Binary<'a> {
    fn decode(term: Term<'a>) -> Result<Self, Error> {
        Binary::from_term(term)
    }
}
impl<'a> Encoder for Binary<'a> {
    fn encode<'b>(&self, env: Env<'b>) -> Term<'b> {
        self.to_term(env)
    }
}
impl Hash for Binary<'_> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.as_slice().hash(state);
    }
}
impl PartialEq for Binary<'_> {
    fn eq(&self, other: &Self) -> bool {
        self.as_slice() == other.as_slice()
    }
}
impl Eq for Binary<'_> {}
impl PartialEq<OwnedBinary> for Binary<'_> {
    fn eq(&self, other: &OwnedBinary) -> bool {
        self.as_slice() == other.as_slice()
    }
}

/// ## Binary terms
impl<'a> Term<'a> {
    pub fn into_binary(self) -> NifResult<Binary<'a>> {
        Binary::from_term(self)
    }
}
<|endoftext|>


<commit_before>extern crate glsl_to_spirv;

#[test]
fn test1() {
    let shader = r#"
#version 330

out vec4 f_color;

void main() {
    f_color = vec4(1.0);
}
"#;

    glsl_to_spirv::compile(Some((shader, glsl_to_spirv::ShaderType::Fragment))).unwrap();
}

#[test]
fn test2() {
    let vertex = r#"
#version 330

in vec2 i_position;
out vec2 v_texcoords;

void main() {
    v_texcoords = i_position;
    gl_Position = vec4(i_position, 0.0, 1.0);
}
"#;

    let fragment = r#"
#version 330

out vec4 f_color;

void main() {
    f_color = vec4(1.0);
}
"#;

    glsl_to_spirv::compile([
        (vertex, glsl_to_spirv::ShaderType::Vertex),
        (fragment, glsl_to_spirv::ShaderType::Fragment)
    ].iter().cloned()).unwrap();
}
<commit_msg>Fix glsl-to-spirv tests<commit_after>extern crate glsl_to_spirv;

#[test]
fn test1() {
    let shader = r#"
#version 330

out vec4 f_color;

void main() {
    f_color = vec4(1.0);
}
"#;

    glsl_to_spirv::compile(shader, glsl_to_spirv::ShaderType::Fragment).unwrap();
}
<|endoftext|>


<commit_before>#![crate_name = "wc"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Boden Garman <bpgarman@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

#![feature(macro_rules)]

extern crate getopts;
extern crate libc;

use std::str::from_utf8;
use std::io::{print, stdin_raw, File, BufferedReader};
use StdResult = std::result::Result;
use getopts::Matches;

#[path = "../common/util.rs"]
mod util;

struct Result {
    filename: String,
    bytes: uint,
    chars: uint,
    lines: uint,
    words: uint,
    max_line_length: uint,
}

static NAME: &'static str = "wc";

pub fn uumain(args: Vec<String>) -> int {
    let program = args.get(0).clone();
    let opts = [
        getopts::optflag("c", "bytes", "print the byte counts"),
        getopts::optflag("m", "chars", "print the character counts"),
        getopts::optflag("l", "lines", "print the newline counts"),
        getopts::optflag("L", "max-line-length", "print the length of the longest line"),
        getopts::optflag("w", "words", "print the word counts"),
        getopts::optflag("h", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(f) => {
            crash!(1, "Invalid options\n{}", f)
        }
    };

    if matches.opt_present("help") {
        println!("Usage:");
        println!("  {0:s} [OPTION]... [FILE]...", program);
        println!("");
        print(getopts::usage("Print newline, word and byte counts for each FILE", opts).as_slice());
        println!("");
        println!("With no FILE, or when FILE is -, read standard input.");
        return 0;
    }

    if matches.opt_present("version") {
        println!("wc 1.0.0");
        return 0;
    }

    let mut files = matches.free.clone();
    if files.is_empty() {
        files = vec!("-".to_string());
    }

    match wc(files, &matches) {
        Ok(()) => ( /* pass */ ),
        Err(e) => return e
    }

    0
}

static CR: u8 = '\r' as u8;
static LF: u8 = '\n' as u8;
static SPACE: u8 = ' ' as u8;
static TAB: u8 = '\t' as u8;
static SYN: u8 = 0x16 as u8;
static FF: u8 = 0x0C as u8;

#[inline(always)]
fn is_word_seperator(byte: u8) -> bool {
    byte == SPACE || byte == TAB || byte == CR || byte == SYN || byte == FF
}

pub fn wc(files: Vec<String>, matches: &Matches) -> StdResult<(), int> {
    let mut total_line_count: uint = 0;
    let mut total_word_count: uint = 0;
    let mut total_char_count: uint = 0;
    let mut total_byte_count: uint = 0;
    let mut total_longest_line_length: uint = 0;

    let mut results = vec!();
    let mut max_str_len: uint = 0;

    for path in files.iter() {
        let mut reader = match open(path.to_string()) {
            Ok(f) => f,
            Err(e) => { return Err(e); }
        };

        let mut line_count: uint = 0;
        let mut word_count: uint = 0;
        let mut byte_count: uint = 0;
        let mut char_count: uint = 0;
        let mut current_char_count: uint = 0;
        let mut longest_line_length: uint = 0;

        loop {
            // reading from a TTY seems to raise a condition on, rather than return Some(0) like a file.
            // hence the option wrapped in a result here
            match reader.read_until(LF) {
                Ok(raw_line) => {
                    // GNU 'wc' only counts lines that end in LF as lines
                    if raw_line.iter().last().unwrap() == &LF {
                        line_count += 1;
                    }

                    byte_count += raw_line.iter().len();

                    // try and convert the bytes to UTF-8 first
                    match from_utf8(raw_line.as_slice()) {
                        Some(line) => {
                            word_count += line.words().count();
                            current_char_count = line.chars().count();
                            char_count += current_char_count;
                        },
                        None => {
                            word_count += raw_line.as_slice().split(|&x| is_word_seperator(x)).count();
                            for byte in raw_line.iter() {
                                match byte.is_ascii() {
                                    true => {
                                        current_char_count += 1;
                                    }
                                    false => { }
                                }
                            }
                            char_count += current_char_count;
                        }
                    }

                    if current_char_count > longest_line_length {
                        // we subtract one here because `line.iter().len()` includes the LF
                        // matches GNU 'wc' behaviour
                        longest_line_length = current_char_count - 1;
                    }
                },
                _ => break
            }

        }

        results.push(Result {
            filename: path.to_string(),
            bytes: byte_count,
            chars: char_count,
            lines: line_count,
            words: word_count,
            max_line_length: longest_line_length,
        });

        total_line_count += line_count;
        total_word_count += word_count;
        total_char_count += char_count;
        total_byte_count += byte_count;

        if longest_line_length > total_longest_line_length {
            total_longest_line_length = longest_line_length;
        }

        // used for formatting
        max_str_len = total_byte_count.to_string().len();
    }

    for result in results.iter() {
        print_stats(result.filename.as_slice(), result.lines, result.words, result.chars, result.bytes, result.max_line_length, matches, max_str_len);
    }

    if files.len() > 1 {
        print_stats("total", total_line_count, total_word_count, total_char_count, total_byte_count, total_longest_line_length, matches, max_str_len);
    }

    Ok(())
}

fn print_stats(filename: &str, line_count: uint, word_count: uint, char_count: uint,
    byte_count: uint, longest_line_length: uint, matches: &Matches, max_str_len: uint) {
    if matches.opt_present("lines") {
        print!("{:1$}", line_count, max_str_len);
    }
    if matches.opt_present("words") {
        print!("{:1$}", word_count, max_str_len);
    }
    if matches.opt_present("bytes") {
        print!("{:1$}", byte_count, max_str_len + 1);
    }
    if matches.opt_present("chars") {
        print!("{:1$}", char_count, max_str_len);
    }
    if matches.opt_present("max-line-length") {
        print!("{:1$}", longest_line_length, max_str_len);
    }

    // defaults
    if !matches.opt_present("bytes")
        && !matches.opt_present("chars")
        && !matches.opt_present("lines")
        && !matches.opt_present("words")
        && !matches.opt_present("max-line-length") {
        print!("{:1$}", line_count, max_str_len);
        print!("{:1$}", word_count, max_str_len + 1);
        print!("{:1$}", byte_count, max_str_len + 1);
    }

    if filename != "-" {
        println!(" {}", filename.as_slice());
    }
    else {
        println!("");
    }
}

fn open(path: String) -> StdResult<BufferedReader<Box<Reader>>, int> {
    if "-" == path.as_slice() {
        let reader = box stdin_raw() as Box<Reader>;
        return Ok(BufferedReader::new(reader));
    }

    match File::open(&std::path::Path::new(path.as_slice())) {
        Ok(fd) => {
            let reader = box fd as Box<Reader>;
            Ok(BufferedReader::new(reader))
        },
        Err(e) => {
            show_error!("wc: {0:s}: {1:s}", path, e.desc.to_string());
            Err(1)
        }
    }
}
<commit_msg>wc: fix build<commit_after>#![crate_name = "wc"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Boden Garman <bpgarman@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

#![feature(macro_rules)]

extern crate getopts;
extern crate libc;

use std::str::from_utf8;
use std::io::{print, File, BufferedReader};
use std::io::stdio::stdin_raw;
use StdResult = std::result::Result;
use getopts::Matches;

#[path = "../common/util.rs"]
mod util;

struct Result {
    filename: String,
    bytes: uint,
    chars: uint,
    lines: uint,
    words: uint,
    max_line_length: uint,
}

static NAME: &'static str = "wc";

pub fn uumain(args: Vec<String>) -> int {
    let program = args.get(0).clone();
    let opts = [
        getopts::optflag("c", "bytes", "print the byte counts"),
        getopts::optflag("m", "chars", "print the character counts"),
        getopts::optflag("l", "lines", "print the newline counts"),
        getopts::optflag("L", "max-line-length", "print the length of the longest line"),
        getopts::optflag("w", "words", "print the word counts"),
        getopts::optflag("h", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(f) => {
            crash!(1, "Invalid options\n{}", f)
        }
    };

    if matches.opt_present("help") {
        println!("Usage:");
        println!("  {0:s} [OPTION]... [FILE]...", program);
        println!("");
        print(getopts::usage("Print newline, word and byte counts for each FILE", opts).as_slice());
        println!("");
        println!("With no FILE, or when FILE is -, read standard input.");
        return 0;
    }

    if matches.opt_present("version") {
        println!("wc 1.0.0");
        return 0;
    }

    let mut files = matches.free.clone();
    if files.is_empty() {
        files = vec!("-".to_string());
    }

    match wc(files, &matches) {
        Ok(()) => ( /* pass */ ),
        Err(e) => return e
    }

    0
}

static CR: u8 = '\r' as u8;
static LF: u8 = '\n' as u8;
static SPACE: u8 = ' ' as u8;
static TAB: u8 = '\t' as u8;
static SYN: u8 = 0x16 as u8;
static FF: u8 = 0x0C as u8;

#[inline(always)]
fn is_word_seperator(byte: u8) -> bool {
    byte == SPACE || byte == TAB || byte == CR || byte == SYN || byte == FF
}

pub fn wc(files: Vec<String>, matches: &Matches) -> StdResult<(), int> {
    let mut total_line_count: uint = 0;
    let mut total_word_count: uint = 0;
    let mut total_char_count: uint = 0;
    let mut total_byte_count: uint = 0;
    let mut total_longest_line_length: uint = 0;

    let mut results = vec!();
    let mut max_str_len: uint = 0;

    for path in files.iter() {
        let mut reader = match open(path.to_string()) {
            Ok(f) => f,
            Err(e) => { return Err(e); }
        };

        let mut line_count: uint = 0;
        let mut word_count: uint = 0;
        let mut byte_count: uint = 0;
        let mut char_count: uint = 0;
        let mut current_char_count: uint = 0;
        let mut longest_line_length: uint = 0;

        loop {
            // reading from a TTY seems to raise a condition on, rather than return Some(0) like a file.
            // hence the option wrapped in a result here
            match reader.read_until(LF) {
                Ok(raw_line) => {
                    // GNU 'wc' only counts lines that end in LF as lines
                    if raw_line.iter().last().unwrap() == &LF {
                        line_count += 1;
                    }

                    byte_count += raw_line.iter().len();

                    // try and convert the bytes to UTF-8 first
                    match from_utf8(raw_line.as_slice()) {
                        Some(line) => {
                            word_count += line.words().count();
                            current_char_count = line.chars().count();
                            char_count += current_char_count;
                        },
                        None => {
                            word_count += raw_line.as_slice().split(|&x| is_word_seperator(x)).count();
                            for byte in raw_line.iter() {
                                match byte.is_ascii() {
                                    true => {
                                        current_char_count += 1;
                                    }
                                    false => { }
                                }
                            }
                            char_count += current_char_count;
                        }
                    }

                    if current_char_count > longest_line_length {
                        // we subtract one here because `line.iter().len()` includes the LF
                        // matches GNU 'wc' behaviour
                        longest_line_length = current_char_count - 1;
                    }
                },
                _ => break
            }

        }

        results.push(Result {
            filename: path.to_string(),
            bytes: byte_count,
            chars: char_count,
            lines: line_count,
            words: word_count,
            max_line_length: longest_line_length,
        });

        total_line_count += line_count;
        total_word_count += word_count;
        total_char_count += char_count;
        total_byte_count += byte_count;

        if longest_line_length > total_longest_line_length {
            total_longest_line_length = longest_line_length;
        }

        // used for formatting
        max_str_len = total_byte_count.to_string().len();
    }

    for result in results.iter() {
        print_stats(result.filename.as_slice(), result.lines, result.words, result.chars, result.bytes, result.max_line_length, matches, max_str_len);
    }

    if files.len() > 1 {
        print_stats("total", total_line_count, total_word_count, total_char_count, total_byte_count, total_longest_line_length, matches, max_str_len);
    }

    Ok(())
}

fn print_stats(filename: &str, line_count: uint, word_count: uint, char_count: uint,
    byte_count: uint, longest_line_length: uint, matches: &Matches, max_str_len: uint) {
    if matches.opt_present("lines") {
        print!("{:1$}", line_count, max_str_len);
    }
    if matches.opt_present("words") {
        print!("{:1$}", word_count, max_str_len);
    }
    if matches.opt_present("bytes") {
        print!("{:1$}", byte_count, max_str_len + 1);
    }
    if matches.opt_present("chars") {
        print!("{:1$}", char_count, max_str_len);
    }
    if matches.opt_present("max-line-length") {
        print!("{:1$}", longest_line_length, max_str_len);
    }

    // defaults
    if !matches.opt_present("bytes")
        && !matches.opt_present("chars")
        && !matches.opt_present("lines")
        && !matches.opt_present("words")
        && !matches.opt_present("max-line-length") {
        print!("{:1$}", line_count, max_str_len);
        print!("{:1$}", word_count, max_str_len + 1);
        print!("{:1$}", byte_count, max_str_len + 1);
    }

    if filename != "-" {
        println!(" {}", filename.as_slice());
    }
    else {
        println!("");
    }
}

fn open(path: String) -> StdResult<BufferedReader<Box<Reader>>, int> {
    if "-" == path.as_slice() {
        let reader = box stdin_raw() as Box<Reader>;
        return Ok(BufferedReader::new(reader));
    }

    match File::open(&std::path::Path::new(path.as_slice())) {
        Ok(fd) => {
            let reader = box fd as Box<Reader>;
            Ok(BufferedReader::new(reader))
        },
        Err(e) => {
            show_error!("wc: {0:s}: {1:s}", path, e.desc.to_string());
            Err(1)
        }
    }
}
<|endoftext|>


<commit_before>#![unstable(issue = "none", feature = "windows_stdio")]

use crate::char::decode_utf16;
use crate::cmp;
use crate::io;
use crate::ptr;
use crate::str;
use crate::sys::c;
use crate::sys::cvt;
use crate::sys::handle::Handle;
use core::str::utf8_char_width;

// Don't cache handles but get them fresh for every read/write. This allows us to track changes to
// the value over time (such as if a process calls `SetStdHandle` while it's running). See #40490.
pub struct Stdin {
    surrogate: u16,
}
pub struct Stdout {
    incomplete_utf8: IncompleteUtf8,
}

pub struct Stderr {
    incomplete_utf8: IncompleteUtf8,
}

struct IncompleteUtf8 {
    bytes: [u8; 4],
    len: u8,
}

// Apparently Windows doesn't handle large reads on stdin or writes to stdout/stderr well (see
// #13304 for details).
//
// From MSDN (2011): "The storage for this buffer is allocated from a shared heap for the
// process that is 64 KB in size. The maximum size of the buffer will depend on heap usage."
//
// We choose the cap at 8 KiB because libuv does the same, and it seems to be acceptable so far.
const MAX_BUFFER_SIZE: usize = 8192;

// The standard buffer size of BufReader for Stdin should be able to hold 3x more bytes than there
// are `u16`'s in MAX_BUFFER_SIZE. This ensures the read data can always be completely decoded from
// UTF-16 to UTF-8.
pub const STDIN_BUF_SIZE: usize = MAX_BUFFER_SIZE / 2 * 3;

pub fn get_handle(handle_id: c::DWORD) -> io::Result<c::HANDLE> {
    let handle = unsafe { c::GetStdHandle(handle_id) };
    if handle == c::INVALID_HANDLE_VALUE {
        Err(io::Error::last_os_error())
    } else if handle.is_null() {
        Err(io::Error::from_raw_os_error(c::ERROR_INVALID_HANDLE as i32))
    } else {
        Ok(handle)
    }
}

fn is_console(handle: c::HANDLE) -> bool {
    // `GetConsoleMode` will return false (0) if this is a pipe (we don't care about the reported
    // mode). This will only detect Windows Console, not other terminals connected to a pipe like
    // MSYS. Which is exactly what we need, as only Windows Console needs a conversion to UTF-16.
    let mut mode = 0;
    unsafe { c::GetConsoleMode(handle, &mut mode) != 0 }
}

fn write(
    handle_id: c::DWORD,
    data: &[u8],
    incomplete_utf8: &mut IncompleteUtf8,
) -> io::Result<usize> {
    if data.is_empty() {
        return Ok(0);
    }

    let handle = get_handle(handle_id)?;
    if !is_console(handle) {
        let handle = Handle::new(handle);
        let ret = handle.write(data);
        handle.into_raw(); // Don't close the handle
        return ret;
    }

    if incomplete_utf8.len > 0 {
        assert!(
            incomplete_utf8.len < 4,
            "Unexpected number of bytes for incomplete UTF-8 codepoint."
        );
        if data[0] >> 6 != 0b10 {
            incomplete_utf8.len = 0;
            // not a continuation byte - reject
            return Err(io::Error::new_const(
                io::ErrorKind::InvalidData,
                &"Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
            ));
        }
        incomplete_utf8.bytes[incomplete_utf8.len as usize] = data[0];
        incomplete_utf8.len += 1;
        let char_width = utf8_char_width(incomplete_utf8.bytes[0]);
        if (incomplete_utf8.len as usize) < char_width {
            // more bytes needed
            return Ok(1);
        }
        let s = str::from_utf8(&incomplete_utf8.bytes[0..incomplete_utf8.len as usize]);
        incomplete_utf8.len = 0;
        match s {
            Ok(s) => {
                assert_eq!(char_width, s.len());
                let written = write_valid_utf8_to_console(handle, s)?;
                assert_eq!(written, s.len()); // guaranteed by write_valid_utf8_to_console() for single codepoint writes
                return Ok(1);
            }
            Err(_) => {
                return Err(io::Error::new_const(
                    io::ErrorKind::InvalidData,
                    &"Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
                ));
            }
        }
    }

    // As the console is meant for presenting text, we assume bytes of `data` are encoded as UTF-8,
    // which needs to be encoded as UTF-16.
    //
    // If the data is not valid UTF-8 we write out as many bytes as are valid.
    // If the first byte is invalid it is either first byte of a multi-byte sequence but the
    // provided byte slice is too short or it is the first byte of an invalide multi-byte sequence.
    let len = cmp::min(data.len(), MAX_BUFFER_SIZE / 2);
    let utf8 = match str::from_utf8(&data[..len]) {
        Ok(s) => s,
        Err(ref e) if e.valid_up_to() == 0 => {
            let first_byte_char_width = utf8_char_width(data[0]);
            if first_byte_char_width > 1 && data.len() < first_byte_char_width {
                incomplete_utf8.bytes[0] = data[0];
                incomplete_utf8.len = 1;
                return Ok(1);
            } else {
                return Err(io::Error::new_const(
                    io::ErrorKind::InvalidData,
                    &"Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
                ));
            }
        }
        Err(e) => str::from_utf8(&data[..e.valid_up_to()]).unwrap(),
    };

    write_valid_utf8_to_console(handle, utf8)
}

fn write_valid_utf8_to_console(handle: c::HANDLE, utf8: &str) -> io::Result<usize> {
    let mut utf16 = [0u16; MAX_BUFFER_SIZE / 2];
    let mut len_utf16 = 0;
    for (chr, dest) in utf8.encode_utf16().zip(utf16.iter_mut()) {
        *dest = chr;
        len_utf16 += 1;
    }
    let utf16 = &utf16[..len_utf16];

    let mut written = write_u16s(handle, &utf16)?;

    // Figure out how many bytes of as UTF-8 were written away as UTF-16.
    if written == utf16.len() {
        Ok(utf8.len())
    } else {
        // Make sure we didn't end up writing only half of a surrogate pair (even though the chance
        // is tiny). Because it is not possible for user code to re-slice `data` in such a way that
        // a missing surrogate can be produced (and also because of the UTF-8 validation above),
        // write the missing surrogate out now.
        // Buffering it would mean we have to lie about the number of bytes written.
        let first_char_remaining = utf16[written];
        if first_char_remaining >= 0xDCEE && first_char_remaining <= 0xDFFF {
            // low surrogate
            // We just hope this works, and give up otherwise
            let _ = write_u16s(handle, &utf16[written..written + 1]);
            written += 1;
        }
        // Calculate the number of bytes of `utf8` that were actually written.
        let mut count = 0;
        for ch in utf16[..written].iter() {
            count += match ch {
                0x0000..=0x007F => 1,
                0x0080..=0x07FF => 2,
                0xDCEE..=0xDFFF => 1, // Low surrogate. We already counted 3 bytes for the other.
                _ => 3,
            };
        }
        debug_assert!(String::from_utf16(&utf16[..written]).unwrap() == utf8[..count]);
        Ok(count)
    }
}

fn write_u16s(handle: c::HANDLE, data: &[u16]) -> io::Result<usize> {
    let mut written = 0;
    cvt(unsafe {
        c::WriteConsoleW(
            handle,
            data.as_ptr() as c::LPCVOID,
            data.len() as u32,
            &mut written,
            ptr::null_mut(),
        )
    })?;
    Ok(written as usize)
}

impl Stdin {
    pub const fn new() -> Stdin {
        Stdin { surrogate: 0 }
    }
}

impl io::Read for Stdin {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let handle = get_handle(c::STD_INPUT_HANDLE)?;
        if !is_console(handle) {
            let handle = Handle::new(handle);
            let ret = handle.read(buf);
            handle.into_raw(); // Don't close the handle
            return ret;
        }

        if buf.len() == 0 {
            return Ok(0);
        } else if buf.len() < 4 {
            return Err(io::Error::new_const(
                io::ErrorKind::InvalidInput,
                &"Windows stdin in console mode does not support a buffer too small to \
                 guarantee holding one arbitrary UTF-8 character (4 bytes)",
            ));
        }

        let mut utf16_buf = [0u16; MAX_BUFFER_SIZE / 2];
        // In the worst case, an UTF-8 string can take 3 bytes for every `u16` of an UTF-16. So
        // we can read at most a third of `buf.len()` chars and uphold the guarantee no data gets
        // lost.
        let amount = cmp::min(buf.len() / 3, utf16_buf.len());
        let read = read_u16s_fixup_surrogates(handle, &mut utf16_buf, amount, &mut self.surrogate)?;

        utf16_to_utf8(&utf16_buf[..read], buf)
    }
}

// We assume that if the last `u16` is an unpaired surrogate they got sliced apart by our
// buffer size, and keep it around for the next read hoping to put them together.
// This is a best effort, and may not work if we are not the only reader on Stdin.
fn read_u16s_fixup_surrogates(
    handle: c::HANDLE,
    buf: &mut [u16],
    mut amount: usize,
    surrogate: &mut u16,
) -> io::Result<usize> {
    // Insert possibly remaining unpaired surrogate from last read.
    let mut start = 0;
    if *surrogate != 0 {
        buf[0] = *surrogate;
        *surrogate = 0;
        start = 1;
        if amount == 1 {
            // Special case: `Stdin::read` guarantees we can always read at least one new `u16`
            // and combine it with an unpaired surrogate, because the UTF-8 buffer is at least
            // 4 bytes.
            amount = 2;
        }
    }
    let mut amount = read_u16s(handle, &mut buf[start..amount])? + start;

    if amount > 0 {
        let last_char = buf[amount - 1];
        if last_char >= 0xD800 && last_char <= 0xDBFF {
            // high surrogate
            *surrogate = last_char;
            amount -= 1;
        }
    }
    Ok(amount)
}

fn read_u16s(handle: c::HANDLE, buf: &mut [u16]) -> io::Result<usize> {
    // Configure the `pInputControl` parameter to not only return on `\r\n` but also Ctrl-Z, the
    // traditional DOS method to indicate end of character stream / user input (SUB).
    // See #38274 and https://stackoverflow.com/questions/43836040/win-api-readconsole.
    const CTRL_Z: u16 = 0x1A;
    const CTRL_Z_MASK: c::ULONG = 1 << CTRL_Z;
    let mut input_control = c::CONSOLE_READCONSOLE_CONTROL {
        nLength: crate::mem::size_of::<c::CONSOLE_READCONSOLE_CONTROL>() as c::ULONG,
        nInitialChars: 0,
        dwCtrlWakeupMask: CTRL_Z_MASK,
        dwControlKeyState: 0,
    };

    let mut amount = 0;
    cvt(unsafe {
        c::ReadConsoleW(
            handle,
            buf.as_mut_ptr() as c::LPVOID,
            buf.len() as u32,
            &mut amount,
            &mut input_control as c::PCONSOLE_READCONSOLE_CONTROL,
        )
    })?;

    if amount > 0 && buf[amount as usize - 1] == CTRL_Z {
        amount -= 1;
    }
    Ok(amount as usize)
}

#[allow(unused)]
fn utf16_to_utf8(utf16: &[u16], utf8: &mut [u8]) -> io::Result<usize> {
    let mut written = 0;
    for chr in decode_utf16(utf16.iter().cloned()) {
        match chr {
            Ok(chr) => {
                chr.encode_utf8(&mut utf8[written..]);
                written += chr.len_utf8();
            }
            Err(_) => {
                // We can't really do any better than forget all data and return an error.
                return Err(io::Error::new_const(
                    io::ErrorKind::InvalidData,
                    &"Windows stdin in console mode does not support non-UTF-16 input; \
                     encountered unpaired surrogate",
                ));
            }
        }
    }
    Ok(written)
}

impl IncompleteUtf8 {
    pub const fn new() -> IncompleteUtf8 {
        IncompleteUtf8 { bytes: [0; 4], len: 0 }
    }
}

impl Stdout {
    pub const fn new() -> Stdout {
        Stdout { incomplete_utf8: IncompleteUtf8::new() }
    }
}

impl io::Write for Stdout {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        write(c::STD_OUTPUT_HANDLE, buf, &mut self.incomplete_utf8)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl Stderr {
    pub const fn new() -> Stderr {
        Stderr { incomplete_utf8: IncompleteUtf8::new() }
    }
}

impl io::Write for Stderr {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        write(c::STD_ERROR_HANDLE, buf, &mut self.incomplete_utf8)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

pub fn is_ebadf(err: &io::Error) -> bool {
    err.raw_os_error() == Some(c::ERROR_INVALID_HANDLE as i32)
}

pub fn panic_output() -> Option<impl io::Write> {
    Some(Stderr::new())
}
<commit_msg>comment pos<commit_after>#![unstable(issue = "none", feature = "windows_stdio")]

use crate::char::decode_utf16;
use crate::cmp;
use crate::io;
use crate::ptr;
use crate::str;
use crate::sys::c;
use crate::sys::cvt;
use crate::sys::handle::Handle;
use core::str::utf8_char_width;

// Don't cache handles but get them fresh for every read/write. This allows us to track changes to
// the value over time (such as if a process calls `SetStdHandle` while it's running). See #40490.
pub struct Stdin {
    surrogate: u16,
}
pub struct Stdout {
    incomplete_utf8: IncompleteUtf8,
}

pub struct Stderr {
    incomplete_utf8: IncompleteUtf8,
}

struct IncompleteUtf8 {
    bytes: [u8; 4],
    len: u8,
}

// Apparently Windows doesn't handle large reads on stdin or writes to stdout/stderr well (see
// #13304 for details).
//
// From MSDN (2011): "The storage for this buffer is allocated from a shared heap for the
// process that is 64 KB in size. The maximum size of the buffer will depend on heap usage."
//
// We choose the cap at 8 KiB because libuv does the same, and it seems to be acceptable so far.
const MAX_BUFFER_SIZE: usize = 8192;

// The standard buffer size of BufReader for Stdin should be able to hold 3x more bytes than there
// are `u16`'s in MAX_BUFFER_SIZE. This ensures the read data can always be completely decoded from
// UTF-16 to UTF-8.
pub const STDIN_BUF_SIZE: usize = MAX_BUFFER_SIZE / 2 * 3;

pub fn get_handle(handle_id: c::DWORD) -> io::Result<c::HANDLE> {
    let handle = unsafe { c::GetStdHandle(handle_id) };
    if handle == c::INVALID_HANDLE_VALUE {
        Err(io::Error::last_os_error())
    } else if handle.is_null() {
        Err(io::Error::from_raw_os_error(c::ERROR_INVALID_HANDLE as i32))
    } else {
        Ok(handle)
    }
}

fn is_console(handle: c::HANDLE) -> bool {
    // `GetConsoleMode` will return false (0) if this is a pipe (we don't care about the reported
    // mode). This will only detect Windows Console, not other terminals connected to a pipe like
    // MSYS. Which is exactly what we need, as only Windows Console needs a conversion to UTF-16.
    let mut mode = 0;
    unsafe { c::GetConsoleMode(handle, &mut mode) != 0 }
}

fn write(
    handle_id: c::DWORD,
    data: &[u8],
    incomplete_utf8: &mut IncompleteUtf8,
) -> io::Result<usize> {
    if data.is_empty() {
        return Ok(0);
    }

    let handle = get_handle(handle_id)?;
    if !is_console(handle) {
        let handle = Handle::new(handle);
        let ret = handle.write(data);
        handle.into_raw(); // Don't close the handle
        return ret;
    }

    if incomplete_utf8.len > 0 {
        assert!(
            incomplete_utf8.len < 4,
            "Unexpected number of bytes for incomplete UTF-8 codepoint."
        );
        if data[0] >> 6 != 0b10 {
            // not a continuation byte - reject
            incomplete_utf8.len = 0;
            return Err(io::Error::new_const(
                io::ErrorKind::InvalidData,
                &"Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
            ));
        }
        incomplete_utf8.bytes[incomplete_utf8.len as usize] = data[0];
        incomplete_utf8.len += 1;
        let char_width = utf8_char_width(incomplete_utf8.bytes[0]);
        if (incomplete_utf8.len as usize) < char_width {
            // more bytes needed
            return Ok(1);
        }
        let s = str::from_utf8(&incomplete_utf8.bytes[0..incomplete_utf8.len as usize]);
        incomplete_utf8.len = 0;
        match s {
            Ok(s) => {
                assert_eq!(char_width, s.len());
                let written = write_valid_utf8_to_console(handle, s)?;
                assert_eq!(written, s.len()); // guaranteed by write_valid_utf8_to_console() for single codepoint writes
                return Ok(1);
            }
            Err(_) => {
                return Err(io::Error::new_const(
                    io::ErrorKind::InvalidData,
                    &"Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
                ));
            }
        }
    }

    // As the console is meant for presenting text, we assume bytes of `data` are encoded as UTF-8,
    // which needs to be encoded as UTF-16.
    //
    // If the data is not valid UTF-8 we write out as many bytes as are valid.
    // If the first byte is invalid it is either first byte of a multi-byte sequence but the
    // provided byte slice is too short or it is the first byte of an invalide multi-byte sequence.
    let len = cmp::min(data.len(), MAX_BUFFER_SIZE / 2);
    let utf8 = match str::from_utf8(&data[..len]) {
        Ok(s) => s,
        Err(ref e) if e.valid_up_to() == 0 => {
            let first_byte_char_width = utf8_char_width(data[0]);
            if first_byte_char_width > 1 && data.len() < first_byte_char_width {
                incomplete_utf8.bytes[0] = data[0];
                incomplete_utf8.len = 1;
                return Ok(1);
            } else {
                return Err(io::Error::new_const(
                    io::ErrorKind::InvalidData,
                    &"Windows stdio in console mode does not support writing non-UTF-8 byte sequences",
                ));
            }
        }
        Err(e) => str::from_utf8(&data[..e.valid_up_to()]).unwrap(),
    };

    write_valid_utf8_to_console(handle, utf8)
}

fn write_valid_utf8_to_console(handle: c::HANDLE, utf8: &str) -> io::Result<usize> {
    let mut utf16 = [0u16; MAX_BUFFER_SIZE / 2];
    let mut len_utf16 = 0;
    for (chr, dest) in utf8.encode_utf16().zip(utf16.iter_mut()) {
        *dest = chr;
        len_utf16 += 1;
    }
    let utf16 = &utf16[..len_utf16];

    let mut written = write_u16s(handle, &utf16)?;

    // Figure out how many bytes of as UTF-8 were written away as UTF-16.
    if written == utf16.len() {
        Ok(utf8.len())
    } else {
        // Make sure we didn't end up writing only half of a surrogate pair (even though the chance
        // is tiny). Because it is not possible for user code to re-slice `data` in such a way that
        // a missing surrogate can be produced (and also because of the UTF-8 validation above),
        // write the missing surrogate out now.
        // Buffering it would mean we have to lie about the number of bytes written.
        let first_char_remaining = utf16[written];
        if first_char_remaining >= 0xDCEE && first_char_remaining <= 0xDFFF {
            // low surrogate
            // We just hope this works, and give up otherwise
            let _ = write_u16s(handle, &utf16[written..written + 1]);
            written += 1;
        }
        // Calculate the number of bytes of `utf8` that were actually written.
        let mut count = 0;
        for ch in utf16[..written].iter() {
            count += match ch {
                0x0000..=0x007F => 1,
                0x0080..=0x07FF => 2,
                0xDCEE..=0xDFFF => 1, // Low surrogate. We already counted 3 bytes for the other.
                _ => 3,
            };
        }
        debug_assert!(String::from_utf16(&utf16[..written]).unwrap() == utf8[..count]);
        Ok(count)
    }
}

fn write_u16s(handle: c::HANDLE, data: &[u16]) -> io::Result<usize> {
    let mut written = 0;
    cvt(unsafe {
        c::WriteConsoleW(
            handle,
            data.as_ptr() as c::LPCVOID,
            data.len() as u32,
            &mut written,
            ptr::null_mut(),
        )
    })?;
    Ok(written as usize)
}

impl Stdin {
    pub const fn new() -> Stdin {
        Stdin { surrogate: 0 }
    }
}

impl io::Read for Stdin {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let handle = get_handle(c::STD_INPUT_HANDLE)?;
        if !is_console(handle) {
            let handle = Handle::new(handle);
            let ret = handle.read(buf);
            handle.into_raw(); // Don't close the handle
            return ret;
        }

        if buf.len() == 0 {
            return Ok(0);
        } else if buf.len() < 4 {
            return Err(io::Error::new_const(
                io::ErrorKind::InvalidInput,
                &"Windows stdin in console mode does not support a buffer too small to \
                 guarantee holding one arbitrary UTF-8 character (4 bytes)",
            ));
        }

        let mut utf16_buf = [0u16; MAX_BUFFER_SIZE / 2];
        // In the worst case, an UTF-8 string can take 3 bytes for every `u16` of an UTF-16. So
        // we can read at most a third of `buf.len()` chars and uphold the guarantee no data gets
        // lost.
        let amount = cmp::min(buf.len() / 3, utf16_buf.len());
        let read = read_u16s_fixup_surrogates(handle, &mut utf16_buf, amount, &mut self.surrogate)?;

        utf16_to_utf8(&utf16_buf[..read], buf)
    }
}

// We assume that if the last `u16` is an unpaired surrogate they got sliced apart by our
// buffer size, and keep it around for the next read hoping to put them together.
// This is a best effort, and may not work if we are not the only reader on Stdin.
fn read_u16s_fixup_surrogates(
    handle: c::HANDLE,
    buf: &mut [u16],
    mut amount: usize,
    surrogate: &mut u16,
) -> io::Result<usize> {
    // Insert possibly remaining unpaired surrogate from last read.
    let mut start = 0;
    if *surrogate != 0 {
        buf[0] = *surrogate;
        *surrogate = 0;
        start = 1;
        if amount == 1 {
            // Special case: `Stdin::read` guarantees we can always read at least one new `u16`
            // and combine it with an unpaired surrogate, because the UTF-8 buffer is at least
            // 4 bytes.
            amount = 2;
        }
    }
    let mut amount = read_u16s(handle, &mut buf[start..amount])? + start;

    if amount > 0 {
        let last_char = buf[amount - 1];
        if last_char >= 0xD800 && last_char <= 0xDBFF {
            // high surrogate
            *surrogate = last_char;
            amount -= 1;
        }
    }
    Ok(amount)
}

fn read_u16s(handle: c::HANDLE, buf: &mut [u16]) -> io::Result<usize> {
    // Configure the `pInputControl` parameter to not only return on `\r\n` but also Ctrl-Z, the
    // traditional DOS method to indicate end of character stream / user input (SUB).
    // See #38274 and https://stackoverflow.com/questions/43836040/win-api-readconsole.
    const CTRL_Z: u16 = 0x1A;
    const CTRL_Z_MASK: c::ULONG = 1 << CTRL_Z;
    let mut input_control = c::CONSOLE_READCONSOLE_CONTROL {
        nLength: crate::mem::size_of::<c::CONSOLE_READCONSOLE_CONTROL>() as c::ULONG,
        nInitialChars: 0,
        dwCtrlWakeupMask: CTRL_Z_MASK,
        dwControlKeyState: 0,
    };

    let mut amount = 0;
    cvt(unsafe {
        c::ReadConsoleW(
            handle,
            buf.as_mut_ptr() as c::LPVOID,
            buf.len() as u32,
            &mut amount,
            &mut input_control as c::PCONSOLE_READCONSOLE_CONTROL,
        )
    })?;

    if amount > 0 && buf[amount as usize - 1] == CTRL_Z {
        amount -= 1;
    }
    Ok(amount as usize)
}

#[allow(unused)]
fn utf16_to_utf8(utf16: &[u16], utf8: &mut [u8]) -> io::Result<usize> {
    let mut written = 0;
    for chr in decode_utf16(utf16.iter().cloned()) {
        match chr {
            Ok(chr) => {
                chr.encode_utf8(&mut utf8[written..]);
                written += chr.len_utf8();
            }
            Err(_) => {
                // We can't really do any better than forget all data and return an error.
                return Err(io::Error::new_const(
                    io::ErrorKind::InvalidData,
                    &"Windows stdin in console mode does not support non-UTF-16 input; \
                     encountered unpaired surrogate",
                ));
            }
        }
    }
    Ok(written)
}

impl IncompleteUtf8 {
    pub const fn new() -> IncompleteUtf8 {
        IncompleteUtf8 { bytes: [0; 4], len: 0 }
    }
}

impl Stdout {
    pub const fn new() -> Stdout {
        Stdout { incomplete_utf8: IncompleteUtf8::new() }
    }
}

impl io::Write for Stdout {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        write(c::STD_OUTPUT_HANDLE, buf, &mut self.incomplete_utf8)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl Stderr {
    pub const fn new() -> Stderr {
        Stderr { incomplete_utf8: IncompleteUtf8::new() }
    }
}

impl io::Write for Stderr {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        write(c::STD_ERROR_HANDLE, buf, &mut self.incomplete_utf8)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

pub fn is_ebadf(err: &io::Error) -> bool {
    err.raw_os_error() == Some(c::ERROR_INVALID_HANDLE as i32)
}

pub fn panic_output() -> Option<impl io::Write> {
    Some(Stderr::new())
}
<|endoftext|>


<commit_before>use std::io::{self, prelude::Write};
use std::time::Duration;

use super::OutputFormatter;
use crate::{
    console::{ConsoleTestState, OutputLocation},
    test_result::TestResult,
    time,
    types::{TestDesc, TestType},
};

pub struct JunitFormatter<T> {
    out: OutputLocation<T>,
    results: Vec<(TestDesc, TestResult, Duration)>,
}

impl<T: Write> JunitFormatter<T> {
    pub fn new(out: OutputLocation<T>) -> Self {
        Self { out, results: Vec::new() }
    }

    fn write_message(&mut self, s: &str) -> io::Result<()> {
        assert!(!s.contains('\n'));

        self.out.write_all(s.as_ref())
    }
}

impl<T: Write> OutputFormatter for JunitFormatter<T> {
    fn write_run_start(&mut self, _test_count: usize) -> io::Result<()> {
        // We write xml header on run start
        self.out.write_all("\n".as_bytes())?;
        self.write_message("<?xml version=\"1.0\" encoding=\"UTF-8\"?>")
    }

    fn write_test_start(&mut self, _desc: &TestDesc) -> io::Result<()> {
        // We do not output anything on test start.
        Ok(())
    }

    fn write_timeout(&mut self, _desc: &TestDesc) -> io::Result<()> {
        // We do not output anything on test timeout.
        Ok(())
    }

    fn write_result(
        &mut self,
        desc: &TestDesc,
        result: &TestResult,
        exec_time: Option<&time::TestExecTime>,
        _stdout: &[u8],
        _state: &ConsoleTestState,
    ) -> io::Result<()> {
        // Because the testsuit node holds some of the information as attributes, we can't write it
        // until all of the tests has ran. Instead of writting every result as they come in, we add
        // them to a Vec and write them all at once when run is complete.
        let duration = exec_time.map(|t| t.0).unwrap_or_default();
        self.results.push((desc.clone(), result.clone(), duration));
        Ok(())
    }
    fn write_run_finish(&mut self, state: &ConsoleTestState) -> io::Result<bool> {
        self.write_message("<testsuites>")?;

        self.write_message(&*format!(
            "<testsuite name=\"test\" package=\"test\" id=\"0\" \
             errors=\"0\" \
             failures=\"{}\" \
             tests=\"{}\" \
             skipped=\"{}\" \
             >",
            state.failed, state.total, state.ignored
        ))?;
        for (desc, result, duration) in std::mem::replace(&mut self.results, Vec::new()) {
            let (class_name, test_name) = parse_class_name(&desc);
            match result {
                TestResult::TrIgnored => { /* no-op */ }
                TestResult::TrFailed => {
                    self.write_message(&*format!(
                        "<testcase classname=\"{}\" \
                         name=\"{}\" time=\"{}\">",
                        class_name,
                        test_name,
                        duration.as_secs_f64()
                    ))?;
                    self.write_message("<failure type=\"assert\"/>")?;
                    self.write_message("</testcase>")?;
                }

                TestResult::TrFailedMsg(ref m) => {
                    self.write_message(&*format!(
                        "<testcase classname=\"{}\" \
                         name=\"{}\" time=\"{}\">",
                        class_name,
                        test_name,
                        duration.as_secs_f64()
                    ))?;
                    self.write_message(&*format!("<failure message=\"{}\" type=\"assert\"/>", m))?;
                    self.write_message("</testcase>")?;
                }

                TestResult::TrTimedFail => {
                    self.write_message(&*format!(
                        "<testcase classname=\"{}\" \
                         name=\"{}\" time=\"{}\">",
                        class_name,
                        test_name,
                        duration.as_secs_f64()
                    ))?;
                    self.write_message("<failure type=\"timeout\"/>")?;
                    self.write_message("</testcase>")?;
                }

                TestResult::TrBench(ref b) => {
                    self.write_message(&*format!(
                        "<testcase classname=\"benchmark::{}\" \
                         name=\"{}\" time=\"{}\" />",
                        class_name, test_name, b.ns_iter_summ.sum
                    ))?;
                }

                TestResult::TrOk | TestResult::TrAllowedFail => {
                    self.write_message(&*format!(
                        "<testcase classname=\"{}\" \
                         name=\"{}\" time=\"{}\"/>",
                        class_name,
                        test_name,
                        duration.as_secs_f64()
                    ))?;
                }
            }
        }
        self.write_message("<system-out/>")?;
        self.write_message("<system-err/>")?;
        self.write_message("</testsuite>")?;
        self.write_message("</testsuites>")?;

        self.out.write_all("\n\n".as_bytes())?;

        Ok(state.failed == 0)
    }
}

fn parse_class_name(desc: &TestDesc) -> (String, String) {
    match desc.test_type {
        TestType::UnitTest => parse_class_name_unit(desc),
        TestType::DocTest => parse_class_name_doc(desc),
        TestType::IntegrationTest => parse_class_name_integration(desc),
        TestType::Unknown => (String::from("unknown"), String::from(desc.name.as_slice())),
    }
}

fn parse_class_name_unit(desc: &TestDesc) -> (String, String) {
    // Module path => classname
    // Function name => name
    let module_segments: Vec<&str> = desc.name.as_slice().split("::").collect();
    let (class_name, test_name) = match module_segments[..] {
        [test] => (String::from("crate"), String::from(test)),
        [ref path @ .., test] => (path.join("::"), String::from(test)),
        [..] => unreachable!(),
    };
    (class_name, test_name)
}

fn parse_class_name_doc(desc: &TestDesc) -> (String, String) {
    // File path => classname
    // Line # => test name
    let segments: Vec<&str> = desc.name.as_slice().split(" - ").collect();
    let (class_name, test_name) = match segments[..] {
        [file, line] => (String::from(file.trim()), String::from(line.trim())),
        [..] => unreachable!(),
    };
    (class_name, test_name)
}

fn parse_class_name_integration(desc: &TestDesc) -> (String, String) {
    (String::from("integration"), String::from(desc.name.as_slice()))
}
<commit_msg>Apply suggestions from code review<commit_after>use std::io::{self, prelude::Write};
use std::time::Duration;

use super::OutputFormatter;
use crate::{
    console::{ConsoleTestState, OutputLocation},
    test_result::TestResult,
    time,
    types::{TestDesc, TestType},
};

pub struct JunitFormatter<T> {
    out: OutputLocation<T>,
    results: Vec<(TestDesc, TestResult, Duration)>,
}

impl<T: Write> JunitFormatter<T> {
    pub fn new(out: OutputLocation<T>) -> Self {
        Self { out, results: Vec::new() }
    }

    fn write_message(&mut self, s: &str) -> io::Result<()> {
        assert!(!s.contains('\n'));

        self.out.write_all(s.as_ref())
    }
}

impl<T: Write> OutputFormatter for JunitFormatter<T> {
    fn write_run_start(&mut self, _test_count: usize) -> io::Result<()> {
        // We write xml header on run start
        self.out.write_all(b"\n")?;
        self.write_message("<?xml version=\"1.0\" encoding=\"UTF-8\"?>")
    }

    fn write_test_start(&mut self, _desc: &TestDesc) -> io::Result<()> {
        // We do not output anything on test start.
        Ok(())
    }

    fn write_timeout(&mut self, _desc: &TestDesc) -> io::Result<()> {
        // We do not output anything on test timeout.
        Ok(())
    }

    fn write_result(
        &mut self,
        desc: &TestDesc,
        result: &TestResult,
        exec_time: Option<&time::TestExecTime>,
        _stdout: &[u8],
        _state: &ConsoleTestState,
    ) -> io::Result<()> {
        // Because the testsuit node holds some of the information as attributes, we can't write it
        // until all of the tests has ran. Instead of writting every result as they come in, we add
        // them to a Vec and write them all at once when run is complete.
        let duration = exec_time.map(|t| t.0).unwrap_or_default();
        self.results.push((desc.clone(), result.clone(), duration));
        Ok(())
    }
    fn write_run_finish(&mut self, state: &ConsoleTestState) -> io::Result<bool> {
        self.write_message("<testsuites>")?;

        self.write_message(&*format!(
            "<testsuite name=\"test\" package=\"test\" id=\"0\" \
             errors=\"0\" \
             failures=\"{}\" \
             tests=\"{}\" \
             skipped=\"{}\" \
             >",
            state.failed, state.total, state.ignored
        ))?;
        for (desc, result, duration) in std::mem::replace(&mut self.results, Vec::new()) {
            let (class_name, test_name) = parse_class_name(&desc);
            match result {
                TestResult::TrIgnored => { /* no-op */ }
                TestResult::TrFailed => {
                    self.write_message(&*format!(
                        "<testcase classname=\"{}\" \
                         name=\"{}\" time=\"{}\">",
                        class_name,
                        test_name,
                        duration.as_secs_f64()
                    ))?;
                    self.write_message("<failure type=\"assert\"/>")?;
                    self.write_message("</testcase>")?;
                }

                TestResult::TrFailedMsg(ref m) => {
                    self.write_message(&*format!(
                        "<testcase classname=\"{}\" \
                         name=\"{}\" time=\"{}\">",
                        class_name,
                        test_name,
                        duration.as_secs_f64()
                    ))?;
                    self.write_message(&*format!("<failure message=\"{}\" type=\"assert\"/>", m))?;
                    self.write_message("</testcase>")?;
                }

                TestResult::TrTimedFail => {
                    self.write_message(&*format!(
                        "<testcase classname=\"{}\" \
                         name=\"{}\" time=\"{}\">",
                        class_name,
                        test_name,
                        duration.as_secs_f64()
                    ))?;
                    self.write_message("<failure type=\"timeout\"/>")?;
                    self.write_message("</testcase>")?;
                }

                TestResult::TrBench(ref b) => {
                    self.write_message(&*format!(
                        "<testcase classname=\"benchmark::{}\" \
                         name=\"{}\" time=\"{}\" />",
                        class_name, test_name, b.ns_iter_summ.sum
                    ))?;
                }

                TestResult::TrOk | TestResult::TrAllowedFail => {
                    self.write_message(&*format!(
                        "<testcase classname=\"{}\" \
                         name=\"{}\" time=\"{}\"/>",
                        class_name,
                        test_name,
                        duration.as_secs_f64()
                    ))?;
                }
            }
        }
        self.write_message("<system-out/>")?;
        self.write_message("<system-err/>")?;
        self.write_message("</testsuite>")?;
        self.write_message("</testsuites>")?;

        self.out.write_all(b"\n\n")?;

        Ok(state.failed == 0)
    }
}

fn parse_class_name(desc: &TestDesc) -> (String, String) {
    match desc.test_type {
        TestType::UnitTest => parse_class_name_unit(desc),
        TestType::DocTest => parse_class_name_doc(desc),
        TestType::IntegrationTest => parse_class_name_integration(desc),
        TestType::Unknown => (String::from("unknown"), String::from(desc.name.as_slice())),
    }
}

fn parse_class_name_unit(desc: &TestDesc) -> (String, String) {
    // Module path => classname
    // Function name => name
    let module_segments: Vec<&str> = desc.name.as_slice().split("::").collect();
    let (class_name, test_name) = match module_segments[..] {
        [test] => (String::from("crate"), String::from(test)),
        [ref path @ .., test] => (path.join("::"), String::from(test)),
        [..] => unreachable!(),
    };
    (class_name, test_name)
}

fn parse_class_name_doc(desc: &TestDesc) -> (String, String) {
    // File path => classname
    // Line # => test name
    let segments: Vec<&str> = desc.name.as_slice().split(" - ").collect();
    let (class_name, test_name) = match segments[..] {
        [file, line] => (String::from(file.trim()), String::from(line.trim())),
        [..] => unreachable!(),
    };
    (class_name, test_name)
}

fn parse_class_name_integration(desc: &TestDesc) -> (String, String) {
    (String::from("integration"), String::from(desc.name.as_slice()))
}
<|endoftext|>


<commit_before>use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
use crossbeam::unbounded;
use lazy_static::lazy_static;
use rosrust;
use std::process::Command;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time;

mod util;

mod msg {
    rosrust::rosmsg_include!(
        roscpp_tutorials / TwoInts,
        rospy_tutorials / AddTwoInts,
        std_msgs / String
    );
}

fn global_init() -> util::ChildProcessTerminator {
    let roscore = util::run_roscore_for(util::Language::None, util::Feature::Benchmarks);
    rosrust::init("benchmarker");
    roscore
}

lazy_static! {
    static ref ROS_CORE: Option<util::ChildProcessTerminator> = Some(global_init());
}

fn setup() {
    assert!(ROS_CORE.is_some());
}

fn subscribe_publish_directly(criterion: &mut Criterion) {
    setup();

    let topic = format!("/topicat{}", line!());

    let publisher = rosrust::publish::<msg::std_msgs::String>(&topic, 2000).unwrap();

    let (tx, rx) = unbounded();

    let _subscriber =
        rosrust::subscribe::<msg::std_msgs::String, _>(&topic, 2000, move |message| {
            tx.send(message.data).unwrap();
        })
        .unwrap();

    loop {
        publisher
            .send(msg::std_msgs::String { data: "".into() })
            .unwrap();
        if !rx.is_empty() {
            break;
        }
    }

    publisher
        .send(msg::std_msgs::String {
            data: "ready".into(),
        })
        .unwrap();

    while rx.recv().unwrap() != "ready" {}

    let inner_publisher = publisher.clone();
    let receiver = rx.clone();
    criterion.bench_function("send and receive single message directly", move |b| {
        let idx = AtomicUsize::new(0);

        b.iter_batched(
            || format!("{}:{}", line!(), idx.fetch_add(1, Ordering::SeqCst)),
            |data| {
                inner_publisher
                    .send(msg::std_msgs::String { data: data.clone() })
                    .unwrap();
                let response = receiver.recv().unwrap();
                assert_eq!(data, response);
            },
            BatchSize::SmallInput,
        );
    });

    let inner_publisher = publisher.clone();
    let receiver = rx.clone();
    criterion.bench_function("send and receive 100 messages directly", move |b| {
        let idx = AtomicUsize::new(0);

        b.iter_batched(
            || format!("{}:{}", line!(), idx.fetch_add(1, Ordering::SeqCst)),
            |data| {
                for _ in 0..100 {
                    inner_publisher
                        .send(msg::std_msgs::String { data: data.clone() })
                        .unwrap();
                }
                for _ in 0..100 {
                    let response = receiver.recv().unwrap();
                    assert_eq!(data, response);
                }
            },
            BatchSize::SmallInput,
        );
    });
}

fn call_service(criterion: &mut Criterion) {
    setup();

    let namespace = format!("/namespaceat{}", line!());

    let _roscpp_service = util::ChildProcessTerminator::spawn(
        Command::new("rosrun")
            .arg("roscpp_tutorials")
            .arg("add_two_ints_server")
            .arg(format!("__ns:={}", namespace))
            .arg("__name:=roscpp_service")
            .arg("add_two_ints:=add_two_ints_cpp"),
    );

    let _roscpp_service = util::ChildProcessTerminator::spawn(
        Command::new("rosrun")
            .arg("rospy_tutorials")
            .arg("add_two_ints_server")
            .arg(format!("__ns:={}", namespace))
            .arg("__name:=rospy_service")
            .arg("add_two_ints:=add_two_ints_py"),
    );

    let _roscpp_service = util::ChildProcessTerminator::spawn_example_bench(
        "../examples/serviceclient",
        Command::new("cargo")
            .arg("run")
            .arg("--release")
            .arg("--bin")
            .arg("service")
            .arg(format!("__ns:={}", namespace))
            .arg("__name:=rosrust_service")
            .arg("add_two_ints:=add_two_ints_rust"),
    );

    let service_name_inline = format!("service_at_line_{}", line!());

    let _service_raii =
        rosrust::service::<msg::roscpp_tutorials::TwoInts, _>(&service_name_inline, move |req| {
            let sum = req.a + req.b;
            Ok(msg::roscpp_tutorials::TwoIntsRes { sum })
        })
        .unwrap();

    let service_name_cpp = format!("{}/add_two_ints_cpp", namespace);
    let service_name_py = format!("{}/add_two_ints_py", namespace);
    let service_name_rust = format!("{}/add_two_ints_rust", namespace);

    rosrust::wait_for_service(&service_name_cpp, Some(time::Duration::from_secs(10))).unwrap();
    rosrust::wait_for_service(&service_name_py, Some(time::Duration::from_secs(10))).unwrap();
    rosrust::wait_for_service(&service_name_rust, Some(time::Duration::from_secs(10))).unwrap();
    rosrust::wait_for_service(&service_name_inline, Some(time::Duration::from_secs(10))).unwrap();

    let client_original =
        rosrust::client::<msg::roscpp_tutorials::TwoInts>(&service_name_cpp).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call roscpp service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::roscpp_tutorials::TwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });
    let client = client_original.clone();
    criterion.bench_function("call roscpp service 50 times in parallel", move |b| {
        b.iter(|| {
            let requests = (0..50)
                .map(|a| client.req_async(msg::roscpp_tutorials::TwoIntsReq { a, b: 5 }))
                .collect::<Vec<_>>();
            for (idx, request) in requests.into_iter().enumerate() {
                assert_eq!(idx as i64 + 5, request.read().unwrap().unwrap().sum);
            }
        });
    });

    let client_original =
        rosrust::client::<msg::rospy_tutorials::AddTwoInts>(&service_name_py).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call rospy service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::rospy_tutorials::AddTwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });

    let client_original =
        rosrust::client::<msg::roscpp_tutorials::TwoInts>(&service_name_rust).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call rosrust service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::roscpp_tutorials::TwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });
    let client = client_original.clone();
    criterion.bench_function("call rosrust service 50 times in parallel", move |b| {
        b.iter(|| {
            let requests = (0..50)
                .map(|a| client.req_async(msg::roscpp_tutorials::TwoIntsReq { a, b: 5 }))
                .collect::<Vec<_>>();
            for (idx, request) in requests.into_iter().enumerate() {
                assert_eq!(idx as i64 + 5, request.read().unwrap().unwrap().sum);
            }
        });
    });

    let client_original =
        rosrust::client::<msg::roscpp_tutorials::TwoInts>(&service_name_inline).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call inline rosrust service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::roscpp_tutorials::TwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });
    let client = client_original.clone();
    criterion.bench_function(
        "call inline rosrust service 50 times in parallel",
        move |b| {
            b.iter(|| {
                let requests = (0..50)
                    .map(|a| client.req_async(msg::roscpp_tutorials::TwoIntsReq { a, b: 5 }))
                    .collect::<Vec<_>>();
                for (idx, request) in requests.into_iter().enumerate() {
                    assert_eq!(idx as i64 + 5, request.read().unwrap().unwrap().sum);
                }
            });
        },
    );
}

criterion_group!(benches, subscribe_publish_directly, call_service);
criterion_main!(benches);
<commit_msg>Add relayed pubsub benchmarks<commit_after>use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
use crossbeam::unbounded;
use lazy_static::lazy_static;
use rosrust;
use std::process::Command;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time;

mod util;

mod msg {
    rosrust::rosmsg_include!(
        roscpp_tutorials / TwoInts,
        rospy_tutorials / AddTwoInts,
        std_msgs / String
    );
}

fn global_init() -> util::ChildProcessTerminator {
    let roscore = util::run_roscore_for(util::Language::None, util::Feature::Benchmarks);
    rosrust::init("benchmarker");
    roscore
}

lazy_static! {
    static ref ROS_CORE: Option<util::ChildProcessTerminator> = Some(global_init());
}

fn setup() {
    assert!(ROS_CORE.is_some());
}

fn subscribe_publish_relayed(criterion: &mut Criterion) {
    setup();

    let pub_topic = format!("/chatter_at_{}", line!());
    let sub_topic = format!("/chatter_at_{}", line!());

    let _subscriber = util::ChildProcessTerminator::spawn(
        Command::new("rosrun")
            .arg("topic_tools")
            .arg("relay")
            .arg(&pub_topic)
            .arg(&sub_topic),
    );

    let (tx, rx) = unbounded();

    let _log_subscriber =
        rosrust::subscribe::<msg::std_msgs::String, _>(&sub_topic, 2000, move |data| {
            tx.send(data.data).unwrap();
        })
        .unwrap();

    let publisher = rosrust::publish::<msg::std_msgs::String>(&pub_topic, 2000).unwrap();

    loop {
        publisher
            .send(msg::std_msgs::String { data: "".into() })
            .unwrap();
        std::thread::sleep(time::Duration::from_millis(100));
        if !rx.is_empty() {
            break;
        }
    }

    publisher
        .send(msg::std_msgs::String {
            data: "ready".into(),
        })
        .unwrap();

    while rx.recv().unwrap() != "ready" {}

    let inner_publisher = publisher.clone();
    let receiver = rx.clone();
    criterion.bench_function("send and receive single message relayed", move |b| {
        let idx = AtomicUsize::new(0);

        b.iter_batched(
            || format!("{}:{}", line!(), idx.fetch_add(1, Ordering::SeqCst)),
            |data| {
                inner_publisher
                    .send(msg::std_msgs::String { data: data.clone() })
                    .unwrap();
                let response = receiver.recv().unwrap();
                assert_eq!(data, response);
            },
            BatchSize::SmallInput,
        );
    });

    let inner_publisher = publisher.clone();
    let receiver = rx.clone();
    // Limit to 8 due to the relay node's buffer limit being 10
    criterion.bench_function("send and receive 8 messages relayed", move |b| {
        let idx = AtomicUsize::new(0);

        b.iter_batched(
            || {
                let id = idx.fetch_add(1, Ordering::SeqCst);
                (0..8)
                    .map(|item| format!("{}:{}:{}", line!(), id, item))
                    .collect::<Vec<String>>()
            },
            |data| {
                for item in &data {
                    inner_publisher
                        .send(msg::std_msgs::String { data: item.clone() })
                        .unwrap();
                }
                for item in &data {
                    let response = receiver.recv().unwrap();
                    assert_eq!(item, &response);
                }
            },
            BatchSize::LargeInput,
        );
    });
}

fn subscribe_publish_directly(criterion: &mut Criterion) {
    setup();

    let topic = format!("/topicat{}", line!());

    let publisher = rosrust::publish::<msg::std_msgs::String>(&topic, 2000).unwrap();

    let (tx, rx) = unbounded();

    let _subscriber =
        rosrust::subscribe::<msg::std_msgs::String, _>(&topic, 2000, move |message| {
            tx.send(message.data).unwrap();
        })
        .unwrap();

    loop {
        publisher
            .send(msg::std_msgs::String { data: "".into() })
            .unwrap();
        if !rx.is_empty() {
            break;
        }
    }

    publisher
        .send(msg::std_msgs::String {
            data: "ready".into(),
        })
        .unwrap();

    while rx.recv().unwrap() != "ready" {}

    let inner_publisher = publisher.clone();
    let receiver = rx.clone();
    criterion.bench_function("send and receive single message directly", move |b| {
        let idx = AtomicUsize::new(0);

        b.iter_batched(
            || format!("{}:{}", line!(), idx.fetch_add(1, Ordering::SeqCst)),
            |data| {
                inner_publisher
                    .send(msg::std_msgs::String { data: data.clone() })
                    .unwrap();
                let response = receiver.recv().unwrap();
                assert_eq!(data, response);
            },
            BatchSize::SmallInput,
        );
    });

    let inner_publisher = publisher.clone();
    let receiver = rx.clone();
    criterion.bench_function("send and receive 100 messages directly", move |b| {
        let idx = AtomicUsize::new(0);

        b.iter_batched(
            || format!("{}:{}", line!(), idx.fetch_add(1, Ordering::SeqCst)),
            |data| {
                for _ in 0..100 {
                    inner_publisher
                        .send(msg::std_msgs::String { data: data.clone() })
                        .unwrap();
                }
                for _ in 0..100 {
                    let response = receiver.recv().unwrap();
                    assert_eq!(data, response);
                }
            },
            BatchSize::SmallInput,
        );
    });
}

fn call_service(criterion: &mut Criterion) {
    setup();

    let namespace = format!("/namespaceat{}", line!());

    let _roscpp_service = util::ChildProcessTerminator::spawn(
        Command::new("rosrun")
            .arg("roscpp_tutorials")
            .arg("add_two_ints_server")
            .arg(format!("__ns:={}", namespace))
            .arg("__name:=roscpp_service")
            .arg("add_two_ints:=add_two_ints_cpp"),
    );

    let _roscpp_service = util::ChildProcessTerminator::spawn(
        Command::new("rosrun")
            .arg("rospy_tutorials")
            .arg("add_two_ints_server")
            .arg(format!("__ns:={}", namespace))
            .arg("__name:=rospy_service")
            .arg("add_two_ints:=add_two_ints_py"),
    );

    let _roscpp_service = util::ChildProcessTerminator::spawn_example_bench(
        "../examples/serviceclient",
        Command::new("cargo")
            .arg("run")
            .arg("--release")
            .arg("--bin")
            .arg("service")
            .arg(format!("__ns:={}", namespace))
            .arg("__name:=rosrust_service")
            .arg("add_two_ints:=add_two_ints_rust"),
    );

    let service_name_inline = format!("service_at_line_{}", line!());

    let _service_raii =
        rosrust::service::<msg::roscpp_tutorials::TwoInts, _>(&service_name_inline, move |req| {
            let sum = req.a + req.b;
            Ok(msg::roscpp_tutorials::TwoIntsRes { sum })
        })
        .unwrap();

    let service_name_cpp = format!("{}/add_two_ints_cpp", namespace);
    let service_name_py = format!("{}/add_two_ints_py", namespace);
    let service_name_rust = format!("{}/add_two_ints_rust", namespace);

    rosrust::wait_for_service(&service_name_cpp, Some(time::Duration::from_secs(10))).unwrap();
    rosrust::wait_for_service(&service_name_py, Some(time::Duration::from_secs(10))).unwrap();
    rosrust::wait_for_service(&service_name_rust, Some(time::Duration::from_secs(10))).unwrap();
    rosrust::wait_for_service(&service_name_inline, Some(time::Duration::from_secs(10))).unwrap();

    let client_original =
        rosrust::client::<msg::roscpp_tutorials::TwoInts>(&service_name_cpp).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call roscpp service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::roscpp_tutorials::TwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });
    let client = client_original.clone();
    criterion.bench_function("call roscpp service 50 times in parallel", move |b| {
        b.iter(|| {
            let requests = (0..50)
                .map(|a| client.req_async(msg::roscpp_tutorials::TwoIntsReq { a, b: 5 }))
                .collect::<Vec<_>>();
            for (idx, request) in requests.into_iter().enumerate() {
                assert_eq!(idx as i64 + 5, request.read().unwrap().unwrap().sum);
            }
        });
    });

    let client_original =
        rosrust::client::<msg::rospy_tutorials::AddTwoInts>(&service_name_py).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call rospy service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::rospy_tutorials::AddTwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });

    let client_original =
        rosrust::client::<msg::roscpp_tutorials::TwoInts>(&service_name_rust).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call rosrust service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::roscpp_tutorials::TwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });
    let client = client_original.clone();
    criterion.bench_function("call rosrust service 50 times in parallel", move |b| {
        b.iter(|| {
            let requests = (0..50)
                .map(|a| client.req_async(msg::roscpp_tutorials::TwoIntsReq { a, b: 5 }))
                .collect::<Vec<_>>();
            for (idx, request) in requests.into_iter().enumerate() {
                assert_eq!(idx as i64 + 5, request.read().unwrap().unwrap().sum);
            }
        });
    });

    let client_original =
        rosrust::client::<msg::roscpp_tutorials::TwoInts>(&service_name_inline).unwrap();
    let client = client_original.clone();
    criterion.bench_function("call inline rosrust service once", move |b| {
        b.iter(|| {
            let sum = client
                .req(&msg::roscpp_tutorials::TwoIntsReq { a: 48, b: 12 })
                .unwrap()
                .unwrap()
                .sum;
            assert_eq!(60, sum);
        });
    });
    let client = client_original.clone();
    criterion.bench_function(
        "call inline rosrust service 50 times in parallel",
        move |b| {
            b.iter(|| {
                let requests = (0..50)
                    .map(|a| client.req_async(msg::roscpp_tutorials::TwoIntsReq { a, b: 5 }))
                    .collect::<Vec<_>>();
                for (idx, request) in requests.into_iter().enumerate() {
                    assert_eq!(idx as i64 + 5, request.read().unwrap().unwrap().sum);
                }
            });
        },
    );
}

criterion_group!(
    benches,
    subscribe_publish_directly,
    subscribe_publish_relayed,
    call_service
);
criterion_main!(benches);
<|endoftext|>