<commit_before>// boolean_expression: a Rust crate for Boolean expressions and BDDs.
//
// Copyright (c) 2016 Chris Fallin <cfallin@c1f.net>. Released under the MIT
// License.
//

use std::collections::HashMap;
use std::collections::hash_map::Entry as HashEntry;
use std::cmp;
use std::fmt::Debug;
use std::hash::Hash;
use std::usize;

use Expr;

/// A `BDDFunc` is a function index within a particular `BDD` index. It must
/// only be used with the `BDD` instance which produced it.
pub type BDDFunc = usize;

/// A special terminal `BDDFunc` which is constant `false` (zero).
pub const BDD_ZERO: BDDFunc = usize::MAX;
/// A special terminal `BDDFunc` which is constant `true` (one).
pub const BDD_ONE: BDDFunc = usize::MAX - 1;

type BDDLabel = usize;

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct BDDNode {
    label: BDDLabel,
    lo: BDDFunc,
    hi: BDDFunc,
}

#[derive(Clone, Debug)]
struct LabelBDD {
    nodes: Vec<BDDNode>,
    dedup_hash: HashMap<BDDNode, BDDFunc>,
}

impl LabelBDD {
    pub fn new() -> LabelBDD {
        LabelBDD {
            nodes: Vec::new(),
            dedup_hash: HashMap::new(),
        }
    }

    fn get_node(&mut self, label: BDDLabel, lo: BDDFunc, hi: BDDFunc) -> BDDFunc {
        let n = BDDNode {
            label: label,
            lo: lo,
            hi: hi,
        };
        match self.dedup_hash.entry(n.clone()) {
            HashEntry::Occupied(o) => *o.get(),
            HashEntry::Vacant(v) => {
                let idx = self.nodes.len() as BDDFunc;
                self.nodes.push(n);
                v.insert(idx);
                idx
            }
        }
    }

    pub fn terminal(&mut self, label: BDDLabel) -> BDDFunc {
        self.get_node(label, BDD_ZERO, BDD_ONE)
    }

    pub fn constant(&mut self, value: bool) -> BDDFunc {
        if value {
            BDD_ONE
        } else {
            BDD_ZERO
        }
    }

    /// Restrict: fundamental building block of logical combinators. Takes a
    /// Shannon cofactor: i.e., returns a new function based on `f` but with the
    /// given label forced to the given value.
    pub fn restrict(&mut self, f: BDDFunc, label: BDDLabel, val: bool) -> BDDFunc {
        if f == BDD_ZERO {
            return BDD_ZERO;
        }
        if f == BDD_ONE {
            return BDD_ONE;
        }

        let node = self.nodes[f].clone();
        if label < node.label {
            f
        } else if label == node.label {
            if val {
                node.hi
            } else {
                node.lo
            }
        } else {
            let lo = self.restrict(node.lo, label, val);
            let hi = self.restrict(node.hi, label, val);
            self.get_node(node.label, lo, hi)
        }
    }

    fn min_label(&self, f: BDDFunc) -> Option<BDDLabel> {
        if f == BDD_ZERO || f == BDD_ONE {
            None
        } else {
            Some(self.nodes[f].label)
        }
    }

    /// If-then-else: fundamental building block of logical combinators. Works
    /// by divide-and-conquer: split on the lowest appearing label, take Shannon
    /// cofactors for the two cases, recurse, and recombine with a new node.
    pub fn ite(&mut self, i: BDDFunc, t: BDDFunc, e: BDDFunc) -> BDDFunc {
        if i == BDD_ONE {
            t
        } else if i == BDD_ZERO {
            e
        } else if t == e {
            t
        } else if t == BDD_ONE && e == BDD_ZERO {
            i
        } else {
            let i_var = self.min_label(i).unwrap_or(usize::MAX);
            let t_var = self.min_label(t).unwrap_or(usize::MAX);
            let e_var = self.min_label(e).unwrap_or(usize::MAX);
            let split = cmp::min(i_var, cmp::min(t_var, e_var));
            assert!(split != usize::MAX);
            let i_lo = self.restrict(i, split, false);
            let t_lo = self.restrict(t, split, false);
            let e_lo = self.restrict(e, split, false);
            let i_hi = self.restrict(i, split, true);
            let t_hi = self.restrict(t, split, true);
            let e_hi = self.restrict(e, split, true);
            let lo = self.ite(i_lo, t_lo, e_lo);
            let hi = self.ite(i_hi, t_hi, e_hi);
            self.get_node(split, lo, hi)
        }
    }

    pub fn not(&mut self, n: BDDFunc) -> BDDFunc {
        self.ite(n, BDD_ZERO, BDD_ONE)
    }

    pub fn and(&mut self, a: BDDFunc, b: BDDFunc) -> BDDFunc {
        self.ite(a, b, BDD_ZERO)
    }

    pub fn or(&mut self, a: BDDFunc, b: BDDFunc) -> BDDFunc {
        self.ite(a, BDD_ONE, b)
    }

    pub fn evaluate(&self, func: BDDFunc, inputs: &[bool]) -> Option<bool> {
        let mut f = func;
        for (i, val) in inputs.iter().enumerate() {
            if f == BDD_ZERO || f == BDD_ONE {
                break;
            }
            let node = &self.nodes[f];
            if node.label > i {
                continue;
            } else if node.label == i {
                f = if *val {
                    node.hi
                } else {
                    node.lo
                };
            }
        }
        match f {
            BDD_ZERO => Some(false),
            BDD_ONE => Some(true),
            _ => None,
        }
    }
}

/// A `BDD` is a Binary Decision Diagram, an efficient way to represent a
/// Boolean function in a canonical way. (It is actually a "Reduced Ordered
/// Binary Decision Diagram", which gives it its canonicity assuming terminals
/// are ordered consistently.)
///
/// A BDD is built up from terminals (free variables) and constants, combined
/// with the logical combinators AND, OR, and NOT. It may be evaluated with
/// certain terminal assignments.
///
/// The major advantage of a BDD is that its logical operations are performed,
/// it will "self-simplify": i.e., taking the OR of `And(a, b)` and `And(a,
/// Not(b))` will produce `a` without any further simplification step. Furthermore,
/// the `BDDFunc` representing this value is canonical: if two different
/// expressions are produced within the same BDD and they both result in
/// (simplify down to) `a`, then the `BDDFunc` values will be equal. The
/// tradeoff is that logical operations may be expensive: they are linear in
/// BDD size, but BDDs may have exponential size (relative to terminal count)
/// in the worst case.
#[derive(Clone, Debug)]
pub struct BDD<T>
    where T: Clone + Debug + Eq + Ord + Hash
{
    bdd: LabelBDD,
    labels: HashMap<T, BDDLabel>,
    next_label: BDDLabel,
}

impl<T> BDD<T>
    where T: Clone + Debug + Eq + Ord + Hash
{
    /// Produce a new, empty, BDD.
    pub fn new() -> BDD<T> {
        BDD {
            bdd: LabelBDD::new(),
            labels: HashMap::new(),
            next_label: 0,
        }
    }

    fn label(&mut self, t: T) -> BDDLabel {
        match self.labels.entry(t.clone()) {
            HashEntry::Occupied(o) => *o.get(),
            HashEntry::Vacant(v) => {
                let next_id = self.next_label;
                self.next_label += 1;
                v.insert(next_id);
                next_id
            }
        }
    }

    /// Produce a function within the BDD representing the terminal `t`. If
    /// this terminal has been used in the BDD before, the same `BDDFunc` will be
    /// returned.
    pub fn terminal(&mut self, t: T) -> BDDFunc {
        let l = self.label(t);
        self.bdd.terminal(l)
    }

    /// Produce a function within the BDD representing the constant value `val`.
    pub fn constant(&mut self, val: bool) -> BDDFunc {
        self.bdd.constant(val)
    }

    /// Produce a function within the BDD representing the logical complement
    /// of the function `n`.
    pub fn not(&mut self, n: BDDFunc) -> BDDFunc {
        self.bdd.not(n)
    }

    /// Produce a function within the BDD representing the logical AND of the
    /// functions `a` and `b`.
    pub fn and(&mut self, a: BDDFunc, b: BDDFunc) -> BDDFunc {
        self.bdd.and(a, b)
    }

    /// Produce a function within the BDD representing the logical OR of the
    /// functions `a` and `b`.
    pub fn or(&mut self, a: BDDFunc, b: BDDFunc) -> BDDFunc {
        self.bdd.or(a, b)
    }

    /// Produce a function within the BDD representing the given expression
    /// `e`, which may contain ANDs, ORs, NOTs, terminals, and constants.
    pub fn expr(&mut self, e: &Expr<T>) -> BDDFunc {
        match e {
            &Expr::Terminal(ref t) => self.terminal(t.clone()),
            &Expr::Const(val) => self.constant(val),
            &Expr::Not(ref x) => {
                let xval = self.expr(&**x);
                self.not(xval)
            }
            &Expr::And(ref a, ref b) => {
                let aval = self.expr(&**a);
                let bval = self.expr(&**b);
                self.and(aval, bval)
            }
            &Expr::Or(ref a, ref b) => {
                let aval = self.expr(&**a);
                let bval = self.expr(&**b);
                self.or(aval, bval)
            }
        }
    }

    /// Evaluate the function `f` in the BDD with the given terminal
    /// assignments. Any terminals not specified in `values` default to `false`.
    pub fn evaluate(&self, f: BDDFunc, values: &HashMap<T, bool>) -> bool {
        let size = self.next_label;
        let mut valarray = Vec::with_capacity(size);
        valarray.resize(size, false);
        for (t, l) in &self.labels {
            valarray[*l as usize] = *values.get(t).unwrap_or(&false);
        }
        self.bdd.evaluate(f, &valarray).unwrap()
    }
}

mod test {
    use super::*;
    use Expr;
    use std::collections::HashMap;
    extern crate rand;
    use self::rand::Rng;

    fn term_hashmap(vals: &[bool], h: &mut HashMap<u32, bool>) {
        h.clear();
        for (i, v) in vals.iter().enumerate() {
            h.insert(i as u32, *v);
        }
    }

    fn test_bdd(b: &BDD<u32>,
                f: BDDFunc,
                h: &mut HashMap<u32, bool>,
                inputs: &[bool],
                expected: bool) {
        term_hashmap(inputs, h);
        assert!(b.evaluate(f, h) == expected);
    }

    #[test]
    fn bdd_eval() {
        let mut h = HashMap::new();
        let mut b = BDD::new();
        let expr = Expr::or(Expr::and(Expr::Terminal(0), Expr::Terminal(1)),
                            Expr::and(Expr::not(Expr::Terminal(2)), Expr::not(Expr::Terminal(3))));
        let f = b.expr(&expr);
        test_bdd(&b, f, &mut h, &[false, false, true, true], false);
        test_bdd(&b, f, &mut h, &[true, false, true, true], false);
        test_bdd(&b, f, &mut h, &[true, true, true, true], true);
        test_bdd(&b, f, &mut h, &[false, false, false, true], false);
        test_bdd(&b, f, &mut h, &[false, false, false, false], true);
    }

    fn bits_to_hashmap(bits: usize, n: usize, h: &mut HashMap<u32, bool>) {
        for b in 0..bits {
            h.insert(b as u32, (n & (1 << b)) != 0);
        }
    }

    fn test_bdd_expr(e: Expr<u32>, nterminals: usize) {
        let mut b = BDD::new();
        let f = b.expr(&e);
        let mut terminal_values = HashMap::new();
        for v in 0..(1 << nterminals) {
            bits_to_hashmap(nterminals, v, &mut terminal_values);
            let expr_val = e.evaluate(&terminal_values);
            let bdd_val = b.evaluate(f, &terminal_values);
            assert!(expr_val == bdd_val);
        }
    }

    fn random_expr(r: &mut rand::XorShiftRng, nterminals: usize) -> Expr<u32> {
        match r.gen_range(0, 5) {
            0 => Expr::Terminal(r.gen_range(0, nterminals) as u32),
            1 => Expr::Const(r.gen_weighted_bool(2)),
            2 => Expr::Not(Box::new(random_expr(r, nterminals))),
            3 => {
                Expr::And(Box::new(random_expr(r, nterminals)),
                          Box::new(random_expr(r, nterminals)))
            }
            4 => {
                Expr::Or(Box::new(random_expr(r, nterminals)),
                         Box::new(random_expr(r, nterminals)))
            }
            _ => unreachable!(),
        }
    }

