<commit_before>//! Parse and write Iridium Short Burst Data (SBD) messages.
//!
//! Iridium is both a [satellite constellation](https://en.wikipedia.org/wiki/Iridium_satellite_constellation)
//! and a [company](https://en.wikipedia.org/wiki/Iridium_Communications) that provides satellite
//! communications. The Iridium network is used by phones, modems, and other communication devices.
//!
//! One mode of transmitting data over the Iridium network is via Short Burst Data (SBD) messages.
//! These messages carry a payload of some small number of bytes, usually less than one thousand.
//! Messages can be Mobile-Originated (MO), meaning that they are sent *from* an Iridium modem, or
//! Mobile-Terminated (MT), meaning that the are sent *to* an Iridium modem. Mobile-originated
//! messages are delivered either to an email address via MIME attachment, or directly to a given
//! IP address and port via TCP; this second method is called DirectIP.
//!
//! This is a simple library for reading mobile originated SBD messages from a stream, decoding
//! their headers and data payloads, and writing them back to a stream. This library does not
//! handle mobile terminated messages.
//!
//! Most of the functionality of this library is exposed by a single executable, named `sbd` that
//! is part of this package. Use the `sbd` executable to inspect raw sbd files stores on a
//! filesystem, discover sbd files on a filesystem, and start a forever-running server to receive
//! Iridium SBD DirectIP messages.

pub mod directip;
mod information_element;
pub mod logger;
pub mod filesystem;
pub mod message;

pub use message::Message;

extern crate byteorder;
extern crate chrono;
#[macro_use] extern crate enum_primitive;
extern crate glob;
#[macro_use] extern crate log;
extern crate num;

use std::result;

#[derive(Debug)]
pub enum Error {
    ByteorderError(byteorder::Error),
    IoError(std::io::Error),
    InvalidImei,
    InvalidProtocolRevisionNumber(u8),
    GlobError(glob::GlobError),
    MissingMobileOriginatedHeader,
    MissingMobileOriginatedPayload,
    Oversized, // Oversized doesn't demand a size since we don't want to find out how much there really is
    PatternError(glob::PatternError),
    Undersized(usize),
}

impl From<byteorder::Error> for Error {
    fn from(err: byteorder::Error) -> Error {
        Error::ByteorderError(err)
    }
}

impl From<glob::PatternError> for Error {
    fn from(err: glob::PatternError) -> Error {
        Error::PatternError(err)
    }
}

impl From<glob::GlobError> for Error {
    fn from(err: glob::GlobError) -> Error {
        Error::GlobError(err)
    }
}

impl From<std::io::Error> for Error {
    fn from(err: std::io::Error) -> Error {
        Error::IoError(err)
    }
}

pub type Result<T> = result::Result<T, Error>;
<commit_msg>Document our errors<commit_after>//! Parse and write Iridium Short Burst Data (SBD) messages.
//!
//! Iridium is both a [satellite constellation](https://en.wikipedia.org/wiki/Iridium_satellite_constellation)
//! and a [company](https://en.wikipedia.org/wiki/Iridium_Communications) that provides satellite
//! communications. The Iridium network is used by phones, modems, and other communication devices.
//!
//! One mode of transmitting data over the Iridium network is via Short Burst Data (SBD) messages.
//! These messages carry a payload of some small number of bytes, usually less than one thousand.
//! Messages can be Mobile-Originated (MO), meaning that they are sent *from* an Iridium modem, or
//! Mobile-Terminated (MT), meaning that the are sent *to* an Iridium modem. Mobile-originated
//! messages are delivered either to an email address via MIME attachment, or directly to a given
//! IP address and port via TCP; this second method is called DirectIP.
//!
//! This is a simple library for reading mobile originated SBD messages from a stream, decoding
//! their headers and data payloads, and writing them back to a stream. This library does not
//! handle mobile terminated messages.
//!
//! Most of the functionality of this library is exposed by a single executable, named `sbd` that
//! is part of this package. Use the `sbd` executable to inspect raw sbd files stores on a
//! filesystem, discover sbd files on a filesystem, and start a forever-running server to receive
//! Iridium SBD DirectIP messages.

pub mod directip;
mod information_element;
pub mod logger;
pub mod filesystem;
pub mod message;

pub use message::Message;

extern crate byteorder;
extern crate chrono;
#[macro_use] extern crate enum_primitive;
extern crate glob;
#[macro_use] extern crate log;
extern crate num;

use std::result;

/// Crate-specific errors
#[derive(Debug)]
pub enum Error {
    /// An error while reading bytes from a stream with the byteorder crate.
    ByteorderError(byteorder::Error),
    /// A wrapper around a std::io::Error.
    IoError(std::io::Error),
    /// Invalid IMEI number.
    InvalidImei,
    /// Invalid protocol revision number.
    InvalidProtocolRevisionNumber(u8),
    /// Wrapper around a glob error.
    GlobError(glob::GlobError),
    /// Missing mobile originated header.
    MissingMobileOriginatedHeader,
    /// Missing mobile originated payload.
    MissingMobileOriginatedPayload,
    /// An oversized message.
    ///
    /// Oversized doesn't demand a size since we don't want to find out how much there really is.
    Oversized,
    /// Wrapper around a glob::PatternError.
    PatternError(glob::PatternError),
    /// An undersized message.
    Undersized(usize),
}

impl From<byteorder::Error> for Error {
    fn from(err: byteorder::Error) -> Error {
        Error::ByteorderError(err)
    }
}

impl From<glob::PatternError> for Error {
    fn from(err: glob::PatternError) -> Error {
        Error::PatternError(err)
    }
}

impl From<glob::GlobError> for Error {
    fn from(err: glob::GlobError) -> Error {
        Error::GlobError(err)
    }
}

impl From<std::io::Error> for Error {
    fn from(err: std::io::Error) -> Error {
        Error::IoError(err)
    }
}

pub type Result<T> = result::Result<T, Error>;
<|endoftext|>


<commit_before>/*
Copyright ⓒ 2016 Daniel Keep.

Licensed under the MIT license (see LICENSE or <http://opensource.org
/licenses/MIT>) or the Apache License, Version 2.0 (see LICENSE of
<http://www.apache.org/licenses/LICENSE-2.0>), at your option. All
files in the project carrying such notice may not be copied, modified,
or distributed except according to those terms.
*/
/*!

Provides the [`Boolinator`](trait.Boolinator.html) trait, which lets you use `Option` and `Result`-style combinators with `bool`s.

**Links**

* [Latest Release](https://crates.io/crates/boolinator/)
* [Latest Docs](https://danielkeep.github.io/rust-boolinator/doc/boolinator/index.html)
* [Repository](https://github.com/DanielKeep/rust-boolinator)

## Compatibility

`boolinator` is tested against Rust 1.0+.  *Exhaustively* so.

*/
// Can't have undocumented APIs!  Nosiree!
#![deny(missing_docs)]

/**
This trait defines a number of combinator-style methods for use with `bool` values.

In general, `true`/`false` map to `Some(_)`/`None` and `Ok(_)`/`Err(_)` respectively.
*/
pub trait Boolinator: Sized {
    /**
    If this value is `true`, returns `Some(())`; `None` otherwise.
    */
    fn as_option(self) -> Option<()>;

    /**
    If this value is `true`, returns `Some(some)`; `None` otherwise.
    */
    fn as_some<T>(self, some: T) -> Option<T>;

    /**
    If this value is `true`, returns `Some(some())`; `None` otherwise.
    */
    fn as_some_from<T, F>(self, some: F) -> Option<T>
    where F: FnOnce() -> T;

    /**
    If this value is `true`, returns `opt`; `None` otherwise.
    */
    fn and_option<T>(self, opt: Option<T>) -> Option<T>;

    /**
    If this value is `true`, returns `opt()`; `None` otherwise.
    */
    fn and_option_from<T, F>(self, opt: F) -> Option<T>
    where F: FnOnce() -> Option<T>;

    /**
    If this value is `true`, returns `Ok(ok)`; `Err(err)` otherwise.
    */
    fn as_result<T, E>(self, ok: T, err: E) -> Result<T, E>;

    /**
    If this value is `true`, returns `Ok(ok())`; `Err(err())` otherwise.
    */
    fn as_result_from<T, E, F, G>(self, ok: F, err: G) -> Result<T, E>
    where F: FnOnce() -> T, G: FnOnce() -> E;

    /**
    If this value is `true`, returns `Ok(())`; `Err(err)` otherwise.
    */
    fn ok_or<E>(self, err: E) -> Result<(), E>;

    /**
    If this value is `true`, returns `Ok(())`; `Err(err())` otherwise.
    */
    fn ok_or_else<E, G>(self, err: G) -> Result<(), E>
    where G: FnOnce() -> E;

    /**
    If this value is `true`, panics with `msg`; does nothing otherwise.
    */
    fn expect(self, msg: &str);
}

impl Boolinator for bool {
    #[inline]
    fn as_option(self) -> Option<()> {
        if self { Some(()) } else { None }
    }

    #[inline]
    fn as_some<T>(self, some: T) -> Option<T> {
        if self { Some(some) } else { None }
    }

    #[inline]
    fn as_some_from<T, F>(self, some: F) -> Option<T>
    where F: FnOnce() -> T {
        if self { Some(some()) } else { None }
    }

    #[inline]
    fn and_option<T>(self, opt: Option<T>) -> Option<T> {
        if self { opt } else { None }
    }

    #[inline]
    fn and_option_from<T, F>(self, opt: F) -> Option<T>
    where F: FnOnce() -> Option<T> {
        if self { opt() } else { None }
    }

    #[inline]
    fn as_result<T, E>(self, ok: T, err: E) -> Result<T, E> {
        if self { Ok(ok) } else { Err(err) }
    }

    #[inline]
    fn as_result_from<T, E, F, G>(self, ok: F, err: G) -> Result<T, E>
    where F: FnOnce() -> T, G: FnOnce() -> E {
        if self { Ok(ok()) } else { Err(err()) }
    }

    #[inline]
    fn ok_or<E>(self, err: E) -> Result<(), E> {
        if self { Ok(()) } else { Err(err) }
    }

    #[inline]
    fn ok_or_else<E, G>(self, err: G) -> Result<(), E>
    where G: FnOnce() -> E {
        if self { Ok(()) } else { Err(err()) }
    }

    #[inline]
    fn expect(self, msg: &str) {
        if self { () } else { panic!("{}", msg) }
    }
}

