<commit_before>use std::io::{Read, BufRead, Error as IoError, ErrorKind};
use std::str::{from_utf8, Utf8Error};
use std::num::{ParseFloatError};

use num::bigint::{BigInt, ToBigInt, Sign, ParseBigIntError};
use byteorder::{ReadBytesExt, LittleEndian, BigEndian, Error as ByteorderError};

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        ReadError(err: ByteorderError) {
            from()
        }
        IoError(err: IoError) {
            from()
        }
        UnknownOpcode(opcode: u8) {}
        InvalidInt {
            from(Utf8Error)
            from(ParseBigIntError)
        }
        InvalidFloat {
            from(ParseFloatError)
        }
        InvalidString
        ExpectedTrailingL
        InvalidLong
        NegativeLength
    }
}

#[derive(Debug)]
pub enum OpCode {
    Int(BigInt),
    BinInt(i32),
    BinInt1(u8),
    BinInt2(u16),
    Long(BigInt),
    Long1(BigInt),
    Long4(BigInt),

    String(Vec<u8>),
    BinString(Vec<u8>),
    ShortBinString(Vec<u8>),

    None,

    NewTrue,
    NewFalse,

    Unicode(String),
    BinUnicode(String),

    Float(f64),
    BinFloat(f64),

    EmptyList,
    Append,
    Appends,
    List,

    EmptyTuple,
    Tuple,
    Tuple1,
    Tuple2,
    Tuple3,

    EmptyDict,
    Dict,
    SetItem,
    SetItems,

    Pop,
    Dup,
    Mark,
    PopMark,

    Get(usize),
    BinGet(usize),
    LongBinGet(usize),
    Put(usize),
    BinPut(usize),
    LongBinPut(usize),

    Ext1(u8),
    Ext2(u16),
    Ext4(i32),

    Global(Vec<u8>, Vec<u8>),
    Reduce,
    Build,
    Inst(Vec<u8>, Vec<u8>),
    Obj,
    NewObj,
    Proto(u8),
    Stop,
    PersId(Vec<u8>),
    BinPersId,
}

fn read_exact<R>(rd: &mut R, mut buf: &mut [u8]) -> Result<(), IoError> where R: Read {
    while !buf.is_empty() {
        match rd.read(buf) {
            Ok(0) => break,
            Ok(n) => { let tmp = buf; buf = &mut tmp[n..]; }
            Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    if !buf.is_empty() {
        Err(IoError::new(ErrorKind::Other,
                       "failed to fill whole buffer"))
    } else {
        Ok(())
    }
}

pub fn read_until_newline<R>(rd: &mut R) -> Result<Vec<u8>, Error> where R: Read + BufRead {
    let mut buf = Vec::new();
    try!(rd.read_until('\n' as u8, &mut buf));

    // Skip last symbol — \n
    match buf.split_last() {
        None => Err(Error::InvalidString),
        Some((_last, init)) => Ok(init.to_vec()),
    }
}

pub fn read_long<R>(rd: &mut R, length: usize) -> Result<BigInt, Error> where R: Read + BufRead {
    let mut buf = vec![0; length];
    try!(read_exact(rd, buf.as_mut()));

    let mut n = BigInt::from_bytes_le(Sign::Plus, &buf);

    let last = match buf.last_mut() {
        None => return Err(Error::InvalidLong),
        Some(last) => last,
    };

    if *last > 127 {
        n = n - (1.to_bigint().unwrap() << (length * 8))
    }

    Ok(n)
}

pub fn read_opcode<R>(rd: &mut R) -> Result<OpCode, Error> where R: Read + BufRead {
    let marker = try!(rd.read_u8());
    return Ok(match marker {
        73 => {
            let s = try!(read_until_newline(rd));
            OpCode::Int(try!(try!(from_utf8(&s)).parse()))  // http://rust-num.github.io/num/num/bigint/struct.BigInt.html#method.parse_bytes
        },
        74 => OpCode::BinInt(try!(rd.read_i32::<LittleEndian>())),
        75 => OpCode::BinInt1(try!(rd.read_u8())),
        77 => OpCode::BinInt2(try!(rd.read_u16::<LittleEndian>())),
        76 => {
            let s = try!(read_until_newline(rd));

            let init = match s.split_last() {
                None => return Err(Error::InvalidString),
                Some((&76, init)) => init,
                Some(_) => return Err(Error::ExpectedTrailingL),
            };

            OpCode::Long(try!(try!(from_utf8(init)).parse()))
        },
        138 => {
            let length = try!(rd.read_u8());
            OpCode::Long1(try!(read_long(rd, length as usize)))

        },
        139 => {
            let length = try!(rd.read_i32::<LittleEndian>());
            if length < 0 {
                return Err(Error::NegativeLength)
            }
            OpCode::Long4(try!(read_long(rd, length as usize)))

        },
        83 => {OpCode::String(try!(read_until_newline(rd)))} // TODO: escaping
        84 => {
            let length = try!(rd.read_i32::<LittleEndian>());
            if length < 0 {
                return Err(Error::NegativeLength)
            }
            let mut buf = vec![0; length as usize];
            try!(read_exact(rd, &mut buf));
            OpCode::BinString(buf)
        }
        85 => {
            let length = try!(rd.read_u8());
            let mut buf = vec![0; length as usize];
            try!(read_exact(rd, &mut buf));
            OpCode::ShortBinString(buf)
        }
        78 => OpCode::None,
        136 => OpCode::NewTrue,
        137 => OpCode::NewFalse,
        86 => unimplemented!(), // Unicode
        88 => unimplemented!(), // BinUnicode
        70 => {
            let s = try!(read_until_newline(rd));
            OpCode::Float(try!(try!(from_utf8(&s)).parse()))
        },
        71 => {
            OpCode::BinFloat(try!(rd.read_f64::<BigEndian>()))
        }
        c => return Err(Error::UnknownOpcode(c)),
    })
}

#[cfg(test)]
mod tests {
    use std::io::{Cursor};

    use num::{FromPrimitive};

    use super::{OpCode, Error, read_opcode};

    macro_rules! t {
        ($buffer: expr, $pat:pat, $result:expr) => ({
            match read_opcode(&mut Cursor::new(&$buffer[..])) {
                $pat => $result,
                other => {
                    println!("ERROR {:?}", other);
                    assert!(false)
                },
            }
        })
    }

    macro_rules! n {
        ($x: expr) => ({FromPrimitive::from_isize($x).unwrap()})
    }

    #[test]
    fn test_int() {
        t!(b"I", Err(Error::InvalidString), assert!(true));
        t!(b"I\n", Err(Error::InvalidInt), assert!(true));
        t!(b"Iabc\n", Err(Error::InvalidInt), assert!(true));
        t!(b"I123\n", Ok(OpCode::Int(n)), assert_eq!(n, n!(123)));
    }

    #[test]
    fn test_bin_int() {
        t!(b"J\x0a", Err(Error::ReadError(_)), assert!(true));
        t!(b"J\x0a\x00\x00\x00", Ok(OpCode::BinInt(n)), assert_eq!(n, 10));
        t!(b"J\x0a\x00\x00\x01", Ok(OpCode::BinInt(n)), assert_eq!(n, 16777226));
    }

    #[test]
    fn test_bin_int1() {
        t!(b"K", Err(Error::ReadError(_)), assert!(true));
        t!(b"K\x0a", Ok(OpCode::BinInt1(n)), assert_eq!(n, 10));
    }

    #[test]
    fn test_bin_int2() {
        t!(b"M\x0a", Err(Error::ReadError(_)), assert!(true));
        t!(b"M\x0a\x00\x00\x00", Ok(OpCode::BinInt2(n)), assert_eq!(n, 10));
        t!(b"M\x0a\x01\x00\x00", Ok(OpCode::BinInt2(n)), assert_eq!(n, 266));
    }

    #[test]
    fn test_long() {
        t!(b"L", Err(Error::InvalidString), assert!(true));
        t!(b"L\n", Err(Error::InvalidString), assert!(true));
        t!(b"Labc\n", Err(Error::ExpectedTrailingL), assert!(true));
        t!(b"LabcL\n", Err(Error::InvalidInt), assert!(true));
        t!(b"L123L\n", Ok(OpCode::Long(n)), assert_eq!(n, n!(123)));
    }

    #[test]
    fn test_long1() {
        t!(b"\x8a", Err(Error::ReadError(_)), assert!(true));
        t!(b"\x8a\x01\x0a", Ok(OpCode::Long1(n)), assert_eq!(n, n!(10)));
        t!(b"\x8a\x01\xf6", Ok(OpCode::Long1(n)), assert_eq!(n, n!(-10)));
        t!(b"\x8a\x02.\xfb", Ok(OpCode::Long1(n)), assert_eq!(n, n!(-1234)));
    }

    #[test]
    fn test_long4() {
        t!(b"\x8b\xff\xff\xff\xff", Err(Error::NegativeLength), assert!(true));
        t!(b"\x8b\x0a", Err(Error::ReadError(_)), assert!(true));
        t!(b"\x8b\x01\x00\x00\x00\x0a", Ok(OpCode::Long4(n)), assert_eq!(n, n!(10)));
        t!(b"\x8b\x01\x00\x00\x00\xf6", Ok(OpCode::Long4(n)), assert_eq!(n, n!(-10)));
        t!(b"\x8b\x02\x00\x00\x00.\xfb", Ok(OpCode::Long4(n)), assert_eq!(n, n!(-1234)));
    }

    #[test]
    fn test_string() {
        t!(b"S", Err(Error::InvalidString), assert!(true));
        t!(b"S\n", Ok(OpCode::String(s)), assert_eq!(s, b""));
        t!(b"Sabc\n", Ok(OpCode::String(s)), assert_eq!(s, b"abc"));
        t!(b"S123\n", Ok(OpCode::String(s)), assert_eq!(s, b"123"));
        t!(b"S\\n\n", Ok(OpCode::String(s)), assert_eq!(s, b"\\n"));
    }

    #[test]
    fn test_bin_string() {
        t!(b"T\xff\xff\xff\xff", Err(Error::NegativeLength), assert!(true));
        t!(b"T\x00\x00\x00\x00", Ok(OpCode::BinString(s)), assert_eq!(s, b""));
        t!(b"T\x03\x00\x00\x00abc", Ok(OpCode::BinString(s)), assert_eq!(s, b"abc"));
        t!(b"T\x03\x00\x00\x00123", Ok(OpCode::BinString(s)), assert_eq!(s, b"123"));
        t!(b"T\x02\x00\x00\x00\\n", Ok(OpCode::BinString(s)), assert_eq!(s, b"\\n"));
    }

    #[test]
    fn test_short_bin_string() {
        t!(b"U\x00", Ok(OpCode::ShortBinString(s)), assert_eq!(s, b""));
        t!(b"U\x03abc", Ok(OpCode::ShortBinString(s)), assert_eq!(s, b"abc"));
        t!(b"U\x03123", Ok(OpCode::ShortBinString(s)), assert_eq!(s, b"123"));
        t!(b"U\x02\\n", Ok(OpCode::ShortBinString(s)), assert_eq!(s, b"\\n"));
    }

    #[test]
    fn test_none() {
        t!(b"N", Ok(OpCode::None), assert!(true));
    }

    #[test]
    fn test_new_true() {
        t!(b"\x88", Ok(OpCode::NewTrue), assert!(true));
    }

    #[test]
    fn test_new_false() {
        t!(b"\x89", Ok(OpCode::NewFalse), assert!(true));
    }

    #[test]
    fn test_unicode() {
    }

    #[test]
    fn test_bin_unicode() {
    }

    #[test]
    fn test_float() {
        t!(b"F", Err(Error::InvalidString), assert!(true));
        t!(b"F\n", Err(Error::InvalidFloat), assert!(true));
        t!(b"Fabc\n", Err(Error::InvalidFloat), assert!(true));
        t!(b"F123\n", Ok(OpCode::Float(n)), assert_eq!(n, 123.0));
        t!(b"F-123\n", Ok(OpCode::Float(n)), assert_eq!(n, -123.0));
        t!(b"F-123.\n", Ok(OpCode::Float(n)), assert_eq!(n, -123.0));
        t!(b"F-123.456\n", Ok(OpCode::Float(n)), assert_eq!(n, -123.456));
    }

    #[test]
    fn test_bin_float() {
        t!(b"G", Err(Error::ReadError(_)), assert!(true));
        t!(b"Gabc", Err(Error::ReadError(_)), assert!(true));
        t!(b"G123", Err(Error::ReadError(_)), assert!(true));
        t!(b"G@^\xc0\x00\x00\x00\x00\x00", Ok(OpCode::BinFloat(n)), assert_eq!(n, 123.0));
        t!(b"G\xc0^\xc0\x00\x00\x00\x00\x00", Ok(OpCode::BinFloat(n)), assert_eq!(n, -123.0));
        t!(b"G\xc0^\xdd/\x1a\x9f\xbew", Ok(OpCode::BinFloat(n)), assert_eq!(n, -123.456));
    }

}
<commit_msg>Add list opcodes<commit_after>use std::io::{Read, BufRead, Error as IoError, ErrorKind};
use std::str::{from_utf8, Utf8Error};
use std::num::{ParseFloatError};

use num::bigint::{BigInt, ToBigInt, Sign, ParseBigIntError};
use byteorder::{ReadBytesExt, LittleEndian, BigEndian, Error as ByteorderError};

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        ReadError(err: ByteorderError) {
            from()
        }
        IoError(err: IoError) {
            from()
        }
        UnknownOpcode(opcode: u8) {}
        InvalidInt {
            from(Utf8Error)
            from(ParseBigIntError)
        }
        InvalidFloat {
            from(ParseFloatError)
        }
        InvalidString
        ExpectedTrailingL
        InvalidLong
        NegativeLength
    }
}

