<commit_before>use std::cell::RefCell;
use std::collections::{HashMap, VecDeque};
use std::ops::Deref;
use std::rc::Rc;

use crate::ast::expressions;
use crate::interpreter::environment;
use crate::utils;

pub enum Type {
    Nil,
    Boolean(bool),
    Number(f64),
    String(String),
    /// Reference to an existing value
    Reference(Rc<RefCell<Type>>),
    Vector(VecDeque<Type>),
    Table {
        /// For comparison
        id: u64,
        map: HashMap<Type, Rc<RefCell<Type>>>,
        metatable: HashMap<String, Type>,
        border: usize,
    },
    Function {
        /// For comparison
        id: u64,
        parameters: Vec<String>,
        varargs: bool,
        body: Rc<Box<expressions::Expression>>,
        // XXX: Capture only vars function needs?
        env: utils::Shared<environment::Environment>,
    },
}

impl Type {
    pub fn call(&self, _arguments: Vec<&Type>) -> Type {
        unimplemented!();
    }

    pub fn as_bool(&self) -> bool {
        match self {
            Type::Nil => false,
            Type::Boolean(false) => false,
            _ => true,
        }
    }

    /// Check if type is nil. We often have special cases for nils
    pub fn is_nil(&self) -> bool {
        match self {
            Type::Nil => true,
            Type::Reference(typeref) => typeref.borrow().is_nil(),
            _ => false,
        }
    }

    /// Create reference to an object or clone reference
    pub fn as_ref(self) -> Rc<RefCell<Self>> {
        match self {
            Type::Reference(typeref) => typeref,
            _ => Rc::new(RefCell::new(self)),
        }
    }
}

#[cfg(test)]
impl ::std::cmp::PartialEq<&'static str> for Type {
    fn eq(&self, other: &&'static str) -> bool {
        format!("{:?}", self) == other.to_string()
    }
}

impl ::std::cmp::PartialEq for Type {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Type::Boolean(left), Type::Boolean(right)) => left == right,
            (Type::Number(left), Type::Number(right)) => left == right,
            (Type::String(left), Type::String(right)) => left == right,
            (Type::Reference(left), right) => right.eq(left.borrow().deref()),
            (left, Type::Reference(right)) => left.eq(right.borrow().deref()),
            (Type::Vector(left), Type::Vector(right)) => left == right,
            (Type::Table { id: left, .. }, Type::Table { id: right, .. }) => left == right,
            (Type::Function { id: left, .. }, Type::Function { id: right, .. }) => left == right,
            _ => false,
        }
    }
}

impl ::std::cmp::Eq for Type {}

impl ::std::hash::Hash for Type {
    fn hash<H: ::std::hash::Hasher>(&self, state: &mut H) {
        match self {
            Type::Nil => 1.hash(state),
            Type::Boolean(value) => value.hash(state),
            Type::Number(value) => value.to_string().hash(state),
            Type::String(value) => value.hash(state),
            Type::Reference(value) => value.borrow().hash(state),
            Type::Vector(vec) => vec.hash(state),
            Type::Table { id, .. } => id.hash(state),
            Type::Function { id, .. } => id.hash(state),
        }
    }
}

/// Display to correctly infor user about runtime errors
impl ::std::fmt::Display for Type {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        match self {
            Type::Function { id, .. } => write!(f, "function ({:x})", id),
            Type::Table { id, .. } => write!(f, "table ({:x})", id),
            Type::Reference(value) => value.borrow().fmt(f),
            _ => write!(f, "{:?}", self),
        }
    }
}

/// Debug, which breaks closured env circular dependency
impl ::std::fmt::Debug for Type {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        match self {
            Type::Nil => write!(f, "Nil"),
            Type::Boolean(value) => write!(f, "Boolean({:?})", value),
            Type::Number(value) => write!(f, "Number({:?})", value),
            Type::String(value) => write!(f, "String({:?})", value),
            Type::Reference(value) => write!(f, "Reference({:?})", value),
            Type::Vector(vec) => write!(f, "Vector({:?})", vec),
            Type::Table {
                id,
                map,
                metatable,
                border,
            } => write!(
                f,
                "Table {{ id: {}, map: {:?}, metatable: {:?}, border: {} }}",
                id, map, metatable, border
            ),
            Type::Function {
                id,
                parameters,
                varargs,
                body,
                env,
            } => write!(
                f,
                "Function {{ id: {:?}, parameters: {:?}, varargs: {:?}, body: {:?}, env: {:?} }}",
                id,
                parameters,
                varargs,
                body,
                env.borrow().id()
            ),
        }
    }
}

/// Macro to use for pattern maching types with respect to type referencing
#[macro_export]
macro_rules! match_type {
    (($($typ:expr),+), $($pat:pat => $result:expr),+) => {{
        let typs = ($(if let $crate::interpreter::types::Type::Reference(value) = $typ { unsafe { &*value.as_ptr() } } else { $typ }),+);

        match typs {$(
            $pat => $result
        ), +}}
    };
    ($typ:expr, $($pat:pat => $result:expr),+) => {{
        let typ = if let $crate::interpreter::types::Type::Reference(value) = $typ { unsafe { &*value.as_ptr() } } else { $typ };

        match typ {$(
            $pat => $result
        ), +}}
    };
}

impl ::std::convert::AsRef<bool> for Type {
    fn as_ref(&self) -> &bool {
        match_type!(&self,
            Type::Boolean(val) => val,
            _ => panic!("Cannot convert vale {} to a boolean", self)
        )
    }
}

impl ::std::convert::AsRef<f64> for Type {
    fn as_ref(&self) -> &f64 {
        match_type!(&self,
            Type::Number(val) => val,
            _ => panic!("Cannot convert vale {} to a number", self)
        )
    }
}

impl ::std::convert::AsRef<String> for Type {
    fn as_ref(&self) -> &String {
        match_type!(&self,
            Type::String(val) => val,
            _ => panic!("Cannot convert vale {} to a string", self)
        )
    }
}

impl ::std::convert::AsRef<VecDeque<Type>> for Type {
    fn as_ref(&self) -> &VecDeque<Type> {
        match_type!(&self,
            Type::Vector(val) => val,
            _ => panic!("Cannot convert vale {} to a string", self)
        )
    }
}

impl ::std::convert::AsRef<HashMap<Type, Rc<RefCell<Type>>>> for Type {
    fn as_ref(&self) -> &HashMap<Type, Rc<RefCell<Type>>> {
        match_type!(&self,
            Type::Table { map, .. } => map,
            _ => panic!("Cannot convert vale {} to a string", self)
        )
    }
}<commit_msg>Minor cleanup<commit_after>use std::cell::RefCell;
use std::collections::{HashMap, VecDeque};
use std::ops::Deref;
use std::rc::Rc;

use crate::ast::expressions;
use crate::interpreter::environment;
use crate::utils;

pub enum Type {
    Nil,
    Boolean(bool),
    Number(f64),
    String(String),
    /// Reference to an existing value
    Reference(Rc<RefCell<Type>>),
    Vector(VecDeque<Type>),
    Table {
        /// For comparison
        id: u64,
        map: HashMap<Type, Rc<RefCell<Type>>>,
        metatable: HashMap<String, Type>,
        border: usize,
    },
    Function {
        /// For comparison
        id: u64,
        parameters: Vec<String>,
        varargs: bool,
        body: Rc<Box<expressions::Expression>>,
        // XXX: Capture only vars function needs?
        env: utils::Shared<environment::Environment>,
    },
}

impl Type {
    pub fn call(&self, _arguments: Vec<&Type>) -> Type {
        unimplemented!();
    }

    pub fn as_bool(&self) -> bool {
        match self {
            Type::Nil => false,
            Type::Boolean(false) => false,
            _ => true,
        }
    }

    /// Check if type is nil. We often have special cases for nils
    pub fn is_nil(&self) -> bool {
        match self {
            Type::Nil => true,
            Type::Reference(typeref) => typeref.borrow().is_nil(),
            _ => false,
        }
    }

    /// Create reference to an object or clone reference
    pub fn as_ref(self) -> Rc<RefCell<Self>> {
        match self {
            Type::Reference(typeref) => typeref,
            _ => Rc::new(RefCell::new(self)),
        }
    }
}

#[cfg(test)]
impl ::std::cmp::PartialEq<&'static str> for Type {
    fn eq(&self, other: &&'static str) -> bool {
        format!("{:?}", self) == other.to_string()
    }
}