    #[test]
    fn bdd_exhaustive_exprs() {
        let mut rng: rand::XorShiftRng = rand::XorShiftRng::new_unseeded();
        for _ in 0..100 {
            let expr = random_expr(&mut rng, 6);
            test_bdd_expr(expr, 6);
        }
    }
}
<commit_msg>Minor optimization: do not create BDD node if hi/lo targets are the same.<commit_after>// boolean_expression: a Rust crate for Boolean expressions and BDDs.
//
// Copyright (c) 2016 Chris Fallin <cfallin@c1f.net>. Released under the MIT
// License.
//

use std::collections::HashMap;
use std::collections::hash_map::Entry as HashEntry;
use std::cmp;
use std::fmt::Debug;
use std::hash::Hash;
use std::usize;

use Expr;

/// A `BDDFunc` is a function index within a particular `BDD` index. It must
/// only be used with the `BDD` instance which produced it.
pub type BDDFunc = usize;

/// A special terminal `BDDFunc` which is constant `false` (zero).
pub const BDD_ZERO: BDDFunc = usize::MAX;
/// A special terminal `BDDFunc` which is constant `true` (one).
pub const BDD_ONE: BDDFunc = usize::MAX - 1;

type BDDLabel = usize;

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct BDDNode {
    label: BDDLabel,
    lo: BDDFunc,
    hi: BDDFunc,
}

#[derive(Clone, Debug)]
struct LabelBDD {
    nodes: Vec<BDDNode>,
    dedup_hash: HashMap<BDDNode, BDDFunc>,
}

impl LabelBDD {
    pub fn new() -> LabelBDD {
        LabelBDD {
            nodes: Vec::new(),
            dedup_hash: HashMap::new(),
        }
    }

    fn get_node(&mut self, label: BDDLabel, lo: BDDFunc, hi: BDDFunc) -> BDDFunc {
        if lo == hi {
            return lo;
        }
        let n = BDDNode {
            label: label,
            lo: lo,
            hi: hi,
        };
        match self.dedup_hash.entry(n.clone()) {
            HashEntry::Occupied(o) => *o.get(),
            HashEntry::Vacant(v) => {
                let idx = self.nodes.len() as BDDFunc;
                self.nodes.push(n);
                v.insert(idx);
                idx
            }
        }
    }

    pub fn terminal(&mut self, label: BDDLabel) -> BDDFunc {
        self.get_node(label, BDD_ZERO, BDD_ONE)
    }

    pub fn constant(&mut self, value: bool) -> BDDFunc {
        if value {
            BDD_ONE
        } else {
            BDD_ZERO
        }
    }

    /// Restrict: fundamental building block of logical combinators. Takes a
    /// Shannon cofactor: i.e., returns a new function based on `f` but with the
    /// given label forced to the given value.
    pub fn restrict(&mut self, f: BDDFunc, label: BDDLabel, val: bool) -> BDDFunc {
        if f == BDD_ZERO {
            return BDD_ZERO;
        }
        if f == BDD_ONE {
            return BDD_ONE;
        }

        let node = self.nodes[f].clone();
        if label < node.label {
            f
        } else if label == node.label {
            if val {
                node.hi
            } else {
                node.lo
            }
        } else {
            let lo = self.restrict(node.lo, label, val);
            let hi = self.restrict(node.hi, label, val);
            self.get_node(node.label, lo, hi)
        }
    }

    fn min_label(&self, f: BDDFunc) -> Option<BDDLabel> {
        if f == BDD_ZERO || f == BDD_ONE {
            None
        } else {
            Some(self.nodes[f].label)
        }
    }

    /// If-then-else: fundamental building block of logical combinators. Works
    /// by divide-and-conquer: split on the lowest appearing label, take Shannon
    /// cofactors for the two cases, recurse, and recombine with a new node.
    pub fn ite(&mut self, i: BDDFunc, t: BDDFunc, e: BDDFunc) -> BDDFunc {
        if i == BDD_ONE {
            t
        } else if i == BDD_ZERO {
            e
        } else if t == e {
            t
        } else if t == BDD_ONE && e == BDD_ZERO {
            i
        } else {
            let i_var = self.min_label(i).unwrap_or(usize::MAX);
            let t_var = self.min_label(t).unwrap_or(usize::MAX);
            let e_var = self.min_label(e).unwrap_or(usize::MAX);
            let split = cmp::min(i_var, cmp::min(t_var, e_var));
            assert!(split != usize::MAX);
            let i_lo = self.restrict(i, split, false);
            let t_lo = self.restrict(t, split, false);
            let e_lo = self.restrict(e, split, false);
            let i_hi = self.restrict(i, split, true);
            let t_hi = self.restrict(t, split, true);
            let e_hi = self.restrict(e, split, true);
            let lo = self.ite(i_lo, t_lo, e_lo);
            let hi = self.ite(i_hi, t_hi, e_hi);
            self.get_node(split, lo, hi)
        }
    }

    pub fn not(&mut self, n: BDDFunc) -> BDDFunc {
        self.ite(n, BDD_ZERO, BDD_ONE)
    }

    pub fn and(&mut self, a: BDDFunc, b: BDDFunc) -> BDDFunc {
        self.ite(a, b, BDD_ZERO)
    }

    pub fn or(&mut self, a: BDDFunc, b: BDDFunc) -> BDDFunc {
        self.ite(a, BDD_ONE, b)
    }

    pub fn evaluate(&self, func: BDDFunc, inputs: &[bool]) -> Option<bool> {
        let mut f = func;
        for (i, val) in inputs.iter().enumerate() {
            if f == BDD_ZERO || f == BDD_ONE {
                break;
            }
            let node = &self.nodes[f];
            if node.label > i {
                continue;
            } else if node.label == i {
                f = if *val {
                    node.hi
                } else {
                    node.lo
                };
            }
        }
        match f {
            BDD_ZERO => Some(false),
            BDD_ONE => Some(true),
            _ => None,
        }
    }
}

/// A `BDD` is a Binary Decision Diagram, an efficient way to represent a
/// Boolean function in a canonical way. (It is actually a "Reduced Ordered
/// Binary Decision Diagram", which gives it its canonicity assuming terminals
/// are ordered consistently.)
///
/// A BDD is built up from terminals (free variables) and constants, combined
/// with the logical combinators AND, OR, and NOT. It may be evaluated with
/// certain terminal assignments.
///
/// The major advantage of a BDD is that its logical operations are performed,
/// it will "self-simplify": i.e., taking the OR of `And(a, b)` and `And(a,
/// Not(b))` will produce `a` without any further simplification step. Furthermore,
/// the `BDDFunc` representing this value is canonical: if two different
/// expressions are produced within the same BDD and they both result in
/// (simplify down to) `a`, then the `BDDFunc` values will be equal. The
/// tradeoff is that logical operations may be expensive: they are linear in
/// BDD size, but BDDs may have exponential size (relative to terminal count)
/// in the worst case.
#[derive(Clone, Debug)]
pub struct BDD<T>
    where T: Clone + Debug + Eq + Ord + Hash
{
    bdd: LabelBDD,
    labels: HashMap<T, BDDLabel>,
    next_label: BDDLabel,
}

impl<T> BDD<T>
    where T: Clone + Debug + Eq + Ord + Hash
{
    /// Produce a new, empty, BDD.
    pub fn new() -> BDD<T> {
        BDD {
            bdd: LabelBDD::new(),
            labels: HashMap::new(),
            next_label: 0,
        }
    }

    fn label(&mut self, t: T) -> BDDLabel {
        match self.labels.entry(t.clone()) {
            HashEntry::Occupied(o) => *o.get(),
            HashEntry::Vacant(v) => {
                let next_id = self.next_label;
                self.next_label += 1;
                v.insert(next_id);
                next_id
            }
        }
    }

    /// Produce a function within the BDD representing the terminal `t`. If
    /// this terminal has been used in the BDD before, the same `BDDFunc` will be
    /// returned.
    pub fn terminal(&mut self, t: T) -> BDDFunc {
        let l = self.label(t);
        self.bdd.terminal(l)
    }

    /// Produce a function within the BDD representing the constant value `val`.
    pub fn constant(&mut self, val: bool) -> BDDFunc {
        self.bdd.constant(val)
    }

    /// Produce a function within the BDD representing the logical complement
    /// of the function `n`.
    pub fn not(&mut self, n: BDDFunc) -> BDDFunc {
        self.bdd.not(n)
    }

    /// Produce a function within the BDD representing the logical AND of the
    /// functions `a` and `b`.
    pub fn and(&mut self, a: BDDFunc, b: BDDFunc) -> BDDFunc {
        self.bdd.and(a, b)
    }

    /// Produce a function within the BDD representing the logical OR of the
    /// functions `a` and `b`.
    pub fn or(&mut self, a: BDDFunc, b: BDDFunc) -> BDDFunc {
        self.bdd.or(a, b)
    }

    /// Produce a function within the BDD representing the given expression
    /// `e`, which may contain ANDs, ORs, NOTs, terminals, and constants.
    pub fn expr(&mut self, e: &Expr<T>) -> BDDFunc {
        match e {
            &Expr::Terminal(ref t) => self.terminal(t.clone()),
            &Expr::Const(val) => self.constant(val),
            &Expr::Not(ref x) => {
                let xval = self.expr(&**x);
                self.not(xval)
            }
            &Expr::And(ref a, ref b) => {
                let aval = self.expr(&**a);
                let bval = self.expr(&**b);
                self.and(aval, bval)
            }
            &Expr::Or(ref a, ref b) => {
                let aval = self.expr(&**a);
                let bval = self.expr(&**b);
                self.or(aval, bval)
            }
        }
    }

    /// Evaluate the function `f` in the BDD with the given terminal
    /// assignments. Any terminals not specified in `values` default to `false`.
    pub fn evaluate(&self, f: BDDFunc, values: &HashMap<T, bool>) -> bool {
        let size = self.next_label;
        let mut valarray = Vec::with_capacity(size);
        valarray.resize(size, false);
        for (t, l) in &self.labels {
            valarray[*l as usize] = *values.get(t).unwrap_or(&false);
        }
        self.bdd.evaluate(f, &valarray).unwrap()
    }
}

mod test {
    use super::*;
    use Expr;
    use std::collections::HashMap;
    extern crate rand;
    use self::rand::Rng;

    fn term_hashmap(vals: &[bool], h: &mut HashMap<u32, bool>) {
        h.clear();
        for (i, v) in vals.iter().enumerate() {
            h.insert(i as u32, *v);
        }
    }

    fn test_bdd(b: &BDD<u32>,
                f: BDDFunc,
                h: &mut HashMap<u32, bool>,
                inputs: &[bool],
                expected: bool) {
        term_hashmap(inputs, h);
        assert!(b.evaluate(f, h) == expected);
    }

    #[test]
    fn bdd_eval() {
        let mut h = HashMap::new();
        let mut b = BDD::new();
        let expr = Expr::or(Expr::and(Expr::Terminal(0), Expr::Terminal(1)),
                            Expr::and(Expr::not(Expr::Terminal(2)), Expr::not(Expr::Terminal(3))));
        let f = b.expr(&expr);
        test_bdd(&b, f, &mut h, &[false, false, true, true], false);
        test_bdd(&b, f, &mut h, &[true, false, true, true], false);
        test_bdd(&b, f, &mut h, &[true, true, true, true], true);
        test_bdd(&b, f, &mut h, &[false, false, false, true], false);
        test_bdd(&b, f, &mut h, &[false, false, false, false], true);
    }

    fn bits_to_hashmap(bits: usize, n: usize, h: &mut HashMap<u32, bool>) {
        for b in 0..bits {
            h.insert(b as u32, (n & (1 << b)) != 0);
        }
    }

    fn test_bdd_expr(e: Expr<u32>, nterminals: usize) {
        let mut b = BDD::new();
        let f = b.expr(&e);
        let mut terminal_values = HashMap::new();
        for v in 0..(1 << nterminals) {
            bits_to_hashmap(nterminals, v, &mut terminal_values);
            let expr_val = e.evaluate(&terminal_values);
            let bdd_val = b.evaluate(f, &terminal_values);
            assert!(expr_val == bdd_val);
        }
    }

    fn random_expr(r: &mut rand::XorShiftRng, nterminals: usize) -> Expr<u32> {
        match r.gen_range(0, 5) {
            0 => Expr::Terminal(r.gen_range(0, nterminals) as u32),
            1 => Expr::Const(r.gen_weighted_bool(2)),
            2 => Expr::Not(Box::new(random_expr(r, nterminals))),
            3 => {
                Expr::And(Box::new(random_expr(r, nterminals)),
                          Box::new(random_expr(r, nterminals)))
            }
            4 => {
                Expr::Or(Box::new(random_expr(r, nterminals)),
                         Box::new(random_expr(r, nterminals)))
            }
            _ => unreachable!(),
        }
    }

