<commit_before>extern crate futures;

use futures::{Async, Poll, Future, Stream};
use std::mem;

pub type BoxStateStream<T, S, E> = Box<StateStream<Item = T, State = S, Error = E> + Send>;

pub enum StreamEvent<I, S> {
    Next(I),
    Done(S),
}

pub trait StateStream {
    type Item;
    type State;
    type Error;

    fn poll(&mut self) -> Poll<StreamEvent<Self::Item, Self::State>, Self::Error>;

    #[inline]
    fn boxed(self) -> BoxStateStream<Self::Item, Self::State, Self::Error>
        where Self: Sized + Send + 'static
    {
        Box::new(self)
    }

    #[inline]
    fn into_future(self) -> IntoFuture<Self>
        where Self: Sized
    {
        IntoFuture(Some(self))
    }

    #[inline]
    fn into_stream(self) -> IntoStream<Self>
        where Self: Sized
    {
        IntoStream(self)
    }

    #[inline]
    fn map<F, B>(self, f: F) -> Map<Self, F>
        where Self: Sized,
              F: FnMut(Self::Item) -> B
    {
        Map {
            stream: self,
            f: f,
        }
    }

    #[inline]
    fn map_err<F, B>(self, f: F) -> MapErr<Self, F>
        where Self: Sized,
              F: FnMut(Self::Error) -> B
    {
        MapErr {
            stream: self,
            f: f,
        }
    }

    #[inline]
    fn map_state<F, B>(self, f: F) -> MapState<Self, F>
        where Self: Sized,
              F: FnOnce(Self::State) -> B
    {
        MapState {
            stream: self,
            f: Some(f),
        }
    }

    #[inline]
    fn collect(self) -> Collect<Self>
        where Self: Sized
    {
        Collect {
            stream: self,
            items: vec![],
        }
    }
}

impl<S: ?Sized> StateStream for Box<S>
    where S: StateStream
{
    type Item = S::Item;
    type State = S::State;
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<S::Item, S::State>, S::Error> {
        S::poll(self)
    }
}

#[inline]
pub fn stream<S>(stream: S) -> FromStream<S>
    where S: Stream
{
    FromStream(stream)
}

#[inline]
pub fn unfold<T, F, Fut, It, St>(init: T, f: F) -> Unfold<T, F, Fut>
     where F: FnMut(T) -> Fut,
           Fut: futures::IntoFuture<Item = StreamEvent<(It, T), St>>
{
    Unfold {
        state: UnfoldState::Ready(init),
        f: f,
    }
}

pub struct FromStream<S>(S);

impl<S> StateStream for FromStream<S>
    where S: Stream
{
    type Item = S::Item;
    type State = ();
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<S::Item, ()>, S::Error> {
        self.0.poll().map(|a| {
            match a {
                Async::Ready(Some(i)) => Async::Ready(StreamEvent::Next(i)),
                Async::Ready(None) => Async::Ready(StreamEvent::Done(())),
                Async::NotReady => Async::NotReady,
            }
        })
    }
}

pub struct IntoFuture<S>(Option<S>);

impl<S> Future for IntoFuture<S>
    where S: StateStream
{
    type Item = (StreamEvent<S::Item, S::State>, S);
    type Error = (S::Error, S);

    #[inline]
    fn poll(&mut self) -> Poll<(StreamEvent<S::Item, S::State>, S), (S::Error, S)> {
        let item = match self.0.as_mut().expect("polling IntoFuture twice").poll() {
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Ok(Async::Ready(i)) => Ok(i),
            Err(e) => Err(e),
        };
        let stream = self.0.take().unwrap();
        match item {
            Ok(i) => Ok(Async::Ready((i, stream))),
            Err(e) => Err((e, stream)),
        }
    }
}

pub struct IntoStream<S>(S);

impl<S> Stream for IntoStream<S>
    where S: StateStream
{
    type Item = S::Item;
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<Option<S::Item>, S::Error> {
        self.0.poll().map(|a| {
            match a {
                Async::Ready(StreamEvent::Next(i)) => Async::Ready(Some(i)),
                Async::Ready(StreamEvent::Done(_)) => Async::Ready(None),
                Async::NotReady => Async::NotReady,
            }
        })
    }
}

pub struct Map<S, F> {
    stream: S,
    f: F,
}

impl<S, F, B> StateStream for Map<S, F>
    where S: StateStream,
          F: FnMut(S::Item) -> B
{
    type Item = B;
    type State = S::State;
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<B, S::State>, S::Error> {
        self.stream.poll().map(|a| {
            match a {
                Async::Ready(StreamEvent::Next(i)) => Async::Ready(StreamEvent::Next((self.f)(i))),
                Async::Ready(StreamEvent::Done(s)) => Async::Ready(StreamEvent::Done(s)),
                Async::NotReady => Async::NotReady,
            }
        })
    }
}

pub struct MapErr<S, F> {
    stream: S,
    f: F,
}

impl<S, F, B> StateStream for MapErr<S, F>
    where S: StateStream,
          F: FnMut(S::Error) -> B
{
    type Item = S::Item;
    type State = S::State;
    type Error = B;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<S::Item, S::State>, B> {
        match self.stream.poll() {
            Ok(a) => Ok(a),
            Err(e) => Err((self.f)(e))
        }
    }
}

pub struct MapState<S, F> {
    stream: S,
    f: Option<F>,
}

impl<S, F, B> StateStream for MapState<S, F>
    where S: StateStream,
          F: FnOnce(S::State) -> B
{
    type Item = S::Item;
    type State = B;
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<S::Item, B>, S::Error> {
        self.stream.poll().map(|a| {
            match a {
                Async::Ready(StreamEvent::Next(i)) => Async::Ready(StreamEvent::Next(i)),
                Async::Ready(StreamEvent::Done(s)) => {
                    let f = self.f.take().expect("polled MapState after completion");
                    Async::Ready(StreamEvent::Done(f(s)))
                }
                Async::NotReady => Async::NotReady,
            }
        })
    }
}

pub struct Collect<S>
    where S: StateStream
{
    stream: S,
    items: Vec<S::Item>,
}

impl<S> Future for Collect<S>
    where S: StateStream
{
    type Item = (Vec<S::Item>, S::State);
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<(Vec<S::Item>, S::State), S::Error> {
        loop {
            match self.stream.poll() {
                Ok(Async::Ready(StreamEvent::Next(i))) => self.items.push(i),
                Ok(Async::Ready(StreamEvent::Done(s))) => {
                    let items = mem::replace(&mut self.items, vec![]);
                    return Ok(Async::Ready((items, s)))
                }
                Ok(Async::NotReady) => return Ok(Async::NotReady),
                Err(e) => return Err(e),
            }
        }
    }
}

pub struct Unfold<T, F, Fut>
    where Fut: futures::IntoFuture
{
    state: UnfoldState<T, Fut::Future>,
    f: F,
}

enum UnfoldState<T, F>
{
    Empty,
    Ready(T),
    Processing(F),
}

impl<T, F, Fut, It, St> StateStream for Unfold<T, F, Fut>
    where F: FnMut(T) -> Fut,
          Fut: futures::IntoFuture<Item = StreamEvent<(It, T), St>>
{
    type Item = It;
    type State = St;
    type Error = Fut::Error;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<It, St>, Fut::Error> {
        loop {
            match mem::replace(&mut self.state, UnfoldState::Empty) {
                UnfoldState::Empty => panic!("polled an Unfold after completion"),
                UnfoldState::Ready(state) => {
                    self.state = UnfoldState::Processing((self.f)(state).into_future())
                }
                UnfoldState::Processing(mut fut) => {
                    match try!(fut.poll()) {
                        Async::Ready(StreamEvent::Next((i, state))) => {
                            self.state = UnfoldState::Ready(state);
                            return Ok(Async::Ready(StreamEvent::Next(i)));
                        }
                        Async::Ready(StreamEvent::Done(s)) => {
                            return Ok(Async::Ready(StreamEvent::Done(s)));
                        }
                        Async::NotReady => {
                            self.state = UnfoldState::Processing(fut);
                            return Ok(Async::NotReady);
                        }
                    }
                }
            }
        }
    }
}
<commit_msg>Add flatten_state_stream to Future<commit_after>extern crate futures;

use futures::{Async, Poll, Future, Stream};
use std::mem;

pub type BoxStateStream<T, S, E> = Box<StateStream<Item = T, State = S, Error = E> + Send>;

pub enum StreamEvent<I, S> {
    Next(I),
    Done(S),
}

pub trait StateStream {
    type Item;
    type State;
    type Error;

    fn poll(&mut self) -> Poll<StreamEvent<Self::Item, Self::State>, Self::Error>;

    #[inline]
    fn boxed(self) -> BoxStateStream<Self::Item, Self::State, Self::Error>
        where Self: Sized + Send + 'static
    {
        Box::new(self)
    }

    #[inline]
    fn into_future(self) -> IntoFuture<Self>
        where Self: Sized
    {
        IntoFuture(Some(self))
    }

    #[inline]
    fn into_stream(self) -> IntoStream<Self>
        where Self: Sized
    {
        IntoStream(self)
    }

    #[inline]
    fn map<F, B>(self, f: F) -> Map<Self, F>
        where Self: Sized,
              F: FnMut(Self::Item) -> B
    {
        Map {
            stream: self,
            f: f,
        }
    }

    #[inline]
    fn map_err<F, B>(self, f: F) -> MapErr<Self, F>
        where Self: Sized,
              F: FnMut(Self::Error) -> B
    {
        MapErr {
            stream: self,
            f: f,
        }
    }

    #[inline]
    fn map_state<F, B>(self, f: F) -> MapState<Self, F>
        where Self: Sized,
              F: FnOnce(Self::State) -> B
    {
        MapState {
            stream: self,
            f: Some(f),
        }
    }

    #[inline]
    fn collect(self) -> Collect<Self>
        where Self: Sized
    {
        Collect {
            stream: self,
            items: vec![],
        }
    }
}

impl<S: ?Sized> StateStream for Box<S>
    where S: StateStream
{
    type Item = S::Item;
    type State = S::State;
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<S::Item, S::State>, S::Error> {
        S::poll(self)
    }
}

