<commit_before>//! The Actual RethinkDB Commands

use errors::*;
use conn::{Connection, ConnectOpts};
use serde_json::{self, Value};
use ql2::proto::{Term_TermType as tt, Query_QueryType as qt};
use std::io::Write;
use byteorder::{WriteBytesExt, LittleEndian};
use std::io::Read;
use std::str;
use protobuf::ProtobufEnum;
use byteorder::ReadBytesExt;
use super::session::Client;
use super::{Result, r};
use std::error::Error as StdError;

#[derive(Debug)]
pub struct RootCommand(Result<String>);
pub struct Command;
pub struct Query;

pub trait IntoCommandArg {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)>;
}

impl<'a> IntoCommandArg for &'a str {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        Ok((Some(format!("{:?}", self)), None))
    }
}

impl IntoCommandArg for String {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        Ok((Some(format!("{:?}", self)), None))
    }
}

impl IntoCommandArg for Value {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        // Arrays are a special case: since ReQL commands are sent as arrays
        // We have to wrap them using MAKE_ARRAY
        let val = wrap_arrays(self.clone());
        match serde_json::to_string(&val) {
            Ok(cmd) => Ok((Some(cmd), None)),
            Err(e) => Err(From::from(DriverError::Json(e))),
        }
    }
}

impl<T> IntoCommandArg for (T, Value)
        where T: IntoCommandArg
{
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        if !self.1.is_object() {
            let msg = String::from("Only objects are allowed as function options. You should use `r.object()` to pass optional arguments in your functions.");
            return Err(From::from(DriverError::Other(msg)));
        }
        let arg = try!(self.0.to_arg());
        let opt = try!(self.1.to_arg());
        Ok((arg.0, opt.0))
    }
}

impl IntoCommandArg for RootCommand {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        match self.0 {
            Ok(ref cmd) => Ok((Some(cmd.to_string()), None)),
            Err(ref e) => Err(From::from(DriverError::Other(e.description().to_string()))),
        }
    }
}

macro_rules! define {
    (impl IntoCommandArg for $T:ty) => {
        impl IntoCommandArg for $T {
            fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
                Ok((Some(self.to_string()), None))
            }
        }
    }
}

define!{ impl IntoCommandArg for bool }
define!{ impl IntoCommandArg for char }
define!{ impl IntoCommandArg for u8 }
define!{ impl IntoCommandArg for u16 }
define!{ impl IntoCommandArg for u32 }
define!{ impl IntoCommandArg for u64 }
define!{ impl IntoCommandArg for usize }
define!{ impl IntoCommandArg for i8 }
define!{ impl IntoCommandArg for i16 }
define!{ impl IntoCommandArg for i32 }
define!{ impl IntoCommandArg for i64 }
define!{ impl IntoCommandArg for isize }
define!{ impl IntoCommandArg for f32 }
define!{ impl IntoCommandArg for f64 }

impl Client {
    pub fn connection(&self) -> ConnectOpts {
        Self::config().read().clone()
    }

    pub fn expr<T>(&self, e: T) -> RootCommand
        where T: IntoCommandArg
    {
        match e.to_arg() {
            Ok(arg) => if let Some(arg) = arg.0 {
                RootCommand(Ok(arg))
            } else {
                RootCommand(Ok(String::new()))
            },
            Err(e) => RootCommand(Err(e)),
        }
    }

    pub fn db_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB_CREATE,
                                     Some(name),
                                     None))
    }

    pub fn db_drop<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB_DROP,
                                     Some(name),
                                     None))
    }

    pub fn db<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB,
                                     Some(name),
                                     None))
    }

    pub fn table_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table_create(name)
    }

    pub fn table<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table(name)
    }

    pub fn table_drop<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table_drop(name)
    }

    pub fn object(&self) -> serde_json::builder::ObjectBuilder {
        serde_json::builder::ObjectBuilder::new()
    }

    pub fn array(&self) -> serde_json::builder::ArrayBuilder {
        serde_json::builder::ArrayBuilder::new()
    }
}

impl RootCommand {
    pub fn table_create<T>(self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::TABLE_CREATE,
                                     Some(name),
                                     Some(&commands)))
    }

    pub fn table<T>(self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::TABLE,
                                     Some(name),
                                     Some(&commands)))
    }

    pub fn table_drop<T>(self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::TABLE_DROP,
                                     Some(name),
                                     Some(&commands)))
    }

    pub fn insert<T>(self, data: T) -> RootCommand
        where T: IntoCommandArg
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::INSERT,
                                     Some(data),
                                     Some(&commands)))
    }

    pub fn delete(self) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::DELETE,
                                  None as Option<&str>,
                                  Some(&commands)))
    }

    pub fn filter<T>(self, filter: T) -> RootCommand
        where T: IntoCommandArg
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::FILTER,
                                     Some(filter),
                                     Some(&commands)))
    }

    pub fn run(self) -> Result<String> {
        let logger = Client::logger().read();
        trace!(logger, "Calling r.run()");
        let commands = try!(self.0);
        let cfg = Client::config().read();
        let query = Query::wrap(qt::START, Some(&commands), None);
        debug!(logger, "{}", query);
        let mut conn = try!(Client::conn());
        // Try sending the query
        {
            let mut i = 0;
            let mut write = true;
            let mut connect = false;
            while i < cfg.retries {
                debug!(logger, "Getting connection...");
                if connect {
                    drop(&mut conn);
                    conn = match Client::conn() {
                        Ok(c) => c,
                        Err(error) => {
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(From::from(error));
                            } else {
                                debug!(logger, "Failed getting a connection. Retrying...");
                                i += 1;
                                continue;
                            }
                        }
                    };
                }
                debug!(logger, "Connection aquired.");
                conn.token += 1;
                if write {
                    debug!(logger, "Writing query...");
                    if let Err(error) = Query::write(&query, &mut conn) {
                        connect = true;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to write query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                    debug!(logger, "Query written...");
                    connect = false;
                }
                debug!(logger, "Reading query...");
                match Query::read(&mut conn) {
                    Ok(resp) => {
                        let result = try!(str::from_utf8(&resp));
                        debug!(logger, "{}", result);
                        // If the write operation failed, retry it
                        // {"t":18,"e":4100000,"r":["Cannot perform write: primary replica for
                        // shard [\"\", +inf) not available"],"b":[]}
                        let msg = r#"{"t":18,"e":4100000,"r":["Cannot perform write: primary replica for shard"#;
                        if result.starts_with(msg) {
                            write = true;
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(
                                From::from(
                                    AvailabilityError::OpFailed(
                                        String::from("Not available")
                                        )));
                            } else {
                                debug!(logger, "Write operation failed. Retrying...");
                                i += 1;
                                continue;
                            }
                        } else {
                            debug!(logger, "Query successfully read.");
                            // This is a successful operation
                            break;
                        }
                    }
                    Err(error) => {
                        write = false;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to read query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                }
            }
        }
        Ok(String::new())
    }
}

impl Command {
    pub fn wrap<T>(command: tt,
                input: Option<T>,
                commands: Option<&str>)
                -> Result<String>
        where T: IntoCommandArg
        {
            let mut cmds = format!("[{},", command.value());
            let mut args = String::new();
            if let Some(commands) = commands {
                args.push_str(commands);
                if input.is_some() {
                    args.push_str(",");
                }
            }
            let mut arguments: Option<String> = None;
            let mut options: Option<String> = None;
            if let Some(input) = input {
                let (args, opts) = try!(input.to_arg());
                arguments = args;
                options = opts;
            }
            if let Some(arguments) = arguments {
                args.push_str(&arguments);
            }
            cmds.push_str(&format!("[{}]", args));
            if let Some(options) = options {
                cmds.push_str(&format!(",{}", options));
            }
            cmds.push(']');
            Ok(cmds)
        }
}

impl Query {
    pub fn wrap(query_type: qt,
                query: Option<&str>,
                options: Option<&str>)
        -> String {
            let mut qry = format!("[{}", query_type.value());
            if let Some(query) = query {
                qry.push_str(&format!(",{}", query));
            }
            if let Some(options) = options {
                qry.push_str(&format!(",{}", options));
            }
            qry.push_str("]");
            qry
        }

    pub fn write(query: &str, conn: &mut Connection) -> Result<()> {
        let query = query.as_bytes();
        let token = conn.token;
        if let Err(error) = conn.stream.write_u64::<LittleEndian>(token) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_u32::<LittleEndian>(query.len() as u32) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_all(query) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.flush() {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(())
    }

    pub fn read(conn: &mut Connection) -> Result<Vec<u8>> {
        // @TODO use response_token to implement parallel reads and writes?
        // let response_token = try!(conn.stream.read_u64::<LittleEndian>());
        let _ = match conn.stream.read_u64::<LittleEndian>() {
            Ok(token) => token,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let len = match conn.stream.read_u32::<LittleEndian>() {
            Ok(len) => len,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let mut resp = vec![0u8; len as usize];
        if let Err(error) = conn.stream.read_exact(&mut resp) {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(resp)
    }
}

fn wrap_arrays(mut val: Value) -> Value {
    if val.is_array() {
        let mut array = Vec::with_capacity(2);
        array.push(Value::I64(tt::MAKE_ARRAY.value() as i64));
        if let Value::Array(vec) = val {
            let mut new_val = Vec::with_capacity(vec.len());
            for v in vec.into_iter() {
                if v.is_array() || v.is_object() {
                    new_val.push(wrap_arrays(v));
                } else {
                    new_val.push(v)
                }
            }
            val = Value::Array(new_val);
        }
        array.push(val);
        val = Value::Array(array);
    } else if val.is_object() {
        if let Value::Object(mut obj) = val {
            for (k, v) in obj.clone() {
                obj.insert(k, wrap_arrays(v));
            }
            val = Value::Object(obj);
        }
    }
    val
}
<commit_msg>Use macros to define commands<commit_after>//! The Actual RethinkDB Commands

use errors::*;
use conn::{Connection, ConnectOpts};
use serde_json::{self, Value};
use ql2::proto::{Term_TermType as tt, Query_QueryType as qt};
use std::io::Write;
use byteorder::{WriteBytesExt, LittleEndian};
use std::io::Read;
use std::str;
use protobuf::ProtobufEnum;
use byteorder::ReadBytesExt;
use super::session::Client;
use super::{Result, r};
use std::error::Error as StdError;

#[derive(Debug)]
pub struct RootCommand(Result<String>);
pub struct Command;
pub struct Query;

pub trait IntoCommandArg {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)>;
}

impl<'a> IntoCommandArg for &'a str {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        Ok((Some(format!("{:?}", self)), None))
    }
}

impl IntoCommandArg for String {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        Ok((Some(format!("{:?}", self)), None))
    }
}

impl IntoCommandArg for Value {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        // Arrays are a special case: since ReQL commands are sent as arrays
        // We have to wrap them using MAKE_ARRAY
        let val = wrap_arrays(self.clone());
        match serde_json::to_string(&val) {
            Ok(cmd) => Ok((Some(cmd), None)),
            Err(e) => Err(From::from(DriverError::Json(e))),
        }
    }
}

