<commit_before>use closure::Closure;

use num::bigint::{BigInt, ToBigInt};
use num::rational::{BigRational, Ratio};
use num::{range, ToPrimitive, Zero};
use std::cell::RefCell;
use std::collections::BTreeMap;
use std::rc::Rc;

#[derive(Clone, Eq, Debug, PartialEq, PartialOrd, Ord)]
pub enum Value {
    Number(BigRational),
    CharString(String),
    Boolean(bool),
    Map(Rc<RefCell<BTreeMap<Value, Value>>>),
    Closure(Rc<Box<Vec<String>>>, Rc<Closure>),
}

type BinopResult = Result<Value, String>;

impl Value {
    pub fn sub(&self, other: Value) -> BinopResult {
        match (self, other) {
            (&Value::Number(ref lratio), Value::Number(ref rratio)) => {
                Ok(Value::Number(lratio - rratio))
            }
            (&Value::CharString(ref lstr), Value::CharString(ref rstr)) => {
                if let Some(index) = lstr.rfind(rstr) {
                    Ok(Value::CharString(
                        lstr[0..index].to_owned() + &lstr[(index + rstr.len())..],
                    ))
                } else {
                    Ok(Value::CharString("".to_owned()))
                }
            }
            (&Value::Closure(_, _), Value::Closure(_, _)) => {
                Err(format!("Subtraction of closures is not supported"))
            }
            (&Value::Boolean(ref lbool), Value::Boolean(ref rbool)) => {
                Ok(Value::Boolean(lbool ^ rbool))
            }
            (&Value::Map(ref lmap), Value::Map(ref rmap)) => {
                let result = lmap
                    .borrow()
                    .clone()
                    .into_iter()
                    .filter(|&(ref key, ref value)| {
                        if let Some(rvalue) = rmap.borrow().get(key) {
                            rvalue != value
                        } else {
                            true
                        }
                    })
                    .collect();
                Ok(Value::Map(Rc::new(RefCell::new(result))))
            }
            (l, r) => Err(format!("Unsupported operation - for {:?} and {:?}", l, r)),
        }
    }

    pub fn mul(&self, right: Value) -> BinopResult {
        match (self, right) {
            (&Value::Number(ref lratio), Value::Number(ref rratio)) => {
                Ok(Value::Number(lratio * rratio))
            }
            (&Value::CharString(ref str), Value::Number(ref ratio)) => {
                let extended: String = range(BigInt::from(0), ratio.ceil().to_integer())
                    .map(|_| str.clone())
                    .collect::<Vec<String>>()
                    .join("");
                let truncation = (ratio.ceil() - ratio.clone())
                    * Ratio::from_integer((str.len() as i64).to_bigint().unwrap());
                let truncation_index =
                    extended.len().to_bigint().unwrap() - truncation.to_integer();
                Ok(Value::CharString(
                    extended[..truncation_index.to_usize().unwrap()].to_owned(),
                ))
            }
            (&Value::Boolean(ref lbool), Value::Boolean(ref rbool)) => {
                Ok(Value::Boolean(lbool & rbool))
            }
            (l, r) => Err(format!("Unsupported operation * for {:?} and {:?}", l, r)),
        }
    }

    pub fn div(&self, right: Value) -> BinopResult {
        match (self, right) {
            (&Value::Number(ref lratio), Value::Number(ref rratio)) => {
                if rratio.is_zero() {
                    Err("Can't divide by zero".to_owned())
                } else {
                    Ok(Value::Number(lratio / rratio))
                }
            }
            (&Value::CharString(_), Value::Number(ref ratio)) => {
                if ratio.is_zero() {
                    Err("Can't divide by zero".to_owned())
                } else {
                    self.mul(Value::Number(Ratio::from_integer(BigInt::from(1)) / ratio))
                }
            }

            (l, r) => Err(format!("Unsupported operation / for {:?} and {:?}", l, r)),
        }
    }

    pub fn add(&self, right: Value) -> BinopResult {
        match (self, right) {
            (&Value::Number(ref lratio), Value::Number(ref rratio)) => {
                Ok(Value::Number(lratio + rratio))
            }
            (&Value::CharString(ref lstr), Value::CharString(ref rstr)) => {
                Ok(Value::CharString((*lstr).clone() + rstr))
            }
            (&Value::Closure(_, _), Value::Closure(_, _)) => {
                Err("Addition of closures is not supported".to_owned())
            }
            (&Value::Boolean(lbool), Value::Boolean(rbool)) => Ok(Value::Boolean(lbool || rbool)),
            (&Value::Map(ref lmap), Value::Map(ref rmap)) => {
                let mut result = (*lmap.borrow()).clone();
                let mut rclone = (*rmap.borrow()).clone();
                result.append(&mut rclone);
                Ok(Value::Map(Rc::new(RefCell::new(result))))
            }
            (l, r) => Err(format!("Unsupported operation + for {:?} and {:?}", l, r)),
        }
    }
}

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

    #[test]
    fn sub() {
        // Numbers
        assert_eq!(Ok(v_number(15, 2)), v_number(8, 1).sub(v_number(1, 2)));
        assert_err!(v_number(8, 1).sub(v_string("toto")));
        // Strings
        assert_eq!(
            Ok(v_string("hello ")),
            v_string("hello world").sub(v_string("world"))
        );
        assert_eq!(
            Ok(v_string("helld")),
            v_string("hello world").sub(v_string("lo wor"))
        );
        assert_eq!(
            Ok(v_string("hello world")),
            v_string("hello world").sub(v_string(""))
        );
        assert_err!(v_string("toto").sub(v_number(1, 1)));
        // Boolean
        assert_eq!(Ok(v_bool(false)), v_bool(true).sub(v_bool(true)));
        assert_eq!(Ok(v_bool(true)), v_bool(true).sub(v_bool(false)));
        assert_eq!(Ok(v_bool(true)), v_bool(false).sub(v_bool(true)));
        assert_eq!(Ok(v_bool(false)), v_bool(false).sub(v_bool(false)));
        assert_err!(v_bool(false).sub(v_number(1, 1)));
        // Map
        assert_eq!(
            Ok(v_map(vec![(v_number(1, 1), v_number(1, 1))])),
            v_map(vec![
                (v_number(1, 1), v_number(1, 1)),
                (v_number(2, 1), v_number(2, 1)),
            ]).sub(v_map(vec![(v_number(2, 1), v_number(2, 1))]))
        );
        // Map
        assert_eq!(
            Ok(v_map(vec![
                (v_number(1, 1), v_number(1, 1)),
                (v_number(2, 1), v_number(2, 1)),
            ])),
            v_map(vec![
                (v_number(1, 1), v_number(1, 1)),
                (v_number(2, 1), v_number(2, 1)),
            ]).sub(v_map(vec![(v_number(2, 1), v_number(2, 2))]))
        );
        assert_err!(v_map(vec![(v_number(1, 1), v_number(1, 1))]).sub(v_number(1, 1)));
    }

    #[test]
    fn mul() {
        // Numbers
        assert_eq!(Ok(v_number(4, 1)), v_number(8, 1).mul(v_number(1, 2)));
        assert_err!(v_number(8, 1).mul(v_string("toto")));
        // Strings
        assert_err!(v_string("hello ").mul(v_string("world")));
        assert_eq!(Ok(v_string("totototo")), v_string("to").mul(v_number(4, 1)));
        assert_eq!(Ok(v_string("tototot")), v_string("to").mul(v_number(7, 2)));
        assert_eq!(Ok(v_string("tot")), v_string("toto").mul(v_number(3, 4)));
        assert_eq!(Ok(v_string("")), v_string("toto").mul(v_number(0, 1)));
        // Boolean
        assert_eq!(Ok(v_bool(true)), v_bool(true).mul(v_bool(true)));
        assert_eq!(Ok(v_bool(false)), v_bool(true).mul(v_bool(false)));
        assert_eq!(Ok(v_bool(false)), v_bool(false).mul(v_bool(true)));
        assert_eq!(Ok(v_bool(false)), v_bool(false).mul(v_bool(false)));
        assert_err!(v_bool(false).mul(v_number(1, 1)));
        // Map can't be multiplied?
        assert_err!(v_map(vec![]).mul(v_map(vec![])));
    }

    #[test]
    fn div() {
        // Numbers
        assert_eq!(Ok(v_number(16, 1)), v_number(8, 1).div(v_number(1, 2)));
        assert_err!(v_number(8, 1).div(v_number(0, 1)));
        assert_err!(v_number(8, 1).div(v_string("toto")));
        // Strings
        assert_err!(v_string("hello ").div(v_string("world")));
        assert_eq!(Ok(v_string("t")), v_string("to").div(v_number(2, 1)));
        assert_eq!(Ok(v_string("tototot")), v_string("to").div(v_number(2, 7)));
        assert_eq!(Ok(v_string("tot")), v_string("toto").div(v_number(4, 3)));
        assert_err!(v_string("toto").div(v_number(0, 1)));
        // Boolean can't be divided
        assert_err!(v_bool(true).div(v_bool(true)));
        // Map can't be div?
        assert_err!(v_map(vec![]).div(v_map(vec![])));
    }

    #[test]
    fn add() {
        // Numbers
        assert_eq!(Ok(v_number(17, 2)), v_number(8, 1).add(v_number(1, 2)));
        assert_err!(v_number(8, 1).add(v_string("toto")));
        // Strings
        assert_eq!(
            Ok(v_string("hello world")),
            v_string("hello ").add(v_string("world"))
        );
        assert_err!(v_string("toto").add(v_number(1, 1)));
        // Boolean
        assert_eq!(Ok(v_bool(true)), v_bool(true).add(v_bool(true)));
        assert_eq!(Ok(v_bool(true)), v_bool(true).add(v_bool(false)));
        assert_eq!(Ok(v_bool(true)), v_bool(false).add(v_bool(true)));
        assert_eq!(Ok(v_bool(false)), v_bool(false).add(v_bool(false)));
        assert_err!(v_bool(false).add(v_number(1, 1)));
        // Map
        assert_eq!(
            Ok(v_map(vec![
                (v_number(1, 1), v_number(1, 1)),
                (v_number(2, 1), v_number(2, 1)),
            ])),
            v_map(vec![(v_number(1, 1), v_number(1, 1))])
                .add(v_map(vec![(v_number(2, 1), v_number(2, 1))]))
        );
        assert_err!(v_map(vec![(v_number(1, 1), v_number(1, 1))]).add(v_number(1, 1)));
    }
}
<commit_msg>Enhance multiplication<commit_after>use closure::Closure;

use num::bigint::{BigInt, ToBigInt};
use num::rational::{BigRational, Ratio};
use num::{range, One, ToPrimitive, Zero};
use std::cell::RefCell;
use std::collections::BTreeMap;
use std::rc::Rc;

#[derive(Clone, Eq, Debug, PartialEq, PartialOrd, Ord)]
pub enum Value {
    Number(BigRational),
    CharString(String),
    Boolean(bool),
    Map(Rc<RefCell<BTreeMap<Value, Value>>>),
    Closure(Rc<Box<Vec<String>>>, Rc<Closure>),
}

type BinopResult = Result<Value, String>;

