<commit_before>// A genetic algorithm in Rust, for testing of stochastic methods for
// accelerating rigorous global maximization.

// Adapted from https://github.com/andschwa/rust-genetic-algorithm
// By Andy Schwartzmeyer (https://github.com/andschwa)

extern crate time;

extern crate rand;
use rand::{Rng};

extern crate gelpia_utils;
use gelpia_utils::{Parameters, INF, Flt};

extern crate gr;
use gr::{GI};

use rand::distributions::{Range, IndependentSample};

use std::cmp::{Eq, PartialEq, Ordering, PartialOrd};

use std::mem;

use std::sync::{Barrier, RwLock, Arc};

use std::sync::atomic::{AtomicBool};
use std::sync::atomic::Ordering as AtOrd;

extern crate function;
use function::FuncObj;

/// A genetic algorithm that searches for convergence to the given
/// tolerance for the problem across the n-dimensional hypercube,
/// using a population of individuals, up to a maximum iterations
/// number of generations.
pub fn ea(x_0: Vec<GI>, params: Parameters,
          population: Arc<RwLock<Vec<Individual>>>,
          f_bestag: Arc<RwLock<Flt>>,
          x_bestbb: Arc<RwLock<Vec<GI>>>,
          b1: Arc<Barrier>, b2: Arc<Barrier>,
          stop: Arc<AtomicBool>, sync: Arc<AtomicBool>, f: FuncObj) {
    // SETUP
    // get thread local random number generator
    let mut rng = rand::thread_rng();

    // initialize population of individuals
    {
        let mut population_w = population.write().unwrap();
        for _ in 0..params.population {
            population_w.push(Individual::new(&x_0, &mut rng, &f));
        }
        assert!(population_w.len() == params.population);
    }

    // let mut running: f64 = 0.0;
    // let mut iters: u64 = 0;

    while !stop.load(AtOrd::Acquire) {
        if sync.load(AtOrd::Acquire) {
            b1.wait();
            b2.wait();
        }

        let mut population = population.write().unwrap();
        {
            let mut fbest = f_bestag.write().unwrap();
            *fbest = match (*population).iter().max() {
                Some(i) => i.fitness,
                None => panic!(),
                }
        }
        // select, mutate, and crossover individuals for next generation
        let mut offspring: Vec<Individual> = Vec::with_capacity(population.len());
        for _ in 0..population.len()/2 {
            let (mut x, mut y) = (select(&population, params.selection, &mut rng),
                                  select(&population, params.selection, &mut rng));
            x.mutate(&x_0, params.mutation, &mut rng);
            y.mutate(&x_0, params.mutation, &mut rng);
            Individual::crossover(&mut x, &mut y, params.crossover, &mut rng, &f);
            offspring.push(x);
            offspring.push(y);
        }

        // replace 2 random individuals with elite of prior generation
        for _ in 0..params.elitism {
            if let Some(x) = population.iter().max() {
                offspring[rng.gen_range(0, population.len())] = x.clone();
            }
        }

        // Replace worst with bestbb
        let mut worst_i = 0;
        let mut worst = INF;
        for i in 0..offspring.len() {
            if offspring[i].fitness < worst {
                worst = offspring[i].fitness;
                worst_i = i;
            }
        }
        {
            let bestbb = x_bestbb.read().unwrap();
            offspring[worst_i] = Individual::new(&bestbb, &mut rng, &f);
        }

        // replace population with next generation
        *population = offspring;

    }
}

/// Tournament selection from n random individuals
fn select<R: Rng>(population: &[Individual], n: usize, rng: &mut R) -> Individual {
    if let Some(selected) = (0..n).map(|_| rng.choose(population)).max() {
        selected.unwrap().clone()
    } else {
        unimplemented!();
    }
}

#[derive(Clone)]
pub struct Individual {
    pub solution: Vec<GI>,
    pub fitness: Flt,
}

impl Individual {
    /// Constructs a new Individual to solve Problem with n random values
    pub fn new<R: Rng>(x: &Vec<GI>, rng: &mut R, f: &FuncObj) -> Self {
        let mut result = Vec::new();
        for i in 0..x.len() {
            let up = x[i].upper();
            let low = x[i].lower();

            let num = if
                up == low {up}
            else
            {Range::new(low,
                        up).ind_sample(rng)};
            result.push(GI::new_d(num, num));
        }
        let fitness = f.call(&result).lower();
        Individual{solution: result, fitness: fitness}
    }

    /// Mutate with chance n a single gene to a new value in the
    /// problem's domain (a "jump" mutation).
    ///
    /// Fitness is NOT evaluated as it is ALWAYS done in `crossover()`
    pub fn mutate<R: Rng>(&mut self, x: &Vec<GI>, chance: f64, rng: &mut R) {
        if rng.gen_range(0_f64, 1_f64) < chance {
            let i = rng.gen_range(0, self.solution.len());
            let up = x[i].upper();
            let low = x[i].lower();

            let num = if up == low {up} else {Range::new(low, up).ind_sample(rng)};

            self.solution[i] = GI::new_d(num, num);
        }
    }

    /// Performs two-point crossover with chance n to swap a random
    /// set of [0, dimension] genes between a pair of individuals.
    ///
    /// Fitness is ALWAYS evaluated because it is NOT done in mutate()
    pub fn crossover<R: Rng>(x: &mut Individual, y: &mut Individual,
                             chance: f64, rng: &mut R, f: &FuncObj) {
        // assert_eq!(x.problem, y.problem);
        if rng.gen_range(0_f64, 1_f64) < chance {
            let len = x.solution.len();
            let (start, n) = (rng.gen_range(0, len), rng.gen_range(0, len));
            for i in start..start + n {
                mem::swap(&mut x.solution[i % len], &mut y.solution[i % len]);
            }
        }
        x.fitness = f.call(&x.solution).lower();
        y.fitness = f.call(&y.solution).lower();
    }
}

impl Eq for Individual {}

impl Ord for Individual {
    /// This dangerously delegates to `partial_cmp`; NaN will panic
    fn cmp(&self, other: &Self) -> Ordering {
        if let Some(result) = self.fitness.partial_cmp(&other.fitness) {
            return result;
        }
        unimplemented!();
    }
}

impl PartialEq<Individual> for Individual {
    /// This doesn't use `fitness.eq()` because it needs to be
    /// consistent with `Eq`
    fn eq(&self, other: &Individual) -> bool {
        if let Some(result) = self.fitness.partial_cmp(&other.fitness) {
            return result == Ordering::Equal;
        }
        unimplemented!();
    }
}

impl PartialOrd for Individual {
    /// This delegates to `fitness.partial_cmp()`
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.fitness.partial_cmp(&other.fitness)
    }
}
<commit_msg>New genetic algorithm<commit_after>
use std::sync::{Barrier, RwLock, Arc, RwLockWriteGuard};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::atomic::Ordering as AtOrd;
extern crate rand;
use rand::{ThreadRng, Rng};
use rand::distributions::{IndependentSample, Range};

extern crate gr;
use gr::{GI, width_box, split_box, midpoint_box, eps_tol};

extern crate gelpia_utils;
use gelpia_utils::{Quple, INF, NINF, Flt, Parameters};

extern crate function;
use function::FuncObj;

#[derive(Clone)]
pub struct Individual {
    pub solution: Vec<GI>,
    pub fitness: Flt,
}

pub fn ea(x_e: Vec<GI>,
          param: Parameters,
          population: Arc<RwLock<Vec<Individual>>>,
          f_bestag: Arc<RwLock<Flt>>,
          x_bestbb: Arc<RwLock<Vec<GI>>>,
          b1: Arc<Barrier>,
          b2: Arc<Barrier>,
          stop: Arc<AtomicBool>,
          sync: Arc<AtomicBool>,
          fo_c: FuncObj) -> (Flt, Vec<GI>, bool) {

    let input = ea_core(&x_e, &param, &stop, &sync, &b1, &b2, &f_bestag,
                        &x_bestbb, population, &fo_c);
    let ans = fo_c.call(&input).upper();

    (ans, input, true)
}

fn ea_core(x_e: &Vec<GI>, param: &Parameters, stop: &Arc<AtomicBool>,
           sync: &Arc<AtomicBool>, b1: &Arc<Barrier>, b2: &Arc<Barrier>,
           f_bestag: &Arc<RwLock<Flt>>,
           x_bestbb: &Arc<RwLock<Vec<GI>>>,
           population: Arc<RwLock<Vec<Individual>>>, fo_c: &FuncObj)
           -> (Vec<GI>) {
    let mut rng = rand::thread_rng();
    let dimension = Range::new(0, x_e.len());
    let mut ranges = Vec::new();
    for g in x_e {
        ranges.push(Range::new(g.lower(), g.upper()));
    }

    while !stop.load(AtOrd::Acquire) {
        if sync.load(AtOrd::Acquire) {
            b1.wait();
            b2.wait();
        }
        let ref mut population = *population.write().unwrap();
        sample(param.population, population, fo_c, &ranges, &mut rng);

        population.sort_by(|a, b| b.fitness.partial_cmp(&a.fitness).unwrap());

        for iteration in 0..100 {
            population.truncate(param.elitism);

            for _ in 0..param.selection {
                population.push(rand_individual(fo_c, &ranges, &mut rng));
            }

            next_generation(param.population, population, fo_c, param.mutation,
                            param.crossover, &dimension, &ranges, &mut rng);
            population.sort_by(|a, b| b.fitness.partial_cmp(&a.fitness).unwrap());

            // Report fittest of the fit.
            {
                let mut fbest = f_bestag.write().unwrap();

                *fbest =
                    if *fbest < population[0].fitness { population[0].fitness }
                else { *fbest };
            }

            // Kill worst of the worst
            let mut ftg = Vec::new();
            {
                let bestbb = x_bestbb.read().unwrap();
                // From The Gods
                for i in 0..bestbb.len() {
                    ftg.push(Range::new(bestbb[i].lower(), bestbb[i].upper()));
                }
            }
            let worst_ind = population.len() - 1;
            population[worst_ind] = rand_individual(fo_c,
                                                    &ftg,
                                                    &mut rng);
            population.sort_by(|a, b| b.fitness.partial_cmp(&a.fitness).unwrap());
        }
    }
    let ref population = *population.read().unwrap();
    let result = population[0].solution.clone();

    result
}

fn sample(population_size: usize, population: &mut Vec<Individual>,
          fo_c: &FuncObj, ranges: &Vec<Range<f64>>, rng: &mut ThreadRng)
          -> () {
    for _ in 0..population_size-population.len() {
        population.push(rand_individual(fo_c, ranges, rng));
    }

    ()
}

fn rand_individual(fo_c: &FuncObj, ranges: &Vec<Range<f64>>, rng: &mut ThreadRng)
                   -> (Individual) {
    let mut new_sol = Vec::new();
    for r in ranges {
        new_sol.push(GI::new_p(r.ind_sample(rng)));
    }
    let fitness = fo_c.call(&new_sol).lower();

    Individual{solution:new_sol, fitness:fitness}
}

