<commit_before>use std::rc::Rc;
use std::cell::{Cell, RefCell};
use std::sync::mpsc;
use std::any::Any;
use helpers::{rc_and_weak, with_weak};
use types::{Callbacks, SumType2};
use signal::Signal;
use maybe_owned::MaybeOwned;

#[cfg(feature="either")]
use either::Either;

/// A source of events that feeds the streams connected to it.
#[derive(Debug)]
pub struct Sink<T>
{
    cbs: Rc<Callbacks<T>>,
}

impl<T> Sink<T>
{
    /// Creates a new sink.
    pub fn new() -> Self
    {
        Sink{ cbs: Default::default() }
    }

    /// Creates a stream that receives the events sent to this sink.
    pub fn stream(&self) -> Stream<T>
    {
        Stream{ cbs: self.cbs.clone(), source: None }
    }

    /// Sends a value into the sink.
    #[inline]
    pub fn send(&self, val: T)
    {
        self.cbs.call(val)
    }

    /// Sends values from an Iterator into the sink.
    #[inline]
    pub fn feed(&self, iter: I)
        where I: IntoIterator<Item=T>
    {
        for val in iter
        {
            self.cbs.call(val)
        }
    }
}

impl<T> Default for Sink<T>
{
    fn default() -> Self
    {
        Sink::new()
    }
}

impl<T> Clone for Sink<T>
{
    fn clone(&self) -> Self
    {
        Sink{ cbs: self.cbs.clone() }
    }
}

/// A stream of discrete events sent over time.
///
/// All the streams returned by the methods below contain an internal reference to it's parent,
/// so dropping intermediate streams won't break the chain.
#[derive(Debug)]
pub struct Stream<T>
{
    cbs: Rc<Callbacks<T>>,
    source: Option<Rc<Any>>,  // strong reference to a parent Stream
}

impl<T: 'static> Stream<T>
{
    /// Creates a stream that never fires.
    pub fn never() -> Self
    {
        Stream{ cbs: Default::default(), source: None }
    }

    /// Maps this stream into another stream using the provided function.
    #[inline]
    pub fn map<F, R>(&self, f: F) -> Stream<R>
        where F: Fn(MaybeOwned<T>) -> R + 'static,
        R: 'static
    {
        self.filter_map(move |arg| Some(f(arg)))
    }

    /// Creates a new stream that only contains the values where the predicate is `true`.
    pub fn filter<F>(&self, pred: F) -> Self
        where F: Fn(&T) -> bool + 'static
    {
        let (new_cbs, weak) = rc_and_weak(Callbacks::new());
        self.cbs.push(move |arg| {
            with_weak(&weak, |cb| if pred(&arg) { cb.call_dyn(arg) })
        });
        Stream{ cbs: new_cbs, source: Some(Rc::new(self.clone())) }
    }

    /// Filter and map a stream simultaneously.
    pub fn filter_map<F, R>(&self, f: F) -> Stream<R>
        where F: Fn(MaybeOwned<T>) -> Option<R> + 'static,
        R: 'static
    {
        let (new_cbs, weak) = rc_and_weak(Callbacks::new());
        self.cbs.push(move |arg| {
            with_weak(&weak, |cb| if let Some(val) = f(arg) { cb.call(val) })
        });
        Stream{ cbs: new_cbs, source: Some(Rc::new(self.clone())) }
    }

    /// Creates a new stream that fires with the events from both streams.
    pub fn merge(&self, other: &Stream<T>) -> Self
    {
        let (new_cbs, weak1) = rc_and_weak(Callbacks::new());
        let weak2 = weak1.clone();
        self.cbs.push(move |arg| {
            with_weak(&weak1, |cb| cb.call_dyn(arg))
        });
        other.cbs.push(move |arg| {
            with_weak(&weak2, |cb| cb.call_dyn(arg))
        });
        Stream{ cbs: new_cbs, source: Some(Rc::new((self.clone(), other.clone()))) }
    }

    /// Merges two streams of different types using the provided function.
    #[cfg(feature="either")]
    pub fn merge_with(&self, other: &Stream, f: F) -> Stream<R>
        where F: Fn(Either<MaybeOwned<T>, MaybeOwned>) -> R + 'static,
        U: 'static, R: 'static
    {
        let (new_cbs, weak1) = rc_and_weak(Callbacks::new());
        let weak2 = weak1.clone();
        let f1 = Rc::new(f);
        let f2 = f1.clone();
        self.cbs.push(move |arg| {
            with_weak(&weak1, |cb| cb.call(f1(Either::Left(arg))))
        });
        other.cbs.push(move |arg| {
            with_weak(&weak2, |cb| cb.call(f2(Either::Right(arg))))
        });
        Stream{ cbs: new_cbs, source: Some(Rc::new((self.clone(), other.clone()))) }
    }

    /// Accumulates the values sent over this stream.
    ///
    /// The fold operation is done by taking the accumulator, consuming it's value, and then
    /// putting back the transformed value. This avoids cloning, but if the closure panics it will
    /// leave the storage empty, and then any sampling attempt on this object will panic until
    /// someone puts back a value on it.
    /// If this is undesirable, use `Stream::fold_clone` instead.
    pub fn fold(&self, initial: A, f: F) -> Signal
        where F: Fn(A, MaybeOwned<T>) -> A + 'static,
        A: 'static
    {
        let (storage, weak) = rc_and_weak(RefCell::new(Some(initial)));
        self.cbs.push(move |arg| {
            weak.upgrade()
                .map(|st| {
                    let old = st.borrow_mut().take().expect("storage empty");
                    let new = f(old, arg);
                    *st.borrow_mut() = Some(new);
                })
                .is_some()
        });

        Signal::from_storage(storage, self.clone())
    }

    /// Folds the stream by cloning the accumulator.
    ///
    /// This will clone the accumulator on every value processed, but if the closure panics, the
    /// storage will remain unchanged and later attempts at sampling will succeed like nothing
    /// happened.
    pub fn fold_clone(&self, initial: A, f: F) -> Signal
        where F: Fn(A, MaybeOwned<T>) -> A + 'static,
        A: Clone + 'static
    {
        let (storage, weak) = rc_and_weak(RefCell::new(Some(initial)));
        self.cbs.push(move |arg| {
            weak.upgrade()
                .map(|st| {
                    let old = st.borrow().clone().expect("storage empty");
                    let new = f(old, arg);
                    *st.borrow_mut() = Some(new);
                })
                .is_some()
        });

        Signal::from_storage(storage, self.clone())
    }

    /// Maps each stream event to `0..N` output values.
    ///
    /// The closure must return it's value by sending it through the provided sink.
    /// Multiple values (or none) can be sent to the output stream this way.
    ///
    /// This primitive is useful to construct asynchronous operations, since you can
    /// store the sink for later usage.
    pub fn map_n<F, R>(&self, f: F) -> Stream<R>
        where F: Fn(MaybeOwned<T>, Sink<R>) + 'static,
        R: 'static
    {
        let (new_cbs, weak) = rc_and_weak(Callbacks::new());
        self.cbs.push(move |arg| {
            with_weak(&weak, |cb| f(arg, Sink{ cbs: cb }))
        });
        Stream{ cbs: new_cbs, source: Some(Rc::new(self.clone())) }
    }

    /// Reads the values without modifying them.
    ///
    /// This is meant to be used as a debugging tool and not to cause side effects.
    pub fn inspect<F>(self, f: F) -> Self
        where F: Fn(MaybeOwned<T>) + 'static
    {
        self.cbs.push(move |arg| { f(arg); true });
        self
    }
}

impl<T: Clone + 'static> Stream<T>
{
    /// Creates a Signal that holds the last value sent to this stream.
    #[inline]
    pub fn hold(&self, initial: T) -> Signal<T>
    {
        self.hold_if(initial, |_| true)
    }

    /// Holds the last value in this stream where the predicate is `true`.
    pub fn hold_if<F>(&self, initial: T, pred: F) -> Signal<T>
        where F: Fn(&T) -> bool + 'static
    {
        let (storage, weak) = rc_and_weak(RefCell::new(Some(initial)));
        self.cbs.push(move |arg| {
            weak.upgrade()
                .map(|st| if pred(&arg) { *st.borrow_mut() = Some(arg.into_owned()) })
                .is_some()
        });

        Signal::from_storage(storage, self.clone())
    }

    /// Creates a channel and sends the stream events through it.
    pub fn as_channel(&self) -> mpsc::Receiver<T>
    {
        let (tx, rx) = mpsc::channel();
        self.cbs.push(move |arg| {
            tx.send(arg.into_owned()).is_ok()
        });
        rx
    }
}

impl<T: Clone + 'static> Stream<Option<T>>
{
    /// Filters a stream of `Option`, returning the unwrapped `Some` values
    #[inline]
    pub fn filter_some(&self) -> Stream<T>
    {
        self.filter_first()
    }
}

impl<T: SumType2 + Clone + 'static> Stream<T>
    where T::Type1: 'static, T::Type2: 'static
{
    /// Creates a stream with only the first element of a sum type
    pub fn filter_first(&self) -> Stream<T::Type1>
    {
        self.filter_map(|res| if res.is_type1() { res.into_owned().into_type1() } else { None })
    }

    /// Creates a stream with only the second element of a sum type
    pub fn filter_second(&self) -> Stream<T::Type2>
    {
        self.filter_map(|res| if res.is_type2() { res.into_owned().into_type2() } else { None })
    }

    /// Splits a two element sum type stream into two streams with the unwrapped values
    pub fn split(&self) -> (Stream<T::Type1>, Stream<T::Type2>)
    {
        let (cbs_1, weak_1) = rc_and_weak(Callbacks::new());
        let (cbs_2, weak_2) = rc_and_weak(Callbacks::new());
        self.cbs.push(move |result| {
            if result.is_type1()
            {
                if let Some(cb) = weak_1.upgrade()
                {
                    cb.call(result.into_owned().into_type1().unwrap());
                    true
                }
                else
                {
                    // drop callback if both output streams dropped
                    weak_2.upgrade().is_some()
                }
            }
            else // if result.is_type2()
            {
                if let Some(cb) = weak_2.upgrade()
                {
                    cb.call(result.into_owned().into_type2().unwrap());
                    true
                }
                else
                {
                    weak_1.upgrade().is_some()
                }
            }
        });
        let source_rc = Rc::new(self.clone());
        let stream_1 = Stream{ cbs: cbs_1, source: Some(source_rc.clone()) };
        let stream_2 = Stream{ cbs: cbs_2, source: Some(source_rc) };
        (stream_1, stream_2)
    }
}

impl<T: 'static> Stream<Stream<T>>
{
    /// Listens to the events from the last stream sent to a nested stream
    pub fn switch(&self) -> Stream<T>
    {
        let (new_cbs, weak) = rc_and_weak(Callbacks::new());
        let id = Rc::new(Cell::new(0u64));  // id of each stream sent
        self.cbs.push(move |stream| {
            if weak.upgrade().is_none() { return false }
            let cbs_w = weak.clone();
            let cur_id = id.clone();
            // increment the id so it will only send to the last stream
            let my_id = id.get() + 1;
            id.set(my_id);
            // redirect the inner stream to the output stream
            stream.cbs.push(move |arg| {
                if my_id != cur_id.get() { return false }
                with_weak(&cbs_w, |cb| cb.call_dyn(arg))
            });
            true
        });
        Stream{ cbs: new_cbs, source: Some(Rc::new(self.clone())) }
    }
}