impl<T> IntoCommandArg for (T, Value)
        where T: IntoCommandArg
{
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        if !self.1.is_object() {
            let msg = String::from("Only objects are allowed as function options. You should use `r.object()` to pass optional arguments in your functions.");
            return Err(From::from(DriverError::Other(msg)));
        }
        let arg = try!(self.0.to_arg());
        let opt = try!(self.1.to_arg());
        Ok((arg.0, opt.0))
    }
}

impl IntoCommandArg for RootCommand {
    fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
        match self.0 {
            Ok(ref cmd) => Ok((Some(cmd.to_string()), None)),
            Err(ref e) => Err(From::from(DriverError::Other(e.description().to_string()))),
        }
    }
}

macro_rules! define {
    (impl IntoCommandArg for $T:ty) => {
        impl IntoCommandArg for $T {
            fn to_arg(&self) -> Result<(Option<String>, Option<String>)> {
                Ok((Some(self.to_string()), None))
            }
        }
    }
}

define!{ impl IntoCommandArg for bool }
define!{ impl IntoCommandArg for char }
define!{ impl IntoCommandArg for u8 }
define!{ impl IntoCommandArg for u16 }
define!{ impl IntoCommandArg for u32 }
define!{ impl IntoCommandArg for u64 }
define!{ impl IntoCommandArg for usize }
define!{ impl IntoCommandArg for i8 }
define!{ impl IntoCommandArg for i16 }
define!{ impl IntoCommandArg for i32 }
define!{ impl IntoCommandArg for i64 }
define!{ impl IntoCommandArg for isize }
define!{ impl IntoCommandArg for f32 }
define!{ impl IntoCommandArg for f64 }

impl Client {
    pub fn connection(&self) -> ConnectOpts {
        Self::config().read().clone()
    }

    pub fn expr<T>(&self, e: T) -> RootCommand
        where T: IntoCommandArg
    {
        match e.to_arg() {
            Ok(arg) => if let Some(arg) = arg.0 {
                RootCommand(Ok(arg))
            } else {
                RootCommand(Ok(String::new()))
            },
            Err(e) => RootCommand(Err(e)),
        }
    }

    pub fn db_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB_CREATE,
                                     Some(name),
                                     None))
    }

    pub fn db_drop<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB_DROP,
                                     Some(name),
                                     None))
    }

    pub fn db<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB,
                                     Some(name),
                                     None))
    }

    pub fn table_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table_create(name)
    }

    pub fn table<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table(name)
    }

    pub fn table_drop<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        let config = Client::config().read();
        r.db(config.db).table_drop(name)
    }

    pub fn object(&self) -> serde_json::builder::ObjectBuilder {
        serde_json::builder::ObjectBuilder::new()
    }

    pub fn array(&self) -> serde_json::builder::ArrayBuilder {
        serde_json::builder::ArrayBuilder::new()
    }
}

macro_rules! command {
    ($name:ident, $cmd:ident) => {
        pub fn $name<T>(self, arg: T) -> RootCommand
            where T: IntoCommandArg
            {
                let commands = match self.0 {
                    Ok(t) => t,
                    Err(e) => return RootCommand(Err(e)),
                };
                RootCommand(Command::wrap(tt::$cmd,
                                          Some(arg),
                                          Some(&commands)))
            }
    }
}

impl RootCommand {
    command!(table_create, TABLE_CREATE);
    command!(table_drop, TABLE_DROP);
    command!(table, TABLE);
    command!(filter, FILTER);
    command!(insert, INSERT);

    pub fn delete(self) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(tt::DELETE,
                                  None as Option<&str>,
                                  Some(&commands)))
    }

    pub fn run(self) -> Result<String> {
        let logger = Client::logger().read();
        trace!(logger, "Calling r.run()");
        let commands = try!(self.0);
        let cfg = Client::config().read();
        let query = Query::wrap(qt::START, Some(&commands), None);
        debug!(logger, "{}", query);
        let mut conn = try!(Client::conn());
        // Try sending the query
        {
            let mut i = 0;
            let mut write = true;
            let mut connect = false;
            while i < cfg.retries {
                debug!(logger, "Getting connection...");
                if connect {
                    drop(&mut conn);
                    conn = match Client::conn() {
                        Ok(c) => c,
                        Err(error) => {
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(From::from(error));
                            } else {
                                debug!(logger, "Failed getting a connection. Retrying...");
                                i += 1;
                                continue;
                            }
                        }
                    };
                }
                debug!(logger, "Connection aquired.");
                conn.token += 1;
                if write {
                    debug!(logger, "Writing query...");
                    if let Err(error) = Query::write(&query, &mut conn) {
                        connect = true;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to write query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                    debug!(logger, "Query written...");
                    connect = false;
                }
                debug!(logger, "Reading query...");
                match Query::read(&mut conn) {
                    Ok(resp) => {
                        let result = try!(str::from_utf8(&resp));
                        debug!(logger, "{}", result);
                        // If the write operation failed, retry it
                        // {"t":18,"e":4100000,"r":["Cannot perform write: primary replica for
                        // shard [\"\", +inf) not available"],"b":[]}
                        let msg = r#"{"t":18,"e":4100000,"r":["Cannot perform write: primary replica for shard"#;
                        if result.starts_with(msg) {
                            write = true;
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(
                                From::from(
                                    AvailabilityError::OpFailed(
                                        String::from("Not available")
                                        )));
                            } else {
                                debug!(logger, "Write operation failed. Retrying...");
                                i += 1;
                                continue;
                            }
                        } else {
                            debug!(logger, "Query successfully read.");
                            // This is a successful operation
                            break;
                        }
                    }
                    Err(error) => {
                        write = false;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to read query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                }
            }
        }
        Ok(String::new())
    }
}

impl Command {
    pub fn wrap<T>(command: tt,
                input: Option<T>,
                commands: Option<&str>)
                -> Result<String>
        where T: IntoCommandArg
        {
            let mut cmds = format!("[{},", command.value());
            let mut args = String::new();
            if let Some(commands) = commands {
                args.push_str(commands);
                if input.is_some() {
                    args.push_str(",");
                }
            }
            let mut arguments: Option<String> = None;
            let mut options: Option<String> = None;
            if let Some(input) = input {
                let (args, opts) = try!(input.to_arg());
                arguments = args;
                options = opts;
            }
            if let Some(arguments) = arguments {
                args.push_str(&arguments);
            }
            cmds.push_str(&format!("[{}]", args));
            if let Some(options) = options {
                cmds.push_str(&format!(",{}", options));
            }
            cmds.push(']');
            Ok(cmds)
        }
}

impl Query {
    pub fn wrap(query_type: qt,
                query: Option<&str>,
                options: Option<&str>)
        -> String {
            let mut qry = format!("[{}", query_type.value());
            if let Some(query) = query {
                qry.push_str(&format!(",{}", query));
            }
            if let Some(options) = options {
                qry.push_str(&format!(",{}", options));
            }
            qry.push_str("]");
            qry
        }