pub trait FutureExt: Future {
    #[inline]
    fn flatten_state_stream(self) -> FlattenStateStream<Self>
        where Self: Sized,
              Self::Item: StateStream<Error = Self::Error>
    {
        FlattenStateStream(FlattenStateStreamState::Future(self))
    }
}

impl<F: ?Sized> FutureExt for F
    where F: Future
{}

#[inline]
pub fn stream<S>(stream: S) -> FromStream<S>
    where S: Stream
{
    FromStream(stream)
}

#[inline]
pub fn unfold<T, F, Fut, It, St>(init: T, f: F) -> Unfold<T, F, Fut>
     where F: FnMut(T) -> Fut,
           Fut: futures::IntoFuture<Item = StreamEvent<(It, T), St>>
{
    Unfold {
        state: UnfoldState::Ready(init),
        f: f,
    }
}

pub struct FromStream<S>(S);

impl<S> StateStream for FromStream<S>
    where S: Stream
{
    type Item = S::Item;
    type State = ();
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<S::Item, ()>, S::Error> {
        self.0.poll().map(|a| {
            match a {
                Async::Ready(Some(i)) => Async::Ready(StreamEvent::Next(i)),
                Async::Ready(None) => Async::Ready(StreamEvent::Done(())),
                Async::NotReady => Async::NotReady,
            }
        })
    }
}

pub struct IntoFuture<S>(Option<S>);

impl<S> Future for IntoFuture<S>
    where S: StateStream
{
    type Item = (StreamEvent<S::Item, S::State>, S);
    type Error = (S::Error, S);

    #[inline]
    fn poll(&mut self) -> Poll<(StreamEvent<S::Item, S::State>, S), (S::Error, S)> {
        let item = match self.0.as_mut().expect("polling IntoFuture twice").poll() {
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Ok(Async::Ready(i)) => Ok(i),
            Err(e) => Err(e),
        };
        let stream = self.0.take().unwrap();
        match item {
            Ok(i) => Ok(Async::Ready((i, stream))),
            Err(e) => Err((e, stream)),
        }
    }
}

pub struct IntoStream<S>(S);

impl<S> Stream for IntoStream<S>
    where S: StateStream
{
    type Item = S::Item;
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<Option<S::Item>, S::Error> {
        self.0.poll().map(|a| {
            match a {
                Async::Ready(StreamEvent::Next(i)) => Async::Ready(Some(i)),
                Async::Ready(StreamEvent::Done(_)) => Async::Ready(None),
                Async::NotReady => Async::NotReady,
            }
        })
    }
}

pub struct Map<S, F> {
    stream: S,
    f: F,
}

impl<S, F, B> StateStream for Map<S, F>
    where S: StateStream,
          F: FnMut(S::Item) -> B
{
    type Item = B;
    type State = S::State;
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<B, S::State>, S::Error> {
        self.stream.poll().map(|a| {
            match a {
                Async::Ready(StreamEvent::Next(i)) => Async::Ready(StreamEvent::Next((self.f)(i))),
                Async::Ready(StreamEvent::Done(s)) => Async::Ready(StreamEvent::Done(s)),
                Async::NotReady => Async::NotReady,
            }
        })
    }
}

pub struct MapErr<S, F> {
    stream: S,
    f: F,
}

impl<S, F, B> StateStream for MapErr<S, F>
    where S: StateStream,
          F: FnMut(S::Error) -> B
{
    type Item = S::Item;
    type State = S::State;
    type Error = B;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<S::Item, S::State>, B> {
        match self.stream.poll() {
            Ok(a) => Ok(a),
            Err(e) => Err((self.f)(e))
        }
    }
}

pub struct MapState<S, F> {
    stream: S,
    f: Option<F>,
}

impl<S, F, B> StateStream for MapState<S, F>
    where S: StateStream,
          F: FnOnce(S::State) -> B
{
    type Item = S::Item;
    type State = B;
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<S::Item, B>, S::Error> {
        self.stream.poll().map(|a| {
            match a {
                Async::Ready(StreamEvent::Next(i)) => Async::Ready(StreamEvent::Next(i)),
                Async::Ready(StreamEvent::Done(s)) => {
                    let f = self.f.take().expect("polled MapState after completion");
                    Async::Ready(StreamEvent::Done(f(s)))
                }
                Async::NotReady => Async::NotReady,
            }
        })
    }
}

pub struct Collect<S>
    where S: StateStream
{
    stream: S,
    items: Vec<S::Item>,
}

impl<S> Future for Collect<S>
    where S: StateStream
{
    type Item = (Vec<S::Item>, S::State);
    type Error = S::Error;

    #[inline]
    fn poll(&mut self) -> Poll<(Vec<S::Item>, S::State), S::Error> {
        loop {
            match self.stream.poll() {
                Ok(Async::Ready(StreamEvent::Next(i))) => self.items.push(i),
                Ok(Async::Ready(StreamEvent::Done(s))) => {
                    let items = mem::replace(&mut self.items, vec![]);
                    return Ok(Async::Ready((items, s)))
                }
                Ok(Async::NotReady) => return Ok(Async::NotReady),
                Err(e) => return Err(e),
            }
        }
    }
}

pub struct Unfold<T, F, Fut>
    where Fut: futures::IntoFuture
{
    state: UnfoldState<T, Fut::Future>,
    f: F,
}

enum UnfoldState<T, F>
{
    Empty,
    Ready(T),
    Processing(F),
}

impl<T, F, Fut, It, St> StateStream for Unfold<T, F, Fut>
    where F: FnMut(T) -> Fut,
          Fut: futures::IntoFuture<Item = StreamEvent<(It, T), St>>
{
    type Item = It;
    type State = St;
    type Error = Fut::Error;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<It, St>, Fut::Error> {
        loop {
            match mem::replace(&mut self.state, UnfoldState::Empty) {
                UnfoldState::Empty => panic!("polled an Unfold after completion"),
                UnfoldState::Ready(state) => {
                    self.state = UnfoldState::Processing((self.f)(state).into_future())
                }
                UnfoldState::Processing(mut fut) => {
                    match try!(fut.poll()) {
                        Async::Ready(StreamEvent::Next((i, state))) => {
                            self.state = UnfoldState::Ready(state);
                            return Ok(Async::Ready(StreamEvent::Next(i)));
                        }
                        Async::Ready(StreamEvent::Done(s)) => {
                            return Ok(Async::Ready(StreamEvent::Done(s)));
                        }
                        Async::NotReady => {
                            self.state = UnfoldState::Processing(fut);
                            return Ok(Async::NotReady);
                        }
                    }
                }
            }
        }
    }
}

enum FlattenStateStreamState<F>
    where F: Future
{
    Future(F),
    Stream(F::Item),
}

pub struct FlattenStateStream<F>(FlattenStateStreamState<F>)
    where F: Future;

impl<F> StateStream for FlattenStateStream<F>
    where F: Future,
          F::Item: StateStream<Error = F::Error>
{
    type Item = <F::Item as StateStream>::Item;
    type State = <F::Item as StateStream>::State;
    type Error = F::Error;

    #[inline]
    fn poll(&mut self) -> Poll<StreamEvent<Self::Item, Self::State>, Self::Error> {
        loop {
            self.0 = match self.0 {
                FlattenStateStreamState::Future(ref mut f) => {
                    match f.poll() {
                        Ok(Async::NotReady) => return Ok(Async::NotReady),
                        Ok(Async::Ready(stream)) => FlattenStateStreamState::Stream(stream),
                        Err(e) => return Err(e),
                    }
                }
                FlattenStateStreamState::Stream(ref mut s) => return s.poll(),
            };
        }
    }
}<|endoftext|>


<commit_before>// Claxon -- A FLAC decoding library in Rust
// Copyright 2014 Ruud van Asseldonk
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// A copy of the License has been included in the root of the repository.

//! Claxon, a FLAC decoding library.
//!
//! Examples
//! ========
//!
//! The following example computes the root mean square (RMS) of an audio file
//! with at most 16 bits per sample.
//!
//! ```
//! use claxon;
//!
//! let mut reader = claxon::FlacReader::open("testsamples/pop.flac").unwrap();
//! let mut sqr_sum = 0.0;
//! let mut count = 0;
//! for sample in reader.samples() {
//!     let s = sample.unwrap() as f64;
//!     sqr_sum += s * s;
//!     count += 1;
//! }
//! println!("RMS is {}", (sqr_sum / count as f64).sqrt());
//! ```
//!
//! TODO: more examples.

#![warn(missing_docs)]

use std::fs;
use std::io;
use std::mem;
use std::path;
use error::fmt_err;
use frame::FrameReader;
use input::{BufferedReader, ReadBytes};
use metadata::{MetadataBlock, MetadataBlockReader, StreamInfo};

mod crc;
mod input;
mod error;
pub mod frame;
pub mod subframe;
pub mod metadata;

pub use error::{Error, Result};
pub use frame::Block;

/// A FLAC decoder that can decode the stream from the underlying reader.
///
/// TODO: Add an example.
pub struct FlacReader<R: io::Read> {
    streaminfo: StreamInfo,
    #[allow(dead_code)] // TODO: Expose metadata nicely.
    metadata_blocks: Vec<MetadataBlock>,
    input: BufferedReader<R>,
}

/// An iterator that yields samples read from a `FlacReader`.
pub struct FlacSamples<R: ReadBytes> {
    frame_reader: FrameReader<R>,
    block: Block,
    sample: u32,
    channel: u32,

    /// If reading ever failed, this flag is set, so that the iterator knows not
    /// to return any new values.
    has_failed: bool,
}

// TODO: Add a `FlacIntoSamples`.

fn read_stream_header<R: ReadBytes>(input: &mut R) -> Result<()> {
    // A FLAC stream starts with a 32-bit header 'fLaC' (big endian).
    const HEADER: u32 = 0x66_4c_61_43;
    let header = try!(input.read_be_u32());
    if header != HEADER {
        fmt_err("invalid stream header")
    } else {
        Ok(())
    }
}

