<commit_before>use super::*;
use crate::{inner::*, norm::*};
use num_traits::{One, Zero};

/// Calc a reflactor `w` from a vector `x`
pub fn calc_reflector(x: &mut ArrayBase<S, Ix1>)
where
    A: Scalar + Lapack,
    S: DataMut<Elem = A>,
{
    let norm = x.norm_l2();
    let alpha = -x[0].mul_real(norm / x[0].abs());
    x[0] -= alpha;
    let inv_rev_norm = A::Real::one() / x.norm_l2();
    azip!(mut a(x) in { *a = a.mul_real(inv_rev_norm)});
}

/// Take a reflection using `w`
pub fn reflect(w: &ArrayBase<S1, Ix1>, a: &mut ArrayBase<S2, Ix1>)
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    assert_eq!(w.len(), a.len());
    let n = a.len();
    let c = A::from(2.0).unwrap() * w.inner(&a);
    for l in 0..n {
        a[l] -= c * w[l];
    }
}

/// Iterative orthogonalizer using Householder reflection
#[derive(Debug, Clone)]
pub struct Householder {
    /// Dimension of orthogonalizer
    dim: usize,

    /// Store Householder reflector.
    ///
    /// The coefficient is copied into another array, and this does not contain
    v: Vec<Array1>,
}

impl Householder {
    /// Create a new orthogonalizer
    pub fn new(dim: usize) -> Self {
        Householder { dim, v: Vec::new() }
    }

    /// Take a Reflection `P = I - 2ww^T`
    fn fundamental_reflection<S>(&self, k: usize, a: &mut ArrayBase<S, Ix1>)
    where
        S: DataMut<Elem = A>,
    {
        assert!(k < self.v.len());
        assert_eq!(a.len(), self.dim, "Input array size mismaches to the dimension");
        reflect(&self.v[k].slice(s![k..]), &mut a.slice_mut(s![k..]));
    }

    /// Take forward reflection `P = P_l ... P_1`
    pub fn forward_reflection<S>(&self, a: &mut ArrayBase<S, Ix1>)
    where
        S: DataMut<Elem = A>,
    {
        assert!(a.len() == self.dim);
        let l = self.v.len();
        for k in 0..l {
            self.fundamental_reflection(k, a);
        }
    }

    /// Take backward reflection `P = P_1 ... P_l`
    pub fn backward_reflection<S>(&self, a: &mut ArrayBase<S, Ix1>)
    where
        S: DataMut<Elem = A>,
    {
        assert!(a.len() == self.dim);
        let l = self.v.len();
        for k in (0..l).rev() {
            self.fundamental_reflection(k, a);
        }
    }

    fn eval_residual<S>(&self, a: &ArrayBase<S, Ix1>) -> A::Real
    where
        S: Data<Elem = A>,
    {
        let l = self.v.len();
        a.slice(s![l..]).norm_l2()
    }
}

impl Orthogonalizer for Householder {
    type Elem = A;

    fn dim(&self) -> usize {
        self.dim
    }

    fn len(&self) -> usize {
        self.v.len()
    }

    fn coeff<S>(&self, a: ArrayBase<S, Ix1>) -> Array1
    where
        S: Data<Elem = A>,
    {
        let mut a = a.into_owned();
        self.forward_reflection(&mut a);
        let res = self.eval_residual(&a);
        let k = self.len();
        let mut c = Array1::zeros(k + 1);
        azip!(mut c(c.slice_mut(s![..k])), a(a.slice(s![..k])) in { *c = a });
        c[k] = A::from_real(res);
        c
    }

    fn append<S>(&mut self, mut a: ArrayBase<S, Ix1>, rtol: A::Real) -> Result<Array1, Array1>
    where
        S: DataMut<Elem = A>,
    {
        assert_eq!(a.len(), self.dim);

        self.forward_reflection(&mut a);

        let k = self.len();
        let alpha = self.eval_residual(&a);
        let alpha = if k  Zero::zero() {
            -a[k].mul_real(alpha / a[k].abs())
        } else {
            A::from_real(alpha)
        };

        let mut coef = Array::zeros(k + 1);
        for i in 0..k {
            coef[i] = a[i];
        }
        coef[k] = alpha;

        if alpha.abs() < rtol {
            // linearly dependent
            return Err(coef);
        }

        assert!(k < a.len()); // this must hold because `alpha == 0` if k >= a.len()

        // Add reflector
        a[k] -= alpha;
        let norm = a.slice(s![k..]).norm_l2();
        azip!(mut a (a.slice_mut(s![..k])) in { *a = A::zero() });
        azip!(mut a (a.slice_mut(s![k..])) in { *a = a.div_real(norm) });
        self.v.push(a.into_owned());
        Ok(coef)
    }

    fn get_q(&self) -> Q {
        assert!(self.len() > 0);
        let mut a = Array::zeros((self.dim(), self.len()));
        for (i, mut col) in a.axis_iter_mut(Axis(1)).enumerate() {
            col[i] = A::one();
            self.backward_reflection(&mut col);
        }
        a
    }
}

/// Online QR decomposition of vectors
pub fn householder(
    iter: impl Iterator<Item = ArrayBase<S, Ix1>>,
    dim: usize,
    rtol: A::Real,
    strategy: Strategy,
) -> (Q, R)
where
    A: Scalar + Lapack,
    S: Data<Elem = A>,
{
    let h = Householder::new(dim);
    qr(iter, h, rtol, strategy)
}

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

    #[test]
    fn check_reflector() {
        let mut a = array![c64::new(1.0, 1.0), c64::new(1.0, 0.0), c64::new(0.0, 1.0)];
        let mut w = a.clone();
        calc_reflector(&mut w);
        reflect(&w, &mut a);
        close_l2(
            &a,
            &array![-c64::new(2.0.sqrt(), 2.0.sqrt()), c64::zero(), c64::zero()],
            1e-9,
        );
    }
}
<commit_msg>Refactoring<commit_after>use super::*;
use crate::{inner::*, norm::*};
use num_traits::One;

/// Calc a reflactor `w` from a vector `x`
pub fn calc_reflector(x: &mut ArrayBase<S, Ix1>) -> A
where
    A: Scalar + Lapack,
    S: DataMut<Elem = A>,
{
    let norm = x.norm_l2();
    let alpha = -x[0].mul_real(norm / x[0].abs());
    x[0] -= alpha;
    let inv_rev_norm = A::Real::one() / x.norm_l2();
    azip!(mut a(x) in { *a = a.mul_real(inv_rev_norm)});
    alpha
}

/// Take a reflection using `w`
pub fn reflect(w: &ArrayBase<S1, Ix1>, a: &mut ArrayBase<S2, Ix1>)
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    assert_eq!(w.len(), a.len());
    let n = a.len();
    let c = A::from(2.0).unwrap() * w.inner(&a);
    for l in 0..n {
        a[l] -= c * w[l];
    }
}

/// Iterative orthogonalizer using Householder reflection
#[derive(Debug, Clone)]
pub struct Householder {
    /// Dimension of orthogonalizer
    dim: usize,

    /// Store Householder reflector.
    ///
    /// The coefficient is copied into another array, and this does not contain
    v: Vec<Array1>,
}

impl Householder {
    /// Create a new orthogonalizer
    pub fn new(dim: usize) -> Self {
        Householder { dim, v: Vec::new() }
    }

    /// Take a Reflection `P = I - 2ww^T`
    fn fundamental_reflection<S>(&self, k: usize, a: &mut ArrayBase<S, Ix1>)
    where
        S: DataMut<Elem = A>,
    {
        assert!(k < self.v.len());
        assert_eq!(a.len(), self.dim, "Input array size mismaches to the dimension");
        reflect(&self.v[k].slice(s![k..]), &mut a.slice_mut(s![k..]));
    }

    /// Take forward reflection `P = P_l ... P_1`
    pub fn forward_reflection<S>(&self, a: &mut ArrayBase<S, Ix1>)
    where
        S: DataMut<Elem = A>,
    {
        assert!(a.len() == self.dim);
        let l = self.v.len();
        for k in 0..l {
            self.fundamental_reflection(k, a);
        }
    }

    /// Take backward reflection `P = P_1 ... P_l`
    pub fn backward_reflection<S>(&self, a: &mut ArrayBase<S, Ix1>)
    where
        S: DataMut<Elem = A>,
    {
        assert!(a.len() == self.dim);
        let l = self.v.len();
        for k in (0..l).rev() {
            self.fundamental_reflection(k, a);
        }
    }

    fn eval_residual<S>(&self, a: &ArrayBase<S, Ix1>) -> A::Real
    where
        S: Data<Elem = A>,
    {
        let l = self.v.len();
        a.slice(s![l..]).norm_l2()
    }
}

impl Orthogonalizer for Householder {
    type Elem = A;

    fn dim(&self) -> usize {
        self.dim
    }

    fn len(&self) -> usize {
        self.v.len()
    }

    fn coeff<S>(&self, a: ArrayBase<S, Ix1>) -> Array1
    where
        S: Data<Elem = A>,
    {
        let mut a = a.into_owned();
        self.forward_reflection(&mut a);
        let res = self.eval_residual(&a);
        let k = self.len();
        let mut c = Array1::zeros(k + 1);
        azip!(mut c(c.slice_mut(s![..k])), a(a.slice(s![..k])) in { *c = a });
        c[k] = A::from_real(res);
        c
    }

    fn append<S>(&mut self, mut a: ArrayBase<S, Ix1>, rtol: A::Real) -> Result<Array1, Array1>
    where
        S: DataMut<Elem = A>,
    {
        assert_eq!(a.len(), self.dim);
        let k = self.len();

        self.forward_reflection(&mut a);
        let mut coef = Array::zeros(k + 1);
        for i in 0..k {
            coef[i] = a[i];
        }
        if self.is_full() {
            return Err(coef); // coef[k] must be zero in this case
        }

        let alpha = calc_reflector(&mut a.slice_mut(s![k..]));
        coef[k] = alpha;

        if alpha.abs() < rtol {
            // linearly dependent
            return Err(coef);
        }
        self.v.push(a.into_owned());
        Ok(coef)
    }

    fn get_q(&self) -> Q {
        assert!(self.len() > 0);
        let mut a = Array::zeros((self.dim(), self.len()));
        for (i, mut col) in a.axis_iter_mut(Axis(1)).enumerate() {
            col[i] = A::one();
            self.backward_reflection(&mut col);
        }
        a
    }
}

/// Online QR decomposition of vectors
pub fn householder(
    iter: impl Iterator<Item = ArrayBase<S, Ix1>>,
    dim: usize,
    rtol: A::Real,
    strategy: Strategy,
) -> (Q, R)
where
    A: Scalar + Lapack,
    S: Data<Elem = A>,
{
    let h = Householder::new(dim);
    qr(iter, h, rtol, strategy)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::assert::*;
    use num_traits::Zero;

    #[test]
    fn check_reflector() {
        let mut a = array![c64::new(1.0, 1.0), c64::new(1.0, 0.0), c64::new(0.0, 1.0)];
        let mut w = a.clone();
        calc_reflector(&mut w);
        reflect(&w, &mut a);
        close_l2(
            &a,
            &array![-c64::new(2.0.sqrt(), 2.0.sqrt()), c64::zero(), c64::zero()],
            1e-9,
        );
    }
}
<|endoftext|>


<commit_before>//! Parsing Expression Grammar for Lea
//!
//! rust-peg generates a parser from this definition at compile time

use std::mem;

use expr_parser::build_tree;
use span::{Span, Spanned, mkspanned};
use op::*;
use parsetree::*;

use lea_core::literal::*;

index -> VarIndex<'input>
    = __* "." __* id:ident {
        VarIndex::DotIndex(id)
    }
    / __* "[" e:expression "]" {
        VarIndex::ExprIndex(Box::new(e))
    }

variable -> Variable<'input>
    = i:ident idxs:index* {
        let mut v = Spanned::new(i.span, VNamed(i.value));
        for idx in idxs {
            let start = v.span.start;
            let end = match idx {
                VarIndex::ExprIndex(ref e) => e.span.start + e.span.len,
                VarIndex::DotIndex(ref id) => id.span.start + id.span.len,
            };

            v = mkspanned(VIndex(Box::new(v), idx), start, end);
        }

        v
    }

