<commit_before>//
// FlightVars
// Copyright (c) 2015, 2016 Alvaro Polo
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::fmt;
use std::io;
use std::str::FromStr;

use byteorder::{BigEndian, ReadBytesExt};
use hex::FromHex;

#[derive(Debug, PartialEq)]
pub struct OffsetAddr(u16);

impl FromHex for OffsetAddr {
    type Error = io::Error;
    fn from_hex(s: &str) -> Result<Self, Self::Error> {
        let error = || io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("invalid FSUIPC offset address in '{}'", s));
        let buf: Vec<u8> = try!(FromHex::from_hex(s).map_err(|_| error()));
        if buf.len() > 2 { Err(error()) }
        else {
            (&buf[..]).read_u16::<BigEndian>()
                .map(|v| OffsetAddr(v))
                .map_err(|_| error())
        }
    }
}

impl fmt::Display for OffsetAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "{:x}", self.0)
    }
}

#[derive(Debug, PartialEq)]
pub enum OffsetLen { Ub, Sb, Uw, Sw, Ud, Sd }

impl fmt::Display for OffsetLen {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            OffsetLen::Ub => write!(f, "UB"),
            OffsetLen::Sb => write!(f, "SB"),
            OffsetLen::Uw => write!(f, "UW"),
            OffsetLen::Sw => write!(f, "SW"),
            OffsetLen::Ud => write!(f, "UD"),
            OffsetLen::Sd => write!(f, "SD"),
        }
    }
}

impl FromStr for OffsetLen {
    type Err = io::Error;
    fn from_str(s: &str) -> Result<OffsetLen, io::Error> {
        match s {
            "UB" => Ok(OffsetLen::Ub),
            "SB" => Ok(OffsetLen::Sb),
            "UW" => Ok(OffsetLen::Uw),
            "SW" => Ok(OffsetLen::Sw),
            "UD" => Ok(OffsetLen::Ud),
            "SD" => Ok(OffsetLen::Sd),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("invalid FSUIPC offset length in '{}'", s))),
        }
    }
}

#[derive(Debug, PartialEq)]
pub struct Offset(OffsetAddr, OffsetLen);

impl fmt::Display for Offset {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "{}:{}", self.0, self.1)
    }
}

impl FromStr for Offset {
    type Err = io::Error;
    fn from_str(s: &str) -> Result<Offset, io::Error> {
        let pair: Vec<&str> = s.split(":").collect();
        if pair.len() == 2 {
            let addr = try!(OffsetAddr::from_hex(pair[0]));
            let len = try!(OffsetLen::from_str(pair[1]));
            Ok(Offset(addr, len))
        } else {
            Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("invalid FSUIPC offset in '{}'", s)))
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum Message {
    Begin { version: u16, client_id: String },
    WriteLvar { lvar: String, value: i64 },
    WriteOffset { offset: Offset, value: i64 },
    ObserveLvar { lvar: String }
}

impl Message {

    pub fn begin(version: u16, client_id: &str) -> Message {
        Message::Begin { version: version, client_id: client_id.to_string() }
    }

    pub fn write_lvar(lvar: &str, value: i64) -> Message {
        Message::WriteLvar { lvar: lvar.to_string(), value: value }
    }

    pub fn write_offset(offset: Offset, value: i64) -> Message {
        Message::WriteOffset { offset: offset, value: value }
    }

    pub fn obs_lvar(lvar: &str) -> Message {
        Message::ObserveLvar { lvar: lvar.to_string() }
    }
}

impl fmt::Display for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match self {
            &Message::Begin { version, ref client_id } =>
                write!(f, "BEGIN {} {}", version, client_id),
            &Message::WriteLvar { ref lvar, value} =>
                write!(f, "WRITE_LVAR {} {}", lvar, value),
            &Message::WriteOffset { ref offset, value } =>
                write!(f, "WRITE_OFFSET {} {}", offset, value),
            &Message::ObserveLvar { ref lvar } =>
                write!(f, "OBS_LVAR {}", lvar),
        }
    }
}

impl FromStr for Message {
    type Err = io::Error;
    fn from_str(s: &str) -> Result<Message, io::Error> {
        let deco = MessageParser::new(s);
        deco.parse()
    }
}

pub struct MessageParser<'a> {
    input: &'a str
}

impl<'a> MessageParser<'a> {
    pub fn new(input: &'a str) -> MessageParser {
        MessageParser { input: input }
    }

    pub fn parse(self) -> io::Result<Message> {
        let args: Vec<&str> = self.input.split_whitespace().collect();
        let cmd = args[0];
        let args = &args[1..];
        match &cmd.to_uppercase()[..] {
            "BEGIN" => self.parse_begin(&args),
            "WRITE_LVAR" => self.parse_write_lvar(&args),
            "WRITE_OFFSET" => self.parse_write_offset(&args),
            "OBS_LVAR" => self.parse_obs_lvar(&args),
            _ => Err(self.input_error()),
        }
    }

    fn parse_begin(self, args: &[&str]) -> io::Result<Message> {
        try!(self.require_argc(args, 2));
        let ver = try!(args[0].parse().map_err(|_| self.input_error()));
        Ok(Message::begin(ver, args[1]))
    }

    fn parse_write_lvar(self, args: &[&str]) -> io::Result<Message> {
        try!(self.require_argc(args, 2));
        let lvar = args[0];
        let value = try!(args[1].parse().map_err(|_| self.input_error()));
        Ok(Message::write_lvar(&lvar, value))
    }

    fn parse_write_offset(self, args: &[&str]) -> io::Result<Message> {
        try!(self.require_argc(args, 2));
        let offset = try!(args[0].parse());
        let value = try!(args[1].parse().map_err(|_| self.input_error()));
        Ok(Message::write_offset(offset, value))
    }

    fn parse_obs_lvar(self, args: &[&str]) -> io::Result<Message> {
        try!(self.require_argc(args, 1));
        Ok(Message::obs_lvar(args[0]))
    }

    fn require_argc(&self, args: &[&str], expected: usize) -> io::Result<()> {
        if args.len() == expected { Ok(()) }
        else { Err(self.input_error()) }
    }

    fn input_error(&self) -> io::Error {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("invalid oacsp message in '{}'", self.input))
    }
}

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

    use super::*;

    #[test]
    fn should_encode_begin_msg() {
        let msg = Message::begin(1, "arduino");
        let buf = format!("{}", msg);
        assert_eq!(buf, "BEGIN 1 arduino")
    }

    #[test]
    fn should_decode_begin_msg() {
        let buf = "BEGIN 1 arduino";
        let msg = Message::from_str(&buf).unwrap();
        assert_eq!(msg, Message::begin(1, "arduino"));
    }

    #[test]
    fn should_encode_write_lvar_msg() {
        let msg = Message::write_lvar("foobar", 42);
        let buf = format!("{}", msg);
        assert_eq!(buf, "WRITE_LVAR foobar 42")
    }

    #[test]
    fn should_decode_write_lvar_msg() {
        let buf = "WRITE_LVAR foobar 42";
        let msg = Message::from_str(&buf).unwrap();
        assert_eq!(msg, Message::write_lvar("foobar", 42));
    }

    #[test]
    fn should_encode_write_offset_msg() {
        let msg = Message::write_offset(Offset(OffsetAddr(0x1234), OffsetLen::Uw), 42);
        let buf = format!("{}", msg);
        assert_eq!(buf, "WRITE_OFFSET 1234:UW 42")
    }

    #[test]
    fn should_decode_write_offset_msg() {
        let buf = "WRITE_OFFSET 1234:UW 42";
        let msg = Message::from_str(&buf).unwrap();
        assert_eq!(msg, Message::write_offset(Offset(OffsetAddr(0x1234), OffsetLen::Uw), 42));
    }

    #[test]
    fn should_encode_obs_lvar_msg() {
        let msg = Message::obs_lvar("foobar");
        let buf = format!("{}", msg);
        assert_eq!(buf, "OBS_LVAR foobar")
    }

    #[test]
    fn should_decode_obs_lvar_msg() {
        let buf = "OBS_LVAR foobar";
        let msg = Message::from_str(&buf).unwrap();
        assert_eq!(msg, Message::obs_lvar("foobar"));
    }
}
<commit_msg>Use appropriate terminology in `oacsp` unit tests<commit_after>//
// FlightVars
// Copyright (c) 2015, 2016 Alvaro Polo
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::fmt;
use std::io;
use std::str::FromStr;

use byteorder::{BigEndian, ReadBytesExt};
use hex::FromHex;

#[derive(Debug, PartialEq)]
pub struct OffsetAddr(u16);

impl FromHex for OffsetAddr {
    type Error = io::Error;
    fn from_hex(s: &str) -> Result<Self, Self::Error> {
        let error = || io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("invalid FSUIPC offset address in '{}'", s));
        let buf: Vec<u8> = try!(FromHex::from_hex(s).map_err(|_| error()));
        if buf.len() > 2 { Err(error()) }
        else {
            (&buf[..]).read_u16::<BigEndian>()
                .map(|v| OffsetAddr(v))
                .map_err(|_| error())
        }
    }
}

impl fmt::Display for OffsetAddr {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "{:x}", self.0)
    }
}