impl<T> Clone for Stream<T>
{
    fn clone(&self) -> Self
    {
        Stream{ cbs: self.cbs.clone(), source: self.source.clone() }
    }
}

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

    #[test]
    fn stream_basic()
    {
        let sink = Sink::new();
        let stream = sink.stream();
        let rx = stream.as_channel();

        sink.send(42);
        sink.send(33);
        sink.feed(0..5);

        let result: Vec<_> = rx.try_iter().collect();
        assert_eq!(result, [42, 33, 0, 1, 2, 3, 4]);
    }

    #[test]
    fn stream_switch()
    {
        let stream_sink = Sink::new();
        let sink1 = Sink::new();
        let sink2 = Sink::new();

        let switched = stream_sink.stream().switch();
        let events = switched.as_channel();

        sink1.send(1);
        sink2.send(2);

        stream_sink.send(sink2.stream());
        sink1.send(3);
        sink2.send(4);
        assert_eq!(events.try_recv(), Ok(4));

        stream_sink.send(sink1.stream());
        sink1.send(5);
        sink2.send(6);
        assert_eq!(events.try_recv(), Ok(5));
    }
}
<commit_msg>added clarification in doc<commit_after>use std::rc::Rc;
use std::cell::{Cell, RefCell};
use std::sync::mpsc;
use std::any::Any;
use helpers::{rc_and_weak, with_weak};
use types::{Callbacks, SumType2};
use signal::Signal;
use maybe_owned::MaybeOwned;

#[cfg(feature="either")]
use either::Either;

/// A source of events that feeds the streams connected to it.
#[derive(Debug)]
pub struct Sink<T>
{
    cbs: Rc<Callbacks<T>>,
}

impl<T> Sink<T>
{
    /// Creates a new sink.
    pub fn new() -> Self
    {
        Sink{ cbs: Default::default() }
    }

    /// Creates a stream that receives the events sent to this sink.
    pub fn stream(&self) -> Stream<T>
    {
        Stream{ cbs: self.cbs.clone(), source: None }
    }

    /// Sends a value into the sink.
    #[inline]
    pub fn send(&self, val: T)
    {
        self.cbs.call(val)
    }

    /// Sends values from an Iterator into the sink.
    #[inline]
    pub fn feed(&self, iter: I)
        where I: IntoIterator<Item=T>
    {
        for val in iter
        {
            self.cbs.call(val)
        }
    }
}

impl<T> Default for Sink<T>
{
    fn default() -> Self
    {
        Sink::new()
    }
}

impl<T> Clone for Sink<T>
{
    fn clone(&self) -> Self
    {
        Sink{ cbs: self.cbs.clone() }
    }
}

/// A stream of discrete events sent over time.
///
/// All the streams returned by the methods below contain an internal reference to it's parent,
/// so dropping intermediate streams won't break the chain.
#[derive(Debug)]
pub struct Stream<T>
{
    cbs: Rc<Callbacks<T>>,
    source: Option<Rc<Any>>,  // strong reference to a parent Stream
}

impl<T: 'static> Stream<T>
{
    /// Creates a stream that never fires.
    pub fn never() -> Self
    {
        Stream{ cbs: Default::default(), source: None }
    }

    /// Maps this stream into another stream using the provided function.
    #[inline]
    pub fn map<F, R>(&self, f: F) -> Stream<R>
        where F: Fn(MaybeOwned<T>) -> R + 'static,
        R: 'static
    {
        self.filter_map(move |arg| Some(f(arg)))
    }

    /// Creates a new stream that only contains the values where the predicate is `true`.
    pub fn filter<F>(&self, pred: F) -> Self
        where F: Fn(&T) -> bool + 'static
    {
        let (new_cbs, weak) = rc_and_weak(Callbacks::new());
        self.cbs.push(move |arg| {
            with_weak(&weak, |cb| if pred(&arg) { cb.call_dyn(arg) })
        });
        Stream{ cbs: new_cbs, source: Some(Rc::new(self.clone())) }
    }

    /// Filter and map a stream simultaneously.
    pub fn filter_map<F, R>(&self, f: F) -> Stream<R>
        where F: Fn(MaybeOwned<T>) -> Option<R> + 'static,
        R: 'static
    {
        let (new_cbs, weak) = rc_and_weak(Callbacks::new());
        self.cbs.push(move |arg| {
            with_weak(&weak, |cb| if let Some(val) = f(arg) { cb.call(val) })
        });
        Stream{ cbs: new_cbs, source: Some(Rc::new(self.clone())) }
    }

    /// Creates a new stream that fires with the events from both streams.
    pub fn merge(&self, other: &Stream<T>) -> Self
    {
        let (new_cbs, weak1) = rc_and_weak(Callbacks::new());
        let weak2 = weak1.clone();
        self.cbs.push(move |arg| {
            with_weak(&weak1, |cb| cb.call_dyn(arg))
        });
        other.cbs.push(move |arg| {
            with_weak(&weak2, |cb| cb.call_dyn(arg))
        });
        Stream{ cbs: new_cbs, source: Some(Rc::new((self.clone(), other.clone()))) }
    }

    /// Merges two streams of different types using the provided function.
    #[cfg(feature="either")]
    pub fn merge_with(&self, other: &Stream, f: F) -> Stream<R>
        where F: Fn(Either<MaybeOwned<T>, MaybeOwned>) -> R + 'static,
        U: 'static, R: 'static
    {
        let (new_cbs, weak1) = rc_and_weak(Callbacks::new());
        let weak2 = weak1.clone();
        let f1 = Rc::new(f);
        let f2 = f1.clone();
        self.cbs.push(move |arg| {
            with_weak(&weak1, |cb| cb.call(f1(Either::Left(arg))))
        });
        other.cbs.push(move |arg| {
            with_weak(&weak2, |cb| cb.call(f2(Either::Right(arg))))
        });
        Stream{ cbs: new_cbs, source: Some(Rc::new((self.clone(), other.clone()))) }
    }

    /// Accumulates the values sent over this stream.
    ///
    /// The fold operation is done by taking the accumulator, consuming it's value, and then
    /// putting back the transformed value. This avoids cloning, but if the closure panics it will
    /// leave the storage empty, and then any sampling attempt on this object will panic until
    /// someone puts back a value on it.
    /// If this is undesirable, use `Stream::fold_clone` instead.
    pub fn fold(&self, initial: A, f: F) -> Signal
        where F: Fn(A, MaybeOwned<T>) -> A + 'static,
        A: 'static
    {
        let (storage, weak) = rc_and_weak(RefCell::new(Some(initial)));
        self.cbs.push(move |arg| {
            weak.upgrade()
                .map(|st| {
                    let old = st.borrow_mut().take().expect("storage empty");
                    let new = f(old, arg);
                    *st.borrow_mut() = Some(new);
                })
                .is_some()
        });

        Signal::from_storage(storage, self.clone())
    }

    /// Folds the stream by cloning the accumulator.
    ///
    /// This will clone the accumulator on every value processed, but if the closure panics, the
    /// storage will remain unchanged and later attempts at sampling will succeed like nothing
    /// happened.
    pub fn fold_clone(&self, initial: A, f: F) -> Signal
        where F: Fn(A, MaybeOwned<T>) -> A + 'static,
        A: Clone + 'static
    {
        let (storage, weak) = rc_and_weak(RefCell::new(Some(initial)));
        self.cbs.push(move |arg| {
            weak.upgrade()
                .map(|st| {
                    let old = st.borrow().clone().expect("storage empty");
                    let new = f(old, arg);
                    *st.borrow_mut() = Some(new);
                })
                .is_some()
        });

        Signal::from_storage(storage, self.clone())
    }

    /// Maps each stream event to `0..N` output values.
    ///
    /// The closure must return it's value by sending it through the provided sink.
    /// Multiple values (or none) can be sent to the output stream this way.
    ///
    /// This primitive is useful to construct asynchronous operations, since you can
    /// store the sink for later usage.
    pub fn map_n<F, R>(&self, f: F) -> Stream<R>
        where F: Fn(MaybeOwned<T>, Sink<R>) + 'static,
        R: 'static
    {
        let (new_cbs, weak) = rc_and_weak(Callbacks::new());
        self.cbs.push(move |arg| {
            with_weak(&weak, |cb| f(arg, Sink{ cbs: cb }))
        });
        Stream{ cbs: new_cbs, source: Some(Rc::new(self.clone())) }
    }

    /// Reads the values without modifying them.
    ///
    /// This is meant to be used as a debugging tool and not to cause side effects.
    pub fn inspect<F>(self, f: F) -> Self
        where F: Fn(MaybeOwned<T>) + 'static
    {
        self.cbs.push(move |arg| { f(arg); true });
        self
    }
}

impl<T: Clone + 'static> Stream<T>
{
    /// Creates a Signal that holds the last value sent to this stream.
    #[inline]
    pub fn hold(&self, initial: T) -> Signal<T>
    {
        self.hold_if(initial, |_| true)
    }

    /// Holds the last value in this stream where the predicate is `true`.
    pub fn hold_if<F>(&self, initial: T, pred: F) -> Signal<T>
        where F: Fn(&T) -> bool + 'static
    {
        let (storage, weak) = rc_and_weak(RefCell::new(Some(initial)));
        self.cbs.push(move |arg| {
            weak.upgrade()
                .map(|st| if pred(&arg) { *st.borrow_mut() = Some(arg.into_owned()) })
                .is_some()
        });

        Signal::from_storage(storage, self.clone())
    }

    /// Creates a channel and sends the stream events through it.
    ///
    /// This doesn't create a strong reference to the parent stream, so the sender will be dropped
    /// when the stream is deleted.
    pub fn as_channel(&self) -> mpsc::Receiver<T>
    {
        let (tx, rx) = mpsc::channel();
        self.cbs.push(move |arg| {
            tx.send(arg.into_owned()).is_ok()
        });
        rx
    }
}

impl<T: Clone + 'static> Stream<Option<T>>
{
    /// Filters a stream of `Option`, returning the unwrapped `Some` values
    #[inline]
    pub fn filter_some(&self) -> Stream<T>
    {
        self.filter_first()
    }
}