    #[test]
    fn bdd_exhaustive_exprs() {
        let mut rng: rand::XorShiftRng = rand::XorShiftRng::new_unseeded();
        for _ in 0..100 {
            let expr = random_expr(&mut rng, 6);
            test_bdd_expr(expr, 6);
        }
    }
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate dent;
extern crate term;
extern crate term_size;

use clap::{App, Arg};
use dent::plot;
use dent::summary::Summary;
use dent::t_test::{SigLevel, TTest, welch_t_test};

use std::error;
use std::fs::File;
use std::path::Path;
use std::io::{self, BufRead, BufReader};

mod fmt;
mod log;

macro_rules! ok {
    ($r: expr) => {
        match $r {
            Ok(t) => t,
            Err(e) => {
                log::error(&format!("{}", e));
                std::process::exit(1);
            }
        }
    }
}

fn print_summary(s: &Summary) {
    let width = 10;
    let size_width = 6;

    println!(
        "{n:>nw$}  {min:>w$}  {max:>w$}  {med:>w$}  {mean:>w$}  {sd:>w$}  {se:>w$}",
        w = width,
        nw = size_width,
        n = "N",
        min = "Min",
        max = "Max",
        med = "Median",
        mean = "Mean",
        sd= "StdDev",
        se = "StdErr",
    );
    println!(
        "{n:>nw$}  {min:>w$}  {max:>w$}  {med:>w$}  {mean:>w$}  {sd:>w$}  {se:>w$}",
        w = width,
        nw = size_width,
        n = fmt::f(s.size(), width),
        min = fmt::f(s.min(), width),
        max = fmt::f(s.max(), width),
        med = fmt::f(s.median(), width),
        mean = fmt::f(s.mean(), width),
        sd = fmt::f(s.standard_deviation(), width),
        se = fmt::f(s.standard_error(), width),
    );
}

fn print_t_test(t_test: &TTest) {
    let width = 12;
    let reject = if t_test.reject { "yes" } else { "no" };

    println!("{l:>w$} = {v}", w = width, l = "Reject H₀?", v = reject);
    println!("{l:>w$} = {v}", w = width, l = "α", v = t_test.alpha);
    println!("{l:>w$} = {v}", w = width, l = "t", v = t_test.t);
    println!("{l:>w$} = {v}", w = width, l = "Crit. v", v = t_test.crit);
    println!("{l:>w$} = {v}", w = width, l = "DF", v = t_test.df);
}

fn summarize_file(path: &str, lax_parsing: bool) -> Result<Summary, Box<error::Error>> {
    let p = Path::new(path);
    let f = File::open(p)?;
    let reader = BufReader::new(f);

    let data = read_data(reader, lax_parsing)?;

    Ok(Summary::new(&data).unwrap())
}

fn read_data<R>(reader: R, lax_parsing: bool) -> Result<Vec<f64>, Box<error::Error>>
    where R: BufRead {
    let mut data: Vec<f64> = vec![];

    for l in reader.lines() {
        let s = l.unwrap().trim().to_string();

        if s.is_empty() {
            continue;
        }

        match s.parse() {
            Ok(d) => data.push(d),
            err => if !lax_parsing { err?; }
        }
    }

    Ok(data)
}

fn parse_alpha(arg: &str) -> SigLevel {
    match arg {
        ".001" => SigLevel::Alpha001,
        ".005" => SigLevel::Alpha005,
        ".01"  => SigLevel::Alpha010,
        ".025" => SigLevel::Alpha025,
        ".05"  => SigLevel::Alpha050,
        ".1"   => SigLevel::Alpha100,
        _ => panic!(),
    }
}

fn summarize_stdin(lax_parsing: bool) -> Result<Summary, Box<error::Error>> {
    let stdin = io::stdin();
    let data = read_data(stdin.lock(), lax_parsing)?;

    Ok(Summary::new(&data).unwrap())
}

fn display_t_test(
    summary1: &Summary,
    summary2: &Summary,
    alpha: SigLevel,
    draw_plot: bool,
    width: usize,
    ascii: bool,
) {
    let t_test = welch_t_test(&summary1, &summary2, alpha);

    if draw_plot {
        let p = plot::comparison_plot(&[summary1, summary2], width, ascii, true);
        println!("{}\n", p);
    }

    print_summary(&summary1);
    println!();
    print_summary(&summary2);
    println!();
    print_t_test(&t_test);
}

fn display_summaries(
    summaries: &[Summary],
    draw_plot: bool,
    width: usize,
    ascii: bool,
) {
    if draw_plot {
        let summary_refs: Vec<&Summary> = summaries
            .iter()
            .collect();

        let plot = plot::comparison_plot(&summary_refs, width, ascii, true);
        println!("{}\n", plot);
    }

    for i in 0..summaries.len() {
        if i > 0 {
            println!();
        }
        print_summary(&summaries[i]);
    }
}

fn main() {
    let matches = App::new("dent")
        .version("0.3.0")
        .author("Joe Ranweiler <joe@lemma.co>")
        .about("A tiny tool for t-tests &c.")
        .arg(Arg::with_name("stdin")
             .short("s")
             .long("stdin")
             .help("Read and summarize data from stdin"))
        .arg(Arg::with_name("files")
             .multiple(true)
             .value_name("FILES")
             .takes_value(true)
             .required_unless("stdin")
             .help("Path to one or more files of sample data"))
        .arg(Arg::with_name("alpha")
             .short("a")
             .long("alpha")
             .value_name("ALPHA")
             .help("Significance level α")
             .takes_value(true)
             .default_value(".05"))
        .arg(Arg::with_name("lax")
             .long("lax")
             .help("Ignore non-numeric input lines"))
        .arg(Arg::with_name("plot")
             .short("p")
             .long("plot")
             .help("Print standard boxplots"))
        .arg(Arg::with_name("ascii")
             .long("ascii")
             .help("Use only ASCII characters in boxplots"))
        .arg(Arg::with_name("width")
             .short("w")
             .long("width")
             .value_name("WIDTH")
             .takes_value(true)
             .help("Width of boxplot"))
        .get_matches();

    let ascii = matches.is_present("ascii");
    let lax_parsing = matches.is_present("lax");
    let draw_plot = matches.is_present("plot");
    let use_stdin = matches.is_present("stdin");

    let width = matches
        .value_of("width")
        .and_then(|w| w.parse::<usize>().ok())
        .or(term_size::dimensions().map(|(w, _)| w))
        .unwrap_or(80);

    let summaries = if use_stdin {
        vec![ok!(summarize_stdin(lax_parsing))]
    } else {
        matches.values_of("files")
            .unwrap()
            .map(|f| ok!(summarize_file(f, lax_parsing)))
            .collect()
    };

    match summaries.len() {
        0 => unreachable!(),
        2 => {
            let alpha = parse_alpha(matches.value_of("alpha").unwrap());

            display_t_test(
                &summaries[0],
                &summaries[1],
                alpha,
                draw_plot,
                width,
                ascii,
            );
        }
        _ => {
            display_summaries(&summaries, draw_plot, width, ascii);
        },
    };
}
<commit_msg>Remove panicking unwraps from exe<commit_after>extern crate clap;
extern crate dent;
extern crate term;
extern crate term_size;

use clap::{App, Arg};
use dent::plot;
use dent::summary::Summary;
use dent::t_test::{SigLevel, TTest, welch_t_test};

use std::error;
use std::fs::File;
use std::path::Path;
use std::io::{self, BufRead, BufReader};

mod fmt;
mod log;

macro_rules! ok {
    ($r: expr) => {
        match $r {
            Ok(t) => t,
            Err(e) => {
                log::error(&format!("{}", e));
                std::process::exit(1);
            }
        }
    }
}

fn print_summary(s: &Summary) {
    let width = 10;
    let size_width = 6;

    println!(
        "{n:>nw$}  {min:>w$}  {max:>w$}  {med:>w$}  {mean:>w$}  {sd:>w$}  {se:>w$}",
        w = width,
        nw = size_width,
        n = "N",
        min = "Min",
        max = "Max",
        med = "Median",
        mean = "Mean",
        sd= "StdDev",
        se = "StdErr",
    );
    println!(
        "{n:>nw$}  {min:>w$}  {max:>w$}  {med:>w$}  {mean:>w$}  {sd:>w$}  {se:>w$}",
        w = width,
        nw = size_width,
        n = fmt::f(s.size(), width),
        min = fmt::f(s.min(), width),
        max = fmt::f(s.max(), width),
        med = fmt::f(s.median(), width),
        mean = fmt::f(s.mean(), width),
        sd = fmt::f(s.standard_deviation(), width),
        se = fmt::f(s.standard_error(), width),
    );
}

fn print_t_test(t_test: &TTest) {
    let width = 12;
    let reject = if t_test.reject { "yes" } else { "no" };

    println!("{l:>w$} = {v}", w = width, l = "Reject H₀?", v = reject);
    println!("{l:>w$} = {v}", w = width, l = "α", v = t_test.alpha);
    println!("{l:>w$} = {v}", w = width, l = "t", v = t_test.t);
    println!("{l:>w$} = {v}", w = width, l = "Crit. v", v = t_test.crit);
    println!("{l:>w$} = {v}", w = width, l = "DF", v = t_test.df);
}

fn summarize_file(path: &str, lax_parsing: bool) -> Result<Summary, Box<error::Error>> {
    let p = Path::new(path);
    let f = File::open(p)?;
    let reader = BufReader::new(f);

    let data = read_data(reader, lax_parsing)?;

    Ok(Summary::new(&data)?)
}

fn read_data<R>(reader: R, lax_parsing: bool) -> Result<Vec<f64>, Box<error::Error>>
    where R: BufRead {
    let mut data: Vec<f64> = vec![];

    for l in reader.lines() {
        let s = l?.trim().to_string();

        if s.is_empty() {
            continue;
        }

        match s.parse() {
            Ok(d) => data.push(d),
            err => if !lax_parsing { err?; }
        }
    }

    Ok(data)
}

fn parse_alpha(arg: &str) -> SigLevel {
    match arg {
        ".001" => SigLevel::Alpha001,
        ".005" => SigLevel::Alpha005,
        ".01"  => SigLevel::Alpha010,
        ".025" => SigLevel::Alpha025,
        ".05"  => SigLevel::Alpha050,
        ".1"   => SigLevel::Alpha100,
        _ => panic!(),
    }
}

fn summarize_stdin(lax_parsing: bool) -> Result<Summary, Box<error::Error>> {
    let stdin = io::stdin();
    let data = read_data(stdin.lock(), lax_parsing)?;

    Ok(Summary::new(&data)?)
}

fn display_t_test(
    summary1: &Summary,
    summary2: &Summary,
    alpha: SigLevel,
    draw_plot: bool,
    width: usize,
    ascii: bool,
) {
    let t_test = welch_t_test(&summary1, &summary2, alpha);

    if draw_plot {
        let p = plot::comparison_plot(&[summary1, summary2], width, ascii, true);
        println!("{}\n", p);
    }

    print_summary(&summary1);
    println!();
    print_summary(&summary2);
    println!();
    print_t_test(&t_test);
}

fn display_summaries(
    summaries: &[Summary],
    draw_plot: bool,
    width: usize,
    ascii: bool,
) {
    if draw_plot {
        let summary_refs: Vec<&Summary> = summaries
            .iter()
            .collect();

        let plot = plot::comparison_plot(&summary_refs, width, ascii, true);
        println!("{}\n", plot);
    }

    for i in 0..summaries.len() {
        if i > 0 {
            println!();
        }
        print_summary(&summaries[i]);
    }
}

fn main() {
    let matches = App::new("dent")
        .version("0.3.0")
        .author("Joe Ranweiler <joe@lemma.co>")
        .about("A tiny tool for t-tests &c.")
        .arg(Arg::with_name("stdin")
             .short("s")
             .long("stdin")
             .help("Read and summarize data from stdin"))
        .arg(Arg::with_name("files")
             .multiple(true)
             .value_name("FILES")
             .takes_value(true)
             .required_unless("stdin")
             .help("Path to one or more files of sample data"))
        .arg(Arg::with_name("alpha")
             .short("a")
             .long("alpha")
             .value_name("ALPHA")
             .help("Significance level α")
             .takes_value(true)
             .default_value(".05"))
        .arg(Arg::with_name("lax")
             .long("lax")
             .help("Ignore non-numeric input lines"))
        .arg(Arg::with_name("plot")
             .short("p")
             .long("plot")
             .help("Print standard boxplots"))
        .arg(Arg::with_name("ascii")
             .long("ascii")
             .help("Use only ASCII characters in boxplots"))
        .arg(Arg::with_name("width")
             .short("w")
             .long("width")
             .value_name("WIDTH")
             .takes_value(true)
             .help("Width of boxplot"))
        .get_matches();

    let ascii = matches.is_present("ascii");
    let lax_parsing = matches.is_present("lax");
    let draw_plot = matches.is_present("plot");
    let use_stdin = matches.is_present("stdin");

    let width = matches
        .value_of("width")
        .and_then(|w| w.parse::<usize>().ok())
        .or(term_size::dimensions().map(|(w, _)| w))
        .unwrap_or(80);

    let summaries = if use_stdin {
        vec![ok!(summarize_stdin(lax_parsing))]
    } else {
        // Required if `stdin` is not present, so we can unwrap.
        matches.values_of("files")
            .unwrap_or_else(|| unreachable!())
            .map(|f| ok!(summarize_file(f, lax_parsing)))
            .collect()
    };

    match summaries.len() {
        0 => unreachable!(),
        2 => {
            // Has a default value, so we can can unwrap.
            let alpha_arg = matches.value_of("alpha").unwrap_or_else(|| unreachable!());
            let alpha = parse_alpha(alpha_arg);

            display_t_test(
                &summaries[0],
                &summaries[1],
                alpha,
                draw_plot,
                width,
                ascii,
            );
        }
        _ => {
            display_summaries(&summaries, draw_plot, width, ascii);
        },
    };
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

mod utils;

use crate::handlers::utils::{
    handle_rejection, reply_with_async_channel_writer, reply_with_lcs_bytes,
    send_size_prefixed_lcs_bytes, unwrap_or_500,
};
use anyhow::Result;
use libra_crypto::hash::HashValue;
use libra_types::transaction::Version;
use libradb::backup::BackupHandler;
use warp::{filters::BoxedFilter, reply::Reply, Filter};

fn get_latest_state_root(backup_handler: &BackupHandler) -> Result<Box<dyn Reply>> {
    reply_with_lcs_bytes(&backup_handler.get_latest_state_root()?)
}

fn get_state_range_proof(
    backup_handler: &BackupHandler,
    version: Version,
    end_key: HashValue,
) -> Result<Box<dyn Reply>> {
    reply_with_lcs_bytes(&backup_handler.get_account_state_range_proof(end_key, version)?)
}

fn get_state_snapshot(backup_handler: &BackupHandler, version: Version) -> Result<Box<dyn Reply>> {
    reply_with_async_channel_writer(backup_handler, |bh, sender| {
        send_size_prefixed_lcs_bytes(bh.get_account_iter(version), sender)
    })
}

fn get_state_root_proof(
    backup_handler: &BackupHandler,
    version: Version,
) -> Result<Box<dyn Reply>> {
    reply_with_lcs_bytes(&backup_handler.get_state_root_proof(version)?)
}

pub(crate) fn get_routes(backup_handler: BackupHandler) -> BoxedFilter<(impl Reply,)> {
    // GET latest_state_root
    let bh = backup_handler.clone();
    let latest_state_root = warp::path::end()
        .map(move || get_latest_state_root(&bh))
        .map(unwrap_or_500)
        .recover(handle_rejection);

    // GET state_range_proof/<version>/<end_key>
    let bh = backup_handler.clone();
    let state_range_proof = warp::path!(Version / HashValue)
        .map(move |version, end_key| get_state_range_proof(&bh, version, end_key))
        .map(unwrap_or_500)
        .recover(handle_rejection);

    // GET state_snapshot/<version>
    let bh = backup_handler.clone();
    let state_snapshot = warp::path!(Version)
        .map(move |version| get_state_snapshot(&bh, version))
        .map(unwrap_or_500)
        .recover(handle_rejection);

    // GET state_root_proof/<version>
    let bh = backup_handler;
    let state_root_proof = warp::path!(Version)
        .map(move |version| get_state_root_proof(&bh, version))
        .map(unwrap_or_500)
        .recover(handle_rejection);

    // Route by endpoint name.
    let routes = warp::any()
        .and(warp::path("latest_state_root").and(latest_state_root))
        .or(warp::path("state_range_proof").and(state_range_proof))
        .or(warp::path("state_snapshot").and(state_snapshot))
        .or(warp::path("state_root_proof").and(state_root_proof));

    // Serve all routes for GET only.
    warp::get().and(routes).boxed()
}
<commit_msg>[backup] service: remove some boilerplates<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

mod utils;

use crate::handlers::utils::{
    handle_rejection, reply_with_async_channel_writer, reply_with_lcs_bytes,
    send_size_prefixed_lcs_bytes, unwrap_or_500,
};
use libra_crypto::hash::HashValue;
use libra_types::transaction::Version;
use libradb::backup::BackupHandler;
use warp::{filters::BoxedFilter, reply::Reply, Filter};

pub(crate) fn get_routes(backup_handler: BackupHandler) -> BoxedFilter<(impl Reply,)> {
    // GET latest_state_root
    let bh = backup_handler.clone();
    let latest_state_root = warp::path::end()
        .map(move || reply_with_lcs_bytes(&bh.get_latest_state_root()?))
        .map(unwrap_or_500)
        .recover(handle_rejection);

    // GET state_range_proof/<version>/<end_key>
    let bh = backup_handler.clone();
    let state_range_proof = warp::path!(Version / HashValue)
        .map(move |version, end_key| {
            reply_with_lcs_bytes(&bh.get_account_state_range_proof(end_key, version)?)
        })
        .map(unwrap_or_500)
        .recover(handle_rejection);

    // GET state_snapshot/<version>
    let bh = backup_handler.clone();
    let state_snapshot = warp::path!(Version)
        .map(move |version| {
            reply_with_async_channel_writer(&bh, |bh, sender| {
                send_size_prefixed_lcs_bytes(bh.get_account_iter(version), sender)
            })
        })
        .map(unwrap_or_500)
        .recover(handle_rejection);

    // GET state_root_proof/<version>
    let bh = backup_handler;
    let state_root_proof = warp::path!(Version)
        .map(move |version| reply_with_lcs_bytes(&bh.get_state_root_proof(version)?))
        .map(unwrap_or_500)
        .recover(handle_rejection);

    // Route by endpoint name.
    let routes = warp::any()
        .and(warp::path("latest_state_root").and(latest_state_root))
        .or(warp::path("state_range_proof").and(state_range_proof))
        .or(warp::path("state_snapshot").and(state_snapshot))
        .or(warp::path("state_root_proof").and(state_root_proof));

    // Serve all routes for GET only.
    warp::get().and(routes).boxed()
}
<|endoftext|>


<commit_before>use crate::backend::UsesAnsiSavepointSyntax;
use crate::connection::{Connection, SimpleConnection};
use crate::result::{DatabaseErrorKind, Error, QueryResult};

/// Manages the internal transaction state for a connection.
///
/// You will not need to interact with this trait, unless you are writing an
/// implementation of [`Connection`](trait.Connection.html).
pub trait TransactionManager<Conn: Connection> {
    /// Begin a new transaction or savepoint
    ///
    /// If the transaction depth is greater than 0,
    /// this should create a savepoint instead.
    /// This function is expected to increment the transaction depth by 1.
    fn begin_transaction(&self, conn: &Conn) -> QueryResult<()>;

    /// Rollback the inner-most transaction or savepoint
    ///
    /// If the transaction depth is greater than 1,
    /// this should rollback to the most recent savepoint.
    /// This function is expected to decrement the transaction depth by 1.
    fn rollback_transaction(&self, conn: &Conn) -> QueryResult<()>;

    /// Commit the inner-most transaction or savepoint
    ///
    /// If the transaction depth is greater than 1,
    /// this should release the most recent savepoint.
    /// This function is expected to decrement the transaction depth by 1.
    fn commit_transaction(&self, conn: &Conn) -> QueryResult<()>;

    /// Fetch the current transaction depth
    ///
    /// Used to ensure that `begin_test_transaction` is not called when already
    /// inside of a transaction.
    fn get_transaction_depth(&self) -> u32;
}

use std::cell::Cell;

/// An implementation of `TransactionManager` which can be used for backends
/// which use ANSI standard syntax for savepoints such as SQLite and PostgreSQL.
#[allow(missing_debug_implementations)]
#[derive(Default)]
pub struct AnsiTransactionManager {
    transaction_depth: Cell<i32>,
}

impl AnsiTransactionManager {
    /// Create a new transaction manager
    pub fn new() -> Self {
        AnsiTransactionManager::default()
    }

    fn change_transaction_depth(&self, by: i32, query: QueryResult<()>) -> QueryResult<()> {
        if query.is_ok() {
            self.transaction_depth
                .set(self.transaction_depth.get() + by)
        }
        query
    }

    /// Begin a transaction with custom SQL
    ///
    /// This is used by connections to implement more complex transaction APIs
    /// to set things such as isolation levels.
    /// Returns an error if already inside of a transaction.
    pub fn begin_transaction_sql<Conn>(&self, conn: &Conn, sql: &str) -> QueryResult<()>
    where
        Conn: SimpleConnection,
    {
        use crate::result::Error::AlreadyInTransaction;

        if self.transaction_depth.get() == 0 {
            self.change_transaction_depth(1, conn.batch_execute(sql))
        } else {
            Err(AlreadyInTransaction)
        }
    }
}

impl<Conn> TransactionManager<Conn> for AnsiTransactionManager
where
    Conn: Connection,
    Conn::Backend: UsesAnsiSavepointSyntax,
{
    fn begin_transaction(&self, conn: &Conn) -> QueryResult<()> {
        let transaction_depth = self.transaction_depth.get();
        self.change_transaction_depth(
            1,
            if transaction_depth == 0 {
                conn.batch_execute("BEGIN")
            } else {
                conn.batch_execute(&format!("SAVEPOINT diesel_savepoint_{}", transaction_depth))
            },
        )
    }

    fn rollback_transaction(&self, conn: &Conn) -> QueryResult<()> {
        let transaction_depth = self.transaction_depth.get();
        self.change_transaction_depth(
            -1,
            if transaction_depth == 1 {
                conn.batch_execute("ROLLBACK")
            } else {
                conn.batch_execute(&format!(
                    "ROLLBACK TO SAVEPOINT diesel_savepoint_{}",
                    transaction_depth - 1
                ))
            },
        )
    }

    /// If the transaction fails to commit due to a `SerializationFailure` or a
    /// `ReadOnlyTransaction` a rollback will be attempted. If the rollback succeeds,
    /// the original error will be returned, otherwise the error generated by the rollback
    /// will be returned. In the second case the connection should be considered broken
    /// as it contains a uncommitted unabortable open transaction.
    fn commit_transaction(&self, conn: &Conn) -> QueryResult<()> {
        let transaction_depth = self.transaction_depth.get();
        if transaction_depth <= 1 {
            match conn.batch_execute("COMMIT") {
                // When any of these kinds of error happen on `COMMIT`, it is expected
                // that a `ROLLBACK` would succeed, leaving the transaction in a non-broken state.
                // If there are other such errors, it is fine to add them here.
                e @ Err(Error::DatabaseError(DatabaseErrorKind::SerializationFailure, _))
                | e @ Err(Error::DatabaseError(DatabaseErrorKind::ReadOnlyTransaction, _)) => {
                    self.change_transaction_depth(-1, conn.batch_execute("ROLLBACK"))?;
                    e
                }
                result => self.change_transaction_depth(-1, result),
            }
        } else {
            self.change_transaction_depth(
                -1,
                conn.batch_execute(&format!(
                    "RELEASE SAVEPOINT diesel_savepoint_{}",
                    transaction_depth - 1
                )),
            )
        }
    }

    fn get_transaction_depth(&self) -> u32 {
        self.transaction_depth.get() as u32
    }
}

#[cfg(test)]
mod test {
    macro_rules! matches {
        ($expression:expr, $( $pattern:pat )|+ $( if $guard: expr )?) => {
            match $expression {
                $( $pattern )|+ $( if $guard )? => true,
                _ => false
            }
        }
    }

    #[test]
    #[cfg(feature = "postgres")]
    fn transaction_depth_is_tracked_properly_on_commit_failure() {
        use crate::result::DatabaseErrorKind::SerializationFailure;
        use crate::result::Error::DatabaseError;
        use crate::*;
        use std::sync::{Arc, Barrier};
        use std::thread;

        table! {
            #[sql_name = "transaction_depth_is_tracked_properly_on_commit_failure"]
            serialization_example {
                id -> Serial,
                class -> Integer,
            }
        }

        let conn = crate::test_helpers::pg_connection_no_transaction();

        sql_query("DROP TABLE IF EXISTS transaction_depth_is_tracked_properly_on_commit_failure;")
            .execute(&conn)
            .unwrap();
        sql_query(
            r#"
            CREATE TABLE transaction_depth_is_tracked_properly_on_commit_failure (
                id SERIAL PRIMARY KEY,
                class INTEGER NOT NULL
            )
        "#,
        )
        .execute(&conn)
        .unwrap();

        insert_into(serialization_example::table)
            .values(&vec![
                serialization_example::class.eq(1),
                serialization_example::class.eq(2),
            ])
            .execute(&conn)
            .unwrap();

        let barrier = Arc::new(Barrier::new(2));
        let threads = (1..3)
            .map(|i| {
                let barrier = barrier.clone();
                thread::spawn(move || {
                    use crate::connection::transaction_manager::AnsiTransactionManager;
                    use crate::connection::transaction_manager::TransactionManager;
                    let conn = crate::test_helpers::pg_connection_no_transaction();
                    assert_eq!(0, <AnsiTransactionManager as TransactionManager<PgConnection>>::get_transaction_depth(&conn.transaction_manager));

                    let result =
                    conn.build_transaction().serializable().run(|| {
                        assert_eq!(1, <AnsiTransactionManager as TransactionManager<PgConnection>>::get_transaction_depth(&conn.transaction_manager));

                        let _ = serialization_example::table
                            .filter(serialization_example::class.eq(i))
                            .count()
                            .execute(&conn)?;

                        barrier.wait();

                        let other_i = if i == 1 { 2 } else { 1 };
                        insert_into(serialization_example::table)
                            .values(serialization_example::class.eq(other_i))
                            .execute(&conn)
                    });

                    assert_eq!(0, <AnsiTransactionManager as TransactionManager<PgConnection>>::get_transaction_depth(&conn.transaction_manager));
                    result
                })
            })
            .collect::<Vec<_>>();

        let mut results = threads
            .into_iter()
            .map(|t| t.join().unwrap())
            .collect::<Vec<_>>();

        results.sort_by_key(|r| r.is_err());

        assert!(matches!(results[0], Ok(_)));
        assert!(matches!(
            results[1],
            Err(DatabaseError(SerializationFailure, _))
        ));
    }
}
<commit_msg>Only define the matches macro when it will be used.<commit_after>use crate::backend::UsesAnsiSavepointSyntax;
use crate::connection::{Connection, SimpleConnection};
use crate::result::{DatabaseErrorKind, Error, QueryResult};

/// Manages the internal transaction state for a connection.
///
/// You will not need to interact with this trait, unless you are writing an
/// implementation of [`Connection`](trait.Connection.html).
pub trait TransactionManager<Conn: Connection> {
    /// Begin a new transaction or savepoint
    ///
    /// If the transaction depth is greater than 0,
    /// this should create a savepoint instead.
    /// This function is expected to increment the transaction depth by 1.
    fn begin_transaction(&self, conn: &Conn) -> QueryResult<()>;

    /// Rollback the inner-most transaction or savepoint
    ///
    /// If the transaction depth is greater than 1,
    /// this should rollback to the most recent savepoint.
    /// This function is expected to decrement the transaction depth by 1.
    fn rollback_transaction(&self, conn: &Conn) -> QueryResult<()>;

    /// Commit the inner-most transaction or savepoint
    ///
    /// If the transaction depth is greater than 1,
    /// this should release the most recent savepoint.
    /// This function is expected to decrement the transaction depth by 1.
    fn commit_transaction(&self, conn: &Conn) -> QueryResult<()>;

    /// Fetch the current transaction depth
    ///
    /// Used to ensure that `begin_test_transaction` is not called when already
    /// inside of a transaction.
    fn get_transaction_depth(&self) -> u32;
}

use std::cell::Cell;

/// An implementation of `TransactionManager` which can be used for backends
/// which use ANSI standard syntax for savepoints such as SQLite and PostgreSQL.
#[allow(missing_debug_implementations)]
#[derive(Default)]
pub struct AnsiTransactionManager {
    transaction_depth: Cell<i32>,
}

impl AnsiTransactionManager {
    /// Create a new transaction manager
    pub fn new() -> Self {
        AnsiTransactionManager::default()
    }

    fn change_transaction_depth(&self, by: i32, query: QueryResult<()>) -> QueryResult<()> {
        if query.is_ok() {
            self.transaction_depth
                .set(self.transaction_depth.get() + by)
        }
        query
    }

    /// Begin a transaction with custom SQL
    ///
    /// This is used by connections to implement more complex transaction APIs
    /// to set things such as isolation levels.
    /// Returns an error if already inside of a transaction.
    pub fn begin_transaction_sql<Conn>(&self, conn: &Conn, sql: &str) -> QueryResult<()>
    where
        Conn: SimpleConnection,
    {
        use crate::result::Error::AlreadyInTransaction;

        if self.transaction_depth.get() == 0 {
            self.change_transaction_depth(1, conn.batch_execute(sql))
        } else {
            Err(AlreadyInTransaction)
        }
    }
}

impl<Conn> TransactionManager<Conn> for AnsiTransactionManager
where
    Conn: Connection,
    Conn::Backend: UsesAnsiSavepointSyntax,
{
    fn begin_transaction(&self, conn: &Conn) -> QueryResult<()> {
        let transaction_depth = self.transaction_depth.get();
        self.change_transaction_depth(
            1,
            if transaction_depth == 0 {
                conn.batch_execute("BEGIN")
            } else {
                conn.batch_execute(&format!("SAVEPOINT diesel_savepoint_{}", transaction_depth))
            },
        )
    }

    fn rollback_transaction(&self, conn: &Conn) -> QueryResult<()> {
        let transaction_depth = self.transaction_depth.get();
        self.change_transaction_depth(
            -1,
            if transaction_depth == 1 {
                conn.batch_execute("ROLLBACK")
            } else {
                conn.batch_execute(&format!(
                    "ROLLBACK TO SAVEPOINT diesel_savepoint_{}",
                    transaction_depth - 1
                ))
            },
        )
    }

    /// If the transaction fails to commit due to a `SerializationFailure` or a
    /// `ReadOnlyTransaction` a rollback will be attempted. If the rollback succeeds,
    /// the original error will be returned, otherwise the error generated by the rollback
    /// will be returned. In the second case the connection should be considered broken
    /// as it contains a uncommitted unabortable open transaction.
    fn commit_transaction(&self, conn: &Conn) -> QueryResult<()> {
        let transaction_depth = self.transaction_depth.get();
        if transaction_depth <= 1 {
            match conn.batch_execute("COMMIT") {
                // When any of these kinds of error happen on `COMMIT`, it is expected
                // that a `ROLLBACK` would succeed, leaving the transaction in a non-broken state.
                // If there are other such errors, it is fine to add them here.
                e @ Err(Error::DatabaseError(DatabaseErrorKind::SerializationFailure, _))
                | e @ Err(Error::DatabaseError(DatabaseErrorKind::ReadOnlyTransaction, _)) => {
                    self.change_transaction_depth(-1, conn.batch_execute("ROLLBACK"))?;
                    e
                }
                result => self.change_transaction_depth(-1, result),
            }
        } else {
            self.change_transaction_depth(
                -1,
                conn.batch_execute(&format!(
                    "RELEASE SAVEPOINT diesel_savepoint_{}",
                    transaction_depth - 1
                )),
            )
        }
    }

    fn get_transaction_depth(&self) -> u32 {
        self.transaction_depth.get() as u32
    }
}

#[cfg(test)]
mod test {
    #[cfg(feature = "postgres")]
    macro_rules! matches {
        ($expression:expr, $( $pattern:pat )|+ $( if $guard: expr )?) => {
            match $expression {
                $( $pattern )|+ $( if $guard )? => true,
                _ => false
            }
        }
    }

    #[test]
    #[cfg(feature = "postgres")]
    fn transaction_depth_is_tracked_properly_on_commit_failure() {
        use crate::result::DatabaseErrorKind::SerializationFailure;
        use crate::result::Error::DatabaseError;
        use crate::*;
        use std::sync::{Arc, Barrier};
        use std::thread;

        table! {
            #[sql_name = "transaction_depth_is_tracked_properly_on_commit_failure"]
            serialization_example {
                id -> Serial,
                class -> Integer,
            }
        }

        let conn = crate::test_helpers::pg_connection_no_transaction();

        sql_query("DROP TABLE IF EXISTS transaction_depth_is_tracked_properly_on_commit_failure;")
            .execute(&conn)
            .unwrap();
        sql_query(
            r#"
            CREATE TABLE transaction_depth_is_tracked_properly_on_commit_failure (
                id SERIAL PRIMARY KEY,
                class INTEGER NOT NULL
            )
        "#,
        )
        .execute(&conn)
        .unwrap();

        insert_into(serialization_example::table)
            .values(&vec![
                serialization_example::class.eq(1),
                serialization_example::class.eq(2),
            ])
            .execute(&conn)
            .unwrap();

        let barrier = Arc::new(Barrier::new(2));
        let threads = (1..3)
            .map(|i| {
                let barrier = barrier.clone();
                thread::spawn(move || {
                    use crate::connection::transaction_manager::AnsiTransactionManager;
                    use crate::connection::transaction_manager::TransactionManager;
                    let conn = crate::test_helpers::pg_connection_no_transaction();
                    assert_eq!(0, <AnsiTransactionManager as TransactionManager<PgConnection>>::get_transaction_depth(&conn.transaction_manager));

                    let result =
                    conn.build_transaction().serializable().run(|| {
                        assert_eq!(1, <AnsiTransactionManager as TransactionManager<PgConnection>>::get_transaction_depth(&conn.transaction_manager));

                        let _ = serialization_example::table
                            .filter(serialization_example::class.eq(i))
                            .count()
                            .execute(&conn)?;

                        barrier.wait();

                        let other_i = if i == 1 { 2 } else { 1 };
                        insert_into(serialization_example::table)
                            .values(serialization_example::class.eq(other_i))
                            .execute(&conn)
                    });

                    assert_eq!(0, <AnsiTransactionManager as TransactionManager<PgConnection>>::get_transaction_depth(&conn.transaction_manager));
                    result
                })
            })
            .collect::<Vec<_>>();

        let mut results = threads
            .into_iter()
            .map(|t| t.join().unwrap())
            .collect::<Vec<_>>();

        results.sort_by_key(|r| r.is_err());

        assert!(matches!(results[0], Ok(_)));
        assert!(matches!(
            results[1],
            Err(DatabaseError(SerializationFailure, _))
        ));
    }
}
<|endoftext|>


<commit_before>extern crate libsqlite3_sys as ffi;

use std::collections::HashMap;
use std::os::raw as libc;
use std::ptr::NonNull;
use std::{slice, str};

use crate::row::*;
use crate::sqlite::Sqlite;

#[allow(missing_debug_implementations, missing_copy_implementations)]
pub struct SqliteValue {
    value: ffi::sqlite3_value,
}

pub struct SqliteRow {
    stmt: NonNull<ffi::sqlite3_stmt>,
    next_col_index: libc::c_int,
}

impl SqliteValue {
    #[allow(clippy::new_ret_no_self)]
    pub(crate) unsafe fn new<'a>(inner: *mut ffi::sqlite3_value) -> Option<&'a Self> {
        (inner as *const _ as *const Self).as_ref().and_then(|v| {
            if v.is_null() {
                None
            } else {
                Some(v)
            }
        })
    }

    pub fn read_text(&self) -> &str {
        unsafe {
            let ptr = ffi::sqlite3_value_text(self.value());
            let len = ffi::sqlite3_value_bytes(self.value());
            let bytes = slice::from_raw_parts(ptr as *const u8, len as usize);
            // The string is guaranteed to be utf8 according to
            // https://www.sqlite.org/c3ref/value_blob.html
            str::from_utf8_unchecked(bytes)
        }
    }

    pub fn read_blob(&self) -> &[u8] {
        unsafe {
            let ptr = ffi::sqlite3_value_blob(self.value());
            let len = ffi::sqlite3_value_bytes(self.value());
            slice::from_raw_parts(ptr as *const u8, len as usize)
        }
    }

    pub fn read_integer(&self) -> i32 {
        unsafe { ffi::sqlite3_value_int(self.value()) as i32 }
    }

    pub fn read_long(&self) -> i64 {
        unsafe { ffi::sqlite3_value_int64(self.value()) as i64 }
    }

    pub fn read_double(&self) -> f64 {
        unsafe { ffi::sqlite3_value_double(self.value()) as f64 }
    }

    pub fn is_null(&self) -> bool {
        let tpe = unsafe { ffi::sqlite3_value_type(self.value()) };
        tpe == ffi::SQLITE_NULL
    }

    fn value(&self) -> *mut ffi::sqlite3_value {
        &self.value as *const _ as _
    }
}

impl SqliteRow {
    pub(crate) fn new(inner_statement: NonNull<ffi::sqlite3_stmt>) -> Self {
        SqliteRow {
            stmt: inner_statement,
            next_col_index: 0,
        }
    }

    pub fn into_named<'a>(self, indices: &'a HashMap<&'a str, usize>) -> SqliteNamedRow<'a> {
        SqliteNamedRow {
            stmt: self.stmt,
            column_indices: indices,
        }
    }
}

impl Row<Sqlite> for SqliteRow {
    fn take(&mut self) -> Option<&SqliteValue> {
        let col_index = self.next_col_index;
        self.next_col_index += 1;

        unsafe {
            let ptr = ffi::sqlite3_column_value(self.stmt.as_ptr(), col_index);
            SqliteValue::new(ptr)
        }
    }

    fn next_is_null(&self, count: usize) -> bool {
        (0..count).all(|i| {
            let idx = self.next_col_index + i as libc::c_int;
            let tpe = unsafe { ffi::sqlite3_column_type(self.stmt.as_ptr(), idx) };
            tpe == ffi::SQLITE_NULL
        })
    }

    fn column_name(&self) -> Option<&str> {
        unsafe {
            let ptr = if self.next_col_index == 0 {
                ffi::sqlite3_column_name(self.stmt.as_ptr(), 0)
            } else {
                ffi::sqlite3_column_name(self.stmt.as_ptr(), self.next_col_index - 1)
            };
            Some(std::ffi::CStr::from_ptr(ptr).to_str().expect(
                "The Sqlite documentation states that this is UTF8. \
                 If you see this error message something has gone \
                 horribliy wrong. Please open an issue at the \
                 diesel repository.",
            ))
        }
    }

    fn column_count(&self) -> usize {
        unsafe { ffi::sqlite3_column_count(self.stmt.as_ptr()) as usize }
    }
}

pub struct SqliteNamedRow<'a> {
    stmt: NonNull<ffi::sqlite3_stmt>,
    column_indices: &'a HashMap<&'a str, usize>,
}

impl<'a> NamedRow<Sqlite> for SqliteNamedRow<'a> {
    fn index_of(&self, column_name: &str) -> Option<usize> {
        self.column_indices.get(column_name).cloned()
    }