#[derive(Debug)]
pub enum OpCode {
    Int(BigInt),
    BinInt(i32),
    BinInt1(u8),
    BinInt2(u16),
    Long(BigInt),
    Long1(BigInt),
    Long4(BigInt),

    String(Vec<u8>),
    BinString(Vec<u8>),
    ShortBinString(Vec<u8>),

    None,

    NewTrue,
    NewFalse,

    Unicode(String),
    BinUnicode(String),

    Float(f64),
    BinFloat(f64),

    EmptyList,
    Append,
    Appends,
    List,

    EmptyTuple,
    Tuple,
    Tuple1,
    Tuple2,
    Tuple3,

    EmptyDict,
    Dict,
    SetItem,
    SetItems,

    Pop,
    Dup,
    Mark,
    PopMark,

    Get(usize),
    BinGet(usize),
    LongBinGet(usize),
    Put(usize),
    BinPut(usize),
    LongBinPut(usize),

    Ext1(u8),
    Ext2(u16),
    Ext4(i32),

    Global(Vec<u8>, Vec<u8>),
    Reduce,
    Build,
    Inst(Vec<u8>, Vec<u8>),
    Obj,
    NewObj,
    Proto(u8),
    Stop,
    PersId(Vec<u8>),
    BinPersId,
}

