<commit_before>#![feature(slicing_syntax, globs)]

use std::error::FromError;

use std::str::Utf8Error;
use std::string::FromUtf8Error;

use std::io::{Reader, BufReader, IoError};
use std::io::net::ip::IpAddr;

#[derive(Show)]
pub enum Error {
    Io(IoError),
    FormatError(FormatError)
}

#[derive(Show, Copy)]
pub enum FormatError {
    UnknownBlock(u32),
    UnknownOption(u16),
    Utf8Error(Utf8Error)
}

impl FromError<IoError> for Error {
    fn from_error(err: IoError) -> Error {
        Error::Io(err)
    }
}

impl FromError<FromUtf8Error> for Error {
    fn from_error(err: FromUtf8Error) -> Error {
        Error::FormatError(FormatError::Utf8Error(err.utf8_error()))
    }
}

impl FromError<FormatError> for Error {
    fn from_error(err: FormatError) -> Error {
        Error::FormatError(err)
    }
}

#[derive(Show)]
pub enum Block {
    SectionHeader(SectionHeaderBlock),
    InterfaceDescription(InterfaceDescriptionBlock),
    InterfaceStatistics(InterfaceStatisticsBlock),
    EnhancedPacket(EnhancedPacketBlock)
}

#[derive(Show)]
pub struct SectionHeaderBlock {
    /// Magic number equal to 0x1A2B3C4D. This field can be used to detect
    /// sections saved on systems with different endianness.
    pub magic: u32,

    pub major_version: u16,
    pub minor_version: u16,

    /// Length in bytes of the following section excluding this block.
    ///
    /// This block can be used to skip the section for faster navigation in
    /// large files. Length of -1u64 means that the length is unspecified.
    pub section_length: u64,
    pub options: Vec<SectionHeaderOption>
}

#[derive(Show)]
pub enum SectionHeaderOption {
    /// Comment associated with the current block
    Comment(String),

    /// Description of the hardware used to create this section
    Hardware(String),

    /// Name of the operating system used to create this section
    OS(String),

    /// Name of the application used to create this section
    UserApplication(String)
}

/// Block type 1
#[derive(Show)]
pub struct InterfaceDescriptionBlock {
    /// Link layer type of the interface
    pub link_type: u16,

    /// Maximum number of bytes stored from each packet
    pub snap_len: u32,

    pub options: Vec<InterfaceDescriptionOption>
}

#[derive(Show)]
pub enum InterfaceDescriptionOption {
    /// Comment associated with the current block
    Comment(String),

    /// Name of the device used to capture data
    Name(String),

    /// Description of the device used to capture data
    Description(String),

    /// IPv4 interface address and netmask
    Ipv4Addr(u32, u32),

    /// IPv6 interface address and prefix length
    Ipv6Addr(IpAddr, u8),

    /// Hardware MAC address (48 bits)
    MacAddr(u64),

    /// Hardware EUI address
    EuiAddr(u64),

    /// Interface speed in bps
    Speed(u64),

    /// Resolution of timestamps.
    ///
    /// If the MSB is equal to zero, the remaining bits indicate the resolution
    /// as a negative power of 10 (e.g. 6 means microsecond resolution,
    /// timestamps are teh number of microseconds since 1/1/1970).
    ///
    /// If the MSB is equal to one, the remaining bits indicate the resolution
    /// as a negative power of 2 (e.g. 10 means 1/1024 of a second).
    ///
    /// If this options is not present, microsecond resolution is assumed.
    TsResolution(u8),

    /// Timezone for GMT support (?)
    Timezone(u32),

    /// The filter used to capture traffic. The first byte specifies the type
    /// of filter used (e.g. libpcap string, BPF bytecode, etc)
    Filter(u8, Vec<u8>),

    /// Operating system of the machine in which this interface is installed.
    /// This can be different from the Section Header Block if the capture
    /// was done remotely.
    OS(String),

    /// Length of Frame Check Sequence in bits for this interface. For link
    /// layers whose FCS length can change during time, the Packet Block
    /// flags be used.
    FcsLen(u8),

    /// Offset pf the timestamp in seconds
    TsOffset(u64)
}

/// Block type 5
#[derive(Show)]
pub struct InterfaceStatisticsBlock {
    pub interface_id: u32,
    pub timestamp: u64,

    pub options: Vec<InterfaceStatisticsOption>
}

#[derive(Show)]
pub enum InterfaceStatisticsOption {
    /// Comment associated with the current block
    Comment(String),

    /// Time when capture started
    StartTime(u64),

    /// Time when capture ended
    EndTime(u64),

    /// Number of packets received from the physical interface
    Received(u64),

    /// Number of packets dropped by the interface due to the lack of resources
    Dropped(u64),

    /// Number of packets accepted by the filter
    FilterAccepted(u64),

    /// Number of packets dropped by the operating system
    OSDropped(u64),

    /// Number of packets delivered to the user starting. The value in this
    /// field can be different from 'Accepted - Dropped' because some packets
    /// could still be in the OS buffers when the capture ended.
    Delivered(u64)
}

#[derive(Show)]
pub struct EnhancedPacketBlock {
    pub interface_id: u32,
    pub timestamp: u64,

    /// Actual length of the packet when it was transmitted on the network
    pub len: u32,

    pub data: Vec<u8>,
    pub options: Vec<EnhancedPacketBlockOption>
}

#[derive(Show)]
pub enum EnhancedPacketBlockOption {
    Comment(String),
    Flags(u32),
    Hash,

    /// Number of packets lost between the last packet
    DropCount(u64)
}

#[inline(always)]
fn dword_aligned(n: uint) -> uint {
    (n + 3) & !3
}

pub fn read_raw_block(r: &mut Reader) -> Result<(u32, Vec<u8>), IoError> {
    let block_type = try!(r.read_le_u32());

    let total_len = try!(r.read_le_u32());
    let data_len = total_len - 12; // 12 = type + 2*length

    let mut data = try!(r.read_exact(dword_aligned(data_len as uint)));
    data.truncate(data_len as uint);

    assert!(total_len == try!(r.read_le_u32()));

    Ok((block_type, data))
}

fn read_option(r: &mut Reader) -> Result<(u16, Vec<u8>), IoError> {
    let code = try!(r.read_le_u16());
    let len = try!(r.read_le_u16()) as uint;

    let mut data = try!(r.read_exact(dword_aligned(len)));
    data.truncate(len as uint);

    Ok((code, data))
}

pub fn read_block(r: &mut Reader) -> Result<Block, Error> {
    use Block::*;

    let (block_type, data) = try!(read_raw_block(r));
    let mut r = BufReader::new(&*data);

    let r = match block_type {
        0x0A0D0D0A => SectionHeader(try!(SectionHeaderBlock::read(&mut r))),
        1 => InterfaceDescription(try!(InterfaceDescriptionBlock::read(&mut r))),
        5 => InterfaceStatistics(try!(InterfaceStatisticsBlock::read(&mut r))),
        6 => EnhancedPacket(try!(EnhancedPacketBlock::read(&mut r))),
        _ => return Err(FromError::from_error(FormatError::UnknownBlock(block_type)))
    };

    Ok(r)
}

pub struct BlockIter<'a> {
    r: &'a mut (Reader + 'a)
}

pub struct PacketIter<'a> {
    r: &'a mut SimpleReader<'a>
}

pub struct SimpleReader<'a> {
    r: &'a mut (Reader + 'a),

    interfaces: Vec<InterfaceDescriptionBlock>,
    if_offset: uint
}

impl<'a> Iterator<Block> for BlockIter<'a> {
    fn next(&mut self) -> Option<Block> {
        match read_block(self.r) {
            Ok(block) => Some(block),
            Err(_) => None
        }
    }
}

pub type IterPacket<'a> = (&'a InterfaceDescriptionBlock, EnhancedPacketBlock);

impl<'a> Iterator<IterPacket<'a>> for PacketIter<'a> {
    fn next(&mut self) -> Option<IterPacket<'a>> {
        while let Ok(block) = read_block(self.r.r) {
            match block {
                Block::SectionHeader(_) => {
                    // Each section is independent from another,
                    // however we might still have live references
                    // to the old interface descriptions so we
                    // can't just clear the vector
                    self.r.if_offset = self.r.interfaces.len()
                },

                Block::InterfaceDescription(id) => self.r.interfaces.push(id),

                Block::EnhancedPacket(ep) => {
                    let iface = &self.r.interfaces[self.r.if_offset + ep.interface_id as uint];
                    unsafe {
                        // The interface description should live as long as
                        // SimpleReader, so this should be safe.
                        let iface = std::mem::transmute(iface);

                        return Some((iface, ep))
                    }
                },
                _ => {}
            }
        }

        None
    }
}

impl<'a> SimpleReader<'a> {
    pub fn new(r: &mut Reader) -> SimpleReader {
        SimpleReader { r: r, interfaces: Vec::new(), if_offset: 0}
    }

    pub fn blocks(&mut self) -> BlockIter {
        BlockIter { r: self.r }
    }

