<commit_before>//
// writer.rs
// Copyright (C) 2015 Adrian Perez <aperez@igalia.com>
// Distributed under terms of the MIT license.
//

use std::io;
use std::collections::BTreeMap;
use super::value::Value;

trait Formatter {
    fn start_compound<W>(&mut self, writer: &mut W, ch: u8) -> io::Result<()>
        where W: io::Write;
    fn end_compound<W>(&mut self, writer: &mut W, ch: u8) -> io::Result<()>
        where W: io::Write;
    fn key_separator<W>(&mut self, writer: &mut W, next: &Value) -> io::Result<()>
        where W: io::Write;
    fn item_separator<W>(&mut self, writer: &mut W, next: &Value) -> io::Result<()>
        where W: io::Write;
}

pub struct CompactFormatter;

impl Formatter for CompactFormatter {
    fn start_compound<W>(&mut self, writer: &mut W, ch: u8) -> io::Result<()>
        where W: io::Write
    {
        writer.write_all(&[ch])
    }

    fn end_compound<W>(&mut self, writer: &mut W, ch: u8) -> io::Result<()>
        where W: io::Write
    {
        writer.write_all(&[ch])
    }

    fn key_separator<W>(&mut self, writer: &mut W, next: &Value) -> io::Result<()>
        where W: io::Write
    {
        match *next {
            Value::Dict(_) | Value::List(_) => Ok(()),
            _ => writer.write_all(b":"),
        }
    }

    fn item_separator<W>(&mut self, writer: &mut W, next: &Value) -> io::Result<()>
        where W: io::Write
    {
        match *next {
            Value::Dict(_) | Value::List(_) => Ok(()),
            _ => writer.write_all(b","),
        }
    }
}

pub struct PrettyFormatter {
    indent: usize,
}

impl PrettyFormatter {
    fn new() -> Self {
        PrettyFormatter { indent: 0 }
    }
}

#[inline]
fn indent<W>(writer: &mut W, indent: usize) -> io::Result<()>
    where W: io::Write
{
    for _ in 0..indent {
        try!(writer.write_all(b"  "));
    }
    Ok(())
}

impl Formatter for PrettyFormatter {
    fn start_compound<W>(&mut self, writer: &mut W, ch: u8) -> io::Result<()>
        where W: io::Write
    {
        self.indent += 1;
        try!(writer.write_all(&[ch, b'\n']));
        indent(writer, self.indent)
    }

    fn end_compound<W>(&mut self, writer: &mut W, ch: u8) -> io::Result<()>
        where W: io::Write
    {
        self.indent -= 1;
        try!(writer.write(b"\n"));
        try!(indent(writer, self.indent));
        writer.write_all(&[ch])
    }

    fn key_separator<W>(&mut self, writer: &mut W, next: &Value) -> io::Result<()>
        where W: io::Write
    {
        match *next {
            Value::Dict(_) | Value::List(_) => writer.write_all(b" "),
            _ => writer.write_all(b": "),
        }
    }

    fn item_separator<W>(&mut self, writer: &mut W, next: &Value) -> io::Result<()>
        where W: io::Write
    {
        try!(writer.write(b"\n"));
        indent(writer, self.indent)
    }
}

pub struct Writer<W, F=PrettyFormatter> {
    writer: W,
    format: F,
}

impl<W> Writer<W, CompactFormatter>
    where W: io::Write
{
    #[inline]
    pub fn compact(writer: W) -> Self {
        Writer::with_formatter(writer, CompactFormatter)
    }
}

impl<W> Writer<W>
    where W: io::Write
{
    #[inline]
    pub fn pretty(writer: W) -> Self {
        Writer::with_formatter(writer, PrettyFormatter::new())
    }
}

impl<W, F> Writer<W, F>
    where W: io::Write,
          F: Formatter,
{
    #[inline]
    fn with_formatter(writer: W, formatter: F) -> Self
        where W: io::Write,
              F: Formatter,
    {
        Writer {
            writer: writer,
            format: formatter,
        }
    }

    fn write_value(&mut self, value: &Value) -> io::Result<()> {
        match *value {
            Value::Bool(value) => self.write_bool(value),
            Value::Float(value) => self.write_float(value),
            Value::Integer(value) => self.write_integer(value),
            Value::String(ref value) => self.write_string(value),
            Value::List(ref value) => self.write_list(value),
            Value::Dict(ref value) => self.write_dict(value),
        }
    }

    #[inline]
    fn write_bool(&mut self, value: bool) -> io::Result<()> {
        if value {
            self.writer.write_all(b"True")
        } else {
            self.writer.write_all(b"False")
        }
    }

    #[inline]
    fn write_float(&mut self, value: f64) -> io::Result<()> {
        if value.is_nan() || value.is_infinite() {
            write!(self.writer, "{}", value)
        } else {
            let s = format!("{}", value);
            try!(self.writer.write_all(s.as_bytes()));
            if !s.contains(".") {
                try!(self.writer.write_all(b".0"));
            }
            Ok(())
        }
    }

    #[inline]
    fn write_integer(&mut self, value: i64) -> io::Result<()> {
        write!(self.writer, "{}", value)
    }

    #[inline]
    fn write_string(&mut self, value: &String) -> io::Result<()> {
        try!(self.writer.write_all(b"\""));
        for ch in value.bytes() {
            try!(match ch {
                0x09 => self.writer.write_all(b"\\t"),
                0x0A => self.writer.write_all(b"\\n"),
                0x0D => self.writer.write_all(b"\\r"),
                0x22 => self.writer.write_all(b"\\\""),
                0x5C => self.writer.write_all(b"\\\\"),
                ch if ch < 0xF => write!(self.writer, "\\0{:X}", ch),
                ch if ch < 0x20 => write!(self.writer, "\\{:X}", ch),
                ch => self.writer.write_all(&[ch]),
            });
        }
        self.writer.write_all(b"\"")
    }

    fn write_list(&mut self, value: &Vec<Value>) -> io::Result<()> {
        if value.is_empty() {
            return self.writer.write_all(b"[]");
        }

        let mut iter = value.iter().peekable();
        try!(self.format.start_compound(&mut self.writer, b'['));
        loop {
            match iter.next() {
                None => break,
                Some(value) => {
                    try!(self.write_value(value));
                    match iter.peek () {
                        Some(_) =>
                            try!(self.format.item_separator(&mut self.writer, value)),
                        None => (),
                    }
                },
            }
        }
        self.format.end_compound(&mut self.writer, b']')
    }

    #[inline]
    fn write_dict(&mut self, dict: &BTreeMap<String, Value>) -> io::Result<()> {
        if dict.is_empty() {
            return self.writer.write_all(b"{}");
        }

        try!(self.format.start_compound(&mut self.writer, b'{'));
        try!(self.write_keyval_items(dict));
        self.format.end_compound(&mut self.writer, b'}')
    }

    fn write_keyval_items(&mut self, dict: &BTreeMap<String, Value>) -> io::Result<()> {
        let mut iter = dict.iter().peekable();
        loop {
            match iter.next() {
                None => break,
                Some((key, value)) => {
                    try!(self.writer.write_all(key.as_bytes()));
                    try!(self.format.key_separator(&mut self.writer, value));
                    try!(self.write_value(value));
                    match iter.peek () {
                        Some(&(_, ref value)) =>
                            try!(self.format.item_separator(&mut self.writer, value)),
                        None => (),
                    }
                },
            }
        }
        Ok(())
    }

    #[inline]
    pub fn write_message(&mut self, value: &BTreeMap<String, Value>) -> io::Result<()> {
        self.write_keyval_items(value)
    }

    #[inline]
    pub fn write_framed_message(&mut self, value: &BTreeMap<String, Value>) -> io::Result<()> {
        self.write_dict(value)
    }
}

#[cfg(test)]
mod tests {
	use super::*;
    use super::super::value::Value;
    use std::io::Cursor;
    use std::collections::BTreeMap;
    use std::f64::{NAN, INFINITY};

    macro_rules! make_write_test {
        ($name:ident, $value:expr, $pretty:expr, $compact:expr) => {
            #[test]
            fn $name() {
                {
                    println!("Pretty");
                    let mut output = Cursor::new(Vec::new());
                    Writer::pretty(&mut output).write_value($value).ok();
                    assert_eq!($pretty, String::from_utf8(output.into_inner()).unwrap());
                }
                {
                    println!("Compact");
                    let mut output = Cursor::new(Vec::new());
                    Writer::compact(&mut output).write_value($value).ok();
                    assert_eq!($compact, String::from_utf8(output.into_inner()).unwrap());
                }
            }
        }
    }

    make_write_test!(bool_true,  &Value::Bool(true),  "True", "True");
    make_write_test!(bool_false, &Value::Bool(false), "False", "False");

    make_write_test!(list_empty, &Value::List(vec![]), "[]", "[]");
    make_write_test!(list_one,   &Value::List(vec![Value::Bool(true)]), "[\n  True\n]", "[True]");
    make_write_test!(list_two,   &Value::List(vec![Value::Bool(true), Value::Bool(false)]),
        "[\n  True\n  False\n]", "[True,False]");
    make_write_test!(list_nested,
                     &Value::List(vec![
                                  Value::Bool(true),
                                  Value::List(vec![Value::Bool(false)])]),
                     "[\n  True\n  [\n    False\n  ]\n]", "[True,[False]]");

    make_write_test!(dict_empty, &Value::Dict(BTreeMap::new()), "{}", "{}");
    make_write_test!(dict_one,   &Value::Dict((|| {
        let mut b = BTreeMap::new();
        b.insert("item".to_string(), Value::Bool(true));
        b
    })()), "{\n  item: True\n}", "{item:True}");
    make_write_test!(dict_two,   &Value::Dict((|| {
        let mut b = BTreeMap::new();
        b.insert("foo".to_string(), Value::Bool(true));
        b.insert("bar".to_string(), Value::Bool(false));
        b
    })()), "{\n  bar: False\n  foo: True\n}", "{bar:False,foo:True}");

    macro_rules! make_write_string_tests {
        ($($name:ident, $value:expr, $expected:expr),+) => {
            $( make_write_test!($name, &Value::String($value.to_string()), $expected, $expected); )*
        }
    }

    make_write_string_tests!(string_empty, "", "\"\"",
                             string_non_empty, "foo bar", "\"foo bar\"",
                             string_unicode, "☺", "\"☺\"",
                             string_escapes, "\n\r\t\\\"", "\"\\n\\r\\t\\\\\\\"\"",
                             string_hexcode, "\0", "\"\\00\"");

    macro_rules! make_write_number_tests {
        ($t:ident, $($name:ident, $value:expr, $expected:expr),+) => {
            $( make_write_test!($name, &Value::$t($value), $expected, $expected); )*
        }
    }

    make_write_number_tests!(Integer,
                             integer_zero, 0, "0",
                             integer_negative, -34, "-34");
    make_write_number_tests!(Float,
                             float_zero, 0.0, "0.0",
                             float_suffix, 1f64, "1.0",
                             float_positive, 4.5, "4.5",
                             float_negative, -3.2, "-3.2",
                             float_nan, NAN, "NaN",
                             float_infinite, INFINITY, "inf");
}
<commit_msg>Silence warning about unused parameter<commit_after>//
// writer.rs
// Copyright (C) 2015 Adrian Perez <aperez@igalia.com>
// Distributed under terms of the MIT license.
//

use std::io;
use std::collections::BTreeMap;
use super::value::Value;

trait Formatter {
    fn start_compound<W>(&mut self, writer: &mut W, ch: u8) -> io::Result<()>
        where W: io::Write;
    fn end_compound<W>(&mut self, writer: &mut W, ch: u8) -> io::Result<()>
        where W: io::Write;
    fn key_separator<W>(&mut self, writer: &mut W, next: &Value) -> io::Result<()>
        where W: io::Write;
    fn item_separator<W>(&mut self, writer: &mut W, next: &Value) -> io::Result<()>
        where W: io::Write;
}

pub struct CompactFormatter;

impl Formatter for CompactFormatter {
    fn start_compound<W>(&mut self, writer: &mut W, ch: u8) -> io::Result<()>
        where W: io::Write
    {
        writer.write_all(&[ch])
    }

    fn end_compound<W>(&mut self, writer: &mut W, ch: u8) -> io::Result<()>
        where W: io::Write
    {
        writer.write_all(&[ch])
    }

    fn key_separator<W>(&mut self, writer: &mut W, next: &Value) -> io::Result<()>
        where W: io::Write
    {
        match *next {
            Value::Dict(_) | Value::List(_) => Ok(()),
            _ => writer.write_all(b":"),
        }
    }

    fn item_separator<W>(&mut self, writer: &mut W, next: &Value) -> io::Result<()>
        where W: io::Write
    {
        match *next {
            Value::Dict(_) | Value::List(_) => Ok(()),
            _ => writer.write_all(b","),
        }
    }
}