/// Returns the expressions to pass as arguments to an invoked function
callargs -> CallArgs<'input>
    = __* "(" args:expression ** listsep __* ")" { CallArgs::Normal(args) }
    / __* strn:string { CallArgs::String(strn) }
    / __* tbl:tablecons { CallArgs::Table(tbl) }

methodname -> Spanned<&'input str>
    = ":" __* name:ident __* { name }

callee -> Expr<'input>
    = v:variable { mkspanned(EVar(v), start_pos, pos) }
    / "(" e:expression ")" { mkspanned(EBraced(Box::new(e)), start_pos, pos) }

call -> Call<'input>
    = callee:callee name:methodname? args:callargs+ {
        let mut it = args.into_iter();
        let first_args = it.next().unwrap();
        let callee_span = callee.span.clone();
        let callee = Box::new(callee);
        let mut what = Spanned::new(callee_span, match name {
            None => SimpleCall(callee, first_args),
            Some(name) => MethodCall(callee, name, first_args),
        });

        for list in it {
            what = Spanned::new(callee_span,
                SimpleCall(Box::new(Spanned::new(what.span, ECall(what.value))), list));
        }

        what.value
    }

atom_inner -> Expr<'input>
    = lit:literal               { Spanned::new(lit.span, ELit(lit.value)) }
    / op:unop a:atom            { mkspanned(EUnOp(op, Box::new(a)), start_pos, pos) }
    / t:tablecons               { mkspanned(ETable(t), start_pos, pos) }
    / a:arraycons               { mkspanned(EArray(a), start_pos, pos) }
    / "function" f:funcbody     { mkspanned(EFunc(f), start_pos, pos) }
    / "..."                     { mkspanned(EVarArgs, start_pos, pos) }
    / c:call                    { mkspanned(ECall(c), start_pos, pos) }
    / callee

// Atomic expression. Either a literal, a unary operator applied to another atom or a full expr
// inside parentheses.
atom -> Expr<'input>
    = __* a:atom_inner __* { a }

kvpair -> (Expr<'input>, Expr<'input>)
    = __* kv:(
        "[" key:expression "]" __* "=" __* val:expression { (key, val) }
        / id:ident __* "=" __* val:expression {
            (Spanned::new(id.span, ELit(TStr(id.value.to_string()))), val)
        }
    ) __* { kv }

tableentry -> TableEntry<'input>
    = pair:kvpair { TableEntry::Pair(pair.0, pair.1) }
    / e:expression { TableEntry::Elem(e) }

tablecons -> TableCons<'input>
    = "{" pairs:tableentry ++ listsep listsep? "}" { pairs }
    / "{" __* "}" { Vec::new() }

arraycons -> Vec<Expr<'input>>
    = "[" vals:expression_list listsep? "]" { vals }
    / "[" __* "]" { Vec::new() }

#[pub]
expression -> Expr<'input>
    = __* left:atom rest:(op:binop r:atom { (op, r) })* __* {
        build_tree(left, rest)
    }

// Comma-separated list of expressions (at least one expression is required)
expression_list -> Vec<Expr<'input>>
    = e:expression ++ listsep { e }

elseif -> Spanned<(Expr<'input>, Block<'input>)>
    = "elseif" e:expression "then" __+ body:block { mkspanned((e, body), start_pos, pos) }

stmt_if -> Stmt<'input>
    = "if" cond:expression "then" __+ body:block elifs:elseif* el:("else" __+ block)? "end" {
        mkspanned(SIf {
            cond: cond,
            body: body,
            elifs: elifs,
            el: el,
        }, start_pos, pos)
    }

// List of comma-separated identifiers with at least one element
identlist -> Vec<Spanned<&'input str>>
    = ids:ident ++ listsep { ids }

// Possibly empty identifier list
eidentlist -> Vec<Spanned<&'input str>>
    = ids:ident ** listsep { ids }

varlist -> Vec<Variable<'input>>
    = vars:variable ++ listsep { vars }