#[derive(Debug, PartialEq)]
pub enum OffsetLen { Ub, Sb, Uw, Sw, Ud, Sd }

impl fmt::Display for OffsetLen {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            OffsetLen::Ub => write!(f, "UB"),
            OffsetLen::Sb => write!(f, "SB"),
            OffsetLen::Uw => write!(f, "UW"),
            OffsetLen::Sw => write!(f, "SW"),
            OffsetLen::Ud => write!(f, "UD"),
            OffsetLen::Sd => write!(f, "SD"),
        }
    }
}

impl FromStr for OffsetLen {
    type Err = io::Error;
    fn from_str(s: &str) -> Result<OffsetLen, io::Error> {
        match s {
            "UB" => Ok(OffsetLen::Ub),
            "SB" => Ok(OffsetLen::Sb),
            "UW" => Ok(OffsetLen::Uw),
            "SW" => Ok(OffsetLen::Sw),
            "UD" => Ok(OffsetLen::Ud),
            "SD" => Ok(OffsetLen::Sd),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("invalid FSUIPC offset length in '{}'", s))),
        }
    }
}

#[derive(Debug, PartialEq)]
pub struct Offset(OffsetAddr, OffsetLen);

impl fmt::Display for Offset {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "{}:{}", self.0, self.1)
    }
}

impl FromStr for Offset {
    type Err = io::Error;
    fn from_str(s: &str) -> Result<Offset, io::Error> {
        let pair: Vec<&str> = s.split(":").collect();
        if pair.len() == 2 {
            let addr = try!(OffsetAddr::from_hex(pair[0]));
            let len = try!(OffsetLen::from_str(pair[1]));
            Ok(Offset(addr, len))
        } else {
            Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("invalid FSUIPC offset in '{}'", s)))
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum Message {
    Begin { version: u16, client_id: String },
    WriteLvar { lvar: String, value: i64 },
    WriteOffset { offset: Offset, value: i64 },
    ObserveLvar { lvar: String }
}

impl Message {

    pub fn begin(version: u16, client_id: &str) -> Message {
        Message::Begin { version: version, client_id: client_id.to_string() }
    }

    pub fn write_lvar(lvar: &str, value: i64) -> Message {
        Message::WriteLvar { lvar: lvar.to_string(), value: value }
    }

    pub fn write_offset(offset: Offset, value: i64) -> Message {
        Message::WriteOffset { offset: offset, value: value }
    }

    pub fn obs_lvar(lvar: &str) -> Message {
        Message::ObserveLvar { lvar: lvar.to_string() }
    }
}

impl fmt::Display for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match self {
            &Message::Begin { version, ref client_id } =>
                write!(f, "BEGIN {} {}", version, client_id),
            &Message::WriteLvar { ref lvar, value} =>
                write!(f, "WRITE_LVAR {} {}", lvar, value),
            &Message::WriteOffset { ref offset, value } =>
                write!(f, "WRITE_OFFSET {} {}", offset, value),
            &Message::ObserveLvar { ref lvar } =>
                write!(f, "OBS_LVAR {}", lvar),
        }
    }
}

impl FromStr for Message {
    type Err = io::Error;
    fn from_str(s: &str) -> Result<Message, io::Error> {
        let deco = MessageParser::new(s);
        deco.parse()
    }
}

pub struct MessageParser<'a> {
    input: &'a str
}

impl<'a> MessageParser<'a> {
    pub fn new(input: &'a str) -> MessageParser {
        MessageParser { input: input }
    }

    pub fn parse(self) -> io::Result<Message> {
        let args: Vec<&str> = self.input.split_whitespace().collect();
        let cmd = args[0];
        let args = &args[1..];
        match &cmd.to_uppercase()[..] {
            "BEGIN" => self.parse_begin(&args),
            "WRITE_LVAR" => self.parse_write_lvar(&args),
            "WRITE_OFFSET" => self.parse_write_offset(&args),
            "OBS_LVAR" => self.parse_obs_lvar(&args),
            _ => Err(self.input_error()),
        }
    }

    fn parse_begin(self, args: &[&str]) -> io::Result<Message> {
        try!(self.require_argc(args, 2));
        let ver = try!(args[0].parse().map_err(|_| self.input_error()));
        Ok(Message::begin(ver, args[1]))
    }

    fn parse_write_lvar(self, args: &[&str]) -> io::Result<Message> {
        try!(self.require_argc(args, 2));
        let lvar = args[0];
        let value = try!(args[1].parse().map_err(|_| self.input_error()));
        Ok(Message::write_lvar(&lvar, value))
    }

    fn parse_write_offset(self, args: &[&str]) -> io::Result<Message> {
        try!(self.require_argc(args, 2));
        let offset = try!(args[0].parse());
        let value = try!(args[1].parse().map_err(|_| self.input_error()));
        Ok(Message::write_offset(offset, value))
    }

    fn parse_obs_lvar(self, args: &[&str]) -> io::Result<Message> {
        try!(self.require_argc(args, 1));
        Ok(Message::obs_lvar(args[0]))
    }

    fn require_argc(&self, args: &[&str], expected: usize) -> io::Result<()> {
        if args.len() == expected { Ok(()) }
        else { Err(self.input_error()) }
    }

    fn input_error(&self) -> io::Error {
        io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("invalid oacsp message in '{}'", self.input))
    }
}

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

    use super::*;

    #[test]
    fn should_display_begin_msg() {
        let msg = Message::begin(1, "arduino");
        let buf = format!("{}", msg);
        assert_eq!(buf, "BEGIN 1 arduino")
    }

    #[test]
    fn should_parse_begin_msg() {
        let buf = "BEGIN 1 arduino";
        let msg = Message::from_str(&buf).unwrap();
        assert_eq!(msg, Message::begin(1, "arduino"));
    }

    #[test]
    fn should_display_write_lvar_msg() {
        let msg = Message::write_lvar("foobar", 42);
        let buf = format!("{}", msg);
        assert_eq!(buf, "WRITE_LVAR foobar 42")
    }

    #[test]
    fn should_parse_write_lvar_msg() {
        let buf = "WRITE_LVAR foobar 42";
        let msg = Message::from_str(&buf).unwrap();
        assert_eq!(msg, Message::write_lvar("foobar", 42));
    }

    #[test]
    fn should_display_write_offset_msg() {
        let msg = Message::write_offset(Offset(OffsetAddr(0x1234), OffsetLen::Uw), 42);
        let buf = format!("{}", msg);
        assert_eq!(buf, "WRITE_OFFSET 1234:UW 42")
    }

    #[test]
    fn should_parse_write_offset_msg() {
        let buf = "WRITE_OFFSET 1234:UW 42";
        let msg = Message::from_str(&buf).unwrap();
        assert_eq!(msg, Message::write_offset(Offset(OffsetAddr(0x1234), OffsetLen::Uw), 42));
    }

    #[test]
    fn should_display_obs_lvar_msg() {
        let msg = Message::obs_lvar("foobar");
        let buf = format!("{}", msg);
        assert_eq!(buf, "OBS_LVAR foobar")
    }

    #[test]
    fn should_parse_obs_lvar_msg() {
        let buf = "OBS_LVAR foobar";
        let msg = Message::from_str(&buf).unwrap();
        assert_eq!(msg, Message::obs_lvar("foobar"));
    }
}
<|endoftext|>


<commit_before>//! Read and write NDArrays as Numpy arrays

use ndarray::prelude::*;
use ndarray as nd;
use std::fs::File;
use std::path::Path;
use std::ptr;
use std::str;
use std::io::{Read, Write};
use regex::bytes::Regex;
use byteorder::{LittleEndian, BigEndian, ReadBytesExt, WriteBytesExt};
use errors::*;
use memmap::{Mmap, Protection};

/// Write an array as a numpy array, fast but only works for native byte order
pub fn write_matrix<S, P>(path: P, arr: &ArrayBase<S, Ix2>) -> Result<()>
    where S: nd::Data<Elem=f64>, P: AsRef<Path> {
    let header_nospace = format!("{{'descr': '<f8', 'fortran_order': False, 'shape': ({},{})}}",
        arr.shape()[0], arr.shape()[1]);
    let virtual_len =
        // Calculating how many bytes we have in the header, so we can get alignment
        header_nospace.len()
        + 6 // The magic string
        + 2 // The version number
        + 2 // An unsigned 2-byte integer for header length
        + 1 ; // Because there will be a \n added
    let padding_needed = (((virtual_len + 15) / 16) * 16) - virtual_len; // to get to the next 16

    // Numpy version 1.0
    let mut writer = File::create(path)?;
    writer.write_all(b"\x93NUMPY\x01\x00")?;
    writer.write_u16::<LittleEndian>((header_nospace.len()
        + padding_needed
        + 1) // newline
        // Magic string, version number and 2 bytes for the header length number not included.
        as u16
        )?;
    write!(writer, "{}{}\n", header_nospace, " ".repeat(padding_needed))?;
    unsafe {
        // O(N) array copy without memory overhead.
        assert!(arr.is_standard_layout(), "To write an array it needs to be in C order (for now).");
        // Treat the incoming data as floating point
        writer.write_all(::std::slice::from_raw_parts(arr.as_ptr() as *const u8, arr.len()*8))?;
    };
    Ok(())
}

