<commit_before>//! Items related to the `osc::Receiver` implementation.

use super::{decode, rosc, CommunicationError, Connected, Packet, Unconnected};
use std;
use std::net::{SocketAddr, SocketAddrV4, ToSocketAddrs, UdpSocket};
use std::sync::atomic::{self, AtomicBool};
use std::sync::Mutex;

/// The default "maximum transmission unit" size as a number of bytes.
///
/// This is a common MTU size for ethernet.
pub const DEFAULT_MTU: usize = rosc::decoder::MTU;
/// By default UDP sockets are blocking so this is the mode in which the receiver is
/// initialised.
pub const DEFAULT_NON_BLOCKING: bool = false;

/// A type used for receiving OSC packets.
pub struct Receiver<M = Unconnected> {
    buffer: Mutex<Vec<u8>>,
    socket: UdpSocket,
    non_blocking: AtomicBool,
    mode: M,
}

/// An iterator that calls `recv` on the inner `Receiver` and yields the results.
///
/// If the `Receiver` is `Connected`, this will yield `Packet`s.
///
/// If the `Receiver` is `Unconnected`, this will yield `Packet`s alongside the source address.
///
/// Each call to `next` will block until the next packet is received or until some error
/// occurs.
pub struct Iter<'a, M = Unconnected>
where
    M: 'a,
{
    receiver: &'a Receiver<M>,
}

/// An iterator that calls `try_recv` on the inner `Receiver` and yields the results.
///
/// If the `Receiver` is `Connected`, this will yield `Packet`s.
///
/// If the `Receiver` is `Unconnected`, this will yield `Packet`s alongside the source address.
///
/// Each call to `next` will only return `Some` while there are pending messages and will
/// return `None` otherwise.
pub struct TryIter<'a, M = Unconnected>
where
    M: 'a,
{
    receiver: &'a Receiver<M>,
}

impl<M> Receiver<M> {
    /// The socket address that this `Receiver`'s socket was created from.
    pub fn local_addr(&self) -> Result<SocketAddr, std::io::Error> {
        self.socket.local_addr()
    }

    // Switch the `Receiver`'s inner socket to blocking mode.
    // This is for internal use only - the `recv` methods will call this automatically.
    fn switch_to_blocking(&self) -> Result<(), std::io::Error> {
        if self.non_blocking.load(atomic::Ordering::Relaxed) {
            self.socket.set_nonblocking(false)?;
            self.non_blocking.store(false, atomic::Ordering::Relaxed);
        }
        Ok(())
    }

    // Switch the `Receiver`'s inner socket to non-blocking mode.
    // This is for internal use only - the `try_recv` methods will call this automatically.
    fn switch_to_non_blocking(&self) -> Result<(), std::io::Error> {
        if !self.non_blocking.load(atomic::Ordering::Relaxed) {
            self.socket.set_nonblocking(true)?;
            self.non_blocking.store(true, atomic::Ordering::Relaxed);
        }
        Ok(())
    }
}

impl Receiver<Unconnected> {
    /// Create a `Receiver` that listen for OSC packets on the given address.
    ///
    /// ```no_run
    /// use nannou_osc::Receiver;
    ///
    /// fn main() {
    ///     let rx = Receiver::bind_to("127.0.0.1:34254").expect("Couldn't bind socket to address");
    /// }
    /// ```
    pub fn bind_to<A>(addr: A) -> Result<Self, std::io::Error>
    where
        A: ToSocketAddrs,
    {
        Self::bind_to_with_mtu(addr, DEFAULT_MTU)
    }

    /// The same as `bind_to`, but allows for manually specifying the MTU (aka "maximum transition
    /// unit").
    ///
    /// The MTU is the maximum UDP packet size in bytes that the `Receiver`'s UDP socket will
    /// receive before returning an error.
    ///
    /// By default this is `DEFAULT_MTU`.
    ///
    /// ```no_run
    /// use nannou_osc::Receiver;
    ///
    /// fn main() {
    ///     let rx = Receiver::bind_to_with_mtu("127.0.0.1:34254", 60_000)
    ///         .expect("Couldn't bind socket to address");
    /// }
    /// ```
    pub fn bind_to_with_mtu<A>(addr: A, mtu: usize) -> Result<Self, std::io::Error>
    where
        A: ToSocketAddrs,
    {
        let buffer = Mutex::new(vec![0; mtu]);
        let socket = UdpSocket::bind(addr)?;
        let non_blocking = AtomicBool::new(DEFAULT_NON_BLOCKING);
        let mode = Unconnected;
        let receiver = Receiver {
            buffer,
            socket,
            non_blocking,
            mode,
        };
        Ok(receiver)
    }

    /// The same as `bind_to`, but assumes that the IP address is `0.0.0.0`.
    ///
    /// The resulting socket address will be `0.0.0.0:<port>`.
    ///
    /// ```no_run
    /// use nannou_osc::Receiver;
    ///
    /// fn main() {
    ///     let rx = Receiver::bind(34254).expect("Couldn't bind socket to default address");
    /// }
    /// ```
    pub fn bind(port: u16) -> Result<Self, std::io::Error> {
        Self::bind_to(SocketAddrV4::new(super::default_ipv4_addr(), port))
    }

    /// The same as `bind_to_with_mtu`, but assumes that the IP address is `0.0.0.0`.
    ///
    /// The resulting socket address will be `0.0.0.0:<port>`.
    ///
    /// ```no_run
    /// use nannou_osc::Receiver;
    ///
    /// fn main() {
    ///     let port = 34254;
    ///     let mtu = 60_000;
    ///     let rx = Receiver::bind_with_mtu(port, mtu).expect("Couldn't bind to default address");
    /// }
    /// ```
    pub fn bind_with_mtu(port: u16, mtu: usize) -> Result<Self, std::io::Error> {
        Self::bind_to_with_mtu(SocketAddrV4::new(super::default_ipv4_addr(), port), mtu)
    }

    /// Connects the `Receiver`'s UDP socket to the given remote address.
    ///
    /// This applies filters so that only data from the given address is received.
    ///
    /// **Panic!**s if the given `addr` cannot resolve to a valid `SocketAddr`.
    ///
    /// ```no_run
    /// use nannou_osc::Receiver;
    ///
    /// fn main() {
    ///     let tx = Receiver::bind(34254)
    ///         .expect("Couldn't bind to default socket")
    ///         .connect("127.0.0.1:34255")
    ///         .expect("Couldn't connect to socket at address");
    /// }
    /// ```
    pub fn connect<A>(self, addr: A) -> Result<Receiver<Connected>, std::io::Error>
    where
        A: ToSocketAddrs,
    {
        let Receiver {
            buffer,
            socket,
            non_blocking,
            ..
        } = self;
        let mut addrs = addr.to_socket_addrs()?;
        let addr = addrs.next().expect("could not resolve any `SocketAddr`s");
        socket.connect(addr)?;
        let mode = Connected { addr };
        Ok(Receiver {
            buffer,
            socket,
            non_blocking,
            mode,
        })
    }

    /// Waits for the next OSC packet to be received and returns it along with the source address.
    ///
    /// If the socket is currently in non-blocking mode, this method will first switch the socket
    /// to blocking.
    ///
    /// This will return a `CommunicationError` if:
    ///
    /// - Switching the socket from "non_blocking" to "blocking" fails,
    /// - The Mutex around the inner buffer (used to collect bytes) was poisoned,
    /// - The MTU was not large enough to receive a UDP packet,
    /// - The inner `UdpSocket::recv` call fails or
    /// - The socket received some bytes that could not be decoded into an OSC `Packet`.
    pub fn recv(&self) -> Result<(Packet, SocketAddr), CommunicationError> {
        self.switch_to_blocking()?;
        let mut buffer = self.buffer.lock()?;
        let (len, addr) = self.socket.recv_from(&mut buffer)?;
        let packet = decode(&buffer[..len])?;
        Ok((packet, addr))
    }

    /// Checks for a pending OSC packet and returns `Ok(Some)` if there is one waiting along with
    /// the source address.
    ///
    /// If there are no packets waiting (or if the inner UDP socket's `recv` method returns an
    /// error) this will immediately return with `Ok(None)`.
    ///
    /// If the socket is currently in blocking mode, this method will first switch the socket to
    /// non-blocking.
    ///
    /// This will return a `CommunicationError` if:
    ///
    /// - Switching the socket from "blocking" to "non_blocking" fails,
    /// - The Mutex around the inner buffer (used to collect bytes) was poisoned,
    /// - The socket received some bytes that could not be decoded into an OSC `Packet`.
    pub fn try_recv(&self) -> Result<Option<(Packet, SocketAddr)>, CommunicationError> {
        self.switch_to_non_blocking()?;
        let mut buffer = self.buffer.lock()?;
        let (len, addr) = match self.socket.recv_from(&mut buffer) {
            Ok(tuple) => tuple,
            // TODO: Don't know how to check for the specific error that is returned when the
            // non_blocking socket has no bytes waiting, so we just always assume that's what the
            // error was. This should probably be fixed somehow to distinguish between errors.
            Err(_) => return Ok(None),
        };
        let packet = decode(&buffer[..len])?;
        Ok(Some((packet, addr)))
    }

    /// An iterator yielding OSC `Packet`s along with their source address.
    ///
    /// Each call to `next` will block until the next packet is received or until some error
    /// occurs.
    pub fn iter(&self) -> Iter<Unconnected> {
        Iter { receiver: self }
    }