impl Value {
    pub fn sub(&self, other: Value) -> BinopResult {
        match (self, other) {
            (&Value::Number(ref lratio), Value::Number(ref rratio)) => {
                Ok(Value::Number(lratio - rratio))
            }
            (&Value::CharString(ref lstr), Value::CharString(ref rstr)) => {
                if let Some(index) = lstr.rfind(rstr) {
                    Ok(Value::CharString(
                        lstr[0..index].to_owned() + &lstr[(index + rstr.len())..],
                    ))
                } else {
                    Ok(Value::CharString("".to_owned()))
                }
            }
            (&Value::Closure(_, _), Value::Closure(_, _)) => {
                Err(format!("Subtraction of closures is not supported"))
            }
            (&Value::Boolean(ref lbool), Value::Boolean(ref rbool)) => {
                Ok(Value::Boolean(lbool ^ rbool))
            }
            (&Value::Map(ref lmap), Value::Map(ref rmap)) => {
                let result = lmap
                    .borrow()
                    .clone()
                    .into_iter()
                    .filter(|&(ref key, ref value)| {
                        if let Some(rvalue) = rmap.borrow().get(key) {
                            rvalue != value
                        } else {
                            true
                        }
                    })
                    .collect();
                Ok(Value::Map(Rc::new(RefCell::new(result))))
            }
            (l, r) => Err(format!("Unsupported operation - for {:?} and {:?}", l, r)),
        }
    }

    pub fn mul(&self, right: Value) -> BinopResult {
        match (self, right) {
            (&Value::Number(ref lratio), Value::Number(ref rratio)) => {
                Ok(Value::Number(lratio * rratio))
            }
            (&Value::CharString(ref str), Value::Number(ref ratio))
            | (&Value::Number(ref ratio), Value::CharString(ref str)) => {
                let extended: String = range(BigInt::from(0), ratio.ceil().to_integer())
                    .map(|_| str.clone())
                    .collect::<Vec<String>>()
                    .join("");
                let truncation = (ratio.ceil() - ratio.clone())
                    * Ratio::from_integer((str.len() as i64).to_bigint().unwrap());
                let truncation_index =
                    extended.len().to_bigint().unwrap() - truncation.to_integer();
                Ok(Value::CharString(
                    extended[..truncation_index.to_usize().unwrap()].to_owned(),
                ))
            }
            (&Value::Map(ref map), Value::Number(ref ratio))
            | (&Value::Number(ref ratio), Value::Map(ref map)) => {
                let result = map
                    .borrow()
                    .clone()
                    .into_iter()
                    .map(|(ref key, ref value)| {
                        (
                            key.mul(Value::Number(ratio.clone()))
                                .expect("everything is multipliable by a number"),
                            value
                                .mul(Value::Number(ratio.clone()))
                                .expect("everything is multipliable by a number"),
                        )
                    })
                    .collect();

                Ok(Value::Map(Rc::new(RefCell::new(result))))
            }
            (&Value::Boolean(ref boolean), Value::Number(ref ratio))
            | (&Value::Number(ref ratio), Value::Boolean(ref boolean)) => match boolean {
                &true => Ok(Value::Number(ratio.clone())),
                &false => Ok(Value::Number(Ratio::new(BigInt::zero(), BigInt::one()))),
            },
            (&Value::CharString(ref str), Value::Boolean(ref boolean))
            | (&Value::Boolean(ref boolean), Value::CharString(ref str)) => match boolean {
                &true => Ok(Value::CharString(str.clone())),
                &false => Ok(Value::CharString("".to_owned())),
            },
            (&Value::Boolean(ref lbool), Value::Boolean(ref rbool)) => {
                Ok(Value::Boolean(lbool & rbool))
            }
            (&Value::Boolean(ref boolean), Value::Map(ref map))
            | (&Value::Map(ref map), Value::Boolean(ref boolean)) => {
                let result = map
                    .borrow()
                    .clone()
                    .into_iter()
                    .map(|(ref key, ref value)| {
                        (
                            key.mul(Value::Boolean(*boolean))
                                .expect("everything is multipliable by a boolean"),
                            value
                                .mul(Value::Boolean(*boolean))
                                .expect("everything is multipliable by a boolean"),
                        )
                    })
                    .collect();

                Ok(Value::Map(Rc::new(RefCell::new(result))))
            }
            (l, r) => Err(format!("Unsupported operation * for {:?} and {:?}", l, r)),
        }
    }

    pub fn div(&self, right: Value) -> BinopResult {
        match (self, right) {
            (&Value::Number(ref lratio), Value::Number(ref rratio)) => {
                if rratio.is_zero() {
                    Err("Can't divide by zero".to_owned())
                } else {
                    Ok(Value::Number(lratio / rratio))
                }
            }
            (&Value::CharString(_), Value::Number(ref ratio)) => {
                if ratio.is_zero() {
                    Err("Can't divide by zero".to_owned())
                } else {
                    self.mul(Value::Number(Ratio::from_integer(BigInt::from(1)) / ratio))
                }
            }

            (l, r) => Err(format!("Unsupported operation / for {:?} and {:?}", l, r)),
        }
    }

    pub fn add(&self, right: Value) -> BinopResult {
        match (self, right) {
            (&Value::Number(ref lratio), Value::Number(ref rratio)) => {
                Ok(Value::Number(lratio + rratio))
            }
            (&Value::CharString(ref lstr), Value::CharString(ref rstr)) => {
                Ok(Value::CharString((*lstr).clone() + rstr))
            }
            (&Value::Closure(_, _), Value::Closure(_, _)) => {
                Err("Addition of closures is not supported".to_owned())
            }
            (&Value::Boolean(lbool), Value::Boolean(rbool)) => Ok(Value::Boolean(lbool || rbool)),
            (&Value::Map(ref lmap), Value::Map(ref rmap)) => {
                let mut result = (*lmap.borrow()).clone();
                let mut rclone = (*rmap.borrow()).clone();
                result.append(&mut rclone);
                Ok(Value::Map(Rc::new(RefCell::new(result))))
            }
            (l, r) => Err(format!("Unsupported operation + for {:?} and {:?}", l, r)),
        }
    }
}

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

    #[test]
    fn sub() {
        // Numbers
        assert_eq!(Ok(v_number(15, 2)), v_number(8, 1).sub(v_number(1, 2)));
        assert_err!(v_number(8, 1).sub(v_string("toto")));
        // Strings
        assert_eq!(
            Ok(v_string("hello ")),
            v_string("hello world").sub(v_string("world"))
        );
        assert_eq!(
            Ok(v_string("helld")),
            v_string("hello world").sub(v_string("lo wor"))
        );
        assert_eq!(
            Ok(v_string("hello world")),
            v_string("hello world").sub(v_string(""))
        );
        assert_err!(v_string("toto").sub(v_number(1, 1)));
        // Boolean
        assert_eq!(Ok(v_bool(false)), v_bool(true).sub(v_bool(true)));
        assert_eq!(Ok(v_bool(true)), v_bool(true).sub(v_bool(false)));
        assert_eq!(Ok(v_bool(true)), v_bool(false).sub(v_bool(true)));
        assert_eq!(Ok(v_bool(false)), v_bool(false).sub(v_bool(false)));
        assert_err!(v_bool(false).sub(v_number(1, 1)));
        // Map
        assert_eq!(
            Ok(v_map(vec![(v_number(1, 1), v_number(1, 1))])),
            v_map(vec![
                (v_number(1, 1), v_number(1, 1)),
                (v_number(2, 1), v_number(2, 1)),
            ]).sub(v_map(vec![(v_number(2, 1), v_number(2, 1))]))
        );
        // Map
        assert_eq!(
            Ok(v_map(vec![
                (v_number(1, 1), v_number(1, 1)),
                (v_number(2, 1), v_number(2, 1)),
            ])),
            v_map(vec![
                (v_number(1, 1), v_number(1, 1)),
                (v_number(2, 1), v_number(2, 1)),
            ]).sub(v_map(vec![(v_number(2, 1), v_number(2, 2))]))
        );
        assert_err!(v_map(vec![(v_number(1, 1), v_number(1, 1))]).sub(v_number(1, 1)));
    }

    #[test]
    fn mul() {
        // Numbers * Number
        assert_eq!(Ok(v_number(4, 1)), v_number(8, 1).mul(v_number(1, 2)));
        // Number * CharString
        assert_eq!(Ok(v_string("totototo")), v_number(4, 1).mul(v_string("to")));
        assert_eq!(Ok(v_string("tototot")), v_number(7, 2).mul(v_string("to")));
        assert_eq!(Ok(v_string("tot")), v_number(3, 4).mul(v_string("toto")));
        assert_eq!(Ok(v_string("")), v_number(0, 1).mul(v_string("toto")));
        // CharString * Number
        assert_eq!(Ok(v_string("totototo")), v_string("to").mul(v_number(4, 1)));
        assert_eq!(Ok(v_string("tototot")), v_string("to").mul(v_number(7, 2)));
        assert_eq!(Ok(v_string("tot")), v_string("toto").mul(v_number(3, 4)));
        assert_eq!(Ok(v_string("")), v_string("toto").mul(v_number(0, 1)));
        // Number * Boolean
        assert_eq!(Ok(v_number(0, 1)), v_number(1, 1).mul(v_bool(false)));
        assert_eq!(Ok(v_number(1, 1)), v_number(1, 1).mul(v_bool(true)));
        // Boolean * Number
        assert_eq!(Ok(v_number(0, 1)), v_bool(false).mul(v_number(1, 1)));
        assert_eq!(Ok(v_number(1, 1)), v_bool(true).mul(v_number(1, 1)));
        // Number * Map
        assert_eq!(
            Ok(v_map(vec![(v_string("tototo"), v_number(3, 1))])),
            v_number(3, 1).mul(v_map(vec![(v_string("to"), v_number(1, 1))]))
        );
        // Map * Number
        assert_eq!(
            Ok(v_map(vec![(v_string("tototo"), v_number(3, 1))])),
            v_map(vec![(v_string("to"), v_number(1, 1))]).mul(v_number(3, 1))
        );
        // TODO: Number * Closure
        // TODO: Closure * Number

        // CharString * CharString
        assert_err!(v_string("hello ").mul(v_string("world")));
        // CharString * Boolean
        assert_eq!(Ok(v_string("")), v_string("toto").mul(v_bool(false)));
        assert_eq!(Ok(v_string("toto")), v_string("toto").mul(v_bool(true)));
        // Boolean * CharString
        assert_eq!(Ok(v_string("")), v_bool(false).mul(v_string("toto")));
        assert_eq!(Ok(v_string("toto")), v_bool(true).mul(v_string("toto")));
        // TODO: CharString * Map
        // TODO: Map * CharString

        // Boolean * Boolean
        assert_eq!(Ok(v_bool(true)), v_bool(true).mul(v_bool(true)));
        assert_eq!(Ok(v_bool(false)), v_bool(true).mul(v_bool(false)));
        assert_eq!(Ok(v_bool(false)), v_bool(false).mul(v_bool(true)));
        assert_eq!(Ok(v_bool(false)), v_bool(false).mul(v_bool(false)));

        // Boolean * Map
        assert_eq!(
            Ok(v_map(vec![(v_string("toto"), v_number(1, 1))])),
            v_bool(true).mul(v_map(vec![(v_string("toto"), v_number(1, 1))]))
        );
        assert_eq!(
            Ok(v_map(vec![(v_string(""), v_number(0, 1))])),
            v_bool(false).mul(v_map(vec![(v_string("toto"), v_number(1, 1))]))
        );
        // Map * Boolean
        assert_eq!(
            Ok(v_map(vec![(v_string("toto"), v_number(1, 1))])),
            v_map(vec![(v_string("toto"), v_number(1, 1))]).mul(v_bool(true))
        );
        assert_eq!(
            Ok(v_map(vec![(v_string(""), v_number(0, 1))])),
            v_map(vec![(v_string("toto"), v_number(1, 1))]).mul(v_bool(false))
        );
        // TODO: Boolean * Closure
        // TODO: Closure * Boolean

        // Map * Map
    }

    #[test]
    fn div() {
        // Number / Number
        assert_eq!(Ok(v_number(16, 1)), v_number(8, 1).div(v_number(1, 2)));

        // Number / CharString
        assert_err!(v_number(8, 1).div(v_number(0, 1)));
        assert_err!(v_number(8, 1).div(v_string("toto")));
        // TODO: Number / Boolean
        // TODO: Number / Map

        // CharString / Number
        assert_eq!(Ok(v_string("t")), v_string("to").div(v_number(2, 1)));
        assert_eq!(Ok(v_string("tototot")), v_string("to").div(v_number(2, 7)));
        assert_eq!(Ok(v_string("tot")), v_string("toto").div(v_number(4, 3)));
        assert_err!(v_string("toto").div(v_number(0, 1)));
        // CharString / CharString
        assert_err!(v_string("hello ").div(v_string("world")));
        // TODO: CharString / Boolean
        // TODO: CharString / Map

        // TODO: Boolean / Number
        // TODO: Boolean / CharStrings
        // Boolean / Boolean
        assert_err!(v_bool(true).div(v_bool(true)));
        // Map / Map
        assert_err!(v_map(vec![]).div(v_map(vec![])));
    }

    #[test]
    fn add() {
        // Numbers
        assert_eq!(Ok(v_number(17, 2)), v_number(8, 1).add(v_number(1, 2)));
        assert_err!(v_number(8, 1).add(v_string("toto")));
        // Strings
        assert_eq!(
            Ok(v_string("hello world")),
            v_string("hello ").add(v_string("world"))
        );
        assert_err!(v_string("toto").add(v_number(1, 1)));
        // Boolean
        assert_eq!(Ok(v_bool(true)), v_bool(true).add(v_bool(true)));
        assert_eq!(Ok(v_bool(true)), v_bool(true).add(v_bool(false)));
        assert_eq!(Ok(v_bool(true)), v_bool(false).add(v_bool(true)));
        assert_eq!(Ok(v_bool(false)), v_bool(false).add(v_bool(false)));
        assert_err!(v_bool(false).add(v_number(1, 1)));
        // Map
        assert_eq!(
            Ok(v_map(vec![
                (v_number(1, 1), v_number(1, 1)),
                (v_number(2, 1), v_number(2, 1)),
            ])),
            v_map(vec![(v_number(1, 1), v_number(1, 1))])
                .add(v_map(vec![(v_number(2, 1), v_number(2, 1))]))
        );
        assert_err!(v_map(vec![(v_number(1, 1), v_number(1, 1))]).add(v_number(1, 1)));
    }
}
<|endoftext|>