pub struct PrettyFormatter {
    indent: usize,
}

impl PrettyFormatter {
    fn new() -> Self {
        PrettyFormatter { indent: 0 }
    }
}

#[inline]
fn indent<W>(writer: &mut W, indent: usize) -> io::Result<()>
    where W: io::Write
{
    for _ in 0..indent {
        try!(writer.write_all(b"  "));
    }
    Ok(())
}

impl Formatter for PrettyFormatter {
    fn start_compound<W>(&mut self, writer: &mut W, ch: u8) -> io::Result<()>
        where W: io::Write
    {
        self.indent += 1;
        try!(writer.write_all(&[ch, b'\n']));
        indent(writer, self.indent)
    }

    fn end_compound<W>(&mut self, writer: &mut W, ch: u8) -> io::Result<()>
        where W: io::Write
    {
        self.indent -= 1;
        try!(writer.write(b"\n"));
        try!(indent(writer, self.indent));
        writer.write_all(&[ch])
    }

    fn key_separator<W>(&mut self, writer: &mut W, next: &Value) -> io::Result<()>
        where W: io::Write
    {
        match *next {
            Value::Dict(_) | Value::List(_) => writer.write_all(b" "),
            _ => writer.write_all(b": "),
        }
    }

    fn item_separator<W>(&mut self, writer: &mut W, _next: &Value) -> io::Result<()>
        where W: io::Write
    {
        try!(writer.write(b"\n"));
        indent(writer, self.indent)
    }
}

pub struct Writer<W, F=PrettyFormatter> {
    writer: W,
    format: F,
}

impl<W> Writer<W, CompactFormatter>
    where W: io::Write
{
    #[inline]
    pub fn compact(writer: W) -> Self {
        Writer::with_formatter(writer, CompactFormatter)
    }
}

impl<W> Writer<W>
    where W: io::Write
{
    #[inline]
    pub fn pretty(writer: W) -> Self {
        Writer::with_formatter(writer, PrettyFormatter::new())
    }
}

impl<W, F> Writer<W, F>
    where W: io::Write,
          F: Formatter,
{
    #[inline]
    fn with_formatter(writer: W, formatter: F) -> Self
        where W: io::Write,
              F: Formatter,
    {
        Writer {
            writer: writer,
            format: formatter,
        }
    }

    fn write_value(&mut self, value: &Value) -> io::Result<()> {
        match *value {
            Value::Bool(value) => self.write_bool(value),
            Value::Float(value) => self.write_float(value),
            Value::Integer(value) => self.write_integer(value),
            Value::String(ref value) => self.write_string(value),
            Value::List(ref value) => self.write_list(value),
            Value::Dict(ref value) => self.write_dict(value),
        }
    }

    #[inline]
    fn write_bool(&mut self, value: bool) -> io::Result<()> {
        if value {
            self.writer.write_all(b"True")
        } else {
            self.writer.write_all(b"False")
        }
    }

    #[inline]
    fn write_float(&mut self, value: f64) -> io::Result<()> {
        if value.is_nan() || value.is_infinite() {
            write!(self.writer, "{}", value)
        } else {
            let s = format!("{}", value);
            try!(self.writer.write_all(s.as_bytes()));
            if !s.contains(".") {
                try!(self.writer.write_all(b".0"));
            }
            Ok(())
        }
    }

    #[inline]
    fn write_integer(&mut self, value: i64) -> io::Result<()> {
        write!(self.writer, "{}", value)
    }

    #[inline]
    fn write_string(&mut self, value: &String) -> io::Result<()> {
        try!(self.writer.write_all(b"\""));
        for ch in value.bytes() {
            try!(match ch {
                0x09 => self.writer.write_all(b"\\t"),
                0x0A => self.writer.write_all(b"\\n"),
                0x0D => self.writer.write_all(b"\\r"),
                0x22 => self.writer.write_all(b"\\\""),
                0x5C => self.writer.write_all(b"\\\\"),
                ch if ch < 0xF => write!(self.writer, "\\0{:X}", ch),
                ch if ch < 0x20 => write!(self.writer, "\\{:X}", ch),
                ch => self.writer.write_all(&[ch]),
            });
        }
        self.writer.write_all(b"\"")
    }

    fn write_list(&mut self, value: &Vec<Value>) -> io::Result<()> {
        if value.is_empty() {
            return self.writer.write_all(b"[]");
        }

        let mut iter = value.iter().peekable();
        try!(self.format.start_compound(&mut self.writer, b'['));
        loop {
            match iter.next() {
                None => break,
                Some(value) => {
                    try!(self.write_value(value));
                    match iter.peek () {
                        Some(_) =>
                            try!(self.format.item_separator(&mut self.writer, value)),
                        None => (),
                    }
                },
            }
        }
        self.format.end_compound(&mut self.writer, b']')
    }

    #[inline]
    fn write_dict(&mut self, dict: &BTreeMap<String, Value>) -> io::Result<()> {
        if dict.is_empty() {
            return self.writer.write_all(b"{}");
        }

        try!(self.format.start_compound(&mut self.writer, b'{'));
        try!(self.write_keyval_items(dict));
        self.format.end_compound(&mut self.writer, b'}')
    }

    fn write_keyval_items(&mut self, dict: &BTreeMap<String, Value>) -> io::Result<()> {
        let mut iter = dict.iter().peekable();
        loop {
            match iter.next() {
                None => break,
                Some((key, value)) => {
                    try!(self.writer.write_all(key.as_bytes()));
                    try!(self.format.key_separator(&mut self.writer, value));
                    try!(self.write_value(value));
                    match iter.peek () {
                        Some(&(_, ref value)) =>
                            try!(self.format.item_separator(&mut self.writer, value)),
                        None => (),
                    }
                },
            }
        }
        Ok(())
    }

    #[inline]
    pub fn write_message(&mut self, value: &BTreeMap<String, Value>) -> io::Result<()> {
        self.write_keyval_items(value)
    }

    #[inline]
    pub fn write_framed_message(&mut self, value: &BTreeMap<String, Value>) -> io::Result<()> {
        self.write_dict(value)
    }
}

#[cfg(test)]
mod tests {
	use super::*;
    use super::super::value::Value;
    use std::io::Cursor;
    use std::collections::BTreeMap;
    use std::f64::{NAN, INFINITY};

    macro_rules! make_write_test {
        ($name:ident, $value:expr, $pretty:expr, $compact:expr) => {
            #[test]
            fn $name() {
                {
                    println!("Pretty");
                    let mut output = Cursor::new(Vec::new());
                    Writer::pretty(&mut output).write_value($value).ok();
                    assert_eq!($pretty, String::from_utf8(output.into_inner()).unwrap());
                }
                {
                    println!("Compact");
                    let mut output = Cursor::new(Vec::new());
                    Writer::compact(&mut output).write_value($value).ok();
                    assert_eq!($compact, String::from_utf8(output.into_inner()).unwrap());
                }
            }
        }
    }

    make_write_test!(bool_true,  &Value::Bool(true),  "True", "True");
    make_write_test!(bool_false, &Value::Bool(false), "False", "False");

    make_write_test!(list_empty, &Value::List(vec![]), "[]", "[]");
    make_write_test!(list_one,   &Value::List(vec![Value::Bool(true)]), "[\n  True\n]", "[True]");
    make_write_test!(list_two,   &Value::List(vec![Value::Bool(true), Value::Bool(false)]),
        "[\n  True\n  False\n]", "[True,False]");
    make_write_test!(list_nested,
                     &Value::List(vec![
                                  Value::Bool(true),
                                  Value::List(vec![Value::Bool(false)])]),
                     "[\n  True\n  [\n    False\n  ]\n]", "[True,[False]]");

    make_write_test!(dict_empty, &Value::Dict(BTreeMap::new()), "{}", "{}");
    make_write_test!(dict_one,   &Value::Dict((|| {
        let mut b = BTreeMap::new();
        b.insert("item".to_string(), Value::Bool(true));
        b
    })()), "{\n  item: True\n}", "{item:True}");
    make_write_test!(dict_two,   &Value::Dict((|| {
        let mut b = BTreeMap::new();
        b.insert("foo".to_string(), Value::Bool(true));
        b.insert("bar".to_string(), Value::Bool(false));
        b
    })()), "{\n  bar: False\n  foo: True\n}", "{bar:False,foo:True}");

    macro_rules! make_write_string_tests {
        ($($name:ident, $value:expr, $expected:expr),+) => {
            $( make_write_test!($name, &Value::String($value.to_string()), $expected, $expected); )*
        }
    }

    make_write_string_tests!(string_empty, "", "\"\"",
                             string_non_empty, "foo bar", "\"foo bar\"",
                             string_unicode, "☺", "\"☺\"",
                             string_escapes, "\n\r\t\\\"", "\"\\n\\r\\t\\\\\\\"\"",
                             string_hexcode, "\0", "\"\\00\"");

    macro_rules! make_write_number_tests {
        ($t:ident, $($name:ident, $value:expr, $expected:expr),+) => {
            $( make_write_test!($name, &Value::$t($value), $expected, $expected); )*
        }
    }

    make_write_number_tests!(Integer,
                             integer_zero, 0, "0",
                             integer_negative, -34, "-34");
    make_write_number_tests!(Float,
                             float_zero, 0.0, "0.0",
                             float_suffix, 1f64, "1.0",
                             float_positive, 4.5, "4.5",
                             float_negative, -3.2, "-3.2",
                             float_nan, NAN, "NaN",
                             float_infinite, INFINITY, "inf");
}
<|endoftext|>


<commit_before>//! Write las points.
//!
//! A `Writer` uses a `Header` for its configuration:
//!
//! ```
//! use std::io::Cursor;
//! use las::{Writer, Header};
//! let mut header = Header::default();
//! header.version = (1, 4).into();
//! let writer = Writer::new(Cursor::new(Vec::new()), header).unwrap();
//! ```
//!
//! The set of optional fields on the point format and the points must match exactly:
//!
//! ```
//! use std::io::Cursor;
//! use las::{Writer, Header, Point};
//! use las::point::Format;
//! use las::Color;
//!
//! let mut header = Header::default();
//! header.point_format = Format::new(1).unwrap();
//! let mut writer = Writer::new(Cursor::new(Vec::new()), header).unwrap();
//!
//! let mut point = Point::default(); // default points don't have any optional attributes
//! assert!(writer.write(point.clone()).is_err());
//!
//! point.gps_time = Some(42.); // point format 1 requires gps time
//! writer.write(point.clone()).unwrap();
//!
//! point.color = Some(Color::new(1, 2, 3));
//! assert!(writer.write(point).is_err()); // the point's color would be lost
//! ```

use {Header, Point, Result};
use point::Format;
use std::collections::HashMap;
use std::fs::File;
use std::io::{BufWriter, Seek, SeekFrom, Write};
use std::path::Path;

quick_error! {
    /// Writer errors.
    #[derive(Debug)]
    pub enum Error {
        /// The writer is closed.
        Closed {
            description("the writer is closed")
        }
        /// The attributes of the point format and point do not match.
        PointAttributes(format: Format, point: Point) {
            description("the attributes of the point format and point do not match")
            display("the attributes of point format {:?} does not match point {:?}", format, point)
        }
    }
}

/// Writes LAS data.
///
/// The LAS header needs to be re-written when the writer closes. For convenience, this is done via
/// the `Drop` implementation of the writer. One consequence is that if the header re-write fails
/// during the drop, a panic will result. If you want to check for errors instead of panicing, use
/// `close` explicitly.
///
/// ```
/// use std::io::Cursor;
/// use las::Writer;
/// {
///     let mut writer = Writer::new(Cursor::new(Vec::new()), Default::default()).unwrap();
///     writer.close().unwrap();
/// } // <- `close` is not called
/// ```
#[derive(Debug)]
pub struct Writer<W: Seek + Write> {
    closed: bool,
    header: Header,
    write: W,
}

impl<W: Seek + Write> Writer<W> {
    /// Creates a new writer.
    ///
    /// The header that is passed in will have various fields zero'd, e.g. bounds, number of
    /// points, etc.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::io::Cursor;
    /// use las::Writer;
    /// let writer = Writer::new(Cursor::new(Vec::new()), Default::default());
    /// ```
    pub fn new(mut write: W, mut header: Header) -> Result<Writer<W>> {
        header.bounds = Default::default();
        header.number_of_points = 0;
        header.number_of_points_by_return = HashMap::new();
        if header.version.requires_point_data_start_signature() {
            // TODO we shouldn't overwrite the old vlr padding
            header.vlr_padding = ::raw::POINT_DATA_START_SIGNATURE.to_vec();
        }
        header.clone().into_raw().and_then(|raw_header| {
            raw_header.write_to(&mut write)
        })?;
        for vlr in header.vlrs().iter() {
            (*vlr).clone().into_raw(false).and_then(|raw_vlr| {
                raw_vlr.write_to(&mut write)
            })?;
        }
        if !header.vlr_padding.is_empty() {
            write.write_all(&header.vlr_padding)?;
        }
        Ok(Writer {
            closed: false,
            header: header,
            write: write,
        })
    }