    /// An iterator yielding OSC `Packet`s along with their source address.
    ///
    /// Each call to `next` will only return `Some` while there are pending packets and will return
    /// `None` otherwise.
    pub fn try_iter(&self) -> TryIter<Unconnected> {
        TryIter { receiver: self }
    }
}

impl Receiver<Connected> {
    /// Returns the address of the socket to which the `Receiver` is `Connected`.
    pub fn remote_addr(&self) -> SocketAddr {
        self.mode.addr
    }

    /// Waits for the next OSC packet to be received and returns it.
    ///
    /// If the socket is currently in non-blocking mode, this method will first switch the socket
    /// to blocking.
    ///
    /// This will return a `CommunicationError` if:
    ///
    /// - Switching the socket from "non_blocking" to "blocking" fails,
    /// - The Mutex around the inner buffer (used to collect bytes) was poisoned,
    /// - The MTU was not large enough to receive a UDP packet,
    /// - The inner `UdpSocket::recv` call fails or
    /// - The socket received some bytes that could not be decoded into an OSC `Packet`.
    pub fn recv(&self) -> Result<Packet, CommunicationError> {
        self.switch_to_blocking()?;
        let mut buffer = self.buffer.lock()?;
        let len = self.socket.recv(&mut buffer)?;
        let packet = decode(&buffer[..len])?;
        Ok(packet)
    }

    /// Checks for a pending OSC packet and returns `Ok(Some)` if there is one waiting.
    ///
    /// If there are no packets waiting (or if the inner UDP socket's `recv` method returns an
    /// error) this will immediately return with `Ok(None)`.
    ///
    /// If the socket is currently in blocking mode, this method will first switch the socket to
    /// non-blocking.
    ///
    /// This will return a `CommunicationError` if:
    ///
    /// - Switching the socket from "blocking" to "non_blocking" fails,
    /// - The Mutex around the inner buffer (used to collect bytes) was poisoned,
    /// - The socket received some bytes that could not be decoded into an OSC `Packet`.
    pub fn try_recv(&self) -> Result<Option<Packet>, CommunicationError> {
        self.switch_to_non_blocking()?;
        let mut buffer = self.buffer.lock()?;
        let len = match self.socket.recv(&mut buffer) {
            Ok(len) => len,
            // TODO: Don't know how to check for the specific error that is returned when the
            // non_blocking socket has no bytes waiting, so we just always assume that's what the
            // error was. This should probably be fixed somehow to distinguish between errors.
            Err(_) => return Ok(None),
        };
        let packet = decode(&buffer[..len])?;
        Ok(Some(packet))
    }

    /// An iterator yielding OSC `Packet`s.
    ///
    /// Each call to `next` will block until the next packet is received or until some error
    /// occurs.
    pub fn iter(&self) -> Iter<Connected> {
        Iter { receiver: self }
    }

    /// An iterator yielding OSC `Packet`s.
    ///
    /// Each call to `next` will only return `Some` while there are pending packets and will return
    /// `None` otherwise.
    pub fn try_iter(&self) -> TryIter<Connected> {
        TryIter { receiver: self }
    }
}

impl<'a> Iterator for Iter<'a, Connected> {
    type Item = Packet;
    fn next(&mut self) -> Option<Self::Item> {
        self.receiver.recv().ok()
    }
}

impl<'a> Iterator for Iter<'a, Unconnected> {
    type Item = (Packet, SocketAddr);
    fn next(&mut self) -> Option<Self::Item> {
        self.receiver.recv().ok()
    }
}

impl<'a> Iterator for TryIter<'a, Connected> {
    type Item = Packet;
    fn next(&mut self) -> Option<Self::Item> {
        self.receiver.try_recv().ok().and_then(|p| p)
    }
}

impl<'a> Iterator for TryIter<'a, Unconnected> {
    type Item = (Packet, SocketAddr);
    fn next(&mut self) -> Option<Self::Item> {
        self.receiver.try_recv().ok().and_then(|p| p)
    }
}
<commit_msg>add note about message being popped<commit_after>//! Items related to the `osc::Receiver` implementation.

use super::{decode, rosc, CommunicationError, Connected, Packet, Unconnected};
use std;
use std::net::{SocketAddr, SocketAddrV4, ToSocketAddrs, UdpSocket};
use std::sync::atomic::{self, AtomicBool};
use std::sync::Mutex;

/// The default "maximum transmission unit" size as a number of bytes.
///
/// This is a common MTU size for ethernet.
pub const DEFAULT_MTU: usize = rosc::decoder::MTU;
/// By default UDP sockets are blocking so this is the mode in which the receiver is
/// initialised.
pub const DEFAULT_NON_BLOCKING: bool = false;

/// A type used for receiving OSC packets.
pub struct Receiver<M = Unconnected> {
    buffer: Mutex<Vec<u8>>,
    socket: UdpSocket,
    non_blocking: AtomicBool,
    mode: M,
}

/// An iterator that calls `recv` on the inner `Receiver` and yields the results.
///
/// If the `Receiver` is `Connected`, this will yield `Packet`s.
///
/// If the `Receiver` is `Unconnected`, this will yield `Packet`s alongside the source address.
///
/// Each call to `next` will block until the next packet is received or until some error
/// occurs.
pub struct Iter<'a, M = Unconnected>
where
    M: 'a,
{
    receiver: &'a Receiver<M>,
}

/// An iterator that calls `try_recv` on the inner `Receiver` and yields the results.
///
/// If the `Receiver` is `Connected`, this will yield `Packet`s.
///
/// If the `Receiver` is `Unconnected`, this will yield `Packet`s alongside the source address.
///
/// Each call to `next` will only return `Some` while there are pending messages and will
/// return `None` otherwise.
pub struct TryIter<'a, M = Unconnected>
where
    M: 'a,
{
    receiver: &'a Receiver<M>,
}

impl<M> Receiver<M> {
    /// The socket address that this `Receiver`'s socket was created from.
    pub fn local_addr(&self) -> Result<SocketAddr, std::io::Error> {
        self.socket.local_addr()
    }

    // Switch the `Receiver`'s inner socket to blocking mode.
    // This is for internal use only - the `recv` methods will call this automatically.
    fn switch_to_blocking(&self) -> Result<(), std::io::Error> {
        if self.non_blocking.load(atomic::Ordering::Relaxed) {
            self.socket.set_nonblocking(false)?;
            self.non_blocking.store(false, atomic::Ordering::Relaxed);
        }
        Ok(())
    }

    // Switch the `Receiver`'s inner socket to non-blocking mode.
    // This is for internal use only - the `try_recv` methods will call this automatically.
    fn switch_to_non_blocking(&self) -> Result<(), std::io::Error> {
        if !self.non_blocking.load(atomic::Ordering::Relaxed) {
            self.socket.set_nonblocking(true)?;
            self.non_blocking.store(true, atomic::Ordering::Relaxed);
        }
        Ok(())
    }
}

impl Receiver<Unconnected> {
    /// Create a `Receiver` that listen for OSC packets on the given address.
    ///
    /// ```no_run
    /// use nannou_osc::Receiver;
    ///
    /// fn main() {
    ///     let rx = Receiver::bind_to("127.0.0.1:34254").expect("Couldn't bind socket to address");
    /// }
    /// ```
    pub fn bind_to<A>(addr: A) -> Result<Self, std::io::Error>
    where
        A: ToSocketAddrs,
    {
        Self::bind_to_with_mtu(addr, DEFAULT_MTU)
    }

    /// The same as `bind_to`, but allows for manually specifying the MTU (aka "maximum transition
    /// unit").
    ///
    /// The MTU is the maximum UDP packet size in bytes that the `Receiver`'s UDP socket will
    /// receive before returning an error.
    ///
    /// By default this is `DEFAULT_MTU`.
    ///
    /// ```no_run
    /// use nannou_osc::Receiver;
    ///
    /// fn main() {
    ///     let rx = Receiver::bind_to_with_mtu("127.0.0.1:34254", 60_000)
    ///         .expect("Couldn't bind socket to address");
    /// }
    /// ```
    pub fn bind_to_with_mtu<A>(addr: A, mtu: usize) -> Result<Self, std::io::Error>
    where
        A: ToSocketAddrs,
    {
        let buffer = Mutex::new(vec![0; mtu]);
        let socket = UdpSocket::bind(addr)?;
        let non_blocking = AtomicBool::new(DEFAULT_NON_BLOCKING);
        let mode = Unconnected;
        let receiver = Receiver {
            buffer,
            socket,
            non_blocking,
            mode,
        };
        Ok(receiver)
    }

    /// The same as `bind_to`, but assumes that the IP address is `0.0.0.0`.
    ///
    /// The resulting socket address will be `0.0.0.0:<port>`.
    ///
    /// ```no_run
    /// use nannou_osc::Receiver;
    ///
    /// fn main() {
    ///     let rx = Receiver::bind(34254).expect("Couldn't bind socket to default address");
    /// }
    /// ```
    pub fn bind(port: u16) -> Result<Self, std::io::Error> {
        Self::bind_to(SocketAddrV4::new(super::default_ipv4_addr(), port))
    }

    /// The same as `bind_to_with_mtu`, but assumes that the IP address is `0.0.0.0`.
    ///
    /// The resulting socket address will be `0.0.0.0:<port>`.
    ///
    /// ```no_run
    /// use nannou_osc::Receiver;
    ///
    /// fn main() {
    ///     let port = 34254;
    ///     let mtu = 60_000;
    ///     let rx = Receiver::bind_with_mtu(port, mtu).expect("Couldn't bind to default address");
    /// }
    /// ```
    pub fn bind_with_mtu(port: u16, mtu: usize) -> Result<Self, std::io::Error> {
        Self::bind_to_with_mtu(SocketAddrV4::new(super::default_ipv4_addr(), port), mtu)
    }