    fn get_raw_value(&self, idx: usize) -> Option<&SqliteValue> {
        unsafe {
            let ptr = ffi::sqlite3_column_value(self.stmt.as_ptr(), idx as libc::c_int);
            SqliteValue::new(ptr)
        }
    }
}
<commit_msg>Add a way to get type information on sqlite<commit_after>extern crate libsqlite3_sys as ffi;

use std::collections::HashMap;
use std::os::raw as libc;
use std::ptr::NonNull;
use std::{slice, str};

use crate::row::*;
use crate::sqlite::{Sqlite, SqliteType};

#[allow(missing_debug_implementations, missing_copy_implementations)]
pub struct SqliteValue {
    value: ffi::sqlite3_value,
}

pub struct SqliteRow {
    stmt: NonNull<ffi::sqlite3_stmt>,
    next_col_index: libc::c_int,
}

impl SqliteValue {
    #[allow(clippy::new_ret_no_self)]
    pub(crate) unsafe fn new<'a>(inner: *mut ffi::sqlite3_value) -> Option<&'a Self> {
        (inner as *const _ as *const Self).as_ref().and_then(|v| {
            if v.is_null() {
                None
            } else {
                Some(v)
            }
        })
    }

    pub fn read_text(&self) -> &str {
        unsafe {
            let ptr = ffi::sqlite3_value_text(self.value());
            let len = ffi::sqlite3_value_bytes(self.value());
            let bytes = slice::from_raw_parts(ptr as *const u8, len as usize);
            // The string is guaranteed to be utf8 according to
            // https://www.sqlite.org/c3ref/value_blob.html
            str::from_utf8_unchecked(bytes)
        }
    }