/*
Serious code must have serious tests, and Boolinator is serious business!
*/
#[cfg(test)]
mod tests {
    use super::Boolinator; // as opposed to the original NES version.

    #[test]
    fn test_as_option() {
        // Very test.
        assert_eq!(true.as_option(), Some(()));
        assert_eq!(false.as_option(), None);
    }

    #[test]
    fn test_as_some() {
        // Much serious.
        let love = true;
        let everybody = love.as_some("body").expect("needs");
        assert_eq!(everybody, "body");

        assert_eq!((!love).as_some("money can buy"), None);
    }

    #[test]
    fn test_as_some_from() {
        // Wow.
        let mothers = vec![true, false, false, true, false, true];
        assert!(mothers.into_iter()
            .map(|e| e.as_some_from(|| Some("em")))
            .filter(Option::is_some)
            .count() > 0);
    }

    #[test]
    fn test_and_option() {
        // Such original.
        assert_eq!(true.and_option(Some("fries with that")), Some("fries with that"));
        assert_eq!(false.and_option(Some("fries with that")), None);
        assert_eq!(true.and_option(None), None::<()>);
        assert_eq!(false.and_option(None), None::<()>);
    }

    #[test]
    fn test_and_option_from() {
        // Such original.
        assert_eq!(true.and_option_from(|| Some("chips too, guv'")), Some("chips too, guv'"));
        assert_eq!(false.and_option_from(|| Some("chips too, guv'")), None);
        assert_eq!(true.and_option_from(|| None), None::<()>);
        assert_eq!(false.and_option_from(|| None), None::<()>);
    }

    #[test]
    fn test_as_result() {
        // Very result.
        assert_eq!(true.as_result("now; ", ", what?"),   Ok("now; "));
        assert_eq!(false.as_result("now; ", ", what?"),  Err(", what?"));
    }

    #[test]
    fn test_as_result_from() {
        // Code good.
        assert_eq!(true.as_result_from(|| "four space indent", || "anything else"), Ok("four space indent"));
        assert_eq!(false.as_result_from(|| "four space indent", || "anything else"), Err("anything else"));
    }

    #[test]
    fn test_ok_or() {
        // Ok.
        let mut annie = true;
        assert_eq!(annie.ok_or("hit back"), Ok(()));
        annie = false;
        assert_eq!(annie.ok_or("hit back"), Err("hit back"));
    }

    #[test]
    fn test_ok_or_else() {
        // Ok.
        let mut annie = true;
        assert_eq!(annie.ok_or_else(|| "hit back"), Ok(()));
        annie = false;
        assert_eq!(annie.ok_or_else(|| "hit back"), Err("hit back"));
    }

    const DREAMS: &'static str = "love and financial security";

    #[test]
    fn test_expect() {
        // Movies lie.
        true.expect(DREAMS);
    }

    #[test]
    #[should_panic]
    fn test_expect_reality() {
        // Send hugs.
        false.expect(DREAMS);
    }
}
<commit_msg>Made `bool::expect` documentation *slightly* more accurate.<commit_after>/*
Copyright ⓒ 2016 Daniel Keep.

Licensed under the MIT license (see LICENSE or <http://opensource.org
/licenses/MIT>) or the Apache License, Version 2.0 (see LICENSE of
<http://www.apache.org/licenses/LICENSE-2.0>), at your option. All
files in the project carrying such notice may not be copied, modified,
or distributed except according to those terms.
*/
/*!

Provides the [`Boolinator`](trait.Boolinator.html) trait, which lets you use `Option` and `Result`-style combinators with `bool`s.

**Links**

* [Latest Release](https://crates.io/crates/boolinator/)
* [Latest Docs](https://danielkeep.github.io/rust-boolinator/doc/boolinator/index.html)
* [Repository](https://github.com/DanielKeep/rust-boolinator)

## Compatibility

`boolinator` is tested against Rust 1.0+.  *Exhaustively* so.

*/
// Can't have undocumented APIs!  Nosiree!
#![deny(missing_docs)]

/**
This trait defines a number of combinator-style methods for use with `bool` values.

In general, `true`/`false` map to `Some(_)`/`None` and `Ok(_)`/`Err(_)` respectively.
*/
pub trait Boolinator: Sized {
    /**
    If this value is `true`, returns `Some(())`; `None` otherwise.
    */
    fn as_option(self) -> Option<()>;

    /**
    If this value is `true`, returns `Some(some)`; `None` otherwise.
    */
    fn as_some<T>(self, some: T) -> Option<T>;

    /**
    If this value is `true`, returns `Some(some())`; `None` otherwise.
    */
    fn as_some_from<T, F>(self, some: F) -> Option<T>
    where F: FnOnce() -> T;

    /**
    If this value is `true`, returns `opt`; `None` otherwise.
    */
    fn and_option<T>(self, opt: Option<T>) -> Option<T>;

    /**
    If this value is `true`, returns `opt()`; `None` otherwise.
    */
    fn and_option_from<T, F>(self, opt: F) -> Option<T>
    where F: FnOnce() -> Option<T>;

    /**
    If this value is `true`, returns `Ok(ok)`; `Err(err)` otherwise.
    */
    fn as_result<T, E>(self, ok: T, err: E) -> Result<T, E>;

    /**
    If this value is `true`, returns `Ok(ok())`; `Err(err())` otherwise.
    */
    fn as_result_from<T, E, F, G>(self, ok: F, err: G) -> Result<T, E>
    where F: FnOnce() -> T, G: FnOnce() -> E;

    /**
    If this value is `true`, returns `Ok(())`; `Err(err)` otherwise.
    */
    fn ok_or<E>(self, err: E) -> Result<(), E>;

    /**
    If this value is `true`, returns `Ok(())`; `Err(err())` otherwise.
    */
    fn ok_or_else<E, G>(self, err: G) -> Result<(), E>
    where G: FnOnce() -> E;

    /**
    If this value is `true`, does nothing; panics with `msg` otherwise.
    */
    fn expect(self, msg: &str);
}

impl Boolinator for bool {
    #[inline]
    fn as_option(self) -> Option<()> {
        if self { Some(()) } else { None }
    }

    #[inline]
    fn as_some<T>(self, some: T) -> Option<T> {
        if self { Some(some) } else { None }
    }

    #[inline]
    fn as_some_from<T, F>(self, some: F) -> Option<T>
    where F: FnOnce() -> T {
        if self { Some(some()) } else { None }
    }

    #[inline]
    fn and_option<T>(self, opt: Option<T>) -> Option<T> {
        if self { opt } else { None }
    }

    #[inline]
    fn and_option_from<T, F>(self, opt: F) -> Option<T>
    where F: FnOnce() -> Option<T> {
        if self { opt() } else { None }
    }

    #[inline]
    fn as_result<T, E>(self, ok: T, err: E) -> Result<T, E> {
        if self { Ok(ok) } else { Err(err) }
    }

    #[inline]
    fn as_result_from<T, E, F, G>(self, ok: F, err: G) -> Result<T, E>
    where F: FnOnce() -> T, G: FnOnce() -> E {
        if self { Ok(ok()) } else { Err(err()) }
    }

    #[inline]
    fn ok_or<E>(self, err: E) -> Result<(), E> {
        if self { Ok(()) } else { Err(err) }
    }

    #[inline]
    fn ok_or_else<E, G>(self, err: G) -> Result<(), E>
    where G: FnOnce() -> E {
        if self { Ok(()) } else { Err(err()) }
    }

    #[inline]
    fn expect(self, msg: &str) {
        if self { () } else { panic!("{}", msg) }
    }
}

/*
Serious code must have serious tests, and Boolinator is serious business!
*/
#[cfg(test)]
mod tests {
    use super::Boolinator; // as opposed to the original NES version.

    #[test]
    fn test_as_option() {
        // Very test.
        assert_eq!(true.as_option(), Some(()));
        assert_eq!(false.as_option(), None);
    }

    #[test]
    fn test_as_some() {
        // Much serious.
        let love = true;
        let everybody = love.as_some("body").expect("needs");
        assert_eq!(everybody, "body");

        assert_eq!((!love).as_some("money can buy"), None);
    }

    #[test]
    fn test_as_some_from() {
        // Wow.
        let mothers = vec![true, false, false, true, false, true];
        assert!(mothers.into_iter()
            .map(|e| e.as_some_from(|| Some("em")))
            .filter(Option::is_some)
            .count() > 0);
    }

    #[test]
    fn test_and_option() {
        // Such original.
        assert_eq!(true.and_option(Some("fries with that")), Some("fries with that"));
        assert_eq!(false.and_option(Some("fries with that")), None);
        assert_eq!(true.and_option(None), None::<()>);
        assert_eq!(false.and_option(None), None::<()>);
    }

    #[test]
    fn test_and_option_from() {
        // Such original.
        assert_eq!(true.and_option_from(|| Some("chips too, guv'")), Some("chips too, guv'"));
        assert_eq!(false.and_option_from(|| Some("chips too, guv'")), None);
        assert_eq!(true.and_option_from(|| None), None::<()>);
        assert_eq!(false.and_option_from(|| None), None::<()>);
    }

    #[test]
    fn test_as_result() {
        // Very result.
        assert_eq!(true.as_result("now; ", ", what?"),   Ok("now; "));
        assert_eq!(false.as_result("now; ", ", what?"),  Err(", what?"));
    }

    #[test]
    fn test_as_result_from() {
        // Code good.
        assert_eq!(true.as_result_from(|| "four space indent", || "anything else"), Ok("four space indent"));
        assert_eq!(false.as_result_from(|| "four space indent", || "anything else"), Err("anything else"));
    }

    #[test]
    fn test_ok_or() {
        // Ok.
        let mut annie = true;
        assert_eq!(annie.ok_or("hit back"), Ok(()));
        annie = false;
        assert_eq!(annie.ok_or("hit back"), Err("hit back"));
    }

    #[test]
    fn test_ok_or_else() {
        // Ok.
        let mut annie = true;
        assert_eq!(annie.ok_or_else(|| "hit back"), Ok(()));
        annie = false;
        assert_eq!(annie.ok_or_else(|| "hit back"), Err("hit back"));
    }

    const DREAMS: &'static str = "love and financial security";

    #[test]
    fn test_expect() {
        // Movies lie.
        true.expect(DREAMS);
    }

    #[test]
    #[should_panic]
    fn test_expect_reality() {
        // Send hugs.
        false.expect(DREAMS);
    }
}
<|endoftext|>


<commit_before>// (c) 2015 Joost Yervante Damad <joost@damad.be>

