<commit_before>extern crate rand;

use std::borrow::Cow;
use std::cmp;
use self::rand::{Rand, Rng};
use parse::{Node, NodeRef, ParsedFile, RangeRef};

#[derive(Debug)]
pub struct FuzzFile<'buf: 'parse, 'parse> {
    root: Cow<'parse, [NodeRef]>,
    nodes: Cow<'parse, [Node<'buf>]>,
    ranges: Cow<'parse, [Vec<NodeRef>]>,
}

#[derive(Clone)]
pub enum FuzzAction {
    DuplicateRange,
    DuplicateRootNode,
    NestDelim,
    RemoveDelim,
    ShuffleRanges,
    SwapDelim,
    SwapRanges,
}

pub struct FuzzConfig {
    pub max_mutations: usize,
    pub max_duplications: usize,
}

impl Rand for FuzzAction {
    fn rand<R: Rng>(rng: &mut R) -> Self {
        let actions = vec![
            FuzzAction::DuplicateRange,
            FuzzAction::DuplicateRootNode,
            FuzzAction::NestDelim,
            FuzzAction::RemoveDelim,
            FuzzAction::ShuffleRanges,
            FuzzAction::SwapDelim,
            FuzzAction::SwapRanges,
        ];
        rng.choose(&actions).unwrap().clone()
    }
}

fn rand_indices<R: Rng, T>(mut rng: &mut Rng, x: &[T]) -> Option<(usize, usize)> {
    if x.len() > 1 {
        let indices = rand::sample(&mut rng, 0..x.len(), 2);
        Some((indices[0], indices[1]))
    } else {
        None
    }
}

fn rand_delim<'buf, R: Rng>(mut rng: &mut R, nodes: &[Node<'buf>]) -> Option<(NodeRef, &'buf [u8], RangeRef, &'buf [u8])> {
    let delims: Vec<_> = nodes.iter().enumerate().filter_map(|item| {
        match item {
            (index, &Node::Delim(start_pattern, rangeref, end_pattern)) => {
                Some((index, start_pattern, rangeref, end_pattern))
            },
            _ => None,
        }
    }).collect();
    rng.choose(&delims[..]).cloned()
}

struct SerializeState {
    have_serialized_range: Vec<bool>,
}

impl SerializeState {
    fn new(ranges: &[Vec<NodeRef>]) -> Self {
        SerializeState {
            have_serialized_range: ranges.iter().map(|_| false).collect()
        }
    }

    fn should_serialize(self: &mut Self, rangeref: RangeRef) -> bool {
        if self.have_serialized_range[rangeref] {
            false
        } else {
            self.have_serialized_range[rangeref] = true;
            true
        }
    }

    fn reset(self: &mut Self, rangeref: RangeRef) {
        self.have_serialized_range[rangeref] = false;
    }
}

pub trait SerializeInto {
    fn push(&mut self, &[u8]);
}

impl SerializeInto for Vec<u8> {
    fn push(self: &mut Self, token: &[u8]) {
        self.extend(token);
    }
}

pub struct SliceSerializer<'buf> {
    slice: &'buf mut [u8],
    cur_offset: usize,
}

impl<'buf> SliceSerializer<'buf> {
    pub fn new(slice: &'buf mut [u8]) -> SliceSerializer {
        SliceSerializer {
            slice: slice,
            cur_offset: 0,
        }
    }

    pub fn bytes_written(self: &Self) -> usize {
        self.cur_offset
    }
}

impl<'buf> SerializeInto for SliceSerializer<'buf> {
    fn push(self: &mut Self, token: &[u8]) {
        let remaining = self.slice.len() - self.cur_offset;
        let num_bytes_to_write = cmp::min(remaining, token.len());
        if num_bytes_to_write > 0 {
            self.slice[self.cur_offset..self.cur_offset+num_bytes_to_write].copy_from_slice(&token[..num_bytes_to_write]);
            self.cur_offset += num_bytes_to_write
        }
    }
}

impl<'buf, 'parse> FuzzFile<'buf, 'parse> {
    pub fn new(parsed: &'parse ParsedFile<'buf>) -> Self {
        FuzzFile {
            root: Cow::from(parsed.root.as_slice()),
            nodes: Cow::from(parsed.nodes.as_slice()),
            ranges: Cow::from(parsed.ranges.as_slice()),
        }
    }

    fn serialize_noderef<S: SerializeInto>(self: &Self, noderef: NodeRef, mut state: &mut SerializeState, mut out: &mut S) {
        match self.nodes[noderef] {
            Node::Delim(prefix, rangeref, postfix) => {
                out.push(prefix);
                if state.should_serialize(rangeref) {
                    for noderef in &self.ranges[rangeref] {
                        self.serialize_noderef(*noderef, &mut state, out)
                    }
                    state.reset(rangeref);
                }
                out.push(postfix);
            },
            Node::Range(rangeref) => {
                if state.should_serialize(rangeref) {
                    for noderef in &self.ranges[rangeref] {
                        self.serialize_noderef(*noderef, &mut state, out)
                    }
                    state.reset(rangeref);
                }
            },
            Node::Token(token) => out.push(token),
        }
    }

    pub fn serialize<S: SerializeInto>(self: &Self, out: &mut S) {
        let mut state = SerializeState::new(&self.ranges[..]);

        for noderef in self.root.iter() {
            self.serialize_noderef(*noderef, &mut state, out);
        }
    }

    pub fn swap_ranges<R: Rng>(self: &mut Self, rng: &mut R) -> bool {
        match rand_indices::<R, _>(rng, &self.ranges[..]) {
            Some((index0, index1)) => {
                let mut ranges = self.ranges.to_mut();
                ranges.swap(index0, index1);
                true
            },
            None => false,
        }
    }

    pub fn shuffle_range<R: Rng>(self: &mut Self, rng: &mut R) -> bool {
        match rng.choose_mut(self.ranges.to_mut()) {
            Some(range) => {
                rng.shuffle(range);
                true
            },
            None => false,
        }
    }

    pub fn duplicate_range<R: Rng>(self: &mut Self, rng: &mut R, max_duplications: usize) -> bool {
        if max_duplications < 1 {
            return false
        }

        match rng.choose_mut(self.ranges.to_mut()) {
            Some(range) => {
                let num_duplications = rng.gen_range(1, max_duplications);
                let mut extension = Vec::new();
                for _ in 0..num_duplications {
                    extension.extend(&range[..])
                }
                range.extend(&extension[..]);
                true
            },
            None => false,
        }
    }

    pub fn duplicate_root_node<R: Rng>(self: &mut Self, rng: &mut R) -> bool {
        match rand_indices::<R, _>(rng, &self.ranges[..]) {
            Some((src_index, dst_index)) => {
                let mut nodes = self.nodes.to_mut();

                let dup_node = nodes[dst_index].clone();
                let noderef = nodes.len();
                nodes.push(dup_node);

                let ranges = self.ranges.to_mut();
                let range = vec![noderef, src_index];
                let rangeref = ranges.len();
                ranges.push(range);

                nodes[dst_index] = Node::Range(rangeref);
                true
            },
            None => false,
        }
    }

    pub fn remove_delim<R: Rng>(self: &mut Self, mut rng: &mut R) -> bool {
        match rand_delim(&mut rng, &self.nodes[..]) {
            Some((index, _, rangeref, _)) => {
                let mut nodes = self.nodes.to_mut();
                nodes[index] = Node::Range(rangeref);
                true
            },
            None => false,
        }
    }

    pub fn swap_delim<R: Rng>(self: &mut Self, mut rng: &mut R) -> bool {
        match rand_delim(&mut rng, &self.nodes[..]) {
            Some((index, start_pattern, rangeref, end_pattern)) => {
                let mut nodes = self.nodes.to_mut();
                nodes[index] = Node::Delim(end_pattern, rangeref, start_pattern);
                true
            },
            None => false,
        }
    }

    pub fn nest_delim<R: Rng>(self: &mut Self, mut rng: &mut R) -> bool {
        match rand_delim(&mut rng, &self.nodes[..]) {
            Some((index, start_pattern, rangeref, end_pattern)) => {
                let mut nodes = self.nodes.to_mut();
                let mut ranges = self.ranges.to_mut();

                let nested_noderef = nodes.len();
                nodes.push(Node::Delim(start_pattern, rangeref, end_pattern));

                let nested_rangeref = ranges.len();
                ranges.push(vec![nested_noderef]);

                nodes[index] = Node::Delim(start_pattern, nested_rangeref, end_pattern);
                true
            },
            None => false,
        }
    }
}

pub fn fuzz_one<'buf, 'parse, R: Rng>(parsed: &'parse ParsedFile<'buf>, mut rng: &mut R, config: &FuzzConfig) -> Option<FuzzFile<'buf, 'parse>> {
    let mut ff = FuzzFile::new(parsed);
    let mut did_mutate = false;
    for _ in 0..config.max_mutations {
        did_mutate |= match rng.gen() {
            FuzzAction::DuplicateRange => ff.duplicate_range(&mut rng, config.max_duplications),
            FuzzAction::DuplicateRootNode => ff.duplicate_root_node(&mut rng),
            FuzzAction::NestDelim => ff.nest_delim(&mut rng),
            FuzzAction::RemoveDelim => ff.remove_delim(&mut rng),
            FuzzAction::ShuffleRanges => ff.shuffle_range(&mut rng),
            FuzzAction::SwapDelim => ff.swap_delim(&mut rng),
            FuzzAction::SwapRanges => ff.swap_ranges(&mut rng),
        }
    }

    if did_mutate {
        Some(ff)
    } else {
        None
    }
}
<commit_msg>Add EmptyDelim mutator.<commit_after>extern crate rand;

use std::borrow::Cow;
use std::cmp;
use self::rand::{Rand, Rng};
use parse::{Node, NodeRef, ParsedFile, RangeRef};

#[derive(Debug)]
pub struct FuzzFile<'buf: 'parse, 'parse> {
    root: Cow<'parse, [NodeRef]>,
    nodes: Cow<'parse, [Node<'buf>]>,
    ranges: Cow<'parse, [Vec<NodeRef>]>,
}

#[derive(Clone)]
pub enum FuzzAction {
    DuplicateRange,
    DuplicateRootNode,
    EmptyDelim,
    NestDelim,
    RemoveDelim,
    ShuffleRanges,
    SwapDelim,
    SwapRanges,
}

pub struct FuzzConfig {
    pub max_mutations: usize,
    pub max_duplications: usize,
}

impl Rand for FuzzAction {
    fn rand<R: Rng>(rng: &mut R) -> Self {
        let actions = vec![
            FuzzAction::DuplicateRange,
            FuzzAction::DuplicateRootNode,
            FuzzAction::EmptyDelim,
            FuzzAction::NestDelim,
            FuzzAction::RemoveDelim,
            FuzzAction::ShuffleRanges,
            FuzzAction::SwapDelim,
            FuzzAction::SwapRanges,
        ];
        rng.choose(&actions).unwrap().clone()
    }
}

fn rand_indices<R: Rng, T>(mut rng: &mut Rng, x: &[T]) -> Option<(usize, usize)> {
    if x.len() > 1 {
        let indices = rand::sample(&mut rng, 0..x.len(), 2);
        Some((indices[0], indices[1]))
    } else {
        None
    }
}