    pub fn read_blob(&self) -> &[u8] {
        unsafe {
            let ptr = ffi::sqlite3_value_blob(self.value());
            let len = ffi::sqlite3_value_bytes(self.value());
            slice::from_raw_parts(ptr as *const u8, len as usize)
        }
    }

    pub fn read_integer(&self) -> i32 {
        unsafe { ffi::sqlite3_value_int(self.value()) as i32 }
    }

    pub fn read_long(&self) -> i64 {
        unsafe { ffi::sqlite3_value_int64(self.value()) as i64 }
    }

    pub fn read_double(&self) -> f64 {
        unsafe { ffi::sqlite3_value_double(self.value()) as f64 }
    }

    pub fn is_null(&self) -> bool {
        self.value_type().is_none()
    }

    pub fn value_type(&self) -> Option<SqliteType> {
        let tpe = unsafe { ffi::sqlite3_value_type(self.value()) };
        match tpe {
            ffi::SQLITE_TEXT => Some(SqliteType::Text),
            ffi::SQLITE_INTEGER => Some(SqliteType::Long),
            ffi::SQLITE_FLOAT => Some(SqliteType::Double),
            ffi::SQLITE_BLOB => Some(SqliteType::Binary),
            ffi::SQLITE_NULL => None,
            _ => unreachable!("Sqlite does saying this is not reachable"),
        }
    }