impl<R: io::Read> FlacReader<R> {
    /// Attempts to create a reader that reads the FLAC format.
    ///
    /// The header and metadata blocks are read immediately. Audio frames will
    /// be read on demand.
    pub fn new(reader: R) -> Result<FlacReader<R>> {
        let mut buf_reader = BufferedReader::new(reader);

        // A flac stream first of all starts with a stream header.
        try!(read_stream_header(&mut buf_reader));

        // Start a new scope, because the input reader must be available again
        // for the frame reader next.
        let (streaminfo, metadata_blocks) = {
            // Next are one or more metadata blocks. The flac specification
            // dictates that the streaminfo block is the first block. The metadata
            // block reader will yield at least one element, so the unwrap is safe.
            let mut metadata_iter = MetadataBlockReader::new(&mut buf_reader);
            let streaminfo_block = try!(metadata_iter.next().unwrap());
            let streaminfo = match streaminfo_block {
                MetadataBlock::StreamInfo(info) => info,
                _ => return fmt_err("streaminfo block missing"),
            };

            // There might be more metadata blocks, read and store them.
            let mut metadata_blocks = Vec::new();
            for block_result in metadata_iter {
                match block_result {
                    Err(error) => return Err(error),
                    Ok(block) => metadata_blocks.push(block),
                }
            }

            (streaminfo, metadata_blocks)
        };

        // The flac reader will contain the reader that will read frames.
        let flac_reader = FlacReader {
            streaminfo: streaminfo,
            metadata_blocks: metadata_blocks,
            input: buf_reader,
        };

        Ok(flac_reader)
    }

    /// Returns the streaminfo metadata.
    ///
    /// This contains information like the sample rate and number of channels.
    pub fn streaminfo(&self) -> StreamInfo {
        self.streaminfo
    }

    /// Returns an iterator that decodes a single frame on every iteration.
    /// TODO: It is not an iterator.
    ///
    /// This is a low-level primitive that gives you control over when decoding
    /// happens. The representation of the decoded audio is somewhat specific to
    /// the FLAC format. For a higher-level interface, see `samples()`.
    pub fn blocks<'r>(&'r mut self) -> FrameReader<&'r mut BufferedReader<R>> {
        FrameReader::new(&mut self.input)
    }

    /// Returns an iterator over all samples.
    ///
    /// This is a user-friendly interface that trades performance for ease of
    /// use. If performance is an issue, consider using `blocks()` instead.
    ///
    /// This is a high-level interface to the decoder. The cost of retrieving
    /// the next sample can vary significantly, as sometimes a new block has to
    /// be decoded. Additionally, there is a cost to every iteration returning a
    /// `Result`. When a block has been decoded, iterating the samples in that
    /// block can never fail, but a match on every sample is required
    /// nonetheless. For more control over when decoding happens, and less error
    /// handling overhead, use `blocks()`.
    ///
    /// The channel data is is interleaved. The iterator is streaming. That is,
    /// if you call this method once, read a few samples, and call this method
    /// again, the second iterator will not start again from the beginning of
    /// the file. It will continue somewhere after where the first iterator
    /// stopped, and it might skip some samples. (This is because FLAC divides
    /// a stream into blocks, which have to be decoded entirely. If you drop the
    /// iterator, you lose the unread samples in that block.)
    pub fn samples<'r>(&'r mut self) -> FlacSamples<&'r mut BufferedReader<R>> {
        FlacSamples {
            frame_reader: frame::FrameReader::new(&mut self.input),
            block: Block::empty(),
            sample: 0,
            channel: 0,
            has_failed: false,
        }
    }

    /// Destroys the FLAC reader and returns the underlying reader.
    ///
    /// Because the reader employs buffering internally, anything in the buffer
    /// will be lost.
    pub fn into_inner(self) -> R {
        self.input.into_inner()
    }
}

impl FlacReader<fs::File> {
    /// Attempts to create a reader that reads from the specified file.
    ///
    /// This is a convenience constructor that opens a `File`, and constructs a
    /// `FlacReader` from it. There is no need to wrap the file in a
    /// `BufReader`, as the `FlacReader` employs buffering already.
    pub fn open<P: AsRef<path::Path>>(filename: P) -> Result<FlacReader<fs::File>> {
        let file = try!(fs::File::open(filename));
        FlacReader::new(file)
    }
}

impl<R: ReadBytes> Iterator for FlacSamples<R> {
    type Item = Result<i32>;

    fn next(&mut self) -> Option<Result<i32>> {
        // If the previous read failed, end iteration.
        if self.has_failed {
            return None;
        }

        // Iterate the samples channel interleaved, so first increment the
        // channel.
        self.channel += 1;

        // If that was the last channel, increment the sample number.
        if self.channel >= self.block.channels() {
            self.channel = 0;
            self.sample += 1;

            // If that was the last sample in the block, decode the next block.
            if self.sample >= self.block.len() {
                self.sample = 0;

                // Replace the current block with an empty one so that we may
                // reuse the current buffer to decode again.
                let current_block = mem::replace(&mut self.block, Block::empty());

                match self.frame_reader.read_next_or_eof(current_block.into_buffer()) {
                    Ok(Some(next_block)) => {
                        self.block = next_block;
                    }
                    Ok(None) => {
                        // The stream ended with EOF.
                        // TODO: If a number of samples was specified in the
                        // streaminfo metadata block, verify that we did not
                        // read more or less samples.
                        return None;
                    }
                    Err(error) => {
                        self.has_failed = true;
                        return Some(Err(error));
                    }
                }
            }
        }

        Some(Ok(self.block.sample(self.channel, self.sample)))
    }
}
<commit_msg>Fix bug in FlacReader::samples() iterator<commit_after>// Claxon -- A FLAC decoding library in Rust
// Copyright 2014 Ruud van Asseldonk
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// A copy of the License has been included in the root of the repository.

//! Claxon, a FLAC decoding library.
//!
//! Examples
//! ========
//!
//! The following example computes the root mean square (RMS) of an audio file
//! with at most 16 bits per sample.
//!
//! ```
//! use claxon;
//!
//! let mut reader = claxon::FlacReader::open("testsamples/pop.flac").unwrap();
//! let mut sqr_sum = 0.0;
//! let mut count = 0;
//! for sample in reader.samples() {
//!     let s = sample.unwrap() as f64;
//!     sqr_sum += s * s;
//!     count += 1;
//! }
//! println!("RMS is {}", (sqr_sum / count as f64).sqrt());
//! ```
//!
//! TODO: more examples.

#![warn(missing_docs)]

use std::fs;
use std::io;
use std::mem;
use std::path;
use error::fmt_err;
use frame::FrameReader;
use input::{BufferedReader, ReadBytes};
use metadata::{MetadataBlock, MetadataBlockReader, StreamInfo};

mod crc;
mod input;
mod error;
pub mod frame;
pub mod subframe;
pub mod metadata;

pub use error::{Error, Result};
pub use frame::Block;

/// A FLAC decoder that can decode the stream from the underlying reader.
///
/// TODO: Add an example.
pub struct FlacReader<R: io::Read> {
    streaminfo: StreamInfo,
    #[allow(dead_code)] // TODO: Expose metadata nicely.
    metadata_blocks: Vec<MetadataBlock>,
    input: BufferedReader<R>,
}

/// An iterator that yields samples read from a `FlacReader`.
pub struct FlacSamples<R: ReadBytes> {
    frame_reader: FrameReader<R>,
    block: Block,
    sample: u32,
    channel: u32,

    /// If reading ever failed, this flag is set, so that the iterator knows not
    /// to return any new values.
    has_failed: bool,
}

// TODO: Add a `FlacIntoSamples`.

fn read_stream_header<R: ReadBytes>(input: &mut R) -> Result<()> {
    // A FLAC stream starts with a 32-bit header 'fLaC' (big endian).
    const HEADER: u32 = 0x66_4c_61_43;
    let header = try!(input.read_be_u32());
    if header != HEADER {
        fmt_err("invalid stream header")
    } else {
        Ok(())
    }
}

impl<R: io::Read> FlacReader<R> {
    /// Attempts to create a reader that reads the FLAC format.
    ///
    /// The header and metadata blocks are read immediately. Audio frames will
    /// be read on demand.
    pub fn new(reader: R) -> Result<FlacReader<R>> {
        let mut buf_reader = BufferedReader::new(reader);

        // A flac stream first of all starts with a stream header.
        try!(read_stream_header(&mut buf_reader));

        // Start a new scope, because the input reader must be available again
        // for the frame reader next.
        let (streaminfo, metadata_blocks) = {
            // Next are one or more metadata blocks. The flac specification
            // dictates that the streaminfo block is the first block. The metadata
            // block reader will yield at least one element, so the unwrap is safe.
            let mut metadata_iter = MetadataBlockReader::new(&mut buf_reader);
            let streaminfo_block = try!(metadata_iter.next().unwrap());
            let streaminfo = match streaminfo_block {
                MetadataBlock::StreamInfo(info) => info,
                _ => return fmt_err("streaminfo block missing"),
            };

            // There might be more metadata blocks, read and store them.
            let mut metadata_blocks = Vec::new();
            for block_result in metadata_iter {
                match block_result {
                    Err(error) => return Err(error),
                    Ok(block) => metadata_blocks.push(block),
                }
            }

            (streaminfo, metadata_blocks)
        };

        // The flac reader will contain the reader that will read frames.
        let flac_reader = FlacReader {
            streaminfo: streaminfo,
            metadata_blocks: metadata_blocks,
            input: buf_reader,
        };

        Ok(flac_reader)
    }

    /// Returns the streaminfo metadata.
    ///
    /// This contains information like the sample rate and number of channels.
    pub fn streaminfo(&self) -> StreamInfo {
        self.streaminfo
    }