impl<T: SumType2 + Clone + 'static> Stream<T>
    where T::Type1: 'static, T::Type2: 'static
{
    /// Creates a stream with only the first element of a sum type
    pub fn filter_first(&self) -> Stream<T::Type1>
    {
        self.filter_map(|res| if res.is_type1() { res.into_owned().into_type1() } else { None })
    }

    /// Creates a stream with only the second element of a sum type
    pub fn filter_second(&self) -> Stream<T::Type2>
    {
        self.filter_map(|res| if res.is_type2() { res.into_owned().into_type2() } else { None })
    }

    /// Splits a two element sum type stream into two streams with the unwrapped values
    pub fn split(&self) -> (Stream<T::Type1>, Stream<T::Type2>)
    {
        let (cbs_1, weak_1) = rc_and_weak(Callbacks::new());
        let (cbs_2, weak_2) = rc_and_weak(Callbacks::new());
        self.cbs.push(move |result| {
            if result.is_type1()
            {
                if let Some(cb) = weak_1.upgrade()
                {
                    cb.call(result.into_owned().into_type1().unwrap());
                    true
                }
                else
                {
                    // drop callback if both output streams dropped
                    weak_2.upgrade().is_some()
                }
            }
            else // if result.is_type2()
            {
                if let Some(cb) = weak_2.upgrade()
                {
                    cb.call(result.into_owned().into_type2().unwrap());
                    true
                }
                else
                {
                    weak_1.upgrade().is_some()
                }
            }
        });
        let source_rc = Rc::new(self.clone());
        let stream_1 = Stream{ cbs: cbs_1, source: Some(source_rc.clone()) };
        let stream_2 = Stream{ cbs: cbs_2, source: Some(source_rc) };
        (stream_1, stream_2)
    }
}

impl<T: 'static> Stream<Stream<T>>
{
    /// Listens to the events from the last stream sent to a nested stream
    pub fn switch(&self) -> Stream<T>
    {
        let (new_cbs, weak) = rc_and_weak(Callbacks::new());
        let id = Rc::new(Cell::new(0u64));  // id of each stream sent
        self.cbs.push(move |stream| {
            if weak.upgrade().is_none() { return false }
            let cbs_w = weak.clone();
            let cur_id = id.clone();
            // increment the id so it will only send to the last stream
            let my_id = id.get() + 1;
            id.set(my_id);
            // redirect the inner stream to the output stream
            stream.cbs.push(move |arg| {
                if my_id != cur_id.get() { return false }
                with_weak(&cbs_w, |cb| cb.call_dyn(arg))
            });
            true
        });
        Stream{ cbs: new_cbs, source: Some(Rc::new(self.clone())) }
    }
}

impl<T> Clone for Stream<T>
{
    fn clone(&self) -> Self
    {
        Stream{ cbs: self.cbs.clone(), source: self.source.clone() }
    }
}

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

    #[test]
    fn stream_basic()
    {
        let sink = Sink::new();
        let stream = sink.stream();
        let rx = stream.as_channel();

        sink.send(42);
        sink.send(33);
        sink.feed(0..5);

        let result: Vec<_> = rx.try_iter().collect();
        assert_eq!(result, [42, 33, 0, 1, 2, 3, 4]);
    }

    #[test]
    fn stream_switch()
    {
        let stream_sink = Sink::new();
        let sink1 = Sink::new();
        let sink2 = Sink::new();

        let switched = stream_sink.stream().switch();
        let events = switched.as_channel();

        sink1.send(1);
        sink2.send(2);

        stream_sink.send(sink2.stream());
        sink1.send(3);
        sink2.send(4);
        assert_eq!(events.try_recv(), Ok(4));

        stream_sink.send(sink1.stream());
        sink1.send(5);
        sink2.send(6);
        assert_eq!(events.try_recv(), Ok(5));
    }
}
<|endoftext|>


<commit_before>/// Streamer describes a "streaming iterator."
///
/// It provides a mechanism for writing code that is generic over streams
/// produced by this crate.
///
/// Note that this is strictly less useful than `Iterator` because the item
/// associated type is bound to a specific lifetime. However, this does permit
/// us to write *some* generic code over streams that produce values tied
/// to the lifetime of the stream.
///
/// Some form of stream abstraction is inherently required for this crate
/// because elements in a finite state transducer are produced *by iterating*
/// over the structure. The alternative would be to create a new allocation
/// for each element iterated over, which would be prohibitively expensive.
///
/// # Usage & motivation
///
/// Streams are hard to use because they don't fit into Rust's current type
/// system very well. They are so hard to use that this author loathes having a
/// publically defined trait for it. Nevertheless, they do just barely provide
/// a means for composing multiple stream abstractions with different concrete
/// types. For example, one might want to take the union of a range query
/// stream with a stream that has been filtered by a regex. These streams have
/// different concrete types. A `Streamer` trait allows us to write code that
/// is generic over these concrete types. (All of the set operations are
/// implemented this way.)
///
/// A problem with streams is that the trait is itself parameterized by a
/// lifetime. In practice, this makes them very unergonomic because specifying
/// a `Streamer` bound generally requires a higher-ranked trait bound. This is
/// necessary because the lifetime can't actually be named in the enclosing
/// function; instead, the lifetime is local to iteration itself. Therefore,
/// one must assert that the bound is valid for *any particular* lifetime.
/// This is the essence of higher-rank trait bounds.
///
/// Because of this, you might expect to see lots of bounds that look like
/// this:
///
/// ```ignore
/// fn takes_stream<T, S>(s: S)
///     where S: for<'a> Streamer<'a, Item=T>
/// {
/// }
/// ```
///
/// There are *three* different problems with this declaration:
///
/// 1. `S` is not bound by any particular lifetime itself, and most streams
///    probably contain a reference to an underlying finite state transducer.
/// 2. It is often convenient to separate the notion of "stream" with
///    "stream constructor." This represents a similar split found in the
///    standard library for `Iterator` and `IntoIterator`, respectively.
/// 3. The `Item=T` is invalid because `Streamer`'s associated type is
///    parameterized by a lifetime and there is no way to parameterize an
///    arbitrary type constructor. (In this context, `T` is the type
///    constructor, because it will invariably require a lifetime to become
///    a concrete type.)
///
/// With that said, we must revise our possibly-workable bounds to a giant
/// scary monster:
///
/// ```ignore
/// fn takes_stream<'f, I, S>(s: S)
///     where I: for<'a> IntoStreamer<'a, Into=S, Item=(&'a [u8], Output)>,
///           S: 'f + for<'a> Streamer<'a, Item=(&'a [u8], Output)>
/// {
/// }
/// ```
///
/// We addressed the above points correspondingly:
///
/// 1. `S` is now bound by `'f`, which corresponds to the lifetime (possibly
///     `'static`) of the underlying stream.
/// 2. The `I` type parameter has been added to refer to a type that knows how
///    to build a stream. Notice that neither of the bounds for `I` or `S`
///    share a lifetime parameter. This is because the higher rank trait bound
///    specifies it works for *any* particular lifetime.
/// 3. `T` has been replaced with specific concrete types. Note that these
///    concrete types are duplicated. With iterators, we could use
///    `Item=S::Item` in the bound for `I`, but one cannot access an associated
///    type through a higher-ranked trait bound. Therefore, we must duplicate
///    the item type.
///
/// As you can see, streams offer little flexibility, little ergonomics and a
/// lot of hard to read trait bounds. The situation is lamentable, but
/// nevertheless, without them, we would not be able to compose streams by
/// leveraging the type system.
///
/// A redeemable quality is that these *same exact* trait bounds (modulo some
/// tweaks in the `Item` associated type) appear in many places in this crate
/// without much variation. Therefore, once you grok it, it's mostly easy to
/// pattern match it with "oh I need a stream." My hope is that clear
/// documentation and examples make these complex bounds easier to burden.
///
/// Stretching this abstraction further with Rust's current type system is not
/// advised.
pub trait Streamer<'a> {
    /// The type of the item emitted by this stream.
    type Item: 'a;

    /// Emits the next element in this stream, or `None` to indicate the stream
    /// has been exhausted.
    ///
    /// It is not specified what a stream does after `None` is emitted. In most
    /// cases, `None` should be emitted on every subsequent call.
    fn next(&'a mut self) -> Option<Self::Item>;
}

/// IntoStreamer describes types that can be converted to streams.
///
/// This is analogous to the `IntoIterator` trait for `Iterator` in
/// `std::iter`.
pub trait IntoStreamer<'a> {
    /// The type of the item emitted by the stream.
    type Item: 'a;
    /// The type of the stream to be constructed.
    type Into: Streamer<'a, Item=Self::Item>;

    /// Construct a stream from `Self`.
    fn into_stream(self) -> Self::Into;
}

impl<'a, S: Streamer<'a>> IntoStreamer<'a> for S {
    type Item = S::Item;
    type Into = S;

    fn into_stream(self) -> S {
        self
    }
}

impl<'a, I: Iterator> Streamer<'a> for I where <I as Iterator>::Item: 'a {
    type Item = I::Item;

    fn next(&'a mut self) -> Option<Self::Item> {
        Iterator::next(self)
    }
}
<commit_msg>fix #1<commit_after>/// Streamer describes a "streaming iterator."
///
/// It provides a mechanism for writing code that is generic over streams
/// produced by this crate.
///
/// Note that this is strictly less useful than `Iterator` because the item
/// associated type is bound to a specific lifetime. However, this does permit
/// us to write *some* generic code over streams that produce values tied
/// to the lifetime of the stream.
///
/// Some form of stream abstraction is inherently required for this crate
/// because elements in a finite state transducer are produced *by iterating*
/// over the structure. The alternative would be to create a new allocation
/// for each element iterated over, which would be prohibitively expensive.
///
/// # Usage & motivation
///
/// Streams are hard to use because they don't fit into Rust's current type
/// system very well. They are so hard to use that this author loathes having a
/// publically defined trait for it. Nevertheless, they do just barely provide
/// a means for composing multiple stream abstractions with different concrete
/// types. For example, one might want to take the union of a range query
/// stream with a stream that has been filtered by a regex. These streams have
/// different concrete types. A `Streamer` trait allows us to write code that
/// is generic over these concrete types. (All of the set operations are
/// implemented this way.)
///
/// A problem with streams is that the trait is itself parameterized by a
/// lifetime. In practice, this makes them very unergonomic because specifying
/// a `Streamer` bound generally requires a higher-ranked trait bound. This is
/// necessary because the lifetime can't actually be named in the enclosing
/// function; instead, the lifetime is local to iteration itself. Therefore,
/// one must assert that the bound is valid for *any particular* lifetime.
/// This is the essence of higher-rank trait bounds.
///
/// Because of this, you might expect to see lots of bounds that look like
/// this:
///
/// ```ignore
/// fn takes_stream<T, S>(s: S)
///     where S: for<'a> Streamer<'a, Item=T>
/// {
/// }
/// ```
///
/// There are *three* different problems with this declaration:
///
/// 1. `S` is not bound by any particular lifetime itself, and most streams
///    probably contain a reference to an underlying finite state transducer.
/// 2. It is often convenient to separate the notion of "stream" with
///    "stream constructor." This represents a similar split found in the
///    standard library for `Iterator` and `IntoIterator`, respectively.
/// 3. The `Item=T` is invalid because `Streamer`'s associated type is
///    parameterized by a lifetime and there is no way to parameterize an
///    arbitrary type constructor. (In this context, `T` is the type
///    constructor, because it will invariably require a lifetime to become
///    a concrete type.)
///
/// With that said, we must revise our possibly-workable bounds to a giant
/// scary monster:
///
/// ```ignore
/// fn takes_stream<'f, I, S>(s: I)
///     where I: for<'a> IntoStreamer<'a, Into=S, Item=(&'a [u8], Output)>,
///           S: 'f + for<'a> Streamer<'a, Item=(&'a [u8], Output)>
/// {
/// }
/// ```
///
/// We addressed the above points correspondingly:
///
/// 1. `S` is now bound by `'f`, which corresponds to the lifetime (possibly
///     `'static`) of the underlying stream.
/// 2. The `I` type parameter has been added to refer to a type that knows how
///    to build a stream. Notice that neither of the bounds for `I` or `S`
///    share a lifetime parameter. This is because the higher rank trait bound
///    specifies it works for *any* particular lifetime.
/// 3. `T` has been replaced with specific concrete types. Note that these
///    concrete types are duplicated. With iterators, we could use
///    `Item=S::Item` in the bound for `I`, but one cannot access an associated
///    type through a higher-ranked trait bound. Therefore, we must duplicate
///    the item type.
///
/// As you can see, streams offer little flexibility, little ergonomics and a
/// lot of hard to read trait bounds. The situation is lamentable, but
/// nevertheless, without them, we would not be able to compose streams by
/// leveraging the type system.
///
/// A redeemable quality is that these *same exact* trait bounds (modulo some
/// tweaks in the `Item` associated type) appear in many places in this crate
/// without much variation. Therefore, once you grok it, it's mostly easy to
/// pattern match it with "oh I need a stream." My hope is that clear
/// documentation and examples make these complex bounds easier to burden.
///
/// Stretching this abstraction further with Rust's current type system is not
/// advised.
pub trait Streamer<'a> {
    /// The type of the item emitted by this stream.
    type Item: 'a;

    /// Emits the next element in this stream, or `None` to indicate the stream
    /// has been exhausted.
    ///
    /// It is not specified what a stream does after `None` is emitted. In most
    /// cases, `None` should be emitted on every subsequent call.
    fn next(&'a mut self) -> Option<Self::Item>;
}