// loosely based on https://github.com/cgaebel/sexp
// latest version can be found at https://github.com/andete/rust_sexp

use std::fmt;
use std::str::FromStr;
use std::f64;

use std::fs::File;
use std::io::prelude::*;

pub enum Atom {
  S(String),
  Q(String),
  I(i64),
  F(f64),
}

impl Atom {
    pub fn f(&self) -> Result<f64,String> {
        match *self {
            Atom::F(f) => Ok(f),
            ref x => Err(format!("not a float: {}", x))
        }
    }
    pub fn i(&self) -> Result<i64,String> {
        match *self {
            Atom::I(i) => Ok(i),
            ref x => Err(format!("not an int: {}", x))
        }
    }
    pub fn string(&self) -> Result<String,String> {
        match *self {
            Atom::S(ref s) => Ok(s.clone()),
            Atom::Q(ref s) => Ok(s.clone()),
            ref x => Err(format!("not a string: {}", x))
        }
    }

}

pub enum Sexp {
  Atom(Atom),
  Empty,
  List(Vec<Sexp>),
}

impl Sexp {
    pub fn atom(&self) -> Result<&Atom,String> {
        match *self {
            Sexp::Atom(ref f) => Ok(f),
            ref x => Err(format!("not an atom: {}", x))
        }
    }
    pub fn list(&self) -> Result<&Vec<Sexp>,String> {
        match *self {
            Sexp::List(ref v) => Ok(v),
            ref x => Err(format!("not a list: {}", x))
        }
    }
}

pub struct ParseError {
    msg: &'static str,
    line: usize,
    lpos: usize,
}

pub enum Error {
    ParseError(ParseError),
    IOError(std::io::Error),
}

pub type ERes<T> = Result<T, Error>;

fn err_parse<T>(msg: &'static str, state: &ParseState) -> ERes<T> {
    Err(Error::ParseError(ParseError { msg: msg, line: state.line, lpos: state.lpos}))
}

fn err_io<T>(err: std::io::Error) -> ERes<T> {
    Err(Error::IOError(err))
}

struct ParseState {
    pos: usize,
    line: usize,
    lpos: usize,
    vec: Vec<char>,
    len: usize,
}

impl ParseState {
    fn peek(&self) -> ERes<char> {
        if self.pos >= self.len {
            return err_parse("end of document reached", self)
        }
        Ok(self.vec[self.pos])
    }
    fn peek_option(&self) -> Option<char> {
        if self.pos >= self.len {
            return None
        }
        Some(self.vec[self.pos])
    }
    fn next(&mut self) -> ERes<()> {
        if self.pos >= self.len {
            return err_parse("end of document reached", self)
        }
        let c = self.vec[self.pos];
        self.pos += 1;
        match c {
            '\r' | '\n' => {
                self.lpos = 0;
                self.line += 1;
            }
            _ => {
                self.lpos += 1;
            }
        }
        Ok(())
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Error::ParseError(ref p) =>
                write!(f, "Parse Error {}:{}: {}", p.line, p.lpos, p.msg),
            Error::IOError(ref i) =>
                write!(f, "{}", i)
        }
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Error::ParseError(ref p) =>
                write!(f, "Parse Error {}:{}: {}", p.line, p.lpos, p.msg),
            Error::IOError(ref i) =>
                write!(f, "{}", i)
        }
    }
}

impl fmt::Display for Atom {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    match *self {
        Atom::S(ref s) => write!(f, "{}", s),
        Atom::Q(ref s) => write!(f, "\"{}\"", s),
        Atom::I(i)     => write!(f, "{}", i),
        Atom::F(d)     => {
            let z = d.floor();
            if d - z < f64::EPSILON {
                write!(f, "{}.0", z)
            } else {
                write!(f, "{}", d)
            }
        }
    }
  }
}

impl fmt::Display for Sexp {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    match *self {
      Sexp::Atom(ref a) => write!(f, "{}", a),
      Sexp::List(ref xs) => {
        try!(write!(f, "("));
        for (i, x) in xs.iter().enumerate() {
          let s = if i == 0 { "" } else { " " };
          try!(write!(f, "{}{}", s, x));
        }
        write!(f, ")")
      },
      Sexp::Empty => write!(f, ""),
    }
  }
}

fn parse_list(state: &mut ParseState) -> ERes<Sexp> {
    //println!("list");
    try!(state.next()); // skip (
    let mut l: Vec<Sexp> = Vec::new();
    loop {
        match try!(state.peek()) {
            ')' => {
                try!(state.next());
                break;
            }
            _ => {
                l.push(try!(parse_sexp(state)));
            }
        }
    }
    //println!("Found list");
    Ok(Sexp::List(l))
}

fn parse_quoted_string(state: &mut ParseState) -> ERes<Atom> {
    //println!("qstring");
    try!(state.next()); // skip "
    let mut s = String::new();
    loop {
        match try!(state.peek()) {
            '"' => {
                try!(state.next());
                break
            }
            x @ '\r' | x @ '\n' => {
                s.push(x);
                try!(state.next());
                }
            x => {
                s.push(x);
                try!(state.next());
            }
        }
    }
    //println!("Found quoted string {}", s);
    Ok(Atom::Q(s))
}

fn parse_string(state: &mut ParseState) -> ERes<Atom> {
    //println!("string");
    let mut s = String::new();
    loop {
        match state.peek_option() {
            Some(x) => {
                match x {
                    ' ' | '\t' | '\r' | '\n' | ')' => break,
                    '"' => return err_parse("unexpected \" in string", state),
                    x => s.push(x),
                }
            }
            None => {
                break;
            }
        }
        try!(state.next())
    }
    //println!("Found string {}", s);
    Ok(Atom::S(s))
}

fn parse_number(state: &mut ParseState) -> ERes<Atom> {
    //println!("number");
    let mut s = String::new();
    loop {
        match state.peek_option() {
            Some(x) => {
                match x {
                    ' ' | '\r' | '\n' | '\t' | ')' => {
                        break
                    },
                    '0' ... '9' | '.' | '-' => {
                        s.push(state.vec[state.pos])
                    },
                    _ => {
                        return err_parse("unexpected char in number", state)
                    },
                }
            }
            None => {
                break
            }
        }
        try!(state.next())
    }
    //println!("Found number {}", s);
    let s2: &str = &s[..];
    if s.contains('.') {
        Ok(Atom::F(f64::from_str(s2).unwrap()))
    } else {
        Ok(Atom::I(i64::from_str(s2).unwrap()))
    }
}

fn parse_atom(state: &mut ParseState) -> ERes<Sexp> {
    //println!("atom");
    let a = match try!(state.peek()) {
        '"' => {
            try!(parse_quoted_string(state))
        }
        '0' ... '9' | '.' | '-' => {
            try!(parse_number(state))
        }
        _ => {
            try!(parse_string(state))
        }
    };
    Ok(Sexp::Atom(a))
}

fn parse_sexp(state: &mut ParseState) -> ERes<Sexp> {
    loop {
        match try!(state.peek()) {
            '(' => {
                return parse_list(state)
            }
            ' ' | '\t' | '\r' | '\n' => {
                try!(state.next());
            }
            ')' => {
                return err_parse("unmatched )", state)
            }
            _ => {
                return parse_atom(state)
            }
        }
    }
}

fn parse(data: &str) -> ERes<Sexp> {
    if data.len() == 0 {
        Ok(Sexp::Empty)
    } else {
        let vec: Vec<char> = data.chars().collect();
        let len = vec.len();
        let state = &mut ParseState { pos: 0, line: 1, lpos: 0, vec: vec, len: len };
        parse_sexp(state)
    }
}

fn read_file(name: &str) -> Result<String, std::io::Error> {
    let mut f = try!(File::open(name));
    let mut s = String::new();
    try!(f.read_to_string(&mut s));
    Ok(s)
}

pub fn parse_str(s: &str) -> ERes<Sexp> {
    parse(s)
}

pub fn parse_file(name: &str) -> ERes<Sexp> {
    let s = try!(match read_file(name) {
        Ok(s) => Ok(s),
        Err(x) => err_io(x),
    });
    parse(&s[..])
}

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

    #[allow(dead_code)]
    fn check_parse(s: &str) {
        let e = parse_str(s).unwrap();
        let t = format!("{}", e);
        assert_eq!(s, t);
    }

    #[allow(dead_code)]
    fn parse_fail(s: &str) {
        parse_str(s).unwrap();
    }

    #[test]
    fn test_empty() { check_parse("") }

    #[test]
    fn test_minimal() { check_parse("()") }

    #[test]
    fn test_string() { check_parse("hello") }

    #[test]
    fn test_qstring() { check_parse("\"hello\"") }

    #[test]
    fn test_number() { check_parse("1.3") }

    #[test]
    fn test_float_vs_int() { check_parse("2.0") }

    #[test]
    fn test_double() { check_parse("(())") }

    #[test]
    fn test_br_string() { check_parse("(world)") }

    #[test]
    fn test_br_qstring() { check_parse("(\"world\")") }

    #[test]
    fn test_br_int() { check_parse("(42)") }

    #[test]
    fn test_br_float() { check_parse("(12.7)") }

    #[test]
    fn test_br_qbrstring() { check_parse("(\"(()\")") }

    #[test]
    #[should_panic(expected="Parse Error 1:1: end of document reached")]
    fn test_invalid1() { parse_fail("(") }

    #[test]
    #[should_panic(expected="Parse Error 1:0: unmatched )")]
    fn test_invalid2() { parse_fail(")") }

    #[test]
    #[should_panic(expected="Parse Error 1:6: end of document reached")]
    fn test_invalid3() { parse_fail("\"hello") }

    #[test]
    fn test_complex() { check_parse("(module SWITCH_3W_SIDE_MMP221-R (layer F.Cu) (descr \"\") (pad 1 thru_hole rect (size 1.2 1.2) (at -2.5 -1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 2 thru_hole rect (size 1.2 1.2) (at 0.0 -1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 3 thru_hole rect (size 1.2 1.2) (at 2.5 -1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 5 thru_hole rect (size 1.2 1.2) (at 0.0 1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 6 thru_hole rect (size 1.2 1.2) (at -2.5 1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 4 thru_hole rect (size 1.2 1.2) (at 2.5 1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (fp_line (start -4.5 -1.75) (end 4.5 -1.75) (layer F.SilkS) (width 0.127)) (fp_line (start 4.5 -1.75) (end 4.5 1.75) (layer F.SilkS) (width 0.127)) (fp_line (start 4.5 1.75) (end -4.5 1.75) (layer F.SilkS) (width 0.127)) (fp_line (start -4.5 1.75) (end -4.5 -1.75) (layer F.SilkS) (width 0.127)))") }
}
<commit_msg>atom as_string<commit_after>// (c) 2015 Joost Yervante Damad <joost@damad.be>