fn rand_delim<'buf, R: Rng>(mut rng: &mut R, nodes: &[Node<'buf>]) -> Option<(NodeRef, &'buf [u8], RangeRef, &'buf [u8])> {
    let delims: Vec<_> = nodes.iter().enumerate().filter_map(|item| {
        match item {
            (index, &Node::Delim(start_pattern, rangeref, end_pattern)) => {
                Some((index, start_pattern, rangeref, end_pattern))
            },
            _ => None,
        }
    }).collect();
    rng.choose(&delims[..]).cloned()
}

struct SerializeState {
    have_serialized_range: Vec<bool>,
}

impl SerializeState {
    fn new(ranges: &[Vec<NodeRef>]) -> Self {
        SerializeState {
            have_serialized_range: ranges.iter().map(|_| false).collect()
        }
    }

    fn should_serialize(self: &mut Self, rangeref: RangeRef) -> bool {
        if self.have_serialized_range[rangeref] {
            false
        } else {
            self.have_serialized_range[rangeref] = true;
            true
        }
    }

    fn reset(self: &mut Self, rangeref: RangeRef) {
        self.have_serialized_range[rangeref] = false;
    }
}

pub trait SerializeInto {
    fn push(&mut self, &[u8]);
}

impl SerializeInto for Vec<u8> {
    fn push(self: &mut Self, token: &[u8]) {
        self.extend(token);
    }
}

pub struct SliceSerializer<'buf> {
    slice: &'buf mut [u8],
    cur_offset: usize,
}

impl<'buf> SliceSerializer<'buf> {
    pub fn new(slice: &'buf mut [u8]) -> SliceSerializer {
        SliceSerializer {
            slice: slice,
            cur_offset: 0,
        }
    }

    pub fn bytes_written(self: &Self) -> usize {
        self.cur_offset
    }
}

impl<'buf> SerializeInto for SliceSerializer<'buf> {
    fn push(self: &mut Self, token: &[u8]) {
        let remaining = self.slice.len() - self.cur_offset;
        let num_bytes_to_write = cmp::min(remaining, token.len());
        if num_bytes_to_write > 0 {
            self.slice[self.cur_offset..self.cur_offset+num_bytes_to_write].copy_from_slice(&token[..num_bytes_to_write]);
            self.cur_offset += num_bytes_to_write
        }
    }
}

impl<'buf, 'parse> FuzzFile<'buf, 'parse> {
    pub fn new(parsed: &'parse ParsedFile<'buf>) -> Self {
        FuzzFile {
            root: Cow::from(parsed.root.as_slice()),
            nodes: Cow::from(parsed.nodes.as_slice()),
            ranges: Cow::from(parsed.ranges.as_slice()),
        }
    }

    fn serialize_noderef<S: SerializeInto>(self: &Self, noderef: NodeRef, mut state: &mut SerializeState, mut out: &mut S) {
        match self.nodes[noderef] {
            Node::Delim(prefix, rangeref, postfix) => {
                out.push(prefix);
                if state.should_serialize(rangeref) {
                    for noderef in &self.ranges[rangeref] {
                        self.serialize_noderef(*noderef, &mut state, out)
                    }
                    state.reset(rangeref);
                }
                out.push(postfix);
            },
            Node::Range(rangeref) => {
                if state.should_serialize(rangeref) {
                    for noderef in &self.ranges[rangeref] {
                        self.serialize_noderef(*noderef, &mut state, out)
                    }
                    state.reset(rangeref);
                }
            },
            Node::Token(token) => out.push(token),
        }
    }

    pub fn serialize<S: SerializeInto>(self: &Self, out: &mut S) {
        let mut state = SerializeState::new(&self.ranges[..]);

        for noderef in self.root.iter() {
            self.serialize_noderef(*noderef, &mut state, out);
        }
    }

    pub fn swap_ranges<R: Rng>(self: &mut Self, rng: &mut R) -> bool {
        match rand_indices::<R, _>(rng, &self.ranges[..]) {
            Some((index0, index1)) => {
                let mut ranges = self.ranges.to_mut();
                ranges.swap(index0, index1);
                true
            },
            None => false,
        }
    }

    pub fn shuffle_range<R: Rng>(self: &mut Self, rng: &mut R) -> bool {
        match rng.choose_mut(self.ranges.to_mut()) {
            Some(range) => {
                rng.shuffle(range);
                true
            },
            None => false,
        }
    }

    pub fn duplicate_range<R: Rng>(self: &mut Self, rng: &mut R, max_duplications: usize) -> bool {
        if max_duplications < 1 {
            return false
        }

        match rng.choose_mut(self.ranges.to_mut()) {
            Some(range) => {
                let num_duplications = rng.gen_range(1, max_duplications);
                let mut extension = Vec::new();
                for _ in 0..num_duplications {
                    extension.extend(&range[..])
                }
                range.extend(&extension[..]);
                true
            },
            None => false,
        }
    }

    pub fn duplicate_root_node<R: Rng>(self: &mut Self, rng: &mut R) -> bool {
        match rand_indices::<R, _>(rng, &self.ranges[..]) {
            Some((src_index, dst_index)) => {
                let mut nodes = self.nodes.to_mut();

                let dup_node = nodes[dst_index].clone();
                let noderef = nodes.len();
                nodes.push(dup_node);

                let ranges = self.ranges.to_mut();
                let range = vec![noderef, src_index];
                let rangeref = ranges.len();
                ranges.push(range);

                nodes[dst_index] = Node::Range(rangeref);
                true
            },
            None => false,
        }
    }

    pub fn remove_delim<R: Rng>(self: &mut Self, mut rng: &mut R) -> bool {
        match rand_delim(&mut rng, &self.nodes[..]) {
            Some((index, _, rangeref, _)) => {
                let mut nodes = self.nodes.to_mut();
                nodes[index] = Node::Range(rangeref);
                true
            },
            None => false,
        }
    }

    pub fn swap_delim<R: Rng>(self: &mut Self, mut rng: &mut R) -> bool {
        match rand_delim(&mut rng, &self.nodes[..]) {
            Some((index, start_pattern, rangeref, end_pattern)) => {
                let mut nodes = self.nodes.to_mut();
                nodes[index] = Node::Delim(end_pattern, rangeref, start_pattern);
                true
            },
            None => false,
        }
    }

    pub fn nest_delim<R: Rng>(self: &mut Self, mut rng: &mut R) -> bool {
        match rand_delim(&mut rng, &self.nodes[..]) {
            Some((index, start_pattern, rangeref, end_pattern)) => {
                let mut nodes = self.nodes.to_mut();
                let mut ranges = self.ranges.to_mut();

                let nested_noderef = nodes.len();
                nodes.push(Node::Delim(start_pattern, rangeref, end_pattern));

                let nested_rangeref = ranges.len();
                ranges.push(vec![nested_noderef]);

                nodes[index] = Node::Delim(start_pattern, nested_rangeref, end_pattern);
                true
            },
            None => false,
        }
    }

    pub fn empty_delim<R: Rng>(self: &mut Self, mut rng: &mut R) -> bool {
        match rand_delim(&mut rng, &self.nodes[..]) {
            Some((index, start_pattern, _, end_pattern)) => {
                let mut nodes = self.nodes.to_mut();
                let mut ranges = self.ranges.to_mut();

                let rangeref = ranges.len();
                ranges.push(vec![]);

                nodes[index] = Node::Delim(start_pattern, rangeref, end_pattern);
                true
            },
            None => false,
        }
    }
}

pub fn fuzz_one<'buf, 'parse, R: Rng>(parsed: &'parse ParsedFile<'buf>, mut rng: &mut R, config: &FuzzConfig) -> Option<FuzzFile<'buf, 'parse>> {
    let mut ff = FuzzFile::new(parsed);
    let mut did_mutate = false;
    for _ in 0..config.max_mutations {
        did_mutate |= match rng.gen() {
            FuzzAction::DuplicateRange => ff.duplicate_range(&mut rng, config.max_duplications),
            FuzzAction::DuplicateRootNode => ff.duplicate_root_node(&mut rng),
            FuzzAction::EmptyDelim => ff.empty_delim(&mut rng),
            FuzzAction::NestDelim => ff.nest_delim(&mut rng),
            FuzzAction::RemoveDelim => ff.remove_delim(&mut rng),
            FuzzAction::ShuffleRanges => ff.shuffle_range(&mut rng),
            FuzzAction::SwapDelim => ff.swap_delim(&mut rng),
            FuzzAction::SwapRanges => ff.swap_ranges(&mut rng),
        }
    }

    if did_mutate {
        Some(ff)
    } else {
        None
    }
}
<|endoftext|>