/// IntoStreamer describes types that can be converted to streams.
///
/// This is analogous to the `IntoIterator` trait for `Iterator` in
/// `std::iter`.
pub trait IntoStreamer<'a> {
    /// The type of the item emitted by the stream.
    type Item: 'a;
    /// The type of the stream to be constructed.
    type Into: Streamer<'a, Item=Self::Item>;

    /// Construct a stream from `Self`.
    fn into_stream(self) -> Self::Into;
}

impl<'a, S: Streamer<'a>> IntoStreamer<'a> for S {
    type Item = S::Item;
    type Into = S;

    fn into_stream(self) -> S {
        self
    }
}

impl<'a, I: Iterator> Streamer<'a> for I where <I as Iterator>::Item: 'a {
    type Item = I::Item;

    fn next(&'a mut self) -> Option<Self::Item> {
        Iterator::next(self)
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::io;
use std::error::Error;
use std::marker::PhantomData;
use std::io::ErrorKind::{WouldBlock, BrokenPipe};

use time::SteadyTime;
use rotor::{Response, Scope, Machine};
use mio::{EventSet, PollOpt};
use netbuf::Buf;
use void::{Void, unreachable};

use substr::find_substr;
use {Expectation, Protocol, StreamSocket, Stream, StreamImpl, Request};
use {Transport, Deadline, Accepted, Exception};

#[derive(Debug)]
enum IoOp {
    Done,
    NoOp,
    Eos,
    Error(io::Error),
}

impl<S: StreamSocket> StreamImpl<S> {
    fn transport(&mut self) -> Transport<S> {
        Transport {
            sock: &mut self.socket,
            inbuf: &mut self.inbuf,
            outbuf: &mut self.outbuf,
        }
    }
    fn _action<C, M>(mut self, req: Request<M>, scope: &mut Scope<C>)
        -> Result<Stream<C, S, M>, ()>
        where M: Protocol<C, S>,
              S: StreamSocket,
    {
        use Expectation::*;
        let mut req = try!(req.ok_or(()));
        let mut can_write = match self.write() {
            IoOp::Done => true,
            IoOp::NoOp => false,
            IoOp::Eos => {
                req = try!(req.0.exception(
                    &mut self.transport(),
                    Exception::WriteError(WriteZero),
                    scope).ok_or(()));
                self.outbuf.remove_range(..);
                false
            }
            IoOp::Error(e) => {
                req = try!(req.0.exception(
                    &mut self.transport(),
                    Exception::WriteError(e),
                    scope).ok_or(()));
                // TODO(tailhook) should we flush buffer here too?
                false
            }
        };
        'outer: loop {
            if can_write {
                can_write = match self.write() {
                    IoOp::Done => true,
                    IoOp::NoOp => false,
                    IoOp::Eos => {
                        req = try!(req.0.exception(
                            &mut self.transport(),
                            Exception::WriteError(WriteZero),
                            scope).ok_or(()));
                        self.outbuf.remove_range(..);
                        false
                    }
                    IoOp::Error(e) => {
                        req = try!(req.0.exception(
                            &mut self.transport(),
                            Exception::WriteError(e),
                            scope).ok_or(()));
                        // TODO(tailhook) should we flush buffer here too?
                        false
                    }
                };
            }
            match req.1 {
                Bytes(num) => {
                    loop {
                        if self.inbuf.len() >= num {
                            req = try!(req.0.bytes_read(&mut self.transport(),
                                num, scope).ok_or(()));
                            continue 'outer;
                        }
                        match self.read() {
                            IoOp::Done => {}
                            IoOp::NoOp => {
                                return Ok(Stream::compose(self, req, scope));
                            }
                            IoOp::Eos => {
                                req = try!(req.0.exception(
                                    &mut self.transport(),
                                    Exception::EndOfStream,
                                    scope).ok_or(()));
                                continue 'outer;
                            }
                            IoOp::Error(e) => {
                                req = try!(req.0.exception(
                                    &mut self.transport(),
                                    Exception::ReadError(e),
                                    scope).ok_or(()));
                                continue 'outer;
                            }
                        }
                    }
                }
                Delimiter(min, delim, max) => {
                    loop {
                        if self.inbuf.len() > min {
                            let opt = find_substr(&self.inbuf[min..], delim);
                            if let Some(num) = opt {
                                req = try!(req.0.bytes_read(
                                    &mut self.transport(),
                                    num, scope).ok_or(()));
                                continue 'outer;
                            }
                        }
                        if self.inbuf.len() > max {
                            return Err(());
                        }
                        match self.read() {
                            IoOp::Done => {}
                            IoOp::NoOp => {
                                return Ok(Stream::compose(self, req, scope));
                            }
                            IoOp::Eos => {
                                req = try!(req.0.exception(
                                    &mut self.transport(),
                                    Exception::EndOfStream,
                                    scope).ok_or(()));
                                continue 'outer;
                            }
                            IoOp::Error(e) => {
                                req = try!(req.0.exception(
                                    &mut self.transport(),
                                    Exception::ReadError(e),
                                    scope).ok_or(()));
                                continue 'outer;
                            }
                        }
                    }
                }
                Flush(num) => {
                    if self.outbuf.len() <= num {
                        req = try!(req.0.bytes_flushed(&mut self.transport(),
                            scope).ok_or(()));
                    } else {
                        return Ok(Stream::compose(self, req, scope));
                    }
                }
                Sleep => {
                    return Ok(Stream::compose(self, req, scope));
                }
            }
        }
    }
    fn action<C, M>(self, req: Request<M>, scope: &mut Scope<C>)
        -> Response<Stream<C, S, M>, Void>
        where M: Protocol<C, S>,
              S: StreamSocket
    {
        let old_timeout = self.timeout;
        match self._action(req, scope) {
            Ok(x) => Response::ok(x),
            Err(()) => {
                scope.clear_timeout(old_timeout);
                Response::done()
            }
        }
    }
    // Returns Ok(true) to if we have read something, does not loop for reading
    // because this might use whole memory, and we may parse and consume the
    // input instead of buffering it whole.
    fn read(&mut self) -> IoOp {
        match self.inbuf.read_from(&mut self.socket) {
            Ok(0) => IoOp::Eos,
            Ok(_) => IoOp::Done,
            Err(ref e) if e.kind() == BrokenPipe => IoOp::Eos,
            Err(ref e) if e.kind() == WouldBlock => IoOp::NoOp,
            Err(e) => IoOp::Error(e),
        }
    }
    fn write(&mut self) -> IoOp {
        loop {
            if self.outbuf.len() == 0 {
                return IoOp::Done;
            }
            match self.outbuf.write_to(&mut self.socket) {
                Ok(0) => return IoOp::Eos,
                Ok(_) => continue,
                Err(ref e) if e.kind() == BrokenPipe => return IoOp::Eos,
                Err(ref e) if e.kind() == WouldBlock => return IoOp::NoOp,
                Err(e) => return IoOp::Error(e),
            }
        }
    }
}

impl<C, S, P> Accepted<C, S> for Stream<C, S, P>
    where S: StreamSocket, P: Protocol<C, S, Seed=()>
{
    fn accepted(sock: S, scope: &mut Scope<C>) -> Result<Self, Box<Error>> {
        Self::new(sock, (), scope)
    }
}

impl<C, S: StreamSocket, P: Protocol<C, S>> Stream<C, S, P> {
    fn decompose(self) -> (P, Expectation, StreamImpl<S>) {
        (self.fsm, self.expectation, StreamImpl {
            socket: self.socket,
            deadline: self.deadline,
            timeout: self.timeout,
            inbuf: self.inbuf,
            outbuf: self.outbuf,
        })
    }
    fn compose(implem: StreamImpl<S>,
        (fsm, exp, dline): (P, Expectation, Deadline),
        scope: &mut Scope<C>)
        -> Stream<C, S, P>
    {
        let mut timeout = implem.timeout;
        if dline != implem.deadline {
            scope.clear_timeout(timeout);
            // Assuming that we can always add timeout since we have just
            // cancelled one. It may be not true if timer is already expired
            // or timeout is too far in future. But I'm not sure that killing
            // state machine here is much better idea than panicking.
            timeout = scope.timeout_ms(
                (dline - SteadyTime::now()).num_milliseconds() as u64)
                // TODO(tailhook) can we process the error somehow?
                .expect("Can't replace timer");
        }
        Stream {
            fsm: fsm,
            socket: implem.socket,
            expectation: exp,
            deadline: dline,
            timeout: timeout,
            inbuf: implem.inbuf,
            outbuf: implem.outbuf,
            phantom: PhantomData,
        }
    }
    pub fn new(mut sock: S, seed: P::Seed, scope: &mut Scope<C>)
        -> Result<Self, Box<Error>>
    {
        // Always register everything in edge-triggered mode.
        // This allows to never reregister socket.
        //
        // The no-reregister strategy is not a goal (although, it's expected
        // to lower number of syscalls for many request-reply-like protocols)
        // but it allows to have single source of truth for
        // readable()/writable() mask (no duplication in kernel space)
        try!(scope.register(&sock,
            EventSet::readable() | EventSet::writable(), PollOpt::edge()));
        match P::create(seed, &mut sock, scope) {
            None => return Err(Box::new(ProtocolStop)),
            Some((m, exp, dline)) => {
                let diff = dline - SteadyTime::now();
                let timeout = scope.timeout_ms(
                    diff.num_milliseconds() as u64)
                    // TODO(tailhook) propagate error carefully
                    .expect("Can't insert timer");
                Ok(Stream {
                    socket: sock,
                    expectation: exp,
                    deadline: dline,
                    timeout: timeout,
                    fsm: m,
                    inbuf: Buf::new(),
                    outbuf: Buf::new(),
                    phantom: PhantomData,
                })
            }
        }
    }
}