    /// Writes a point.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::io::Cursor;
    /// use las::Writer;
    ///
    /// let mut writer = Writer::new(Cursor::new(Vec::new()), Default::default()).unwrap();
    /// writer.write(Default::default()).unwrap();
    /// ```
    pub fn write(&mut self, point: Point) -> Result<()> {
        if self.closed {
            return Err(Error::Closed.into());
        }
        if !point.matches(self.header.point_format) {
            return Err(
                Error::PointAttributes(self.header.point_format, point).into(),
            );
        }
        self.header.number_of_points += 1;
        {
            let entry = self.header
                .number_of_points_by_return
                .entry(point.return_number)
                .or_insert(0);
            *entry += 1;
        }
        self.header.bounds.grow(&point);
        point.into_raw(self.header.transforms).and_then(
            |raw_point| {
                raw_point.write_to(&mut self.write, self.header.point_format)
            },
        )?;
        Ok(())
    }

    /// Close this writer.
    ///
    /// A second call to close is a no-op.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::io::Cursor;
    /// use las::Writer;
    /// let mut writer = Writer::new(Cursor::new(Vec::new()), Default::default()).unwrap();
    /// writer.close().unwrap();
    /// writer.close().unwrap(); // <- no-op
    ///
    pub fn close(&mut self) -> Result<()> {
        if !self.closed {
            for raw_evlr in self.header.evlrs().into_iter().map(|evlr| {
                evlr.clone().into_raw(true)
            })
            {
                raw_evlr?.write_to(&mut self.write)?;
            }
            self.write.seek(SeekFrom::Start(0))?;
            self.header.clone().into_raw().and_then(|raw_header| {
                raw_header.write_to(&mut self.write)
            })?;
            self.closed = true;
        }
        Ok(())
    }
}

impl Writer<BufWriter<File>> {
    /// Creates a new writer for a path.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::Writer;
    /// let writer = Writer::from_path("/dev/null", Default::default());
    /// ```
    pub fn from_path<P: AsRef<Path>>(path: P, header: Header) -> Result<Writer<BufWriter<File>>> {
        File::create(path).map_err(::Error::from).and_then(|file| {
            Writer::new(BufWriter::new(file), header)
        })
    }
}