    /// Returns an iterator that decodes a single frame on every iteration.
    /// TODO: It is not an iterator.
    ///
    /// This is a low-level primitive that gives you control over when decoding
    /// happens. The representation of the decoded audio is somewhat specific to
    /// the FLAC format. For a higher-level interface, see `samples()`.
    pub fn blocks<'r>(&'r mut self) -> FrameReader<&'r mut BufferedReader<R>> {
        FrameReader::new(&mut self.input)
    }

    /// Returns an iterator over all samples.
    ///
    /// This is a user-friendly interface that trades performance for ease of
    /// use. If performance is an issue, consider using `blocks()` instead.
    ///
    /// This is a high-level interface to the decoder. The cost of retrieving
    /// the next sample can vary significantly, as sometimes a new block has to
    /// be decoded. Additionally, there is a cost to every iteration returning a
    /// `Result`. When a block has been decoded, iterating the samples in that
    /// block can never fail, but a match on every sample is required
    /// nonetheless. For more control over when decoding happens, and less error
    /// handling overhead, use `blocks()`.
    ///
    /// The channel data is is interleaved. The iterator is streaming. That is,
    /// if you call this method once, read a few samples, and call this method
    /// again, the second iterator will not start again from the beginning of
    /// the file. It will continue somewhere after where the first iterator
    /// stopped, and it might skip some samples. (This is because FLAC divides
    /// a stream into blocks, which have to be decoded entirely. If you drop the
    /// iterator, you lose the unread samples in that block.)
    pub fn samples<'r>(&'r mut self) -> FlacSamples<&'r mut BufferedReader<R>> {
        FlacSamples {
            frame_reader: frame::FrameReader::new(&mut self.input),
            block: Block::empty(),
            sample: 0,
            channel: 0,
            has_failed: false,
        }
    }

    /// Destroys the FLAC reader and returns the underlying reader.
    ///
    /// Because the reader employs buffering internally, anything in the buffer
    /// will be lost.
    pub fn into_inner(self) -> R {
        self.input.into_inner()
    }
}

impl FlacReader<fs::File> {
    /// Attempts to create a reader that reads from the specified file.
    ///
    /// This is a convenience constructor that opens a `File`, and constructs a
    /// `FlacReader` from it. There is no need to wrap the file in a
    /// `BufReader`, as the `FlacReader` employs buffering already.
    pub fn open<P: AsRef<path::Path>>(filename: P) -> Result<FlacReader<fs::File>> {
        let file = try!(fs::File::open(filename));
        FlacReader::new(file)
    }
}

impl<R: ReadBytes> Iterator for FlacSamples<R> {
    type Item = Result<i32>;

    fn next(&mut self) -> Option<Result<i32>> {
        // If the previous read failed, end iteration.
        if self.has_failed {
            return None;
        }

        // Iterate the samples channel interleaved, so first increment the
        // channel.
        self.channel += 1;

        // If that was the last channel, increment the sample number.
        if self.channel >= self.block.channels() {
            self.channel = 0;
            self.sample += 1;

            // If that was the last sample in the block, decode the next block.
            if self.sample >= self.block.duration() {
                self.sample = 0;

                // Replace the current block with an empty one so that we may
                // reuse the current buffer to decode again.
                let current_block = mem::replace(&mut self.block, Block::empty());

                match self.frame_reader.read_next_or_eof(current_block.into_buffer()) {
                    Ok(Some(next_block)) => {
                        self.block = next_block;
                    }
                    Ok(None) => {
                        // The stream ended with EOF.
                        // TODO: If a number of samples was specified in the
                        // streaminfo metadata block, verify that we did not
                        // read more or less samples.
                        return None;
                    }
                    Err(error) => {
                        self.has_failed = true;
                        return Some(Err(error));
                    }
                }
            }
        }

        Some(Ok(self.block.sample(self.channel, self.sample)))
    }
}
<|endoftext|>


<commit_before>#![feature(core)]
#![cfg_attr(test, feature(test))]

mod number_streams;

use number_streams::NumberStreams;

const CHUNK_SIZE: usize = 32;

const PRIME_1: u64 = 11400714785074694791;
const PRIME_2: u64 = 14029467366897019727;
const PRIME_3: u64 = 1609587929392839161;
const PRIME_4: u64 = 9650029242287828579;
const PRIME_5: u64 = 2870177450012600261;

#[derive(Copy,Clone,PartialEq)]
struct XxCore {
    v1: u64,
    v2: u64,
    v3: u64,
    v4: u64,
}

#[derive(Debug,Copy,Clone)]
pub struct XxHash {
    total_len: u64,
    seed: u64,
    core: XxCore,
    buffer: [u8; CHUNK_SIZE],
    buffer_usage: usize,
}

impl XxCore {
    fn from_seed(seed: u64) -> XxCore {
        XxCore {
            v1: seed.wrapping_add(PRIME_1).wrapping_add(PRIME_2),
            v2: seed.wrapping_add(PRIME_2),
            v3: seed,
            v4: seed.wrapping_sub(PRIME_1),
        }
    }

    #[inline(always)]
    fn ingest_chunks<I>(&mut self, values: I)
        where I: Iterator<Item=u64>
    {
        #[inline(always)]
        fn ingest_one_number(mut current_value: u64, mut value: u64) -> u64 {
            value = value.wrapping_mul(PRIME_2);
            current_value = current_value.wrapping_add(value);
            current_value = current_value.rotate_left(31);
            current_value.wrapping_mul(PRIME_1)
        };

        // By drawing these out, we can avoid going back and forth to
        // memory. It only really helps for large files, when we need
        // to iterate multiple times here.

        let mut v1 = self.v1;
        let mut v2 = self.v2;
        let mut v3 = self.v3;
        let mut v4 = self.v4;

        let mut values = values.peekable();

        while values.peek().is_some() {
            v1 = ingest_one_number(v1, values.next().unwrap());
            v2 = ingest_one_number(v2, values.next().unwrap());
            v3 = ingest_one_number(v3, values.next().unwrap());
            v4 = ingest_one_number(v4, values.next().unwrap());
        }

        self.v1 = v1;
        self.v2 = v2;
        self.v3 = v3;
        self.v4 = v4;
    }

    #[inline(always)]
    fn finish(&self) -> u64 {
        // TODO: The original code pulls out local vars for
        // v[1234], presumably for performance
        // reasons. Investigate.

        let mut hash;

        hash =                   self.v1.rotate_left( 1);
        hash = hash.wrapping_add(self.v2.rotate_left( 7));
        hash = hash.wrapping_add(self.v3.rotate_left(12));
        hash = hash.wrapping_add(self.v4.rotate_left(18));

        #[inline(always)]
        fn mix_one(mut hash: u64, mut value: u64) -> u64 {
            value = value.wrapping_mul(PRIME_2);
            value = value.rotate_left(31);
            value = value.wrapping_mul(PRIME_1);
            hash ^= value;
            hash = hash.wrapping_mul(PRIME_1);
            hash.wrapping_add(PRIME_4)
        }

        hash = mix_one(hash, self.v1);
        hash = mix_one(hash, self.v2);
        hash = mix_one(hash, self.v3);
        hash = mix_one(hash, self.v4);

        hash
    }
}

impl std::fmt::Debug for XxCore {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        write!(
            f, "XxCore {{ {:016x} {:016x} {:016x} {:016x} }}",
            self.v1, self.v2, self.v3, self.v4
        )
    }
}

impl XxHash {
    pub fn from_seed(seed: u64) -> XxHash {
        XxHash {
            total_len: 0,
            seed: seed,
            core: XxCore::from_seed(seed),
            buffer: [0; CHUNK_SIZE],
            buffer_usage: 0,
        }
    }
}

impl XxHash {
    pub fn write(&mut self, bytes: &[u8]) {
        let mut bytes = bytes;

        self.total_len += bytes.len() as u64;

        // Even with new data, we still don't have a full buffer. Wait
        // until we have a full buffer.
        if self.buffer_usage + bytes.len() < self.buffer.len() {
            unsafe {
                let tail = self.buffer.as_mut_ptr().offset(self.buffer_usage as isize);
                std::ptr::copy_nonoverlapping(bytes.as_ptr(), tail, bytes.len());
            }
            self.buffer_usage += bytes.len();
            return;
        }

        // Some data left from previous update. Fill the buffer and
        // consume it first.
        if self.buffer_usage > 0 {
            let bytes_to_use = self.buffer.len() - self.buffer_usage;
            let (to_use, leftover) = bytes.split_at(bytes_to_use);

            unsafe {
                let tail = self.buffer.as_mut_ptr().offset(self.buffer_usage as isize);
                std::ptr::copy_nonoverlapping(to_use.as_ptr(), tail, bytes_to_use);
            }

            let (iter, _) = self.buffer.u64_stream();

            self.core.ingest_chunks(iter);

            bytes = leftover;
            self.buffer_usage = 0;
        }

        // Consume the input data in large chunks
        if bytes.len() >= CHUNK_SIZE {
            let (iter, leftover) = bytes.u64_stream();
            self.core.ingest_chunks(iter);
            bytes = leftover;
        }

        // Save any leftover data for the next call
        if bytes.len() > 0 {
            unsafe {
                std::ptr::copy_nonoverlapping(bytes.as_ptr(), self.buffer.as_mut_ptr(), bytes.len());
            }
            self.buffer_usage = bytes.len();
        }
    }