/// Read a Numpy matrix into memory. Be careful if it's large. You could run out of memory.
///
/// You need to know the number of dimensions at compile time so for convenience, we assume you
/// need a matrix. Also, this method is fast but only works for native byte order.
pub fn read_matrix<P: AsRef<Path>>(path: P) -> Result<Array2<f64>> {
    let header_match = Regex::new(r"NUMPY\x01\x00(?s:..)\{'descr': ?'<f8', ?'fortran_order': ?False, ?'shape': ?\((\d+), ?(\d+)\)\} *\n").unwrap();
    let mut reader = File::open(path.as_ref())?;
    let mut content = vec![];
    reader.read_to_end(&mut content)?;
    // The skip and the nested context here is so that we can regex parse (with a borrow)
    // and then reuse the buffer as the array. This way there are never two copies in memory.
    let res: Result<(usize, usize, usize)> = {
        let captures = header_match.captures(&content)
            .ok_or(helpful_complaint(path.as_ref(), &content))?;
        Ok((// where the full match ends
            captures.get(0).unwrap().end(),
            // The shape of the array as described in the metadata
            str::from_utf8(&captures[1]).unwrap().parse().unwrap(),
            str::from_utf8(&captures[2]).unwrap().parse().unwrap()
        ))
    };
    let (skip, h, w) = res?;
    let content_f64 = unsafe {
        // O(N) array copy without memory overhead.
        let length = content.len() - skip;
        assert_eq!(length, h*w*8,
            "The numpy file's array is the wrong length for a {}x{} array. \
            It should be {} elements, ({} bytes), but it is actually {} bytes.",
            h, w, h*w, h*w*8, length
        );
        ptr::copy((&content[skip..]).as_ptr(),
                  (&mut content[..length]).as_mut_ptr(),
                  length);
        content.set_len(length);
        // Treat the incoming data as floating point
        let mut out : Vec<f64> = ::std::mem::transmute(content);
        out.set_len(length/8);
        out
    };
    Ok(ArrayBase::from_shape_vec([h, w], content_f64)?)
}

pub struct MatFile(usize, usize, Mmap);

/// Load a Numpy matrix as an mmap. This only consumes address space. (Part 1)
///
/// This is a two-step process because the Mmap needs to outlive the matrix.
pub fn open_matrix_mmap<P: AsRef<Path>>(path: P) -> Result<MatFile> {
    let header_match = Regex::new(r"NUMPY\x01\x00(?s:..)\{'descr': ?'<f8', ?'fortran_order': ?False, ?'shape': ?\((\d+), ?(\d+)\)\} *\n").unwrap();
    let mut reader = File::open(path.as_ref())?;
    let mut content = [0u8; 128];
    let bytes_read = reader.read(&mut content)?;
    assert!(bytes_read > 0, format!("The numpy file {} seems to be empty.", path.as_ref().display()));
    // The skip and the nested context here is so that we can regex parse (with a borrow)
    // and then reuse the buffer as the array. This way there are never two copies in memory.
    let captures = header_match.captures(&content)
        .ok_or(helpful_complaint(path.as_ref(), &content))?;
    // where the full match ends
    let skip = captures.get(0).unwrap().end();
    // The shape of the array as described in the metadata
    let h = str::from_utf8(&captures[1]).unwrap().parse().unwrap();
    let w = str::from_utf8(&captures[2]).unwrap().parse().unwrap();

    Ok(MatFile(h, w, Mmap::open_with_offset(&reader, Protection::ReadWrite, skip, h*w*8)?))
}

/// Load a Numpy matrix as an mmap (Part 2)
///
/// You need to know the number of dimensions at compile time so for convenience, we assume you
/// need a matrix. Also, this method is even faster than read_matrix but only works for native byte
/// order.
///
/// This is different than you might expect. The file you read from must outlive the array,
/// because the array is based on an mmap of that file (so that it doesn't read it into memory).
pub fn read_matrix_mmap<'t>(mmap: &'t MatFile) -> Result<ArrayViewMut2<'t, f64>> {
    unsafe {
        let new_slice= ::std::slice::from_raw_parts_mut(mmap.2.ptr() as *mut f64, mmap.2.len()/8);
        Ok(ArrayViewMut2::from_shape([mmap.0, mmap.1], new_slice)?)
    }
}

/// Tell the user more info about the file
///
/// It seems verbose but you can see this error often so it save you time.
fn helpful_complaint(p: &Path, header: &[u8]) -> Error {
    let cap = ::std::cmp::min(header.len(), 100);
    let complaint = format!(
        "Expected {} to be an uncompressed numpy (.npy) file, but couldn't \
        parse the header. The first hundred bytes look like:

        {}

        As bytes, the header is as follows:

        {:?}

        It should look something like this example, where . are non-printable characters: \
        NUMPY..{{'descr': '<f8', 'fortran_order': False, 'shape': (34, 27)}}\
        Note: Cabarrus only supports 2D big-endian 64-bit float matrices in C order (for \
        simplicity). You may need to change the dtype accordingly.",
        p.display(),
        String::from_utf8_lossy(&header[..cap]),
        &header[..cap]);
    Error::Other(complaint)
}<commit_msg>Mismatched File and Mmap permissions<commit_after>//! Read and write NDArrays as Numpy arrays

use ndarray::prelude::*;
use ndarray as nd;
use std::fs::{OpenOption, File};
use std::path::Path;
use std::ptr;
use std::str;
use std::io::{Read, Write};
use regex::bytes::Regex;
use byteorder::{LittleEndian, BigEndian, ReadBytesExt, WriteBytesExt};
use errors::*;
use memmap::{Mmap, Protection};

/// Write an array as a numpy array, fast but only works for native byte order
pub fn write_matrix<S, P>(path: P, arr: &ArrayBase<S, Ix2>) -> Result<()>
    where S: nd::Data<Elem=f64>, P: AsRef<Path> {
    let header_nospace = format!("{{'descr': '<f8', 'fortran_order': False, 'shape': ({},{})}}",
        arr.shape()[0], arr.shape()[1]);
    let virtual_len =
        // Calculating how many bytes we have in the header, so we can get alignment
        header_nospace.len()
        + 6 // The magic string
        + 2 // The version number
        + 2 // An unsigned 2-byte integer for header length
        + 1 ; // Because there will be a \n added
    let padding_needed = (((virtual_len + 15) / 16) * 16) - virtual_len; // to get to the next 16

    // Numpy version 1.0
    let mut writer = File::create(path)?;
    writer.write_all(b"\x93NUMPY\x01\x00")?;
    writer.write_u16::<LittleEndian>((header_nospace.len()
        + padding_needed
        + 1) // newline
        // Magic string, version number and 2 bytes for the header length number not included.
        as u16
        )?;
    write!(writer, "{}{}\n", header_nospace, " ".repeat(padding_needed))?;
    unsafe {
        // O(N) array copy without memory overhead.
        assert!(arr.is_standard_layout(), "To write an array it needs to be in C order (for now).");
        // Treat the incoming data as floating point
        writer.write_all(::std::slice::from_raw_parts(arr.as_ptr() as *const u8, arr.len()*8))?;
    };
    Ok(())
}

/// Read a Numpy matrix into memory. Be careful if it's large. You could run out of memory.
///
/// You need to know the number of dimensions at compile time so for convenience, we assume you
/// need a matrix. Also, this method is fast but only works for native byte order.
pub fn read_matrix<P: AsRef<Path>>(path: P) -> Result<Array2<f64>> {
    let header_match = Regex::new(r"NUMPY\x01\x00(?s:..)\{'descr': ?'<f8', ?'fortran_order': ?False, ?'shape': ?\((\d+), ?(\d+)\)\} *\n").unwrap();
    let mut reader = File::open(path.as_ref())?;
    let mut content = vec![];
    reader.read_to_end(&mut content)?;
    // The skip and the nested context here is so that we can regex parse (with a borrow)
    // and then reuse the buffer as the array. This way there are never two copies in memory.
    let res: Result<(usize, usize, usize)> = {
        let captures = header_match.captures(&content)
            .ok_or(helpful_complaint(path.as_ref(), &content))?;
        Ok((// where the full match ends
            captures.get(0).unwrap().end(),
            // The shape of the array as described in the metadata
            str::from_utf8(&captures[1]).unwrap().parse().unwrap(),
            str::from_utf8(&captures[2]).unwrap().parse().unwrap()
        ))
    };
    let (skip, h, w) = res?;
    let content_f64 = unsafe {
        // O(N) array copy without memory overhead.
        let length = content.len() - skip;
        assert_eq!(length, h*w*8,
            "The numpy file's array is the wrong length for a {}x{} array. \
            It should be {} elements, ({} bytes), but it is actually {} bytes.",
            h, w, h*w, h*w*8, length
        );
        ptr::copy((&content[skip..]).as_ptr(),
                  (&mut content[..length]).as_mut_ptr(),
                  length);
        content.set_len(length);
        // Treat the incoming data as floating point
        let mut out : Vec<f64> = ::std::mem::transmute(content);
        out.set_len(length/8);
        out
    };
    Ok(ArrayBase::from_shape_vec([h, w], content_f64)?)
}

pub struct MatFile(usize, usize, Mmap);