// loosely based on https://github.com/cgaebel/sexp
// latest version can be found at https://github.com/andete/rust_sexp

use std::fmt;
use std::str::FromStr;
use std::f64;

use std::fs::File;
use std::io::prelude::*;

pub enum Atom {
  S(String),
  Q(String),
  I(i64),
  F(f64),
}

impl Atom {
    pub fn f(&self) -> Result<f64,String> {
        match *self {
            Atom::F(f) => Ok(f),
            ref x => Err(format!("not a float: {}", x))
        }
    }
    pub fn i(&self) -> Result<i64,String> {
        match *self {
            Atom::I(i) => Ok(i),
            ref x => Err(format!("not an int: {}", x))
        }
    }
    pub fn string(&self) -> Result<String,String> {
        match *self {
            Atom::S(ref s) => Ok(s.clone()),
            Atom::Q(ref s) => Ok(s.clone()),
            ref x => Err(format!("not a string: {}", x))
        }
    }
    pub fn as_string(&self) -> Result<String,String> {
        match *self {
            Atom::S(ref s) => Ok(s.clone()),
            Atom::Q(ref s) => Ok(s.clone()),
            Atom::F(ref s) => Ok(format!("{}", s)),
            Atom::I(ref s) => Ok(format!("{}", s)),
        }
    }
}

pub enum Sexp {
  Atom(Atom),
  Empty,
  List(Vec<Sexp>),
}

impl Sexp {
    pub fn atom(&self) -> Result<&Atom,String> {
        match *self {
            Sexp::Atom(ref f) => Ok(f),
            ref x => Err(format!("not an atom: {}", x))
        }
    }
    pub fn list(&self) -> Result<&Vec<Sexp>,String> {
        match *self {
            Sexp::List(ref v) => Ok(v),
            ref x => Err(format!("not a list: {}", x))
        }
    }
}

pub struct ParseError {
    msg: &'static str,
    line: usize,
    lpos: usize,
}

pub enum Error {
    ParseError(ParseError),
    IOError(std::io::Error),
}

pub type ERes<T> = Result<T, Error>;

fn err_parse<T>(msg: &'static str, state: &ParseState) -> ERes<T> {
    Err(Error::ParseError(ParseError { msg: msg, line: state.line, lpos: state.lpos}))
}

fn err_io<T>(err: std::io::Error) -> ERes<T> {
    Err(Error::IOError(err))
}

struct ParseState {
    pos: usize,
    line: usize,
    lpos: usize,
    vec: Vec<char>,
    len: usize,
}

impl ParseState {
    fn peek(&self) -> ERes<char> {
        if self.pos >= self.len {
            return err_parse("end of document reached", self)
        }
        Ok(self.vec[self.pos])
    }
    fn peek_option(&self) -> Option<char> {
        if self.pos >= self.len {
            return None
        }
        Some(self.vec[self.pos])
    }
    fn next(&mut self) -> ERes<()> {
        if self.pos >= self.len {
            return err_parse("end of document reached", self)
        }
        let c = self.vec[self.pos];
        self.pos += 1;
        match c {
            '\r' | '\n' => {
                self.lpos = 0;
                self.line += 1;
            }
            _ => {
                self.lpos += 1;
            }
        }
        Ok(())
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Error::ParseError(ref p) =>
                write!(f, "Parse Error {}:{}: {}", p.line, p.lpos, p.msg),
            Error::IOError(ref i) =>
                write!(f, "{}", i)
        }
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Error::ParseError(ref p) =>
                write!(f, "Parse Error {}:{}: {}", p.line, p.lpos, p.msg),
            Error::IOError(ref i) =>
                write!(f, "{}", i)
        }
    }
}

impl fmt::Display for Atom {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    match *self {
        Atom::S(ref s) => write!(f, "{}", s),
        Atom::Q(ref s) => write!(f, "\"{}\"", s),
        Atom::I(i)     => write!(f, "{}", i),
        Atom::F(d)     => {
            let z = d.floor();
            if d - z < f64::EPSILON {
                write!(f, "{}.0", z)
            } else {
                write!(f, "{}", d)
            }
        }
    }
  }
}

impl fmt::Display for Sexp {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    match *self {
      Sexp::Atom(ref a) => write!(f, "{}", a),
      Sexp::List(ref xs) => {
        try!(write!(f, "("));
        for (i, x) in xs.iter().enumerate() {
          let s = if i == 0 { "" } else { " " };
          try!(write!(f, "{}{}", s, x));
        }
        write!(f, ")")
      },
      Sexp::Empty => write!(f, ""),
    }
  }
}

fn parse_list(state: &mut ParseState) -> ERes<Sexp> {
    //println!("list");
    try!(state.next()); // skip (
    let mut l: Vec<Sexp> = Vec::new();
    loop {
        match try!(state.peek()) {
            ')' => {
                try!(state.next());
                break;
            }
            _ => {
                l.push(try!(parse_sexp(state)));
            }
        }
    }
    //println!("Found list");
    Ok(Sexp::List(l))
}

fn parse_quoted_string(state: &mut ParseState) -> ERes<Atom> {
    //println!("qstring");
    try!(state.next()); // skip "
    let mut s = String::new();
    loop {
        match try!(state.peek()) {
            '"' => {
                try!(state.next());
                break
            }
            x @ '\r' | x @ '\n' => {
                s.push(x);
                try!(state.next());
                }
            x => {
                s.push(x);
                try!(state.next());
            }
        }
    }
    //println!("Found quoted string {}", s);
    Ok(Atom::Q(s))
}

fn parse_string(state: &mut ParseState) -> ERes<Atom> {
    //println!("string");
    let mut s = String::new();
    loop {
        match state.peek_option() {
            Some(x) => {
                match x {
                    ' ' | '\t' | '\r' | '\n' | ')' => break,
                    '"' => return err_parse("unexpected \" in string", state),
                    x => s.push(x),
                }
            }
            None => {
                break;
            }
        }
        try!(state.next())
    }
    //println!("Found string {}", s);
    Ok(Atom::S(s))
}

fn parse_number(state: &mut ParseState) -> ERes<Atom> {
    //println!("number");
    let mut s = String::new();
    loop {
        match state.peek_option() {
            Some(x) => {
                match x {
                    ' ' | '\r' | '\n' | '\t' | ')' => {
                        break
                    },
                    '0' ... '9' | '.' | '-' => {
                        s.push(state.vec[state.pos])
                    },
                    _ => {
                        return err_parse("unexpected char in number", state)
                    },
                }
            }
            None => {
                break
            }
        }
        try!(state.next())
    }
    //println!("Found number {}", s);
    let s2: &str = &s[..];
    if s.contains('.') {
        Ok(Atom::F(f64::from_str(s2).unwrap()))
    } else {
        Ok(Atom::I(i64::from_str(s2).unwrap()))
    }
}

fn parse_atom(state: &mut ParseState) -> ERes<Sexp> {
    //println!("atom");
    let a = match try!(state.peek()) {
        '"' => {
            try!(parse_quoted_string(state))
        }
        '0' ... '9' | '.' | '-' => {
            try!(parse_number(state))
        }
        _ => {
            try!(parse_string(state))
        }
    };
    Ok(Sexp::Atom(a))
}

fn parse_sexp(state: &mut ParseState) -> ERes<Sexp> {
    loop {
        match try!(state.peek()) {
            '(' => {
                return parse_list(state)
            }
            ' ' | '\t' | '\r' | '\n' => {
                try!(state.next());
            }
            ')' => {
                return err_parse("unmatched )", state)
            }
            _ => {
                return parse_atom(state)
            }
        }
    }
}

fn parse(data: &str) -> ERes<Sexp> {
    if data.len() == 0 {
        Ok(Sexp::Empty)
    } else {
        let vec: Vec<char> = data.chars().collect();
        let len = vec.len();
        let state = &mut ParseState { pos: 0, line: 1, lpos: 0, vec: vec, len: len };
        parse_sexp(state)
    }
}

fn read_file(name: &str) -> Result<String, std::io::Error> {
    let mut f = try!(File::open(name));
    let mut s = String::new();
    try!(f.read_to_string(&mut s));
    Ok(s)
}

pub fn parse_str(s: &str) -> ERes<Sexp> {
    parse(s)
}

pub fn parse_file(name: &str) -> ERes<Sexp> {
    let s = try!(match read_file(name) {
        Ok(s) => Ok(s),
        Err(x) => err_io(x),
    });
    parse(&s[..])
}

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

    #[allow(dead_code)]
    fn check_parse(s: &str) {
        let e = parse_str(s).unwrap();
        let t = format!("{}", e);
        assert_eq!(s, t);
    }

    #[allow(dead_code)]
    fn parse_fail(s: &str) {
        parse_str(s).unwrap();
    }

    #[test]
    fn test_empty() { check_parse("") }

    #[test]
    fn test_minimal() { check_parse("()") }

    #[test]
    fn test_string() { check_parse("hello") }

    #[test]
    fn test_qstring() { check_parse("\"hello\"") }

    #[test]
    fn test_number() { check_parse("1.3") }

    #[test]
    fn test_float_vs_int() { check_parse("2.0") }

    #[test]
    fn test_double() { check_parse("(())") }

    #[test]
    fn test_br_string() { check_parse("(world)") }

    #[test]
    fn test_br_qstring() { check_parse("(\"world\")") }

    #[test]
    fn test_br_int() { check_parse("(42)") }

    #[test]
    fn test_br_float() { check_parse("(12.7)") }

    #[test]
    fn test_br_qbrstring() { check_parse("(\"(()\")") }

    #[test]
    #[should_panic(expected="Parse Error 1:1: end of document reached")]
    fn test_invalid1() { parse_fail("(") }

    #[test]
    #[should_panic(expected="Parse Error 1:0: unmatched )")]
    fn test_invalid2() { parse_fail(")") }

    #[test]
    #[should_panic(expected="Parse Error 1:6: end of document reached")]
    fn test_invalid3() { parse_fail("\"hello") }

    #[test]
    fn test_complex() { check_parse("(module SWITCH_3W_SIDE_MMP221-R (layer F.Cu) (descr \"\") (pad 1 thru_hole rect (size 1.2 1.2) (at -2.5 -1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 2 thru_hole rect (size 1.2 1.2) (at 0.0 -1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 3 thru_hole rect (size 1.2 1.2) (at 2.5 -1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 5 thru_hole rect (size 1.2 1.2) (at 0.0 1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 6 thru_hole rect (size 1.2 1.2) (at -2.5 1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 4 thru_hole rect (size 1.2 1.2) (at 2.5 1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (fp_line (start -4.5 -1.75) (end 4.5 -1.75) (layer F.SilkS) (width 0.127)) (fp_line (start 4.5 -1.75) (end 4.5 1.75) (layer F.SilkS) (width 0.127)) (fp_line (start 4.5 1.75) (end -4.5 1.75) (layer F.SilkS) (width 0.127)) (fp_line (start -4.5 1.75) (end -4.5 -1.75) (layer F.SilkS) (width 0.127)))") }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]