impl<W: Seek + Write> Drop for Writer<W> {
    fn drop(&mut self) {
        if !self.closed {
            self.close().expect("Error when dropping the writer");
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use {Header, Version};
    use byteorder::{LittleEndian, ReadBytesExt};
    use point::Format;
    use std::io::Cursor;

    fn writer(format: Format, version: Version) -> Writer<Cursor<Vec<u8>>> {
        let mut header = Header::default();
        header.point_format = format;
        header.version = version;
        Writer::new(Cursor::new(Vec::new()), header).unwrap()
    }

    #[test]
    fn las_1_0_point_data_start_signature() {
        let mut cursor = Cursor::new(Vec::new());
        {
            let mut header = Header::default();
            header.version = (1, 0).into();
            header.push_vlr(Default::default());
            let mut writer = Writer::new(&mut cursor, header).unwrap();
            writer.write(Default::default()).unwrap();
        }
        cursor.set_position(281);
        assert_eq!(0xCCDD, cursor.read_u16::<LittleEndian>().unwrap());
    }

    #[test]
    fn already_closed() {
        let mut writer = Writer::new(Cursor::new(Vec::new()), Default::default()).unwrap();
        writer.close().unwrap();
        assert!(writer.write(Default::default()).is_err());
    }

    #[test]
    fn missing_extra_bytes() {
        let format = Format {
            extra_bytes: 1,
            ..Default::default()
        };
        let mut writer = writer(format, Version::new(1, 4));
        assert!(writer.write(Default::default()).is_err());
    }

    #[test]
    fn missing_gps_time() {
        let format = Format::new(1).unwrap();
        let mut writer = writer(format, Version::new(1, 2));
        assert!(writer.write(Default::default()).is_err());
    }

    #[test]
    fn missing_color() {
        let format = Format::new(2).unwrap();
        let mut writer = writer(format, Version::new(1, 2));
        assert!(writer.write(Default::default()).is_err());
    }

    #[test]
    fn missing_nir() {
        let format = Format::new(8).unwrap();
        let mut writer = writer(format, Version::new(1, 4));
        let point = Point {
            gps_time: Some(0.),
            color: Some(Default::default()),
            ..Default::default()
        };
        assert!(writer.write(point).is_err());
    }

    #[test]
    fn missing_waveform() {
        let format = Format::new(4).unwrap();
        let mut writer = writer(format, Version::new(1, 4));
        assert!(writer.write(Default::default()).is_err());
    }
}
<commit_msg>A second close for a writer is an error<commit_after>//! Write las points.
//!
//! A `Writer` uses a `Header` for its configuration:
//!
//! ```
//! use std::io::Cursor;
//! use las::{Writer, Header};
//! let mut header = Header::default();
//! header.version = (1, 4).into();
//! let writer = Writer::new(Cursor::new(Vec::new()), header).unwrap();
//! ```
//!
//! The set of optional fields on the point format and the points must match exactly:
//!
//! ```
//! use std::io::Cursor;
//! use las::{Writer, Header, Point};
//! use las::point::Format;
//! use las::Color;
//!
//! let mut header = Header::default();
//! header.point_format = Format::new(1).unwrap();
//! let mut writer = Writer::new(Cursor::new(Vec::new()), header).unwrap();
//!
//! let mut point = Point::default(); // default points don't have any optional attributes
//! assert!(writer.write(point.clone()).is_err());
//!
//! point.gps_time = Some(42.); // point format 1 requires gps time
//! writer.write(point.clone()).unwrap();
//!
//! point.color = Some(Color::new(1, 2, 3));
//! assert!(writer.write(point).is_err()); // the point's color would be lost
//! ```

use {Header, Point, Result};
use point::Format;
use std::collections::HashMap;
use std::fs::File;
use std::io::{BufWriter, Seek, SeekFrom, Write};
use std::path::Path;

quick_error! {
    /// Writer errors.
    #[derive(Debug)]
    pub enum Error {
        /// The writer is closed.
        Closed {
            description("the writer is closed")
        }
        /// The attributes of the point format and point do not match.
        PointAttributes(format: Format, point: Point) {
            description("the attributes of the point format and point do not match")
            display("the attributes of point format {:?} does not match point {:?}", format, point)
        }
    }
}

/// Writes LAS data.
///
/// The LAS header needs to be re-written when the writer closes. For convenience, this is done via
/// the `Drop` implementation of the writer. One consequence is that if the header re-write fails
/// during the drop, a panic will result. If you want to check for errors instead of panicing, use
/// `close` explicitly.
///
/// ```
/// use std::io::Cursor;
/// use las::Writer;
/// {
///     let mut writer = Writer::new(Cursor::new(Vec::new()), Default::default()).unwrap();
///     writer.close().unwrap();
/// } // <- `close` is not called
/// ```
#[derive(Debug)]
pub struct Writer<W: Seek + Write> {
    closed: bool,
    header: Header,
    write: W,
}

impl<W: Seek + Write> Writer<W> {
    /// Creates a new writer.
    ///
    /// The header that is passed in will have various fields zero'd, e.g. bounds, number of
    /// points, etc.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::io::Cursor;
    /// use las::Writer;
    /// let writer = Writer::new(Cursor::new(Vec::new()), Default::default());
    /// ```
    pub fn new(mut write: W, mut header: Header) -> Result<Writer<W>> {
        header.bounds = Default::default();
        header.number_of_points = 0;
        header.number_of_points_by_return = HashMap::new();
        if header.version.requires_point_data_start_signature() {
            // TODO we shouldn't overwrite the old vlr padding
            header.vlr_padding = ::raw::POINT_DATA_START_SIGNATURE.to_vec();
        }
        header.clone().into_raw().and_then(|raw_header| {
            raw_header.write_to(&mut write)
        })?;
        for vlr in header.vlrs().iter() {
            (*vlr).clone().into_raw(false).and_then(|raw_vlr| {
                raw_vlr.write_to(&mut write)
            })?;
        }
        if !header.vlr_padding.is_empty() {
            write.write_all(&header.vlr_padding)?;
        }
        Ok(Writer {
            closed: false,
            header: header,
            write: write,
        })
    }

    /// Writes a point.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::io::Cursor;
    /// use las::Writer;
    ///
    /// let mut writer = Writer::new(Cursor::new(Vec::new()), Default::default()).unwrap();
    /// writer.write(Default::default()).unwrap();
    /// ```
    pub fn write(&mut self, point: Point) -> Result<()> {
        if self.closed {
            return Err(Error::Closed.into());
        }
        if !point.matches(self.header.point_format) {
            return Err(
                Error::PointAttributes(self.header.point_format, point).into(),
            );
        }
        self.header.number_of_points += 1;
        {
            let entry = self.header
                .number_of_points_by_return
                .entry(point.return_number)
                .or_insert(0);
            *entry += 1;
        }
        self.header.bounds.grow(&point);
        point.into_raw(self.header.transforms).and_then(
            |raw_point| {
                raw_point.write_to(&mut self.write, self.header.point_format)
            },
        )?;
        Ok(())
    }

    /// Close this writer.
    ///
    /// A second call to close is a no-op.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::io::Cursor;
    /// use las::Writer;
    /// let mut writer = Writer::new(Cursor::new(Vec::new()), Default::default()).unwrap();
    /// writer.close().unwrap();
    /// assert!(writer.close().is_err());
    /// ```
    pub fn close(&mut self) -> Result<()> {
        if self.closed {
            return Err(Error::Closed.into());
        }
        for raw_evlr in self.header.evlrs().into_iter().map(|evlr| {
            evlr.clone().into_raw(true)
        })
        {
            raw_evlr?.write_to(&mut self.write)?;
        }
        self.write.seek(SeekFrom::Start(0))?;
        self.header.clone().into_raw().and_then(|raw_header| {
            raw_header.write_to(&mut self.write)
        })?;
        self.closed = true;
        Ok(())
    }
}

impl Writer<BufWriter<File>> {
    /// Creates a new writer for a path.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::Writer;
    /// let writer = Writer::from_path("/dev/null", Default::default());
    /// ```
    pub fn from_path<P: AsRef<Path>>(path: P, header: Header) -> Result<Writer<BufWriter<File>>> {
        File::create(path).map_err(::Error::from).and_then(|file| {
            Writer::new(BufWriter::new(file), header)
        })
    }
}

impl<W: Seek + Write> Drop for Writer<W> {
    fn drop(&mut self) {
        if !self.closed {
            self.close().expect("Error when dropping the writer");
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use {Header, Version};
    use byteorder::{LittleEndian, ReadBytesExt};
    use point::Format;
    use std::io::Cursor;

    fn writer(format: Format, version: Version) -> Writer<Cursor<Vec<u8>>> {
        let mut header = Header::default();
        header.point_format = format;
        header.version = version;
        Writer::new(Cursor::new(Vec::new()), header).unwrap()
    }

    #[test]
    fn las_1_0_point_data_start_signature() {
        let mut cursor = Cursor::new(Vec::new());
        {
            let mut header = Header::default();
            header.version = (1, 0).into();
            header.push_vlr(Default::default());
            let mut writer = Writer::new(&mut cursor, header).unwrap();
            writer.write(Default::default()).unwrap();
        }
        cursor.set_position(281);
        assert_eq!(0xCCDD, cursor.read_u16::<LittleEndian>().unwrap());
    }

    #[test]
    fn already_closed() {
        let mut writer = Writer::new(Cursor::new(Vec::new()), Default::default()).unwrap();
        writer.close().unwrap();
        assert!(writer.close().is_err());
        assert!(writer.write(Default::default()).is_err());
    }

    #[test]
    fn missing_extra_bytes() {
        let format = Format {
            extra_bytes: 1,
            ..Default::default()
        };
        let mut writer = writer(format, Version::new(1, 4));
        assert!(writer.write(Default::default()).is_err());
    }

    #[test]
    fn missing_gps_time() {
        let format = Format::new(1).unwrap();
        let mut writer = writer(format, Version::new(1, 2));
        assert!(writer.write(Default::default()).is_err());
    }

    #[test]
    fn missing_color() {
        let format = Format::new(2).unwrap();
        let mut writer = writer(format, Version::new(1, 2));
        assert!(writer.write(Default::default()).is_err());
    }

    #[test]
    fn missing_nir() {
        let format = Format::new(8).unwrap();
        let mut writer = writer(format, Version::new(1, 4));
        let point = Point {
            gps_time: Some(0.),
            color: Some(Default::default()),
            ..Default::default()
        };
        assert!(writer.write(point).is_err());
    }

    #[test]
    fn missing_waveform() {
        let format = Format::new(4).unwrap();
        let mut writer = writer(format, Version::new(1, 4));
        assert!(writer.write(Default::default()).is_err());
    }
}
<|endoftext|>


<commit_before>use common::types::Ty;

#[derive(Clone)]
pub struct Expr
{
	out_ty: Ty,
	spec: ExprSpec
}

impl Expr
{
	pub fn spec(&self) -> &ExprSpec
	{
		&self.spec
	}
}

/// Expression Specific Element
#[derive(Clone)]
pub enum ExprSpec {
  // Unary Expressions
  Not(Box<Expr>),
  IsNull(Box<Expr>, bool), // bool - 'is null' if true. 'is not null' if false.
  Sign(Box<Expr>, bool), // true - Plus, false - Minus
  Cast(Box<Expr>, Box<Ty>, Box<Ty>),

  // Binary Arithmetic Expressions
  And(Box<Expr>,Box<Expr>),
  Or(Box<Expr>,Box<Expr>),
  Comp(CompOp, Box<Expr>, Box<Expr>),
  Arithm(ArithmOp, Box<Expr>, Box<Expr>),

  // Functions
  Fn(Box<FnDecl>, Vec<Box<Expr>>),

  // Other predicates
  Between(Box<Expr>,Box<Expr>,Box<Expr>), // predicand, begin, end

  // condition
  Switch(Vec<Box<Expr>>, Box<Expr>), // multiple conditions, else return value
  Case  (Box<Expr>, Box<Expr>),      // condition, return value

  // values
  Field(String),
  Const(Literal)
}

/// Comparison Operator Type
#[derive(Clone, Copy)]
pub enum CompOp {
  Eq,
  Ne,
  Lt,
  Le,
  Gt,
  Ge
}

/// Arithmetic Operator Type
#[derive(Clone, Copy)]
pub enum ArithmOp {
  Plus,
  Sub,
  Mul,
  Div,
  Rem,
}

/// Function Declaration
#[derive(Clone)]
pub struct FnDecl {
  signature: String
}

/// Representation for a single value
#[derive(Clone)]
pub enum Literal
{
  Bool(bool),
  Int1(i8),
  Int2(i16),
  Int4(i32),
  Int8(i64),
  Float4(f32),
  Float8(f64),
  Time(i64),
  Date(i32),
  Timestamp(i64),
  Interval(i64, i32),
  Char(String),
  Text(String),
  Varchar(String),
  Blob(Vec<u8>)
}

pub mod visitor {
	use common::types::Ty;
	use super::*;

	/// Visitor for Expr Tree
	pub trait Visitor<'v>: Sized {
	  fn visit_not(&mut self, child: &'v Expr) {
	    walk_expr(self, child);
	  }
	  fn visit_is_null(&mut self, child: &'v Expr, not: bool) {
	    walk_expr(self, child);
	  }
	  fn visit_sign(&mut self, child: &'v Expr, not: bool) {
	    walk_expr(self, child);
	  }
	  fn visit_cast(&mut self, expr: &'v Expr, from: &'v Ty, to: &'v Ty) {
	    walk_expr(self, expr);
	  }

	  fn visit_and(&mut self, lhs: &'v Expr, rhs: &'v Expr) {
	    walk_expr(self, lhs);
	    walk_expr(self, rhs);
	  }
	  fn visit_or(&mut self, lhs: &'v Expr, rhs: &'v Expr) {
	    walk_expr(self, lhs);
	    walk_expr(self, rhs);
	  }
	  fn visit_comp(&mut self, op: &CompOp, lhs: &'v Expr, rhs: &'v Expr) {
	    walk_expr(self, lhs);
	    walk_expr(self, rhs);
	  }
	  fn visit_arithm(&mut self, op: &ArithmOp, lhs: &'v Expr, rhs: &'v Expr) {
	    walk_expr(self, lhs);
	    walk_expr(self, rhs);
	  }
	  fn visit_concat(&mut self, lhs: &'v Expr, rhs: &'v Expr) {
	    walk_expr(self, lhs);
	    walk_expr(self, rhs);
	  }

	  fn visit_fn(&mut self, fd: &'v FnDecl, args: &'v Vec<Box<Expr>>) {
	    for ref arg in args {
	      walk_expr(self, arg);
	    }
	  }

	  fn visit_between(&mut self, predicand: &'v Expr, begin: &'v Expr, end: &'v Expr) {
	    walk_expr(self, predicand);
	    walk_expr(self, begin);
	    walk_expr(self, end);
	  }

	  fn visit_switch(&mut self, cases: &'v Vec<Box<Expr>>, default_value: &'v Expr) {
	    for ref each_case in cases {
	      walk_expr(self, each_case);
	    }
	    walk_expr(self, default_value);
	  }
	  fn visit_case(&mut self, condition: &'v Expr, result: &'v Expr) {
	    walk_expr(self, condition);
	    walk_expr(self, result);
	  }

	  fn visit_field(&mut self, field: &'v String) {}
	  fn visit_const(&mut self, literal: &'v Literal) {}
	}

	/// Walker for Expr Tree
	pub fn walk_expr<'v, V: Visitor<'v>>(visitor: &mut V, expr: &'v Expr) {

	  match *expr.spec() {

	    ExprSpec::Not(ref child) => visitor.visit_not(child),

	    ExprSpec::IsNull(ref child, positive) => {
	      visitor.visit_is_null(child, positive)
	    }

	    ExprSpec::Sign(ref child, positive) => {
	      visitor.visit_sign(child, positive)
	    }

	    ExprSpec::Cast(ref value, ref from, ref to) => {
	      visitor.visit_cast(value, from, to)
	    }

	    ExprSpec::And(ref lhs, ref rhs) => visitor.visit_and(lhs, rhs),

	    ExprSpec::Or(ref lhs, ref rhs) => visitor.visit_or(lhs, rhs),

	    ExprSpec::Comp(ref op, ref lhs, ref rhs) => {
	      visitor.visit_comp(op, lhs, rhs);
	    }

	    ExprSpec::Arithm(ref op, ref lhs, ref rhs) => {
	      visitor.visit_arithm(op, lhs, rhs);
	    },

	    ExprSpec::Fn(ref fd, ref args) => visitor.visit_fn(fd, args),

	    ExprSpec::Between(ref predicand, ref begin, ref end) => {
	      visitor.visit_between(predicand, begin, end)
	    },

	    ExprSpec::Switch(ref cases, ref default_value) => {
	      visitor.visit_switch(cases, default_value)
	    }
	    ExprSpec::Case(ref condition, ref result) => {
	      visitor.visit_case(condition, result)
	    },

	    ExprSpec::Field(ref name) => visitor.visit_field(name),

	    ExprSpec::Const(ref datum) => visitor.visit_const(datum),
	  }
	}

	pub trait SimpleVisitor {
	  fn accept(&self, op: &Expr) {
	    self.accept_by_default(op);
	  }

	  fn accept_by_default(&self, e: &Expr) {
	    match *e.spec() {
	      ExprSpec::Not      (ref c)               => self.accept(c),
	      ExprSpec::IsNull   (ref c, _)            => self.accept(c),
	      ExprSpec::Sign     (ref c, _)            => self.accept(c),
	      ExprSpec::Cast     (ref c, _, _)         => self.accept(c),

	      ExprSpec::Const    (_)                   => {},
	      ExprSpec::And      (ref l, ref r)        => { self.accept(l); self.accept(r) },
	      _ => unimplemented!()
	    }
	  }
	}

}<commit_msg>Implemented SimpleVisitor for Expr.<commit_after>use common::types::Ty;

#[derive(Clone)]
pub struct Expr
{
	out_ty: Ty,
	spec: ExprSpec
}

impl Expr
{
	pub fn spec(&self) -> &ExprSpec
	{
		&self.spec
	}
}

/// Expression Specific Element
#[derive(Clone)]
pub enum ExprSpec {
  // Unary Expressions
  Not(Box<Expr>),
  IsNull(Box<Expr>, bool), // bool - 'is null' if true. 'is not null' if false.
  Sign(Box<Expr>, bool), // true - Plus, false - Minus
  Cast(Box<Expr>, Box<Ty>, Box<Ty>),

  // Binary Arithmetic Expressions
  And(Box<Expr>,Box<Expr>),
  Or(Box<Expr>,Box<Expr>),
  Comp(CompOp, Box<Expr>, Box<Expr>),
  Arithm(ArithmOp, Box<Expr>, Box<Expr>),

  // Functions
  Fn(Box<FnDecl>, Vec<Box<Expr>>),

  // Other predicates
  Between(Box<Expr>,Box<Expr>,Box<Expr>), // predicand, begin, end

  // condition
  Switch(Vec<Box<Expr>>, Box<Expr>), // multiple conditions, else return value
  Case  (Box<Expr>, Box<Expr>),      // condition, return value

  // values
  Field(String),
  Const(Literal)
}

/// Comparison Operator Type
#[derive(Clone, Copy)]
pub enum CompOp {
  Eq,
  Ne,
  Lt,
  Le,
  Gt,
  Ge
}

/// Arithmetic Operator Type
#[derive(Clone, Copy)]
pub enum ArithmOp {
  Plus,
  Sub,
  Mul,
  Div,
  Rem,
}

/// Function Declaration
#[derive(Clone)]
pub struct FnDecl {
  signature: String
}

/// Representation for a single value
#[derive(Clone)]
pub enum Literal
{
  Bool(bool),
  Int1(i8),
  Int2(i16),
  Int4(i32),
  Int8(i64),
  Float4(f32),
  Float8(f64),
  Time(i64),
  Date(i32),
  Timestamp(i64),
  Interval(i64, i32),
  Char(String),
  Text(String),
  Varchar(String),
  Blob(Vec<u8>)
}

pub mod visitor {
	use common::types::Ty;
	use super::*;

	/// Visitor for Expr Tree
	pub trait Visitor<'v>: Sized {
	  fn visit_not(&mut self, child: &'v Expr) {
	    walk_expr(self, child);
	  }
	  fn visit_is_null(&mut self, child: &'v Expr, not: bool) {
	    walk_expr(self, child);
	  }
	  fn visit_sign(&mut self, child: &'v Expr, not: bool) {
	    walk_expr(self, child);
	  }
	  fn visit_cast(&mut self, expr: &'v Expr, from: &'v Ty, to: &'v Ty) {
	    walk_expr(self, expr);
	  }

	  fn visit_and(&mut self, lhs: &'v Expr, rhs: &'v Expr) {
	    walk_expr(self, lhs);
	    walk_expr(self, rhs);
	  }
	  fn visit_or(&mut self, lhs: &'v Expr, rhs: &'v Expr) {
	    walk_expr(self, lhs);
	    walk_expr(self, rhs);
	  }
	  fn visit_comp(&mut self, op: &CompOp, lhs: &'v Expr, rhs: &'v Expr) {
	    walk_expr(self, lhs);
	    walk_expr(self, rhs);
	  }
	  fn visit_arithm(&mut self, op: &ArithmOp, lhs: &'v Expr, rhs: &'v Expr) {
	    walk_expr(self, lhs);
	    walk_expr(self, rhs);
	  }
	  fn visit_concat(&mut self, lhs: &'v Expr, rhs: &'v Expr) {
	    walk_expr(self, lhs);
	    walk_expr(self, rhs);
	  }

	  fn visit_fn(&mut self, fd: &'v FnDecl, args: &'v Vec<Box<Expr>>) {
	    for ref arg in args {
	      walk_expr(self, arg);
	    }
	  }

	  fn visit_between(&mut self, predicand: &'v Expr, begin: &'v Expr, end: &'v Expr) {
	    walk_expr(self, predicand);
	    walk_expr(self, begin);
	    walk_expr(self, end);
	  }

	  fn visit_switch(&mut self, cases: &'v Vec<Box<Expr>>, default_value: &'v Expr) {
	    for ref each_case in cases {
	      walk_expr(self, each_case);
	    }
	    walk_expr(self, default_value);
	  }
	  fn visit_case(&mut self, condition: &'v Expr, result: &'v Expr) {
	    walk_expr(self, condition);
	    walk_expr(self, result);
	  }

	  fn visit_field(&mut self, field: &'v String) {}
	  fn visit_const(&mut self, literal: &'v Literal) {}
	}

	/// Walker for Expr Tree
	pub fn walk_expr<'v, V: Visitor<'v>>(visitor: &mut V, expr: &'v Expr) {

	  match *expr.spec() {

	    ExprSpec::Not(ref child) => visitor.visit_not(child),

	    ExprSpec::IsNull(ref child, positive) => {
	      visitor.visit_is_null(child, positive)
	    }

	    ExprSpec::Sign(ref child, positive) => {
	      visitor.visit_sign(child, positive)
	    }

	    ExprSpec::Cast(ref value, ref from, ref to) => {
	      visitor.visit_cast(value, from, to)
	    }

	    ExprSpec::And(ref lhs, ref rhs) => visitor.visit_and(lhs, rhs),

	    ExprSpec::Or(ref lhs, ref rhs) => visitor.visit_or(lhs, rhs),

	    ExprSpec::Comp(ref op, ref lhs, ref rhs) => {
	      visitor.visit_comp(op, lhs, rhs);
	    }

	    ExprSpec::Arithm(ref op, ref lhs, ref rhs) => {
	      visitor.visit_arithm(op, lhs, rhs);
	    },

	    ExprSpec::Fn(ref fd, ref args) => visitor.visit_fn(fd, args),

	    ExprSpec::Between(ref predicand, ref begin, ref end) => {
	      visitor.visit_between(predicand, begin, end)
	    },

	    ExprSpec::Switch(ref cases, ref default_value) => {
	      visitor.visit_switch(cases, default_value)
	    }
	    ExprSpec::Case(ref condition, ref result) => {
	      visitor.visit_case(condition, result)
	    },

	    ExprSpec::Field(ref name) => visitor.visit_field(name),

	    ExprSpec::Const(ref datum) => visitor.visit_const(datum),
	  }
	}

	pub trait SimpleVisitor {
	  fn accept(&self, op: &Expr) {
	    self.accept_by_default(op);
	  }

	  fn accept_by_default(&self, e: &Expr) {
	    match *e.spec() {
	      ExprSpec::Not      (ref c)               => self.accept(c),
	      ExprSpec::IsNull   (ref c, _)            => self.accept(c),
	      ExprSpec::Sign     (ref c, _)            => self.accept(c),
	      ExprSpec::Cast     (ref c, _, _)         => self.accept(c),

	      ExprSpec::And      (ref l, ref r)        => { self.accept(l); self.accept(r) },
	      ExprSpec::Or       (ref l, ref r)        => { self.accept(l); self.accept(r) },
	      ExprSpec::Comp     (_, ref l, ref r)     => { self.accept(l); self.accept(r) },
	      ExprSpec::Arithm   (_, ref l, ref r)     => { self.accept(l); self.accept(r) },

	      ExprSpec::Fn       (_, ref args)         => { for e in args.iter() { self.accept(e) } },

	      ExprSpec::Between  (ref p, ref b, ref e)     => {
	      	self.accept(p); self.accept(b); self.accept(e);
	     	},

	      ExprSpec::Switch   (ref cases, ref default)  => {
	      	for c in cases.iter() {
	      		self.accept(c);
	      	}

	      	self.accept(default);
	      },
	      ExprSpec::Case      (ref l, ref r)        => { self.accept(l); self.accept(r) },

				ExprSpec::Field     (_) 	                => {},
	      ExprSpec::Const     (_)                   => {}
	    }
	  }
	}

}<|endoftext|>


<commit_before>//! Parse the PE SafeSEH table for references to valid exception handler
//! functions.
//!
//! This table referenced by the Load Config directory.
//! The reference consists of: (offset: VA, count: u32/u64).
//! The table is simply an array of RVAs to function start addresses.
//!
//! Therefore, we can follow the table reference and walk the array, pulling out
//! functions.
//!
//! Bail if the table reference, table, or any of its entries don't make sense.
//!
//! references:
//!   - https://docs.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-image_load_config_directory32
//!   - https://docs.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-image_load_config_directory64

use anyhow::Result;
use log::debug;

use crate::{
    aspace::AddressSpace,
    loader::pe::PE,
    module::{Arch, Permissions},
    RVA, VA,
};

pub fn find_pe_safeseh_handlers(pe: &PE) -> Result<Vec<VA>> {
    let mut ret = vec![];

    let load_config_directory_rva: RVA = {
        let opt_header = match pe.header.optional_header {
            Some(opt_header) => opt_header,
            _ => return Ok(ret),
        };

        let load_config_directory = match opt_header.data_directories.get_load_config_table() {
            Some(load_config_directory) => load_config_directory,
            _ => return Ok(ret),
        };

        load_config_directory.virtual_address as RVA
    };

    debug!(
        "load config directory: {:#x}",
        pe.module.address_space.base_address + load_config_directory_rva
    );

    // offsets defined here:
    // https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory

    // according to IDA, the first DWORD is `Size` not `Characteristics` (unused).
    let size = pe.module.address_space.relative.read_u32(load_config_directory_rva)?;

    // max offset into the config directory that we'll read.
    let max_config_directory_offset = match pe.module.arch {
        // CFG flags
        Arch::X32 => 0x48,
        Arch::X64 => 0x70,
    };

    if max_config_directory_offset > size {
        debug!("no SafeSEH table: load config directory too small");
        return Ok(ret);
    }

    let sehandler_table_va: VA = match pe.module.arch {
        Arch::X32 => pe.module.read_va_at_rva(load_config_directory_rva + 0x40)?,
        Arch::X64 => pe.module.read_va_at_rva(load_config_directory_rva + 0x60)?,
    };
    if sehandler_table_va == 0 {
        debug!("SafeSEH table empty");
        return Ok(ret);
    };
    debug!("SafeSEH table: {:#x}", sehandler_table_va);

    let sehandler_table_count = match pe.module.arch {
        Arch::X32 => pe
            .module
            .address_space
            .relative
            .read_u32(load_config_directory_rva + 0x44)? as u64,
        Arch::X64 => pe
            .module
            .address_space
            .relative
            .read_u64(load_config_directory_rva + 0x68)? as u64,
    };
    debug!("SafeSEH table count: {:#x}", sehandler_table_count);

    let mut offset = sehandler_table_va;
    for _ in 0..sehandler_table_count {
        let target = pe.module.read_rva_at_va(offset)?;
        let target = target as u64 + pe.module.address_space.base_address;

        if pe.module.probe_va(target, Permissions::X) {
            ret.push(target);
        } else {
            debug!("unexpected non-executable SafeSEH target: {:#x}", target);
            break;
        }
        offset += pe.module.arch.pointer_size() as u64
    }

    Ok(ret)
}

#[cfg(test)]
mod tests {
    use crate::rsrc::*;
    use anyhow::Result;

    #[test]
    fn k32() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let fns = crate::analysis::pe::safeseh::find_pe_safeseh_handlers(&pe)?;
        assert_eq!(0, fns.len());

        Ok(())
    }

    #[test]
    fn tiny() -> Result<()> {
        let buf = get_buf(Rsrc::TINY);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let fns = crate::analysis::pe::safeseh::find_pe_safeseh_handlers(&pe)?;
        assert_eq!(0, fns.len());

        Ok(())
    }

    #[test]
    fn nop() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let fns = crate::analysis::pe::safeseh::find_pe_safeseh_handlers(&pe)?;
        assert_eq!(2, fns.len());

        Ok(())
    }

    #[test]
    fn mimi() -> Result<()> {
        let buf = get_buf(Rsrc::MIMI);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let fns = crate::analysis::pe::safeseh::find_pe_safeseh_handlers(&pe)?;
        assert_eq!(0, fns.len());

        Ok(())
    }
}
<commit_msg>pe: safeseh: reuse data directory accessor<commit_after>//! Parse the PE SafeSEH table for references to valid exception handler
//! functions.
//!
//! This table referenced by the Load Config directory.
//! The reference consists of: (offset: VA, count: u32/u64).
//! The table is simply an array of RVAs to function start addresses.
//!
//! Therefore, we can follow the table reference and walk the array, pulling out
//! functions.
//!
//! Bail if the table reference, table, or any of its entries don't make sense.
//!
//! references:
//!   - https://docs.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-image_load_config_directory32
//!   - https://docs.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-image_load_config_directory64

use anyhow::Result;
use log::debug;

use crate::{
    aspace::AddressSpace,
    loader::{pe, pe::PE},
    module::{Arch, Permissions},
    VA,
};

pub fn find_pe_safeseh_handlers(pe: &PE) -> Result<Vec<VA>> {
    let mut ret = vec![];

    if let Ok(Some(load_config_directory)) = pe.get_data_directory(pe::IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG) {
        debug!("load config directory: {:#x}", load_config_directory.address);

        // offsets defined here:
        // https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory

        // according to IDA, the first DWORD is `Size` not `Characteristics` (unused).
        let size = pe.module.address_space.read_u32(load_config_directory.address)?;

        // max offset into the config directory that we'll read.
        let max_config_directory_offset = match pe.module.arch {
            // CFG flags
            Arch::X32 => 0x48,
            Arch::X64 => 0x70,
        };

        if max_config_directory_offset > size {
            debug!("no SafeSEH table: load config directory too small");
            return Ok(ret);
        }

        let sehandler_table_va: VA = match pe.module.arch {
            Arch::X32 => pe.module.read_va_at_va(load_config_directory.address + 0x40)?,
            Arch::X64 => pe.module.read_va_at_va(load_config_directory.address + 0x60)?,
        };
        if sehandler_table_va == 0 {
            debug!("SafeSEH table empty");
            return Ok(ret);
        };
        debug!("SafeSEH table: {:#x}", sehandler_table_va);

        let sehandler_table_count = match pe.module.arch {
            Arch::X32 => pe.module.address_space.read_u32(load_config_directory.address + 0x44)? as u64,
            Arch::X64 => pe.module.address_space.read_u64(load_config_directory.address + 0x68)? as u64,
        };
        debug!("SafeSEH table count: {:#x}", sehandler_table_count);

        let mut offset = sehandler_table_va;
        for _ in 0..sehandler_table_count {
            let target = pe.module.read_rva_at_va(offset)?;
            let target = target as u64 + pe.module.address_space.base_address;

            if pe.module.probe_va(target, Permissions::X) {
                ret.push(target);
            } else {
                debug!("unexpected non-executable SafeSEH target: {:#x}", target);
                break;
            }
            offset += pe.module.arch.pointer_size() as u64
        }
    }

    Ok(ret)
}

#[cfg(test)]
mod tests {
    use crate::rsrc::*;
    use anyhow::Result;

    #[test]
    fn k32() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let fns = crate::analysis::pe::safeseh::find_pe_safeseh_handlers(&pe)?;
        assert_eq!(0, fns.len());

        Ok(())
    }

    #[test]
    fn tiny() -> Result<()> {
        let buf = get_buf(Rsrc::TINY);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let fns = crate::analysis::pe::safeseh::find_pe_safeseh_handlers(&pe)?;
        assert_eq!(0, fns.len());

        Ok(())
    }

    #[test]
    fn nop() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let fns = crate::analysis::pe::safeseh::find_pe_safeseh_handlers(&pe)?;
        assert_eq!(2, fns.len());

        Ok(())
    }

    #[test]
    fn mimi() -> Result<()> {
        let buf = get_buf(Rsrc::MIMI);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let fns = crate::analysis::pe::safeseh::find_pe_safeseh_handlers(&pe)?;
        assert_eq!(0, fns.len());

        Ok(())
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code, unused_macros)]

use utils::constants::WALLET_CREDENTIALS;

pub mod callback;

#[macro_use]
#[path = "../../src/utils/memzeroize.rs"]
pub mod zeroize;

#[path = "../../src/utils/environment.rs"]
pub mod environment;

pub mod pool;
pub mod crypto;
pub mod did;
pub mod wallet;
pub mod ledger;
pub mod anoncreds;
pub mod types;
pub mod pairwise;
pub mod constants;
pub mod blob_storage;
pub mod non_secrets;
pub mod results;
pub mod payments;
pub mod rand_utils;
pub mod logger;

#[macro_use]
#[allow(unused_macros)]
#[path = "../../src/utils/test.rs"]
pub mod test;

pub mod timeout;

#[path = "../../src/utils/sequence.rs"]
pub mod sequence;

#[macro_use]
#[allow(unused_macros)]
#[path = "../../src/utils/ctypes.rs"]
pub mod ctypes;

#[path = "../../src/utils/inmem_wallet.rs"]
pub mod inmem_wallet;

#[path = "../../src/domain/mod.rs"]
pub mod domain;

pub fn setup(name: &str) {
    test::cleanup_storage(name);
    logger::set_default_logger();
}

pub fn tear_down_delete_wallet(wallet_config: &str) {
    wallet::delete_wallet(wallet_config, WALLET_CREDENTIALS).unwrap();
}

pub fn tear_down(name: &str) {
    test::cleanup_storage(name);
}

pub fn setup_with_wallet(name: &str) -> (i32, String) {
    setup(name);
    wallet::create_and_open_default_wallet(name).unwrap()
}

pub fn setup_with_plugged_wallet(name: &str) -> (i32, String) {
    setup(name);
    wallet::create_and_open_plugged_wallet().unwrap()
}

pub fn tear_down_with_wallet(wallet_handle: i32, name: &str, wallet_config: &str) {
    wallet::close_wallet(wallet_handle).unwrap();
    wallet::delete_wallet(wallet_config, WALLET_CREDENTIALS).unwrap();
    tear_down(name);
}

pub fn setup_with_pool(name: &str) -> i32 {
    setup(name);
    pool::create_and_open_pool_ledger(name).unwrap()
}

pub fn tear_down_with_pool(pool_handle: i32, name: &str) {
    pool::close(pool_handle).unwrap();
    tear_down(name);
}

pub fn setup_with_wallet_and_pool(name: &str) -> (i32, i32, String) {
    let (wallet_handle, config) = setup_with_wallet(name);
    let pool_handle = pool::create_and_open_pool_ledger(name).unwrap();
    (wallet_handle, pool_handle, config)
}

pub fn tear_down_with_wallet_and_pool(wallet_handle: i32, pool_handle: i32, name: &str, wallet_config: &str) {
    pool::close(pool_handle).unwrap();
    tear_down_with_wallet(wallet_handle, name, wallet_config);
}

pub fn setup_trustee(name: &str) -> (i32, i32, String, String) {
    let (wallet_handle, pool_handle, config) = setup_with_wallet_and_pool(name);
    let (did, _) = did::create_and_store_my_did(wallet_handle, Some(constants::TRUSTEE_SEED)).unwrap();
    (wallet_handle, pool_handle, did, config)
}

pub fn setup_steward(name: &str) -> (i32, i32, String, String) {
    let (wallet_handle, pool_handle, config) = setup_with_wallet_and_pool(name);
    let (did, _) = did::create_and_store_my_did(wallet_handle, Some(constants::STEWARD_SEED)).unwrap();
    (wallet_handle, pool_handle, did, config)
}

pub fn setup_did(name: &str) -> (i32, String, String) {
    let (wallet_handle, config) = setup_with_wallet(name);
    let (did, _) = did::create_and_store_my_did(wallet_handle, None).unwrap();
    (wallet_handle, did, config)
}

pub fn setup_new_identity(name: &str) -> (i32, i32, String, String, String) {
    let (wallet_handle, pool_handle, trustee_did, config) = setup_trustee(name);

    let (my_did, my_vk) = did::create_and_store_my_did(wallet_handle, None).unwrap();
    let nym = ledger::build_nym_request(&trustee_did, &my_did, Some(&my_vk), None, Some("TRUSTEE")).unwrap();
    let response = ledger::sign_and_submit_request(pool_handle, wallet_handle, &trustee_did, &nym).unwrap();
    pool::check_response_type(&response, types::ResponseType::REPLY);

    (wallet_handle, pool_handle, my_did, my_vk, config)
}<commit_msg>use the correct credentials to delete the wallets<commit_after>#![allow(dead_code, unused_macros)]

use utils::constants::WALLET_CREDENTIALS;

pub mod callback;

#[macro_use]
#[path = "../../src/utils/memzeroize.rs"]
pub mod zeroize;

#[path = "../../src/utils/environment.rs"]
pub mod environment;

pub mod pool;
pub mod crypto;
pub mod did;
pub mod wallet;
pub mod ledger;
pub mod anoncreds;
pub mod types;
pub mod pairwise;
pub mod constants;
pub mod blob_storage;
pub mod non_secrets;
pub mod results;
pub mod payments;
pub mod rand_utils;
pub mod logger;

#[macro_use]
#[allow(unused_macros)]
#[path = "../../src/utils/test.rs"]
pub mod test;

pub mod timeout;

#[path = "../../src/utils/sequence.rs"]
pub mod sequence;

#[macro_use]
#[allow(unused_macros)]
#[path = "../../src/utils/ctypes.rs"]
pub mod ctypes;

#[path = "../../src/utils/inmem_wallet.rs"]
pub mod inmem_wallet;

#[path = "../../src/domain/mod.rs"]
pub mod domain;

pub fn setup(name: &str) {
    test::cleanup_storage(name);
    logger::set_default_logger();
}

pub fn tear_down_delete_wallet(wallet_config: &str) {
    wallet::delete_wallet(wallet_config, WALLET_CREDENTIALS).unwrap();
}

pub fn tear_down_delete_wallet_with_credentials(wallet_config: &str, wallet_credentials: &str) {
    wallet::delete_wallet(wallet_config, wallet_credentials).unwrap();
}

pub fn tear_down(name: &str) {
    test::cleanup_storage(name);
}

pub fn setup_with_wallet(name: &str) -> (i32, String) {
    setup(name);
    wallet::create_and_open_default_wallet(name).unwrap()
}

pub fn setup_with_plugged_wallet(name: &str) -> (i32, String) {
    setup(name);
    wallet::create_and_open_plugged_wallet().unwrap()
}

pub fn tear_down_with_wallet(wallet_handle: i32, name: &str, wallet_config: &str) {
    wallet::close_wallet(wallet_handle).unwrap();
    wallet::delete_wallet(wallet_config, WALLET_CREDENTIALS).unwrap();
    tear_down(name);
}

pub fn setup_with_pool(name: &str) -> i32 {
    setup(name);
    pool::create_and_open_pool_ledger(name).unwrap()
}

pub fn tear_down_with_pool(pool_handle: i32, name: &str) {
    pool::close(pool_handle).unwrap();
    tear_down(name);
}

pub fn setup_with_wallet_and_pool(name: &str) -> (i32, i32, String) {
    let (wallet_handle, config) = setup_with_wallet(name);
    let pool_handle = pool::create_and_open_pool_ledger(name).unwrap();
    (wallet_handle, pool_handle, config)
}

pub fn tear_down_with_wallet_and_pool(wallet_handle: i32, pool_handle: i32, name: &str, wallet_config: &str) {
    pool::close(pool_handle).unwrap();
    tear_down_with_wallet(wallet_handle, name, wallet_config);
}

pub fn setup_trustee(name: &str) -> (i32, i32, String, String) {
    let (wallet_handle, pool_handle, config) = setup_with_wallet_and_pool(name);
    let (did, _) = did::create_and_store_my_did(wallet_handle, Some(constants::TRUSTEE_SEED)).unwrap();
    (wallet_handle, pool_handle, did, config)
}

pub fn setup_steward(name: &str) -> (i32, i32, String, String) {
    let (wallet_handle, pool_handle, config) = setup_with_wallet_and_pool(name);
    let (did, _) = did::create_and_store_my_did(wallet_handle, Some(constants::STEWARD_SEED)).unwrap();
    (wallet_handle, pool_handle, did, config)
}

pub fn setup_did(name: &str) -> (i32, String, String) {
    let (wallet_handle, config) = setup_with_wallet(name);
    let (did, _) = did::create_and_store_my_did(wallet_handle, None).unwrap();
    (wallet_handle, did, config)
}

pub fn setup_new_identity(name: &str) -> (i32, i32, String, String, String) {
    let (wallet_handle, pool_handle, trustee_did, config) = setup_trustee(name);

    let (my_did, my_vk) = did::create_and_store_my_did(wallet_handle, None).unwrap();
    let nym = ledger::build_nym_request(&trustee_did, &my_did, Some(&my_vk), None, Some("TRUSTEE")).unwrap();
    let response = ledger::sign_and_submit_request(pool_handle, wallet_handle, &trustee_did, &nym).unwrap();
    pool::check_response_type(&response, types::ResponseType::REPLY);

    (wallet_handle, pool_handle, my_did, my_vk, config)
}<|endoftext|>


<commit_before>use std::sync::Mutex;

use fall_tree::{File, AstNode};
use fall_tree::visitor::{Visitor, NodeVisitor};

use {FallFile, RefExpr, CallExpr};
use editor_api::Diagnostic;

mod diagnostics;
mod db;
mod query;

use self::diagnostics::DiagnosticSink;
pub use self::query::{CallKind, RefKind};

pub struct Analysis<'f> {
    db: db::DB<'f>,
    file: FallFile<'f>,

    diagnostics: Mutex<Vec<Diagnostic>>,
}

impl<'f> Analysis<'f> {
    pub fn new(file: FallFile) -> Analysis {
        Analysis {
            db: db::DB::new(file),
            file,
            diagnostics: Default::default(),
        }
    }

