<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),

    pub interfaces: Vec<InterfaceDescriptionBlock>
}

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

impl<'a> Iterator<EnhancedPacketBlock> for PacketIter<'a> {
    fn next(&mut self) -> Option<EnhancedPacketBlock> {
        while let Ok(block) = read_block(self.r.r) {
            match block {
                Block::SectionHeader(_) => self.r.interfaces.clear(),
                Block::InterfaceDescription(id) => self.r.interfaces.push(id),
                Block::EnhancedPacket(ep) => return Some(ep),
                _ => {}
            }
        }

        None
    }
}

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

    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>Provide interface descriptor while iterating packets<commit_after>#![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
        })
    }
}
<|endoftext|>


<commit_before>//! BleepScript interpreter
//!
//! # Loading and Executing a Script
//!
//! ```
//! let mut bleep = Bleep::new();
//! bleep.load_script("test.tst").unwrap();
//! bleep.call_function("main", &[]).unwrap();
//! ```
//!
//! # Example Script
//!
//! ```text
//! function main() {
//!     printf("Hello, world!\n");
//!     return 42;
//! }
//! ```
//!

mod errors;
mod parser;
mod ast;
mod exec;
mod src_loc;
mod env;
mod sym_tab;
mod value;
mod native;

use std::rc::Rc;
use std::path;

pub use self::env::Env;
pub use self::errors::RunError;
pub use self::parser::ParseError;
pub use self::value::Value;

use self::src_loc::SrcLoc;
use self::sym_tab::SymTab;
use self::parser::{ops, Parser};

/// Script loader and executor.
///
/// # Example
///
/// ```
/// let mut bleep = Bleep::new();
///
/// bleep.load_script("test.tst").unwrap();
///
/// let ret = bleep.call_function("main", &[]).unwrap();
/// println!("function returned {}", ret);
/// ```
///
/// # Adding User Functions
///
/// To make a Rust function callable from a script, first create your Rust
/// function with a signature like this:
///
/// ```
/// fn my_function(args : &[Value], _env : &Rc<Env>) -> Result<Value,RunError> {
///     Ok(Value::Number(42.0))
/// }
/// ```
///
/// Then, use `Value::new_native_func()` to create a value for your function,
/// and `add_var()` to add it to the script global environment:
///
/// ```
/// fn main() {
///     let mut bleep = Bleep::new();
///
///     // add your function to the script environment
///     bleep.add_var("my_function", Value::new_native_func(my_function));
///
///     // load a script that calls your function
///     bleep.load_script("test.tst").unwrap();
///
///     // ...
/// }
/// ```
///
/// It's important to add your functions *before* loading a script that will
/// use them, or `load_script()` will complain about the function not existing.
///
pub struct Bleep {
    env : Rc<Env>,
    sym_tab : Rc<SymTab>,
    funcs : Vec<Rc<exec::FuncDef>>,
}

impl Bleep {
    /// Constructs a new script loader.
    pub fn new() -> Bleep {
        let mut bleep = Bleep {
            env : Rc::new(Env::new_global()),
            sym_tab : Rc::new(SymTab::new_global()),
            funcs : Vec::new(),
        };

        bleep.init_env();
        bleep
    }

    /// Loads a script file, adding its functions to the global environment.
    pub fn load_script<P: AsRef<path::Path>>(&mut self, filename : P) -> Result<(), ParseError> {
        let mut parser = self.create_parser();

        let ast_funcs = try!(parser.parse(filename));

        // add all function names to the environment, so the order of
        // function definitions doesn't matter
        for ast_func in &ast_funcs {
            self.set_var(&*ast_func.name, Value::Null);
            //println!("{:?}", ast_func);
        }

        // analyze each function and add the the result to the environment
        for ast_func in ast_funcs {
            let func = Rc::new(try!(ast_func.analyze(&self.sym_tab, &mut ast::analysis::State::new())));
            let closure = exec::FuncDef::eval(func.clone(), &self.env);
            self.set_var(&*ast_func.name, closure);
            self.funcs.push(func);
        }
        Ok(())
    }

    /// Calls the specified script function with the given arguments.
    ///
    /// The number of arguments must match the function definition, or an error will be returned.
    ///
    /// Returns the value returned by the function (`Value::Null` if the function didn't return
    /// a value), or any error encountered during execution.
    pub fn call_function(&self, func_name : &str, args : &[Value]) -> Result<Value, RunError> {
        let loc = SrcLoc::new("(no file)", 0, 0);
        let func = match self.get_var(func_name) {
            Some(f) => f,
            None => return Err(RunError::new_script(loc, &format!("function not found: '{}'", func_name))),
        };
        func.call(args, &self.env, &loc)
    }

    /// Sets the value of the given global variable to the given value.
    /// If the variable doesn't exist, it will be created.
    pub fn set_var(&mut self, var : &str, val : Value) {
        let index = self.sym_tab.add_name(var);
        if index >= self.env.size() {
            self.env.grow();
        }
        self.env.set_value(index, 0, val);
    }

    /// Returns the value of the given global variable, or
    /// `None` if the variable doesn't exist.
    pub fn get_var(&self, var : &str) -> Option<Value> {
        match self.sym_tab.get_name(var) {
            Some((vi, ei)) => Some(self.env.get_value(vi, ei)),
            None => None,
        }
    }

    fn init_env(&mut self) {
        self.set_var("null", Value::Null);
        self.set_var("true", Value::Bool(true));
        self.set_var("false", Value::Bool(false));

        self.set_var("printf", Value::new_native_func(native::func_printf));
        self.set_var("error", Value::new_native_func(native::func_error));
        self.set_var("dump_env", Value::new_native_func(native::func_dump_env));

        self.set_var("!",  Value::new_native_func(native::func_logic_not));
        self.set_var("==", Value::new_native_func(native::func_cmp_eq));
        self.set_var("!=", Value::new_native_func(native::func_cmp_ne));
        self.set_var("<",  Value::new_native_func(native::func_cmp_lt));
        self.set_var("<=", Value::new_native_func(native::func_cmp_le));
        self.set_var(">",  Value::new_native_func(native::func_cmp_gt));
        self.set_var(">=", Value::new_native_func(native::func_cmp_ge));
        self.set_var("+",  Value::new_native_func(native::func_num_add));
        self.set_var("-",  Value::new_native_func(native::func_num_sub));
        self.set_var("*",  Value::new_native_func(native::func_num_mul));
        self.set_var("/",  Value::new_native_func(native::func_num_div));
        self.set_var("^",  Value::new_native_func(native::func_num_pow));
        self.set_var("%",  Value::new_native_func(native::func_num_mod));
    }

    fn create_parser(&self) -> parser::Parser {
        let mut parser = Parser::new();

        parser.add_op("=",   10, ops::Assoc::Right);
        parser.add_op("||",  20, ops::Assoc::Left);
        parser.add_op("&&",  30, ops::Assoc::Left);
        parser.add_op("==",  40, ops::Assoc::Right);
        parser.add_op("!=",  40, ops::Assoc::Right);
        parser.add_op("<",   50, ops::Assoc::Left);
        parser.add_op(">",   50, ops::Assoc::Left);
        parser.add_op("<=",  50, ops::Assoc::Left);
        parser.add_op(">=",  50, ops::Assoc::Left);
        parser.add_op("+",   60, ops::Assoc::Left);
        parser.add_op("-",   60, ops::Assoc::Left);
        parser.add_op("*",   70, ops::Assoc::Left);
        parser.add_op("/",   70, ops::Assoc::Left);
        parser.add_op("%",   70, ops::Assoc::Left);
        parser.add_op("-",   80, ops::Assoc::Prefix);
        parser.add_op("!",   80, ops::Assoc::Prefix);
        parser.add_op("^",   90, ops::Assoc::Right);

        parser
    }

    /// Dumps the global environment (used for debugging).
    pub fn dump_env(&self) {
        println!("--- global environment ------------------------");
        self.sym_tab.dump_env(&self.env);
        println!("-----------------------------------------------");
    }

    /// Dumps all loaded functions (used for debugging).
    pub fn dump_funcs(&self) {
        println!("--- functions ---------------------------------");
        for func in &self.funcs {
            println!("{:?}", func);
            println!("");
        }
        println!("-----------------------------------------------");
    }

}
<commit_msg>Documentation fixes<commit_after>//! BleepScript interpreter
//!
//! # Loading and Executing a Script
//!
//! ```
//! let mut bleep = Bleep::new();
//! bleep.load_script("test.tst").unwrap();
//! bleep.call_function("main", &[]).unwrap();
//! ```
//!
//! # Example Script
//!
//! ```text
//! function main() {
//!     printf("Hello, world!\n");
//!     return 42;
//! }
//! ```
//!