#![feature(slicing_syntax)]

use error::{FlacError, FlacResult};
use decode::{FrameReader};
use metadata::{MetadataBlock, MetadataBlockReader, StreamInfo};

mod checksum;
pub mod decode;
pub mod error;
pub mod metadata;

pub struct FlacStream<'r> {
    streaminfo: StreamInfo,
    metadata_blocks: Vec<MetadataBlock>,
    frame_reader: FrameReader<'r>
}

fn read_stream_header(input: &mut Reader) -> FlacResult<()> {
    // A FLAC stream starts with a 32-bit header 'fLaC' (big endian).
    const HEADER: u32 = 0x66_4c_61_43;
    let header = try!(input.read_be_u32());
    if header != HEADER { return Err(FlacError::InvalidStreamHeader); }
    Ok(())
}

impl<'r> FlacStream<'r> {
    /// Constructs a flac stream from the given input.
    ///
    /// This will read all metadata and stop at the first audio frame.
    pub fn new<R>(input: &mut R) -> FlacResult<FlacStream> where R: Reader {
        // A flac stream first of all starts with a stream header.
        try!(read_stream_header(input));

        // Start a new scope, because the input reader must be available again
        // for the frame reader next.
        let (streaminfo, metadata_blocks) = {
            // Next are one or more metadata blocks. The flac specification
            // dictates that the streaminfo block is the first block. The metadata
            // block reader will yield at least one element, so the unwrap is safe.
            let mut metadata_iter = MetadataBlockReader::new(input);
            let streaminfo_block = try!(metadata_iter.next().unwrap());
            let streaminfo = match streaminfo_block {
                MetadataBlock::StreamInfo(info) => info,
                _ => return Err(FlacError::MissingStreamInfoBlock)
            };

            // There might be more metadata blocks, read and store them.
            let mut metadata_blocks = Vec::new();
            for block_result in metadata_iter {
                match block_result {
                    Err(error) => return Err(error),
                    Ok(block) => metadata_blocks.push(block)
                }
            }

            (streaminfo, metadata_blocks)
        };

        // The flac stream will contain the reader that will read frames.
        let flac_stream = FlacStream {
            streaminfo: streaminfo,
            metadata_blocks: metadata_blocks,
            frame_reader: FrameReader::new(input)
        };

        Ok(flac_stream)
    }

    /// Returns the streaminfo metadata.
    pub fn streaminfo(&self) -> &StreamInfo {
        &self.streaminfo
    }

    /// Returns an iterator that decodes frames while iterating.
    pub fn frames(&'r self) -> &'r FrameReader<'r> {
        &self.frame_reader
    }
}

#[test]
fn test_open_stream() {
    use std::io::File;
    let mut input = File::open(&Path::new("foo.flac")).unwrap();
    let flac_stream = FlacStream::new(&mut input).unwrap();
}
<commit_msg>test reading frame header<commit_after>#![allow(dead_code)]
#![feature(slicing_syntax)]

use error::{FlacError, FlacResult};
use decode::{FrameReader};
use metadata::{MetadataBlock, MetadataBlockReader, StreamInfo};

mod checksum;
pub mod decode;
pub mod error;
pub mod metadata;

pub struct FlacStream<'r> {
    streaminfo: StreamInfo,
    metadata_blocks: Vec<MetadataBlock>,
    frame_reader: FrameReader<'r>
}

fn read_stream_header(input: &mut Reader) -> FlacResult<()> {
    // A FLAC stream starts with a 32-bit header 'fLaC' (big endian).
    const HEADER: u32 = 0x66_4c_61_43;
    let header = try!(input.read_be_u32());
    if header != HEADER { return Err(FlacError::InvalidStreamHeader); }
    Ok(())
}

impl<'r> FlacStream<'r> {
    /// Constructs a flac stream from the given input.
    ///
    /// This will read all metadata and stop at the first audio frame.
    pub fn new<R>(input: &mut R) -> FlacResult<FlacStream> where R: Reader {
        // A flac stream first of all starts with a stream header.
        try!(read_stream_header(input));

        // Start a new scope, because the input reader must be available again
        // for the frame reader next.
        let (streaminfo, metadata_blocks) = {
            // Next are one or more metadata blocks. The flac specification
            // dictates that the streaminfo block is the first block. The metadata
            // block reader will yield at least one element, so the unwrap is safe.
            let mut metadata_iter = MetadataBlockReader::new(input);
            let streaminfo_block = try!(metadata_iter.next().unwrap());
            let streaminfo = match streaminfo_block {
                MetadataBlock::StreamInfo(info) => info,
                _ => return Err(FlacError::MissingStreamInfoBlock)
            };

            // There might be more metadata blocks, read and store them.
            let mut metadata_blocks = Vec::new();
            for block_result in metadata_iter {
                match block_result {
                    Err(error) => return Err(error),
                    Ok(block) => metadata_blocks.push(block)
                }
            }

            (streaminfo, metadata_blocks)
        };

        // The flac stream will contain the reader that will read frames.
        let flac_stream = FlacStream {
            streaminfo: streaminfo,
            metadata_blocks: metadata_blocks,
            frame_reader: FrameReader::new(input)
        };

        Ok(flac_stream)
    }

    /// Returns the streaminfo metadata.
    pub fn streaminfo(&self) -> &StreamInfo {
        &self.streaminfo
    }

    /// Returns an iterator that decodes frames while iterating.
    pub fn frames(&'r mut self) -> &'r mut FrameReader<'r> {
        &mut self.frame_reader
    }
}

#[test]
fn test_open_stream() {
    use std::io::File;
    let mut input = File::open(&Path::new("foo.flac")).unwrap();
    let mut stream = FlacStream::new(&mut input).unwrap();
    let frame = stream.frames().next().unwrap();
}
<|endoftext|>


<commit_before>//! Serenity is a Rust library for the Discord API.
//!
//! View the [examples] on how to make and structure a bot.
//!
//! Serenity supports bot user authentication via the use of [`Client::builder`].
//!
//! Once logged in, you may add handlers to your client to dispatch [`Event`]s,
//! such as [`EventHandler::message`]. This will cause your handler to be called
//! when a [`Event::MessageCreate`] is received. Each handler is given a
//! [`Context`], giving information about the event. See the
//! [client's module-level documentation].
//!
//! The [`Shard`] is transparently handled by the library, removing
//! unnecessary complexity. Sharded connections are automatically handled for
//! you. See the [gateway's documentation][gateway docs] for more information.
//!
//! A [`Cache`] is also provided for you. This will be updated automatically for
//! you as data is received from the Discord API via events. When calling a
//! method on a [`Context`], the cache will first be searched for relevant data
//! to avoid unnecessary HTTP requests to the Discord API. For more information,
//! see the [cache's module-level documentation][cache docs].
//!
//! Note that, although this documentation will try to be as up-to-date and
//! accurate as possible, Discord hosts [official documentation][docs]. If you
//! need to be sure that some information piece is sanctioned by Discord, refer
//! to their own documentation.
//!
//! ### Full Examples
//!
//! Full examples, detailing and explaining usage of the basic functionality of the
//! library, can be found in the [`examples`] directory.
//!
//! # Installation
//!
//! Add the following to your `Cargo.toml` file:
//!
//! ```toml
//! [dependencies]
//! serenity = "0.11"
//! ```
//!
//! [`Context`]: crate::client::Context
//! [`EventHandler::message`]: crate::client::EventHandler::message
//! [`Event`]: crate::model::event::Event
//! [`Event::MessageCreate`]: crate::model::event::Event::MessageCreate
//! [`Shard`]: crate::gateway::Shard
//! [`examples`]: https://github.com/serenity-rs/serenity/blob/current/examples
//! [cache docs]: crate::cache
//! [client's module-level documentation]: crate::client
//! [docs]: https://discord.com/developers/docs/intro
//! [examples]: https://github.com/serenity-rs/serenity/tree/current/examples
//! [gateway docs]: crate::gateway
#![doc(html_root_url = "https://docs.rs/serenity/*")]
#![cfg_attr(docsrs, feature(doc_auto_cfg))]
#![deny(rust_2018_idioms)]
#![warn(clippy::pedantic)]
#![allow(
    // Allowed to avoid breaking changes.
    clippy::module_name_repetitions,
    clippy::needless_pass_by_value,
    clippy::struct_excessive_bools,
    clippy::ref_option_ref,
    // Allowed as they are too pedantic
    clippy::cast_possible_truncation,
    clippy::unreadable_literal,
    clippy::cast_possible_wrap,
    clippy::wildcard_imports,
    clippy::cast_sign_loss,
    clippy::too_many_lines,
    clippy::doc_markdown,
    clippy::cast_lossless,
    // Covered by other lints
    clippy::missing_panics_doc, // clippy::unwrap_used
)]
#![deny(rustdoc::broken_intra_doc_links, rustdoc::private_intra_doc_links)]
#![deny(
    clippy::unwrap_used,
    clippy::non_ascii_literal,
    clippy::missing_errors_doc,
    clippy::let_underscore_must_use,
    clippy::unused_async,
    unused
)]
#![cfg_attr(test, allow(clippy::unwrap_used))]
#![type_length_limit = "3294819"] // needed so ShardRunner::run compiles with instrument.

#[macro_use]
extern crate serde;

#[macro_use]
mod internal;

pub mod constants;
pub mod json;
pub mod model;
pub mod prelude;