    /// Connects the `Receiver`'s UDP socket to the given remote address.
    ///
    /// This applies filters so that only data from the given address is received.
    ///
    /// **Panic!**s if the given `addr` cannot resolve to a valid `SocketAddr`.
    ///
    /// ```no_run
    /// use nannou_osc::Receiver;
    ///
    /// fn main() {
    ///     let tx = Receiver::bind(34254)
    ///         .expect("Couldn't bind to default socket")
    ///         .connect("127.0.0.1:34255")
    ///         .expect("Couldn't connect to socket at address");
    /// }
    /// ```
    pub fn connect<A>(self, addr: A) -> Result<Receiver<Connected>, std::io::Error>
    where
        A: ToSocketAddrs,
    {
        let Receiver {
            buffer,
            socket,
            non_blocking,
            ..
        } = self;
        let mut addrs = addr.to_socket_addrs()?;
        let addr = addrs.next().expect("could not resolve any `SocketAddr`s");
        socket.connect(addr)?;
        let mode = Connected { addr };
        Ok(Receiver {
            buffer,
            socket,
            non_blocking,
            mode,
        })
    }

    /// Waits for the next OSC packet to be received and returns it along with the source address.
    ///
    /// If the socket is currently in non-blocking mode, this method will first switch the socket
    /// to blocking.
    ///
    /// This will return a `CommunicationError` if:
    ///
    /// - Switching the socket from "non_blocking" to "blocking" fails,
    /// - The Mutex around the inner buffer (used to collect bytes) was poisoned,
    /// - The MTU was not large enough to receive a UDP packet,
    /// - The inner `UdpSocket::recv` call fails or
    /// - The socket received some bytes that could not be decoded into an OSC `Packet`.
    pub fn recv(&self) -> Result<(Packet, SocketAddr), CommunicationError> {
        self.switch_to_blocking()?;
        let mut buffer = self.buffer.lock()?;
        let (len, addr) = self.socket.recv_from(&mut buffer)?;
        let packet = decode(&buffer[..len])?;
        Ok((packet, addr))
    }

    /// Checks for a pending OSC packet and returns `Ok(Some)` if there is one waiting along with
    /// the source address.
    ///
    /// If there are no packets waiting (or if the inner UDP socket's `recv` method returns an
    /// error) this will immediately return with `Ok(None)`.
    ///
    /// When called, this method will pop a message from the receiver's queue.
    ///
    ///
    /// If the socket is currently in blocking mode, this method will first switch the socket to
    /// non-blocking.
    ///
    /// This will return a `CommunicationError` if:
    ///
    /// - Switching the socket from "blocking" to "non_blocking" fails,
    /// - The Mutex around the inner buffer (used to collect bytes) was poisoned,
    /// - The socket received some bytes that could not be decoded into an OSC `Packet`.
    pub fn try_recv(&self) -> Result<Option<(Packet, SocketAddr)>, CommunicationError> {
        self.switch_to_non_blocking()?;
        let mut buffer = self.buffer.lock()?;
        let (len, addr) = match self.socket.recv_from(&mut buffer) {
            Ok(tuple) => tuple,
            // TODO: Don't know how to check for the specific error that is returned when the
            // non_blocking socket has no bytes waiting, so we just always assume that's what the
            // error was. This should probably be fixed somehow to distinguish between errors.
            Err(_) => return Ok(None),
        };
        let packet = decode(&buffer[..len])?;
        Ok(Some((packet, addr)))
    }

    /// An iterator yielding OSC `Packet`s along with their source address.
    ///
    /// Each call to `next` will block until the next packet is received or until some error
    /// occurs.
    pub fn iter(&self) -> Iter<Unconnected> {
        Iter { receiver: self }
    }

    /// An iterator yielding OSC `Packet`s along with their source address.
    ///
    /// Each call to `next` will only return `Some` while there are pending packets and will return
    /// `None` otherwise.
    pub fn try_iter(&self) -> TryIter<Unconnected> {
        TryIter { receiver: self }
    }
}

impl Receiver<Connected> {
    /// Returns the address of the socket to which the `Receiver` is `Connected`.
    pub fn remote_addr(&self) -> SocketAddr {
        self.mode.addr
    }

    /// Waits for the next OSC packet to be received and returns it.
    ///
    /// If the socket is currently in non-blocking mode, this method will first switch the socket
    /// to blocking.
    ///
    /// This will return a `CommunicationError` if:
    ///
    /// - Switching the socket from "non_blocking" to "blocking" fails,
    /// - The Mutex around the inner buffer (used to collect bytes) was poisoned,
    /// - The MTU was not large enough to receive a UDP packet,
    /// - The inner `UdpSocket::recv` call fails or
    /// - The socket received some bytes that could not be decoded into an OSC `Packet`.
    pub fn recv(&self) -> Result<Packet, CommunicationError> {
        self.switch_to_blocking()?;
        let mut buffer = self.buffer.lock()?;
        let len = self.socket.recv(&mut buffer)?;
        let packet = decode(&buffer[..len])?;
        Ok(packet)
    }

    /// Checks for a pending OSC packet and returns `Ok(Some)` if there is one waiting.
    ///
    /// If there are no packets waiting (or if the inner UDP socket's `recv` method returns an
    /// error) this will immediately return with `Ok(None)`.
    ///
    /// If the socket is currently in blocking mode, this method will first switch the socket to
    /// non-blocking.
    ///
    /// This will return a `CommunicationError` if:
    ///
    /// - Switching the socket from "blocking" to "non_blocking" fails,
    /// - The Mutex around the inner buffer (used to collect bytes) was poisoned,
    /// - The socket received some bytes that could not be decoded into an OSC `Packet`.
    pub fn try_recv(&self) -> Result<Option<Packet>, CommunicationError> {
        self.switch_to_non_blocking()?;
        let mut buffer = self.buffer.lock()?;
        let len = match self.socket.recv(&mut buffer) {
            Ok(len) => len,
            // TODO: Don't know how to check for the specific error that is returned when the
            // non_blocking socket has no bytes waiting, so we just always assume that's what the
            // error was. This should probably be fixed somehow to distinguish between errors.
            Err(_) => return Ok(None),
        };
        let packet = decode(&buffer[..len])?;
        Ok(Some(packet))
    }

    /// An iterator yielding OSC `Packet`s.
    ///
    /// Each call to `next` will block until the next packet is received or until some error
    /// occurs.
    pub fn iter(&self) -> Iter<Connected> {
        Iter { receiver: self }
    }

    /// An iterator yielding OSC `Packet`s.
    ///
    /// Each call to `next` will only return `Some` while there are pending packets and will return
    /// `None` otherwise.
    pub fn try_iter(&self) -> TryIter<Connected> {
        TryIter { receiver: self }
    }
}

impl<'a> Iterator for Iter<'a, Connected> {
    type Item = Packet;
    fn next(&mut self) -> Option<Self::Item> {
        self.receiver.recv().ok()
    }
}

impl<'a> Iterator for Iter<'a, Unconnected> {
    type Item = (Packet, SocketAddr);
    fn next(&mut self) -> Option<Self::Item> {
        self.receiver.recv().ok()
    }
}

impl<'a> Iterator for TryIter<'a, Connected> {
    type Item = Packet;
    fn next(&mut self) -> Option<Self::Item> {
        self.receiver.try_recv().ok().and_then(|p| p)
    }
}

impl<'a> Iterator for TryIter<'a, Unconnected> {
    type Item = (Packet, SocketAddr);
    fn next(&mut self) -> Option<Self::Item> {
        self.receiver.try_recv().ok().and_then(|p| p)
    }
}
<|endoftext|>


<commit_before>#[derive(Debug)]
#[repr(C)]
pub struct MemoryMapTag {
    typ: u32,
    size: u32,
    entry_size: u32,
    entry_version: u32,
    first_area: MemoryArea,
}

impl MemoryMapTag {
    pub fn memory_areas(&self) -> MemoryAreaIter {
        let self_ptr = self as *const MemoryMapTag;
        let start_area = (&self.first_area) as *const MemoryArea;
        MemoryAreaIter {
            current_area: start_area as u64,
            last_area: (self_ptr as u64 + (self.size - self.entry_size) as u64),
            entry_size: self.entry_size,
        }
    }
}

#[derive(Debug)]
#[repr(C)]
pub struct MemoryArea {
    base_addr: u64,
    length: u64,
    typ: u32,
    _reserved: u32,
}

impl MemoryArea {
    pub fn start_address(&self) -> u64 {
        self.base_addr
    }

    pub fn end_address(&self) -> u64 {
        (self.base_addr + self.length)
    }

    pub fn size(&self) -> u64 {
        self.length
    }
}

#[derive(Clone, Debug)]
pub struct MemoryAreaIter {
    current_area: u64,
    last_area: u64,
    entry_size: u32,
}

impl Iterator for MemoryAreaIter {
    type Item = &'static MemoryArea;
    fn next(&mut self) -> Option<&'static MemoryArea> {
        if self.current_area > self.last_area {
            None
        } else {
            let area = unsafe{&*(self.current_area as *const MemoryArea)};
            self.current_area = self.current_area + (self.entry_size as u64);
            if area.typ == 1 {
                Some(area)
            } else {self.next()}
        }
    }
}
<commit_msg>Use 'a instead of 'static for MemoryAreaIter<commit_after>use core::marker::PhantomData;