    pub fn packets(&'a mut self) -> PacketIter<'a> {
        PacketIter { r: self }
    }
}

impl SectionHeaderBlock {
    pub fn read(r: &mut BufReader) -> Result<SectionHeaderBlock, Error> {
        let magic = try!(r.read_le_u32());

        assert!(magic == 0x1A2B3C4D, "unsupported endianness");

        let major_version = try!(r.read_le_u16());
        let minor_version = try!(r.read_le_u16());
        let section_length = try!(r.read_le_u64());

        let mut options = Vec::new();

        if !r.eof() {
            loop {
                use SectionHeaderOption::*;

                let (code, data) = try!(read_option(r));

                let opt = match code {
                    0 => break,

                    1...4 => {
                        let s = try!(String::from_utf8(data));

                        match code {
                            1 => Comment(s),
                            2 => Hardware(s),
                            3 => OS(s),
                            4 => UserApplication(s),
                            _ => unreachable!()
                        }
                    }

                    _ => return Err(FromError::from_error(FormatError::UnknownOption(code)))
                };

                options.push(opt);
            }
        }

        Ok(SectionHeaderBlock {
            magic: magic,
            major_version: major_version,
            minor_version: minor_version,
            section_length: section_length,

            options: options
        })
    }
}

impl InterfaceDescriptionBlock {
    pub fn read(r: &mut BufReader) -> Result<InterfaceDescriptionBlock, Error> {
        let link_type = try!(r.read_le_u16());
        try!(r.read_le_u16()); // reserved
        let snap_len = try!(r.read_le_u32());

        let mut options = Vec::new();

        if !r.eof() {
            loop {
                use InterfaceDescriptionOption::*;
                let (code, data) = try!(read_option(r));

                let mut d = &*data;

                let opt = match code {
                    0 => break,
                    1...3 | 12 => {
                        let s = try!(String::from_utf8(data.clone()));

                        match code {
                            1 => Comment(s),
                            2 => Name(s),
                            3 => Description(s),
                            12 => OS(s),
                            _ => unreachable!()
                        }
                    }

                    4 => {
                        let ip = try!(d.read_le_u32());
                        let mask = try!(d.read_le_u32());

                        Ipv4Addr(ip, mask)
                    }

                    6 => MacAddr(try!(d.read_le_uint_n(6))),
                    9 => TsResolution(try!(d.read_byte())),

                    _ => return Err(FromError::from_error(FormatError::UnknownOption(code)))
                };

                options.push(opt);
            }
        }

        Ok(InterfaceDescriptionBlock {
            link_type: link_type,
            snap_len: snap_len,
            options: options
        })
    }
}

impl EnhancedPacketBlock {
    pub fn read(r: &mut BufReader) -> Result<EnhancedPacketBlock, Error> {
        let interface_id = try!(r.read_le_u32());
        let ts = try!(r.read_le_u64());
        let cap_len = try!(r.read_le_u32());
        let len = try!(r.read_le_u32());

        let aligned_len = dword_aligned(cap_len as uint);

        let mut packet_data = try!(r.read_exact(aligned_len));
        packet_data.truncate(cap_len as uint);

        let mut options = Vec::new();

        if !r.eof() {
            loop {
                use EnhancedPacketBlockOption::*;
                let (code, data) = try!(read_option(r));

                let opt = match code {
                    0 => break,
                    1 => {
                        match String::from_utf8(data.clone()) {
                            Ok(r) => Comment(r),
                            Err(err) => return Err(FromError::from_error(err))
                        }
                    }

                    // TODO: rest of the options
                    _ => return Err(FromError::from_error(FormatError::UnknownOption(code)))
                };

                options.push(opt);
            }
        }

        Ok(EnhancedPacketBlock {
            interface_id: interface_id,
            timestamp: ts,
            len: len,
            data: packet_data,
            options: options
        })
    }
}

impl InterfaceStatisticsBlock {
    pub fn read(r: &mut BufReader) -> Result<InterfaceStatisticsBlock, Error> {
        let interface_id = try!(r.read_le_u32());
        let ts = try!(r.read_le_u64());

        let mut options = Vec::new();

        if !r.eof() {
            loop {
                use InterfaceStatisticsOption::*;
                let (code, data) = try!(read_option(r));

                let opt = match code {
                    0 => break,
                    1 => {
                        match String::from_utf8(data.clone()) {
                            Ok(r) => Comment(r),
                            Err(err) => return Err(FromError::from_error(err))
                        }
                    }

                    2...8 => {
                        let data = try!((&*data).read_le_u64());

                        match code {
                            2 => StartTime(data),
                            3 => EndTime(data),
                            4 => Received(data),
                            5 => Dropped(data),
                            6 => FilterAccepted(data),
                            7 => OSDropped(data),
                            8 => Delivered(data),
                            _ => unreachable!()
                        }
                    }

                    _ => return Err(FromError::from_error(FormatError::UnknownOption(code)))
                };

                options.push(opt);
            }
        }

        Ok(InterfaceStatisticsBlock {
            interface_id: interface_id,
            timestamp: ts,
            options: options
        })
    }
}
<commit_msg>Use associated types for iterators<commit_after>#![feature(slicing_syntax, globs, associated_types)]

use std::error::FromError;

use std::str::Utf8Error;
use std::string::FromUtf8Error;

use std::io::{Reader, BufReader, IoError};
use std::io::net::ip::IpAddr;

#[derive(Show)]
pub enum Error {
    Io(IoError),
    FormatError(FormatError)
}

#[derive(Show, Copy)]
pub enum FormatError {
    UnknownBlock(u32),
    UnknownOption(u16),
    Utf8Error(Utf8Error)
}

impl FromError<IoError> for Error {
    fn from_error(err: IoError) -> Error {
        Error::Io(err)
    }
}

impl FromError<FromUtf8Error> for Error {
    fn from_error(err: FromUtf8Error) -> Error {
        Error::FormatError(FormatError::Utf8Error(err.utf8_error()))
    }
}

impl FromError<FormatError> for Error {
    fn from_error(err: FormatError) -> Error {
        Error::FormatError(err)
    }
}

#[derive(Show)]
pub enum Block {
    SectionHeader(SectionHeaderBlock),
    InterfaceDescription(InterfaceDescriptionBlock),
    InterfaceStatistics(InterfaceStatisticsBlock),
    EnhancedPacket(EnhancedPacketBlock)
}

#[derive(Show)]
pub struct SectionHeaderBlock {
    /// Magic number equal to 0x1A2B3C4D. This field can be used to detect
    /// sections saved on systems with different endianness.
    pub magic: u32,

    pub major_version: u16,
    pub minor_version: u16,

    /// Length in bytes of the following section excluding this block.
    ///
    /// This block can be used to skip the section for faster navigation in
    /// large files. Length of -1u64 means that the length is unspecified.
    pub section_length: u64,
    pub options: Vec<SectionHeaderOption>
}

#[derive(Show)]
pub enum SectionHeaderOption {
    /// Comment associated with the current block
    Comment(String),

    /// Description of the hardware used to create this section
    Hardware(String),

    /// Name of the operating system used to create this section
    OS(String),

    /// Name of the application used to create this section
    UserApplication(String)
}

/// Block type 1
#[derive(Show)]
pub struct InterfaceDescriptionBlock {
    /// Link layer type of the interface
    pub link_type: u16,

    /// Maximum number of bytes stored from each packet
    pub snap_len: u32,

    pub options: Vec<InterfaceDescriptionOption>
}

#[derive(Show)]
pub enum InterfaceDescriptionOption {
    /// Comment associated with the current block
    Comment(String),

    /// Name of the device used to capture data
    Name(String),

    /// Description of the device used to capture data
    Description(String),

    /// IPv4 interface address and netmask
    Ipv4Addr(u32, u32),

    /// IPv6 interface address and prefix length
    Ipv6Addr(IpAddr, u8),

    /// Hardware MAC address (48 bits)
    MacAddr(u64),

    /// Hardware EUI address
    EuiAddr(u64),

    /// Interface speed in bps
    Speed(u64),

    /// Resolution of timestamps.
    ///
    /// If the MSB is equal to zero, the remaining bits indicate the resolution
    /// as a negative power of 10 (e.g. 6 means microsecond resolution,
    /// timestamps are teh number of microseconds since 1/1/1970).
    ///
    /// If the MSB is equal to one, the remaining bits indicate the resolution
    /// as a negative power of 2 (e.g. 10 means 1/1024 of a second).
    ///
    /// If this options is not present, microsecond resolution is assumed.
    TsResolution(u8),

    /// Timezone for GMT support (?)
    Timezone(u32),

    /// The filter used to capture traffic. The first byte specifies the type
    /// of filter used (e.g. libpcap string, BPF bytecode, etc)
    Filter(u8, Vec<u8>),

    /// Operating system of the machine in which this interface is installed.
    /// This can be different from the Section Header Block if the capture
    /// was done remotely.
    OS(String),

    /// Length of Frame Check Sequence in bits for this interface. For link
    /// layers whose FCS length can change during time, the Packet Block
    /// flags be used.
    FcsLen(u8),

    /// Offset pf the timestamp in seconds
    TsOffset(u64)
}

/// Block type 5
#[derive(Show)]
pub struct InterfaceStatisticsBlock {
    pub interface_id: u32,
    pub timestamp: u64,

    pub options: Vec<InterfaceStatisticsOption>
}

#[derive(Show)]
pub enum InterfaceStatisticsOption {
    /// Comment associated with the current block
    Comment(String),

    /// Time when capture started
    StartTime(u64),

    /// Time when capture ended
    EndTime(u64),

    /// Number of packets received from the physical interface
    Received(u64),

    /// Number of packets dropped by the interface due to the lack of resources
    Dropped(u64),

    /// Number of packets accepted by the filter
    FilterAccepted(u64),

    /// Number of packets dropped by the operating system
    OSDropped(u64),

    /// Number of packets delivered to the user starting. The value in this
    /// field can be different from 'Accepted - Dropped' because some packets
    /// could still be in the OS buffers when the capture ended.
    Delivered(u64)
}

#[derive(Show)]
pub struct EnhancedPacketBlock {
    pub interface_id: u32,
    pub timestamp: u64,

    /// Actual length of the packet when it was transmitted on the network
    pub len: u32,

    pub data: Vec<u8>,
    pub options: Vec<EnhancedPacketBlockOption>
}

#[derive(Show)]
pub enum EnhancedPacketBlockOption {
    Comment(String),
    Flags(u32),
    Hash,

    /// Number of packets lost between the last packet
    DropCount(u64)
}

#[inline(always)]
fn dword_aligned(n: uint) -> uint {
    (n + 3) & !3
}

pub fn read_raw_block(r: &mut Reader) -> Result<(u32, Vec<u8>), IoError> {
    let block_type = try!(r.read_le_u32());

    let total_len = try!(r.read_le_u32());
    let data_len = total_len - 12; // 12 = type + 2*length

    let mut data = try!(r.read_exact(dword_aligned(data_len as uint)));
    data.truncate(data_len as uint);

    assert!(total_len == try!(r.read_le_u32()));

    Ok((block_type, data))
}

fn read_option(r: &mut Reader) -> Result<(u16, Vec<u8>), IoError> {
    let code = try!(r.read_le_u16());
    let len = try!(r.read_le_u16()) as uint;

    let mut data = try!(r.read_exact(dword_aligned(len)));
    data.truncate(len as uint);

    Ok((code, data))
}

pub fn read_block(r: &mut Reader) -> Result<Block, Error> {
    use Block::*;

    let (block_type, data) = try!(read_raw_block(r));
    let mut r = BufReader::new(&*data);

    let r = match block_type {
        0x0A0D0D0A => SectionHeader(try!(SectionHeaderBlock::read(&mut r))),
        1 => InterfaceDescription(try!(InterfaceDescriptionBlock::read(&mut r))),
        5 => InterfaceStatistics(try!(InterfaceStatisticsBlock::read(&mut r))),
        6 => EnhancedPacket(try!(EnhancedPacketBlock::read(&mut r))),
        _ => return Err(FromError::from_error(FormatError::UnknownBlock(block_type)))
    };

    Ok(r)
}

pub struct BlockIter<'a> {
    r: &'a mut (Reader + 'a)
}

pub struct PacketIter<'a> {
    r: &'a mut SimpleReader<'a>
}

pub struct SimpleReader<'a> {
    r: &'a mut (Reader + 'a),

    interfaces: Vec<InterfaceDescriptionBlock>,
    if_offset: uint
}

impl<'a> Iterator for BlockIter<'a> {
    type Item = Block;

    fn next(&mut self) -> Option<Block> {
        match read_block(self.r) {
            Ok(block) => Some(block),
            Err(_) => None
        }
    }
}

type IterPacket<'a> = (&'a InterfaceDescriptionBlock, EnhancedPacketBlock);

impl<'a> Iterator for PacketIter<'a> {
    type Item = IterPacket<'a>;

    fn next(&mut self) -> Option<IterPacket<'a>> {
        while let Ok(block) = read_block(self.r.r) {
            match block {
                Block::SectionHeader(_) => {
                    // Each section is independent from another,
                    // however we might still have live references
                    // to the old interface descriptions so we
                    // can't just clear the vector
                    self.r.if_offset = self.r.interfaces.len()
                },

                Block::InterfaceDescription(id) => self.r.interfaces.push(id),

                Block::EnhancedPacket(ep) => {
                    let iface = &self.r.interfaces[self.r.if_offset + ep.interface_id as uint];
                    unsafe {
                        // The interface description should live as long as
                        // SimpleReader, so this should be safe.
                        let iface = std::mem::transmute(iface);

                        return Some((iface, ep))
                    }
                },
                _ => {}
            }
        }

        None
    }
}

impl<'a> SimpleReader<'a> {
    pub fn new(r: &mut Reader) -> SimpleReader {
        SimpleReader { r: r, interfaces: Vec::new(), if_offset: 0}
    }

    pub fn blocks(&mut self) -> BlockIter {
        BlockIter { r: self.r }
    }