impl ::std::cmp::PartialEq for Type {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Type::Boolean(left), Type::Boolean(right)) => left == right,
            (Type::Number(left), Type::Number(right)) => left == right,
            (Type::String(left), Type::String(right)) => left == right,
            (Type::Reference(left), right) => right.eq(left.borrow().deref()),
            (left, Type::Reference(right)) => left.eq(right.borrow().deref()),
            (Type::Vector(left), Type::Vector(right)) => left == right,
            (Type::Table { id: left, .. }, Type::Table { id: right, .. }) => left == right,
            (Type::Function { id: left, .. }, Type::Function { id: right, .. }) => left == right,
            _ => false,
        }
    }
}

impl ::std::cmp::Eq for Type {}

impl ::std::hash::Hash for Type {
    fn hash<H: ::std::hash::Hasher>(&self, state: &mut H) {
        match self {
            Type::Nil => 1.hash(state),
            Type::Boolean(value) => value.hash(state),
            Type::Number(value) => value.to_string().hash(state),
            Type::String(value) => value.hash(state),
            Type::Reference(value) => value.borrow().hash(state),
            Type::Vector(vec) => vec.hash(state),
            Type::Table { id, .. } => id.hash(state),
            Type::Function { id, .. } => id.hash(state),
        }
    }
}

/// Display to correctly infor user about runtime errors
impl ::std::fmt::Display for Type {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        match self {
            Type::Function { id, .. } => write!(f, "function ({:x})", id),
            Type::Table { id, .. } => write!(f, "table ({:x})", id),
            Type::Reference(value) => value.borrow().fmt(f),
            _ => write!(f, "{:?}", self),
        }
    }
}

/// Debug, which breaks closured env circular dependency
impl ::std::fmt::Debug for Type {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        match self {
            Type::Nil => write!(f, "Nil"),
            Type::Boolean(value) => write!(f, "Boolean({:?})", value),
            Type::Number(value) => write!(f, "Number({:?})", value),
            Type::String(value) => write!(f, "String({:?})", value),
            Type::Reference(value) => write!(f, "Reference({:?})", value),
            Type::Vector(vec) => write!(f, "Vector({:?})", vec),
            Type::Table {
                id,
                map,
                metatable,
                border,
            } => write!(
                f,
                "Table {{ id: {}, map: {:?}, metatable: {:?}, border: {} }}",
                id, map, metatable, border
            ),
            Type::Function {
                id,
                parameters,
                varargs,
                body,
                env,
            } => write!(
                f,
                "Function {{ id: {:?}, parameters: {:?}, varargs: {:?}, body: {:?}, env: {:?} }}",
                id,
                parameters,
                varargs,
                body,
                env.borrow().id()
            ),
        }
    }
}

/// Macro to use for pattern maching types with respect to type referencing
#[macro_export]
macro_rules! match_type {
    (($($typ:expr),+), $($pat:pat => $result:expr),+) => {{
        let typs = ($(if let $crate::interpreter::types::Type::Reference(value) = $typ { unsafe { &*value.as_ptr() } } else { $typ }),+);

        match typs {$(
            $pat => $result
        ), +}}
    };
    ($typ:expr, $($pat:pat => $result:expr),+) => {{
        let typ = if let $crate::interpreter::types::Type::Reference(value) = $typ { unsafe { &*value.as_ptr() } } else { $typ };

        match typ {$(
            $pat => $result
        ), +}}
    };
}

impl ::std::convert::AsRef<bool> for Type {
    fn as_ref(&self) -> &bool {
        match_type!(&self,
            Type::Boolean(val) => val,
            _ => panic!("Cannot convert lua value {} to a boolean", self)
        )
    }
}

impl ::std::convert::AsRef<f64> for Type {
    fn as_ref(&self) -> &f64 {
        match_type!(&self,
            Type::Number(val) => val,
            _ => panic!("Cannot convert lua value {} to a number", self)
        )
    }
}

impl ::std::convert::AsRef<String> for Type {
    fn as_ref(&self) -> &String {
        match_type!(&self,
            Type::String(val) => val,
            _ => panic!("Cannot convert lua value {} to a string", self)
        )
    }
}

impl ::std::convert::AsRef<VecDeque<Type>> for Type {
    fn as_ref(&self) -> &VecDeque<Type> {
        match_type!(&self,
            Type::Vector(val) => val,
            _ => panic!("Cannot convert lua value {} to a deque", self)
        )
    }
}

impl ::std::convert::AsRef<HashMap<Type, Rc<RefCell<Type>>>> for Type {
    fn as_ref(&self) -> &HashMap<Type, Rc<RefCell<Type>>> {
        match_type!(&self,
            Type::Table { map, .. } => map,
            _ => panic!("Cannot convert lua value {} to a hashmap", self)
        )
    }
}<|endoftext|>


<commit_before>//! The Actual RethinkDB Commands

use errors::*;
use conn::{Connection, ConnectOpts};
use serde_json;
use ql2::proto;
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::fmt;
use std::error::Error as StdError;

pub struct RootCommand(Result<String>);

impl fmt::Display for RootCommand {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.0 {
            Ok(ref cmd) => write!(f, "{}", cmd),
            Err(_) => Err(fmt::Error),
        }
    }
}

pub struct Command;
pub struct Query;

pub trait IntoCommandArg : fmt::Display {
    fn to_arg(&self) -> Result<String> {
        Ok(format!("{}", self))
    }
}

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

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

impl IntoCommandArg for serde_json::Value {
    fn to_arg(&self) -> Result<String> {
        match serde_json::to_string(self) {
            Ok(cmd) => Ok(cmd),
            Err(e) => Err(From::from(DriverError::Json(e))),
        }
    }
}

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

impl IntoCommandArg for bool { }
impl IntoCommandArg for char { }
impl IntoCommandArg for u8 { }
impl IntoCommandArg for u16 { }
impl IntoCommandArg for u32 { }
impl IntoCommandArg for u64 { }
impl IntoCommandArg for usize { }
impl IntoCommandArg for i8 { }
impl IntoCommandArg for i16 { }
impl IntoCommandArg for i32 { }
impl IntoCommandArg for i64 { }
impl IntoCommandArg for isize { }
impl IntoCommandArg for f32 { }
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
    {
        RootCommand(e.to_arg())
    }

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

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

    pub fn db<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(proto::Term_TermType::DB,
                                     Some(name),
                                     None,
                                     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 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(proto::Term_TermType::TABLE_CREATE,
                                     Some(name),
                                     None,
                                     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(proto::Term_TermType::TABLE,
                                     Some(name),
                                     None,
                                     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(proto::Term_TermType::INSERT,
                                     Some(data),
                                     None,
                                     Some(&commands)))
    }

    pub fn delete(self) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Command::wrap(proto::Term_TermType::DELETE, None as Option<&str>, 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(proto::Term_TermType::FILTER,
                                     Some(filter),
                                     None,
                                     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(proto::Query_QueryType::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: proto::Term_TermType,
                arguments: Option<T>,
                options: 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 arguments.is_some() {
                args.push_str(",");
            }
        }
        if let Some(arguments) = arguments {
            args.push_str(&try!(arguments.to_arg()));
        }
        cmds.push_str(&format!("[{}]", args));
        if let Some(options) = options {
            cmds.push_str(&format!(",{{{}}}", try!(options.to_arg())));
        }
        cmds.push(']');
        Ok(cmds)
    }
}

impl Query {
    pub fn wrap(query_type: proto::Query_QueryType,
                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)
    }
}
<commit_msg>Abbreviate protobuf types<commit_after>//! The Actual RethinkDB Commands

use errors::*;
use conn::{Connection, ConnectOpts};
use serde_json;
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::fmt;
use std::error::Error as StdError;

pub struct RootCommand(Result<String>);

impl fmt::Display for RootCommand {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.0 {
            Ok(ref cmd) => write!(f, "{}", cmd),
            Err(_) => Err(fmt::Error),
        }
    }
}

pub struct Command;
pub struct Query;

pub trait IntoCommandArg : fmt::Display {
    fn to_arg(&self) -> Result<String> {
        Ok(format!("{}", self))
    }
}

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

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

impl IntoCommandArg for serde_json::Value {
    fn to_arg(&self) -> Result<String> {
        match serde_json::to_string(self) {
            Ok(cmd) => Ok(cmd),
            Err(e) => Err(From::from(DriverError::Json(e))),
        }
    }
}

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