    pub fn write(query: &str, conn: &mut Connection) -> Result<()> {
        let query = query.as_bytes();
        let token = conn.token;
        if let Err(error) = conn.stream.write_u64::<LittleEndian>(token) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_u32::<LittleEndian>(query.len() as u32) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_all(query) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.flush() {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(())
    }

    pub fn read(conn: &mut Connection) -> Result<Vec<u8>> {
        // @TODO use response_token to implement parallel reads and writes?
        // let response_token = try!(conn.stream.read_u64::<LittleEndian>());
        let _ = match conn.stream.read_u64::<LittleEndian>() {
            Ok(token) => token,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let len = match conn.stream.read_u32::<LittleEndian>() {
            Ok(len) => len,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let mut resp = vec![0u8; len as usize];
        if let Err(error) = conn.stream.read_exact(&mut resp) {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(resp)
    }
}

fn wrap_arrays(mut val: Value) -> Value {
    if val.is_array() {
        let mut array = Vec::with_capacity(2);
        array.push(Value::I64(tt::MAKE_ARRAY.value() as i64));
        if let Value::Array(vec) = val {
            let mut new_val = Vec::with_capacity(vec.len());
            for v in vec.into_iter() {
                if v.is_array() || v.is_object() {
                    new_val.push(wrap_arrays(v));
                } else {
                    new_val.push(v)
                }
            }
            val = Value::Array(new_val);
        }
        array.push(val);
        val = Value::Array(array);
    } else if val.is_object() {
        if let Value::Object(mut obj) = val {
            for (k, v) in obj.clone() {
                obj.insert(k, wrap_arrays(v));
            }
            val = Value::Object(obj);
        }
    }
    val
}
<|endoftext|>


<commit_before>use std::path::PathBuf;

pub trait Completer {
    fn completions(&self, start: &str) -> Vec<String>;
}

pub struct BasicCompleter {
    prefixes: Vec<String>,
}

impl BasicCompleter {
    pub fn new<T: Into<String>>(prefixes: Vec<T>) -> BasicCompleter {
        BasicCompleter { prefixes: prefixes.into_iter().map(|s| s.into()).collect() }
    }
}

impl Completer for BasicCompleter {
    fn completions(&self, start: &str) -> Vec<String> {
        self.prefixes.iter().filter(|s| s.starts_with(start)).cloned().collect()
    }
}

pub struct FilenameCompleter {
    working_dir: Option<PathBuf>,
}

impl FilenameCompleter {
    pub fn new<T: Into<PathBuf>>(working_dir: Option<T>) -> Self {
        FilenameCompleter { working_dir: working_dir.map(|p| p.into()) }
    }
}

impl Completer for FilenameCompleter {
    fn completions(&self, mut start: &str) -> Vec<String> {
        // XXX: this function is really bad, TODO rewrite

        let start_owned;
        if start.starts_with("\"") || start.starts_with("'") {
            start = &start[1..];
            if start.len() >= 1 {
                start = &start[..start.len() - 1];
            }
            start_owned = start.into();
        } else {
            start_owned = start.replace("\\ ", " ");
        }

        let full_path;
        let start_path = PathBuf::from(&start_owned[..]);

        if let Some(ref wd) = self.working_dir {
            let mut fp = PathBuf::from(wd);
            fp.push(start_owned.clone());
            full_path = fp;
        } else {
            full_path = PathBuf::from(start_owned.clone());
        }

        if full_path.is_relative() {
            return vec![];
        }

        let p;
        let start_name;
        let completing_dir;
        match full_path.parent() {
            // XXX non-unix separaor
            Some(parent) if start != "" && !start_owned.ends_with("/") &&
                            !full_path.ends_with("..") => {
                p = PathBuf::from(parent);
                start_name = full_path.file_name().unwrap().to_string_lossy().into_owned();
                completing_dir = false;
            }
            _ => {
                p = full_path.clone();
                start_name = "".into();
                completing_dir = start == "" || start.ends_with("/") || full_path.ends_with("..");
            }
        }

        let read_dir = match p.read_dir() {
            Ok(x) => x,
            Err(_) => return vec![],
        };

        let mut matches = vec![];
        for dir in read_dir {
            let dir = match dir {
                Ok(x) => x,
                Err(_) => continue,
            };
            let file_name = dir.file_name();
            let file_name = file_name.to_string_lossy();

            if start_name == "" || file_name.starts_with(&*start_name) {
                let mut a = start_path.clone();
                if !a.is_absolute() {
                    a = PathBuf::new();
                } else if !completing_dir && !a.pop() {
                    return vec![];
                }
                a.push(dir.file_name());
                let mut s = a.to_string_lossy().into_owned();
                if dir.path().is_dir() {
                    s = s + "/";
                }

                let mut b = PathBuf::from(start_owned.clone());
                if !completing_dir {
                    b.pop();
                }
                b.push(s);

                matches.push(b.to_string_lossy().to_owned().replace(" ", "\\ "));
            }
        }

        matches
    }
}
<commit_msg>Patch for Redox - is_relative does not work yet<commit_after>use std::path::PathBuf;

pub trait Completer {
    fn completions(&self, start: &str) -> Vec<String>;
}

pub struct BasicCompleter {
    prefixes: Vec<String>,
}

impl BasicCompleter {
    pub fn new<T: Into<String>>(prefixes: Vec<T>) -> BasicCompleter {
        BasicCompleter { prefixes: prefixes.into_iter().map(|s| s.into()).collect() }
    }
}

impl Completer for BasicCompleter {
    fn completions(&self, start: &str) -> Vec<String> {
        self.prefixes.iter().filter(|s| s.starts_with(start)).cloned().collect()
    }
}

pub struct FilenameCompleter {
    working_dir: Option<PathBuf>,
}

impl FilenameCompleter {
    pub fn new<T: Into<PathBuf>>(working_dir: Option<T>) -> Self {
        FilenameCompleter { working_dir: working_dir.map(|p| p.into()) }
    }
}

impl Completer for FilenameCompleter {
    fn completions(&self, mut start: &str) -> Vec<String> {
        // XXX: this function is really bad, TODO rewrite

        let start_owned;
        if start.starts_with("\"") || start.starts_with("'") {
            start = &start[1..];
            if start.len() >= 1 {
                start = &start[..start.len() - 1];
            }
            start_owned = start.into();
        } else {
            start_owned = start.replace("\\ ", " ");
        }

        let full_path;
        let start_path = PathBuf::from(&start_owned[..]);

        if let Some(ref wd) = self.working_dir {
            let mut fp = PathBuf::from(wd);
            fp.push(start_owned.clone());
            full_path = fp;
        } else {
            full_path = PathBuf::from(start_owned.clone());
        }

        let p;
        let start_name;
        let completing_dir;
        match full_path.parent() {
            // XXX non-unix separaor
            Some(parent) if start != "" && !start_owned.ends_with("/") &&
                            !full_path.ends_with("..") => {
                p = PathBuf::from(parent);
                start_name = full_path.file_name().unwrap().to_string_lossy().into_owned();
                completing_dir = false;
            }
            _ => {
                p = full_path.clone();
                start_name = "".into();
                completing_dir = start == "" || start.ends_with("/") || full_path.ends_with("..");
            }
        }

        let read_dir = match p.read_dir() {
            Ok(x) => x,
            Err(_) => return vec![],
        };

        let mut matches = vec![];
        for dir in read_dir {
            let dir = match dir {
                Ok(x) => x,
                Err(_) => continue,
            };
            let file_name = dir.file_name();
            let file_name = file_name.to_string_lossy();

            if start_name == "" || file_name.starts_with(&*start_name) {
                let mut a = start_path.clone();
                if !a.is_absolute() {
                    a = PathBuf::new();
                } else if !completing_dir && !a.pop() {
                    return vec![];
                }
                a.push(dir.file_name());
                let mut s = a.to_string_lossy().into_owned();
                if dir.path().is_dir() {
                    s = s + "/";
                }

                let mut b = PathBuf::from(start_owned.clone());
                if !completing_dir {
                    b.pop();
                }
                b.push(s);

                matches.push(b.to_string_lossy().to_owned().replace(" ", "\\ "));
            }
        }

        matches
    }
}
<|endoftext|>


<commit_before>//!
//! Hold the `Game` struct and the `play()` function
//!

// Tcod
extern crate tcod;
use self::tcod::{console, Console};
use core::tcod::input;

// core::world
pub mod world;
use core::world::World;

// core::object
pub mod object;
// Game state depends on what actions the player can do
use core::object::actions::Actions;

// core::renderer
pub mod renderer;
use self::renderer::Renderer;
use self::renderer::RGB;

// core::log
pub mod log;
use self::log::GlobalLog;

// Use init here privately
// After all, what game function requires that something in the initalizer be ran?
mod init;

///
/// Enum representing the state of the game
///
pub enum State {
  // Game just created
  New,
  // Player acted
  Act(Actions),
  // Debug
  Debug
}

///
/// Engine struct. Holds a world, state, and renderer
///
pub struct Engine {
  world: World,
  state: State,
  ren: Renderer,
  root: console::Root,

  // Debug
  noclip: bool

}

impl Engine {

  ///
  /// Capture keyboard input from tcod and update player state
  ///
  fn process_keypress(&mut self, keypress: input::Key) {

    match keypress.code {

      // If the keycode isn't escape we continue checking for important keys
      input::KeyCode::Escape => panic!("Bye"),
      _ => {

        if keypress.printable != ' ' {

          let oldpos = self.world.player.actor.pos.clone();

          // Wipe the state
          self.state = State::New;

          match keypress.printable {

            // Movement keys are bound to vim-like controls (hjklyubn)
            'h' => {
              self.world.player.actor.move_cart(-1, 0);
              self.world.player.state = Actions::Move;
            },
            'j' => {
              self.world.player.actor.move_cart(0, 1);
              self.world.player.state = Actions::Move;
            },
            'k' => {
              self.world.player.actor.move_cart(0, -1);
              self.world.player.state = Actions::Move;
            },
            'l' => {
              self.world.player.actor.move_cart(1, 0);
              self.world.player.state = Actions::Move;
            },
            'y' => {
              self.world.player.actor.move_cart(-1, -1);
              self.world.player.state = Actions::Move;
            },
            'u' => {
              self.world.player.actor.move_cart(1, -1);
              self.world.player.state = Actions::Move;
            },
            'b' => {
              self.world.player.actor.move_cart(-1, 1);
              self.world.player.state = Actions::Move;
            },
            'n' => {
              self.world.player.actor.move_cart(1, 1);
              self.world.player.state = Actions::Move;
            },
            // "Regenerate" current level
            'w' => {
              self.world.test_traverse();
              self.state = State::Act(Actions::Unknown);
            },
            // Create an empty level for testing
            'q' => {
              self.world.test_empty();
              self.state = State::Act(Actions::Unknown);
            },
            // Wait
            '.' => {
              self.world.player.state = Actions::Wait;
            },
            // Go downstars (if possible)
            '>' => { self.world.player.state = Actions::DownStair },
            // Go upstairs (if possible)
            '<' => { self.world.player.state = Actions::UpStair },
            // Unbound key, so we just say we don't know what the player did

            // Debug keypresses

            // Toggle scent
            'r' => {
              self.ren.show_scent = !self.ren.show_scent;
              self.ren.draw_all(&mut self.root, &mut self.world);
              self.state = State::Debug;
            },

            // Toggle sound
            't' => {
              self.ren.show_sound = !self.ren.show_sound;
              self.ren.draw_all(&mut self.root, &mut self.world);
              self.state = State::Debug;
            },

            // Toggle FoV
            'f' => {
              self.ren.fov = !self.ren.fov;
              self.ren.draw_all(&mut self.root, &mut self.world);
              self.state = State::Debug;
            },

            // Toggle noclip
            'z' => {
              self.noclip = !self.noclip;
              self.state = State::Debug;
            },

            _ => { self.world.player.state = Actions::Unknown }

          }

          // If state is Debug, don't override
          match self.state {
            State::Debug => (),
            _ => {
              // Set game state to player state
              self.state = State::Act(self.world.player.state.clone());

              match self.world.player.state {
                Actions::Move => {
                  // Make sure player doesn't do anything dumb
                  if !self.world.is_valid(self.world.player.actor.pos.x, self.world.player.actor.pos.y) && !self.noclip {
                    self.world.player.actor.pos = oldpos;
                    self.world.player.state = Actions::Unknown;
                  }
                }
                _ => ()
              }

            }
          }

        }

        /*
        // Prints keycode to console in case if you're trying to find a key that isn't intutive
        else {
          println!("{:?}", keypress.code);
        }
        */

      }

    }

  }

  ///
  /// Return a new `Game`
  ///
  /// This function assumes you will just be passing in tcod::console::Root.width() and height(),
  /// so inputs are i32s instead of usizes (they get converted)
  ///
  pub fn new() -> Engine {

    // Get map height
    let map_dim = init::map_dimensions();

    // Get root console
    let root = init::root();

    Engine {
      world: World::new(map_dim),
      state: State::New,
      ren: Renderer::new(
        map_dim,
        (root.width() as isize, root.height() as isize),
        init::console_height(),
        init::panel_width()
      ),
      root: root,

      // Debug
      noclip: false
    }

  }

  ///
  /// Update the game state, then update depending on the new state
  ///
  fn update(&mut self) {

    match &self.state {
      // Moving or waiting prompts a world update
      &State::Act(Actions::Move) | &State::Act(Actions::Wait) => self.world.update(),

      // Trying to go up and downstairs prompts the respective response from world
      &State::Act(Actions::DownStair) => {
        if self.world.can_go_down() {
          self.world.go_down();
        }
      },
      &State::Act(Actions::UpStair) => {
        if self.world.can_go_up() {
          self.world.go_up();
        }
      }

      // All other variants are dropped
      _ => ()
    }

  }

  ///
  /// Play the game.
  ///
  pub fn play(&mut self) {

    // Some starting messages, will be removed in later versions
    let mut log = GlobalLog.lock().unwrap();
    log.push(("Welcome to Edgequest", RGB(255, 0, 255)));
    log.push(("Move with vim keys", RGB(255, 255, 255)));
    log.push(("esc to quit, w to regenerate", RGB(255, 255, 255)));
    log.push(("r to toggle scent, t to toggle sound", RGB(255, 255, 255)));
    log.push(("f to toggle FoV, z to toggle noclip", RGB(255, 255, 255)));
    drop(log);

    // Initial update
    self.update();

    // Draw all and capture keypresses
    while !self.root.window_closed() {

      // Draw what the camera sees
      self.ren.draw_all(&mut self.root, &mut self.world);

      // Capture game keys (Keys that change the state of the player)
      // This is what gives it the turn based nature, i.e. waits for player input before
      // doing anything
      let keypress = self.root.wait_for_keypress(true);
      self.process_keypress(keypress);

      // Update game
      self.update();

    }

  }

}<commit_msg>Update<commit_after>//!
//! Hold the `Game` struct and the `play()` function
//!

// Tcod
extern crate tcod;
use self::tcod::{console, Console};
use core::tcod::input;

// core::world
pub mod world;
use core::world::World;

// core::object
pub mod object;
// Game state depends on what actions the player can do
use core::object::actions::Actions;

// core::renderer
pub mod renderer;
use self::renderer::Renderer;
use self::renderer::RGB;

// core::log
pub mod log;
use self::log::GlobalLog;

// Use init here privately
// After all, what game function requires that something in the initalizer be ran?
mod init;

///
/// Enum representing the state of the game
///
pub enum State {
  // Game just created
  New,
  // Player acted
  Act(Actions),
  // Debug
  Debug
}

///
/// Engine struct. Holds a world, state, and renderer
///
pub struct Engine {
  world: World,
  state: State,
  ren: Renderer,
  root: console::Root,

  // Debug
  noclip: bool

}

impl Engine {

  ///
  /// Capture keyboard input from tcod and update player state
  ///
  fn process_keypress(&mut self, keypress: input::Key) {

    match keypress.code {

      // If the keycode isn't escape we continue checking for important keys
      input::KeyCode::Escape => panic!("Bye"),
      _ => {

        if keypress.printable != ' ' {

          let oldpos = self.world.player.actor.pos.clone();

          // Wipe the state
          self.state = State::New;

          match keypress.printable {

            // Movement keys are bound to vim-like controls (hjklyubn)
            'h' => {
              self.world.player.actor.move_cart(-1, 0);
              self.world.player.state = Actions::Move;
            },
            'j' => {
              self.world.player.actor.move_cart(0, 1);
              self.world.player.state = Actions::Move;
            },
            'k' => {
              self.world.player.actor.move_cart(0, -1);
              self.world.player.state = Actions::Move;
            },
            'l' => {
              self.world.player.actor.move_cart(1, 0);
              self.world.player.state = Actions::Move;
            },
            'y' => {
              self.world.player.actor.move_cart(-1, -1);
              self.world.player.state = Actions::Move;
            },
            'u' => {
              self.world.player.actor.move_cart(1, -1);
              self.world.player.state = Actions::Move;
            },
            'b' => {
              self.world.player.actor.move_cart(-1, 1);
              self.world.player.state = Actions::Move;
            },
            'n' => {
              self.world.player.actor.move_cart(1, 1);
              self.world.player.state = Actions::Move;
            },
            // "Regenerate" current level
            'w' => {
              self.world.test_traverse();
              self.state = State::Act(Actions::Unknown);
            },
            // Create an empty level for testing
            'q' => {
              self.world.test_empty();
              self.state = State::Act(Actions::Unknown);
            },
            // Wait
            '.' => {
              self.world.player.state = Actions::Wait;
            },
            // Go downstars (if possible)
            '>' => { self.world.player.state = Actions::DownStair },
            // Go upstairs (if possible)
            '<' => { self.world.player.state = Actions::UpStair },
            // Unbound key, so we just say we don't know what the player did

            // Debug keypresses

            // Toggle scent
            'r' => {
              self.ren.show_scent = !self.ren.show_scent;
              self.ren.draw_all(&mut self.root, &mut self.world);
              self.state = State::Debug;
            },

            // Toggle sound
            't' => {
              self.ren.show_sound = !self.ren.show_sound;
              self.ren.draw_all(&mut self.root, &mut self.world);
              self.state = State::Debug;
            },

            // Toggle FoV
            'f' => {
              self.ren.fov = !self.ren.fov;
              self.ren.draw_all(&mut self.root, &mut self.world);
              self.state = State::Debug;
            },

            // Toggle noclip
            'z' => {
              self.noclip = !self.noclip;
              self.state = State::Debug;
            },

            _ => { self.world.player.state = Actions::Unknown }

          }

          // If state is Debug, don't override
          match self.state {
            State::Debug => (),
            _ => {
              // Set game state to player state
              self.state = State::Act(self.world.player.state.clone());

              match self.world.player.state {
                Actions::Move => {
                  // Make sure player doesn't do anything dumb
                  if !self.world.is_valid(self.world.player.actor.pos.x, self.world.player.actor.pos.y) && !self.noclip {
                    self.world.player.actor.pos = oldpos;
                    self.world.player.state = Actions::Unknown;
                  }
                }
                _ => ()
              }

            }
          }

        }

        /*
        // Prints keycode to console in case if you're trying to find a key that isn't intutive
        else {
          println!("{:?}", keypress.code);
        }
        */

      }

    }

  }

  ///
  /// Return a new `Game`
  ///
  /// This function assumes you will just be passing in tcod::console::Root.width() and height(),
  /// so inputs are i32s instead of usizes (they get converted)
  ///
  pub fn new() -> Engine {

    // Get map height
    let map_dim = init::map_dimensions();

    // Get root console
    let root = init::root();

    Engine {
      world: World::new(map_dim),
      state: State::New,
      ren: Renderer::new(
        map_dim,
        (root.width() as isize, root.height() as isize),
        init::console_height(),
        init::panel_width()
      ),
      root: root,

      // Debug
      noclip: false
    }

  }

  ///
  /// Update the game state, then update depending on the new state
  ///
  fn update(&mut self) {

    match self.state {
      // Moving or waiting prompts a world update
      State::Act(Actions::Move) | State::Act(Actions::Wait) => self.world.update(),

      // Trying to go up and downstairs prompts the respective response from world
      State::Act(Actions::DownStair) => {
        if self.world.can_go_down() {
          self.world.go_down();
        }
      },
      State::Act(Actions::UpStair) => {
        if self.world.can_go_up() {
          self.world.go_up();
        }
      }

      // All other variants are dropped
      _ => ()
    }

  }

  ///
  /// Play the game.
  ///
  pub fn play(&mut self) {

    // Some starting messages, will be removed in later versions
    let mut log = GlobalLog.lock().unwrap();
    log.push(("Welcome to Edgequest", RGB(255, 0, 255)));
    log.push(("Move with vim keys", RGB(255, 255, 255)));
    log.push(("esc to quit, w to regenerate", RGB(255, 255, 255)));
    log.push(("r to toggle scent, t to toggle sound", RGB(255, 255, 255)));
    log.push(("f to toggle FoV, z to toggle noclip", RGB(255, 255, 255)));
    drop(log);

    // Initial update
    self.update();

    // Draw all and capture keypresses
    while !self.root.window_closed() {

      // Draw what the camera sees
      self.ren.draw_all(&mut self.root, &mut self.world);

      // Capture game keys (Keys that change the state of the player)
      // This is what gives it the turn based nature, i.e. waits for player input before
      // doing anything
      let keypress = self.root.wait_for_keypress(true);
      self.process_keypress(keypress);

      // Update game
      self.update();

    }

  }

}<|endoftext|>


<commit_before>use super::content::Content;
use super::encodings::{self, bytes_to_string, string_to_bytes};
use super::{Dictionary, Object, ObjectId};
use crate::xref::Xref;
use crate::{Error, Result};
use encoding::all::UTF_16BE;
use encoding::types::{DecoderTrap, EncoderTrap, Encoding};
use log::info;
use std::collections::BTreeMap;
use std::io::Write;
use std::str;

/// PDF document.
#[derive(Debug, Clone)]
pub struct Document {
	/// The version of the PDF specification to which the file conforms.
	pub version: String,

	/// The trailer gives the location of the cross-reference table and of certain special objects.
	pub trailer: Dictionary,

	/// The cross-reference table contains locations of the indirect objects.
	pub reference_table: Xref,

	/// The objects that make up the document contained in the file.
	pub objects: BTreeMap<ObjectId, Object>,

	/// Current maximum object id within the document.
	pub max_id: u32,
}

impl Document {
	/// Create new PDF document.
	pub fn new() -> Document {
		Document {
			version: "1.4".to_string(),
			trailer: Dictionary::new(),
			reference_table: Xref::new(0),
			objects: BTreeMap::new(),
			max_id: 0,
		}
	}

	/// Get object by object id, will recursively dereference a referenced object.
	pub fn get_object(&self, id: ObjectId) -> Result<&Object> {
		if let Some(object) = self.objects.get(&id) {
			if let Ok(id) = object.as_reference() {
				return self.get_object(id);
			} else {
				return Ok(object);
			}
		}
		Err(Error::ObjectNotFound)
	}

	/// Get mutable reference to object by object id, will recursively dereference a referenced object.
	pub fn get_object_mut(&mut self, id: ObjectId) -> Result<&mut Object> {
		unsafe {
			let s = self as *mut Self;
			if let Some(object) = (*s).objects.get_mut(&id) {
				if let Ok(id) = object.as_reference() {
					return (*s).get_object_mut(id);
				} else {
					return Ok(object);
				}
			}
			Err(Error::ObjectNotFound)
		}
	}

	/// Get dictionary object by id.
	pub fn get_dictionary(&self, id: ObjectId) -> Result<&Dictionary> {
		self.get_object(id).and_then(Object::as_dict)
	}

	/// Traverse objects from trailer recursively, return all referenced object IDs.
	pub fn traverse_objects<A: Fn(&mut Object) -> ()>(&mut self, action: A) -> Vec<ObjectId> {
		fn traverse_array<A: Fn(&mut Object) -> ()>(array: &mut Vec<Object>, action: &A, refs: &mut Vec<ObjectId>) {
			for item in array.iter_mut() {
				traverse_object(item, action, refs);
			}
		}
		fn traverse_dictionary<A: Fn(&mut Object) -> ()>(dict: &mut Dictionary, action: &A, refs: &mut Vec<ObjectId>) {
			for (_, v) in dict.iter_mut() {
				traverse_object(v, action, refs);
			}
		}
		fn traverse_object<A: Fn(&mut Object) -> ()>(object: &mut Object, action: &A, refs: &mut Vec<ObjectId>) {
			action(object);
			match *object {
				Object::Array(ref mut array) => traverse_array(array, action, refs),
				Object::Dictionary(ref mut dict) => traverse_dictionary(dict, action, refs),
				Object::Stream(ref mut stream) => traverse_dictionary(&mut stream.dict, action, refs),
				Object::Reference(id) => {
					if !refs.contains(&id) {
						refs.push(id);
					}
				}
				_ => {}
			}
		}
		let mut refs = vec![];
		traverse_dictionary(&mut self.trailer, &action, &mut refs);
		let mut index = 0;
		while index < refs.len() {
			if let Some(object) = self.objects.get_mut(&refs[index]) {
				traverse_object(object, &action, &mut refs);
			}
			index += 1;
		}
		refs
	}

	/// Get catalog dictionary.
	pub fn catalog(&self) -> Result<&Dictionary> {
		self.trailer.get(b"Root").and_then(Object::as_reference).and_then(|id| self.get_dictionary(id))
	}

	/// Get page numbers and corresponding object ids.
	pub fn get_pages(&self) -> BTreeMap<u32, ObjectId> {
		fn collect_pages(doc: &Document, page_tree_id: ObjectId, page_number: &mut u32, pages: &mut BTreeMap<u32, ObjectId>) {
			if let Ok(kids) = doc.get_dictionary(page_tree_id).and_then(|page_tree| page_tree.get(b"Kids")).and_then(Object::as_array) {
				for kid in kids {
					if let Ok(kid_id) = kid.as_reference() {
						if let Ok(type_name) = doc.get_dictionary(kid_id).and_then(Dictionary::type_name) {
							match type_name {
								"Page" => {
									pages.insert(*page_number, kid_id);
									*page_number += 1;
								}
								"Pages" => {
									collect_pages(doc, kid_id, page_number, pages);
								}
								_ => {}
							}
						}
					}
				}
			}
		}

		let mut pages = BTreeMap::new();
		let mut page_number = 1;
		if let Ok(page_tree_id) = self.catalog().and_then(|cat| cat.get(b"Pages")).and_then(Object::as_reference) {
			collect_pages(self, page_tree_id, &mut page_number, &mut pages);
		}
		pages
	}

	/// Get content stream object ids of a page.
	pub fn get_page_contents(&self, page_id: ObjectId) -> Vec<ObjectId> {
		let mut streams = vec![];
		if let Ok(page) = self.get_dictionary(page_id) {
			if let Ok(contents) = page.get(b"Contents") {
				match *contents {
					Object::Reference(ref id) => {
						streams.push(*id);
					}
					Object::Array(ref arr) => {
						for content in arr {
							if let Ok(id) = content.as_reference() { streams.push(id) }
						}
					}
					_ => {}
				}
			}
		}
		streams
	}

	/// Get content of a page.
	pub fn get_page_content(&self, page_id: ObjectId) -> Result<Vec<u8>> {
		let mut content = Vec::new();
		let content_streams = self.get_page_contents(page_id);
		for object_id in content_streams {
			if let Ok(content_stream) = self.get_object(object_id).and_then(Object::as_stream) {
				match content_stream.decompressed_content() {
					Ok(data) => content.write_all(&data)?,
					Err(_) => content.write_all(&content_stream.content)?,
				};
			}
		}
		Ok(content)
	}

	/// Get decoded page content;
	pub fn get_and_decode_page_content(&self, page_id: ObjectId) -> Result<Content> {
		let content_data = self.get_page_content(page_id)?;
		Content::decode(&content_data)
	}

	/// Get resources used by a page.
	pub fn get_page_resources(&self, page_id: ObjectId) -> (Option<&Dictionary>, Vec<ObjectId>) {
		fn collect_resources(page_node: &Dictionary, resource_ids: &mut Vec<ObjectId>, doc: &Document) {
			if let Ok(resources_id) = page_node.get(b"Resources").and_then(Object::as_reference) {
				resource_ids.push(resources_id);
			}
			if let Ok(page_tree) = page_node.get(b"Parent").and_then(Object::as_reference).and_then(|id| doc.get_dictionary(id)) {
				collect_resources(page_tree, resource_ids, doc);
			}
		};

		let mut resource_dict = None;
		let mut resource_ids = Vec::new();
		if let Ok(page) = self.get_dictionary(page_id) {
			resource_dict = page.get(b"Resources").and_then(Object::as_dict).ok();
			collect_resources(page, &mut resource_ids, self);
		}
		(resource_dict, resource_ids)
	}

	/// Get fonts used by a page.
	pub fn get_page_fonts(&self, page_id: ObjectId) -> BTreeMap<Vec<u8>, &Dictionary> {
		fn collect_fonts_from_resources<'a>(resources: &'a Dictionary, fonts: &mut BTreeMap<Vec<u8>, &'a Dictionary>, doc: &'a Document) {
			if let Ok(font_dict) = resources.get(b"Font").and_then(Object::as_dict) {
				for (name, value) in font_dict.iter() {
					let font = match *value {
						Object::Reference(id) => doc.get_dictionary(id).ok(),
						Object::Dictionary(ref dict) => Some(dict),
						_ => None,
					};
					if !fonts.contains_key(name) {
						font.map(|font| fonts.insert(name.clone(), font));
					}
				}
			}
		};

		let mut fonts = BTreeMap::new();
		let (resource_dict, resource_ids) = self.get_page_resources(page_id);
		if let Some(resources) = resource_dict {
			collect_fonts_from_resources(resources, &mut fonts, self);
		}
		for resource_id in resource_ids {
			if let Ok(resources) = self.get_dictionary(resource_id) {
				collect_fonts_from_resources(resources, &mut fonts, self);
			}
		}
		fonts
	}

	pub fn decode_text(encoding: Option<&str>, bytes: &[u8]) -> String {
		if let Some(encoding) = encoding {
			info!("{}", encoding);
			match encoding {
				"StandardEncoding" => bytes_to_string(encodings::STANDARD_ENCODING, bytes),
				"MacRomanEncoding" => bytes_to_string(encodings::MAC_ROMAN_ENCODING, bytes),
				"MacExpertEncoding" => bytes_to_string(encodings::MAC_EXPERT_ENCODING, bytes),
				"WinAnsiEncoding" => bytes_to_string(encodings::WIN_ANSI_ENCODING, bytes),
				"UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.decode(bytes, DecoderTrap::Ignore).unwrap(),
				"Identity-H" => "?Identity-H Unimplemented?".to_string(), // Unimplemented
				_ => String::from_utf8_lossy(bytes).to_string(),
			}
		} else {
			bytes_to_string(encodings::STANDARD_ENCODING, bytes)
		}
	}

	pub fn encode_text(encoding: Option<&str>, text: &str) -> Vec<u8> {
		if let Some(encoding) = encoding {
			match encoding {
				"StandardEncoding" => string_to_bytes(encodings::STANDARD_ENCODING, text),
				"MacRomanEncoding" => string_to_bytes(encodings::MAC_ROMAN_ENCODING, text),
				"MacExpertEncoding" => string_to_bytes(encodings::MAC_EXPERT_ENCODING, text),
				"WinAnsiEncoding" => string_to_bytes(encodings::WIN_ANSI_ENCODING, text),
				"UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.encode(text, EncoderTrap::Ignore).unwrap(),
				"Identity-H" => vec![], // Unimplemented
				_ => text.as_bytes().to_vec(),
			}
		} else {
			string_to_bytes(encodings::STANDARD_ENCODING, text)
		}
	}
}

impl Default for Document {
	fn default() -> Self {
		Self::new()
	}
}
<commit_msg>Remove unsafe code on get_object_mut.<commit_after>use super::content::Content;
use super::encodings::{self, bytes_to_string, string_to_bytes};
use super::{Dictionary, Object, ObjectId};
use crate::xref::Xref;
use crate::{Error, Result};
use encoding::all::UTF_16BE;
use encoding::types::{DecoderTrap, EncoderTrap, Encoding};
use log::info;
use std::collections::BTreeMap;
use std::io::Write;
use std::str;

/// PDF document.
#[derive(Debug, Clone)]
pub struct Document {
	/// The version of the PDF specification to which the file conforms.
	pub version: String,

	/// The trailer gives the location of the cross-reference table and of certain special objects.
	pub trailer: Dictionary,

	/// The cross-reference table contains locations of the indirect objects.
	pub reference_table: Xref,

	/// The objects that make up the document contained in the file.
	pub objects: BTreeMap<ObjectId, Object>,

	/// Current maximum object id within the document.
	pub max_id: u32,
}

impl Document {
	/// Create new PDF document.
	pub fn new() -> Document {
		Document {
			version: "1.4".to_string(),
			trailer: Dictionary::new(),
			reference_table: Xref::new(0),
			objects: BTreeMap::new(),
			max_id: 0,
		}
	}

	/// Get object by object id, will recursively dereference a referenced object.
	pub fn get_object(&self, id: ObjectId) -> Result<&Object> {
		if let Some(object) = self.objects.get(&id) {
			if let Ok(id) = object.as_reference() {
				return self.get_object(id);
			} else {
				return Ok(object);
			}
		}
		Err(Error::ObjectNotFound)
	}

	/// Get mutable reference to object by object id, will iteratively dereference a referenced object.
	pub fn get_object_mut(&mut self, id: ObjectId) -> Result<&mut Object> {
		let mut object = self.objects.get(&id).ok_or(Error::ObjectNotFound)?;
		let mut cur_id = id;

		while let Ok(id) = object.as_reference() {
			cur_id = id;
			object = self.objects.get(&cur_id).ok_or(Error::ObjectNotFound)?;
		}

		Ok(self.objects.get_mut(&cur_id).unwrap())
	}

	/// Get dictionary object by id.
	pub fn get_dictionary(&self, id: ObjectId) -> Result<&Dictionary> {
		self.get_object(id).and_then(Object::as_dict)
	}

	/// Traverse objects from trailer recursively, return all referenced object IDs.
	pub fn traverse_objects<A: Fn(&mut Object) -> ()>(&mut self, action: A) -> Vec<ObjectId> {
		fn traverse_array<A: Fn(&mut Object) -> ()>(array: &mut Vec<Object>, action: &A, refs: &mut Vec<ObjectId>) {
			for item in array.iter_mut() {
				traverse_object(item, action, refs);
			}
		}
		fn traverse_dictionary<A: Fn(&mut Object) -> ()>(dict: &mut Dictionary, action: &A, refs: &mut Vec<ObjectId>) {
			for (_, v) in dict.iter_mut() {
				traverse_object(v, action, refs);
			}
		}
		fn traverse_object<A: Fn(&mut Object) -> ()>(object: &mut Object, action: &A, refs: &mut Vec<ObjectId>) {
			action(object);
			match *object {
				Object::Array(ref mut array) => traverse_array(array, action, refs),
				Object::Dictionary(ref mut dict) => traverse_dictionary(dict, action, refs),
				Object::Stream(ref mut stream) => traverse_dictionary(&mut stream.dict, action, refs),
				Object::Reference(id) => {
					if !refs.contains(&id) {
						refs.push(id);
					}
				}
				_ => {}
			}
		}
		let mut refs = vec![];
		traverse_dictionary(&mut self.trailer, &action, &mut refs);
		let mut index = 0;
		while index < refs.len() {
			if let Some(object) = self.objects.get_mut(&refs[index]) {
				traverse_object(object, &action, &mut refs);
			}
			index += 1;
		}
		refs
	}

	/// Get catalog dictionary.
	pub fn catalog(&self) -> Result<&Dictionary> {
		self.trailer.get(b"Root").and_then(Object::as_reference).and_then(|id| self.get_dictionary(id))
	}

	/// Get page numbers and corresponding object ids.
	pub fn get_pages(&self) -> BTreeMap<u32, ObjectId> {
		fn collect_pages(doc: &Document, page_tree_id: ObjectId, page_number: &mut u32, pages: &mut BTreeMap<u32, ObjectId>) {
			if let Ok(kids) = doc.get_dictionary(page_tree_id).and_then(|page_tree| page_tree.get(b"Kids")).and_then(Object::as_array) {
				for kid in kids {
					if let Ok(kid_id) = kid.as_reference() {
						if let Ok(type_name) = doc.get_dictionary(kid_id).and_then(Dictionary::type_name) {
							match type_name {
								"Page" => {
									pages.insert(*page_number, kid_id);
									*page_number += 1;
								}
								"Pages" => {
									collect_pages(doc, kid_id, page_number, pages);
								}
								_ => {}
							}
						}
					}
				}
			}
		}

		let mut pages = BTreeMap::new();
		let mut page_number = 1;
		if let Ok(page_tree_id) = self.catalog().and_then(|cat| cat.get(b"Pages")).and_then(Object::as_reference) {
			collect_pages(self, page_tree_id, &mut page_number, &mut pages);
		}
		pages
	}

	/// Get content stream object ids of a page.
	pub fn get_page_contents(&self, page_id: ObjectId) -> Vec<ObjectId> {
		let mut streams = vec![];
		if let Ok(page) = self.get_dictionary(page_id) {
			if let Ok(contents) = page.get(b"Contents") {
				match *contents {
					Object::Reference(ref id) => {
						streams.push(*id);
					}
					Object::Array(ref arr) => {
						for content in arr {
							if let Ok(id) = content.as_reference() { streams.push(id) }
						}
					}
					_ => {}
				}
			}
		}
		streams
	}

	/// Get content of a page.
	pub fn get_page_content(&self, page_id: ObjectId) -> Result<Vec<u8>> {
		let mut content = Vec::new();
		let content_streams = self.get_page_contents(page_id);
		for object_id in content_streams {
			if let Ok(content_stream) = self.get_object(object_id).and_then(Object::as_stream) {
				match content_stream.decompressed_content() {
					Ok(data) => content.write_all(&data)?,
					Err(_) => content.write_all(&content_stream.content)?,
				};
			}
		}
		Ok(content)
	}

	/// Get decoded page content;
	pub fn get_and_decode_page_content(&self, page_id: ObjectId) -> Result<Content> {
		let content_data = self.get_page_content(page_id)?;
		Content::decode(&content_data)
	}

	/// Get resources used by a page.
	pub fn get_page_resources(&self, page_id: ObjectId) -> (Option<&Dictionary>, Vec<ObjectId>) {
		fn collect_resources(page_node: &Dictionary, resource_ids: &mut Vec<ObjectId>, doc: &Document) {
			if let Ok(resources_id) = page_node.get(b"Resources").and_then(Object::as_reference) {
				resource_ids.push(resources_id);
			}
			if let Ok(page_tree) = page_node.get(b"Parent").and_then(Object::as_reference).and_then(|id| doc.get_dictionary(id)) {
				collect_resources(page_tree, resource_ids, doc);
			}
		};

		let mut resource_dict = None;
		let mut resource_ids = Vec::new();
		if let Ok(page) = self.get_dictionary(page_id) {
			resource_dict = page.get(b"Resources").and_then(Object::as_dict).ok();
			collect_resources(page, &mut resource_ids, self);
		}
		(resource_dict, resource_ids)
	}

	/// Get fonts used by a page.
	pub fn get_page_fonts(&self, page_id: ObjectId) -> BTreeMap<Vec<u8>, &Dictionary> {
		fn collect_fonts_from_resources<'a>(resources: &'a Dictionary, fonts: &mut BTreeMap<Vec<u8>, &'a Dictionary>, doc: &'a Document) {
			if let Ok(font_dict) = resources.get(b"Font").and_then(Object::as_dict) {
				for (name, value) in font_dict.iter() {
					let font = match *value {
						Object::Reference(id) => doc.get_dictionary(id).ok(),
						Object::Dictionary(ref dict) => Some(dict),
						_ => None,
					};
					if !fonts.contains_key(name) {
						font.map(|font| fonts.insert(name.clone(), font));
					}
				}
			}
		};

		let mut fonts = BTreeMap::new();
		let (resource_dict, resource_ids) = self.get_page_resources(page_id);
		if let Some(resources) = resource_dict {
			collect_fonts_from_resources(resources, &mut fonts, self);
		}
		for resource_id in resource_ids {
			if let Ok(resources) = self.get_dictionary(resource_id) {
				collect_fonts_from_resources(resources, &mut fonts, self);
			}
		}
		fonts
	}

	pub fn decode_text(encoding: Option<&str>, bytes: &[u8]) -> String {
		if let Some(encoding) = encoding {
			info!("{}", encoding);
			match encoding {
				"StandardEncoding" => bytes_to_string(encodings::STANDARD_ENCODING, bytes),
				"MacRomanEncoding" => bytes_to_string(encodings::MAC_ROMAN_ENCODING, bytes),
				"MacExpertEncoding" => bytes_to_string(encodings::MAC_EXPERT_ENCODING, bytes),
				"WinAnsiEncoding" => bytes_to_string(encodings::WIN_ANSI_ENCODING, bytes),
				"UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.decode(bytes, DecoderTrap::Ignore).unwrap(),
				"Identity-H" => "?Identity-H Unimplemented?".to_string(), // Unimplemented
				_ => String::from_utf8_lossy(bytes).to_string(),
			}
		} else {
			bytes_to_string(encodings::STANDARD_ENCODING, bytes)
		}
	}

	pub fn encode_text(encoding: Option<&str>, text: &str) -> Vec<u8> {
		if let Some(encoding) = encoding {
			match encoding {
				"StandardEncoding" => string_to_bytes(encodings::STANDARD_ENCODING, text),
				"MacRomanEncoding" => string_to_bytes(encodings::MAC_ROMAN_ENCODING, text),
				"MacExpertEncoding" => string_to_bytes(encodings::MAC_EXPERT_ENCODING, text),
				"WinAnsiEncoding" => string_to_bytes(encodings::WIN_ANSI_ENCODING, text),
				"UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.encode(text, EncoderTrap::Ignore).unwrap(),
				"Identity-H" => vec![], // Unimplemented
				_ => text.as_bytes().to_vec(),
			}
		} else {
			string_to_bytes(encodings::STANDARD_ENCODING, text)
		}
	}
}

impl Default for Document {
	fn default() -> Self {
		Self::new()
	}
}
<|endoftext|>


<commit_before>use std::time;
use std::fmt;
use std::ops;
use std::convert;
use std::f64;

use chrono;

use super::error;
use super::error::DurationError;

pub const SECS_PER_DAY: f64 = 60.0*60.0*24.0;
pub const SECS_PER_HOUR: f64 = 60.0*60.0;
pub const SECS_PER_MINUTE: f64 = 60.0;

pub trait TimePoint<Rhs=Self> {
    fn duration_between(&self, rhs: &Rhs) -> FloatDuration;
}

#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
pub struct FloatDuration {
    secs: f64,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct DecomposedTime {
    pub days: u32,
    pub hours: u32,
    pub minutes: u32,
    pub seconds: u32,
    pub fractional_seconds: f64,
    pub sign: i8,
}

impl FloatDuration {
    pub fn new(secs: f64) -> FloatDuration {
        FloatDuration {secs}
    }

    pub fn days(days: f64) -> FloatDuration {
        FloatDuration {secs: days * (3600.0 * 24.0)}
    }
    pub fn hours(hours: f64) -> FloatDuration {
        FloatDuration {secs: hours * 3600.0}
    }
    pub fn minutes(mins: f64) -> FloatDuration {
        FloatDuration {secs: mins * 60.0}
    }
    pub fn seconds(secs: f64) -> FloatDuration {
        FloatDuration {secs}
    }
    pub fn milliseconds(millis: f64) -> FloatDuration {
        FloatDuration {secs: millis / 1000.0}
    }
    pub fn microseconds(micros: f64) -> FloatDuration {
        FloatDuration {secs: micros / 1.0e6}
    }
    pub fn nanoseconds(nanos: f64) -> FloatDuration {
        FloatDuration {secs: nanos / 1.0e9}
    }

    pub fn decompose(&self) -> DecomposedTime {
        let mut secs_left = self.secs;

        let days = (secs_left / SECS_PER_DAY).trunc();
        secs_left -= days*SECS_PER_DAY;
        let hours = (secs_left / SECS_PER_HOUR).trunc();
        secs_left -= hours*SECS_PER_HOUR;
        let minutes = (secs_left / SECS_PER_MINUTE).trunc();
        secs_left -= minutes*SECS_PER_MINUTE;
        let seconds = secs_left.trunc();
        let fract = secs_left.fract();

        DecomposedTime {
            days: days as u32,
            hours: hours as u32,
            minutes: minutes as u32,
            seconds: seconds as u32,
            fractional_seconds: fract,
            sign: self.secs.signum() as i8,
        }
    }

    pub fn as_days(&self) -> f64 {
        self.secs / SECS_PER_DAY
    }
    pub fn as_hours(&self) -> f64 {
        self.secs / SECS_PER_HOUR
    }
    pub fn as_minutes(&self) -> f64 {
        self.secs / SECS_PER_MINUTE
    }
    pub fn as_seconds(&self) -> f64 {
        self.secs
    }
    pub fn as_milliseconds(&self) -> f64 {
        self.secs * 1.0e3
    }
    pub fn as_microseconds(&self) -> f64 {
        self.secs * 1.0e6
    }
    pub fn as_nanoseconds(&self) -> f64 {
        self.secs * 1.0e9
    }

    pub fn zero() -> FloatDuration {
        FloatDuration {secs: 0.0}
    }
    pub fn is_zero(&self) -> bool {
        self.secs == 0.0
    }
    pub fn is_positive(&self) -> bool {
        self.secs.is_sign_positive()
    }
    pub fn is_negative(&self) -> bool {
        self.secs.is_sign_negative()
    }

    pub fn min_value() -> FloatDuration {
        FloatDuration { secs: f64::MIN }
    }
    pub fn max_value() -> FloatDuration {
        FloatDuration { secs: f64::MAX }
    }

    pub fn as_std(&self) -> error::Result<time::Duration> {
        if self.secs.is_sign_negative() {
            Err(DurationError::StdOutOfRange)
        } else {
            let seconds = self.secs.trunc();
            let nanos = self.secs.fract() * 1e9;

            Ok(time::Duration::new(seconds as u64, nanos as u32))
        }
    }

    pub fn from_std(duration: time::Duration) -> error::Result<FloatDuration> {
        Ok(FloatDuration::seconds(
            (duration.as_secs() as f64) + (duration.subsec_nanos() as f64)*1.0e-9))
    }

    pub fn from_decomposed(decomposed: &DecomposedTime) -> FloatDuration {
        let seconds = (decomposed.days as f64)*SECS_PER_DAY
            + (decomposed.hours as f64)*SECS_PER_HOUR
            + (decomposed.minutes as f64)*SECS_PER_MINUTE
            + (decomposed.seconds as f64)
            + decomposed.fractional_seconds;

        FloatDuration {
            secs: seconds
        }
    }
}

impl FloatDuration {
    pub fn as_chrono_duration(&self) -> error::Result<chrono::Duration> {
        let std_duration = self.as_std()?;
        chrono::Duration::from_std(std_duration).map_err(|e| convert::From::from(e))
    }

    pub fn from_chrono_duration(duration: &chrono::Duration)
            -> error::Result<FloatDuration>
    {
        let std_duration = duration.to_std()?;
        FloatDuration::from_std(std_duration)
    }
}

impl DecomposedTime {
    pub fn new() -> DecomposedTime {
        DecomposedTime { days: 0, hours: 0, minutes: 0, seconds: 0,
            fractional_seconds: 0.0, sign: 1 }
    }

    pub fn from_components(days: u32, hours: u32, minutes: u32,
               seconds: u32, fractional_seconds: f64) -> DecomposedTime
    {
        DecomposedTime {
            days, hours, minutes, seconds, fractional_seconds, sign: 1}
    }

    pub fn negate(mut self) -> DecomposedTime {
        self.sign *= -1;
        self
    }
}

impl fmt::Display for DecomposedTime {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        if self.days > 0 {
            write!(fmt, "{}d ", self.days as u64)?;
        }

        if self.sign.is_negative() {
            write!(fmt, "-")?;
        }
        write!(fmt, "{:02}:{:02}:{:02}.{}",
            self.hours,
            self.minutes,
            self.seconds,
            self.fractional_seconds)
    }
}

impl fmt::Display for FloatDuration {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        if self.secs > SECS_PER_DAY {
            write!(fmt, "{} days", self.as_days())
        } else if self.secs > SECS_PER_HOUR {
            write!(fmt, "{} hours", self.as_hours())
        } else if self.secs > SECS_PER_MINUTE {
            write!(fmt, "{} minutes", self.as_minutes())
        } else if self.secs > 1.0 {
            write!(fmt, "{} seconds", self.as_seconds())
        } else if self.secs > 1.0e-3 {
            write!(fmt, "{} milliseconds", self.as_milliseconds())
        } else if self.secs > 1.0e-6 {
            write!(fmt, "{} microseconds", self.as_microseconds())
        } else {
            write!(fmt, "{} nanoseconds", self.as_nanoseconds())
        }
    }
}

impl ops::Neg for FloatDuration {
    type Output = FloatDuration;