// Returns the parameter list and a boolean indicating whether the function takes variable args
funcparams -> (Vec<Spanned<&'input str>>, bool)
    = "(" __* params:identlist va:(__* ',' __* "...")? __* ")" {
        (params, va.is_some())
    }
    / "(" __* "..." __* ")" { (Vec::new(), true) }
    / "(" __* ")" { (Vec::new(), false) }

funcbody -> Function<'input>
    = __* params:funcparams __* body:block "end" {
        let (pars, varargs) = params;
        Function::<'input>::new(pars, varargs, body)
    }

stmt_inner -> Stmt<'input>
    = stmt_if
    / "while" cond:expression "do" __+ body:block "end" {
        mkspanned(SWhile { cond: cond, body: body }, start_pos, pos)
    }
    / "repeat" __+ body:block "until" cond:expression {
        mkspanned(SRepeat { abort_on: cond, body: body }, start_pos, pos)
    }
    / "for" __+ vars:identlist __+ "in" it:expression_list "do" __+ body:block "end" {
        mkspanned(SForIn {
            vars: vars,
            iter: it,
            body: body,
        }, start_pos, pos)
    }
    / "for" __+ var:ident __* "=" start:expression "," end:expression step:("," expression)? "do" __+ body:block "end" {
        mkspanned(SFor {
            var: var,
            start: start,
            step: step,
            end: end,
            body: body,
        }, start_pos, pos)
    }
    / "function" __+ var:variable method:methodname? f:funcbody {
        mkspanned(match method {
            None => SFunc(var, f),
            Some(method) => SMethod(var, method, f),
        }, start_pos, pos)
    }
    / "local" __+ "function" __+ name:ident f:funcbody {
        mkspanned(SLFunc(name, f), start_pos, pos)
    }
    / "do" __+ b:block "end" { mkspanned(SDo(b), start_pos, pos) }
    / "break" { mkspanned(SBreak, start_pos, pos) }
    / ";" { mkspanned(SSemi, start_pos, pos) }
    / "return" vals:expression_list { mkspanned(SReturn(vals), start_pos, pos) }
    / "return" { mkspanned(SReturn(vec![]), start_pos, pos) }
    / "local" __+ locals:identlist __* "=" exprs:expression_list {
        mkspanned(SDecl(locals, exprs), start_pos, pos)
    }
    / "local" __+ locals:identlist { mkspanned(SDecl(locals, vec![]), start_pos, pos) }
    / c:call { mkspanned(SCall(c), start_pos, pos) }
    / vars:varlist __* "=" vals:expression_list {
        mkspanned(SAssign(vars, vals), start_pos, pos)
    }

#[pub]
statement -> Stmt<'input>
    = __* s:stmt_inner __* { s }

#[pub]
block -> Block<'input>
    = s:statement* {
        Block {
            stmts: s,
            span: Span::new(start_pos, pos),
        }
    }

/// Lexical elements

/// Parses Lua's "long bracket" and returns the content. A long bracket is opened with '[',
/// followed by any number of '=' (including 0), followed by '[' and is closed with ']', followed
/// by the same number of '=' used in the opening bracket, followed by ']'.
longbracket -> &'input str
    = "[" open:"="* "[" ( !("]" close:"="* "]" {?
        if open.len() == close.len() { Ok(()) } else { Err("") }
    }) . )* "]" "="* "]" { &match_str[open.len()+2 .. match_str.len() - open.len() - 2] }

comment -> ()
    = "--" longbracket {()}
    / "--" [^\n]*

real_whitespace = " " / "\t" / "\n" / "\r"

__ = real_whitespace / comment

#[pub]
ident -> Spanned<&'input str>
    = [_a-zA-Z][_a-zA-Z0-9]* { mkspanned(match_str, start_pos, pos) }

decstr -> &'input str
    = [0-9]+ { match_str }

hexdigit -> char
    = [0-9a-fA-F] { match_str.chars().next().unwrap() }

hexstr -> &'input str
    = hexdigit+ { match_str }

octstr -> &'input str
    = [0-7]+ { match_str }

// integer without preceding sign
uinteger -> i64
    = "0x" hex:hexstr { i64::from_str_radix(hex, 16).unwrap() }
    / "0o" oct:octstr { i64::from_str_radix(oct, 8).unwrap() }
    / [0-9]+ { match_str.parse().unwrap() }

integer -> i64
    = minus:"-"? i:uinteger {
        match minus {
            None => i,
            Some(_) => -i,
        }
    }

/// C-like float (only one part before/after the decimal point is required)
cfloat -> &'input str
    = [0-9]* "." decstr { match_str }
    / decstr "." ! "."  { match_str }   // no second "." allowed, since that's the concat operator

/// Float in exponential notation
efloat -> &'input str
    = ( cfloat / decstr ) "e" [+-]? decstr { match_str }

float -> f64
    = "-"? ( efloat / cfloat ) { match_str.parse().unwrap() }

/// Matches the decimal notation of a single byte: 0-255
/// (leading '0's are ignored, but bytes are limited to at most 3 digits, so "0001" for example
/// will return 0 here and parse the "000", but stop at the "1")
dec_byte -> u8
    = [0-9]{,3} {?
        match match_str.parse() {
            Ok(num) => Ok(num),
            Err(_) => Err("decimal byte"),
        }
    }

esc_seq -> char
    = "\\a" { '\x07' }  // bell
    / "\\b" { '\x08' }  // backspace
    / "\\f" { '\x0c' }  // form feed
    / "\\n" { '\n' }
    / "\\r" { '\r' }
    / "\\t" { '\t' }
    / "\\v" { '\x0b' }  // vertical tab
    / "\\\\" { '\\' }
    / "\\\"" { '\"' }
    / "\\'" { '\'' }
    / "\\\n" { '\n' }   // '\' at the end of a line continues the string and embeds a "\n"
    / "\\\r\n" { '\n' } // There's a bug in Windows where a newline is preceded by a "\r"
    / "\\x" digits:hexdigit{2} {
        let string: String = digits.into_iter().collect();
        let i = u8::from_str_radix(string.as_ref(), 16).unwrap();

        i as char
    }
    / "\\0" { 0 as char }
    / "\\" val:dec_byte { val as char }

/// \z will skip all whitespace that follows it. Useful for breaking string literals into multiple
/// lines: "bla\z
///         bla" == "blabla"
blank_esc -> ()
    = "\\z" real_whitespace*

_double_quote_content -> char
    = esc_seq
    / blank_esc _double_quote_content
    / [^\\"\n]  { match_str.chars().next().unwrap() }   //" to fix syntax highlighting

_single_quote_content -> char
    = esc_seq
    / blank_esc _single_quote_content
    / [^\\'\n]  { match_str.chars().next().unwrap() }

#[pub]
string -> String
    = '"' chars:_double_quote_content* blank_esc? '"' { chars.into_iter().collect() }
    / '\'' chars:_single_quote_content* blank_esc? '\'' { chars.into_iter().collect() }
    / content:longbracket { content.to_string() }

boolean -> bool
    = "true" { true } / "false" { false }

_literal -> Literal
    = f:float   { TFloat(f) }
    / i:integer { TInt(i) }
    / s:string  { TStr(s) }
    / b:boolean { TBool(b) }
    / "nil"     { TNil }

#[pub]
literal -> Spanned<Literal>
    = l:_literal { mkspanned(l, start_pos, pos) }

binop -> BinOp
    = "==" {BinOp::Eq} / "!=" {BinOp::NEq} / "~=" {BinOp::NEqLua} / ">=" {BinOp::GEq}
    / "<=" {BinOp::LEq} / ">>" {BinOp::ShiftR} / "<<" {BinOp::ShiftL} / "<" {BinOp::Less}
    / ">" {BinOp::Greater} / "&&" {BinOp::LAnd} / "and" {BinOp::LAndLua} / "||" {BinOp::LOr}
    / "or" {BinOp::LOrLua} / "&" {BinOp::BAnd} / "~" {BinOp::BXor} / "|" {BinOp::BOr}
    / "+" {BinOp::Add} / "-" {BinOp::Sub} / "*" {BinOp::Mul} / "/" {BinOp::Div} / "%" {BinOp::Mod}
    / "^" {BinOp::Pow} / ".." {BinOp::Concat}

unop -> UnOp
    = "-" {UnOp::Negate} / "!" {UnOp::LNot} / "not" {UnOp::LNotLua} / "~" {UnOp::BNot}
    / "#" {UnOp::Len}

listsep = __* ',' __*
<commit_msg>grammar formatting<commit_after>//! Parsing Expression Grammar for Lea
//!
//! rust-peg generates a parser from this definition at compile time

use std::mem;

use expr_parser::build_tree;
use span::{Span, Spanned, mkspanned};
use op::*;
use parsetree::*;

use lea_core::literal::*;

index -> VarIndex<'input>
    = __* "." __* id:ident      { VarIndex::DotIndex(id) }
    / __* "[" e:expression "]"  { VarIndex::ExprIndex(Box::new(e)) }

variable -> Variable<'input>
    = i:ident idxs:index* {
        let mut v = Spanned::new(i.span, VNamed(i.value));
        for idx in idxs {
            let start = v.span.start;
            let end = match idx {
                VarIndex::ExprIndex(ref e) => e.span.start + e.span.len,
                VarIndex::DotIndex(ref id) => id.span.start + id.span.len,
            };

            v = mkspanned(VIndex(Box::new(v), idx), start, end);
        }

        v
    }

/// Returns the expressions to pass as arguments to an invoked function
callargs -> CallArgs<'input>
    = __* "(" args:expression ** listsep __* ")"    { CallArgs::Normal(args) }
    / __* strn:string                               { CallArgs::String(strn) }
    / __* tbl:tablecons                             { CallArgs::Table(tbl) }

methodname -> Spanned<&'input str>
    = ":" __* name:ident __* { name }

callee -> Expr<'input>
    = v:variable            { mkspanned(EVar(v), start_pos, pos) }
    / "(" e:expression ")"  { mkspanned(EBraced(Box::new(e)), start_pos, pos) }

call -> Call<'input>
    = callee:callee name:methodname? args:callargs+ {
        let mut it = args.into_iter();
        let first_args = it.next().unwrap();
        let callee_span = callee.span.clone();
        let callee = Box::new(callee);
        let mut what = Spanned::new(callee_span, match name {
            None => SimpleCall(callee, first_args),
            Some(name) => MethodCall(callee, name, first_args),
        });

        for list in it {
            what = Spanned::new(callee_span,
                SimpleCall(Box::new(Spanned::new(what.span, ECall(what.value))), list));
        }

        what.value
    }

atom_inner -> Expr<'input>
    = lit:literal           { Spanned::new(lit.span, ELit(lit.value)) }
    / op:unop a:atom        { mkspanned(EUnOp(op, Box::new(a)), start_pos, pos) }
    / t:tablecons           { mkspanned(ETable(t), start_pos, pos) }
    / a:arraycons           { mkspanned(EArray(a), start_pos, pos) }
    / "function" f:funcbody { mkspanned(EFunc(f), start_pos, pos) }
    / "..."                 { mkspanned(EVarArgs, start_pos, pos) }
    / c:call                { mkspanned(ECall(c), start_pos, pos) }
    / callee

// Atomic expression. Either a literal, a unary operator applied to another atom or a full expr
// inside parentheses.
atom -> Expr<'input>
    = __* a:atom_inner __* { a }

kvpair -> (Expr<'input>, Expr<'input>)
    = __* kv:(
        "[" key:expression "]" __* "=" __* val:expression { (key, val) }
        / id:ident __* "=" __* val:expression {
            (Spanned::new(id.span, ELit(TStr(id.value.to_string()))), val)
        }
    ) __* { kv }

tableentry -> TableEntry<'input>
    = pair:kvpair   { TableEntry::Pair(pair.0, pair.1) }
    / e:expression  { TableEntry::Elem(e) }

tablecons -> TableCons<'input>
    = "{" pairs:tableentry ++ listsep listsep? "}"  { pairs }
    / "{" __* "}"                                   { Vec::new() }

arraycons -> Vec<Expr<'input>>
    = "[" vals:expression_list listsep? "]" { vals }
    / "[" __* "]"                           { Vec::new() }

#[pub]
expression -> Expr<'input>
    = __* left:atom rest:(op:binop r:atom { (op, r) })* __* {
        build_tree(left, rest)
    }

// Comma-separated list of expressions (at least one expression is required)
expression_list -> Vec<Expr<'input>>
    = e:expression ++ listsep { e }

elseif -> Spanned<(Expr<'input>, Block<'input>)>
    = "elseif" e:expression "then" __+ body:block { mkspanned((e, body), start_pos, pos) }

stmt_if -> Stmt<'input>
    = "if" cond:expression "then" __+ body:block elifs:elseif* el:("else" __+ block)? "end" {
        mkspanned(SIf {
            cond: cond,
            body: body,
            elifs: elifs,
            el: el,
        }, start_pos, pos)
    }

// List of comma-separated identifiers with at least one element
identlist -> Vec<Spanned<&'input str>>
    = ids:ident ++ listsep { ids }

// Possibly empty identifier list
eidentlist -> Vec<Spanned<&'input str>>
    = ids:ident ** listsep { ids }

varlist -> Vec<Variable<'input>>
    = vars:variable ++ listsep { vars }

// Returns the parameter list and a boolean indicating whether the function takes variable args
funcparams -> (Vec<Spanned<&'input str>>, bool)
    = "(" __* params:identlist va:(__* ',' __* "...")? __* ")" {
        (params, va.is_some())
    }
    / "(" __* "..." __* ")" { (Vec::new(), true) }
    / "(" __* ")" { (Vec::new(), false) }

funcbody -> Function<'input>
    = __* params:funcparams __* body:block "end" {
        let (pars, varargs) = params;
        Function::<'input>::new(pars, varargs, body)
    }

stmt_inner -> Stmt<'input>
    = stmt_if
    / "while" cond:expression "do" __+ body:block "end" {
        mkspanned(SWhile { cond: cond, body: body }, start_pos, pos)
    }
    / "repeat" __+ body:block "until" cond:expression {
        mkspanned(SRepeat { abort_on: cond, body: body }, start_pos, pos)
    }
    / "for" __+ vars:identlist __+ "in" it:expression_list "do" __+ body:block "end" {
        mkspanned(SForIn {
            vars: vars,
            iter: it,
            body: body,
        }, start_pos, pos)
    }
    / "for" __+ var:ident __* "=" start:expression "," end:expression step:("," expression)? "do" __+ body:block "end" {
        mkspanned(SFor {
            var: var,
            start: start,
            step: step,
            end: end,
            body: body,
        }, start_pos, pos)
    }
    / "function" __+ var:variable method:methodname? f:funcbody {
        mkspanned(match method {
            None => SFunc(var, f),
            Some(method) => SMethod(var, method, f),
        }, start_pos, pos)
    }
    / "local" __+ "function" __+ name:ident f:funcbody {
        mkspanned(SLFunc(name, f), start_pos, pos)
    }
    / "do" __+ b:block "end" { mkspanned(SDo(b), start_pos, pos) }
    / "break" { mkspanned(SBreak, start_pos, pos) }
    / ";" { mkspanned(SSemi, start_pos, pos) }
    / "return" vals:expression_list { mkspanned(SReturn(vals), start_pos, pos) }
    / "return" { mkspanned(SReturn(vec![]), start_pos, pos) }
    / "local" __+ locals:identlist __* "=" exprs:expression_list {
        mkspanned(SDecl(locals, exprs), start_pos, pos)
    }
    / "local" __+ locals:identlist { mkspanned(SDecl(locals, vec![]), start_pos, pos) }
    / c:call { mkspanned(SCall(c), start_pos, pos) }
    / vars:varlist __* "=" vals:expression_list {
        mkspanned(SAssign(vars, vals), start_pos, pos)
    }

#[pub]
statement -> Stmt<'input>
    = __* s:stmt_inner __* { s }

#[pub]
block -> Block<'input>
    = s:statement* {
        Block {
            stmts: s,
            span: Span::new(start_pos, pos),
        }
    }

/// Lexical elements

/// Parses Lua's "long bracket" and returns the content. A long bracket is opened with '[',
/// followed by any number of '=' (including 0), followed by '[' and is closed with ']', followed
/// by the same number of '=' used in the opening bracket, followed by ']'.
longbracket -> &'input str
    = "[" open:"="* "[" ( !("]" close:"="* "]" {?
        if open.len() == close.len() { Ok(()) } else { Err("") }
    }) . )* "]" "="* "]" { &match_str[open.len()+2 .. match_str.len() - open.len() - 2] }

comment -> ()
    = "--" longbracket { () }
    / "--" [^\n]*

real_whitespace = " " / "\t" / "\n" / "\r"

__ = real_whitespace / comment

#[pub]
ident -> Spanned<&'input str>
    = [_a-zA-Z][_a-zA-Z0-9]* { mkspanned(match_str, start_pos, pos) }

decstr -> &'input str
    = [0-9]+ { match_str }

hexdigit -> char
    = [0-9a-fA-F] { match_str.chars().next().unwrap() }

hexstr -> &'input str
    = hexdigit+ { match_str }

octstr -> &'input str
    = [0-7]+ { match_str }

// integer without preceding sign
uinteger -> i64
    = "0x" hex:hexstr   { i64::from_str_radix(hex, 16).unwrap() }
    / "0o" oct:octstr   { i64::from_str_radix(oct, 8).unwrap() }
    / [0-9]+            { match_str.parse().unwrap() }

integer -> i64
    = minus:"-"? i:uinteger {
        match minus {
            None => i,
            Some(_) => -i,
        }
    }

/// C-like float (only one part before/after the decimal point is required)
cfloat -> &'input str
    = [0-9]* "." decstr { match_str }
    / decstr "." ! "."  { match_str }   // no second "." allowed, since that's the concat operator

/// Float in exponential notation
efloat -> &'input str
    = ( cfloat / decstr ) "e" [+-]? decstr { match_str }

float -> f64
    = "-"? ( efloat / cfloat ) { match_str.parse().unwrap() }

/// Matches the decimal notation of a single byte: 0-255
/// (leading '0's are ignored, but bytes are limited to at most 3 digits, so "0001" for example
/// will return 0 here and parse the "000", but stop at the "1")
dec_byte -> u8
    = [0-9]{,3} {?
        match match_str.parse() {
            Ok(num) => Ok(num),
            Err(_) => Err("decimal byte"),
        }
    }

esc_seq -> char
    = "\\a"     { '\x07' }  // bell
    / "\\b"     { '\x08' }  // backspace
    / "\\f"     { '\x0c' }  // form feed
    / "\\n"     { '\n' }
    / "\\r"     { '\r' }
    / "\\t"     { '\t' }
    / "\\v"     { '\x0b' }  // vertical tab
    / "\\\\"    { '\\' }
    / "\\\""    { '\"' }
    / "\\'"     { '\'' }
    / "\\\n"    { '\n' }    // '\' at the end of a line continues the string and embeds a "\n"
    / "\\\r\n"  { '\n' }    // There's a bug in Windows where a newline is preceded by a "\r"
    / "\\x" digits:hexdigit{2} {
        let string: String = digits.into_iter().collect();
        let i = u8::from_str_radix(string.as_ref(), 16).unwrap();

        i as char
    }
    / "\\0"     { 0 as char }
    / "\\" val:dec_byte { val as char }

/// \z will skip all whitespace that follows it. Useful for breaking string literals into multiple
/// lines: "bla\z
///         bla" == "blabla"
blank_esc -> ()
    = "\\z" real_whitespace*

_double_quote_content -> char
    = esc_seq
    / blank_esc _double_quote_content
    / [^\\"\n]  { match_str.chars().next().unwrap() }   //" to fix syntax highlighting

_single_quote_content -> char
    = esc_seq
    / blank_esc _single_quote_content
    / [^\\'\n]  { match_str.chars().next().unwrap() }

#[pub]
string -> String
    = '"' chars:_double_quote_content* blank_esc? '"'   { chars.into_iter().collect() }
    / '\'' chars:_single_quote_content* blank_esc? '\'' { chars.into_iter().collect() }
    / content:longbracket                               { content.to_string() }

boolean -> bool
    = "true" { true } / "false" { false }

_literal -> Literal
    = f:float   { TFloat(f) }
    / i:integer { TInt(i) }
    / s:string  { TStr(s) }
    / b:boolean { TBool(b) }
    / "nil"     { TNil }

#[pub]
literal -> Spanned<Literal>
    = l:_literal { mkspanned(l, start_pos, pos) }

binop -> BinOp
    = "==" {BinOp::Eq} / "!=" {BinOp::NEq} / "~=" {BinOp::NEqLua} / ">=" {BinOp::GEq}
    / "<=" {BinOp::LEq} / ">>" {BinOp::ShiftR} / "<<" {BinOp::ShiftL} / "<" {BinOp::Less}
    / ">" {BinOp::Greater} / "&&" {BinOp::LAnd} / "and" {BinOp::LAndLua} / "||" {BinOp::LOr}
    / "or" {BinOp::LOrLua} / "&" {BinOp::BAnd} / "~" {BinOp::BXor} / "|" {BinOp::BOr}
    / "+" {BinOp::Add} / "-" {BinOp::Sub} / "*" {BinOp::Mul} / "/" {BinOp::Div} / "%" {BinOp::Mod}
    / "^" {BinOp::Pow} / ".." {BinOp::Concat}

unop -> UnOp
    = "-" {UnOp::Negate} / "!" {UnOp::LNot} / "not" {UnOp::LNotLua} / "~" {UnOp::BNot}
    / "#" {UnOp::Len}

listsep = __* ',' __*
<|endoftext|>


<commit_before>//! A palette for use with Objects

use serde;
use rustc_serialize; // needed for toml
use na;

use std::error::Error;
use std::io::Read;

use std::collections::HashMap;
use std::iter::{repeat,FromIterator};

use kiss3d::scene::SceneNode;

use objects::ObjectID;

/// An RGB color
#[derive(Eq,PartialEq,Ord,PartialOrd,Hash,Copy,Clone)]
pub struct Color(pub u8,pub u8,pub u8);

/// Default colors to use, when no others are specified.
pub static DEFAULT_COLORS : [(u8, u8, u8); 11] = [
    ( 77, 175,  74), // Green
    (152,  78, 163), // Purple
    (255, 127,   0), // Orange
    (228,  26,  28), // Red
    ( 55, 126, 184), // Blue
    (166,  86,  40), // Brown
    (247, 129, 191), // Pink
    (153, 153, 153), // Gray
    (255, 255,  51), // Yellow
    (255, 255, 255), // White
    (  0,   0,   0), // Black
];

impl Color {
    /// Convert to a 3-tuple, from 0 to 1
    pub fn to_floats(self) -> (f32, f32, f32) {
        let Color(r,g,b) = self;
        (
            (r as f32) / (u8::max_value() as f32),
            (g as f32) / (u8::max_value() as f32),
            (b as f32) / (u8::max_value() as f32),
        )
    }

    /// Convert to `na::Pnt3`, needed for `kiss3d::window::Window::draw_text`
    pub fn to_pnt3(self) -> na::Pnt3<f32> {
        let (r,g,b) = self.to_floats();

        na::Pnt3::new(r,g,b)
    }
}

/// A [bool] for keeping track of which part of the objectID should be used
#[derive(Debug,Eq,PartialEq,Ord,PartialOrd,Hash,Clone)]
#[derive(Serialize,Deserialize)]
#[derive(RustcDecodable,RustcEncodable)]
pub struct PartialIDer {
    /// Which sections of an ObjectID should be converted
    pub bools : Vec<bool>,
}

impl PartialIDer {
    /// New, with all set to given boolean value
    pub fn new(n : usize, value : bool) -> Self {
        PartialIDer {
            bools : repeat(value).take(n).collect(),
        }
    }

    /// Get the portion of an ID
    pub fn partial<'a>(&self, name: &'a ObjectID) -> Vec<&'a str> {
        let &ObjectID(ref names) = name;
        names.iter()
            .zip(self.bools.iter())
            .filter_map(|(n, &b)| if b {Some(&n[..])} else {None})
            .collect()
    }

    /// Get the portion of an ID
    pub fn as_id(&self, name: &ObjectID) -> ObjectID {
        let &ObjectID(ref names) = name;
        ObjectID(
            names
            .iter()
            .zip(self.bools.iter())
            .filter_map(|(n, &b)| if b {Some(n.clone())} else {None})
            .collect()
        )
    }

    /// Get a string representing the current state
    pub fn as_string(&self) -> String {
        let mut s = String::with_capacity(self.bools.len());
        let pieces : Vec<String> = self.bools.iter().enumerate().map(
            |(n, &b)| if b {format!("{}", (n+1) % 10)} else {String::from("_")}
        ).collect();
        s.extend(pieces.iter().map(|p| &p[0..1]));
        s
    }
}

/// A way to convert string names to colors
#[derive(Clone)]
pub struct Palette {
    default_colors : Vec<Color>,
    partials : PartialIDer,

    /// Mapping of names -> colors, when found. names not found will be given default colors.
    pub assigned : HashMap<ObjectID, Color>,
    next_color : usize
}

impl Palette {
    /// Set the color of an object to match what the palette says
    pub fn set_color(&mut self, name: &ObjectID, node: &mut SceneNode) {
        let (r, g, b) = self.get_color(name).to_floats();
        node.set_color(r, g, b);
    }

    /// Get the color for a particular ID (using partials mask)
    pub fn get_color(&mut self, name: &ObjectID) -> Color {
        let partial = self.partials.as_id(name);
        let (next_color, default_colors, assigned) =
            (&mut self.next_color, &self.default_colors, &mut self.assigned);
        let color = *assigned.entry(partial).or_insert_with(|| {
            let col = default_colors[*next_color];
            *next_color = (*next_color + 1) % default_colors.len();
            col
        });
        return color
    }

    /// Toggle whether or not to use a certain partial value
    pub fn toggle_partial(&mut self, n : usize){
        if n > self.partials.bools.len() {return;};
        self.partials.bools[n] = !self.partials.bools[n];
        self.next_color = 0;
    }

    /// Check value of a certain partial
    pub fn get_partial(&mut self, n : usize) -> bool {
        self.partials.bools[n]
    }

    /// Check value of a certain partial
    pub fn set_partial(&mut self, n : usize, value : bool){
        self.partials.bools[n] = value;
        self.next_color = 0;
    }

    /// Check value of a certain partial
    pub fn set_all_partial(&mut self, value : bool){
        let n = self.partials.bools.len();
        for i in 0..n {
            self.partials.bools[i] = value;
        }
        self.next_color = 0;
    }

    /// Get a string like '_23____' for which parts of the ObjectID we are coloring with respect to
    pub fn partials_string(&self) -> String {
        self.partials.as_string()
    }
}

impl Default for Palette {
    fn default() -> Self {
        Palette {
            default_colors : FromIterator::from_iter(DEFAULT_COLORS.iter().map(
                |&(r,g,b)|{Color(r,g,b)}
            )),
            partials : PartialIDer::new(8, true),
            assigned : HashMap::new(),
            next_color : 0
        }
    }
}

#[derive(Serialize, RustcEncodable, Copy, Clone)]
struct AssignmentRef<'a> {
    names : &'a ObjectID,
    color : &'a Color
}

fn to_assignments<'a>(assigned : &'a HashMap<ObjectID, Color>) -> Vec<AssignmentRef<'a>> {
    assigned.into_iter().map(|(k, v)| {
        AssignmentRef{
            names : k,
            color : v
        }
    }).collect()
}

/// A reference to a palette, with minimal data, for serialization purposes
#[derive(Serialize, RustcEncodable, Clone)]
pub struct PaletteRef<'a> {
    defaults : &'a Vec<Color>,
    partials : &'a PartialIDer,
    assigned : Vec<AssignmentRef<'a>>
}