    fn value(&self) -> *mut ffi::sqlite3_value {
        &self.value as *const _ as _
    }
}

impl SqliteRow {
    pub(crate) fn new(inner_statement: NonNull<ffi::sqlite3_stmt>) -> Self {
        SqliteRow {
            stmt: inner_statement,
            next_col_index: 0,
        }
    }

    pub fn into_named<'a>(self, indices: &'a HashMap<&'a str, usize>) -> SqliteNamedRow<'a> {
        SqliteNamedRow {
            stmt: self.stmt,
            column_indices: indices,
        }
    }
}

impl Row<Sqlite> for SqliteRow {
    fn take(&mut self) -> Option<&SqliteValue> {
        let col_index = self.next_col_index;
        self.next_col_index += 1;

        unsafe {
            let ptr = ffi::sqlite3_column_value(self.stmt.as_ptr(), col_index);
            SqliteValue::new(ptr)
        }
    }

    fn next_is_null(&self, count: usize) -> bool {
        (0..count).all(|i| {
            let idx = self.next_col_index + i as libc::c_int;
            let tpe = unsafe { ffi::sqlite3_column_type(self.stmt.as_ptr(), idx) };
            tpe == ffi::SQLITE_NULL
        })
    }

    fn column_name(&self) -> Option<&str> {
        unsafe {
            let ptr = if self.next_col_index == 0 {
                ffi::sqlite3_column_name(self.stmt.as_ptr(), 0)
            } else {
                ffi::sqlite3_column_name(self.stmt.as_ptr(), self.next_col_index - 1)
            };
            Some(std::ffi::CStr::from_ptr(ptr).to_str().expect(
                "The Sqlite documentation states that this is UTF8. \
                 If you see this error message something has gone \
                 horribliy wrong. Please open an issue at the \
                 diesel repository.",
            ))
        }
    }

    fn column_count(&self) -> usize {
        unsafe { ffi::sqlite3_column_count(self.stmt.as_ptr()) as usize }
    }
}

pub struct SqliteNamedRow<'a> {
    stmt: NonNull<ffi::sqlite3_stmt>,
    column_indices: &'a HashMap<&'a str, usize>,
}

impl<'a> NamedRow<Sqlite> for SqliteNamedRow<'a> {
    fn index_of(&self, column_name: &str) -> Option<usize> {
        self.column_indices.get(column_name).cloned()
    }

    fn get_raw_value(&self, idx: usize) -> Option<&SqliteValue> {
        unsafe {
            let ptr = ffi::sqlite3_column_value(self.stmt.as_ptr(), idx as libc::c_int);
            SqliteValue::new(ptr)
        }
    }
}
<|endoftext|>


<commit_before>use std::io::IoResult;

use collections::HashMap;

use backend::Backend;
use colors;

pub struct LatexBackend;

impl LatexBackend {
    pub fn new() -> LatexBackend {
        LatexBackend
    }
}

static HEADER: &'static str = "\
\\usepackage{xcolor}
\\usepackage{fancyvrb}
\\newcommand{\\VerbBar}{|}
\\newcommand{\\VERB}{\\Verb[commandchars=\\\\\\{\\}]}
\\DefineVerbatimEnvironment{Highlighting}{Verbatim}{commandchars=\\\\\\{\\}}
% Add ',fontsize=\\small' for more characters per line
\\newenvironment{Shaded}{}{}
";

impl Backend for LatexBackend {
    fn configure(&mut self, _vars: &HashMap<~str, ~str>) -> Result<(), ~str> {
        Ok(())
    }

    fn header(&mut self, w: &mut Writer) -> IoResult<()> {
        try!(w.write_str(HEADER));

        for (ty, color) in colors::get_colors().iter() {
            try!(writeln!(w, "\\\\definecolor\\{{}\\}\\{HTML\\}\\{{}\\}", ty, color));
        }

        Ok(())
    }

    fn code_start(&mut self, w: &mut Writer) -> IoResult<()> {
        try!(w.write_line("\\begin{Shaded}"));
        try!(w.write_line("\\begin{Highlighting}[]"));

        Ok(())
    }

    fn code_end(&mut self, w: &mut Writer) -> IoResult<()> {
        try!(w.write_line("\\end{Highlighting}"));
        try!(w.write_line("\\end{Shaded}"));

        Ok(())
    }

    fn start(&mut self, w: &mut Writer, ty: &str) -> IoResult<()> {
        if colors::get_types().contains(&ty.to_owned()) {
            try!(write!(w, "\\\\textcolor\\{{}\\}\\{", ty));
        }

        if ty == "attribute" {
            try!(w.write_str("#"));
        }

        Ok(())
    }

    fn end(&mut self, w: &mut Writer, ty: &str) -> IoResult<()> {
        if ty == "attribute" {
            try!(w.write_str("]"));
        }

        if colors::get_types().contains(&ty.to_owned()) {
            try!(w.write_str("}"));
        }

        Ok(())
    }

    fn text(&mut self, w: &mut Writer, text: &str) -> IoResult<()> {
        fn escape_latex(text: &str) -> ~str {
            let mut result = ~"";

            let mut escape = false;
            for c in text.chars() {
                if escape {
                    result.push_str("\\textbackslash{");
                }

                if c == '{' || c == '}' {
                    result.push_char('\\');
                }

                if c == '\\' {
                    escape = true;
                } else {
                    result.push_char(c);
                }

                if escape && c != '\\' {
                    result.push_char('}');
                    escape = false;
                }
            }

            result
        }
        let text = escape_latex(text);

        try!(w.write_str(text));

        Ok(())
    }
}
<commit_msg>Fix comments handling in LaTeX backend.<commit_after>use std::io::IoResult;

use collections::HashMap;

use backend::Backend;
use colors;

pub struct LatexBackend {
    priv contexts: Vec<~str>,
}

impl LatexBackend {
    pub fn new() -> LatexBackend {
        LatexBackend {
            contexts: Vec::new(),
        }
    }
}

static HEADER: &'static str = "\
\\usepackage{xcolor}
\\usepackage{fancyvrb}
\\newcommand{\\VerbBar}{|}
\\newcommand{\\VERB}{\\Verb[commandchars=\\\\\\{\\}]}
\\DefineVerbatimEnvironment{Highlighting}{Verbatim}{commandchars=\\\\\\{\\}}
% Add ',fontsize=\\small' for more characters per line
\\newenvironment{Shaded}{}{}
";

impl Backend for LatexBackend {
    fn configure(&mut self, _vars: &HashMap<~str, ~str>) -> Result<(), ~str> {
        Ok(())
    }