/// Load a Numpy matrix as an mmap. This only consumes address space. (Part 1)
///
/// This is a two-step process because the Mmap needs to outlive the matrix.
pub fn open_matrix_mmap<P: AsRef<Path>>(path: P) -> Result<MatFile> {
    let header_match = Regex::new(r"NUMPY\x01\x00(?s:..)\{'descr': ?'<f8', ?'fortran_order': ?False, ?'shape': ?\((\d+), ?(\d+)\)\} *\n").unwrap();
    let mut reader = OpenOption::new()
        .read(true)
        .write(true)
        .create(true)
        .open(path.as_ref())?;
    let mut content = [0u8; 128];
    let bytes_read = reader.read(&mut content)?;
    assert!(bytes_read > 0, format!("The numpy file {} seems to be empty.", path.as_ref().display()));
    // The skip and the nested context here is so that we can regex parse (with a borrow)
    // and then reuse the buffer as the array. This way there are never two copies in memory.
    let captures = header_match.captures(&content)
        .ok_or(helpful_complaint(path.as_ref(), &content))?;
    // where the full match ends
    let skip = captures.get(0).unwrap().end();
    // The shape of the array as described in the metadata
    let h = str::from_utf8(&captures[1]).unwrap().parse().unwrap();
    let w = str::from_utf8(&captures[2]).unwrap().parse().unwrap();

    Ok(MatFile(h, w, Mmap::open_with_offset(&reader, Protection::ReadWrite, skip, h*w*8)?))
}

/// Load a Numpy matrix as an mmap (Part 2)
///
/// You need to know the number of dimensions at compile time so for convenience, we assume you
/// need a matrix. Also, this method is even faster than read_matrix but only works for native byte
/// order.
///
/// This is different than you might expect. The file you read from must outlive the array,
/// because the array is based on an mmap of that file (so that it doesn't read it into memory).
pub fn read_matrix_mmap<'t>(mmap: &'t MatFile) -> Result<ArrayViewMut2<'t, f64>> {
    unsafe {
        let new_slice= ::std::slice::from_raw_parts_mut(mmap.2.ptr() as *mut f64, mmap.2.len()/8);
        Ok(ArrayViewMut2::from_shape([mmap.0, mmap.1], new_slice)?)
    }
}

/// Tell the user more info about the file
///
/// It seems verbose but you can see this error often so it save you time.
fn helpful_complaint(p: &Path, header: &[u8]) -> Error {
    let cap = ::std::cmp::min(header.len(), 100);
    let complaint = format!(
        "Expected {} to be an uncompressed numpy (.npy) file, but couldn't \
        parse the header. The first hundred bytes look like:

        {}

        As bytes, the header is as follows:

        {:?}

        It should look something like this example, where . are non-printable characters: \
        NUMPY..{{'descr': '<f8', 'fortran_order': False, 'shape': (34, 27)}}\
        Note: Cabarrus only supports 2D big-endian 64-bit float matrices in C order (for \
        simplicity). You may need to change the dtype accordingly.",
        p.display(),
        String::from_utf8_lossy(&header[..cap]),
        &header[..cap]);
    Error::Other(complaint)
}<|endoftext|>


<commit_before>pub type c_char = i8;
pub type c_long = i64;
pub type c_ulong = u64;

pub type wchar_t = i16;

pub type off_t = usize;
pub type mode_t = u16;
pub type time_t = i64;
pub type pid_t = usize;
pub type gid_t = usize;
pub type uid_t = usize;

pub type in_addr_t = u32;
pub type in_port_t = u16;

pub type socklen_t = u32;
pub type sa_family_t = u16;

s! {
    pub struct in_addr {
        pub s_addr: in_addr_t,
    }

    pub struct in6_addr {
        pub s6_addr: [u8; 16],
        __align: [u32; 0],
    }

    pub struct sockaddr {
        pub sa_family: sa_family_t,
        pub sa_data: [::c_char; 14],
    }

    pub struct sockaddr_in {
        pub sin_family: sa_family_t,
        pub sin_port: ::in_port_t,
        pub sin_addr: ::in_addr,
        pub sin_zero: [u8; 8],
    }

    pub struct sockaddr_in6 {
        pub sin6_family: sa_family_t,
        pub sin6_port: in_port_t,
        pub sin6_flowinfo: u32,
        pub sin6_addr: ::in6_addr,
        pub sin6_scope_id: u32,
    }
}

pub const STDIN_FILENO: ::c_int = 0;
pub const STDOUT_FILENO: ::c_int = 1;
pub const STDERR_FILENO: ::c_int = 2;

extern {
    pub fn memalign(align: ::size_t, size: ::size_t) -> *mut ::c_void;
    pub fn read(fd: ::c_int, buf: *mut ::c_void, count: ::size_t) -> ::ssize_t;
    pub fn write(fd: ::c_int, buf: *const ::c_void, count: ::size_t) -> ::ssize_t;
}

#[link(name = "c")]
#[link(name = "m")]
extern {}
<commit_msg>Update redox.rs<commit_after>pub type c_char = i8;
pub type c_long = i64;
pub type c_ulong = u64;

pub type wchar_t = i16;

pub type off_t = usize;
pub type mode_t = u16;
pub type time_t = i64;
pub type pid_t = usize;
pub type gid_t = usize;
pub type uid_t = usize;

pub type in_addr_t = u32;
pub type in_port_t = u16;

pub type socklen_t = u32;
pub type sa_family_t = u16;

s! {
    pub struct in_addr {
        pub s_addr: in_addr_t,
    }

    pub struct in6_addr {
        pub s6_addr: [u8; 16],
        __align: [u32; 0],
    }

    pub struct sockaddr {
        pub sa_family: sa_family_t,
        pub sa_data: [::c_char; 14],
    }

    pub struct sockaddr_in {
        pub sin_family: sa_family_t,
        pub sin_port: ::in_port_t,
        pub sin_addr: ::in_addr,
        pub sin_zero: [u8; 8],
    }

    pub struct sockaddr_in6 {
        pub sin6_family: sa_family_t,
        pub sin6_port: in_port_t,
        pub sin6_flowinfo: u32,
        pub sin6_addr: ::in6_addr,
        pub sin6_scope_id: u32,
    }
}

pub const STDIN_FILENO: ::c_int = 0;
pub const STDOUT_FILENO: ::c_int = 1;
pub const STDERR_FILENO: ::c_int = 2;

extern {
    pub fn memalign(align: ::size_t, size: ::size_t) -> *mut ::c_void;
    pub fn read(fd: ::c_int, buf: *mut ::c_void, count: ::size_t)
                -> ::ssize_t;
    pub fn write(fd: ::c_int, buf: *const ::c_void, count: ::size_t)
                 -> ::ssize_t;
}

#[link(name = "c")]
#[link(name = "m")]
extern {}
<|endoftext|>


<commit_before>use std::io::Read;
use std::io::Write;
use std::io::Result;
use std::io::Error;
use std::io::ErrorKind;

const BUFFER_SIZE: usize = 2048;

const VALUE_BITS: usize = 8;
const VALUE_EOF: usize = 1 << VALUE_BITS;
const VALUE_COUNT: usize = VALUE_EOF + 2;

const CODE_BITS: usize = 16;
const CODE_MIN: u64 = 0;
const CODE_MAX: u64 = (1 << CODE_BITS) - 1;
const ONE_FOURTH: u64 = 1 << (CODE_BITS - 2);
const ONE_HALF: u64 = 2 * ONE_FOURTH;
const THREE_FOURTHS: u64 = 3 * ONE_FOURTH;

const FREQ_BITS: usize = 14;
const FREQ_MAX: u64 = (1 << FREQ_BITS) - 1;

struct Buffer {
    data: [u8; BUFFER_SIZE],
    len: usize,
    cur: usize,
    bit: u8,
    sum: u64
}

impl Buffer {
    fn create() -> Buffer {
        return Buffer {
            data: [0u8; BUFFER_SIZE],
            len: 0us,
            cur: 0us,
            bit: 7u8,
            sum: 0u64
        };
    }

    fn fill_buffer(&mut self, istream: &mut Read) -> Result<bool> {
        match istream.read(&mut self.data[0..]) {
            Ok(0) => Ok(true),
            Ok(n) => {
                self.sum += n as u64;
                self.len = n;
                self.cur = 0;
                Ok(false)
            },
            Err(e) => Err(e)
        }
    }

    fn get_byte(&mut self, istream: &mut Read) -> Result<(u8, bool)> {
        let value = self.data[self.cur];
        self.cur += 1;
        if self.cur == self.len {
            return match self.fill_buffer(istream) {
                Ok(b) => Ok((value, b)),
                Err(e) => Err(e)
            }
        } else {
            return Ok((value, false));
        }
    }

    fn get_bit(&mut self, istream: &mut Read) -> Result<(u8, bool)> {
        let value = (self.data[self.cur] & (1 << self.bit)) >> self.bit;

        if self.bit == 0 {
            self.bit = 7;
            self.cur += 1;
            if self.cur == self.len {
                return match self.fill_buffer(istream) {
                    Ok(b) => Ok((value, b)),
                    Err(e) => Err(e)
                }
            }
        } else {
            self.bit -= 1;
        }
        return Ok((value, false));
    }