impl<'a> From<&'a Palette> for PaletteRef<'a> {
    fn from(palette: &'a Palette) -> Self {
        PaletteRef {
            defaults : &palette.default_colors,
            partials : &palette.partials,
            assigned : to_assignments(&palette.assigned),
        }
    }
}

#[derive(Deserialize, RustcDecodable, Clone)]
struct Assignment {
    names : ObjectID,
    color : Color
}

fn from_assignments(assigned : Vec<Assignment>) -> HashMap<ObjectID, Color> {
    assigned.into_iter().map(|a| {
        (a.names, a.color)
    }).collect()
}

/// A reference to a palette, with minimal data
#[derive(Deserialize, RustcDecodable, Clone)]
pub struct PaletteBasic {
    defaults : Vec<Color>,
    partials : PartialIDer,
    assigned : Option<Vec<Assignment>>,
    next_color : Option<usize>
}

impl From<PaletteBasic> for Palette {
    fn from(p : PaletteBasic) -> Self {
        Palette {
            default_colors : p.defaults,
            partials : p.partials,
            assigned : p.assigned.map(from_assignments).unwrap_or_default(),
            next_color : p.next_color.unwrap_or(0)
        }
    }
}

//TODO: add test for round-trip through toml and json

impl serde::Serialize for Color {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error> where S: serde::Serializer {
        let &Color(r,g,b) = self;
        serde::Serialize::serialize(&(r,g,b), serializer)
    }
}

impl serde::Deserialize for Color {
    fn deserialize<D : serde::Deserializer>(deserializer: &mut D) -> Result<Self, D::Error> {
        let (r,g,b) = try!(serde::Deserialize::deserialize(deserializer));
        Ok(Color(r,g,b))
    }
}

impl rustc_serialize::Encodable for Color {
    fn encode<S: rustc_serialize::Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        let &Color(r,g,b) = self;
        rustc_serialize::Encodable::encode(&(r,g,b), s)
    }
}

impl rustc_serialize::Decodable for Color {
    fn decode<D: rustc_serialize::Decoder>(d: &mut D) -> Result<Self, D::Error> {
        let (r,g,b) = try!(rustc_serialize::Decodable::decode(d));
        Ok(Color(r,g,b))
    }
}

impl serde::Serialize for Palette {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error> where S: serde::Serializer {
        PaletteRef::from(self).serialize(serializer)
    }
}

impl serde::Deserialize for Palette {
    fn deserialize<D : serde::Deserializer>(deserializer: &mut D) -> Result<Self, D::Error>  {
        let p = try!(PaletteBasic::deserialize(deserializer));
        Ok(Palette::from(p))
    }
}

impl rustc_serialize::Encodable for Palette {
    fn encode<S: rustc_serialize::Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        PaletteRef::from(self).encode(s)
    }
}

impl rustc_serialize::Decodable for Palette {
    fn decode<D: rustc_serialize::Decoder>(d: &mut D) -> Result<Self, D::Error> {
        let p = try!(PaletteBasic::decode(d));
        Ok(Palette::from(p))
    }
}
<commit_msg>Added serialization round-trip test for Palette<commit_after>//! A palette for use with Objects

use serde;
use rustc_serialize; // needed for toml
use na;

#[cfg(test)]
use toml;
#[cfg(test)]
use serde_json;

use std::error::Error;
use std::io::Read;

use std::collections::HashMap;
use std::iter::{repeat,FromIterator};

use kiss3d::scene::SceneNode;

use objects::ObjectID;

/// An RGB color
#[derive(Eq,PartialEq,Ord,PartialOrd,Hash,Copy,Clone,Debug)]
pub struct Color(pub u8,pub u8,pub u8);

/// Default colors to use, when no others are specified.
pub static DEFAULT_COLORS : [(u8, u8, u8); 11] = [
    ( 77, 175,  74), // Green
    (152,  78, 163), // Purple
    (255, 127,   0), // Orange
    (228,  26,  28), // Red
    ( 55, 126, 184), // Blue
    (166,  86,  40), // Brown
    (247, 129, 191), // Pink
    (153, 153, 153), // Gray
    (255, 255,  51), // Yellow
    (255, 255, 255), // White
    (  0,   0,   0), // Black
];