mod errors;
mod parser;
mod ast;
mod exec;
mod src_loc;
mod env;
mod sym_tab;
mod value;
mod native;

use std::rc::Rc;
use std::path;

pub use self::env::Env;
pub use self::errors::RunError;
pub use self::parser::ParseError;
pub use self::value::Value;

use self::src_loc::SrcLoc;
use self::sym_tab::SymTab;
use self::parser::{ops, Parser};

/// Script loader and executor.
///
/// # Example
///
/// ```
/// let mut bleep = Bleep::new();
///
/// bleep.load_script("test.tst").unwrap();
///
/// let ret = bleep.call_function("main", &[]).unwrap();
/// println!("function returned {}", ret);
/// ```
///
/// # Adding User Functions
///
/// To make a Rust function callable from a script, first create your Rust
/// function with a signature like this:
///
/// ```
/// fn my_function(_args : &[Value], _env : &Rc<Env>) -> Result<Value,RunError> {
///     Ok(Value::Number(42.0))
/// }
/// ```
///
/// Then, use `Value::new_native_func()` to create a value for your function,
/// and `set_var()` to add it to the script global environment:
///
/// ```
/// fn main() {
///     let mut bleep = Bleep::new();
///
///     // add your function to the script environment
///     bleep.set_var("my_function", Value::new_native_func(my_function));
///
///     // load a script that calls your function
///     bleep.load_script("test.tst").unwrap();
///
///     // ...
/// }
/// ```
///
/// It's important to add your functions *before* loading a script that will
/// use them, or `load_script()` will complain about the function not existing.
///
pub struct Bleep {
    env : Rc<Env>,
    sym_tab : Rc<SymTab>,
    funcs : Vec<Rc<exec::FuncDef>>,
}

impl Bleep {
    /// Constructs a new script loader.
    pub fn new() -> Bleep {
        let mut bleep = Bleep {
            env : Rc::new(Env::new_global()),
            sym_tab : Rc::new(SymTab::new_global()),
            funcs : Vec::new(),
        };

        bleep.init_env();
        bleep
    }

    /// Loads a script file, adding its functions to the global environment.
    pub fn load_script<P: AsRef<path::Path>>(&mut self, filename : P) -> Result<(), ParseError> {
        let mut parser = self.create_parser();

        let ast_funcs = try!(parser.parse(filename));

        // add all function names to the environment, so the order of
        // function definitions doesn't matter
        for ast_func in &ast_funcs {
            self.set_var(&*ast_func.name, Value::Null);
            //println!("{:?}", ast_func);
        }

        // analyze each function and add the the result to the environment
        for ast_func in ast_funcs {
            let func = Rc::new(try!(ast_func.analyze(&self.sym_tab, &mut ast::analysis::State::new())));
            let closure = exec::FuncDef::eval(func.clone(), &self.env);
            self.set_var(&*ast_func.name, closure);
            self.funcs.push(func);
        }
        Ok(())
    }

    /// Calls the specified script function with the given arguments.
    ///
    /// The number of arguments must match the function definition, or an error will be returned.
    ///
    /// Returns the value returned by the function (`Value::Null` if the function didn't return
    /// a value), or any error encountered during execution.
    pub fn call_function(&self, func_name : &str, args : &[Value]) -> Result<Value, RunError> {
        let loc = SrcLoc::new("(no file)", 0, 0);
        let func = match self.get_var(func_name) {
            Some(f) => f,
            None => return Err(RunError::new_script(loc, &format!("function not found: '{}'", func_name))),
        };
        func.call(args, &self.env, &loc)
    }

    /// Sets the value of the given global variable to the given value.
    /// If the variable doesn't exist, it will be created.
    pub fn set_var(&mut self, var : &str, val : Value) {
        let index = self.sym_tab.add_name(var);
        if index >= self.env.size() {
            self.env.grow();
        }
        self.env.set_value(index, 0, val);
    }

    /// Returns the value of the given global variable, or
    /// `None` if the variable doesn't exist.
    pub fn get_var(&self, var : &str) -> Option<Value> {
        match self.sym_tab.get_name(var) {
            Some((vi, ei)) => Some(self.env.get_value(vi, ei)),
            None => None,
        }
    }

    fn init_env(&mut self) {
        self.set_var("null", Value::Null);
        self.set_var("true", Value::Bool(true));
        self.set_var("false", Value::Bool(false));

        self.set_var("printf", Value::new_native_func(native::func_printf));
        self.set_var("error", Value::new_native_func(native::func_error));
        self.set_var("dump_env", Value::new_native_func(native::func_dump_env));

        self.set_var("!",  Value::new_native_func(native::func_logic_not));
        self.set_var("==", Value::new_native_func(native::func_cmp_eq));
        self.set_var("!=", Value::new_native_func(native::func_cmp_ne));
        self.set_var("<",  Value::new_native_func(native::func_cmp_lt));
        self.set_var("<=", Value::new_native_func(native::func_cmp_le));
        self.set_var(">",  Value::new_native_func(native::func_cmp_gt));
        self.set_var(">=", Value::new_native_func(native::func_cmp_ge));
        self.set_var("+",  Value::new_native_func(native::func_num_add));
        self.set_var("-",  Value::new_native_func(native::func_num_sub));
        self.set_var("*",  Value::new_native_func(native::func_num_mul));
        self.set_var("/",  Value::new_native_func(native::func_num_div));
        self.set_var("^",  Value::new_native_func(native::func_num_pow));
        self.set_var("%",  Value::new_native_func(native::func_num_mod));
    }

    fn create_parser(&self) -> parser::Parser {
        let mut parser = Parser::new();

        parser.add_op("=",   10, ops::Assoc::Right);
        parser.add_op("||",  20, ops::Assoc::Left);
        parser.add_op("&&",  30, ops::Assoc::Left);
        parser.add_op("==",  40, ops::Assoc::Right);
        parser.add_op("!=",  40, ops::Assoc::Right);
        parser.add_op("<",   50, ops::Assoc::Left);
        parser.add_op(">",   50, ops::Assoc::Left);
        parser.add_op("<=",  50, ops::Assoc::Left);
        parser.add_op(">=",  50, ops::Assoc::Left);
        parser.add_op("+",   60, ops::Assoc::Left);
        parser.add_op("-",   60, ops::Assoc::Left);
        parser.add_op("*",   70, ops::Assoc::Left);
        parser.add_op("/",   70, ops::Assoc::Left);
        parser.add_op("%",   70, ops::Assoc::Left);
        parser.add_op("-",   80, ops::Assoc::Prefix);
        parser.add_op("!",   80, ops::Assoc::Prefix);
        parser.add_op("^",   90, ops::Assoc::Right);

        parser
    }

    /// Dumps the global environment (used for debugging).
    pub fn dump_env(&self) {
        println!("--- global environment ------------------------");
        self.sym_tab.dump_env(&self.env);
        println!("-----------------------------------------------");
    }

    /// Dumps all loaded functions (used for debugging).
    pub fn dump_funcs(&self) {
        println!("--- functions ---------------------------------");
        for func in &self.funcs {
            println!("{:?}", func);
            println!("");
        }
        println!("-----------------------------------------------");
    }

}
<|endoftext|>


<commit_before>#![crate_name = "neovim"]
#![crate_type = "lib"]
#![crate_type = "rlib"]

extern crate libc;

use std::c_str::CString;
use std::vec::Vec;

#[cfg(target_os="macos")]
mod platform {
    #[link(name = "nvim")]
    #[link(name = "uv")]
    #[link(name = "msgpack")]
    #[link(name = "curses")]
    #[link(name = "iconv")]
    extern{}
}

#[cfg(target_os="linux")]
mod platform {
    #[link(name = "nvim")]
    #[link(name = "uv")]
    #[link(name = "msgpack")]
    #[link(name = "curses")]
    extern{}
}

#[cfg(target_os="windows")]
mod platform {
    #[link(name = "nvim")]
    #[link(name = "uv")]
    #[link(name = "msgpack")]
    #[link(name = "curses")]
    extern{}
}

mod ffi {
    pub use libc::{c_char, c_int, uint64_t};

    extern "C" {
        pub fn nvim_main (argc: c_int, argv: *const *const c_char) -> c_int;
        pub fn channel_from_fds (read_fd: c_int, write_fd: c_int) -> uint64_t;
        pub fn channel_subscribe (id: uint64_t, event: *const c_char);
    }
}

pub fn nvim_main(args: &[&str]) -> i32 {
    let v: Vec<CString> = args.iter().map(|s| s.to_c_str()).collect();
    let vp: Vec<*const ffi::c_char> = v.iter().map(|s| s.as_ptr()).collect();
    let p_vp: *const *const ffi::c_char = vp.as_ptr();

    unsafe { ffi::nvim_main(vp.len() as i32, p_vp) }
}