    pub fn file(&self) -> FallFile<'f> {
        self.file
    }

    pub fn resolve_call(&self, call: CallExpr<'f>) -> Option<CallKind<'f>> {
        self.db.get(query::ResolveCall(call))
    }

    pub fn resolve_reference(&self, ref_: RefExpr<'f>) -> Option<RefKind<'f>> {
        self.db.get(query::ResolveRefExpr(ref_))
    }

    pub fn collect_all_diagnostics(&self) -> Vec<Diagnostic> {
        let mut result = Visitor(Vec::new())
            .visit::<RefExpr, _>(|_, ref_| { self.resolve_reference(ref_); })
            .visit::<CallExpr, _>(|_, call| { self.resolve_call(call); })
            .walk_recursively_children_first(self.file().node());

        self.db.get(query::UnusedRules);

        result.extend(self.diagnostics());
        result.sort_by_key(|d| (d.range.start(), d.range.end()));
        result
    }

    #[allow(unused)] // for debugging
    pub fn debug_diagnostics(&self) -> String {
        diagnostics::debug_diagnostics(&self.diagnostics(), self.file().node().text())
    }

    fn diagnostics(&self) -> Vec<Diagnostic> {
        let mut result = self.diagnostics.lock().unwrap().clone();
        result.extend(self.db.diagnostics.lock().unwrap().clone());
        result
    }
}

pub struct FileWithAnalysis {
    rent: rent::R
}

impl FileWithAnalysis {
    pub fn new(text: String) -> FileWithAnalysis {
        FileWithAnalysis::from_file(::parse(text))
    }

    fn from_file(file: File) -> FileWithAnalysis {
        FileWithAnalysis {
            rent: rent::R::new(Box::new(file), |file| {
                Analysis::new(FallFile::new(file.root()))
            })
        }
    }

    pub fn file(&self) -> &File {
        self.rent.head()
    }

    pub fn analyse<T, F: FnOnce(&Analysis) -> T>(&self, f: F) -> T {
        self.rent.rent(|a: &Analysis| f(a))
    }
}

rental! {
    mod rent {
        use super::*;

        #[rental]
        pub struct R {
            file: Box<File>,
            analysis: Analysis<'file>
        }
    }
}

#[test]
fn analysis_is_sync() {
    fn assert_is_sync<T: Sync>() {}
    assert_is_sync::<Analysis>();
}

#[cfg(test)]
fn check_diagnostics(code: &str, expected_diagnostics: &str) {
    use fall_tree::test_util::report_diff;
    let file = ::editor_api::analyse(code.to_string());

    file.analyse(|a| {
        let d = a.collect_all_diagnostics();
        let actual = d.into_iter().map(|d| {
            let s = match d.severity {
                Severity::Error => 'E',
                Severity::Warning => 'W',
            };
            format!("{}: {} {}", a.file().node().text().slice(d.range), s, d.message)
        }).collect::<Vec<_>>().join("\n");

        report_diff(expected_diagnostics, &actual);
    })
}

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

    use fall_tree::search;
    use fall_tree::search::ast;

    use ::test_util::parse_with_caret;

    #[test]
    fn test_diagnostics() {
        check_diagnostics(r"
       pub rule foo { <eof x> }
       rule bar { foo <abracadabra>}
       rule baz { <prev_is foo> <prev_is bar> <prev_is {foo}>}
       rule dupe { dupe }
       rule dupe { dupe }
    ", "\
<eof x>: E Wrong number of arguments, expected 0, got 1
x: E Unresolved reference
abracadabra: E Unresolved reference
baz: W Unused rule
<prev_is bar>: E <prev_is> arguments must be public rules
<prev_is {foo}>: E <prev_is> arguments must be public rules
dupe: W Unused rule
dupe: E Duplicate rule");
    }

    #[test]
    fn test_resolve_no_args() {
        check_resolved("rule foo { <^eof> }", "Eof");
        check_resolved("rule foo { <^any> }", "Any");
        check_resolved("rule foo { <^commit> }", "Commit");
    }

    #[test]
    fn test_resolve_simple_args() {
        check_resolved("rule foo { <^not a> }", "Not(RefExpr@[16; 17))");
        check_resolved("rule foo { <^with_skip a b> }", "WithSkip(RefExpr@[22; 23), RefExpr@[24; 25))");
        check_resolved("rule foo { <^layer a b> }", "Layer(RefExpr@[18; 19), RefExpr@[20; 21))");
    }

    #[test]
    fn test_resolve_layer() {
        check_resolved("rule foo { <^enter 'ctx' a> <is_in 'ctx' >}", "Enter(0, RefExpr@[24; 25))");
        check_resolved("rule foo { <^exit  'ctx' a> <is_in 'ctx' >}", "Exit(0, RefExpr@[24; 25))");

        check_resolved("rule foo { <is_in 'ctx1'> <^enter 'ctx2' a> <is_in 'ctx1'> <enter 'ctx2' a> }", "Enter(1, RefExpr@[40; 41))");
        check_resolved("rule foo { <is_in 'ctx1'> <^exit  'ctx2' a> <is_in 'ctx1'> <exit  'ctx2' a> }", "Exit(1, RefExpr@[40; 41))");

        check_resolved(
            "rule foo { <enter 'ctx1' a> <enter 'ctx2' b> <^is_in 'ctx2'> <is_in 'ctx1'> }",
            "IsIn(1)"
        );

        check(
            r" rule foo { <^is_in foo>}",
            None,
            r#"[(foo, "Context should be a single quoted string")]"#,
        )
    }

    #[test]
    fn test_not() {
        resolve_call("rule foo { <^not bar>}", |c, _| {
            let c = c.unwrap();
            assert_eq!(format!("{:?}", c), "Not(RefExpr@[16; 19))");
        })
    }

    #[test]
    fn test_unresolved_call() {
        check(
            r"rule foo { <^abracadabra>}",
            None,
            r#"[(abracadabra, "Unresolved reference")]"#
        );
    }

    #[test]
    fn wrong_arity() {
        check(
            "rule foo { <^eof foo> }",
            Some(CallKind::Eof),
            r#"[(<eof foo>, "Wrong number of arguments, expected 0, got 1")]"#
        );

        check(
            "rule foo { <^any foo bar> }",
            Some(CallKind::Any),
            r#"[(<any foo bar>, "Wrong number of arguments, expected 0, got 2")]"#
        );

        check(
            "rule foo { <^commit foo bar baz> }",
            Some(CallKind::Commit),
            r#"[(<commit foo bar baz>, "Wrong number of arguments, expected 0, got 3")]"#
        );

        check(
            "rule foo { <^not> }",
            None,
            r#"[(<not>, "Wrong number of arguments, expected 1, got 0")]"#
        );
    }

    #[test]
    fn check_prev_is() {
        check_resolved(
            "pub rule foo {} pub rule bar {} rule baz { <^prev_is foo bar> }",
            "PrevIs([1, 2])"
        )
    }

    #[test]
    fn call_custom_rule() {
        check_resolved(
            "rule foo { <^bar a b>} rule bar(x, y) { }",
            "RuleCall(SynRule@[22; 40), [(0, RefExpr@[16; 17)), (1, RefExpr@[18; 19))])"
        );
        ::analysis::check_diagnostics(
            "rule foo { <bar <eof>>} rule bar(x, y) { }",
            "<bar <eof>>: E Expected 2 arguments, got 1"
        )
    }

    fn check(text: &str, kind: Option<CallKind>, diagnostics: &str) {
        resolve_call(text, |c, a| {
            assert_eq!(c, kind);
            assert_eq!(a.debug_diagnostics(), diagnostics);
        })
    }

    fn check_resolved(text: &str, kind: &str) {
        resolve_call(text, |c, a| {
            let c = c.unwrap();
            assert_eq!(format!("{:?}", c), kind);
            assert_eq!(a.debug_diagnostics(), "[]");
        })
    }

    fn resolve_call<F: FnOnce(Option<CallKind>, Analysis)>(text: &str, f: F) {
        let (file, caret) = parse_with_caret(&text);
        let call = {
            let leaf = search::find_leaf_at_offset(file.root(), caret)
                .left_biased().unwrap();
            ast::ancestor_exn(leaf)
        };
        let a = Analysis::new(::ast(&file));
        let c = a.resolve_call(call);
        f(c, a)
    }
}

<commit_msg>Kill old diagnostics<commit_after>use fall_tree::{File, AstNode};
use fall_tree::visitor::{Visitor, NodeVisitor};

use {FallFile, RefExpr, CallExpr};
use editor_api::Diagnostic;

mod diagnostics;
mod db;
mod query;

use self::diagnostics::DiagnosticSink;
pub use self::query::{CallKind, RefKind};

pub struct Analysis<'f> {
    db: db::DB<'f>,
    file: FallFile<'f>,
}

impl<'f> Analysis<'f> {
    pub fn new(file: FallFile) -> Analysis {
        Analysis {
            db: db::DB::new(file),
            file,
        }
    }