    pub fn finish(&self) -> u64 {
        let mut hash;

        // We have processed at least one full chunk
        if self.total_len >= CHUNK_SIZE as u64 {
            hash = self.core.finish();
        } else {
            hash = self.seed.wrapping_add(PRIME_5);
        }

        hash = hash.wrapping_add(self.total_len);

        let buffered = &self.buffer[..self.buffer_usage];
        let (buffered_u64s, buffered) = buffered.u64_stream();

        for mut k1 in buffered_u64s {
            k1 = k1.wrapping_mul(PRIME_2);
            k1 = k1.rotate_left(31);
            k1 = k1.wrapping_mul(PRIME_1);
            hash ^= k1;
            hash = hash.rotate_left(27);
            hash = hash.wrapping_mul(PRIME_1);
            hash = hash.wrapping_add(PRIME_4);
        }

        let (buffered_u32s, buffered) = buffered.u32_stream();

        for k1 in buffered_u32s {
            let k1 = (k1 as u64).wrapping_mul(PRIME_1);
            hash ^= k1;
            hash = hash.rotate_left(23);
            hash = hash.wrapping_mul(PRIME_2);
            hash = hash.wrapping_add(PRIME_3);
        }

        for buffered_u8 in buffered {
            let k1 = (*buffered_u8 as u64).wrapping_mul(PRIME_5);
            hash ^= k1;
            hash = hash.rotate_left(11);
            hash = hash.wrapping_mul(PRIME_1);
        }

        // The final intermixing
        hash ^= hash.wrapping_shr(33);
        hash = hash.wrapping_mul(PRIME_2);
        hash ^= hash.wrapping_shr(29);
        hash = hash.wrapping_mul(PRIME_3);
        hash ^= hash.wrapping_shr(32);

        hash
    }
}

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

    #[test]
    fn ingesting_byte_by_byte_is_equivalent_to_large_chunks() {
        let bytes: Vec<_> = (0..32).map(|_| 0).collect();

        let mut byte_by_byte = XxHash::from_seed(0);
        for byte in bytes.chunks(1) {
            byte_by_byte.write(byte);
        }

        let mut one_chunk = XxHash::from_seed(0);
        one_chunk.write(&bytes);

        assert_eq!(byte_by_byte.core, one_chunk.core);
    }

    #[test]
    fn hash_of_nothing_matches_c_implementation() {
        let mut hasher = XxHash::from_seed(0);
        hasher.write(&[]);
        assert_eq!(hasher.finish(), 0xef46db3751d8e999);
    }

    #[test]
    fn hash_of_single_byte_matches_c_implementation() {
        let mut hasher = XxHash::from_seed(0);
        hasher.write(&[42]);
        assert_eq!(hasher.finish(), 0x0a9edecebeb03ae4);
    }

    #[test]
    fn hash_of_multiple_bytes_matches_c_implementation() {
        let mut hasher = XxHash::from_seed(0);
        hasher.write(b"Hello, world!\0");
        assert_eq!(hasher.finish(), 0x7b06c531ea43e89f);
    }

    #[test]
    fn hash_of_multiple_chunks_matches_c_implementation() {
        let bytes: Vec<_> = (0..100).collect();
        let mut hasher = XxHash::from_seed(0);
        hasher.write(&bytes);
        assert_eq!(hasher.finish(), 0x6ac1e58032166597);
    }

    #[test]
    fn hash_with_different_seed_matches_c_implementation() {
        let mut hasher = XxHash::from_seed(0xae0543311b702d91);
        hasher.write(&[]);
        assert_eq!(hasher.finish(), 0x4b6a04fcdf7a4672);
    }

    #[test]
    fn hash_with_different_seed_and_multiple_chunks_matches_c_implementation() {
        let bytes: Vec<_> = (0..100).collect();
        let mut hasher = XxHash::from_seed(0xae0543311b702d91);
        hasher.write(&bytes);
        assert_eq!(hasher.finish(), 0x567e355e0682e1f1);
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use super::XxHash;

    #[inline(always)]
    fn straight_line_slice_bench(b: &mut test::Bencher, len: usize) {
        let bytes: Vec<_> = (0..100).cycle().take(len).collect();
        b.bytes = bytes.len() as u64;
        b.iter(|| {
            let mut hasher = XxHash::from_seed(0);
            hasher.write(&bytes);
            hasher.finish()
        });
    }

    #[bench]
    fn straight_line_megabyte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 1024*1024);
    }

    #[bench]
    fn straight_line_1024_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 1024);
    }

    #[bench]
    fn straight_line_512_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 512);
    }

    #[bench]
    fn straight_line_256_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 256);
    }

    #[bench]
    fn straight_line_128_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 128);
    }

    #[bench]
    fn straight_line_32_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 32);
    }

    #[bench]
    fn straight_line_16_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 16);
    }

    #[bench]
    fn straight_line_4_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 4);
    }

    #[bench]
    fn straight_line_1_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 1);
    }

    #[bench]
    fn straight_line_0_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 0);
    }
}
<commit_msg>Update comment with performance note<commit_after>#![feature(core)]
#![cfg_attr(test, feature(test))]

mod number_streams;

use number_streams::NumberStreams;

const CHUNK_SIZE: usize = 32;

const PRIME_1: u64 = 11400714785074694791;
const PRIME_2: u64 = 14029467366897019727;
const PRIME_3: u64 = 1609587929392839161;
const PRIME_4: u64 = 9650029242287828579;
const PRIME_5: u64 = 2870177450012600261;

#[derive(Copy,Clone,PartialEq)]
struct XxCore {
    v1: u64,
    v2: u64,
    v3: u64,
    v4: u64,
}

#[derive(Debug,Copy,Clone)]
pub struct XxHash {
    total_len: u64,
    seed: u64,
    core: XxCore,
    buffer: [u8; CHUNK_SIZE],
    buffer_usage: usize,
}

impl XxCore {
    fn from_seed(seed: u64) -> XxCore {
        XxCore {
            v1: seed.wrapping_add(PRIME_1).wrapping_add(PRIME_2),
            v2: seed.wrapping_add(PRIME_2),
            v3: seed,
            v4: seed.wrapping_sub(PRIME_1),
        }
    }

    #[inline(always)]
    fn ingest_chunks<I>(&mut self, values: I)
        where I: Iterator<Item=u64>
    {
        #[inline(always)]
        fn ingest_one_number(mut current_value: u64, mut value: u64) -> u64 {
            value = value.wrapping_mul(PRIME_2);
            current_value = current_value.wrapping_add(value);
            current_value = current_value.rotate_left(31);
            current_value.wrapping_mul(PRIME_1)
        };

        // By drawing these out, we can avoid going back and forth to
        // memory. It only really helps for large files, when we need
        // to iterate multiple times here.

        let mut v1 = self.v1;
        let mut v2 = self.v2;
        let mut v3 = self.v3;
        let mut v4 = self.v4;

        let mut values = values.peekable();

        while values.peek().is_some() {
            v1 = ingest_one_number(v1, values.next().unwrap());
            v2 = ingest_one_number(v2, values.next().unwrap());
            v3 = ingest_one_number(v3, values.next().unwrap());
            v4 = ingest_one_number(v4, values.next().unwrap());
        }

        self.v1 = v1;
        self.v2 = v2;
        self.v3 = v3;
        self.v4 = v4;
    }

    #[inline(always)]
    fn finish(&self) -> u64 {
        // The original code pulls out local vars for v[1234]
        // here. Performance tests did not show that to be effective
        // here, presumably because this method is not called in a
        // tight loop.

        let mut hash;

        hash =                   self.v1.rotate_left( 1);
        hash = hash.wrapping_add(self.v2.rotate_left( 7));
        hash = hash.wrapping_add(self.v3.rotate_left(12));
        hash = hash.wrapping_add(self.v4.rotate_left(18));

        #[inline(always)]
        fn mix_one(mut hash: u64, mut value: u64) -> u64 {
            value = value.wrapping_mul(PRIME_2);
            value = value.rotate_left(31);
            value = value.wrapping_mul(PRIME_1);
            hash ^= value;
            hash = hash.wrapping_mul(PRIME_1);
            hash.wrapping_add(PRIME_4)
        }

        hash = mix_one(hash, self.v1);
        hash = mix_one(hash, self.v2);
        hash = mix_one(hash, self.v3);
        hash = mix_one(hash, self.v4);

        hash
    }
}

impl std::fmt::Debug for XxCore {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        write!(
            f, "XxCore {{ {:016x} {:016x} {:016x} {:016x} }}",
            self.v1, self.v2, self.v3, self.v4
        )
    }
}

impl XxHash {
    pub fn from_seed(seed: u64) -> XxHash {
        XxHash {
            total_len: 0,
            seed: seed,
            core: XxCore::from_seed(seed),
            buffer: [0; CHUNK_SIZE],
            buffer_usage: 0,
        }
    }
}

impl XxHash {
    pub fn write(&mut self, bytes: &[u8]) {
        let mut bytes = bytes;

        self.total_len += bytes.len() as u64;

        // Even with new data, we still don't have a full buffer. Wait
        // until we have a full buffer.
        if self.buffer_usage + bytes.len() < self.buffer.len() {
            unsafe {
                let tail = self.buffer.as_mut_ptr().offset(self.buffer_usage as isize);
                std::ptr::copy_nonoverlapping(bytes.as_ptr(), tail, bytes.len());
            }
            self.buffer_usage += bytes.len();
            return;
        }

        // Some data left from previous update. Fill the buffer and
        // consume it first.
        if self.buffer_usage > 0 {
            let bytes_to_use = self.buffer.len() - self.buffer_usage;
            let (to_use, leftover) = bytes.split_at(bytes_to_use);

            unsafe {
                let tail = self.buffer.as_mut_ptr().offset(self.buffer_usage as isize);
                std::ptr::copy_nonoverlapping(to_use.as_ptr(), tail, bytes_to_use);
            }

            let (iter, _) = self.buffer.u64_stream();

            self.core.ingest_chunks(iter);

            bytes = leftover;
            self.buffer_usage = 0;
        }

        // Consume the input data in large chunks
        if bytes.len() >= CHUNK_SIZE {
            let (iter, leftover) = bytes.u64_stream();
            self.core.ingest_chunks(iter);
            bytes = leftover;
        }

        // Save any leftover data for the next call
        if bytes.len() > 0 {
            unsafe {
                std::ptr::copy_nonoverlapping(bytes.as_ptr(), self.buffer.as_mut_ptr(), bytes.len());
            }
            self.buffer_usage = bytes.len();
        }
    }