pub struct Channel {
    id: ffi::uint64_t
}

impl Channel {
    pub fn new(read_fd: i32, write_fd: i32) -> Channel {
        Channel {
            id: unsafe { ffi::channel_from_fds(read_fd, write_fd) }
        }
    }

    pub fn subscribe(&mut self, event: &str) {
        unsafe { ffi::channel_subscribe(self.id, event.to_c_str().as_ptr()) }
    }
}
<commit_msg>Make nvim_main accept Vec<String> instead of &[&str] so you can pass os::args() into it<commit_after>#![crate_name = "neovim"]
#![crate_type = "lib"]
#![crate_type = "rlib"]

extern crate libc;

use std::c_str::CString;
use std::string::String;
use std::vec::Vec;

#[cfg(target_os="macos")]
mod platform {
    #[link(name = "nvim")]
    #[link(name = "uv")]
    #[link(name = "msgpack")]
    #[link(name = "curses")]
    #[link(name = "iconv")]
    extern{}
}

#[cfg(target_os="linux")]
mod platform {
    #[link(name = "nvim")]
    #[link(name = "uv")]
    #[link(name = "msgpack")]
    #[link(name = "curses")]
    extern{}
}

#[cfg(target_os="windows")]
mod platform {
    #[link(name = "nvim")]
    #[link(name = "uv")]
    #[link(name = "msgpack")]
    #[link(name = "curses")]
    extern{}
}

mod ffi {
    pub use libc::{c_char, c_int, uint64_t};

    extern "C" {
        pub fn channel_from_fds (read_fd: c_int, write_fd: c_int) -> uint64_t;
        pub fn channel_subscribe (id: uint64_t, event: *const c_char);
        pub fn nvim_main (argc: c_int, argv: *const *const c_char) -> c_int;
    }
}

pub struct Channel {
    id: ffi::uint64_t
}

impl Channel {
    pub fn new(read_fd: i32, write_fd: i32) -> Channel {
        Channel {
            id: unsafe { ffi::channel_from_fds(read_fd, write_fd) }
        }
    }

    pub fn subscribe(&mut self, event: &str) {
        unsafe { ffi::channel_subscribe(self.id, event.to_c_str().as_ptr()) }
    }
}

pub fn nvim_main(args: Vec<String>) -> i32 {
    let v: Vec<CString> = args.iter().map(|s| s.as_slice().to_c_str()).collect();
    let vp: Vec<*const ffi::c_char> = v.iter().map(|s| s.as_ptr()).collect();
    let p_vp: *const *const ffi::c_char = vp.as_ptr();

    unsafe { ffi::nvim_main(vp.len() as i32, p_vp) }
}
<|endoftext|>


<commit_before>#![feature(test)]
#[cfg(test)]
extern crate test;
#[cfg(test)]
extern crate rand;
#[cfg(test)]
extern crate quickcheck;

pub mod afsort {

    pub fn sort(vec: &mut [String]) {
        sort_req(vec, 0);
    }

    fn sort_req(vec: &mut [String], depth: usize) {
        if vec.len() <= 32 {
            vec.sort_unstable();
            return;
        }
        let mut min = u16::max_value();
        let mut max = 0u16;
        for elem in vec.iter() {
            if elem.len() > depth {
                let val = elem.as_bytes()[depth] as u16;
                if val < min {
                    min = val;
                }
                if val > max {
                    max = val;
                }
            }
        }

        let num_items = (max - min + 2) as u16;
        let mut counts: Vec<usize> = vec![0usize; num_items as usize];

        for elem in vec.iter() {
            let radix_val = radix_for_str(elem, depth, min);
            counts[radix_val as usize] += 1;
        }

        let mut sum = 0usize;
        let mut offsets: Vec<usize> = vec![0usize; num_items as usize];
        for i in 0..num_items {
            offsets[i as usize] = sum;
            sum += counts[i as usize];
        }
        let mut next_free = offsets.clone();

        let mut block = 0u16;
        let mut i = 0usize;
        while block < num_items - 1 {
            if i >= offsets[block as usize + 1] as usize {
                block += 1;
            } else {
                let radix_val = radix_for_str(&vec[i as usize], depth, min);
                if radix_val == block {
                    i += 1;
                } else {
                    vec.swap(i as usize, next_free[radix_val as usize] as usize);
                    next_free[radix_val as usize] += 1;
                }
            }
        }
        for i in 1..num_items - 1 {
            sort_req(
                &mut vec[offsets[i as usize] as usize..offsets[i as usize + 1] as usize],
                depth + 1,
            );
        }
        sort_req(&mut vec[offsets[num_items as usize - 1]..], depth + 1);
    }

    fn radix_for_str(s: &str, d: usize, base: u16) -> u16 {
        if s.as_bytes().len() > d {
            (s.as_bytes()[d] as u16 + 1 - base)
        } else {
            0
        }
    }
}

#[cfg(test)]
mod tests {

    use afsort;
    use quickcheck::QuickCheck;
    use test::Bencher;
    use std::fs::File;
    use std::io::{BufRead, BufReader};
    use rand::{self, Rng};

    #[test]
    fn sorts_same_as_unstable() {
        QuickCheck::new().tests(1000).quickcheck(
            compare_sort as
                fn(Vec<String>) -> bool,
        );
    }

    #[test]
    fn foobar() {
        assert!(compare_sort(
            vec![String::from("\u{40001}"), String::from("\u{80000}")],
        ))
    }

    fn compare_sort(mut strings: Vec<String>) -> bool {
        let mut copy = strings.clone();
        copy.sort_unstable();
        afsort::sort(&mut strings);
        strings == copy
    }

    #[bench]
    fn std_sort(b: &mut Bencher) {
        let strings = strings_en();
        b.iter(|| strings.clone().sort_unstable())
    }

    #[bench]
    fn af_sort(b: &mut Bencher) {
        let strings = strings_en();
        b.iter(|| afsort::sort(&mut strings.clone()))
    }

    fn strings_en() -> Vec<String> {
        let f = File::open("/usr/share/dict/american-english").unwrap();
        let b = BufReader::new(f);
        let mut strings = b.lines().map(|l| l.unwrap()).collect::<Vec<String>>();
        let mut rng = rand::thread_rng();
        rng.shuffle(&mut strings);
        strings.into_iter().take(10000).collect::<Vec<String>>()
    }

}
<commit_msg>Cleaned up tests, once again run 100k quickcheks.<commit_after>#![feature(test)]
#[cfg(test)]
extern crate test;
#[cfg(test)]
extern crate rand;
#[cfg(test)]
extern crate quickcheck;

pub mod afsort {

    pub fn sort(vec: &mut [String]) {
        sort_req(vec, 0);
    }

    fn sort_req(vec: &mut [String], depth: usize) {
        if vec.len() <= 32 {
            vec.sort_unstable();
            return;
        }
        let mut min = u16::max_value();
        let mut max = 0u16;
        for elem in vec.iter() {
            if elem.len() > depth {
                let val = elem.as_bytes()[depth] as u16;
                if val < min {
                    min = val;
                }
                if val > max {
                    max = val;
                }
            }
        }

        let num_items = (max - min + 2) as u16;
        let mut counts: Vec<usize> = vec![0usize; num_items as usize];

        for elem in vec.iter() {
            let radix_val = radix_for_str(elem, depth, min);
            counts[radix_val as usize] += 1;
        }

        let mut sum = 0usize;
        let mut offsets: Vec<usize> = vec![0usize; num_items as usize];
        for i in 0..num_items {
            offsets[i as usize] = sum;
            sum += counts[i as usize];
        }
        let mut next_free = offsets.clone();

        let mut block = 0u16;
        let mut i = 0usize;
        while block < num_items - 1 {
            if i >= offsets[block as usize + 1] as usize {
                block += 1;
            } else {
                let radix_val = radix_for_str(&vec[i as usize], depth, min);
                if radix_val == block {
                    i += 1;
                } else {
                    vec.swap(i as usize, next_free[radix_val as usize] as usize);
                    next_free[radix_val as usize] += 1;
                }
            }
        }
        for i in 1..num_items - 1 {
            sort_req(
                &mut vec[offsets[i as usize] as usize..offsets[i as usize + 1] as usize],
                depth + 1,
            );
        }
        sort_req(&mut vec[offsets[num_items as usize - 1]..], depth + 1);
    }

    fn radix_for_str(s: &str, d: usize, base: u16) -> u16 {
        if s.as_bytes().len() > d {
            (s.as_bytes()[d] as u16 + 1 - base)
        } else {
            0
        }
    }
}