impl<C, S: StreamSocket, P: Protocol<C, S>> Machine<C> for Stream<C, S, P> {
    type Seed = Void;
    fn create(void: Void, _scope: &mut Scope<C>)
        -> Result<Self, Box<Error>>
    {
        unreachable(void);
    }
    fn ready(self, _events: EventSet, scope: &mut Scope<C>)
        -> Response<Self, Self::Seed>
    {
        // TODO(tailhook) use `events` to optimize reading
        let (fsm, exp, imp) = self.decompose();
        let deadline = imp.deadline;
        imp.action(Some((fsm, exp, deadline)), scope)
    }
    fn spawned(self, _scope: &mut Scope<C>) -> Response<Self, Self::Seed> {
        unreachable!();
    }
    fn timeout(self, scope: &mut Scope<C>) -> Response<Self, Self::Seed> {
        if Deadline::now() >= self.deadline {
            let (fsm, _exp, mut imp) = self.decompose();
            let res = fsm.timeout(&mut imp.transport(), scope);
            imp.action(res, scope)
        } else {
            // Spurious timeouts are possible for the couple of reasons
            Response::ok(self)
        }
    }
    fn wakeup(self, scope: &mut Scope<C>) -> Response<Self, Self::Seed> {
        let (fsm, _exp, mut imp) = self.decompose();
        let res = fsm.wakeup(&mut imp.transport(), scope);
        imp.action(res, scope)
    }
}

/// Protocol returned None right at the start of the stream processing
#[derive(Debug)]
pub struct ProtocolStop;

impl fmt::Display for ProtocolStop {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "ProtocolStop")
    }
}

impl Error for ProtocolStop {
    fn cause(&self) -> Option<&Error> { None }
    fn description(&self) -> &'static str {
        r#"Protocol returned None (which means "stop") at start"#
    }
}

/// Can't insert timer, so can't create a connection
///
/// It's may be because there are too many timers in mio event loop or
/// because timer is too far away in the future (this is a limitation of
/// mio event loop too)
#[derive(Debug)]
pub struct TimerError;

impl fmt::Display for TimerError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "TimerError")
    }
}

impl Error for TimerError {
    fn cause(&self) -> Option<&Error> { None }
    fn description(&self) -> &'static str {
        "Can't insert timer, probably too much timers or time is too far away"
    }
}
<commit_msg>Fixup: proper creation of `io::Error`s<commit_after>use std::fmt;
use std::io;
use std::error::Error;
use std::marker::PhantomData;
use std::io::ErrorKind::{WouldBlock, BrokenPipe, WriteZero};

use time::SteadyTime;
use rotor::{Response, Scope, Machine};
use mio::{EventSet, PollOpt};
use netbuf::Buf;
use void::{Void, unreachable};

use substr::find_substr;
use {Expectation, Protocol, StreamSocket, Stream, StreamImpl, Request};
use {Transport, Deadline, Accepted, Exception};

#[derive(Debug)]
enum IoOp {
    Done,
    NoOp,
    Eos,
    Error(io::Error),
}

impl<S: StreamSocket> StreamImpl<S> {
    fn transport(&mut self) -> Transport<S> {
        Transport {
            sock: &mut self.socket,
            inbuf: &mut self.inbuf,
            outbuf: &mut self.outbuf,
        }
    }
    fn _action<C, M>(mut self, req: Request<M>, scope: &mut Scope<C>)
        -> Result<Stream<C, S, M>, ()>
        where M: Protocol<C, S>,
              S: StreamSocket,
    {
        use Expectation::*;
        let mut req = try!(req.ok_or(()));
        let mut can_write = match self.write() {
            IoOp::Done => true,
            IoOp::NoOp => false,
            IoOp::Eos => {
                req = try!(req.0.exception(
                    &mut self.transport(),
                    Exception::WriteError(io::Error::new(
                        WriteZero, "failed to write whole buffer")),
                    scope).ok_or(()));
                self.outbuf.remove_range(..);
                false
            }
            IoOp::Error(e) => {
                req = try!(req.0.exception(
                    &mut self.transport(),
                    Exception::WriteError(e),
                    scope).ok_or(()));
                // TODO(tailhook) should we flush buffer here too?
                false
            }
        };
        'outer: loop {
            if can_write {
                can_write = match self.write() {
                    IoOp::Done => true,
                    IoOp::NoOp => false,
                    IoOp::Eos => {
                        req = try!(req.0.exception(
                            &mut self.transport(),
                            Exception::WriteError(io::Error::new(
                                WriteZero, "failed to write whole buffer")),
                            scope).ok_or(()));
                        self.outbuf.remove_range(..);
                        false
                    }
                    IoOp::Error(e) => {
                        req = try!(req.0.exception(
                            &mut self.transport(),
                            Exception::WriteError(e),
                            scope).ok_or(()));
                        // TODO(tailhook) should we flush buffer here too?
                        false
                    }
                };
            }
            match req.1 {
                Bytes(num) => {
                    loop {
                        if self.inbuf.len() >= num {
                            req = try!(req.0.bytes_read(&mut self.transport(),
                                num, scope).ok_or(()));
                            continue 'outer;
                        }
                        match self.read() {
                            IoOp::Done => {}
                            IoOp::NoOp => {
                                return Ok(Stream::compose(self, req, scope));
                            }
                            IoOp::Eos => {
                                req = try!(req.0.exception(
                                    &mut self.transport(),
                                    Exception::EndOfStream,
                                    scope).ok_or(()));
                                continue 'outer;
                            }
                            IoOp::Error(e) => {
                                req = try!(req.0.exception(
                                    &mut self.transport(),
                                    Exception::ReadError(e),
                                    scope).ok_or(()));
                                continue 'outer;
                            }
                        }
                    }
                }
                Delimiter(min, delim, max) => {
                    loop {
                        if self.inbuf.len() > min {
                            let opt = find_substr(&self.inbuf[min..], delim);
                            if let Some(num) = opt {
                                req = try!(req.0.bytes_read(
                                    &mut self.transport(),
                                    num, scope).ok_or(()));
                                continue 'outer;
                            }
                        }
                        if self.inbuf.len() > max {
                            return Err(());
                        }
                        match self.read() {
                            IoOp::Done => {}
                            IoOp::NoOp => {
                                return Ok(Stream::compose(self, req, scope));
                            }
                            IoOp::Eos => {
                                req = try!(req.0.exception(
                                    &mut self.transport(),
                                    Exception::EndOfStream,
                                    scope).ok_or(()));
                                continue 'outer;
                            }
                            IoOp::Error(e) => {
                                req = try!(req.0.exception(
                                    &mut self.transport(),
                                    Exception::ReadError(e),
                                    scope).ok_or(()));
                                continue 'outer;
                            }
                        }
                    }
                }
                Flush(num) => {
                    if self.outbuf.len() <= num {
                        req = try!(req.0.bytes_flushed(&mut self.transport(),
                            scope).ok_or(()));
                    } else {
                        return Ok(Stream::compose(self, req, scope));
                    }
                }
                Sleep => {
                    return Ok(Stream::compose(self, req, scope));
                }
            }
        }
    }
    fn action<C, M>(self, req: Request<M>, scope: &mut Scope<C>)
        -> Response<Stream<C, S, M>, Void>
        where M: Protocol<C, S>,
              S: StreamSocket
    {
        let old_timeout = self.timeout;
        match self._action(req, scope) {
            Ok(x) => Response::ok(x),
            Err(()) => {
                scope.clear_timeout(old_timeout);
                Response::done()
            }
        }
    }
    // Returns Ok(true) to if we have read something, does not loop for reading
    // because this might use whole memory, and we may parse and consume the
    // input instead of buffering it whole.
    fn read(&mut self) -> IoOp {
        match self.inbuf.read_from(&mut self.socket) {
            Ok(0) => IoOp::Eos,
            Ok(_) => IoOp::Done,
            Err(ref e) if e.kind() == BrokenPipe => IoOp::Eos,
            Err(ref e) if e.kind() == WouldBlock => IoOp::NoOp,
            Err(e) => IoOp::Error(e),
        }
    }
    fn write(&mut self) -> IoOp {
        loop {
            if self.outbuf.len() == 0 {
                return IoOp::Done;
            }
            match self.outbuf.write_to(&mut self.socket) {
                Ok(0) => return IoOp::Eos,
                Ok(_) => continue,
                Err(ref e) if e.kind() == BrokenPipe => return IoOp::Eos,
                Err(ref e) if e.kind() == WouldBlock => return IoOp::NoOp,
                Err(e) => return IoOp::Error(e),
            }
        }
    }
}

impl<C, S, P> Accepted<C, S> for Stream<C, S, P>
    where S: StreamSocket, P: Protocol<C, S, Seed=()>
{
    fn accepted(sock: S, scope: &mut Scope<C>) -> Result<Self, Box<Error>> {
        Self::new(sock, (), scope)
    }
}

impl<C, S: StreamSocket, P: Protocol<C, S>> Stream<C, S, P> {
    fn decompose(self) -> (P, Expectation, StreamImpl<S>) {
        (self.fsm, self.expectation, StreamImpl {
            socket: self.socket,
            deadline: self.deadline,
            timeout: self.timeout,
            inbuf: self.inbuf,
            outbuf: self.outbuf,
        })
    }
    fn compose(implem: StreamImpl<S>,
        (fsm, exp, dline): (P, Expectation, Deadline),
        scope: &mut Scope<C>)
        -> Stream<C, S, P>
    {
        let mut timeout = implem.timeout;
        if dline != implem.deadline {
            scope.clear_timeout(timeout);
            // Assuming that we can always add timeout since we have just
            // cancelled one. It may be not true if timer is already expired
            // or timeout is too far in future. But I'm not sure that killing
            // state machine here is much better idea than panicking.
            timeout = scope.timeout_ms(
                (dline - SteadyTime::now()).num_milliseconds() as u64)
                // TODO(tailhook) can we process the error somehow?
                .expect("Can't replace timer");
        }
        Stream {
            fsm: fsm,
            socket: implem.socket,
            expectation: exp,
            deadline: dline,
            timeout: timeout,
            inbuf: implem.inbuf,
            outbuf: implem.outbuf,
            phantom: PhantomData,
        }
    }
    pub fn new(mut sock: S, seed: P::Seed, scope: &mut Scope<C>)
        -> Result<Self, Box<Error>>
    {
        // Always register everything in edge-triggered mode.
        // This allows to never reregister socket.
        //
        // The no-reregister strategy is not a goal (although, it's expected
        // to lower number of syscalls for many request-reply-like protocols)
        // but it allows to have single source of truth for
        // readable()/writable() mask (no duplication in kernel space)
        try!(scope.register(&sock,
            EventSet::readable() | EventSet::writable(), PollOpt::edge()));
        match P::create(seed, &mut sock, scope) {
            None => return Err(Box::new(ProtocolStop)),
            Some((m, exp, dline)) => {
                let diff = dline - SteadyTime::now();
                let timeout = scope.timeout_ms(
                    diff.num_milliseconds() as u64)
                    // TODO(tailhook) propagate error carefully
                    .expect("Can't insert timer");
                Ok(Stream {
                    socket: sock,
                    expectation: exp,
                    deadline: dline,
                    timeout: timeout,
                    fsm: m,
                    inbuf: Buf::new(),
                    outbuf: Buf::new(),
                    phantom: PhantomData,
                })
            }
        }
    }
}