    fn flush_buffer(&mut self, ostream: &mut Write) -> Result<()> {
        // Flush the remaining incomplete byte as well
        if self.bit < 7 {
            self.len += 1;
        }

        while self.cur < self.len {
            match ostream.write(&self.data[self.cur..self.len]) {
                Ok(0) => { return Err(Error::new(ErrorKind::WriteZero, "Zero bytes written to output")); }
                Ok(n) => {
                    self.cur += n;
                    self.sum += n as u64;
                },
                Err(e) => { return Err(e); }
            }
        }

        self.len = 0;
        self.cur = 0;
        return Ok(());
    }

    fn put_byte(&mut self, value: u8, ostream: &mut Write) -> Result<()> {
        self.data[self.len] = value;
        self.len += 1;
        if self.len == self.data.len() {
            if let Err(e) = self.flush_buffer(ostream) {
                return Err(e);
            }
        }
        return Ok(());
    }

    fn put_bit(&mut self, value: u8, ostream: &mut Write) -> Result<()> {
        self.data[self.len] |= value << self.bit;

        if self.bit == 0 {
            self.bit = 7;
            self.len += 1;
            if self.len == self.data.len() {
                if let Err(e) = self.flush_buffer(ostream) {
                    return Err(e);
                }
            }
            self.data[self.len] = 0;
        } else {
            self.bit -= 1;
        }
        return Ok(());
    }
}

struct Status {
    low: u64,
    high: u64,
    range: u64,
    pending: u8,
    freq: [u64; VALUE_COUNT],
    input: Buffer,
    output: Buffer
}

impl Status {
    fn init() -> Status {
        let mut res = Status {
            low: CODE_MIN,
            high: CODE_MAX,
            range: 0u64,
            pending: 0u8,
            freq: [0u64; VALUE_COUNT],
            input: Buffer::create(),
            output: Buffer::create()
        };
        for i in 1..258 {
            res.freq[i] = i as u64;
        }
        return res;
    }

    fn get_count(&self) -> u64 {
        self.freq[self.freq.len() - 1]
    }

    fn get_probability(&mut self, chr: usize) -> (u64, u64, u64) {
        let res = (self.freq[chr], self.freq[chr + 1], self.get_count());
        self.update(chr);
        return res;
    }

    fn get_char(&mut self, scaled: u64) -> (usize, u64, u64) {
        let chr = match self.freq.binary_search(&scaled) {
            Ok(i) => i,
            Err(i) => i - 1
        };
        let res = (chr, self.freq[chr], self.freq[chr + 1]);
        self.update(chr);
        return res;
    }

    fn update(&mut self, chr: usize) {
        if self.get_count() < FREQ_MAX {
            for index in chr+1..self.freq.len() {
                self.freq[index] += 1;
            }
        }
    }

    fn put_pending_bits(&mut self, bit: u8, ostream: &mut Write) -> Result<()> {
        try!(self.output.put_bit(bit, ostream));
        while self.pending > 0 {
            try!(self.output.put_bit(!bit & 1u8, ostream));
            self.pending -= 1;
        }
        return Ok(());
    }

    fn metrics(&self) -> (u64, u64) {
        (self.input.sum, self.output.sum)
    }
}

pub fn compress(istream: &mut Read, ostream: &mut Write) -> Result<(u64, u64)> {
    let mut status = Status::init();

    let mut eof = try!(status.input.fill_buffer(istream));
    loop {
        let chr = if eof {
            VALUE_EOF
        } else {
            match status.input.get_byte(istream) {
                Ok((c, e)) => { eof = e; c as usize },
                Err(e) => { return Err(e); }
            }
        };

        {
            let (low, high, count) = status.get_probability(chr);
            status.range = status.high - status.low + 1;
            status.high = status.low + (status.range * high / count) - 1;
            status.low = status.low + (status.range * low / count);
        }

        loop {
           if status.high < ONE_HALF {
               try!(status.put_pending_bits(0, ostream));
           } else if status.low >= ONE_HALF {
               try!(status.put_pending_bits(1, ostream));
           } else if status.low >= ONE_FOURTH && status.high < THREE_FOURTHS {
                status.pending += 1;
                status.low -= ONE_FOURTH;
                status.high -= ONE_FOURTH;
           } else {
               break;
           }

           status.high = ((status.high << 1) + 1) & CODE_MAX;
           status.low = (status.low << 1) & CODE_MAX;
        }

        if chr == VALUE_EOF {
            break;
        }
    }

    status.pending += 1;
    if status.low < ONE_FOURTH {
        try!(status.put_pending_bits(0, ostream));
    } else {
        try!(status.put_pending_bits(1, ostream));
    }
    try!(status.output.flush_buffer(ostream));

    return Ok(status.metrics());
}

pub fn decompress(istream: &mut Read, ostream: &mut Write) -> Result<(u64, u64)> {
    let mut status = Status::init();
    let mut value = 0u64;

    if try!(status.input.fill_buffer(istream)) {
        return Err(Error::new(ErrorKind::InvalidInput, "Empty input stream"));
    }
    let mut eof = false;
    for _ in 0..CODE_BITS {
        value = if eof {
            (value << 1) | (value & 1u64)
        } else {
            match status.input.get_bit(istream) {
                Ok((b, e)) => { eof = e; (value << 1) | (b as u64) },
                Err(e) => { return Err(e); }
            }
        }
    }

    loop {
        {
            status.range = status.high - status.low + 1;
            let count = status.get_count();
            let scaled = ((value - status.low + 1) * count - 1) / status.range;
            let (chr, low, high) = status.get_char(scaled);

            if chr == VALUE_EOF {
                break;
            } else {
                try!(status.output.put_byte(chr as u8, ostream));
            }

            status.high = status.low + (status.range * high / count) - 1;
            status.low = status.low + (status.range * low / count);
        }

        loop {
            if status.high < ONE_HALF {
                // do nothing
            } else if status.low >= ONE_HALF {
                value -= ONE_HALF;
                status.low -= ONE_HALF;
                status.high -= ONE_HALF;
            } else if status.low >= ONE_FOURTH && status.high < THREE_FOURTHS {
                value -= ONE_FOURTH;
                status.low -= ONE_FOURTH;
                status.high -= ONE_FOURTH;
            } else {
                break;
            }

            status.low = status.low << 1;
            status.high = (status.high << 1) + 1;

            value = if eof {
                return Err(Error::new(ErrorKind::InvalidInput, "Unexpected end of input stream"));
            } else {
                match status.input.get_bit(istream) {
                    Ok((b, e)) => { eof = e; (value << 1) | (b as u64) },
                    Err(e) => { return Err(e); }
                }
            }
        }
    }

    try!(status.output.flush_buffer(ostream));
    return Ok(status.metrics());
}

<commit_msg>Moved de/compressor methods inside the impl<commit_after>use std::io::Read;
use std::io::Write;
use std::io::Result;
use std::io::Error;
use std::io::ErrorKind;

const BUFFER_SIZE: usize = 16384;

const VALUE_BITS: usize = 8;
const VALUE_EOF: usize = 1 << VALUE_BITS;
const VALUE_COUNT: usize = VALUE_EOF + 2;

const CODE_BITS: usize = 16; // 33;
const CODE_MIN: u64 = 0;
const CODE_MAX: u64 = (1 << CODE_BITS) - 1;
const ONE_FOURTH: u64 = 1 << (CODE_BITS - 2);
const ONE_HALF: u64 = 2 * ONE_FOURTH;
const THREE_FOURTHS: u64 = 3 * ONE_FOURTH;

const FREQ_BITS: usize = 14; // 64 - CODE_BITS;
const FREQ_MAX: u64 = (1 << FREQ_BITS) - 1;

struct Buffer {
    data: [u8; BUFFER_SIZE],
    len: usize,
    cur: usize,
    bit: u8,
    sum: u64
}

impl Buffer {
    fn create() -> Buffer {
        return Buffer {
            data: [0u8; BUFFER_SIZE],
            len: 0us,
            cur: 0us,
            bit: 7u8,
            sum: 0u64
        };
    }

    fn fill_buffer(&mut self, istream: &mut Read) -> Result<bool> {
        match istream.read(&mut self.data[0..]) {
            Ok(0) => Ok(true),
            Ok(n) => {
                self.sum += n as u64;
                self.len = n;
                self.cur = 0;
                Ok(false)
            },
            Err(e) => Err(e)
        }
    }

    fn get_byte(&mut self, istream: &mut Read) -> Result<(u8, bool)> {
        let value = self.data[self.cur];
        self.cur += 1;
        if self.cur == self.len {
            return match self.fill_buffer(istream) {
                Ok(b) => Ok((value, b)),
                Err(e) => Err(e)
            }
        } else {
            return Ok((value, false));
        }
    }

    fn get_bit(&mut self, istream: &mut Read) -> Result<(u8, bool)> {
        let value = (self.data[self.cur] & (1 << self.bit)) >> self.bit;

        if self.bit == 0 {
            self.bit = 7;
            self.cur += 1;
            if self.cur == self.len {
                return match self.fill_buffer(istream) {
                    Ok(b) => Ok((value, b)),
                    Err(e) => Err(e)
                }
            }
        } else {
            self.bit -= 1;
        }
        return Ok((value, false));
    }