    fn header(&mut self, w: &mut Writer) -> IoResult<()> {
        try!(w.write_str(HEADER));

        for (ty, color) in colors::get_colors().iter() {
            try!(writeln!(w, "\\\\definecolor\\{{}\\}\\{HTML\\}\\{{}\\}", ty, color));
        }

        Ok(())
    }

    fn code_start(&mut self, w: &mut Writer) -> IoResult<()> {
        try!(w.write_line("\\begin{Shaded}"));
        try!(w.write_line("\\begin{Highlighting}[]"));

        Ok(())
    }

    fn code_end(&mut self, w: &mut Writer) -> IoResult<()> {
        try!(w.write_line("\\end{Highlighting}"));
        try!(w.write_line("\\end{Shaded}"));

        Ok(())
    }

    fn start(&mut self, w: &mut Writer, ty: &str) -> IoResult<()> {
        if ty != "comment" {
            if colors::get_types().contains(&ty.to_owned()) {
                try!(write!(w, "\\\\textcolor\\{{}\\}\\{", ty));
            }

            if ty == "attribute" {
                try!(w.write_str("#"));
            }
        }

        self.contexts.push(ty.to_owned());

        Ok(())
    }

    fn end(&mut self, w: &mut Writer, ty: &str) -> IoResult<()> {
        if ty != "comment" {
            if ty == "attribute" {
                try!(w.write_str("]"));
            }

            if colors::get_types().contains(&ty.to_owned()) {
                try!(w.write_str("}"));
            }
        }

        self.contexts.pop();

        Ok(())
    }

    fn text(&mut self, w: &mut Writer, text: &str) -> IoResult<()> {
        fn escape_latex(text: &str) -> ~str {
            let mut result = ~"";

            let mut escape = false;
            for c in text.chars() {
                if escape {
                    result.push_str("\\textbackslash{");
                }

                if c == '{' || c == '}' {
                    result.push_char('\\');
                }

                if c == '\\' {
                    escape = true;
                } else {
                    result.push_char(c);
                }

                if escape && c != '\\' {
                    result.push_char('}');
                    escape = false;
                }
            }

            result
        }

        fn escape_comment(text: &str, has_color: bool) -> ~str {
            let mut result = ~"";

            let mut first = true;
            for line in text.lines() {
                if !first {
                    result.push_str("\n");
                }

                if line.len() > 0 && has_color {
                    result.push_str("\\textcolor{comment}{");
                }

                result.push_str(line);

                if line.len() > 0 && has_color {
                    result.push_str("}");
                }

                first = false;
            }

            let old_len = text.len();
            let text = text.trim_right_chars(&'\n').to_owned();
            let new_len = text.len();

            range(0, old_len - new_len).advance(|_| {
                result.push_char('\n');

                true
            });

            result
        }

        let context = self.contexts.last().unwrap();
        let has_color = colors::get_types().contains(context);
        let context = context.as_slice();

        let text = if context == "comment" {
            escape_comment(text, has_color)
        } else {
            escape_latex(text)
        };
        try!(w.write_str(text));

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


<commit_before>use std::rc::{Rc, Weak};
use std::cell::RefCell;
use std::borrow::Borrow;
use dovahkiin::types::custom_types::id::Id;
use smallvec::SmallVec;
use std::sync::Arc;
use ram::chunk::Chunks;
use ram::cell::Cell;
use dovahkiin::types;
use dovahkiin::types::custom_types::map::key_hash;
use dovahkiin::types::value::{ValueIter};
use dovahkiin::types::{Value, PrimitiveArray};
use utils::lru_cache::LRUCache;
use parking_lot::{Mutex, MutexGuard};
use std::io::Cursor;
use std::ops::Index;
use std::cmp::Ordering;

lazy_static! {
    pub static ref ENTRIES_KEY_HASH : u64 = key_hash("entries");
    pub static ref KEYS_KEY_HASH : u64 = key_hash("keys");
    pub static ref PREV_PAGE_KEY_HASH : u64 = key_hash("prev");
    pub static ref NEXT_PAGE_KEY_HASH : u64 = key_hash("next");
}

type EntryKey = SmallVec<[u8; 32]>;
type CachedExtNodeRef = Arc<CachedExtNode>;

pub struct LSMTree {
    num_levels: u8,
    levels: Vec<Rc<BPlusTree>>,
}

 trait BPlusTree {
      fn root(&self) -> &Node;
      fn root_mut(&self) -> &mut Node;
      fn get_height(&self) -> u32;
      fn set_height(&mut self) -> u32;
      fn get_num_nodes(&self) -> u32;
      fn set_num_nodes(&self) -> u32;
      fn page_size(&self) -> u32;
      fn chunks(&self) -> &Arc<Chunks>;
      fn page_cache(&self) -> MutexGuard<LRUCache<Id, CachedExtNodeRef>>;
      fn get_page_direct(&self, id: &Id) -> CachedExtNodeRef {
          let cell = self.chunks().read_cell(id).unwrap(); // should crash if not exists
          let keys = &cell.data[*KEYS_KEY_HASH];
          let keys_len = keys.len().unwrap();
          let mut v_node = CachedExtNode {
              keys: Vec::new(),
              ids: Vec::new(),
              cap: self.page_size()
          };
          for key in keys.iter_value().unwrap() {
              let key = if let Value::PrimArray(PrimitiveArray::U8(ref array)) = key {
                  EntryKey::from_slice(array.as_slice())
              } else { panic!("invalid entry") };
              let id = id_from_key(&key);
              v_node.keys.push(EntryKey::from_slice(key.as_slice()));
              v_node.ids.push(Delimiter::External(id));
          }
          return Arc::new(v_node);
      }
      fn get_page(&self, id: &Id) -> CachedExtNodeRef {
          self.page_cache().get_or_fetch(id).unwrap().clone()
      }
      fn get(&self, key: &EntryKey) -> Option<Id> {
          self.search(self.root(), key, self.get_height())
      }
      fn search<'a>(&self, node: &'a Node, key: &EntryKey, ht: u32) -> Option<Id> {
          if ht == 0 { // persist page
              let search_result = node.keys()
                  .search(key);
              let index = if let Ok(i) = search_result { i } else { return None };
              match &node.delimiters().get(index) {
                  Some(Delimiter::External(id)) => Some(*id),
                  Some(Delimiter::Internal(_)) => panic!("Error on leaf delimiter type"),
                  None => panic!("Error on leaf delimiter is none")
              }
          } else {
              let index = node.keys()
                  .search(key)
                  .map(|i| i + 1)
                  .unwrap_or_else(|i| i);
              match &node.delimiters().get(index) {
                  Some(Delimiter::External(ref id)) => {
                      assert_eq!(ht, 0);
                      // search in leaf page
                      let page = self.get_page(id);
                      return self.search(&*page, key, ht - 1);
                  },
                  Some(Delimiter::Internal(node)) => return self.search(node.borrow(), key, ht - 1),
                  None => return None
              }
          }
      }
     fn insert(&self, key: &EntryKey, value: &Id) {
         unimplemented!()
     }
     fn put<'a>(&self, node: &'a Node, key: &EntryKey, ht: u32) {
         let keys = node.keys();
         let index = keys
             .search(key)
             .map(|i| i + 1)
             .unwrap_or_else(|i| i);
         if ht == 0 {

         }
     }
 }

enum Delimiter {
    External(Id), // to leaf
    Internal(Box<Node>) // to higher level node
}

trait Array<T> where T: Ord {
    #[inline]
    fn index_of(&self, index: usize) -> &T;
    #[inline]
    fn size(&self) -> usize;
    #[inline]
    fn search(&self, item: &T) -> Result<usize, usize>;
    #[inline]
    fn get(&self, index: usize) -> Option<&T>;
}

trait Keys : Array<EntryKey> {}
trait Delimiters : Array<Delimiter> {}
impl Keys for [EntryKey] {}
impl Delimiters for [Delimiter] {}
impl <T> Array<T> for [T] where T: Ord {
    #[inline]
    fn index_of(&self, index: usize) -> &T { &self[index] }

    #[inline]
    fn size(&self) -> usize { self.len() }

    #[inline]
    fn search(&self, item: &T) -> Result<usize, usize> { self.binary_search(item) }

    #[inline]
    fn get(&self, index: usize) -> Option<&T> {
        if index >= self.size() { None } else { Some(&self[index]) }
    }
}
impl <T> Array<T> for Vec<T> where T: Ord {
    #[inline]
    fn index_of(&self, index: usize) -> &T {
        self.get(index).unwrap()
    }

    #[inline]
    fn size(&self) -> usize {
        self.len()
    }

    #[inline]
    fn search(&self, item: &T) -> Result<usize, usize> {
        self.as_slice().binary_search(item)
    }

    #[inline]
    fn get(&self, index: usize) -> Option<&T> {
        if index >= self.len() { None } else { Some(&self[index]) }
    }
}
impl Delimiters for Vec<Delimiter> {}
impl Keys for Vec<EntryKey> {}

trait Node {
    #[inline]
    fn keys(&self) -> &Keys;
    #[inline]
    fn delimiters(&self) -> &Delimiters;
    #[inline]
    fn add(&mut self, key: EntryKey) -> Option<None>;
    #[inline]
    fn del(&mut self, key: &EntryKey);
    #[inline]
    fn merge(&mut self, x: None);
}

struct CachedExtNode {
    keys: Vec<EntryKey>,
    ids: Vec<Delimiter>,
    cap: u32,
}

impl Node for CachedExtNode {
    #[inline]
    fn keys(&self) -> &Keys {
        &self.keys
    }

    #[inline]
    fn delimiters(&self) -> &Delimiters {
        &self.ids
    }

    fn add(&mut self, key: EntryKey) -> Option<Self> {
        let insert_pos = match self.keys.binary_search(&key) {
            Ok(_) => return None, // existed
            Err(i) => i
        };
        let id = id_from_key(&key);
        self.keys.insert(i, key);
        self.ids.insert(i, Delimiter::External(id));
        if self.keys.len() > self.cap as usize {
            // need to split
            let mid = self.keys.len() / 2;
            let keys_2: Vec<_> = self.keys.drain(mid..).collect();
            let ids_2: Vec<_> = self.ids.drain(mid..).collect();
            return Some(CachedExtNode {
                keys: keys_2,
                ids: ids_2,
                cap: self.cap
            });
        } else {
            return None;
        }
    }

    fn del(&mut self, key: &EntryKey) {
        let pos = match self.keys.binary_search(&key) {
            Ok(i) => i, // existed
            Err(i) => return
        };
        self.keys.remove(pos);
        self.ids.remove(pos);
    }

    fn merge(&mut self, mut x: Self) {
        self.keys.append(&mut x.keys);
        self.ids.append(&mut x.ids);
    }
}

trait InterNode {

}

impl <T> Index<usize> for Array<T> where T: Ord {
    type Output = T;

    fn index(&self, index: usize) -> &<Self as Index<usize>>::Output {
        self.index_of(index)
    }
}

macro_rules! impl_nodes {
    ($(($level: ident, $entry_size: expr, $delimiter_size: expr)),+) => {
        $(
            mod $level {
                use super::*;

                type LEntryKeys = [EntryKey; $entry_size];

                impl Array<EntryKey> for LEntryKeys {
                    #[inline]
                    fn size(&self) -> usize { $entry_size }
                    #[inline]
                    fn index_of(&self, index: usize) -> &EntryKey { &self[index] }
                    #[inline]
                    fn search(&self, item: &EntryKey) -> Result<usize, usize> {
                        self.binary_search(item)
                    }
                    #[inline]
                    fn get(&self, index: usize) -> Option<&EntryKey> {
                        if index >= $entry_size { None } else { Some(&self[index]) }
                    }
                }

                type LDelimiter = [Delimiter; $delimiter_size];

                impl Array<Delimiter> for LDelimiter {
                    #[inline]
                    fn size(&self) -> usize { $delimiter_size }
                    #[inline]
                    fn index_of(&self, index: usize) -> &Delimiter { &self[index] }
                    #[inline]
                    fn search(&self, item: &Delimiter) -> Result<usize, usize> {
                        self.binary_search(item)
                    }
                    #[inline]
                    fn get(&self, index: usize) -> Option<&Delimiter> {
                        if index >= $entry_size { None } else { Some(&self[index]) }
                    }
                }

                impl Delimiters for LDelimiter {}
                impl Keys for LEntryKeys {}

                struct LNode {
                    keys: LEntryKeys,
                    delimeters: LDelimiter,
                    pos: u32
                }

                impl Node for LNode {
                    #[inline]
                    fn keys(&self) -> &Keys {
                        &self.keys
                    }
                    #[inline]
                    fn delimiters(&self) -> &Delimiters {
                        &self.delimeters
                    }
                    #[inline]
                    fn add(&mut self, key: EntryKey) -> Option<Node> {
                        unimplemented!()
                    }
                    #[inline]
                    fn del(&mut self, key: &EntryKey) {
                        unimplemented!()
                    }
                    #[inline]
                    fn merge(&mut self, x: Self) {
                        unimplemented!()
                    }
                }
            }
        )+
    };
}