    pub fn packets(&'a mut self) -> PacketIter<'a> {
        PacketIter { r: self }
    }
}

impl SectionHeaderBlock {
    pub fn read(r: &mut BufReader) -> Result<SectionHeaderBlock, Error> {
        let magic = try!(r.read_le_u32());

        assert!(magic == 0x1A2B3C4D, "unsupported endianness");

        let major_version = try!(r.read_le_u16());
        let minor_version = try!(r.read_le_u16());
        let section_length = try!(r.read_le_u64());

        let mut options = Vec::new();

        if !r.eof() {
            loop {
                use SectionHeaderOption::*;

                let (code, data) = try!(read_option(r));

                let opt = match code {
                    0 => break,

                    1...4 => {
                        let s = try!(String::from_utf8(data));

                        match code {
                            1 => Comment(s),
                            2 => Hardware(s),
                            3 => OS(s),
                            4 => UserApplication(s),
                            _ => unreachable!()
                        }
                    }

                    _ => return Err(FromError::from_error(FormatError::UnknownOption(code)))
                };

                options.push(opt);
            }
        }

        Ok(SectionHeaderBlock {
            magic: magic,
            major_version: major_version,
            minor_version: minor_version,
            section_length: section_length,

            options: options
        })
    }
}

impl InterfaceDescriptionBlock {
    pub fn read(r: &mut BufReader) -> Result<InterfaceDescriptionBlock, Error> {
        let link_type = try!(r.read_le_u16());
        try!(r.read_le_u16()); // reserved
        let snap_len = try!(r.read_le_u32());

        let mut options = Vec::new();

        if !r.eof() {
            loop {
                use InterfaceDescriptionOption::*;
                let (code, data) = try!(read_option(r));

                let mut d = &*data;

                let opt = match code {
                    0 => break,
                    1...3 | 12 => {
                        let s = try!(String::from_utf8(data.clone()));

                        match code {
                            1 => Comment(s),
                            2 => Name(s),
                            3 => Description(s),
                            12 => OS(s),
                            _ => unreachable!()
                        }
                    }

                    4 => {
                        let ip = try!(d.read_le_u32());
                        let mask = try!(d.read_le_u32());

                        Ipv4Addr(ip, mask)
                    }

                    6 => MacAddr(try!(d.read_le_uint_n(6))),
                    9 => TsResolution(try!(d.read_byte())),

                    _ => return Err(FromError::from_error(FormatError::UnknownOption(code)))
                };

                options.push(opt);
            }
        }

        Ok(InterfaceDescriptionBlock {
            link_type: link_type,
            snap_len: snap_len,
            options: options
        })
    }
}

impl EnhancedPacketBlock {
    pub fn read(r: &mut BufReader) -> Result<EnhancedPacketBlock, Error> {
        let interface_id = try!(r.read_le_u32());
        let ts = try!(r.read_le_u64());
        let cap_len = try!(r.read_le_u32());
        let len = try!(r.read_le_u32());

        let aligned_len = dword_aligned(cap_len as uint);

        let mut packet_data = try!(r.read_exact(aligned_len));
        packet_data.truncate(cap_len as uint);

        let mut options = Vec::new();

        if !r.eof() {
            loop {
                use EnhancedPacketBlockOption::*;
                let (code, data) = try!(read_option(r));

                let opt = match code {
                    0 => break,
                    1 => {
                        match String::from_utf8(data.clone()) {
                            Ok(r) => Comment(r),
                            Err(err) => return Err(FromError::from_error(err))
                        }
                    }

                    // TODO: rest of the options
                    _ => return Err(FromError::from_error(FormatError::UnknownOption(code)))
                };

                options.push(opt);
            }
        }

        Ok(EnhancedPacketBlock {
            interface_id: interface_id,
            timestamp: ts,
            len: len,
            data: packet_data,
            options: options
        })
    }
}

impl InterfaceStatisticsBlock {
    pub fn read(r: &mut BufReader) -> Result<InterfaceStatisticsBlock, Error> {
        let interface_id = try!(r.read_le_u32());
        let ts = try!(r.read_le_u64());

        let mut options = Vec::new();

        if !r.eof() {
            loop {
                use InterfaceStatisticsOption::*;
                let (code, data) = try!(read_option(r));

                let opt = match code {
                    0 => break,
                    1 => {
                        match String::from_utf8(data.clone()) {
                            Ok(r) => Comment(r),
                            Err(err) => return Err(FromError::from_error(err))
                        }
                    }

                    2...8 => {
                        let data = try!((&*data).read_le_u64());

                        match code {
                            2 => StartTime(data),
                            3 => EndTime(data),
                            4 => Received(data),
                            5 => Dropped(data),
                            6 => FilterAccepted(data),
                            7 => OSDropped(data),
                            8 => Delivered(data),
                            _ => unreachable!()
                        }
                    }

                    _ => return Err(FromError::from_error(FormatError::UnknownOption(code)))
                };

                options.push(opt);
            }
        }

        Ok(InterfaceStatisticsBlock {
            interface_id: interface_id,
            timestamp: ts,
            options: options
        })
    }
}
<|endoftext|>


<commit_before>extern crate crypto;
extern crate rustc_serialize;

use crypto::{buffer, blockmodes};
use crypto::aes::{cbc_decryptor, KeySize as AesKeySize};
use crypto::buffer::{WriteBuffer, ReadBuffer, BufferResult};
use crypto::hmac::Hmac;
use crypto::mac::{Mac};
use crypto::sha1::Sha1;
use crypto::pbkdf2::pbkdf2;
use crypto::util::fixed_time_eq;

use rustc_serialize::hex::FromHex;
use rustc_serialize::base64::FromBase64;

use std::str::from_utf8;

pub struct Verifier {
    secret_key: Vec<u8>
}

#[derive(Debug, PartialEq)]
pub enum VerifierError {
    MessageParse,
    InvalidSignature,
    DecodePayload
}

pub fn create_derived_keys(secret: &str, salts: &Vec<&str>, key_sz: u32, key_iters: u32) -> Vec<Vec<u8>> {
    let mut mac = Hmac::new(Sha1::new(), secret.as_bytes());

    salts.iter().map(|salt| {
        let mut result: Vec<u8> = vec![0; key_sz as usize];

        pbkdf2(&mut mac, salt.as_bytes(), key_iters, &mut result);

        result
    }).collect()
}

fn split_by_dashes(message: &str) -> Option<(&str, &str)> {
    let split: Vec<&str> = message.splitn(2, "--").collect();

    match split.len() {
        2 => Some((split.first().unwrap(), split.last().unwrap())),
        _ => None
    }
}

fn split_by_dashes_from_u8_slice(slice: &[u8]) -> Option<(&str, &str)> {
    match from_utf8(slice) {
        Ok(string) => split_by_dashes(string),
        Err(_) => None
    }
}

impl Verifier {
    pub fn new(secret: &str) -> Verifier {
        Verifier {
            secret_key: secret.bytes().collect()
        }
    }

    pub fn verify(&self, message: &str) -> Result<Vec<u8>, VerifierError> {
        match split_by_dashes(&message) {
            Some((encoded_data, signature)) => {
                match self.is_valid_message(encoded_data, signature) {
                    true => {
                        match encoded_data.from_base64() {
                            Ok(decoded) => Ok(decoded),
                            Err(_) => Err(VerifierError::DecodePayload)
                        }
                    },

                    false => Err(VerifierError::InvalidSignature)
                }
            },

            None => Err(VerifierError::MessageParse)
        }
    }

    pub fn is_valid_message(&self, encoded_data: &str, signature: &str) -> bool {
        match signature.from_hex() {
            Ok(sig_bytes) => {
                let mut mac = Hmac::new(Sha1::new(), &self.secret_key);

                mac.input(encoded_data.as_bytes());

                fixed_time_eq(mac.result().code(), &sig_bytes)
            },

            Err(_) => false
        }
    }
}

pub struct Encryptor {
    secret_key: Vec<u8>,
    verifier: Verifier
}

#[derive(Debug, PartialEq)]
pub enum EncryptorError {
    KeyDerivationFailed,
    InvalidSignature,
    MessageParse,
    MessageDecode,
    DecryptError
}

impl Encryptor {
    pub fn new(secret: &str, salt: &str, sign_salt: &str, key_sz: u32, key_iters: u32) -> Result<Encryptor, EncryptorError> {
        let salts = vec![salt, sign_salt];
        let keys = create_derived_keys(secret, &salts, key_sz, key_iters);

        match (keys.first(), keys.last()) {
            (Some(cipher_key), Some(sig_key)) => {
                Ok(Encryptor {
                    secret_key: cipher_key.to_vec(),
                    verifier: Verifier {
                        secret_key: sig_key.to_vec()
                    }
                })
            }

            _ => Err(EncryptorError::KeyDerivationFailed)
        }
    }

    pub fn decrypt_and_verify(&self, message: &str) -> Result<Vec<u8>, EncryptorError>  {
        match self.verifier.verify(message) {
            Ok(decoded) => {
                match split_by_dashes_from_u8_slice(&decoded) {
                    Some((encoded_cipher_text, encoded_iv)) => {
                        match (encoded_cipher_text.from_base64(), encoded_iv.from_base64()) {
                            (Ok(cipher_text), Ok(iv)) => {
                                let key_sz = AesKeySize::KeySize256; // TODO make configurable
                                let mut decryptor = cbc_decryptor(key_sz, &self.secret_key, &iv, blockmodes::PkcsPadding);

                                let mut final_result: Vec<u8> = Vec::new();
                                let mut buffer = [0; 4096];

                                let mut read_buffer = buffer::RefReadBuffer::new(&cipher_text);
                                let mut write_buffer = buffer::RefWriteBuffer::new(&mut buffer);

                                loop {
                                    match decryptor.decrypt(&mut read_buffer, &mut write_buffer, true) {
                                        Ok(buffer_result) => {
                                            final_result.extend(write_buffer.take_read_buffer().take_remaining().iter().map(|&i| i));

                                            match buffer_result {
                                                BufferResult::BufferUnderflow => break,
                                                BufferResult::BufferOverflow => continue
                                            }
                                        }

                                        Err(_) => return Err(EncryptorError::DecryptError)
                                    }
                                }

                                Ok(final_result)
                            },

                            _ => Err(EncryptorError::MessageDecode)
                        }
                    },

                    None => Err(EncryptorError::MessageParse)
                }
            }

            Err(_) => Err(EncryptorError::InvalidSignature)
        }

    }
}

#[cfg(test)]
mod tests {
    use std::str;

    use Verifier;
    use VerifierError;

    use Encryptor;
    use EncryptorError;

    #[test]
    fn is_valid_message_returns_true_for_valid_signatures() {
        let data = "eyJrZXkiOiJ2YWx1ZSJ9";
        let sig  = "fa115453dbb4a28277b1ba07ef4c7437621f5d72";

        let verifier = Verifier::new("helloworld");

        assert_eq!(verifier.is_valid_message(data, sig), true);
    }

    #[test]
    fn is_valid_message_returns_false_for_invalid_signatures() {
        let data = "eyJrZXkiOiJ2YWx1ZSJ9";
        let sig  = "05330518df0e21fb9beec7a71a5f5f951c3f5254";

        let verifier = Verifier::new("helloworld");

        assert_eq!(verifier.is_valid_message(data, sig), false);
    }

    #[test]
    fn is_valid_message_returns_false_for_invalid_messages() {
        let data = "baddata";
        let sig  = "badsig";

        let verifier = Verifier::new("helloworld");

        assert_eq!(verifier.is_valid_message(data, sig), false);
    }

    #[test]
    fn verify_returns_decoded_message_for_valid_signatures() {
        let msg = "eyJrZXkiOiJ2YWx1ZSJ9--fa115453dbb4a28277b1ba07ef4c7437621f5d72";

        let verifier = Verifier::new("helloworld");

        assert_eq!(verifier.verify(msg).unwrap(), "{\"key\":\"value\"}".as_bytes());
    }

    #[test]
    fn verify_returns_invalid_signature_error_for_wrong_key() {
        let msg = "eyJrZXkiOiJ2YWx1ZSJ9--05330518df0e21fb9beec7a71a5f5f951c3f5254";

        let verifier = Verifier::new("helloworld");

        assert_eq!(verifier.verify(msg).unwrap_err(), VerifierError::InvalidSignature);
    }