#[derive(Debug)]
#[repr(C)]
pub struct MemoryMapTag {
    typ: u32,
    size: u32,
    entry_size: u32,
    entry_version: u32,
    first_area: MemoryArea,
}

impl MemoryMapTag {
    pub fn memory_areas(&self) -> MemoryAreaIter {
        let self_ptr = self as *const MemoryMapTag;
        let start_area = (&self.first_area) as *const MemoryArea;
        MemoryAreaIter {
            current_area: start_area as u64,
            last_area: (self_ptr as u64 + (self.size - self.entry_size) as u64),
            entry_size: self.entry_size,
            phantom: PhantomData,
        }
    }
}

#[derive(Debug)]
#[repr(C)]
pub struct MemoryArea {
    base_addr: u64,
    length: u64,
    typ: u32,
    _reserved: u32,
}

impl MemoryArea {
    pub fn start_address(&self) -> u64 {
        self.base_addr
    }

    pub fn end_address(&self) -> u64 {
        (self.base_addr + self.length)
    }

    pub fn size(&self) -> u64 {
        self.length
    }
}

#[derive(Clone, Debug)]
pub struct MemoryAreaIter<'a> {
    current_area: u64,
    last_area: u64,
    entry_size: u32,
    phantom: PhantomData<&'a MemoryArea>,
}

impl<'a> Iterator for MemoryAreaIter<'a> {
    type Item = &'a MemoryArea;
    fn next(&mut self) -> Option<&'a MemoryArea> {
        if self.current_area > self.last_area {
            None
        } else {
            let area = unsafe{&*(self.current_area as *const MemoryArea)};
            self.current_area = self.current_area + (self.entry_size as u64);
            if area.typ == 1 {
                Some(area)
            } else {self.next()}
        }
    }
}
<|endoftext|>


<commit_before>use crate::kernel_data::with_kernel_data;
use crate::upcall::{UpcallId, UpcallQueueEntry};
use std::cell::Cell;

/// A reference used by a `fake::SyscallDriver` to access data shared between it
/// and the `fake::Kernel`, such as upcalls.
///
/// A `Default`-initialized `DriverShareRef` links to an empty share. It can be
/// used as normal, but is generally a no-op. This allows `fake::SyscallDriver`
/// implementations to store a `DriverShareRef` directly, rather than having to
/// contain a `Cell<Option<DriverShareRef>>`.
#[derive(Default)]
pub struct DriverShareRef {
    pub(crate) driver_num: Cell<u32>,
}

impl DriverShareRef {
    /// Replaces this DriverShareRef with another. `fake:SyscallDrivers` can use
    /// `replace` to implement `register` to avoid having to store their
    /// `DriverShareRef` inside a `Cell`.
    pub fn replace(&self, new: Self) {
        self.driver_num.set(new.driver_num.get());
    }

    /// Schedules the upcall with the specified subscribe number. Like the real
    /// kernel, this does nothing if there is no upcall with number
    /// `subscribe_num` or the upcall is the null upcall.
    pub fn schedule_upcall(
        &self,
        subscribe_num: u32,
        args: (u32, u32, u32),
    ) -> Result<(), InvalidSubscribeNum> {
        with_kernel_data(|kernel_data| {
            let kernel_data = match kernel_data {
                Some(kernel_data) => kernel_data,
                None => return Ok(()),
            };
            let driver_data = kernel_data
                .drivers
                .get(&self.driver_num.get())
                .expect("DriverShareRef: registered but nonexistent?");
            if subscribe_num >= driver_data.num_upcalls {
                return Err(InvalidSubscribeNum {
                    upcall_count: driver_data.num_upcalls,
                    requested: subscribe_num,
                });
            }
            let upcall = match driver_data.upcalls.get(&subscribe_num) {
                Some(&upcall) => upcall,
                None => return Ok(()),
            };
            // Don't bother queueing a null upcall, as they don't do anything
            // when invoked anyway, and the core kernel does not queue them
            // either.
            if upcall.is_null() {
                return Ok(());
            }
            kernel_data.upcall_queue.push_back(UpcallQueueEntry {
                args,
                id: UpcallId {
                    driver_num: self.driver_num.get(),
                    subscribe_num,
                },
                upcall,
            });
            Ok(())
        })
    }
}

#[derive(Debug, Eq, PartialEq, thiserror::Error)]
#[error("Upcall number {requested} too large, expected < {upcall_count}.")]
pub struct InvalidSubscribeNum {
    requested: u32,
    upcall_count: u32,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::upcall::Upcall;
    use crate::DriverInfo;
    use std::rc::Rc;

    #[derive(Default)]
    struct MockDriver {
        share_ref: DriverShareRef,
    }
    impl crate::fake::SyscallDriver for MockDriver {
        fn info(&self) -> DriverInfo {
            DriverInfo::new(1).upcall_count(10)
        }
        fn register(&self, share_ref: DriverShareRef) {
            self.share_ref.replace(share_ref);
        }
        fn command(&self, _: u32, _: u32, _: u32) -> libtock_platform::CommandReturn {
            crate::command_return::failure(libtock_platform::ErrorCode::NoSupport)
        }
    }

    #[test]
    fn schedule_errors() {
        let mock_driver = Rc::new(MockDriver::default());
        let kernel = crate::fake::Kernel::new();
        kernel.add_driver(&mock_driver);
        assert_eq!(
            mock_driver.share_ref.schedule_upcall(10, (3, 4, 5)),
            Err(InvalidSubscribeNum {
                upcall_count: 10,
                requested: 10
            })
        );
    }

    #[test]
    fn schedule_success() {
        let mock_driver = Rc::new(MockDriver::default());
        let kernel = crate::fake::Kernel::new();
        kernel.add_driver(&mock_driver);

        // Call schedule with no registered upcall.
        assert_eq!(mock_driver.share_ref.schedule_upcall(2, (3, 4, 5)), Ok(()));
        crate::kernel_data::with_kernel_data(|kernel_data| {
            let kernel_data = kernel_data.unwrap();

            // There was no upcall to schedule, so the queue should still be
            // empty.
            assert!(kernel_data.upcall_queue.is_empty());

            // Register a null upcall.
            kernel_data.drivers.get_mut(&1).unwrap().upcalls.insert(
                2,
                Upcall {
                    fn_pointer: None,
                    data: 1234u32.into(),
                },
            );
        });
        // Call schedule again. This should still do nothing, because the upcall
        // is a null upcall.
        assert_eq!(mock_driver.share_ref.schedule_upcall(2, (3, 4, 5)), Ok(()));
        unsafe extern "C" fn upcall(_: u32, _: u32, _: u32, _: libtock_platform::Register) {}
        crate::kernel_data::with_kernel_data(|kernel_data| {
            let kernel_data = kernel_data.unwrap();

            // Very the upcall was not queued.
            assert!(kernel_data.upcall_queue.is_empty());

            // Register a non-null upcall.
            kernel_data.drivers.get_mut(&1).unwrap().upcalls.insert(
                2,
                Upcall {
                    fn_pointer: Some(upcall),
                    data: 1111usize.into(),
                },
            );
        });
        // Call schedule again. This should schedule the upcall.
        assert_eq!(mock_driver.share_ref.schedule_upcall(2, (3, 4, 5)), Ok(()));
        crate::kernel_data::with_kernel_data(|kernel_data| {
            let kernel_data = kernel_data.unwrap();

            // Verify the upcall was queued.
            assert_eq!(kernel_data.upcall_queue.len(), 1);
            let upcall_queue_entry = kernel_data.upcall_queue.front().expect("Upcall not queued");
            assert_eq!(upcall_queue_entry.args, (3, 4, 5));
            assert_eq!(
                upcall_queue_entry.id,
                UpcallId {
                    driver_num: 1,
                    subscribe_num: 2
                }
            );
            assert!(upcall_queue_entry.upcall.fn_pointer == Some(upcall));
            let data: usize = upcall_queue_entry.upcall.data.into();
            assert_eq!(data, 1111);

            // Register a non-null upcall.
            kernel_data.drivers.get_mut(&1).unwrap().upcalls.insert(
                2,
                Upcall {
                    fn_pointer: Some(upcall),
                    data: 2222u32.into(),
                },
            );
        });
        // Call schedule again. This should schedule another upcall, after the
        // first.
        assert_eq!(
            mock_driver.share_ref.schedule_upcall(2, (30, 40, 50)),
            Ok(())
        );
        crate::kernel_data::with_kernel_data(|kernel_data| {
            let kernel_data = kernel_data.unwrap();

            // Very the upcall was queued.
            assert_eq!(kernel_data.upcall_queue.len(), 2);
            let front_queue_entry = kernel_data.upcall_queue.front().expect("Upcall not queued");
            assert_eq!(front_queue_entry.args, (3, 4, 5));
            assert_eq!(
                front_queue_entry.id,
                UpcallId {
                    driver_num: 1,
                    subscribe_num: 2
                }
            );
            assert!(front_queue_entry.upcall.fn_pointer == Some(upcall));
            let front_data: usize = front_queue_entry.upcall.data.into();
            assert_eq!(front_data, 1111);
            let back_queue_entry = kernel_data.upcall_queue.back().expect("Upcall not queued");
            assert_eq!(back_queue_entry.args, (30, 40, 50));
            assert_eq!(
                back_queue_entry.id,
                UpcallId {
                    driver_num: 1,
                    subscribe_num: 2
                }
            );
            assert!(back_queue_entry.upcall.fn_pointer == Some(upcall));
            let back_data: usize = back_queue_entry.upcall.data.into();
            assert_eq!(back_data, 2222);
        });
    }
}
<commit_msg>Typo correction: "Very the" -> "Verify the".<commit_after>use crate::kernel_data::with_kernel_data;
use crate::upcall::{UpcallId, UpcallQueueEntry};
use std::cell::Cell;