impl<C, S: StreamSocket, P: Protocol<C, S>> Machine<C> for Stream<C, S, P> {
    type Seed = Void;
    fn create(void: Void, _scope: &mut Scope<C>)
        -> Result<Self, Box<Error>>
    {
        unreachable(void);
    }
    fn ready(self, _events: EventSet, scope: &mut Scope<C>)
        -> Response<Self, Self::Seed>
    {
        // TODO(tailhook) use `events` to optimize reading
        let (fsm, exp, imp) = self.decompose();
        let deadline = imp.deadline;
        imp.action(Some((fsm, exp, deadline)), scope)
    }
    fn spawned(self, _scope: &mut Scope<C>) -> Response<Self, Self::Seed> {
        unreachable!();
    }
    fn timeout(self, scope: &mut Scope<C>) -> Response<Self, Self::Seed> {
        if Deadline::now() >= self.deadline {
            let (fsm, _exp, mut imp) = self.decompose();
            let res = fsm.timeout(&mut imp.transport(), scope);
            imp.action(res, scope)
        } else {
            // Spurious timeouts are possible for the couple of reasons
            Response::ok(self)
        }
    }
    fn wakeup(self, scope: &mut Scope<C>) -> Response<Self, Self::Seed> {
        let (fsm, _exp, mut imp) = self.decompose();
        let res = fsm.wakeup(&mut imp.transport(), scope);
        imp.action(res, scope)
    }
}

/// Protocol returned None right at the start of the stream processing
#[derive(Debug)]
pub struct ProtocolStop;

impl fmt::Display for ProtocolStop {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "ProtocolStop")
    }
}

impl Error for ProtocolStop {
    fn cause(&self) -> Option<&Error> { None }
    fn description(&self) -> &'static str {
        r#"Protocol returned None (which means "stop") at start"#
    }
}

/// Can't insert timer, so can't create a connection
///
/// It's may be because there are too many timers in mio event loop or
/// because timer is too far away in the future (this is a limitation of
/// mio event loop too)
#[derive(Debug)]
pub struct TimerError;

impl fmt::Display for TimerError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "TimerError")
    }
}

impl Error for TimerError {
    fn cause(&self) -> Option<&Error> { None }
    fn description(&self) -> &'static str {
        "Can't insert timer, probably too much timers or time is too far away"
    }
}
<|endoftext|>


<commit_before>// TODO: Maybe stop handing out slices and instead use a wrapper type that supports copying etc.
// Invalid pages will cause page faults, which can be handled so that they are caught and EFAULT is
// returned. This will also make SMAP much, much, easier. c.f. Linux's copy_from_user, copy_to_user
// which are written in assembly and handle page faults.
use core::{mem, slice, str};

use crate::context;
use crate::paging::{Page, TableKind, VirtualAddress};
use crate::syscall::error::*;

use alloc::sync::Arc;

fn validate(address: usize, size: usize, writable: bool) -> Result<()> {
    if VirtualAddress::new(address.saturating_add(size)).kind() != TableKind::User {
        return Err(Error::new(EFAULT));
    }

    let end_offset = size.checked_sub(1).ok_or(Error::new(EFAULT))?;
    let end_address = address.checked_add(end_offset).ok_or(Error::new(EFAULT))?;

    let addr_space = Arc::clone(context::current()?.read().addr_space()?);
    let addr_space = addr_space.read();

    let start_page = Page::containing_address(VirtualAddress::new(address));
    let end_page = Page::containing_address(VirtualAddress::new(end_address));
    for page in Page::range_inclusive(start_page, end_page) {
        if let Some((_, flags)) = addr_space.table.utable.translate(page.start_address()) {
            if !flags.has_user() {
                // println!("{:X}: Not usermode", page.start_address().data());
                return Err(Error::new(EFAULT));
            }

            if writable && !flags.has_write() {
                // println!("{:X}: Not writable {}", page.start_address().data(), writable);
                return Err(Error::new(EFAULT));
            }
        } else {
            // println!("{:X}: Not found", page.start_address().data());
            return Err(Error::new(EFAULT));
        }
    }

    Ok(())
}

/// Convert a pointer and length to reference, if valid
pub unsafe fn validate_ref<T>(ptr: *const T, size: usize) -> Result<&'static T> {
    if size == mem::size_of::<T>() {
        validate(ptr as usize, mem::size_of::<T>(), false)?;
        Ok(&*ptr)
    } else {
        Err(Error::new(EINVAL))
    }
}

/// Convert a pointer and length to mutable reference, if valid
pub unsafe fn validate_ref_mut<T>(ptr: *mut T, size: usize) -> Result<&'static mut T> {
    if size == mem::size_of::<T>() {
        validate(ptr as usize, mem::size_of::<T>(), false)?;
        Ok(&mut *ptr)
    } else {
        Err(Error::new(EINVAL))
    }
}

/// Convert a pointer and length to slice, if valid
//TODO: Mark unsafe
pub fn validate_slice<T>(ptr: *const T, len: usize) -> Result<&'static [T]> {
    if len == 0 {
        Ok(&[])
    } else {
        validate(ptr as usize, len * mem::size_of::<T>(), false)?;
        Ok(unsafe { slice::from_raw_parts(ptr, len) })
    }
}
/// Convert a pointer with fixed static length to a reference to an array, if valid.
// TODO: This is probably also quite unsafe, mainly because we have no idea unless we do very
// careful checking, that this upholds the rules that LLVM relies with shared references, namely
// that the value cannot change by others. Atomic volatile.
pub unsafe fn validate_array<'a, T, const N: usize>(ptr: *const T) -> Result<&'a [T; N]> {
    validate(ptr as usize, mem::size_of::<T>() * N, false)?;
    Ok(&*ptr.cast::<[T; N]>())
}
pub unsafe fn validate_array_mut<'a, T, const N: usize>(ptr: *mut T) -> Result<&'a mut [T; N]> {
    validate(ptr as usize, mem::size_of::<T>() * N, true)?;
    Ok(&mut *ptr.cast::<[T; N]>())
}

/// Convert a pointer and length to slice, if valid
// TODO: Mark unsafe
//
// FIXME: This is probably never ever safe, except under very special circumstances. Any &mut
// reference will allow LLVM to assume that nobody else will ever modify this value, which is
// certainly not the case for multithreaded userspace programs. Instead, we will want something
// like atomic volatile.
pub fn validate_slice_mut<T>(ptr: *mut T, len: usize) -> Result<&'static mut [T]> {
    if len == 0 {
        Ok(&mut [])
    } else {
        validate(ptr as usize, len * mem::size_of::<T>(), true)?;
        Ok(unsafe { slice::from_raw_parts_mut(ptr, len) })
    }
}

/// Convert a pointer and length to str, if valid
//TODO: Mark unsafe
pub fn validate_str(ptr: *const u8, len: usize) -> Result<&'static str> {
    let slice = validate_slice(ptr, len)?;
    str::from_utf8(slice).map_err(|_| Error::new(EINVAL))
}
<commit_msg>Add fn for validating user memory range bounds.<commit_after>// TODO: Maybe stop handing out slices and instead use a wrapper type that supports copying etc.
// Invalid pages will cause page faults, which can be handled so that they are caught and EFAULT is
// returned. This will also make SMAP much, much, easier. c.f. Linux's copy_from_user, copy_to_user
// which are written in assembly and handle page faults.
use core::{mem, slice, str};

use crate::context;
use crate::memory::PAGE_SIZE;
use crate::paging::{Page, TableKind, VirtualAddress};
use crate::syscall::error::*;

use alloc::sync::Arc;

fn validate(address: usize, size: usize, writable: bool) -> Result<()> {
    if VirtualAddress::new(address.saturating_add(size)).kind() != TableKind::User {
        return Err(Error::new(EFAULT));
    }

    let end_offset = size.checked_sub(1).ok_or(Error::new(EFAULT))?;
    let end_address = address.checked_add(end_offset).ok_or(Error::new(EFAULT))?;

    let addr_space = Arc::clone(context::current()?.read().addr_space()?);
    let addr_space = addr_space.read();

    let start_page = Page::containing_address(VirtualAddress::new(address));
    let end_page = Page::containing_address(VirtualAddress::new(end_address));
    for page in Page::range_inclusive(start_page, end_page) {
        if let Some((_, flags)) = addr_space.table.utable.translate(page.start_address()) {
            if !flags.has_user() {
                // println!("{:X}: Not usermode", page.start_address().data());
                return Err(Error::new(EFAULT));
            }

            if writable && !flags.has_write() {
                // println!("{:X}: Not writable {}", page.start_address().data(), writable);
                return Err(Error::new(EFAULT));
            }
        } else {
            // println!("{:X}: Not found", page.start_address().data());
            return Err(Error::new(EFAULT));
        }
    }

    Ok(())
}

/// Convert a pointer and length to reference, if valid
pub unsafe fn validate_ref<T>(ptr: *const T, size: usize) -> Result<&'static T> {
    if size == mem::size_of::<T>() {
        validate(ptr as usize, mem::size_of::<T>(), false)?;
        Ok(&*ptr)
    } else {
        Err(Error::new(EINVAL))
    }
}

/// Convert a pointer and length to mutable reference, if valid
pub unsafe fn validate_ref_mut<T>(ptr: *mut T, size: usize) -> Result<&'static mut T> {
    if size == mem::size_of::<T>() {
        validate(ptr as usize, mem::size_of::<T>(), false)?;
        Ok(&mut *ptr)
    } else {
        Err(Error::new(EINVAL))
    }
}

/// Convert a pointer and length to slice, if valid
//TODO: Mark unsafe
pub fn validate_slice<T>(ptr: *const T, len: usize) -> Result<&'static [T]> {
    if len == 0 {
        Ok(&[])
    } else {
        validate(ptr as usize, len * mem::size_of::<T>(), false)?;
        Ok(unsafe { slice::from_raw_parts(ptr, len) })
    }
}
/// Convert a pointer with fixed static length to a reference to an array, if valid.
// TODO: This is probably also quite unsafe, mainly because we have no idea unless we do very
// careful checking, that this upholds the rules that LLVM relies with shared references, namely
// that the value cannot change by others. Atomic volatile.
pub unsafe fn validate_array<'a, T, const N: usize>(ptr: *const T) -> Result<&'a [T; N]> {
    validate(ptr as usize, mem::size_of::<T>() * N, false)?;
    Ok(&*ptr.cast::<[T; N]>())
}
pub unsafe fn validate_array_mut<'a, T, const N: usize>(ptr: *mut T) -> Result<&'a mut [T; N]> {
    validate(ptr as usize, mem::size_of::<T>() * N, true)?;
    Ok(&mut *ptr.cast::<[T; N]>())
}