    #[test]
    fn verify_returns_message_parse_error_for_empty_message() {
        let msg = "";

        let verifier = Verifier::new("helloworld");

        assert_eq!(verifier.verify(msg).unwrap_err(), VerifierError::MessageParse);
    }

    #[test]
    fn decrypt_and_verify_returns_decoded_message_for_valid_messages() {
        let msg = "c20wSnp6Z1o1U2MyWDVjU3BPeWNNQT09LS1JOWNyR25LdDRpZUUvcmoxVTdoSTNRPT0=--a79c9522355e55bf8e4302c66d8bf5638f1a50ec";

        let encryptor = Encryptor::new("helloworld", "test salt", "test signed salt", 64, 1000).unwrap();

        assert_eq!(encryptor.decrypt_and_verify(msg).unwrap(), "{\"key\":\"value\"}".as_bytes());
    }

    #[test]
    fn decrypt_and_verify_returns_invalid_signature_error_for_wrong_key(){
        let msg = "SnRXQXFhOE9WSGg2QmVGUDdHdkhNZz09LS1vcjFWcm53VU40YmV0SVcwdWFlK2NRPT0=--c879b51cbd92559d4d684c406b3aaebfbc958e9d";

        let encryptor = Encryptor::new("helloworld", "test salt", "test signed salt", 64, 1000).unwrap();

        assert_eq!(encryptor.decrypt_and_verify(msg).unwrap_err(), EncryptorError::InvalidSignature);
    }

    #[test]
    fn decrypt_and_verify_returns_invalid_signature_for_empty_message(){
        let msg = "";

        let encryptor = Encryptor::new("helloworld", "test salt", "test signed salt", 64, 1000).unwrap();

        assert_eq!(encryptor.decrypt_and_verify(msg).unwrap_err(), EncryptorError::InvalidSignature);
    }
}
<commit_msg>Add Verifier#generate function<commit_after>extern crate crypto;
extern crate rustc_serialize;

use crypto::{buffer, blockmodes};
use crypto::aes::{cbc_decryptor, KeySize as AesKeySize};
use crypto::buffer::{WriteBuffer, ReadBuffer, BufferResult};
use crypto::hmac::Hmac;
use crypto::mac::{Mac};
use crypto::sha1::Sha1;
use crypto::pbkdf2::pbkdf2;
use crypto::util::fixed_time_eq;

use rustc_serialize::hex::{FromHex, ToHex};
use rustc_serialize::base64::{FromBase64, ToBase64, STANDARD};

use std::str::from_utf8;

pub struct Verifier {
    secret_key: Vec<u8>
}

#[derive(Debug, PartialEq)]
pub enum VerifierError {
    MessageParse,
    InvalidSignature,
    DecodePayload
}

pub fn create_derived_keys(secret: &str, salts: &Vec<&str>, key_sz: u32, key_iters: u32) -> Vec<Vec<u8>> {
    let mut mac = Hmac::new(Sha1::new(), secret.as_bytes());

    salts.iter().map(|salt| {
        let mut result: Vec<u8> = vec![0; key_sz as usize];

        pbkdf2(&mut mac, salt.as_bytes(), key_iters, &mut result);

        result
    }).collect()
}

fn split_by_dashes(message: &str) -> Option<(&str, &str)> {
    let split: Vec<&str> = message.splitn(2, "--").collect();

    match split.len() {
        2 => Some((split.first().unwrap(), split.last().unwrap())),
        _ => None
    }
}

fn split_by_dashes_from_u8_slice(slice: &[u8]) -> Option<(&str, &str)> {
    match from_utf8(slice) {
        Ok(string) => split_by_dashes(string),
        Err(_) => None
    }
}

impl Verifier {
    pub fn new(secret: &str) -> Verifier {
        Verifier {
            secret_key: secret.bytes().collect()
        }
    }

    pub fn verify(&self, message: &str) -> Result<Vec<u8>, VerifierError> {
        match split_by_dashes(&message) {
            Some((encoded_data, signature)) => {
                match self.is_valid_message(encoded_data, signature) {
                    true => {
                        match encoded_data.from_base64() {
                            Ok(decoded) => Ok(decoded),
                            Err(_) => Err(VerifierError::DecodePayload)
                        }
                    },

                    false => Err(VerifierError::InvalidSignature)
                }
            },

            None => Err(VerifierError::MessageParse)
        }
    }

    pub fn is_valid_message(&self, encoded_data: &str, signature: &str) -> bool {
        match signature.from_hex() {
            Ok(sig_bytes) => {
                let mut mac = Hmac::new(Sha1::new(), &self.secret_key);

                mac.input(encoded_data.as_bytes());

                fixed_time_eq(mac.result().code(), &sig_bytes)
            },

            Err(_) => false
        }
    }

    pub fn generate(&self, message: &str) -> Result<String, VerifierError> {
        let mut mac = Hmac::new(Sha1::new(), &self.secret_key);
        let encoded_data = message.as_bytes().to_base64(STANDARD);

        mac.input(encoded_data.as_bytes());

        let signature = mac.result();
        let result = format!("{}--{}", encoded_data, signature.code().to_hex());

        Ok(result.clone())
    }
}

pub struct Encryptor {
    secret_key: Vec<u8>,
    verifier: Verifier
}

#[derive(Debug, PartialEq)]
pub enum EncryptorError {
    KeyDerivationFailed,
    InvalidSignature,
    MessageParse,
    MessageDecode,
    DecryptError
}

impl Encryptor {
    pub fn new(secret: &str, salt: &str, sign_salt: &str, key_sz: u32, key_iters: u32) -> Result<Encryptor, EncryptorError> {
        let salts = vec![salt, sign_salt];
        let keys = create_derived_keys(secret, &salts, key_sz, key_iters);

        match (keys.first(), keys.last()) {
            (Some(cipher_key), Some(sig_key)) => {
                Ok(Encryptor {
                    secret_key: cipher_key.to_vec(),
                    verifier: Verifier {
                        secret_key: sig_key.to_vec()
                    }
                })
            }

            _ => Err(EncryptorError::KeyDerivationFailed)
        }
    }

    pub fn decrypt_and_verify(&self, message: &str) -> Result<Vec<u8>, EncryptorError>  {
        match self.verifier.verify(message) {
            Ok(decoded) => {
                match split_by_dashes_from_u8_slice(&decoded) {
                    Some((encoded_cipher_text, encoded_iv)) => {
                        match (encoded_cipher_text.from_base64(), encoded_iv.from_base64()) {
                            (Ok(cipher_text), Ok(iv)) => {
                                let key_sz = AesKeySize::KeySize256; // TODO make configurable
                                let mut decryptor = cbc_decryptor(key_sz, &self.secret_key, &iv, blockmodes::PkcsPadding);

                                let mut final_result: Vec<u8> = Vec::new();
                                let mut buffer = [0; 4096];

                                let mut read_buffer = buffer::RefReadBuffer::new(&cipher_text);
                                let mut write_buffer = buffer::RefWriteBuffer::new(&mut buffer);

                                loop {
                                    match decryptor.decrypt(&mut read_buffer, &mut write_buffer, true) {
                                        Ok(buffer_result) => {
                                            final_result.extend(write_buffer.take_read_buffer().take_remaining().iter().map(|&i| i));

                                            match buffer_result {
                                                BufferResult::BufferUnderflow => break,
                                                BufferResult::BufferOverflow => continue
                                            }
                                        }

                                        Err(_) => return Err(EncryptorError::DecryptError)
                                    }
                                }

                                Ok(final_result)
                            },

                            _ => Err(EncryptorError::MessageDecode)
                        }
                    },

                    None => Err(EncryptorError::MessageParse)
                }
            }

            Err(_) => Err(EncryptorError::InvalidSignature)
        }
    }
}

#[cfg(test)]
mod tests {
    use std::str;

    use Verifier;
    use VerifierError;

    use Encryptor;
    use EncryptorError;

    #[test]
    fn is_valid_message_returns_true_for_valid_signatures() {
        let data = "eyJrZXkiOiJ2YWx1ZSJ9";
        let sig  = "fa115453dbb4a28277b1ba07ef4c7437621f5d72";

        let verifier = Verifier::new("helloworld");

        assert_eq!(verifier.is_valid_message(data, sig), true);
    }

    #[test]
    fn is_valid_message_returns_false_for_invalid_signatures() {
        let data = "eyJrZXkiOiJ2YWx1ZSJ9";
        let sig  = "05330518df0e21fb9beec7a71a5f5f951c3f5254";

        let verifier = Verifier::new("helloworld");

        assert_eq!(verifier.is_valid_message(data, sig), false);
    }

    #[test]
    fn is_valid_message_returns_false_for_invalid_messages() {
        let data = "baddata";
        let sig  = "badsig";

        let verifier = Verifier::new("helloworld");

        assert_eq!(verifier.is_valid_message(data, sig), false);
    }

    #[test]
    fn verify_returns_decoded_message_for_valid_signatures() {
        let msg = "eyJrZXkiOiJ2YWx1ZSJ9--fa115453dbb4a28277b1ba07ef4c7437621f5d72";

        let verifier = Verifier::new("helloworld");

        assert_eq!(verifier.verify(msg).unwrap(), "{\"key\":\"value\"}".as_bytes());
    }

    #[test]
    fn verify_returns_invalid_signature_error_for_wrong_key() {
        let msg = "eyJrZXkiOiJ2YWx1ZSJ9--05330518df0e21fb9beec7a71a5f5f951c3f5254";

        let verifier = Verifier::new("helloworld");

        assert_eq!(verifier.verify(msg).unwrap_err(), VerifierError::InvalidSignature);
    }

    #[test]
    fn verify_returns_message_parse_error_for_empty_message() {
        let msg = "";

        let verifier = Verifier::new("helloworld");

        assert_eq!(verifier.verify(msg).unwrap_err(), VerifierError::MessageParse);
    }

    #[test]
    fn decrypt_and_verify_returns_decoded_message_for_valid_messages() {
        let msg = "c20wSnp6Z1o1U2MyWDVjU3BPeWNNQT09LS1JOWNyR25LdDRpZUUvcmoxVTdoSTNRPT0=--a79c9522355e55bf8e4302c66d8bf5638f1a50ec";

        let encryptor = Encryptor::new("helloworld", "test salt", "test signed salt", 64, 1000).unwrap();

        assert_eq!(encryptor.decrypt_and_verify(msg).unwrap(), "{\"key\":\"value\"}".as_bytes());
    }

    #[test]
    fn decrypt_and_verify_returns_invalid_signature_error_for_wrong_key(){
        let msg = "SnRXQXFhOE9WSGg2QmVGUDdHdkhNZz09LS1vcjFWcm53VU40YmV0SVcwdWFlK2NRPT0=--c879b51cbd92559d4d684c406b3aaebfbc958e9d";

        let encryptor = Encryptor::new("helloworld", "test salt", "test signed salt", 64, 1000).unwrap();

        assert_eq!(encryptor.decrypt_and_verify(msg).unwrap_err(), EncryptorError::InvalidSignature);
    }

    #[test]
    fn decrypt_and_verify_returns_invalid_signature_for_empty_message(){
        let msg = "";

        let encryptor = Encryptor::new("helloworld", "test salt", "test signed salt", 64, 1000).unwrap();

        assert_eq!(encryptor.decrypt_and_verify(msg).unwrap_err(), EncryptorError::InvalidSignature);
    }

