<commit_before>#![feature(generators)]

#![allow(non_camel_case_types)]
#![allow(dead_code)]
#![allow(unreachable_code)]
#![allow(unused_parens)]

#![recursion_limit = "256"]

use std::cell::Cell;
use std::mem::swap;

// Just a grab bag of stuff that you wouldn't want to actually write.

fn strange() -> bool { let _x: bool = return true; }

fn funny() {
    fn f(_x: ()) { }
    f(return);
}

fn what() {
    fn the(x: &Cell<bool>) {
        return while !x.get() { x.set(true); };
    }
    let i = &Cell::new(false);
    let dont = {||the(i)};
    dont();
    assert!((i.get()));
}

fn zombiejesus() {
    loop {
        while (return) {
            if (return) {
                match (return) {
                    1 => {
                        if (return) {
                            return
                        } else {
                            return
                        }
                    }
                    _ => { return }
                };
            } else if (return) {
                return;
            }
        }
        if (return) { break; }
    }
}

fn notsure() {
    let mut _x: isize;
    let mut _y = (_x = 0) == (_x = 0);
    let mut _z = (_x = 0) < (_x = 0);
    let _a = (_x += 0) == (_x = 0);
    let _b = swap(&mut _y, &mut _z) == swap(&mut _y, &mut _z);
}

fn canttouchthis() -> usize {
    fn p() -> bool { true }
    let _a = (assert!((true)) == (assert!(p())));
    let _c = (assert!((p())) == ());
    let _b: bool = (println!("{}", 0) == (return 0));
}

fn angrydome() {
    loop { if break { } }
    let mut i = 0;
    loop { i += 1; if i == 1 { match (continue) { 1 => { }, _ => panic!("wat") } }
      break; }
}

fn evil_lincoln() { let _evil = println!("lincoln"); }

fn dots() {
    assert_eq!(String::from(".................................................."),
               format!("{:?}", .. .. .. .. .. .. .. .. .. .. .. .. ..
                               .. .. .. .. .. .. .. .. .. .. .. ..));
}

fn u8(u8: u8) {
    if u8 != 0u8 {
        assert_eq!(8u8, {
            macro_rules! u8 {
                (u8) => {
                    mod u8 {
                        pub fn u8<'u8: 'u8 + 'u8>(u8: &'u8 u8) -> &'u8 u8 {
                            "u8";
                            u8
                        }
                    }
                };
            }

            u8!(u8);
            let &u8: &u8 = u8::u8(&8u8);
            ::u8(0u8);
            u8
        });
    }
}

fn fishy() {
    assert_eq!(String::from("><>"),
               String::<>::from::<>("><>").chars::<>().rev::<>().collect::<String>());
}

fn union() {
    union union<'union> { union: &'union union<'union>, }
}

fn special_characters() {
    let val = !((|(..):(_,_),__@_|__)((&*"\\",'🤔')/**/,{})=={&[..=..][..];})//
    ;
    assert!(!val);
}

fn punch_card() -> impl std::fmt::Debug {
    ..=..=.. ..    .. .. .. ..    .. .. .. ..    .. ..=.. ..
    ..=.. ..=..    .. .. .. ..    .. .. .. ..    ..=..=..=..
    ..=.. ..=..    ..=.. ..=..    .. ..=..=..    .. ..=.. ..
    ..=..=.. ..    ..=.. ..=..    ..=.. .. ..    .. ..=.. ..
    ..=.. ..=..    ..=.. ..=..    .. ..=.. ..    .. ..=.. ..
    ..=.. ..=..    ..=.. ..=..    .. .. ..=..    .. ..=.. ..
    ..=.. ..=..    .. ..=..=..    ..=..=.. ..    .. ..=.. ..
}

fn r#match() {
    let val = match match match match match () {
        () => ()
    } {
        () => ()
    } {
        () => ()
    } {
        () => ()
    } {
        () => ()
    };
    assert_eq!(val, ());
}

fn i_yield() {
    static || {
        yield yield yield yield yield yield yield yield yield;
    };
}

pub fn main() {
    strange();
    funny();
    what();
    zombiejesus();
    notsure();
    canttouchthis();
    angrydome();
    evil_lincoln();
    dots();
    u8(8u8);
    fishy();
    union();
    special_characters();
    punch_card();
    r#match();
    i_yield();
}
<commit_msg>Auto merge of #61892 - rijenkii:master, r=Centril<commit_after>#![feature(generators)]

#![allow(non_camel_case_types)]
#![allow(dead_code)]
#![allow(unreachable_code)]
#![allow(unused_parens)]

#![recursion_limit = "256"]

use std::cell::Cell;
use std::mem::swap;

// Just a grab bag of stuff that you wouldn't want to actually write.

fn strange() -> bool { let _x: bool = return true; }

fn funny() {
    fn f(_x: ()) { }
    f(return);
}

fn what() {
    fn the(x: &Cell<bool>) {
        return while !x.get() { x.set(true); };
    }
    let i = &Cell::new(false);
    let dont = {||the(i)};
    dont();
    assert!((i.get()));
}

fn zombiejesus() {
    loop {
        while (return) {
            if (return) {
                match (return) {
                    1 => {
                        if (return) {
                            return
                        } else {
                            return
                        }
                    }
                    _ => { return }
                };
            } else if (return) {
                return;
            }
        }
        if (return) { break; }
    }
}

fn notsure() {
    let mut _x: isize;
    let mut _y = (_x = 0) == (_x = 0);
    let mut _z = (_x = 0) < (_x = 0);
    let _a = (_x += 0) == (_x = 0);
    let _b = swap(&mut _y, &mut _z) == swap(&mut _y, &mut _z);
}

fn canttouchthis() -> usize {
    fn p() -> bool { true }
    let _a = (assert!((true)) == (assert!(p())));
    let _c = (assert!((p())) == ());
    let _b: bool = (println!("{}", 0) == (return 0));
}

fn angrydome() {
    loop { if break { } }
    let mut i = 0;
    loop { i += 1; if i == 1 { match (continue) { 1 => { }, _ => panic!("wat") } }
      break; }
}

fn evil_lincoln() { let _evil = println!("lincoln"); }

fn dots() {
    assert_eq!(String::from(".................................................."),
               format!("{:?}", .. .. .. .. .. .. .. .. .. .. .. .. ..
                               .. .. .. .. .. .. .. .. .. .. .. ..));
}

fn u8(u8: u8) {
    if u8 != 0u8 {
        assert_eq!(8u8, {
            macro_rules! u8 {
                (u8) => {
                    mod u8 {
                        pub fn u8<'u8: 'u8 + 'u8>(u8: &'u8 u8) -> &'u8 u8 {
                            "u8";
                            u8
                        }
                    }
                };
            }

            u8!(u8);
            let &u8: &u8 = u8::u8(&8u8);
            ::u8(0u8);
            u8
        });
    }
}

fn fishy() {
    assert_eq!(String::from("><>"),
               String::<>::from::<>("><>").chars::<>().rev::<>().collect::<String>());
}

fn union() {
    union union<'union> { union: &'union union<'union>, }
}

fn special_characters() {
    let val = !((|(..):(_,_),__@_|__)((&*"\\",'🤔')/**/,{})=={&[..=..][..];})//
    ;
    assert!(!val);
}

fn punch_card() -> impl std::fmt::Debug {
    ..=..=.. ..    .. .. .. ..    .. .. .. ..    .. ..=.. ..
    ..=.. ..=..    .. .. .. ..    .. .. .. ..    ..=..=..=..
    ..=.. ..=..    ..=.. ..=..    .. ..=..=..    .. ..=.. ..
    ..=..=.. ..    ..=.. ..=..    ..=.. .. ..    .. ..=.. ..
    ..=.. ..=..    ..=.. ..=..    .. ..=.. ..    .. ..=.. ..
    ..=.. ..=..    ..=.. ..=..    .. .. ..=..    .. ..=.. ..
    ..=.. ..=..    .. ..=..=..    ..=..=.. ..    .. ..=.. ..
}

fn r#match() {
    let val = match match match match match () {
        () => ()
    } {
        () => ()
    } {
        () => ()
    } {
        () => ()
    } {
        () => ()
    };
    assert_eq!(val, ());
}

fn i_yield() {
    static || {
        yield yield yield yield yield yield yield yield yield;
    };
}

fn match_nested_if() {
    let val = match () {
        () if if if if true {true} else {false} {true} else {false} {true} else {false} => true,
        _ => false,
    };
    assert!(val);
}

pub fn main() {
    strange();
    funny();
    what();
    zombiejesus();
    notsure();
    canttouchthis();
    angrydome();
    evil_lincoln();
    dots();
    u8(8u8);
    fishy();
    union();
    special_characters();
    punch_card();
    r#match();
    i_yield();
    match_nested_if();
}
<|endoftext|>


<commit_before>use std::str::FromStr;
use std::collections::HashMap;

type Stack<T> = Vec<T>;

#[derive(Default)]
pub struct Interpreter<'a, 'b, T: 'b> {
    data_stack: Stack<T>,
    vocabulary: HashMap<&'a str, &'b Word<T>>,
}

impl<'a, 'b, T: 'b> Interpreter<'a, 'b, T> where T: FromStr {

    fn push(&mut self, value: T) {
        self.data_stack.push(value)
    }

    fn pop(&mut self) -> Option<T> {
        self.data_stack.pop()
    }

    fn parse(&mut self, input: &str) {
        for token in input.split_whitespace() {
            self.eval_token(token)
        }
    }

    fn eval_token(&mut self, token: &str) {
        if let Ok(literal) = token.parse::<T>() {
            self.push(literal);
        } else {
            if let Some(&word) = self.lookup(token) {
                word.eval_within(self);
            }
        }
    }

    fn define_word(&mut self, name: &'a str, word: &'b Word<T>) {
        self.vocabulary.insert(name, word);
    }

    fn lookup(&self, token: &str) -> Option<&&'b Word<T>> {
        self.vocabulary.get(token)
    }
}

struct Word<T> {
    entry: Box<Fn(&mut Interpreter<T>) -> ()>,
}