impl Color {
    /// Convert to a 3-tuple, from 0 to 1
    pub fn to_floats(self) -> (f32, f32, f32) {
        let Color(r,g,b) = self;
        (
            (r as f32) / (u8::max_value() as f32),
            (g as f32) / (u8::max_value() as f32),
            (b as f32) / (u8::max_value() as f32),
        )
    }

    /// Convert to `na::Pnt3`, needed for `kiss3d::window::Window::draw_text`
    pub fn to_pnt3(self) -> na::Pnt3<f32> {
        let (r,g,b) = self.to_floats();

        na::Pnt3::new(r,g,b)
    }
}

/// A [bool] for keeping track of which part of the objectID should be used
#[derive(Debug,Eq,PartialEq,Ord,PartialOrd,Hash,Clone)]
#[derive(Serialize,Deserialize)]
#[derive(RustcDecodable,RustcEncodable)]
pub struct PartialIDer {
    /// Which sections of an ObjectID should be converted
    pub bools : Vec<bool>,
}

impl PartialIDer {
    /// New, with all set to given boolean value
    pub fn new(n : usize, value : bool) -> Self {
        PartialIDer {
            bools : repeat(value).take(n).collect(),
        }
    }

    /// Get the portion of an ID
    pub fn partial<'a>(&self, name: &'a ObjectID) -> Vec<&'a str> {
        let &ObjectID(ref names) = name;
        names.iter()
            .zip(self.bools.iter())
            .filter_map(|(n, &b)| if b {Some(&n[..])} else {None})
            .collect()
    }

    /// Get the portion of an ID
    pub fn as_id(&self, name: &ObjectID) -> ObjectID {
        let &ObjectID(ref names) = name;
        ObjectID(
            names
            .iter()
            .zip(self.bools.iter())
            .filter_map(|(n, &b)| if b {Some(n.clone())} else {None})
            .collect()
        )
    }

    /// Get a string representing the current state
    pub fn as_string(&self) -> String {
        let mut s = String::with_capacity(self.bools.len());
        let pieces : Vec<String> = self.bools.iter().enumerate().map(
            |(n, &b)| if b {format!("{}", (n+1) % 10)} else {String::from("_")}
        ).collect();
        s.extend(pieces.iter().map(|p| &p[0..1]));
        s
    }
}

/// A way to convert string names to colors
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Palette {
    default_colors : Vec<Color>,
    partials : PartialIDer,

    /// Mapping of names -> colors, when found. names not found will be given default colors.
    pub assigned : HashMap<ObjectID, Color>,
    next_color : usize
}

impl Palette {
    /// Set the color of an object to match what the palette says
    pub fn set_color(&mut self, name: &ObjectID, node: &mut SceneNode) {
        let (r, g, b) = self.get_color(name).to_floats();
        node.set_color(r, g, b);
    }

    /// Get the color for a particular ID (using partials mask)
    pub fn get_color(&mut self, name: &ObjectID) -> Color {
        let partial = self.partials.as_id(name);
        let (next_color, default_colors, assigned) =
            (&mut self.next_color, &self.default_colors, &mut self.assigned);
        let color = *assigned.entry(partial).or_insert_with(|| {
            let col = default_colors[*next_color];
            *next_color = (*next_color + 1) % default_colors.len();
            col
        });
        return color
    }

    /// Toggle whether or not to use a certain partial value
    pub fn toggle_partial(&mut self, n : usize){
        if n > self.partials.bools.len() {return;};
        self.partials.bools[n] = !self.partials.bools[n];
        self.next_color = 0;
    }

    /// Check value of a certain partial
    pub fn get_partial(&mut self, n : usize) -> bool {
        self.partials.bools[n]
    }

    /// Check value of a certain partial
    pub fn set_partial(&mut self, n : usize, value : bool){
        self.partials.bools[n] = value;
        self.next_color = 0;
    }

    /// Check value of a certain partial
    pub fn set_all_partial(&mut self, value : bool){
        let n = self.partials.bools.len();
        for i in 0..n {
            self.partials.bools[i] = value;
        }
        self.next_color = 0;
    }

    /// Get a string like '_23____' for which parts of the ObjectID we are coloring with respect to
    pub fn partials_string(&self) -> String {
        self.partials.as_string()
    }
}

impl Default for Palette {
    fn default() -> Self {
        Palette {
            default_colors : FromIterator::from_iter(DEFAULT_COLORS.iter().map(
                |&(r,g,b)|{Color(r,g,b)}
            )),
            partials : PartialIDer::new(8, true),
            assigned : HashMap::new(),
            next_color : 0
        }
    }
}

#[derive(Serialize, RustcEncodable, Copy, Clone)]
struct AssignmentRef<'a> {
    names : &'a ObjectID,
    color : &'a Color
}

fn to_assignments<'a>(assigned : &'a HashMap<ObjectID, Color>) -> Vec<AssignmentRef<'a>> {
    assigned.into_iter().map(|(k, v)| {
        AssignmentRef{
            names : k,
            color : v
        }
    }).collect()
}

/// A reference to a palette, with minimal data, for serialization purposes
#[derive(Serialize, RustcEncodable, Clone)]
pub struct PaletteRef<'a> {
    defaults : &'a Vec<Color>,
    partials : &'a PartialIDer,
    assigned : Vec<AssignmentRef<'a>>
}

impl<'a> From<&'a Palette> for PaletteRef<'a> {
    fn from(palette: &'a Palette) -> Self {
        PaletteRef {
            defaults : &palette.default_colors,
            partials : &palette.partials,
            assigned : to_assignments(&palette.assigned),
        }
    }
}

#[derive(Deserialize, RustcDecodable, Clone)]
struct Assignment {
    names : ObjectID,
    color : Color
}

fn from_assignments(assigned : Vec<Assignment>) -> HashMap<ObjectID, Color> {
    assigned.into_iter().map(|a| {
        (a.names, a.color)
    }).collect()
}

/// A reference to a palette, with minimal data
#[derive(Deserialize, RustcDecodable, Clone)]
pub struct PaletteBasic {
    defaults : Vec<Color>,
    partials : PartialIDer,
    assigned : Option<Vec<Assignment>>,
    next_color : Option<usize>
}

impl From<PaletteBasic> for Palette {
    fn from(p : PaletteBasic) -> Self {
        Palette {
            default_colors : p.defaults,
            partials : p.partials,
            assigned : p.assigned.map(from_assignments).unwrap_or_default(),
            next_color : p.next_color.unwrap_or(0)
        }
    }
}

impl serde::Serialize for Color {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error> where S: serde::Serializer {
        let &Color(r,g,b) = self;
        serde::Serialize::serialize(&(r,g,b), serializer)
    }
}

impl serde::Deserialize for Color {
    fn deserialize<D : serde::Deserializer>(deserializer: &mut D) -> Result<Self, D::Error> {
        let (r,g,b) = try!(serde::Deserialize::deserialize(deserializer));
        Ok(Color(r,g,b))
    }
}

impl rustc_serialize::Encodable for Color {
    fn encode<S: rustc_serialize::Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        let &Color(r,g,b) = self;
        rustc_serialize::Encodable::encode(&(r,g,b), s)
    }
}

impl rustc_serialize::Decodable for Color {
    fn decode<D: rustc_serialize::Decoder>(d: &mut D) -> Result<Self, D::Error> {
        let (r,g,b) = try!(rustc_serialize::Decodable::decode(d));
        Ok(Color(r,g,b))
    }
}

impl serde::Serialize for Palette {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error> where S: serde::Serializer {
        PaletteRef::from(self).serialize(serializer)
    }
}

impl serde::Deserialize for Palette {
    fn deserialize<D : serde::Deserializer>(deserializer: &mut D) -> Result<Self, D::Error>  {
        let p = try!(PaletteBasic::deserialize(deserializer));
        Ok(Palette::from(p))
    }
}

impl rustc_serialize::Encodable for Palette {
    fn encode<S: rustc_serialize::Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        PaletteRef::from(self).encode(s)
    }
}

impl rustc_serialize::Decodable for Palette {
    fn decode<D: rustc_serialize::Decoder>(d: &mut D) -> Result<Self, D::Error> {
        let p = try!(PaletteBasic::decode(d));
        Ok(Palette::from(p))
    }
}

#[test]
fn palette_toml_round_trip() {
    let p = Palette::default();

    let s = toml::encode_str(&p);
    let p2 = toml::decode_str(&*s).unwrap();
    assert_eq!(p, p2);
}

#[test]
fn palette_json_round_trip() {
    let p = Palette::default();

    let s = serde_json::to_string_pretty(&p).unwrap();
    let p2 = serde_json::from_str(&*s).unwrap();
    assert_eq!(p, p2);
}
<|endoftext|>


<commit_before>use std::{fmt, cmp};
use std::rc::{Rc, Weak};
use std::cell::RefCell;
use std::collections::VecDeque;
use bytes::{Bytes, BytesMut};
use futures::{Async, Poll, Stream};
use futures::task::{Task, current as current_task};

use actix::ResponseType;
use error::PayloadError;

pub(crate) const DEFAULT_BUFFER_SIZE: usize = 65_536; // max buffer size 64k

/// Just Bytes object
pub struct PayloadItem(pub Bytes);

impl ResponseType for PayloadItem {
    type Item = ();
    type Error = ();
}

impl fmt::Debug for PayloadItem {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, f)
    }
}

/// Stream of byte chunks
///
/// Payload stores chunks in vector. First chunk can be received with `.readany()` method.
#[derive(Debug)]
pub struct Payload {
    inner: Rc<RefCell<Inner>>,
}

impl Payload {

    pub(crate) fn new(eof: bool) -> (PayloadSender, Payload) {
        let shared = Rc::new(RefCell::new(Inner::new(eof)));

        (PayloadSender{inner: Rc::downgrade(&shared)}, Payload{inner: shared})
    }

    /// Indicates EOF of payload
    pub fn eof(&self) -> bool {
        self.inner.borrow().eof()
    }

    /// Length of the data in this payload
    pub fn len(&self) -> usize {
        self.inner.borrow().len()
    }

    /// Is payload empty
    pub fn is_empty(&self) -> bool {
        self.inner.borrow().len() == 0
    }

    /// Get first available chunk of data.
    /// Returns Some(PayloadItem) as chunk, `None` indicates eof.
    pub fn readany(&mut self) -> Poll<Option<PayloadItem>, PayloadError> {
        self.inner.borrow_mut().readany()
    }

    /// Get exactly number of bytes
    /// Returns Some(PayloadItem) as chunk, `None` indicates eof.
    pub fn readexactly(&mut self, size: usize) -> Result<Async<Bytes>, PayloadError> {
        self.inner.borrow_mut().readexactly(size)
    }

    /// Read until `\n`
    /// Returns Some(PayloadItem) as line, `None` indicates eof.
    pub fn readline(&mut self) -> Result<Async<Bytes>, PayloadError> {
        self.inner.borrow_mut().readline()
    }

    /// Read until match line
    /// Returns Some(PayloadItem) as line, `None` indicates eof.
    pub fn readuntil(&mut self, line: &[u8]) -> Result<Async<Bytes>, PayloadError> {
        self.inner.borrow_mut().readuntil(line)
    }

    #[doc(hidden)]
    pub fn readall(&mut self) -> Option<Bytes> {
        self.inner.borrow_mut().readall()
    }

    /// Put unused data back to payload
    pub fn unread_data(&mut self, data: Bytes) {
        self.inner.borrow_mut().unread_data(data);
    }

    /// Get size of payload buffer
    pub fn buffer_size(&self) -> usize {
        self.inner.borrow().buffer_size()
    }

    /// Set size of payload buffer
    pub fn set_buffer_size(&self, size: usize) {
        self.inner.borrow_mut().set_buffer_size(size)
    }
}

impl Stream for Payload {
    type Item = PayloadItem;
    type Error = PayloadError;

    fn poll(&mut self) -> Poll<Option<PayloadItem>, PayloadError> {
        self.readany()
    }
}

pub(crate) trait PayloadWriter {
    fn set_error(&mut self, err: PayloadError);

    fn feed_eof(&mut self);

    fn feed_data(&mut self, data: Bytes);

    fn capacity(&self) -> usize;
}

pub(crate) struct PayloadSender {
    inner: Weak<RefCell<Inner>>,
}