    #[test]
    fn generate_returns_signed_and_encoded_string(){
        let verifier = Verifier::new("helloworld");
        let expected = "eyJrZXkiOiJ2YWx1ZSJ9--fa115453dbb4a28277b1ba07ef4c7437621f5d72";

        assert_eq!(verifier.generate("{\"key\":\"value\"}").unwrap(), expected.to_string());
    }
}
<|endoftext|>


<commit_before>//! This crate provides (at this time) a single function, `take()`.
//!
//! `take()` allows for taking `T` out of a `&mut T`, doing anything with it including consuming it, and producing another `T` to put back in the `&mut T`.
//!
//! During `take()`, if a panic occurs, the entire process will be exited, as there's no valid `T` to put back into the `&mut T`.
//! Use `take_or_recover()` to replace the `&mut T` with a recovery value before continuing the panic.
//!
//! Contrast with `std::mem::replace()`, which allows for putting a different `T` into a `&mut T`, but requiring the new `T` to be available before being able to consume the old `T`.

mod exit_on_panic;

use exit_on_panic::exit_on_panic;
use std::panic;

/// Allows use of a value pointed to by `&mut T` as though it was owned, as long as a `T` is made available afterwards.
///
/// The closure must return a valid T.
/// # Important
/// Will exit the program (with status code 101) if the closure panics.
///
/// # Example
/// ```
/// struct Foo;
/// let mut foo = Foo;
/// take_mut::take(&mut foo, |foo| {
///     // Can now consume the Foo, and provide a new value later
///     drop(foo);
///     // Do more stuff
///     Foo // Return new Foo from closure, which goes back into the &mut Foo
/// });
/// ```
pub fn take<T, F>(mut_ref: &mut T, closure: F)
  where F: FnOnce(T) -> T {
    use std::ptr;
    exit_on_panic(|| {
        unsafe {
            let old_t = ptr::read(mut_ref);
            let new_t = closure(old_t);
            ptr::write(mut_ref, new_t);
        }
    });
}

#[test]
fn it_works() {
    #[derive(PartialEq, Eq, Debug)]
    enum Foo {A, B};
    impl Drop for Foo {
        fn drop(&mut self) {
            match *self {
                Foo::A => println!("Foo::A dropped"),
                Foo::B => println!("Foo::B dropped")
            }
        }
    }
    let mut foo = Foo::A;
    take(&mut foo, |f| {
       drop(f);
       Foo::B
    });
    assert_eq!(&foo, &Foo::B);
}

/// Allows use of a value pointed to by `&mut T` as though it was owned, as long as a `T` is made available afterwards.
///
/// The closure must return a valid T.
/// # Important
/// Will replace `&mut T` with `recover` if the closure panics, then continues the panic.
///
/// # Example
/// ```
/// struct Foo;
/// let mut foo = Foo;
/// take_mut::take_or_recover(&mut foo, Foo, |foo| {
///     // Can now consume the Foo, and provide a new value later
///     drop(foo);
///     // Do more stuff
///     Foo // Return new Foo from closure, which goes back into the &mut Foo
/// });
/// ```
pub fn take_or_recover<T, F>(mut_ref: &mut T, recover: T, closure: F)
  where F: FnOnce(T) -> T {
    use std::ptr;
    unsafe {
        let old_t = ptr::read(mut_ref);
        let new_t = panic::catch_unwind(panic::AssertUnwindSafe(|| closure(old_t)));
        match new_t {
            Err(err) => {
                ptr::write(mut_ref, recover);
                panic::resume_unwind(err);
            }
            Ok(new_t) => ptr::write(mut_ref, new_t),
        }
    }
}

#[test]
fn it_works_recover() {
    #[derive(PartialEq, Eq, Debug)]
    enum Foo {A, B};
    impl Drop for Foo {
        fn drop(&mut self) {
            match *self {
                Foo::A => println!("Foo::A dropped"),
                Foo::B => println!("Foo::B dropped")
            }
        }
    }
    let mut foo = Foo::A;
    take_or_recover(&mut foo, Foo::A, |f| {
       drop(f);
       Foo::B
    });
    assert_eq!(&foo, &Foo::B);
}

#[test]
fn it_works_recover_panic() {
    #[derive(PartialEq, Eq, Debug)]
    enum Foo {A, B, C};
    impl Drop for Foo {
        fn drop(&mut self) {
            match *self {
                Foo::A => println!("Foo::A dropped"),
                Foo::B => println!("Foo::B dropped"),
                Foo::C => println!("Foo::C dropped")
            }
        }
    }
    let mut foo = Foo::A;

    let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
        take_or_recover(&mut foo, Foo::C, |f| {
            drop(f);
            panic!("panic");
            Foo::B
        });
    }));

    assert!(res.is_err());
    assert_eq!(&foo, &Foo::C);
}

<commit_msg>Use a closure for the recovery value for take_or_recover()<commit_after>//! This crate provides (at this time) a single function, `take()`.
//!
//! `take()` allows for taking `T` out of a `&mut T`, doing anything with it including consuming it, and producing another `T` to put back in the `&mut T`.
//!
//! During `take()`, if a panic occurs, the entire process will be exited, as there's no valid `T` to put back into the `&mut T`.
//! Use `take_or_recover()` to replace the `&mut T` with a recovery value before continuing the panic.
//!
//! Contrast with `std::mem::replace()`, which allows for putting a different `T` into a `&mut T`, but requiring the new `T` to be available before being able to consume the old `T`.

mod exit_on_panic;

use exit_on_panic::exit_on_panic;
use std::panic;

/// Allows use of a value pointed to by `&mut T` as though it was owned, as long as a `T` is made available afterwards.
///
/// The closure must return a valid T.
/// # Important
/// Will exit the program (with status code 101) if the closure panics.
///
/// # Example
/// ```
/// struct Foo;
/// let mut foo = Foo;
/// take_mut::take(&mut foo, |foo| {
///     // Can now consume the Foo, and provide a new value later
///     drop(foo);
///     // Do more stuff
///     Foo // Return new Foo from closure, which goes back into the &mut Foo
/// });
/// ```
pub fn take<T, F>(mut_ref: &mut T, closure: F)
  where F: FnOnce(T) -> T {
    use std::ptr;
    exit_on_panic(|| {
        unsafe {
            let old_t = ptr::read(mut_ref);
            let new_t = closure(old_t);
            ptr::write(mut_ref, new_t);
        }
    });
}

#[test]
fn it_works() {
    #[derive(PartialEq, Eq, Debug)]
    enum Foo {A, B};
    impl Drop for Foo {
        fn drop(&mut self) {
            match *self {
                Foo::A => println!("Foo::A dropped"),
                Foo::B => println!("Foo::B dropped")
            }
        }
    }
    let mut foo = Foo::A;
    take(&mut foo, |f| {
       drop(f);
       Foo::B
    });
    assert_eq!(&foo, &Foo::B);
}

/// Allows use of a value pointed to by `&mut T` as though it was owned, as long as a `T` is made available afterwards.
///
/// The closure must return a valid T.
/// # Important
/// Will replace `&mut T` with `recover` if the closure panics, then continues the panic.
///
/// # Example
/// ```
/// struct Foo;
/// let mut foo = Foo;
/// take_mut::take_or_recover(&mut foo, || Foo, |foo| {
///     // Can now consume the Foo, and provide a new value later
///     drop(foo);
///     // Do more stuff
///     Foo // Return new Foo from closure, which goes back into the &mut Foo
/// });
/// ```
pub fn take_or_recover<T, F, R>(mut_ref: &mut T, recover: R, closure: F)
  where F: FnOnce(T) -> T, R: FnOnce() -> T {
    use std::ptr;
    unsafe {
        let old_t = ptr::read(mut_ref);
        let new_t = panic::catch_unwind(panic::AssertUnwindSafe(|| closure(old_t)));
        match new_t {
            Err(err) => {
                let r = panic::catch_unwind(panic::AssertUnwindSafe(|| recover()))
                    .unwrap_or_else(|_| ::std::process::exit(101));
                ptr::write(mut_ref, r);
                panic::resume_unwind(err);
            }
            Ok(new_t) => ptr::write(mut_ref, new_t),
        }
    }
}

#[test]
fn it_works_recover() {
    #[derive(PartialEq, Eq, Debug)]
    enum Foo {A, B};
    impl Drop for Foo {
        fn drop(&mut self) {
            match *self {
                Foo::A => println!("Foo::A dropped"),
                Foo::B => println!("Foo::B dropped")
            }
        }
    }
    let mut foo = Foo::A;
    take_or_recover(&mut foo, || Foo::A, |f| {
       drop(f);
       Foo::B
    });
    assert_eq!(&foo, &Foo::B);
}

#[test]
fn it_works_recover_panic() {
    #[derive(PartialEq, Eq, Debug)]
    enum Foo {A, B, C};
    impl Drop for Foo {
        fn drop(&mut self) {
            match *self {
                Foo::A => println!("Foo::A dropped"),
                Foo::B => println!("Foo::B dropped"),
                Foo::C => println!("Foo::C dropped")
            }
        }
    }
    let mut foo = Foo::A;

    let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
        take_or_recover(&mut foo, || Foo::C, |f| {
            drop(f);
            panic!("panic");
            Foo::B
        });
    }));

    assert!(res.is_err());
    assert_eq!(&foo, &Foo::C);
}

<|endoftext|>


<commit_before>extern crate nix;
extern crate docopt;
extern crate cargo;
extern crate rustc_serialize;
extern crate gimli;
extern crate object;
extern crate memmap;
extern crate fallible_iterator;
extern crate rustc_demangle;
extern crate regex;

use std::io;
use std::io::{Error, ErrorKind};
use std::ffi::CString;
use std::path::Path;
use std::collections::HashMap;
use nix::unistd::*;
use nix::libc::pid_t;
use nix::sys::signal;
use nix::sys::wait::*;

pub mod tracer;
pub mod collectors;
pub mod breakpoint;
/// Should be unnecessary with a future nix crate release.
mod personality;
mod ptrace_control;

use tracer::*;
use breakpoint::*;
use ptrace_control::*;

/// Returns the coverage statistics for a test executable in the given workspace
pub fn get_test_coverage(root: &Path, test: &Path) {
    match fork() {
        Ok(ForkResult::Parent{ child }) => {
            match collect_coverage(root, test, child) {
                Ok(_) => println!("Coverage successful"),
                Err(e) => println!("Error occurred: {}", e),
            }
        }
        Ok(ForkResult::Child) => {
            execute_test(test, true);
        }
        Err(err) => {
            println!("Failed to run {}", test.display());
            println!("Error {}", err);
        }
    }
}
/// Collects the coverage data from the launched test
fn collect_coverage(project_path: &Path,
                    test_path: &Path,
                    test: pid_t) -> io::Result<()> {
    let mut traces = generate_tracer_data(project_path, test_path)?;
    let mut bps: HashMap<u64, Breakpoint> = HashMap::new();

    match waitpid(test, None) {
        Ok(WaitStatus::Stopped(child, signal::SIGTRAP)) => {
            for trace in traces.iter() {
                let file = trace.path.file_name().unwrap().to_str().unwrap();
                println!("Instrumenting {}:{} - {:X}", file, trace.line, trace.address);
                match Breakpoint::new(child, trace.address) {
                    Ok(bp) => {
                        let _ = bps.insert(trace.address, bp);
                    },
                    Err(e) => println!("Failed to instrument {}", e),
                }
            }
        },
        Ok(_) => println!("Unexpected grab"),
        Err(err) => println!("{}", err)
    }
    println!("Test process: {}", test);
    // Now we start hitting lines!

    let e = run_function(test, u64::max_value(), &mut traces, &mut bps);
    println!("{:?}", e);

    for t in traces.iter() {
        let file = t.path.file_name().unwrap().to_str().unwrap();
        println!("{}:{} - hits: {}", file, t.line, t.hits);
    }
    Ok(())
}