<commit_before>
use ::std;
use ::std::ffi::{OsString, OsStr};
use ::std::os::windows::ffi::{OsStringExt, OsStrExt};

use ::WCString;
use ::NoNulError;
use ::error;

/// Representation of a borrowed Win32 style "wide" string.
#[derive(PartialEq, PartialOrd, Eq, Ord, Hash)]
pub struct WCStr {
    inner: [u16]
}

impl WCStr {
    /// Create a ```&WCStr``` from a raw pointer and a length.
    ///
    /// This function is unsafe for the reasons mentioned below.
    ///
    /// This function assumes that the pointer passed in has these properties:
    ///
    /// * It is not null.
    /// * It is a valid pointer.
    /// * It points to an array of ```u16```'s that does not contain any ```nul``` values.
    /// * It points to an array of ```u16```'s that is terminated with a ```nul``` at exactly the offset "```len```".
    ///
    /// This function will assert/panic when ```nul``` is not found at offset "```len```".
    ///
    /// The lifetime of the ```&WCStr``` returned from this function is not guranteed to be correct and
    /// it is up to the caller to determine the appropriate lifetime.
    ///
    /// # ```from_raw_parts()``` example
    ///
    ///     use wcstr::WCStr;
    ///     static a : &'static [u16] = &[116u16, 101u16, 115u16, 116u16, 0];
    ///     let s = unsafe { WCStr::from_raw_parts(a.as_ptr(), a.len() - 1) };
    ///     assert!(s.len() == (a.len() - 1));
    pub unsafe fn from_raw_parts<'a>(ptr: *const u16, len: usize) -> &'a WCStr {
        assert!(*ptr.offset(len as isize) == 0u16);
        std::mem::transmute(std::slice::from_raw_parts(ptr, len + 1))
    }

    /// Create a ```&WCStr``` from a slice of ```u16```'s.
    /// This function will scan the slice for ```nul``` and assume that ```nul``` terminates the string.
    /// If no ```nul``` is found in the slice, it will return ```Err(NoNulError(None))```
    /// # ```frm_slice_with_nul()``` example
    ///
    ///     use wcstr::WCStr;
    ///     static a : &'static [u16] = &[116u16, 101u16, 115u16, 116u16, 0];
    ///     let s = WCStr::from_slice_with_nul(a).unwrap();
    ///     assert!(s.len() == (a.len() - 1));
    pub fn from_slice_with_nul<'a>(slice: &'a [u16]) -> Result<&'a WCStr, NoNulError> {
        match slice.iter().position(|x| *x == 0) {
            None => Err(error::no_nul(None)),
            Some(i) => Ok(unsafe { std::mem::transmute(&slice[..i + 1]) }),
        }
    }

    /// length of the string in u16 units
    pub fn len(&self) -> usize {
        self.inner.len() - 1
    }

    /// Return a raw pointer to this "wide" string.
    ///
    ///  * The pointer remains valid only as long as this string is valid.
    ///  * The pointer points to a contiguous region of memory terminated with ```nul```.
    pub fn as_ptr(&self) -> *const u16 {
        self.inner.as_ptr()
    }

    /// Return this "wide" string as a slice of ```u16```s without a ```nul``` terminator.
    pub fn to_slice(&self) -> &[u16] {
        &self.inner[..self.len()]
    }

    /// Return this "wide" string as a slice of ```u16```s with a ```nul``` terminator.
    pub fn to_slice_with_nul(&self) -> &[u16] {
        &self.inner
    }

    /// Convert this "wide" string to a ```String``` by using ```String::from_utf16```
    pub fn to_string(&self) -> Result<String, std::string::FromUtf16Error> {
        String::from_utf16(self.to_slice())
    }

    /// Convert this "wide" string to a ```String``` by using ```String::from_utf16_lossy```
    pub fn to_string_lossy(&self) -> String {
        String::from_utf16_lossy(self.to_slice())
    }

    /// Convert this "wide" string to an ```OsString``` by using ```OsString::from_wide```
    pub fn to_os_string(&self) -> OsString {
        OsString::from_wide(self.to_slice())
    }

    /// starts with a string
    pub fn starts_with<T>(&self, s: T) -> bool
        where T: AsRef<WCStr> {
        let s = s.as_ref();
        let len = self.len();
        if s.len() > len {
            return false;
        }

        self.to_slice().iter().zip(s.to_slice().iter()).all(|(&a, &b)| a == b)
    }

    /// starts with a string
    pub fn starts_with_str<T>(&self, s: T) -> bool
        where T: AsRef<OsStr> {
        let s = s.as_ref();
        let mut s_iter = s.encode_wide();
        let mut t_iter = self.to_slice().iter();

        while let Some(a) = s_iter.next() {
            match t_iter.next() {
                Some(&b) if a == b => (),
                _ => return false,
            }
        }

        true
    }
}

impl std::fmt::Debug for WCStr {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        try!(write!(f, "\""));
        for &w in self.to_slice().iter() {
            if w < 0xD800 || w >= 0xE000 {
                for c in std::char::from_u32(w as u32).unwrap().escape_default() {
                    use std::fmt::Write;
                    try!(f.write_char(c));
                }
            }
            else {
                try!(write!(f, "\\u{{{:X}}}", w));
            }
        }
        write!(f, "\"")
    }
}

impl AsRef<WCStr> for WCStr {
    fn as_ref(&self) -> &WCStr {
        self
    }
}

impl AsRef<[u16]> for WCStr {
    fn as_ref(&self) -> &[u16] {
        &self.inner[..self.len()]
    }
}

impl ToOwned for WCStr {
    type Owned = WCString;
    fn to_owned(&self) -> WCString {
        unsafe {
            WCString::from_vec_with_nul_unchecked(self.inner.to_owned())
        }
    }
}

<commit_msg>add examples<commit_after>
use ::std;
use ::std::ffi::{OsString, OsStr};
use ::std::os::windows::ffi::{OsStringExt, OsStrExt};

use ::WCString;
use ::NoNulError;
use ::error;

/// Representation of a borrowed Win32 style "wide" string.
#[derive(PartialEq, PartialOrd, Eq, Ord, Hash)]
pub struct WCStr {
    inner: [u16]
}

impl WCStr {
    /// Create a ```&WCStr``` from a raw pointer and a length.
    ///
    /// This function is unsafe for the reasons mentioned below.
    ///
    /// This function assumes that the pointer passed in has these properties:
    ///
    /// * It is not null.
    /// * It is a valid pointer.
    /// * It points to an array of ```u16```'s that does not contain any ```nul``` values.
    /// * It points to an array of ```u16```'s that is terminated with a ```nul``` at exactly the offset "```len```".
    ///
    /// This function will assert/panic when ```nul``` is not found at offset "```len```".
    ///
    /// The lifetime of the ```&WCStr``` returned from this function is not guranteed to be correct and
    /// it is up to the caller to determine the appropriate lifetime.
    ///
    /// # ```from_raw_parts()``` example
    ///
    ///     use wcstr::WCStr;
    ///     static a : &'static [u16] = &[116u16, 101u16, 115u16, 116u16, 0];
    ///     let s = unsafe { WCStr::from_raw_parts(a.as_ptr(), a.len() - 1) };
    ///     assert!(s.len() == (a.len() - 1));
    pub unsafe fn from_raw_parts<'a>(ptr: *const u16, len: usize) -> &'a WCStr {
        assert!(*ptr.offset(len as isize) == 0u16);
        std::mem::transmute(std::slice::from_raw_parts(ptr, len + 1))
    }

    /// Create a ```&WCStr``` from a slice of ```u16```'s.
    /// This function will scan the slice for ```nul``` and assume that ```nul``` terminates the string.
    /// If no ```nul``` is found in the slice, it will return ```Err(NoNulError(None))```
    /// # ```frm_slice_with_nul()``` example
    ///
    ///     use wcstr::WCStr;
    ///     static a : &'static [u16] = &[116u16, 101u16, 115u16, 116u16, 0];
    ///     let s = WCStr::from_slice_with_nul(a).unwrap();
    ///     assert!(s.len() == (a.len() - 1));
    pub fn from_slice_with_nul<'a>(slice: &'a [u16]) -> Result<&'a WCStr, NoNulError> {
        match slice.iter().position(|x| *x == 0) {
            None => Err(error::no_nul(None)),
            Some(i) => Ok(unsafe { std::mem::transmute(&slice[..i + 1]) }),
        }
    }

    /// length of the string in u16 units
    pub fn len(&self) -> usize {
        self.inner.len() - 1
    }

    /// Return a raw pointer to this "wide" string.
    ///
    ///  * The pointer remains valid only as long as this string is valid.
    ///  * The pointer points to a contiguous region of memory terminated with ```nul```.
    pub fn as_ptr(&self) -> *const u16 {
        self.inner.as_ptr()
    }

    /// Return this "wide" string as a slice of ```u16```s without a ```nul``` terminator.
    pub fn to_slice(&self) -> &[u16] {
        &self.inner[..self.len()]
    }

    /// Return this "wide" string as a slice of ```u16```s with a ```nul``` terminator.
    pub fn to_slice_with_nul(&self) -> &[u16] {
        &self.inner
    }

    /// Convert this "wide" string to a ```String``` by using ```String::from_utf16```
    pub fn to_string(&self) -> Result<String, std::string::FromUtf16Error> {
        String::from_utf16(self.to_slice())
    }

    /// Convert this "wide" string to a ```String``` by using ```String::from_utf16_lossy```
    pub fn to_string_lossy(&self) -> String {
        String::from_utf16_lossy(self.to_slice())
    }

    /// Convert this "wide" string to an ```OsString``` by using ```OsString::from_wide```
    pub fn to_os_string(&self) -> OsString {
        OsString::from_wide(self.to_slice())
    }

    /// starts with a string.
    ///
    /// # ```starts_with()``` example
    ///
    ///     use wcstr::{WCStr, WCString};
    ///     let s = WCString::from_str("abcefg").unwrap();
    ///     let t = WCString::from_str("abc").unwrap();
    ///     let u = WCString::from_str("efg").unwrap();
    ///     let v = WCString::from_str("abcefgh").unwrap();
    ///     assert!(s.starts_with(t));
    ///     assert!(!s.starts_with(u));
    ///     assert!(!s.starts_with(v));
    pub fn starts_with<T>(&self, s: T) -> bool
        where T: AsRef<WCStr> {
        let s = s.as_ref();
        let len = self.len();
        if s.len() > len {
            return false;
        }

        self.to_slice().iter().zip(s.to_slice().iter()).all(|(&a, &b)| a == b)
    }

    /// starts with a string.
    ///
    /// # ```starts_with()``` example
    ///
    ///     use wcstr::{WCStr, WCString};
    ///     let s = WCString::from_str("abcefg").unwrap();
    ///     assert!(s.starts_with_str("abc"));
    ///     assert!(!s.starts_with_str("efg"));
    ///     assert!(!s.starts_with_str("abcefgh"));
    pub fn starts_with_str<T>(&self, s: T) -> bool
        where T: AsRef<OsStr> {
        let s = s.as_ref();
        let mut s_iter = s.encode_wide();
        let mut t_iter = self.to_slice().iter();

        while let Some(a) = s_iter.next() {
            match t_iter.next() {
                Some(&b) if a == b => (),
                _ => return false,
            }
        }

        true
    }
}