impl PayloadWriter for PayloadSender {

    fn set_error(&mut self, err: PayloadError) {
        if let Some(shared) = self.inner.upgrade() {
            shared.borrow_mut().set_error(err)
        }
    }

    fn feed_eof(&mut self) {
        if let Some(shared) = self.inner.upgrade() {
            shared.borrow_mut().feed_eof()
        }
    }

    fn feed_data(&mut self, data: Bytes) {
        if let Some(shared) = self.inner.upgrade() {
            shared.borrow_mut().feed_data(data)
        }
    }

    fn capacity(&self) -> usize {
        if let Some(shared) = self.inner.upgrade() {
            shared.borrow().capacity()
        } else {
            0
        }
    }
}

#[derive(Debug)]
struct Inner {
    len: usize,
    eof: bool,
    err: Option<PayloadError>,
    task: Option<Task>,
    items: VecDeque<Bytes>,
    buf_size: usize,
}

impl Inner {

    fn new(eof: bool) -> Self {
        Inner {
            len: 0,
            eof: eof,
            err: None,
            task: None,
            items: VecDeque::new(),
            buf_size: DEFAULT_BUFFER_SIZE,
        }
    }

    fn set_error(&mut self, err: PayloadError) {
        self.err = Some(err);
        if let Some(task) = self.task.take() {
            task.notify()
        }
    }

    fn feed_eof(&mut self) {
        self.eof = true;
        if let Some(task) = self.task.take() {
            task.notify()
        }
    }

    fn feed_data(&mut self, data: Bytes) {
        self.len += data.len();
        self.items.push_back(data);
        if let Some(task) = self.task.take() {
            task.notify()
        }
    }

    fn eof(&self) -> bool {
        self.items.is_empty() && self.eof
    }

    fn len(&self) -> usize {
        self.len
    }

    fn readany(&mut self) -> Poll<Option<PayloadItem>, PayloadError> {
        if let Some(data) = self.items.pop_front() {
            self.len -= data.len();
            Ok(Async::Ready(Some(PayloadItem(data))))
        } else if self.eof {
            Ok(Async::Ready(None))
        } else if let Some(err) = self.err.take() {
            Err(err)
        } else {
            self.task = Some(current_task());
            Ok(Async::NotReady)
        }
    }

    fn readexactly(&mut self, size: usize) -> Result<Async<Bytes>, PayloadError> {
        if size <= self.len {
            let mut buf = BytesMut::with_capacity(size);
            while buf.len() < size {
                let mut chunk = self.items.pop_front().unwrap();
                let rem = cmp::min(size - buf.len(), chunk.len());
                self.len -= rem;
                buf.extend(&chunk.split_to(rem));
                if !chunk.is_empty() {
                    self.items.push_front(chunk);
                    return Ok(Async::Ready(buf.freeze()))
                }
            }
        }

        if let Some(err) = self.err.take() {
            Err(err)
        } else {
            self.task = Some(current_task());
            Ok(Async::NotReady)
        }
    }

    fn readuntil(&mut self, line: &[u8]) -> Result<Async<Bytes>, PayloadError> {
        let mut idx = 0;
        let mut num = 0;
        let mut offset = 0;
        let mut found = false;
        let mut length = 0;

        for no in 0..self.items.len() {
            {
                let chunk = &self.items[no];
                for (pos, ch) in chunk.iter().enumerate() {
                    if *ch == line[idx] {
                        idx += 1;
                        if idx == line.len() {
                            num = no;
                            offset = pos+1;
                            length += pos+1;
                            found = true;
                            break;
                        }
                    } else {
                        idx = 0
                    }
                }
                if !found {
                    length += chunk.len()
                }
            }

            if found {
                let mut buf = BytesMut::with_capacity(length);
                if num > 0 {
                    for _ in 0..num {
                        buf.extend(self.items.pop_front().unwrap());
                    }
                }
                if offset > 0 {
                    let mut chunk = self.items.pop_front().unwrap();
                    buf.extend(chunk.split_to(offset));
                    if !chunk.is_empty() {
                        self.items.push_front(chunk)
                    }
                }
                self.len -= length;
                return Ok(Async::Ready(buf.freeze()))
            }
        }
        if let Some(err) = self.err.take() {
            Err(err)
        } else {
            self.task = Some(current_task());
            Ok(Async::NotReady)
        }
    }

    fn readline(&mut self) -> Result<Async<Bytes>, PayloadError> {
        self.readuntil(b"\n")
    }

    pub fn readall(&mut self) -> Option<Bytes> {
        let len = self.items.iter().fold(0, |cur, item| cur + item.len());
        if len > 0 {
            let mut buf = BytesMut::with_capacity(len);
            for item in &self.items {
                buf.extend(item);
            }
            self.items = VecDeque::new();
            self.len = 0;
            Some(buf.take().freeze())
        } else {
            None
        }
    }

    fn unread_data(&mut self, data: Bytes) {
        self.len += data.len();
        self.items.push_front(data)
    }

    fn capacity(&self) -> usize {
        if self.len > self.buf_size {
            0
        } else {
            self.buf_size - self.len
        }
    }

    fn buffer_size(&self) -> usize {
        self.buf_size
    }