fn run_function(pid: pid_t,
                end: u64,
                mut traces: &mut Vec<TracerData>,
                mut breakpoints: &mut HashMap<u64, Breakpoint>) -> Result<i8, Error> {
    let mut res = 0i8;
    loop {
        continue_exec(pid)?;
        match waitpid(pid, None) {
            Ok(WaitStatus::Exited(_, sig)) => {
                res = sig;
                break;
            },
            Ok(WaitStatus::Stopped(child, signal::SIGTRAP)) |
                Ok(WaitStatus::Signaled(child, signal::SIGTRAP, true))=> {
                if let Ok(rip) = current_instruction_pointer(child) {
                    let rip = (rip - 1) as u64;
                    if let Some(ref mut bp) = breakpoints.get_mut(&rip) {
                        for mut t in traces.iter_mut()
                                       .filter(|ref x| x.address == rip) {
                            (*t).hits += 1;
                        }
                        if rip != end {
                            let _ = bp.step();
                        } else {
                            // Hit function end.
                            break;
                        }
                    }
                }

            },
            Ok(WaitStatus::Stopped(_, sig)) => {
                println!("Stopped due to {:?}", sig);
                break;
            },
            Ok(x) => println!("Unexpected: {:?}", x),
            Err(e) => {
                println!("e {}", e);
                return Err(Error::new(ErrorKind::Other, e))
            },
        }
    }
    Ok(res)
}

fn tests_mod_coverage(pid: pid_t,
                      mut traces: &mut Vec<TracerData>,
                      mut breakpoints: &mut HashMap<u64, Breakpoint>) {

    let is_test = | t: &TracerData | {
        match t.trace_type {
            LineType::TestEntry(_) => true,
            _ => false,
        }
    };

    let test_entries = traces.iter()
                             .filter(|t| is_test(t))
                             .map(|t| (t.address, t.trace_type))
                             .collect::<Vec<(u64, LineType)>>();

    println!("{:?}", test_entries);
    for te in test_entries.iter() {
        println!("Setting RIP");
        set_instruction_pointer(pid, te.0);

        let func_length:u64 = match te.1 {
            LineType::TestEntry(len) => len as u64,
            _ => 0,
        };
        println!("Running");
        match run_function(pid, te.0 + func_length, &mut traces, &mut breakpoints) {
            Ok(r) => println!("Ran function return {}", r),
            Err(e) => println!("Error running function: {}", e),
        }

        // Disable breakpoints in that function. It's a top level test and we've
        // covered it. This is likely unnecessary. But it gives me some level of
        // verification later.
        for bp in breakpoints.values() {
            if bp.pc >= te.0 && bp.pc < (te.0 + func_length) {
                let _ = bp.disable();
            }
        }
    }
}

/// Launches the test executable
fn execute_test(test: &Path, backtrace_on: bool) {
    let exec_path = CString::new(test.to_str().unwrap()).unwrap();
    match personality::disable_aslr() {
        Ok(_) => {},
        Err(e) => println!("ASLR disable failed: {}", e),
    }
    request_trace().ok()
                   .expect("Failed to trace");

    let mut envars: Vec<CString> = vec![CString::new("RUST_TEST_THREADS=1").unwrap()];
    if backtrace_on {
        envars.push(CString::new("RUST_BACKTRACE=1").unwrap());
    }
    execve(&exec_path, &[exec_path.clone()], envars.as_slice())
        .unwrap();
}
<commit_msg>Cleaned up some of lib.rs<commit_after>extern crate nix;
extern crate docopt;
extern crate cargo;
extern crate rustc_serialize;
extern crate gimli;
extern crate object;
extern crate memmap;
extern crate fallible_iterator;
extern crate rustc_demangle;
extern crate regex;

use std::io;
use std::io::{Error, ErrorKind};
use std::ffi::CString;
use std::path::Path;
use std::collections::HashMap;
use nix::unistd::*;
use nix::libc::pid_t;
use nix::sys::signal;
use nix::sys::wait::*;

pub mod tracer;
pub mod collectors;
pub mod breakpoint;
/// Should be unnecessary with a future nix crate release.
mod personality;
mod ptrace_control;

use tracer::*;
use breakpoint::*;
use ptrace_control::*;

/// Returns the coverage statistics for a test executable in the given workspace
pub fn get_test_coverage(root: &Path, test: &Path) {
    match fork() {
        Ok(ForkResult::Parent{ child }) => {
            match collect_coverage(root, test, child) {
                Ok(_) => println!("Coverage successful"),
                Err(e) => println!("Error occurred: {}", e),
            }
        }
        Ok(ForkResult::Child) => {
            execute_test(test, true);
        }
        Err(err) => {
            println!("Failed to run {}", test.display());
            println!("Error {}", err);
        }
    }
}
/// Collects the coverage data from the launched test
fn collect_coverage(project_path: &Path,
                    test_path: &Path,
                    test: pid_t) -> io::Result<()> {
    let mut traces = generate_tracer_data(project_path, test_path)?;
    let mut bps: HashMap<u64, Breakpoint> = HashMap::new();

    match waitpid(test, None) {
        Ok(WaitStatus::Stopped(child, signal::SIGTRAP)) => {
            for trace in traces.iter() {
                let file = trace.path.file_name().unwrap().to_str().unwrap();
                println!("Instrumenting {}:{} - {:X}", file, trace.line, trace.address);
                match Breakpoint::new(child, trace.address) {
                    Ok(bp) => {
                        let _ = bps.insert(trace.address, bp);
                    },
                    Err(e) => println!("Failed to instrument {}", e),
                }
            }
        },
        Ok(_) => println!("Unexpected grab"),
        Err(err) => println!("{}", err)
    }
    println!("Test process: {}", test);
    // Now we start hitting lines!
    let e = run_function(test, u64::max_value(), &mut traces, &mut bps);
    println!("{:?}", e);

    for t in traces.iter() {
        let file = t.path.file_name().unwrap().to_str().unwrap();
        println!("{}:{} - hits: {}", file, t.line, t.hits);
    }
    Ok(())
}

fn run_function(pid: pid_t,
                end: u64,
                mut traces: &mut Vec<TracerData>,
                mut breakpoints: &mut HashMap<u64, Breakpoint>) -> Result<i8, Error> {
    let mut res = 0i8;
    loop {
        continue_exec(pid)?;
        match waitpid(pid, None) {
            Ok(WaitStatus::Exited(_, sig)) => {
                res = sig;
                break;
            },
            Ok(WaitStatus::Stopped(child, signal::SIGTRAP)) |
                Ok(WaitStatus::Signaled(child, signal::SIGTRAP, true))=> {
                if let Ok(rip) = current_instruction_pointer(child) {
                    let rip = (rip - 1) as u64;
                    if let Some(ref mut bp) = breakpoints.get_mut(&rip) {
                        for mut t in traces.iter_mut()
                                       .filter(|ref x| x.address == rip) {
                            (*t).hits += 1;
                        }
                        if rip != end {
                            let _ = bp.step();
                        } else {
                            // Hit function end.
                            break;
                        }
                    }
                }

            },
            Ok(WaitStatus::Stopped(_, sig)) => {
                println!("Stopped due to {:?}", sig);
                break;
            },
            Ok(x) => println!("Unexpected: {:?}", x),
            Err(e) => {
                println!("e {}", e);
                return Err(Error::new(ErrorKind::Other, e))
            },
        }
    }
    Ok(res)
}

fn tests_mod_coverage(pid: pid_t,
                      mut traces: &mut Vec<TracerData>,
                      mut breakpoints: &mut HashMap<u64, Breakpoint>) {

    let is_test = | t: &TracerData | {
        match t.trace_type {
            LineType::TestEntry(_) => true,
            _ => false,
        }
    };

    let test_entries = traces.iter()
                             .filter(|t| is_test(t))
                             .map(|t| (t.address, t.trace_type))
                             .collect::<Vec<(u64, LineType)>>();

    println!("{:?}", test_entries);
    for te in test_entries.iter() {
        println!("Setting RIP");
        let _ = set_instruction_pointer(pid, te.0);

        let func_length:u64 = match te.1 {
            LineType::TestEntry(len) => len as u64,
            _ => 0,
        };
        println!("Running");
        match run_function(pid, te.0 + func_length, &mut traces, &mut breakpoints) {
            Ok(r) => println!("Ran function return {}", r),
            Err(e) => println!("Error running function: {}", e),
        }

        // Disable breakpoints in that function. It's a top level test and we've
        // covered it. This is likely unnecessary. But it gives me some level of
        // verification later.
        for bp in breakpoints.values() {
            if bp.pc >= te.0 && bp.pc < (te.0 + func_length) {
                let _ = bp.disable();
            }
        }
    }
}

/// Launches the test executable
fn execute_test(test: &Path, backtrace_on: bool) {
    let exec_path = CString::new(test.to_str().unwrap()).unwrap();
    match personality::disable_aslr() {
        Ok(_) => {},
        Err(e) => println!("ASLR disable failed: {}", e),
    }
    request_trace().ok()
                   .expect("Failed to trace");

    let mut envars: Vec<CString> = vec![CString::new("RUST_TEST_THREADS=1").unwrap()];
    if backtrace_on {
        envars.push(CString::new("RUST_BACKTRACE=1").unwrap());
    }
    execve(&exec_path, &[exec_path.clone()], envars.as_slice())
        .unwrap();
}
<|endoftext|>


<commit_before>//! An implementation of [FLAC](https://xiph.org/flac), free lossless audio
//! codec, written in Rust.
//!
//! The code is available on [GitHub](https://github.com/sourrust/flac).
//!
//! # Examples
//!
//! Basic decoding from a file.
//!
//! ```
//! use flac::StreamReader;
//! use std::fs::File;
//!
//! match StreamReader::<File>::from_file("path/to/file.flac") {
//!   Ok(mut stream) => {
//!     // Copy of `StreamInfo` to help convert to a different audio format.
//!     let info = stream.info();
//!
//!     // The explicit type needs to be one bit size up from the goal
//!     // size. (ie `i8` -> `i16`, `i16` -> `i32`, `i32` -> `i64`)
//!     for sample in stream.iter::<i32>() {
//!       // Iterate over each decoded sample
//!     }
//!   }
//!   Err(error)     => println!("{:?}", error),
//! }
//! ```

#[macro_use]
extern crate nom;

#[macro_use]
mod utility;
mod frame;
mod subframe;
pub mod metadata;
pub mod stream;

pub use metadata::Metadata;
pub use stream::{Stream, StreamBuffer, StreamReader};
pub use utility::{Sample, StreamProducer, ReadStream, ByteStream, ErrorKind};
<commit_msg>Export the SampleSize trait<commit_after>//! An implementation of [FLAC](https://xiph.org/flac), free lossless audio
//! codec, written in Rust.
//!
//! The code is available on [GitHub](https://github.com/sourrust/flac).
//!
//! # Examples
//!
//! Basic decoding from a file.
//!
//! ```
//! use flac::StreamReader;
//! use std::fs::File;
//!
//! match StreamReader::<File>::from_file("path/to/file.flac") {
//!   Ok(mut stream) => {
//!     // Copy of `StreamInfo` to help convert to a different audio format.
//!     let info = stream.info();
//!
//!     // The explicit type needs to be one bit size up from the goal
//!     // size. (ie `i8` -> `i16`, `i16` -> `i32`, `i32` -> `i64`)
//!     for sample in stream.iter::<i32>() {
//!       // Iterate over each decoded sample
//!     }
//!   }
//!   Err(error)     => println!("{:?}", error),
//! }
//! ```

#[macro_use]
extern crate nom;

#[macro_use]
mod utility;
mod frame;
mod subframe;
pub mod metadata;
pub mod stream;

pub use metadata::Metadata;
pub use stream::{Stream, StreamBuffer, StreamReader};
pub use utility::{
  Sample, SampleSize,
  StreamProducer, ReadStream, ByteStream,
  ErrorKind
};
<|endoftext|>