/// Convert a pointer and length to slice, if valid
// TODO: Mark unsafe
//
// FIXME: This is probably never ever safe, except under very special circumstances. Any &mut
// reference will allow LLVM to assume that nobody else will ever modify this value, which is
// certainly not the case for multithreaded userspace programs. Instead, we will want something
// like atomic volatile.
pub fn validate_slice_mut<T>(ptr: *mut T, len: usize) -> Result<&'static mut [T]> {
    if len == 0 {
        Ok(&mut [])
    } else {
        validate(ptr as usize, len * mem::size_of::<T>(), true)?;
        Ok(unsafe { slice::from_raw_parts_mut(ptr, len) })
    }
}

/// Convert a pointer and length to str, if valid
//TODO: Mark unsafe
pub fn validate_str(ptr: *const u8, len: usize) -> Result<&'static str> {
    let slice = validate_slice(ptr, len)?;
    str::from_utf8(slice).map_err(|_| Error::new(EINVAL))
}

pub fn validate_region(address: usize, size: usize) -> Result<(Page, usize)> {
    if address % PAGE_SIZE != 0 || size % PAGE_SIZE != 0 || size == 0 {
        return Err(Error::new(EINVAL));
    }
    if address.saturating_add(size) > crate::USER_END_OFFSET {
        return Err(Error::new(EFAULT));
    }
    Ok((Page::containing_address(VirtualAddress::new(address)), size / PAGE_SIZE))
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::collections::btree_map::Iter;
use std::path::{PathBuf, Path};
use std::fmt::{Display, Formatter, Result};
use std::ops::Add;
use std::cmp::{Ord, Ordering};

/// Used to track the state of logical conditions
#[derive(Debug, Clone, Copy, Default, Hash, PartialEq, Eq, PartialOrd)]
pub struct LogicState {
    /// Whether the condition has been observed as true
    pub been_true: bool,
    /// Whether the condition has been observed as false
    pub been_false: bool,
}

impl<'a> Add for &'a LogicState {
    type Output = LogicState;

    fn add(self, other: &'a LogicState) -> LogicState {
        LogicState {
            been_true: self.been_true | other.been_true,
            been_false: self.been_false | other.been_false,
        }
    }
}

/// Shows what type of coverage data is being collected by a given trace
#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd)]
pub enum CoverageStat {
    /// Line coverage data (whether line has been hit)
    Line(u64),
    /// Branch coverage data (whether branch has been true and false
    Branch(LogicState),
    /// Condition coverage data (each boolean subcondition true and false)
    Condition(Vec<LogicState>),
}

impl Add for CoverageStat {
    type Output = CoverageStat;

    fn add(self, other: CoverageStat) -> CoverageStat {
        match (self, other) {
            (CoverageStat::Line(ref l), CoverageStat::Line(ref r)) => {
                CoverageStat::Line(l+r)
            },
            (CoverageStat::Branch(ref l), CoverageStat::Branch(ref r)) => {
                CoverageStat::Branch(l + r)
            },
            t @ _ => t.0,
        }
    }
}

impl Display for CoverageStat {
    fn fmt(&self, f: &mut Formatter) -> Result {
        match self {
            &CoverageStat::Line(x) => {
                write!(f, "hits: {}", x)
            },
            _ => write!(f, ""),
        }
    }
}

#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd)]
pub struct Trace {
    /// Line the trace is on in the file
    pub line: u64,
    /// Optional address showing location in the test artefact
    pub address: Option<u64>,
    /// Length of the instruction (useful to get entire condition/branch)
    pub length: usize,
    /// Coverage stats
    pub stats: CoverageStat,
}

/// Implemented to allow Traces to be sorted by line number
impl Ord for Trace {
    fn cmp(&self, other: &Trace) -> Ordering {
        self.line.cmp(&other.line)
    }
    fn max(self, other: Trace) -> Trace {
        if self.line > other.line {
            self
        } else {
            other
        }
    }
    fn min(self, other: Trace) -> Trace {
        if self.line < other.line {
            self
        } else {
            other
        }
    }
}

/// Stores all the program traces mapped to files and provides an interface to
/// add, query and change traces.
#[derive(Debug)]
pub struct TraceMap {
    /// Traces in the program mapped to the given file
    traces: BTreeMap<PathBuf, Vec<Trace>>,
}

impl TraceMap {
    /// Create a new TraceMap
    pub fn new() -> TraceMap {
        TraceMap {
            traces: BTreeMap::new(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.traces.is_empty()
    }

    pub fn iter(&self) -> Iter<PathBuf, Vec<Trace>> {
        self.traces.iter()
    }

    pub fn merge(&mut self, other: &TraceMap) {
        for ( k,  values) in other.iter() {
            if !self.traces.contains_key(k) {
                self.traces.insert(k.to_path_buf(), values.to_vec());
            } else {
                let mut existing = self.traces.get_mut(k).unwrap();
                for ref v in values.iter() {
                    let mut added = false;
                    if let Some(ref mut t) = existing.iter_mut().find(|ref x| x.line == v.line) {
                        t.stats = t.stats.clone() + v.stats.clone();
                        added = true;
                    }
                    if !added {
                        existing.push((*v).clone());
                        existing.sort_unstable();
                    }
                }
            }
        }
    }

    /// Add a trace to the tracemap for the given file
    pub fn add_trace(&mut self, file: &Path, trace: Trace) {
        if self.traces.contains_key(file) {
            if let Some(trace_vec) = self.traces.get_mut(file) {
                trace_vec.push(trace);
                trace_vec.sort_unstable();
            }
        } else {
            self.traces.insert(file.to_path_buf(), vec![trace]);
        }
    }

    /// Gets an immutable reference to a trace from an address. Returns None if
    /// there is no trace at that address
    pub fn get_trace(&self, address: u64) -> Option<&Trace> {
        self.all_traces()
            .iter()
            .find(|x| x.address == Some(address))
            .map(|x| *x)
    }

    /// Gets a mutable reference to a trace at a given address
    /// Returns None if there is no trace at that address
    pub fn get_trace_mut(&mut self, address: u64) -> Option<&mut Trace> {
        for val in self.all_traces_mut() {
            if val.address == Some(address) {
                return Some(val);
            }
        }
        None
    }

    /// Returns true if the location described by file and line number is present
    /// in the tracemap
    pub fn contains_location(&self, file: &Path, line: u64) -> bool {
        match self.traces.get(file) {
            Some(traces) => traces.iter().any(|x| x.line == line),
            None => false,
        }
    }

    /// Gets all traces below a certain path
    pub fn get_child_traces(&self, root: &Path) -> Vec<&Trace> {
        self.traces.iter()
                   .filter(|&(ref k, _)| k.starts_with(root))
                   .flat_map(|(_, ref v)| v.iter())
                   .collect()
    }

    /// Gets all traces
    pub fn all_traces(&self) -> Vec<&Trace> {
        self.traces.values().flat_map(|ref x| x.iter()).collect()
    }

    /// Gets a vector of all the traces to mutate
    fn all_traces_mut(&mut self) -> Vec<&mut Trace> {
        self.traces.values_mut().flat_map(|x| x.iter_mut()).collect()
    }

    pub fn files(&self) -> Vec<&PathBuf> {
        self.traces.keys().collect()
    }

    pub fn coverable_in_path(&self, path: &Path) -> usize {
        let mut result = 0usize;
        for t in self.get_child_traces(path) {
            result += match t.stats {
                CoverageStat::Branch(_) => {
                    2usize
                },
                CoverageStat::Condition(ref x) => {
                    x.len() * 2usize
                }
                _ => 1usize
            };
        }
        result
    }

    pub fn covered_in_path(&self, path: &Path) -> usize {
        let mut result = 0usize;
        for t in self.get_child_traces(path) {
            result += match t.stats {
                CoverageStat::Branch(ref x) => {
                    (x.been_true as usize) + (x.been_false as usize)
                },
                CoverageStat::Condition(ref x) => {
                    x.iter()
                     .fold(0, |acc, ref x| acc + (x.been_true as usize) + (x.been_false as usize))
                }
                CoverageStat::Line(ref x) => {
                    (x > &0) as usize
                }
            };
        }
        result
    }

    /// Give the total amount of coverable points in the code. This will vary
    /// based on the statistics available for line coverage it will be total
    /// line whereas for condition or decision it will count the number of
    /// conditions available
    pub fn total_coverable(&self) -> usize {
        let mut result = 0usize;
        for t in self.all_traces() {
            result += match t.stats {
                CoverageStat::Branch(_) => {
                    2usize
                },
                CoverageStat::Condition(ref x) => {
                    x.len() * 2usize
                }
                _ => 1usize
            };
        }
        result
    }

    /// From all the coverable data return the amount covered
    pub fn total_covered(&self) -> usize {
        let mut result = 0usize;
        for t in self.all_traces() {
            result += match t.stats {
                CoverageStat::Branch(ref x) => {
                    (x.been_true as usize) + (x.been_false as usize)
                },
                CoverageStat::Condition(ref x) => {
                    x.iter()
                     .fold(0, |acc, ref x| acc + (x.been_true as usize) + (x.been_false as usize))
                }
                CoverageStat::Line(ref x) => {
                    (x > &0) as usize
                }
            };
        }
        result
    }

    /// Returns coverage percentage ranging from 0.0-1.0
    pub fn coverage_percentage(&self) -> f64 {
        (self.total_covered() as f64) / (self.total_coverable() as f64)
    }

}

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

    #[test]
    fn stat_addition() {
        let x = CoverageStat::Line(0);
        let y = CoverageStat::Line(5);
        let z = CoverageStat::Line(7);
        let xy = x.clone() + y.clone();
        let yx = y.clone() + x.clone();
        let yy = y.clone() + y.clone();
        let zy = z.clone() + y.clone();
        assert_eq!(&xy, &CoverageStat::Line(5));
        assert_eq!(&yx, &xy);
        assert_eq!(&yy, &CoverageStat::Line(10));
        assert_eq!(&zy, &CoverageStat::Line(12));

        let tf = LogicState{been_true:true, been_false:true};
        let t = LogicState{been_true:true, been_false:false};
        let f = LogicState{been_true:false, been_false:true};
        let n = LogicState{been_true:false, been_false:false};

        assert_eq!(&t+&f, tf);
        assert_eq!(&t+&t, t);
        assert_eq!(&tf+&f, tf);
        assert_eq!(&tf+&t, tf);
        assert_eq!(&t+&n, t);
        assert_eq!(&n+&f, f);
        assert_eq!(&n+&n,n);
    }
}
<commit_msg>Started work on dedup<commit_after>use std::collections::{BTreeMap, HashMap};
use std::collections::btree_map::Iter;
use std::path::{PathBuf, Path};
use std::fmt::{Display, Formatter, Result};
use std::ops::Add;
use std::cmp::{Ord, Ordering};

/// Used to track the state of logical conditions
#[derive(Debug, Clone, Copy, Default, Hash, PartialEq, Eq, PartialOrd)]
pub struct LogicState {
    /// Whether the condition has been observed as true
    pub been_true: bool,
    /// Whether the condition has been observed as false
    pub been_false: bool,
}

impl<'a> Add for &'a LogicState {
    type Output = LogicState;

    fn add(self, other: &'a LogicState) -> LogicState {
        LogicState {
            been_true: self.been_true | other.been_true,
            been_false: self.been_false | other.been_false,
        }
    }
}

/// Shows what type of coverage data is being collected by a given trace
#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd)]
pub enum CoverageStat {
    /// Line coverage data (whether line has been hit)
    Line(u64),
    /// Branch coverage data (whether branch has been true and false
    Branch(LogicState),
    /// Condition coverage data (each boolean subcondition true and false)
    Condition(Vec<LogicState>),
}

impl Add for CoverageStat {
    type Output = CoverageStat;

    fn add(self, other: CoverageStat) -> CoverageStat {
        match (self, other) {
            (CoverageStat::Line(ref l), CoverageStat::Line(ref r)) => {
                CoverageStat::Line(l+r)
            },
            (CoverageStat::Branch(ref l), CoverageStat::Branch(ref r)) => {
                CoverageStat::Branch(l + r)
            },
            t @ _ => t.0,
        }
    }
}

impl Display for CoverageStat {
    fn fmt(&self, f: &mut Formatter) -> Result {
        match self {
            &CoverageStat::Line(x) => {
                write!(f, "hits: {}", x)
            },
            _ => write!(f, ""),
        }
    }
}

#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd)]
pub struct Trace {
    /// Line the trace is on in the file
    pub line: u64,
    /// Optional address showing location in the test artefact
    pub address: Option<u64>,
    /// Length of the instruction (useful to get entire condition/branch)
    pub length: usize,
    /// Coverage stats
    pub stats: CoverageStat,
}

/// Implemented to allow Traces to be sorted by line number
impl Ord for Trace {
    fn cmp(&self, other: &Trace) -> Ordering {
        self.line.cmp(&other.line)
    }
    fn max(self, other: Trace) -> Trace {
        if self.line > other.line {
            self
        } else {
            other
        }
    }
    fn min(self, other: Trace) -> Trace {
        if self.line < other.line {
            self
        } else {
            other
        }
    }
}

/// Stores all the program traces mapped to files and provides an interface to
/// add, query and change traces.
#[derive(Debug)]
pub struct TraceMap {
    /// Traces in the program mapped to the given file
    traces: BTreeMap<PathBuf, Vec<Trace>>,
}

impl TraceMap {
    /// Create a new TraceMap
    pub fn new() -> TraceMap {
        TraceMap {
            traces: BTreeMap::new(),
        }
    }

    /// Returns true if there are no traces
    pub fn is_empty(&self) -> bool {
        self.traces.is_empty()
    }

    /// Provides an interator to the underlying map of PathBufs to Vec<Trace>
    pub fn iter(&self) -> Iter<PathBuf, Vec<Trace>> {
        self.traces.iter()
    }

    /// Merges the results of one tracemap into the current one.
    /// This adds records which are missing and adds the statistics gathered to
    /// existing records
    pub fn merge(&mut self, other: &TraceMap) {
        for (k,  values) in other.iter() {
            if !self.traces.contains_key(k) {
                self.traces.insert(k.to_path_buf(), values.to_vec());
            } else {
                let mut existing = self.traces.get_mut(k).unwrap();
                for ref v in values.iter() {
                    let mut added = false;
                    if let Some(ref mut t) = existing.iter_mut().find(|ref x| x.line == v.line) {
                        t.stats = t.stats.clone() + v.stats.clone();
                        added = true;
                    }
                    if !added {
                        existing.push((*v).clone());
                        existing.sort_unstable();
                    }
                }
            }
        }
    }

    /// This will collapse duplicate Traces into a single trace. Warning this
    /// will lose the addresses of the duplicate traces but increment the results
    /// should be called only if you don't need those addresses from then on
    /// TODO possibly not the cleanest solution
    pub fn dedup(&mut self) {
        for (_, values) in self.traces.iter_mut() {
            let mut lines: HashMap<u64, CoverageStat> = HashMap::new();
            for v in values.iter() {
                if lines.contains_key(&v.line) {
                    let s = lines.get(&v.line).unwrap().clone() + v.stats.clone();
                    lines.insert(v.line, s);
                } else {
                    lines.insert(v.line, v.stats.clone());
                }
            }
        }
    }

    /// Add a trace to the tracemap for the given file
    pub fn add_trace(&mut self, file: &Path, trace: Trace) {
        if self.traces.contains_key(file) {
            if let Some(trace_vec) = self.traces.get_mut(file) {
                trace_vec.push(trace);
                trace_vec.sort_unstable();
            }
        } else {
            self.traces.insert(file.to_path_buf(), vec![trace]);
        }
    }

    /// Gets an immutable reference to a trace from an address. Returns None if
    /// there is no trace at that address
    pub fn get_trace(&self, address: u64) -> Option<&Trace> {
        self.all_traces()
            .iter()
            .find(|x| x.address == Some(address))
            .map(|x| *x)
    }

    /// Gets a mutable reference to a trace at a given address
    /// Returns None if there is no trace at that address
    pub fn get_trace_mut(&mut self, address: u64) -> Option<&mut Trace> {
        for val in self.all_traces_mut() {
            if val.address == Some(address) {
                return Some(val);
            }
        }
        None
    }

    /// Returns true if the location described by file and line number is present
    /// in the tracemap
    pub fn contains_location(&self, file: &Path, line: u64) -> bool {
        match self.traces.get(file) {
            Some(traces) => traces.iter().any(|x| x.line == line),
            None => false,
        }
    }

    /// Gets all traces below a certain path
    pub fn get_child_traces(&self, root: &Path) -> Vec<&Trace> {
        self.traces.iter()
                   .filter(|&(ref k, _)| k.starts_with(root))
                   .flat_map(|(_, ref v)| v.iter())
                   .collect()
    }

    /// Gets all traces
    pub fn all_traces(&self) -> Vec<&Trace> {
        self.traces.values().flat_map(|ref x| x.iter()).collect()
    }

    /// Gets a vector of all the traces to mutate
    fn all_traces_mut(&mut self) -> Vec<&mut Trace> {
        self.traces.values_mut().flat_map(|x| x.iter_mut()).collect()
    }

    pub fn files(&self) -> Vec<&PathBuf> {
        self.traces.keys().collect()
    }

    pub fn coverable_in_path(&self, path: &Path) -> usize {
        let mut result = 0usize;
        for t in self.get_child_traces(path) {
            result += match t.stats {
                CoverageStat::Branch(_) => {
                    2usize
                },
                CoverageStat::Condition(ref x) => {
                    x.len() * 2usize
                }
                _ => 1usize
            };
        }
        result
    }

    pub fn covered_in_path(&self, path: &Path) -> usize {
        let mut result = 0usize;
        for t in self.get_child_traces(path) {
            result += match t.stats {
                CoverageStat::Branch(ref x) => {
                    (x.been_true as usize) + (x.been_false as usize)
                },
                CoverageStat::Condition(ref x) => {
                    x.iter()
                     .fold(0, |acc, ref x| acc + (x.been_true as usize) + (x.been_false as usize))
                }
                CoverageStat::Line(ref x) => {
                    (x > &0) as usize
                }
            };
        }
        result
    }

    /// Give the total amount of coverable points in the code. This will vary
    /// based on the statistics available for line coverage it will be total
    /// line whereas for condition or decision it will count the number of
    /// conditions available
    pub fn total_coverable(&self) -> usize {
        let mut result = 0usize;
        for t in self.all_traces() {
            result += match t.stats {
                CoverageStat::Branch(_) => {
                    2usize
                },
                CoverageStat::Condition(ref x) => {
                    x.len() * 2usize
                }
                _ => 1usize
            };
        }
        result
    }

    /// From all the coverable data return the amount covered
    pub fn total_covered(&self) -> usize {
        let mut result = 0usize;
        for t in self.all_traces() {
            result += match t.stats {
                CoverageStat::Branch(ref x) => {
                    (x.been_true as usize) + (x.been_false as usize)
                },
                CoverageStat::Condition(ref x) => {
                    x.iter()
                     .fold(0, |acc, ref x| acc + (x.been_true as usize) + (x.been_false as usize))
                }
                CoverageStat::Line(ref x) => {
                    (x > &0) as usize
                }
            };
        }
        result
    }

    /// Returns coverage percentage ranging from 0.0-1.0
    pub fn coverage_percentage(&self) -> f64 {
        (self.total_covered() as f64) / (self.total_coverable() as f64)
    }

}

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

    #[test]
    fn stat_addition() {
        let x = CoverageStat::Line(0);
        let y = CoverageStat::Line(5);
        let z = CoverageStat::Line(7);
        let xy = x.clone() + y.clone();
        let yx = y.clone() + x.clone();
        let yy = y.clone() + y.clone();
        let zy = z.clone() + y.clone();
        assert_eq!(&xy, &CoverageStat::Line(5));
        assert_eq!(&yx, &xy);
        assert_eq!(&yy, &CoverageStat::Line(10));
        assert_eq!(&zy, &CoverageStat::Line(12));

        let tf = LogicState{been_true:true, been_false:true};
        let t = LogicState{been_true:true, been_false:false};
        let f = LogicState{been_true:false, been_false:true};
        let n = LogicState{been_true:false, been_false:false};

        assert_eq!(&t+&f, tf);
        assert_eq!(&t+&t, t);
        assert_eq!(&tf+&f, tf);
        assert_eq!(&tf+&t, tf);
        assert_eq!(&t+&n, t);
        assert_eq!(&n+&f, f);
        assert_eq!(&n+&n,n);
    }
}
<|endoftext|>


<commit_before>use super::Result;

pub trait R {
    type Connection;
    fn connect<T: IntoConnectOpts>(&self, opts: T) -> Result<Self::Connection>;
}

pub trait IntoConnectOpts {
    fn into(self) -> ::ConnectOpts;
}

pub trait Connector {
    type Connection;
    fn close(&self, noreply_wait: bool);
    fn reconnect(&self, noreply_wait: bool) -> Self::Connection;
    // use is a reserved keyword in Rust
    fn use_db(&self, db_name: &str) -> Self::Connection;
}
<commit_msg>Do not return a connection<commit_after>use super::Result;

pub trait R {
    fn connect<T: IntoConnectOpts>(&self, opts: T) -> Result<()>;
}

pub trait IntoConnectOpts {
    fn into(self) -> ::ConnectOpts;
}

pub trait Connector {
    type Connection;
    fn close(&self, noreply_wait: bool);
    fn reconnect(&self, noreply_wait: bool) -> Self::Connection;
    // use is a reserved keyword in Rust
    fn use_db(&self, db_name: &str) -> Self::Connection;
}
<|endoftext|>