    fn set_buffer_size(&mut self, size: usize) {
        self.buf_size = size
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io;
    use failure::Fail;
    use futures::future::{lazy, result};
    use tokio_core::reactor::Core;

    #[test]
    fn test_error() {
        let err: PayloadError = io::Error::new(io::ErrorKind::Other, "ParseError").into();
        assert_eq!(format!("{}", err), "ParseError");
        assert_eq!(format!("{}", err.cause().unwrap()), "ParseError");

        let err = PayloadError::Incomplete;
        assert_eq!(format!("{}", err), "A payload reached EOF, but is not complete.");
    }

    #[test]
    fn test_basic() {
        Core::new().unwrap().run(lazy(|| {
            let (_, mut payload) = Payload::new(false);

            assert!(!payload.eof());
            assert!(payload.is_empty());
            assert_eq!(payload.len(), 0);

            match payload.readany() {
                Ok(Async::NotReady) => (),
                _ => panic!("error"),
            }

            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }

    #[test]
    fn test_eof() {
        Core::new().unwrap().run(lazy(|| {
            let (mut sender, mut payload) = Payload::new(false);

            match payload.readany() {
                Ok(Async::NotReady) => (),
                _ => panic!("error"),
            }

            assert!(!payload.eof());

            sender.feed_data(Bytes::from("data"));
            sender.feed_eof();

            assert!(!payload.eof());

            match payload.readany() {
                Ok(Async::Ready(Some(data))) => assert_eq!(&data.0, "data"),
                _ => panic!("error"),
            }
            assert!(payload.is_empty());
            assert!(payload.eof());
            assert_eq!(payload.len(), 0);

            match payload.readany() {
                Ok(Async::Ready(None)) => (),
                _ => panic!("error"),
            }
            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }

    #[test]
    fn test_err() {
        Core::new().unwrap().run(lazy(|| {
            let (mut sender, mut payload) = Payload::new(false);

            match payload.readany() {
                Ok(Async::NotReady) => (),
                _ => panic!("error"),
            }

            sender.set_error(PayloadError::Incomplete);
            match payload.readany() {
                Err(_) => (),
                _ => panic!("error"),
            }
            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }

    #[test]
    fn test_readany() {
        Core::new().unwrap().run(lazy(|| {
            let (mut sender, mut payload) = Payload::new(false);

            sender.feed_data(Bytes::from("line1"));

            assert!(!payload.is_empty());
            assert_eq!(payload.len(), 5);

            sender.feed_data(Bytes::from("line2"));
            assert!(!payload.is_empty());
            assert_eq!(payload.len(), 10);

            match payload.readany() {
                Ok(Async::Ready(Some(data))) => assert_eq!(&data.0, "line1"),
                _ => panic!("error"),
            }
            assert!(!payload.is_empty());
            assert_eq!(payload.len(), 5);

            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }

    #[test]
    fn test_readexactly() {
        Core::new().unwrap().run(lazy(|| {
            let (mut sender, mut payload) = Payload::new(false);

            match payload.readexactly(2) {
                Ok(Async::NotReady) => (),
                _ => panic!("error"),
            }

            sender.feed_data(Bytes::from("line1"));
            sender.feed_data(Bytes::from("line2"));
            assert_eq!(payload.len(), 10);

            match payload.readexactly(2) {
                Ok(Async::Ready(data)) => assert_eq!(&data, "li"),
                _ => panic!("error"),
            }
            assert_eq!(payload.len(), 8);

            match payload.readexactly(4) {
                Ok(Async::Ready(data)) => assert_eq!(&data, "ne1l"),
                _ => panic!("error"),
            }
            assert_eq!(payload.len(), 4);

            sender.set_error(PayloadError::Incomplete);
            match payload.readexactly(10) {
                Err(_) => (),
                _ => panic!("error"),
            }

            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }

    #[test]
    fn test_readuntil() {
        Core::new().unwrap().run(lazy(|| {
            let (mut sender, mut payload) = Payload::new(false);

            match payload.readuntil(b"ne") {
                Ok(Async::NotReady) => (),
                _ => panic!("error"),
            }

            sender.feed_data(Bytes::from("line1"));
            sender.feed_data(Bytes::from("line2"));
            assert_eq!(payload.len(), 10);

            match payload.readuntil(b"ne") {
                Ok(Async::Ready(data)) => assert_eq!(&data, "line"),
                _ => panic!("error"),
            }
            assert_eq!(payload.len(), 6);

            match payload.readuntil(b"2") {
                Ok(Async::Ready(data)) => assert_eq!(&data, "1line2"),
                _ => panic!("error"),
            }
            assert_eq!(payload.len(), 0);

            sender.set_error(PayloadError::Incomplete);
            match payload.readuntil(b"b") {
                Err(_) => (),
                _ => panic!("error"),
            }

            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }

    #[test]
    fn test_unread_data() {
        Core::new().unwrap().run(lazy(|| {
            let (_, mut payload) = Payload::new(false);

            payload.unread_data(Bytes::from("data"));
            assert!(!payload.is_empty());
            assert_eq!(payload.len(), 4);

            match payload.readany() {
                Ok(Async::Ready(Some(data))) => assert_eq!(&data.0, "data"),
                _ => panic!("error"),
            }

            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }
}
<commit_msg>update payload tests<commit_after>use std::{fmt, cmp};
use std::rc::{Rc, Weak};
use std::cell::RefCell;
use std::collections::VecDeque;
use bytes::{Bytes, BytesMut};
use futures::{Async, Poll, Stream};
use futures::task::{Task, current as current_task};

use actix::ResponseType;
use error::PayloadError;

pub(crate) const DEFAULT_BUFFER_SIZE: usize = 65_536; // max buffer size 64k

/// Just Bytes object
#[derive(PartialEq)]
pub struct PayloadItem(pub Bytes);

impl ResponseType for PayloadItem {
    type Item = ();
    type Error = ();
}

impl fmt::Debug for PayloadItem {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, f)
    }
}

/// Stream of byte chunks
///
/// Payload stores chunks in vector. First chunk can be received with `.readany()` method.
#[derive(Debug)]
pub struct Payload {
    inner: Rc<RefCell<Inner>>,
}

impl Payload {

    pub(crate) fn new(eof: bool) -> (PayloadSender, Payload) {
        let shared = Rc::new(RefCell::new(Inner::new(eof)));

        (PayloadSender{inner: Rc::downgrade(&shared)}, Payload{inner: shared})
    }

    /// Indicates EOF of payload
    pub fn eof(&self) -> bool {
        self.inner.borrow().eof()
    }

    /// Length of the data in this payload
    pub fn len(&self) -> usize {
        self.inner.borrow().len()
    }

    /// Is payload empty
    pub fn is_empty(&self) -> bool {
        self.inner.borrow().len() == 0
    }

    /// Get first available chunk of data.
    /// Returns Some(PayloadItem) as chunk, `None` indicates eof.
    pub fn readany(&mut self) -> Poll<Option<PayloadItem>, PayloadError> {
        self.inner.borrow_mut().readany()
    }

    /// Get exactly number of bytes
    /// Returns Some(PayloadItem) as chunk, `None` indicates eof.
    pub fn readexactly(&mut self, size: usize) -> Result<Async<Bytes>, PayloadError> {
        self.inner.borrow_mut().readexactly(size)
    }

    /// Read until `\n`
    /// Returns Some(PayloadItem) as line, `None` indicates eof.
    pub fn readline(&mut self) -> Result<Async<Bytes>, PayloadError> {
        self.inner.borrow_mut().readline()
    }

    /// Read until match line
    /// Returns Some(PayloadItem) as line, `None` indicates eof.
    pub fn readuntil(&mut self, line: &[u8]) -> Result<Async<Bytes>, PayloadError> {
        self.inner.borrow_mut().readuntil(line)
    }

    #[doc(hidden)]
    pub fn readall(&mut self) -> Option<Bytes> {
        self.inner.borrow_mut().readall()
    }

    /// Put unused data back to payload
    pub fn unread_data(&mut self, data: Bytes) {
        self.inner.borrow_mut().unread_data(data);
    }

    /// Get size of payload buffer
    pub fn buffer_size(&self) -> usize {
        self.inner.borrow().buffer_size()
    }

    /// Set size of payload buffer
    pub fn set_buffer_size(&self, size: usize) {
        self.inner.borrow_mut().set_buffer_size(size)
    }
}

impl Stream for Payload {
    type Item = PayloadItem;
    type Error = PayloadError;

    fn poll(&mut self) -> Poll<Option<PayloadItem>, PayloadError> {
        self.readany()
    }
}

pub(crate) trait PayloadWriter {
    fn set_error(&mut self, err: PayloadError);

    fn feed_eof(&mut self);

    fn feed_data(&mut self, data: Bytes);

    fn capacity(&self) -> usize;
}

pub(crate) struct PayloadSender {
    inner: Weak<RefCell<Inner>>,
}

impl PayloadWriter for PayloadSender {

    fn set_error(&mut self, err: PayloadError) {
        if let Some(shared) = self.inner.upgrade() {
            shared.borrow_mut().set_error(err)
        }
    }

    fn feed_eof(&mut self) {
        if let Some(shared) = self.inner.upgrade() {
            shared.borrow_mut().feed_eof()
        }
    }

    fn feed_data(&mut self, data: Bytes) {
        if let Some(shared) = self.inner.upgrade() {
            shared.borrow_mut().feed_data(data)
        }
    }

    fn capacity(&self) -> usize {
        if let Some(shared) = self.inner.upgrade() {
            shared.borrow().capacity()
        } else {
            0
        }
    }
}

#[derive(Debug)]
struct Inner {
    len: usize,
    eof: bool,
    err: Option<PayloadError>,
    task: Option<Task>,
    items: VecDeque<Bytes>,
    buf_size: usize,
}

impl Inner {

    fn new(eof: bool) -> Self {
        Inner {
            len: 0,
            eof: eof,
            err: None,
            task: None,
            items: VecDeque::new(),
            buf_size: DEFAULT_BUFFER_SIZE,
        }
    }

    fn set_error(&mut self, err: PayloadError) {
        self.err = Some(err);
        if let Some(task) = self.task.take() {
            task.notify()
        }
    }

    fn feed_eof(&mut self) {
        self.eof = true;
        if let Some(task) = self.task.take() {
            task.notify()
        }
    }

    fn feed_data(&mut self, data: Bytes) {
        self.len += data.len();
        self.items.push_back(data);
        if let Some(task) = self.task.take() {
            task.notify()
        }
    }

    fn eof(&self) -> bool {
        self.items.is_empty() && self.eof
    }

    fn len(&self) -> usize {
        self.len
    }

    fn readany(&mut self) -> Poll<Option<PayloadItem>, PayloadError> {
        if let Some(data) = self.items.pop_front() {
            self.len -= data.len();
            Ok(Async::Ready(Some(PayloadItem(data))))
        } else if self.eof {
            Ok(Async::Ready(None))
        } else if let Some(err) = self.err.take() {
            Err(err)
        } else {
            self.task = Some(current_task());
            Ok(Async::NotReady)
        }
    }

    fn readexactly(&mut self, size: usize) -> Result<Async<Bytes>, PayloadError> {
        if size <= self.len {
            let mut buf = BytesMut::with_capacity(size);
            while buf.len() < size {
                let mut chunk = self.items.pop_front().unwrap();
                let rem = cmp::min(size - buf.len(), chunk.len());
                self.len -= rem;
                buf.extend(&chunk.split_to(rem));
                if !chunk.is_empty() {
                    self.items.push_front(chunk);
                    return Ok(Async::Ready(buf.freeze()))
                }
            }
        }

        if let Some(err) = self.err.take() {
            Err(err)
        } else {
            self.task = Some(current_task());
            Ok(Async::NotReady)
        }
    }

    fn readuntil(&mut self, line: &[u8]) -> Result<Async<Bytes>, PayloadError> {
        let mut idx = 0;
        let mut num = 0;
        let mut offset = 0;
        let mut found = false;
        let mut length = 0;

        for no in 0..self.items.len() {
            {
                let chunk = &self.items[no];
                for (pos, ch) in chunk.iter().enumerate() {
                    if *ch == line[idx] {
                        idx += 1;
                        if idx == line.len() {
                            num = no;
                            offset = pos+1;
                            length += pos+1;
                            found = true;
                            break;
                        }
                    } else {
                        idx = 0
                    }
                }
                if !found {
                    length += chunk.len()
                }
            }

            if found {
                let mut buf = BytesMut::with_capacity(length);
                if num > 0 {
                    for _ in 0..num {
                        buf.extend(self.items.pop_front().unwrap());
                    }
                }
                if offset > 0 {
                    let mut chunk = self.items.pop_front().unwrap();
                    buf.extend(chunk.split_to(offset));
                    if !chunk.is_empty() {
                        self.items.push_front(chunk)
                    }
                }
                self.len -= length;
                return Ok(Async::Ready(buf.freeze()))
            }
        }
        if let Some(err) = self.err.take() {
            Err(err)
        } else {
            self.task = Some(current_task());
            Ok(Async::NotReady)
        }
    }

    fn readline(&mut self) -> Result<Async<Bytes>, PayloadError> {
        self.readuntil(b"\n")
    }

    pub fn readall(&mut self) -> Option<Bytes> {
        let len = self.items.iter().fold(0, |cur, item| cur + item.len());
        if len > 0 {
            let mut buf = BytesMut::with_capacity(len);
            for item in &self.items {
                buf.extend(item);
            }
            self.items = VecDeque::new();
            self.len = 0;
            Some(buf.take().freeze())
        } else {
            None
        }
    }

    fn unread_data(&mut self, data: Bytes) {
        self.len += data.len();
        self.items.push_front(data)
    }

    fn capacity(&self) -> usize {
        if self.len > self.buf_size {
            0
        } else {
            self.buf_size - self.len
        }
    }

    fn buffer_size(&self) -> usize {
        self.buf_size
    }

    fn set_buffer_size(&mut self, size: usize) {
        self.buf_size = size
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io;
    use failure::Fail;
    use futures::future::{lazy, result};
    use tokio_core::reactor::Core;

    #[test]
    fn test_error() {
        let err: PayloadError = io::Error::new(io::ErrorKind::Other, "ParseError").into();
        assert_eq!(format!("{}", err), "ParseError");
        assert_eq!(format!("{}", err.cause().unwrap()), "ParseError");

        let err = PayloadError::Incomplete;
        assert_eq!(format!("{}", err), "A payload reached EOF, but is not complete.");
    }

    #[test]
    fn test_basic() {
        Core::new().unwrap().run(lazy(|| {
            let (_, mut payload) = Payload::new(false);

            assert!(!payload.eof());
            assert!(payload.is_empty());
            assert_eq!(payload.len(), 0);
            assert_eq!(Async::NotReady, payload.readany().ok().unwrap());

            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }

    #[test]
    fn test_eof() {
        Core::new().unwrap().run(lazy(|| {
            let (mut sender, mut payload) = Payload::new(false);

            assert_eq!(Async::NotReady, payload.readany().ok().unwrap());
            assert!(!payload.eof());

            sender.feed_data(Bytes::from("data"));
            sender.feed_eof();

            assert!(!payload.eof());

            assert_eq!(Async::Ready(Some(PayloadItem(Bytes::from("data")))),
                       payload.readany().ok().unwrap());
            assert!(payload.is_empty());
            assert!(payload.eof());
            assert_eq!(payload.len(), 0);

            assert_eq!(Async::Ready(None), payload.readany().ok().unwrap());
            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }

    #[test]
    fn test_err() {
        Core::new().unwrap().run(lazy(|| {
            let (mut sender, mut payload) = Payload::new(false);

            assert_eq!(Async::NotReady, payload.readany().ok().unwrap());

            sender.set_error(PayloadError::Incomplete);
            payload.readany().err().unwrap();
            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }

    #[test]
    fn test_readany() {
        Core::new().unwrap().run(lazy(|| {
            let (mut sender, mut payload) = Payload::new(false);

            sender.feed_data(Bytes::from("line1"));

            assert!(!payload.is_empty());
            assert_eq!(payload.len(), 5);

            sender.feed_data(Bytes::from("line2"));
            assert!(!payload.is_empty());
            assert_eq!(payload.len(), 10);

            assert_eq!(Async::Ready(Some(PayloadItem(Bytes::from("line1")))),
                       payload.readany().ok().unwrap());
            assert!(!payload.is_empty());
            assert_eq!(payload.len(), 5);

            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }

    #[test]
    fn test_readexactly() {
        Core::new().unwrap().run(lazy(|| {
            let (mut sender, mut payload) = Payload::new(false);

            assert_eq!(Async::NotReady, payload.readexactly(2).ok().unwrap());

            sender.feed_data(Bytes::from("line1"));
            sender.feed_data(Bytes::from("line2"));
            assert_eq!(payload.len(), 10);

            assert_eq!(Async::Ready(Bytes::from("li")), payload.readexactly(2).ok().unwrap());
            assert_eq!(payload.len(), 8);

            assert_eq!(Async::Ready(Bytes::from("ne1l")), payload.readexactly(4).ok().unwrap());
            assert_eq!(payload.len(), 4);

            sender.set_error(PayloadError::Incomplete);
            payload.readexactly(10).err().unwrap();

            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }

    #[test]
    fn test_readuntil() {
        Core::new().unwrap().run(lazy(|| {
            let (mut sender, mut payload) = Payload::new(false);

            assert_eq!(Async::NotReady, payload.readuntil(b"ne").ok().unwrap());

            sender.feed_data(Bytes::from("line1"));
            sender.feed_data(Bytes::from("line2"));
            assert_eq!(payload.len(), 10);

            assert_eq!(Async::Ready(Bytes::from("line")),
                       payload.readuntil(b"ne").ok().unwrap());
            assert_eq!(payload.len(), 6);

            assert_eq!(Async::Ready(Bytes::from("1line2")),
                       payload.readuntil(b"2").ok().unwrap());
            assert_eq!(payload.len(), 0);

            sender.set_error(PayloadError::Incomplete);
            payload.readuntil(b"b").err().unwrap();

            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }

    #[test]
    fn test_unread_data() {
        Core::new().unwrap().run(lazy(|| {
            let (_, mut payload) = Payload::new(false);

            payload.unread_data(Bytes::from("data"));
            assert!(!payload.is_empty());
            assert_eq!(payload.len(), 4);

            assert_eq!(Async::Ready(Some(PayloadItem(Bytes::from("data")))),
                       payload.readany().ok().unwrap());

            let res: Result<(), ()> = Ok(());
            result(res)
        })).unwrap();
    }
}
<|endoftext|>


<commit_before>//! Process-specific information from `/proc/[pid]/`.

mod cwd;
mod stat;
mod statm;
mod status;

pub use pid::cwd::{cwd, cwd_self};
pub use pid::statm::{Statm, statm, statm_self};
pub use pid::status::{SeccompMode, Status, status, status_self};
pub use pid::stat::{Stat};

/// The state of a process.
#[derive(Debug, PartialEq, Eq, Hash)]
pub enum State {
    /// Running.
    Running,
    /// Sleeping in an interruptible wait.
    Sleeping,
    /// Waiting in uninterruptible disk sleep.
    Waiting,
    /// Zombie.
    Zombie,
    /// Stopped (on a signal) or (before Linux 2.6.33) trace stopped.
    Stopped,
    /// trace stopped.
    ///
    /// Linux 2.6.33 onward.
    TraceStopped,
    /// Paging.
    ///
    /// Only before linux 2.6.0.
    Paging,
    /// Dead.
    ///
    /// Linux 2.6.33 to 3.13 only.
    Dead,
    /// Wakekill.
    ///
    /// Linux 2.6.33 to 3.13 only.
    Wakekill,
    /// Waking.
    ///
    /// Linux 2.6.33 to 3.13 only.
    Waking,
    /// Parked.
    ///
    /// Linux 3.9 to 3.13 only.
    Parked,
}

impl Default for State {
    fn default() -> State {
        State::Running
    }
}
<commit_msg>Export stat and stat_self<commit_after>//! Process-specific information from `/proc/[pid]/`.

mod cwd;
mod stat;
mod statm;
mod status;

pub use pid::cwd::{cwd, cwd_self};
pub use pid::statm::{Statm, statm, statm_self};
pub use pid::status::{SeccompMode, Status, status, status_self};
pub use pid::stat::{Stat, stat, stat_self};

/// The state of a process.
#[derive(Debug, PartialEq, Eq, Hash)]
pub enum State {
    /// Running.
    Running,
    /// Sleeping in an interruptible wait.
    Sleeping,
    /// Waiting in uninterruptible disk sleep.
    Waiting,
    /// Zombie.
    Zombie,
    /// Stopped (on a signal) or (before Linux 2.6.33) trace stopped.
    Stopped,
    /// trace stopped.
    ///
    /// Linux 2.6.33 onward.
    TraceStopped,
    /// Paging.
    ///
    /// Only before linux 2.6.0.
    Paging,
    /// Dead.
    ///
    /// Linux 2.6.33 to 3.13 only.
    Dead,
    /// Wakekill.
    ///
    /// Linux 2.6.33 to 3.13 only.
    Wakekill,
    /// Waking.
    ///
    /// Linux 2.6.33 to 3.13 only.
    Waking,
    /// Parked.
    ///
    /// Linux 3.9 to 3.13 only.
    Parked,
}

impl Default for State {
    fn default() -> State {
        State::Running
    }
}
<|endoftext|>


<commit_before>use command::Command;
use vm::VM;

pub struct Program {
    instructions: Vec<Command>,
    instruction_pointer: Option<usize>,
    current_depth: u64,
    pub status: ProgramStatus,
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum ProgramStatus {
    Normal,
    Seeking(u64),
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct ProgramPatch {
    pub instruction_pointer: usize,
    pub current_depth: u64,
    pub status: ProgramStatus,
}

fn find_by<T, F>(list: &[T], start: Option<usize>, mut predicate: F) -> Option<usize>
where
    F: FnMut(&T) -> bool,
{
    let start = start.unwrap_or(0);
    for (offset, element) in list.iter().enumerate() {
        if predicate(element) {
            return Some(start + offset);
        }
    }

    return None;
}

impl Program {
    pub fn new() -> Program {
        Program {
            instructions: Vec::new(),
            instruction_pointer: None,
            current_depth: 0,
            status: ProgramStatus::Normal,
        }
    }

    pub fn append(&mut self, instructions: &[Command]) {
        self.instructions.extend(instructions.iter().cloned());
        if self.instruction_pointer.is_none() {
            self.instruction_pointer = Some(0);
        }
    }

    pub fn execute(&mut self, vm: &mut VM) {
        match (self.instruction_pointer, self.status) {
            (Some(index), ProgramStatus::Seeking(goal)) if index < self.instructions.len() => {
                let patch = Program::attempt_forward_jump(
                    &self.instructions,
                    index,
                    goal,
                    self.current_depth,
                );
                println!("SEEKING BRANCH\n{:?}\n", patch);
                return self.update(patch).execute(vm);
            }

            (Some(start), ProgramStatus::Normal) if start < self.instructions.len() => {
                for index in start..self.instructions.len() {
                    match self.instructions[index] {
                        Command::JumpBackward => {
                            if vm.cells[vm.data_pointer] == 0 {
                                self.current_depth -= 1;
                                continue;
                            }

                            let patch = Program::backward_jump(
                                &self.instructions,
                                index,
                                self.current_depth,
                            );
                            println!("BACKWARD BRANCH\n{:?}\n", patch);
                            return self.update(patch).execute(vm);
                        }

                        Command::JumpForward => {
                            if vm.cells[vm.data_pointer] != 0 {
                                self.current_depth += 1;
                                continue;
                            }

                            let patch = Program::attempt_forward_jump(
                                &self.instructions,
                                index,
                                self.current_depth,
                                self.current_depth,
                            );
                            println!("FORWARD BRANCH\n{:?}\n", patch);
                            return self.update(patch).execute(vm);
                        }

                        command if start < self.instructions.len() => vm.apply(command),

                        _ => {}
                    }
                }
            }

            _ => {}
        }
    }

    fn update(&mut self, patch: ProgramPatch) -> &mut Self {
        self.instruction_pointer = Some(patch.instruction_pointer);
        self.current_depth = patch.current_depth;
        self.status = patch.status;
        return self;
    }

    fn backward_jump(commands: &Vec<Command>, start: usize, depth: u64) -> ProgramPatch {
        let goal_depth = depth - 1;
        let mut current_depth = depth;
        for index in (0..start).rev() {
            match commands[index] {
                Command::JumpBackward => current_depth += 1,
                Command::JumpForward => current_depth -= 1,
                _ => {}
            }
            if goal_depth == current_depth {
                return ProgramPatch {
                    instruction_pointer: index + 1,
                    status: ProgramStatus::Normal,
                    current_depth: current_depth + 1,
                };
            }
        }
        panic!("No opening bracket found?");
    }

    fn attempt_forward_jump(
        commands: &Vec<Command>,
        index: usize,
        goal_depth: u64,
        mut current_depth: u64,
    ) -> ProgramPatch {
        let index = find_by(&commands[index..], Some(index), |command| {
            match command {
                &Command::JumpBackward => current_depth -= 1,
                &Command::JumpForward => current_depth += 1,
                _ => {}
            };
            return current_depth == goal_depth;
        });

        return match index {
            Some(index) => ProgramPatch {
                instruction_pointer: index + 1,
                status: ProgramStatus::Normal,
                current_depth,
            },

            None => ProgramPatch {
                instruction_pointer: commands.len(),
                status: ProgramStatus::Seeking(goal_depth),
                current_depth,
            },
        };
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn attempt_forward_jump() {
        let commands = vec![
            Command::JumpForward,
            Command::JumpForward,
            Command::JumpBackward,
        ];
        let index = 0;
        let goal_depth = 0;
        let current_depth = 0;
        let actual = Program::attempt_forward_jump(&commands, index, goal_depth, current_depth);
        assert_eq!(
            actual,
            ProgramPatch {
                instruction_pointer: 3,
                current_depth: 1,
                status: ProgramStatus::Seeking(0),
            }
        )
    }
}
<commit_msg>Rewrite execute function to utilize ProgramPatch more<commit_after>use command::Command;
use vm::VM;

pub struct Program {
    instructions: Vec<Command>,
    instruction_pointer: Option<usize>,
    current_depth: u64,
    pub status: ProgramStatus,
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum ProgramStatus {
    Normal,
    Seeking(u64),
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct ProgramPatch {
    pub instruction_pointer: usize,
    pub current_depth: u64,
    pub status: ProgramStatus,
}

fn find_by<T, F>(list: &[T], start: Option<usize>, mut predicate: F) -> Option<usize>
where
    F: FnMut(&T) -> bool,
{
    let start = start.unwrap_or(0);
    for (offset, element) in list.iter().enumerate() {
        if predicate(element) {
            return Some(start + offset);
        }
    }

    return None;
}

impl Program {
    pub fn new() -> Program {
        Program {
            instructions: Vec::new(),
            instruction_pointer: None,
            current_depth: 0,
            status: ProgramStatus::Normal,
        }
    }

    pub fn append(&mut self, instructions: &[Command]) {
        self.instructions.extend(instructions.iter().cloned());
        if self.instruction_pointer.is_none() {
            self.instruction_pointer = Some(0);
        }
    }

    pub fn execute(&mut self, vm: &mut VM) {
        if let Some(mut index) = self.instruction_pointer {
            while index < self.instructions.len() {
                let current_cell = vm.cells[vm.data_pointer];
                let patch = match (self.instructions[index], self.status) {
                    (_, ProgramStatus::Seeking(goal)) => {
                        Program::attempt_forward_jump(
                            &self.instructions,
                            index,
                            goal,
                            self.current_depth,
                        )
                    }

                    (Command::JumpForward, _) => {
                        match current_cell {
                            0 => {
                                Program::attempt_forward_jump(
                                    &self.instructions,
                                    index,
                                    self.current_depth,
                                    self.current_depth,
                                )
                            }

                            _ => ProgramPatch {
                                instruction_pointer: index + 1,
                                status: ProgramStatus::Normal,
                                current_depth: self.current_depth + 1,
                            },
                        }
                    }

                    (Command::JumpBackward, _) => {
                        match current_cell {
                            0 => ProgramPatch {
                                instruction_pointer: index + 1,
                                status: ProgramStatus::Normal,
                                current_depth: self.current_depth - 1,
                            },

                            _ => {
                                Program::backward_jump(
                                    &self.instructions,
                                    index,
                                    self.current_depth,
                                )
                            }
                        }
                    }

                    (command, _) => {
                        vm.apply(command);
                        ProgramPatch {
                            status: ProgramStatus::Normal,
                            current_depth: self.current_depth,
                            instruction_pointer: index + 1,
                        }
                    }
                };

                self.update(patch);
                index = patch.instruction_pointer;
            }
        }
    }

    fn update(&mut self, patch: ProgramPatch) -> &mut Self {
        self.instruction_pointer = Some(patch.instruction_pointer);
        self.current_depth = patch.current_depth;
        self.status = patch.status;
        return self;
    }

    fn backward_jump(commands: &Vec<Command>, start: usize, depth: u64) -> ProgramPatch {
        let goal_depth = depth - 1;
        let mut current_depth = depth;
        for index in (0..start).rev() {
            match commands[index] {
                Command::JumpBackward => current_depth += 1,
                Command::JumpForward => current_depth -= 1,
                _ => {}
            }
            if goal_depth == current_depth {
                return ProgramPatch {
                    instruction_pointer: index + 1,
                    status: ProgramStatus::Normal,
                    current_depth: current_depth + 1,
                };
            }
        }
        panic!("No opening bracket found?");
    }

    fn attempt_forward_jump(
        commands: &Vec<Command>,
        index: usize,
        goal_depth: u64,
        mut current_depth: u64,
    ) -> ProgramPatch {
        let index = find_by(&commands[index..], Some(index), |command| {
            match command {
                &Command::JumpBackward => current_depth -= 1,
                &Command::JumpForward => current_depth += 1,
                _ => {}
            };
            return current_depth == goal_depth;
        });

        return match index {
            Some(index) => ProgramPatch {
                instruction_pointer: index + 1,
                status: ProgramStatus::Normal,
                current_depth,
            },

            None => ProgramPatch {
                instruction_pointer: commands.len(),
                status: ProgramStatus::Seeking(goal_depth),
                current_depth,
            },
        };
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn attempt_forward_jump() {
        let commands = vec![
            Command::JumpForward,
            Command::JumpForward,
            Command::JumpBackward,
        ];
        let index = 0;
        let goal_depth = 0;
        let current_depth = 0;
        let actual = Program::attempt_forward_jump(&commands, index, goal_depth, current_depth);
        assert_eq!(
            actual,
            ProgramPatch {
                instruction_pointer: 3,
                current_depth: 1,
                status: ProgramStatus::Seeking(0),
            }
        )
    }
}
<|endoftext|>