<commit_before>#![allow(unstable)]

use std::ops::Index;

#[derive(Copy, Clone, PartialEq)]
pub struct Matrix4([[f32; 4]; 4]);

impl Index<usize> for Matrix4 {
    type Output = [f32; 4];

    fn index<'a>(&'a self, row: &usize) -> &'a [f32; 4] {
        &self.0[*row]
    }
}

impl Matrix4 {
    pub fn as_array(&self) -> &[[f32; 4]; 4] {
        &self.0
    }

    pub fn from_array(array: [[f32; 4]; 4]) -> Self {
        Matrix4(array)
    }

    pub fn identity() -> Self {
        Matrix4::from_array([[1.0, 0.0, 0.0, 0.0],
                          [0.0, 1.0, 0.0, 0.0],
                          [0.0, 0.0, 1.0, 0.0],
                          [0.0, 0.0, 0.0, 1.0]])
    }

    pub fn operate<F: Fn(f32) -> f32>(&mut self, f: F) {
        self.0[0][0] = f(self.0[0][0]);
        self.0[0][1] = f(self.0[0][1]);
        self.0[0][2] = f(self.0[0][2]);
        self.0[0][3] = f(self.0[0][3]);

        self.0[1][0] = f(self.0[1][0]);
        self.0[1][1] = f(self.0[1][1]);
        self.0[1][2] = f(self.0[1][2]);
        self.0[1][3] = f(self.0[1][3]);

        self.0[2][0] = f(self.0[2][0]);
        self.0[2][1] = f(self.0[2][1]);
        self.0[2][2] = f(self.0[2][2]);
        self.0[2][3] = f(self.0[2][3]);

        self.0[3][0] = f(self.0[3][0]);
        self.0[3][1] = f(self.0[3][1]);
        self.0[3][2] = f(self.0[3][2]);
        self.0[3][3] = f(self.0[3][3]);
    }

    pub fn scale(&mut self, scalar: f32) {
        self.operate( |x| { x * scalar } )
    }
}

#[derive(Copy, Clone, PartialEq)]
pub struct Vector4 {
    pub x: f32,
    pub y: f32,
    pub z: f32,
    pub w: f32,
}

#[test]
fn test_matrix_to_array() {
    let array = {
        let foo = Matrix4::from_array([[1.0, 0.0, 0.0, 0.0],
                                       [0.0, 1.0, 0.0, 0.0],
                                       [0.0, 0.0, 1.0, 0.0],
                                       [0.0, 0.0, 0.0, 1.0]]);
        assert_eq!(foo.as_array()[0][0], 1.0);
        *foo.as_array()
    };
    assert_eq!(array[1][1], 1.0);
}

#[test]
fn test_scale() {
    let mut m = Matrix4::identity();
    assert_eq!(m[2][2], 1.0);
    m.scale(3.0);
    assert_eq!(m[2][2], 3.0);
}
<commit_msg>add translate API<commit_after>#![allow(unstable)]

use std::ops::Index;

#[derive(Copy, Clone, PartialEq)]
pub struct Matrix4([[f32; 4]; 4]);

impl Index<usize> for Matrix4 {
    type Output = [f32; 4];

    fn index<'a>(&'a self, row: &usize) -> &'a [f32; 4] {
        &self.0[*row]
    }
}

impl Matrix4 {
    pub fn as_array(&self) -> &[[f32; 4]; 4] {
        &self.0
    }

    pub fn from_array(array: [[f32; 4]; 4]) -> Self {
        Matrix4(array)
    }

    pub fn identity() -> Self {
        Matrix4::from_array([[1.0, 0.0, 0.0, 0.0],
                             [0.0, 1.0, 0.0, 0.0],
                             [0.0, 0.0, 1.0, 0.0],
                             [0.0, 0.0, 0.0, 1.0]])
    }

    pub fn scale(&mut self, x: f32, y: f32, z: f32) {
        self.0[0][0] *= x;
        self.0[1][1] *= y;
        self.0[2][2] *= z;
    }

    pub fn scale_by_vector(&mut self, v: &Vector4) {
        self.scale(v[0], v[1], v[2])
    }

    pub fn translate(&mut self, x: f32, y: f32, z: f32) {
        self.0[0][2] += x;
        self.0[1][2] += y;
        self.0[2][2] += z;
    }

    pub fn translate_by_vector(&mut self, v: &Vector4) {
        self.translate(v[0], v[1], v[2])
    }
}

#[derive(Copy, Clone, PartialEq)]
pub struct Vector4([f32; 4]);

impl Index<usize> for Vector4 {
    type Output = f32;

    fn index<'a>(&'a self, index: &usize) -> &'a f32 {
        &self.0[*index]
    }
}

impl Vector4 {
    pub fn as_array(&self) -> &[f32; 4] {
        &self.0
    }

    pub fn from_array(array: [f32; 4]) -> Self {
        Vector4(array)
    }

    pub fn new(x: f32, y: f32, z: f32) -> Self {
        Vector4::from_array([x, y, z, 1.0])
    }
}

#[test]
fn test_matrix_to_array() {
    let array = {
        let foo = Matrix4::from_array([[1.0, 0.0, 0.0, 0.0],
                                       [0.0, 1.0, 0.0, 0.0],
                                       [0.0, 0.0, 1.0, 0.0],
                                       [0.0, 0.0, 0.0, 1.0]]);
        assert_eq!(foo.as_array()[0][0], 1.0);
        *foo.as_array()
    };
    assert_eq!(array[1][1], 1.0);
}

#[test]
fn test_scale() {
    let mut m = Matrix4::identity();
    assert_eq!(m[2][2], 1.0);
    m.scale(3.0, 3.0, 3.0);
    assert_eq!(m[2][2], 3.0);
    m.scale_by_vector(&Vector4::new(2.0, 4.0, 3.0));
    assert_eq!(m[2][2], 9.0);
}

#[test]
fn test_translate() {
    let mut m = Matrix4::identity();
    m.translate(1.0, 2.0, 3.0);
    assert_eq!(m[0][2], 1.0);
}
<|endoftext|>


<commit_before>#![warn(missing_docs)]

//! `mudi` is a multi-dimensional array library.
//!
//! `mudi` provides multi-dimensional arrays to stora data.
//! The aim of this library is not to provide complex algorithms, but a
//! powerful indexing scheme and iterators.
//! For example, with `mudi` it is possible to use range dimensions
//! (`-42..42`) and the corresponding negative indexing (`array[-23]`).
//!
//! The `mudi` array type is closer to a Fortran array than to Numpy's
//! `ndarray`. The fact that negative indexing is supported means that python
//! style *indexing from the end* will not be supported.
//! If the dimension does not contain a negative range then negative indexing
//! is not permitted:
//!
//! ```no_run
//! use mudi::Array;
//! let array = Array::from_element(0.0, (2, 6));
//! // This fails to compile
//! array[(1, -1)];
//! ```
//!
//! If the dimensions contain a negative range, negative indexing
//! will yield the corresponding value
//!
//! ```
//! use mudi::Array;
//! let array = Array::from_element(0.0, (-42..42));
//! assert_eq!(array[-12], 0.0);
//! ```
//!
//! The basic type is `Array`, which provides an owned array for any kind of
//! data. The size of the array is fixed at initial construction.

mod dimensions;
pub use dimensions::Dimensions;

mod storage;
pub use storage::Storage;

mod arrays;
pub use arrays::{Array, ArrayBase};
<commit_msg>Fix doctest which should not compile<commit_after>#![warn(missing_docs)]

//! `mudi` is a multi-dimensional array library.
//!
//! `mudi` provides multi-dimensional arrays to stora data.
//! The aim of this library is not to provide complex algorithms, but a
//! powerful indexing scheme and iterators.
//! For example, with `mudi` it is possible to use range dimensions
//! (`-42..42`) and the corresponding negative indexing (`array[-23]`).
//!
//! The `mudi` array type is closer to a Fortran array than to Numpy's
//! `ndarray`. The fact that negative indexing is supported means that python
//! style *indexing from the end* will not be supported.
//! If the dimension does not contain a negative range then negative indexing
//! is not permitted:
//!
//! ```
//! use mudi::Array;
//! let array = Array::from_element(0.0, (2, 6));
//! // This fails to compile
//! // array[(1, -1)];
//! ```
//!
//! If the dimensions contain a negative range, negative indexing
//! will yield the corresponding value
//!
//! ```
//! use mudi::Array;
//! let array = Array::from_element(0.0, (-42..42));
//! assert_eq!(array[-12], 0.0);
//! ```
//!
//! The basic type is `Array`, which provides an owned array for any kind of
//! data. The size of the array is fixed at initial construction.

mod dimensions;
pub use dimensions::Dimensions;

mod storage;
pub use storage::Storage;

mod arrays;
pub use arrays::{Array, ArrayBase};
<|endoftext|>


<commit_before>#![cfg_attr(all(feature = "bench", test), feature(test))]

extern crate num;

mod butterflies;

use num::{Complex, Zero, One, Float, Num, FromPrimitive, Signed};
use num::traits::cast;
use std::f32;

use butterflies::{butterfly_2, butterfly_3, butterfly_4, butterfly_5};

pub struct FFT<T> {
    factors: Vec<(usize, usize)>,
    twiddles: Vec<Complex<T>>,
    scratch: Vec<Complex<T>>,
    inverse: bool,
}

impl<T> FFT<T> where T: Signed + FromPrimitive + Copy {
    /// Creates a new FFT context that will process signal of length
    /// `len`. If `inverse` is `true`, then this struct will run inverse
    /// FFTs. This implementation of the FFT doesn't do any scaling on both
    /// the forward and backward transforms, so doing a forward then backward
    /// FFT on a signal will scale the signal by its length.
    pub fn new(len: usize, inverse: bool) -> Self {
        let dir = if inverse { 1 } else { -1 };
        let factors = factor(len);
        let max_fft_len = factors.iter().map(|&(a, _)| a).max();
        let scratch = match max_fft_len {
            None | Some(0...5) => vec![Zero::zero(); 0],
            Some(l) => vec![Zero::zero(); l],
        };
        FFT {
            factors: factor(len),
            twiddles: (0..len)
                      .map(|i| dir as f32 * i as f32 * 2.0 * f32::consts::PI / len as f32)
                      .map(|phase| Complex::from_polar(&1.0, &phase))
                      .map(|c| Complex {re: FromPrimitive::from_f32(c.re).unwrap(),
                                        im: FromPrimitive::from_f32(c.im).unwrap()})
                      .collect(),
            scratch: scratch,
            inverse: inverse,
        }
    }

    /// Runs the FFT on the input `signal` buffer, and places the output in the
    /// `spectrum` buffer.
    ///
    /// # Panics
    /// This method will panic if `signal` and `spectrum` are not the length
    /// specified in the struct's constructor.
    pub fn process(&mut self, signal: &[Complex<T>], spectrum: &mut [Complex<T>]) {
        assert!(signal.len() == spectrum.len());
        assert!(signal.len() == self.twiddles.len());
        cooley_tukey(signal, spectrum, 1, &self.twiddles[..], &self.factors[..],
                     &mut self.scratch, self.inverse);
    }
}