impl<T> Word<T> {
    fn eval_within(&self, i: &mut Interpreter<T>) {
      (self.entry)(i)
    }
}

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

    fn fixture_plus() -> Word<i32> {
        Word {
            entry: Box::new(|ref mut i| {
                if let (Some(b), Some(a)) = (i.pop(), i.pop()) {
                    i.push(a + b)
                } else { panic!("Stack underflow") }
            })
        }
    }

    fn fixture_dup() -> Word<i32> {
        Word {
            entry: Box::new(|ref mut i| {
                if let Some(x) = i.pop() {
                    i.push(x);
                    i.push(x);
                } else { panic!("Stack underflow") }
            })
        }
    }

    #[test]
    fn eval_expression() {
        let plus = fixture_plus();
        let dup = fixture_dup();
        let mut interpreter = Interpreter::<i32>::default();
        interpreter.define_word("+", &plus);
        interpreter.define_word("dup", &dup);
        interpreter.parse("2 3 + dup +");
        assert_eq!(Some(10), interpreter.pop())
    }

    #[test]
    fn eval_successive_words() {
        let plus = fixture_plus();
        let dup = fixture_dup();
        let mut interpreter = Interpreter::<i32>::default();
        interpreter.define_word("+", &plus);
        interpreter.define_word("dup", &dup);
        interpreter.eval_token("2");
        interpreter.eval_token("dup");
        interpreter.eval_token("+");
        assert_eq!(Some(4), interpreter.pop())
    }

    #[test]
    fn register_word() {
        let plus = fixture_plus(); // FIXME needs to have longer lifetime than the interpreter…
        let mut interpreter = Interpreter::<i32>::default();
        interpreter.define_word("+", &plus);
        interpreter.push(51);
        interpreter.push(42);
        interpreter.eval_token("+");
        assert_eq!(Some(93), interpreter.pop())
    }

    #[test]
    fn plus_word() {
        let mut interpreter = Interpreter::<i32>::default();
        let plus = fixture_plus();
        interpreter.push(42);
        interpreter.push(51);
        plus.eval_within(&mut interpreter);
        assert_eq!(Some(93), interpreter.pop())
    }

    #[test]
    fn push_literal() {
        let mut interpreter = Interpreter::<i32>::default();
        interpreter.eval_token("42");
        assert_eq!(Some(42), interpreter.pop());
    }

    #[test]
    fn pops_nothing() {
        let mut interpreter = Interpreter::<i32>::default();
        assert_eq!(None, interpreter.pop())
    }

    #[test]
    fn pops_as_pushed() {
        let mut interpreter = Interpreter::<i32>::default();
        interpreter.push(42);
        assert_eq!(Some(42), interpreter.pop());
        assert_eq!(None, interpreter.pop())
    }

    #[test]
    fn pops_in_reverse_order() {
        let mut interpreter = Interpreter::<i32>::default();
        interpreter.push(42);
        interpreter.push(51);
        assert_eq!(Some(51), interpreter.pop());
        assert_eq!(Some(42), interpreter.pop());
        assert_eq!(None, interpreter.pop())
    }
}
<commit_msg>Allocate words on heap via Rc's<commit_after>use std::str::FromStr;
use std::collections::HashMap;
use std::rc::Rc;

type Stack<T> = Vec<T>;

#[derive(Default)]
pub struct Interpreter<'a, T> {
    data_stack: Stack<T>,
    vocabulary: HashMap<&'a str, Rc<Word<T>>>,
}

impl<'a, T> Interpreter<'a, T> where T: FromStr {

    fn push(&mut self, value: T) {
        self.data_stack.push(value)
    }

    fn pop(&mut self) -> Option<T> {
        self.data_stack.pop()
    }

    fn parse(&mut self, input: &str) {
        for token in input.split_whitespace() {
            self.eval_token(token)
        }
    }

    fn eval_token(&mut self, token: &str) {
        if let Ok(literal) = token.parse::<T>() {
            self.push(literal);
        } else {
            if let Some(word) = self.lookup(token) {
                word.eval_within(self);
            }
        }
    }

    fn define_word(&mut self, name: &'a str, word: Rc<Word<T>>) {
        self.vocabulary.insert(name, word);
    }

    fn lookup(&self, token: &str) -> Option<Rc<Word<T>>> {
        if let Some(word) = self.vocabulary.get(token) {
            Some(word.clone())
        } else { None }
    }
}

struct Word<T> {
    entry: Box<Fn(&mut Interpreter<T>) -> ()>,
}

impl<T> Word<T> {
    fn eval_within(&self, i: &mut Interpreter<T>) {
      (self.entry)(i)
    }
}

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

    fn fixture_plus() -> Rc<Word<i32>> {
        Rc::new(Word {
            entry: Box::new(|ref mut i| {
                if let (Some(b), Some(a)) = (i.pop(), i.pop()) {
                    i.push(a + b)
                } else { panic!("Stack underflow") }
            })
        })
    }

    fn fixture_dup() -> Rc<Word<i32>> {
        Rc::new(Word {
            entry: Box::new(|ref mut i| {
                if let Some(x) = i.pop() {
                    i.push(x);
                    i.push(x);
                } else { panic!("Stack underflow") }
            })
        })
    }

    #[test]
    fn eval_interleaved_expressions() {
        let mut interpreter = Interpreter::<i32>::default();
        interpreter.parse("2 3");
        interpreter.define_word("+", fixture_plus());
        interpreter.define_word("dup", fixture_dup());
        interpreter.parse("+ dup +");
        assert_eq!(Some(10), interpreter.pop())
    }

    #[test]
    fn eval_successive_words() {
        let mut interpreter = Interpreter::<i32>::default();
        interpreter.define_word("+", fixture_plus());
        interpreter.define_word("dup", fixture_dup());
        interpreter.eval_token("2");
        interpreter.eval_token("dup");
        interpreter.eval_token("+");
        assert_eq!(Some(4), interpreter.pop())
    }

    #[test]
    fn register_word() {
        let mut interpreter = Interpreter::<i32>::default();
        interpreter.define_word("+", fixture_plus());
        interpreter.push(51);
        interpreter.push(42);
        interpreter.eval_token("+");
        assert_eq!(Some(93), interpreter.pop())
    }

    #[test]
    fn plus_word() {
        let mut interpreter = Interpreter::<i32>::default();
        let plus = fixture_plus();
        interpreter.push(42);
        interpreter.push(51);
        plus.eval_within(&mut interpreter);
        assert_eq!(Some(93), interpreter.pop())
    }

    #[test]
    fn push_literal() {
        let mut interpreter = Interpreter::<i32>::default();
        interpreter.eval_token("42");
        assert_eq!(Some(42), interpreter.pop());
    }

    #[test]
    fn pops_nothing() {
        let mut interpreter = Interpreter::<i32>::default();
        assert_eq!(None, interpreter.pop())
    }

    #[test]
    fn pops_as_pushed() {
        let mut interpreter = Interpreter::<i32>::default();
        interpreter.push(42);
        assert_eq!(Some(42), interpreter.pop());
        assert_eq!(None, interpreter.pop())
    }

    #[test]
    fn pops_in_reverse_order() {
        let mut interpreter = Interpreter::<i32>::default();
        interpreter.push(42);
        interpreter.push(51);
        assert_eq!(Some(51), interpreter.pop());
        assert_eq!(Some(42), interpreter.pop());
        assert_eq!(None, interpreter.pop())
    }
}
<|endoftext|>


<commit_before>use ast::*;
use std::collections::HashMap;
use std::io;
use std::io::prelude::*;

#[derive(Debug, PartialEq, Clone)]
pub enum Value {
    Nil,
    Boolean(bool),
    Number(f64),
    String(String),
}

impl Value {
    fn is_true(&self) -> bool {
        match *self {
            Value::Nil => false,
            Value::Boolean(b) => b,
            _ => true,
        }
    }

    fn to_string(&self) -> String {
        match *self {
            Value::Nil => "nil".into(),
            Value::Boolean(ref b) => b.to_string(),
            Value::Number(ref n) => n.to_string(),
            Value::String(ref s) => s.clone(),
        }
    }
}

struct Environment {
    values: Vec<HashMap<Identifier, Value>>,
}

impl Environment {
    fn new() -> Environment {
        Environment { values: vec![HashMap::new()] }
    }

    fn push(&mut self) {
        self.values.push(HashMap::new());
    }

    fn pop(&mut self) {
        self.values.pop();
    }

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

    fn define(&mut self, identifier: Identifier, value: Value) {
        // NOTE that this allow for variable redefinition. See the chapter.
        let index = self.top();
        self.values[index].insert(identifier, value);
    }

    fn try_set(&mut self, identifier: Identifier, value: Value) -> bool {
        for index in (0..self.values.len()).rev() {
            if self.values[index].contains_key(&identifier) {
                self.values[index].insert(identifier, value);
                return true;
            }
        }
        false
    }

    fn get(&self, identifier: &Identifier) -> Option<&Value> {
        for index in (0..self.values.len()).rev() {
            if let Some(value) = self.values[index].get(identifier) {
                return Some(value);
            }
        }
        None
    }
}

pub trait Interpreter {
    fn execute(&mut self, statement: &Statement) -> Result<Option<Value>, RuntimeError>;
}

pub struct StatementInterpreter {
    environment: Environment,
}

impl StatementInterpreter {
    pub fn new() -> StatementInterpreter {
        StatementInterpreter { environment: Environment::new() }
    }
}
impl Interpreter for StatementInterpreter {
    fn execute(&mut self, statement: &Statement) -> Result<Option<Value>, RuntimeError> {
        statement.execute(&mut self.environment)
    }
}

#[derive(Debug)]
pub enum RuntimeError {
    UnaryMinusTypeMismatch(Value),
    BinaryOperatorTypeMismatch(BinaryOperator, Value, Value),
    UndefinedIdentifier(Identifier),
}

trait Interpret {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError>;
}

trait Execute {
    fn execute(&self, environment: &mut Environment) -> Result<Option<Value>, RuntimeError>;
}

impl Interpret for Expr {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError> {
        match *self {
            Expr::Literal(ref l) => l.interpret(environment),
            Expr::Unary(ref u) => u.interpret(environment),
            Expr::Binary(ref b) => b.interpret(environment),
            Expr::Grouping(ref g) => g.interpret(environment),
            Expr::Identifier(ref i) => i.interpret(environment),
            Expr::Assignment(ref a) => a.interpret(environment),
        }
    }
}