impl std::fmt::Debug for WCStr {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        try!(write!(f, "\""));
        for &w in self.to_slice().iter() {
            if w < 0xD800 || w >= 0xE000 {
                for c in std::char::from_u32(w as u32).unwrap().escape_default() {
                    use std::fmt::Write;
                    try!(f.write_char(c));
                }
            }
            else {
                try!(write!(f, "\\u{{{:X}}}", w));
            }
        }
        write!(f, "\"")
    }
}

impl AsRef<WCStr> for WCStr {
    fn as_ref(&self) -> &WCStr {
        self
    }
}

impl AsRef<[u16]> for WCStr {
    fn as_ref(&self) -> &[u16] {
        &self.inner[..self.len()]
    }
}

impl ToOwned for WCStr {
    type Owned = WCString;
    fn to_owned(&self) -> WCString {
        unsafe {
            WCString::from_vec_with_nul_unchecked(self.inner.to_owned())
        }
    }
}

<|endoftext|>


<commit_before>//! Structs for creating a new zip archive

use compression::CompressionMethod;
use types::ZipFileData;
use spec;
use crc32;
use result::{ZipResult, ZipError};
use std::default::Default;
use std::io;
use std::io::prelude::*;
use std::mem;
use std::ascii::AsciiExt;
use time;
use flate2;
use flate2::FlateWriteExt;
use flate2::write::DeflateEncoder;
use bzip2;
use bzip2::writer::BzCompressor;
use util;
use podio::{WritePodExt, LittleEndian};

enum GenericZipWriter<W: Write + io::Seek>
{
    Closed,
    Storer(W),
    Deflater(DeflateEncoder<W>),
    Bzip2(BzCompressor<W>),
}

/// Generator for ZIP files.
///
/// ```
/// fn doit() -> zip::result::ZipResult<()>
/// {
///     use std::io::Write;
///
///     // For this example we write to a buffer, but normally you should use a File
///     let mut buf: &mut [u8] = &mut [0u8; 65536];
///     let mut w = std::io::Cursor::new(buf);
///     let mut zip = zip::ZipWriter::new(w);
///
///     try!(zip.start_file("hello_world.txt", zip::CompressionMethod::Stored));
///     try!(zip.write(b"Hello, World!"));
///
///     // Optionally finish the zip. (this is also done on drop)
///     try!(zip.finish());
///
///     Ok(())
/// }
///
/// println!("Result: {:?}", doit());
/// ```
pub struct ZipWriter<W: Write + io::Seek>
{
    inner: GenericZipWriter<W>,
    files: Vec<ZipFileData>,
    stats: ZipWriterStats,
}

#[derive(Default)]
struct ZipWriterStats
{
    crc32: u32,
    start: u64,
    bytes_written: u64,
}

impl<W: Write+io::Seek> Write for ZipWriter<W>
{
    fn write(&mut self, buf: &[u8]) -> io::Result<usize>
    {
        if self.files.len() == 0 { return Err(io::Error::new(io::ErrorKind::Other, "No file has been started", None)) }
        self.stats.update(buf);
        match self.inner
        {
            GenericZipWriter::Storer(ref mut w) => w.write(buf),
            GenericZipWriter::Deflater(ref mut w) => w.write(buf),
            GenericZipWriter::Bzip2(ref mut w) => w.write(buf),
            GenericZipWriter::Closed => Err(io::Error::new(io::ErrorKind::BrokenPipe, "ZipWriter was already closed", None)),
        }
    }

    fn flush(&mut self) -> io::Result<()>
    {
        match self.inner
        {
            GenericZipWriter::Storer(ref mut w) => w.flush(),
            GenericZipWriter::Deflater(ref mut w) => w.flush(),
            GenericZipWriter::Bzip2(ref mut w) => w.flush(),
            GenericZipWriter::Closed => Err(io::Error::new(io::ErrorKind::BrokenPipe, "ZipWriter was already closed", None)),
        }
    }
}

impl ZipWriterStats
{
    fn update(&mut self, buf: &[u8])
    {
        self.crc32 = crc32::update(self.crc32, buf);
        self.bytes_written += buf.len() as u64;
    }
}

impl<W: Write+io::Seek> ZipWriter<W>
{
    /// Initializes the ZipWriter.
    ///
    /// Before writing to this object, the start_file command should be called.
    pub fn new(inner: W) -> ZipWriter<W>
    {
        ZipWriter
        {
            inner: GenericZipWriter::Storer(inner),
            files: Vec::new(),
            stats: Default::default(),
        }
    }

    /// Start a new file for with the requested compression method.
    pub fn start_file(&mut self, name: &str, compression: CompressionMethod) -> ZipResult<()>
    {
        try!(self.finish_file());

        {
            let writer = self.inner.get_plain();
            let header_start = try!(writer.seek(io::SeekFrom::Current(0)));

            let mut file = ZipFileData
            {
                encrypted: false,
                compression_method: compression,
                last_modified_time: time::now(),
                crc32: 0,
                compressed_size: 0,
                uncompressed_size: 0,
                file_name: name.to_string(),
                file_comment: String::new(),
                header_start: header_start,
                data_start: 0,
            };
            try!(write_local_file_header(writer, &file));

            let header_end = try!(writer.seek(io::SeekFrom::Current(0)));
            self.stats.start = header_end;
            file.data_start = header_end;

            self.stats.bytes_written = 0;
            self.stats.crc32 = 0;

            self.files.push(file);
        }

        try!(self.inner.switch_to(compression));

        Ok(())
    }

    fn finish_file(&mut self) -> ZipResult<()>
    {
        try!(self.inner.switch_to(CompressionMethod::Stored));
        let writer = self.inner.get_plain();

        let file = match self.files.last_mut()
        {
            None => return Ok(()),
            Some(f) => f,
        };
        file.crc32 = self.stats.crc32;
        file.uncompressed_size = self.stats.bytes_written;
        file.compressed_size = try!(writer.seek(io::SeekFrom::Current(0))) - self.stats.start;

        try!(update_local_file_header(writer, file));
        try!(writer.seek(io::SeekFrom::End(0)));
        Ok(())
    }

    /// Finish the last file and write all other zip-structures
    ///
    /// This will return the writer, but one should normally not append any data to the end of the file.
    /// Note that the zipfile will also be finished on drop.
    pub fn finish(mut self) -> ZipResult<W>
    {
        try!(self.finalize());
        let inner = mem::replace(&mut self.inner, GenericZipWriter::Closed);
        Ok(inner.unwrap())
    }

    fn finalize(&mut self) -> ZipResult<()>
    {
        try!(self.finish_file());

        {
            let writer = self.inner.get_plain();

            let central_start = try!(writer.seek(io::SeekFrom::Current(0)));
            for file in self.files.iter()
            {
                try!(write_central_directory_header(writer, file));
            }
            let central_size = try!(writer.seek(io::SeekFrom::Current(0))) - central_start;

            let footer = spec::CentralDirectoryEnd
            {
                disk_number: 0,
                disk_with_central_directory: 0,
                number_of_files_on_this_disk: self.files.len() as u16,
                number_of_files: self.files.len() as u16,
                central_directory_size: central_size as u32,
                central_directory_offset: central_start as u32,
                zip_file_comment: b"zip-rs".to_vec(),
            };

            try!(footer.write(writer));
        }

        Ok(())
    }
}

#[unsafe_destructor]
impl<W: Write+io::Seek> Drop for ZipWriter<W>
{
    fn drop(&mut self)
    {
        if !self.inner.is_closed()
        {
            if let Err(e) = self.finalize() {
                let _ = write!(&mut io::stderr(), "ZipWriter drop failed: {:?}", e);
            }
        }
    }
}

impl<W: Write+io::Seek> GenericZipWriter<W>
{
    fn switch_to(&mut self, compression: CompressionMethod) -> ZipResult<()>
    {
        let bare = match mem::replace(self, GenericZipWriter::Closed)
        {
            GenericZipWriter::Storer(w) => w,
            GenericZipWriter::Deflater(w) => try!(w.finish()),
            GenericZipWriter::Bzip2(w) => match w.into_inner() { Ok(r) => r, Err((_, err)) => try!(Err(err)) },
            GenericZipWriter::Closed => try!(Err(io::Error::new(io::ErrorKind::BrokenPipe, "ZipWriter was already closed", None))),
        };

        *self = match compression
        {
            CompressionMethod::Stored => GenericZipWriter::Storer(bare),
            CompressionMethod::Deflated => GenericZipWriter::Deflater(bare.deflate_encode(flate2::Compression::Default)),
            CompressionMethod::Bzip2 => GenericZipWriter::Bzip2(BzCompressor::new(bare, bzip2::Compress::Default)),
            _ => return Err(ZipError::UnsupportedArchive("Unsupported compression")),
        };

        Ok(())
    }

    fn is_closed(&self) -> bool
    {
        match *self
        {
            GenericZipWriter::Closed => true,
            _ => false,
        }
    }

    fn get_plain(&mut self) -> &mut W
    {
        match *self
        {
            GenericZipWriter::Storer(ref mut w) => w,
            _ => panic!("Should have switched to stored beforehand"),
        }
    }

    fn unwrap(self) -> W
    {
        match self
        {
            GenericZipWriter::Storer(w) => w,
            _ => panic!("Should have switched to stored beforehand"),
        }
    }
}

fn write_local_file_header<T: Write>(writer: &mut T, file: &ZipFileData) -> ZipResult<()>
{
    try!(writer.write_u32::<LittleEndian>(spec::LOCAL_FILE_HEADER_SIGNATURE));
    try!(writer.write_u16::<LittleEndian>(20));
    let flag = if !file.file_name.is_ascii() { 1u16 << 11 } else { 0 };
    try!(writer.write_u16::<LittleEndian>(flag));
    try!(writer.write_u16::<LittleEndian>(file.compression_method.to_u16()));
    try!(writer.write_u16::<LittleEndian>(util::tm_to_msdos_time(file.last_modified_time)));
    try!(writer.write_u16::<LittleEndian>(util::tm_to_msdos_date(file.last_modified_time)));
    try!(writer.write_u32::<LittleEndian>(file.crc32));
    try!(writer.write_u32::<LittleEndian>(file.compressed_size as u32));
    try!(writer.write_u32::<LittleEndian>(file.uncompressed_size as u32));
    try!(writer.write_u16::<LittleEndian>(file.file_name.as_bytes().len() as u16));
    let extra_field = try!(build_extra_field(file));
    try!(writer.write_u16::<LittleEndian>(extra_field.len() as u16));
    try!(writer.write_all(file.file_name.as_bytes()));
    try!(writer.write_all(&extra_field));

    Ok(())
}