impl IntoCommandArg for bool { }
impl IntoCommandArg for char { }
impl IntoCommandArg for u8 { }
impl IntoCommandArg for u16 { }
impl IntoCommandArg for u32 { }
impl IntoCommandArg for u64 { }
impl IntoCommandArg for usize { }
impl IntoCommandArg for i8 { }
impl IntoCommandArg for i16 { }
impl IntoCommandArg for i32 { }
impl IntoCommandArg for i64 { }
impl IntoCommandArg for isize { }
impl IntoCommandArg for f32 { }
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
    {
        RootCommand(e.to_arg())
    }

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

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

    pub fn db<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg
    {
        RootCommand(Command::wrap(tt::DB,
                                     Some(name),
                                     None,
                                     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 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),
                                     None,
                                     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),
                                     None,
                                     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),
                                     None,
                                     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>, 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),
                                     None,
                                     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,
                arguments: Option<T>,
                options: 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 arguments.is_some() {
                args.push_str(",");
            }
        }
        if let Some(arguments) = arguments {
            args.push_str(&try!(arguments.to_arg()));
        }
        cmds.push_str(&format!("[{}]", args));
        if let Some(options) = options {
            cmds.push_str(&format!(",{{{}}}", try!(options.to_arg())));
        }
        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)
    }
}
<|endoftext|>


<commit_before>use super::V2_COOKIE;
use super::super::{Counter, Histogram, RestatState};
use super::super::num::ToPrimitive;
use std::io::{self, Cursor, ErrorKind, Read};
use std::marker::PhantomData;
use std;
use super::byteorder::{BigEndian, ReadBytesExt};

/// Errors that can happen during deserialization.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum DeserializeError {
    /// An i/o operation failed.
    IoError(ErrorKind),
    /// The cookie (first 4 bytes) did not match that for any supported format.
    InvalidCookie,
    /// The histogram uses features that this implementation doesn't support (yet), so it cannot
    /// be deserialized correctly.
    UnsupportedFeature,
    /// A count exceeded what can be represented in the chosen counter type.
    UnsuitableCounterType,
    /// The histogram instance could not be created because the serialized parameters were invalid
    /// (e.g. lowest value, highest value, etc.)
    InvalidParameters,
    /// The current system's pointer width cannot represent the encoded histogram.
    UsizeTypeTooSmall,
    /// The encoded array is longer than it should be for the histogram's value range.
    EncodedArrayTooLong
}

impl std::convert::From<std::io::Error> for DeserializeError {
    fn from(e: std::io::Error) -> Self {
        DeserializeError::IoError(e.kind())
    }
}

/// Deserializer for all supported formats.
///
/// Since the serialization formats all include some magic bytes that allow reliable identification
/// of the different formats, only one Deserializer implementation is needed.
pub struct Deserializer {
    payload_buf: Vec<u8>
}

impl Deserializer {
    /// Create a new deserializer.
    pub fn new() -> Deserializer {
        Deserializer {
            payload_buf: Vec::new()
        }
    }

    /// Deserialize an encoded histogram from the provided reader.
    ///
    /// Note that `&[u8]` and `Cursor` are convenient implementations of `Read` if you have some
    /// bytes already in slice or `Vec` form.
    pub fn deserialize<T: Counter, R: Read>(&mut self, reader: &mut R)
                                            -> Result<Histogram<T>, DeserializeError> {
        // TODO benchmark minimizing read calls by reading into a fixed-size header buffer

        let cookie = reader.read_u32::<BigEndian>()?;

        if cookie != V2_COOKIE {
            return Err(DeserializeError::InvalidCookie);
        }

        let payload_len = reader.read_u32::<BigEndian>()?.to_usize()
            .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        let normalizing_offset = reader.read_u32::<BigEndian>()?;
        if normalizing_offset != 0 {
            return Err(DeserializeError::UnsupportedFeature);
        }
        let num_digits = reader.read_u32::<BigEndian>()?.to_u8()
            .ok_or(DeserializeError::InvalidParameters)?;
        let low = reader.read_u64::<BigEndian>()?;
        let high = reader.read_u64::<BigEndian>()?;
        let int_double_ratio = reader.read_f64::<BigEndian>()?;
        if int_double_ratio != 1.0 {
            return Err(DeserializeError::UnsupportedFeature);
        }

        let mut h = Histogram::new_with_bounds(low, high, num_digits)
            .map_err(|_| DeserializeError::InvalidParameters)?;

        if payload_len > self.payload_buf.len() {
            self.payload_buf.resize(payload_len, 0);
        }

        let mut payload_slice = &mut self.payload_buf[0..payload_len];
        reader.read_exact(&mut payload_slice)?;

        let mut payload_index: usize = 0;
        let mut restat_state = RestatState::new();
        let mut decode_state = DecodeLoopState::new();

        while payload_index < payload_len.saturating_sub(9) {
            // Read with fast loop until we are within 9 of the end. Fast loop can't handle EOF,
            // so bail to slow version for the last few bytes.

            let (zz_num, bytes_read) = varint_read_slice(
                &payload_slice[payload_index..(payload_index + 9)]);
            payload_index += bytes_read;

            let count_or_zeros = zig_zag_decode(zz_num);

            decode_state.on_decoded_num(count_or_zeros, &mut restat_state, &mut h)?;
        }

        // Now read the leftovers
        let leftover_slice = &payload_slice[payload_index..];
        let mut cursor = Cursor::new(&leftover_slice);
        while cursor.position() < leftover_slice.len() as u64 {
            let count_or_zeros = zig_zag_decode(varint_read(&mut cursor)?);

            decode_state.on_decoded_num(count_or_zeros, &mut restat_state, &mut h)?;
        }

        restat_state.update_histogram(&mut h);

        Ok(h)
    }
}

// Only public for testing.
/// Read from a slice that must be 9 bytes long or longer. Returns the decoded number and how many
/// bytes were consumed.
#[inline]
pub fn varint_read_slice(slice: &[u8]) -> (u64, usize) {
    let mut b = slice[0];

    // take low 7 bits
    let mut value: u64 = low_7_bits(b);
    if !is_high_bit_set(b) {
        return (value, 1);
    }
    // high bit set, keep reading
    b = slice[1];
    value |= low_7_bits(b) << 7;
    if !is_high_bit_set(b) {
        return (value, 2);
    }
    b = slice[2];
    value |= low_7_bits(b) << 7 * 2;
    if !is_high_bit_set(b) {
        return (value, 3);
    }
    b = slice[3];
    value |= low_7_bits(b) << 7 * 3;
    if !is_high_bit_set(b) {
        return (value, 4);
    }
    b = slice[4];
    value |= low_7_bits(b) << 7 * 4;
    if !is_high_bit_set(b) {
        return (value, 5);
    }
    b = slice[5];
    value |= low_7_bits(b) << 7 * 5;
    if !is_high_bit_set(b) {
        return (value, 6);
    }
    b = slice[6];
    value |= low_7_bits(b) << 7 * 6;
    if !is_high_bit_set(b) {
        return (value, 7);
    }
    b = slice[7];
    value |= low_7_bits(b) << 7 * 7;
    if !is_high_bit_set(b) {
        return (value, 8);
    }

    b = slice[8];
    // special case: use last byte as is
    value |= (b as u64) << 7 * 8;

    (value, 9)
}