impl Interpret for Literal {
    fn interpret(&self, _: &mut Environment) -> Result<Value, RuntimeError> {
        match *self {
            Literal::NilLiteral => Ok(Value::Nil),
            Literal::BoolLiteral(b) => Ok(Value::Boolean(b)),
            Literal::StringLiteral(ref s) => Ok(Value::String(s.clone())),
            Literal::NumberLiteral(n) => Ok(Value::Number(n)),
        }
    }
}

impl Interpret for Identifier {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError> {
        match environment.get(self) {
            Some(value) => Ok(value.clone()),
            None => Err(RuntimeError::UndefinedIdentifier(self.clone())),
        }
    }
}

impl Interpret for Assignment {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError> {
        let target = match self.lvalue {
            Target::Identifier(ref i) => Identifier { name: i.name.clone() },
        };
        match self.rvalue.interpret(environment) {
            Ok(value) => {
                if environment.try_set(target.clone(), value.clone()) {
                    Ok(value.clone())
                } else {
                    Err(RuntimeError::UndefinedIdentifier(target))
                }
            }
            Err(error) => Err(error),
        }
    }
}

impl Interpret for Grouping {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError> {
        self.expr.interpret(environment)
    }
}

impl Interpret for UnaryExpr {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError> {
        let value = try!(self.right.interpret(environment));
        match self.operator {
            UnaryOperator::Bang => Ok(Value::Boolean(!value.is_true())),
            UnaryOperator::Minus => {
                match value {
                    Value::Number(n) => Ok(Value::Number(-n)),
                    _ => Err(RuntimeError::UnaryMinusTypeMismatch(value)),
                }
            }
        }
    }
}

impl Interpret for BinaryExpr {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError> {
        let left = try!(self.left.interpret(environment));
        let right = try!(self.right.interpret(environment));
        match (&self.operator, &left, &right) {
            (&BinaryOperator::Minus, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Number(l - r))
            }
            (&BinaryOperator::Slash, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Number(l / r))
            }
            (&BinaryOperator::Star, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Number(l * r))
            }
            (&BinaryOperator::Plus, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Number(l + r))
            }
            (&BinaryOperator::Plus, &Value::String(ref l), &Value::String(ref r)) => {
                let mut result = String::new();
                result.push_str(l);
                result.push_str(r);
                Ok(Value::String(result))
            }
            (&BinaryOperator::Greater, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Boolean(l > r))
            }
            (&BinaryOperator::Less, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Boolean(l < r))
            }
            (&BinaryOperator::GreaterEqual, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Boolean(l >= r))
            }
            (&BinaryOperator::LessEqual, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Boolean(l <= r))
            }
            (&BinaryOperator::NotEqual, l, r) => Ok(Value::Boolean(l != r)),
            (&BinaryOperator::Equal, l, r) => Ok(Value::Boolean(l == r)),
            _ => {
                Err(RuntimeError::BinaryOperatorTypeMismatch(self.operator,
                                                             left.clone(),
                                                             right.clone()))
            }
        }
    }
}

impl Execute for Statement {
    fn execute(&self, mut environment: &mut Environment) -> Result<Option<Value>, RuntimeError> {
        match *self {
            Statement::Expression(ref e) => e.interpret(environment).map(Some),
            Statement::Print(ref e) => {
                e.interpret(environment)
                    .map(|value| {
                             println!("{}", value.to_string());
                             let _ = io::stdout().flush(); //TODO: is this okay?
                             None
                         })
            }
            Statement::VariableDefinition(ref identifier) => {
                let identifier = Identifier { name: identifier.name.clone() };
                environment.define(identifier, Value::Nil);
                Ok(None)
            }
            Statement::VariableDefinitionWithInitalizer(ref identifier, ref expression) => {
                let identifier = Identifier { name: identifier.name.clone() };
                expression
                    .interpret(environment)
                    .map(|initializer| {
                             environment.define(identifier, initializer);
                             None
                         })
            }
            Statement::Block(ref b) => {
                environment.push();
                let mut result = Ok(None);
                for statement in &b.statements {
                    result = statement.execute(environment);
                }
                environment.pop();
                result
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use ast::*;
    use interpreter::{Interpret, Execute, Environment, Value};

    #[test]
    fn literal() {
        let mut environment = Environment::new();
        let string = String::from("abc");
        let expr = Expr::Literal(Literal::StringLiteral(string.clone()));
        assert_eq!(Value::String(string),
                   expr.interpret(&mut environment).unwrap());
    }

    #[test]
    fn grouping() {
        let mut environment = Environment::new();
        let expr = Grouping { expr: Expr::Literal(Literal::NumberLiteral(45.67f64)) };
        assert_eq!(Value::Number(45.67f64),
                   expr.interpret(&mut environment).unwrap());
    }

    #[test]
    fn unary() {
        let mut environment = Environment::new();
        let expr = UnaryExpr {
            operator: UnaryOperator::Bang,
            right: Expr::Literal(Literal::BoolLiteral(false)),
        };
        assert_eq!(Value::Boolean(true),
                   expr.interpret(&mut environment).unwrap());
    }

    #[test]
    fn binary() {
        let mut environment = Environment::new();
        let expr = BinaryExpr {
            operator: BinaryOperator::Plus,
            left: Expr::Literal(Literal::NumberLiteral(1.0f64)),
            right: Expr::Literal(Literal::NumberLiteral(1.0f64)),
        };
        assert_eq!(Value::Number(2.0f64),
                   expr.interpret(&mut environment).unwrap());
    }

    #[test]
    fn expression_statement() {
        let mut environment = Environment::new();
        let expr = Expr::Literal(Literal::NumberLiteral(1.0f64));
        let statement = Statement::Expression(expr);
        assert_eq!(Some(Value::Number(1.0f64)),
                   statement.execute(&mut environment).unwrap());
    }

    #[test]
    fn complex_expression_statement() {
        let mut environment = Environment::new();
        let subexpr1 = UnaryExpr {
            operator: UnaryOperator::Minus,
            right: Expr::Literal(Literal::NumberLiteral(2f64)),
        };
        let subexpr2 = Grouping { expr: Expr::Literal(Literal::NumberLiteral(5f64)) };
        let binary_expr = BinaryExpr {
            left: Expr::Unary(Box::new(subexpr1)),
            operator: BinaryOperator::Star,
            right: Expr::Grouping(Box::new(subexpr2)),
        };
        let expr = Expr::Binary(Box::new(binary_expr));
        let statement = Statement::Expression(expr);
        assert_eq!(Some(Value::Number(-10.0f64)),
                   statement.execute(&mut environment).unwrap());
    }

    #[test]
    fn variable_definition() {
        let mut environment = Environment::new();
        let identifier = Identifier { name: "x".into() };
        let statement = Statement::VariableDefinition(identifier.clone());
        assert_eq!(None, statement.execute(&mut environment).unwrap());
        assert_eq!(&Value::Nil, environment.get(&identifier).unwrap());
    }

    #[test]
    fn error_accessing_undefined_variable() {
        let mut environment = Environment::new();
        let identifier = Identifier { name: "x".into() };
        let statement = Statement::Expression(Expr::Identifier(identifier));
        assert!(statement.execute(&mut environment).is_err());
    }

    #[test]
    fn error_assigning_undefined_variable() {
        let mut environment = Environment::new();
        let identifier = Identifier { name: "x".into() };
        let statement = Statement::Expression(
            Expr::Assignment(Box::new(Assignment{
                lvalue: Target::Identifier(identifier),
                rvalue: Expr::Literal(Literal::BoolLiteral(true))}))
            );
        assert!(statement.execute(&mut environment).is_err());
    }

    #[test]
    fn variable_definition_with_initializer() {
        let mut environment = Environment::new();
        let identifier = Identifier { name: "x".into() };
        let expr = Expr::Literal(Literal::NumberLiteral(1.0f64));
        let statement = Statement::VariableDefinitionWithInitalizer(identifier.clone(), expr);
        assert_eq!(None, statement.execute(&mut environment).unwrap());
        assert_eq!(&Value::Number(1.0f64),
                   environment.get(&identifier).unwrap());
    }

    #[test]
    fn block_affects_outer_scope() {
        let mut environment = Environment::new();
        let identifier = Identifier { name: "x".into() };
        let expr = Expr::Literal(Literal::NumberLiteral(1.0f64));
        let outer_statement = Statement::VariableDefinitionWithInitalizer(identifier.clone(), expr);
        assert_eq!(None, outer_statement.execute(&mut environment).unwrap());

        let statements = vec![
            Statement::Expression(
                Expr::Assignment(Box::new(Assignment{
                    lvalue: Target::Identifier(Identifier{name: "x".into()}),
                    rvalue: Expr::Literal(Literal::BoolLiteral(false))})))];
        let block = Statement::Block(Box::new(Block { statements: statements }));
        assert!(block.execute(&mut environment).is_ok());
        assert_eq!(&Value::Boolean(false),
                   environment.get(&identifier).unwrap());
    }
    #[test]
    fn block_variable_dont_escape_scope() {
        let mut environment = Environment::new();
        let identifier = Identifier { name: "x".into() };
        let expr = Expr::Literal(Literal::NumberLiteral(1.0f64));
        let statements = vec![Statement::VariableDefinitionWithInitalizer(identifier.clone(),
                                                                          expr)];
        let block = Statement::Block(Box::new(Block { statements: statements }));
        assert_eq!(None, block.execute(&mut environment).unwrap());
        // The variable declaration gets lost when we exit the scope
        assert_eq!(None, environment.get(&identifier));
    }

}
<commit_msg>Test for truthiness<commit_after>use ast::*;
use std::collections::HashMap;
use std::io;
use std::io::prelude::*;

#[derive(Debug, PartialEq, Clone)]
pub enum Value {
    Nil,
    Boolean(bool),
    Number(f64),
    String(String),
}

impl Value {
    fn is_true(&self) -> bool {
        match *self {
            Value::Nil => false,
            Value::Boolean(b) => b,
            _ => true,
        }
    }

    fn to_string(&self) -> String {
        match *self {
            Value::Nil => "nil".into(),
            Value::Boolean(ref b) => b.to_string(),
            Value::Number(ref n) => n.to_string(),
            Value::String(ref s) => s.clone(),
        }
    }
}

struct Environment {
    values: Vec<HashMap<Identifier, Value>>,
}

impl Environment {
    fn new() -> Environment {
        Environment { values: vec![HashMap::new()] }
    }

    fn push(&mut self) {
        self.values.push(HashMap::new());
    }

    fn pop(&mut self) {
        self.values.pop();
    }

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

    fn define(&mut self, identifier: Identifier, value: Value) {
        // NOTE that this allow for variable redefinition. See the chapter.
        let index = self.top();
        self.values[index].insert(identifier, value);
    }

    fn try_set(&mut self, identifier: Identifier, value: Value) -> bool {
        for index in (0..self.values.len()).rev() {
            if self.values[index].contains_key(&identifier) {
                self.values[index].insert(identifier, value);
                return true;
            }
        }
        false
    }

    fn get(&self, identifier: &Identifier) -> Option<&Value> {
        for index in (0..self.values.len()).rev() {
            if let Some(value) = self.values[index].get(identifier) {
                return Some(value);
            }
        }
        None
    }
}

pub trait Interpreter {
    fn execute(&mut self, statement: &Statement) -> Result<Option<Value>, RuntimeError>;
}

pub struct StatementInterpreter {
    environment: Environment,
}

impl StatementInterpreter {
    pub fn new() -> StatementInterpreter {
        StatementInterpreter { environment: Environment::new() }
    }
}
impl Interpreter for StatementInterpreter {
    fn execute(&mut self, statement: &Statement) -> Result<Option<Value>, RuntimeError> {
        statement.execute(&mut self.environment)
    }
}

#[derive(Debug)]
pub enum RuntimeError {
    UnaryMinusTypeMismatch(Value),
    BinaryOperatorTypeMismatch(BinaryOperator, Value, Value),
    UndefinedIdentifier(Identifier),
}

trait Interpret {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError>;
}

trait Execute {
    fn execute(&self, environment: &mut Environment) -> Result<Option<Value>, RuntimeError>;
}

impl Interpret for Expr {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError> {
        match *self {
            Expr::Literal(ref l) => l.interpret(environment),
            Expr::Unary(ref u) => u.interpret(environment),
            Expr::Binary(ref b) => b.interpret(environment),
            Expr::Grouping(ref g) => g.interpret(environment),
            Expr::Identifier(ref i) => i.interpret(environment),
            Expr::Assignment(ref a) => a.interpret(environment),
        }
    }
}

impl Interpret for Literal {
    fn interpret(&self, _: &mut Environment) -> Result<Value, RuntimeError> {
        match *self {
            Literal::NilLiteral => Ok(Value::Nil),
            Literal::BoolLiteral(b) => Ok(Value::Boolean(b)),
            Literal::StringLiteral(ref s) => Ok(Value::String(s.clone())),
            Literal::NumberLiteral(n) => Ok(Value::Number(n)),
        }
    }
}

impl Interpret for Identifier {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError> {
        match environment.get(self) {
            Some(value) => Ok(value.clone()),
            None => Err(RuntimeError::UndefinedIdentifier(self.clone())),
        }
    }
}

impl Interpret for Assignment {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError> {
        let target = match self.lvalue {
            Target::Identifier(ref i) => Identifier { name: i.name.clone() },
        };
        match self.rvalue.interpret(environment) {
            Ok(value) => {
                if environment.try_set(target.clone(), value.clone()) {
                    Ok(value.clone())
                } else {
                    Err(RuntimeError::UndefinedIdentifier(target))
                }
            }
            Err(error) => Err(error),
        }
    }
}

impl Interpret for Grouping {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError> {
        self.expr.interpret(environment)
    }
}

impl Interpret for UnaryExpr {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError> {
        let value = try!(self.right.interpret(environment));
        match self.operator {
            UnaryOperator::Bang => Ok(Value::Boolean(!value.is_true())),
            UnaryOperator::Minus => {
                match value {
                    Value::Number(n) => Ok(Value::Number(-n)),
                    _ => Err(RuntimeError::UnaryMinusTypeMismatch(value)),
                }
            }
        }
    }
}

impl Interpret for BinaryExpr {
    fn interpret(&self, environment: &mut Environment) -> Result<Value, RuntimeError> {
        let left = try!(self.left.interpret(environment));
        let right = try!(self.right.interpret(environment));
        match (&self.operator, &left, &right) {
            (&BinaryOperator::Minus, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Number(l - r))
            }
            (&BinaryOperator::Slash, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Number(l / r))
            }
            (&BinaryOperator::Star, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Number(l * r))
            }
            (&BinaryOperator::Plus, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Number(l + r))
            }
            (&BinaryOperator::Plus, &Value::String(ref l), &Value::String(ref r)) => {
                let mut result = String::new();
                result.push_str(l);
                result.push_str(r);
                Ok(Value::String(result))
            }
            (&BinaryOperator::Greater, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Boolean(l > r))
            }
            (&BinaryOperator::Less, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Boolean(l < r))
            }
            (&BinaryOperator::GreaterEqual, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Boolean(l >= r))
            }
            (&BinaryOperator::LessEqual, &Value::Number(l), &Value::Number(r)) => {
                Ok(Value::Boolean(l <= r))
            }
            (&BinaryOperator::NotEqual, l, r) => Ok(Value::Boolean(l != r)),
            (&BinaryOperator::Equal, l, r) => Ok(Value::Boolean(l == r)),
            _ => {
                Err(RuntimeError::BinaryOperatorTypeMismatch(self.operator,
                                                             left.clone(),
                                                             right.clone()))
            }
        }
    }
}