    pub fn file(&self) -> FallFile<'f> {
        self.file
    }

    pub fn resolve_call(&self, call: CallExpr<'f>) -> Option<CallKind<'f>> {
        self.db.get(query::ResolveCall(call))
    }

    pub fn resolve_reference(&self, ref_: RefExpr<'f>) -> Option<RefKind<'f>> {
        self.db.get(query::ResolveRefExpr(ref_))
    }

    pub fn collect_all_diagnostics(&self) -> Vec<Diagnostic> {
        Visitor(())
            .visit::<RefExpr, _>(|_, ref_| { self.resolve_reference(ref_); })
            .visit::<CallExpr, _>(|_, call| { self.resolve_call(call); })
            .walk_recursively_children_first(self.file().node());
        self.db.get(query::UnusedRules);
        self.db.diagnostics.lock().unwrap().clone()
    }

    #[allow(unused)] // for debugging
    pub fn debug_diagnostics(&self) -> String {
        diagnostics::debug_diagnostics(&self.db.diagnostics.lock().unwrap(), self.file().node().text())
    }
}

pub struct FileWithAnalysis {
    rent: rent::R
}

impl FileWithAnalysis {
    pub fn new(text: String) -> FileWithAnalysis {
        FileWithAnalysis::from_file(::parse(text))
    }

    fn from_file(file: File) -> FileWithAnalysis {
        FileWithAnalysis {
            rent: rent::R::new(Box::new(file), |file| {
                Analysis::new(FallFile::new(file.root()))
            })
        }
    }

    pub fn file(&self) -> &File {
        self.rent.head()
    }

    pub fn analyse<T, F: FnOnce(&Analysis) -> T>(&self, f: F) -> T {
        self.rent.rent(|a: &Analysis| f(a))
    }
}

rental! {
    mod rent {
        use super::*;

        #[rental]
        pub struct R {
            file: Box<File>,
            analysis: Analysis<'file>
        }
    }
}

#[test]
fn analysis_is_sync() {
    fn assert_is_sync<T: Sync>() {}
    assert_is_sync::<Analysis>();
}

#[cfg(test)]
fn check_diagnostics(code: &str, expected_diagnostics: &str) {
    use fall_tree::test_util::report_diff;
    let file = ::editor_api::analyse(code.to_string());

    file.analyse(|a| {
        let d = a.collect_all_diagnostics();
        let actual = d.into_iter().map(|d| {
            let s = match d.severity {
                Severity::Error => 'E',
                Severity::Warning => 'W',
            };
            format!("{}: {} {}", a.file().node().text().slice(d.range), s, d.message)
        }).collect::<Vec<_>>().join("\n");

        report_diff(expected_diagnostics, &actual);
    })
}

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

    use fall_tree::search;
    use fall_tree::search::ast;

    use ::test_util::parse_with_caret;

    #[test]
    fn test_diagnostics() {
        check_diagnostics(r"
       pub rule foo { <eof x> }
       rule bar { foo <abracadabra>}
       rule baz { <prev_is foo> <prev_is bar> <prev_is {foo}>}
       rule dupe { dupe }
       rule dupe { dupe }
    ", "\
<eof x>: E Wrong number of arguments, expected 0, got 1
x: E Unresolved reference
abracadabra: E Unresolved reference
baz: W Unused rule
<prev_is bar>: E <prev_is> arguments must be public rules
<prev_is {foo}>: E <prev_is> arguments must be public rules
dupe: W Unused rule
dupe: E Duplicate rule");
    }

    #[test]
    fn test_resolve_no_args() {
        check_resolved("rule foo { <^eof> }", "Eof");
        check_resolved("rule foo { <^any> }", "Any");
        check_resolved("rule foo { <^commit> }", "Commit");
    }

    #[test]
    fn test_resolve_simple_args() {
        check_resolved("rule foo { <^not a> }", "Not(RefExpr@[16; 17))");
        check_resolved("rule foo { <^with_skip a b> }", "WithSkip(RefExpr@[22; 23), RefExpr@[24; 25))");
        check_resolved("rule foo { <^layer a b> }", "Layer(RefExpr@[18; 19), RefExpr@[20; 21))");
    }

    #[test]
    fn test_resolve_layer() {
        check_resolved("rule foo { <^enter 'ctx' a> <is_in 'ctx' >}", "Enter(0, RefExpr@[24; 25))");
        check_resolved("rule foo { <^exit  'ctx' a> <is_in 'ctx' >}", "Exit(0, RefExpr@[24; 25))");

        check_resolved("rule foo { <is_in 'ctx1'> <^enter 'ctx2' a> <is_in 'ctx1'> <enter 'ctx2' a> }", "Enter(1, RefExpr@[40; 41))");
        check_resolved("rule foo { <is_in 'ctx1'> <^exit  'ctx2' a> <is_in 'ctx1'> <exit  'ctx2' a> }", "Exit(1, RefExpr@[40; 41))");

        check_resolved(
            "rule foo { <enter 'ctx1' a> <enter 'ctx2' b> <^is_in 'ctx2'> <is_in 'ctx1'> }",
            "IsIn(1)"
        );

        check(
            r" rule foo { <^is_in foo>}",
            None,
            r#"[(foo, "Context should be a single quoted string")]"#,
        )
    }

    #[test]
    fn test_not() {
        resolve_call("rule foo { <^not bar>}", |c, _| {
            let c = c.unwrap();
            assert_eq!(format!("{:?}", c), "Not(RefExpr@[16; 19))");
        })
    }

    #[test]
    fn test_unresolved_call() {
        check(
            r"rule foo { <^abracadabra>}",
            None,
            r#"[(abracadabra, "Unresolved reference")]"#
        );
    }

    #[test]
    fn wrong_arity() {
        check(
            "rule foo { <^eof foo> }",
            Some(CallKind::Eof),
            r#"[(<eof foo>, "Wrong number of arguments, expected 0, got 1")]"#
        );

        check(
            "rule foo { <^any foo bar> }",
            Some(CallKind::Any),
            r#"[(<any foo bar>, "Wrong number of arguments, expected 0, got 2")]"#
        );

        check(
            "rule foo { <^commit foo bar baz> }",
            Some(CallKind::Commit),
            r#"[(<commit foo bar baz>, "Wrong number of arguments, expected 0, got 3")]"#
        );

        check(
            "rule foo { <^not> }",
            None,
            r#"[(<not>, "Wrong number of arguments, expected 1, got 0")]"#
        );
    }

    #[test]
    fn check_prev_is() {
        check_resolved(
            "pub rule foo {} pub rule bar {} rule baz { <^prev_is foo bar> }",
            "PrevIs([1, 2])"
        )
    }

    #[test]
    fn call_custom_rule() {
        check_resolved(
            "rule foo { <^bar a b>} rule bar(x, y) { }",
            "RuleCall(SynRule@[22; 40), [(0, RefExpr@[16; 17)), (1, RefExpr@[18; 19))])"
        );
        ::analysis::check_diagnostics(
            "rule foo { <bar <eof>>} rule bar(x, y) { }",
            "<bar <eof>>: E Expected 2 arguments, got 1"
        )
    }

    fn check(text: &str, kind: Option<CallKind>, diagnostics: &str) {
        resolve_call(text, |c, a| {
            assert_eq!(c, kind);
            assert_eq!(a.debug_diagnostics(), diagnostics);
        })
    }

    fn check_resolved(text: &str, kind: &str) {
        resolve_call(text, |c, a| {
            let c = c.unwrap();
            assert_eq!(format!("{:?}", c), kind);
            assert_eq!(a.debug_diagnostics(), "[]");
        })
    }

    fn resolve_call<F: FnOnce(Option<CallKind>, Analysis)>(text: &str, f: F) {
        let (file, caret) = parse_with_caret(&text);
        let call = {
            let leaf = search::find_leaf_at_offset(file.root(), caret)
                .left_biased().unwrap();
            ast::ancestor_exn(leaf)
        };
        let a = Analysis::new(::ast(&file));
        let c = a.resolve_call(call);
        f(c, a)
    }
}

<|endoftext|>


<commit_before>// xfail-fast
// xfail-test

// Broken on 32 bit, I guess?

trait Foo {
    fn f();
}

trait Bar : Foo {
    fn g();
}

struct A {
    x: int
}

impl A : Bar {
    fn g() { io::println("in g"); }
    fn f() { io::println("in f"); }
}

fn h<T:Foo>(a: &T) {
    a.f();
}

fn main() {
    let a = A { x: 3 };
    h(&a);
}

<commit_msg>Enable trait-inheritance-simple.rs<commit_after>trait Foo {
    fn f();
}

trait Bar : Foo {
    fn g();
}

struct A {
    x: int
}

impl A : Bar {
    fn g() { io::println("in g"); }
    fn f() { io::println("in f"); }
}

fn h<T:Foo>(a: &T) {
    a.f();
}

fn main() {
    let a = A { x: 3 };
    h(&a);
}

<|endoftext|>


<commit_before>use std::io::Write;
use std::io;
use common::serialize::BinarySerializable;
use std::mem;
use std::ops::Deref;

/// Computes the number of bits that will be used for bitpacking.
///
/// In general the target is the minimum number of bits
/// required to express the amplitude given in argument.
///
/// e.g. If the amplitude is 10, we can store all ints on simply 4bits.
///
/// The logic is slightly more convoluted here as for optimization
/// reasons, we want to ensure that a value spawns over at most 8 bytes
/// of aligns bytes.
///
/// Spawning over 9 bytes is possible for instance, if we do
/// bitpacking with an amplitude of 63 bits.
/// In this case, the second int will start on bit
/// 63 (which belongs to byte 7) and ends at byte 15;
/// Hence 9 bytes (from byte 7 to byte 15 included).
///
/// To avoid this, we force the number of bits to 64bits
/// when the result is greater than `64-8 = 56 bits`.
///
/// Note that this only affects rare use cases spawning over
/// a very large range of values. Even in this case, it results
/// in an extra cost of at most 12% compared to the optimal
/// number of bits.
pub fn compute_num_bits(amplitude: u64) -> u8 {
    let amplitude = (64u32 - amplitude.leading_zeros()) as u8;
    if amplitude <= 64 - 8 { amplitude } else { 64 }
}

pub struct BitPacker {
    mini_buffer: u64,
    mini_buffer_written: usize,
    num_bits: usize,
    written_size: usize,
}

impl BitPacker {
    pub fn new(num_bits: usize) -> BitPacker {
        BitPacker {
            mini_buffer: 0u64,
            mini_buffer_written: 0,
            num_bits: num_bits,
            written_size: 0,
        }
    }

    pub fn write<TWrite: Write>(&mut self, val: u64, output: &mut TWrite) -> io::Result<()> {
        let val_u64 = val as u64;
        if self.mini_buffer_written + self.num_bits > 64 {
            self.mini_buffer |= val_u64.wrapping_shl(self.mini_buffer_written as u32);
            self.written_size += self.mini_buffer.serialize(output)?;
            self.mini_buffer = val_u64.wrapping_shr((64 - self.mini_buffer_written) as u32);
            self.mini_buffer_written = self.mini_buffer_written + (self.num_bits as usize) - 64;
        } else {
            self.mini_buffer |= val_u64 << self.mini_buffer_written;
            self.mini_buffer_written += self.num_bits;
            if self.mini_buffer_written == 64 {
                self.written_size += self.mini_buffer.serialize(output)?;
                self.mini_buffer_written = 0;
                self.mini_buffer = 0u64;
            }
        }
        Ok(())
    }

    fn flush<TWrite: Write>(&mut self, output: &mut TWrite) -> io::Result<()> {
        if self.mini_buffer_written > 0 {
            let num_bytes = (self.mini_buffer_written + 7) / 8;
            let arr: [u8; 8] = unsafe { mem::transmute::<u64, [u8; 8]>(self.mini_buffer) };
            output.write_all(&arr[..num_bytes])?;
            self.written_size += num_bytes;
            self.mini_buffer_written = 0;
        }
        Ok(())
    }

    pub fn close<TWrite: Write>(&mut self, output: &mut TWrite) -> io::Result<usize> {
        self.flush(output)?;
        // Padding the write file to simplify reads.
        output.write_all(&[0u8; 7])?;
        self.written_size += 7;
        Ok(self.written_size)
    }
}

pub struct BitUnpacker<Data>
    where Data: Deref<Target = [u8]>
{
    num_bits: usize,
    mask: u64,
    data: Data,
}

impl<Data> BitUnpacker<Data>
    where Data: Deref<Target = [u8]>
{
    pub fn new(data: Data, num_bits: usize) -> BitUnpacker<Data> {
        let mask: u64 = if num_bits == 64 {
            !0u64
        } else {
            (1u64 << num_bits) - 1u64
        };
        BitUnpacker {
            num_bits: num_bits,
            mask: mask,
            data: data,
        }
    }

    pub fn get(&self, idx: usize) -> u64 {
        if self.num_bits == 0 {
            return 0;
        }
        let addr = (idx * self.num_bits) / 8;
        let bit_shift = idx * self.num_bits - addr * 8;
        let val_unshifted_unmasked: u64;
        debug_assert!(addr + 8 <= self.data.len(),
                      "The fast field field should have been padded with 7 bytes.");
        val_unshifted_unmasked =
            unsafe { *(self.data.as_ptr().offset(addr as isize) as *const u64) };
        let val_shifted = (val_unshifted_unmasked >> bit_shift) as u64;
        (val_shifted & self.mask)
    }
}

#[cfg(test)]
mod test {
    use super::{BitPacker, BitUnpacker, compute_num_bits};

    #[test]
    fn test_compute_num_bits() {
        assert_eq!(compute_num_bits(1), 1u8);
        assert_eq!(compute_num_bits(0), 0u8);
        assert_eq!(compute_num_bits(2), 2u8);
        assert_eq!(compute_num_bits(3), 2u8);
        assert_eq!(compute_num_bits(4), 3u8);
        assert_eq!(compute_num_bits(255), 8u8);
        assert_eq!(compute_num_bits(256), 9u8);
        assert_eq!(compute_num_bits(5_000_000_000), 33u8);
    }

    fn test_bitpacker_util(len: usize, num_bits: usize) {
        let mut data = Vec::new();
        let mut bitpacker = BitPacker::new(num_bits);
        let max_val: u64 = (1 << num_bits) - 1;
        let vals: Vec<u64> = (0u64..len as u64)
            .map(|i| if max_val == 0 { 0 } else { i % max_val })
            .collect();
        for &val in &vals {
            bitpacker.write(val, &mut data).unwrap();
        }
        let num_bytes = bitpacker.close(&mut data).unwrap();
        assert_eq!(num_bytes, (num_bits * len + 7) / 8 + 7);
        assert_eq!(data.len(), num_bytes);
        let bitunpacker = BitUnpacker::new(data, num_bits);
        for (i, val) in vals.iter().enumerate() {
            assert_eq!(bitunpacker.get(i), *val);
        }
    }

    #[test]
    fn test_bitpacker() {
        test_bitpacker_util(10, 3);
        test_bitpacker_util(10, 0);
        test_bitpacker_util(10, 1);
        test_bitpacker_util(6, 14);
        test_bitpacker_util(1000, 14);
    }
}
<commit_msg>Optimization in bitpacker<commit_after>use std::io::Write;
use std::io;
use common::serialize::BinarySerializable;
use std::mem;
use std::ops::Deref;

/// Computes the number of bits that will be used for bitpacking.
///
/// In general the target is the minimum number of bits
/// required to express the amplitude given in argument.
///
/// e.g. If the amplitude is 10, we can store all ints on simply 4bits.
///
/// The logic is slightly more convoluted here as for optimization
/// reasons, we want to ensure that a value spawns over at most 8 bytes
/// of aligns bytes.
///
/// Spawning over 9 bytes is possible for instance, if we do
/// bitpacking with an amplitude of 63 bits.
/// In this case, the second int will start on bit
/// 63 (which belongs to byte 7) and ends at byte 15;
/// Hence 9 bytes (from byte 7 to byte 15 included).
///
/// To avoid this, we force the number of bits to 64bits
/// when the result is greater than `64-8 = 56 bits`.
///
/// Note that this only affects rare use cases spawning over
/// a very large range of values. Even in this case, it results
/// in an extra cost of at most 12% compared to the optimal
/// number of bits.
pub fn compute_num_bits(amplitude: u64) -> u8 {
    let amplitude = (64u32 - amplitude.leading_zeros()) as u8;
    if amplitude <= 64 - 8 { amplitude } else { 64 }
}

pub struct BitPacker {
    mini_buffer: u64,
    mini_buffer_written: usize,
    num_bits: usize,
    written_size: usize,
}

impl BitPacker {
    pub fn new(num_bits: usize) -> BitPacker {
        BitPacker {
            mini_buffer: 0u64,
            mini_buffer_written: 0,
            num_bits: num_bits,
            written_size: 0,
        }
    }

    pub fn write<TWrite: Write>(&mut self, val: u64, output: &mut TWrite) -> io::Result<()> {
        let val_u64 = val as u64;
        if self.mini_buffer_written + self.num_bits > 64 {
            self.mini_buffer |= val_u64.wrapping_shl(self.mini_buffer_written as u32);
            self.written_size += self.mini_buffer.serialize(output)?;
            self.mini_buffer = val_u64.wrapping_shr((64 - self.mini_buffer_written) as u32);
            self.mini_buffer_written = self.mini_buffer_written + (self.num_bits as usize) - 64;
        } else {
            self.mini_buffer |= val_u64 << self.mini_buffer_written;
            self.mini_buffer_written += self.num_bits;
            if self.mini_buffer_written == 64 {
                self.written_size += self.mini_buffer.serialize(output)?;
                self.mini_buffer_written = 0;
                self.mini_buffer = 0u64;
            }
        }
        Ok(())
    }

    fn flush<TWrite: Write>(&mut self, output: &mut TWrite) -> io::Result<()> {
        if self.mini_buffer_written > 0 {
            let num_bytes = (self.mini_buffer_written + 7) / 8;
            let arr: [u8; 8] = unsafe { mem::transmute::<u64, [u8; 8]>(self.mini_buffer) };
            output.write_all(&arr[..num_bytes])?;
            self.written_size += num_bytes;
            self.mini_buffer_written = 0;
        }
        Ok(())
    }

    pub fn close<TWrite: Write>(&mut self, output: &mut TWrite) -> io::Result<usize> {
        self.flush(output)?;
        // Padding the write file to simplify reads.
        output.write_all(&[0u8; 7])?;
        self.written_size += 7;
        Ok(self.written_size)
    }
}

pub struct BitUnpacker<Data>
    where Data: Deref<Target = [u8]>
{
    num_bits: usize,
    mask: u64,
    data: Data,
}

impl<Data> BitUnpacker<Data>
    where Data: Deref<Target = [u8]>
{
    pub fn new(data: Data, num_bits: usize) -> BitUnpacker<Data> {
        let mask: u64 = if num_bits == 64 {
            !0u64
        } else {
            (1u64 << num_bits) - 1u64
        };
        BitUnpacker {
            num_bits: num_bits,
            mask: mask,
            data: data,
        }
    }

    pub fn get(&self, idx: usize) -> u64 {
        if self.num_bits == 0 {
            return 0;
        }
        let data: &[u8] = &*self.data;
        let num_bits = self.num_bits;
        let mask = self.mask;
        let addr_in_bits = idx * num_bits;
        let addr = addr_in_bits >> 3;
        let bit_shift = addr_in_bits & 7;
        debug_assert!(addr + 8 <= data.len(),
                      "The fast field field should have been padded with 7 bytes.");
        let val_unshifted_unmasked: u64 =
            unsafe { *(data.as_ptr().offset(addr as isize) as *const u64) };
        let val_shifted = (val_unshifted_unmasked >> bit_shift) as u64;
        (val_shifted & mask)
    }
}

#[cfg(test)]
mod test {
    use super::{BitPacker, BitUnpacker, compute_num_bits};

    #[test]
    fn test_compute_num_bits() {
        assert_eq!(compute_num_bits(1), 1u8);
        assert_eq!(compute_num_bits(0), 0u8);
        assert_eq!(compute_num_bits(2), 2u8);
        assert_eq!(compute_num_bits(3), 2u8);
        assert_eq!(compute_num_bits(4), 3u8);
        assert_eq!(compute_num_bits(255), 8u8);
        assert_eq!(compute_num_bits(256), 9u8);
        assert_eq!(compute_num_bits(5_000_000_000), 33u8);
    }

    fn test_bitpacker_util(len: usize, num_bits: usize) {
        let mut data = Vec::new();
        let mut bitpacker = BitPacker::new(num_bits);
        let max_val: u64 = (1 << num_bits) - 1;
        let vals: Vec<u64> = (0u64..len as u64)
            .map(|i| if max_val == 0 { 0 } else { i % max_val })
            .collect();
        for &val in &vals {
            bitpacker.write(val, &mut data).unwrap();
        }
        let num_bytes = bitpacker.close(&mut data).unwrap();
        assert_eq!(num_bytes, (num_bits * len + 7) / 8 + 7);
        assert_eq!(data.len(), num_bytes);
        let bitunpacker = BitUnpacker::new(data, num_bits);
        for (i, val) in vals.iter().enumerate() {
            assert_eq!(bitunpacker.get(i), *val);
        }
    }

    #[test]
    fn test_bitpacker() {
        test_bitpacker_util(10, 3);
        test_bitpacker_util(10, 0);
        test_bitpacker_util(10, 1);
        test_bitpacker_util(6, 14);
        test_bitpacker_util(1000, 14);
    }
}
<|endoftext|>