    fn neg(self) -> FloatDuration {
        FloatDuration {secs: -self.secs}
    }
}

impl ops::Add<FloatDuration> for FloatDuration {
    type Output = FloatDuration;

    fn add(self, rhs: FloatDuration) -> FloatDuration {
        FloatDuration {secs: self.secs + rhs.secs}
    }
}
impl ops::Sub<FloatDuration> for FloatDuration {
    type Output = FloatDuration;

    fn sub(self, rhs: FloatDuration) -> FloatDuration {
        FloatDuration {secs: self.secs - rhs.secs}
    }
}

impl ops::Mul<f64> for FloatDuration {
    type Output = FloatDuration;

    fn mul(self, rhs: f64) -> FloatDuration {
        FloatDuration {secs: self.secs * rhs}
    }
}
impl ops::Div<f64> for FloatDuration {
    type Output = FloatDuration;

    fn div(self, rhs: f64) -> FloatDuration {
        FloatDuration {secs: self.secs / rhs}
    }
}

impl ops::AddAssign<FloatDuration> for FloatDuration {
    fn add_assign(&mut self, rhs: FloatDuration) {
        self.secs += rhs.secs;
    }
}
impl ops::SubAssign<FloatDuration> for FloatDuration {
    fn sub_assign(&mut self, rhs: FloatDuration) {
        self.secs -= rhs.secs;
    }
}

impl ops::MulAssign<f64> for FloatDuration {
    fn mul_assign(&mut self, rhs: f64) {
        self.secs *= rhs;
    }
}
impl ops::DivAssign<f64> for FloatDuration {
    fn div_assign(&mut self, rhs: f64) {
        self.secs /= rhs;
    }
}

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