impl Execute for Statement {
    fn execute(&self, mut environment: &mut Environment) -> Result<Option<Value>, RuntimeError> {
        match *self {
            Statement::Expression(ref e) => e.interpret(environment).map(Some),
            Statement::Print(ref e) => {
                e.interpret(environment)
                    .map(|value| {
                             println!("{}", value.to_string());
                             let _ = io::stdout().flush(); //TODO: is this okay?
                             None
                         })
            }
            Statement::VariableDefinition(ref identifier) => {
                let identifier = Identifier { name: identifier.name.clone() };
                environment.define(identifier, Value::Nil);
                Ok(None)
            }
            Statement::VariableDefinitionWithInitalizer(ref identifier, ref expression) => {
                let identifier = Identifier { name: identifier.name.clone() };
                expression
                    .interpret(environment)
                    .map(|initializer| {
                             environment.define(identifier, initializer);
                             None
                         })
            }
            Statement::Block(ref b) => {
                environment.push();
                let mut result = Ok(None);
                for statement in &b.statements {
                    result = statement.execute(environment);
                }
                environment.pop();
                result
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use ast::*;
    use interpreter::{Interpret, Execute, Environment, Value};

    #[test]
    fn literal() {
        let mut environment = Environment::new();
        let string = String::from("abc");
        let expr = Expr::Literal(Literal::StringLiteral(string.clone()));
        assert_eq!(Value::String(string),
                   expr.interpret(&mut environment).unwrap());
    }

    #[test]
    fn grouping() {
        let mut environment = Environment::new();
        let expr = Grouping { expr: Expr::Literal(Literal::NumberLiteral(45.67f64)) };
        assert_eq!(Value::Number(45.67f64),
                   expr.interpret(&mut environment).unwrap());
    }

    #[test]
    fn unary() {
        let mut environment = Environment::new();
        let expr = UnaryExpr {
            operator: UnaryOperator::Bang,
            right: Expr::Literal(Literal::BoolLiteral(false)),
        };
        assert_eq!(Value::Boolean(true),
                   expr.interpret(&mut environment).unwrap());
    }

    #[test]
    fn binary() {
        let mut environment = Environment::new();
        let expr = BinaryExpr {
            operator: BinaryOperator::Plus,
            left: Expr::Literal(Literal::NumberLiteral(1.0f64)),
            right: Expr::Literal(Literal::NumberLiteral(1.0f64)),
        };
        assert_eq!(Value::Number(2.0f64),
                   expr.interpret(&mut environment).unwrap());
    }

    #[test]
    fn expression_statement() {
        let mut environment = Environment::new();
        let expr = Expr::Literal(Literal::NumberLiteral(1.0f64));
        let statement = Statement::Expression(expr);
        assert_eq!(Some(Value::Number(1.0f64)),
                   statement.execute(&mut environment).unwrap());
    }

    #[test]
    fn complex_expression_statement() {
        let mut environment = Environment::new();
        let subexpr1 = UnaryExpr {
            operator: UnaryOperator::Minus,
            right: Expr::Literal(Literal::NumberLiteral(2f64)),
        };
        let subexpr2 = Grouping { expr: Expr::Literal(Literal::NumberLiteral(5f64)) };
        let binary_expr = BinaryExpr {
            left: Expr::Unary(Box::new(subexpr1)),
            operator: BinaryOperator::Star,
            right: Expr::Grouping(Box::new(subexpr2)),
        };
        let expr = Expr::Binary(Box::new(binary_expr));
        let statement = Statement::Expression(expr);
        assert_eq!(Some(Value::Number(-10.0f64)),
                   statement.execute(&mut environment).unwrap());
    }

    #[test]
    fn variable_definition() {
        let mut environment = Environment::new();
        let identifier = Identifier { name: "x".into() };
        let statement = Statement::VariableDefinition(identifier.clone());
        assert_eq!(None, statement.execute(&mut environment).unwrap());
        assert_eq!(&Value::Nil, environment.get(&identifier).unwrap());
    }

    #[test]
    fn error_accessing_undefined_variable() {
        let mut environment = Environment::new();
        let identifier = Identifier { name: "x".into() };
        let statement = Statement::Expression(Expr::Identifier(identifier));
        assert!(statement.execute(&mut environment).is_err());
    }

    #[test]
    fn error_assigning_undefined_variable() {
        let mut environment = Environment::new();
        let identifier = Identifier { name: "x".into() };
        let statement = Statement::Expression(
            Expr::Assignment(Box::new(Assignment{
                lvalue: Target::Identifier(identifier),
                rvalue: Expr::Literal(Literal::BoolLiteral(true))}))
            );
        assert!(statement.execute(&mut environment).is_err());
    }

    #[test]
    fn variable_definition_with_initializer() {
        let mut environment = Environment::new();
        let identifier = Identifier { name: "x".into() };
        let expr = Expr::Literal(Literal::NumberLiteral(1.0f64));
        let statement = Statement::VariableDefinitionWithInitalizer(identifier.clone(), expr);
        assert_eq!(None, statement.execute(&mut environment).unwrap());
        assert_eq!(&Value::Number(1.0f64),
                   environment.get(&identifier).unwrap());
    }

    #[test]
    fn block_affects_outer_scope() {
        let mut environment = Environment::new();
        let identifier = Identifier { name: "x".into() };
        let expr = Expr::Literal(Literal::NumberLiteral(1.0f64));
        let outer_statement = Statement::VariableDefinitionWithInitalizer(identifier.clone(), expr);
        assert_eq!(None, outer_statement.execute(&mut environment).unwrap());

        let statements = vec![
            Statement::Expression(
                Expr::Assignment(Box::new(Assignment{
                    lvalue: Target::Identifier(Identifier{name: "x".into()}),
                    rvalue: Expr::Literal(Literal::BoolLiteral(false))})))];
        let block = Statement::Block(Box::new(Block { statements: statements }));
        assert!(block.execute(&mut environment).is_ok());
        assert_eq!(&Value::Boolean(false),
                   environment.get(&identifier).unwrap());
    }
    #[test]
    fn block_variable_dont_escape_scope() {
        let mut environment = Environment::new();
        let identifier = Identifier { name: "x".into() };
        let expr = Expr::Literal(Literal::NumberLiteral(1.0f64));
        let statements = vec![Statement::VariableDefinitionWithInitalizer(identifier.clone(),
                                                                          expr)];
        let block = Statement::Block(Box::new(Block { statements: statements }));
        assert_eq!(None, block.execute(&mut environment).unwrap());
        // The variable declaration gets lost when we exit the scope
        assert_eq!(None, environment.get(&identifier));
    }

    #[test]
    fn truth_table() {
        // Note that lox follows Ruby's rules, not Javascript's
        assert_eq!(false, Value::Nil.is_true());
        assert_eq!(true, Value::Boolean(true).is_true());
        assert_eq!(false, Value::Boolean(false).is_true());
        assert_eq!(true, Value::Number(1.0f64).is_true());
        assert_eq!(true, Value::Number(0.0f64).is_true());
        assert_eq!(true, Value::String("".into()).is_true());
        assert_eq!(true, Value::String("Hello World!".into()).is_true());
    }

}
<|endoftext|>


<commit_before>use std::cell::UnsafeCell;
use std::io;

use buf::{ByteBuf, Appender};
use io::{AsyncRead, AsyncWrite};

const RECV_BUF_SIZE: usize = 128 * 1024;

struct RecvBuf {
    inner: UnsafeCell<ByteBuf>,
}

impl RecvBuf {
    #[inline]
    fn new() -> Self {
        RecvBuf {
            inner: UnsafeCell::new(ByteBuf::with_capacity(RECV_BUF_SIZE)),
        }
    }

    #[inline]
    fn append(v: usize, chain: &mut Appender) -> io::Result<usize> {
        if let Some(block) = chain.last_mut() {
            if block.appendable() >= v {
                return Ok(0);
            }
        }
        chain.append(v);
        Ok(0)
    }

    #[inline]
    fn get_mut(&self) -> &mut ByteBuf {
        let buf = unsafe { &mut *self.inner.get() };
        buf.append(RECV_BUF_SIZE, Self::append).unwrap();
        buf
    }
}

thread_local!(static RECV_BUF: RecvBuf = RecvBuf::new());

#[inline]
fn read<R>(r: &mut R) -> io::Result<Option<(ByteBuf, usize)>>
where
    R: AsyncRead,
{
    RECV_BUF.with(|recv_buf| {
        let buf = recv_buf.get_mut();
        let n = buf.read_in(r)?;
        let data = if n > 0 {
            Some((buf.drain_to(n)?, n))
        } else {
            None
        };
        Ok(data)
    })
}

pub mod read;
pub mod write;
pub mod copy;
pub mod split;

use futures::{Stream, Sink, Poll, Async, StartSend, AsyncSink};
use sink::IntoSink;
use stream::IntoStream;

#[derive(Debug)]
pub struct IStream<R> {
    r: Option<R>,
    would_block: bool,
}

impl<R> IStream<R>
where
    R: AsyncRead,
{
    #[inline]
    pub fn get_ref(&self) -> &R {
        self.r
            .as_ref()
            .expect("Attempted IStream::get_ref after completion")
    }

    #[inline]
    pub fn get_mut(&mut self) -> &mut R {
        self.r
            .as_mut()
            .expect("Attempted IStream::get_mut after completion")
    }

    #[inline]
    pub fn into_inner(self) -> R {
        self.r
            .expect("Attempted IStream::into_inner after completion")
    }

    #[inline]
    fn poll_read(&mut self) -> Poll<Option<ByteBuf>, io::Error> {
        match self.r {
            Some(ref mut r) => {
                if self.would_block {
                    self.would_block = false;
                    r.need_read()?;
                    return Ok(Async::NotReady);
                }
                if !r.is_readable() {
                    return Ok(Async::NotReady);
                }
                match read(r) {
                    Ok(Some((data, _))) => {
                        self.would_block = true;
                        Ok(Async::Ready(Some(data)))
                    }
                    Ok(None) => Ok(Async::Ready(None)),
                    Err(e) => {
                        match e.kind() {
                            io::ErrorKind::WouldBlock => {
                                r.need_read()?;
                                Ok(Async::NotReady)
                            }
                            _ => Err(e),
                        }
                    }
                }
            }
            None => panic!("Attempted to poll IStream after completion"),
        }
    }
}

impl<R> Stream for IStream<R>
where
    R: AsyncRead,
{
    type Item = ByteBuf;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        match self.poll_read() {
            Ok(o) => Ok(o),
            Err(e) => {
                self.r = None;
                Err(e)
            }
        }
    }
}

#[derive(Debug)]
pub struct OStream<W> {
    w: Option<W>,
    buf: Option<ByteBuf>,
}

impl<W> OStream<W>
where
    W: AsyncWrite,
{
    #[inline]
    pub fn get_ref(&self) -> &W {
        self.w
            .as_ref()
            .expect("Attempted OStream::get_ref after completion")
    }

    #[inline]
    pub fn get_mut(&mut self) -> &mut W {
        self.w
            .as_mut()
            .expect("Attempted OStream::get_mut after completion")
    }

    #[inline]
    pub fn into_inner(self) -> W {
        self.w
            .expect("Attempted OStream::into_inner after completion")
    }

    #[inline]
    fn poll_write(&mut self) -> Poll<(), io::Error> {
        match self.w {
            Some(ref mut w) => {
                if w.is_writable() {
                    if let Some(ref mut data) = self.buf {
                        if let Err(e) = data.write_out(w) {
                            return match e.kind() {
                                io::ErrorKind::WouldBlock => {
                                    w.need_write()?;
                                    Ok(Async::NotReady)
                                }
                                _ => Err(e),
                            };
                        }
                        if !data.is_empty() {
                            w.need_write()?;
                            data.compact();
                            return Ok(Async::NotReady);
                        }
                    }
                    w.no_need_write()?;
                    Ok(Async::Ready(()))
                } else {
                    Ok(Async::NotReady)
                }
            }
            None => panic!("Attempted to poll OStream after completion"),
        }
    }
}

impl<W> Sink for OStream<W>
where
    W: AsyncWrite,
{
    type SinkItem = ByteBuf;
    type SinkError = io::Error;

    fn start_send(&mut self, bytes: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        let buf = match self.buf.take() {
            Some(mut buf) => {
                buf.extend(bytes);
                buf
            }
            None => bytes,
        };
        self.buf = Some(buf);
        self.poll_complete()?;
        Ok(AsyncSink::Ready)
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        match self.poll_write() {
            Ok(Async::Ready(())) => {
                self.buf = None;
                Ok(Async::Ready(()))
            }
            Ok(Async::NotReady) => Ok(Async::NotReady),
            Err(e) => {
                self.w = None;
                Err(e)
            }
        }
    }

    #[inline]
    fn close(&mut self) -> Poll<(), Self::SinkError> {
        self.poll_complete()
    }
}

impl<R: AsyncRead> IntoStream for R {
    type Stream = IStream<Self>;

    #[inline]
    fn into_stream(self) -> Self::Stream {
        IStream {
            r: Some(self),
            would_block: false,
        }
    }
}

impl<W: AsyncWrite> IntoSink for W {
    type Sink = OStream<Self>;

    #[inline]
    fn into_sink(self) -> Self::Sink {
        OStream {
            w: Some(self),
            buf: None,
        }
    }
}
<commit_msg>Clean up<commit_after>use std::cell::UnsafeCell;
use std::io;

use buf::{ByteBuf, Appender};
use nio::ReadV;
use io::{AsyncRead, AsyncWrite};

const RECV_BUF_SIZE: usize = 128 * 1024;

struct RecvBuf {
    inner: UnsafeCell<ByteBuf>,
}

impl RecvBuf {
    #[inline]
    fn new() -> Self {
        RecvBuf {
            inner: UnsafeCell::new(ByteBuf::with_capacity(RECV_BUF_SIZE)),
        }
    }

    #[inline]
    fn append(v: usize, chain: &mut Appender) -> io::Result<usize> {
        if let Some(block) = chain.last_mut() {
            if block.appendable() >= v {
                return Ok(0);
            }
        }
        chain.append(v);
        Ok(0)
    }

    #[inline]
    fn get_mut(&self) -> &mut ByteBuf {
        let buf = unsafe { &mut *self.inner.get() };
        buf.append(RECV_BUF_SIZE, Self::append).unwrap();
        buf
    }
}

thread_local!(static RECV_BUF: RecvBuf = RecvBuf::new());

#[inline]
fn read<R>(r: &mut R) -> io::Result<Option<(ByteBuf, usize)>>
where
    R: io::Read + ReadV,
{
    RECV_BUF.with(|recv_buf| {
        let buf = recv_buf.get_mut();
        let n = buf.read_in(r)?;
        let data = if n > 0 {
            Some((buf.drain_to(n)?, n))
        } else {
            None
        };
        Ok(data)
    })
}

pub mod read;
pub mod write;
pub mod copy;
pub mod split;

use futures::{Stream, Sink, Poll, Async, StartSend, AsyncSink};
use sink::IntoSink;
use stream::IntoStream;

#[derive(Debug)]
pub struct IStream<R> {
    r: Option<R>,
    would_block: bool,
}

impl<R> IStream<R>
where
    R: AsyncRead,
{
    #[inline]
    pub fn get_ref(&self) -> &R {
        self.r
            .as_ref()
            .expect("Attempted IStream::get_ref after completion")
    }

    #[inline]
    pub fn get_mut(&mut self) -> &mut R {
        self.r
            .as_mut()
            .expect("Attempted IStream::get_mut after completion")
    }

    #[inline]
    pub fn into_inner(self) -> R {
        self.r
            .expect("Attempted IStream::into_inner after completion")
    }

    #[inline]
    fn poll_read(&mut self) -> Poll<Option<ByteBuf>, io::Error> {
        match self.r {
            Some(ref mut r) => {
                if self.would_block {
                    self.would_block = false;
                    r.need_read()?;
                    return Ok(Async::NotReady);
                }
                if !r.is_readable() {
                    return Ok(Async::NotReady);
                }
                match read(r) {
                    Ok(Some((data, _))) => {
                        self.would_block = true;
                        Ok(Async::Ready(Some(data)))
                    }
                    Ok(None) => Ok(Async::Ready(None)),
                    Err(e) => {
                        match e.kind() {
                            io::ErrorKind::WouldBlock => {
                                r.need_read()?;
                                Ok(Async::NotReady)
                            }
                            _ => Err(e),
                        }
                    }
                }
            }
            None => panic!("Attempted to poll IStream after completion"),
        }
    }
}

impl<R> Stream for IStream<R>
where
    R: AsyncRead,
{
    type Item = ByteBuf;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        match self.poll_read() {
            Ok(o) => Ok(o),
            Err(e) => {
                self.r = None;
                Err(e)
            }
        }
    }
}

#[derive(Debug)]
pub struct OStream<W> {
    w: Option<W>,
    buf: Option<ByteBuf>,
}

impl<W> OStream<W>
where
    W: AsyncWrite,
{
    #[inline]
    pub fn get_ref(&self) -> &W {
        self.w
            .as_ref()
            .expect("Attempted OStream::get_ref after completion")
    }

    #[inline]
    pub fn get_mut(&mut self) -> &mut W {
        self.w
            .as_mut()
            .expect("Attempted OStream::get_mut after completion")
    }

    #[inline]
    pub fn into_inner(self) -> W {
        self.w
            .expect("Attempted OStream::into_inner after completion")
    }

    #[inline]
    fn poll_write(&mut self) -> Poll<(), io::Error> {
        match self.w {
            Some(ref mut w) => {
                if w.is_writable() {
                    if let Some(ref mut data) = self.buf {
                        if let Err(e) = data.write_out(w) {
                            return match e.kind() {
                                io::ErrorKind::WouldBlock => {
                                    w.need_write()?;
                                    Ok(Async::NotReady)
                                }
                                _ => Err(e),
                            };
                        }
                        if !data.is_empty() {
                            w.need_write()?;
                            data.compact();
                            return Ok(Async::NotReady);
                        }
                    }
                    w.no_need_write()?;
                    Ok(Async::Ready(()))
                } else {
                    Ok(Async::NotReady)
                }
            }
            None => panic!("Attempted to poll OStream after completion"),
        }
    }
}

impl<W> Sink for OStream<W>
where
    W: AsyncWrite,
{
    type SinkItem = ByteBuf;
    type SinkError = io::Error;

    fn start_send(&mut self, bytes: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        let buf = match self.buf.take() {
            Some(mut buf) => {
                buf.extend(bytes);
                buf
            }
            None => bytes,
        };
        self.buf = Some(buf);
        self.poll_complete()?;
        Ok(AsyncSink::Ready)
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        match self.poll_write() {
            Ok(Async::Ready(())) => {
                self.buf = None;
                Ok(Async::Ready(()))
            }
            Ok(Async::NotReady) => Ok(Async::NotReady),
            Err(e) => {
                self.w = None;
                Err(e)
            }
        }
    }

    #[inline]
    fn close(&mut self) -> Poll<(), Self::SinkError> {
        self.poll_complete()
    }
}

impl<R: AsyncRead> IntoStream for R {
    type Stream = IStream<Self>;