fn update_local_file_header<T: Write+io::Seek>(writer: &mut T, file: &ZipFileData) -> ZipResult<()>
{
    static CRC32_OFFSET : u64 = 14;
    try!(writer.seek(io::SeekFrom::Start(file.header_start + CRC32_OFFSET)));
    try!(writer.write_u32::<LittleEndian>(file.crc32));
    try!(writer.write_u32::<LittleEndian>(file.compressed_size as u32));
    try!(writer.write_u32::<LittleEndian>(file.uncompressed_size as u32));
    Ok(())
}

fn write_central_directory_header<T: Write>(writer: &mut T, file: &ZipFileData) -> ZipResult<()>
{
    try!(writer.write_u32::<LittleEndian>(spec::CENTRAL_DIRECTORY_HEADER_SIGNATURE));
    try!(writer.write_u16::<LittleEndian>(0x14FF));
    try!(writer.write_u16::<LittleEndian>(20));
    let flag = if !file.file_name.is_ascii() { 1u16 << 11 } else { 0 };
    try!(writer.write_u16::<LittleEndian>(flag));
    try!(writer.write_u16::<LittleEndian>(file.compression_method.to_u16()));
    try!(writer.write_u16::<LittleEndian>(util::tm_to_msdos_time(file.last_modified_time)));
    try!(writer.write_u16::<LittleEndian>(util::tm_to_msdos_date(file.last_modified_time)));
    try!(writer.write_u32::<LittleEndian>(file.crc32));
    try!(writer.write_u32::<LittleEndian>(file.compressed_size as u32));
    try!(writer.write_u32::<LittleEndian>(file.uncompressed_size as u32));
    try!(writer.write_u16::<LittleEndian>(file.file_name.as_bytes().len() as u16));
    let extra_field = try!(build_extra_field(file));
    try!(writer.write_u16::<LittleEndian>(extra_field.len() as u16));
    try!(writer.write_u16::<LittleEndian>(0));
    try!(writer.write_u16::<LittleEndian>(0));
    try!(writer.write_u16::<LittleEndian>(0));
    try!(writer.write_u32::<LittleEndian>(0));
    try!(writer.write_u32::<LittleEndian>(file.header_start as u32));
    try!(writer.write_all(file.file_name.as_bytes()));
    try!(writer.write_all(&extra_field));

    Ok(())
}

fn build_extra_field(_file: &ZipFileData) -> ZipResult<Vec<u8>>
{
    let writer = Vec::new();
    // Future work
    Ok(writer)
}
<commit_msg>Move the pattern matching of genericzipwriter<commit_after>//! Structs for creating a new zip archive

use compression::CompressionMethod;
use types::ZipFileData;
use spec;
use crc32;
use result::{ZipResult, ZipError};
use std::default::Default;
use std::io;
use std::io::prelude::*;
use std::mem;
use std::ascii::AsciiExt;
use time;
use flate2;
use flate2::FlateWriteExt;
use flate2::write::DeflateEncoder;
use bzip2;
use bzip2::writer::BzCompressor;
use util;
use podio::{WritePodExt, LittleEndian};

enum GenericZipWriter<W: Write + io::Seek>
{
    Closed,
    Storer(W),
    Deflater(DeflateEncoder<W>),
    Bzip2(BzCompressor<W>),
}

/// Generator for ZIP files.
///
/// ```
/// fn doit() -> zip::result::ZipResult<()>
/// {
///     use std::io::Write;
///
///     // For this example we write to a buffer, but normally you should use a File
///     let mut buf: &mut [u8] = &mut [0u8; 65536];
///     let mut w = std::io::Cursor::new(buf);
///     let mut zip = zip::ZipWriter::new(w);
///
///     try!(zip.start_file("hello_world.txt", zip::CompressionMethod::Stored));
///     try!(zip.write(b"Hello, World!"));
///
///     // Optionally finish the zip. (this is also done on drop)
///     try!(zip.finish());
///
///     Ok(())
/// }
///
/// println!("Result: {:?}", doit());
/// ```
pub struct ZipWriter<W: Write + io::Seek>
{
    inner: GenericZipWriter<W>,
    files: Vec<ZipFileData>,
    stats: ZipWriterStats,
}

#[derive(Default)]
struct ZipWriterStats
{
    crc32: u32,
    start: u64,
    bytes_written: u64,
}

impl<W: Write+io::Seek> Write for ZipWriter<W>
{
    fn write(&mut self, buf: &[u8]) -> io::Result<usize>
    {
        if self.files.len() == 0 { return Err(io::Error::new(io::ErrorKind::Other, "No file has been started", None)) }
        self.stats.update(buf);
        match self.inner.ref_mut()
        {
            Some(ref mut w) => w.write(buf),
            None => Err(io::Error::new(io::ErrorKind::BrokenPipe, "ZipWriter was already closed", None)),
        }
    }

    fn flush(&mut self) -> io::Result<()>
    {
        match self.inner.ref_mut()
        {
            Some(ref mut w) => w.flush(),
            None => Err(io::Error::new(io::ErrorKind::BrokenPipe, "ZipWriter was already closed", None)),
        }
    }
}

impl ZipWriterStats
{
    fn update(&mut self, buf: &[u8])
    {
        self.crc32 = crc32::update(self.crc32, buf);
        self.bytes_written += buf.len() as u64;
    }
}

impl<W: Write+io::Seek> ZipWriter<W>
{
    /// Initializes the ZipWriter.
    ///
    /// Before writing to this object, the start_file command should be called.
    pub fn new(inner: W) -> ZipWriter<W>
    {
        ZipWriter
        {
            inner: GenericZipWriter::Storer(inner),
            files: Vec::new(),
            stats: Default::default(),
        }
    }

    /// Start a new file for with the requested compression method.
    pub fn start_file(&mut self, name: &str, compression: CompressionMethod) -> ZipResult<()>
    {
        try!(self.finish_file());

        {
            let writer = self.inner.get_plain();
            let header_start = try!(writer.seek(io::SeekFrom::Current(0)));

            let mut file = ZipFileData
            {
                encrypted: false,
                compression_method: compression,
                last_modified_time: time::now(),
                crc32: 0,
                compressed_size: 0,
                uncompressed_size: 0,
                file_name: name.to_string(),
                file_comment: String::new(),
                header_start: header_start,
                data_start: 0,
            };
            try!(write_local_file_header(writer, &file));

            let header_end = try!(writer.seek(io::SeekFrom::Current(0)));
            self.stats.start = header_end;
            file.data_start = header_end;

            self.stats.bytes_written = 0;
            self.stats.crc32 = 0;

            self.files.push(file);
        }

        try!(self.inner.switch_to(compression));

        Ok(())
    }

    fn finish_file(&mut self) -> ZipResult<()>
    {
        try!(self.inner.switch_to(CompressionMethod::Stored));
        let writer = self.inner.get_plain();

        let file = match self.files.last_mut()
        {
            None => return Ok(()),
            Some(f) => f,
        };
        file.crc32 = self.stats.crc32;
        file.uncompressed_size = self.stats.bytes_written;
        file.compressed_size = try!(writer.seek(io::SeekFrom::Current(0))) - self.stats.start;

        try!(update_local_file_header(writer, file));
        try!(writer.seek(io::SeekFrom::End(0)));
        Ok(())
    }

    /// Finish the last file and write all other zip-structures
    ///
    /// This will return the writer, but one should normally not append any data to the end of the file.
    /// Note that the zipfile will also be finished on drop.
    pub fn finish(mut self) -> ZipResult<W>
    {
        try!(self.finalize());
        let inner = mem::replace(&mut self.inner, GenericZipWriter::Closed);
        Ok(inner.unwrap())
    }

    fn finalize(&mut self) -> ZipResult<()>
    {
        try!(self.finish_file());

        {
            let writer = self.inner.get_plain();

            let central_start = try!(writer.seek(io::SeekFrom::Current(0)));
            for file in self.files.iter()
            {
                try!(write_central_directory_header(writer, file));
            }
            let central_size = try!(writer.seek(io::SeekFrom::Current(0))) - central_start;

            let footer = spec::CentralDirectoryEnd
            {
                disk_number: 0,
                disk_with_central_directory: 0,
                number_of_files_on_this_disk: self.files.len() as u16,
                number_of_files: self.files.len() as u16,
                central_directory_size: central_size as u32,
                central_directory_offset: central_start as u32,
                zip_file_comment: b"zip-rs".to_vec(),
            };

            try!(footer.write(writer));
        }

        Ok(())
    }
}

#[unsafe_destructor]
impl<W: Write+io::Seek> Drop for ZipWriter<W>
{
    fn drop(&mut self)
    {
        if !self.inner.is_closed()
        {
            if let Err(e) = self.finalize() {
                let _ = write!(&mut io::stderr(), "ZipWriter drop failed: {:?}", e);
            }
        }
    }
}

impl<W: Write+io::Seek> GenericZipWriter<W>
{
    fn switch_to(&mut self, compression: CompressionMethod) -> ZipResult<()>
    {
        let bare = match mem::replace(self, GenericZipWriter::Closed)
        {
            GenericZipWriter::Storer(w) => w,
            GenericZipWriter::Deflater(w) => try!(w.finish()),
            GenericZipWriter::Bzip2(w) => match w.into_inner() { Ok(r) => r, Err((_, err)) => try!(Err(err)) },
            GenericZipWriter::Closed => try!(Err(io::Error::new(io::ErrorKind::BrokenPipe, "ZipWriter was already closed", None))),
        };

        *self = match compression
        {
            CompressionMethod::Stored => GenericZipWriter::Storer(bare),
            CompressionMethod::Deflated => GenericZipWriter::Deflater(bare.deflate_encode(flate2::Compression::Default)),
            CompressionMethod::Bzip2 => GenericZipWriter::Bzip2(BzCompressor::new(bare, bzip2::Compress::Default)),
            _ => return Err(ZipError::UnsupportedArchive("Unsupported compression")),
        };

        Ok(())
    }

    fn ref_mut(&mut self) -> Option<&mut Write> {
        match *self {
            GenericZipWriter::Storer(ref mut w) => Some(w as &mut Write),
            GenericZipWriter::Deflater(ref mut w) => Some(w as &mut Write),
            GenericZipWriter::Bzip2(ref mut w) => Some(w as &mut Write),
            GenericZipWriter::Closed => None,
        }
    }

    fn is_closed(&self) -> bool
    {
        match *self
        {
            GenericZipWriter::Closed => true,
            _ => false,
        }
    }

    fn get_plain(&mut self) -> &mut W
    {
        match *self
        {
            GenericZipWriter::Storer(ref mut w) => w,
            _ => panic!("Should have switched to stored beforehand"),
        }
    }

    fn unwrap(self) -> W
    {
        match self
        {
            GenericZipWriter::Storer(w) => w,
            _ => panic!("Should have switched to stored beforehand"),
        }
    }
}