fn next_generation(population_size:usize, population: &mut Vec<Individual>,
                   fo_c: &FuncObj, mut_rate: f64, crossover: f64,
                   dimension: &Range<usize>, ranges: &Vec<Range<f64>>,
                   rng: &mut ThreadRng)
                   -> () {

    let elites = population.clone();

    for _ in 0..population_size-population.len() {
        if rng.gen::<f64>() < crossover {
            population.push(breed(rng.choose(&elites).unwrap(),
                                  rng.choose(&elites).unwrap(),
                                  fo_c,
                                  dimension, rng));
        } else {
            population.push(mutate(rng.choose(&elites).unwrap(), fo_c, mut_rate,
                                   dimension, ranges, rng));
        }
    }

    ()
}

fn mutate(input: &Individual, fo_c: &FuncObj, mut_rate: f64,
          dimension: &Range<usize>, ranges: &Vec<Range<f64>>, rng: &mut ThreadRng)
          -> (Individual) {
    let mut output_sol = Vec::new();

    for (r, &ind) in ranges.iter().zip(input.solution.iter()) {
        output_sol.push(
            if rng.gen::<f64>() < mut_rate {
                ind
            } else {
                GI::new_p(r.ind_sample(rng))
            });
    }

    let fitness = fo_c.call(&output_sol).lower();

    Individual{solution: output_sol, fitness: fitness}
}

fn breed(parent1: &Individual, parent2: &Individual, fo_c: &FuncObj,
         dimention: &Range<usize>, rng: &mut ThreadRng) -> (Individual) {
    let mut child = parent1.clone();
    let crossover_point = dimention.ind_sample(rng);
    child.solution.truncate(crossover_point);
    let mut rest = parent2.clone().solution.split_off(crossover_point);
    child.solution.append(&mut rest);
    child.fitness = fo_c.call(&child.solution).lower();
    child
}
<|endoftext|>


<commit_before>use std::vec::Vec;

pub struct Machine {
    memory: Vec<u32>,
    counter: u32,
}

impl Machine {
    pub fn new() -> Machine {
        Machine {
            memory: vec![0; 2097152],
            counter: 0,
        }
    }

    pub fn with_data(data: Vec<u32>) -> Machine {
        Machine {
            memory: data,
            counter: 0,
        }
    }

    pub fn loc(&self) -> u32 {
        self.counter
    }

    pub fn dump(&self) -> &Vec<u32> {
        &self.memory
    }

    pub fn load(&mut self, data: Vec<u32>) {
        self.memory = data;
    }

    pub fn step(&mut self) {
        let counter = self.counter as usize;
        let opcode = self.memory[counter];
        let a = self.memory[counter + 1];
        let b = self.memory[counter + 2];
        let c = self.memory[counter + 3];
        self.exec(opcode, a, b, c);
    }

    pub fn exec(&mut self, opcode: u32, a: u32, b: u32, c: u32) {
        match opcode {
            // PC <- M[A]
            1 => {
                self.counter = self.memory[a as usize];
            },

            // If M[B] = 0, then PC <- M[A]
            2 => {
                if 0 == self.memory[b as usize] {
                  self.counter = self.memory[a as usize];
                }
            },

            // M[A] <- PC
            3 => {
                self.memory[a as usize] = self.counter;
                self.counter += 4;
            },

            // M[A] <- M[B]
            4 => {
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[A] <- M[M[B]]
            5 => {
                let b = self.memory[b as usize];
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[M[B]] <- M[A]
            6 => {
                let b = self.memory[b as usize];
                self.memory[b as usize] = self.memory[a as usize];
                self.counter += 4;
            },

            // M[A] <- M[B] + M[C]
            7 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b + c;
                self.counter += 4;
            },

            // M[A] <- M[B] - M[C]
            8 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b - c;
                self.counter += 4;
            },

            // M[A] <- M[B] x M[C]
            9 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b * c;
                self.counter += 4;
            },

            // M[A] <- M[B] / M[C]
            10 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b / c;
                self.counter += 4;
            },