    #[inline]
    fn into_stream(self) -> Self::Stream {
        IStream {
            r: Some(self),
            would_block: false,
        }
    }
}

impl<W: AsyncWrite> IntoSink for W {
    type Sink = OStream<Self>;

    #[inline]
    fn into_sink(self) -> Self::Sink {
        OStream {
            w: Some(self),
            buf: None,
        }
    }
}
<|endoftext|>


<commit_before>use std::ascii::AsciiExt;

use ::{ NifTerm, NifEnv, NifResult, NifError, NifEncoder };
use ::wrapper::nif_interface::NIF_TERM;
use ::wrapper::atom;

// Atoms are a special case of a term. They can be stored and used on all envs regardless of where
// it lives and when it is created.
#[derive(PartialEq, Eq, Clone, Copy)]
pub struct NifAtom {
    term: NIF_TERM,
}

impl NifAtom {
    pub fn as_c_arg(&self) -> NIF_TERM {
        self.term
    }

    pub fn to_term<'a>(self, env: NifEnv<'a>) -> NifTerm<'a> {
        // Safe because atoms are not associated with any environment.
        unsafe { NifTerm::new(env, self.term) }
    }

    unsafe fn from_nif_term(term: NIF_TERM) -> Self {
        NifAtom {
            term: term
        }
    }

    pub fn from_term(term: NifTerm) -> NifResult<Self> {
        match term.is_atom() {
            true => Ok(unsafe { NifAtom::from_nif_term(term.as_c_arg()) }),
            false => Err(NifError::BadArg)
        }
    }

    /// Return the atom whose text representation is `bytes`, like `erlang:binary_to_atom/2`.
    ///
    /// # Errors
    /// `NifError::BadArg` if `bytes.len() > 255`.
    pub fn from_bytes<'a>(env: NifEnv<'a>, bytes: &[u8]) -> NifResult<NifAtom> {
        if bytes.len() > 255 {
            return Err(NifError::BadArg);
        }
        unsafe {
            Ok(NifAtom::from_nif_term(atom::make_atom(env.as_c_arg(), bytes)))
        }
    }

    /// Return the atom whose text representation is the given `string`, like `erlang:list_to_atom/2`.
    ///
    /// # Errors
    /// `NifError::BadArg` if `string` contains characters that aren't in Latin-1, or if it's too
    /// long. The maximum length is 255 characters.
    pub fn from_str<'a>(env: NifEnv<'a>, string: &str) -> NifResult<NifAtom> {
        if string.is_ascii() {
            // Fast path.
            NifAtom::from_bytes(env, string.as_bytes())
        } else {
            // Convert from Rust UTF-8 to Latin-1.
            let mut bytes = Vec::with_capacity(string.len());
            for c in string.chars() {
                if (c as u32) >= 256 {
                    return Err(NifError::BadArg);
                }
                bytes.push(c as u8);
            }
            NifAtom::from_bytes(env, &bytes)
        }
    }
}