#[cfg(test)]
mod tests {

    use afsort;
    use quickcheck::QuickCheck;
    use test::Bencher;
    use std::fs::File;
    use std::io::{BufRead, BufReader};
    use rand::{self, Rng};

    #[test]
    fn sorts_same_as_unstable() {
        QuickCheck::new().tests(100000).quickcheck(
            compare_sort as
                fn(Vec<String>) -> bool,
        );
    }

    fn compare_sort(mut strings: Vec<String>) -> bool {
        let mut copy = strings.clone();
        copy.sort_unstable();
        afsort::sort(&mut strings);
        strings == copy
    }

    #[bench]
    fn std_sort(b: &mut Bencher) {
        let strings = strings_en();
        b.iter(|| strings.clone().sort_unstable())
    }

    #[bench]
    fn af_sort(b: &mut Bencher) {
        let strings = strings_en();
        b.iter(|| afsort::sort(&mut strings.clone()))
    }

    fn strings_en() -> Vec<String> {
        let f = File::open("/usr/share/dict/american-english").unwrap();
        let b = BufReader::new(f);
        let mut strings = b.lines().map(|l| l.unwrap()).collect::<Vec<String>>();
        let mut rng = rand::thread_rng();
        rng.shuffle(&mut strings);
        strings.into_iter().take(10000).collect::<Vec<String>>()
    }

}
<|endoftext|>


<commit_before>#[macro_use(bson, doc)]
extern crate bson;
extern crate mongodb;
extern crate chrono;

use std::fmt;
use std::result;
use std::error;
use mongodb::{Client, ThreadedClient};
use mongodb::cursor::Cursor;
use mongodb::db::ThreadedDatabase;
use mongodb::coll::options::{FindOptions, CursorType};
use chrono::{DateTime, UTC, TimeZone};

#[derive(Debug)]
pub enum Error {
    MissingField(bson::ValueAccessError),
    Database(mongodb::Error),
    UnknownOperation(String),
}

impl From<bson::ValueAccessError> for Error {
    fn from(original: bson::ValueAccessError) -> Error {
        Error::MissingField(original)
    }
}

impl From<mongodb::Error> for Error {
    fn from(original: mongodb::Error) -> Error {
        Error::Database(original)
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::MissingField(ref err) => err.description(),
            Error::Database(ref err) => err.description(),
            Error::UnknownOperation(_) => "unknown operation type",
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::MissingField(ref err) => err.fmt(f),
            Error::Database(ref err) => err.fmt(f),
            Error::UnknownOperation(ref op) => write!(f, "Unknown operation type found: {}", op),
        }
    }
}

type Result<T> = result::Result<T, Error>;

#[derive(PartialEq, Debug)]
pub enum Operation {
    Noop {
        id: i64,
        timestamp: DateTime<UTC>,
        message: String,
    },
    Insert {
        id: i64,
        timestamp: DateTime<UTC>,
        namespace: String,
        document: bson::Document,
    },
    Update {
        id: i64,
        timestamp: DateTime<UTC>,
        namespace: String,
        query: bson::Document,
        update: bson::Document,
    },
    Delete {
        id: i64,
        timestamp: DateTime<UTC>,
        namespace: String,
        query: bson::Document,
    },
    Command {
        id: i64,
        timestamp: DateTime<UTC>,
        namespace: String,
        command: bson::Document,
    },
}

impl Operation {
    pub fn new(document: &bson::Document) -> Result<Operation> {
        let op = try!(document.get_str("op"));

        match op {
            "n" => noop(document),
            "i" => insert(document),
            "u" => update(document),
            "d" => delete(document),
            "c" => command(document),
            _ => Err(Error::UnknownOperation(op.to_owned())),
        }
    }
}

impl fmt::Display for Operation {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Operation::Noop { id, timestamp, ref message } => {
                write!(f, "No-op #{} at {}: {}", id, timestamp, message)
            }
            Operation::Insert { id, timestamp, ref namespace, ref document } => {
                write!(f,
                       "Insert #{} into {} at {}: {}",
                       id,
                       namespace,
                       timestamp,
                       document)
            }
            Operation::Update { id, timestamp, ref namespace, ref query, ref update } => {
                write!(f,
                       "Update #{} {} with {} at {}: {}",
                       id,
                       namespace,
                       query,
                       timestamp,
                       update)
            }
            Operation::Delete { id, timestamp, ref namespace, ref query } => {
                write!(f,
                       "Delete #{} from {} at {}: {}",
                       id,
                       namespace,
                       timestamp,
                       query)
            }
            Operation::Command { id, timestamp, ref namespace, ref command } => {
                write!(f,
                       "Command #{} {} at {}: {}",
                       id,
                       namespace,
                       timestamp,
                       command)
            }
        }
    }
}

fn noop(document: &bson::Document) -> Result<Operation> {
    let h = try!(document.get_i64("h"));
    let ts = try!(document.get_time_stamp("ts"));
    let o = try!(document.get_document("o"));
    let msg = try!(o.get_str("msg"));

    Ok(Operation::Noop {
        id: h,
        timestamp: timestamp_to_datetime(ts),
        message: msg.to_owned(),
    })
}

fn insert(document: &bson::Document) -> Result<Operation> {
    let h = try!(document.get_i64("h"));
    let ts = try!(document.get_time_stamp("ts"));
    let ns = try!(document.get_str("ns"));
    let o = try!(document.get_document("o"));

    Ok(Operation::Insert {
        id: h,
        timestamp: timestamp_to_datetime(ts),
        namespace: ns.to_owned(),
        document: o.to_owned(),
    })
}

fn update(document: &bson::Document) -> Result<Operation> {
    let h = try!(document.get_i64("h"));
    let ts = try!(document.get_time_stamp("ts"));
    let ns = try!(document.get_str("ns"));
    let o = try!(document.get_document("o"));
    let o2 = try!(document.get_document("o2"));

    Ok(Operation::Update {
        id: h,
        timestamp: timestamp_to_datetime(ts),
        namespace: ns.to_owned(),
        query: o2.to_owned(),
        update: o.to_owned(),
    })
}

fn delete(document: &bson::Document) -> Result<Operation> {
    let h = try!(document.get_i64("h"));
    let ts = try!(document.get_time_stamp("ts"));
    let ns = try!(document.get_str("ns"));
    let o = try!(document.get_document("o"));

    Ok(Operation::Delete {
        id: h,
        timestamp: timestamp_to_datetime(ts),
        namespace: ns.to_owned(),
        query: o.to_owned(),
    })
}

fn command(document: &bson::Document) -> Result<Operation> {
    let h = try!(document.get_i64("h"));
    let ts = try!(document.get_time_stamp("ts"));
    let ns = try!(document.get_str("ns"));
    let o = try!(document.get_document("o"));

    Ok(Operation::Command {
        id: h,
        timestamp: timestamp_to_datetime(ts),
        namespace: ns.to_owned(),
        command: o.to_owned(),
    })
}

fn timestamp_to_datetime(timestamp: i64) -> DateTime<UTC> {
    let seconds = timestamp >> 32;
    let nanoseconds = ((timestamp & 0xFFFFFFFF) * 1000000) as u32;

    UTC.timestamp(seconds, nanoseconds)
}

pub struct Oplog {
    cursor: Cursor,
}

impl Iterator for Oplog {
    type Item = Operation;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.cursor.next() {
                Some(Ok(document)) => return Operation::new(&document).ok(),
                _ => continue,
            }
        }
    }
}