// Only public for testing.
/// Read a LEB128-64b9B from the buffer
pub fn varint_read<R: Read>(reader: &mut R) -> io::Result<u64> {
    let mut b = reader.read_u8()?;

    // take low 7 bits
    let mut value: u64 = low_7_bits(b);

    if is_high_bit_set(b) {
        // high bit set, keep reading
        b = reader.read_u8()?;
        value |= low_7_bits(b) << 7;
        if is_high_bit_set(b) {
            b = reader.read_u8()?;
            value |= low_7_bits(b) << 7 * 2;
            if is_high_bit_set(b) {
                b = reader.read_u8()?;
                value |= low_7_bits(b) << 7 * 3;
                if is_high_bit_set(b) {
                    b = reader.read_u8()?;
                    value |= low_7_bits(b) << 7 * 4;
                    if is_high_bit_set(b) {
                        b = reader.read_u8()?;
                        value |= low_7_bits(b) << 7 * 5;
                        if is_high_bit_set(b) {
                            b = reader.read_u8()?;
                            value |= low_7_bits(b) << 7 * 6;
                            if is_high_bit_set(b) {
                                b = reader.read_u8()?;
                                value |= low_7_bits(b) << 7 * 7;
                                if is_high_bit_set(b) {
                                    b = reader.read_u8()?;
                                    // special case: use last byte as is
                                    value |= (b as u64) << 7 * 8;
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    Ok(value)
}

/// truncate byte to low 7 bits, cast to u64
#[inline]
fn low_7_bits(b: u8) -> u64 {
    (b & 0x7F) as u64
}

#[inline]
fn is_high_bit_set(b: u8) -> bool {
    (b & 0x80) != 0
}

// Only public for testing.
#[inline]
pub fn zig_zag_decode(encoded: u64) -> i64 {
    ((encoded >> 1) as i64) ^ -((encoded & 1) as i64)
}

/// We need to perform the same logic in two different decode loops while carrying over a modicum
/// of state.
struct DecodeLoopState<T: Counter> {
    dest_index:usize,
    phantom: PhantomData<T>
}

impl <T: Counter> DecodeLoopState<T> {

    fn new() -> DecodeLoopState<T> {
        DecodeLoopState {
            dest_index: 0,
            phantom: PhantomData
        }
    }

    #[inline]
    fn on_decoded_num(&mut self, count_or_zeros: i64, restat_state: &mut RestatState<T>,
                      h: &mut Histogram<T>) -> Result<(), DeserializeError> {
        if count_or_zeros < 0 {
            let zero_count = (-count_or_zeros).to_usize()
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
            // skip the zeros
            self.dest_index = self.dest_index.checked_add(zero_count)
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        } else {
            let count: T = T::from_i64(count_or_zeros)
                .ok_or(DeserializeError::UnsuitableCounterType)?;

            h.set_count_at_index(self.dest_index, count)
                .map_err(|_| DeserializeError::EncodedArrayTooLong)?;

            restat_state.on_nonzero_count(self.dest_index, count);

            self.dest_index = self.dest_index.checked_add(1)
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        }

        Ok(())
    }
}
<commit_msg>Remove TODO about header i/o: didn't show up anywhere in profiling<commit_after>use super::V2_COOKIE;
use super::super::{Counter, Histogram, RestatState};
use super::super::num::ToPrimitive;
use std::io::{self, Cursor, ErrorKind, Read};
use std::marker::PhantomData;
use std;
use super::byteorder::{BigEndian, ReadBytesExt};

/// Errors that can happen during deserialization.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum DeserializeError {
    /// An i/o operation failed.
    IoError(ErrorKind),
    /// The cookie (first 4 bytes) did not match that for any supported format.
    InvalidCookie,
    /// The histogram uses features that this implementation doesn't support (yet), so it cannot
    /// be deserialized correctly.
    UnsupportedFeature,
    /// A count exceeded what can be represented in the chosen counter type.
    UnsuitableCounterType,
    /// The histogram instance could not be created because the serialized parameters were invalid
    /// (e.g. lowest value, highest value, etc.)
    InvalidParameters,
    /// The current system's pointer width cannot represent the encoded histogram.
    UsizeTypeTooSmall,
    /// The encoded array is longer than it should be for the histogram's value range.
    EncodedArrayTooLong
}

impl std::convert::From<std::io::Error> for DeserializeError {
    fn from(e: std::io::Error) -> Self {
        DeserializeError::IoError(e.kind())
    }
}

/// Deserializer for all supported formats.
///
/// Since the serialization formats all include some magic bytes that allow reliable identification
/// of the different formats, only one Deserializer implementation is needed.
pub struct Deserializer {
    payload_buf: Vec<u8>
}

impl Deserializer {
    /// Create a new deserializer.
    pub fn new() -> Deserializer {
        Deserializer {
            payload_buf: Vec::new()
        }
    }

    /// Deserialize an encoded histogram from the provided reader.
    ///
    /// Note that `&[u8]` and `Cursor` are convenient implementations of `Read` if you have some
    /// bytes already in slice or `Vec` form.
    pub fn deserialize<T: Counter, R: Read>(&mut self, reader: &mut R)
                                            -> Result<Histogram<T>, DeserializeError> {
        let cookie = reader.read_u32::<BigEndian>()?;

        if cookie != V2_COOKIE {
            return Err(DeserializeError::InvalidCookie);
        }

        let payload_len = reader.read_u32::<BigEndian>()?.to_usize()
            .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        let normalizing_offset = reader.read_u32::<BigEndian>()?;
        if normalizing_offset != 0 {
            return Err(DeserializeError::UnsupportedFeature);
        }
        let num_digits = reader.read_u32::<BigEndian>()?.to_u8()
            .ok_or(DeserializeError::InvalidParameters)?;
        let low = reader.read_u64::<BigEndian>()?;
        let high = reader.read_u64::<BigEndian>()?;
        let int_double_ratio = reader.read_f64::<BigEndian>()?;
        if int_double_ratio != 1.0 {
            return Err(DeserializeError::UnsupportedFeature);
        }

        let mut h = Histogram::new_with_bounds(low, high, num_digits)
            .map_err(|_| DeserializeError::InvalidParameters)?;

        if payload_len > self.payload_buf.len() {
            self.payload_buf.resize(payload_len, 0);
        }

        let mut payload_slice = &mut self.payload_buf[0..payload_len];
        reader.read_exact(&mut payload_slice)?;

        let mut payload_index: usize = 0;
        let mut restat_state = RestatState::new();
        let mut decode_state = DecodeLoopState::new();

        while payload_index < payload_len.saturating_sub(9) {
            // Read with fast loop until we are within 9 of the end. Fast loop can't handle EOF,
            // so bail to slow version for the last few bytes.

            let (zz_num, bytes_read) = varint_read_slice(
                &payload_slice[payload_index..(payload_index + 9)]);
            payload_index += bytes_read;

            let count_or_zeros = zig_zag_decode(zz_num);

            decode_state.on_decoded_num(count_or_zeros, &mut restat_state, &mut h)?;
        }

        // Now read the leftovers
        let leftover_slice = &payload_slice[payload_index..];
        let mut cursor = Cursor::new(&leftover_slice);
        while cursor.position() < leftover_slice.len() as u64 {
            let count_or_zeros = zig_zag_decode(varint_read(&mut cursor)?);

            decode_state.on_decoded_num(count_or_zeros, &mut restat_state, &mut h)?;
        }

        restat_state.update_histogram(&mut h);

        Ok(h)
    }
}

// Only public for testing.
/// Read from a slice that must be 9 bytes long or longer. Returns the decoded number and how many
/// bytes were consumed.
#[inline]
pub fn varint_read_slice(slice: &[u8]) -> (u64, usize) {
    let mut b = slice[0];

    // take low 7 bits
    let mut value: u64 = low_7_bits(b);
    if !is_high_bit_set(b) {
        return (value, 1);
    }
    // high bit set, keep reading
    b = slice[1];
    value |= low_7_bits(b) << 7;
    if !is_high_bit_set(b) {
        return (value, 2);
    }
    b = slice[2];
    value |= low_7_bits(b) << 7 * 2;
    if !is_high_bit_set(b) {
        return (value, 3);
    }
    b = slice[3];
    value |= low_7_bits(b) << 7 * 3;
    if !is_high_bit_set(b) {
        return (value, 4);
    }
    b = slice[4];
    value |= low_7_bits(b) << 7 * 4;
    if !is_high_bit_set(b) {
        return (value, 5);
    }
    b = slice[5];
    value |= low_7_bits(b) << 7 * 5;
    if !is_high_bit_set(b) {
        return (value, 6);
    }
    b = slice[6];
    value |= low_7_bits(b) << 7 * 6;
    if !is_high_bit_set(b) {
        return (value, 7);
    }
    b = slice[7];
    value |= low_7_bits(b) << 7 * 7;
    if !is_high_bit_set(b) {
        return (value, 8);
    }

    b = slice[8];
    // special case: use last byte as is
    value |= (b as u64) << 7 * 8;

    (value, 9)
}

// Only public for testing.
/// Read a LEB128-64b9B from the buffer
pub fn varint_read<R: Read>(reader: &mut R) -> io::Result<u64> {
    let mut b = reader.read_u8()?;

    // take low 7 bits
    let mut value: u64 = low_7_bits(b);

    if is_high_bit_set(b) {
        // high bit set, keep reading
        b = reader.read_u8()?;
        value |= low_7_bits(b) << 7;
        if is_high_bit_set(b) {
            b = reader.read_u8()?;
            value |= low_7_bits(b) << 7 * 2;
            if is_high_bit_set(b) {
                b = reader.read_u8()?;
                value |= low_7_bits(b) << 7 * 3;
                if is_high_bit_set(b) {
                    b = reader.read_u8()?;
                    value |= low_7_bits(b) << 7 * 4;
                    if is_high_bit_set(b) {
                        b = reader.read_u8()?;
                        value |= low_7_bits(b) << 7 * 5;
                        if is_high_bit_set(b) {
                            b = reader.read_u8()?;
                            value |= low_7_bits(b) << 7 * 6;
                            if is_high_bit_set(b) {
                                b = reader.read_u8()?;
                                value |= low_7_bits(b) << 7 * 7;
                                if is_high_bit_set(b) {
                                    b = reader.read_u8()?;
                                    // special case: use last byte as is
                                    value |= (b as u64) << 7 * 8;
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    Ok(value)
}

/// truncate byte to low 7 bits, cast to u64
#[inline]
fn low_7_bits(b: u8) -> u64 {
    (b & 0x7F) as u64
}

#[inline]
fn is_high_bit_set(b: u8) -> bool {
    (b & 0x80) != 0
}

// Only public for testing.
#[inline]
pub fn zig_zag_decode(encoded: u64) -> i64 {
    ((encoded >> 1) as i64) ^ -((encoded & 1) as i64)
}

/// We need to perform the same logic in two different decode loops while carrying over a modicum
/// of state.
struct DecodeLoopState<T: Counter> {
    dest_index:usize,
    phantom: PhantomData<T>
}

impl <T: Counter> DecodeLoopState<T> {

    fn new() -> DecodeLoopState<T> {
        DecodeLoopState {
            dest_index: 0,
            phantom: PhantomData
        }
    }

    #[inline]
    fn on_decoded_num(&mut self, count_or_zeros: i64, restat_state: &mut RestatState<T>,
                      h: &mut Histogram<T>) -> Result<(), DeserializeError> {
        if count_or_zeros < 0 {
            let zero_count = (-count_or_zeros).to_usize()
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
            // skip the zeros
            self.dest_index = self.dest_index.checked_add(zero_count)
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        } else {
            let count: T = T::from_i64(count_or_zeros)
                .ok_or(DeserializeError::UnsuitableCounterType)?;

            h.set_count_at_index(self.dest_index, count)
                .map_err(|_| DeserializeError::EncodedArrayTooLong)?;

            restat_state.on_nonzero_count(self.dest_index, count);

            self.dest_index = self.dest_index.checked_add(1)
                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
        }

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


<commit_before>//! Data consumers
//!
//! The goal of data producers is to parse data depending on the previous result.
//! It can be used to selectively seek in a file.
//!
//! ## Example
//!
//! This consumer will take 4 samples from the input, print them, then stop
//!
//! ```rust
//!  use nom::{MemProducer,Consumer,ConsumerState};
//!  use std::str;
//!  use std::io::SeekFrom;
//!
//!  struct TestPrintConsumer {
//!    counter: usize
//!  }
//!
//!  impl TestPrintConsumer {
//!    fn new() -> TestPrintConsumer {
//!      TestPrintConsumer { counter: 0 }
//!    }
//!  }
//!
//!  // Return ConsumerState::Await if it needs more data, or ConsumerDone when it ends
//!  impl Consumer for TestPrintConsumer {
//!    fn consume(&mut self, input: &[u8]) -> ConsumerState {
//!      println!("{} -> {}", self.counter, str::from_utf8(input).unwrap());
//!      self.counter = self.counter + 1;
//!      if self.counter <=4 {
//!        ConsumerState::Await(0, 0)
//!      } else {
//!        ConsumerState::ConsumerDone
//!      }
//!    }
//!
//!    fn end(&mut self) {
//!      println!("finished");
//!    }
//!  }
//!
//!  // It can consume data directly from a producer
//!  let mut p = MemProducer::new("abcdefghijklmnopqrstuvwx".as_bytes(), 4);
//!  let mut c = TestPrintConsumer::new();
//!  c.run(&mut p);
//! ```

use self::ConsumerState::*;
use producer::Producer;
use producer::ProducerState::*;
use internal::Err;
use std::io::SeekFrom;
use util::HexDisplay;

/// Holds the current state of the consumer
///
/// * Await if more data is needed
///
/// * ConsumerDone if the consumer does not need anymore data to be parsed
///
/// * ConsumerError when something went wrong
#[derive(Debug,PartialEq,Eq,Copy)]
pub enum ConsumerState {
  Await(
    usize,    // consumed
    usize     // needed buffer size
  ),
  Seek(
    usize,    // consumed
    SeekFrom, // new position
    usize     // needed buffer size
  ),
  Incomplete,
  ConsumerDone,
  ConsumerError(Err)
}

/// Implement the consume method, taking a byte array as input and returning a consumer state
///
/// The run function takes care of continuing or not
pub trait Consumer {
  fn consume(&mut self, input: &[u8]) -> ConsumerState;
  fn end(&mut self);

  fn run(&mut self, producer: &mut Producer) {
    let mut acc: Vec<u8>      = Vec::new();
    let mut position          = 0;
    let mut shouldSeek        = false;
    let mut consumed:usize    = 0;
    let mut needed:usize      = 0;
    let mut seekFrom:SeekFrom = SeekFrom::Current(0);
    let mut eof = false;
    let mut end = false;

    loop {
      //self.getDataFromProducer(producer, seekFrom, needed, acc);
      if !shouldSeek && acc.len() - consumed >= needed {
        //println!("buffer is large enough, skipping");
        let mut tmp = Vec::new();
        //println!("before:\n{}", acc.to_hex(16));
        //println!("after:\n{}", (&acc[consumed..acc.len()]).to_hex(16));
        tmp.push_all(&acc[consumed..acc.len()]);
        acc.clear();
        acc = tmp;
      } else {
        if shouldSeek {
          let pos = producer.seek(seekFrom);
          //println!("seeking: {:?}", pos);
          shouldSeek = false;
          acc.clear();
        } else {
          let mut tmp = Vec::new();
          tmp.push_all(&acc[consumed..acc.len()]);
          acc.clear();
          acc = tmp;
        }

        loop {
          let state   = producer.produce();
          match state {
            Data(v) => {
              //println!("got data: {} bytes", v.len());
              acc.push_all(v);
              position = position + v.len();
            },
            Eof([])  => {
              //println!("eof empty");
              eof = true;
              self.end();
              return
            }
            Eof(v) => {
              //println!("eof with {} bytes", v.len());
              eof = true;
              acc.push_all(v);
              position = position + v.len();
              break;
            }
            ProducerError(e) => {break;}
            Continue => { continue; }
          }
          //println!("acc size: {}", acc.len());
          if acc.len() >= needed { break; }
        }
      }

      //println!("full:\n{}", acc.to_hex(8));
      //println!("truncated:\n{}", (&acc[0..needed]).to_hex(16));
      match self.consume(&acc[0..needed]) {
        ConsumerError(e) => {
          //println!("consumer error, stopping: {}", e);
        },
        ConsumerDone => {
          //println!("data, done");
          end = true;
        },
        Seek(consumed_bytes, sf, needed_bytes) => {
          //println!("Seek: consumed {} bytes, got {:?} and asked {} bytes", consumed_bytes, sf, needed_bytes);
          seekFrom = match sf {
            SeekFrom::Current(i) => SeekFrom::Current(i - (acc.len() - needed) as i64),
            a => a
          };
          shouldSeek = true;
          consumed = consumed_bytes;
          needed   = needed_bytes;
        },
        Await(consumed_bytes, needed_bytes) => {
          //println!("consumed: {} bytes | needed: {} bytes", consumed_bytes, needed_bytes);
          consumed = consumed_bytes;
          needed   = needed_bytes;
        },
        Incomplete => {
          //println!("incomplete");
        }
      }
      if eof {
        self.end();
        break;
      }
      if end {
        self.end();
        break;
      }
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use super::ConsumerState::*;
  use producer::MemProducer;
  use std::str;
  use std::io::SeekFrom;

  struct TestPrintConsumer {
    counter: usize,
    ended:   bool
  }

  impl TestPrintConsumer {
    fn new() -> TestPrintConsumer {
      TestPrintConsumer { counter: 0, ended: false }
    }
  }

  impl Consumer for TestPrintConsumer {
    fn consume(&mut self, input: &[u8]) -> ConsumerState {
      println!("{} -> {}", self.counter, str::from_utf8(input).unwrap());
      self.counter = self.counter + 1;
      if self.counter <= 4 {
        Await(0, 0)
      } else {
        ConsumerDone
      }
    }

    fn end(&mut self) {
      assert_eq!(self.counter, 5);
      self.ended = true;
    }
  }

  #[test]
  fn pull() {
    let mut p = MemProducer::new("abcdefghijklmnopqrstuvwx".as_bytes(), 4);
    let mut c = TestPrintConsumer::new();
    c.run(&mut p);

    assert!(c.ended);
  }
}
<commit_msg>Better examples for the consumer<commit_after>//! Data consumers
//!
//! The goal of data producers is to parse data depending on the previous result.
//! It can be used to selectively seek in a file.
//!
//! ## Example
//!
//! This consumer will take 4 samples from the input, print them, then stop
//!
//! ```rust
//!  use nom::{IResult,MemProducer,Consumer,ConsumerState};
//!  use std::str;
//!  use std::io::SeekFrom;
//!
//!  struct TestPrintConsumer {
//!    counter: usize
//!  }
//!
//!  impl TestPrintConsumer {
//!    fn new() -> TestPrintConsumer {
//!      TestPrintConsumer { counter: 0 }
//!    }
//!  }
//!
//!  fn take4(i:&[u8]) -> IResult<&[u8], &[u8]>{
//!    if i.len() < 4 {
//!      IResult::Incomplete(4)
//!    } else {
//!      IResult::Done(&i[4..],&i[0..4])
//!    }
//!  }
//!
//!  // Return ConsumerState::Await if it needs more data, or ConsumerDone when it ends
//!  impl Consumer for TestPrintConsumer {
//!    fn consume(&mut self, input: &[u8]) -> ConsumerState {
//!      match take4(input) {
//!        IResult::Error(a)      => ConsumerState::ConsumerError(a),
//!        IResult::Incomplete(a) => ConsumerState::Await(0, 4),
//!        IResult::Done(i, o)    => {
//!          println!("{} -> {}", self.counter, str::from_utf8(o).unwrap());
//!          self.counter = self.counter + 1;
//!          if self.counter <= 4 {
//!            ConsumerState::Await(4, 4)
//!          } else {
//!            ConsumerState::ConsumerDone
//!          }
//!        }
//!      }
//!    }
//!
//!    fn end(&mut self) {
//!      println!("finished");
//!    }
//!  }
//!
//!  // It can consume data directly from a producer
//!  let mut p = MemProducer::new("abcdefghijklmnopqrstuvwx".as_bytes(), 4);
//!  let mut c = TestPrintConsumer::new();
//!  c.run(&mut p);
//! ```

use self::ConsumerState::*;
use producer::Producer;
use producer::ProducerState::*;
use internal::Err;
use std::io::SeekFrom;
use util::HexDisplay;

/// Holds the current state of the consumer
///
/// * Await if more data is needed
///
/// * ConsumerDone if the consumer does not need anymore data to be parsed
///
/// * ConsumerError when something went wrong
#[derive(Debug,PartialEq,Eq,Copy)]
pub enum ConsumerState {
  Await(
    usize,    // consumed
    usize     // needed buffer size
  ),
  Seek(
    usize,    // consumed
    SeekFrom, // new position
    usize     // needed buffer size
  ),
  Incomplete,
  ConsumerDone,
  ConsumerError(Err)
}

/// Implement the consume method, taking a byte array as input and returning a consumer state
///
/// The run function takes care of continuing or not
pub trait Consumer {
  fn consume(&mut self, input: &[u8]) -> ConsumerState;
  fn end(&mut self);

  fn run(&mut self, producer: &mut Producer) {
    let mut acc: Vec<u8>      = Vec::new();
    let mut position          = 0;
    let mut shouldSeek        = false;
    let mut consumed:usize    = 0;
    let mut needed:usize      = 0;
    let mut seekFrom:SeekFrom = SeekFrom::Current(0);
    let mut eof = false;
    let mut end = false;

    loop {
      //self.getDataFromProducer(producer, seekFrom, needed, acc);
      if !shouldSeek && acc.len() - consumed >= needed {
        //println!("buffer is large enough, skipping");
        let mut tmp = Vec::new();
        //println!("before:\n{}", acc.to_hex(16));
        //println!("after:\n{}", (&acc[consumed..acc.len()]).to_hex(16));
        tmp.push_all(&acc[consumed..acc.len()]);
        acc.clear();
        acc = tmp;
      } else {
        if shouldSeek {
          let pos = producer.seek(seekFrom);
          //println!("seeking: {:?}", pos);
          shouldSeek = false;
          acc.clear();
        } else {
          let mut tmp = Vec::new();
          tmp.push_all(&acc[consumed..acc.len()]);
          acc.clear();
          acc = tmp;
        }

        loop {
          let state   = producer.produce();
          match state {
            Data(v) => {
              //println!("got data: {} bytes", v.len());
              acc.push_all(v);
              position = position + v.len();
            },
            Eof([])  => {
              //println!("eof empty");
              eof = true;
              self.end();
              return
            }
            Eof(v) => {
              //println!("eof with {} bytes", v.len());
              eof = true;
              acc.push_all(v);
              position = position + v.len();
              break;
            }
            ProducerError(e) => {break;}
            Continue => { continue; }
          }
          //println!("acc size: {}", acc.len());
          if acc.len() >= needed { break; }
        }
      }

      //println!("full:\n{}", acc.to_hex(8));
      //println!("truncated:\n{}", (&acc[0..needed]).to_hex(16));
      match self.consume(&acc[0..needed]) {
        ConsumerError(e) => {
          //println!("consumer error, stopping: {}", e);
        },
        ConsumerDone => {
          //println!("data, done");
          end = true;
        },
        Seek(consumed_bytes, sf, needed_bytes) => {
          //println!("Seek: consumed {} bytes, got {:?} and asked {} bytes", consumed_bytes, sf, needed_bytes);
          seekFrom = match sf {
            SeekFrom::Current(i) => SeekFrom::Current(i - (acc.len() - needed) as i64),
            a => a
          };
          shouldSeek = true;
          consumed = consumed_bytes;
          needed   = needed_bytes;
        },
        Await(consumed_bytes, needed_bytes) => {
          //println!("consumed: {} bytes | needed: {} bytes", consumed_bytes, needed_bytes);
          consumed = consumed_bytes;
          needed   = needed_bytes;
        },
        Incomplete => {
          //println!("incomplete");
        }
      }
      if eof {
        self.end();
        break;
      }
      if end {
        self.end();
        break;
      }
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use super::ConsumerState::*;
  use producer::MemProducer;
  use internal::IResult;
  use std::str;
  use std::io::SeekFrom;

#[macro_export]
macro_rules! take(
  ($name:ident $count:expr) => (
    fn $name(i:&[u8]) -> IResult<&[u8], &[u8]>{
      if i.len() < $count {
        IResult::Incomplete(0)
      } else {
        IResult::Done(&i[$count..],&i[0..$count])
      }
    }
  )
);

  struct TestPrintConsumer {
    counter: usize,
    ended:   bool
  }

  impl TestPrintConsumer {
    fn new() -> TestPrintConsumer {
      TestPrintConsumer { counter: 0, ended: false }
    }
  }

  take!(take4 4);

  impl Consumer for TestPrintConsumer {
    fn consume(&mut self, input: &[u8]) -> ConsumerState {
      match take4(input) {
        IResult::Error(a)      => ConsumerError(a),
        IResult::Incomplete(a) => Await(0, 4),
        IResult::Done(i, o)    => {
          println!("{} -> {}", self.counter, str::from_utf8(o).unwrap());
          self.counter = self.counter + 1;
          if self.counter <= 4 {
            Await(4, 4)
          } else {
            ConsumerDone
          }
        }
      }
    }

    fn end(&mut self) {
      assert_eq!(self.counter, 5);
      self.ended = true;
    }
  }

  #[test]
  fn pull() {
    let mut p = MemProducer::new("abcdefghijklmnopqrstuvwx".as_bytes(), 4);
    let mut c = TestPrintConsumer::new();
    c.run(&mut p);

    assert!(c.ended);
  }
}
<|endoftext|>


<commit_before>//! The lookups.

use truetype::GlyphID;

use {Result, Tape, Value};

mod class;
mod coverage;

pub use self::class::{Class, Class1, Class2};
pub use self::coverage::{Coverage, Coverage1, Coverage2};

table! {
    @define
    #[doc = "A lookup list."]
    pub Lookups {
        count   (u16        ), // LookupCount
        offsets (Vec<u16>   ), // Lookup
        records (Vec<Record>),
    }
}

table! {
    @define
    #[doc = "A lookup record."]
    pub Record {
        kind               (Kind       ), // LookupType
        flags              (Flags      ), // LookupFlag
        table_count        (u16        ), // SubTableCount
        table_offsets      (Vec<u16>   ), // SubTable
        mark_filtering_set (Option<u16>), // MarkFilteringSet
    }
}

/// The type of a lookup.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Kind {
    SingleAdjustment = 1,
    PairAdjustment = 2,
    CursiveAttachment = 3,
    MarkToBaseAttachment = 4,
    MarkToLigatureAttachment = 5,
    MarkToMarkAttachment = 6,
    ContextPositioning = 7,
    ChainedContextPositioning = 8,
    ExtensionPositioning = 9,
}

flags! {
    #[doc = "Flags of a lookup record."]
    pub Flags(u16) {
        0b0000_0000_0000_0001 => is_right_to_left,
        0b0000_0000_0000_0010 => should_ignore_base_glyphs,
        0b0000_0000_0000_0100 => should_ignore_ligature,
        0b0000_0000_0000_1000 => should_ignore_marks,
        0b0000_0000_0001_0000 => has_mark_filtering,
        0b0000_0000_1110_0000 => is_invalid,
    }
}

table! {
    #[doc = "A glyph range."]
    #[derive(Copy)]
    pub Range {
        start (GlyphID), // Start
        end   (GlyphID), // End
        index (u16    ), // Class or StartCoverageIndex
    }
}

impl Value for Lookups {
    fn read<T: Tape>(tape: &mut T) -> Result<Self> {
        let position = try!(tape.position());
        let count = try!(tape.take::<u16>());
        let offsets: Vec<u16> = try!(tape.take_given(count as usize));
        let mut records: Vec<Record> = Vec::with_capacity(count as usize);
        for i in 0..(count as usize) {
            try!(tape.jump(position + offsets[i] as u64));
            records.push(try!(tape.take()));
        }
        Ok(Lookups { count: count, offsets: offsets, records: records })
    }
}

impl Value for Record {
    fn read<T: Tape>(tape: &mut T) -> Result<Self> {
        let kind = try!(tape.take());
        let flags = try!(tape.take::<Flags>());
        if flags.is_invalid() {
            raise!("found a malformed lookup record");
        }
        let table_count = try!(tape.take::<u16>());
        let table_offsets: Vec<u16> = try!(tape.take_given(table_count as usize));
        let mark_filtering_set = if flags.has_mark_filtering() {
            Some(try!(tape.take()))
        } else {
            None
        };
        Ok(Record {
            kind: kind,
            flags: flags,
            table_count: table_count,
            table_offsets: table_offsets,
            mark_filtering_set: mark_filtering_set,
        })
    }
}

impl Value for Kind {
    fn read<T: Tape>(tape: &mut T) -> Result<Self> {
        let kind = try!(tape.take::<u16>());
        if kind < 1 || kind > 9 {
            raise!("found an unknown lookup type");
        }
        Ok(unsafe { ::std::mem::transmute(kind as u8) })
    }
}
<commit_msg>Adjust a couple of descriptions<commit_after>//! The lookups.

use truetype::GlyphID;

use {Result, Tape, Value};

mod class;
mod coverage;

pub use self::class::{Class, Class1, Class2};
pub use self::coverage::{Coverage, Coverage1, Coverage2};

table! {
    @define
    #[doc = "A lookup list."]
    pub Lookups {
        count   (u16        ), // LookupCount
        offsets (Vec<u16>   ), // Lookup
        records (Vec<Record>),
    }
}

table! {
    @define
    #[doc = "A lookup record."]
    pub Record {
        kind               (Kind       ), // LookupType
        flags              (Flags      ), // LookupFlag
        table_count        (u16        ), // SubTableCount
        table_offsets      (Vec<u16>   ), // SubTable
        mark_filtering_set (Option<u16>), // MarkFilteringSet
    }
}

/// A lookup type.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Kind {
    SingleAdjustment = 1,
    PairAdjustment = 2,
    CursiveAttachment = 3,
    MarkToBaseAttachment = 4,
    MarkToLigatureAttachment = 5,
    MarkToMarkAttachment = 6,
    ContextPositioning = 7,
    ChainedContextPositioning = 8,
    ExtensionPositioning = 9,
}

flags! {
    #[doc = "Lookup flags."]
    pub Flags(u16) {
        0b0000_0000_0000_0001 => is_right_to_left,
        0b0000_0000_0000_0010 => should_ignore_base_glyphs,
        0b0000_0000_0000_0100 => should_ignore_ligature,
        0b0000_0000_0000_1000 => should_ignore_marks,
        0b0000_0000_0001_0000 => has_mark_filtering,
        0b0000_0000_1110_0000 => is_invalid,
    }
}

table! {
    #[doc = "A glyph range."]
    #[derive(Copy)]
    pub Range {
        start (GlyphID), // Start
        end   (GlyphID), // End
        index (u16    ), // Class or StartCoverageIndex
    }
}

impl Value for Lookups {
    fn read<T: Tape>(tape: &mut T) -> Result<Self> {
        let position = try!(tape.position());
        let count = try!(tape.take::<u16>());
        let offsets: Vec<u16> = try!(tape.take_given(count as usize));
        let mut records: Vec<Record> = Vec::with_capacity(count as usize);
        for i in 0..(count as usize) {
            try!(tape.jump(position + offsets[i] as u64));
            records.push(try!(tape.take()));
        }
        Ok(Lookups { count: count, offsets: offsets, records: records })
    }
}

impl Value for Record {
    fn read<T: Tape>(tape: &mut T) -> Result<Self> {
        let kind = try!(tape.take());
        let flags = try!(tape.take::<Flags>());
        if flags.is_invalid() {
            raise!("found a malformed lookup record");
        }
        let table_count = try!(tape.take::<u16>());
        let table_offsets: Vec<u16> = try!(tape.take_given(table_count as usize));
        let mark_filtering_set = if flags.has_mark_filtering() {
            Some(try!(tape.take()))
        } else {
            None
        };
        Ok(Record {
            kind: kind,
            flags: flags,
            table_count: table_count,
            table_offsets: table_offsets,
            mark_filtering_set: mark_filtering_set,
        })
    }
}

impl Value for Kind {
    fn read<T: Tape>(tape: &mut T) -> Result<Self> {
        let kind = try!(tape.take::<u16>());
        if kind < 1 || kind > 9 {
            raise!("found an unknown lookup type");
        }
        Ok(unsafe { ::std::mem::transmute(kind as u8) })
    }
}
<|endoftext|>


<commit_before>// helper for transmutation, shown below.
type RustClosure = (int,int);

struct Condition<T, U:Copy> {
    key: task::local_data::LocalDataKey<Handler<T,U>>
}

struct Handler<T, U:Copy> {
    handle: RustClosure,
    prev: Option<@Handler<T, U>>
}

struct ProtectBlock<T, U:Copy> {
    cond: &Condition<T, U>,
    inner: RustClosure
}

struct PopHandler<T, U:Copy> {
    cond: &Condition<T,U>,
    drop {
        unsafe {
            debug!("PopHandler: popping handler from TLS");
            match task::local_data::local_data_pop(self.cond.key) {
                None => (),
                Some(h) => {
                    match h.prev {
                        None => (),
                        Some(p) =>
                        task::local_data::local_data_set(self.cond.key, p)
                    }
                }
            }
        }
    }
}

struct HandleBlock<T, U:Copy> {
    pb: &ProtectBlock<T,U>,
    handler: @Handler<T,U>,
    drop {
        unsafe {
            debug!("HandleBlock: pushing handler to TLS");
            task::local_data::local_data_set(self.pb.cond.key,
                                             self.handler);
            let _pop = PopHandler { cond: self.pb.cond };
            // transmutation to avoid copying non-copyable, should
            // be fixable by tracking closure pointees in regionck.
            let f : &fn() = ::cast::transmute(self.pb.inner);
            debug!("HandleBlock: invoking protected code");
            f();
            debug!("HandleBlock: returned from protected code");
        }
    }
}

impl<T, U: Copy> ProtectBlock<T,U> {
    fn handle(&self, h: &self/fn(&T) ->U) -> HandleBlock/&self<T,U> {
        unsafe {
            debug!("ProtectBlock.handle: setting up handler block");
            let p : *RustClosure = ::cast::transmute(&h);
            let prev = task::local_data::local_data_get(self.cond.key);
            HandleBlock { pb: self,
                          handler: @Handler{handle: *p, prev: prev} }
        }
    }
}

impl<T, U: Copy>  Condition<T,U> {

    fn protect(&self, inner: &self/fn()) -> ProtectBlock/&self<T,U> {
        unsafe {
            // transmutation to avoid copying non-copyable, should
            // be fixable by tracking closure pointees in regionck.
            debug!("Condition.protect: setting up protected block");
            let p : *RustClosure = ::cast::transmute(&inner);
            ProtectBlock { cond: self,
                           inner: *p }
        }
    }

    fn raise(t:&T) -> U {
        unsafe {
            match task::local_data::local_data_get(self.key) {
                None => {
                    debug!("Condition.raise: found no handler");
                    fail
                }

                Some(handler) => {
                    debug!("Condition.raise: found handler");
                    let f : &fn(&T) -> U = ::cast::transmute(handler.handle);
                    f(t)
                }
            }
        }
    }
}

#[cfg(test)]
fn sadness_key(_x: @Handler<int,int>) { }

#[cfg(test)]
fn trouble(i: int) {
    // Condition should work as a const, just limitations in consts.
    let sadness_condition : Condition<int,int> =
        Condition { key: sadness_key };
    debug!("trouble: raising conition");
    let j = sadness_condition.raise(&i);
    debug!("trouble: handler recovered with %d", j);
}

#[test]
fn test1() {

    let sadness_condition : Condition<int,int> =
        Condition { key: sadness_key };

    let mut i = 10;

    let b = do sadness_condition.protect {
        debug!("test1: in protected block");
        trouble(1);
        trouble(2);
        trouble(3);
    };

    do b.handle |j| {
        debug!("test1: in handler");
        i += *j;
        i
    };

    assert i == 16;
}
#[cfg(test)]
fn nested_test_inner() {
    let sadness_condition : Condition<int,int> =
        Condition { key: sadness_key };

    let mut inner_trapped = false;

    let b = do sadness_condition.protect {
        debug!("nested_test_inner: in protected block");
        trouble(1);
    };

    do b.handle |_j:&int| {
        debug!("nested_test_inner: in handler");
        inner_trapped = true;
        0
    };

    assert inner_trapped;
}

#[test]
fn nested_test_outer() {

    let sadness_condition : Condition<int,int> =
        Condition { key: sadness_key };

    let mut outer_trapped = false;

    let b = do sadness_condition.protect {
        debug!("nested_test_outer: in protected block");
        nested_test_inner();
        trouble(1);
    };

    do b.handle |_j:&int| {
        debug!("nested_test_outer: in handler");
        outer_trapped = true;
        0
    };

    assert outer_trapped;
}
<commit_msg>core: add potential "simpler interface" to core::condition, also reduce TLS hits.<commit_after>// helper for transmutation, shown below.
type RustClosure = (int,int);

struct Condition<T, U:Copy> {
    key: task::local_data::LocalDataKey<Handler<T,U>>
}

struct Handler<T, U:Copy> {
    handle: RustClosure
}

struct ProtectBlock<T, U:Copy> {
    cond: &Condition<T, U>,
    inner: RustClosure
}

struct Guard<T, U:Copy> {
    cond: &Condition<T,U>,
    prev: Option<@Handler<T, U>>,
    drop {
        match self.prev {
            None => (),
            Some(p) =>
            unsafe {
                debug!("Guard: popping handler from TLS");
                task::local_data::local_data_set(self.cond.key, p)
            }
        }
    }
}

struct HandleBlock<T, U:Copy> {
    pb: &ProtectBlock<T,U>,
    prev: Option<@Handler<T,U>>,
    handler: @Handler<T,U>,
    drop {
        unsafe {
            debug!("HandleBlock: pushing handler to TLS");
            let _g = Guard { cond: self.pb.cond,
                             prev: self.prev };
            task::local_data::local_data_set(self.pb.cond.key,
                                             self.handler);
            // transmutation to avoid copying non-copyable, should
            // be fixable by tracking closure pointees in regionck.
            let f : &fn() = ::cast::transmute(self.pb.inner);
            debug!("HandleBlock: invoking protected code");
            f();
            debug!("HandleBlock: returned from protected code");
        }
    }
}

impl<T, U: Copy> ProtectBlock<T,U> {
    fn handle(&self, h: &self/fn(&T) ->U) -> HandleBlock/&self<T,U> {
        unsafe {
            debug!("ProtectBlock.handle: setting up handler block");
            let p : *RustClosure = ::cast::transmute(&h);
            let prev = task::local_data::local_data_get(self.cond.key);
            HandleBlock { pb: self,
                          prev: prev,
                          handler: @Handler{handle: *p} }
        }
    }
}

impl<T, U: Copy>  Condition<T,U> {

    fn guard(&self, h: &self/fn(&T) ->U) -> Guard/&self<T,U> {
        unsafe {
            let prev = task::local_data::local_data_get(self.key);
            let g = Guard { cond: self, prev: prev };
            debug!("Guard: pushing handler to TLS");
            let p : *RustClosure = ::cast::transmute(&h);
            let h = @Handler{handle: *p};
            task::local_data::local_data_set(self.key, h);
            move g
        }
    }

    fn protect(&self, inner: &self/fn()) -> ProtectBlock/&self<T,U> {
        unsafe {
            // transmutation to avoid copying non-copyable, should
            // be fixable by tracking closure pointees in regionck.
            debug!("Condition.protect: setting up protected block");
            let p : *RustClosure = ::cast::transmute(&inner);
            ProtectBlock { cond: self,
                           inner: *p }
        }
    }

    fn raise(t:&T) -> U {
        unsafe {
            match task::local_data::local_data_get(self.key) {
                None => {
                    debug!("Condition.raise: found no handler");
                    fail
                }

                Some(handler) => {
                    debug!("Condition.raise: found handler");
                    let f : &fn(&T) -> U = ::cast::transmute(handler.handle);
                    f(t)
                }
            }
        }
    }
}

#[cfg(test)]
fn sadness_key(_x: @Handler<int,int>) { }

#[cfg(test)]
fn trouble(i: int) {
    // Condition should work as a const, just limitations in consts.
    let sadness_condition : Condition<int,int> =
        Condition { key: sadness_key };
    debug!("trouble: raising conition");
    let j = sadness_condition.raise(&i);
    debug!("trouble: handler recovered with %d", j);
}

#[test]
fn test1() {

    let sadness_condition : Condition<int,int> =
        Condition { key: sadness_key };

    let mut i = 10;

    let b = do sadness_condition.protect {
        debug!("test1: in protected block");
        trouble(1);
        trouble(2);
        trouble(3);
    };

    do b.handle |j| {
        debug!("test1: in handler");
        i += *j;
        i
    };

    assert i == 16;
}
#[cfg(test)]
fn nested_test_inner() {
    let sadness_condition : Condition<int,int> =
        Condition { key: sadness_key };

    let mut inner_trapped = false;

    let b = do sadness_condition.protect {
        debug!("nested_test_inner: in protected block");
        trouble(1);
    };

    do b.handle |_j| {
        debug!("nested_test_inner: in handler");
        inner_trapped = true;
        0
    };

    assert inner_trapped;
}

#[test]
fn nested_test_outer() {

    let sadness_condition : Condition<int,int> =
        Condition { key: sadness_key };

    let mut outer_trapped = false;

    let b = do sadness_condition.protect {
        debug!("nested_test_outer: in protected block");
        nested_test_inner();
        trouble(1);
    };

    do b.handle |_j| {
        debug!("nested_test_outer: in handler");
        outer_trapped = true;
        0
    };

    assert outer_trapped;
}

#[cfg(test)]
fn nested_guard_test_inner() {
    let sadness_condition : Condition<int,int> =
        Condition { key: sadness_key };

    let mut inner_trapped = false;

    let _g = do sadness_condition.guard |_j| {
        debug!("nested_guard_test_inner: in handler");
        inner_trapped = true;
        0
    };

    debug!("nested_guard_test_inner: in protected block");
    trouble(1);

    assert inner_trapped;
}

#[test]
fn nested_guard_test_outer() {

    let sadness_condition : Condition<int,int> =
        Condition { key: sadness_key };

    let mut outer_trapped = false;

    let _g = do sadness_condition.guard |_j| {
        debug!("nested_guard_test_outer: in handler");
        outer_trapped = true;
        0
    };

    debug!("nested_guard_test_outer: in protected block");
    nested_guard_test_inner();
    trouble(1);

    assert outer_trapped;
}
<|endoftext|>