use ::std::fmt;
impl fmt::Debug for NifAtom {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        ::wrapper::term::fmt(self.as_c_arg(), f)
    }
}

impl NifEncoder for NifAtom {
    fn encode<'a>(&self, env: NifEnv<'a>) -> NifTerm<'a> {
        self.to_term(env)
    }
}

impl<'a> PartialEq<NifTerm<'a>> for NifAtom {
    fn eq(&self, other: &NifTerm<'a>) -> bool {
        self.as_c_arg() == other.as_c_arg()
    }
}

/// ## Atom terms
impl<'a> NifTerm<'a> {

    /// When the term is an atom, this method will return the string
    /// representation of it.
    ///
    /// If you only need to test for equality, comparing the terms directly
    /// is much faster.
    ///
    /// Will return None if the term is not an atom.
    pub fn atom_to_string(&self) -> NifResult<String> {
        unsafe { atom::get_atom(self.get_env().as_c_arg(), self.as_c_arg()) }
    }

}

pub fn is_truthy(term: NifTerm) -> bool {
    !((term.as_c_arg() == false_().as_c_arg()) || (term.as_c_arg() == nil().as_c_arg()))
}

// This is safe because atoms are never removed/changed once they are created.
unsafe impl Sync for NifAtom {}
unsafe impl Send for NifAtom {}