#[cfg(feature = "builder")]
pub mod builder;
#[cfg(feature = "cache")]
pub mod cache;
#[cfg(feature = "client")]
pub mod client;
#[cfg(feature = "collector")]
pub mod collector;
#[cfg(feature = "framework")]
pub mod framework;
#[cfg(feature = "gateway")]
pub mod gateway;
#[cfg(feature = "http")]
pub mod http;
#[cfg(feature = "utils")]
pub mod utils;

mod error;

#[cfg(feature = "client")]
use std::sync::Arc;

#[cfg(all(feature = "client", feature = "cache"))]
use crate::cache::Cache;
#[cfg(all(feature = "client", feature = "gateway"))]
pub use crate::client::Client;
pub use crate::error::{Error, Result};
#[cfg(feature = "client")]
use crate::http::Http;

#[cfg(feature = "client")]
#[non_exhaustive]
pub struct CacheAndHttp {
    #[cfg(feature = "cache")]
    pub cache: Arc<Cache>,
    pub http: Arc<Http>,
}

// For the procedural macros in `command_attr`.
pub use async_trait::async_trait;
pub use futures;
pub use futures::future::FutureExt;
#[cfg(feature = "standard_framework")]
#[doc(hidden)]
pub use static_assertions;
<commit_msg>Implement `AsRef<Cache>` and `AsRef<Http>` for `CacheAndHttp` (#2114)<commit_after>//! Serenity is a Rust library for the Discord API.
//!
//! View the [examples] on how to make and structure a bot.
//!
//! Serenity supports bot user authentication via the use of [`Client::builder`].
//!
//! Once logged in, you may add handlers to your client to dispatch [`Event`]s,
//! such as [`EventHandler::message`]. This will cause your handler to be called
//! when a [`Event::MessageCreate`] is received. Each handler is given a
//! [`Context`], giving information about the event. See the
//! [client's module-level documentation].
//!
//! The [`Shard`] is transparently handled by the library, removing
//! unnecessary complexity. Sharded connections are automatically handled for
//! you. See the [gateway's documentation][gateway docs] for more information.
//!
//! A [`Cache`] is also provided for you. This will be updated automatically for
//! you as data is received from the Discord API via events. When calling a
//! method on a [`Context`], the cache will first be searched for relevant data
//! to avoid unnecessary HTTP requests to the Discord API. For more information,
//! see the [cache's module-level documentation][cache docs].
//!
//! Note that, although this documentation will try to be as up-to-date and
//! accurate as possible, Discord hosts [official documentation][docs]. If you
//! need to be sure that some information piece is sanctioned by Discord, refer
//! to their own documentation.
//!
//! ### Full Examples
//!
//! Full examples, detailing and explaining usage of the basic functionality of the
//! library, can be found in the [`examples`] directory.
//!
//! # Installation
//!
//! Add the following to your `Cargo.toml` file:
//!
//! ```toml
//! [dependencies]
//! serenity = "0.11"
//! ```
//!
//! [`Context`]: crate::client::Context
//! [`EventHandler::message`]: crate::client::EventHandler::message
//! [`Event`]: crate::model::event::Event
//! [`Event::MessageCreate`]: crate::model::event::Event::MessageCreate
//! [`Shard`]: crate::gateway::Shard
//! [`examples`]: https://github.com/serenity-rs/serenity/blob/current/examples
//! [cache docs]: crate::cache
//! [client's module-level documentation]: crate::client
//! [docs]: https://discord.com/developers/docs/intro
//! [examples]: https://github.com/serenity-rs/serenity/tree/current/examples
//! [gateway docs]: crate::gateway
#![doc(html_root_url = "https://docs.rs/serenity/*")]
#![cfg_attr(docsrs, feature(doc_auto_cfg))]
#![deny(rust_2018_idioms)]
#![warn(clippy::pedantic)]
#![allow(
    // Allowed to avoid breaking changes.
    clippy::module_name_repetitions,
    clippy::needless_pass_by_value,
    clippy::struct_excessive_bools,
    clippy::ref_option_ref,
    // Allowed as they are too pedantic
    clippy::cast_possible_truncation,
    clippy::unreadable_literal,
    clippy::cast_possible_wrap,
    clippy::wildcard_imports,
    clippy::cast_sign_loss,
    clippy::too_many_lines,
    clippy::doc_markdown,
    clippy::cast_lossless,
    // Covered by other lints
    clippy::missing_panics_doc, // clippy::unwrap_used
)]
#![deny(rustdoc::broken_intra_doc_links, rustdoc::private_intra_doc_links)]
#![deny(
    clippy::unwrap_used,
    clippy::non_ascii_literal,
    clippy::missing_errors_doc,
    clippy::let_underscore_must_use,
    clippy::unused_async,
    unused
)]
#![cfg_attr(test, allow(clippy::unwrap_used))]
#![type_length_limit = "3294819"] // needed so ShardRunner::run compiles with instrument.

#[macro_use]
extern crate serde;

#[macro_use]
mod internal;

pub mod constants;
pub mod json;
pub mod model;
pub mod prelude;

#[cfg(feature = "builder")]
pub mod builder;
#[cfg(feature = "cache")]
pub mod cache;
#[cfg(feature = "client")]
pub mod client;
#[cfg(feature = "collector")]
pub mod collector;
#[cfg(feature = "framework")]
pub mod framework;
#[cfg(feature = "gateway")]
pub mod gateway;
#[cfg(feature = "http")]
pub mod http;
#[cfg(feature = "utils")]
pub mod utils;

mod error;

#[cfg(feature = "client")]
use std::sync::Arc;

#[cfg(all(feature = "client", feature = "cache"))]
use crate::cache::Cache;
#[cfg(all(feature = "client", feature = "gateway"))]
pub use crate::client::Client;
pub use crate::error::{Error, Result};
#[cfg(feature = "client")]
use crate::http::Http;

#[cfg(feature = "client")]
#[non_exhaustive]
pub struct CacheAndHttp {
    #[cfg(feature = "cache")]
    pub cache: Arc<Cache>,
    pub http: Arc<Http>,
}

#[cfg(all(feature = "client", feature = "cache"))]
impl AsRef<Cache> for CacheAndHttp {
    fn as_ref(&self) -> &Cache {
        &self.cache
    }
}

#[cfg(feature = "client")]
impl AsRef<Http> for CacheAndHttp {
    fn as_ref(&self) -> &Http {
        &self.http
    }
}

// For the procedural macros in `command_attr`.
pub use async_trait::async_trait;
pub use futures;
pub use futures::future::FutureExt;
#[cfg(feature = "standard_framework")]
#[doc(hidden)]
pub use static_assertions;
<|endoftext|>


<commit_before>//! Lua functionality

use hlua;
use hlua::Lua;
use std::thread;
use std::time::Duration;
use std::fs::{File};
use std::io::Write;
use std::sync::Mutex;
use std::sync::mpsc::{channel, Sender, Receiver};

lazy_static! {
    static ref SENDER: Mutex<Sender<LuaQuery>> = {
        let (tx, rx) = channel::<LuaQuery>();
        Mutex::new(tx)
    };
    static ref RECEIVER: Mutex<Receiver<LuaResponse>> = {
        let (tx, rx) = channel::<LuaResponse>();
        Mutex::new(rx)
    };
}

/// Messages sent to the lua thread
pub enum LuaQuery {
    CallFunction(String),
    GetVariable(String)
}

/// Messages received from lua thread
pub enum LuaResponse {
    Error(hlua::LuaError),
    Function(hlua::functions_read::LuaFunction<String>)
}

unsafe impl Send for LuaQuery { }
unsafe impl Send for LuaResponse { }
unsafe impl Sync for LuaQuery { }
unsafe impl Sync for LuaResponse { }

pub fn init() {
    print!("[lua] Initializing...");
    let (query_tx, query_rx) = channel::<LuaQuery>();
    let (answer_tx, answer_rx) = channel::<LuaResponse>();
    {
        let mut sender = SENDER.lock().unwrap();
        let mut receiver = RECEIVER.lock().unwrap();

        *sender = query_tx;
        *receiver = answer_rx;
    }

    thread::spawn(move || {
        println!("[lua] Inside thread!");
        let receiver = query_rx;
        let sender = answer_tx;
        let mut LUA = Lua::new();
        print!("[lua] Loading libs...");
        LUA.openlibs();
        println!(" done!");
        println!("[lua] Creating init file");
        let mut file = File::create("/tmp/init.lua").unwrap();
        file.write(b"print('Hello world!')").unwrap();
        println!("[lua] Created hello world file!");
        println!("[lua] Executing init.lua...");
        LUA.execute_from_reader::<(), File>(File::open("/tmp/init.lua").unwrap());
        println!("[lua] Done!");
        thread::sleep(Duration::from_secs(10));
        println!("[lua] Exiting thread...");
    });
    println!(" created thread.")
}
<commit_msg>Add variants to LuaQuery to allow better communication<commit_after>//! Lua functionality

use hlua;
use hlua::Lua;
use std::thread;
use std::time::Duration;
use std::fs::{File};
use std::path::Path;
use std::io::Write;
use std::sync::Mutex;
use std::sync::mpsc::{channel, Sender, Receiver};

lazy_static! {
    static ref SENDER: Mutex<Sender<LuaQuery>> = {
        let (tx, rx) = channel::<LuaQuery>();
        Mutex::new(tx)
    };
    static ref RECEIVER: Mutex<Receiver<LuaResponse>> = {
        let (tx, rx) = channel::<LuaResponse>();
        Mutex::new(rx)
    };
}

/// Messages sent to the lua thread
pub enum LuaQuery {
    /// Execute a string
    Execute(String),
    /// Execute a file
    ExecuteFile(String),
    /// Get a variable
    GetVariable(String),
    /// Set a value
    SetValue {
        name: Box<::std::borrow::Borrow<str> + Sized>,
        val: Box<hlua::Push<&'static mut Lua<'static>> + Sized>
    },
    /// Create a new array
    EmptyArray(String),
}

/// Messages received from lua thread
pub enum LuaResponse {
    Error(hlua::LuaError),
    Function(hlua::functions_read::LuaFunction<String>)
}

unsafe impl Send for LuaQuery { }
unsafe impl Send for LuaResponse { }
unsafe impl Sync for LuaQuery { }
unsafe impl Sync for LuaResponse { }