    pub fn finish(&self) -> u64 {
        let mut hash;

        // We have processed at least one full chunk
        if self.total_len >= CHUNK_SIZE as u64 {
            hash = self.core.finish();
        } else {
            hash = self.seed.wrapping_add(PRIME_5);
        }

        hash = hash.wrapping_add(self.total_len);

        let buffered = &self.buffer[..self.buffer_usage];
        let (buffered_u64s, buffered) = buffered.u64_stream();

        for mut k1 in buffered_u64s {
            k1 = k1.wrapping_mul(PRIME_2);
            k1 = k1.rotate_left(31);
            k1 = k1.wrapping_mul(PRIME_1);
            hash ^= k1;
            hash = hash.rotate_left(27);
            hash = hash.wrapping_mul(PRIME_1);
            hash = hash.wrapping_add(PRIME_4);
        }

        let (buffered_u32s, buffered) = buffered.u32_stream();

        for k1 in buffered_u32s {
            let k1 = (k1 as u64).wrapping_mul(PRIME_1);
            hash ^= k1;
            hash = hash.rotate_left(23);
            hash = hash.wrapping_mul(PRIME_2);
            hash = hash.wrapping_add(PRIME_3);
        }

        for buffered_u8 in buffered {
            let k1 = (*buffered_u8 as u64).wrapping_mul(PRIME_5);
            hash ^= k1;
            hash = hash.rotate_left(11);
            hash = hash.wrapping_mul(PRIME_1);
        }

        // The final intermixing
        hash ^= hash.wrapping_shr(33);
        hash = hash.wrapping_mul(PRIME_2);
        hash ^= hash.wrapping_shr(29);
        hash = hash.wrapping_mul(PRIME_3);
        hash ^= hash.wrapping_shr(32);

        hash
    }
}

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

    #[test]
    fn ingesting_byte_by_byte_is_equivalent_to_large_chunks() {
        let bytes: Vec<_> = (0..32).map(|_| 0).collect();

        let mut byte_by_byte = XxHash::from_seed(0);
        for byte in bytes.chunks(1) {
            byte_by_byte.write(byte);
        }

        let mut one_chunk = XxHash::from_seed(0);
        one_chunk.write(&bytes);

        assert_eq!(byte_by_byte.core, one_chunk.core);
    }

    #[test]
    fn hash_of_nothing_matches_c_implementation() {
        let mut hasher = XxHash::from_seed(0);
        hasher.write(&[]);
        assert_eq!(hasher.finish(), 0xef46db3751d8e999);
    }

    #[test]
    fn hash_of_single_byte_matches_c_implementation() {
        let mut hasher = XxHash::from_seed(0);
        hasher.write(&[42]);
        assert_eq!(hasher.finish(), 0x0a9edecebeb03ae4);
    }

    #[test]
    fn hash_of_multiple_bytes_matches_c_implementation() {
        let mut hasher = XxHash::from_seed(0);
        hasher.write(b"Hello, world!\0");
        assert_eq!(hasher.finish(), 0x7b06c531ea43e89f);
    }

    #[test]
    fn hash_of_multiple_chunks_matches_c_implementation() {
        let bytes: Vec<_> = (0..100).collect();
        let mut hasher = XxHash::from_seed(0);
        hasher.write(&bytes);
        assert_eq!(hasher.finish(), 0x6ac1e58032166597);
    }

    #[test]
    fn hash_with_different_seed_matches_c_implementation() {
        let mut hasher = XxHash::from_seed(0xae0543311b702d91);
        hasher.write(&[]);
        assert_eq!(hasher.finish(), 0x4b6a04fcdf7a4672);
    }

    #[test]
    fn hash_with_different_seed_and_multiple_chunks_matches_c_implementation() {
        let bytes: Vec<_> = (0..100).collect();
        let mut hasher = XxHash::from_seed(0xae0543311b702d91);
        hasher.write(&bytes);
        assert_eq!(hasher.finish(), 0x567e355e0682e1f1);
    }
}

#[cfg(test)]
mod bench {
    extern crate test;

    use super::XxHash;

    #[inline(always)]
    fn straight_line_slice_bench(b: &mut test::Bencher, len: usize) {
        let bytes: Vec<_> = (0..100).cycle().take(len).collect();
        b.bytes = bytes.len() as u64;
        b.iter(|| {
            let mut hasher = XxHash::from_seed(0);
            hasher.write(&bytes);
            hasher.finish()
        });
    }

    #[bench]
    fn straight_line_megabyte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 1024*1024);
    }

    #[bench]
    fn straight_line_1024_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 1024);
    }

    #[bench]
    fn straight_line_512_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 512);
    }

    #[bench]
    fn straight_line_256_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 256);
    }

    #[bench]
    fn straight_line_128_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 128);
    }

    #[bench]
    fn straight_line_32_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 32);
    }

    #[bench]
    fn straight_line_16_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 16);
    }

    #[bench]
    fn straight_line_4_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 4);
    }

    #[bench]
    fn straight_line_1_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 1);
    }

    #[bench]
    fn straight_line_0_byte(b: &mut test::Bencher) {
        straight_line_slice_bench(b, 0);
    }
}
<|endoftext|>


<commit_before>//! Does all the magic to have you potentially long output piped through the
//! external pager. Similar to what git does for its output.
//!
//! # Quick Start
//!
//! ```rust
//! extern crate pager;
//! use pager::Pager;
//! fn main() {
//!     Pager::new().setup();
//!     // The rest of your program goes here
//! }
//! ```
//!
//! Under the hood this forks the current process, connects child' stdout
//! to parent's stdin, and then replaces the parent with the pager of choice
//! (environment variable PAGER). The child just continues as normal. If PAGER
//! environment variable is not present `Pager` probes current PATH for `more`.
//! If found it is used as a default pager.
//!
//! You can control pager to a limited degree. For example you can change the
//! environment variable used for finding pager executable.
//!
//! ```rust
//! extern crate pager;
//! use pager::Pager;
//! fn main() {
//!     Pager::env("MY_PAGER").setup();
//!     // The rest of your program goes here
//! }
//! ```
//!
//! If no suitable pager found `setup()` does nothing and your executable keeps
//! running as usual. `Pager` cleans after itself and doesn't leak resources in
//! case of setup failure.
//!
//! If you need to disable pager altogether set environment variable `NOPAGER` and `Pager::setup()`
//! will skip initialization. The host application will continue as normal. `Pager::ok()` will
//! reflect the fact that no Pager is active.

extern crate errno;
extern crate libc;

mod utils;

use std::ffi::OsString;

const DEFAULT_PAGER_ENV: &'static str = "PAGER";

#[derive(Debug, Default)]
pub struct Pager {
    pager: Option<OsString>,
    env: String,
    ok: bool,
}

impl Pager {
    /// Creates new instance of pager with default settings
    pub fn new() -> Self {
        Pager::env(DEFAULT_PAGER_ENV)
    }

    /// Creates new instance of pager using `env` environment variable instead of PAGER
    pub fn env(env: &str) -> Self {
        let pager = utils::find_pager(env);

        Pager {
            pager: pager,
            env: env.into(),
            ok: true,
        }
    }

    /// Gives quick assessment of successful Pager setup
    pub fn ok(&self) -> bool {
        self.ok
    }

    /// Initiates Pager framework and sets up all the necessary environment for sending standard
    /// output to the activated pager.
    pub fn setup(&mut self) {
        if let Some(ref pager) = self.pager {
            let (pager_stdin, main_stdout) = utils::pipe();
            let pid = utils::fork();
            match pid {
                -1 => {
                    // Fork failed
                    utils::close(pager_stdin);
                    utils::close(main_stdout);
                    self.ok = false
                }
                0 => {
                    // I am child
                    utils::dup2(main_stdout, libc::STDOUT_FILENO);
                    utils::close(pager_stdin);
                }
                _ => {
                    // I am parent
                    utils::dup2(pager_stdin, libc::STDIN_FILENO);
                    utils::close(main_stdout);
                    utils::execvp(pager);
                }
            }
        }
    }
}
<commit_msg>Drop 'static lifetime - it is not needed anymore<commit_after>//! Does all the magic to have you potentially long output piped through the
//! external pager. Similar to what git does for its output.
//!
//! # Quick Start
//!
//! ```rust
//! extern crate pager;
//! use pager::Pager;
//! fn main() {
//!     Pager::new().setup();
//!     // The rest of your program goes here
//! }
//! ```
//!
//! Under the hood this forks the current process, connects child' stdout
//! to parent's stdin, and then replaces the parent with the pager of choice
//! (environment variable PAGER). The child just continues as normal. If PAGER
//! environment variable is not present `Pager` probes current PATH for `more`.
//! If found it is used as a default pager.
//!
//! You can control pager to a limited degree. For example you can change the
//! environment variable used for finding pager executable.
//!
//! ```rust
//! extern crate pager;
//! use pager::Pager;
//! fn main() {
//!     Pager::env("MY_PAGER").setup();
//!     // The rest of your program goes here
//! }
//! ```
//!
//! If no suitable pager found `setup()` does nothing and your executable keeps
//! running as usual. `Pager` cleans after itself and doesn't leak resources in
//! case of setup failure.
//!
//! If you need to disable pager altogether set environment variable `NOPAGER` and `Pager::setup()`
//! will skip initialization. The host application will continue as normal. `Pager::ok()` will
//! reflect the fact that no Pager is active.

extern crate errno;
extern crate libc;

mod utils;

use std::ffi::OsString;

const DEFAULT_PAGER_ENV: &str = "PAGER";

#[derive(Debug, Default)]
pub struct Pager {
    pager: Option<OsString>,
    env: String,
    ok: bool,
}

impl Pager {
    /// Creates new instance of pager with default settings
    pub fn new() -> Self {
        Pager::env(DEFAULT_PAGER_ENV)
    }

    /// Creates new instance of pager using `env` environment variable instead of PAGER
    pub fn env(env: &str) -> Self {
        let pager = utils::find_pager(env);

        Pager {
            pager: pager,
            env: env.into(),
            ok: true,
        }
    }

    /// Gives quick assessment of successful Pager setup
    pub fn ok(&self) -> bool {
        self.ok
    }