/// Macro for defining Rust functions that return Erlang atoms.
/// To use this macro, you must also import the `lazy_static` crate.
///
/// For example, this code:
///
///     #[macro_use] extern crate rustler;
///     #[macro_use] extern crate lazy_static;
///
///     mod my_atoms {
///         rustler_atoms! {
///             atom jpeg;
///         }
///     }
///     # fn main() {}
///
/// defines a public function `my_atoms::jpeg()` that returns the `NifAtom` for the `jpeg` atom.
///
/// Multiple atoms can be defined. Each one can have its own doc comment and other attributes.
///
///     # #[macro_use] extern crate rustler;
///     # #[macro_use] extern crate lazy_static;
///     rustler_atoms! {
///         /// The `jpeg` atom.
///         atom jpeg;
///
///         /// The `png` atom.
///         atom png;
///
///         #[allow(non_snake_case)]
///         atom WebP;
///     }
///     # fn main() {}
///
/// When you need an atom that's not a legal Rust function name, write `atom NAME = "ATOM"`, like
/// this:
///
///     # #[macro_use] extern crate rustler;
///     # #[macro_use] extern crate lazy_static;
///     rustler_atoms! {
///         /// The `mod` atom. The function isn't called `mod` because that's
///         /// a Rust keyword.
///         atom mod_atom = "mod";
///
///         /// The atom `'hello world'`. Obviously this function can't be
///         /// called `hello world` because there's a space in it.
///         atom hello_world = "hello world";
///     }
///     # fn main() {}
///
/// # Performance
///
/// These functions are faster than `get_atom` and `get_atom_init`. The first time you call one, it
/// creates atoms for all its sibling functions and caches them, so that all later calls are fast.
/// The only overhead is checking that the atoms have been created (an atomic integer read).
///
#[macro_export]
macro_rules! rustler_atoms {
    {
        $(
            $( #[$attr:meta] )*
            atom $name:ident $( = $str:expr )*;
        )*
    } => {
        #[allow(non_snake_case)]
        struct RustlerAtoms {
            $( $name : $crate::types::atom::NifAtom ),*
        }
        lazy_static! {
            static ref RUSTLER_ATOMS: RustlerAtoms = $crate::env::OwnedEnv::new().run(|env| {
                RustlerAtoms {
                    $( $name: rustler_atoms!(@internal_make_atom(env, $name $( = $str)* )) ),*
                }
            });
        }
        $(
            $( #[$attr] )*
            pub fn $name() -> $crate::types::atom::NifAtom {
                RUSTLER_ATOMS.$name
            }
        )*
    };

    // Internal helper macros.
    { @internal_make_atom($env:ident, $name:ident) } => {
        rustler_atoms!(@internal_make_atom($env, $name = stringify!($name)))
    };
    { @internal_make_atom($env:ident, $name:ident = $str:expr) } => {
        $crate::types::atom::NifAtom::from_str($env, $str)
            .ok().expect("rustler_atoms: bad atom string")
    };
}

rustler_atoms! {
    /// The `nil` atom.
    atom nil;

    /// The `ok` atom, commonly used in success tuples.
    atom ok;

    /// The `error` atom, commonly used in error tuples.
    atom error;

    /// The `badarg` atom, which Rustler sometimes returns to indicate that a function was
    /// called with incorrect arguments.
    atom badarg;

    /// The `false` atom. (Trailing underscore because `false` is a keyword in Rust.)
    ///
    /// If you're looking to convert between Erlang terms and Rust `bool`
    /// values, use `NifEncoder` and `NifDecoder` instead.
    atom false_ = "false";

    /// The `true` atom. (Trailing underscore because `true` is a keyword in Rust.)
    ///
    /// If you're looking to convert between Erlang terms and Rust `bool`
    /// values, use `NifEncoder` and `NifDecoder` instead.
    atom true_ = "true";

    /// The `__struct__` atom used by Elixir.
    atom __struct__;
}
<commit_msg>implement NifDecoder for NifAtom<commit_after>use std::ascii::AsciiExt;

use ::{ NifTerm, NifEnv, NifResult, NifError, NifEncoder, NifDecoder };
use ::wrapper::nif_interface::NIF_TERM;
use ::wrapper::atom;

// Atoms are a special case of a term. They can be stored and used on all envs regardless of where
// it lives and when it is created.
#[derive(PartialEq, Eq, Clone, Copy)]
pub struct NifAtom {
    term: NIF_TERM,
}

impl NifAtom {
    pub fn as_c_arg(&self) -> NIF_TERM {
        self.term
    }

    pub fn to_term<'a>(self, env: NifEnv<'a>) -> NifTerm<'a> {
        // Safe because atoms are not associated with any environment.
        unsafe { NifTerm::new(env, self.term) }
    }

    unsafe fn from_nif_term(term: NIF_TERM) -> Self {
        NifAtom {
            term: term
        }
    }

    pub fn from_term(term: NifTerm) -> NifResult<Self> {
        match term.is_atom() {
            true => Ok(unsafe { NifAtom::from_nif_term(term.as_c_arg()) }),
            false => Err(NifError::BadArg)
        }
    }

    /// Return the atom whose text representation is `bytes`, like `erlang:binary_to_atom/2`.
    ///
    /// # Errors
    /// `NifError::BadArg` if `bytes.len() > 255`.
    pub fn from_bytes<'a>(env: NifEnv<'a>, bytes: &[u8]) -> NifResult<NifAtom> {
        if bytes.len() > 255 {
            return Err(NifError::BadArg);
        }
        unsafe {
            Ok(NifAtom::from_nif_term(atom::make_atom(env.as_c_arg(), bytes)))
        }
    }

    /// Return the atom whose text representation is the given `string`, like `erlang:list_to_atom/2`.
    ///
    /// # Errors
    /// `NifError::BadArg` if `string` contains characters that aren't in Latin-1, or if it's too
    /// long. The maximum length is 255 characters.
    pub fn from_str<'a>(env: NifEnv<'a>, string: &str) -> NifResult<NifAtom> {
        if string.is_ascii() {
            // Fast path.
            NifAtom::from_bytes(env, string.as_bytes())
        } else {
            // Convert from Rust UTF-8 to Latin-1.
            let mut bytes = Vec::with_capacity(string.len());
            for c in string.chars() {
                if (c as u32) >= 256 {
                    return Err(NifError::BadArg);
                }
                bytes.push(c as u8);
            }
            NifAtom::from_bytes(env, &bytes)
        }
    }
}

use ::std::fmt;
impl fmt::Debug for NifAtom {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        ::wrapper::term::fmt(self.as_c_arg(), f)
    }
}

impl NifEncoder for NifAtom {
    fn encode<'a>(&self, env: NifEnv<'a>) -> NifTerm<'a> {
        self.to_term(env)
    }
}
impl<'a> NifDecoder<'a> for NifAtom {
    fn decode(term: NifTerm<'a>) -> NifResult<NifAtom> {
        NifAtom::from_term(term)
    }
}

impl<'a> PartialEq<NifTerm<'a>> for NifAtom {
    fn eq(&self, other: &NifTerm<'a>) -> bool {
        self.as_c_arg() == other.as_c_arg()
    }
}

/// ## Atom terms
impl<'a> NifTerm<'a> {

    /// When the term is an atom, this method will return the string
    /// representation of it.
    ///
    /// If you only need to test for equality, comparing the terms directly
    /// is much faster.
    ///
    /// Will return None if the term is not an atom.
    pub fn atom_to_string(&self) -> NifResult<String> {
        unsafe { atom::get_atom(self.get_env().as_c_arg(), self.as_c_arg()) }
    }

}

pub fn is_truthy(term: NifTerm) -> bool {
    !((term.as_c_arg() == false_().as_c_arg()) || (term.as_c_arg() == nil().as_c_arg()))
}

// This is safe because atoms are never removed/changed once they are created.
unsafe impl Sync for NifAtom {}
unsafe impl Send for NifAtom {}