impl Oplog {
    pub fn new(client: &Client) -> Result<Oplog> {
        let coll = client.db("local").collection("oplog.rs");

        let mut opts = FindOptions::new();
        opts.cursor_type = CursorType::TailableAwait;
        opts.no_cursor_timeout = true;

        let cursor = try!(coll.find(None, Some(opts)));

        Ok(Oplog { cursor: cursor })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bson::Bson;
    use bson::oid::ObjectId;
    use chrono::{UTC, TimeZone};

    #[test]
    fn operation_converts_noops() {
        let doc = doc! {
            "ts" => (Bson::TimeStamp(1479419535 << 32)),
            "h" => (-2135725856567446411i64),
            "v" => 2,
            "op" => "n",
            "ns" => "",
            "o" => {
                "msg" => "initiating set"
            }
        };

        match Operation::new(&doc) {
            Ok(Operation::Noop { id, timestamp, message }) => {
                assert_eq!(-2135725856567446411i64, id);
                assert_eq!(UTC.timestamp(1479419535, 0), timestamp);
                assert_eq!("initiating set", message);
            }
            _ => panic!("Unexpected operation"),
        }
    }

    #[test]
    fn operation_converts_inserts() {
        let oid = ObjectId::with_string("583050b26813716e505a5bf2").unwrap();
        let doc = doc! {
            "ts" => (Bson::TimeStamp(1479561394 << 32)),
            "h" => (-1742072865587022793i64),
            "v" => 2,
            "op" => "i",
            "ns" => "foo.bar",
            "o" => {
                "_id" => (Bson::ObjectId(oid)),
                "foo" => "bar"
            }
        };

        match Operation::new(&doc) {
            Ok(Operation::Insert { id, timestamp, namespace, document }) => {
                assert_eq!(-1742072865587022793i64, id);
                assert_eq!(UTC.timestamp(1479561394, 0), timestamp);
                assert_eq!("foo.bar", namespace);
                assert_eq!("bar", document.get_str("foo").expect("foo missing"));
            }
            _ => panic!("Unexpected type of operation"),
        }
    }

    #[test]
    fn operation_converts_updates() {
        let oid = ObjectId::with_string("583033a3643431ab5be6ec35").unwrap();
        let doc = doc! {
            "ts" => (Bson::TimeStamp(1479561033 << 32)),
            "h" => (3511341713062188019i64),
            "v" => 2,
            "op" => "u",
            "ns" => "foo.bar",
            "o2" => {
                "_id" => (Bson::ObjectId(oid))
            },
            "o" => {
                "$set" => {
                    "foo" => "baz"
                }
            }
        };

        match Operation::new(&doc) {
            Ok(Operation::Update { id, timestamp, namespace, query, update }) => {
                assert_eq!(3511341713062188019i64, id);
                assert_eq!(UTC.timestamp(1479561033, 0), timestamp);
                assert_eq!("foo.bar", namespace);
                assert_eq!(ObjectId::with_string("583033a3643431ab5be6ec35").unwrap(),
                           query.get_object_id("_id").expect("_id missing").to_owned());
                assert_eq!("baz",
                           update.get_document("$set").and_then(|o| o.get_str("foo")).unwrap());
            }
            _ => panic!("Unexpected type of operation"),
        }
    }

    #[test]
    fn operation_converts_deletes() {
        let oid = ObjectId::with_string("582e287cfedf6fb051b2efdf").unwrap();
        let doc = doc! {
            "ts" => (Bson::TimeStamp(1479421186 << 32)),
            "h" => (-5457382347563537847i64),
            "v" => 2,
            "op" => "d",
            "ns" => "foo.bar",
            "o" => {
                "_id" => (Bson::ObjectId(oid))
            }
        };

        match Operation::new(&doc) {
            Ok(Operation::Delete { id, timestamp, namespace, query }) => {
                assert_eq!(-5457382347563537847i64, id);
                assert_eq!(UTC.timestamp(1479421186, 0), timestamp);
                assert_eq!("foo.bar", namespace);
                assert_eq!(ObjectId::with_string("582e287cfedf6fb051b2efdf").unwrap(),
                           query.get_object_id("_id").expect("_id missing").to_owned());
            }
            _ => panic!("Unexpected type of operation"),
        }
    }

    #[test]
    fn operation_converts_commands() {
        let doc = doc! {
            "ts" => (Bson::TimeStamp(1479553955 << 32)),
            "h" => (-7222343681970774929i64),
            "v" => 2,
            "op" => "c",
            "ns" => "test.$cmd",
            "o" => {
                "create" => "foo"
            }
        };

        match Operation::new(&doc) {
            Ok(Operation::Command { id, timestamp, namespace, command }) => {
                assert_eq!(-7222343681970774929i64, id);
                assert_eq!(UTC.timestamp(1479553955, 0), timestamp);
                assert_eq!("test.$cmd", namespace);
                assert_eq!("foo", command.get_str("create").expect("create missing"));
            }
            _ => panic!("Unexpected type of operation"),
        }
    }
}
<commit_msg>Move the document when converting into Operations<commit_after>#[macro_use(bson, doc)]
extern crate bson;
extern crate mongodb;
extern crate chrono;

use std::fmt;
use std::result;
use std::error;
use mongodb::{Client, ThreadedClient};
use mongodb::cursor::Cursor;
use mongodb::db::ThreadedDatabase;
use mongodb::coll::options::{FindOptions, CursorType};
use chrono::{DateTime, UTC, TimeZone};

#[derive(Debug)]
pub enum Error {
    MissingField(bson::ValueAccessError),
    Database(mongodb::Error),
    UnknownOperation(String),
    InvalidOperation,
}

impl From<bson::ValueAccessError> for Error {
    fn from(original: bson::ValueAccessError) -> Error {
        Error::MissingField(original)
    }
}

impl From<mongodb::Error> for Error {
    fn from(original: mongodb::Error) -> Error {
        Error::Database(original)
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::MissingField(ref err) => err.description(),
            Error::Database(ref err) => err.description(),
            Error::UnknownOperation(_) => "unknown operation type",
            Error::InvalidOperation => "invalid operation",
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::MissingField(ref err) => err.fmt(f),
            Error::Database(ref err) => err.fmt(f),
            Error::UnknownOperation(ref op) => write!(f, "Unknown operation type found: {}", op),
            Error::InvalidOperation => write!(f, "Invalid operation"),
        }
    }
}

type Result<T> = result::Result<T, Error>;

#[derive(PartialEq, Debug)]
pub enum Operation {
    Noop {
        id: i64,
        timestamp: DateTime<UTC>,
        message: String,
    },
    Insert {
        id: i64,
        timestamp: DateTime<UTC>,
        namespace: String,
        document: bson::Document,
    },
    Update {
        id: i64,
        timestamp: DateTime<UTC>,
        namespace: String,
        query: bson::Document,
        update: bson::Document,
    },
    Delete {
        id: i64,
        timestamp: DateTime<UTC>,
        namespace: String,
        query: bson::Document,
    },
    Command {
        id: i64,
        timestamp: DateTime<UTC>,
        namespace: String,
        command: bson::Document,
    },
}

impl Operation {
    pub fn new(document: bson::Document) -> Result<Operation> {
        let op = operation(&document);

        match op {
            Some('n') => noop(document),
            Some('i') => insert(document),
            Some('u') => update(document),
            Some('d') => delete(document),
            Some('c') => command(document),
            Some(unknown) => Err(Error::UnknownOperation(unknown.to_string())),
            None => Err(Error::InvalidOperation),
        }
    }
}

fn operation(document: &bson::Document) -> Option<char> {
    document.get_str("op").ok().and_then(|op| op.chars().next())
}

impl fmt::Display for Operation {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Operation::Noop { id, timestamp, ref message } => {
                write!(f, "No-op #{} at {}: {}", id, timestamp, message)
            }
            Operation::Insert { id, timestamp, ref namespace, ref document } => {
                write!(f,
                       "Insert #{} into {} at {}: {}",
                       id,
                       namespace,
                       timestamp,
                       document)
            }
            Operation::Update { id, timestamp, ref namespace, ref query, ref update } => {
                write!(f,
                       "Update #{} {} with {} at {}: {}",
                       id,
                       namespace,
                       query,
                       timestamp,
                       update)
            }
            Operation::Delete { id, timestamp, ref namespace, ref query } => {
                write!(f,
                       "Delete #{} from {} at {}: {}",
                       id,
                       namespace,
                       timestamp,
                       query)
            }
            Operation::Command { id, timestamp, ref namespace, ref command } => {
                write!(f,
                       "Command #{} {} at {}: {}",
                       id,
                       namespace,
                       timestamp,
                       command)
            }
        }
    }
}

fn noop(document: bson::Document) -> Result<Operation> {
    let h = try!(document.get_i64("h"));
    let ts = try!(document.get_time_stamp("ts"));
    let o = try!(document.get_document("o"));
    let msg = try!(o.get_str("msg"));

    Ok(Operation::Noop {
        id: h,
        timestamp: timestamp_to_datetime(ts),
        message: msg.to_owned(),
    })
}

fn insert(document: bson::Document) -> Result<Operation> {
    let h = try!(document.get_i64("h"));
    let ts = try!(document.get_time_stamp("ts"));
    let ns = try!(document.get_str("ns"));
    let o = try!(document.get_document("o"));

    Ok(Operation::Insert {
        id: h,
        timestamp: timestamp_to_datetime(ts),
        namespace: ns.to_owned(),
        document: o.to_owned(),
    })
}

fn update(document: bson::Document) -> Result<Operation> {
    let h = try!(document.get_i64("h"));
    let ts = try!(document.get_time_stamp("ts"));
    let ns = try!(document.get_str("ns"));
    let o = try!(document.get_document("o"));
    let o2 = try!(document.get_document("o2"));

    Ok(Operation::Update {
        id: h,
        timestamp: timestamp_to_datetime(ts),
        namespace: ns.to_owned(),
        query: o2.to_owned(),
        update: o.to_owned(),
    })
}