impl_nodes!((level_0, 2, 3), (level_1, 20, 21), (level_2, 200, 201), (level_3, 2000, 2001), (level_4, 20000, 20001));

impl Ord for Delimiter {
    fn cmp(&self, other: &Self) -> Ordering {
        panic!()
    }
}
impl PartialOrd for Delimiter {
    fn partial_cmp(&self, other: &Delimiter) -> Option<Ordering> {
        panic!()
    }
}
impl Eq for Delimiter {}
impl PartialEq for Delimiter {
    fn eq(&self, other: &Delimiter) -> bool {
        panic!()
    }
}

fn id_from_key(key: &EntryKey) -> Id {
    let mut id_cursor = Cursor::new(&key[key.len() - 16 ..]);
    return Id::from_binary(&mut id_cursor).unwrap(); // read id from tailing 128 bits
}<commit_msg>fix typo<commit_after>use std::rc::{Rc, Weak};
use std::cell::RefCell;
use std::borrow::Borrow;
use dovahkiin::types::custom_types::id::Id;
use smallvec::SmallVec;
use std::sync::Arc;
use ram::chunk::Chunks;
use ram::cell::Cell;
use dovahkiin::types;
use dovahkiin::types::custom_types::map::key_hash;
use dovahkiin::types::value::{ValueIter};
use dovahkiin::types::{Value, PrimitiveArray};
use utils::lru_cache::LRUCache;
use parking_lot::{Mutex, MutexGuard};
use std::io::Cursor;
use std::ops::Index;
use std::cmp::Ordering;

lazy_static! {
    pub static ref ENTRIES_KEY_HASH : u64 = key_hash("entries");
    pub static ref KEYS_KEY_HASH : u64 = key_hash("keys");
    pub static ref PREV_PAGE_KEY_HASH : u64 = key_hash("prev");
    pub static ref NEXT_PAGE_KEY_HASH : u64 = key_hash("next");
}

type EntryKey = SmallVec<[u8; 32]>;
type CachedExtNodeRef = Arc<CachedExtNode>;

pub struct LSMTree {
    num_levels: u8,
    levels: Vec<Rc<BPlusTree>>,
}

 trait BPlusTree {
      fn root(&self) -> &Node;
      fn root_mut(&self) -> &mut Node;
      fn get_height(&self) -> u32;
      fn set_height(&mut self) -> u32;
      fn get_num_nodes(&self) -> u32;
      fn set_num_nodes(&self) -> u32;
      fn page_size(&self) -> u32;
      fn chunks(&self) -> &Arc<Chunks>;
      fn page_cache(&self) -> MutexGuard<LRUCache<Id, CachedExtNodeRef>>;
      fn get_page_direct(&self, id: &Id) -> CachedExtNodeRef {
          let cell = self.chunks().read_cell(id).unwrap(); // should crash if not exists
          let keys = &cell.data[*KEYS_KEY_HASH];
          let keys_len = keys.len().unwrap();
          let mut v_node = CachedExtNode {
              keys: Vec::new(),
              ids: Vec::new(),
              cap: self.page_size()
          };
          for key in keys.iter_value().unwrap() {
              let key = if let Value::PrimArray(PrimitiveArray::U8(ref array)) = key {
                  EntryKey::from_slice(array.as_slice())
              } else { panic!("invalid entry") };
              let id = id_from_key(&key);
              v_node.keys.push(EntryKey::from_slice(key.as_slice()));
              v_node.ids.push(Delimiter::External(id));
          }
          return Arc::new(v_node);
      }
      fn get_page(&self, id: &Id) -> CachedExtNodeRef {
          self.page_cache().get_or_fetch(id).unwrap().clone()
      }
      fn get(&self, key: &EntryKey) -> Option<Id> {
          self.search(self.root(), key, self.get_height())
      }
      fn search<'a>(&self, node: &'a Node, key: &EntryKey, ht: u32) -> Option<Id> {
          if ht == 0 { // persist page
              let search_result = node.keys()
                  .search(key);
              let index = if let Ok(i) = search_result { i } else { return None };
              match &node.delimiters().get(index) {
                  Some(Delimiter::External(id)) => Some(*id),
                  Some(Delimiter::Internal(_)) => panic!("Error on leaf delimiter type"),
                  None => panic!("Error on leaf delimiter is none")
              }
          } else {
              let index = node.keys()
                  .search(key)
                  .map(|i| i + 1)
                  .unwrap_or_else(|i| i);
              match &node.delimiters().get(index) {
                  Some(Delimiter::External(ref id)) => {
                      assert_eq!(ht, 0);
                      // search in leaf page
                      let page = self.get_page(id);
                      return self.search(&*page, key, ht - 1);
                  },
                  Some(Delimiter::Internal(node)) => return self.search(node.borrow(), key, ht - 1),
                  None => return None
              }
          }
      }
     fn insert(&self, key: &EntryKey, value: &Id) {
         unimplemented!()
     }
     fn put<'a>(&self, node: &'a Node, key: &EntryKey, ht: u32) {
         let keys = node.keys();
         let index = keys
             .search(key)
             .map(|i| i + 1)
             .unwrap_or_else(|i| i);
         if ht == 0 {

         }
     }
 }

enum Delimiter {
    External(Id), // to leaf
    Internal(Box<Node>) // to higher level node
}

trait Array<T> where T: Ord {
    #[inline]
    fn index_of(&self, index: usize) -> &T;
    #[inline]
    fn size(&self) -> usize;
    #[inline]
    fn search(&self, item: &T) -> Result<usize, usize>;
    #[inline]
    fn get(&self, index: usize) -> Option<&T>;
}

trait Keys : Array<EntryKey> {}
trait Delimiters : Array<Delimiter> {}
impl Keys for [EntryKey] {}
impl Delimiters for [Delimiter] {}
impl <T> Array<T> for [T] where T: Ord {
    #[inline]
    fn index_of(&self, index: usize) -> &T { &self[index] }

    #[inline]
    fn size(&self) -> usize { self.len() }

    #[inline]
    fn search(&self, item: &T) -> Result<usize, usize> { self.binary_search(item) }

    #[inline]
    fn get(&self, index: usize) -> Option<&T> {
        if index >= self.size() { None } else { Some(&self[index]) }
    }
}
impl <T> Array<T> for Vec<T> where T: Ord {
    #[inline]
    fn index_of(&self, index: usize) -> &T {
        self.get(index).unwrap()
    }

    #[inline]
    fn size(&self) -> usize {
        self.len()
    }

    #[inline]
    fn search(&self, item: &T) -> Result<usize, usize> {
        self.as_slice().binary_search(item)
    }

    #[inline]
    fn get(&self, index: usize) -> Option<&T> {
        if index >= self.len() { None } else { Some(&self[index]) }
    }
}
impl Delimiters for Vec<Delimiter> {}
impl Keys for Vec<EntryKey> {}

trait Node {
    #[inline]
    fn keys(&self) -> &Keys;
    #[inline]
    fn delimiters(&self) -> &Delimiters;
    #[inline]
    fn add(&mut self, key: EntryKey) -> Option<Node>;
    #[inline]
    fn del(&mut self, key: &EntryKey);
    #[inline]
    fn merge(&mut self, x: Node);
}

struct CachedExtNode {
    keys: Vec<EntryKey>,
    ids: Vec<Delimiter>,
    cap: u32,
}

impl Node for CachedExtNode {
    #[inline]
    fn keys(&self) -> &Keys {
        &self.keys
    }

    #[inline]
    fn delimiters(&self) -> &Delimiters {
        &self.ids
    }

    fn add(&mut self, key: EntryKey) -> Option<Self> {
        let insert_pos = match self.keys.binary_search(&key) {
            Ok(_) => return None, // existed
            Err(i) => i
        };
        let id = id_from_key(&key);
        self.keys.insert(insert_pos, key);
        self.ids.insert(insert_pos, Delimiter::External(id));
        if self.keys.len() > self.cap as usize {
            // need to split
            let mid = self.keys.len() / 2;
            let keys_2: Vec<_> = self.keys.drain(mid..).collect();
            let ids_2: Vec<_> = self.ids.drain(mid..).collect();
            return Some(CachedExtNode {
                keys: keys_2,
                ids: ids_2,
                cap: self.cap
            });
        } else {
            return None;
        }
    }

    fn del(&mut self, key: &EntryKey) {
        let pos = match self.keys.binary_search(&key) {
            Ok(i) => i, // existed
            Err(i) => return
        };
        self.keys.remove(pos);
        self.ids.remove(pos);
    }

    fn merge(&mut self, mut x: Self) {
        self.keys.append(&mut x.keys);
        self.ids.append(&mut x.ids);
    }
}

trait InterNode {

}

impl <T> Index<usize> for Array<T> where T: Ord {
    type Output = T;

    fn index(&self, index: usize) -> &<Self as Index<usize>>::Output {
        self.index_of(index)
    }
}

macro_rules! impl_nodes {
    ($(($level: ident, $entry_size: expr, $delimiter_size: expr)),+) => {
        $(
            mod $level {
                use super::*;

                type LEntryKeys = [EntryKey; $entry_size];

                impl Array<EntryKey> for LEntryKeys {
                    #[inline]
                    fn size(&self) -> usize { $entry_size }
                    #[inline]
                    fn index_of(&self, index: usize) -> &EntryKey { &self[index] }
                    #[inline]
                    fn search(&self, item: &EntryKey) -> Result<usize, usize> {
                        self.binary_search(item)
                    }
                    #[inline]
                    fn get(&self, index: usize) -> Option<&EntryKey> {
                        if index >= $entry_size { None } else { Some(&self[index]) }
                    }
                }

                type LDelimiter = [Delimiter; $delimiter_size];

                impl Array<Delimiter> for LDelimiter {
                    #[inline]
                    fn size(&self) -> usize { $delimiter_size }
                    #[inline]
                    fn index_of(&self, index: usize) -> &Delimiter { &self[index] }
                    #[inline]
                    fn search(&self, item: &Delimiter) -> Result<usize, usize> {
                        self.binary_search(item)
                    }
                    #[inline]
                    fn get(&self, index: usize) -> Option<&Delimiter> {
                        if index >= $entry_size { None } else { Some(&self[index]) }
                    }
                }

                impl Delimiters for LDelimiter {}
                impl Keys for LEntryKeys {}

                struct LNode {
                    keys: LEntryKeys,
                    delimeters: LDelimiter,
                    pos: u32
                }

                impl Node for LNode {
                    #[inline]
                    fn keys(&self) -> &Keys {
                        &self.keys
                    }
                    #[inline]
                    fn delimiters(&self) -> &Delimiters {
                        &self.delimeters
                    }
                    #[inline]
                    fn add(&mut self, key: EntryKey) -> Option<Node> {
                        unimplemented!()
                    }
                    #[inline]
                    fn del(&mut self, key: &EntryKey) {
                        unimplemented!()
                    }
                    #[inline]
                    fn merge(&mut self, x: Self) {
                        unimplemented!()
                    }
                }
            }
        )+
    };
}

impl_nodes!((level_0, 2, 3), (level_1, 20, 21), (level_2, 200, 201), (level_3, 2000, 2001), (level_4, 20000, 20001));

impl Ord for Delimiter {
    fn cmp(&self, other: &Self) -> Ordering {
        panic!()
    }
}
impl PartialOrd for Delimiter {
    fn partial_cmp(&self, other: &Delimiter) -> Option<Ordering> {
        panic!()
    }
}
impl Eq for Delimiter {}
impl PartialEq for Delimiter {
    fn eq(&self, other: &Delimiter) -> bool {
        panic!()
    }
}

fn id_from_key(key: &EntryKey) -> Id {
    let mut id_cursor = Cursor::new(&key[key.len() - 16 ..]);
    return Id::from_binary(&mut id_cursor).unwrap(); // read id from tailing 128 bits
}<|endoftext|>