/// Macro for defining Rust functions that return Erlang atoms.
/// To use this macro, you must also import the `lazy_static` crate.
///
/// For example, this code:
///
///     #[macro_use] extern crate rustler;
///     #[macro_use] extern crate lazy_static;
///
///     mod my_atoms {
///         rustler_atoms! {
///             atom jpeg;
///         }
///     }
///     # fn main() {}
///
/// defines a public function `my_atoms::jpeg()` that returns the `NifAtom` for the `jpeg` atom.
///
/// Multiple atoms can be defined. Each one can have its own doc comment and other attributes.
///
///     # #[macro_use] extern crate rustler;
///     # #[macro_use] extern crate lazy_static;
///     rustler_atoms! {
///         /// The `jpeg` atom.
///         atom jpeg;
///
///         /// The `png` atom.
///         atom png;
///
///         #[allow(non_snake_case)]
///         atom WebP;
///     }
///     # fn main() {}
///
/// When you need an atom that's not a legal Rust function name, write `atom NAME = "ATOM"`, like
/// this:
///
///     # #[macro_use] extern crate rustler;
///     # #[macro_use] extern crate lazy_static;
///     rustler_atoms! {
///         /// The `mod` atom. The function isn't called `mod` because that's
///         /// a Rust keyword.
///         atom mod_atom = "mod";
///
///         /// The atom `'hello world'`. Obviously this function can't be
///         /// called `hello world` because there's a space in it.
///         atom hello_world = "hello world";
///     }
///     # fn main() {}
///
/// # Performance
///
/// These functions are faster than `get_atom` and `get_atom_init`. The first time you call one, it
/// creates atoms for all its sibling functions and caches them, so that all later calls are fast.
/// The only overhead is checking that the atoms have been created (an atomic integer read).
///
#[macro_export]
macro_rules! rustler_atoms {
    {
        $(
            $( #[$attr:meta] )*
            atom $name:ident $( = $str:expr )*;
        )*
    } => {
        #[allow(non_snake_case)]
        struct RustlerAtoms {
            $( $name : $crate::types::atom::NifAtom ),*
        }
        lazy_static! {
            static ref RUSTLER_ATOMS: RustlerAtoms = $crate::env::OwnedEnv::new().run(|env| {
                RustlerAtoms {
                    $( $name: rustler_atoms!(@internal_make_atom(env, $name $( = $str)* )) ),*
                }
            });
        }
        $(
            $( #[$attr] )*
            pub fn $name() -> $crate::types::atom::NifAtom {
                RUSTLER_ATOMS.$name
            }
        )*
    };

    // Internal helper macros.
    { @internal_make_atom($env:ident, $name:ident) } => {
        rustler_atoms!(@internal_make_atom($env, $name = stringify!($name)))
    };
    { @internal_make_atom($env:ident, $name:ident = $str:expr) } => {
        $crate::types::atom::NifAtom::from_str($env, $str)
            .ok().expect("rustler_atoms: bad atom string")
    };
}

rustler_atoms! {
    /// The `nil` atom.
    atom nil;

    /// The `ok` atom, commonly used in success tuples.
    atom ok;

    /// The `error` atom, commonly used in error tuples.
    atom error;

    /// The `badarg` atom, which Rustler sometimes returns to indicate that a function was
    /// called with incorrect arguments.
    atom badarg;

    /// The `false` atom. (Trailing underscore because `false` is a keyword in Rust.)
    ///
    /// If you're looking to convert between Erlang terms and Rust `bool`
    /// values, use `NifEncoder` and `NifDecoder` instead.
    atom false_ = "false";

    /// The `true` atom. (Trailing underscore because `true` is a keyword in Rust.)
    ///
    /// If you're looking to convert between Erlang terms and Rust `bool`
    /// values, use `NifEncoder` and `NifDecoder` instead.
    atom true_ = "true";

    /// The `__struct__` atom used by Elixir.
    atom __struct__;
}
<|endoftext|>


<commit_before>use rand::{distributions::Uniform, rngs::OsRng, Rng};
use sha2::{Digest, Sha256};

const TOKEN_LENGTH: usize = 32;

pub(crate) struct SecureToken {
    sha256: Vec<u8>,
}

impl SecureToken {
    pub(crate) fn generate(kind: SecureTokenKind) -> NewSecureToken {
        let plaintext = format!(
            "{}{}",
            kind.prefix(),
            generate_secure_alphanumeric_string(TOKEN_LENGTH)
        );
        let sha256 = Sha256::digest(plaintext.as_bytes()).as_slice().to_vec();

        NewSecureToken {
            plaintext,
            inner: Self { sha256 },
        }
    }

    pub(crate) fn parse(kind: SecureTokenKind, plaintext: &str) -> Option<Self> {
        // This will both reject tokens without a prefix and tokens of the wrong kind.
        if SecureTokenKind::from_token(plaintext) != Some(kind) {
            return None;
        }

        let sha256 = Sha256::digest(plaintext.as_bytes()).as_slice().to_vec();
        Some(Self { sha256 })
    }

    pub(crate) fn sha256(&self) -> &[u8] {
        &self.sha256
    }
}

pub(crate) struct NewSecureToken {
    plaintext: String,
    inner: SecureToken,
}

impl NewSecureToken {
    pub(crate) fn plaintext(&self) -> &str {
        &self.plaintext
    }
}

impl std::ops::Deref for NewSecureToken {
    type Target = SecureToken;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

fn generate_secure_alphanumeric_string(len: usize) -> String {
    const CHARS: &[u8] = b"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";

    OsRng
        .sample_iter(Uniform::from(0..CHARS.len()))
        .map(|idx| CHARS[idx] as char)
        .take(len)
        .collect()
}

macro_rules! secure_token_kind {
    ($(#[$attr:meta])* $vis:vis enum $name:ident { $($key:ident => $repr:expr,)* }) => {
        $(#[$attr])*
        $vis enum $name {
            $($key,)*
        }

        impl $name {
            const VARIANTS: &'static [Self] = &[$(Self::$key,)*];

            fn prefix(&self) -> &'static str {
                match self {
                    $(Self::$key => $repr,)*
                }
            }
        }
    }
}

secure_token_kind! {
    /// Represents every kind of secure token generated by crates.io. When you need to generate a
    /// new kind of token you should also add its own kind with its own unique prefix.
    ///
    /// NEVER CHANGE THE PREFIX OF EXISTING TOKEN TYPES!!! Doing so will implicitly revoke all the
    /// tokens of that kind, distrupting production users.
    #[derive(Debug, Copy, Clone, Eq, PartialEq)]
    pub(crate) enum SecureTokenKind {
        API => "cio", // Crates.IO
    }
}

impl SecureTokenKind {
    fn from_token(token: &str) -> Option<Self> {
        Self::VARIANTS
            .iter()
            .find(|v| token.starts_with(v.prefix()))
            .copied()
    }
}

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

    #[test]
    fn test_generated_and_parse() {
        const KIND: SecureTokenKind = SecureTokenKind::API;

        let token = SecureToken::generate(KIND);
        assert!(token.plaintext().starts_with(KIND.prefix()));
        assert_eq!(
            token.sha256(),
            Sha256::digest(token.plaintext().as_bytes()).as_slice()
        );

        let parsed =
            SecureToken::parse(KIND, &token.plaintext()).expect("failed to parse back the token");
        assert_eq!(parsed.sha256(), token.sha256());
    }

    #[test]
    fn test_parse_no_kind() {
        assert!(SecureToken::parse(SecureTokenKind::API, "nokind").is_none());
    }
}
<commit_msg>util: add more tests to the token generation<commit_after>use rand::{distributions::Uniform, rngs::OsRng, Rng};
use sha2::{Digest, Sha256};

const TOKEN_LENGTH: usize = 32;

pub(crate) struct SecureToken {
    sha256: Vec<u8>,
}

impl SecureToken {
    pub(crate) fn generate(kind: SecureTokenKind) -> NewSecureToken {
        let plaintext = format!(
            "{}{}",
            kind.prefix(),
            generate_secure_alphanumeric_string(TOKEN_LENGTH)
        );
        let sha256 = Sha256::digest(plaintext.as_bytes()).as_slice().to_vec();

        NewSecureToken {
            plaintext,
            inner: Self { sha256 },
        }
    }

    pub(crate) fn parse(kind: SecureTokenKind, plaintext: &str) -> Option<Self> {
        // This will both reject tokens without a prefix and tokens of the wrong kind.
        if SecureTokenKind::from_token(plaintext) != Some(kind) {
            return None;
        }

        let sha256 = Sha256::digest(plaintext.as_bytes()).as_slice().to_vec();
        Some(Self { sha256 })
    }

    pub(crate) fn sha256(&self) -> &[u8] {
        &self.sha256
    }
}

pub(crate) struct NewSecureToken {
    plaintext: String,
    inner: SecureToken,
}

impl NewSecureToken {
    pub(crate) fn plaintext(&self) -> &str {
        &self.plaintext
    }
}

impl std::ops::Deref for NewSecureToken {
    type Target = SecureToken;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

fn generate_secure_alphanumeric_string(len: usize) -> String {
    const CHARS: &[u8] = b"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";

    OsRng
        .sample_iter(Uniform::from(0..CHARS.len()))
        .map(|idx| CHARS[idx] as char)
        .take(len)
        .collect()
}

macro_rules! secure_token_kind {
    ($(#[$attr:meta])* $vis:vis enum $name:ident { $($key:ident => $repr:expr,)* }) => {
        $(#[$attr])*
        $vis enum $name {
            $($key,)*
        }

        impl $name {
            const VARIANTS: &'static [Self] = &[$(Self::$key,)*];

            fn prefix(&self) -> &'static str {
                match self {
                    $(Self::$key => $repr,)*
                }
            }
        }
    }
}

secure_token_kind! {
    /// Represents every kind of secure token generated by crates.io. When you need to generate a
    /// new kind of token you should also add its own kind with its own unique prefix.
    ///
    /// NEVER CHANGE THE PREFIX OF EXISTING TOKEN TYPES!!! Doing so will implicitly revoke all the
    /// tokens of that kind, distrupting production users.
    #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
    pub(crate) enum SecureTokenKind {
        API => "cio", // Crates.IO
    }
}

impl SecureTokenKind {
    fn from_token(token: &str) -> Option<Self> {
        Self::VARIANTS
            .iter()
            .find(|v| token.starts_with(v.prefix()))
            .copied()
    }
}

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

    #[test]
    fn test_generated_and_parse() {
        const KIND: SecureTokenKind = SecureTokenKind::API;

        let token = SecureToken::generate(KIND);
        assert!(token.plaintext().starts_with(KIND.prefix()));
        assert_eq!(
            token.sha256(),
            Sha256::digest(token.plaintext().as_bytes()).as_slice()
        );

        let parsed =
            SecureToken::parse(KIND, &token.plaintext()).expect("failed to parse back the token");
        assert_eq!(parsed.sha256(), token.sha256());
    }

    #[test]
    fn test_parse_no_kind() {
        assert!(SecureToken::parse(SecureTokenKind::API, "nokind").is_none());
    }

    #[test]
    fn test_persistent_prefixes() {
        // Changing prefixes will implicitly revoke all the tokens of that kind, disrupting users.
        // This test serves as a reminder for future maintainers not to change the prefixes, and
        // to ensure all the variants are tested by this test.
        let mut remaining: HashSet<_> = SecureTokenKind::VARIANTS.iter().copied().collect();
        let mut ensure = |kind: SecureTokenKind, prefix| {
            assert_eq!(kind.prefix(), prefix);
            remaining.remove(&kind);
        };

        ensure(SecureTokenKind::API, "cio");

        assert!(
            remaining.is_empty(),
            "not all variants have a test to check the prefix"
        );
    }

    #[test]
    fn test_conflicting_prefixes() {
        // This sanity check prevents multiple tokens from starting with the same prefix, which
        // would mess up the token kind detection. If this test fails after adding another variant
        // do not change the test, choose another prefix instead.
        for variant in SecureTokenKind::VARIANTS {
            for other in SecureTokenKind::VARIANTS {
                if variant == other {
                    continue;
                }
                if variant.prefix().starts_with(other.prefix()) {
                    panic!("variants {:?} and {:?} share a prefix", variant, other);
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use ::serde_json;
use ::zip;
use std::collections::HashMap;

#[derive(Serialize, Debug)]
struct MultiArchiveJsonWriter<'a> {
    archives: HashMap<&'a str, &'a ZipArchiveJsonWriter<'a>>,
}

#[derive(Serialize, Debug)]
struct ZipArchiveJsonWriter<'a> {
    objects: HashMap<&'a str, &'a ZipObjectJsonWriter>,
}

#[derive(Serialize, Debug)]
struct ZipObjectJsonWriter {
    compression_type: String,
    original_size: u64,
    compressed_size: u64,
    compression_rate: String,
}

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

    #[test]
    fn test_serialize_json_writer() {
        let zip_object = ZipObjectJsonWriter {
            compression_type: format!("{}", zip::CompressionMethod::Deflated),
            original_size: 100,
            compressed_size: 50,
            compression_rate: String::from("50%"),
        };

        let zip_object_serialized = serde_json::to_string(&zip_object)
            .unwrap();

        println!("zip_object_serialized: {}", zip_object_serialized);
    }
}
<commit_msg>Add test for serializing ZipObjectJsonWriter<commit_after>use ::serde_json;
use ::zip;
use std::collections::HashMap;

#[derive(Serialize, Debug)]
struct MultiArchiveJsonWriter<'a> {
    archives: HashMap<&'a str, &'a ZipArchiveJsonWriter<'a>>,
}

#[derive(Serialize, Debug)]
struct ZipArchiveJsonWriter<'a> {
    objects: HashMap<&'a str, &'a ZipObjectJsonWriter>,
}

#[derive(Serialize, Deserialize, Debug, PartialEq)]
struct ZipObjectJsonWriter {
    compression_type: String,
    original_size: u64,
    compressed_size: u64,
    compression_rate: String,
}

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

    #[test]
    fn test_serialize_json_writer() {
        let zip_object = ZipObjectJsonWriter {
            compression_type: format!("{}", zip::CompressionMethod::Deflated),
            original_size: 100,
            compressed_size: 50,
            compression_rate: String::from("50%"),
        };

        let zip_object_serialized = serde_json::to_string(&zip_object)
            .unwrap();

        let zip_object_pretty = serde_json::
            to_string_pretty(&zip_object).unwrap();

        let zip_object_deserialized: ZipObjectJsonWriter =
            serde_json::from_str(zip_object_serialized.as_str()).unwrap();

        let zip_object_depretty: ZipObjectJsonWriter =
            serde_json::from_str(zip_object_pretty.as_str()).unwrap();

        assert_eq!(zip_object, zip_object_deserialized);
        assert_eq!(zip_object, zip_object_depretty);
    }
}
<|endoftext|>