fn read_exact<R>(rd: &mut R, mut buf: &mut [u8]) -> Result<(), IoError> where R: Read {
    while !buf.is_empty() {
        match rd.read(buf) {
            Ok(0) => break,
            Ok(n) => { let tmp = buf; buf = &mut tmp[n..]; }
            Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    if !buf.is_empty() {
        Err(IoError::new(ErrorKind::Other,
                       "failed to fill whole buffer"))
    } else {
        Ok(())
    }
}

pub fn read_until_newline<R>(rd: &mut R) -> Result<Vec<u8>, Error> where R: Read + BufRead {
    let mut buf = Vec::new();
    try!(rd.read_until('\n' as u8, &mut buf));

    // Skip last symbol — \n
    match buf.split_last() {
        None => Err(Error::InvalidString),
        Some((_last, init)) => Ok(init.to_vec()),
    }
}

pub fn read_long<R>(rd: &mut R, length: usize) -> Result<BigInt, Error> where R: Read + BufRead {
    let mut buf = vec![0; length];
    try!(read_exact(rd, buf.as_mut()));

    let mut n = BigInt::from_bytes_le(Sign::Plus, &buf);

    let last = match buf.last_mut() {
        None => return Err(Error::InvalidLong),
        Some(last) => last,
    };

    if *last > 127 {
        n = n - (1.to_bigint().unwrap() << (length * 8))
    }

    Ok(n)
}

pub fn read_opcode<R>(rd: &mut R) -> Result<OpCode, Error> where R: Read + BufRead {
    let marker = try!(rd.read_u8());
    return Ok(match marker {
        73 => {
            let s = try!(read_until_newline(rd));
            OpCode::Int(try!(try!(from_utf8(&s)).parse()))  // http://rust-num.github.io/num/num/bigint/struct.BigInt.html#method.parse_bytes
        },
        74 => OpCode::BinInt(try!(rd.read_i32::<LittleEndian>())),
        75 => OpCode::BinInt1(try!(rd.read_u8())),
        77 => OpCode::BinInt2(try!(rd.read_u16::<LittleEndian>())),
        76 => {
            let s = try!(read_until_newline(rd));

            let init = match s.split_last() {
                None => return Err(Error::InvalidString),
                Some((&76, init)) => init,
                Some(_) => return Err(Error::ExpectedTrailingL),
            };

            OpCode::Long(try!(try!(from_utf8(init)).parse()))
        },
        138 => {
            let length = try!(rd.read_u8());
            OpCode::Long1(try!(read_long(rd, length as usize)))

        },
        139 => {
            let length = try!(rd.read_i32::<LittleEndian>());
            if length < 0 {
                return Err(Error::NegativeLength)
            }
            OpCode::Long4(try!(read_long(rd, length as usize)))

        },
        83 => {OpCode::String(try!(read_until_newline(rd)))} // TODO: escaping
        84 => {
            let length = try!(rd.read_i32::<LittleEndian>());
            if length < 0 {
                return Err(Error::NegativeLength)
            }
            let mut buf = vec![0; length as usize];
            try!(read_exact(rd, &mut buf));
            OpCode::BinString(buf)
        }
        85 => {
            let length = try!(rd.read_u8());
            let mut buf = vec![0; length as usize];
            try!(read_exact(rd, &mut buf));
            OpCode::ShortBinString(buf)
        }
        78 => OpCode::None,
        136 => OpCode::NewTrue,
        137 => OpCode::NewFalse,
        86 => unimplemented!(), // Unicode
        88 => unimplemented!(), // BinUnicode
        70 => {
            let s = try!(read_until_newline(rd));
            OpCode::Float(try!(try!(from_utf8(&s)).parse()))
        },
        71 => {
            OpCode::BinFloat(try!(rd.read_f64::<BigEndian>()))
        }
        93 => OpCode::EmptyList,
        97 => OpCode::Append,
        101 => OpCode::Appends,
        108 => OpCode::List,
        c => return Err(Error::UnknownOpcode(c)),
    })
}

#[cfg(test)]
mod tests {
    use std::io::{Cursor};

    use num::{FromPrimitive};

    use super::{OpCode, Error, read_opcode};

    macro_rules! t {
        ($buffer: expr, $pat:pat, $result:expr) => ({
            match read_opcode(&mut Cursor::new(&$buffer[..])) {
                $pat => $result,
                other => {
                    println!("ERROR {:?}", other);
                    assert!(false)
                },
            }
        })
    }

    macro_rules! n {
        ($x: expr) => ({FromPrimitive::from_isize($x).unwrap()})
    }

    #[test]
    fn test_int() {
        t!(b"I", Err(Error::InvalidString), assert!(true));
        t!(b"I\n", Err(Error::InvalidInt), assert!(true));
        t!(b"Iabc\n", Err(Error::InvalidInt), assert!(true));
        t!(b"I123\n", Ok(OpCode::Int(n)), assert_eq!(n, n!(123)));
    }

    #[test]
    fn test_bin_int() {
        t!(b"J\x0a", Err(Error::ReadError(_)), assert!(true));
        t!(b"J\x0a\x00\x00\x00", Ok(OpCode::BinInt(n)), assert_eq!(n, 10));
        t!(b"J\x0a\x00\x00\x01", Ok(OpCode::BinInt(n)), assert_eq!(n, 16777226));
    }

    #[test]
    fn test_bin_int1() {
        t!(b"K", Err(Error::ReadError(_)), assert!(true));
        t!(b"K\x0a", Ok(OpCode::BinInt1(n)), assert_eq!(n, 10));
    }

    #[test]
    fn test_bin_int2() {
        t!(b"M\x0a", Err(Error::ReadError(_)), assert!(true));
        t!(b"M\x0a\x00\x00\x00", Ok(OpCode::BinInt2(n)), assert_eq!(n, 10));
        t!(b"M\x0a\x01\x00\x00", Ok(OpCode::BinInt2(n)), assert_eq!(n, 266));
    }

    #[test]
    fn test_long() {
        t!(b"L", Err(Error::InvalidString), assert!(true));
        t!(b"L\n", Err(Error::InvalidString), assert!(true));
        t!(b"Labc\n", Err(Error::ExpectedTrailingL), assert!(true));
        t!(b"LabcL\n", Err(Error::InvalidInt), assert!(true));
        t!(b"L123L\n", Ok(OpCode::Long(n)), assert_eq!(n, n!(123)));
    }

    #[test]
    fn test_long1() {
        t!(b"\x8a", Err(Error::ReadError(_)), assert!(true));
        t!(b"\x8a\x01\x0a", Ok(OpCode::Long1(n)), assert_eq!(n, n!(10)));
        t!(b"\x8a\x01\xf6", Ok(OpCode::Long1(n)), assert_eq!(n, n!(-10)));
        t!(b"\x8a\x02.\xfb", Ok(OpCode::Long1(n)), assert_eq!(n, n!(-1234)));
    }

    #[test]
    fn test_long4() {
        t!(b"\x8b\xff\xff\xff\xff", Err(Error::NegativeLength), assert!(true));
        t!(b"\x8b\x0a", Err(Error::ReadError(_)), assert!(true));
        t!(b"\x8b\x01\x00\x00\x00\x0a", Ok(OpCode::Long4(n)), assert_eq!(n, n!(10)));
        t!(b"\x8b\x01\x00\x00\x00\xf6", Ok(OpCode::Long4(n)), assert_eq!(n, n!(-10)));
        t!(b"\x8b\x02\x00\x00\x00.\xfb", Ok(OpCode::Long4(n)), assert_eq!(n, n!(-1234)));
    }

    #[test]
    fn test_string() {
        t!(b"S", Err(Error::InvalidString), assert!(true));
        t!(b"S\n", Ok(OpCode::String(s)), assert_eq!(s, b""));
        t!(b"Sabc\n", Ok(OpCode::String(s)), assert_eq!(s, b"abc"));
        t!(b"S123\n", Ok(OpCode::String(s)), assert_eq!(s, b"123"));
        t!(b"S\\n\n", Ok(OpCode::String(s)), assert_eq!(s, b"\\n"));
    }

    #[test]
    fn test_bin_string() {
        t!(b"T\xff\xff\xff\xff", Err(Error::NegativeLength), assert!(true));
        t!(b"T\x00\x00\x00\x00", Ok(OpCode::BinString(s)), assert_eq!(s, b""));
        t!(b"T\x03\x00\x00\x00abc", Ok(OpCode::BinString(s)), assert_eq!(s, b"abc"));
        t!(b"T\x03\x00\x00\x00123", Ok(OpCode::BinString(s)), assert_eq!(s, b"123"));
        t!(b"T\x02\x00\x00\x00\\n", Ok(OpCode::BinString(s)), assert_eq!(s, b"\\n"));
    }

    #[test]
    fn test_short_bin_string() {
        t!(b"U\x00", Ok(OpCode::ShortBinString(s)), assert_eq!(s, b""));
        t!(b"U\x03abc", Ok(OpCode::ShortBinString(s)), assert_eq!(s, b"abc"));
        t!(b"U\x03123", Ok(OpCode::ShortBinString(s)), assert_eq!(s, b"123"));
        t!(b"U\x02\\n", Ok(OpCode::ShortBinString(s)), assert_eq!(s, b"\\n"));
    }

    #[test]
    fn test_none() {
        t!(b"N", Ok(OpCode::None), assert!(true));
    }

    #[test]
    fn test_new_true() {
        t!(b"\x88", Ok(OpCode::NewTrue), assert!(true));
    }

    #[test]
    fn test_new_false() {
        t!(b"\x89", Ok(OpCode::NewFalse), assert!(true));
    }

    #[test]
    fn test_unicode() {
    }

    #[test]
    fn test_bin_unicode() {
    }

    #[test]
    fn test_float() {
        t!(b"F", Err(Error::InvalidString), assert!(true));
        t!(b"F\n", Err(Error::InvalidFloat), assert!(true));
        t!(b"Fabc\n", Err(Error::InvalidFloat), assert!(true));
        t!(b"F123\n", Ok(OpCode::Float(n)), assert_eq!(n, 123.0));
        t!(b"F-123\n", Ok(OpCode::Float(n)), assert_eq!(n, -123.0));
        t!(b"F-123.\n", Ok(OpCode::Float(n)), assert_eq!(n, -123.0));
        t!(b"F-123.456\n", Ok(OpCode::Float(n)), assert_eq!(n, -123.456));
    }

    #[test]
    fn test_bin_float() {
        t!(b"G", Err(Error::ReadError(_)), assert!(true));
        t!(b"Gabc", Err(Error::ReadError(_)), assert!(true));
        t!(b"G123", Err(Error::ReadError(_)), assert!(true));
        t!(b"G@^\xc0\x00\x00\x00\x00\x00", Ok(OpCode::BinFloat(n)), assert_eq!(n, 123.0));
        t!(b"G\xc0^\xc0\x00\x00\x00\x00\x00", Ok(OpCode::BinFloat(n)), assert_eq!(n, -123.0));
        t!(b"G\xc0^\xdd/\x1a\x9f\xbew", Ok(OpCode::BinFloat(n)), assert_eq!(n, -123.456));
    }

    #[test]
    fn test_empty_list() {
        t!(b"]", Ok(OpCode::EmptyList), assert!(true));
    }

    #[test]
    fn test_append() {
        t!(b"a", Ok(OpCode::Append), assert!(true));
    }

    #[test]
    fn test_appends() {
        t!(b"e", Ok(OpCode::Appends), assert!(true));
    }

    #[test]
    fn test_list() {
        t!(b"l", Ok(OpCode::List), assert!(true));
    }

}
<|endoftext|>


<commit_before>use super::helpers::allocations;
use super::helpers::fixtures::{get_language, get_language_queries_path};
use std::ffi::CStr;
use std::ffi::CString;
use std::{fs, ptr, slice, str};
use tree_sitter::Point;
use tree_sitter_tags::c_lib as c;
use tree_sitter_tags::{Error, TagsConfiguration, TagsContext};

const PYTHON_TAG_QUERY: &'static str = r#"
(
  (function_definition
    name: (identifier) @name
    body: (block . (expression_statement (string) @doc))) @definition.function
  (#strip! @doc "(^['\"\\s]*)|(['\"\\s]*$)")
)

(function_definition
  name: (identifier) @name) @definition.function

(
  (class_definition
    name: (identifier) @name
    body: (block
      . (expression_statement (string) @doc))) @definition.class
  (#strip! @doc "(^['\"\\s]*)|(['\"\\s]*$)")
)

(class_definition
  name: (identifier) @name) @definition.class

(call
  function: (identifier) @name) @reference.call

(call
  function: (attribute
    attribute: (identifier) @name)) @reference.call
"#;

const JS_TAG_QUERY: &'static str = r#"
(
    (comment)* @doc .
    (class_declaration
        name: (identifier) @name) @definition.class
    (#select-adjacent! @doc @definition.class)
    (#strip! @doc "(^[/\\*\\s]*)|([/\\*\\s]*$)")
)

(
    (comment)* @doc .
    (method_definition
        name: (property_identifier) @name) @definition.method
    (#select-adjacent! @doc @definition.method)
    (#strip! @doc "(^[/\\*\\s]*)|([/\\*\\s]*$)")
)

(
    (comment)* @doc .
    (function_declaration
        name: (identifier) @name) @definition.function
    (#select-adjacent! @doc @definition.function)
    (#strip! @doc "(^[/\\*\\s]*)|([/\\*\\s]*$)")
)

(call_expression
    function: (identifier) @name) @reference.call
"#;

const RUBY_TAG_QUERY: &'static str = r#"
(method
    name: (_) @name) @definition.method

(method_call
    method: (identifier) @name) @reference.call

(setter (identifier) @ignore)

((identifier) @name @reference.call
 (#is-not? local))
"#;

#[test]
fn test_tags_python() {
    let language = get_language("python");
    let tags_config = TagsConfiguration::new(language, PYTHON_TAG_QUERY, "").unwrap();
    let mut tag_context = TagsContext::new();

    let source = br#"
    class Customer:
        """
        Data about a customer
        """

        def age(self):
            '''
            Get the customer's age
            '''
            compute_age(self.id)
    }
    "#;

    let tags = tag_context
        .generate_tags(&tags_config, source, None)
        .unwrap()
        .0
        .collect::<Result<Vec<_>, _>>()
        .unwrap();

    assert_eq!(
        tags.iter()
            .map(|t| (
                substr(source, &t.name_range),
                tags_config.syntax_type_name(t.syntax_type_id)
            ))
            .collect::<Vec<_>>(),
        &[
            ("Customer", "class"),
            ("age", "function"),
            ("compute_age", "call"),
        ]
    );

    assert_eq!(substr(source, &tags[0].line_range), "class Customer:");
    assert_eq!(substr(source, &tags[1].line_range), "def age(self):");
    assert_eq!(tags[0].docs.as_ref().unwrap(), "Data about a customer");
    assert_eq!(tags[1].docs.as_ref().unwrap(), "Get the customer's age");
}

#[test]
fn test_tags_javascript() {
    let language = get_language("javascript");
    let tags_config = TagsConfiguration::new(language, JS_TAG_QUERY, "").unwrap();
    let source = br#"
    // hi

    // Data about a customer.
    // bla bla bla
    class Customer {
        /*
         * Get the customer's age
         */
        getAge() {
        }
    }

    // ok

    class Agent {

    }
    "#;

    let mut tag_context = TagsContext::new();
    let tags = tag_context
        .generate_tags(&tags_config, source, None)
        .unwrap()
        .0
        .collect::<Result<Vec<_>, _>>()
        .unwrap();

    assert_eq!(
        tags.iter()
            .map(|t| (
                substr(source, &t.name_range),
                t.span.clone(),
                tags_config.syntax_type_name(t.syntax_type_id)
            ))
            .collect::<Vec<_>>(),
        &[
            ("Customer", Point::new(5, 10)..Point::new(5, 18), "class",),
            ("getAge", Point::new(9, 8)..Point::new(9, 14), "method",),
            ("Agent", Point::new(15, 10)..Point::new(15, 15), "class",)
        ]
    );
    assert_eq!(
        tags[0].docs.as_ref().unwrap(),
        "Data about a customer.\nbla bla bla"
    );
    assert_eq!(tags[1].docs.as_ref().unwrap(), "Get the customer's age");
    assert_eq!(tags[2].docs, None);
}

#[test]
fn test_tags_columns_measured_in_utf16_code_units() {
    let language = get_language("python");
    let tags_config = TagsConfiguration::new(language, PYTHON_TAG_QUERY, "").unwrap();
    let mut tag_context = TagsContext::new();

    let source = r#""❤❤❤".hello_α_ω()"#.as_bytes();

    let tag = tag_context
        .generate_tags(&tags_config, source, None)
        .unwrap()
        .0
        .next()
        .unwrap()
        .unwrap();

    assert_eq!(substr(source, &tag.name_range), "hello_α_ω");
    assert_eq!(tag.span, Point::new(0, 21)..Point::new(0, 32));
    assert_eq!(tag.utf16_column_range, 9..18);
}

#[test]
fn test_tags_ruby() {
    let language = get_language("ruby");
    let locals_query =
        fs::read_to_string(get_language_queries_path("ruby").join("locals.scm")).unwrap();
    let tags_config = TagsConfiguration::new(language, RUBY_TAG_QUERY, &locals_query).unwrap();
    let source = strip_whitespace(
        8,
        "
        b = 1

        def foo=()
            c = 1

            # a is a method because it is not in scope
            # b is a method because `b` doesn't capture variables from its containing scope
            bar a, b, c

            [1, 2, 3].each do |a|
                # a is a parameter
                # b is a method
                # c is a variable, because the block captures variables from its containing scope.
                baz a, b, c
            end
        end",
    );

    let mut tag_context = TagsContext::new();
    let tags = tag_context
        .generate_tags(&tags_config, source.as_bytes(), None)
        .unwrap()
        .0
        .collect::<Result<Vec<_>, _>>()
        .unwrap();

    assert_eq!(
        tags.iter()
            .map(|t| (
                substr(source.as_bytes(), &t.name_range),
                tags_config.syntax_type_name(t.syntax_type_id),
                (t.span.start.row, t.span.start.column),
            ))
            .collect::<Vec<_>>(),
        &[
            ("foo=", "method", (2, 4)),
            ("bar", "call", (7, 4)),
            ("a", "call", (7, 8)),
            ("b", "call", (7, 11)),
            ("each", "call", (9, 14)),
            ("baz", "call", (13, 8)),
            ("b", "call", (13, 15),),
        ]
    );
}

#[test]
fn test_tags_cancellation() {
    use std::sync::atomic::{AtomicUsize, Ordering};

    allocations::record(|| {
        // Large javascript document
        let source = (0..500)
            .map(|_| "/* hi */ class A { /* ok */ b() {} }\n")
            .collect::<String>();

        let cancellation_flag = AtomicUsize::new(0);
        let language = get_language("javascript");
        let tags_config = TagsConfiguration::new(language, JS_TAG_QUERY, "").unwrap();

        let mut tag_context = TagsContext::new();
        let tags = tag_context
            .generate_tags(&tags_config, source.as_bytes(), Some(&cancellation_flag))
            .unwrap();

        for (i, tag) in tags.0.enumerate() {
            if i == 150 {
                cancellation_flag.store(1, Ordering::SeqCst);
            }
            if let Err(e) = tag {
                assert_eq!(e, Error::Cancelled);
                return;
            }
        }

        panic!("Expected to halt tagging with an error");
    });
}

#[test]
fn test_invalid_capture() {
    let language = get_language("python");
    let e = TagsConfiguration::new(language, "(identifier) @method", "")
        .expect_err("expected InvalidCapture error");
    assert_eq!(e, Error::InvalidCapture("method".to_string()));
}

#[test]
fn test_tags_via_c_api() {
    allocations::record(|| {
        let tagger = c::ts_tagger_new();
        let buffer = c::ts_tags_buffer_new();
        let scope_name = "source.js";
        let language = get_language("javascript");

        let source_code = strip_whitespace(
            12,
            "
            var a = 1;

            // one
            // two
            // three
            function b() {
            }

            // four
            // five
            class C extends D {

            }

            b(a);",
        );

        let c_scope_name = CString::new(scope_name).unwrap();
        let result = c::ts_tagger_add_language(
            tagger,
            c_scope_name.as_ptr(),
            language,
            JS_TAG_QUERY.as_ptr(),
            ptr::null(),
            JS_TAG_QUERY.len() as u32,
            0,
        );
        assert_eq!(result, c::TSTagsError::Ok);

        let result = c::ts_tagger_tag(
            tagger,
            c_scope_name.as_ptr(),
            source_code.as_ptr(),
            source_code.len() as u32,
            buffer,
            ptr::null(),
        );
        assert_eq!(result, c::TSTagsError::Ok);
        let tags = unsafe {
            slice::from_raw_parts(
                c::ts_tags_buffer_tags(buffer),
                c::ts_tags_buffer_tags_len(buffer) as usize,
            )
        };
        let docs = str::from_utf8(unsafe {
            slice::from_raw_parts(
                c::ts_tags_buffer_docs(buffer) as *const u8,
                c::ts_tags_buffer_docs_len(buffer) as usize,
            )
        })
        .unwrap();

        let syntax_types: Vec<&str> = unsafe {
            let mut len: u32 = 0;
            let ptr =
                c::ts_tagger_syntax_kinds_for_scope_name(tagger, c_scope_name.as_ptr(), &mut len);
            slice::from_raw_parts(ptr, len as usize)
                .iter()
                .map(|i| CStr::from_ptr(*i).to_str().unwrap())
                .collect()
        };

        assert_eq!(
            tags.iter()
                .map(|tag| (
                    syntax_types[tag.syntax_type_id as usize],
                    &source_code[tag.name_start_byte as usize..tag.name_end_byte as usize],
                    &source_code[tag.line_start_byte as usize..tag.line_end_byte as usize],
                    &docs[tag.docs_start_byte as usize..tag.docs_end_byte as usize],
                ))
                .collect::<Vec<_>>(),
            &[
                ("function", "b", "function b() {", "one\ntwo\nthree"),
                ("class", "C", "class C extends D {", "four\nfive"),
                ("call", "b", "b(a);", "")
            ]
        );

        c::ts_tags_buffer_delete(buffer);
        c::ts_tagger_delete(tagger);
    });
}

fn substr<'a>(source: &'a [u8], range: &std::ops::Range<usize>) -> &'a str {
    std::str::from_utf8(&source[range.clone()]).unwrap()
}

fn strip_whitespace(indent: usize, s: &str) -> String {
    s.lines()
        .skip(1)
        .map(|line| &line[line.len().min(indent)..])
        .collect::<Vec<_>>()
        .join("\n")
}
<commit_msg>Add a test.<commit_after>use super::helpers::allocations;
use super::helpers::fixtures::{get_language, get_language_queries_path};
use std::ffi::CStr;
use std::ffi::CString;
use std::{fs, ptr, slice, str};
use tree_sitter::Point;
use tree_sitter_tags::c_lib as c;
use tree_sitter_tags::{Error, TagsConfiguration, TagsContext};

const PYTHON_TAG_QUERY: &'static str = r#"
(
  (function_definition
    name: (identifier) @name
    body: (block . (expression_statement (string) @doc))) @definition.function
  (#strip! @doc "(^['\"\\s]*)|(['\"\\s]*$)")
)

(function_definition
  name: (identifier) @name) @definition.function

(
  (class_definition
    name: (identifier) @name
    body: (block
      . (expression_statement (string) @doc))) @definition.class
  (#strip! @doc "(^['\"\\s]*)|(['\"\\s]*$)")
)

(class_definition
  name: (identifier) @name) @definition.class

(call
  function: (identifier) @name) @reference.call

(call
  function: (attribute
    attribute: (identifier) @name)) @reference.call
"#;

const JS_TAG_QUERY: &'static str = r#"
(
    (comment)* @doc .
    (class_declaration
        name: (identifier) @name) @definition.class
    (#select-adjacent! @doc @definition.class)
    (#strip! @doc "(^[/\\*\\s]*)|([/\\*\\s]*$)")
)

(
    (comment)* @doc .
    (method_definition
        name: (property_identifier) @name) @definition.method
    (#select-adjacent! @doc @definition.method)
    (#strip! @doc "(^[/\\*\\s]*)|([/\\*\\s]*$)")
)

(
    (comment)* @doc .
    (function_declaration
        name: (identifier) @name) @definition.function
    (#select-adjacent! @doc @definition.function)
    (#strip! @doc "(^[/\\*\\s]*)|([/\\*\\s]*$)")
)

(call_expression
    function: (identifier) @name) @reference.call
"#;

const RUBY_TAG_QUERY: &'static str = r#"
(method
    name: (_) @name) @definition.method

(method_call
    method: (identifier) @name) @reference.call

(setter (identifier) @ignore)

((identifier) @name @reference.call
 (#is-not? local))
"#;

#[test]
fn test_tags_python() {
    let language = get_language("python");
    let tags_config = TagsConfiguration::new(language, PYTHON_TAG_QUERY, "").unwrap();
    let mut tag_context = TagsContext::new();

    let source = br#"
    class Customer:
        """
        Data about a customer
        """

        def age(self):
            '''
            Get the customer's age
            '''
            compute_age(self.id)
    }
    "#;

    let tags = tag_context
        .generate_tags(&tags_config, source, None)
        .unwrap()
        .0
        .collect::<Result<Vec<_>, _>>()
        .unwrap();

    assert_eq!(
        tags.iter()
            .map(|t| (
                substr(source, &t.name_range),
                tags_config.syntax_type_name(t.syntax_type_id)
            ))
            .collect::<Vec<_>>(),
        &[
            ("Customer", "class"),
            ("age", "function"),
            ("compute_age", "call"),
        ]
    );

    assert_eq!(substr(source, &tags[0].line_range), "class Customer:");
    assert_eq!(substr(source, &tags[1].line_range), "def age(self):");
    assert_eq!(tags[0].docs.as_ref().unwrap(), "Data about a customer");
    assert_eq!(tags[1].docs.as_ref().unwrap(), "Get the customer's age");
}

#[test]
fn test_tags_javascript() {
    let language = get_language("javascript");
    let tags_config = TagsConfiguration::new(language, JS_TAG_QUERY, "").unwrap();
    let source = br#"
    // hi

    // Data about a customer.
    // bla bla bla
    class Customer {
        /*
         * Get the customer's age
         */
        getAge() {
        }
    }

    // ok

    class Agent {

    }
    "#;

    let mut tag_context = TagsContext::new();
    let tags = tag_context
        .generate_tags(&tags_config, source, None)
        .unwrap()
        .0
        .collect::<Result<Vec<_>, _>>()
        .unwrap();

    assert_eq!(
        tags.iter()
            .map(|t| (
                substr(source, &t.name_range),
                t.span.clone(),
                tags_config.syntax_type_name(t.syntax_type_id)
            ))
            .collect::<Vec<_>>(),
        &[
            ("Customer", Point::new(5, 10)..Point::new(5, 18), "class",),
            ("getAge", Point::new(9, 8)..Point::new(9, 14), "method",),
            ("Agent", Point::new(15, 10)..Point::new(15, 15), "class",)
        ]
    );
    assert_eq!(
        tags[0].docs.as_ref().unwrap(),
        "Data about a customer.\nbla bla bla"
    );
    assert_eq!(tags[1].docs.as_ref().unwrap(), "Get the customer's age");
    assert_eq!(tags[2].docs, None);
}

#[test]
fn test_tags_columns_measured_in_utf16_code_units() {
    let language = get_language("python");
    let tags_config = TagsConfiguration::new(language, PYTHON_TAG_QUERY, "").unwrap();
    let mut tag_context = TagsContext::new();

    let source = r#""❤❤❤".hello_α_ω()"#.as_bytes();

    let tag = tag_context
        .generate_tags(&tags_config, source, None)
        .unwrap()
        .0
        .next()
        .unwrap()
        .unwrap();

    assert_eq!(substr(source, &tag.name_range), "hello_α_ω");
    assert_eq!(tag.span, Point::new(0, 21)..Point::new(0, 32));
    assert_eq!(tag.utf16_column_range, 9..18);
}

#[test]
fn test_tags_ruby() {
    let language = get_language("ruby");
    let locals_query =
        fs::read_to_string(get_language_queries_path("ruby").join("locals.scm")).unwrap();
    let tags_config = TagsConfiguration::new(language, RUBY_TAG_QUERY, &locals_query).unwrap();
    let source = strip_whitespace(
        8,
        "
        b = 1

        def foo=()
            c = 1

            # a is a method because it is not in scope
            # b is a method because `b` doesn't capture variables from its containing scope
            bar a, b, c

            [1, 2, 3].each do |a|
                # a is a parameter
                # b is a method
                # c is a variable, because the block captures variables from its containing scope.
                baz a, b, c
            end
        end",
    );

    let mut tag_context = TagsContext::new();
    let tags = tag_context
        .generate_tags(&tags_config, source.as_bytes(), None)
        .unwrap()
        .0
        .collect::<Result<Vec<_>, _>>()
        .unwrap();

    assert_eq!(
        tags.iter()
            .map(|t| (
                substr(source.as_bytes(), &t.name_range),
                tags_config.syntax_type_name(t.syntax_type_id),
                (t.span.start.row, t.span.start.column),
            ))
            .collect::<Vec<_>>(),
        &[
            ("foo=", "method", (2, 4)),
            ("bar", "call", (7, 4)),
            ("a", "call", (7, 8)),
            ("b", "call", (7, 11)),
            ("each", "call", (9, 14)),
            ("baz", "call", (13, 8)),
            ("b", "call", (13, 15),),
        ]
    );
}

#[test]
fn test_tags_cancellation() {
    use std::sync::atomic::{AtomicUsize, Ordering};

    allocations::record(|| {
        // Large javascript document
        let source = (0..500)
            .map(|_| "/* hi */ class A { /* ok */ b() {} }\n")
            .collect::<String>();

        let cancellation_flag = AtomicUsize::new(0);
        let language = get_language("javascript");
        let tags_config = TagsConfiguration::new(language, JS_TAG_QUERY, "").unwrap();

        let mut tag_context = TagsContext::new();
        let tags = tag_context
            .generate_tags(&tags_config, source.as_bytes(), Some(&cancellation_flag))
            .unwrap();

        for (i, tag) in tags.0.enumerate() {
            if i == 150 {
                cancellation_flag.store(1, Ordering::SeqCst);
            }
            if let Err(e) = tag {
                assert_eq!(e, Error::Cancelled);
                return;
            }
        }

        panic!("Expected to halt tagging with an error");
    });
}

#[test]
fn test_invalid_capture() {
    let language = get_language("python");
    let e = TagsConfiguration::new(language, "(identifier) @method", "")
        .expect_err("expected InvalidCapture error");
    assert_eq!(e, Error::InvalidCapture("method".to_string()));
}

#[test]
fn test_tags_with_parse_error() {
    let language = get_language("python");
    let tags_config = TagsConfiguration::new(language, PYTHON_TAG_QUERY, "").unwrap();
    let mut tag_context = TagsContext::new();

    let source = br#"
    class Fine: pass
    class Bad
    "#;

    let (tags, failed) = tag_context
        .generate_tags(&tags_config, source, None)
        .unwrap();

    let newtags = tags.collect::<Result<Vec<_>, _>>().unwrap();

    assert!(failed, "syntax error should have been detected");

    assert_eq!(
        newtags.iter()
            .map(|t| (
                substr(source, &t.name_range),
                tags_config.syntax_type_name(t.syntax_type_id)
            ))
            .collect::<Vec<_>>(),
        &[
            ("Fine", "class"),
        ]
    );
}

#[test]
fn test_tags_via_c_api() {
    allocations::record(|| {
        let tagger = c::ts_tagger_new();
        let buffer = c::ts_tags_buffer_new();
        let scope_name = "source.js";
        let language = get_language("javascript");

        let source_code = strip_whitespace(
            12,
            "
            var a = 1;

            // one
            // two
            // three
            function b() {
            }

            // four
            // five
            class C extends D {

            }

            b(a);",
        );

        let c_scope_name = CString::new(scope_name).unwrap();
        let result = c::ts_tagger_add_language(
            tagger,
            c_scope_name.as_ptr(),
            language,
            JS_TAG_QUERY.as_ptr(),
            ptr::null(),
            JS_TAG_QUERY.len() as u32,
            0,
        );
        assert_eq!(result, c::TSTagsError::Ok);

        let result = c::ts_tagger_tag(
            tagger,
            c_scope_name.as_ptr(),
            source_code.as_ptr(),
            source_code.len() as u32,
            buffer,
            ptr::null(),
        );
        assert_eq!(result, c::TSTagsError::Ok);
        let tags = unsafe {
            slice::from_raw_parts(
                c::ts_tags_buffer_tags(buffer),
                c::ts_tags_buffer_tags_len(buffer) as usize,
            )
        };
        let docs = str::from_utf8(unsafe {
            slice::from_raw_parts(
                c::ts_tags_buffer_docs(buffer) as *const u8,
                c::ts_tags_buffer_docs_len(buffer) as usize,
            )
        })
        .unwrap();

        let syntax_types: Vec<&str> = unsafe {
            let mut len: u32 = 0;
            let ptr =
                c::ts_tagger_syntax_kinds_for_scope_name(tagger, c_scope_name.as_ptr(), &mut len);
            slice::from_raw_parts(ptr, len as usize)
                .iter()
                .map(|i| CStr::from_ptr(*i).to_str().unwrap())
                .collect()
        };

        assert_eq!(
            tags.iter()
                .map(|tag| (
                    syntax_types[tag.syntax_type_id as usize],
                    &source_code[tag.name_start_byte as usize..tag.name_end_byte as usize],
                    &source_code[tag.line_start_byte as usize..tag.line_end_byte as usize],
                    &docs[tag.docs_start_byte as usize..tag.docs_end_byte as usize],
                ))
                .collect::<Vec<_>>(),
            &[
                ("function", "b", "function b() {", "one\ntwo\nthree"),
                ("class", "C", "class C extends D {", "four\nfive"),
                ("call", "b", "b(a);", "")
            ]
        );

        c::ts_tags_buffer_delete(buffer);
        c::ts_tagger_delete(tagger);
    });
}

fn substr<'a>(source: &'a [u8], range: &std::ops::Range<usize>) -> &'a str {
    std::str::from_utf8(&source[range.clone()]).unwrap()
}

fn strip_whitespace(indent: usize, s: &str) -> String {
    s.lines()
        .skip(1)
        .map(|line| &line[line.len().min(indent)..])
        .collect::<Vec<_>>()
        .join("\n")
}
<|endoftext|>


<commit_before>//! Write systems to disk.

use error::Result;
use system::System;

use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::PathBuf;

/// Output a system to disk as a GROMOS formatted file.
/// The filename extension is adjusted to .gro.
///
/// # Errors
/// Returns an error if the file could not be written to.
pub fn write_gromos(system: &System, output_file: &str, title: &str) -> Result<()> {
    let path = PathBuf::from(output_file).with_extension("gro");
    let file = File::create(&path)?;
    let mut writer = BufWriter::new(file);

    writer.write_fmt(format_args!("{}\n", title))?;
    writer.write_fmt(format_args!("{}\n", system.num_atoms()))?;

    for (i, residue) in system.residues.iter().enumerate() {
        for (j, atom) in residue.atoms.iter().enumerate() {
            let residue_number = (i + 1) % 100_000;
            let atom_number = (j + 1) % 100_000;

            writer.write_fmt(format_args!("{:>5}{:<5}{:>5}{:>5}{:>8.3}{:>8.3}{:>8.3}\n",
                                        residue_number,
                                        residue.code,
                                        atom.code,
                                        atom_number,
                                        atom.position.x,
                                        atom.position.y,
                                        atom.position.z))?;

        }
    }

    writer.write_fmt(format_args!("{:12.8} {:12.8} {:12.8}\n",
                                system.dimensions.x,
                                system.dimensions.y,
                                system.dimensions.z))?;

    Ok(())
}
<commit_msg>Atom relative positions are now added to the residue position<commit_after>//! Write systems to disk.

use error::Result;
use system::System;

use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::PathBuf;

/// Output a system to disk as a GROMOS formatted file.
/// The filename extension is adjusted to .gro.
///
/// # Errors
/// Returns an error if the file could not be written to.
pub fn write_gromos(system: &System, output_file: &str, title: &str) -> Result<()> {
    let path = PathBuf::from(output_file).with_extension("gro");
    let file = File::create(&path)?;
    let mut writer = BufWriter::new(file);

    writer.write_fmt(format_args!("{}\n", title))?;
    writer.write_fmt(format_args!("{}\n", system.num_atoms()))?;

    for (i, residue) in system.residues.iter().enumerate() {
        for (j, atom) in residue.atoms.iter().enumerate() {
            let residue_number = (i + 1) % 100_000;
            let atom_number = (j + 1) % 100_000;

            let position = residue.position.add(&atom.position);

            writer.write_fmt(format_args!("{:>5}{:<5}{:>5}{:>5}{:>8.3}{:>8.3}{:>8.3}\n",
                                        residue_number,
                                        residue.code,
                                        atom.code,
                                        atom_number,
                                        position.x,
                                        position.y,
                                        position.z))?;

        }
    }

    writer.write_fmt(format_args!("{:12.8} {:12.8} {:12.8}\n",
                                system.dimensions.x,
                                system.dimensions.y,
                                system.dimensions.z))?;

    Ok(())
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use lexer::{Token, Tokens};
use error::ParserError;

// XXX: Change operator type from &'static str to String
type Operators = HashMap<&'static str, Operator>;

#[derive(PartialEq, Debug)]
enum Associativity {
    LeftToRight,
    RightToLeft,
}

#[derive(PartialEq, Debug)]
struct Operator(usize, Associativity);

impl Operator {
    fn new(precedence: usize, associativity: Associativity) -> Operator {
        Operator(precedence, associativity)
    }
}

pub struct Parser {
    operators: Operators,
    stack: Tokens,
    output: Tokens,
}

impl Parser {
    fn new() -> Parser {
        Parser {
            operators: Operators::new(),
            stack: Tokens::with_capacity(3),
            output: Tokens::with_capacity(3),
        }
    }

    fn lower_precedence(&self, new_token: &Token, top_token: &Token) -> bool {
        let &Operator(new_token_prec, ref new_token_assoc) = match *new_token {
            Token::Operator(ref name) => self.operators.get::<str>(&name).unwrap(),
            _ => unreachable!(),
        };

        let &Operator(top_token_prec, _) = match *top_token {
            Token::Operator(ref name) => self.operators.get::<str>(&name).unwrap(),
            _ => unreachable!(),
        };

        (*new_token_assoc == Associativity::LeftToRight && new_token_prec <= top_token_prec) ||
        (*new_token_assoc == Associativity::RightToLeft && new_token_prec < top_token_prec)
    }

    pub fn parse(&mut self, tokens: &Tokens) -> Result<&Tokens, ParserError> {
        self.reset_parser();

        let mut iter = tokens.iter().peekable();
        while let Some(&(position, ref token)) = iter.next() {
            match *token {
                Token::Decimal(_) => self.output.push((position, token.clone())),

                // Token::Identifier(ref id) => {
                //     if self.is_keyword(id) {
                //         return Err(ParserError::KeywordError(position, id.clone()));
                //     } else {
                //         self.stack.push((position, token.clone()));
                //     }
                // }
                Token::Operator(ref name) => {
                    // if the token is an operator, o1, then:
                    // while there is an operator token o2, at the top of the
                    // operator stack and either o1 is left-associative and its
                    // precedence is less than or equal to that of o2, or o1 is
                    // right associative, and has precedence less than that of
                    // o2, pop o2 off the operator stack, onto the output queue
                    if !self.operators.contains_key::<str>(&name) {
                        return Err(ParserError::IllegalOperator(position));
                    }

                    loop {
                        match self.stack.last() {
                            Some(&(_, Token::Operator(_))) => {
                                if self.lower_precedence(&token, &self.stack.last().unwrap().1) {
                                    self.output.push(self.stack.pop().unwrap());
                                } else {
                                    break;
                                }
                            }
                            _ => break,
                        }
                    }

                    self.stack.push((position, token.clone()));
                }
                Token::OpenBracket => self.stack.push((position, token.clone())),
                Token::CloseBracket => {
                    let mut found = false;

                    loop {
                        match self.stack.last() {
                            Some(&(_, Token::OpenBracket)) => {
                                found = true;
                                self.stack.pop();
                                break;
                            }
                            None => break,
                            _ => self.output.push(self.stack.pop().unwrap()),
                        }
                    }
                    if !found {
                        return Err(ParserError::MissingOpeningBracket(position));
                    }
                }
                Token::Unknown(_) => {
                    return Err(ParserError::IllegalOperator(position));
                }
                _ => continue,

            }
        }
        loop {
            match self.stack.last() {
                Some(&(position, Token::OpenBracket)) => {
                    return Err(ParserError::MissingClosingBracket(position))
                }
                Some(_) => self.output.push(self.stack.pop().unwrap()),
                None => break,
            }
        }

        Ok((&self.output))
    }

    fn reset_parser(&mut self) {
        self.stack.clear();
        self.output.clear();
    }
}

impl Default for Parser {
    fn default() -> Parser {
        let mut parser = Parser::new();

        // XXX: Try to convert operators to String with .to_string()
        parser.operators.insert("(", Operator::new(1, Associativity::LeftToRight));
        parser.operators.insert(")", Operator::new(1, Associativity::LeftToRight));
        parser.operators.insert("~", Operator::new(2, Associativity::RightToLeft));
        // TODO: Unary Minus
        // parser.operators.insert("-", Operator::new(2, Associativity::RightToLeft));
        parser.operators.insert(">>", Operator::new(3, Associativity::LeftToRight));
        parser.operators.insert("<<", Operator::new(3, Associativity::LeftToRight));
        parser.operators.insert("&", Operator::new(4, Associativity::LeftToRight));
        parser.operators.insert("^", Operator::new(5, Associativity::LeftToRight));
        parser.operators.insert("|", Operator::new(6, Associativity::LeftToRight));

        parser
    }
}
<commit_msg>Remove comments about operator strings<commit_after>use std::collections::HashMap;
use lexer::{Token, Tokens};
use error::ParserError;

type Operators = HashMap<&'static str, Operator>;

#[derive(PartialEq, Debug)]
enum Associativity {
    LeftToRight,
    RightToLeft,
}

#[derive(PartialEq, Debug)]
struct Operator(usize, Associativity);

impl Operator {
    fn new(precedence: usize, associativity: Associativity) -> Operator {
        Operator(precedence, associativity)
    }
}

pub struct Parser {
    operators: Operators,
    stack: Tokens,
    output: Tokens,
}

impl Parser {
    fn new() -> Parser {
        Parser {
            operators: Operators::new(),
            stack: Tokens::with_capacity(3),
            output: Tokens::with_capacity(3),
        }
    }

    fn lower_precedence(&self, new_token: &Token, top_token: &Token) -> bool {
        let &Operator(new_token_prec, ref new_token_assoc) = match *new_token {
            Token::Operator(ref name) => self.operators.get::<str>(&name).unwrap(),
            _ => unreachable!(),
        };

        let &Operator(top_token_prec, _) = match *top_token {
            Token::Operator(ref name) => self.operators.get::<str>(&name).unwrap(),
            _ => unreachable!(),
        };

        (*new_token_assoc == Associativity::LeftToRight && new_token_prec <= top_token_prec) ||
        (*new_token_assoc == Associativity::RightToLeft && new_token_prec < top_token_prec)
    }

    pub fn parse(&mut self, tokens: &Tokens) -> Result<&Tokens, ParserError> {
        self.reset_parser();

        let mut iter = tokens.iter().peekable();
        while let Some(&(position, ref token)) = iter.next() {
            match *token {
                Token::Decimal(_) => self.output.push((position, token.clone())),

                // Token::Identifier(ref id) => {
                //     if self.is_keyword(id) {
                //         return Err(ParserError::KeywordError(position, id.clone()));
                //     } else {
                //         self.stack.push((position, token.clone()));
                //     }
                // }
                Token::Operator(ref name) => {
                    // if the token is an operator, o1, then:
                    // while there is an operator token o2, at the top of the
                    // operator stack and either o1 is left-associative and its
                    // precedence is less than or equal to that of o2, or o1 is
                    // right associative, and has precedence less than that of
                    // o2, pop o2 off the operator stack, onto the output queue
                    if !self.operators.contains_key::<str>(&name) {
                        return Err(ParserError::IllegalOperator(position));
                    }

                    loop {
                        match self.stack.last() {
                            Some(&(_, Token::Operator(_))) => {
                                if self.lower_precedence(&token, &self.stack.last().unwrap().1) {
                                    self.output.push(self.stack.pop().unwrap());
                                } else {
                                    break;
                                }
                            }
                            _ => break,
                        }
                    }

                    self.stack.push((position, token.clone()));
                }
                Token::OpenBracket => self.stack.push((position, token.clone())),
                Token::CloseBracket => {
                    let mut found = false;

                    loop {
                        match self.stack.last() {
                            Some(&(_, Token::OpenBracket)) => {
                                found = true;
                                self.stack.pop();
                                break;
                            }
                            None => break,
                            _ => self.output.push(self.stack.pop().unwrap()),
                        }
                    }
                    if !found {
                        return Err(ParserError::MissingOpeningBracket(position));
                    }
                }
                Token::Unknown(_) => {
                    return Err(ParserError::IllegalOperator(position));
                }
                _ => continue,

            }
        }
        loop {
            match self.stack.last() {
                Some(&(position, Token::OpenBracket)) => {
                    return Err(ParserError::MissingClosingBracket(position))
                }
                Some(_) => self.output.push(self.stack.pop().unwrap()),
                None => break,
            }
        }

        Ok((&self.output))
    }

    fn reset_parser(&mut self) {
        self.stack.clear();
        self.output.clear();
    }
}

impl Default for Parser {
    fn default() -> Parser {
        let mut parser = Parser::new();

        parser.operators.insert("(", Operator::new(1, Associativity::LeftToRight));
        parser.operators.insert(")", Operator::new(1, Associativity::LeftToRight));
        parser.operators.insert("~", Operator::new(2, Associativity::RightToLeft));
        // TODO: Unary Minus
        // parser.operators.insert("-", Operator::new(2, Associativity::RightToLeft));
        parser.operators.insert(">>", Operator::new(3, Associativity::LeftToRight));
        parser.operators.insert("<<", Operator::new(3, Associativity::LeftToRight));
        parser.operators.insert("&", Operator::new(4, Associativity::LeftToRight));
        parser.operators.insert("^", Operator::new(5, Associativity::LeftToRight));
        parser.operators.insert("|", Operator::new(6, Associativity::LeftToRight));

        parser
    }
}
<|endoftext|>


<commit_before>use super::bit;
use super::error::{Error, ErrorKind};
use super::query::Query;
use super::result::Result;
use super::utf8::{COMMA, CR, HT, LF, RIGHT_BRACE, SPACE};
use fnv::{FnvHashMap, FnvHashSet};

#[inline]
pub fn basic_parse<'a>(
    rec: &'a [u8],
    index: &[Vec<u64>],
    queries: &mut FnvHashMap<&[u8], Query>,
    start: usize,
    end: usize,
    level: usize,
    stats: &mut Vec<FnvHashSet<usize>>,
    set_stats: bool,
    results: &mut Vec<Option<&'a [u8]>>,
    b_quote: &[u64],
    colon_positions: &mut Vec<Vec<usize>>,
) -> Result<()> {
    let query_num = queries.len();
    let mut found_num = 0;
    generate_colon_positions(index, start, end, level, colon_positions);
    let mut vei = end;
    let cp_len = &colon_positions[level].len();
    for i in (0..*cp_len).rev() {
        let (fsi, fei) = search_pre_field_indices(
            b_quote,
            if i > 0 {
                colon_positions[level][i - 1]
            } else {
                start
            },
            colon_positions[level][i],
        )?;
        let field = &rec[fsi + 1..fei];
        if let Some(query) = queries.get_mut(field) {
            let (vsi, vei) = search_post_value_indices(
                rec,
                colon_positions[level][i] + 1,
                vei,
                if i == *cp_len - 1 { RIGHT_BRACE } else { COMMA },
            )?;
            found_num += 1;
            if set_stats && !stats[query.i].contains(&i) {
                stats[query.i].insert(i);
            }
            if let Some(ref mut children) = query.children {
                basic_parse(
                    rec,
                    index,
                    children,
                    vsi,
                    vei,
                    level + 1,
                    stats,
                    set_stats,
                    results,
                    b_quote,
                    colon_positions,
                )?;
            }
            if query.target {
                results[query.ri] = Some(&rec[vsi..vei + 1]);
            }
            if found_num == query_num {
                return Ok(());
            }
        }
        vei = fsi - 1;
    }
    Ok(())
}

#[inline]
pub fn speculative_parse<'a>(rec: &'a [u8], index: &[Vec<u64>], queries: &FnvHashMap<&[u8], Query>, start: usize, end: usize, level: usize, stats: &[FnvHashSet<usize>], results: &mut Vec<Option<&'a [u8]>>, b_quote: &[u64], colon_positions: &mut Vec<Vec<usize>>) -> Result<bool> {
    generate_colon_positions(index, start, end, level, colon_positions);
    for (s, q) in queries.iter() {
        let mut found = false;
        for i in &stats[q.i] {
            let cp_len = &colon_positions[level].len();
            if *i >= *cp_len {
                continue;
            }
            let (fsi, fei) = search_pre_field_indices(
                b_quote,
                if *i > 0 {
                    colon_positions[level][*i - 1]
                } else {
                    start
                },
                colon_positions[level][*i],
            )?;
            let field = &rec[fsi + 1..fei];
            if s == &field {
                let vei = if *i < *cp_len - 1 {
                    let (nfsi, _) = search_pre_field_indices(
                        b_quote,
                        colon_positions[level][*i],
                        colon_positions[level][*i + 1],
                    )?;
                    nfsi - 1
                } else {
                    end
                };
                let (vsi, vei) = search_post_value_indices(
                    rec,
                    colon_positions[level][*i] + 1,
                    vei,
                    if *i == *cp_len - 1 {
                        RIGHT_BRACE
                    } else {
                        COMMA
                    },
                )?;
                if let Some(ref children) = q.children {
                    found = speculative_parse(
                        rec,
                        index,
                        children,
                        vsi,
                        vei,
                        level + 1,
                        stats,
                        results,
                        b_quote,
                        colon_positions,
                    )?;
                } else {
                    found = true;
                }
                if q.target {
                    results[q.ri] = Some(&rec[vsi..vei + 1]);
                }
                break;
            }
        }
        if !found {
            return Ok(false);
        }
    }
    Ok(true)
}

#[inline]
fn generate_colon_positions(index: &[Vec<u64>], start: usize, end: usize, level: usize, colon_positions: &mut Vec<Vec<usize>>) {
    let cp = &mut colon_positions[level];
    cp.clear();
    for i in start / 64..(end + 63) / 64 {
        let mut m_colon = index[level][i];
        while m_colon != 0 {
            let m_bit = bit::e(m_colon);
            let offset = i * 64 + (m_bit.trailing_zeros() as usize);
            if start <= offset && offset <= end {
                cp.push(offset);
            }
            m_colon = bit::r(m_colon);
        }
    }
}

#[inline]
fn search_pre_field_indices(b_quote: &[u64], start: usize, end: usize) -> Result<(usize, usize)> {
    let mut si = 0;
    let mut ei = 0;
    let mut ei_set = false;
    let mut n_quote = 0;
    for i in (start / 64..(end + 63) / 64).rev() {
        let mut m_quote = b_quote[i];
        while m_quote != 0 {
            let m_bit = bit::e(m_quote);
            let offset = i * 64 + (m_bit.trailing_zeros() as usize);
            if end <= offset {
                break;
            }
            if start < offset {
                if ei_set {
                    si = offset;
                } else {
                    si = ei;
                    ei = offset;
                }
                n_quote += 1;
            }
            m_quote = bit::r(m_quote);
        }
        if n_quote >= 2 {
            break;
        }
        if n_quote == 1 && !ei_set {
            ei_set = true;
        }
    }
    if n_quote >= 2 {
        Ok((si, ei))
    } else {
        Err(Error::from(ErrorKind::InvalidRecord))
    }
}

#[inline]
fn search_post_value_indices(rec: &[u8], si: usize, ei: usize, ignore_once_char: u8) -> Result<(usize, usize)> {
    let mut si = si;
    let mut ei = ei;
    let mut ignore_once_char_ignored = false;
    let n = rec.len();
    while si < n {
        match rec[si] {
            SPACE | HT | LF | CR => {
                si += 1;
            }
            _ => {
                break;
            }
        }
    }
    if si == n {
        return Err(Error::from(ErrorKind::InvalidRecord));
    }
    while si <= ei {
        match rec[ei] {
            SPACE | HT | LF | CR => {
                ei -= 1;
            }
            char if char == ignore_once_char => {
                if ignore_once_char_ignored {
                    break;
                }
                ignore_once_char_ignored = true;
                ei -= 1;
            }
            _ => {
                break;
            }
        }
    }
    if ei < si {
        return Err(Error::from(ErrorKind::InvalidRecord));
    }
    Ok((si, ei))
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::avx;
    use super::super::index_builder;
    use super::super::utf8::{BACKSLASH, COLON, LEFT_BRACE, QUOTE, RIGHT_BRACE};

    #[test]
    fn test_basic_parse() {
        let json_rec_str = r#"{ "aaa" : "AAA", "bbb" : 111, "ccc": ["C1", "C2"], "ddd" : { "d1" : "D1", "d2" : "D2", "d3": 333 }, "eee": { "e1": "EEE" } } "#;
        let json_rec = json_rec_str.as_bytes();
        let mut b_backslash = Vec::with_capacity((json_rec.len() + 63) / 64);
        let mut b_quote = Vec::with_capacity((json_rec.len() + 63) / 64);
        let mut b_colon = Vec::with_capacity((json_rec.len() + 63) / 64);
        let mut b_left = Vec::with_capacity((json_rec.len() + 63) / 64);
        let mut b_right = Vec::with_capacity((json_rec.len() + 63) / 64);
        index_builder::build_structural_character_bitmap(
            &json_rec,
            &mut b_backslash,
            &mut b_quote,
            &mut b_colon,
            &mut b_left,
            &mut b_right,
            &avx::mm256i(BACKSLASH as i8),
            &avx::mm256i(QUOTE as i8),
            &avx::mm256i(COLON as i8),
            &avx::mm256i(LEFT_BRACE as i8),
            &avx::mm256i(RIGHT_BRACE as i8),
        );
        index_builder::build_structural_quote_bitmap(&b_backslash, &mut b_quote);
        let mut b_string_mask = Vec::with_capacity(b_quote.len());
        index_builder::build_string_mask_bitmap(&b_quote, &mut b_string_mask);
        for i in 0..b_string_mask.len() {
            let b = b_string_mask[i];
            b_colon[i] &= b;
            b_left[i] &= b;
            b_right[i] &= b;
        }
        let l = 10;
        let mut index = vec![Vec::new(); l];
        let mut s_left = Vec::new();
        let r = index_builder::build_leveled_colon_bitmap(&b_colon, &b_left, &b_right, l, &mut s_left, &mut index);
        assert_eq!(Ok(()), r);
        let mut children = FnvHashMap::default();
        children.insert(
            "d1".as_bytes(),
            Query {
                i: 0,
                ri: 0,
                target: true,
                children: None,
            },
        );
        children.insert(
            "d3".as_bytes(),
            Query {
                i: 1,
                ri: 1,
                target: true,
                children: None,
            },
        );
        let mut queries = FnvHashMap::default();
        queries.insert(
            "aaa".as_bytes(),
            Query {
                i: 2,
                ri: 2,
                target: true,
                children: None,
            },
        );
        queries.insert(
            "bbb".as_bytes(),
            Query {
                i: 3,
                ri: 3,
                target: true,
                children: None,
            },
        );
        queries.insert(
            "ccc".as_bytes(),
            Query {
                i: 4,
                ri: 4,
                target: true,
                children: None,
            },
        );
        queries.insert(
            "ddd".as_bytes(),
            Query {
                i: 5,
                ri: 0,
                target: false,
                children: Some(children),
            },
        );
        queries.insert(
            "eee".as_bytes(),
            Query {
                i: 6,
                ri: 5,
                target: true,
                children: None,
            },
        );

        let mut stats = vec![
            FnvHashSet::default(),
            FnvHashSet::default(),
            FnvHashSet::default(),
            FnvHashSet::default(),
            FnvHashSet::default(),
            FnvHashSet::default(),
            FnvHashSet::default(),
        ];

        let mut results = vec![None, None, None, None, None, None];
        let mut colon_positions = vec![Vec::new(); l];
        let result = basic_parse(
            json_rec,
            &index,
            &mut queries,
            0,
            json_rec.len() - 1,
            0,
            &mut stats,
            true,
            &mut results,
            &b_quote,
            &mut colon_positions,
        );

        assert_eq!(Ok(()), result);
        assert_eq!(Some(r#""D1""#.as_bytes()), results[0]);
        assert_eq!(Some(r#"333"#.as_bytes()), results[1]);
        assert_eq!(Some(r#""AAA""#.as_bytes()), results[2]);
        assert_eq!(Some(r#"111"#.as_bytes()), results[3]);
        assert_eq!(Some(r#"["C1", "C2"]"#.as_bytes()), results[4]);
        assert_eq!(Some(r#"{ "e1": "EEE" }"#.as_bytes()), results[5]);
    }

    #[test]
    fn test_generate_colon_positions() {
        struct TestCase {
            index: Vec<Vec<u64>>,
            start: usize,
            end: usize,
            level: usize,
            want: Vec<usize>,
        }
        let test_cases = vec![
            TestCase {
                index: vec![vec![]],
                start: 0,
                end: 0,
                level: 0,
                want: vec![],
            },
            TestCase {
                index: vec![
                    vec![
                        0b00000001_11000000_00000000_00000000_00001000_00000000_00000011_10000000,
                    ],
                ],
                start: 8,
                end: 55,
                level: 0,
                want: vec![8, 9, 27, 54, 55],
            },
            TestCase {
                index: vec![
                    vec![
                        0b10000000_00000000_00000000_00000000_00001000_00000000_00000011_10000000,
                        0b00000001_11000001_00000000_00000011_00000000_00000000_00000000_00000001,
                    ],
                ],
                start: 8,
                end: 119,
                level: 0,
                want: vec![8, 9, 27, 63, 64, 96, 97, 112, 118, 119],
            },
        ];
        for t in test_cases {
            let mut cp = vec![Vec::new(); t.level + 1];
            generate_colon_positions(&t.index, t.start, t.end, t.level, &mut cp);
            assert_eq!(t.want, cp[0]);
        }
    }
}
<commit_msg>chore(parser): cleanup<commit_after>use super::bit;
use super::error::{Error, ErrorKind};
use super::query::Query;
use super::result::Result;
use super::utf8::{COMMA, CR, HT, LF, RIGHT_BRACE, SPACE};
use fnv::{FnvHashMap, FnvHashSet};

#[inline]
pub fn basic_parse<'a>(
    rec: &'a [u8],
    index: &[Vec<u64>],
    queries: &mut FnvHashMap<&[u8], Query>,
    start: usize,
    end: usize,
    level: usize,
    stats: &mut Vec<FnvHashSet<usize>>,
    set_stats: bool,
    results: &mut Vec<Option<&'a [u8]>>,
    b_quote: &[u64],
    colon_positions: &mut Vec<Vec<usize>>,
) -> Result<()> {
    generate_colon_positions(index, start, end, level, colon_positions);

    let query_num = queries.len();
    let mut found_num = 0;
    let mut vei = end;
    let cp_len = colon_positions[level].len();
    for i in (0..cp_len).rev() {
        let (fsi, fei) = search_pre_field_indices(
            b_quote,
            if i > 0 {
                colon_positions[level][i - 1]
            } else {
                start
            },
            colon_positions[level][i],
        )?;
        let field = &rec[fsi + 1..fei];
        if let Some(query) = queries.get_mut(field) {
            let (vsi, vei) = search_post_value_indices(
                rec,
                colon_positions[level][i] + 1,
                vei,
                if i == cp_len - 1 { RIGHT_BRACE } else { COMMA },
            )?;
            found_num += 1;
            if set_stats && !stats[query.i].contains(&i) {
                stats[query.i].insert(i);
            }
            if let Some(ref mut children) = query.children {
                basic_parse(
                    rec,
                    index,
                    children,
                    vsi,
                    vei,
                    level + 1,
                    stats,
                    set_stats,
                    results,
                    b_quote,
                    colon_positions,
                )?;
            }
            if query.target {
                results[query.ri] = Some(&rec[vsi..vei + 1]);
            }
            if found_num == query_num {
                return Ok(());
            }
        }
        vei = fsi - 1;
    }
    Ok(())
}

#[inline]
pub fn speculative_parse<'a>(rec: &'a [u8], index: &[Vec<u64>], queries: &FnvHashMap<&[u8], Query>, start: usize, end: usize, level: usize, stats: &[FnvHashSet<usize>], results: &mut Vec<Option<&'a [u8]>>, b_quote: &[u64], colon_positions: &mut Vec<Vec<usize>>) -> Result<bool> {
    generate_colon_positions(index, start, end, level, colon_positions);

    for (&s, q) in queries {
        let mut found = false;
        for &i in &stats[q.i] {
            let cp_len = colon_positions[level].len();
            if i >= cp_len {
                continue;
            }
            let (fsi, fei) = search_pre_field_indices(
                b_quote,
                if i > 0 {
                    colon_positions[level][i - 1]
                } else {
                    start
                },
                colon_positions[level][i],
            )?;
            let field = &rec[fsi + 1..fei];
            if s == field {
                let vei = if i < cp_len - 1 {
                    let (nfsi, _) = search_pre_field_indices(
                        b_quote,
                        colon_positions[level][i],
                        colon_positions[level][i + 1],
                    )?;
                    nfsi - 1
                } else {
                    end
                };
                let (vsi, vei) = search_post_value_indices(
                    rec,
                    colon_positions[level][i] + 1,
                    vei,
                    if i == cp_len - 1 { RIGHT_BRACE } else { COMMA },
                )?;
                if let Some(ref children) = q.children {
                    found = speculative_parse(
                        rec,
                        index,
                        children,
                        vsi,
                        vei,
                        level + 1,
                        stats,
                        results,
                        b_quote,
                        colon_positions,
                    )?;
                } else {
                    found = true;
                }
                if q.target {
                    results[q.ri] = Some(&rec[vsi..vei + 1]);
                }
                break;
            }
        }
        if !found {
            return Ok(false);
        }
    }
    Ok(true)
}

#[inline]
fn generate_colon_positions(index: &[Vec<u64>], start: usize, end: usize, level: usize, colon_positions: &mut Vec<Vec<usize>>) {
    let cp = &mut colon_positions[level];
    cp.clear();
    for i in start / 64..(end + 63) / 64 {
        let mut m_colon = index[level][i];
        while m_colon != 0 {
            let m_bit = bit::e(m_colon);
            let offset = i * 64 + (m_bit.trailing_zeros() as usize);
            if start <= offset && offset <= end {
                cp.push(offset);
            }
            m_colon = bit::r(m_colon);
        }
    }
}

#[inline]
fn search_pre_field_indices(b_quote: &[u64], start: usize, end: usize) -> Result<(usize, usize)> {
    let mut si = 0;
    let mut ei = 0;
    let mut ei_set = false;
    let mut n_quote = 0;
    for i in (start / 64..(end + 63) / 64).rev() {
        let mut m_quote = b_quote[i];
        while m_quote != 0 {
            let m_bit = bit::e(m_quote);
            let offset = i * 64 + (m_bit.trailing_zeros() as usize);
            if end <= offset {
                break;
            }
            if start < offset {
                if ei_set {
                    si = offset;
                } else {
                    si = ei;
                    ei = offset;
                }
                n_quote += 1;
            }
            m_quote = bit::r(m_quote);
        }
        if n_quote >= 2 {
            break;
        }
        if n_quote == 1 && !ei_set {
            ei_set = true;
        }
    }
    if n_quote >= 2 {
        Ok((si, ei))
    } else {
        Err(Error::from(ErrorKind::InvalidRecord))
    }
}

#[inline]
fn search_post_value_indices(rec: &[u8], si: usize, ei: usize, ignore_once_char: u8) -> Result<(usize, usize)> {
    let mut si = si;
    let mut ei = ei;
    let mut ignore_once_char_ignored = false;
    let n = rec.len();
    while si < n {
        match rec[si] {
            SPACE | HT | LF | CR => {
                si += 1;
            }
            _ => {
                break;
            }
        }
    }
    if si == n {
        return Err(Error::from(ErrorKind::InvalidRecord));
    }
    while si <= ei {
        match rec[ei] {
            SPACE | HT | LF | CR => {
                ei -= 1;
            }
            char if char == ignore_once_char => {
                if ignore_once_char_ignored {
                    break;
                }
                ignore_once_char_ignored = true;
                ei -= 1;
            }
            _ => {
                break;
            }
        }
    }
    if ei < si {
        return Err(Error::from(ErrorKind::InvalidRecord));
    }
    Ok((si, ei))
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::avx;
    use super::super::index_builder;
    use super::super::utf8::{BACKSLASH, COLON, LEFT_BRACE, QUOTE, RIGHT_BRACE};

    #[test]
    fn test_basic_parse() {
        let json_rec_str = r#"{ "aaa" : "AAA", "bbb" : 111, "ccc": ["C1", "C2"], "ddd" : { "d1" : "D1", "d2" : "D2", "d3": 333 }, "eee": { "e1": "EEE" } } "#;
        let json_rec = json_rec_str.as_bytes();
        let mut b_backslash = Vec::with_capacity((json_rec.len() + 63) / 64);
        let mut b_quote = Vec::with_capacity((json_rec.len() + 63) / 64);
        let mut b_colon = Vec::with_capacity((json_rec.len() + 63) / 64);
        let mut b_left = Vec::with_capacity((json_rec.len() + 63) / 64);
        let mut b_right = Vec::with_capacity((json_rec.len() + 63) / 64);
        index_builder::build_structural_character_bitmap(
            &json_rec,
            &mut b_backslash,
            &mut b_quote,
            &mut b_colon,
            &mut b_left,
            &mut b_right,
            &avx::mm256i(BACKSLASH as i8),
            &avx::mm256i(QUOTE as i8),
            &avx::mm256i(COLON as i8),
            &avx::mm256i(LEFT_BRACE as i8),
            &avx::mm256i(RIGHT_BRACE as i8),
        );
        index_builder::build_structural_quote_bitmap(&b_backslash, &mut b_quote);
        let mut b_string_mask = Vec::with_capacity(b_quote.len());
        index_builder::build_string_mask_bitmap(&b_quote, &mut b_string_mask);
        for i in 0..b_string_mask.len() {
            let b = b_string_mask[i];
            b_colon[i] &= b;
            b_left[i] &= b;
            b_right[i] &= b;
        }
        let l = 10;
        let mut index = vec![Vec::new(); l];
        let mut s_left = Vec::new();
        let r = index_builder::build_leveled_colon_bitmap(&b_colon, &b_left, &b_right, l, &mut s_left, &mut index);
        assert_eq!(Ok(()), r);
        let mut children = FnvHashMap::default();
        children.insert(
            "d1".as_bytes(),
            Query {
                i: 0,
                ri: 0,
                target: true,
                children: None,
            },
        );
        children.insert(
            "d3".as_bytes(),
            Query {
                i: 1,
                ri: 1,
                target: true,
                children: None,
            },
        );
        let mut queries = FnvHashMap::default();
        queries.insert(
            "aaa".as_bytes(),
            Query {
                i: 2,
                ri: 2,
                target: true,
                children: None,
            },
        );
        queries.insert(
            "bbb".as_bytes(),
            Query {
                i: 3,
                ri: 3,
                target: true,
                children: None,
            },
        );
        queries.insert(
            "ccc".as_bytes(),
            Query {
                i: 4,
                ri: 4,
                target: true,
                children: None,
            },
        );
        queries.insert(
            "ddd".as_bytes(),
            Query {
                i: 5,
                ri: 0,
                target: false,
                children: Some(children),
            },
        );
        queries.insert(
            "eee".as_bytes(),
            Query {
                i: 6,
                ri: 5,
                target: true,
                children: None,
            },
        );

        let mut stats = vec![
            FnvHashSet::default(),
            FnvHashSet::default(),
            FnvHashSet::default(),
            FnvHashSet::default(),
            FnvHashSet::default(),
            FnvHashSet::default(),
            FnvHashSet::default(),
        ];

        let mut results = vec![None, None, None, None, None, None];
        let mut colon_positions = vec![Vec::new(); l];
        let result = basic_parse(
            json_rec,
            &index,
            &mut queries,
            0,
            json_rec.len() - 1,
            0,
            &mut stats,
            true,
            &mut results,
            &b_quote,
            &mut colon_positions,
        );

        assert_eq!(Ok(()), result);
        assert_eq!(Some(r#""D1""#.as_bytes()), results[0]);
        assert_eq!(Some(r#"333"#.as_bytes()), results[1]);
        assert_eq!(Some(r#""AAA""#.as_bytes()), results[2]);
        assert_eq!(Some(r#"111"#.as_bytes()), results[3]);
        assert_eq!(Some(r#"["C1", "C2"]"#.as_bytes()), results[4]);
        assert_eq!(Some(r#"{ "e1": "EEE" }"#.as_bytes()), results[5]);
    }

    #[test]
    fn test_generate_colon_positions() {
        struct TestCase {
            index: Vec<Vec<u64>>,
            start: usize,
            end: usize,
            level: usize,
            want: Vec<usize>,
        }
        let test_cases = vec![
            TestCase {
                index: vec![vec![]],
                start: 0,
                end: 0,
                level: 0,
                want: vec![],
            },
            TestCase {
                index: vec![
                    vec![
                        0b00000001_11000000_00000000_00000000_00001000_00000000_00000011_10000000,
                    ],
                ],
                start: 8,
                end: 55,
                level: 0,
                want: vec![8, 9, 27, 54, 55],
            },
            TestCase {
                index: vec![
                    vec![
                        0b10000000_00000000_00000000_00000000_00001000_00000000_00000011_10000000,
                        0b00000001_11000001_00000000_00000011_00000000_00000000_00000000_00000001,
                    ],
                ],
                start: 8,
                end: 119,
                level: 0,
                want: vec![8, 9, 27, 63, 64, 96, 97, 112, 118, 119],
            },
        ];
        for t in test_cases {
            let mut cp = vec![Vec::new(); t.level + 1];
            generate_colon_positions(&t.index, t.start, t.end, t.level, &mut cp);
            assert_eq!(t.want, cp[0]);
        }
    }
}
<|endoftext|>


<commit_before>use std::str::from_utf8;
use nom::IResult;

#[derive(Debug,PartialEq,Eq)]
pub struct PosixHeader<'a> {
    pub name:     & 'a str,
    pub mode:     & 'a str,
    pub uid:      u32,
    pub gid:      u32,
    pub size:     u32,
    pub mtime:    u32,
    pub chksum:   & 'a str,
    pub typeflag: char, /* TODO: enum */
    pub linkname: & 'a str,
    pub ustar:    Option<UStarHeader<'a>>
}

#[derive(Debug,PartialEq,Eq)]
pub struct UStarHeader<'a> {
    pub magic:    & 'a str,
    pub version:  & 'a str,
    pub uname:    & 'a str,
    pub gname:    & 'a str,
    pub devmajor: & 'a str,
    pub devminor: & 'a str,
    pub prefix:   & 'a str,
}

#[derive(Debug,PartialEq,Eq)]
pub struct TarEntry<'a> {
    pub header:   PosixHeader<'a>,
    pub contents: & 'a str
}

/* TODO: validation */
fn str_to_u32(s: &str, base: u32) -> u32 {
    let mut u = 0;
    let mut f = 1;

    for c in s.chars().rev() {
        u += f * ((c as u32) - ('0' as u32));
        f *= base;
    }

    u
}

pub fn octal_to_u32(o: &str) -> u32 {
    str_to_u32(o, 8)
}

pub fn decimal_to_u32(d: &str) -> u32 {
    str_to_u32(d, 10)
}

fn parse_ustar(i: &[u8]) -> IResult<&[u8], Option<UStarHeader>> {
    chain!(i,
        magic:    map_res!(take!(6),   from_utf8) ~
        version:  map_res!(take!(2),   from_utf8) ~
        uname:    map_res!(take!(32),  from_utf8) ~
        gname:    map_res!(take!(32),  from_utf8) ~
        devmajor: map_res!(take!(8),   from_utf8) ~
        devminor: map_res!(take!(8),   from_utf8) ~
        prefix:   map_res!(take!(155), from_utf8),
        ||{
            match magic {
                "ustar" => Some(UStarHeader {
                    magic:    magic,
                    version:  version,
                    uname:    uname,
                    gname:    gname,
                    devmajor: devmajor,
                    devminor: devminor,
                    prefix:   prefix
                }),
                _ => None,
            }
        }
    )
}

fn parse_header(i: &[u8]) -> IResult<&[u8], PosixHeader> {
    chain!(i,
        name:     map_res!(take!(100), from_utf8) ~
        mode:     map_res!(take!(8),   from_utf8) ~
        uid:      map_res!(take!(8),   from_utf8) ~
        gid:      map_res!(take!(8),   from_utf8) ~
        size:     map_res!(take!(12),  from_utf8) ~
        mtime:    map_res!(take!(12),  from_utf8) ~
        chksum:   map_res!(take!(8),   from_utf8) ~
        typeflag: take!(1)                        ~
        linkname: map_res!(take!(100), from_utf8) ~
        ustar:    parse_ustar                     ~
        take!(12), /* padding to 512 */
        ||{
            PosixHeader {
                name:     name,
                mode:     mode,
                uid:      decimal_to_u32(uid),
                gid:      decimal_to_u32(gid),
                size:     octal_to_u32(size),
                mtime:    octal_to_u32(mtime),
                chksum:   chksum,
                typeflag: typeflag[0] as char,
                linkname: linkname,
                ustar:    ustar
            }
        }
    )
}

fn parse_contents(i: &[u8], size: u32) -> IResult<&[u8], &str> {
    let trailing = size % 512;
    let padding = match trailing {
        0 => 0,
        t => 512 - t
    };
    chain!(i,
        contents: map_res!(take!(size as usize), from_utf8) ~
        take!(padding as usize),
        ||{
            contents
        }
    )
}

macro_rules! apply (
 ($i:expr, $fun:expr, $arg:expr ) => ( $fun( $i, $arg ) );
);

fn parse_entry(i: &[u8]) -> IResult<&[u8], TarEntry> {
    chain!(i,
        header:   parse_header ~
        contents: apply!(parse_contents, header.size),
        ||{
            TarEntry {
                header: header,
                contents: contents
            }
        }
    )
}

pub fn parse_tar(i: &[u8]) -> IResult<&[u8], Vec<TarEntry>> {
    many0!(i, parse_entry)
}

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

    #[test]
    fn octal_to_u32_test() {
        match octal_to_u32("756") {
            494 => {},
            o => panic!("octal_to_u32 failed, expected 494 but got {}", o)
        }
        match octal_to_u32("") {
            0 => {},
            o => panic!("octal_to_u32 failed, expected 0 but got {}", o)
        }
    }

    #[test]
    fn decimal_to_u32_test() {
        match decimal_to_u32("756") {
            756 => {},
            d => panic!("decimal_to_u32 failed, expected 756 but got {}", d)
        }
        match decimal_to_u32("") {
            0 => {},
            d => panic!("decimal_to_u32 failed, expected 0 but got {}", d)
        }
    }
}
<commit_msg>use u64<commit_after>use std::str::from_utf8;
use nom::IResult;

#[derive(Debug,PartialEq,Eq)]
pub struct PosixHeader<'a> {
    pub name:     & 'a str,
    pub mode:     & 'a str,
    pub uid:      u64,
    pub gid:      u64,
    pub size:     u64,
    pub mtime:    u64,
    pub chksum:   & 'a str,
    pub typeflag: char, /* TODO: enum */
    pub linkname: & 'a str,
    pub ustar:    Option<UStarHeader<'a>>
}

#[derive(Debug,PartialEq,Eq)]
pub struct UStarHeader<'a> {
    pub magic:    & 'a str,
    pub version:  & 'a str,
    pub uname:    & 'a str,
    pub gname:    & 'a str,
    pub devmajor: & 'a str,
    pub devminor: & 'a str,
    pub prefix:   & 'a str,
}

#[derive(Debug,PartialEq,Eq)]
pub struct TarEntry<'a> {
    pub header:   PosixHeader<'a>,
    pub contents: & 'a str
}

/* TODO: validation */
fn str_to_u64(s: &str, base: u64) -> u64 {
    let mut u = 0;
    let mut f = 1;

    for c in s.chars().rev().skip_while(|&c| c == '\0') {
        u += f * ((c as u64) - ('0' as u64));
        f *= base;
    }

    u
}

pub fn octal_to_u64(o: &str) -> u64 {
    str_to_u64(o, 8)
}

pub fn decimal_to_u64(d: &str) -> u64 {
    str_to_u64(d, 10)
}

fn parse_ustar(i: &[u8]) -> IResult<&[u8], Option<UStarHeader>> {
    chain!(i,
        magic:    map_res!(take!(6),   from_utf8) ~
        version:  map_res!(take!(2),   from_utf8) ~
        uname:    map_res!(take!(32),  from_utf8) ~
        gname:    map_res!(take!(32),  from_utf8) ~
        devmajor: map_res!(take!(8),   from_utf8) ~
        devminor: map_res!(take!(8),   from_utf8) ~
        prefix:   map_res!(take!(155), from_utf8),
        ||{
            match magic {
                "ustar" => Some(UStarHeader {
                    magic:    magic,
                    version:  version,
                    uname:    uname,
                    gname:    gname,
                    devmajor: devmajor,
                    devminor: devminor,
                    prefix:   prefix
                }),
                _ => None,
            }
        }
    )
}

fn parse_header(i: &[u8]) -> IResult<&[u8], PosixHeader> {
    chain!(i,
        name:     map_res!(take!(100), from_utf8) ~
        mode:     map_res!(take!(8),   from_utf8) ~
        uid:      map_res!(take!(8),   from_utf8) ~
        gid:      map_res!(take!(8),   from_utf8) ~
        size:     map_res!(take!(12),  from_utf8) ~
        mtime:    map_res!(take!(12),  from_utf8) ~
        chksum:   map_res!(take!(8),   from_utf8) ~
        typeflag: take!(1)                        ~
        linkname: map_res!(take!(100), from_utf8) ~
        ustar:    parse_ustar                     ~
        take!(12), /* padding to 512 */
        ||{
            PosixHeader {
                name:     name,
                mode:     mode,
                uid:      decimal_to_u64(uid),
                gid:      decimal_to_u64(gid),
                size:     octal_to_u64(size),
                mtime:    octal_to_u64(mtime), /* TODO: u64 */
                chksum:   chksum,
                typeflag: typeflag[0] as char,
                linkname: linkname,
                ustar:    ustar
            }
        }
    )
}

fn parse_contents(i: &[u8], size: u64) -> IResult<&[u8], &str> {
    let trailing = size % 512;
    let padding = match trailing {
        0 => 0,
        t => 512 - t
    };
    chain!(i,
        contents: map_res!(take!(size as usize), from_utf8) ~
        take!(padding as usize),
        ||{
            contents
        }
    )
}

macro_rules! apply (
 ($i:expr, $fun:expr, $arg:expr ) => ( $fun( $i, $arg ) );
);

fn parse_entry(i: &[u8]) -> IResult<&[u8], TarEntry> {
    chain!(i,
        header:   parse_header ~
        contents: apply!(parse_contents, header.size),
        ||{
            TarEntry {
                header: header,
                contents: contents
            }
        }
    )
}

pub fn parse_tar(i: &[u8]) -> IResult<&[u8], Vec<TarEntry>> {
    many0!(i, parse_entry)
}

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

    #[test]
    fn octal_to_u32_test() {
        match octal_to_u32("756") {
            494 => {},
            o => panic!("octal_to_u32 failed, expected 494 but got {}", o)
        }
        match octal_to_u32("") {
            0 => {},
            o => panic!("octal_to_u32 failed, expected 0 but got {}", o)
        }
    }

    #[test]
    fn decimal_to_u32_test() {
        match decimal_to_u32("756") {
            756 => {},
            d => panic!("decimal_to_u32 failed, expected 756 but got {}", d)
        }
        match decimal_to_u32("") {
            0 => {},
            d => panic!("decimal_to_u32 failed, expected 0 but got {}", d)
        }
    }
}
<|endoftext|>