pub fn init() {
    print!("[lua] Initializing...");
    let (query_tx, query_rx) = channel::<LuaQuery>();
    let (answer_tx, answer_rx) = channel::<LuaResponse>();
    {
        let mut sender = SENDER.lock().unwrap();
        let mut receiver = RECEIVER.lock().unwrap();

        *sender = query_tx;
        *receiver = answer_rx;
    }

    thread::spawn(move || {
        println!("[lua] Inside thread!");
        let receiver = query_rx;
        let sender = answer_tx;
        let mut lua = Lua::new();
        print!("[lua] Loading libs...");
        lua.openlibs();
        println!(" done!");
        println!("[lua] Creating init file");
        let mut file = File::create("/tmp/init.lua").unwrap();
        file.write(b"print('Hello world!')").unwrap();
        println!("[lua] Created hello world file!");
        println!("[lua] Executing init.lua...");
        lua.execute_from_reader::<(), File>(File::open("/tmp/init.lua")
                                            .unwrap()).unwrap();
        println!("[lua] Done!");
        thread::sleep(Duration::from_secs(10));
        println!("[lua] Exiting thread...");
    });
    println!(" created thread.")
}
<|endoftext|>


<commit_before>//! Support for binary-format `COPY` query execution with rust-postgres.
//!
//! # Example
//!
//! ```rust,no_run
//! extern crate postgres;
//! extern crate postgres_binary_copy;
//!
//! use postgres::{Connection, SslMode};
//! use postgres::types::{Type, ToSql};
//! use postgres_binary_copy::BinaryCopyReader;
//!
//! fn main() {
//!     let conn = Connection::connect("postgres://postgres@localhost",
//!                                    &SslMode::None).unwrap();
//!
//!     conn.execute("CREATE TABLE foo (id INT PRIMARY KEY, bar VARCHAR)", &[])
//!         .unwrap();
//!
//!     let types = &[Type::Int4, Type::Varchar];
//!     let data: Vec<Box<ToSql>> = vec![Box::new(1i32), Box::new("hello"),
//!                                      Box::new(2i32), Box::new("world")];
//!     let data = data.iter().map(|v| &**v);
//!     let mut reader = BinaryCopyReader::new(types, data);
//!
//!     let stmt = conn.prepare("COPY foo (id, bar) FROM STDIN BINARY").unwrap();
//!     stmt.copy_in(&[], &mut reader).unwrap();
//! }
//! ```
extern crate byteorder;
extern crate postgres;

use byteorder::{BigEndian, WriteBytesExt};
use postgres::error::Error;
use postgres::types::{Type, ToSql, SessionInfo, IsNull, ReadWithInfo};
use std::error;
use std::io::prelude::*;
use std::io::{self, Cursor};

/// Like `Iterator`, except that it returns borrowed values.
///
/// In constrast to `Iterator<Item = &T>`, a type implementing
/// `StreamingIterator<Item = T>` does not need to have all of the values it
/// returns in memory at the same time.
///
/// All `Iterator`s over `&T` are also `StreamingIterator`s over `T`.
pub trait StreamingIterator {
    /// The type of elements being iterated.
    type Item: ?Sized;

    /// Advances the iterator and returns the next value.
    ///
    /// Returns `None` when the end is reached.
    fn next(&mut self) -> Option<&Self::Item>;
}

impl<'a, T: 'a + ?Sized, I: Iterator<Item = &'a T>> StreamingIterator for I {
    type Item = T;

    fn next(&mut self) -> Option<&T> {
        unsafe { std::mem::transmute(Iterator::next(self)) }
    }
}

#[derive(Debug, Copy, Clone)]
enum ReadState {
    Header,
    Body(usize),
    Footer,
}

/// A `ReadWithInfo` implementation that generates binary-formatted output
/// for use with `COPY ... FROM STDIN BINARY` statements.
#[derive(Debug)]
pub struct BinaryCopyReader<'a, I> {
    types: &'a [Type],
    state: ReadState,
    it: I,
    buf: Cursor<Vec<u8>>,
}

impl<'a, I> BinaryCopyReader<'a, I> where I: StreamingIterator<Item = ToSql> {
    /// Creates a new `BinaryCopyReader`.
    ///
    /// The reader will output tuples with a structure described by `types` and
    /// values from `it`. `it` should return values in row-major order.
    pub fn new(types: &'a [Type], it: I) -> BinaryCopyReader<'a, I> {
        let mut buf = vec![];
        let _ = buf.write(b"PGCOPY\n\xff\r\n\0");
        let _ = buf.write_i32::<BigEndian>(0);
        let _ = buf.write_i32::<BigEndian>(0);

        BinaryCopyReader {
            types: types,
            state: ReadState::Header,
            it: it,
            buf: Cursor::new(buf),
        }
    }

    fn fill_buf(&mut self, info: &SessionInfo) -> io::Result<()> {
        enum Op<'a> {
            Value(usize, &'a ToSql),
            Footer,
            Nothing,
        }

        let op = match (self.state, self.it.next()) {
            (ReadState::Header, Some(value)) => {
                self.state = ReadState::Body(0);
                Op::Value(0, value)
            }
            (ReadState::Body(old_idx), Some(value)) => {
                let idx = (old_idx + 1) % self.types.len();
                self.state = ReadState::Body(idx);
                Op::Value(idx, value)
            }
            (ReadState::Header, None) | (ReadState::Body(_), None) => {
                self.state = ReadState::Footer;
                Op::Footer
            }
            (ReadState::Footer, _) => Op::Nothing,
        };

        self.buf.set_position(0);
        self.buf.get_mut().clear();

        match op {
            Op::Value(idx, value) => {
                if idx == 0 {
                    let len = self.types.len();
                    let len = if len > i16::max_value() as usize {
                        let err: Box<error::Error+Sync+Send> =
                            "value too large to transmit".into();
                        return Err(io::Error::new(io::ErrorKind::InvalidInput,
                                                  Error::Conversion(err)));
                    } else {
                        len as i16
                    };
                    let _ = self.buf.write_i16::<BigEndian>(len);
                }

                let len_pos = self.buf.position();
                let _ = self.buf.write_i32::<BigEndian>(0); // space for length
                let len = match value.to_sql_checked(&self.types[idx], &mut self.buf, info) {
                    Ok(IsNull::Yes) => -1,
                    Ok(IsNull::No) => {
                        let len = self.buf.position() - 4 - len_pos;
                        if len > i32::max_value() as u64 {
                            let err: Box<error::Error+Sync+Send> =
                                "value too large to transmit".into();
                            return Err(io::Error::new(io::ErrorKind::InvalidInput,
                                                      Error::Conversion(err)));
                        } else {
                            len as i32
                        }
                    }
                    Err(e) => return Err(io::Error::new(io::ErrorKind::InvalidInput, e)),
                };
                self.buf.set_position(len_pos);
                let _ = self.buf.write_i32::<BigEndian>(len);
            }
            Op::Footer => {
                let _ = self.buf.write_i16::<BigEndian>(-1);
            }
            Op::Nothing => {}
        }

        self.buf.set_position(0);
        Ok(())
    }
}

impl<'a, I> ReadWithInfo for BinaryCopyReader<'a, I> where I: StreamingIterator<Item = ToSql> {
    fn read_with_info(&mut self, buf: &mut [u8], info: &SessionInfo) -> io::Result<usize> {
        if self.buf.position() == self.buf.get_ref().len() as u64 {
            try!(self.fill_buf(info));
        }
        self.buf.read(buf)
    }
}

#[cfg(test)]
mod test {
    use super::BinaryCopyReader;
    use postgres::{Connection, SslMode};
    use postgres::types::{Type, ToSql};

    #[test]
    fn basic() {
        let conn = Connection::connect("postgres://postgres@localhost", &SslMode::None).unwrap();
        conn.execute("CREATE TEMPORARY TABLE foo (id INT PRIMARY KEY, bar VARCHAR)", &[]).unwrap();

        let stmt = conn.prepare("COPY foo (id, bar) FROM STDIN BINARY").unwrap();

        let types = &[Type::Int4, Type::Varchar];
        let values: Vec<Box<ToSql>> = vec![Box::new(1i32), Box::new("foobar"),
                                           Box::new(2i32), Box::new(None::<String>)];
        let values = values.iter().map(|e| &**e);
        let mut reader = BinaryCopyReader::new(types, values);

        stmt.copy_in(&[], &mut reader).unwrap();

        let stmt = conn.prepare("SELECT id, bar FROM foo ORDER BY id").unwrap();
        assert_eq!(vec![(1i32, Some("foobar".to_string())), (2i32, None)],
                   stmt.query(&[])
                        .unwrap()
                        .into_iter()
                        .map(|r| (r.get(0), r.get(1)))
                        .collect::<Vec<(i32, Option<String>)>>());
    }

    #[test]
    fn many_rows() {
        let conn = Connection::connect("postgres://postgres@localhost", &SslMode::None).unwrap();
        conn.execute("CREATE TEMPORARY TABLE foo (id INT PRIMARY KEY, bar VARCHAR)", &[]).unwrap();

        let stmt = conn.prepare("COPY foo (id, bar) FROM STDIN BINARY").unwrap();

        let types = &[Type::Int4, Type::Varchar];
        let mut values: Vec<Box<ToSql>> = vec![];
        for i in 0..10000i32 {
            values.push(Box::new(i));
            values.push(Box::new(format!("the value for {}", i)));
        }

        let values = values.iter().map(|e| &**e);
        let mut reader = BinaryCopyReader::new(types, values);

        stmt.copy_in(&[], &mut reader).unwrap();

        let stmt = conn.prepare("SELECT id, bar FROM foo ORDER BY id").unwrap();
        let result = stmt.query(&[]).unwrap();
        assert_eq!(10000, result.len());
        for (i, row) in result.into_iter().enumerate() {
            assert_eq!(i as i32, row.get(0));
            assert_eq!(format!("the value for {}", i), row.get::<_, String>(1));
        }
    }