/// A reference used by a `fake::SyscallDriver` to access data shared between it
/// and the `fake::Kernel`, such as upcalls.
///
/// A `Default`-initialized `DriverShareRef` links to an empty share. It can be
/// used as normal, but is generally a no-op. This allows `fake::SyscallDriver`
/// implementations to store a `DriverShareRef` directly, rather than having to
/// contain a `Cell<Option<DriverShareRef>>`.
#[derive(Default)]
pub struct DriverShareRef {
    pub(crate) driver_num: Cell<u32>,
}

impl DriverShareRef {
    /// Replaces this DriverShareRef with another. `fake:SyscallDrivers` can use
    /// `replace` to implement `register` to avoid having to store their
    /// `DriverShareRef` inside a `Cell`.
    pub fn replace(&self, new: Self) {
        self.driver_num.set(new.driver_num.get());
    }

    /// Schedules the upcall with the specified subscribe number. Like the real
    /// kernel, this does nothing if there is no upcall with number
    /// `subscribe_num` or the upcall is the null upcall.
    pub fn schedule_upcall(
        &self,
        subscribe_num: u32,
        args: (u32, u32, u32),
    ) -> Result<(), InvalidSubscribeNum> {
        with_kernel_data(|kernel_data| {
            let kernel_data = match kernel_data {
                Some(kernel_data) => kernel_data,
                None => return Ok(()),
            };
            let driver_data = kernel_data
                .drivers
                .get(&self.driver_num.get())
                .expect("DriverShareRef: registered but nonexistent?");
            if subscribe_num >= driver_data.num_upcalls {
                return Err(InvalidSubscribeNum {
                    upcall_count: driver_data.num_upcalls,
                    requested: subscribe_num,
                });
            }
            let upcall = match driver_data.upcalls.get(&subscribe_num) {
                Some(&upcall) => upcall,
                None => return Ok(()),
            };
            // Don't bother queueing a null upcall, as they don't do anything
            // when invoked anyway, and the core kernel does not queue them
            // either.
            if upcall.is_null() {
                return Ok(());
            }
            kernel_data.upcall_queue.push_back(UpcallQueueEntry {
                args,
                id: UpcallId {
                    driver_num: self.driver_num.get(),
                    subscribe_num,
                },
                upcall,
            });
            Ok(())
        })
    }
}

#[derive(Debug, Eq, PartialEq, thiserror::Error)]
#[error("Upcall number {requested} too large, expected < {upcall_count}.")]
pub struct InvalidSubscribeNum {
    requested: u32,
    upcall_count: u32,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::upcall::Upcall;
    use crate::DriverInfo;
    use std::rc::Rc;

    #[derive(Default)]
    struct MockDriver {
        share_ref: DriverShareRef,
    }
    impl crate::fake::SyscallDriver for MockDriver {
        fn info(&self) -> DriverInfo {
            DriverInfo::new(1).upcall_count(10)
        }
        fn register(&self, share_ref: DriverShareRef) {
            self.share_ref.replace(share_ref);
        }
        fn command(&self, _: u32, _: u32, _: u32) -> libtock_platform::CommandReturn {
            crate::command_return::failure(libtock_platform::ErrorCode::NoSupport)
        }
    }

    #[test]
    fn schedule_errors() {
        let mock_driver = Rc::new(MockDriver::default());
        let kernel = crate::fake::Kernel::new();
        kernel.add_driver(&mock_driver);
        assert_eq!(
            mock_driver.share_ref.schedule_upcall(10, (3, 4, 5)),
            Err(InvalidSubscribeNum {
                upcall_count: 10,
                requested: 10
            })
        );
    }

    #[test]
    fn schedule_success() {
        let mock_driver = Rc::new(MockDriver::default());
        let kernel = crate::fake::Kernel::new();
        kernel.add_driver(&mock_driver);

        // Call schedule with no registered upcall.
        assert_eq!(mock_driver.share_ref.schedule_upcall(2, (3, 4, 5)), Ok(()));
        crate::kernel_data::with_kernel_data(|kernel_data| {
            let kernel_data = kernel_data.unwrap();

            // There was no upcall to schedule, so the queue should still be
            // empty.
            assert!(kernel_data.upcall_queue.is_empty());

            // Register a null upcall.
            kernel_data.drivers.get_mut(&1).unwrap().upcalls.insert(
                2,
                Upcall {
                    fn_pointer: None,
                    data: 1234u32.into(),
                },
            );
        });
        // Call schedule again. This should still do nothing, because the upcall
        // is a null upcall.
        assert_eq!(mock_driver.share_ref.schedule_upcall(2, (3, 4, 5)), Ok(()));
        unsafe extern "C" fn upcall(_: u32, _: u32, _: u32, _: libtock_platform::Register) {}
        crate::kernel_data::with_kernel_data(|kernel_data| {
            let kernel_data = kernel_data.unwrap();

            // Verify the upcall was not queued.
            assert!(kernel_data.upcall_queue.is_empty());

            // Register a non-null upcall.
            kernel_data.drivers.get_mut(&1).unwrap().upcalls.insert(
                2,
                Upcall {
                    fn_pointer: Some(upcall),
                    data: 1111usize.into(),
                },
            );
        });
        // Call schedule again. This should schedule the upcall.
        assert_eq!(mock_driver.share_ref.schedule_upcall(2, (3, 4, 5)), Ok(()));
        crate::kernel_data::with_kernel_data(|kernel_data| {
            let kernel_data = kernel_data.unwrap();

            // Verify the upcall was queued.
            assert_eq!(kernel_data.upcall_queue.len(), 1);
            let upcall_queue_entry = kernel_data.upcall_queue.front().expect("Upcall not queued");
            assert_eq!(upcall_queue_entry.args, (3, 4, 5));
            assert_eq!(
                upcall_queue_entry.id,
                UpcallId {
                    driver_num: 1,
                    subscribe_num: 2
                }
            );
            assert!(upcall_queue_entry.upcall.fn_pointer == Some(upcall));
            let data: usize = upcall_queue_entry.upcall.data.into();
            assert_eq!(data, 1111);

            // Register a non-null upcall.
            kernel_data.drivers.get_mut(&1).unwrap().upcalls.insert(
                2,
                Upcall {
                    fn_pointer: Some(upcall),
                    data: 2222u32.into(),
                },
            );
        });
        // Call schedule again. This should schedule another upcall, after the
        // first.
        assert_eq!(
            mock_driver.share_ref.schedule_upcall(2, (30, 40, 50)),
            Ok(())
        );
        crate::kernel_data::with_kernel_data(|kernel_data| {
            let kernel_data = kernel_data.unwrap();

            // Very the upcall was queued.
            assert_eq!(kernel_data.upcall_queue.len(), 2);
            let front_queue_entry = kernel_data.upcall_queue.front().expect("Upcall not queued");
            assert_eq!(front_queue_entry.args, (3, 4, 5));
            assert_eq!(
                front_queue_entry.id,
                UpcallId {
                    driver_num: 1,
                    subscribe_num: 2
                }
            );
            assert!(front_queue_entry.upcall.fn_pointer == Some(upcall));
            let front_data: usize = front_queue_entry.upcall.data.into();
            assert_eq!(front_data, 1111);
            let back_queue_entry = kernel_data.upcall_queue.back().expect("Upcall not queued");
            assert_eq!(back_queue_entry.args, (30, 40, 50));
            assert_eq!(
                back_queue_entry.id,
                UpcallId {
                    driver_num: 1,
                    subscribe_num: 2
                }
            );
            assert!(back_queue_entry.upcall.fn_pointer == Some(upcall));
            let back_data: usize = back_queue_entry.upcall.data.into();
            assert_eq!(back_data, 2222);
        });
    }
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;
use std::iter::Fuse;
use std::fmt;

use super::adaptors::{PutBack, put_back};
use crate::either_or_both::EitherOrBoth;

/// Return an iterator adaptor that merge-joins items from the two base iterators in ascending order.
///
/// See [`.merge_join_by()`](crate::Itertools::merge_join_by) for more information.
pub fn merge_join_by<I, J, F>(left: I, right: J, cmp_fn: F)
    -> MergeJoinBy<I::IntoIter, J::IntoIter, F>
    where I: IntoIterator,
          J: IntoIterator,
          F: FnMut(&I::Item, &J::Item) -> Ordering
{
    MergeJoinBy {
        left: put_back(left.into_iter().fuse()),
        right: put_back(right.into_iter().fuse()),
        cmp_fn,
    }
}

/// An iterator adaptor that merge-joins items from the two base iterators in ascending order.
///
/// See [`.merge_join_by()`](crate::Itertools::merge_join_by) for more information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct MergeJoinBy<I: Iterator, J: Iterator, F> {
    left: PutBack<Fuse<I>>,
    right: PutBack<Fuse<J>>,
    cmp_fn: F
}

impl<I, J, F> Clone for MergeJoinBy<I, J, F>
    where I: Iterator,
          J: Iterator,
          PutBack<Fuse<I>>: Clone,
          PutBack<Fuse<J>>: Clone,
          F: Clone,
{
    clone_fields!(left, right, cmp_fn);
}