    #[test]
    fn test_construct() {
        let duration1 = FloatDuration::hours(3.0);
        assert_eq!(duration1.as_hours(), 3.0);
        assert_eq!(duration1.as_minutes(), 180.0);
        assert_eq!(duration1.as_seconds(), 180.0*60.0);
        assert_eq!(duration1.as_days(), 3.0/24.0);
        assert_eq!(duration1.as_milliseconds(), 180.0*60.0*1000.0);

        let duration2 = FloatDuration::milliseconds(55.0);
        assert_eq!(duration2.as_seconds(), 0.055);
        assert_eq!(duration2.as_milliseconds(), 55.0);
        assert_eq!(duration2.as_microseconds(), 55000.0);
        assert_eq!(duration2.as_nanoseconds(), 55000000.0);
        assert!(!duration2.is_zero());

        let duration3 = FloatDuration::zero();
        assert!(duration3.is_zero());

        assert_eq!(FloatDuration::days(1.5), FloatDuration::hours(36.0));
        assert_eq!(FloatDuration::minutes(30.0), FloatDuration::hours(0.5));
        assert_eq!(FloatDuration::seconds(180.0), FloatDuration::minutes(3.0));
        assert_eq!(FloatDuration::seconds(3.5), FloatDuration::milliseconds(3500.0));
        assert_eq!(FloatDuration::microseconds(300.0), FloatDuration::milliseconds(0.30));
        assert_eq!(FloatDuration::nanoseconds(1000.0), FloatDuration::microseconds(1.0));
    }