    #[test]
    fn big_rows() {
        let conn = Connection::connect("postgres://postgres@localhost", &SslMode::None).unwrap();
        conn.execute("CREATE TEMPORARY TABLE foo (id INT PRIMARY KEY, bar BYTEA)", &[]).unwrap();

        let stmt = conn.prepare("COPY foo (id, bar) FROM STDIN BINARY").unwrap();

        let types = &[Type::Int4, Type::Bytea];
        let mut values: Vec<Box<ToSql>> = vec![];
        for i in 0..2i32 {
            values.push(Box::new(i));
            values.push(Box::new(vec![i as u8; 128 * 1024]));
        }

        let values = values.iter().map(|e| &**e);
        let mut reader = BinaryCopyReader::new(types, values);

        stmt.copy_in(&[], &mut reader).unwrap();

        let stmt = conn.prepare("SELECT id, bar FROM foo ORDER BY id").unwrap();
        let result = stmt.query(&[]).unwrap();
        assert_eq!(2, result.len());
        for (i, row) in result.into_iter().enumerate() {
            assert_eq!(i as i32, row.get(0));
            assert_eq!(vec![i as u8; 128 * 1024], row.get::<_, Vec<u8>>(1));
        }
    }
}
<commit_msg>Doc tweak<commit_after>//! Support for binary-format `COPY` query execution with rust-postgres.
//!
//! # Example
//!
//! ```rust,no_run
//! extern crate postgres;
//! extern crate postgres_binary_copy;
//!
//! use postgres::{Connection, SslMode};
//! use postgres::types::{Type, ToSql};
//! use postgres_binary_copy::BinaryCopyReader;
//!
//! fn main() {
//!     let conn = Connection::connect("postgres://postgres@localhost",
//!                                    &SslMode::None).unwrap();
//!
//!     conn.execute("CREATE TABLE foo (id INT PRIMARY KEY, bar VARCHAR)", &[])
//!         .unwrap();
//!
//!     let types = &[Type::Int4, Type::Varchar];
//!     let data: Vec<Box<ToSql>> = vec![Box::new(1i32), Box::new("hello"),
//!                                      Box::new(2i32), Box::new("world")];
//!     let data = data.iter().map(|v| &**v);
//!     let mut reader = BinaryCopyReader::new(types, data);
//!
//!     let stmt = conn.prepare("COPY foo (id, bar) FROM STDIN BINARY").unwrap();
//!     stmt.copy_in(&[], &mut reader).unwrap();
//! }
//! ```
extern crate byteorder;
extern crate postgres;

use byteorder::{BigEndian, WriteBytesExt};
use postgres::error::Error;
use postgres::types::{Type, ToSql, SessionInfo, IsNull, ReadWithInfo};
use std::error;
use std::io::prelude::*;
use std::io::{self, Cursor};

/// Like `Iterator`, except that it returns borrowed values.
///
/// In constrast to `Iterator<Item = &T>`, a type implementing
/// `StreamingIterator<Item = T>` does not need to have all of the values it
/// returns in memory at the same time.
///
/// All `Iterator`s over `&T` are also `StreamingIterator`s over `T`.
pub trait StreamingIterator {
    /// The type of elements being iterated.
    type Item: ?Sized;

    /// Advances the iterator and returns the next value.
    ///
    /// Returns `None` when the end is reached.
    fn next(&mut self) -> Option<&Self::Item>;
}

impl<'a, T: 'a + ?Sized, I: Iterator<Item = &'a T>> StreamingIterator for I {
    type Item = T;

    fn next(&mut self) -> Option<&T> {
        unsafe { std::mem::transmute(Iterator::next(self)) }
    }
}

#[derive(Debug, Copy, Clone)]
enum ReadState {
    Header,
    Body(usize),
    Footer,
}

/// A `ReadWithInfo` implementation that generates binary-formatted output
/// for use with binary-format `COPY FROM STDIN` statements.
#[derive(Debug)]
pub struct BinaryCopyReader<'a, I> {
    types: &'a [Type],
    state: ReadState,
    it: I,
    buf: Cursor<Vec<u8>>,
}

impl<'a, I> BinaryCopyReader<'a, I> where I: StreamingIterator<Item = ToSql> {
    /// Creates a new `BinaryCopyReader`.
    ///
    /// The reader will output tuples with a structure described by `types` and
    /// values from `it`. `it` should return values in row-major order.
    pub fn new(types: &'a [Type], it: I) -> BinaryCopyReader<'a, I> {
        let mut buf = vec![];
        let _ = buf.write(b"PGCOPY\n\xff\r\n\0");
        let _ = buf.write_i32::<BigEndian>(0);
        let _ = buf.write_i32::<BigEndian>(0);

        BinaryCopyReader {
            types: types,
            state: ReadState::Header,
            it: it,
            buf: Cursor::new(buf),
        }
    }

    fn fill_buf(&mut self, info: &SessionInfo) -> io::Result<()> {
        enum Op<'a> {
            Value(usize, &'a ToSql),
            Footer,
            Nothing,
        }

        let op = match (self.state, self.it.next()) {
            (ReadState::Header, Some(value)) => {
                self.state = ReadState::Body(0);
                Op::Value(0, value)
            }
            (ReadState::Body(old_idx), Some(value)) => {
                let idx = (old_idx + 1) % self.types.len();
                self.state = ReadState::Body(idx);
                Op::Value(idx, value)
            }
            (ReadState::Header, None) | (ReadState::Body(_), None) => {
                self.state = ReadState::Footer;
                Op::Footer
            }
            (ReadState::Footer, _) => Op::Nothing,
        };

        self.buf.set_position(0);
        self.buf.get_mut().clear();

        match op {
            Op::Value(idx, value) => {
                if idx == 0 {
                    let len = self.types.len();
                    let len = if len > i16::max_value() as usize {
                        let err: Box<error::Error+Sync+Send> =
                            "value too large to transmit".into();
                        return Err(io::Error::new(io::ErrorKind::InvalidInput,
                                                  Error::Conversion(err)));
                    } else {
                        len as i16
                    };
                    let _ = self.buf.write_i16::<BigEndian>(len);
                }

                let len_pos = self.buf.position();
                let _ = self.buf.write_i32::<BigEndian>(0); // space for length
                let len = match value.to_sql_checked(&self.types[idx], &mut self.buf, info) {
                    Ok(IsNull::Yes) => -1,
                    Ok(IsNull::No) => {
                        let len = self.buf.position() - 4 - len_pos;
                        if len > i32::max_value() as u64 {
                            let err: Box<error::Error+Sync+Send> =
                                "value too large to transmit".into();
                            return Err(io::Error::new(io::ErrorKind::InvalidInput,
                                                      Error::Conversion(err)));
                        } else {
                            len as i32
                        }
                    }
                    Err(e) => return Err(io::Error::new(io::ErrorKind::InvalidInput, e)),
                };
                self.buf.set_position(len_pos);
                let _ = self.buf.write_i32::<BigEndian>(len);
            }
            Op::Footer => {
                let _ = self.buf.write_i16::<BigEndian>(-1);
            }
            Op::Nothing => {}
        }

        self.buf.set_position(0);
        Ok(())
    }
}

impl<'a, I> ReadWithInfo for BinaryCopyReader<'a, I> where I: StreamingIterator<Item = ToSql> {
    fn read_with_info(&mut self, buf: &mut [u8], info: &SessionInfo) -> io::Result<usize> {
        if self.buf.position() == self.buf.get_ref().len() as u64 {
            try!(self.fill_buf(info));
        }
        self.buf.read(buf)
    }
}

#[cfg(test)]
mod test {
    use super::BinaryCopyReader;
    use postgres::{Connection, SslMode};
    use postgres::types::{Type, ToSql};

    #[test]
    fn basic() {
        let conn = Connection::connect("postgres://postgres@localhost", &SslMode::None).unwrap();
        conn.execute("CREATE TEMPORARY TABLE foo (id INT PRIMARY KEY, bar VARCHAR)", &[]).unwrap();

        let stmt = conn.prepare("COPY foo (id, bar) FROM STDIN BINARY").unwrap();

        let types = &[Type::Int4, Type::Varchar];
        let values: Vec<Box<ToSql>> = vec![Box::new(1i32), Box::new("foobar"),
                                           Box::new(2i32), Box::new(None::<String>)];
        let values = values.iter().map(|e| &**e);
        let mut reader = BinaryCopyReader::new(types, values);

        stmt.copy_in(&[], &mut reader).unwrap();

        let stmt = conn.prepare("SELECT id, bar FROM foo ORDER BY id").unwrap();
        assert_eq!(vec![(1i32, Some("foobar".to_string())), (2i32, None)],
                   stmt.query(&[])
                        .unwrap()
                        .into_iter()
                        .map(|r| (r.get(0), r.get(1)))
                        .collect::<Vec<(i32, Option<String>)>>());
    }

    #[test]
    fn many_rows() {
        let conn = Connection::connect("postgres://postgres@localhost", &SslMode::None).unwrap();
        conn.execute("CREATE TEMPORARY TABLE foo (id INT PRIMARY KEY, bar VARCHAR)", &[]).unwrap();

        let stmt = conn.prepare("COPY foo (id, bar) FROM STDIN BINARY").unwrap();

        let types = &[Type::Int4, Type::Varchar];
        let mut values: Vec<Box<ToSql>> = vec![];
        for i in 0..10000i32 {
            values.push(Box::new(i));
            values.push(Box::new(format!("the value for {}", i)));
        }

        let values = values.iter().map(|e| &**e);
        let mut reader = BinaryCopyReader::new(types, values);

        stmt.copy_in(&[], &mut reader).unwrap();

        let stmt = conn.prepare("SELECT id, bar FROM foo ORDER BY id").unwrap();
        let result = stmt.query(&[]).unwrap();
        assert_eq!(10000, result.len());
        for (i, row) in result.into_iter().enumerate() {
            assert_eq!(i as i32, row.get(0));
            assert_eq!(format!("the value for {}", i), row.get::<_, String>(1));
        }
    }

    #[test]
    fn big_rows() {
        let conn = Connection::connect("postgres://postgres@localhost", &SslMode::None).unwrap();
        conn.execute("CREATE TEMPORARY TABLE foo (id INT PRIMARY KEY, bar BYTEA)", &[]).unwrap();

        let stmt = conn.prepare("COPY foo (id, bar) FROM STDIN BINARY").unwrap();

        let types = &[Type::Int4, Type::Bytea];
        let mut values: Vec<Box<ToSql>> = vec![];
        for i in 0..2i32 {
            values.push(Box::new(i));
            values.push(Box::new(vec![i as u8; 128 * 1024]));
        }

        let values = values.iter().map(|e| &**e);
        let mut reader = BinaryCopyReader::new(types, values);

        stmt.copy_in(&[], &mut reader).unwrap();

        let stmt = conn.prepare("SELECT id, bar FROM foo ORDER BY id").unwrap();
        let result = stmt.query(&[]).unwrap();
        assert_eq!(2, result.len());
        for (i, row) in result.into_iter().enumerate() {
            assert_eq!(i as i32, row.get(0));
            assert_eq!(vec![i as u8; 128 * 1024], row.get::<_, Vec<u8>>(1));
        }
    }
}
<|endoftext|>