impl<I, J, F> fmt::Debug for MergeJoinBy<I, J, F>
    where I: Iterator + fmt::Debug,
          I::Item: fmt::Debug,
          J: Iterator + fmt::Debug,
          J::Item: fmt::Debug,
{
    debug_fmt_fields!(MergeJoinBy, left, right);
}

impl<I, J, F> Iterator for MergeJoinBy<I, J, F>
    where I: Iterator,
          J: Iterator,
          F: FnMut(&I::Item, &J::Item) -> Ordering
{
    type Item = EitherOrBoth<I::Item, J::Item>;

    fn next(&mut self) -> Option<Self::Item> {
        match (self.left.next(), self.right.next()) {
            (None, None) => None,
            (Some(left), None) =>
                Some(EitherOrBoth::Left(left)),
            (None, Some(right)) =>
                Some(EitherOrBoth::Right(right)),
            (Some(left), Some(right)) => {
                match (self.cmp_fn)(&left, &right) {
                    Ordering::Equal =>
                        Some(EitherOrBoth::Both(left, right)),
                    Ordering::Less => {
                        self.right.put_back(right);
                        Some(EitherOrBoth::Left(left))
                    },
                    Ordering::Greater => {
                        self.left.put_back(left);
                        Some(EitherOrBoth::Right(right))
                    }
                }
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let (a_lower, a_upper) = self.left.size_hint();
        let (b_lower, b_upper) = self.right.size_hint();

        let lower = ::std::cmp::max(a_lower, b_lower);

        let upper = match (a_upper, b_upper) {
            (Some(x), Some(y)) => x.checked_add(y),
            _ => None,
        };

        (lower, upper)
    }

    fn count(mut self) -> usize {
        let mut count = 0;
        loop {
            match (self.left.next(), self.right.next()) {
                (None, None) => break count,
                (Some(_left), None) => break count + 1 + self.left.into_parts().1.count(),
                (None, Some(_right)) => break count + 1 + self.right.into_parts().1.count(),
                (Some(left), Some(right)) => {
                    count += 1;
                    match (self.cmp_fn)(&left, &right) {
                        Ordering::Equal => {}
                        Ordering::Less => self.right.put_back(right),
                        Ordering::Greater => self.left.put_back(left),
                    }
                }
            }
        }
    }

    fn last(mut self) -> Option<Self::Item> {
        let mut previous_element = None;
        loop {
            match (self.left.next(), self.right.next()) {
                (None, None) => break previous_element,
                (Some(left), None) => {
                    break Some(EitherOrBoth::Left(
                        self.left.into_parts().1.last().unwrap_or(left),
                    ))
                }
                (None, Some(right)) => {
                    break Some(EitherOrBoth::Right(
                        self.right.into_parts().1.last().unwrap_or(right),
                    ))
                }
                (Some(left), Some(right)) => {
                    previous_element = match (self.cmp_fn)(&left, &right) {
                        Ordering::Equal => Some(EitherOrBoth::Both(left, right)),
                        Ordering::Less => {
                            self.right.put_back(right);
                            Some(EitherOrBoth::Left(left))
                        }
                        Ordering::Greater => {
                            self.left.put_back(left);
                            Some(EitherOrBoth::Right(right))
                        }
                    }
                }
            }
        }
    }

    fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
        loop {
            if n == 0 {
                break self.next();
            }
            n -= 1;
            match (self.left.next(), self.right.next()) {
                (None, None) => break None,
                (Some(_left), None) => break self.left.nth(n).map(EitherOrBoth::Left),
                (None, Some(_right)) => break self.right.nth(n).map(EitherOrBoth::Right),
                (Some(left), Some(right)) => match (self.cmp_fn)(&left, &right) {
                    Ordering::Equal => {}
                    Ordering::Less => self.right.put_back(right),
                    Ordering::Greater => self.left.put_back(left),
                },
            }
        }
    }
}
<commit_msg>Canonicalize some free functions (4) merge_join_by: documentation<commit_after>use std::cmp::Ordering;
use std::iter::Fuse;
use std::fmt;

use super::adaptors::{PutBack, put_back};
use crate::either_or_both::EitherOrBoth;
#[cfg(doc)]
use crate::Itertools;

/// Return an iterator adaptor that merge-joins items from the two base iterators in ascending order.
///
/// [`IntoIterator`] enabled version of [`Itertools::merge_join_by`].
pub fn merge_join_by<I, J, F>(left: I, right: J, cmp_fn: F)
    -> MergeJoinBy<I::IntoIter, J::IntoIter, F>
    where I: IntoIterator,
          J: IntoIterator,
          F: FnMut(&I::Item, &J::Item) -> Ordering
{
    MergeJoinBy {
        left: put_back(left.into_iter().fuse()),
        right: put_back(right.into_iter().fuse()),
        cmp_fn,
    }
}

/// An iterator adaptor that merge-joins items from the two base iterators in ascending order.
///
/// See [`.merge_join_by()`](crate::Itertools::merge_join_by) for more information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct MergeJoinBy<I: Iterator, J: Iterator, F> {
    left: PutBack<Fuse<I>>,
    right: PutBack<Fuse<J>>,
    cmp_fn: F
}

impl<I, J, F> Clone for MergeJoinBy<I, J, F>
    where I: Iterator,
          J: Iterator,
          PutBack<Fuse<I>>: Clone,
          PutBack<Fuse<J>>: Clone,
          F: Clone,
{
    clone_fields!(left, right, cmp_fn);
}

impl<I, J, F> fmt::Debug for MergeJoinBy<I, J, F>
    where I: Iterator + fmt::Debug,
          I::Item: fmt::Debug,
          J: Iterator + fmt::Debug,
          J::Item: fmt::Debug,
{
    debug_fmt_fields!(MergeJoinBy, left, right);
}

impl<I, J, F> Iterator for MergeJoinBy<I, J, F>
    where I: Iterator,
          J: Iterator,
          F: FnMut(&I::Item, &J::Item) -> Ordering
{
    type Item = EitherOrBoth<I::Item, J::Item>;

    fn next(&mut self) -> Option<Self::Item> {
        match (self.left.next(), self.right.next()) {
            (None, None) => None,
            (Some(left), None) =>
                Some(EitherOrBoth::Left(left)),
            (None, Some(right)) =>
                Some(EitherOrBoth::Right(right)),
            (Some(left), Some(right)) => {
                match (self.cmp_fn)(&left, &right) {
                    Ordering::Equal =>
                        Some(EitherOrBoth::Both(left, right)),
                    Ordering::Less => {
                        self.right.put_back(right);
                        Some(EitherOrBoth::Left(left))
                    },
                    Ordering::Greater => {
                        self.left.put_back(left);
                        Some(EitherOrBoth::Right(right))
                    }
                }
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let (a_lower, a_upper) = self.left.size_hint();
        let (b_lower, b_upper) = self.right.size_hint();

        let lower = ::std::cmp::max(a_lower, b_lower);

        let upper = match (a_upper, b_upper) {
            (Some(x), Some(y)) => x.checked_add(y),
            _ => None,
        };

        (lower, upper)
    }

    fn count(mut self) -> usize {
        let mut count = 0;
        loop {
            match (self.left.next(), self.right.next()) {
                (None, None) => break count,
                (Some(_left), None) => break count + 1 + self.left.into_parts().1.count(),
                (None, Some(_right)) => break count + 1 + self.right.into_parts().1.count(),
                (Some(left), Some(right)) => {
                    count += 1;
                    match (self.cmp_fn)(&left, &right) {
                        Ordering::Equal => {}
                        Ordering::Less => self.right.put_back(right),
                        Ordering::Greater => self.left.put_back(left),
                    }
                }
            }
        }
    }

    fn last(mut self) -> Option<Self::Item> {
        let mut previous_element = None;
        loop {
            match (self.left.next(), self.right.next()) {
                (None, None) => break previous_element,
                (Some(left), None) => {
                    break Some(EitherOrBoth::Left(
                        self.left.into_parts().1.last().unwrap_or(left),
                    ))
                }
                (None, Some(right)) => {
                    break Some(EitherOrBoth::Right(
                        self.right.into_parts().1.last().unwrap_or(right),
                    ))
                }
                (Some(left), Some(right)) => {
                    previous_element = match (self.cmp_fn)(&left, &right) {
                        Ordering::Equal => Some(EitherOrBoth::Both(left, right)),
                        Ordering::Less => {
                            self.right.put_back(right);
                            Some(EitherOrBoth::Left(left))
                        }
                        Ordering::Greater => {
                            self.left.put_back(left);
                            Some(EitherOrBoth::Right(right))
                        }
                    }
                }
            }
        }
    }

    fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
        loop {
            if n == 0 {
                break self.next();
            }
            n -= 1;
            match (self.left.next(), self.right.next()) {
                (None, None) => break None,
                (Some(_left), None) => break self.left.nth(n).map(EitherOrBoth::Left),
                (None, Some(_right)) => break self.right.nth(n).map(EitherOrBoth::Right),
                (Some(left), Some(right)) => match (self.cmp_fn)(&left, &right) {
                    Ordering::Equal => {}
                    Ordering::Less => self.right.put_back(right),
                    Ordering::Greater => self.left.put_back(left),
                },
            }
        }
    }
}
<|endoftext|>


<commit_before>// pest. The Elegant Parser
// Copyright (C) 2017  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::fmt;
use std::rc::Rc;

use super::pair::{self, Pair};
use super::queueable_token::QueueableToken;
use super::token_iterator::{self, TokenIterator};
use super::super::inputs::Input;
use super::super::RuleType;

/// A `struct` containing `Pairs`. It is created by
/// [`Pairs::flatten`](struct.Pairs.html#method.flatten).
pub struct FlatPairs<R, I: Input> {
    queue: Rc<Vec<QueueableToken<R>>>,
    input: Rc<I>,
    start: usize,
    end: usize
}

pub fn new<R: RuleType, I: Input>(
    queue: Rc<Vec<QueueableToken<R>>>,
    input: Rc<I>,
    start: usize,
    end: usize
) -> FlatPairs<R, I> {
    FlatPairs {
        queue,
        input,
        start,
        end
    }
}

impl<R: RuleType, I: Input> FlatPairs<R, I> {
    /// Converts the `FlatPairs` into a `TokenIterator`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::rc::Rc;
    /// # use pest;
    /// # use pest::inputs::StringInput;
    /// # #[allow(non_camel_case_types)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// enum Rule {
    ///     a
    /// }
    ///
    /// let input = Rc::new(StringInput::new("".to_owned()));
    /// let pairs = pest::state(input, |state, pos| {
    ///     // generating Token pair with Rule::a ...
    /// #     state.rule(Rule::a, pos, |_, p| Ok(p))
    /// }).unwrap();
    /// let tokens: Vec<_> = pairs.flatten().into_iter().collect();
    ///
    /// assert_eq!(tokens.len(), 2);
    /// ```
    #[inline]
    pub fn into_iter(self) -> TokenIterator<R, I> {
        token_iterator::new(
            self.queue,
            self.input,
            self.start,
            self.end
        )
    }

    fn next_start(&mut self) {
        self.start += 1;

        while !self.is_start() {
            if self.start >= self.end {
                break;
            }

            self.start += 1;
        }
    }

    fn is_start(&self) -> bool {
        match self.queue[self.start] {
            QueueableToken::Start { .. } => true,
            QueueableToken::End { .. } => false
        }
    }
}

impl<R: RuleType, I: Input> Iterator for FlatPairs<R, I> {
    type Item = Pair<R, I>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.start >= self.end {
            return None;
        }

        let pair = pair::new(
            self.queue.clone(),
            self.input.clone(),
            self.start
        );

        self.next_start();

        Some(pair)
    }
}