    fn flush_buffer(&mut self, ostream: &mut Write) -> Result<()> {
        // Flush the remaining incomplete byte as well
        if self.bit < 7 {
            self.len += 1;
        }

        while self.cur < self.len {
            match ostream.write(&self.data[self.cur..self.len]) {
                Ok(0) => { return Err(Error::new(ErrorKind::WriteZero, "Zero bytes written to output")); }
                Ok(n) => {
                    self.cur += n;
                    self.sum += n as u64;
                },
                Err(e) => { return Err(e); }
            }
        }

        self.len = 0;
        self.cur = 0;
        return Ok(());
    }

    fn put_byte(&mut self, value: u8, ostream: &mut Write) -> Result<()> {
        self.data[self.len] = value;
        self.len += 1;
        if self.len == self.data.len() {
            if let Err(e) = self.flush_buffer(ostream) {
                return Err(e);
            }
        }
        return Ok(());
    }

    fn put_bit(&mut self, value: u8, ostream: &mut Write) -> Result<()> {
        self.data[self.len] |= value << self.bit;

        if self.bit == 0 {
            self.bit = 7;
            self.len += 1;
            if self.len == self.data.len() {
                if let Err(e) = self.flush_buffer(ostream) {
                    return Err(e);
                }
            }
            self.data[self.len] = 0;
        } else {
            self.bit -= 1;
        }
        return Ok(());
    }
}

struct Status {
    low: u64,
    high: u64,
    pending: u64,
    freq: [u64; VALUE_COUNT],
    input: Buffer,
    output: Buffer
}

impl Status {
    fn init() -> Status {
        let mut res = Status {
            low: CODE_MIN,
            high: CODE_MAX,
            pending: 0u64,
            freq: [0u64; VALUE_COUNT],
            input: Buffer::create(),
            output: Buffer::create()
        };
        for i in 1..res.freq.len() {
            res.freq[i] = i as u64;
        }
        return res;
    }

    fn update(&mut self, chr: usize) {
        if self.get_count() < FREQ_MAX {
            for index in chr+1..self.freq.len() {
                self.freq[index] += 1;
            }
        }
    }

    fn get_count(&self) -> u64 {
        self.freq[self.freq.len() - 1]
    }

    fn get_probability(&mut self, chr: usize) -> (u64, u64, u64) {
        let res = (self.freq[chr], self.freq[chr + 1], self.get_count());
        self.update(chr);
        return res;
    }

    fn get_char(&mut self, scaled: u64) -> (usize, u64, u64) {
        let chr = match self.freq.binary_search(&scaled) {
            Ok(i) => i,
            Err(i) => i - 1
        };
        let res = (chr, self.freq[chr], self.freq[chr + 1]);
        self.update(chr);
        return res;
    }

    fn get_metrics(&self) -> (u64, u64) {
        (self.input.sum, self.output.sum)
    }

    fn put_bits(&mut self, bit: u8, ostream: &mut Write) -> Result<()> {
        try!(self.output.put_bit(bit, ostream));
        while self.pending > 0 {
            try!(self.output.put_bit(!bit & 1u8, ostream));
            self.pending -= 1;
        }
        return Ok(());
    }

    fn compress(&mut self, istream: &mut Read, ostream: &mut Write) -> Result<()> {
        let mut eof = try!(self.input.fill_buffer(istream));
        loop {
            let chr = if eof {
                VALUE_EOF
            } else {
                match self.input.get_byte(istream) {
                    Ok((c, e)) => { eof = e; c as usize },
                    Err(e) => { return Err(e); }
                }
            };

            {
                let (low, high, count) = self.get_probability(chr);
                let range = self.high - self.low + 1;
                self.high = self.low + (range * high / count) - 1;
                self.low = self.low + (range * low / count);
            }

            loop {
               if self.high < ONE_HALF {
                   try!(self.put_bits(0, ostream));
               } else if self.low >= ONE_HALF {
                   try!(self.put_bits(1, ostream));
               } else if self.low >= ONE_FOURTH && self.high < THREE_FOURTHS {
                    self.pending += 1;
                    self.low -= ONE_FOURTH;
                    self.high -= ONE_FOURTH;
               } else {
                   break;
               }

               self.high = ((self.high << 1) + 1) & CODE_MAX;
               self.low = (self.low << 1) & CODE_MAX;
            }

            if chr == VALUE_EOF {
                break;
            }
        }

        self.pending += 1;
        if self.low < ONE_FOURTH {
            try!(self.put_bits(0, ostream));
        } else {
            try!(self.put_bits(1, ostream));
        }
        try!(self.output.flush_buffer(ostream));

        return Ok(());
    }

    fn decompress(&mut self, istream: &mut Read, ostream: &mut Write) -> Result<()> {
        if try!(self.input.fill_buffer(istream)) {
            return Err(Error::new(ErrorKind::InvalidInput, "Empty input stream"));
        }
        let mut eof = false;
        for _ in 0..CODE_BITS {
            self.pending = if eof {
                (self.pending << 1) | (self.pending & 1u64)
            } else {
                match self.input.get_bit(istream) {
                    Ok((b, e)) => { eof = e; (self.pending << 1) | (b as u64) },
                    Err(e) => { return Err(e); }
                }
            }
        }

        loop {
            {
                let range = self.high - self.low + 1;
                let count = self.get_count();
                let scaled = ((self.pending - self.low + 1) * count - 1) / range;
                let (chr, low, high) = self.get_char(scaled);

                if chr == VALUE_EOF {
                    break;
                } else {
                    try!(self.output.put_byte(chr as u8, ostream));
                }

                self.high = self.low + (range * high / count) - 1;
                self.low = self.low + (range * low / count);
            }

            loop {
                if self.high < ONE_HALF {
                    // do nothing
                } else if self.low >= ONE_HALF {
                    self.pending -= ONE_HALF;
                    self.low -= ONE_HALF;
                    self.high -= ONE_HALF;
                } else if self.low >= ONE_FOURTH && self.high < THREE_FOURTHS {
                    self.pending -= ONE_FOURTH;
                    self.low -= ONE_FOURTH;
                    self.high -= ONE_FOURTH;
                } else {
                    break;
                }

                self.low = self.low << 1;
                self.high = (self.high << 1) + 1;

                self.pending = if eof {
                    return Err(Error::new(ErrorKind::InvalidInput, "Unexpected end of input stream"));
                } else {
                    match self.input.get_bit(istream) {
                        Ok((b, e)) => { eof = e; (self.pending << 1) | (b as u64) },
                        Err(e) => { return Err(e); }
                    }
                }
            }
        }

        try!(self.output.flush_buffer(ostream));
        return Ok(());
    }
}

pub fn compress(istream: &mut Read, ostream: &mut Write) -> Result<(u64, u64)> {
    let mut status = Status::init();
    try!(status.compress(istream, ostream));
    return Ok(status.get_metrics());
}

pub fn decompress(istream: &mut Read, ostream: &mut Write) -> Result<(u64, u64)> {
    let mut status = Status::init();
    try!(status.decompress(istream, ostream));
    return Ok(status.get_metrics());
}

<|endoftext|>


<commit_before>use super::namepart::{Location, NamePart};
use super::suffix;
use super::surname;
use super::title;
use super::utils::is_mixed_case;
use smallvec::SmallVec;

struct ParseOp<'a> {
    words: SmallVec<[NamePart<'a>; 7]>,
    surname_index: usize,
    generation_from_suffix: Option<usize>,
    possible_false_prefix: Option<NamePart<'a>>,
    possible_false_postfix: Option<NamePart<'a>>,
    use_capitalization: bool,
}

pub fn parse(name: &str) -> Option<(SmallVec<[NamePart; 7]>, usize, Option<usize>)> {
    let mut op = ParseOp {
        words: SmallVec::new(),
        surname_index: 0,
        generation_from_suffix: None,
        possible_false_prefix: None,
        possible_false_postfix: None,
        use_capitalization: is_mixed_case(name),
    };

    if op.run(name) {
        Some((op.words, op.surname_index, op.generation_from_suffix))
    } else {
        None
    }
}

impl<'a> ParseOp<'a> {
    fn run(&mut self, name: &'a str) -> bool {
        // Separate comma-separated titles and suffixes, then flip remaining words
        // around remaining comma, if any
        if name.contains(',') {
            for part in name.split(',') {
                if self.words.is_empty() {
                    // We're in the surname part (if the format is "Smith, John"),
                    // or the only actual name part (if the format is "John Smith,
                    // esq." or just "John Smith")
                    self.handle_before_comma(part, true);
                } else if self.surname_index == 0 {
                    // We already processed one comma-separated part, but we think
                    // it was just the surname, so this might be the given name or
                    // initials
                    self.handle_after_comma(part);
                } else {
                    // We already found the full name, so this is a comma-separated
                    // postfix title or suffix
                    self.handle_after_surname(part);
                }
            }
        } else {
            self.handle_before_comma(name, false);
        }

        // If there are two or fewer words, e.g. "JOHN MA", we treat
        // ambiguous strings like "MA" as surnames rather than titles
        // (not initials, which should only be at the end of the input
        // if they are comma-separated, and we already handled that case)
        if !self.valid() {
            if let Some(ref postfix) = self.possible_false_postfix {
                self.words.push(postfix.clone());
            } else if let Some(ref prefix) = self.possible_false_prefix {
                self.words.insert(0, prefix.clone());
            }
        }

        // Anything trailing that looks like initials is probably a stray postfix
        while self.words.last().filter(|w| !w.is_namelike()).is_some() {
            let removed = self.words.pop().unwrap();

            // If we guessed the surname previously, our guess is no longer valid
            self.surname_index = 0;

            // If the alternative is having less than two words and giving up,
            // check if we can treat the last word as a name if we just ignore
            // case; this handles the not-quite-rare-enough case of an all-caps
            // last name (e.g.Neto John SMITH), among others
            if self.use_capitalization && !self.valid() {
                let word = NamePart::from_word(removed.word, false, Location::End);
                if word.is_namelike() {
                    self.words.push(word);
                    break;
                }
            }
        }

        // Handle case where we thought the whole before-comma part was a surname,
        // but we never found a plausible given name or initial afterwards,
        // as well as the reset just above
        if self.surname_index == 0 && self.words.len() > 1 {
            self.surname_index = surname::find_surname_index(&self.words[1..]) + 1;
        }

        debug_assert!(self
            .words
            .iter()
            .all(|w| w.is_namelike() || w.is_initials()));

        // Check the plausibility of what we've found
        self.valid()
    }