    #[test]
    fn test_arithmetic() {
        let duration1 = FloatDuration::minutes(5.0) + FloatDuration::seconds(30.0);
        assert_eq!(duration1, FloatDuration::seconds(330.0));

        let duration2 = FloatDuration::hours(3.0) * 2.5;
        assert_eq!(duration2, FloatDuration::hours(7.5));

        let duration3 = FloatDuration::days(3.0) / 3.0 - FloatDuration::hours(2.0);
        assert_eq!(duration3, FloatDuration::hours(22.0));

        let duration4 = FloatDuration::zero() + FloatDuration::milliseconds(500.0)
            + FloatDuration::microseconds(500.0);
        assert_eq!(duration4, FloatDuration::microseconds(500500.0));

    }

    #[test]
    fn test_std_conversion() {
        let duration1 = FloatDuration::minutes(5.0);
        let std_duration1 = duration1.as_std().unwrap();
        assert!(duration1.is_positive());
        assert_eq!(std_duration1, time::Duration::new(300, 0));
        assert_eq!(FloatDuration::from_std(std_duration1).unwrap(), duration1);

        let duration2 = FloatDuration::hours(-2.0);
        assert!(duration2.is_negative());
        assert!(!duration2.as_std().is_ok());
        let std_duration2 = (-duration2).as_std().unwrap();
        assert_eq!(std_duration2, time::Duration::new(3600*2, 0));
        assert_eq!(FloatDuration::from_std(std_duration2).unwrap(), -duration2);
    }