impl<R: RuleType, I: Input> fmt::Debug for FlatPairs<R, I> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "FlatPairs {{ pairs: {:?} }}", self.clone().collect::<Vec<_>>())
    }
}

impl<R: Clone, I: Input> Clone for FlatPairs<R, I> {
    fn clone(&self) -> FlatPairs<R, I> {
        FlatPairs {
            queue: self.queue.clone(),
            input: self.input.clone(),
            start: self.start,
            end: self.end
        }
    }
}
<commit_msg>Fixed FlatPairs index ouf of bounds. Fixes #144.<commit_after>// pest. The Elegant Parser
// Copyright (C) 2017  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::fmt;
use std::rc::Rc;

use super::pair::{self, Pair};
use super::queueable_token::QueueableToken;
use super::token_iterator::{self, TokenIterator};
use super::super::inputs::Input;
use super::super::RuleType;

/// A `struct` containing `Pairs`. It is created by
/// [`Pairs::flatten`](struct.Pairs.html#method.flatten).
pub struct FlatPairs<R, I: Input> {
    queue: Rc<Vec<QueueableToken<R>>>,
    input: Rc<I>,
    start: usize,
    end: usize
}

pub fn new<R: RuleType, I: Input>(
    queue: Rc<Vec<QueueableToken<R>>>,
    input: Rc<I>,
    start: usize,
    end: usize
) -> FlatPairs<R, I> {
    FlatPairs {
        queue,
        input,
        start,
        end
    }
}

impl<R: RuleType, I: Input> FlatPairs<R, I> {
    /// Converts the `FlatPairs` into a `TokenIterator`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::rc::Rc;
    /// # use pest;
    /// # use pest::inputs::StringInput;
    /// # #[allow(non_camel_case_types)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// enum Rule {
    ///     a
    /// }
    ///
    /// let input = Rc::new(StringInput::new("".to_owned()));
    /// let pairs = pest::state(input, |state, pos| {
    ///     // generating Token pair with Rule::a ...
    /// #     state.rule(Rule::a, pos, |_, p| Ok(p))
    /// }).unwrap();
    /// let tokens: Vec<_> = pairs.flatten().into_iter().collect();
    ///
    /// assert_eq!(tokens.len(), 2);
    /// ```
    #[inline]
    pub fn into_iter(self) -> TokenIterator<R, I> {
        token_iterator::new(
            self.queue,
            self.input,
            self.start,
            self.end
        )
    }

    fn next_start(&mut self) {
        self.start += 1;

        while self.start < self.end && !self.is_start() {
            self.start += 1;
        }
    }

    fn is_start(&self) -> bool {
        match self.queue[self.start] {
            QueueableToken::Start { .. } => true,
            QueueableToken::End { .. } => false
        }
    }
}

impl<R: RuleType, I: Input> Iterator for FlatPairs<R, I> {
    type Item = Pair<R, I>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.start >= self.end {
            return None;
        }

        let pair = pair::new(
            self.queue.clone(),
            self.input.clone(),
            self.start
        );

        self.next_start();

        Some(pair)
    }
}

impl<R: RuleType, I: Input> fmt::Debug for FlatPairs<R, I> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "FlatPairs {{ pairs: {:?} }}", self.clone().collect::<Vec<_>>())
    }
}

impl<R: Clone, I: Input> Clone for FlatPairs<R, I> {
    fn clone(&self) -> FlatPairs<R, I> {
        FlatPairs {
            queue: self.queue.clone(),
            input: self.input.clone(),
            start: self.start,
            end: self.end
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

use std::f64;
use std::iter;
use std::slice;

use num::traits::{cast, NumCast};
use itertools::Itertools;

use bio::stats::logprobs::{self, LogProb};

/// Implementation of a cumulative distribution function.
#[derive(Debug, Clone)]
pub struct CDF<T: PartialOrd> {
    inner: Vec<(T, LogProb)>
}

impl<T: PartialOrd> CDF<T> {
    /// Create CDF from given probability mass function (PMF). The PMF may contain duplicate values
    /// the probabilities of which are summed during generation of the CDF.
    pub fn from_pmf(mut entries: Vec<(T, LogProb)>) -> Self {
        entries.sort_by(|&(ref a, _), &(ref b, _)| a.partial_cmp(b).unwrap());
        let mut inner: Vec<(T, LogProb)> = Vec::new();
        for mut e in entries.into_iter() {
            let p = logprobs::add(inner.last().map_or(f64::NEG_INFINITY, |e| e.1), e.1);
            if !inner.is_empty() && inner.last().unwrap().0 == e.0 {
                inner.last_mut().unwrap().1 = p;
            }
            else {
                e.1 = p;
                inner.push(e);
            }
        }
        let mut cdf = CDF {
            inner: inner
        };

        if relative_eq!(cdf.total_prob(), 0.0) && cdf.total_prob() > 0.0 {
            cdf.inner.last_mut().unwrap().1 = 0.0;
        }

        cdf
    }

    /// Create CDF from iterator. This can be used to replace the values of a CDF.
    pub fn from_cdf<I: Iterator<Item = (T, LogProb)>>(entries: I) -> Self {
        CDF { inner: entries.collect_vec() }
    }

    pub fn reduce(self) -> Self {
        let mut inner = Vec::new();
        let mut last = f64::NEG_INFINITY;
        for e in self.inner.into_iter() {
            if last != e.1 {
                last = e.1;
                inner.push(e);
            }
        }
        CDF { inner: inner }
    }

    pub fn sample(mut self, n: usize) -> Self {
        if self.inner.len() <= n {
            self
        }
        else {
            let s = self.inner.len() / (n - 1);
            let last = self.inner.pop().unwrap();
            let mut inner = self.inner.into_iter().step(s).collect_vec();
            inner.push(last);
            CDF {
                inner: inner
            }
        }
    }

    pub fn iter(&self) -> slice::Iter<(T, f64)>{
        self.inner.iter()
    }

    pub fn iter_pmf<'a>(&'a self) -> CDFPMFIter<'a, T> {
        fn cdf_to_pmf<'a, G>(last_prob: &mut LogProb, e: &'a (G, LogProb)) -> Option<(&'a G, LogProb)> {
            let &(ref value, cdf_prob) = e;
            let prob = logprobs::sub(cdf_prob, *last_prob);
            *last_prob = cdf_prob;
            Some((value, prob))
        }
        self.inner.iter().scan(f64::NEG_INFINITY, cdf_to_pmf)
    }

    /// Get cumulative probability for a given value. If the value is not present,
    /// return the probability of the previous value.
    pub fn get(&self, value: &T) -> Option<LogProb> {
        if self.inner.is_empty() {
            None
        }
        else {
            Some(match self.inner.binary_search_by(|e| e.0.partial_cmp(value).unwrap()) {
                Ok(i) => self.inner[i].1,
                Err(i) => if i > 0 { self.inner[i - 1].1 } else { f64::NEG_INFINITY }
            })
        }
    }

    pub fn get_pmf(&self, value: &T) -> Option<LogProb> {
        if self.inner.is_empty() {
            None
        }
        else {
            Some(match self.inner.binary_search_by(|e| e.0.partial_cmp(value).unwrap()) {
                Ok(i) => if i > 0 { logprobs::sub(self.inner[i].1, self.inner[i - 1].1) } else { self.inner[0].1 },
                Err(i) => if i > 0 { self.inner[i - 1].1 } else { f64::NEG_INFINITY }
            })
        }
    }