fn write_local_file_header<T: Write>(writer: &mut T, file: &ZipFileData) -> ZipResult<()>
{
    try!(writer.write_u32::<LittleEndian>(spec::LOCAL_FILE_HEADER_SIGNATURE));
    try!(writer.write_u16::<LittleEndian>(20));
    let flag = if !file.file_name.is_ascii() { 1u16 << 11 } else { 0 };
    try!(writer.write_u16::<LittleEndian>(flag));
    try!(writer.write_u16::<LittleEndian>(file.compression_method.to_u16()));
    try!(writer.write_u16::<LittleEndian>(util::tm_to_msdos_time(file.last_modified_time)));
    try!(writer.write_u16::<LittleEndian>(util::tm_to_msdos_date(file.last_modified_time)));
    try!(writer.write_u32::<LittleEndian>(file.crc32));
    try!(writer.write_u32::<LittleEndian>(file.compressed_size as u32));
    try!(writer.write_u32::<LittleEndian>(file.uncompressed_size as u32));
    try!(writer.write_u16::<LittleEndian>(file.file_name.as_bytes().len() as u16));
    let extra_field = try!(build_extra_field(file));
    try!(writer.write_u16::<LittleEndian>(extra_field.len() as u16));
    try!(writer.write_all(file.file_name.as_bytes()));
    try!(writer.write_all(&extra_field));

    Ok(())
}

fn update_local_file_header<T: Write+io::Seek>(writer: &mut T, file: &ZipFileData) -> ZipResult<()>
{
    static CRC32_OFFSET : u64 = 14;
    try!(writer.seek(io::SeekFrom::Start(file.header_start + CRC32_OFFSET)));
    try!(writer.write_u32::<LittleEndian>(file.crc32));
    try!(writer.write_u32::<LittleEndian>(file.compressed_size as u32));
    try!(writer.write_u32::<LittleEndian>(file.uncompressed_size as u32));
    Ok(())
}

fn write_central_directory_header<T: Write>(writer: &mut T, file: &ZipFileData) -> ZipResult<()>
{
    try!(writer.write_u32::<LittleEndian>(spec::CENTRAL_DIRECTORY_HEADER_SIGNATURE));
    try!(writer.write_u16::<LittleEndian>(0x14FF));
    try!(writer.write_u16::<LittleEndian>(20));
    let flag = if !file.file_name.is_ascii() { 1u16 << 11 } else { 0 };
    try!(writer.write_u16::<LittleEndian>(flag));
    try!(writer.write_u16::<LittleEndian>(file.compression_method.to_u16()));
    try!(writer.write_u16::<LittleEndian>(util::tm_to_msdos_time(file.last_modified_time)));
    try!(writer.write_u16::<LittleEndian>(util::tm_to_msdos_date(file.last_modified_time)));
    try!(writer.write_u32::<LittleEndian>(file.crc32));
    try!(writer.write_u32::<LittleEndian>(file.compressed_size as u32));
    try!(writer.write_u32::<LittleEndian>(file.uncompressed_size as u32));
    try!(writer.write_u16::<LittleEndian>(file.file_name.as_bytes().len() as u16));
    let extra_field = try!(build_extra_field(file));
    try!(writer.write_u16::<LittleEndian>(extra_field.len() as u16));
    try!(writer.write_u16::<LittleEndian>(0));
    try!(writer.write_u16::<LittleEndian>(0));
    try!(writer.write_u16::<LittleEndian>(0));
    try!(writer.write_u32::<LittleEndian>(0));
    try!(writer.write_u32::<LittleEndian>(file.header_start as u32));
    try!(writer.write_all(file.file_name.as_bytes()));
    try!(writer.write_all(&extra_field));

    Ok(())
}

fn build_extra_field(_file: &ZipFileData) -> ZipResult<Vec<u8>>
{
    let writer = Vec::new();
    // Future work
    Ok(writer)
}
<|endoftext|>


<commit_before>
extern crate ndarray;
extern crate ndarray_linalg;
extern crate ndarray_rand;
extern crate rand;

use ndarray::prelude::*;
use ndarray_linalg::prelude::*;
use rand::distributions::*;
use ndarray_rand::RandomExt;