fn cooley_tukey<T>(signal: &[Complex<T>],
                   spectrum: &mut [Complex<T>],
                   stride: usize,
                   twiddles: &[Complex<T>],
                   factors: &[(usize, usize)],
                   scratch: &mut [Complex<T>],
                   inverse: bool) where T: Signed + FromPrimitive + Copy {
    if let Some(&(n1, n2)) = factors.first() {
        if n2 == 1 {
            // An FFT of length 1 is just the identity operator
            let mut spectrum_idx = 0usize;
            let mut signal_idx = 0usize;
            while signal_idx < signal.len() {
                unsafe { *spectrum.get_unchecked_mut(spectrum_idx) =
                    *signal.get_unchecked(signal_idx); }
                spectrum_idx += 1;
                signal_idx += stride;
            }
        } else {
            // Recursive call to perform n1 ffts of length n2
            for i in 0..n1 {
                cooley_tukey(&signal[i * stride..],
                             &mut spectrum[i * n2..],
                             stride * n1, twiddles, &factors[1..],
                             scratch, inverse);
            }
        }

        match n1 {
            5 => unsafe { butterfly_5(spectrum, stride, twiddles, n2) },
            4 => unsafe { butterfly_4(spectrum, stride, twiddles, n2, inverse) },
            3 => unsafe { butterfly_3(spectrum, stride, twiddles, n2) },
            2 => unsafe { butterfly_2(spectrum, stride, twiddles, n2) },
            _ => butterfly(spectrum, stride, twiddles, n2, n1, &mut scratch[..n1]),
        }
    }
}

fn butterfly<T: Num + Copy>(data: &mut [Complex<T>], stride: usize,
                            twiddles: &[Complex<T>], num_ffts: usize,
                            fft_len: usize, scratch: &mut [Complex<T>]) {
    // for each fft we have to perform...
    for fft_idx in 0..num_ffts {

        // copy over data into scratch space
        let mut data_idx = fft_idx;
        for s in scratch.iter_mut() {
            *s = unsafe { *data.get_unchecked(data_idx) };
            data_idx += num_ffts;
        }

        // perfom the butterfly from the scratch space into the original buffer
        let mut data_idx = fft_idx;
        while data_idx < fft_len * num_ffts {
            let out_sample = unsafe { data.get_unchecked_mut(data_idx) };
            *out_sample = Zero::zero();
            let mut twiddle_idx = 0usize;
            for in_sample in scratch.iter() {
                let twiddle = unsafe { twiddles.get_unchecked(twiddle_idx) };
                *out_sample = *out_sample + in_sample * twiddle;
                twiddle_idx += stride * data_idx;
                if twiddle_idx >= twiddles.len() { twiddle_idx -= twiddles.len() }
            }
            data_idx += num_ffts;
        }

    }
}

pub fn dft<T: Float>(signal: &[Complex<T>], spectrum: &mut [Complex<T>]) {
    for (k, spec_bin) in spectrum.iter_mut().enumerate() {
        let mut sum = Zero::zero();
        for (i, &x) in signal.iter().enumerate() {
            let angle = cast::<_, T>(-1 * (i * k) as isize).unwrap()
                * cast(2.0 * f32::consts::PI).unwrap()
                / cast(signal.len()).unwrap();
            let twiddle = Complex::from_polar(&One::one(), &angle);
            sum = sum + twiddle * x;
        }
        *spec_bin = sum;
    }
}

/// Factors an integer into its prime factors.
fn factor(n: usize) -> Vec<(usize, usize)> {
    let mut factors = Vec::new();
    let mut next = n;
    while next > 1 {
        for div in 2..next + 1 {
            if next % div == 0 {
                next = next / div;
                factors.push((div, next));
                break;
            }
        }
    }
    return factors;
}
<commit_msg>Use radix-4 algorithm on power of 2 sizes<commit_after>#![cfg_attr(all(feature = "bench", test), feature(test))]

extern crate num;

mod butterflies;

use num::{Complex, Zero, One, Float, Num, FromPrimitive, Signed};
use num::traits::cast;
use std::f32;

use butterflies::{butterfly_2, butterfly_3, butterfly_4, butterfly_5};

enum Algorithm<T> {
    MixedRadix(Vec<(usize, usize)>, Vec<Complex<T>>),
    Radix4,
    Noop,
}

pub struct FFT<T> {
    algorithm: Algorithm<T>,
    twiddles: Vec<Complex<T>>,
    inverse: bool,
}

impl<T> FFT<T> where T: Signed + FromPrimitive + Copy {
    /// Creates a new FFT context that will process signal of length
    /// `len`. If `inverse` is `true`, then this struct will run inverse
    /// FFTs. This implementation of the FFT doesn't do any scaling on both
    /// the forward and backward transforms, so doing a forward then backward
    /// FFT on a signal will scale the signal by its length.
    pub fn new(len: usize, inverse: bool) -> Self {
        let dir = if inverse { 1 } else { -1 };

        let algorithm = if len < 2 {
            Algorithm::Noop
        } else if is_power_of_two(len) {
            Algorithm::Radix4
        } else {
            let factors = factor(len);
            let max_fft_len = factors.iter().map(|&(a, _)| a).max();
            let scratch = match max_fft_len {
                None | Some(0...5) => vec![Zero::zero(); 0],
                Some(l) => vec![Zero::zero(); l],
            };

            Algorithm::MixedRadix(factors, scratch)
        };

        FFT {
            algorithm: algorithm,
            twiddles: (0..len)
                      .map(|i| dir as f32 * i as f32 * 2.0 * f32::consts::PI / len as f32)
                      .map(|phase| Complex::from_polar(&1.0, &phase))
                      .map(|c| Complex {re: FromPrimitive::from_f32(c.re).unwrap(),
                                        im: FromPrimitive::from_f32(c.im).unwrap()})
                      .collect(),
            inverse: inverse,
        }
    }

    /// Runs the FFT on the input `signal` buffer, and places the output in the
    /// `spectrum` buffer.
    ///
    /// # Panics
    /// This method will panic if `signal` and `spectrum` are not the length
    /// specified in the struct's constructor.
    pub fn process(&mut self, signal: &[Complex<T>], spectrum: &mut [Complex<T>]) {
        assert!(signal.len() == spectrum.len());
        assert!(signal.len() == self.twiddles.len());

        match self.algorithm {
            Algorithm::Radix4 => radix_4(signal.len(), signal, spectrum, 1, &self.twiddles[..], self.inverse),
            Algorithm::MixedRadix(ref factors, ref mut scratch) =>
                cooley_tukey(signal, spectrum, 1, &self.twiddles[..], factors, scratch, self.inverse),
            Algorithm::Noop => copy_data(signal, spectrum, 1),
        }
    }
}

fn cooley_tukey<T>(signal: &[Complex<T>],
                   spectrum: &mut [Complex<T>],
                   stride: usize,
                   twiddles: &[Complex<T>],
                   factors: &[(usize, usize)],
                   scratch: &mut [Complex<T>],
                   inverse: bool) where T: Signed + FromPrimitive + Copy {
    if let Some(&(n1, n2)) = factors.first() {
        if n2 == 1 {
            // we theoretically need to compute n1 ffts of size n2
            // but n2 is 1, and a fft of size 1 is just a copy
            copy_data(signal, spectrum, stride);
        } else {
            // Recursive call to perform n1 ffts of length n2
            for i in 0..n1 {
                cooley_tukey(&signal[i * stride..],
                             &mut spectrum[i * n2..],
                             stride * n1, twiddles, &factors[1..],
                             scratch, inverse);
            }
        }

        match n1 {
            5 => unsafe { butterfly_5(spectrum, stride, twiddles, n2) },
            4 => unsafe { butterfly_4(spectrum, stride, twiddles, n2, inverse) },
            3 => unsafe { butterfly_3(spectrum, stride, twiddles, n2) },
            2 => unsafe { butterfly_2(spectrum, stride, twiddles, n2) },
            _ => butterfly(spectrum, stride, twiddles, n2, n1, &mut scratch[..n1]),
        }
    }
}

fn radix_4<T>(  size: usize,
                signal: &[Complex<T>],
                spectrum: &mut [Complex<T>],
                stride: usize,
                twiddles: &[Complex<T>],
                inverse: bool) where T: Signed + FromPrimitive + Copy {
    match size {
        4 => unsafe {
            copy_data_4(signal, spectrum, stride);
            butterfly_4(spectrum, stride, twiddles, 1, inverse)
        },
        2 => unsafe {
            copy_data_2(signal, spectrum, stride);
            butterfly_2(spectrum, stride, twiddles, 1)
        },
        _ => {
            for i in 0..4 {
                radix_4(
                    size / 4,
                    &signal[i * stride..],
                    &mut spectrum[i * (size / 4)..],
                    stride * 4, twiddles, inverse
                    );
            }
            unsafe { butterfly_4(spectrum, stride, twiddles, size/4, inverse) };
        }
    }
}

fn butterfly<T: Num + Copy>(data: &mut [Complex<T>], stride: usize,
                            twiddles: &[Complex<T>], num_ffts: usize,
                            fft_len: usize, scratch: &mut [Complex<T>]) {
    // for each fft we have to perform...
    for fft_idx in 0..num_ffts {

        // copy over data into scratch space
        let mut data_idx = fft_idx;
        for s in scratch.iter_mut() {
            *s = unsafe { *data.get_unchecked(data_idx) };
            data_idx += num_ffts;
        }

        // perfom the butterfly from the scratch space into the original buffer
        let mut data_idx = fft_idx;
        while data_idx < fft_len * num_ffts {
            let out_sample = unsafe { data.get_unchecked_mut(data_idx) };
            *out_sample = Zero::zero();
            let mut twiddle_idx = 0usize;
            for in_sample in scratch.iter() {
                let twiddle = unsafe { twiddles.get_unchecked(twiddle_idx) };
                *out_sample = *out_sample + in_sample * twiddle;
                twiddle_idx += stride * data_idx;
                if twiddle_idx >= twiddles.len() { twiddle_idx -= twiddles.len() }
            }
            data_idx += num_ffts;
        }

    }
}

fn copy_data<T>(signal: &[Complex<T>],
                spectrum: &mut [Complex<T>],
                stride: usize) where T: Copy  {
    let mut spectrum_idx = 0usize;
    let mut signal_idx = 0usize;
    while signal_idx < signal.len() {
        unsafe { *spectrum.get_unchecked_mut(spectrum_idx) =
            *signal.get_unchecked(signal_idx); }
        spectrum_idx += 1;
        signal_idx += stride;
    }
}

fn copy_data_2<T>(signal: &[Complex<T>],
                spectrum: &mut [Complex<T>],
                stride: usize) where T: Copy  {
    unsafe {
        *spectrum.get_unchecked_mut(0) = *signal.get_unchecked(0);
        *spectrum.get_unchecked_mut(1) = *signal.get_unchecked(stride);
    }
}
fn copy_data_4<T>(signal: &[Complex<T>],
                spectrum: &mut [Complex<T>],
                stride: usize) where T: Signed + FromPrimitive + Copy  {
    unsafe {
        for i in 0..4 {
            *spectrum.get_unchecked_mut(i) = *signal.get_unchecked(i*stride);
        }
    }
}

pub fn dft<T: Float>(signal: &[Complex<T>], spectrum: &mut [Complex<T>]) {
    for (k, spec_bin) in spectrum.iter_mut().enumerate() {
        let mut sum = Zero::zero();
        for (i, &x) in signal.iter().enumerate() {
            let angle = cast::<_, T>(-1 * (i * k) as isize).unwrap()
                * cast(2.0 * f32::consts::PI).unwrap()
                / cast(signal.len()).unwrap();
            let twiddle = Complex::from_polar(&One::one(), &angle);
            sum = sum + twiddle * x;
        }
        *spec_bin = sum;
    }
}

/// Factors an integer into its prime factors.
fn factor(n: usize) -> Vec<(usize, usize)> {
    let mut factors = Vec::new();
    let mut next = n;
    while next > 1 {
        for div in 2..next + 1 {
            if next % div == 0 {
                next = next / div;
                factors.push((div, next));
                break;
            }
        }
    }
    return factors;
}

//returns true if n is a power of 2, false otherwise
fn is_power_of_two(n: usize) -> bool {
    return n & n - 1 == 0;
}
<|endoftext|>