            // M[A] <- M[B] modulo M[C]
            11 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b % c;
                self.counter += 4;
            },

            // If M[B] < M[C], then M[A] <- 1, else M[A] <- 0
            12 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                if b < c {
                  self.memory[a as usize] = 1;
                } else {
                  self.memory[a as usize] = 0;
                }
                self.counter += 4;
            },

            // Unknown opcode
            _ => {
            },
        }
    }
}

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

    #[test]
    fn it_loads_data() {
        let mut m = Machine::new();
        assert_eq!(Some(&0), m.dump().first());

        m.load(vec![1]);
        assert_eq!(Some(&1), m.dump().first());
    }

    #[test]
    fn it_runs_opcode_1() {
        // PC <- M[A]
        let mut m = Machine::with_data(vec![1, 3, 0, 2]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_then_branch() {
        // If M[B] = 0, then PC <- M[A]
        let mut m = Machine::with_data(vec![2, 4, 5, 0, 2, 0]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_else_branch() {
        // If M[B] = 0, then PC <- M[A]
        let mut m = Machine::with_data(vec![2, 4, 5, 0, 2, 1]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(0, m.loc());
    }

    #[test]
    fn it_runs_opcode_3() {
        let mut m = Machine::with_data(vec![1]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(3, 0, 0, 0);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_4() {
        let mut m = Machine::with_data(vec![1, 0]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(4, 0, 1, 0);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_5() {
        let mut m = Machine::with_data(vec![1, 2, 0]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(5, 0, 1, 0);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_6() {
        let mut m = Machine::with_data(vec![1, 0]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(6, 1, 1, 0);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_7() {
        let mut m = Machine::with_data(vec![1, 2, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(7, 0, 1, 2);
        assert_eq!(Some(&5), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_8() {
        let mut m = Machine::with_data(vec![0, 3, 2]);

        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(8, 0, 1, 2);
        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_9() {
        let mut m = Machine::with_data(vec![1, 2, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(9, 0, 1, 2);
        assert_eq!(Some(&6), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_10() {
        let mut m = Machine::with_data(vec![1, 6, 2]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(10, 0, 1, 2);
        assert_eq!(Some(&3), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_11() {
        let mut m = Machine::with_data(vec![1, 11, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(11, 0, 1, 2);
        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_12_then_branch() {
        let mut m = Machine::with_data(vec![2, 3, 4]);

        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(12, 0, 1, 2);
        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_12_else_branch() {
        let mut m = Machine::with_data(vec![2, 4, 3]);

        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(12, 0, 1, 2);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }
}
<commit_msg>Use four words for happy path opcode tests.<commit_after>use std::vec::Vec;

pub struct Machine {
    memory: Vec<u32>,
    counter: u32,
}

impl Machine {
    pub fn new() -> Machine {
        Machine {
            memory: vec![0; 2097152],
            counter: 0,
        }
    }

    pub fn with_data(data: Vec<u32>) -> Machine {
        Machine {
            memory: data,
            counter: 0,
        }
    }

    pub fn loc(&self) -> u32 {
        self.counter
    }

    pub fn dump(&self) -> &Vec<u32> {
        &self.memory
    }

    pub fn load(&mut self, data: Vec<u32>) {
        self.memory = data;
    }

    pub fn step(&mut self) {
        let counter = self.counter as usize;
        let opcode = self.memory[counter];
        let a = self.memory[counter + 1];
        let b = self.memory[counter + 2];
        let c = self.memory[counter + 3];
        self.exec(opcode, a, b, c);
    }

    pub fn exec(&mut self, opcode: u32, a: u32, b: u32, c: u32) {
        match opcode {
            // PC <- M[A]
            1 => {
                self.counter = self.memory[a as usize];
            },

            // If M[B] = 0, then PC <- M[A]
            2 => {
                if 0 == self.memory[b as usize] {
                  self.counter = self.memory[a as usize];
                }
            },

            // M[A] <- PC
            3 => {
                self.memory[a as usize] = self.counter;
                self.counter += 4;
            },

            // M[A] <- M[B]
            4 => {
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[A] <- M[M[B]]
            5 => {
                let b = self.memory[b as usize];
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[M[B]] <- M[A]
            6 => {
                let b = self.memory[b as usize];
                self.memory[b as usize] = self.memory[a as usize];
                self.counter += 4;
            },

            // M[A] <- M[B] + M[C]
            7 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b + c;
                self.counter += 4;
            },

            // M[A] <- M[B] - M[C]
            8 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b - c;
                self.counter += 4;
            },

            // M[A] <- M[B] x M[C]
            9 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b * c;
                self.counter += 4;
            },

            // M[A] <- M[B] / M[C]
            10 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b / c;
                self.counter += 4;
            },

            // M[A] <- M[B] modulo M[C]
            11 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b % c;
                self.counter += 4;
            },

            // If M[B] < M[C], then M[A] <- 1, else M[A] <- 0
            12 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                if b < c {
                  self.memory[a as usize] = 1;
                } else {
                  self.memory[a as usize] = 0;
                }
                self.counter += 4;
            },

            // Unknown opcode
            _ => {
            },
        }
    }
}

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

    #[test]
    fn it_loads_data() {
        let mut m = Machine::new();
        assert_eq!(Some(&0), m.dump().first());

        m.load(vec![1]);
        assert_eq!(Some(&1), m.dump().first());
    }

    #[test]
    fn it_runs_opcode_1() {
        // PC <- M[A]
        let mut m = Machine::with_data(vec![1, 4, 0, 0, 2]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_then_branch() {
        // If M[B] = 0, then PC <- M[A]
        let mut m = Machine::with_data(vec![2, 4, 5, 0, 2, 0]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_else_branch() {
        // If M[B] = 0, then PC <- M[A]
        let mut m = Machine::with_data(vec![2, 4, 5, 0, 2, 1]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(0, m.loc());
    }

    #[test]
    fn it_runs_opcode_3() {
        let mut m = Machine::with_data(vec![1]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(3, 0, 0, 0);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_4() {
        let mut m = Machine::with_data(vec![1, 0]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(4, 0, 1, 0);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_5() {
        let mut m = Machine::with_data(vec![1, 2, 0]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(5, 0, 1, 0);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_6() {
        let mut m = Machine::with_data(vec![1, 0]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(6, 1, 1, 0);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_7() {
        let mut m = Machine::with_data(vec![1, 2, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(7, 0, 1, 2);
        assert_eq!(Some(&5), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_8() {
        let mut m = Machine::with_data(vec![0, 3, 2]);

        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(8, 0, 1, 2);
        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_9() {
        let mut m = Machine::with_data(vec![1, 2, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(9, 0, 1, 2);
        assert_eq!(Some(&6), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_10() {
        let mut m = Machine::with_data(vec![1, 6, 2]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(10, 0, 1, 2);
        assert_eq!(Some(&3), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_11() {
        let mut m = Machine::with_data(vec![1, 11, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(11, 0, 1, 2);
        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_12_then_branch() {
        let mut m = Machine::with_data(vec![2, 3, 4]);

        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(12, 0, 1, 2);
        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_12_else_branch() {
        let mut m = Machine::with_data(vec![2, 4, 3]);

        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(12, 0, 1, 2);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }
}
<|endoftext|>


<commit_before>use ops::{CommutativeOperation, Inverse};
use std::marker::PhantomData;

use std::default::Default;
use std::hash::{Hash, Hasher};

/// This data structure allows prefix queries and single element modification.
///
/// This tree allocates `n * sizeof(N)` bytes of memory, and can be resized.
///
/// This data structure is implemented using a fenwick tree, which is also known as a binary
/// indexed tree.
///
///# Examples
///
/// Showcase of functionality:
///
///```rust
///use segment_tree::ops::Add;
///use segment_tree::PrefixPoint;
///
///let buf = vec![10, 5, 30, 40];
///
///let mut pp: PrefixPoint<_, Add> = PrefixPoint::build(buf);
///
/// // If we query, we get the sum up until the specified value.
///assert_eq!(pp.query(0), 10);
///assert_eq!(pp.query(1), 15);
///assert_eq!(pp.query(2), 45);
///assert_eq!(pp.query(3), 85);
///
/// // Add five to the second value.
///pp.modify(1, 5);
///assert_eq!(pp.query(0), 10);
///assert_eq!(pp.query(1), 20);
///assert_eq!(pp.query(2), 50);
///assert_eq!(pp.query(3), 90);
///
/// // Multiply every value with 2.
///pp.map(|v| *v *= 2);
///assert_eq!(pp.query(0), 20);
///assert_eq!(pp.query(1), 40);
///assert_eq!(pp.query(2), 100);
///assert_eq!(pp.query(3), 180);
///
/// // Divide with two to undo.
///pp.map(|v| *v /= 2);
/// // Add some more values.
///pp.append(vec![0, 10]);
///assert_eq!(pp.query(0), 10);
///assert_eq!(pp.query(1), 20);
///assert_eq!(pp.query(2), 50);
///assert_eq!(pp.query(3), 90);
///assert_eq!(pp.query(4), 90);
///assert_eq!(pp.query(5), 100);
///
/// // Get the values.
///assert_eq!(pp.get(0), 10);
///assert_eq!(pp.get(1), 10);
///assert_eq!(pp.get(2), 30);
///assert_eq!(pp.get(3), 40);
///assert_eq!(pp.get(4), 0);
///assert_eq!(pp.get(5), 10);
///
/// // Remove the last value
///pp.truncate(5);
///assert_eq!(pp.get(0), 10);
///assert_eq!(pp.get(1), 10);
///assert_eq!(pp.get(2), 30);
///assert_eq!(pp.get(3), 40);
///assert_eq!(pp.get(4), 0);
///
/// // Get back the original values.
///assert_eq!(pp.unwrap(), vec![10, 10, 30, 40, 0]);
///```
///
/// You can also use other operators:
///
///```rust
///use segment_tree::ops::Mul;
///use segment_tree::PrefixPoint;
///
///let buf = vec![10, 5, 30, 40];
///
///let mut pp: PrefixPoint<_, Mul> = PrefixPoint::build(buf);
///
///assert_eq!(pp.query(0), 10);
///assert_eq!(pp.query(1), 50);
///assert_eq!(pp.query(2), 1500);
///assert_eq!(pp.query(3), 60000);
///```
pub struct PrefixPoint<N, O> where O: CommutativeOperation<N> {
    buf: Vec<N>,
    op: PhantomData<O>
}

/// Returns the least significant bit that is one.
#[inline(always)]
fn lsb(i: usize) -> usize {
    i & (1 + !i)
}

/// Could also be done with slice_at_mut, but that's a giant pain
#[inline(always)]
unsafe fn combine_mut<N, O: CommutativeOperation<N>>(buf: &mut Vec<N>, i: usize, j: usize) {
    let ptr1 = &mut buf[i] as *mut N;
    let ptr2 = &buf[j] as *const N;
    O::combine_mut(&mut *ptr1, &*ptr2);
}
/// Could also be done with slice_at_mut, but that's a giant pain
#[inline(always)]
unsafe fn uncombine_mut<N, O: Inverse<N>>(buf: &mut Vec<N>, i: usize, j: usize) {
    let ptr1 = &mut buf[i] as *mut N;
    let ptr2 = &buf[j] as *const N;
    O::uncombine(&mut *ptr1, &*ptr2);
}

impl<N, O: CommutativeOperation<N>> PrefixPoint<N, O> {
    /// Creates a `PrefixPoint` containing the given values.
    /// Uses `O(len)` time.
    pub fn build(mut buf: Vec<N>) -> PrefixPoint<N, O> {
        let len = buf.len();
        for i in 0..len {
            let j = i + lsb(i+1);
            if j < len {
                unsafe {
                    combine_mut::<N, O>(&mut buf, j, i);
                }
            }
        }
        PrefixPoint { buf: buf, op: PhantomData }
    }
    /// Returns the number of values in this tree.
    /// Uses `O(1)` time.
    pub fn len(&self) -> usize {
        self.buf.len()
    }
    /// Computes `a[0] * a[1] * ... * a[i]`.  Note that `i` is inclusive.
    /// Uses `O(log(i))` time.
    #[inline]
    pub fn query(&self, mut i: usize) -> N where N: Clone {
        let mut sum = self.buf[i].clone();
        i -= lsb(1+i) - 1;
        while i > 0 {
            sum = O::combine_left(sum, &self.buf[i-1]);
            i -= lsb(i);
        }
        sum
    }
    /// Combine the value at `i` with `delta`.
    /// Uses `O(log(len))` time.
    #[inline]
    pub fn modify(&mut self, mut i: usize, delta: N) {
        let len = self.len();
        while i < len {
            O::combine_mut(&mut self.buf[i], &delta);
            i += lsb(i+1);
        }
    }
    /// Truncates the `PrefixPoint` to the given size.  If `size >= len`, this method does nothing.
    /// Uses `O(1)` time.
    #[inline(always)]
    pub fn truncate(&mut self, size: usize) {
        self.buf.truncate(size);
    }
    /// Replace every value in the type with `f(value)`.
    /// This function assumes that `f(a) * f(b) = f(a * b)`.
    /// Applies the function `len` times.
    #[inline]
    pub fn map<F: FnMut(&mut N)>(&mut self, mut f: F) {
        for val in &mut self.buf {
            f(val);
        }
    }
    /// Adds the given values to the `PrefixPoint`, increasing its size.
    /// Uses `O(len)` time.
    #[inline]
    pub fn append(&mut self, mut values: Vec<N>) {
        self.extend(values.drain(..))
    }
}
impl<N, O: CommutativeOperation<N>> Extend<N> for PrefixPoint<N, O> {
    /// Adds the given values to the `PrefixPoint`, increasing its size.
    /// Uses `O(len)` time.
    fn extend<I: IntoIterator<Item=N>>(&mut self, values: I) {
        let oldlen = self.len();
        self.buf.extend(values);
        let len = self.len();
        for i in 0..len {
            let j = i + lsb(i+1);
            if oldlen <= j && j < len {
                unsafe {
                    combine_mut::<N, O>(&mut self.buf, j, i);
                }
            }
        }
    }
}
impl<N, O: CommutativeOperation<N> + Inverse<N>> PrefixPoint<N, O> {
    /// Returns the value at `i`.
    /// Uses `O(log(i))` time.
    /// Store your own copy of the array if you want constant time.
    pub fn get(&self, mut i: usize) -> N where N: Clone {
        let mut sum = self.buf[i].clone();
        let z = 1 + i - lsb(i+1);
        while i != z {
            O::uncombine(&mut sum, &self.buf[i-1]);
            i -= lsb(i);
        }
        sum
    }
    /// Compute the underlying array of values.
    /// Uses `O(len)` time.
    pub fn unwrap(self) -> Vec<N> {
        let mut buf = self.buf;
        let len = buf.len();
        for i in (0..len).rev() {
            let j = i + lsb(i+1);
            if j < len {
                unsafe {
                    uncombine_mut::<N, O>(&mut buf, j, i);
                }
            }
        }
        buf
    }
}

impl<N: Clone, O: CommutativeOperation<N>> Clone for PrefixPoint<N, O> {
    fn clone(&self) -> PrefixPoint<N, O> {
        PrefixPoint {
            buf: self.buf.clone(), op: PhantomData
        }
    }
}
impl<N, O: CommutativeOperation<N>> Default for PrefixPoint<N, O> {
    #[inline]
    fn default() -> PrefixPoint<N, O> {
        PrefixPoint { buf: Vec::new(), op: PhantomData }
    }
}
impl<'a, N: 'a + Hash, O: CommutativeOperation<N>> Hash for PrefixPoint<N, O> {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.buf.hash(state);
    }
}

#[cfg(test)]
mod tests {

    /// Modifies the given slice such that the n'th element becomes the sum of the first n elements.
    pub fn compute_prefix_sum<N: ::std::ops::Add<Output=N> + Copy>(buf: &mut[N]) {
        let mut iter = buf.iter_mut();
        match iter.next() {
            None => {},
            Some(s) => {
                let mut sum = *s;
                for item in iter {
                    sum = sum + *item;
                    *item = sum;
                }
            }
        }
    }

    use super::*;
    use rand::{Rng, thread_rng};
    use std::num::Wrapping;
    use ops::Add;

    #[test]
    fn fenwick_query() {
        let mut rng = thread_rng();
        for n in 0..130 {
            let mut vec: Vec<Wrapping<i32>> = rng.gen_iter::<i32>().take(n).map(|i| Wrapping(i)).collect();
            let fenwick: PrefixPoint<_, Add> = PrefixPoint::build(vec.clone());
            compute_prefix_sum(&mut vec);
            for i in 0..vec.len() {
                assert_eq!(vec[i], fenwick.query(i));
            }
        }
    }
    #[test]
    fn fenwick_map() {
        let mut rng = thread_rng();
        for n in 0..130 {
            let mut vec: Vec<Wrapping<i32>> = rng.gen_iter::<i32>().take(n).map(|i| Wrapping(i)).collect();
            let mut fenwick: PrefixPoint<_, Add> = PrefixPoint::build(vec.clone());
            assert_eq!(fenwick.clone().unwrap(), vec);
            compute_prefix_sum(&mut vec);
            fenwick.map(|n| *n = Wrapping(12) * *n);
            for i in 0..vec.len() {
                assert_eq!(vec[i]*Wrapping(12), fenwick.query(i));
            }
        }
    }
    #[test]
    fn fenwick_modify() {
        let mut rng = thread_rng();
        for n in 0..130 {
            let mut vec: Vec<Wrapping<i32>> = rng.gen_iter::<i32>().take(n).map(|i| Wrapping(i)).collect();
            let diff: Vec<Wrapping<i32>> = rng.gen_iter::<i32>().take(n).map(|i| Wrapping(i)).collect();
            let mut fenwick: PrefixPoint<_, Add> = PrefixPoint::build(vec.clone());
            for i in 0..diff.len() {
                let mut ps: Vec<Wrapping<i32>> = vec.clone();
                compute_prefix_sum(&mut ps);
                assert_eq!(fenwick.clone().unwrap(), vec);
                for j in 0..vec.len() {
                    assert_eq!(ps[j], fenwick.query(j));
                    assert_eq!(vec[j], fenwick.get(j));
                }
                vec[i] += diff[i];
                fenwick.modify(i, diff[i]);
            }
        }
    }
    #[test]
    fn fenwick_extend() {
        let mut rng = thread_rng();
        for n in 0..130 {
            let vec: Vec<Wrapping<i32>> = rng.gen_iter::<i32>().take(n).map(|i| Wrapping(i)).collect();
            let mut sum = vec.clone();
            compute_prefix_sum(&mut sum);
            for i in 0..sum.len() {
                let mut fenwick: PrefixPoint<_, Add> = PrefixPoint::build(vec.iter().take(i/2).map(|&i| i).collect());
                fenwick.extend(vec.iter().skip(i/2).take(i - i/2).map(|&i| i));
                for j in 0..i {
                    assert_eq!(sum[j], fenwick.query(j));
                }
            }
        }
    }
}
<commit_msg>Fenwick is a name, so it is with a capital letter<commit_after>use ops::{CommutativeOperation, Inverse};
use std::marker::PhantomData;

use std::default::Default;
use std::hash::{Hash, Hasher};

/// This data structure allows prefix queries and single element modification.
///
/// This tree allocates `n * sizeof(N)` bytes of memory, and can be resized.
///
/// This data structure is implemented using a Fenwick tree, which is also known as a binary
/// indexed tree.
///
///# Examples
///
/// Showcase of functionality:
///
///```rust
///use segment_tree::ops::Add;
///use segment_tree::PrefixPoint;
///
///let buf = vec![10, 5, 30, 40];
///
///let mut pp: PrefixPoint<_, Add> = PrefixPoint::build(buf);
///
/// // If we query, we get the sum up until the specified value.
///assert_eq!(pp.query(0), 10);
///assert_eq!(pp.query(1), 15);
///assert_eq!(pp.query(2), 45);
///assert_eq!(pp.query(3), 85);
///
/// // Add five to the second value.
///pp.modify(1, 5);
///assert_eq!(pp.query(0), 10);
///assert_eq!(pp.query(1), 20);
///assert_eq!(pp.query(2), 50);
///assert_eq!(pp.query(3), 90);
///
/// // Multiply every value with 2.
///pp.map(|v| *v *= 2);
///assert_eq!(pp.query(0), 20);
///assert_eq!(pp.query(1), 40);
///assert_eq!(pp.query(2), 100);
///assert_eq!(pp.query(3), 180);
///
/// // Divide with two to undo.
///pp.map(|v| *v /= 2);
/// // Add some more values.
///pp.append(vec![0, 10]);
///assert_eq!(pp.query(0), 10);
///assert_eq!(pp.query(1), 20);
///assert_eq!(pp.query(2), 50);
///assert_eq!(pp.query(3), 90);
///assert_eq!(pp.query(4), 90);
///assert_eq!(pp.query(5), 100);
///
/// // Get the values.
///assert_eq!(pp.get(0), 10);
///assert_eq!(pp.get(1), 10);
///assert_eq!(pp.get(2), 30);
///assert_eq!(pp.get(3), 40);
///assert_eq!(pp.get(4), 0);
///assert_eq!(pp.get(5), 10);
///
/// // Remove the last value
///pp.truncate(5);
///assert_eq!(pp.get(0), 10);
///assert_eq!(pp.get(1), 10);
///assert_eq!(pp.get(2), 30);
///assert_eq!(pp.get(3), 40);
///assert_eq!(pp.get(4), 0);
///
/// // Get back the original values.
///assert_eq!(pp.unwrap(), vec![10, 10, 30, 40, 0]);
///```
///
/// You can also use other operators:
///
///```rust
///use segment_tree::ops::Mul;
///use segment_tree::PrefixPoint;
///
///let buf = vec![10, 5, 30, 40];
///
///let mut pp: PrefixPoint<_, Mul> = PrefixPoint::build(buf);
///
///assert_eq!(pp.query(0), 10);
///assert_eq!(pp.query(1), 50);
///assert_eq!(pp.query(2), 1500);
///assert_eq!(pp.query(3), 60000);
///```
pub struct PrefixPoint<N, O> where O: CommutativeOperation<N> {
    buf: Vec<N>,
    op: PhantomData<O>
}

/// Returns the least significant bit that is one.
#[inline(always)]
fn lsb(i: usize) -> usize {
    i & (1 + !i)
}

/// Could also be done with slice_at_mut, but that's a giant pain
#[inline(always)]
unsafe fn combine_mut<N, O: CommutativeOperation<N>>(buf: &mut Vec<N>, i: usize, j: usize) {
    let ptr1 = &mut buf[i] as *mut N;
    let ptr2 = &buf[j] as *const N;
    O::combine_mut(&mut *ptr1, &*ptr2);
}
/// Could also be done with slice_at_mut, but that's a giant pain
#[inline(always)]
unsafe fn uncombine_mut<N, O: Inverse<N>>(buf: &mut Vec<N>, i: usize, j: usize) {
    let ptr1 = &mut buf[i] as *mut N;
    let ptr2 = &buf[j] as *const N;
    O::uncombine(&mut *ptr1, &*ptr2);
}

impl<N, O: CommutativeOperation<N>> PrefixPoint<N, O> {
    /// Creates a `PrefixPoint` containing the given values.
    /// Uses `O(len)` time.
    pub fn build(mut buf: Vec<N>) -> PrefixPoint<N, O> {
        let len = buf.len();
        for i in 0..len {
            let j = i + lsb(i+1);
            if j < len {
                unsafe {
                    combine_mut::<N, O>(&mut buf, j, i);
                }
            }
        }
        PrefixPoint { buf: buf, op: PhantomData }
    }
    /// Returns the number of values in this tree.
    /// Uses `O(1)` time.
    pub fn len(&self) -> usize {
        self.buf.len()
    }
    /// Computes `a[0] * a[1] * ... * a[i]`.  Note that `i` is inclusive.
    /// Uses `O(log(i))` time.
    #[inline]
    pub fn query(&self, mut i: usize) -> N where N: Clone {
        let mut sum = self.buf[i].clone();
        i -= lsb(1+i) - 1;
        while i > 0 {
            sum = O::combine_left(sum, &self.buf[i-1]);
            i -= lsb(i);
        }
        sum
    }
    /// Combine the value at `i` with `delta`.
    /// Uses `O(log(len))` time.
    #[inline]
    pub fn modify(&mut self, mut i: usize, delta: N) {
        let len = self.len();
        while i < len {
            O::combine_mut(&mut self.buf[i], &delta);
            i += lsb(i+1);
        }
    }
    /// Truncates the `PrefixPoint` to the given size.  If `size >= len`, this method does nothing.
    /// Uses `O(1)` time.
    #[inline(always)]
    pub fn truncate(&mut self, size: usize) {
        self.buf.truncate(size);
    }
    /// Replace every value in the type with `f(value)`.
    /// This function assumes that `f(a) * f(b) = f(a * b)`.
    /// Applies the function `len` times.
    #[inline]
    pub fn map<F: FnMut(&mut N)>(&mut self, mut f: F) {
        for val in &mut self.buf {
            f(val);
        }
    }
    /// Adds the given values to the `PrefixPoint`, increasing its size.
    /// Uses `O(len)` time.
    #[inline]
    pub fn append(&mut self, mut values: Vec<N>) {
        self.extend(values.drain(..))
    }
}
impl<N, O: CommutativeOperation<N>> Extend<N> for PrefixPoint<N, O> {
    /// Adds the given values to the `PrefixPoint`, increasing its size.
    /// Uses `O(len)` time.
    fn extend<I: IntoIterator<Item=N>>(&mut self, values: I) {
        let oldlen = self.len();
        self.buf.extend(values);
        let len = self.len();
        for i in 0..len {
            let j = i + lsb(i+1);
            if oldlen <= j && j < len {
                unsafe {
                    combine_mut::<N, O>(&mut self.buf, j, i);
                }
            }
        }
    }
}
impl<N, O: CommutativeOperation<N> + Inverse<N>> PrefixPoint<N, O> {
    /// Returns the value at `i`.
    /// Uses `O(log(i))` time.
    /// Store your own copy of the array if you want constant time.
    pub fn get(&self, mut i: usize) -> N where N: Clone {
        let mut sum = self.buf[i].clone();
        let z = 1 + i - lsb(i+1);
        while i != z {
            O::uncombine(&mut sum, &self.buf[i-1]);
            i -= lsb(i);
        }
        sum
    }
    /// Compute the underlying array of values.
    /// Uses `O(len)` time.
    pub fn unwrap(self) -> Vec<N> {
        let mut buf = self.buf;
        let len = buf.len();
        for i in (0..len).rev() {
            let j = i + lsb(i+1);
            if j < len {
                unsafe {
                    uncombine_mut::<N, O>(&mut buf, j, i);
                }
            }
        }
        buf
    }
}

impl<N: Clone, O: CommutativeOperation<N>> Clone for PrefixPoint<N, O> {
    fn clone(&self) -> PrefixPoint<N, O> {
        PrefixPoint {
            buf: self.buf.clone(), op: PhantomData
        }
    }
}
impl<N, O: CommutativeOperation<N>> Default for PrefixPoint<N, O> {
    #[inline]
    fn default() -> PrefixPoint<N, O> {
        PrefixPoint { buf: Vec::new(), op: PhantomData }
    }
}
impl<'a, N: 'a + Hash, O: CommutativeOperation<N>> Hash for PrefixPoint<N, O> {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.buf.hash(state);
    }
}

#[cfg(test)]
mod tests {

    /// Modifies the given slice such that the n'th element becomes the sum of the first n elements.
    pub fn compute_prefix_sum<N: ::std::ops::Add<Output=N> + Copy>(buf: &mut[N]) {
        let mut iter = buf.iter_mut();
        match iter.next() {
            None => {},
            Some(s) => {
                let mut sum = *s;
                for item in iter {
                    sum = sum + *item;
                    *item = sum;
                }
            }
        }
    }

    use super::*;
    use rand::{Rng, thread_rng};
    use std::num::Wrapping;
    use ops::Add;

    #[test]
    fn fenwick_query() {
        let mut rng = thread_rng();
        for n in 0..130 {
            let mut vec: Vec<Wrapping<i32>> = rng.gen_iter::<i32>().take(n).map(|i| Wrapping(i)).collect();
            let fenwick: PrefixPoint<_, Add> = PrefixPoint::build(vec.clone());
            compute_prefix_sum(&mut vec);
            for i in 0..vec.len() {
                assert_eq!(vec[i], fenwick.query(i));
            }
        }
    }
    #[test]
    fn fenwick_map() {
        let mut rng = thread_rng();
        for n in 0..130 {
            let mut vec: Vec<Wrapping<i32>> = rng.gen_iter::<i32>().take(n).map(|i| Wrapping(i)).collect();
            let mut fenwick: PrefixPoint<_, Add> = PrefixPoint::build(vec.clone());
            assert_eq!(fenwick.clone().unwrap(), vec);
            compute_prefix_sum(&mut vec);
            fenwick.map(|n| *n = Wrapping(12) * *n);
            for i in 0..vec.len() {
                assert_eq!(vec[i]*Wrapping(12), fenwick.query(i));
            }
        }
    }
    #[test]
    fn fenwick_modify() {
        let mut rng = thread_rng();
        for n in 0..130 {
            let mut vec: Vec<Wrapping<i32>> = rng.gen_iter::<i32>().take(n).map(|i| Wrapping(i)).collect();
            let diff: Vec<Wrapping<i32>> = rng.gen_iter::<i32>().take(n).map(|i| Wrapping(i)).collect();
            let mut fenwick: PrefixPoint<_, Add> = PrefixPoint::build(vec.clone());
            for i in 0..diff.len() {
                let mut ps: Vec<Wrapping<i32>> = vec.clone();
                compute_prefix_sum(&mut ps);
                assert_eq!(fenwick.clone().unwrap(), vec);
                for j in 0..vec.len() {
                    assert_eq!(ps[j], fenwick.query(j));
                    assert_eq!(vec[j], fenwick.get(j));
                }
                vec[i] += diff[i];
                fenwick.modify(i, diff[i]);
            }
        }
    }
    #[test]
    fn fenwick_extend() {
        let mut rng = thread_rng();
        for n in 0..130 {
            let vec: Vec<Wrapping<i32>> = rng.gen_iter::<i32>().take(n).map(|i| Wrapping(i)).collect();
            let mut sum = vec.clone();
            compute_prefix_sum(&mut sum);
            for i in 0..sum.len() {
                let mut fenwick: PrefixPoint<_, Add> = PrefixPoint::build(vec.iter().take(i/2).map(|&i| i).collect());
                fenwick.extend(vec.iter().skip(i/2).take(i - i/2).map(|&i| i));
                for j in 0..i {
                    assert_eq!(sum[j], fenwick.query(j));
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Compiler plugin for Rust-PHF
//!
//! See the documentation for the `phf` crate for more details.
#[crate_id="github.com/sfackler/rust-phf/phf_mac"];
#[crate_type="dylib"];
#[doc(html_root_url="http://www.rust-ci.org/sfackler/rust-phf/doc")];
#[feature(managed_boxes, macro_registrar, quote)];

extern crate extra;
extern crate syntax;
extern crate phf;

use extra::time;
use std::hashmap::HashMap;
use std::rand;
use std::os;
use std::vec;
use syntax::ast;
use syntax::ast::{Name, TokenTree, LitStr, MutImmutable, Expr, ExprVec, ExprLit};
use syntax::codemap::Span;
use syntax::ext::base::{SyntaxExtension,
                        ExtCtxt,
                        MacResult,
                        MRExpr,
                        NormalTT,
                        BasicMacroExpander};
use syntax::parse;
use syntax::parse::token;
use syntax::parse::token::{InternedString, COMMA, EOF, FAT_ARROW};

static DEFAULT_LAMBDA: uint = 5;

#[macro_registrar]
#[doc(hidden)]
pub fn macro_registrar(register: |Name, SyntaxExtension|) {
    register(token::intern("phf_map"),
             NormalTT(~BasicMacroExpander {
                expander: expand_mphf_map,
                span: None
             },
             None));
}

struct Entry {
    key_str: InternedString,
    key: @Expr,
    value: @Expr
}

fn expand_mphf_map(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> MacResult {
    let mut entries = match parse_entries(cx, tts) {
        Some(entries) => entries,
        None => return MacResult::dummy_expr(sp)
    };

    entries.sort_by(|a, b| a.key_str.cmp(&b.key_str));
    if has_duplicates(cx, sp, entries) {
        return MacResult::dummy_expr(sp);
    }

    let start = time::precise_time_s();
    let state;
    loop {
        match generate_hash(entries) {
            Some(s) => {
                state = s;
                break;
            }
            None => {}
        }
    }
    let time = time::precise_time_s() - start;
    if os::getenv("PHF_STATS").is_some() {
        cx.span_note(sp, format!("PHF generation took {} seconds", time));
    }

    create_map(cx, sp, entries, state)
}

fn parse_entries(cx: &mut ExtCtxt, tts: &[TokenTree]) -> Option<~[Entry]> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(),
                                                tts.to_owned());
    let mut entries = ~[];

    let mut bad = false;
    while parser.token != EOF {
        let key = parser.parse_expr();

        let key_str = match key.node {
            ExprLit(lit) => {
                match lit.node {
                    LitStr(ref s, _) => s.clone(),
                    _ => {
                        cx.span_err(key.span, "expected string literal");
                        bad = true;
                        InternedString::new("")
                    }
                }
            }
            _ => {
                cx.span_err(key.span, "expected string literal");
                bad = true;
                InternedString::new("")
            }
        };

        if !parser.eat(&FAT_ARROW) {
            cx.span_err(parser.span, "expected `=>`");
            return None;
        }

        let value = parser.parse_expr();

        entries.push(Entry {
            key_str: key_str,
            key: key,
            value: value
        });

        if !parser.eat(&COMMA) && parser.token != EOF {
            cx.span_err(parser.span, "expected `,`");
            return None;
        }
    }

    if entries.len() > phf::MAX_SIZE {
        cx.span_err(parser.span,
                    format!("maps with more than {} entries are not supported",
                            phf::MAX_SIZE));
        return None;
    }

    if bad {
        return None;
    }

    Some(entries)
}

fn has_duplicates(cx: &mut ExtCtxt, sp: Span, entries: &[Entry]) -> bool {
    let mut dups = false;
    let mut in_dup = false;
    for window in entries.windows(2) {
        let ref a = window[0];
        let ref b = window[1];
        if a.key_str == b.key_str {
            dups = true;
            if !in_dup {
                cx.span_err(sp, format!("duplicate key \"{}\"", a.key_str));
                cx.span_note(a.key.span, "one occurrence here");
                in_dup = true;
            }
            cx.span_note(b.key.span, "one occurrence here");
        } else {
            in_dup = false;
        }
    }

    dups
}

struct HashState {
    k1: u64,
    k2: u64,
    disps: ~[(uint, uint)],
    map: ~[Option<uint>],
}

fn generate_hash(entries: &[Entry]) -> Option<HashState> {
    struct Bucket {
        idx: uint,
        keys: ~[uint],
    }

    struct Hashes {
        g: uint,
        f1: uint,
        f2: uint,
    }

    let k1 = rand::random();
    let k2 = rand::random();

    let hashes = entries.iter().map(|entry| {
            let (g, f1, f2) = phf::hash(entry.key_str.get(), k1, k2);
            Hashes {
                g: g,
                f1: f1,
                f2: f2
            }
        }).to_owned_vec();

    let buckets_len = (entries.len() + DEFAULT_LAMBDA - 1) / DEFAULT_LAMBDA;
    let mut buckets = vec::from_fn(buckets_len,
                                   |i| Bucket { idx: i, keys: ~[] });

    for (i, hash) in hashes.iter().enumerate() {
        buckets[hash.g % buckets_len].keys.push(i);
    }

    // Sort descending
    buckets.sort_by(|a, b| b.keys.len().cmp(&a.keys.len()));

    let table_len = entries.len();
    let mut map = vec::from_elem(table_len, None);
    let mut disps = vec::from_elem(buckets_len, None);
    let mut try_map = HashMap::new();
    'buckets: for bucket in buckets.iter() {
        for d1 in range(0, table_len) {
            'disps: for d2 in range(0, table_len) {
                try_map.clear();
                for &key in bucket.keys.iter() {
                    let idx = phf::displace(hashes[key].f1, hashes[key].f2,
                                            d1, d2) % table_len;
                    if try_map.find(&idx).is_some() || map[idx].is_some() {
                        continue 'disps;
                    }
                    try_map.insert(idx, key);
                }

                // We've picked a good set of disps
                disps[bucket.idx] = Some((d1, d2));
                for (&idx, &key) in try_map.iter() {
                    map[idx] = Some(key);
                }
                continue 'buckets;
            }
        }

        // Unable to find displacements for a bucket
        return None;
    }

    let disps = disps.move_iter().map(|i| i.expect("should have a bucket"))
            .collect();

    Some(HashState {
        k1: k1,
        k2: k2,
        disps: disps,
        map: map,
    })
}

fn create_map(cx: &mut ExtCtxt, sp: Span, entries: ~[Entry], state: HashState)
              -> MacResult {
    let len = entries.len();
    let k1 = state.k1;
    let k2 = state.k2;
    let disps = state.disps.iter().map(|&(d1, d2)| {
            quote_expr!(&*cx, ($d1, $d2))
        }).collect();
    let disps = @Expr {
        id: ast::DUMMY_NODE_ID,
        node: ExprVec(disps, MutImmutable),
        span: sp,
    };
    let entries = state.map.iter().map(|&idx| {
            match idx {
                Some(idx) => {
                    let Entry { key, value, .. } = entries[idx];
                    quote_expr!(&*cx, Some(($key, $value)))
                }
                None => quote_expr!(&*cx, None),
            }
        }).collect();
    let entries = @Expr {
        id: ast::DUMMY_NODE_ID,
        node: ExprVec(entries, MutImmutable),
        span: sp,
    };

    MRExpr(quote_expr!(cx, PhfMap {
        len: $len,
        k1: $k1,
        k2: $k2,
        disps: &'static $disps,
        entries: &'static $entries,
    }))
}
<commit_msg>Update time move<commit_after>//! Compiler plugin for Rust-PHF
//!
//! See the documentation for the `phf` crate for more details.
#[crate_id="github.com/sfackler/rust-phf/phf_mac"];
#[crate_type="dylib"];
#[doc(html_root_url="http://www.rust-ci.org/sfackler/rust-phf/doc")];
#[feature(managed_boxes, macro_registrar, quote)];

extern crate extra;
extern crate syntax;
extern crate time;
extern crate phf;

use std::hashmap::HashMap;
use std::rand;
use std::os;
use std::vec;
use syntax::ast;
use syntax::ast::{Name, TokenTree, LitStr, MutImmutable, Expr, ExprVec, ExprLit};
use syntax::codemap::Span;
use syntax::ext::base::{SyntaxExtension,
                        ExtCtxt,
                        MacResult,
                        MRExpr,
                        NormalTT,
                        BasicMacroExpander};
use syntax::parse;
use syntax::parse::token;
use syntax::parse::token::{InternedString, COMMA, EOF, FAT_ARROW};

static DEFAULT_LAMBDA: uint = 5;

#[macro_registrar]
#[doc(hidden)]
pub fn macro_registrar(register: |Name, SyntaxExtension|) {
    register(token::intern("phf_map"),
             NormalTT(~BasicMacroExpander {
                expander: expand_mphf_map,
                span: None
             },
             None));
}

struct Entry {
    key_str: InternedString,
    key: @Expr,
    value: @Expr
}

fn expand_mphf_map(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> MacResult {
    let mut entries = match parse_entries(cx, tts) {
        Some(entries) => entries,
        None => return MacResult::dummy_expr(sp)
    };

    entries.sort_by(|a, b| a.key_str.cmp(&b.key_str));
    if has_duplicates(cx, sp, entries) {
        return MacResult::dummy_expr(sp);
    }

    let start = time::precise_time_s();
    let state;
    loop {
        match generate_hash(entries) {
            Some(s) => {
                state = s;
                break;
            }
            None => {}
        }
    }
    let time = time::precise_time_s() - start;
    if os::getenv("PHF_STATS").is_some() {
        cx.span_note(sp, format!("PHF generation took {} seconds", time));
    }

    create_map(cx, sp, entries, state)
}

fn parse_entries(cx: &mut ExtCtxt, tts: &[TokenTree]) -> Option<~[Entry]> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(),
                                                tts.to_owned());
    let mut entries = ~[];

    let mut bad = false;
    while parser.token != EOF {
        let key = parser.parse_expr();

        let key_str = match key.node {
            ExprLit(lit) => {
                match lit.node {
                    LitStr(ref s, _) => s.clone(),
                    _ => {
                        cx.span_err(key.span, "expected string literal");
                        bad = true;
                        InternedString::new("")
                    }
                }
            }
            _ => {
                cx.span_err(key.span, "expected string literal");
                bad = true;
                InternedString::new("")
            }
        };

        if !parser.eat(&FAT_ARROW) {
            cx.span_err(parser.span, "expected `=>`");
            return None;
        }

        let value = parser.parse_expr();

        entries.push(Entry {
            key_str: key_str,
            key: key,
            value: value
        });

        if !parser.eat(&COMMA) && parser.token != EOF {
            cx.span_err(parser.span, "expected `,`");
            return None;
        }
    }

    if entries.len() > phf::MAX_SIZE {
        cx.span_err(parser.span,
                    format!("maps with more than {} entries are not supported",
                            phf::MAX_SIZE));
        return None;
    }

    if bad {
        return None;
    }

    Some(entries)
}

fn has_duplicates(cx: &mut ExtCtxt, sp: Span, entries: &[Entry]) -> bool {
    let mut dups = false;
    let mut in_dup = false;
    for window in entries.windows(2) {
        let ref a = window[0];
        let ref b = window[1];
        if a.key_str == b.key_str {
            dups = true;
            if !in_dup {
                cx.span_err(sp, format!("duplicate key \"{}\"", a.key_str));
                cx.span_note(a.key.span, "one occurrence here");
                in_dup = true;
            }
            cx.span_note(b.key.span, "one occurrence here");
        } else {
            in_dup = false;
        }
    }

    dups
}

struct HashState {
    k1: u64,
    k2: u64,
    disps: ~[(uint, uint)],
    map: ~[Option<uint>],
}

fn generate_hash(entries: &[Entry]) -> Option<HashState> {
    struct Bucket {
        idx: uint,
        keys: ~[uint],
    }

    struct Hashes {
        g: uint,
        f1: uint,
        f2: uint,
    }

    let k1 = rand::random();
    let k2 = rand::random();

    let hashes = entries.iter().map(|entry| {
            let (g, f1, f2) = phf::hash(entry.key_str.get(), k1, k2);
            Hashes {
                g: g,
                f1: f1,
                f2: f2
            }
        }).to_owned_vec();

    let buckets_len = (entries.len() + DEFAULT_LAMBDA - 1) / DEFAULT_LAMBDA;
    let mut buckets = vec::from_fn(buckets_len,
                                   |i| Bucket { idx: i, keys: ~[] });

    for (i, hash) in hashes.iter().enumerate() {
        buckets[hash.g % buckets_len].keys.push(i);
    }

    // Sort descending
    buckets.sort_by(|a, b| b.keys.len().cmp(&a.keys.len()));

    let table_len = entries.len();
    let mut map = vec::from_elem(table_len, None);
    let mut disps = vec::from_elem(buckets_len, None);
    let mut try_map = HashMap::new();
    'buckets: for bucket in buckets.iter() {
        for d1 in range(0, table_len) {
            'disps: for d2 in range(0, table_len) {
                try_map.clear();
                for &key in bucket.keys.iter() {
                    let idx = phf::displace(hashes[key].f1, hashes[key].f2,
                                            d1, d2) % table_len;
                    if try_map.find(&idx).is_some() || map[idx].is_some() {
                        continue 'disps;
                    }
                    try_map.insert(idx, key);
                }

                // We've picked a good set of disps
                disps[bucket.idx] = Some((d1, d2));
                for (&idx, &key) in try_map.iter() {
                    map[idx] = Some(key);
                }
                continue 'buckets;
            }
        }

        // Unable to find displacements for a bucket
        return None;
    }

    let disps = disps.move_iter().map(|i| i.expect("should have a bucket"))
            .collect();

    Some(HashState {
        k1: k1,
        k2: k2,
        disps: disps,
        map: map,
    })
}

fn create_map(cx: &mut ExtCtxt, sp: Span, entries: ~[Entry], state: HashState)
              -> MacResult {
    let len = entries.len();
    let k1 = state.k1;
    let k2 = state.k2;
    let disps = state.disps.iter().map(|&(d1, d2)| {
            quote_expr!(&*cx, ($d1, $d2))
        }).collect();
    let disps = @Expr {
        id: ast::DUMMY_NODE_ID,
        node: ExprVec(disps, MutImmutable),
        span: sp,
    };
    let entries = state.map.iter().map(|&idx| {
            match idx {
                Some(idx) => {
                    let Entry { key, value, .. } = entries[idx];
                    quote_expr!(&*cx, Some(($key, $value)))
                }
                None => quote_expr!(&*cx, None),
            }
        }).collect();
    let entries = @Expr {
        id: ast::DUMMY_NODE_ID,
        node: ExprVec(entries, MutImmutable),
        span: sp,
    };

    MRExpr(quote_expr!(cx, PhfMap {
        len: $len,
        k1: $k1,
        k2: $k2,
        disps: &'static $disps,
        entries: &'static $entries,
    }))
}
<|endoftext|>


<commit_before>extern crate rustc_serialize;
extern crate rand;

use rand::Rng;
use rustc_serialize::{Encodable, Encoder, Decodable, Decoder};

use options::{Options, Details};
use transport::Serializer;

use std::result;
use std::collections::HashMap;

use WampError;
use WampResult;

/// All WAMP events and message types and their numeric counterparts
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum MessageType {
    HELLO = 1,
    WELCOME,
    ABORT,
    CHALLENGE,
    AUTHENTICATE,
    GOODBYE,
    HEARTBEAT,
    ERROR,
    PUBLISH = 16,
    PUBLISHED,
    SUBSCRIBE = 32,
    SUBSCRIBED,
    UNSUBSCRIBE,
    UNSUBSCRIBED,
    EVENT,
    CALL = 48,
    CANCEL,
    RESULT,
    REGISTER = 64,
    REGISTERED,
    UNREGISTER,
    UNREGISTERED,
    INVOCATION,
    INTERRUPT,
    YIELD,
    /// The message type value was not found, usually this means a protocol violation occured.
    NONE = 0
}

impl From<u32> for MessageType {
    fn from(value: u32) -> Self {
        match value {
            1 => MessageType::HELLO,
            2 => MessageType::WELCOME,
            3 => MessageType::ABORT,
            4 => MessageType:: CHALLENGE,
            5 => MessageType::AUTHENTICATE,
            6 => MessageType::GOODBYE,
            7 => MessageType::HEARTBEAT,
            8 => MessageType::ERROR,
            16 => MessageType::PUBLISH,
            17 => MessageType::PUBLISHED,
            32 => MessageType::SUBSCRIBE,
            33 => MessageType::SUBSCRIBED,
            34 => MessageType::UNSUBSCRIBE,
            35 => MessageType::UNSUBSCRIBED,
            36 => MessageType::EVENT,
            48 => MessageType::CALL,
            49 => MessageType::CANCEL,
            50 => MessageType::RESULT,
            64 => MessageType::REGISTER,
            65 => MessageType::REGISTERED,
            66 => MessageType::UNREGISTER,
            67 => MessageType::UNREGISTERED,
            68 => MessageType::INVOCATION,
            69 => MessageType::INTERRUPT,
            70 => MessageType::YIELD,
            _ => MessageType::NONE,
        }
    }
}

//#[derive(Debug)]
pub enum WampEvent {
    Subscribed {
        message_type: MessageType,
        event_id: u64,
        topic_id: u64,
    },
    Event {
        message_type: MessageType,
        topic_id: u64,
        event_id: u64,
        options: Options,
        has_kwargs: bool,
    },
}

impl Decodable for WampEvent {
    fn decode<D: Decoder>(d: &mut D) -> Result<WampEvent, D::Error> {
        d.read_seq(|d, len| {
            if len > 0 {
                let message_type = try!(d.read_seq_elt(0, |d| d.read_u32()));
                let message_type = MessageType::from(message_type);
                match message_type {
                    MessageType::SUBSCRIBED => {
                        if len != 3 {
                            Err(d.error("unexpected len != 3 for SUBSCRIBED message"))
                        } else {
                            let event_id = try!(d.read_seq_elt(1, |d| d.read_u64()));
                            let topic_id = try!(d.read_seq_elt(2, |d| d.read_u64()));
                            Ok(WampEvent::Subscribed {
                                message_type: message_type,
                                event_id: event_id,
                                topic_id: topic_id,
                            })
                        }
                    },

                    MessageType::EVENT => {
                        if len != 5 && len != 6 {
                            Err(d.error("unexpected len != {5, 6} for EVENT message"))
                        } else {
                            let topic_id = try!(d.read_seq_elt(1, |d| d.read_u64()));
                            let event_id = try!(d.read_seq_elt(2, |d| d.read_u64()));
                            //TODO: Actually fulfill these options!
                            let options = Options::Empty;
                            let has_kwargs = len == 6;

                            // Read positional arguments
                            //let args = try!(d.read_seq_elt(4, |d| d.read_seq(|d, len| {
                            //   let mut v : Vec<WampType> = Vec::with_capacity(len);
                            //   for i in 0..len {
                            //    v.push(try!(d.read_seq_elt(i, |d| Decodable::decode(d))));
                            //   }
                            //   Ok(v)
                            //})));

                           Ok(WampEvent::Event {
                               message_type: message_type,
                               topic_id: topic_id,
                               event_id: event_id,
                               options: options,
                               has_kwargs: has_kwargs,
                           })
                        }
                    }
                    _ => Err(d.error(&*format!("protocol violation: no message type {} exists",
                                               message_type as u32)))
                }
            } else {
                // bad things happen
                Err(d.error("empty message received"))
            }
        })
    }
}

#[derive(Debug)]
#[allow(dead_code)]
pub struct Payload {
    args: String,
    kwargs: Option<String>,
    serializer: Serializer
}

impl Payload {
    fn capture_braces(raw: &str, braces: (char, char)) -> Option<(usize, usize)> {
        let left_brace = braces.0;
        let right_brace = braces.1;
        assert!(left_brace != right_brace);

        let mut first_index : usize  = 0;
        let mut final_index : usize  = 0;
        let mut found_initial_brace = false;
        let mut count = -1;

        for (i, c) in raw.char_indices() {

            if c == left_brace {
                count = count + 1;
                if !found_initial_brace {
                    first_index = i;
                    found_initial_brace = true;
                }
            } else if c == right_brace {
                if found_initial_brace {
                    if count == 0 {
                        final_index = i;
                        break;
                    } else {
                        count = count - 1;
                    }
                }
            }
        }

        if count == 0 && found_initial_brace && final_index > first_index {
            Some((first_index, final_index))
        } else {
            None
        }
    }

    /// Extract args and kwargs from a Wamp Message that is known to have a payload
    pub fn from_str(raw: &str) -> WampResult<Payload> {
        if let Some((_, end)) = Self::capture_braces(raw, ('{', '}')) {
            let (_, next) = raw.split_at(end + 1);
            if let Some((args_l, args_r)) = Self::capture_braces(next, ('[', ']')) {
                let args = &next[args_l .. args_r + 1];
                let (_, kw_next) = next.split_at(args_r);
                let kwargs = Self::capture_braces(kw_next, ('{', '}'))
                    .map(|(kwargs_l, kwargs_r)| kw_next[kwargs_l .. kwargs_r + 1].to_string());

                return Ok(Payload {
                    args: args.to_string(),
                    kwargs: kwargs,
                    serializer: Serializer::json()
                });
            }
        }

        Err(WampError::ProtocolError)
    }

    #[allow(dead_code)]
    pub fn decode_args<T: Decodable>(&self) -> WampResult<T> {
        self.serializer.decode(&*self.args)
    }

    #[allow(dead_code)]
    pub fn decode_kwargs<T: Decodable>(&self) -> WampResult<T> {
        self.serializer.decode(&**self.kwargs.as_ref().unwrap_or(&String::from("")))
    }

    #[allow(dead_code)]
    pub fn has_args(&self) -> bool {
        self.args.len() > 0
    }

    #[allow(dead_code)]
    pub fn has_kwargs(&self) -> bool {
        self.kwargs.is_some()
    }
}

impl Encodable for MessageType {
    fn encode<S: Encoder>(&self, s: &mut S) -> result::Result<(), S::Error> {
        s.emit_u32(*self as u32)
    }
}

/// Generates a new event_id to track messages sent to and from the WAMP Router
pub fn new_event_id() -> u64 {
    rand::thread_rng().next_u32() as u64
}

//TODO: Currently every message type has its own struct,
// perhaps there is a better way to abstract this with an enum
// however, message types should never be directly exposed to users of the library
//

#[derive(Debug, Clone)]
pub struct EventPublish<A: Encodable, K: Encodable> {
    pub message_type: MessageType,
    pub id: u64,
    pub options: Options,
    pub topic: String,
    pub args: Vec<A>,
    pub kwargs: K,
}

impl<A: Encodable, K: Encodable> Encodable for EventPublish<A, K> {
    fn encode<S: Encoder>(&self, s: &mut S) -> result::Result<(), S::Error> {
        // [self.message_type, self.id, self.options, self.topic, self.args, self.kwargs];
        s.emit_seq(6, |s| {
            try!(s.emit_seq_elt(0, |s| self.message_type.encode(s)));
            try!(s.emit_seq_elt(1, |s| self.id.encode(s)));
            try!(s.emit_seq_elt(2, |s| self.options.encode(s)));
            try!(s.emit_seq_elt(3, |s| self.topic.encode(s)));
            try!(s.emit_seq_elt(4, |s| self.args.encode(s)));
            try!(s.emit_seq_elt(5, |s| self.kwargs.encode(s)));
            Ok(())
        })
    }
}

#[derive(Debug, Clone)]
pub struct EventSubscribe {
    message_type: MessageType,
    id: u64,
    options: Options,
    topic: String,
}

impl EventSubscribe {
    pub fn new(topic: String) -> Self {
        EventSubscribe {
            message_type: MessageType::SUBSCRIBE,
            id: new_event_id(),
            topic: topic,
            options: Options::Empty
        }
    }

    pub fn get_id(&self) -> u64 {
        self.id
    }
}

impl Encodable for EventSubscribe {
    fn encode<S: Encoder>(&self, s: &mut S) -> result::Result<(), S::Error> {
        // [self.message_type, self.id, self.options, self.topic, self.args, self.kwargs];
        s.emit_seq(4, |s| {
            try!(s.emit_seq_elt(0, |s| self.message_type.encode(s)));
            try!(s.emit_seq_elt(1, |s| self.id.encode(s)));
            try!(s.emit_seq_elt(2, |s| self.options.encode(s)));
            try!(s.emit_seq_elt(3, |s| self.topic.encode(s)));
            Ok(())
        })
    }
}

#[derive(Debug, Clone)]
pub struct EventJoin {
    pub message_type: MessageType,
    pub realm: String,
    pub details: Details,
}

impl EventJoin {
    pub fn new (realm: String) {
        unimplemented!();
    }
}

impl Encodable for EventJoin {
    fn encode<S: Encoder>(&self, s: &mut S) -> result::Result<(), S::Error> {
        // [self.message_type, self.id, self.options, self.topic, self.args, self.kwargs];
        s.emit_seq(3, |s| {
            try!(s.emit_seq_elt(0, |s| self.message_type.encode(s)));
            try!(s.emit_seq_elt(1, |s| self.realm.encode(s)));
            try!(s.emit_seq_elt(2, |s| self.details.encode(s)));
            Ok(())
        })
    }
}

macro_rules! wamp_type {
    ($($t:ident),+) => {
        /// All types that can be sent to/from a WAMP Router
        /// Used to publish hetereogenous arguments
        #[allow(dead_code, non_camel_case_types)]
        #[derive(Debug, Clone, PartialEq)]
        pub enum WampType {
            $($t($t),)*
                Vec(Vec<WampType>),
                Map(HashMap<String, WampType>),
                /// Used to send an empty struct or keymap value "{}"
                None,
        }

        impl Encodable for WampType {
            fn encode<S: Encoder>(&self, s: &mut S) -> result::Result<(), S::Error> {
                match self {
                    $(&WampType::$t(ref value) => value.encode(s),)+
                        &WampType::Vec(ref value) => value.encode(s),
                        &WampType::Map(ref value) => value.encode(s),
                        &WampType::None => s.emit_map(0, |_| Ok(())),
                }
            }
        }

        // impl Decodable for WampType {
        //     fn decode<D: Decoder>(d: &mut D) -> Result<Self, D::Error> {
        //         unimplemented!();
        //     }
        // }

        // Other clients can have this simplifier type alias
        // type WampEncodable = WampEncodable<()>;
    }
}

// TODO: add types per request
wamp_type!(usize, u8, u16, u32, u64, isize, i8, i16, i32, i64, String, f32, f64, bool, char);

#[test]
fn message_enum_value() {
    assert!(MessageType::HELLO as u32 == 1);
    assert!(MessageType::SUBSCRIBE as u32 == 32);
}

#[test]
fn message_payload_braces() {
    let simple_brace = "[hello, world]";
    assert!(Payload::capture_braces(simple_brace, ('[', ']')) == Some((0, simple_brace.len()-1)));
    assert!(Payload::capture_braces("[[hello], [test, [thing]]], [other, [stuff]]", ('[',']'))
                                    == Some((0, 25)));
}

#[test]
fn message_extract_payload() {
    let sample_message_nokwargs = "[36,1232131,64713717171,{},[42, \"yup\"]]";
    let sample_payload = Payload::from_str(sample_message_nokwargs);
    let (number, yup) : (u32, String) = sample_payload.unwrap().decode_args().unwrap();
    assert!(number == 42);
    assert!(yup == "yup".to_string());

    let message_kwargs_only = "[42, 12415261, 16171, {},[],{\"field\": 42, \"binary\": false, \"word\": \"hello world\"}]";
    #[derive(PartialEq, RustcDecodable)]
    struct TestStruct {
        field: u32,
        binary: bool,
        word: String,
    }
    let payload2 = Payload::from_str(message_kwargs_only);
    let test_struct: TestStruct= payload2.unwrap().decode_kwargs().unwrap();
    assert!(test_struct == TestStruct{field: 42, binary: false, word: "hello world".to_string()});
}
<commit_msg>added base doc to message::Payload<commit_after>extern crate rustc_serialize;
extern crate rand;

use rand::Rng;
use rustc_serialize::{Encodable, Encoder, Decodable, Decoder};

use options::{Options, Details};
use transport::Serializer;

use std::result;
use std::collections::HashMap;

use WampError;
use WampResult;

/// All WAMP events and message types and their numeric counterparts
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum MessageType {
    HELLO = 1,
    WELCOME,
    ABORT,
    CHALLENGE,
    AUTHENTICATE,
    GOODBYE,
    HEARTBEAT,
    ERROR,
    PUBLISH = 16,
    PUBLISHED,
    SUBSCRIBE = 32,
    SUBSCRIBED,
    UNSUBSCRIBE,
    UNSUBSCRIBED,
    EVENT,
    CALL = 48,
    CANCEL,
    RESULT,
    REGISTER = 64,
    REGISTERED,
    UNREGISTER,
    UNREGISTERED,
    INVOCATION,
    INTERRUPT,
    YIELD,
    /// The message type value was not found, usually this means a protocol violation occured.
    NONE = 0
}

impl From<u32> for MessageType {
    fn from(value: u32) -> Self {
        match value {
            1 => MessageType::HELLO,
            2 => MessageType::WELCOME,
            3 => MessageType::ABORT,
            4 => MessageType:: CHALLENGE,
            5 => MessageType::AUTHENTICATE,
            6 => MessageType::GOODBYE,
            7 => MessageType::HEARTBEAT,
            8 => MessageType::ERROR,
            16 => MessageType::PUBLISH,
            17 => MessageType::PUBLISHED,
            32 => MessageType::SUBSCRIBE,
            33 => MessageType::SUBSCRIBED,
            34 => MessageType::UNSUBSCRIBE,
            35 => MessageType::UNSUBSCRIBED,
            36 => MessageType::EVENT,
            48 => MessageType::CALL,
            49 => MessageType::CANCEL,
            50 => MessageType::RESULT,
            64 => MessageType::REGISTER,
            65 => MessageType::REGISTERED,
            66 => MessageType::UNREGISTER,
            67 => MessageType::UNREGISTERED,
            68 => MessageType::INVOCATION,
            69 => MessageType::INTERRUPT,
            70 => MessageType::YIELD,
            _ => MessageType::NONE,
        }
    }
}

//#[derive(Debug)]
pub enum WampEvent {
    Subscribed {
        message_type: MessageType,
        event_id: u64,
        topic_id: u64,
    },
    Event {
        message_type: MessageType,
        topic_id: u64,
        event_id: u64,
        options: Options,
        has_kwargs: bool,
    },
}

impl Decodable for WampEvent {
    fn decode<D: Decoder>(d: &mut D) -> Result<WampEvent, D::Error> {
        d.read_seq(|d, len| {
            if len > 0 {
                let message_type = try!(d.read_seq_elt(0, |d| d.read_u32()));
                let message_type = MessageType::from(message_type);
                match message_type {
                    MessageType::SUBSCRIBED => {
                        if len != 3 {
                            Err(d.error("unexpected len != 3 for SUBSCRIBED message"))
                        } else {
                            let event_id = try!(d.read_seq_elt(1, |d| d.read_u64()));
                            let topic_id = try!(d.read_seq_elt(2, |d| d.read_u64()));
                            Ok(WampEvent::Subscribed {
                                message_type: message_type,
                                event_id: event_id,
                                topic_id: topic_id,
                            })
                        }
                    },

                    MessageType::EVENT => {
                        if len != 5 && len != 6 {
                            Err(d.error("unexpected len != {5, 6} for EVENT message"))
                        } else {
                            let topic_id = try!(d.read_seq_elt(1, |d| d.read_u64()));
                            let event_id = try!(d.read_seq_elt(2, |d| d.read_u64()));
                            //TODO: Actually fulfill these options!
                            let options = Options::Empty;
                            let has_kwargs = len == 6;

                            // Read positional arguments
                            //let args = try!(d.read_seq_elt(4, |d| d.read_seq(|d, len| {
                            //   let mut v : Vec<WampType> = Vec::with_capacity(len);
                            //   for i in 0..len {
                            //    v.push(try!(d.read_seq_elt(i, |d| Decodable::decode(d))));
                            //   }
                            //   Ok(v)
                            //})));

                           Ok(WampEvent::Event {
                               message_type: message_type,
                               topic_id: topic_id,
                               event_id: event_id,
                               options: options,
                               has_kwargs: has_kwargs,
                           })
                        }
                    }
                    _ => Err(d.error(&*format!("protocol violation: no message type {} exists",
                                               message_type as u32)))
                }
            } else {
                // bad things happen
                Err(d.error("empty message received"))
            }
        })
    }
}

#[derive(Debug)]
#[allow(dead_code)]
/// A struct representing the payload that's received from a WAMP event.
pub struct Payload {
    args: String,
    kwargs: Option<String>,
    serializer: Serializer
}

impl Payload {
    fn capture_braces(raw: &str, braces: (char, char)) -> Option<(usize, usize)> {
        let left_brace = braces.0;
        let right_brace = braces.1;
        assert!(left_brace != right_brace);

        let mut first_index : usize  = 0;
        let mut final_index : usize  = 0;
        let mut found_initial_brace = false;
        let mut count = -1;

        for (i, c) in raw.char_indices() {

            if c == left_brace {
                count = count + 1;
                if !found_initial_brace {
                    first_index = i;
                    found_initial_brace = true;
                }
            } else if c == right_brace {
                if found_initial_brace {
                    if count == 0 {
                        final_index = i;
                        break;
                    } else {
                        count = count - 1;
                    }
                }
            }
        }

        if count == 0 && found_initial_brace && final_index > first_index {
            Some((first_index, final_index))
        } else {
            None
        }
    }

    /// Parse args and kwargs from a WAMP Message that is known to have a payload
    pub fn from_str(raw: &str) -> WampResult<Payload> {
        if let Some((_, end)) = Self::capture_braces(raw, ('{', '}')) {
            let (_, next) = raw.split_at(end + 1);
            if let Some((args_l, args_r)) = Self::capture_braces(next, ('[', ']')) {
                let args = &next[args_l .. args_r + 1];
                let (_, kw_next) = next.split_at(args_r);
                let kwargs = Self::capture_braces(kw_next, ('{', '}'))
                    .map(|(kwargs_l, kwargs_r)| kw_next[kwargs_l .. kwargs_r + 1].to_string());

                return Ok(Payload {
                    args: args.to_string(),
                    kwargs: kwargs,
                    serializer: Serializer::json()
                });
            }
        }

        Err(WampError::ProtocolError)
    }

    #[allow(dead_code)]
    /// Extract positional arguments from the payload.
    pub fn decode_args<T: Decodable>(&self) -> WampResult<T> {
        self.serializer.decode(&*self.args)
    }

    #[allow(dead_code)]
    /// Extract keyword arguments from the payload.
    pub fn decode_kwargs<T: Decodable>(&self) -> WampResult<T> {
        self.serializer.decode(&**self.kwargs.as_ref().unwrap_or(&String::from("")))
    }

    #[allow(dead_code)]
    pub fn has_args(&self) -> bool {
        self.args.len() > 0
    }

    #[allow(dead_code)]
    pub fn has_kwargs(&self) -> bool {
        self.kwargs.is_some()
    }
}

impl Encodable for MessageType {
    fn encode<S: Encoder>(&self, s: &mut S) -> result::Result<(), S::Error> {
        s.emit_u32(*self as u32)
    }
}

/// Generates a new event_id to track messages sent to and from the WAMP Router
pub fn new_event_id() -> u64 {
    rand::thread_rng().next_u32() as u64
}

//TODO: Currently every message type has its own struct,
// perhaps there is a better way to abstract this with an enum
// however, message types should never be directly exposed to users of the library
//

#[derive(Debug, Clone)]
pub struct EventPublish<A: Encodable, K: Encodable> {
    pub message_type: MessageType,
    pub id: u64,
    pub options: Options,
    pub topic: String,
    pub args: Vec<A>,
    pub kwargs: K,
}

impl<A: Encodable, K: Encodable> Encodable for EventPublish<A, K> {
    fn encode<S: Encoder>(&self, s: &mut S) -> result::Result<(), S::Error> {
        // [self.message_type, self.id, self.options, self.topic, self.args, self.kwargs];
        s.emit_seq(6, |s| {
            try!(s.emit_seq_elt(0, |s| self.message_type.encode(s)));
            try!(s.emit_seq_elt(1, |s| self.id.encode(s)));
            try!(s.emit_seq_elt(2, |s| self.options.encode(s)));
            try!(s.emit_seq_elt(3, |s| self.topic.encode(s)));
            try!(s.emit_seq_elt(4, |s| self.args.encode(s)));
            try!(s.emit_seq_elt(5, |s| self.kwargs.encode(s)));
            Ok(())
        })
    }
}

#[derive(Debug, Clone)]
pub struct EventSubscribe {
    message_type: MessageType,
    id: u64,
    options: Options,
    topic: String,
}

impl EventSubscribe {
    pub fn new(topic: String) -> Self {
        EventSubscribe {
            message_type: MessageType::SUBSCRIBE,
            id: new_event_id(),
            topic: topic,
            options: Options::Empty
        }
    }

    pub fn get_id(&self) -> u64 {
        self.id
    }
}

impl Encodable for EventSubscribe {
    fn encode<S: Encoder>(&self, s: &mut S) -> result::Result<(), S::Error> {
        // [self.message_type, self.id, self.options, self.topic, self.args, self.kwargs];
        s.emit_seq(4, |s| {
            try!(s.emit_seq_elt(0, |s| self.message_type.encode(s)));
            try!(s.emit_seq_elt(1, |s| self.id.encode(s)));
            try!(s.emit_seq_elt(2, |s| self.options.encode(s)));
            try!(s.emit_seq_elt(3, |s| self.topic.encode(s)));
            Ok(())
        })
    }
}

#[derive(Debug, Clone)]
pub struct EventJoin {
    pub message_type: MessageType,
    pub realm: String,
    pub details: Details,
}

impl EventJoin {
    pub fn new (realm: String) {
        unimplemented!();
    }
}

impl Encodable for EventJoin {
    fn encode<S: Encoder>(&self, s: &mut S) -> result::Result<(), S::Error> {
        // [self.message_type, self.id, self.options, self.topic, self.args, self.kwargs];
        s.emit_seq(3, |s| {
            try!(s.emit_seq_elt(0, |s| self.message_type.encode(s)));
            try!(s.emit_seq_elt(1, |s| self.realm.encode(s)));
            try!(s.emit_seq_elt(2, |s| self.details.encode(s)));
            Ok(())
        })
    }
}

macro_rules! wamp_type {
    ($($t:ident),+) => {
        /// All types that can be sent to/from a WAMP Router
        /// Used to publish hetereogenous arguments
        #[allow(dead_code, non_camel_case_types)]
        #[derive(Debug, Clone, PartialEq)]
        pub enum WampType {
            $($t($t),)*
                Vec(Vec<WampType>),
                Map(HashMap<String, WampType>),
                /// Used to send an empty struct or keymap value "{}"
                None,
        }

        impl Encodable for WampType {
            fn encode<S: Encoder>(&self, s: &mut S) -> result::Result<(), S::Error> {
                match self {
                    $(&WampType::$t(ref value) => value.encode(s),)+
                        &WampType::Vec(ref value) => value.encode(s),
                        &WampType::Map(ref value) => value.encode(s),
                        &WampType::None => s.emit_map(0, |_| Ok(())),
                }
            }
        }

        // impl Decodable for WampType {
        //     fn decode<D: Decoder>(d: &mut D) -> Result<Self, D::Error> {
        //         unimplemented!();
        //     }
        // }

        // Other clients can have this simplifier type alias
        // type WampEncodable = WampEncodable<()>;
    }
}

// TODO: add types per request
wamp_type!(usize, u8, u16, u32, u64, isize, i8, i16, i32, i64, String, f32, f64, bool, char);

#[test]
fn message_enum_value() {
    assert!(MessageType::HELLO as u32 == 1);
    assert!(MessageType::SUBSCRIBE as u32 == 32);
}

#[test]
fn message_payload_braces() {
    let simple_brace = "[hello, world]";
    assert!(Payload::capture_braces(simple_brace, ('[', ']')) == Some((0, simple_brace.len()-1)));
    assert!(Payload::capture_braces("[[hello], [test, [thing]]], [other, [stuff]]", ('[',']'))
                                    == Some((0, 25)));
}

#[test]
fn message_extract_payload() {
    let sample_message_nokwargs = "[36,1232131,64713717171,{},[42, \"yup\"]]";
    let sample_payload = Payload::from_str(sample_message_nokwargs);
    let (number, yup) : (u32, String) = sample_payload.unwrap().decode_args().unwrap();
    assert!(number == 42);
    assert!(yup == "yup".to_string());

    let message_kwargs_only = "[42, 12415261, 16171, {},[],{\"field\": 42, \"binary\": false, \"word\": \"hello world\"}]";
    #[derive(PartialEq, RustcDecodable)]
    struct TestStruct {
        field: u32,
        binary: bool,
        word: String,
    }
    let payload2 = Payload::from_str(message_kwargs_only);
    let test_struct: TestStruct= payload2.unwrap().decode_kwargs().unwrap();
    assert!(test_struct == TestStruct{field: 42, binary: false, word: "hello world".to_string()});
}
<|endoftext|>