fn all_close(a: Array<f64, (Ix, Ix)>, b: Array<f64, (Ix, Ix)>) {
    if !a.all_close(&b, 1.0e-7) {
        panic!("\nTwo matrices are not equal:\na = \n{:?}\nb = \n{:?}\n",
               a,
               b);
    }
}
#[test]
fn permutate() {
    let a = arr2(&[[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]);
    println!("a= \n{:?}", &a);
    let p = vec![2, 2, 3]; // replace 1-2
    let pa = a.permutate_column(&p);
    println!("permutated = \n{:?}", &pa);
    all_close(pa, arr2(&[[4., 5., 6.], [1., 2., 3.], [7., 8., 9.]]))
}

#[test]
fn permutate_t() {
    let a = arr2(&[[1., 4., 7.], [2., 5., 8.], [3., 6., 9.]]).reversed_axes();
    println!("a= \n{:?}", &a);
    let p = vec![2, 2, 3]; // replace 1-2
    let pa = a.permutate_column(&p);
    println!("permutated = \n{:?}", &pa);
    all_close(pa, arr2(&[[4., 5., 6.], [1., 2., 3.], [7., 8., 9.]]))
}

#[test]
fn permutate_3x4_t() {
    let a = arr2(&[[1., 5., 9.], [2., 6., 10.], [3., 7., 11.], [4., 8., 12.]]).reversed_axes();
    println!("a= \n{:?}", &a);
    let p = vec![1, 3, 3]; // replace 2-3
    println!("permutation = \n{:?}", &p);
    let pa = a.permutate_column(&p);
    println!("permutated = \n{:?}", &pa);
    all_close(pa,
              arr2(&[[1., 2., 3., 4.], [9., 10., 11., 12.], [5., 6., 7., 8.]]));
}

#[test]
fn permutate_4x3_t() {
    let a = arr2(&[[1., 4., 7., 10.], [2., 5., 8., 11.], [3., 6., 9., 12.]]).reversed_axes();
    println!("a= \n{:?}", &a);
    let p = vec![4, 2, 3, 4]; // replace 1-4
    println!("permutation = \n{:?}", &p);
    let pa = a.permutate_column(&p);
    println!("permutated = \n{:?}", &pa);
    all_close(pa,
              arr2(&[[10., 11., 12.], [4., 5., 6.], [7., 8., 9.], [1., 2., 3.]]))
}

#[test]
fn lu_square_upper() {
    let r_dist = Range::new(0., 1.);
    let mut a = Array::<f64, _>::random((3, 3), r_dist);
    for ((i, j), val) in a.indexed_iter_mut() {
        if i > j {
            *val = 0.0;
        }
    }
    println!("a = \n{:?}", &a);
    let (p, l, u) = a.clone().lu().unwrap();
    println!("P = \n{:?}", &p);
    println!("L = \n{:?}", &l);
    println!("U = \n{:?}", &u);
    all_close(l.dot(&u).permutate_column(&p), a);
}

#[test]
fn lu_square_upper_t() {
    let r_dist = Range::new(0., 1.);
    let mut a = Array::<f64, _>::random((3, 3), r_dist).reversed_axes();
    for ((i, j), val) in a.indexed_iter_mut() {
        if i > j {
            *val = 0.0;
        }
    }
    println!("a = \n{:?}", &a);
    let (p, l, u) = a.clone().lu().unwrap();
    println!("P = \n{:?}", &p);
    println!("L = \n{:?}", &l);
    println!("U = \n{:?}", &u);
    all_close(l.dot(&u).permutate_column(&p), a);
}

#[test]
fn lu_square_lower_t() {
    let r_dist = Range::new(0., 1.);
    let mut a = Array::<f64, _>::random((3, 3), r_dist).reversed_axes();
    for ((i, j), val) in a.indexed_iter_mut() {
        if i < j {
            *val = 0.0;
        }
    }
    println!("a = \n{:?}", &a);
    let (p, l, u) = a.clone().lu().unwrap();
    println!("P = \n{:?}", &p);
    println!("L = \n{:?}", &l);
    println!("U = \n{:?}", &u);
    println!("LU = \n{:?}", l.dot(&u));
    all_close(l.dot(&u).permutate_column(&p), a);
}

#[test]
fn lu_square_t() {
    let r_dist = Range::new(0., 1.);
    let a = Array::<f64, _>::random((3, 3), r_dist).reversed_axes();
    println!("a = \n{:?}", &a);
    let (p, l, u) = a.clone().lu().unwrap();
    println!("P = \n{:?}", &p);
    println!("L = \n{:?}", &l);
    println!("U = \n{:?}", &u);
    all_close(l.dot(&u).permutate_column(&p), a);
}
<commit_msg>Add test for Row major<commit_after>
extern crate ndarray;
extern crate ndarray_linalg;
extern crate ndarray_rand;
extern crate rand;

use ndarray::prelude::*;
use ndarray_linalg::prelude::*;
use rand::distributions::*;
use ndarray_rand::RandomExt;

fn all_close(a: Array<f64, (Ix, Ix)>, b: Array<f64, (Ix, Ix)>) {
    if !a.all_close(&b, 1.0e-7) {
        panic!("\nTwo matrices are not equal:\na = \n{:?}\nb = \n{:?}\n",
               a,
               b);
    }
}
#[test]
fn permutate() {
    let a = arr2(&[[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]);
    println!("a= \n{:?}", &a);
    let p = vec![2, 2, 3]; // replace 1-2
    let pa = a.permutate_column(&p);
    println!("permutated = \n{:?}", &pa);
    all_close(pa, arr2(&[[4., 5., 6.], [1., 2., 3.], [7., 8., 9.]]))
}

#[test]
fn permutate_t() {
    let a = arr2(&[[1., 4., 7.], [2., 5., 8.], [3., 6., 9.]]).reversed_axes();
    println!("a= \n{:?}", &a);
    let p = vec![2, 2, 3]; // replace 1-2
    let pa = a.permutate_column(&p);
    println!("permutated = \n{:?}", &pa);
    all_close(pa, arr2(&[[4., 5., 6.], [1., 2., 3.], [7., 8., 9.]]))
}

#[test]
fn permutate_3x4_t() {
    let a = arr2(&[[1., 5., 9.], [2., 6., 10.], [3., 7., 11.], [4., 8., 12.]]).reversed_axes();
    println!("a= \n{:?}", &a);
    let p = vec![1, 3, 3]; // replace 2-3
    println!("permutation = \n{:?}", &p);
    let pa = a.permutate_column(&p);
    println!("permutated = \n{:?}", &pa);
    all_close(pa,
              arr2(&[[1., 2., 3., 4.], [9., 10., 11., 12.], [5., 6., 7., 8.]]));
}

#[test]
fn permutate_4x3_t() {
    let a = arr2(&[[1., 4., 7., 10.], [2., 5., 8., 11.], [3., 6., 9., 12.]]).reversed_axes();
    println!("a= \n{:?}", &a);
    let p = vec![4, 2, 3, 4]; // replace 1-4
    println!("permutation = \n{:?}", &p);
    let pa = a.permutate_column(&p);
    println!("permutated = \n{:?}", &pa);
    all_close(pa,
              arr2(&[[10., 11., 12.], [4., 5., 6.], [7., 8., 9.], [1., 2., 3.]]))
}

#[test]
fn lu_square_upper() {
    let r_dist = Range::new(0., 1.);
    let mut a = Array::<f64, _>::random((3, 3), r_dist);
    for ((i, j), val) in a.indexed_iter_mut() {
        if i > j {
            *val = 0.0;
        }
    }
    println!("a = \n{:?}", &a);
    let (p, l, u) = a.clone().lu().unwrap();
    println!("P = \n{:?}", &p);
    println!("L = \n{:?}", &l);
    println!("U = \n{:?}", &u);
    all_close(l.dot(&u).permutate_column(&p), a);
}

#[test]
fn lu_square_upper_t() {
    let r_dist = Range::new(0., 1.);
    let mut a = Array::<f64, _>::random((3, 3), r_dist).reversed_axes();
    for ((i, j), val) in a.indexed_iter_mut() {
        if i > j {
            *val = 0.0;
        }
    }
    println!("a = \n{:?}", &a);
    let (p, l, u) = a.clone().lu().unwrap();
    println!("P = \n{:?}", &p);
    println!("L = \n{:?}", &l);
    println!("U = \n{:?}", &u);
    all_close(l.dot(&u).permutate_column(&p), a);
}

#[test]
fn lu_square_lower() {
    let r_dist = Range::new(0., 1.);
    let mut a = Array::<f64, _>::random((3, 3), r_dist);
    for ((i, j), val) in a.indexed_iter_mut() {
        if i < j {
            *val = 0.0;
        }
    }
    println!("a = \n{:?}", &a);
    let (p, l, u) = a.clone().lu().unwrap();
    println!("P = \n{:?}", &p);
    println!("L = \n{:?}", &l);
    println!("U = \n{:?}", &u);
    println!("LU = \n{:?}", l.dot(&u));
    all_close(l.dot(&u).permutate_column(&p), a);
}

#[test]
fn lu_square_lower_t() {
    let r_dist = Range::new(0., 1.);
    let mut a = Array::<f64, _>::random((3, 3), r_dist).reversed_axes();
    for ((i, j), val) in a.indexed_iter_mut() {
        if i < j {
            *val = 0.0;
        }
    }
    println!("a = \n{:?}", &a);
    let (p, l, u) = a.clone().lu().unwrap();
    println!("P = \n{:?}", &p);
    println!("L = \n{:?}", &l);
    println!("U = \n{:?}", &u);
    println!("LU = \n{:?}", l.dot(&u));
    all_close(l.dot(&u).permutate_column(&p), a);
}

#[test]
fn lu_square() {
    let r_dist = Range::new(0., 1.);
    let a = Array::<f64, _>::random((3, 3), r_dist);
    println!("a = \n{:?}", &a);
    let (p, l, u) = a.clone().lu().unwrap();
    println!("P = \n{:?}", &p);
    println!("L = \n{:?}", &l);
    println!("U = \n{:?}", &u);
    println!("LU = \n{:?}", l.dot(&u));
    all_close(l.dot(&u).permutate_column(&p), a);
}

#[test]
fn lu_square_t() {
    let r_dist = Range::new(0., 1.);
    let a = Array::<f64, _>::random((3, 3), r_dist).reversed_axes();
    println!("a = \n{:?}", &a);
    let (p, l, u) = a.clone().lu().unwrap();
    println!("P = \n{:?}", &p);
    println!("L = \n{:?}", &l);
    println!("U = \n{:?}", &u);
    all_close(l.dot(&u).permutate_column(&p), a);
}
<|endoftext|>


<commit_before>/* Copyright (c) 2010-2011 Dmitry Vyukov. All rights reserved.
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *    1. Redistributions of source code must retain the above copyright notice,
 *       this list of conditions and the following disclaimer.
 *
 *    2. Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY DMITRY VYUKOV "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
 * SHALL DMITRY VYUKOV OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * The views and conclusions contained in the software and documentation are
 * those of the authors and should not be interpreted as representing official
 * policies, either expressed or implied, of Dmitry Vyukov.
 */

// http://www.1024cores.net/home/lock-free-algorithms/queues/unbounded-spsc-queue

//! A single-producer single-consumer concurrent queue
//!
//! This module contains the implementation of an SPSC queue which can be used
//! concurrently between two threads. This data structure is safe to use and
//! enforces the semantics that there is one pusher and one popper.

#![unstable(feature = "std_misc")]

use core::prelude::*;

use alloc::boxed;
use alloc::boxed::Box;
use core::ptr;
use core::cell::UnsafeCell;

use sync::atomic::{AtomicPtr, AtomicUsize, Ordering};

// Node within the linked list queue of messages to send
struct Node<T> {
    // FIXME: this could be an uninitialized T if we're careful enough, and
    //      that would reduce memory usage (and be a bit faster).
    //      is it worth it?
    value: Option<T>,           // nullable for re-use of nodes
    next: AtomicPtr<Node<T>>,   // next node in the queue
}

/// The single-producer single-consumer queue. This structure is not cloneable,
/// but it can be safely shared in an Arc if it is guaranteed that there
/// is only one popper and one pusher touching the queue at any one point in
/// time.
pub struct Queue<T> {
    // consumer fields
    tail: UnsafeCell<*mut Node<T>>, // where to pop from
    tail_prev: AtomicPtr<Node<T>>, // where to pop from

    // producer fields
    head: UnsafeCell<*mut Node<T>>,      // where to push to
    first: UnsafeCell<*mut Node<T>>,     // where to get new nodes from
    tail_copy: UnsafeCell<*mut Node<T>>, // between first/tail

    // Cache maintenance fields. Additions and subtractions are stored
    // separately in order to allow them to use nonatomic addition/subtraction.
    cache_bound: usize,
    cache_additions: AtomicUsize,
    cache_subtractions: AtomicUsize,
}

unsafe impl<T: Send> Send for Queue<T> { }

unsafe impl<T: Send> Sync for Queue<T> { }

impl<T> Node<T> {
    fn new() -> *mut Node<T> {
        unsafe {
            boxed::into_raw(box Node {
                value: None,
                next: AtomicPtr::new(ptr::null_mut::<Node<T>>()),
            })
        }
    }
}

impl<T> Queue<T> {
    /// Creates a new queue.
    ///
    /// This is unsafe as the type system doesn't enforce a single
    /// consumer-producer relationship. It also allows the consumer to `pop`
    /// items while there is a `peek` active due to all methods having a
    /// non-mutable receiver.
    ///
    /// # Arguments
    ///
    ///   * `bound` - This queue implementation is implemented with a linked
    ///               list, and this means that a push is always a malloc. In
    ///               order to amortize this cost, an internal cache of nodes is
    ///               maintained to prevent a malloc from always being
    ///               necessary. This bound is the limit on the size of the
    ///               cache (if desired). If the value is 0, then the cache has
    ///               no bound. Otherwise, the cache will never grow larger than
    ///               `bound` (although the queue itself could be much larger.
    pub unsafe fn new(bound: usize) -> Queue<T> {
        let n1 = Node::new();
        let n2 = Node::new();
        (*n1).next.store(n2, Ordering::Relaxed);
        Queue {
            tail: UnsafeCell::new(n2),
            tail_prev: AtomicPtr::new(n1),
            head: UnsafeCell::new(n2),
            first: UnsafeCell::new(n1),
            tail_copy: UnsafeCell::new(n1),
            cache_bound: bound,
            cache_additions: AtomicUsize::new(0),
            cache_subtractions: AtomicUsize::new(0),
        }
    }

    /// Pushes a new value onto this queue. Note that to use this function
    /// safely, it must be externally guaranteed that there is only one pusher.
    pub fn push(&self, t: T) {
        unsafe {
            // Acquire a node (which either uses a cached one or allocates a new
            // one), and then append this to the 'head' node.
            let n = self.alloc();
            assert!((*n).value.is_none());
            (*n).value = Some(t);
            (*n).next.store(ptr::null_mut(), Ordering::Relaxed);
            (**self.head.get()).next.store(n, Ordering::Release);
            *self.head.get() = n;
        }
    }

    unsafe fn alloc(&self) -> *mut Node<T> {
        // First try to see if we can consume the 'first' node for our uses.
        // We try to avoid as many atomic instructions as possible here, so
        // the addition to cache_subtractions is not atomic (plus we're the
        // only one subtracting from the cache).
        if *self.first.get() != *self.tail_copy.get() {
            if self.cache_bound > 0 {
                let b = self.cache_subtractions.load(Ordering::Relaxed);
                self.cache_subtractions.store(b + 1, Ordering::Relaxed);
            }
            let ret = *self.first.get();
            *self.first.get() = (*ret).next.load(Ordering::Relaxed);
            return ret;
        }
        // If the above fails, then update our copy of the tail and try
        // again.
        *self.tail_copy.get() = self.tail_prev.load(Ordering::Acquire);
        if *self.first.get() != *self.tail_copy.get() {
            if self.cache_bound > 0 {
                let b = self.cache_subtractions.load(Ordering::Relaxed);
                self.cache_subtractions.store(b + 1, Ordering::Relaxed);
            }
            let ret = *self.first.get();
            *self.first.get() = (*ret).next.load(Ordering::Relaxed);
            return ret;
        }
        // If all of that fails, then we have to allocate a new node
        // (there's nothing in the node cache).
        Node::new()
    }

    /// Attempts to pop a value from this queue. Remember that to use this type
    /// safely you must ensure that there is only one popper at a time.
    pub fn pop(&self) -> Option<T> {
        unsafe {
            // The `tail` node is not actually a used node, but rather a
            // sentinel from where we should start popping from. Hence, look at
            // tail's next field and see if we can use it. If we do a pop, then
            // the current tail node is a candidate for going into the cache.
            let tail = *self.tail.get();
            let next = (*tail).next.load(Ordering::Acquire);
            if next.is_null() { return None }
            assert!((*next).value.is_some());
            let ret = (*next).value.take();

            *self.tail.get() = next;
            if self.cache_bound == 0 {
                self.tail_prev.store(tail, Ordering::Release);
            } else {
                // FIXME: this is dubious with overflow.
                let additions = self.cache_additions.load(Ordering::Relaxed);
                let subtractions = self.cache_subtractions.load(Ordering::Relaxed);
                let size = additions - subtractions;

                if size < self.cache_bound {
                    self.tail_prev.store(tail, Ordering::Release);
                    self.cache_additions.store(additions + 1, Ordering::Relaxed);
                } else {
                    (*self.tail_prev.load(Ordering::Relaxed))
                          .next.store(next, Ordering::Relaxed);
                    // We have successfully erased all references to 'tail', so
                    // now we can safely drop it.
                    let _: Box<Node<T>> = Box::from_raw(tail);
                }
            }
            return ret;
        }
    }

    /// Attempts to peek at the head of the queue, returning `None` if the queue
    /// has no data currently
    ///
    /// # Warning
    /// The reference returned is invalid if it is not used before the consumer
    /// pops the value off the queue. If the producer then pushes another value
    /// onto the queue, it will overwrite the value pointed to by the reference.
    pub fn peek<'a>(&'a self) -> Option<&'a mut T> {
        // This is essentially the same as above with all the popping bits
        // stripped out.
        unsafe {
            let tail = *self.tail.get();
            let next = (*tail).next.load(Ordering::Acquire);
            if next.is_null() { return None }
            return (*next).value.as_mut();
        }
    }
}