fn delete(document: bson::Document) -> Result<Operation> {
    let h = try!(document.get_i64("h"));
    let ts = try!(document.get_time_stamp("ts"));
    let ns = try!(document.get_str("ns"));
    let o = try!(document.get_document("o"));

    Ok(Operation::Delete {
        id: h,
        timestamp: timestamp_to_datetime(ts),
        namespace: ns.to_owned(),
        query: o.to_owned(),
    })
}

fn command(document: bson::Document) -> Result<Operation> {
    let h = try!(document.get_i64("h"));
    let ts = try!(document.get_time_stamp("ts"));
    let ns = try!(document.get_str("ns"));
    let o = try!(document.get_document("o"));

    Ok(Operation::Command {
        id: h,
        timestamp: timestamp_to_datetime(ts),
        namespace: ns.to_owned(),
        command: o.to_owned(),
    })
}

fn timestamp_to_datetime(timestamp: i64) -> DateTime<UTC> {
    let seconds = timestamp >> 32;
    let nanoseconds = ((timestamp & 0xFFFFFFFF) * 1000000) as u32;

    UTC.timestamp(seconds, nanoseconds)
}

pub struct Oplog {
    cursor: Cursor,
}

impl Iterator for Oplog {
    type Item = Operation;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.cursor.next() {
                Some(Ok(document)) => return Operation::new(document).ok(),
                _ => continue,
            }
        }
    }
}

impl Oplog {
    pub fn new(client: &Client) -> Result<Oplog> {
        let coll = client.db("local").collection("oplog.rs");

        let mut opts = FindOptions::new();
        opts.cursor_type = CursorType::TailableAwait;
        opts.no_cursor_timeout = true;

        let cursor = try!(coll.find(None, Some(opts)));

        Ok(Oplog { cursor: cursor })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bson::Bson;
    use bson::oid::ObjectId;
    use chrono::{UTC, TimeZone};

    #[test]
    fn operation_converts_noops() {
        let doc = doc! {
            "ts" => (Bson::TimeStamp(1479419535 << 32)),
            "h" => (-2135725856567446411i64),
            "v" => 2,
            "op" => "n",
            "ns" => "",
            "o" => {
                "msg" => "initiating set"
            }
        };

        match Operation::new(doc) {
            Ok(Operation::Noop { id, timestamp, message }) => {
                assert_eq!(-2135725856567446411i64, id);
                assert_eq!(UTC.timestamp(1479419535, 0), timestamp);
                assert_eq!("initiating set", message);
            }
            _ => panic!("Unexpected operation"),
        }
    }

    #[test]
    fn operation_converts_inserts() {
        let oid = ObjectId::with_string("583050b26813716e505a5bf2").unwrap();
        let doc = doc! {
            "ts" => (Bson::TimeStamp(1479561394 << 32)),
            "h" => (-1742072865587022793i64),
            "v" => 2,
            "op" => "i",
            "ns" => "foo.bar",
            "o" => {
                "_id" => (Bson::ObjectId(oid)),
                "foo" => "bar"
            }
        };

        match Operation::new(doc) {
            Ok(Operation::Insert { id, timestamp, namespace, document }) => {
                assert_eq!(-1742072865587022793i64, id);
                assert_eq!(UTC.timestamp(1479561394, 0), timestamp);
                assert_eq!("foo.bar", namespace);
                assert_eq!("bar", document.get_str("foo").expect("foo missing"));
            }
            _ => panic!("Unexpected type of operation"),
        }
    }

    #[test]
    fn operation_converts_updates() {
        let oid = ObjectId::with_string("583033a3643431ab5be6ec35").unwrap();
        let doc = doc! {
            "ts" => (Bson::TimeStamp(1479561033 << 32)),
            "h" => (3511341713062188019i64),
            "v" => 2,
            "op" => "u",
            "ns" => "foo.bar",
            "o2" => {
                "_id" => (Bson::ObjectId(oid))
            },
            "o" => {
                "$set" => {
                    "foo" => "baz"
                }
            }
        };

        match Operation::new(doc) {
            Ok(Operation::Update { id, timestamp, namespace, query, update }) => {
                assert_eq!(3511341713062188019i64, id);
                assert_eq!(UTC.timestamp(1479561033, 0), timestamp);
                assert_eq!("foo.bar", namespace);
                assert_eq!(ObjectId::with_string("583033a3643431ab5be6ec35").unwrap(),
                           query.get_object_id("_id").expect("_id missing").to_owned());
                assert_eq!("baz",
                           update.get_document("$set").and_then(|o| o.get_str("foo")).unwrap());
            }
            _ => panic!("Unexpected type of operation"),
        }
    }

    #[test]
    fn operation_converts_deletes() {
        let oid = ObjectId::with_string("582e287cfedf6fb051b2efdf").unwrap();
        let doc = doc! {
            "ts" => (Bson::TimeStamp(1479421186 << 32)),
            "h" => (-5457382347563537847i64),
            "v" => 2,
            "op" => "d",
            "ns" => "foo.bar",
            "o" => {
                "_id" => (Bson::ObjectId(oid))
            }
        };

        match Operation::new(doc) {
            Ok(Operation::Delete { id, timestamp, namespace, query }) => {
                assert_eq!(-5457382347563537847i64, id);
                assert_eq!(UTC.timestamp(1479421186, 0), timestamp);
                assert_eq!("foo.bar", namespace);
                assert_eq!(ObjectId::with_string("582e287cfedf6fb051b2efdf").unwrap(),
                           query.get_object_id("_id").expect("_id missing").to_owned());
            }
            _ => panic!("Unexpected type of operation"),
        }
    }

    #[test]
    fn operation_converts_commands() {
        let doc = doc! {
            "ts" => (Bson::TimeStamp(1479553955 << 32)),
            "h" => (-7222343681970774929i64),
            "v" => 2,
            "op" => "c",
            "ns" => "test.$cmd",
            "o" => {
                "create" => "foo"
            }
        };

        match Operation::new(doc) {
            Ok(Operation::Command { id, timestamp, namespace, command }) => {
                assert_eq!(-7222343681970774929i64, id);
                assert_eq!(UTC.timestamp(1479553955, 0), timestamp);
                assert_eq!("test.$cmd", namespace);
                assert_eq!("foo", command.get_str("create").expect("create missing"));
            }
            _ => panic!("Unexpected type of operation"),
        }
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]
pub mod p1;
pub mod p2;
pub mod p3;
pub mod p4;
<commit_msg>p5 ok<commit_after>#![allow(dead_code)]
pub mod p1;
pub mod p2;
pub mod p3;
pub mod p4;
pub mod p5;
<|endoftext|>


<commit_before>//! Setup and control loop devices.
//!
//! Provides rust interface with similar functionality to the Linux utility `losetup`.
//!
//! # Examples
//!
//! ```rust
//! use loopdev::LoopControl;
//! let lc = LoopControl::open().unwrap();
//! let ld = lc.next_free().unwrap();
//!
//! println!("{}", ld.path().unwrap().display());
//!
//! ld.attach_file("test.img").unwrap();
//! // ...
//! ld.detach().unwrap();
//! ```

extern crate libc;

use std::fs::File;
use std::fs::OpenOptions;

use libc::{c_int, ioctl};
use std::default::Default;
use std::io;
use std::os::unix::prelude::*;
use std::path::{Path, PathBuf};

// TODO support missing operations
const LOOP_SET_FD: u16 = 0x4C00;
const LOOP_CLR_FD: u16 = 0x4C01;
const LOOP_SET_STATUS64: u16 = 0x4C04;
//const LOOP_GET_STATUS64: u16 = 0x4C05;
const LOOP_SET_CAPACITY: u16 = 0x4C07;
//const LOOP_SET_DIRECT_IO: u16 = 0x4C08;
//const LOOP_SET_BLOCK_SIZE: u16 = 0x4C09;

//const LOOP_CTL_ADD: u16 = 0x4C80;
//const LOOP_CTL_REMOVE: u16 = 0x4C81;
const LOOP_CTL_GET_FREE: u16 = 0x4C82;

const LOOP_CONTROL: &str = "/dev/loop-control";
const LOOP_PREFIX: &str = "/dev/loop";

/// Interface to the loop control device: `/dev/loop-control`.
#[derive(Debug)]
pub struct LoopControl {
    dev_file: File,
}

impl LoopControl {
    /// Opens the loop control device.
    pub fn open() -> io::Result<Self> {
        Ok(Self {
            dev_file: OpenOptions::new()
                .read(true)
                .write(true)
                .open(LOOP_CONTROL)?,
        })
    }

    /// Finds and opens the next available loop device.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use loopdev::LoopControl;
    /// let lc = LoopControl::open().unwrap();
    /// let ld = lc.next_free().unwrap();
    /// println!("{}", ld.path().unwrap().display());
    /// ```
    pub fn next_free(&self) -> io::Result<LoopDevice> {
        let dev_num = ioctl_to_error(unsafe {
            ioctl(self.dev_file.as_raw_fd() as c_int, LOOP_CTL_GET_FREE.into())
        })?;
        LoopDevice::open(&format!("{}{}", LOOP_PREFIX, dev_num))
    }
}

/// Interface to a loop device ie `/dev/loop0`.
#[derive(Debug)]
pub struct LoopDevice {
    device: File,
}

impl AsRawFd for LoopDevice {
    fn as_raw_fd(&self) -> RawFd {
        self.device.as_raw_fd()
    }
}

impl LoopDevice {
    /// Opens a loop device.
    pub fn open<P: AsRef<Path>>(dev: P) -> io::Result<Self> {
        // TODO create dev if it does not exist and begins with LOOP_PREFIX
        Ok(Self {
            device: OpenOptions::new().read(true).write(true).open(dev)?,
        })
    }

    /// Attach the loop device to a file starting at offset into the file.
    #[deprecated(
        since = "0.2.0",
        note = "use `attach_file` or `attach_with_offset` instead"
    )]
    pub fn attach<P: AsRef<Path>>(&self, backing_file: P, offset: u64) -> io::Result<()> {
        self.attach_with_sizelimit(backing_file, offset, 0)
    }