    /// Initiates Pager framework and sets up all the necessary environment for sending standard
    /// output to the activated pager.
    pub fn setup(&mut self) {
        if let Some(ref pager) = self.pager {
            let (pager_stdin, main_stdout) = utils::pipe();
            let pid = utils::fork();
            match pid {
                -1 => {
                    // Fork failed
                    utils::close(pager_stdin);
                    utils::close(main_stdout);
                    self.ok = false
                }
                0 => {
                    // I am child
                    utils::dup2(main_stdout, libc::STDOUT_FILENO);
                    utils::close(pager_stdin);
                }
                _ => {
                    // I am parent
                    utils::dup2(pager_stdin, libc::STDIN_FILENO);
                    utils::close(main_stdout);
                    utils::execvp(pager);
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Hroi Sigurdsson
//
// Licensed under the MIT license <LICENSE-MIT or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed except according to those terms.

//! # Fast IP lookup table for IPv4/IPv6 prefixes
//!
//! This crate provides a datastructure for fast IP address lookups.
//! It aims at fast lookup times, and a small memory footprint.
//! A full IPv4 BGP table of more than 600k entries fits in less than 5 MB. A
//! full IPv6 BGP table of more than 25k entries fits in less than 1 MB.
//!
//! Longest match lookups on full BGP IP tables take on the order of 100ns.
//!
//! The internal datastructure is based on the Tree-bitmap algorithm described
//! by W. Eatherton, Z. Dittia, G. Varghes.
//!

use std::net::{Ipv4Addr, Ipv6Addr};
use std::marker::PhantomData;

mod tree_bitmap;
use tree_bitmap::TreeBitmap;

mod address;
use address::Address;

///The operations defined on the lookup table.
pub trait IpLookupTableOps<Addr, T> {
    /// Insert a value for the prefix designated by ip and masklen. If prefix
    /// existed previously, the old value is returned.
    /// # Example
    /// ```
    /// use treebitmap::{IpLookupTable, IpLookupTableOps};
    /// use std::net::Ipv6Addr;
    ///
    /// let mut table: IpLookupTable<Ipv6Addr,&str> = IpLookupTable::new();
    /// let prefix = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0);
    /// let masklen = 32;
    ///
    /// assert_eq!(table.insert(prefix, masklen, "foo"), None);
    /// // Insert duplicate
    /// assert_eq!(table.insert(prefix, masklen, "bar"), Some("foo"));
    /// ```
    fn insert(&mut self, ip: Addr, masklen: u32, value: T) -> Option<T>;

    /// Remove an entry from the lookup table. If the prefix existed previously,
    /// the value is returned.
    /// # Example
    /// ```
    /// use treebitmap::{IpLookupTable, IpLookupTableOps};
    /// use std::net::Ipv6Addr;
    ///
    /// let mut table: IpLookupTable<Ipv6Addr,&str> = IpLookupTable::new();
    /// let prefix = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0);
    /// let masklen = 32;
    /// table.insert(prefix, masklen, "foo");
    ///
    /// assert_eq!(table.remove(prefix, masklen), Some("foo"));
    /// // Remove non-existant
    /// assert_eq!(table.remove(prefix, masklen), None);
    /// ```
    fn remove(&mut self, ip: Addr, masklen: u32) -> Option<T>;

    /// Perform exact match lookup of ```ip```/```masklen``` and return the
    /// value.
    /// # Example
    /// ```
    /// use treebitmap::{IpLookupTable, IpLookupTableOps};
    /// use std::net::Ipv6Addr;
    ///
    /// let mut table: IpLookupTable<Ipv6Addr,&str> = IpLookupTable::new();
    /// let prefix = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0);
    /// let masklen = 32;
    /// table.insert(prefix, masklen, "foo");
    ///
    /// assert_eq!(table.exact_match(prefix, masklen), Some(&"foo"));
    /// // differing mask
    /// assert_eq!(table.exact_match(prefix, 48), None);
    /// ```
    fn exact_match(&self, ip: Addr, masklen: u32) -> Option<&T>;

    /// Perform longest match lookup of ```ip``` and return the best matching
    /// prefix, designated by ip, masklen, along with its value.
    /// # Example
    /// ```
    /// use treebitmap::{IpLookupTable, IpLookupTableOps};
    /// use std::net::Ipv6Addr;
    ///
    /// let mut table: IpLookupTable<Ipv6Addr,&str> = IpLookupTable::new();
    /// let less_specific = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0);
    /// let more_specific = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0, 0, 0, 0, 0);
    /// table.insert(less_specific, 32, "foo");
    /// table.insert(more_specific, 48, "bar");
    ///
    /// let lookupip = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef,
    ///                              0xcafe, 0xbabe, 0, 1);
    /// let result = table.longest_match(lookupip);
    /// assert_eq!(result, Some((more_specific, 48, &"bar")));
    ///
    /// let lookupip = Ipv6Addr::new(0x2001, 0xdb8, 0xcafe, 0xf00,
    ///                              0xf00, 0xf00, 0, 1);
    /// let result = table.longest_match(lookupip);
    /// assert_eq!(result, Some((less_specific, 32, &"foo")));
    /// ```
    fn longest_match(&self, ip: Addr) -> Option<(Addr, u32, &T)>;

    /// Returns iterator over prefixes and values.
    /// # Example
    /// ```
    /// use treebitmap::{IpLookupTable, IpLookupTableOps};
    /// use std::net::Ipv6Addr;
    ///
    /// let mut table: IpLookupTable<Ipv6Addr,&str> = IpLookupTable::new();
    /// let less_specific = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0);
    /// let more_specific = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0, 0, 0, 0, 0);
    /// table.insert(less_specific, 32, "foo");
    /// table.insert(more_specific, 48, "bar");
    ///
    /// let mut iter = table.iter();
    /// assert_eq!(iter.next(), Some((less_specific, 32, &"foo")));
    /// assert_eq!(iter.next(), Some((more_specific, 48, &"bar")));
    /// assert_eq!(iter.next(), None);
    /// ```
    fn iter(&self) -> Iter<Addr, T>;

    /// Mutable version of iter()
    fn iter_mut(&self) -> IterMut<Addr, T>;
}

/// A fast, compressed IP lookup table.
pub struct IpLookupTable<A, T> {
    inner: TreeBitmap<T>,
    _addrtype: PhantomData<A>,
}

impl<A, T> IpLookupTable<A, T> {
    /// Initialize an empty lookup table with no preallocation.
    pub fn new() -> Self {
        IpLookupTable {
            inner: TreeBitmap::new(),
            _addrtype: PhantomData,
        }
    }

    /// Initialize an empty lookup table with pre-allocated buffers.
    pub fn with_capacity(n: usize) -> Self {
        IpLookupTable {
            inner: TreeBitmap::with_capacity(n),
            _addrtype: PhantomData,
        }
    }

    /// Return the bytes used by nodes and results.
    pub fn mem_usage(&self) -> (usize, usize) {
        self.inner.mem_usage()
    }
}

/// Iterator over prefixes and associated values. The prefixes are returned in
/// "tree"-order.
pub struct Iter<'a, A, T: 'a> {
    inner: tree_bitmap::Iter<'a, T>,
    _addrtype: PhantomData<A>,
}

/// Mutable iterator over prefixes and associated values. The prefixes are
/// returned in "tree"-order.
pub struct IterMut<'a, A, T: 'a> {
    inner: tree_bitmap::IterMut<'a, T>,
    _addrtype: PhantomData<A>,
}

/// Converts ```IpLookupTable``` into an iterator. The prefixes are returned in
/// "tree"-order.
pub struct IntoIter<A, T> {
    inner: tree_bitmap::IntoIter<T>,
    _addrtype: PhantomData<A>,
}

macro_rules! impl_ops {
    ($addr_type:ty) => {
        impl<T: Sized> IpLookupTableOps<$addr_type, T> for IpLookupTable<$addr_type, T> {

            fn insert(&mut self, ip: $addr_type, masklen: u32, value: T) -> Option<T>{
                self.inner.insert(&ip.nibbles(), masklen, value)
            }

            fn remove(&mut self, ip: $addr_type, masklen: u32) -> Option<T>{
                self.inner.remove(&ip.nibbles(), masklen)
            }

            fn exact_match(&self, ip: $addr_type, masklen: u32) -> Option<&T> {
                self.inner.exact_match(&ip.nibbles(), masklen)
            }

            fn longest_match(&self, ip: $addr_type) -> Option<($addr_type, u32, &T)> {
                match self.inner.longest_match(&ip.nibbles()) {
                    Some((bits_matched,value)) => Some((ip.mask(bits_matched),
                                                        bits_matched, value)),
                    None => None
                }
            }

            fn iter(&self) -> Iter<$addr_type,T> {
                Iter{
                    inner: self.inner.iter(),
                    _addrtype: PhantomData,
                }
            }

            fn iter_mut(&self) -> IterMut<$addr_type,T> {
                IterMut{
                    inner: self.inner.iter_mut(),
                    _addrtype: PhantomData,
                }
            }
        }

        impl<'a, T: 'a> Iterator for Iter<'a, $addr_type, T> {
            type Item = ($addr_type, u32, &'a T);

            fn next(&mut self) -> Option<Self::Item> {
                match self.inner.next() {
                    Some((nibbles, masklen, value)) => {
                        Some((Address::from_nibbles(&nibbles[..]), masklen, value))
                    },
                    None => None,
                }
            }
        }

        impl<'a, T: 'a> Iterator for IterMut<'a, $addr_type, T> {
            type Item = ($addr_type, u32, &'a mut T);

            fn next(&mut self) -> Option<Self::Item> {
                match self.inner.next() {
                    Some((nibbles, masklen, value)) => {
                        Some((Address::from_nibbles(&nibbles[..]), masklen, value))
                    },
                    None => None,
                }
            }
        }

        impl<'a, T: 'a> Iterator for IntoIter<$addr_type, T> {
            type Item = ($addr_type, u32, T);

            fn next(&mut self) -> Option<Self::Item> {
                match self.inner.next() {
                    Some((nibbles, masklen, value)) => {
                        Some((Address::from_nibbles(&nibbles[..]), masklen, value))
                    },
                    None => None,
                }
            }
        }

        impl<T> IntoIterator for IpLookupTable<$addr_type, T> {
            type Item = ($addr_type, u32, T);
            type IntoIter = IntoIter<$addr_type,T>;

            fn into_iter(self) -> IntoIter<$addr_type,T> {
                IntoIter {
                    inner: self.inner.into_iter(),
                    _addrtype: PhantomData,
                }
            }
        }
    }
}

impl_ops!(Ipv4Addr);
impl_ops!(Ipv6Addr);

#[cfg(test)]
mod tests;
#[cfg(all(test, feature = "full-bgp-tests"))]
mod full_bgp_tests;
<commit_msg>Fixed documentation code style<commit_after>// Copyright 2016 Hroi Sigurdsson
//
// Licensed under the MIT license <LICENSE-MIT or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed except according to those terms.

//! # Fast IP lookup table for IPv4/IPv6 prefixes
//!
//! This crate provides a datastructure for fast IP address lookups.
//! It aims at fast lookup times, and a small memory footprint.
//! A full IPv4 BGP table of more than 600k entries fits in less than 5 MB. A
//! full IPv6 BGP table of more than 25k entries fits in less than 1 MB.
//!
//! Longest match lookups on full BGP IP tables take on the order of 100ns.
//!
//! The internal datastructure is based on the Tree-bitmap algorithm described
//! by W. Eatherton, Z. Dittia, G. Varghes.
//!

use std::net::{Ipv4Addr, Ipv6Addr};
use std::marker::PhantomData;

mod tree_bitmap;
use tree_bitmap::TreeBitmap;

mod address;
use address::Address;

/// The operations defined on the lookup table.
pub trait IpLookupTableOps<Addr, T> {
    /// Insert a value for the prefix designated by ip and masklen. If prefix
    /// existed previously, the old value is returned.
    ///
    /// # Examples
    ///
    /// ```
    /// use treebitmap::{IpLookupTable, IpLookupTableOps};
    /// use std::net::Ipv6Addr;
    ///
    /// let mut table: IpLookupTable<Ipv6Addr,&str> = IpLookupTable::new();
    /// let prefix = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0);
    /// let masklen = 32;
    ///
    /// assert_eq!(table.insert(prefix, masklen, "foo"), None);
    /// // Insert duplicate
    /// assert_eq!(table.insert(prefix, masklen, "bar"), Some("foo"));
    /// ```
    fn insert(&mut self, ip: Addr, masklen: u32, value: T) -> Option<T>;

    /// Remove an entry from the lookup table. If the prefix existed previously,
    /// the value is returned.
    ///
    /// # Examples
    ///
    /// ```
    /// use treebitmap::{IpLookupTable, IpLookupTableOps};
    /// use std::net::Ipv6Addr;
    ///
    /// let mut table: IpLookupTable<Ipv6Addr,&str> = IpLookupTable::new();
    /// let prefix = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef, 0, 0, 0, 0);
    /// let masklen = 32;
    /// table.insert(prefix, masklen, "foo");
    ///
    /// assert_eq!(table.remove(prefix, masklen), Some("foo"));
    /// // Remove non-existant
    /// assert_eq!(table.remove(prefix, masklen), None);
    /// ```
    fn remove(&mut self, ip: Addr, masklen: u32) -> Option<T>;

    /// Perform exact match lookup of `ip`/`masklen` and return the
    /// value.
    ///
    /// # Examples
    ///
    /// ```
    /// use treebitmap::{IpLookupTable, IpLookupTableOps};
    /// use std::net::Ipv6Addr;
    ///
    /// let mut table: IpLookupTable<Ipv6Addr,&str> = IpLookupTable::new();
    /// let prefix = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0);
    /// let masklen = 32;
    /// table.insert(prefix, masklen, "foo");
    ///
    /// assert_eq!(table.exact_match(prefix, masklen), Some(&"foo"));
    /// // differing mask
    /// assert_eq!(table.exact_match(prefix, 48), None);
    /// ```
    fn exact_match(&self, ip: Addr, masklen: u32) -> Option<&T>;

    /// Perform longest match lookup of `ip` and return the best matching
    /// prefix, designated by ip, masklen, along with its value.
    ///
    /// # Example
    ///
    /// ```
    /// use treebitmap::{IpLookupTable, IpLookupTableOps};
    /// use std::net::Ipv6Addr;
    ///
    /// let mut table: IpLookupTable<Ipv6Addr,&str> = IpLookupTable::new();
    /// let less_specific = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0);
    /// let more_specific = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0, 0, 0, 0, 0);
    /// table.insert(less_specific, 32, "foo");
    /// table.insert(more_specific, 48, "bar");
    ///
    /// let lookupip = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0xbeef,
    ///                              0xcafe, 0xbabe, 0, 1);
    /// let result = table.longest_match(lookupip);
    /// assert_eq!(result, Some((more_specific, 48, &"bar")));
    ///
    /// let lookupip = Ipv6Addr::new(0x2001, 0xdb8, 0xcafe, 0xf00,
    ///                              0xf00, 0xf00, 0, 1);
    /// let result = table.longest_match(lookupip);
    /// assert_eq!(result, Some((less_specific, 32, &"foo")));
    /// ```
    fn longest_match(&self, ip: Addr) -> Option<(Addr, u32, &T)>;

    /// Returns iterator over prefixes and values.
    ///
    /// # Examples
    ///
    /// ```
    /// use treebitmap::{IpLookupTable, IpLookupTableOps};
    /// use std::net::Ipv6Addr;
    ///
    /// let mut table: IpLookupTable<Ipv6Addr,&str> = IpLookupTable::new();
    /// let less_specific = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0);
    /// let more_specific = Ipv6Addr::new(0x2001, 0xdb8, 0xdead, 0, 0, 0, 0, 0);
    /// table.insert(less_specific, 32, "foo");
    /// table.insert(more_specific, 48, "bar");
    ///
    /// let mut iter = table.iter();
    /// assert_eq!(iter.next(), Some((less_specific, 32, &"foo")));
    /// assert_eq!(iter.next(), Some((more_specific, 48, &"bar")));
    /// assert_eq!(iter.next(), None);
    /// ```
    fn iter(&self) -> Iter<Addr, T>;

    /// Mutable version of iter()
    fn iter_mut(&self) -> IterMut<Addr, T>;
}

/// A fast, compressed IP lookup table.
pub struct IpLookupTable<A, T> {
    inner: TreeBitmap<T>,
    _addrtype: PhantomData<A>,
}

impl<A, T> IpLookupTable<A, T> {
    /// Initialize an empty lookup table with no preallocation.
    pub fn new() -> Self {
        IpLookupTable {
            inner: TreeBitmap::new(),
            _addrtype: PhantomData,
        }
    }

    /// Initialize an empty lookup table with pre-allocated buffers.
    pub fn with_capacity(n: usize) -> Self {
        IpLookupTable {
            inner: TreeBitmap::with_capacity(n),
            _addrtype: PhantomData,
        }
    }

    /// Return the bytes used by nodes and results.
    pub fn mem_usage(&self) -> (usize, usize) {
        self.inner.mem_usage()
    }
}

/// Iterator over prefixes and associated values. The prefixes are returned in
/// "tree"-order.
pub struct Iter<'a, A, T: 'a> {
    inner: tree_bitmap::Iter<'a, T>,
    _addrtype: PhantomData<A>,
}

/// Mutable iterator over prefixes and associated values. The prefixes are
/// returned in "tree"-order.
pub struct IterMut<'a, A, T: 'a> {
    inner: tree_bitmap::IterMut<'a, T>,
    _addrtype: PhantomData<A>,
}

/// Converts ```IpLookupTable``` into an iterator. The prefixes are returned in
/// "tree"-order.
pub struct IntoIter<A, T> {
    inner: tree_bitmap::IntoIter<T>,
    _addrtype: PhantomData<A>,
}

macro_rules! impl_ops {
    ($addr_type:ty) => {
        impl<T: Sized> IpLookupTableOps<$addr_type, T> for IpLookupTable<$addr_type, T> {

            fn insert(&mut self, ip: $addr_type, masklen: u32, value: T) -> Option<T>{
                self.inner.insert(&ip.nibbles(), masklen, value)
            }

            fn remove(&mut self, ip: $addr_type, masklen: u32) -> Option<T>{
                self.inner.remove(&ip.nibbles(), masklen)
            }

            fn exact_match(&self, ip: $addr_type, masklen: u32) -> Option<&T> {
                self.inner.exact_match(&ip.nibbles(), masklen)
            }

            fn longest_match(&self, ip: $addr_type) -> Option<($addr_type, u32, &T)> {
                match self.inner.longest_match(&ip.nibbles()) {
                    Some((bits_matched,value)) => Some((ip.mask(bits_matched),
                                                        bits_matched, value)),
                    None => None
                }
            }

            fn iter(&self) -> Iter<$addr_type,T> {
                Iter{
                    inner: self.inner.iter(),
                    _addrtype: PhantomData,
                }
            }

            fn iter_mut(&self) -> IterMut<$addr_type,T> {
                IterMut{
                    inner: self.inner.iter_mut(),
                    _addrtype: PhantomData,
                }
            }
        }

        impl<'a, T: 'a> Iterator for Iter<'a, $addr_type, T> {
            type Item = ($addr_type, u32, &'a T);

            fn next(&mut self) -> Option<Self::Item> {
                match self.inner.next() {
                    Some((nibbles, masklen, value)) => {
                        Some((Address::from_nibbles(&nibbles[..]), masklen, value))
                    },
                    None => None,
                }
            }
        }

        impl<'a, T: 'a> Iterator for IterMut<'a, $addr_type, T> {
            type Item = ($addr_type, u32, &'a mut T);

            fn next(&mut self) -> Option<Self::Item> {
                match self.inner.next() {
                    Some((nibbles, masklen, value)) => {
                        Some((Address::from_nibbles(&nibbles[..]), masklen, value))
                    },
                    None => None,
                }
            }
        }

        impl<'a, T: 'a> Iterator for IntoIter<$addr_type, T> {
            type Item = ($addr_type, u32, T);

            fn next(&mut self) -> Option<Self::Item> {
                match self.inner.next() {
                    Some((nibbles, masklen, value)) => {
                        Some((Address::from_nibbles(&nibbles[..]), masklen, value))
                    },
                    None => None,
                }
            }
        }

        impl<T> IntoIterator for IpLookupTable<$addr_type, T> {
            type Item = ($addr_type, u32, T);
            type IntoIter = IntoIter<$addr_type,T>;

            fn into_iter(self) -> IntoIter<$addr_type,T> {
                IntoIter {
                    inner: self.inner.into_iter(),
                    _addrtype: PhantomData,
                }
            }
        }
    }
}

impl_ops!(Ipv4Addr);
impl_ops!(Ipv6Addr);

#[cfg(test)]
mod tests;
#[cfg(all(test, feature = "full-bgp-tests"))]
mod full_bgp_tests;
<|endoftext|>