    #[test]
    fn test_display() {
        use std::fmt::Write;

        let mut buffer1 = "".to_string();
        let mut buffer1_2 = "".to_string();
        let duration1 = FloatDuration::minutes(3.5);

        write!(buffer1, "{}", duration1).unwrap();
        write!(buffer1_2, "{}", duration1.decompose()).unwrap();
        assert_eq!(buffer1, "3.5 minutes");
        assert_eq!(buffer1_2, "00:03:30.0");

        let mut buffer2 = "".to_string();
        let mut buffer2_2 = "".to_string();
        let duration2 = FloatDuration::days(3.0) + FloatDuration::hours(12.0);

        write!(buffer2, "{}", duration2).unwrap();
        write!(buffer2_2, "{}", duration2.decompose()).unwrap();
        assert_eq!(buffer2, "3.5 days");
        assert_eq!(buffer2_2, "3d 12:00:00.0");
    }
}
<commit_msg>Added interoperability between FloatDuration and std::time, chrono.<commit_after>use std::time;
use std::fmt;
use std::ops;
use std::f64;

use chrono;

use super::error;
use super::error::DurationError;

pub const SECS_PER_DAY: f64 = 60.0*60.0*24.0;
pub const SECS_PER_HOUR: f64 = 60.0*60.0;
pub const SECS_PER_MINUTE: f64 = 60.0;

pub trait TimePoint<Rhs=Self> {
    fn float_duration_since(self, rhs: Rhs) -> error::Result<FloatDuration>;
}

#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
pub struct FloatDuration {
    secs: f64,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct DecomposedTime {
    pub days: u32,
    pub hours: u32,
    pub minutes: u32,
    pub seconds: u32,
    pub fractional_seconds: f64,
    pub sign: i8,
}

impl FloatDuration {
    pub fn new(secs: f64) -> FloatDuration {
        FloatDuration {secs}
    }

    pub fn days(days: f64) -> FloatDuration {
        FloatDuration {secs: days * (3600.0 * 24.0)}
    }
    pub fn hours(hours: f64) -> FloatDuration {
        FloatDuration {secs: hours * 3600.0}
    }
    pub fn minutes(mins: f64) -> FloatDuration {
        FloatDuration {secs: mins * 60.0}
    }
    pub fn seconds(secs: f64) -> FloatDuration {
        FloatDuration {secs}
    }
    pub fn milliseconds(millis: f64) -> FloatDuration {
        FloatDuration {secs: millis / 1000.0}
    }
    pub fn microseconds(micros: f64) -> FloatDuration {
        FloatDuration {secs: micros / 1.0e6}
    }
    pub fn nanoseconds(nanos: f64) -> FloatDuration {
        FloatDuration {secs: nanos / 1.0e9}
    }

    pub fn decompose(&self) -> DecomposedTime {
        let mut secs_left = self.secs;

        let days = (secs_left / SECS_PER_DAY).trunc();
        secs_left -= days*SECS_PER_DAY;
        let hours = (secs_left / SECS_PER_HOUR).trunc();
        secs_left -= hours*SECS_PER_HOUR;
        let minutes = (secs_left / SECS_PER_MINUTE).trunc();
        secs_left -= minutes*SECS_PER_MINUTE;
        let seconds = secs_left.trunc();
        let fract = secs_left.fract();

        DecomposedTime {
            days: days as u32,
            hours: hours as u32,
            minutes: minutes as u32,
            seconds: seconds as u32,
            fractional_seconds: fract,
            sign: self.secs.signum() as i8,
        }
    }

    pub fn as_days(&self) -> f64 {
        self.secs / SECS_PER_DAY
    }
    pub fn as_hours(&self) -> f64 {
        self.secs / SECS_PER_HOUR
    }
    pub fn as_minutes(&self) -> f64 {
        self.secs / SECS_PER_MINUTE
    }
    pub fn as_seconds(&self) -> f64 {
        self.secs
    }
    pub fn as_milliseconds(&self) -> f64 {
        self.secs * 1.0e3
    }
    pub fn as_microseconds(&self) -> f64 {
        self.secs * 1.0e6
    }
    pub fn as_nanoseconds(&self) -> f64 {
        self.secs * 1.0e9
    }