    pub fn total_prob(&self) -> LogProb {
        let &(_, prob) = self.inner.last().unwrap();
        prob
    }

    /// Return maximum a posteriori probability estimate (MAP).
    pub fn map(&self) -> &T {
        let mut max = self.iter_pmf().next().unwrap();
        for e in self.iter_pmf() {
            if e.1 >= max.1 {
                max = e;
            }
        }
        &max.0
    }

    /// Return the 95% credible interval.
    pub fn credible_interval(&self) -> (&T, &T) {
        let lower = self.inner.binary_search_by(|&(_, p)| p.partial_cmp(&0.025f64.ln()).unwrap()).unwrap_or_else(|i| if i > 0 { i - 1 } else { 0 });
        let upper = self.inner.binary_search_by(|&(_, p)| p.partial_cmp(&0.975f64.ln()).unwrap()).unwrap_or_else(|i| i);

        (&self.inner[lower].0, &self.inner[upper].0)
    }

    pub fn len(&self) -> usize {
        self.inner.len()
    }
}

impl<T: NumCast + Clone + PartialOrd> CDF<T> {
    pub fn expected_value(&self) -> f64 {
        self.iter_pmf().map(|(value, prob)| {
            cast::<T, f64>(value.clone()).unwrap() * prob.exp()
        }).fold(0.0f64, |s, e| s + e)
    }

    pub fn variance(&self) -> f64 {
        let ev = self.expected_value();
        self.iter_pmf().map(|(value, prob)| {
                (cast::<T, f64>(value.clone()).unwrap() - ev).powi(2) * prob.exp()
        }).fold(0.0, |s, e| s + e)
    }

    pub fn standard_deviation(&self) -> f64 {
        self.variance().sqrt()
    }
}

pub type CDFPMFIter<'a, T> = iter::Scan<slice::Iter<'a, (T, LogProb)>, LogProb, fn(&mut LogProb, &'a (T, LogProb)) -> Option<(&'a T, LogProb)>>;

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

    #[test]
    fn test_cdf() {
        let mut pmf = vec![(0.0, 0.1f64.ln())];
        for i in 0..9 {
            pmf.push((i as f64, 0.1f64.ln()));
        }
        println!("{:?}", pmf);

        let cdf = CDF::from_pmf(pmf.clone());
        println!("{:?}", cdf);
        for &(value, prob) in pmf.iter().skip(2) {
            assert_ulps_eq!(prob, cdf.get_pmf(&value).unwrap(), epsilon = 0.0000000000001);
        }
        assert_relative_eq!(cdf.total_prob(), 1.0f64.ln());
        assert_relative_eq!(cdf.get(&1.0).unwrap(), 0.3f64.ln(), epsilon = 0.00000001);
    }
}
<commit_msg>Try to generate even tighter credible intervals.<commit_after>// Copyright 2016 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

use std::f64;
use std::iter;
use std::slice;

use num::traits::{cast, NumCast};
use itertools::Itertools;

use bio::stats::logprobs::{self, LogProb};

/// Implementation of a cumulative distribution function.
#[derive(Debug, Clone)]
pub struct CDF<T: PartialOrd> {
    inner: Vec<(T, LogProb)>
}

impl<T: PartialOrd> CDF<T> {
    /// Create CDF from given probability mass function (PMF). The PMF may contain duplicate values
    /// the probabilities of which are summed during generation of the CDF.
    pub fn from_pmf(mut entries: Vec<(T, LogProb)>) -> Self {
        entries.sort_by(|&(ref a, _), &(ref b, _)| a.partial_cmp(b).unwrap());
        let mut inner: Vec<(T, LogProb)> = Vec::new();
        for mut e in entries.into_iter() {
            let p = logprobs::add(inner.last().map_or(f64::NEG_INFINITY, |e| e.1), e.1);
            if !inner.is_empty() && inner.last().unwrap().0 == e.0 {
                inner.last_mut().unwrap().1 = p;
            }
            else {
                e.1 = p;
                inner.push(e);
            }
        }
        let mut cdf = CDF {
            inner: inner
        };

        if relative_eq!(cdf.total_prob(), 0.0) && cdf.total_prob() > 0.0 {
            cdf.inner.last_mut().unwrap().1 = 0.0;
        }

        cdf
    }

    /// Create CDF from iterator. This can be used to replace the values of a CDF.
    pub fn from_cdf<I: Iterator<Item = (T, LogProb)>>(entries: I) -> Self {
        CDF { inner: entries.collect_vec() }
    }

    pub fn reduce(self) -> Self {
        let mut inner = Vec::new();
        let mut last = f64::NEG_INFINITY;
        for e in self.inner.into_iter() {
            if last != e.1 {
                last = e.1;
                inner.push(e);
            }
        }
        CDF { inner: inner }
    }

    pub fn sample(mut self, n: usize) -> Self {
        if self.inner.len() <= n {
            self
        }
        else {
            let s = self.inner.len() / (n - 1);
            let last = self.inner.pop().unwrap();
            let mut inner = self.inner.into_iter().step(s).collect_vec();
            inner.push(last);
            CDF {
                inner: inner
            }
        }
    }

    pub fn iter(&self) -> slice::Iter<(T, f64)>{
        self.inner.iter()
    }

    pub fn iter_pmf<'a>(&'a self) -> CDFPMFIter<'a, T> {
        fn cdf_to_pmf<'a, G>(last_prob: &mut LogProb, e: &'a (G, LogProb)) -> Option<(&'a G, LogProb)> {
            let &(ref value, cdf_prob) = e;
            let prob = logprobs::sub(cdf_prob, *last_prob);
            *last_prob = cdf_prob;
            Some((value, prob))
        }
        self.inner.iter().scan(f64::NEG_INFINITY, cdf_to_pmf)
    }

    /// Get cumulative probability for a given value. If the value is not present,
    /// return the probability of the previous value.
    pub fn get(&self, value: &T) -> Option<LogProb> {
        if self.inner.is_empty() {
            None
        }
        else {
            Some(match self.inner.binary_search_by(|e| e.0.partial_cmp(value).unwrap()) {
                Ok(i) => self.inner[i].1,
                Err(i) => if i > 0 { self.inner[i - 1].1 } else { f64::NEG_INFINITY }
            })
        }
    }

    pub fn get_pmf(&self, value: &T) -> Option<LogProb> {
        if self.inner.is_empty() {
            None
        }
        else {
            Some(match self.inner.binary_search_by(|e| e.0.partial_cmp(value).unwrap()) {
                Ok(i) => if i > 0 { logprobs::sub(self.inner[i].1, self.inner[i - 1].1) } else { self.inner[0].1 },
                Err(i) => if i > 0 { self.inner[i - 1].1 } else { f64::NEG_INFINITY }
            })
        }
    }

    pub fn total_prob(&self) -> LogProb {
        let &(_, prob) = self.inner.last().unwrap();
        prob
    }

    /// Return maximum a posteriori probability estimate (MAP).
    pub fn map(&self) -> &T {
        let mut max = self.iter_pmf().next().unwrap();
        for e in self.iter_pmf() {
            if e.1 >= max.1 {
                max = e;
            }
        }
        &max.0
    }

    /// Return the 95% credible interval.
    pub fn credible_interval(&self) -> (&T, &T) {
        let lower = self.inner.binary_search_by(|&(_, p)| p.partial_cmp(&0.025f64.ln()).unwrap()).unwrap_or_else(|i| if i > 0 { i - 1 } else { 0 });
        let lower_p = self.inner[lower].1;
        let upper = self.inner.binary_search_by(|&(_, p)| (p - lower_p).partial_cmp(&0.95f64.ln()).unwrap()).unwrap_or_else(|i| i);

        (&self.inner[lower].0, &self.inner[upper].0)
    }

    pub fn len(&self) -> usize {
        self.inner.len()
    }
}

impl<T: NumCast + Clone + PartialOrd> CDF<T> {
    pub fn expected_value(&self) -> f64 {
        self.iter_pmf().map(|(value, prob)| {
            cast::<T, f64>(value.clone()).unwrap() * prob.exp()
        }).fold(0.0f64, |s, e| s + e)
    }

    pub fn variance(&self) -> f64 {
        let ev = self.expected_value();
        self.iter_pmf().map(|(value, prob)| {
                (cast::<T, f64>(value.clone()).unwrap() - ev).powi(2) * prob.exp()
        }).fold(0.0, |s, e| s + e)
    }

    pub fn standard_deviation(&self) -> f64 {
        self.variance().sqrt()
    }
}

pub type CDFPMFIter<'a, T> = iter::Scan<slice::Iter<'a, (T, LogProb)>, LogProb, fn(&mut LogProb, &'a (T, LogProb)) -> Option<(&'a T, LogProb)>>;

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

    #[test]
    fn test_cdf() {
        let mut pmf = vec![(0.0, 0.1f64.ln())];
        for i in 0..9 {
            pmf.push((i as f64, 0.1f64.ln()));
        }
        println!("{:?}", pmf);

        let cdf = CDF::from_pmf(pmf.clone());
        println!("{:?}", cdf);
        for &(value, prob) in pmf.iter().skip(2) {
            assert_ulps_eq!(prob, cdf.get_pmf(&value).unwrap(), epsilon = 0.0000000000001);
        }
        assert_relative_eq!(cdf.total_prob(), 1.0f64.ln());
        assert_relative_eq!(cdf.get(&1.0).unwrap(), 0.3f64.ln(), epsilon = 0.00000001);
    }
}
<|endoftext|>