    fn valid(&self) -> bool {
        self.words.len() >= 2
            && self.words[self.surname_index..]
                .iter()
                .any(|w| w.is_namelike())
    }

    // Called only until any words are found
    fn handle_before_comma(&mut self, part: &'a str, any_after: bool) {
        debug_assert!(
            self.words.is_empty()
                && self.surname_index == 0
                && self.possible_false_prefix.is_none(),
            "Invalid state for handle_before_comma!"
        );

        self.words.extend(NamePart::all_from_text(
            part,
            self.use_capitalization,
            if any_after {
                Location::End
            } else {
                Location::Start
            },
        ));

        if self.words.is_empty() {
            return;
        }

        // Check for title as prefix (e.g. "Dr. John Smith" or "Right Hon.
        // John Smith")
        let prefix_title_len = title::find_prefix_len(&self.words);
        for i in (0..prefix_title_len).rev() {
            let word = self.words.remove(i);
            self.found_prefix(word);
        }

        // Strip non-comma-separated titles & suffixes (e.g. "John Smith Jr.")
        let first_postfix_index = title::find_postfix_index(&self.words[1..], false) + 1;
        if first_postfix_index < self.words.len() {
            let postfix = self.words.swap_remove(first_postfix_index);
            self.found_suffix_or_postfix(postfix, false);
            self.words.truncate(first_postfix_index);
        }

        if prefix_title_len > 0 || !any_after {
            // Finding a prefix title means the next word is a first name or
            // initial (we don't support "Dr. Smith, John")
            self.surname_index = surname::find_surname_index(&self.words[1..]) + 1;
        } else {
            // Have to guess whether this is just the surname (as in "Smith, John")
            // or the full name (as in "John Smith")
            //
            // Note we might be wrong, and have to go back, if we think the given
            // name is coming after a comma, but it never does
            self.surname_index = surname::find_surname_index(&self.words);
        }
    }