<commit_before>// vim: tw=80
/// Generate Expectation and Expectations types with specific signatures
///
/// # Arguments
///
/// * `expectation`:    Name of the generated Expectation type
/// * `module`:         Name of the private module to create
/// * `o`:              Output type.  Must be static
/// * `methty`:         Comma-delimited sequence of arguments types for the
///                     method being mocked.
/// * `matchcall`:      Comma-delimited sequence of expressions that produce
///                     values of type `matchty` from values of type `methty`.
/// * `args`:           comma-delimited sequence of argument names for each
///                     argument.  Ideally this wouldn't need to be specified,
///                     but it is due to Rust's macro hygiene rules.
/// * `predargs`:       Comma-delimited sequence of identifiers of the same
///                     length as `args`, but distinct.
/// * `matchty`:        comma-delimited sequence of types for each match
///                     argument.  Must all be `'static`.
///
/// # Examples
///
/// ```
/// // Mock a method like foo(&self, x: u32, y: &i16) -> u32
/// omnimock!{(u32, &i16), u32, i, [&i.0, i.1], [i0, i1], [u32, i16]}
#[macro_export]
macro_rules! omnimock {(
        $expectation:ident,
        $module:ident,
        $o:ty,
        [ $( $methty:ty ),* ],
        [ $( $matchcall:expr ),* ],
        [ $( $args:ident ),* ],
        [ $( $predargs:ident ),* ],
        [ $( $matchty:ty ),* ]) =>
    {
        mod $module {

        enum Rfunc {
            Default,
            // Indicates that a `return_once` expectation has already returned
            Expired,
            Mut(Box<dyn FnMut($( $methty, )*) -> $o + Send>),
            // Should be Box<dyn FnOnce> once that feature is stabilized
            // https://github.com/rust-lang/rust/issues/28796
            Once(Box<dyn FnMut($( $methty, )*) -> $o + Send>),
        }

        impl  Rfunc {
            fn call_mut(&mut self, $( $args: $methty, )* ) -> $o {
                match self {
                    Rfunc::Default => {
                        unimplemented!()
                        //Self::return_default()
                    },
                    Rfunc::Expired => {
                        panic!("Called a method twice that was expected only once")
                    },
                    Rfunc::Mut(f) => {
                        f( $( $args, )* )
                    },
                    Rfunc::Once(_) => {
                        let fo = ::std::mem::replace(self, Rfunc::Expired);
                        if let Rfunc::Once(mut f) = fo {
                            f( $( $args, )* )
                        } else {
                            unreachable!()
                        }
                    },
                }
            }
        }

        impl std::default::Default for Rfunc {
            fn default() -> Self {
                Rfunc::Default
            }
        }

        enum Matcher {
            Func(Box<Fn($( &$matchty, )* ) -> bool>),
            Pred( $( Box<::mockall::Predicate<$matchty>>, )* )
        }

        impl Matcher {
            fn eval(&self, $( $args: &$matchty, )*) -> bool {
                match self {
                    Matcher::Func(f) => f($( $args, )*),
                    Matcher::Pred($( $predargs, )*) =>
                        [$( $predargs.eval($args), )*]
                        .into_iter()
                        .all(|x| *x),
                }
            }

            fn verify(&self, $( $args: &$matchty, )* ) {
                use predicates_tree::CaseTreeExt;
                match self {
                    Matcher::Func(f) => assert!(f($( $args, )*),
                        "Expectation didn't match arguments"),
                    Matcher::Pred($( $predargs, )*) => {
                        $(if let Some(c) = $predargs.find_case(false, $args) {
                            panic!("Expectation didn't match arguments:\n{}",
                                   c.tree());
                        })*
                    }
                }
            }
        }

        impl Default for Matcher {
            #[allow(unused_variables)]
            fn default() -> Self {
                Matcher::Func(Box::new(|$( $args, )*| true))
            }
        }

        #[derive(Default)]
        struct Common {
            matcher: ::std::sync::Mutex<Matcher>,
            seq_handle: Option<::mockall::SeqHandle>,
            times: ::mockall::Times
        }

        impl Common {
            fn call(&self, $( $args: &$matchty, )* ) {
                self.matcher.lock().unwrap().verify($( $args, )*);
                self.times.call();
                self.verify_sequence();
                if self.times.is_satisfied() {
                    self.satisfy_sequence()
                }
            }

            fn in_sequence(&mut self, seq: &mut ::mockall::Sequence)
                -> &mut Self
            {
                assert!(self.times.is_exact(),
                    "Only Expectations with an exact call count have sequences");
                self.seq_handle = Some(seq.next());
                self
            }

            /// Forbid this expectation from ever being called
            fn never(&mut self) {
                self.times.never();
            }

            fn satisfy_sequence(&self) {
                if let Some(handle) = &self.seq_handle {
                    handle.satisfy()
                }
            }

            /// Require this expectation to be called exactly `n` times.
            fn times(&mut self, n: usize) {
                self.times.n(n);
            }

            /// Allow this expectation to be called any number of times
            fn times_any(&mut self) {
                self.times.any();
            }

            /// Allow this expectation to be called any number of times within a given
            /// range
            fn times_range(&mut self, range: ::std::ops::Range<usize>) {
                self.times.range(range);
            }

            fn verify_sequence(&self) {
                if let Some(handle) = &self.seq_handle {
                    handle.verify()
                }
            }

            #[allow(non_camel_case_types)]  // Repurpose $predargs for generics
            fn with<$( $predargs: ::mockall::Predicate<$matchty> + 'static,)*>
                (&mut self, $( $args: $predargs,)*)
            {
                use ::std::ops::DerefMut;
                let mut guard = self.matcher.lock().unwrap();
                let m = Matcher::Pred($( Box::new($args), )*);
                ::std::mem::replace(guard.deref_mut(), m);
            }

            fn withf<F>(&mut self, f: F)
                where F: Fn($( &$matchty, )* ) -> bool + Send + 'static
            {
                use ::std::ops::DerefMut;
                let mut guard = self.matcher.lock().unwrap();
                let m = Matcher::Func(Box::new(f));
                ::std::mem::replace(guard.deref_mut(), m);
            }
        }

        #[derive(Default)]
        pub struct Expectation {
            common: Common,
            rfunc: ::std::sync::Mutex<Rfunc>,
        }

        impl Expectation {
            pub fn call(&self, $( $args: $methty, )* ) -> $o {
                self.common.call($( $matchcall, )*);
                self.rfunc.lock().unwrap().call_mut($( $args, )*)
            }

            pub fn in_sequence(&mut self, seq: &mut ::mockall::Sequence)
                -> &mut Self
            {
                self.common.in_sequence(seq);
                self
            }

            /// Forbid this expectation from ever being called
            pub fn never(&mut self) -> &mut Self {
                self.common.never();
                self
            }

            pub fn returning<F>(&mut self, f: F) -> &mut Self
                where F: FnMut($( $methty, )*) -> $o + Send + 'static
            {
                {
                    use ::std::ops::DerefMut;
                    let mut guard = self.rfunc.lock().unwrap();
                    ::std::mem::replace(guard.deref_mut(), Rfunc::Mut(Box::new(f)));
                }
                self
            }

            /// Require this expectation to be called exactly `n` times.
            pub fn times(&mut self, n: usize) -> &mut Self {
                self.common.times(n);
                self
            }

            /// Allow this expectation to be called any number of times
            ///
            /// This behavior is the default, but the method is provided in case the
            /// default behavior changes.
            pub fn times_any(&mut self) -> &mut Self {
                self.common.times_any();
                self
            }

            /// Allow this expectation to be called any number of times within a
            /// given range
            pub fn times_range(&mut self, range: ::std::ops::Range<usize>)
                -> &mut Self
            {
                self.common.times_range(range);
                self
            }

            #[allow(non_camel_case_types)]  // Repurpose $predargs for generics
            pub fn with<$( $predargs: ::mockall::Predicate<$matchty> + 'static,)*>
                (&mut self, $( $args: $predargs,)*) -> &mut Self
            {
                self.common.with($( $args, )*);
                self
            }

            pub fn withf<F>(&mut self, f: F) -> &mut Self
                where F: Fn($( &$matchty, )* ) -> bool + Send + 'static
            {
                self.common.withf(f);
                self
            }
        }

        }
        use $module::Expectation as $expectation;
    }
}
<commit_msg>omnimock: add return_once<commit_after>// vim: tw=80
/// Generate Expectation and Expectations types with specific signatures
///
/// # Arguments
///
/// * `expectation`:    Name of the generated Expectation type
/// * `module`:         Name of the private module to create
/// * `o`:              Output type.  Must be static
/// * `methty`:         Comma-delimited sequence of arguments types for the
///                     method being mocked.
/// * `matchcall`:      Comma-delimited sequence of expressions that produce
///                     values of type `matchty` from values of type `methty`.
/// * `args`:           comma-delimited sequence of argument names for each
///                     argument.  Ideally this wouldn't need to be specified,
///                     but it is due to Rust's macro hygiene rules.
/// * `predargs`:       Comma-delimited sequence of identifiers of the same
///                     length as `args`, but distinct.
/// * `matchty`:        comma-delimited sequence of types for each match
///                     argument.  Must all be `'static`.
///
/// # Examples
///
/// ```
/// // Mock a method like foo(&self, x: u32, y: &i16) -> u32
/// omnimock!{(u32, &i16), u32, i, [&i.0, i.1], [i0, i1], [u32, i16]}
#[macro_export]
macro_rules! omnimock {(
        $expectation:ident,
        $module:ident,
        $o:ty,
        [ $( $methty:ty ),* ],
        [ $( $matchcall:expr ),* ],
        [ $( $args:ident ),* ],
        [ $( $predargs:ident ),* ],
        [ $( $matchty:ty ),* ]) =>
    {
        mod $module {
            use super::*;

        enum Rfunc {
            Default,
            // Indicates that a `return_once` expectation has already returned
            Expired,
            Mut(Box<dyn FnMut($( $methty, )*) -> $o + Send>),
            // Should be Box<dyn FnOnce> once that feature is stabilized
            // https://github.com/rust-lang/rust/issues/28796
            Once(Box<dyn FnMut($( $methty, )*) -> $o + Send>),
        }

        impl  Rfunc {
            fn call_mut(&mut self, $( $args: $methty, )* ) -> $o {
                match self {
                    Rfunc::Default => {
                        unimplemented!()
                        //Self::return_default()
                    },
                    Rfunc::Expired => {
                        panic!("Called a method twice that was expected only once")
                    },
                    Rfunc::Mut(f) => {
                        f( $( $args, )* )
                    },
                    Rfunc::Once(_) => {
                        let fo = ::std::mem::replace(self, Rfunc::Expired);
                        if let Rfunc::Once(mut f) = fo {
                            f( $( $args, )* )
                        } else {
                            unreachable!()
                        }
                    },
                }
            }
        }

        impl std::default::Default for Rfunc {
            fn default() -> Self {
                Rfunc::Default
            }
        }

        enum Matcher {
            Func(Box<Fn($( &$matchty, )* ) -> bool>),
            Pred( $( Box<::mockall::Predicate<$matchty>>, )* )
        }

        impl Matcher {
            fn eval(&self, $( $args: &$matchty, )*) -> bool {
                match self {
                    Matcher::Func(f) => f($( $args, )*),
                    Matcher::Pred($( $predargs, )*) =>
                        [$( $predargs.eval($args), )*]
                        .into_iter()
                        .all(|x| *x),
                }
            }

            fn verify(&self, $( $args: &$matchty, )* ) {
                use predicates_tree::CaseTreeExt;
                match self {
                    Matcher::Func(f) => assert!(f($( $args, )*),
                        "Expectation didn't match arguments"),
                    Matcher::Pred($( $predargs, )*) => {
                        $(if let Some(c) = $predargs.find_case(false, $args) {
                            panic!("Expectation didn't match arguments:\n{}",
                                   c.tree());
                        })*
                    }
                }
            }
        }

        impl Default for Matcher {
            #[allow(unused_variables)]
            fn default() -> Self {
                Matcher::Func(Box::new(|$( $args, )*| true))
            }
        }

        #[derive(Default)]
        struct Common {
            matcher: ::std::sync::Mutex<Matcher>,
            seq_handle: Option<::mockall::SeqHandle>,
            times: ::mockall::Times
        }

        impl Common {
            fn call(&self, $( $args: &$matchty, )* ) {
                self.matcher.lock().unwrap().verify($( $args, )*);
                self.times.call();
                self.verify_sequence();
                if self.times.is_satisfied() {
                    self.satisfy_sequence()
                }
            }

            fn in_sequence(&mut self, seq: &mut ::mockall::Sequence)
                -> &mut Self
            {
                assert!(self.times.is_exact(),
                    "Only Expectations with an exact call count have sequences");
                self.seq_handle = Some(seq.next());
                self
            }

            /// Forbid this expectation from ever being called
            fn never(&mut self) {
                self.times.never();
            }

            fn satisfy_sequence(&self) {
                if let Some(handle) = &self.seq_handle {
                    handle.satisfy()
                }
            }

            /// Require this expectation to be called exactly `n` times.
            fn times(&mut self, n: usize) {
                self.times.n(n);
            }

            /// Allow this expectation to be called any number of times
            fn times_any(&mut self) {
                self.times.any();
            }

            /// Allow this expectation to be called any number of times within a given
            /// range
            fn times_range(&mut self, range: ::std::ops::Range<usize>) {
                self.times.range(range);
            }

            fn verify_sequence(&self) {
                if let Some(handle) = &self.seq_handle {
                    handle.verify()
                }
            }

            #[allow(non_camel_case_types)]  // Repurpose $predargs for generics
            fn with<$( $predargs: ::mockall::Predicate<$matchty> + 'static,)*>
                (&mut self, $( $args: $predargs,)*)
            {
                use ::std::ops::DerefMut;
                let mut guard = self.matcher.lock().unwrap();
                let m = Matcher::Pred($( Box::new($args), )*);
                ::std::mem::replace(guard.deref_mut(), m);
            }

            fn withf<F>(&mut self, f: F)
                where F: Fn($( &$matchty, )* ) -> bool + Send + 'static
            {
                use ::std::ops::DerefMut;
                let mut guard = self.matcher.lock().unwrap();
                let m = Matcher::Func(Box::new(f));
                ::std::mem::replace(guard.deref_mut(), m);
            }
        }

        #[derive(Default)]
        pub struct Expectation {
            common: Common,
            rfunc: ::std::sync::Mutex<Rfunc>,
        }

        impl Expectation {
            pub fn call(&self, $( $args: $methty, )* ) -> $o {
                self.common.call($( $matchcall, )*);
                self.rfunc.lock().unwrap().call_mut($( $args, )*)
            }

            pub fn in_sequence(&mut self, seq: &mut ::mockall::Sequence)
                -> &mut Self
            {
                self.common.in_sequence(seq);
                self
            }

            /// Forbid this expectation from ever being called
            pub fn never(&mut self) -> &mut Self {
                self.common.never();
                self
            }

            /// Supply an `FnOnce` closure that will provide the return value
            /// for this Expectation.  This is useful for return types that
            /// aren't `Clone`.  It will be an error to call this Expectation
            /// multiple times.
            pub fn return_once<F>(&mut self, f: F) -> &mut Self
                where F: FnOnce($( $methty, )*) -> $o + Send + 'static
            {
                use ::std::ops::DerefMut;
                let mut fopt = Some(f);
                let fmut = move |$( $args, )*| {
                    if let Some(f) = fopt.take() {
                        f($( $args, )*)
                    } else {
                        panic!("Called a method twice that was expected only once")
                    }
                };
                {
                    let mut guard = self.rfunc.lock().unwrap();
                    ::std::mem::replace(guard.deref_mut(),
                                        Rfunc::Once(Box::new(fmut)));
                }
                self
            }

            pub fn returning<F>(&mut self, f: F) -> &mut Self
                where F: FnMut($( $methty, )*) -> $o + Send + 'static
            {
                {
                    use ::std::ops::DerefMut;
                    let mut guard = self.rfunc.lock().unwrap();
                    ::std::mem::replace(guard.deref_mut(), Rfunc::Mut(Box::new(f)));
                }
                self
            }

            /// Require this expectation to be called exactly `n` times.
            pub fn times(&mut self, n: usize) -> &mut Self {
                self.common.times(n);
                self
            }

            /// Allow this expectation to be called any number of times
            ///
            /// This behavior is the default, but the method is provided in case the
            /// default behavior changes.
            pub fn times_any(&mut self) -> &mut Self {
                self.common.times_any();
                self
            }

            /// Allow this expectation to be called any number of times within a
            /// given range
            pub fn times_range(&mut self, range: ::std::ops::Range<usize>)
                -> &mut Self
            {
                self.common.times_range(range);
                self
            }

            #[allow(non_camel_case_types)]  // Repurpose $predargs for generics
            pub fn with<$( $predargs: ::mockall::Predicate<$matchty> + 'static,)*>
                (&mut self, $( $args: $predargs,)*) -> &mut Self
            {
                self.common.with($( $args, )*);
                self
            }

            pub fn withf<F>(&mut self, f: F) -> &mut Self
                where F: Fn($( &$matchty, )* ) -> bool + Send + 'static
            {
                self.common.withf(f);
                self
            }
        }

        }
        use $module::Expectation as $expectation;
    }
}
<|endoftext|>


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

use proc_macro2::TokenStream;
use syn::{self, DeriveInput};

type Fields = syn::punctuated::Punctuated<syn::Field, syn::token::Comma>;

fn derive_trait(
    input: &DeriveInput,
    trait_name: TokenStream,
    generics: &syn::Generics,
    body: impl FnOnce() -> TokenStream,
) -> TokenStream {
    let struct_name = &input.ident;

    let (impl_generics, _, where_clause) = generics.split_for_impl();
    let (_, ty_generics, _) = input.generics.split_for_impl();

    let body = body();
    quote! {
        impl #impl_generics #trait_name for #struct_name #ty_generics #where_clause {
            #body
        }
    }
}

fn derive_simple_trait(
    input: &DeriveInput,
    trait_name: TokenStream,
    t: &syn::TypeParam,
    body: impl FnOnce() -> TokenStream,
) -> TokenStream {
    let mut generics = input.generics.clone();
    generics
        .make_where_clause()
        .predicates
        .push(parse_quote!(#t: #trait_name));
    derive_trait(input, trait_name, &generics, body)
}

fn each_field_except_unit(
    fields: &Fields,
    unit: &syn::Field,
    mut field_expr: impl FnMut(&syn::Ident) -> TokenStream,
) -> TokenStream {
    fields.iter().filter(|f| f.ident != unit.ident).fold(quote! {}, |body, field| {
        let name = field.ident.as_ref().unwrap();
        let expr = field_expr(name);
        quote! {
            #body
            #expr
        }
    })
}

fn derive_struct_body(
    fields: &Fields,
    unit: &syn::Field,
    mut field_expr: impl FnMut(&syn::Ident) -> TokenStream,
) -> TokenStream {
    let body = each_field_except_unit(fields, unit, |name| {
        let expr = field_expr(name);
        quote! {
            #name: #expr,
        }
    });

    let unit_name = unit.ident.as_ref().unwrap();
    quote! {
        Self {
            #body
            #unit_name: PhantomData,
        }
    }
}

fn clone_impl(input: &DeriveInput, fields: &Fields, unit: &syn::Field, t: &syn::TypeParam) -> TokenStream {
    derive_simple_trait(input, quote! { Clone }, t, || {
        let body = derive_struct_body(fields, unit, |name| {
            quote! { self.#name.clone() }
        });
        quote! {
            fn clone(&self) -> Self {
                #body
            }
        }
    })
}

fn copy_impl(input: &DeriveInput, t: &syn::TypeParam) -> TokenStream {
    derive_simple_trait(input, quote!{ Copy }, t, || quote! {})
}

fn eq_impl(input: &DeriveInput, t: &syn::TypeParam) -> TokenStream {
    derive_simple_trait(input, quote!{ ::core::cmp::Eq }, t, || quote! {})
}

fn partialeq_impl(input: &DeriveInput, fields: &Fields, unit: &syn::Field, t: &syn::TypeParam) -> TokenStream {
    derive_simple_trait(input, quote!{ ::core::cmp::PartialEq }, t, || {
        let body = each_field_except_unit(fields, unit, |name| {
            quote! { && self.#name == other.#name }
        });

        quote! {
            fn eq(&self, other: &Self) -> bool {
                true #body
            }
        }
    })
}

fn hash_impl(input: &DeriveInput, fields: &Fields, unit: &syn::Field, t: &syn::TypeParam) -> TokenStream {
    derive_simple_trait(input, quote!{ ::core::hash::Hash }, t, || {
        let body = each_field_except_unit(fields, unit, |name| {
            quote! { self.#name.hash(h); }
        });

        quote! {
            fn hash<H: ::core::hash::Hasher>(&self, h: &mut H) {
                #body
            }
        }
    })
}

fn serde_impl(
    input: &DeriveInput,
    fields: &Fields,
    unit: &syn::Field,
    t: &syn::TypeParam,
) -> TokenStream {
    let deserialize_impl = {
        let mut generics = input.generics.clone();
        generics.params.insert(0, parse_quote!('de));
        generics
            .make_where_clause()
            .predicates
            .push(parse_quote!(#t: ::serde::Deserialize<'de>));
        derive_trait(input, quote!{ ::serde::Deserialize<'de> }, &generics, || {
            let tuple = each_field_except_unit(fields, unit, |name| {
                quote! { #name, }
            });
            let body = derive_struct_body(fields, unit, |name| quote! { #name });
            quote! {
                fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
                where
                    D: ::serde::Deserializer<'de>,
                {
                    let (#tuple) = ::serde::Deserialize::deserialize(deserializer)?;
                    Ok(#body)
                }
            }
        })
    };

    let serialize_impl = derive_simple_trait(input, quote! { ::serde::Serialize }, t, || {
        let tuple = each_field_except_unit(fields, unit, |name| {
            quote! { &self.#name, }
        });
        quote! {
            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
            where
                S: ::serde::Serializer,
            {
                (#tuple).serialize(serializer)
            }
        }
    });

    quote! {
        #[cfg(feature = "serde")]
        #serialize_impl
        #[cfg(feature = "serde")]
        #deserialize_impl
    }
}

pub fn derive(input: DeriveInput) -> TokenStream {
    let s = match input.data {
        syn::Data::Struct(ref s) => s,
        _ => panic!("Need to derive this on a struct"),
    };

    let fields = match s.fields {
        syn::Fields::Named(ref named) => &named.named,
        _ => panic!("Need to use named fields"),
    };

    assert!(!fields.is_empty());

    let unit_field = fields.last().unwrap();
    assert_eq!(
        unit_field.value().ident.as_ref().unwrap().to_string(),
        "_unit",
        "You need to have a _unit field to derive this trait",
    );

    let type_param =
        input.generics.type_params().next().cloned().expect("Need a T");

    let clone = clone_impl(&input, fields, unit_field.value(), &type_param);
    let copy = copy_impl(&input, &type_param);
    let serde = serde_impl(&input, fields, unit_field.value(), &type_param);
    let eq = eq_impl(&input, &type_param);
    let partialeq = partialeq_impl(&input, fields, unit_field.value(), &type_param);
    let hash = hash_impl(&input, fields, unit_field.value(), &type_param);

    quote! {
        #clone
        #copy
        #serde
        #eq
        #partialeq
        #hash
    }
}
<commit_msg>macros: Ensure the structs are repr(C) and the unit field is public.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

use proc_macro2::TokenStream;
use syn::{self, DeriveInput};

type Fields = syn::punctuated::Punctuated<syn::Field, syn::token::Comma>;

fn derive_trait(
    input: &DeriveInput,
    trait_name: TokenStream,
    generics: &syn::Generics,
    body: impl FnOnce() -> TokenStream,
) -> TokenStream {
    let struct_name = &input.ident;

    let (impl_generics, _, where_clause) = generics.split_for_impl();
    let (_, ty_generics, _) = input.generics.split_for_impl();

    let body = body();
    quote! {
        impl #impl_generics #trait_name for #struct_name #ty_generics #where_clause {
            #body
        }
    }
}

fn derive_simple_trait(
    input: &DeriveInput,
    trait_name: TokenStream,
    t: &syn::TypeParam,
    body: impl FnOnce() -> TokenStream,
) -> TokenStream {
    let mut generics = input.generics.clone();
    generics
        .make_where_clause()
        .predicates
        .push(parse_quote!(#t: #trait_name));
    derive_trait(input, trait_name, &generics, body)
}

fn each_field_except_unit(
    fields: &Fields,
    unit: &syn::Field,
    mut field_expr: impl FnMut(&syn::Ident) -> TokenStream,
) -> TokenStream {
    fields.iter().filter(|f| f.ident != unit.ident).fold(quote! {}, |body, field| {
        let name = field.ident.as_ref().unwrap();
        let expr = field_expr(name);
        quote! {
            #body
            #expr
        }
    })
}

fn derive_struct_body(
    fields: &Fields,
    unit: &syn::Field,
    mut field_expr: impl FnMut(&syn::Ident) -> TokenStream,
) -> TokenStream {
    let body = each_field_except_unit(fields, unit, |name| {
        let expr = field_expr(name);
        quote! {
            #name: #expr,
        }
    });

    let unit_name = unit.ident.as_ref().unwrap();
    quote! {
        Self {
            #body
            #unit_name: PhantomData,
        }
    }
}

fn clone_impl(input: &DeriveInput, fields: &Fields, unit: &syn::Field, t: &syn::TypeParam) -> TokenStream {
    derive_simple_trait(input, quote! { Clone }, t, || {
        let body = derive_struct_body(fields, unit, |name| {
            quote! { self.#name.clone() }
        });
        quote! {
            fn clone(&self) -> Self {
                #body
            }
        }
    })
}

fn copy_impl(input: &DeriveInput, t: &syn::TypeParam) -> TokenStream {
    derive_simple_trait(input, quote!{ Copy }, t, || quote! {})
}

fn eq_impl(input: &DeriveInput, t: &syn::TypeParam) -> TokenStream {
    derive_simple_trait(input, quote!{ ::core::cmp::Eq }, t, || quote! {})
}

fn partialeq_impl(input: &DeriveInput, fields: &Fields, unit: &syn::Field, t: &syn::TypeParam) -> TokenStream {
    derive_simple_trait(input, quote!{ ::core::cmp::PartialEq }, t, || {
        let body = each_field_except_unit(fields, unit, |name| {
            quote! { && self.#name == other.#name }
        });

        quote! {
            fn eq(&self, other: &Self) -> bool {
                true #body
            }
        }
    })
}

fn hash_impl(input: &DeriveInput, fields: &Fields, unit: &syn::Field, t: &syn::TypeParam) -> TokenStream {
    derive_simple_trait(input, quote!{ ::core::hash::Hash }, t, || {
        let body = each_field_except_unit(fields, unit, |name| {
            quote! { self.#name.hash(h); }
        });

        quote! {
            fn hash<H: ::core::hash::Hasher>(&self, h: &mut H) {
                #body
            }
        }
    })
}

fn serde_impl(
    input: &DeriveInput,
    fields: &Fields,
    unit: &syn::Field,
    t: &syn::TypeParam,
) -> TokenStream {
    let deserialize_impl = {
        let mut generics = input.generics.clone();
        generics.params.insert(0, parse_quote!('de));
        generics
            .make_where_clause()
            .predicates
            .push(parse_quote!(#t: ::serde::Deserialize<'de>));
        derive_trait(input, quote!{ ::serde::Deserialize<'de> }, &generics, || {
            let tuple = each_field_except_unit(fields, unit, |name| {
                quote! { #name, }
            });
            let body = derive_struct_body(fields, unit, |name| quote! { #name });
            quote! {
                fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
                where
                    D: ::serde::Deserializer<'de>,
                {
                    let (#tuple) = ::serde::Deserialize::deserialize(deserializer)?;
                    Ok(#body)
                }
            }
        })
    };

    let serialize_impl = derive_simple_trait(input, quote! { ::serde::Serialize }, t, || {
        let tuple = each_field_except_unit(fields, unit, |name| {
            quote! { &self.#name, }
        });
        quote! {
            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
            where
                S: ::serde::Serializer,
            {
                (#tuple).serialize(serializer)
            }
        }
    });

    quote! {
        #[cfg(feature = "serde")]
        #serialize_impl
        #[cfg(feature = "serde")]
        #deserialize_impl
    }
}

pub fn derive(input: DeriveInput) -> TokenStream {
    let s = match input.data {
        syn::Data::Struct(ref s) => s,
        _ => panic!("Need to derive this on a struct"),
    };

    let fields = match s.fields {
        syn::Fields::Named(ref named) => &named.named,
        _ => panic!("Need to use named fields"),
    };

    assert!(!fields.is_empty());

    let unit_field = fields.last().unwrap();
    assert_eq!(
        unit_field.value().ident.as_ref().unwrap().to_string(),
        "_unit",
        "You need to have a _unit field to derive this trait",
    );

    assert!(match unit_field.value().vis {
        syn::Visibility::Public(..) => true,
        _ => false,
    }, "Unit field should be public");

    assert!(input.attrs.iter().filter_map(|attr| attr.interpret_meta()).any(|attr| {
        match attr {
            syn::Meta::Word(..) |
            syn::Meta::NameValue(..) => false,
            syn::Meta::List(ref list) => {
                list.ident == "repr" && list.nested.iter().any(|meta| {
                    match meta {
                        syn::NestedMeta::Meta(syn::Meta::Word(ref w)) => w == "C",
                        _ => false,
                    }
                })
            }
        }
    }), "struct should be #[repr(C)]");

    let type_param =
        input.generics.type_params().next().cloned().expect("Need a T");

    let clone = clone_impl(&input, fields, unit_field.value(), &type_param);
    let copy = copy_impl(&input, &type_param);
    let serde = serde_impl(&input, fields, unit_field.value(), &type_param);
    let eq = eq_impl(&input, &type_param);
    let partialeq = partialeq_impl(&input, fields, unit_field.value(), &type_param);
    let hash = hash_impl(&input, fields, unit_field.value(), &type_param);

    quote! {
        #clone
        #copy
        #serde
        #eq
        #partialeq
        #hash
    }
}
<|endoftext|>


<commit_before>//! Encoding of AVIF images.
///
/// The [AVIF] specification defines an image derivative of the AV1 bitstream, an open video codec.
///
/// [AVIF]: https://aomediacodec.github.io/av1-avif/
use std::borrow::Cow;
use std::io::Write;
use std::cmp::{min, max};

use crate::{ColorType, ImageBuffer, ImageFormat, Pixel};
use crate::{ImageError, ImageResult};
use crate::buffer::ConvertBuffer;
use crate::color::{FromColor, Luma, LumaA, Bgr, Bgra, Rgb, Rgba};
use crate::error::{EncodingError, ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind};

use bytemuck::{Pod, PodCastError, try_cast_slice, try_cast_slice_mut};
use num_traits::Zero;
use ravif::{Img, ColorSpace, Config, RGBA8, encode_rgba};
use rgb::AsPixels;

/// AVIF Encoder.
///
/// Writes one image into the chosen output.
pub struct AvifEncoder<W> {
    inner: W,
    fallback: Vec<u8>,
}

impl<W: Write> AvifEncoder<W> {
    /// Create a new encoder that writes its output to `w`.
    pub fn new(w: W) -> Self {
        AvifEncoder { inner: w, fallback: vec![] }
    }

    /// Encode image data with the indicated color type.
    ///
    /// The encoder currently requires all data to be RGBA8, it will be converted internally if
    /// necessary. When data is suitably aligned, i.e. u16 channels to two bytes, then the
    /// conversion may be more efficient.
    pub fn write_image(mut self, data: &[u8], width: u32, height: u32, color: ColorType) -> ImageResult<()> {
        self.write_image_with_speed_quality(data, width, height, color, 100, 1)
    }

    /// Encode image data with the indicated color type, quality and speed.
    ///
    /// The encoder currently requires all data to be RGBA8, it will be converted internally if
    /// necessary. When data is suitably aligned, i.e. u16 channels to two bytes, then the
    /// conversion may be more efficient.
    pub fn write_image_with_speed_quality(mut self, data: &[u8], width: u32, height: u32, color: ColorType, quality: u8, speed: u8) -> ImageResult<()> {
        // Clamp quality and speed to range
        let quality = max(1, min(quality, 100));
        let speed = max(1, min(speed, 10));

        let config = self.config(color, quality, speed);
        // `ravif` needs strongly typed data so let's convert. We can either use a temporarily
        // owned version in our own buffer or zero-copy if possible by using the input buffer.
        // This requires going through `rgb`.
        let buffer = self.encode_as_img(data, width, height, color)?;
        let (data, _color_size, _alpha_size) = encode_rgba(buffer, &config)
            .map_err(|err| ImageError::Encoding(
                EncodingError::new(ImageFormat::Avif.into(), err)
            ))?;
        self.inner.write_all(&data)?;
        Ok(())
    }

    fn config(&self, _color: ColorType, quality: u8, speed: u8) -> Config {
        Config {
            quality: quality,
            alpha_quality: quality,
            speed: speed,
            premultiplied_alpha: false,
            color_space: ColorSpace::RGB,
        }
    }

    fn encode_as_img<'buf>(&'buf mut self, data: &'buf [u8], width: u32, height: u32, color: ColorType)
        -> ImageResult<Img<&'buf [RGBA8]>>
    {
        // Error wrapping utility for color dependent buffer dimensions.
        fn try_from_raw<P: Pixel + 'static>(data: &[P::Subpixel], width: u32, height: u32)
            -> ImageResult<ImageBuffer<P, &[P::Subpixel]>>
        {
            ImageBuffer::from_raw(width, height, data).ok_or_else(|| {
                ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch))
            })
        };

        // Convert to target color type using few buffer allocations.
        fn convert_into<'buf, P>(buf: &'buf mut Vec<u8>, image: ImageBuffer<P, &[P::Subpixel]>)
            -> Img<&'buf [RGBA8]>
        where
            P: Pixel + 'static,
            Rgba<u8>: FromColor<P>,
        {
            let (width, height) = image.dimensions();
            // TODO: conversion re-using the target buffer?
            let image: ImageBuffer<Rgba<u8>, _> = image.convert();
            *buf = image.into_raw();
            Img::new(buf.as_pixels(), width as usize, height as usize)
        }

        // Cast the input slice using few buffer allocations if possible.
        // In particular try not to allocate if the caller did the infallible reverse.
        fn cast_buffer<Channel>(buf: &[u8]) -> ImageResult<Cow<[Channel]>>
        where
            Channel: Pod + Zero,
        {
            match try_cast_slice(buf) {
                Ok(slice) => Ok(Cow::Borrowed(slice)),
                Err(PodCastError::OutputSliceWouldHaveSlop) => {
                    Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                }
                Err(PodCastError::TargetAlignmentGreaterAndInputNotAligned) => {
                    // Sad, but let's allocate.
                    // bytemuck checks alignment _before_ slop but size mismatch before this..
                    if buf.len() % std::mem::size_of::<Channel>() != 0 {
                        Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                    } else {
                        let len = buf.len() / std::mem::size_of::<Channel>();
                        let mut data = vec![Channel::zero(); len];
                        let view = try_cast_slice_mut::<_, u8>(data.as_mut_slice()).unwrap();
                        view.copy_from_slice(buf);
                        Ok(Cow::Owned(data))
                    }
                }
                Err(err) => { // Are you trying to encode a ZST??
                    Err(ImageError::Parameter(ParameterError::from_kind(
                        ParameterErrorKind::Generic(format!("{:?}", err))
                    )))
                }
            }
        }

        match color {
            ColorType::Rgba8 => {
                // ravif doesn't do any checks but has some asserts, so we do the checks.
                let img = try_from_raw::<Rgba<u8>>(data, width, height)?;
                // Now, internally ravif uses u32 but it takes usize. We could do some checked
                // conversion but instead we use that a non-empty image must be addressable.
                if img.pixels().len() == 0 {
                    return Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)));
                }

                Ok(Img::new(rgb::AsPixels::as_pixels(data), width as usize, height as usize))
            },
            // we need a separate buffer..
            ColorType::L8 => {
                let image = try_from_raw::<Luma<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La8 => {
                let image = try_from_raw::<LumaA<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb8 => {
                let image = try_from_raw::<Rgb<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgr8 => {
                let image = try_from_raw::<Bgr<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgra8 => {
                let image = try_from_raw::<Bgra<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // we need to really convert data..
            ColorType::L16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Luma<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<LumaA<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgb<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgba16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgba<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // for cases we do not support at all?
            _ => Err(ImageError::Unsupported(UnsupportedError::from_format_and_kind(
                    ImageFormat::Avif.into(),
                    UnsupportedErrorKind::Color(color.into()),
                )))
        }
    }
}
<commit_msg>Set quality minimum to 0.<commit_after>//! Encoding of AVIF images.
///
/// The [AVIF] specification defines an image derivative of the AV1 bitstream, an open video codec.
///
/// [AVIF]: https://aomediacodec.github.io/av1-avif/
use std::borrow::Cow;
use std::io::Write;
use std::cmp::{min, max};

use crate::{ColorType, ImageBuffer, ImageFormat, Pixel};
use crate::{ImageError, ImageResult};
use crate::buffer::ConvertBuffer;
use crate::color::{FromColor, Luma, LumaA, Bgr, Bgra, Rgb, Rgba};
use crate::error::{EncodingError, ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind};

use bytemuck::{Pod, PodCastError, try_cast_slice, try_cast_slice_mut};
use num_traits::Zero;
use ravif::{Img, ColorSpace, Config, RGBA8, encode_rgba};
use rgb::AsPixels;

/// AVIF Encoder.
///
/// Writes one image into the chosen output.
pub struct AvifEncoder<W> {
    inner: W,
    fallback: Vec<u8>,
}

impl<W: Write> AvifEncoder<W> {
    /// Create a new encoder that writes its output to `w`.
    pub fn new(w: W) -> Self {
        AvifEncoder { inner: w, fallback: vec![] }
    }

    /// Encode image data with the indicated color type.
    ///
    /// The encoder currently requires all data to be RGBA8, it will be converted internally if
    /// necessary. When data is suitably aligned, i.e. u16 channels to two bytes, then the
    /// conversion may be more efficient.
    pub fn write_image(mut self, data: &[u8], width: u32, height: u32, color: ColorType) -> ImageResult<()> {
        self.write_image_with_speed_quality(data, width, height, color, 100, 1)
    }

    /// Encode image data with the indicated color type, quality and speed.
    ///
    /// The encoder currently requires all data to be RGBA8, it will be converted internally if
    /// necessary. When data is suitably aligned, i.e. u16 channels to two bytes, then the
    /// conversion may be more efficient.
    pub fn write_image_with_speed_quality(mut self, data: &[u8], width: u32, height: u32, color: ColorType, quality: u8, speed: u8) -> ImageResult<()> {
        // Clamp quality and speed to range
        let quality = max(0, min(quality, 100));
        let speed = max(1, min(speed, 10));

        let config = self.config(color, quality, speed);
        // `ravif` needs strongly typed data so let's convert. We can either use a temporarily
        // owned version in our own buffer or zero-copy if possible by using the input buffer.
        // This requires going through `rgb`.
        let buffer = self.encode_as_img(data, width, height, color)?;
        let (data, _color_size, _alpha_size) = encode_rgba(buffer, &config)
            .map_err(|err| ImageError::Encoding(
                EncodingError::new(ImageFormat::Avif.into(), err)
            ))?;
        self.inner.write_all(&data)?;
        Ok(())
    }

    fn config(&self, _color: ColorType, quality: u8, speed: u8) -> Config {
        Config {
            quality: quality,
            alpha_quality: quality,
            speed: speed,
            premultiplied_alpha: false,
            color_space: ColorSpace::RGB,
        }
    }

    fn encode_as_img<'buf>(&'buf mut self, data: &'buf [u8], width: u32, height: u32, color: ColorType)
        -> ImageResult<Img<&'buf [RGBA8]>>
    {
        // Error wrapping utility for color dependent buffer dimensions.
        fn try_from_raw<P: Pixel + 'static>(data: &[P::Subpixel], width: u32, height: u32)
            -> ImageResult<ImageBuffer<P, &[P::Subpixel]>>
        {
            ImageBuffer::from_raw(width, height, data).ok_or_else(|| {
                ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch))
            })
        };

        // Convert to target color type using few buffer allocations.
        fn convert_into<'buf, P>(buf: &'buf mut Vec<u8>, image: ImageBuffer<P, &[P::Subpixel]>)
            -> Img<&'buf [RGBA8]>
        where
            P: Pixel + 'static,
            Rgba<u8>: FromColor<P>,
        {
            let (width, height) = image.dimensions();
            // TODO: conversion re-using the target buffer?
            let image: ImageBuffer<Rgba<u8>, _> = image.convert();
            *buf = image.into_raw();
            Img::new(buf.as_pixels(), width as usize, height as usize)
        }

        // Cast the input slice using few buffer allocations if possible.
        // In particular try not to allocate if the caller did the infallible reverse.
        fn cast_buffer<Channel>(buf: &[u8]) -> ImageResult<Cow<[Channel]>>
        where
            Channel: Pod + Zero,
        {
            match try_cast_slice(buf) {
                Ok(slice) => Ok(Cow::Borrowed(slice)),
                Err(PodCastError::OutputSliceWouldHaveSlop) => {
                    Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                }
                Err(PodCastError::TargetAlignmentGreaterAndInputNotAligned) => {
                    // Sad, but let's allocate.
                    // bytemuck checks alignment _before_ slop but size mismatch before this..
                    if buf.len() % std::mem::size_of::<Channel>() != 0 {
                        Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                    } else {
                        let len = buf.len() / std::mem::size_of::<Channel>();
                        let mut data = vec![Channel::zero(); len];
                        let view = try_cast_slice_mut::<_, u8>(data.as_mut_slice()).unwrap();
                        view.copy_from_slice(buf);
                        Ok(Cow::Owned(data))
                    }
                }
                Err(err) => { // Are you trying to encode a ZST??
                    Err(ImageError::Parameter(ParameterError::from_kind(
                        ParameterErrorKind::Generic(format!("{:?}", err))
                    )))
                }
            }
        }

        match color {
            ColorType::Rgba8 => {
                // ravif doesn't do any checks but has some asserts, so we do the checks.
                let img = try_from_raw::<Rgba<u8>>(data, width, height)?;
                // Now, internally ravif uses u32 but it takes usize. We could do some checked
                // conversion but instead we use that a non-empty image must be addressable.
                if img.pixels().len() == 0 {
                    return Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)));
                }

                Ok(Img::new(rgb::AsPixels::as_pixels(data), width as usize, height as usize))
            },
            // we need a separate buffer..
            ColorType::L8 => {
                let image = try_from_raw::<Luma<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La8 => {
                let image = try_from_raw::<LumaA<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb8 => {
                let image = try_from_raw::<Rgb<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgr8 => {
                let image = try_from_raw::<Bgr<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgra8 => {
                let image = try_from_raw::<Bgra<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // we need to really convert data..
            ColorType::L16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Luma<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<LumaA<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgb<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgba16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgba<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // for cases we do not support at all?
            _ => Err(ImageError::Unsupported(UnsupportedError::from_format_and_kind(
                    ImageFormat::Avif.into(),
                    UnsupportedErrorKind::Color(color.into()),
                )))
        }
    }
}
<|endoftext|>


<commit_before>use std::{mem, ptr};
use libc::{c_int, socklen_t};
use fcntl::{Fd, fcntl, F_SETFL, F_SETFD, FD_CLOEXEC, O_NONBLOCK};
use errno::{SysResult, SysError, from_ffi};
use features;

pub use libc::{in_addr, sockaddr_in, sockaddr_in6, sockaddr_un, sa_family_t, ip_mreq};

pub use self::consts::*;

mod ffi {
    use libc::{c_int, c_void, socklen_t};
    pub use libc::{socket, listen, bind, accept, connect, setsockopt};

    extern {
        pub fn getsockopt(
            sockfd: c_int,
            level: c_int,
            optname: c_int,
            optval: *mut c_void,
            optlen: *mut socklen_t) -> c_int;
    }
}

// Extra flags - Supported by Linux 2.6.27, normalized on other platforms
bitflags!(
    flags SockFlag: c_int {
        const SOCK_NONBLOCK = 0o0004000,
        const SOCK_CLOEXEC  = 0o2000000
    }
)

pub enum SockAddr {
    SockIpV4(sockaddr_in),
    SockIpV6(sockaddr_in6),
    SockUnix(sockaddr_un)
}

#[cfg(target_os = "linux")]
mod consts {
    use libc::{c_int};

    pub type AddressFamily = c_int;

    pub const AF_UNIX: AddressFamily  = 1;
    pub const AF_LOCAL: AddressFamily = AF_UNIX;
    pub const AF_INET: AddressFamily  = 2;
    pub const AF_INET6: AddressFamily = 10;

    pub type SockType = c_int;

    pub const SOCK_STREAM: SockType = 1;
    pub const SOCK_DGRAM: SockType = 2;
    pub const SOCK_SEQPACKET: SockType = 5;
    pub const SOCK_RAW: SockType = 3;
    pub const SOCK_RDM: SockType = 4;

    pub type SockLevel = c_int;

    pub const SOL_IP: SockLevel     = 0;
    pub const SOL_SOCKET: SockLevel = 1;
    pub const SOL_TCP: SockLevel    = 6;
    pub static IPPROTO_TCP: SockLevel = SOL_TCP;
    pub const SOL_UDP: SockLevel    = 17;
    pub const SOL_IPV6: SockLevel   = 41;

    pub type SockOpt = c_int;

    pub const SO_ACCEPTCONN: SockOpt = 30;
    pub const SO_BINDTODEVICE: SockOpt = 25;
    pub const SO_BROADCAST: SockOpt = 6;
    pub const SO_BSDCOMPAT: SockOpt = 14;
    pub const SO_DEBUG: SockOpt = 1;
    pub const SO_DOMAIN: SockOpt = 39;
    pub const SO_ERROR: SockOpt = 4;
    pub const SO_DONTROUTE: SockOpt = 5;
    pub const SO_KEEPALIVE: SockOpt = 9;
    pub const SO_LINGER: SockOpt = 13;
    pub const SO_MARK: SockOpt = 36;
    pub const SO_OOBINLINE: SockOpt = 10;
    pub const SO_PASSCRED: SockOpt = 16;
    pub const SO_PEEK_OFF: SockOpt = 42;
    pub const SO_PEERCRED: SockOpt = 17;
    pub const SO_PRIORITY: SockOpt = 12;
    pub const SO_PROTOCOL: SockOpt = 38;
    pub const SO_RCVBUF: SockOpt = 8;
    pub const SO_RCVBUFFORCE: SockOpt = 33;
    pub const SO_RCVLOWAT: SockOpt = 18;
    pub const SO_SNDLOWAT: SockOpt = 19;
    pub const SO_RCVTIMEO: SockOpt = 20;
    pub const SO_SNDTIMEO: SockOpt = 21;
    pub const SO_REUSEADDR: SockOpt = 2;
    pub const SO_REUSEPORT: SockOpt = 15;
    pub const SO_RXQ_OVFL: SockOpt = 40;
    pub const SO_SNDBUF: SockOpt = 7;
    pub const SO_SNDBUFFORCE: SockOpt = 32;
    pub const SO_TIMESTAMP: SockOpt = 29;
    pub const SO_TYPE: SockOpt = 3;
    pub const SO_BUSY_POLL: SockOpt = 46;

    // Socket options for TCP sockets
    pub const TCP_NODELAY: SockOpt = 1;
    pub const TCP_MAXSEG: SockOpt = 2;
    pub const TCP_CORK: SockOpt = 3;

    // Socket options for the IP layer of the socket
    pub const IP_MULTICAST_IF: SockOpt = 32;
    pub const IP_MULTICAST_TTL: SockOpt = 33;
    pub const IP_MULTICAST_LOOP: SockOpt = 34;
    pub const IP_ADD_MEMBERSHIP: SockOpt = 35;
    pub const IP_DROP_MEMBERSHIP: SockOpt = 36;
}

#[cfg(any(target_os = "macos", target_os = "ios"))]
mod consts {
    use libc::{c_int};

    pub type AddressFamily = c_int;

    pub const AF_UNIX: AddressFamily  = 1;
    pub const AF_LOCAL: AddressFamily = AF_UNIX;
    pub const AF_INET: AddressFamily  = 2;
    pub const AF_INET6: AddressFamily = 30;

    pub type SockType = c_int;

    pub const SOCK_STREAM: SockType = 1;
    pub const SOCK_DGRAM: SockType = 2;
    pub const SOCK_SEQPACKET: SockType = 5;
    pub const SOCK_RAW: SockType = 3;
    pub const SOCK_RDM: SockType = 4;

    pub type SockLevel = c_int;

    pub const SOL_SOCKET: SockLevel = 0xffff;

    pub type SockOpt = c_int;

    pub const SO_ACCEPTCONN: SockOpt          = 0x0002;
    pub const SO_BROADCAST: SockOpt           = 0x0020;
    pub const SO_DEBUG: SockOpt               = 0x0001;
    pub const SO_DONTTRUNC: SockOpt           = 0x2000;
    pub const SO_ERROR: SockOpt               = 0x1007;
    pub const SO_DONTROUTE: SockOpt           = 0x0010;
    pub const SO_KEEPALIVE: SockOpt           = 0x0008;
    pub const SO_LABEL: SockOpt               = 0x1010;
    pub const SO_LINGER: SockOpt              = 0x0080;
    pub const SO_NREAD: SockOpt               = 0x1020;
    pub const SO_NKE: SockOpt                 = 0x1021;
    pub const SO_NOSIGPIPE: SockOpt           = 0x1022;
    pub const SO_NOADDRERR: SockOpt           = 0x1023;
    pub const SO_NOTIFYCONFLICT: SockOpt      = 0x1026;
    pub const SO_NP_EXTENSIONS: SockOpt       = 0x1083;
    pub const SO_NWRITE: SockOpt              = 0x1024;
    pub const SO_OOBINLINE: SockOpt           = 0x0100;
    pub const SO_PEERLABEL: SockOpt           = 0x1011;
    pub const SO_RCVBUF: SockOpt              = 0x1002;
    pub const SO_RCVLOWAT: SockOpt            = 0x1004;
    pub const SO_SNDLOWAT: SockOpt            = 0x1003;
    pub const SO_RCVTIMEO: SockOpt            = 0x1006;
    pub const SO_SNDTIMEO: SockOpt            = 0x1005;
    pub const SO_RANDOMPORT: SockOpt          = 0x1082;
    pub const SO_RESTRICTIONS: SockOpt        = 0x1081;
    pub const SO_RESTRICT_DENYIN: SockOpt     = 0x00000001;
    pub const SO_RESTRICT_DENYOUT: SockOpt    = 0x00000002;
    pub const SO_REUSEADDR: SockOpt           = 0x0004;
    pub const SO_REUSESHAREUID: SockOpt       = 0x1025;
    pub const SO_SNDBUF: SockOpt              = 0x1001;
    pub const SO_TIMESTAMP: SockOpt           = 0x0400;
    pub const SO_TIMESTAMP_MONOTONIC: SockOpt = 0x0800;
    pub const SO_TYPE: SockOpt                = 0x1008;
    pub const SO_WANTMORE: SockOpt            = 0x4000;
    pub const SO_WANTOOBFLAG: SockOpt         = 0x8000;
    #[allow(type_overflow)]
    pub const SO_RESTRICT_DENYSET: SockOpt    = 0x80000000;
}

pub fn socket(domain: AddressFamily, mut ty: SockType, flags: SockFlag) -> SysResult<Fd> {
    let feat_atomic = features::socket_atomic_cloexec();

    if feat_atomic {
        ty = ty | flags.bits();
    }

    // TODO: Check the kernel version
    let res = unsafe { ffi::socket(domain, ty, 0) };

    if res < 0 {
        return Err(SysError::last());
    }

    if !feat_atomic {
        if flags.contains(SOCK_CLOEXEC) {
            try!(fcntl(res, F_SETFD(FD_CLOEXEC)));
        }

        if flags.contains(SOCK_NONBLOCK) {
            try!(fcntl(res, F_SETFL(O_NONBLOCK)));
        }
    }

    Ok(res)
}

pub fn listen(sockfd: Fd, backlog: uint) -> SysResult<()> {
    let res = unsafe { ffi::listen(sockfd, backlog as c_int) };
    from_ffi(res)
}

pub fn bind(sockfd: Fd, addr: &SockAddr) -> SysResult<()> {
    let res = unsafe {
        match *addr {
            SockIpV4(ref addr) => ffi::bind(sockfd, mem::transmute(addr), mem::size_of::<sockaddr_in>() as socklen_t),
            SockIpV6(ref addr) => ffi::bind(sockfd, mem::transmute(addr), mem::size_of::<sockaddr_in6>() as socklen_t),
            SockUnix(ref addr) => ffi::bind(sockfd, mem::transmute(addr), mem::size_of::<sockaddr_un>() as socklen_t)
        }
    };

    from_ffi(res)
}

pub fn accept(sockfd: Fd) -> SysResult<Fd> {
    let res = unsafe { ffi::accept(sockfd, ptr::null_mut(), ptr::null_mut()) };

    if res < 0 {
        return Err(SysError::last());
    }

    Ok(res)
}

#[cfg(not(any(target_os = "macos", target_os = "ios")))]
pub fn accept4(sockfd: Fd, flags: SockFlag) -> SysResult<Fd> {
    use libc::sockaddr;

    type F = unsafe extern "C" fn(c_int, *mut sockaddr, *mut socklen_t, c_int) -> c_int;

    extern {
        #[linkage = "extern_weak"]
        static accept4: *const ();
    }

    if !accept4.is_null() {
        let res = unsafe {
            mem::transmute::<*const (), F>(accept4)(
                sockfd, ptr::null_mut(), ptr::null_mut(), flags.bits)
        };

        if res < 0 {
            return Err(SysError::last());
        }

        Ok(res)
    } else {
        accept4_polyfill(sockfd, flags)
    }
}

#[cfg(any(target_os = "macos", target_os = "ios"))]
pub fn accept4(sockfd: Fd, flags: SockFlag) -> SysResult<Fd> {
    accept4_polyfill(sockfd, flags)
}

#[inline]
fn accept4_polyfill(sockfd: Fd, flags: SockFlag) -> SysResult<Fd> {
    let res =  unsafe { ffi::accept(sockfd, ptr::null_mut(), ptr::null_mut()) };

    if res < 0 {
        return Err(SysError::last());
    }

    if flags.contains(SOCK_CLOEXEC) {
        try!(fcntl(res, F_SETFD(FD_CLOEXEC)));
    }

    if flags.contains(SOCK_NONBLOCK) {
        try!(fcntl(res, F_SETFL(O_NONBLOCK)));
    }

    Ok(res)
}

pub fn connect(sockfd: Fd, addr: &SockAddr) -> SysResult<()> {
    let res = unsafe {
        match *addr {
            SockIpV4(ref addr) => ffi::connect(sockfd, mem::transmute(addr), mem::size_of::<sockaddr_in>() as socklen_t),
            SockIpV6(ref addr) => ffi::connect(sockfd, mem::transmute(addr), mem::size_of::<sockaddr_in6>() as socklen_t),
            SockUnix(ref addr) => ffi::connect(sockfd, mem::transmute(addr), mem::size_of::<sockaddr_un>() as socklen_t)
        }
    };

    from_ffi(res)
}

#[repr(C)]
pub struct linger {
    pub l_onoff: c_int,
    pub l_linger: c_int
}

pub fn getsockopt<T>(fd: Fd, level: SockLevel, opt: SockOpt, val: &mut T) -> SysResult<uint> {
    let mut len = mem::size_of::<T>() as socklen_t;

    let res = unsafe {
        ffi::getsockopt(
            fd, level, opt,
            mem::transmute(val),
            &mut len as *mut socklen_t)
    };

    if res < 0 {
        return Err(SysError::last());
    }

    Ok(len as uint)
}

pub fn setsockopt<T>(fd: Fd, level: SockLevel, opt: SockOpt, val: &T) -> SysResult<()> {
    let len = mem::size_of::<T>() as socklen_t;

    let res = unsafe {
            ffi::setsockopt(
            fd, level, opt,
            mem::transmute(val),
            len)
    };

    from_ffi(res)
}
<commit_msg>Add IPPROTO_IP for level, based on existing SOL_IP<commit_after>use std::{mem, ptr};
use libc::{c_int, socklen_t};
use fcntl::{Fd, fcntl, F_SETFL, F_SETFD, FD_CLOEXEC, O_NONBLOCK};
use errno::{SysResult, SysError, from_ffi};
use features;

pub use libc::{in_addr, sockaddr_in, sockaddr_in6, sockaddr_un, sa_family_t, ip_mreq};

pub use self::consts::*;

mod ffi {
    use libc::{c_int, c_void, socklen_t};
    pub use libc::{socket, listen, bind, accept, connect, setsockopt};

    extern {
        pub fn getsockopt(
            sockfd: c_int,
            level: c_int,
            optname: c_int,
            optval: *mut c_void,
            optlen: *mut socklen_t) -> c_int;
    }
}

// Extra flags - Supported by Linux 2.6.27, normalized on other platforms
bitflags!(
    flags SockFlag: c_int {
        const SOCK_NONBLOCK = 0o0004000,
        const SOCK_CLOEXEC  = 0o2000000
    }
)

pub enum SockAddr {
    SockIpV4(sockaddr_in),
    SockIpV6(sockaddr_in6),
    SockUnix(sockaddr_un)
}

#[cfg(target_os = "linux")]
mod consts {
    use libc::{c_int};

    pub type AddressFamily = c_int;

    pub const AF_UNIX: AddressFamily  = 1;
    pub const AF_LOCAL: AddressFamily = AF_UNIX;
    pub const AF_INET: AddressFamily  = 2;
    pub const AF_INET6: AddressFamily = 10;

    pub type SockType = c_int;

    pub const SOCK_STREAM: SockType = 1;
    pub const SOCK_DGRAM: SockType = 2;
    pub const SOCK_SEQPACKET: SockType = 5;
    pub const SOCK_RAW: SockType = 3;
    pub const SOCK_RDM: SockType = 4;

    pub type SockLevel = c_int;

    pub const SOL_IP: SockLevel     = 0;
    pub const IPPROTO_IP: SockLevel = SOL_IP;
    pub const SOL_SOCKET: SockLevel = 1;
    pub const SOL_TCP: SockLevel    = 6;
    pub const IPPROTO_TCP: SockLevel = SOL_TCP;
    pub const SOL_UDP: SockLevel    = 17;
    pub const SOL_IPV6: SockLevel   = 41;

    pub type SockOpt = c_int;

    pub const SO_ACCEPTCONN: SockOpt = 30;
    pub const SO_BINDTODEVICE: SockOpt = 25;
    pub const SO_BROADCAST: SockOpt = 6;
    pub const SO_BSDCOMPAT: SockOpt = 14;
    pub const SO_DEBUG: SockOpt = 1;
    pub const SO_DOMAIN: SockOpt = 39;
    pub const SO_ERROR: SockOpt = 4;
    pub const SO_DONTROUTE: SockOpt = 5;
    pub const SO_KEEPALIVE: SockOpt = 9;
    pub const SO_LINGER: SockOpt = 13;
    pub const SO_MARK: SockOpt = 36;
    pub const SO_OOBINLINE: SockOpt = 10;
    pub const SO_PASSCRED: SockOpt = 16;
    pub const SO_PEEK_OFF: SockOpt = 42;
    pub const SO_PEERCRED: SockOpt = 17;
    pub const SO_PRIORITY: SockOpt = 12;
    pub const SO_PROTOCOL: SockOpt = 38;
    pub const SO_RCVBUF: SockOpt = 8;
    pub const SO_RCVBUFFORCE: SockOpt = 33;
    pub const SO_RCVLOWAT: SockOpt = 18;
    pub const SO_SNDLOWAT: SockOpt = 19;
    pub const SO_RCVTIMEO: SockOpt = 20;
    pub const SO_SNDTIMEO: SockOpt = 21;
    pub const SO_REUSEADDR: SockOpt = 2;
    pub const SO_REUSEPORT: SockOpt = 15;
    pub const SO_RXQ_OVFL: SockOpt = 40;
    pub const SO_SNDBUF: SockOpt = 7;
    pub const SO_SNDBUFFORCE: SockOpt = 32;
    pub const SO_TIMESTAMP: SockOpt = 29;
    pub const SO_TYPE: SockOpt = 3;
    pub const SO_BUSY_POLL: SockOpt = 46;

    // Socket options for TCP sockets
    pub const TCP_NODELAY: SockOpt = 1;
    pub const TCP_MAXSEG: SockOpt = 2;
    pub const TCP_CORK: SockOpt = 3;

    // Socket options for the IP layer of the socket
    pub const IP_MULTICAST_IF: SockOpt = 32;
    pub const IP_MULTICAST_TTL: SockOpt = 33;
    pub const IP_MULTICAST_LOOP: SockOpt = 34;
    pub const IP_ADD_MEMBERSHIP: SockOpt = 35;
    pub const IP_DROP_MEMBERSHIP: SockOpt = 36;
}

#[cfg(any(target_os = "macos", target_os = "ios"))]
mod consts {
    use libc::{c_int};

    pub type AddressFamily = c_int;

    pub const AF_UNIX: AddressFamily  = 1;
    pub const AF_LOCAL: AddressFamily = AF_UNIX;
    pub const AF_INET: AddressFamily  = 2;
    pub const AF_INET6: AddressFamily = 30;

    pub type SockType = c_int;

    pub const SOCK_STREAM: SockType = 1;
    pub const SOCK_DGRAM: SockType = 2;
    pub const SOCK_SEQPACKET: SockType = 5;
    pub const SOCK_RAW: SockType = 3;
    pub const SOCK_RDM: SockType = 4;

    pub type SockLevel = c_int;

    pub const SOL_SOCKET: SockLevel = 0xffff;

    pub type SockOpt = c_int;

    pub const SO_ACCEPTCONN: SockOpt          = 0x0002;
    pub const SO_BROADCAST: SockOpt           = 0x0020;
    pub const SO_DEBUG: SockOpt               = 0x0001;
    pub const SO_DONTTRUNC: SockOpt           = 0x2000;
    pub const SO_ERROR: SockOpt               = 0x1007;
    pub const SO_DONTROUTE: SockOpt           = 0x0010;
    pub const SO_KEEPALIVE: SockOpt           = 0x0008;
    pub const SO_LABEL: SockOpt               = 0x1010;
    pub const SO_LINGER: SockOpt              = 0x0080;
    pub const SO_NREAD: SockOpt               = 0x1020;
    pub const SO_NKE: SockOpt                 = 0x1021;
    pub const SO_NOSIGPIPE: SockOpt           = 0x1022;
    pub const SO_NOADDRERR: SockOpt           = 0x1023;
    pub const SO_NOTIFYCONFLICT: SockOpt      = 0x1026;
    pub const SO_NP_EXTENSIONS: SockOpt       = 0x1083;
    pub const SO_NWRITE: SockOpt              = 0x1024;
    pub const SO_OOBINLINE: SockOpt           = 0x0100;
    pub const SO_PEERLABEL: SockOpt           = 0x1011;
    pub const SO_RCVBUF: SockOpt              = 0x1002;
    pub const SO_RCVLOWAT: SockOpt            = 0x1004;
    pub const SO_SNDLOWAT: SockOpt            = 0x1003;
    pub const SO_RCVTIMEO: SockOpt            = 0x1006;
    pub const SO_SNDTIMEO: SockOpt            = 0x1005;
    pub const SO_RANDOMPORT: SockOpt          = 0x1082;
    pub const SO_RESTRICTIONS: SockOpt        = 0x1081;
    pub const SO_RESTRICT_DENYIN: SockOpt     = 0x00000001;
    pub const SO_RESTRICT_DENYOUT: SockOpt    = 0x00000002;
    pub const SO_REUSEADDR: SockOpt           = 0x0004;
    pub const SO_REUSESHAREUID: SockOpt       = 0x1025;
    pub const SO_SNDBUF: SockOpt              = 0x1001;
    pub const SO_TIMESTAMP: SockOpt           = 0x0400;
    pub const SO_TIMESTAMP_MONOTONIC: SockOpt = 0x0800;
    pub const SO_TYPE: SockOpt                = 0x1008;
    pub const SO_WANTMORE: SockOpt            = 0x4000;
    pub const SO_WANTOOBFLAG: SockOpt         = 0x8000;
    #[allow(type_overflow)]
    pub const SO_RESTRICT_DENYSET: SockOpt    = 0x80000000;
}

pub fn socket(domain: AddressFamily, mut ty: SockType, flags: SockFlag) -> SysResult<Fd> {
    let feat_atomic = features::socket_atomic_cloexec();

    if feat_atomic {
        ty = ty | flags.bits();
    }

    // TODO: Check the kernel version
    let res = unsafe { ffi::socket(domain, ty, 0) };

    if res < 0 {
        return Err(SysError::last());
    }

    if !feat_atomic {
        if flags.contains(SOCK_CLOEXEC) {
            try!(fcntl(res, F_SETFD(FD_CLOEXEC)));
        }

        if flags.contains(SOCK_NONBLOCK) {
            try!(fcntl(res, F_SETFL(O_NONBLOCK)));
        }
    }

    Ok(res)
}

pub fn listen(sockfd: Fd, backlog: uint) -> SysResult<()> {
    let res = unsafe { ffi::listen(sockfd, backlog as c_int) };
    from_ffi(res)
}

pub fn bind(sockfd: Fd, addr: &SockAddr) -> SysResult<()> {
    let res = unsafe {
        match *addr {
            SockIpV4(ref addr) => ffi::bind(sockfd, mem::transmute(addr), mem::size_of::<sockaddr_in>() as socklen_t),
            SockIpV6(ref addr) => ffi::bind(sockfd, mem::transmute(addr), mem::size_of::<sockaddr_in6>() as socklen_t),
            SockUnix(ref addr) => ffi::bind(sockfd, mem::transmute(addr), mem::size_of::<sockaddr_un>() as socklen_t)
        }
    };

    from_ffi(res)
}

pub fn accept(sockfd: Fd) -> SysResult<Fd> {
    let res = unsafe { ffi::accept(sockfd, ptr::null_mut(), ptr::null_mut()) };

    if res < 0 {
        return Err(SysError::last());
    }

    Ok(res)
}

#[cfg(not(any(target_os = "macos", target_os = "ios")))]
pub fn accept4(sockfd: Fd, flags: SockFlag) -> SysResult<Fd> {
    use libc::sockaddr;

    type F = unsafe extern "C" fn(c_int, *mut sockaddr, *mut socklen_t, c_int) -> c_int;

    extern {
        #[linkage = "extern_weak"]
        static accept4: *const ();
    }

    if !accept4.is_null() {
        let res = unsafe {
            mem::transmute::<*const (), F>(accept4)(
                sockfd, ptr::null_mut(), ptr::null_mut(), flags.bits)
        };

        if res < 0 {
            return Err(SysError::last());
        }

        Ok(res)
    } else {
        accept4_polyfill(sockfd, flags)
    }
}

#[cfg(any(target_os = "macos", target_os = "ios"))]
pub fn accept4(sockfd: Fd, flags: SockFlag) -> SysResult<Fd> {
    accept4_polyfill(sockfd, flags)
}

#[inline]
fn accept4_polyfill(sockfd: Fd, flags: SockFlag) -> SysResult<Fd> {
    let res =  unsafe { ffi::accept(sockfd, ptr::null_mut(), ptr::null_mut()) };

    if res < 0 {
        return Err(SysError::last());
    }

    if flags.contains(SOCK_CLOEXEC) {
        try!(fcntl(res, F_SETFD(FD_CLOEXEC)));
    }

    if flags.contains(SOCK_NONBLOCK) {
        try!(fcntl(res, F_SETFL(O_NONBLOCK)));
    }

    Ok(res)
}

pub fn connect(sockfd: Fd, addr: &SockAddr) -> SysResult<()> {
    let res = unsafe {
        match *addr {
            SockIpV4(ref addr) => ffi::connect(sockfd, mem::transmute(addr), mem::size_of::<sockaddr_in>() as socklen_t),
            SockIpV6(ref addr) => ffi::connect(sockfd, mem::transmute(addr), mem::size_of::<sockaddr_in6>() as socklen_t),
            SockUnix(ref addr) => ffi::connect(sockfd, mem::transmute(addr), mem::size_of::<sockaddr_un>() as socklen_t)
        }
    };

    from_ffi(res)
}

#[repr(C)]
pub struct linger {
    pub l_onoff: c_int,
    pub l_linger: c_int
}

pub fn getsockopt<T>(fd: Fd, level: SockLevel, opt: SockOpt, val: &mut T) -> SysResult<uint> {
    let mut len = mem::size_of::<T>() as socklen_t;

    let res = unsafe {
        ffi::getsockopt(
            fd, level, opt,
            mem::transmute(val),
            &mut len as *mut socklen_t)
    };

    if res < 0 {
        return Err(SysError::last());
    }

    Ok(len as uint)
}

pub fn setsockopt<T>(fd: Fd, level: SockLevel, opt: SockOpt, val: &T) -> SysResult<()> {
    let len = mem::size_of::<T>() as socklen_t;

    let res = unsafe {
            ffi::setsockopt(
            fd, level, opt,
            mem::transmute(val),
            len)
    };

    from_ffi(res)
}
<|endoftext|>