    /// Attach the loop device to a file that maps to the whole file.
    ///
    /// # Examples
    ///
    /// Attach the device to a file.
    ///
    /// ```rust
    /// use loopdev::LoopDevice;
    /// let ld = LoopDevice::open("/dev/loop4").unwrap();
    /// ld.attach_file("test.img").unwrap();
    /// # ld.detach().unwrap();
    /// ```
    pub fn attach_file<P: AsRef<Path>>(&self, backing_file: P) -> io::Result<()> {
        self.attach_with_sizelimit(backing_file, 0, 0)
    }

    /// Attach the loop device to a file starting at offset into the file.
    ///
    /// # Examples
    ///
    /// Attach the device to the start of a file.
    ///
    /// ```rust
    /// use loopdev::LoopDevice;
    /// let ld = LoopDevice::open("/dev/loop5").unwrap();
    /// ld.attach_with_offset("test.img", 0).unwrap();
    /// # ld.detach().unwrap();
    /// ```
    pub fn attach_with_offset<P: AsRef<Path>>(
        &self,
        backing_file: P,
        offset: u64,
    ) -> io::Result<()> {
        self.attach_with_sizelimit(backing_file, offset, 0)
    }

    /// Attach the loop device to a file starting at offset into the file and a the given sizelimit.
    ///
    /// # Examples
    ///
    /// Attach the device to the start of a file with a maximum size of 1024 bytes.
    ///
    /// ```rust
    /// use loopdev::LoopDevice;
    /// let ld = LoopDevice::open("/dev/loop6").unwrap();
    /// ld.attach_with_sizelimit("test.img", 0, 1024).unwrap();
    /// # ld.detach().unwrap();
    /// ```
    pub fn attach_with_sizelimit<P: AsRef<Path>>(
        &self,
        backing_file: P,
        offset: u64,
        sizelimit: u64,
    ) -> io::Result<()> {
        let bf = OpenOptions::new()
            .read(true)
            .write(true)
            .open(backing_file)?;

        // Attach the file
        ioctl_to_error(unsafe {
            ioctl(
                self.device.as_raw_fd() as c_int,
                LOOP_SET_FD.into(),
                bf.as_raw_fd() as c_int,
            )
        })?;

        // Set offset for backing_file
        let mut info = loop_info64::default();
        info.lo_offset = offset;
        info.lo_sizelimit = sizelimit;

        if let Err(err) = ioctl_to_error(unsafe {
            ioctl(
                self.device.as_raw_fd() as c_int,
                LOOP_SET_STATUS64.into(),
                &mut info,
            )
        }) {
            // Ignore the error to preserve the original error
            let _ = self.detach();
            return Err(err);
        }
        Ok(())
    }

    /// Get the path of the loop device.
    #[deprecated(since = "0.2.0", note = "use `path` instead")]
    pub fn get_path(&self) -> Option<PathBuf> {
        self.path()
    }

    /// Get the path of the loop device.
    pub fn path(&self) -> Option<PathBuf> {
        let mut p = PathBuf::from("/proc/self/fd");
        p.push(self.device.as_raw_fd().to_string());
        std::fs::read_link(&p).ok()
    }

    /// Detach a loop device from its backing file.
    ///
    /// Note that the device won't fully detach until a short delay after the underling device file
    /// gets closed. This happens when LoopDev goes out of scope so you should ensure the LoopDev
    /// lives for a short a time as possible.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use loopdev::LoopDevice;
    /// let ld = LoopDevice::open("/dev/loop7").unwrap();
    /// # ld.attach_file("test.img").unwrap();
    /// ld.detach().unwrap();
    /// ```
    pub fn detach(&self) -> io::Result<()> {
        ioctl_to_error(unsafe { ioctl(self.device.as_raw_fd() as c_int, LOOP_CLR_FD.into(), 0) })?;
        Ok(())
    }

    /// Resize a live loop device. If the size of the backing file changes this can be called to
    /// inform the loop driver about the new size.
    pub fn set_capacity(&self) -> io::Result<()> {
        ioctl_to_error(unsafe {
            ioctl(
                self.device.as_raw_fd() as c_int,
                LOOP_SET_CAPACITY.into(),
                0,
            )
        })?;
        Ok(())
    }
}

#[repr(C)]
struct loop_info64 {
    pub lo_device: u64,
    pub lo_inode: u64,
    pub lo_rdevice: u64,
    pub lo_offset: u64,
    pub lo_sizelimit: u64,
    pub lo_number: u32,
    pub lo_encrypt_type: u32,
    pub lo_encrypt_key_size: u32,
    pub lo_flags: u32,
    pub lo_file_name: [u8; 64],
    pub lo_crypt_name: [u8; 64],
    pub lo_encrypt_key: [u8; 32],
    pub lo_init: [u64; 2],
}

impl Default for loop_info64 {
    fn default() -> Self {
        Self {
            lo_device: 0,
            lo_inode: 0,
            lo_rdevice: 0,
            lo_offset: 0,
            lo_sizelimit: 0,
            lo_number: 0,
            lo_encrypt_type: 0,
            lo_encrypt_key_size: 0,
            lo_flags: 0,
            lo_file_name: [0; 64],
            lo_crypt_name: [0; 64],
            lo_encrypt_key: [0; 32],
            lo_init: [0; 2],
        }
    }
}

fn ioctl_to_error(ret: i32) -> io::Result<i32> {
    if ret < 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(ret)
    }
}
<commit_msg>Adds automatic partition scanning function<commit_after>//! Setup and control loop devices.
//!
//! Provides rust interface with similar functionality to the Linux utility `losetup`.
//!
//! # Examples
//!
//! ```rust
//! use loopdev::LoopControl;
//! let lc = LoopControl::open().unwrap();
//! let ld = lc.next_free().unwrap();
//!
//! println!("{}", ld.path().unwrap().display());
//!
//! ld.attach_file("test.img").unwrap();
//! // ...
//! ld.detach().unwrap();
//! ```

extern crate libc;

use std::fs::File;
use std::fs::OpenOptions;

use libc::{c_int, ioctl};
use std::default::Default;
use std::io;
use std::os::unix::prelude::*;
use std::path::{Path, PathBuf};

// TODO support missing operations
const LOOP_SET_FD: u16 = 0x4C00;
const LOOP_CLR_FD: u16 = 0x4C01;
const LOOP_SET_STATUS64: u16 = 0x4C04;
//const LOOP_GET_STATUS64: u16 = 0x4C05;
const LOOP_SET_CAPACITY: u16 = 0x4C07;
//const LOOP_SET_DIRECT_IO: u16 = 0x4C08;
//const LOOP_SET_BLOCK_SIZE: u16 = 0x4C09;

//const LOOP_CTL_ADD: u16 = 0x4C80;
//const LOOP_CTL_REMOVE: u16 = 0x4C81;
const LOOP_CTL_GET_FREE: u16 = 0x4C82;

const LOOP_CONTROL: &str = "/dev/loop-control";
const LOOP_PREFIX: &str = "/dev/loop";