    pub fn abs(self) -> FloatDuration {
        FloatDuration { secs: self.secs.abs() }
    }
    pub fn zero() -> FloatDuration {
        FloatDuration {secs: 0.0}
    }
    pub fn is_zero(&self) -> bool {
        self.secs == 0.0
    }
    pub fn is_positive(&self) -> bool {
        self.secs.is_sign_positive()
    }
    pub fn is_negative(&self) -> bool {
        self.secs.is_sign_negative()
    }

    pub fn min_value() -> FloatDuration {
        FloatDuration { secs: f64::MIN }
    }
    pub fn max_value() -> FloatDuration {
        FloatDuration { secs: f64::MAX }
    }

    pub fn as_std(&self) -> error::Result<time::Duration> {
        if self.secs.is_sign_negative() {
            Err(DurationError::StdOutOfRange)
        } else {
            let seconds = self.secs.trunc();
            let nanos = self.secs.fract() * 1e9;

            Ok(time::Duration::new(seconds as u64, nanos as u32))
        }
    }

    pub fn from_std(duration: time::Duration) -> error::Result<FloatDuration> {
        Ok(FloatDuration::seconds(
            (duration.as_secs() as f64) + (duration.subsec_nanos() as f64)*1.0e-9))
    }

    pub fn from_decomposed(decomposed: &DecomposedTime) -> FloatDuration {
        let seconds = (decomposed.days as f64)*SECS_PER_DAY
            + (decomposed.hours as f64)*SECS_PER_HOUR
            + (decomposed.minutes as f64)*SECS_PER_MINUTE
            + (decomposed.seconds as f64)
            + decomposed.fractional_seconds;

        FloatDuration {
            secs: seconds
        }
    }
}

impl FloatDuration {
    pub fn as_chrono_duration(&self) -> error::Result<chrono::Duration> {
        let is_negative = self.is_negative();
        let std_duration = self.abs().as_std()?;
        let chrono_duration = chrono::Duration::from_std(std_duration)?;
        if is_negative {
            Ok(-chrono_duration)
        } else {
            Ok(chrono_duration)
        }
    }

    pub fn from_chrono_duration(duration: &chrono::Duration)
            -> error::Result<FloatDuration>
    {
        let is_negative = duration.num_seconds() < 0;

        let std_duration =
            if is_negative {
                (-*duration).to_std()?
            } else {
                duration.to_std()?
            };

        let float_duration = FloatDuration::from_std(std_duration)?;
        if is_negative {
            Ok(-float_duration)
        } else {
            Ok(float_duration)
        }
    }
}

impl DecomposedTime {
    pub fn new() -> DecomposedTime {
        DecomposedTime { days: 0, hours: 0, minutes: 0, seconds: 0,
            fractional_seconds: 0.0, sign: 1 }
    }

    pub fn from_components(days: u32, hours: u32, minutes: u32,
               seconds: u32, fractional_seconds: f64) -> DecomposedTime
    {
        DecomposedTime {
            days, hours, minutes, seconds, fractional_seconds, sign: 1}
    }

    pub fn negate(mut self) -> DecomposedTime {
        self.sign *= -1;
        self
    }
}

impl<Tz: chrono::TimeZone> TimePoint for chrono::DateTime<Tz> {
    fn float_duration_since(self, since: chrono::DateTime<Tz>)
            -> error::Result<FloatDuration>
    {
        let chrono_duration = self.signed_duration_since(since);
        FloatDuration::from_chrono_duration(&chrono_duration)
    }
}
impl TimePoint for time::Instant {
    fn float_duration_since(self, since: time::Instant) -> error::Result<FloatDuration> {
        let std_duration = self.duration_since(since);
        FloatDuration::from_std(std_duration)
    }
}
impl TimePoint for time::SystemTime {
    fn float_duration_since(self, since: time::SystemTime) -> error::Result<FloatDuration> {
        let std_duration = self.duration_since(since)?;
        FloatDuration::from_std(std_duration)
    }
}

impl fmt::Display for DecomposedTime {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        if self.days > 0 {
            write!(fmt, "{}d ", self.days as u64)?;
        }

        if self.sign.is_negative() {
            write!(fmt, "-")?;
        }
        write!(fmt, "{:02}:{:02}:{:02}.{}",
            self.hours,
            self.minutes,
            self.seconds,
            self.fractional_seconds)
    }
}

impl fmt::Display for FloatDuration {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        if self.secs > SECS_PER_DAY {
            write!(fmt, "{} days", self.as_days())
        } else if self.secs > SECS_PER_HOUR {
            write!(fmt, "{} hours", self.as_hours())
        } else if self.secs > SECS_PER_MINUTE {
            write!(fmt, "{} minutes", self.as_minutes())
        } else if self.secs > 1.0 {
            write!(fmt, "{} seconds", self.as_seconds())
        } else if self.secs > 1.0e-3 {
            write!(fmt, "{} milliseconds", self.as_milliseconds())
        } else if self.secs > 1.0e-6 {
            write!(fmt, "{} microseconds", self.as_microseconds())
        } else {
            write!(fmt, "{} nanoseconds", self.as_nanoseconds())
        }
    }
}

impl ops::Neg for FloatDuration {
    type Output = FloatDuration;

    fn neg(self) -> FloatDuration {
        FloatDuration {secs: -self.secs}
    }
}

impl ops::Add<FloatDuration> for FloatDuration {
    type Output = FloatDuration;

    fn add(self, rhs: FloatDuration) -> FloatDuration {
        FloatDuration {secs: self.secs + rhs.secs}
    }
}
impl ops::Sub<FloatDuration> for FloatDuration {
    type Output = FloatDuration;

    fn sub(self, rhs: FloatDuration) -> FloatDuration {
        FloatDuration {secs: self.secs - rhs.secs}
    }
}

impl ops::Mul<f64> for FloatDuration {
    type Output = FloatDuration;

    fn mul(self, rhs: f64) -> FloatDuration {
        FloatDuration {secs: self.secs * rhs}
    }
}
impl ops::Div<f64> for FloatDuration {
    type Output = FloatDuration;

    fn div(self, rhs: f64) -> FloatDuration {
        FloatDuration {secs: self.secs / rhs}
    }
}

impl ops::AddAssign<FloatDuration> for FloatDuration {
    fn add_assign(&mut self, rhs: FloatDuration) {
        self.secs += rhs.secs;
    }
}
impl ops::SubAssign<FloatDuration> for FloatDuration {
    fn sub_assign(&mut self, rhs: FloatDuration) {
        self.secs -= rhs.secs;
    }
}

impl ops::MulAssign<f64> for FloatDuration {
    fn mul_assign(&mut self, rhs: f64) {
        self.secs *= rhs;
    }
}
impl ops::DivAssign<f64> for FloatDuration {
    fn div_assign(&mut self, rhs: f64) {
        self.secs /= rhs;
    }
}

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

    #[test]
    fn test_construct() {
        let duration1 = FloatDuration::hours(3.0);
        assert_eq!(duration1.as_hours(), 3.0);
        assert_eq!(duration1.as_minutes(), 180.0);
        assert_eq!(duration1.as_seconds(), 180.0*60.0);
        assert_eq!(duration1.as_days(), 3.0/24.0);
        assert_eq!(duration1.as_milliseconds(), 180.0*60.0*1000.0);

        let duration2 = FloatDuration::milliseconds(55.0);
        assert_eq!(duration2.as_seconds(), 0.055);
        assert_eq!(duration2.as_milliseconds(), 55.0);
        assert_eq!(duration2.as_microseconds(), 55000.0);
        assert_eq!(duration2.as_nanoseconds(), 55000000.0);
        assert!(!duration2.is_zero());

        let duration3 = FloatDuration::zero();
        assert!(duration3.is_zero());

        assert_eq!(FloatDuration::days(1.5), FloatDuration::hours(36.0));
        assert_eq!(FloatDuration::minutes(30.0), FloatDuration::hours(0.5));
        assert_eq!(FloatDuration::seconds(180.0), FloatDuration::minutes(3.0));
        assert_eq!(FloatDuration::seconds(3.5), FloatDuration::milliseconds(3500.0));
        assert_eq!(FloatDuration::microseconds(300.0), FloatDuration::milliseconds(0.30));
        assert_eq!(FloatDuration::nanoseconds(1000.0), FloatDuration::microseconds(1.0));
    }

    #[test]
    fn test_arithmetic() {
        let duration1 = FloatDuration::minutes(5.0) + FloatDuration::seconds(30.0);
        assert_eq!(duration1, FloatDuration::seconds(330.0));

        let duration2 = FloatDuration::hours(3.0) * 2.5;
        assert_eq!(duration2, FloatDuration::hours(7.5));

        let duration3 = FloatDuration::days(3.0) / 3.0 - FloatDuration::hours(2.0);
        assert_eq!(duration3, FloatDuration::hours(22.0));

        let duration4 = FloatDuration::zero() + FloatDuration::milliseconds(500.0)
            + FloatDuration::microseconds(500.0);
        assert_eq!(duration4, FloatDuration::microseconds(500500.0));

    }

    #[test]
    fn test_std_conversion() {
        let duration1 = FloatDuration::minutes(5.0);
        let std_duration1 = duration1.as_std().unwrap();
        assert!(duration1.is_positive());
        assert_eq!(std_duration1, time::Duration::new(300, 0));
        assert_eq!(FloatDuration::from_std(std_duration1).unwrap(), duration1);

        let duration2 = FloatDuration::hours(-2.0);
        assert!(duration2.is_negative());
        assert!(!duration2.as_std().is_ok());
        let std_duration2 = (-duration2).as_std().unwrap();
        assert_eq!(std_duration2, time::Duration::new(3600*2, 0));
        assert_eq!(FloatDuration::from_std(std_duration2).unwrap(), -duration2);
    }

    #[test]
    fn test_display() {
        use std::fmt::Write;

        let mut buffer1 = "".to_string();
        let mut buffer1_2 = "".to_string();
        let duration1 = FloatDuration::minutes(3.5);

        write!(buffer1, "{}", duration1).unwrap();
        write!(buffer1_2, "{}", duration1.decompose()).unwrap();
        assert_eq!(buffer1, "3.5 minutes");
        assert_eq!(buffer1_2, "00:03:30.0");

        let mut buffer2 = "".to_string();
        let mut buffer2_2 = "".to_string();
        let duration2 = FloatDuration::days(3.0) + FloatDuration::hours(12.0);

        write!(buffer2, "{}", duration2).unwrap();
        write!(buffer2_2, "{}", duration2.decompose()).unwrap();
        assert_eq!(buffer2, "3.5 days");
        assert_eq!(buffer2_2, "3d 12:00:00.0");
    }
}
<|endoftext|>