impl<T> Drop for Queue<T> {
    fn drop(&mut self) {
        unsafe {
            let mut cur = *self.first.get();
            while !cur.is_null() {
                let next = (*cur).next.load(Ordering::Relaxed);
                let _n: Box<Node<T>> = Box::from_raw(cur);
                cur = next;
            }
        }
    }
}

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

    use sync::Arc;
    use super::Queue;
    use thread;
    use sync::mpsc::channel;

    #[test]
    fn smoke() {
        unsafe {
            let queue = Queue::new(0);
            queue.push(1);
            queue.push(2);
            assert_eq!(queue.pop(), Some(1));
            assert_eq!(queue.pop(), Some(2));
            assert_eq!(queue.pop(), None);
            queue.push(3);
            queue.push(4);
            assert_eq!(queue.pop(), Some(3));
            assert_eq!(queue.pop(), Some(4));
            assert_eq!(queue.pop(), None);
        }
    }

    #[test]
    fn peek() {
        unsafe {
            let queue = Queue::new(0);
            queue.push(vec![1]);

            // Ensure the borrowchecker works
            match queue.peek() {
                Some(vec) => match &**vec {
                    // Note that `pop` is not allowed here due to borrow
                    [1] => {}
                    _ => return
                },
                None => unreachable!()
            }

            queue.pop();
        }
    }

    #[test]
    fn drop_full() {
        unsafe {
            let q: Queue<Box<_>> = Queue::new(0);
            q.push(box 1);
            q.push(box 2);
        }
    }

    #[test]
    fn smoke_bound() {
        unsafe {
            let q = Queue::new(0);
            q.push(1);
            q.push(2);
            assert_eq!(q.pop(), Some(1));
            assert_eq!(q.pop(), Some(2));
            assert_eq!(q.pop(), None);
            q.push(3);
            q.push(4);
            assert_eq!(q.pop(), Some(3));
            assert_eq!(q.pop(), Some(4));
            assert_eq!(q.pop(), None);
        }
    }

    #[test]
    fn stress() {
        unsafe {
            stress_bound(0);
            stress_bound(1);
        }

        unsafe fn stress_bound(bound: usize) {
            let q = Arc::new(Queue::new(bound));

            let (tx, rx) = channel();
            let q2 = q.clone();
            let _t = thread::spawn(move|| {
                for _ in 0..100000 {
                    loop {
                        match q2.pop() {
                            Some(1) => break,
                            Some(_) => panic!(),
                            None => {}
                        }
                    }
                }
                tx.send(()).unwrap();
            });
            for _ in 0..100000 {
                q.push(1);
            }
            rx.recv().unwrap();
        }
    }
}
<commit_msg>Remove unsafe block around boxed::into_raw() as it is now safe<commit_after>/* Copyright (c) 2010-2011 Dmitry Vyukov. All rights reserved.
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *    1. Redistributions of source code must retain the above copyright notice,
 *       this list of conditions and the following disclaimer.
 *
 *    2. Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY DMITRY VYUKOV "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
 * SHALL DMITRY VYUKOV OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * The views and conclusions contained in the software and documentation are
 * those of the authors and should not be interpreted as representing official
 * policies, either expressed or implied, of Dmitry Vyukov.
 */

// http://www.1024cores.net/home/lock-free-algorithms/queues/unbounded-spsc-queue

//! A single-producer single-consumer concurrent queue
//!
//! This module contains the implementation of an SPSC queue which can be used
//! concurrently between two threads. This data structure is safe to use and
//! enforces the semantics that there is one pusher and one popper.

#![unstable(feature = "std_misc")]

use core::prelude::*;

use alloc::boxed;
use alloc::boxed::Box;
use core::ptr;
use core::cell::UnsafeCell;

use sync::atomic::{AtomicPtr, AtomicUsize, Ordering};

// Node within the linked list queue of messages to send
struct Node<T> {
    // FIXME: this could be an uninitialized T if we're careful enough, and
    //      that would reduce memory usage (and be a bit faster).
    //      is it worth it?
    value: Option<T>,           // nullable for re-use of nodes
    next: AtomicPtr<Node<T>>,   // next node in the queue
}

/// The single-producer single-consumer queue. This structure is not cloneable,
/// but it can be safely shared in an Arc if it is guaranteed that there
/// is only one popper and one pusher touching the queue at any one point in
/// time.
pub struct Queue<T> {
    // consumer fields
    tail: UnsafeCell<*mut Node<T>>, // where to pop from
    tail_prev: AtomicPtr<Node<T>>, // where to pop from

    // producer fields
    head: UnsafeCell<*mut Node<T>>,      // where to push to
    first: UnsafeCell<*mut Node<T>>,     // where to get new nodes from
    tail_copy: UnsafeCell<*mut Node<T>>, // between first/tail

    // Cache maintenance fields. Additions and subtractions are stored
    // separately in order to allow them to use nonatomic addition/subtraction.
    cache_bound: usize,
    cache_additions: AtomicUsize,
    cache_subtractions: AtomicUsize,
}

unsafe impl<T: Send> Send for Queue<T> { }

unsafe impl<T: Send> Sync for Queue<T> { }

impl<T> Node<T> {
    fn new() -> *mut Node<T> {
        boxed::into_raw(box Node {
            value: None,
            next: AtomicPtr::new(ptr::null_mut::<Node<T>>()),
        })
    }
}

impl<T> Queue<T> {
    /// Creates a new queue.
    ///
    /// This is unsafe as the type system doesn't enforce a single
    /// consumer-producer relationship. It also allows the consumer to `pop`
    /// items while there is a `peek` active due to all methods having a
    /// non-mutable receiver.
    ///
    /// # Arguments
    ///
    ///   * `bound` - This queue implementation is implemented with a linked
    ///               list, and this means that a push is always a malloc. In
    ///               order to amortize this cost, an internal cache of nodes is
    ///               maintained to prevent a malloc from always being
    ///               necessary. This bound is the limit on the size of the
    ///               cache (if desired). If the value is 0, then the cache has
    ///               no bound. Otherwise, the cache will never grow larger than
    ///               `bound` (although the queue itself could be much larger.
    pub unsafe fn new(bound: usize) -> Queue<T> {
        let n1 = Node::new();
        let n2 = Node::new();
        (*n1).next.store(n2, Ordering::Relaxed);
        Queue {
            tail: UnsafeCell::new(n2),
            tail_prev: AtomicPtr::new(n1),
            head: UnsafeCell::new(n2),
            first: UnsafeCell::new(n1),
            tail_copy: UnsafeCell::new(n1),
            cache_bound: bound,
            cache_additions: AtomicUsize::new(0),
            cache_subtractions: AtomicUsize::new(0),
        }
    }

    /// Pushes a new value onto this queue. Note that to use this function
    /// safely, it must be externally guaranteed that there is only one pusher.
    pub fn push(&self, t: T) {
        unsafe {
            // Acquire a node (which either uses a cached one or allocates a new
            // one), and then append this to the 'head' node.
            let n = self.alloc();
            assert!((*n).value.is_none());
            (*n).value = Some(t);
            (*n).next.store(ptr::null_mut(), Ordering::Relaxed);
            (**self.head.get()).next.store(n, Ordering::Release);
            *self.head.get() = n;
        }
    }

    unsafe fn alloc(&self) -> *mut Node<T> {
        // First try to see if we can consume the 'first' node for our uses.
        // We try to avoid as many atomic instructions as possible here, so
        // the addition to cache_subtractions is not atomic (plus we're the
        // only one subtracting from the cache).
        if *self.first.get() != *self.tail_copy.get() {
            if self.cache_bound > 0 {
                let b = self.cache_subtractions.load(Ordering::Relaxed);
                self.cache_subtractions.store(b + 1, Ordering::Relaxed);
            }
            let ret = *self.first.get();
            *self.first.get() = (*ret).next.load(Ordering::Relaxed);
            return ret;
        }
        // If the above fails, then update our copy of the tail and try
        // again.
        *self.tail_copy.get() = self.tail_prev.load(Ordering::Acquire);
        if *self.first.get() != *self.tail_copy.get() {
            if self.cache_bound > 0 {
                let b = self.cache_subtractions.load(Ordering::Relaxed);
                self.cache_subtractions.store(b + 1, Ordering::Relaxed);
            }
            let ret = *self.first.get();
            *self.first.get() = (*ret).next.load(Ordering::Relaxed);
            return ret;
        }
        // If all of that fails, then we have to allocate a new node
        // (there's nothing in the node cache).
        Node::new()
    }

    /// Attempts to pop a value from this queue. Remember that to use this type
    /// safely you must ensure that there is only one popper at a time.
    pub fn pop(&self) -> Option<T> {
        unsafe {
            // The `tail` node is not actually a used node, but rather a
            // sentinel from where we should start popping from. Hence, look at
            // tail's next field and see if we can use it. If we do a pop, then
            // the current tail node is a candidate for going into the cache.
            let tail = *self.tail.get();
            let next = (*tail).next.load(Ordering::Acquire);
            if next.is_null() { return None }
            assert!((*next).value.is_some());
            let ret = (*next).value.take();

            *self.tail.get() = next;
            if self.cache_bound == 0 {
                self.tail_prev.store(tail, Ordering::Release);
            } else {
                // FIXME: this is dubious with overflow.
                let additions = self.cache_additions.load(Ordering::Relaxed);
                let subtractions = self.cache_subtractions.load(Ordering::Relaxed);
                let size = additions - subtractions;

                if size < self.cache_bound {
                    self.tail_prev.store(tail, Ordering::Release);
                    self.cache_additions.store(additions + 1, Ordering::Relaxed);
                } else {
                    (*self.tail_prev.load(Ordering::Relaxed))
                          .next.store(next, Ordering::Relaxed);
                    // We have successfully erased all references to 'tail', so
                    // now we can safely drop it.
                    let _: Box<Node<T>> = Box::from_raw(tail);
                }
            }
            return ret;
        }
    }

    /// Attempts to peek at the head of the queue, returning `None` if the queue
    /// has no data currently
    ///
    /// # Warning
    /// The reference returned is invalid if it is not used before the consumer
    /// pops the value off the queue. If the producer then pushes another value
    /// onto the queue, it will overwrite the value pointed to by the reference.
    pub fn peek<'a>(&'a self) -> Option<&'a mut T> {
        // This is essentially the same as above with all the popping bits
        // stripped out.
        unsafe {
            let tail = *self.tail.get();
            let next = (*tail).next.load(Ordering::Acquire);
            if next.is_null() { return None }
            return (*next).value.as_mut();
        }
    }
}

impl<T> Drop for Queue<T> {
    fn drop(&mut self) {
        unsafe {
            let mut cur = *self.first.get();
            while !cur.is_null() {
                let next = (*cur).next.load(Ordering::Relaxed);
                let _n: Box<Node<T>> = Box::from_raw(cur);
                cur = next;
            }
        }
    }
}

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

    use sync::Arc;
    use super::Queue;
    use thread;
    use sync::mpsc::channel;

    #[test]
    fn smoke() {
        unsafe {
            let queue = Queue::new(0);
            queue.push(1);
            queue.push(2);
            assert_eq!(queue.pop(), Some(1));
            assert_eq!(queue.pop(), Some(2));
            assert_eq!(queue.pop(), None);
            queue.push(3);
            queue.push(4);
            assert_eq!(queue.pop(), Some(3));
            assert_eq!(queue.pop(), Some(4));
            assert_eq!(queue.pop(), None);
        }
    }

    #[test]
    fn peek() {
        unsafe {
            let queue = Queue::new(0);
            queue.push(vec![1]);

            // Ensure the borrowchecker works
            match queue.peek() {
                Some(vec) => match &**vec {
                    // Note that `pop` is not allowed here due to borrow
                    [1] => {}
                    _ => return
                },
                None => unreachable!()
            }

            queue.pop();
        }
    }

    #[test]
    fn drop_full() {
        unsafe {
            let q: Queue<Box<_>> = Queue::new(0);
            q.push(box 1);
            q.push(box 2);
        }
    }

    #[test]
    fn smoke_bound() {
        unsafe {
            let q = Queue::new(0);
            q.push(1);
            q.push(2);
            assert_eq!(q.pop(), Some(1));
            assert_eq!(q.pop(), Some(2));
            assert_eq!(q.pop(), None);
            q.push(3);
            q.push(4);
            assert_eq!(q.pop(), Some(3));
            assert_eq!(q.pop(), Some(4));
            assert_eq!(q.pop(), None);
        }
    }

    #[test]
    fn stress() {
        unsafe {
            stress_bound(0);
            stress_bound(1);
        }

        unsafe fn stress_bound(bound: usize) {
            let q = Arc::new(Queue::new(bound));

            let (tx, rx) = channel();
            let q2 = q.clone();
            let _t = thread::spawn(move|| {
                for _ in 0..100000 {
                    loop {
                        match q2.pop() {
                            Some(1) => break,
                            Some(_) => panic!(),
                            None => {}
                        }
                    }
                }
                tx.send(()).unwrap();
            });
            for _ in 0..100000 {
                q.push(1);
            }
            rx.recv().unwrap();
        }
    }
}
<|endoftext|>