    // Called after the first comma, until we find a given name or first initial
    fn handle_after_comma(&mut self, part: &'a str) {
        debug_assert!(
            !self.words.is_empty() && self.surname_index == 0,
            "Invalid state for handle_after_comma!"
        );

        let mut given_middle_or_postfix_words: SmallVec<[NamePart<'a>; 5]> =
            NamePart::all_from_text(part, self.use_capitalization, Location::Start).collect();

        if given_middle_or_postfix_words.is_empty() {
            return;
        }

        // Handle (unusual) formats like "Smith, Dr. John M."
        if given_middle_or_postfix_words.len() > 1 {
            let prefix_title_len = title::find_prefix_len(&given_middle_or_postfix_words);
            for i in (0..prefix_title_len).rev() {
                self.found_prefix(given_middle_or_postfix_words.remove(i));
            }
        }

        // Handle isolated suffixes or titles as well as (unusual) formats like
        // "Smith, John Jr."
        let first_postfix_index = title::find_postfix_index(&given_middle_or_postfix_words, true);
        if first_postfix_index < given_middle_or_postfix_words.len() {
            let postfix = given_middle_or_postfix_words.swap_remove(first_postfix_index);
            self.found_suffix_or_postfix(postfix, true);
            given_middle_or_postfix_words.truncate(first_postfix_index);
        }

        // Now if there are any words left, they include the given name or first
        // initial (in a format like "Smith, John" or "Smith, J. M."), so we put
        // them in front
        if !given_middle_or_postfix_words.is_empty() {
            self.surname_index = given_middle_or_postfix_words.len();

            self.words.reserve(given_middle_or_postfix_words.len());
            for word in given_middle_or_postfix_words.into_iter().rev() {
                self.words.insert(0, word);
            }
        }
    }

    // Called on any parts remaining after full name is found
    fn handle_after_surname(&mut self, part: &'a str) {
        debug_assert!(
            self.surname_index > 0,
            "Invalid state for handle_after_surname!"
        );

        if self.possible_false_postfix.is_some() && self.generation_from_suffix.is_some() {
            return;
        }

        let mut postfix_words =
            NamePart::all_from_text(part, self.use_capitalization, Location::End);
        while self.possible_false_postfix.is_none() || self.generation_from_suffix.is_none() {
            if let Some(word) = postfix_words.next() {
                self.found_suffix_or_postfix(word, false);
            } else {
                break;
            }
        }
    }

    fn found_prefix(&mut self, prefix: NamePart<'a>) {
        // We drop prefixes, but keep the last word that's namelike,
        // just in case we make a mistake and it turns out by process of
        // elimination that this must actually be a given name
        if self.possible_false_prefix.is_none() && (prefix.is_namelike() || prefix.is_initials()) {
            self.possible_false_prefix = Some(prefix);
        }
    }

    fn found_suffix_or_postfix(&mut self, postfix: NamePart<'a>, expect_initials: bool) {
        if self.generation_from_suffix.is_none() {
            self.generation_from_suffix = suffix::generation_from_suffix(&postfix, expect_initials);
        }

        // We throw away most postfix titles, but keep the first one that's namelike,
        // just in case we make a mistake and it turns out by process of elimination
        // that this must actually be a surname
        if self.possible_false_postfix.is_none() && (postfix.is_namelike() || postfix.is_initials())
        {
            self.possible_false_postfix = Some(postfix);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use test::{black_box, Bencher};

    #[test]
    fn first_last() {
        let (parts, surname_index, generation) = parse("John Doe").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn initial_last() {
        let (parts, surname_index, generation) = parse("J. Doe").unwrap();
        assert_eq!("J.", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn last_first() {
        let (parts, surname_index, generation) = parse("Doe, John").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn last_initial() {
        let (parts, surname_index, generation) = parse("Doe, J.").unwrap();
        assert_eq!("J.", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn suffix() {
        let (parts, surname_index, generation) = parse("John Doe III").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(Some(3), generation);
    }

    #[test]
    fn suffix_comma() {
        let (parts, surname_index, generation) = parse("Doe, John III").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(Some(3), generation);
    }

    #[bench]
    fn parse_simple(b: &mut Bencher) {
        b.iter(|| black_box(parse("John Doe").is_some()))
    }

    #[bench]
    fn parse_nonascii(b: &mut Bencher) {
        b.iter(|| black_box(parse("이용희").is_some()))
    }

    #[bench]
    fn parse_comma(b: &mut Bencher) {
        b.iter(|| black_box(parse("Doe, John").is_some()))
    }

    #[bench]
    fn parse_all_caps(b: &mut Bencher) {
        b.iter(|| black_box(parse("JOHN DOE").is_some()))
    }

    #[bench]
    fn parse_complex(b: &mut Bencher) {
        b.iter(|| black_box(parse("James S. Brown MD, FRCS, FDSRCS").is_some()))
    }
}
<commit_msg>More fast pathing<commit_after>use super::namepart::{Location, NamePart};
use super::suffix;
use super::surname;
use super::title;
use super::utils::is_mixed_case;
use smallvec::SmallVec;

struct ParseOp<'a> {
    words: SmallVec<[NamePart<'a>; 7]>,
    surname_index: usize,
    generation_from_suffix: Option<usize>,
    possible_false_prefix: Option<NamePart<'a>>,
    possible_false_postfix: Option<NamePart<'a>>,
    use_capitalization: bool,
}

pub fn parse(name: &str) -> Option<(SmallVec<[NamePart; 7]>, usize, Option<usize>)> {
    let mut op = ParseOp {
        words: SmallVec::new(),
        surname_index: 0,
        generation_from_suffix: None,
        possible_false_prefix: None,
        possible_false_postfix: None,
        use_capitalization: is_mixed_case(name),
    };

    if op.run(name) {
        Some((op.words, op.surname_index, op.generation_from_suffix))
    } else {
        None
    }
}

impl<'a> ParseOp<'a> {
    fn run(&mut self, name: &'a str) -> bool {
        // Separate comma-separated titles and suffixes, then flip remaining words
        // around remaining comma, if any
        if name.contains(',') {
            for part in name.split(',') {
                if self.words.is_empty() {
                    // We're in the surname part (if the format is "Smith, John"),
                    // or the only actual name part (if the format is "John Smith,
                    // esq." or just "John Smith")
                    self.handle_before_comma(part, true);
                } else if self.surname_index == 0 {
                    // We already processed one comma-separated part, but we think
                    // it was just the surname, so this might be the given name or
                    // initials
                    self.handle_after_comma(part);
                } else {
                    // We already found the full name, so this is a comma-separated
                    // postfix title or suffix
                    self.handle_after_surname(part);
                }
            }
        } else {
            self.handle_before_comma(name, false);
        }

        // If there are two or fewer words, e.g. "JOHN MA", we treat
        // ambiguous strings like "MA" as surnames rather than titles
        // (not initials, which should only be at the end of the input
        // if they are comma-separated, and we already handled that case)
        if !self.valid() {
            if let Some(ref postfix) = self.possible_false_postfix {
                self.words.push(postfix.clone());
            } else if let Some(ref prefix) = self.possible_false_prefix {
                self.words.insert(0, prefix.clone());
            }
        }

        // Anything trailing that looks like initials is probably a stray postfix
        while self.words.last().filter(|w| !w.is_namelike()).is_some() {
            let removed = self.words.pop().unwrap();

            // If we guessed the surname previously, our guess is no longer valid
            self.surname_index = 0;

            // If the alternative is having less than two words and giving up,
            // check if we can treat the last word as a name if we just ignore
            // case; this handles the not-quite-rare-enough case of an all-caps
            // last name (e.g.Neto John SMITH), among others
            if self.use_capitalization && !self.valid() {
                let word = NamePart::from_word(removed.word, false, Location::End);
                if word.is_namelike() {
                    self.words.push(word);
                    break;
                }
            }
        }

        // Handle case where we thought the whole before-comma part was a surname,
        // but we never found a plausible given name or initial afterwards,
        // as well as the reset just above
        if self.surname_index == 0 && self.words.len() > 1 {
            self.surname_index = surname::find_surname_index(&self.words[1..]) + 1;
        }

        // Check the plausibility of what we've found
        self.valid()
    }

    fn valid(&self) -> bool {
        self.words.len() >= 2
            && self
                .words
                .iter()
                .all(|w| w.is_namelike() || w.is_initials())
            && self.words[self.surname_index..]
                .iter()
                .any(|w| w.is_namelike())
    }

    // Called only until any words are found
    fn handle_before_comma(&mut self, part: &'a str, any_after: bool) {
        debug_assert!(
            self.words.is_empty()
                && self.surname_index == 0
                && self.possible_false_prefix.is_none(),
            "Invalid state for handle_before_comma!"
        );

        self.words.extend(NamePart::all_from_text(
            part,
            self.use_capitalization,
            if any_after {
                Location::End
            } else {
                Location::Start
            },
        ));

        if self.words.is_empty() {
            return;
        }

        // Check for title as prefix (e.g. "Dr. John Smith" or "Right Hon.
        // John Smith")
        let prefix_title_len = if any_after || self.words.len() > 2 {
            title::find_prefix_len(&self.words)
        } else {
            0
        };
        for i in (0..prefix_title_len).rev() {
            let word = self.words.remove(i);
            self.found_prefix(word);
        }

        // Strip non-comma-separated titles & suffixes (e.g. "John Smith Jr.")
        let first_postfix_index =
            if any_after || self.words.len() + self.possible_false_prefix.iter().count() > 2 {
                title::find_postfix_index(&self.words[1..], false) + 1
            } else {
                self.words.len()
            };
        if first_postfix_index < self.words.len() {
            let postfix = self.words.swap_remove(first_postfix_index);
            self.found_suffix_or_postfix(postfix, false);
            self.words.truncate(first_postfix_index);
        }

        if prefix_title_len > 0 || !any_after {
            // Finding a prefix title means the next word is a first name or
            // initial (we don't support "Dr. Smith, John")
            self.surname_index = surname::find_surname_index(&self.words[1..]) + 1;
        } else {
            // Have to guess whether this is just the surname (as in "Smith, John")
            // or the full name (as in "John Smith")
            //
            // Note we might be wrong, and have to go back, if we think the given
            // name is coming after a comma, but it never does
            self.surname_index = surname::find_surname_index(&self.words);
        }
    }

    // Called after the first comma, until we find a given name or first initial
    fn handle_after_comma(&mut self, part: &'a str) {
        debug_assert!(
            !self.words.is_empty() && self.surname_index == 0,
            "Invalid state for handle_after_comma!"
        );

        let mut given_middle_or_postfix_words: SmallVec<[NamePart<'a>; 5]> =
            NamePart::all_from_text(part, self.use_capitalization, Location::Start).collect();

        if given_middle_or_postfix_words.is_empty() {
            return;
        }

        // Handle (unusual) formats like "Smith, Dr. John M."
        if given_middle_or_postfix_words.len() > 1 {
            let prefix_title_len = title::find_prefix_len(&given_middle_or_postfix_words);
            for i in (0..prefix_title_len).rev() {
                self.found_prefix(given_middle_or_postfix_words.remove(i));
            }
        }

        // Handle isolated suffixes or titles as well as (unusual) formats like
        // "Smith, John Jr."
        let first_postfix_index = title::find_postfix_index(&given_middle_or_postfix_words, true);
        if first_postfix_index < given_middle_or_postfix_words.len() {
            let postfix = given_middle_or_postfix_words.swap_remove(first_postfix_index);
            self.found_suffix_or_postfix(postfix, true);
            given_middle_or_postfix_words.truncate(first_postfix_index);
        }

        // Now if there are any words left, they include the given name or first
        // initial (in a format like "Smith, John" or "Smith, J. M."), so we put
        // them in front
        if !given_middle_or_postfix_words.is_empty() {
            self.surname_index = given_middle_or_postfix_words.len();

            self.words.reserve(given_middle_or_postfix_words.len());
            for word in given_middle_or_postfix_words.into_iter().rev() {
                self.words.insert(0, word);
            }
        }
    }

    // Called on any parts remaining after full name is found
    fn handle_after_surname(&mut self, part: &'a str) {
        debug_assert!(
            self.surname_index > 0,
            "Invalid state for handle_after_surname!"
        );

        if self.possible_false_postfix.is_some() && self.generation_from_suffix.is_some() {
            return;
        }

        let mut postfix_words =
            NamePart::all_from_text(part, self.use_capitalization, Location::End);
        while self.possible_false_postfix.is_none() || self.generation_from_suffix.is_none() {
            if let Some(word) = postfix_words.next() {
                self.found_suffix_or_postfix(word, false);
            } else {
                break;
            }
        }
    }

    fn found_prefix(&mut self, prefix: NamePart<'a>) {
        // We drop prefixes, but keep the last word that's namelike,
        // just in case we make a mistake and it turns out by process of
        // elimination that this must actually be a given name
        if self.possible_false_prefix.is_none() && (prefix.is_namelike() || prefix.is_initials()) {
            self.possible_false_prefix = Some(prefix);
        }
    }

    fn found_suffix_or_postfix(&mut self, postfix: NamePart<'a>, expect_initials: bool) {
        if self.generation_from_suffix.is_none() {
            self.generation_from_suffix = suffix::generation_from_suffix(&postfix, expect_initials);
        }

        // We throw away most postfix titles, but keep the first one that's namelike,
        // just in case we make a mistake and it turns out by process of elimination
        // that this must actually be a surname
        if self.possible_false_postfix.is_none() && (postfix.is_namelike() || postfix.is_initials())
        {
            self.possible_false_postfix = Some(postfix);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use test::{black_box, Bencher};

    #[test]
    fn first_last() {
        let (parts, surname_index, generation) = parse("John Doe").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn initial_last() {
        let (parts, surname_index, generation) = parse("J. Doe").unwrap();
        assert_eq!("J.", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn last_first() {
        let (parts, surname_index, generation) = parse("Doe, John").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn last_initial() {
        let (parts, surname_index, generation) = parse("Doe, J.").unwrap();
        assert_eq!("J.", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(None, generation);
    }

    #[test]
    fn suffix() {
        let (parts, surname_index, generation) = parse("John Doe III").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(Some(3), generation);
    }

    #[test]
    fn suffix_comma() {
        let (parts, surname_index, generation) = parse("Doe, John III").unwrap();
        assert_eq!("John", parts[0].word);
        assert_eq!("Doe", parts[1].word);
        assert_eq!(1, surname_index);
        assert_eq!(Some(3), generation);
    }

    #[bench]
    fn parse_simple(b: &mut Bencher) {
        b.iter(|| black_box(parse("John Doe").is_some()))
    }

    #[bench]
    fn parse_nonascii(b: &mut Bencher) {
        b.iter(|| black_box(parse("이용희").is_some()))
    }

    #[bench]
    fn parse_comma(b: &mut Bencher) {
        b.iter(|| black_box(parse("Doe, John").is_some()))
    }

    #[bench]
    fn parse_all_caps(b: &mut Bencher) {
        b.iter(|| black_box(parse("JOHN DOE").is_some()))
    }

    #[bench]
    fn parse_complex(b: &mut Bencher) {
        b.iter(|| black_box(parse("James S. Brown MD, FRCS, FDSRCS").is_some()))
    }
}
<|endoftext|>