/// Interface to the loop control device: `/dev/loop-control`.
#[derive(Debug)]
pub struct LoopControl {
    dev_file: File,
}

impl LoopControl {
    /// Opens the loop control device.
    pub fn open() -> io::Result<Self> {
        Ok(Self {
            dev_file: OpenOptions::new()
                .read(true)
                .write(true)
                .open(LOOP_CONTROL)?,
        })
    }

    /// Finds and opens the next available loop device.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use loopdev::LoopControl;
    /// let lc = LoopControl::open().unwrap();
    /// let ld = lc.next_free().unwrap();
    /// println!("{}", ld.path().unwrap().display());
    /// ```
    pub fn next_free(&self) -> io::Result<LoopDevice> {
        let dev_num = ioctl_to_error(unsafe {
            ioctl(self.dev_file.as_raw_fd() as c_int, LOOP_CTL_GET_FREE.into())
        })?;
        LoopDevice::open(&format!("{}{}", LOOP_PREFIX, dev_num))
    }
}

/// Interface to a loop device ie `/dev/loop0`.
#[derive(Debug)]
pub struct LoopDevice {
    device: File,
}

impl AsRawFd for LoopDevice {
    fn as_raw_fd(&self) -> RawFd {
        self.device.as_raw_fd()
    }
}

impl LoopDevice {
    /// Opens a loop device.
    pub fn open<P: AsRef<Path>>(dev: P) -> io::Result<Self> {
        // TODO create dev if it does not exist and begins with LOOP_PREFIX
        Ok(Self {
            device: OpenOptions::new().read(true).write(true).open(dev)?,
        })
    }

    /// Attach the loop device to a file starting at offset into the file.
    #[deprecated(
        since = "0.2.0",
        note = "use `attach_file` or `attach_with_offset` instead"
    )]
    pub fn attach<P: AsRef<Path>>(&self, backing_file: P, offset: u64) -> io::Result<()> {
        self.attach_with_sizelimit(backing_file, offset, 0)
    }

    /// Attach the loop device to a file that maps to the whole file.
    ///
    /// # Examples
    ///
    /// Attach the device to a file.
    ///
    /// ```rust
    /// use loopdev::LoopDevice;
    /// let ld = LoopDevice::open("/dev/loop4").unwrap();
    /// ld.attach_file("test.img").unwrap();
    /// # ld.detach().unwrap();
    /// ```
    pub fn attach_file<P: AsRef<Path>>(&self, backing_file: P) -> io::Result<()> {
        self.attach_with_sizelimit(backing_file, 0, 0)
    }

    /// Attach the loop device to a file starting at offset into the file.
    ///
    /// # Examples
    ///
    /// Attach the device to the start of a file.
    ///
    /// ```rust
    /// use loopdev::LoopDevice;
    /// let ld = LoopDevice::open("/dev/loop5").unwrap();
    /// ld.attach_with_offset("test.img", 0).unwrap();
    /// # ld.detach().unwrap();
    /// ```
    pub fn attach_with_offset<P: AsRef<Path>>(
        &self,
        backing_file: P,
        offset: u64,
    ) -> io::Result<()> {
        self.attach_with_sizelimit(backing_file, offset, 0)
    }

    /// Attach the loop device to a file starting at offset into the file and a the given sizelimit.
    ///
    /// # Examples
    ///
    /// Attach the device to the start of a file with a maximum size of 1024 bytes.
    ///
    /// ```rust
    /// use loopdev::LoopDevice;
    /// let ld = LoopDevice::open("/dev/loop6").unwrap();
    /// ld.attach_with_sizelimit("test.img", 0, 1024).unwrap();
    /// # ld.detach().unwrap();
    /// ```
    pub fn attach_with_sizelimit<P: AsRef<Path>>(
        &self,
        backing_file: P,
        offset: u64,
        sizelimit: u64,
    ) -> io::Result<()> {
        // Set offset for backing_file
        let mut info = loop_info64::default();
        info.lo_offset = offset;
        info.lo_sizelimit = sizelimit;

        Self::attach_with_loop_info(self, backing_file, &mut info)
    }

    /// Attach the loop device with automatic partition scanning.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use loopdev::LoopDevice;
    /// let ld = LoopDevice::open("/dev/loop5").unwrap();
    /// ld.attach_with_partscan("test.img").unwrap();
    /// # ld.detach().unwrap();
    /// ```
    pub fn attach_with_partscan<P: AsRef<Path>>(
        &self,
        backing_file: P,
    ) -> io::Result<()> {
        // Set lo_flags to LO_FLAGS_PARTSCAN value
        let mut info = loop_info64::default();
        info.lo_flags = 8;

        Self::attach_with_loop_info(self, backing_file, &mut info)
    }

    /// Attach the loop device to a file with loop_info.
    ///
    /// # Examples
    ///
    /// Attach the device with default loop_info values.
    ///
    /// ```rust
    /// use loopdev::LoopDevice;
    /// let ld = LoopDevice::open("/dev/loop5").unwrap();
    /// ld.attach_with_loop_info("test.img", loop_info64::default()).unwrap();
    /// # ld.detach().unwrap();
    /// ```
    pub fn attach_with_loop_info<P: AsRef<Path>>(
        &self,
        backing_file: P,
        loop_info: &mut loop_info64,
    ) -> io::Result<()> {
        let bf = OpenOptions::new()
            .read(true)
            .write(true)
            .open(backing_file)?;

        // Attach the file
        ioctl_to_error(unsafe {
            ioctl(
                self.device.as_raw_fd() as c_int,
                LOOP_SET_FD.into(),
                bf.as_raw_fd() as c_int,
            )
        })?;

        if let Err(err) = ioctl_to_error(unsafe {
            ioctl(
                self.device.as_raw_fd() as c_int,
                LOOP_SET_STATUS64.into(),
                loop_info,
            )
        }) {
            // Ignore the error to preserve the original error
            let _ = self.detach();
            return Err(err);
        }
        Ok(())
    }

    /// Get the path of the loop device.
    #[deprecated(since = "0.2.0", note = "use `path` instead")]
    pub fn get_path(&self) -> Option<PathBuf> {
        self.path()
    }

    /// Get the path of the loop device.
    pub fn path(&self) -> Option<PathBuf> {
        let mut p = PathBuf::from("/proc/self/fd");
        p.push(self.device.as_raw_fd().to_string());
        std::fs::read_link(&p).ok()
    }

    /// Detach a loop device from its backing file.
    ///
    /// Note that the device won't fully detach until a short delay after the underling device file
    /// gets closed. This happens when LoopDev goes out of scope so you should ensure the LoopDev
    /// lives for a short a time as possible.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use loopdev::LoopDevice;
    /// let ld = LoopDevice::open("/dev/loop7").unwrap();
    /// # ld.attach_file("test.img").unwrap();
    /// ld.detach().unwrap();
    /// ```
    pub fn detach(&self) -> io::Result<()> {
        ioctl_to_error(unsafe { ioctl(self.device.as_raw_fd() as c_int, LOOP_CLR_FD.into(), 0) })?;
        Ok(())
    }

    /// Resize a live loop device. If the size of the backing file changes this can be called to
    /// inform the loop driver about the new size.
    pub fn set_capacity(&self) -> io::Result<()> {
        ioctl_to_error(unsafe {
            ioctl(
                self.device.as_raw_fd() as c_int,
                LOOP_SET_CAPACITY.into(),
                0,
            )
        })?;
        Ok(())
    }
}

#[repr(C)]
pub struct loop_info64 {
    pub lo_device: u64,
    pub lo_inode: u64,
    pub lo_rdevice: u64,
    pub lo_offset: u64,
    pub lo_sizelimit: u64,
    pub lo_number: u32,
    pub lo_encrypt_type: u32,
    pub lo_encrypt_key_size: u32,
    pub lo_flags: u32,
    pub lo_file_name: [u8; 64],
    pub lo_crypt_name: [u8; 64],
    pub lo_encrypt_key: [u8; 32],
    pub lo_init: [u64; 2],
}

impl Default for loop_info64 {
    fn default() -> Self {
        Self {
            lo_device: 0,
            lo_inode: 0,
            lo_rdevice: 0,
            lo_offset: 0,
            lo_sizelimit: 0,
            lo_number: 0,
            lo_encrypt_type: 0,
            lo_encrypt_key_size: 0,
            lo_flags: 0,
            lo_file_name: [0; 64],
            lo_crypt_name: [0; 64],
            lo_encrypt_key: [0; 32],
            lo_init: [0; 2],
        }
    }
}

fn ioctl_to_error(ret: i32) -> io::Result<i32> {
    if ret < 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(ret)
    }
}
<|endoftext|>