<commit_before>use super::fitness::Fitness;
use super::traits::{Genotype, Distance};
use super::selection;
use rand::{Rng, Closed01};
use super::traits::Mate;
use std::marker::PhantomData;
use std::num::Zero;
use std::fmt::Debug;
use std::cmp;
use rayon::par_iter::*;
use super::prob::probabilistic_round;

#[derive(Debug)]
pub struct Individual<T: Debug + Genotype> {
    fitness: Fitness,
    genome: Box<T>,
}

pub trait Rating {}
#[derive(Debug)]
pub struct Rated;
#[derive(Debug)]
pub struct Unrated;

impl Rating for Rated {}
impl Rating for Unrated {}

#[derive(Debug)]
pub struct Population<T: Genotype + Debug, R: Rating> {
    individuals: Vec<Individual<T>>,
    _marker: PhantomData<R>,
}

impl<T: Genotype + Debug, R: Rating> Population<T, R> {
    pub fn new() -> Population<T, R> {
        Population {
            individuals: Vec::new(),
            _marker: PhantomData,
        }
    }

    pub fn new_from_vec(v: Vec<Individual<T>>) -> Population<T, R> {
        Population {
            individuals: v,
            _marker: PhantomData,
        }
    }

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

impl<T: Genotype + Debug> Population<T, Unrated> {
    pub fn add_genome(&mut self, genome: Box<T>) {
        self.individuals.push(Individual {
            fitness: Zero::zero(),
            genome: genome,
        });
    }

    pub fn rate<F>(mut self, f: &F) -> Population<T, Rated>
        where F: Fn(&T) -> Fitness
    {
        for ind in self.individuals.iter_mut() {
            let fitness = f(&ind.genome);
            ind.fitness = fitness;
        }
        Population {
            individuals: self.individuals,
            _marker: PhantomData,
        }
    }

    pub fn rate_par<F>(mut self, f: &F) -> Population<T, Rated>
        where F: Sync + Fn(&T) -> Fitness
    {
        self.individuals.par_iter_mut().for_each(|ind| {
            let fitness = f(&ind.genome);
            ind.fitness = fitness;
        });

        Population {
            individuals: self.individuals,
            _marker: PhantomData,
        }
    }
}

impl<T: Genotype + Debug> Population<T, Rated> {
    pub fn add_individual(&mut self, ind: Individual<T>) {
        self.individuals.push(ind);
    }

    // Return true if genome at position `i` is fitter that `j`
    fn is_fitter(&self, i: usize, j: usize) -> bool {
        self.individuals[i].fitness > self.individuals[j].fitness
    }

    fn mean_fitness(&self) -> Fitness {
        let sum: Fitness = self.individuals.iter().map(|ind| ind.fitness).sum();
        sum / Fitness::new(self.len() as f64)
    }

    pub fn max_fitness(&self) -> Option<Fitness> {
        self.individuals.iter().max_by_key(|ind| ind.fitness).map(|i| i.fitness)
    }

    // higher value of fitness means that the individual is fitter.
    pub fn sort(mut self) -> Self {
        (&mut self.individuals).sort_by(|a, b| a.fitness.cmp(&b.fitness).reverse());
        self
    }

    /// Merge `self` with the first `n` individuals from population `other`.
    pub fn merge(&mut self, other: Population<T, Rated>, n: Option<usize>) {
        let len = other.individuals.len();
        self.individuals.extend(other.individuals.into_iter().take(n.unwrap_or(len)));
    }

    // We want to create a population of approximately `pop_size`. We do not care if there are
    // slightly more or slightly less individuals than that.
    //
    // 1. sort whole popluation into niches
    // 2. calculate the number of offspring for each niche.
    // 3. each niche produces offspring.
    //    - sort each niche according to the fitness value.
    //    - r% of the best genomes produce offspring
    //    - determine the elitist size. those are always copied into the new generation.
    //
    pub fn produce_offspring<C, M, R>(self,
                                      pop_size: usize,
                                      // how many of the best individuals of a niche are copied as-is into the
                                      // new population?
                                      elite_percentage: Closed01<f64>,
                                      // how many of the best individuals of a niche are selected for
                                      // reproduction?
                                      selection_percentage: Closed01<f64>,
                                      tournament_k: usize,
                                      compatibility_threshold: f64,
                                      compatibility: &C,
                                      mate: &mut M,
                                      rng: &mut R)
                                      -> (Population<T, Rated>, Population<T, Unrated>)
        where R: Rng,
              C: Distance<T>,
              M: Mate<T>
    {
        debug!("population size: {}", self.len());

        assert!(elite_percentage.0 <= selection_percentage.0); // XXX
        assert!(self.len() > 0);
        let niches = self.partition(rng, compatibility_threshold, compatibility);
        let num_niches = niches.len();

        debug!("number of niches: {}", num_niches);

        assert!(num_niches > 0);
        let total_mean: Fitness = niches.iter().map(|ind| ind.mean_fitness()).sum();
        assert!(total_mean.get() >= 0.0);

        let mut new_unrated_population: Population<T, Unrated> = Population::new();
        let mut new_rated_population: Population<T, Rated> = Population::new();

        for niche in niches.into_iter() {
            let percentage_of_population: f64 = if total_mean.get() == 0.0 {
                // all individuals have a fitness of 0.0.
                // we will equally allow each niche to procude offspring.
                1.0 / (num_niches as f64)
            } else {
                (niche.mean_fitness() / total_mean).get()
            };

            // calculate new size of niche, and size of elites, and selection size.
            assert!(percentage_of_population >= 0.0 && percentage_of_population <= 1.0);

            let niche_size = pop_size as f64 * percentage_of_population;

            // number of elitary individuals to copy from the old niche generation into the new.
            let elite_size =
                cmp::max(1,
                         probabilistic_round(niche_size * elite_percentage.0, rng) as usize);

            // number of offspring to produce.
            let offspring_size = probabilistic_round(niche_size * (1.0 - elite_percentage.0),
                                                     rng) as usize;

            // number of the best individuals to use for mating.
            let select_size = probabilistic_round(niche_size *
                                                  selection_percentage.0,
                                                  rng) as usize;

            let sorted_niche = niche.sort();

            let select_size = cmp::min(select_size, sorted_niche.len());

            // at first produce `offspring_size` individuals from the top `select_size`
            // individuals.
            if select_size > 0 {
                let mut n = offspring_size;
                while n > 0 {
                    let (parent1, parent2) =
                        selection::tournament_selection_fast2(rng,
                                                              &|i, j| sorted_niche.is_fitter(i, j),
                                                              select_size,
                                                              cmp::min(select_size, tournament_k),
                                                              3);

                    // `mate` assumes that the left parent performs better.
                    let (left, right) = if sorted_niche.is_fitter(parent1, parent2) {
                        (parent1, parent2)
                    } else {
                        (parent2, parent1)
                    };
                    let offspring = mate.mate(&sorted_niche.individuals[left].genome,
                                              &sorted_niche.individuals[right].genome,
                                              left == right,
                                              rng);
                    new_unrated_population.add_genome(Box::new(offspring));
                    n -= 1;
                }
            }

            // now copy the elites
            new_rated_population.merge(sorted_niche, Some(elite_size));
        }

        return (new_rated_population, new_unrated_population);
    }

    // Partitions the whole population into species (niches)
    fn partition<R, C>(self,
                       rng: &mut R,
                       compatibility_threshold: f64,
                       compatibility: &C)
                       -> Vec<Population<T, Rated>>
        where R: Rng,
              C: Distance<T>
    {
        let mut niches: Vec<Population<T, Rated>> = Vec::new();

        'outer: for ind in self.individuals.into_iter() {
            for niche in niches.iter_mut() {
                // Is this genome compatible with this niche? Test against a random genome.
                let compatible = match rng.choose(&niche.individuals) {
                    Some(probe) => {
                        compatibility.distance(&probe.genome, &ind.genome) < compatibility_threshold
                    }
                    // If a niche is empty, a genome always is compatible (note that a niche can't be empyt)
                    None => true,
                };
                if compatible {
                    niche.add_individual(ind);
                    continue 'outer;
                }
            }
            // if no compatible niche was found, create a new niche containing this genome.
            let mut new_niche = Population::new();
            new_niche.add_individual(ind);
            niches.push(new_niche);
        }

        niches
    }
}

pub struct Runner<'a, T, C, M, F>
    where T: Genotype + Debug,
          C: Distance<T> + 'a,
          M: Mate<T> + 'a,
          F: Sync + Fn(&T) -> Fitness + 'a
{
    // anticipated population size
    pub pop_size: usize,
    // how many of the best individuals of a niche are copied as-is into the
    // new population?
    pub elite_percentage: Closed01<f64>,
    // how many of the best individuals of a niche are selected for
    // reproduction?
    pub selection_percentage: Closed01<f64>,
    pub tournament_k: usize,
    pub compatibility_threshold: f64,
    pub compatibility: &'a C,
    pub mate: &'a mut M,
    pub fitness: &'a F,
    pub _marker: PhantomData<T>,
}

impl<'a, T, C, M, F> Runner<'a, T, C, M, F>
    where T: Genotype + Debug,
          C: Distance<T> + 'a,
          M: Mate<T> + 'a,
          F: Sync + Fn(&T) -> Fitness + 'a
{
    pub fn run<R, G>(&mut self,
                     initial_pop: Population<T, Unrated>,
                     goal_condition: &G,
                     rng: &mut R)
                     -> (usize, Population<T, Rated>)
        where R: Rng,
              G: Fn(usize, &Population<T, Rated>) -> bool
    {
        let mut iteration: usize = 0;
        let mut current_rated_pop = initial_pop.rate_par(self.fitness);

        while !goal_condition(iteration, &current_rated_pop) {
            let (new_rated, new_unrated) =
                current_rated_pop.produce_offspring(self.pop_size,
                                                    Closed01(self.elite_percentage.0),
                                                    Closed01(self.selection_percentage.0),
                                                    self.tournament_k,
                                                    self.compatibility_threshold,
                                                    self.compatibility,
                                                    self.mate,
                                                    rng);

            current_rated_pop = new_rated;
            current_rated_pop.merge(new_unrated.rate_par(self.fitness), None);
            iteration += 1;
        }

        return (iteration, current_rated_pop);
    }
}
<commit_msg>Add best_individual() method to Population<commit_after>use super::fitness::Fitness;
use super::traits::{Genotype, Distance};
use super::selection;
use rand::{Rng, Closed01};
use super::traits::Mate;
use std::marker::PhantomData;
use std::num::Zero;
use std::fmt::Debug;
use std::cmp;
use rayon::par_iter::*;
use super::prob::probabilistic_round;

#[derive(Debug)]
pub struct Individual<T: Debug + Genotype> {
    fitness: Fitness,
    genome: Box<T>,
}

impl<T: Debug + Genotype> Individual<T> {
    pub fn fitness(&self) -> Fitness {
        self.fitness
    }

    pub fn genome(&self) -> &T {
        &self.genome
    }
}

pub trait Rating {}

#[derive(Debug)]
pub struct Unrated;

#[derive(Debug)]
pub struct Rated;

impl Rating for Unrated {}
impl Rating for Rated {}

#[derive(Debug)]
pub struct Population<T: Genotype + Debug, R: Rating> {
    individuals: Vec<Individual<T>>,
    _marker: PhantomData<R>,
}

impl<T: Genotype + Debug, R: Rating> Population<T, R> {
    pub fn new() -> Population<T, R> {
        Population {
            individuals: Vec::new(),
            _marker: PhantomData,
        }
    }

    pub fn new_from_vec(v: Vec<Individual<T>>) -> Population<T, R> {
        Population {
            individuals: v,
            _marker: PhantomData,
        }
    }

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

impl<T: Genotype + Debug> Population<T, Unrated> {
    pub fn add_genome(&mut self, genome: Box<T>) {
        self.individuals.push(Individual {
            fitness: Zero::zero(),
            genome: genome,
        });
    }

    pub fn rate<F>(mut self, f: &F) -> Population<T, Rated>
        where F: Fn(&T) -> Fitness
    {
        for ind in self.individuals.iter_mut() {
            let fitness = f(&ind.genome);
            ind.fitness = fitness;
        }
        Population {
            individuals: self.individuals,
            _marker: PhantomData,
        }
    }

    pub fn rate_par<F>(mut self, f: &F) -> Population<T, Rated>
        where F: Sync + Fn(&T) -> Fitness
    {
        self.individuals.par_iter_mut().for_each(|ind| {
            let fitness = f(&ind.genome);
            ind.fitness = fitness;
        });

        Population {
            individuals: self.individuals,
            _marker: PhantomData,
        }
    }
}

impl<T: Genotype + Debug> Population<T, Rated> {
    pub fn add_individual(&mut self, ind: Individual<T>) {
        self.individuals.push(ind);
    }

    // Return true if genome at position `i` is fitter that `j`
    // XXX: If population is sorted, we only have to check the indices!
    fn is_fitter(&self, i: usize, j: usize) -> bool {
        self.individuals[i].fitness > self.individuals[j].fitness
    }

    fn mean_fitness(&self) -> Fitness {
        let sum: Fitness = self.individuals.iter().map(|ind| ind.fitness).sum();
        sum / Fitness::new(self.len() as f64)
    }

    pub fn max_fitness(&self) -> Option<Fitness> {
        self.individuals.iter().max_by_key(|ind| ind.fitness).map(|i| i.fitness)
    }

    pub fn best_individual(&self) -> Option<&Individual<T>> {
        self.individuals.iter().max_by_key(|ind| ind.fitness)
    }

    // higher value of fitness means that the individual is fitter.
    pub fn sort(mut self) -> Self {
        (&mut self.individuals).sort_by(|a, b| a.fitness.cmp(&b.fitness).reverse());
        self
    }

    /// Merge `self` with the first `n` individuals from population `other`.
    pub fn merge(&mut self, other: Population<T, Rated>, n: Option<usize>) {
        let len = other.individuals.len();
        self.individuals.extend(other.individuals.into_iter().take(n.unwrap_or(len)));
    }

    // We want to create a population of approximately `pop_size`. We do not care if there are
    // slightly more or slightly less individuals than that.
    //
    // 1. sort whole popluation into niches
    // 2. calculate the number of offspring for each niche.
    // 3. each niche produces offspring.
    //    - sort each niche according to the fitness value.
    //    - r% of the best genomes produce offspring
    //    - determine the elitist size. those are always copied into the new generation.
    //
    pub fn produce_offspring<C, M, R>(self,
                                      pop_size: usize,
                                      // how many of the best individuals of a niche are copied as-is into the
                                      // new population?
                                      elite_percentage: Closed01<f64>,
                                      // how many of the best individuals of a niche are selected for
                                      // reproduction?
                                      selection_percentage: Closed01<f64>,
                                      tournament_k: usize,
                                      compatibility_threshold: f64,
                                      compatibility: &C,
                                      mate: &mut M,
                                      rng: &mut R)
                                      -> (Population<T, Rated>, Population<T, Unrated>)
        where R: Rng,
              C: Distance<T>,
              M: Mate<T>
    {
        debug!("population size: {}", self.len());

        assert!(elite_percentage.0 <= selection_percentage.0); // XXX
        assert!(self.len() > 0);
        let niches = self.partition(rng, compatibility_threshold, compatibility);
        let num_niches = niches.len();

        debug!("number of niches: {}", num_niches);

        assert!(num_niches > 0);
        let total_mean: Fitness = niches.iter().map(|ind| ind.mean_fitness()).sum();
        assert!(total_mean.get() >= 0.0);

        let mut new_unrated_population: Population<T, Unrated> = Population::new();
        let mut new_rated_population: Population<T, Rated> = Population::new();

        for niche in niches.into_iter() {
            let percentage_of_population: f64 = if total_mean.get() == 0.0 {
                // all individuals have a fitness of 0.0.
                // we will equally allow each niche to procude offspring.
                1.0 / (num_niches as f64)
            } else {
                (niche.mean_fitness() / total_mean).get()
            };

            // calculate new size of niche, and size of elites, and selection size.
            assert!(percentage_of_population >= 0.0 && percentage_of_population <= 1.0);

            let niche_size = pop_size as f64 * percentage_of_population;

            // number of elitary individuals to copy from the old niche generation into the new.
            let elite_size =
                cmp::max(1,
                         probabilistic_round(niche_size * elite_percentage.0, rng) as usize);

            // number of offspring to produce.
            let offspring_size = probabilistic_round(niche_size * (1.0 - elite_percentage.0),
                                                     rng) as usize;

            // number of the best individuals to use for mating.
            let select_size = probabilistic_round(niche_size *
                                                  selection_percentage.0,
                                                  rng) as usize;

            let sorted_niche = niche.sort();

            let select_size = cmp::min(select_size, sorted_niche.len());

            // at first produce `offspring_size` individuals from the top `select_size`
            // individuals.
            if select_size > 0 {
                let mut n = offspring_size;
                while n > 0 {
                    let (parent1, parent2) =
                        selection::tournament_selection_fast2(rng,
                                                              // XXX: in a sorted niche, the
                                                              // individuals with the lower index
                                                              // has better fitness.
                                                              &|i, j| sorted_niche.is_fitter(i, j),
                                                              select_size,
                                                              cmp::min(select_size, tournament_k),
                                                              3);

                    // `mate` assumes that the left parent performs better.
                    // XXX: `is_fitter` same as above.
                    let (left, right) = if sorted_niche.is_fitter(parent1, parent2) {
                        (parent1, parent2)
                    } else {
                        (parent2, parent1)
                    };
                    let offspring = mate.mate(&sorted_niche.individuals[left].genome,
                                              &sorted_niche.individuals[right].genome,
                                              left == right,
                                              rng);
                    new_unrated_population.add_genome(Box::new(offspring));
                    n -= 1;
                }
            }

            // now copy the elites
            new_rated_population.merge(sorted_niche, Some(elite_size));
        }

        return (new_rated_population, new_unrated_population);
    }

    // Partitions the whole population into species (niches)
    fn partition<R, C>(self,
                       rng: &mut R,
                       compatibility_threshold: f64,
                       compatibility: &C)
                       -> Vec<Population<T, Rated>>
        where R: Rng,
              C: Distance<T>
    {
        let mut niches: Vec<Population<T, Rated>> = Vec::new();

        'outer: for ind in self.individuals.into_iter() {
            for niche in niches.iter_mut() {
                // Is this genome compatible with this niche? Test against a random genome.
                let compatible = match rng.choose(&niche.individuals) {
                    Some(probe) => {
                        compatibility.distance(&probe.genome, &ind.genome) < compatibility_threshold
                    }
                    // If a niche is empty, a genome always is compatible (note that a niche can't be empyt)
                    None => true,
                };
                if compatible {
                    niche.add_individual(ind);
                    continue 'outer;
                }
            }
            // if no compatible niche was found, create a new niche containing this genome.
            let mut new_niche = Population::new();
            new_niche.add_individual(ind);
            niches.push(new_niche);
        }

        niches
    }
}

pub struct Runner<'a, T, C, M, F>
    where T: Genotype + Debug,
          C: Distance<T> + 'a,
          M: Mate<T> + 'a,
          F: Sync + Fn(&T) -> Fitness + 'a
{
    // anticipated population size
    pub pop_size: usize,
    // how many of the best individuals of a niche are copied as-is into the
    // new population?
    pub elite_percentage: Closed01<f64>,
    // how many of the best individuals of a niche are selected for
    // reproduction?
    pub selection_percentage: Closed01<f64>,
    pub tournament_k: usize,
    pub compatibility_threshold: f64,
    pub compatibility: &'a C,
    pub mate: &'a mut M,
    pub fitness: &'a F,
    pub _marker: PhantomData<T>,
}

impl<'a, T, C, M, F> Runner<'a, T, C, M, F>
    where T: Genotype + Debug,
          C: Distance<T> + 'a,
          M: Mate<T> + 'a,
          F: Sync + Fn(&T) -> Fitness + 'a
{
    pub fn run<R, G>(&mut self,
                     initial_pop: Population<T, Unrated>,
                     goal_condition: &G,
                     rng: &mut R)
                     -> (usize, Population<T, Rated>)
        where R: Rng,
              G: Fn(usize, &Population<T, Rated>) -> bool
    {
        let mut iteration: usize = 0;
        let mut current_rated_pop = initial_pop.rate_par(self.fitness);

        while !goal_condition(iteration, &current_rated_pop) {
            let (new_rated, new_unrated) =
                current_rated_pop.produce_offspring(self.pop_size,
                                                    Closed01(self.elite_percentage.0),
                                                    Closed01(self.selection_percentage.0),
                                                    self.tournament_k,
                                                    self.compatibility_threshold,
                                                    self.compatibility,
                                                    self.mate,
                                                    rng);

            current_rated_pop = new_rated;
            current_rated_pop.merge(new_unrated.rate_par(self.fitness), None);
            iteration += 1;
        }

        return (iteration, current_rated_pop);
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::prelude::*;
use std::io::BufReader;
use std::fs::{self, File, ReadDir, DirEntry};
use std::path::Path;
use std::collections::HashMap;
use std::cmp::Ordering;

pub type TaskId = u32;

fn err_str<T: ToString>(err: T) -> String {
  err.to_string()
}

fn parse_taskid(taskid_str: String) -> Result<TaskId, String> {
  taskid_str.parse().map_err(err_str)
}

#[derive(Debug)]
pub struct Proc {
  pub status: ProcStatus,
  pub cmdline: Vec<String>
}

impl Proc {
  pub fn new(pid: TaskId) -> Result<Self, String> {
    let proc_dir = format!("/proc/{}", pid);
    let proc_status = try!(ProcStatus::new(&proc_dir));
    let cmdline = try!(Self::read_cmdline(&proc_dir));

    let proc_struct = Proc{
      status: proc_status,
      cmdline: cmdline
    };

    let _ = ProcStat::new(&proc_dir).unwrap();

    Ok(proc_struct)
  }

  fn read_cmdline(proc_dir: &str) -> Result<Vec<String>, String> {
    File::open(Path::new(proc_dir).join("cmdline"))
      .map_err(err_str)
      .and_then(|mut file| {
        let mut contents = Vec::new();
        try!(
          file.read_to_end(&mut contents)
            .map_err(err_str)
        );
        if contents.ends_with(&['\0' as u8]) {
          let _ = contents.pop();
        }
        Ok(contents)
      }).and_then(|contents| {
        String::from_utf8(contents)
          .map_err(err_str)
      }).map(|contents|
        contents
          .split('\0')
          .map(|a| a.to_string())
          .collect()
      )
  }

  // Return true if query matches this process
  fn query(&self, query: &ProcQuery) -> bool {
    match *query {
      ProcQuery::PidQuery(q) => taskid_query(self.status.pid, q),
      ProcQuery::PpidQuery(q) => taskid_query(self.status.ppid, q),
      ProcQuery::NameQuery(ref q) => string_query(&self.status.name, &q),
      ProcQuery::CmdlineQuery(ref q) => string_s_query(&self.cmdline, &q),
      ProcQuery::NoneQuery => true
    }
  }
}

pub enum ProcState {
  Running,
  Sleeping,
  Waiting,
  Zombie,
  Stopped,
  Tracing,
  Paging,
  Dead,
  Wakekill,
  Waking,
  Parked
}

fn get_procstate(state: char) -> Option<ProcState> {
  match state {
    'R' => Some(ProcState::Running),
    'S' => Some(ProcState::Sleeping),
    'D' => Some(ProcState::Waiting),
    'Z' => Some(ProcState::Zombie),
    'T' => Some(ProcState::Stopped),
    't' => Some(ProcState::Tracing),
    'X' | 'x' => Some(ProcState::Dead),
    'K' => Some(ProcState::Wakekill),
    'W' => Some(ProcState::Waking),
    'P' => Some(ProcState::Parked),
     _  => None
  }
}

#[derive(Debug)]
pub struct ProcStat {
  pub pid: TaskId,
//  pub comm: String,
//  pub state: ProcState,
//  pub ppid: TaskId,
//  pub pgrp: i32,
//  pub session: i32,
//  pub tty_nr: i32,
//  pub tpgid: i32,
//  pub flags: u32,
//  pub minflt: u64,
//  pub cminflt: u64,
//  pub majflt: u64,
//  pub cmajflt: u64,
//  pub utime: u64,
//  pub stime: u64,
//  pub cutime: i64,
//  pub cstime: i64,
//  pub priority: i64,
//  pub nice: i64,
//  pub num_threads: i64,
//  pub itrealvalue: i64,
//  pub starttime: u64,
//  pub vsize: u64,
//  pub rss: i64,
//  pub rsslim: u64,
//  pub startcode: u64,
//  pub endcode: u64,
//  pub startstack: u64,
//  pub kstkesp: u64,
//  pub kstkeip: u64,
//  pub signal: u64,
//  pub blocked: u64,
//  pub sigignore: u64,
//  pub sigcatch: u64,
//  pub wchan: u64,
//  pub nswap: u64,
//  pub cnswap: u64,
//  pub exit_signal: i32,
//  pub processor: i32,
//  pub rt_priority: u32,
//  pub policy: u32,
//  pub delayacct_blkio_ticks: u64,
//  pub guest_time: u64,
//  pub cguest_time: i64,
//  pub start_data: u64,
//  pub end_data: u64,
//  pub start_brk: u64,
//  pub arg_start: u64,
//  pub arg_end: u64,
//  pub env_start: u64,
//  pub env_end: u64,
  pub exit_code: i32
}

const READ_ERROR: &'static str = "Error parsing /proc/../stat file";

impl ProcStat {
  // Generate ProcStat struct given a process directory
  fn new(proc_dir: &str) -> Result<Self, String> {
    let file = try!(
      File::open(Path::new(proc_dir).join("stat"))
        .map_err(err_str)
    );
    let mut contents = try!(
      file.bytes().collect::<Result<Vec<_>, _>>()
        .map_err(err_str)
        .and_then(|s|
          String::from_utf8(s)
          .map_err(err_str)
        )
    );

    // Remove new-line at end
    let _ = contents.pop();

    let mut split = contents.split(' ');

    Ok(ProcStat {
      pid: try!(
        split.next().ok_or(READ_ERROR.to_owned())
        .and_then(|s|
           s.parse()
           .map_err(err_str)
        )
      ),
      // From here parse from back, since arbitrary data can be in program name
      exit_code: try!(
        split.next_back().ok_or(READ_ERROR.to_owned())
        .and_then(|s|
           s.parse()
           .map_err(err_str)
        )
      )
    })
  }
}

#[derive(Debug)]
pub struct ProcStatus {
  pub pid: TaskId,
  pub ppid: TaskId,
  pub tgid: TaskId,
  pub name: String,
}

macro_rules! extract_key {
  ($map:expr, $key:expr, $func:expr) =>
    (try!(
      $map.remove($key)
        .ok_or(format!("Key '{}' not found", $key))
        .and_then($func)
    ))
}

impl ProcStatus {
  // Generate ProcStatus struct given a process directory
  fn new(proc_dir: &str) -> Result<Self, String> {
    // Try opening file
    let status_file = try!(
      File::open(Path::new(proc_dir).join("status"))
        .map_err(err_str)
    );

    let mut status: HashMap<String, String> =
      BufReader::new(status_file).lines().filter_map(
        |line|
          line
            .map_err(err_str)
            .and_then(|line| {
              let split = line.splitn(2, ':').collect::<Vec<&str>>();

              match (split.get(0), split.get(1)) {
                (Some(key), Some(value)) =>
                  Ok((key.trim().to_owned(),
                     value.trim().to_owned())),
                _ => Err("Error reading status line".to_owned())
              }
             }).ok()
      ).collect();

    let pid = extract_key!(status, "Pid", parse_taskid);
    let ppid = extract_key!(status, "PPid", parse_taskid);
    let tgid = extract_key!(status, "Tgid", parse_taskid);
    let name = extract_key!(status, "Name", |a| Ok(a));

    Ok(ProcStatus{
      pid: pid,
      ppid: ppid,
      tgid: tgid,
      name: name
    })
  }
}

impl PartialEq for Proc {
  fn eq(&self, other: &Self) -> bool {
    self.status.pid.eq(&other.status.pid)
  }
}

impl Eq for Proc {}
impl PartialOrd for Proc {
  fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
    Some(self.cmp(other))
  }
}
impl Ord for Proc {
  fn cmp(&self, other: &Self) -> Ordering {
    self.status.pid.cmp(&other.status.pid)
  }
}

pub struct ProcIter {
  dir_iter: ReadDir,
  query: ProcQuery
}

impl ProcIter {
  pub fn new() -> Result<Self, String> {
    Self::new_query(ProcQuery::NoneQuery)
  }

  pub fn new_query(query: ProcQuery) -> Result<Self, String> {
    let proc_dir = Path::new("/proc");
    let dir_iter = try!(fs::read_dir(proc_dir).map_err(err_str));
    Ok(ProcIter{
      dir_iter: dir_iter,
      query: query
    })
  }

  fn proc_dir_filter(entry_opt: Result<DirEntry, io::Error>, query: &ProcQuery)
    -> Option<Proc> {
    // TODO: This sucks, find a better way
    entry_opt.ok()
      .and_then(|entry|
        entry.file_name().into_string().ok()
      ).and_then(|name|
        name.parse().ok()
      ).and_then(|pid|
        Proc::new(pid).ok()
      ).and_then(|proc_struct|
        if proc_struct.query(query) {
          Some(proc_struct)
        } else {
          None
        }
      )
  }
}

impl Iterator for ProcIter {
  type Item = Proc;

  fn next(&mut self) -> Option<Self::Item> {
    for entry in self.dir_iter.by_ref() {
      match Self::proc_dir_filter(entry, &self.query) {
        Some(p) => return Some(p),
        None => continue
      }
    }
    None
  }

  // Size may be anywhere from 0 to number of dirs
  fn size_hint(&self) -> (usize, Option<usize>) {
    (0, self.dir_iter.size_hint().1)
  }
}

pub type ProcMap = HashMap<TaskId, Proc>;

pub fn get_proc_map() -> Result<ProcMap, String> {
  let iter = try!(ProcIter::new());
  Ok(iter.map(|proc_struct|
    (proc_struct.status.pid, proc_struct)
  ).collect())
}

pub enum ProcQuery {
  PidQuery(TaskId),
  PpidQuery(TaskId),
  NameQuery(String),
  CmdlineQuery(String),
  NoneQuery
}

pub fn create_query(query: &str) -> Result<ProcQuery, String> {
  let splits: Vec<_> = query.splitn(2, '=').collect();

  match splits.len() {
    0 => Ok(ProcQuery::NoneQuery),
    1 => Ok(match query.parse().ok() {
      Some(tid) => ProcQuery::PidQuery(tid),
      None => ProcQuery::NameQuery(query.to_owned())
    }),
    _ => {
      let q_text = splits[1].to_owned();
      let q_tid = q_text.parse();
      match &*splits[0].to_lowercase() {
        "pid" => q_tid.map(|q| ProcQuery::PidQuery(q))
          .or(Err("Query value for type 'pid' not valid".to_owned())),
        "ppid" => q_tid.map(|q| ProcQuery::PpidQuery(q))
          .or(Err("Query value for type 'ppid' not valid".to_owned())),
        "name" => Ok(ProcQuery::NameQuery(q_text)),
        "cmdline" => Ok(ProcQuery::CmdlineQuery(q_text)),
        _ => Err("Invalid query type".to_owned())
      }
    }
  }
}

pub fn taskid_query(tid: TaskId, query: TaskId) -> bool {
  tid == query
}

pub fn string_query(text: &str, query: &str) -> bool {
  text == query
}

pub fn string_s_query(text_vec: &Vec<String>, query: &str) -> bool {
  for text in text_vec {
    if string_query(text, query) {
      return true;
    }
  }
  false
}

<commit_msg>Add parsing of /proc/.../stat file<commit_after>use std::io;
use std::io::prelude::*;
use std::io::BufReader;
use std::fs::{self, File, ReadDir, DirEntry};
use std::path::Path;
use std::collections::HashMap;
use std::cmp::Ordering;

pub type TaskId = u32;

fn err_str<T: ToString>(err: T) -> String {
  err.to_string()
}

fn parse_taskid(taskid_str: String) -> Result<TaskId, String> {
  taskid_str.parse().map_err(err_str)
}

#[derive(Debug)]
pub struct Proc {
  pub stat: Box<ProcStat>,
  pub status: Box<ProcStatus>,
  pub cmdline: Vec<String>
}

impl Proc {
  pub fn new(pid: TaskId) -> Result<Self, String> {
    let proc_dir = format!("/proc/{}", pid);
    let proc_stat = try!(ProcStat::new(&proc_dir));
    let proc_status = try!(ProcStatus::new(&proc_dir));
    let cmdline = try!(Self::read_cmdline(&proc_dir));

    let proc_struct = Proc{
      stat: Box::new(proc_stat),
      status: Box::new(proc_status),
      cmdline: cmdline
    };

    Ok(proc_struct)
  }

  fn read_cmdline(proc_dir: &str) -> Result<Vec<String>, String> {
    File::open(Path::new(proc_dir).join("cmdline"))
      .map_err(err_str)
      .and_then(|mut file| {
        let mut contents = Vec::new();
        try!(
          file.read_to_end(&mut contents)
            .map_err(err_str)
        );
        if contents.ends_with(&['\0' as u8]) {
          let _ = contents.pop();
        }
        Ok(contents)
      }).and_then(|contents| {
        String::from_utf8(contents)
          .map_err(err_str)
      }).map(|contents|
        contents
          .split('\0')
          .map(|a| a.to_string())
          .collect()
      )
  }

  // Return true if query matches this process
  fn query(&self, query: &ProcQuery) -> bool {
    match *query {
      ProcQuery::PidQuery(q) => taskid_query(self.status.pid, q),
      ProcQuery::PpidQuery(q) => taskid_query(self.status.ppid, q),
      ProcQuery::NameQuery(ref q) => string_query(&self.status.name, &q),
      ProcQuery::CmdlineQuery(ref q) => string_s_query(&self.cmdline, &q),
      ProcQuery::NoneQuery => true
    }
  }
}

#[derive(Debug)]
pub enum ProcState {
  Running,
  Sleeping,
  Waiting,
  Zombie,
  Stopped,
  Tracing,
  Paging,
  Dead,
  Wakekill,
  Waking,
  Parked
}

fn get_procstate(state: &str) -> Option<ProcState> {
  match state {
    "R" => Some(ProcState::Running),
    "S" => Some(ProcState::Sleeping),
    "D" => Some(ProcState::Waiting),
    "Z" => Some(ProcState::Zombie),
    "T" => Some(ProcState::Stopped),
    "t" => Some(ProcState::Tracing),
    "X" | "x" => Some(ProcState::Dead),
    "K" => Some(ProcState::Wakekill),
    "W" => Some(ProcState::Waking),
    "P" => Some(ProcState::Parked),
     _  => None
  }
}

#[derive(Debug)]
pub struct ProcStat {
  pub pid: TaskId,
  pub comm: String,
  pub state: ProcState,
  pub ppid: TaskId,
  pub pgrp: i32,
  pub session: i32,
  pub tty_nr: i32,
  pub tpgid: i32,
  pub flags: u32,
  pub minflt: u64,
  pub cminflt: u64,
  pub majflt: u64,
  pub cmajflt: u64,
  pub utime: u64,
  pub stime: u64,
  pub cutime: i64,
  pub cstime: i64,
  pub priority: i64,
  pub nice: i64,
  pub num_threads: i64,
  pub itrealvalue: i64,
  pub starttime: u64,
  pub vsize: u64,
  pub rss: i64,
  pub rsslim: u64,
  pub startcode: u64,
  pub endcode: u64,
  pub startstack: u64,
  pub kstkesp: u64,
  pub kstkeip: u64,
  pub signal: u64,
  pub blocked: u64,
  pub sigignore: u64,
  pub sigcatch: u64,
  pub wchan: u64,
  pub nswap: u64,
  pub cnswap: u64,
  // These fields depend on kernel version (linux 2.1 -> 3.5), so wrap in Option
  pub exit_signal: Option<i32>,
  pub processor: Option<i32>,
  pub rt_priority: Option<u32>,
  pub policy: Option<u32>,
  pub delayacct_blkio_ticks: Option<u64>,
  pub guest_time: Option<u64>,
  pub cguest_time: Option<i64>,
  pub start_data: Option<u64>,
  pub end_data: Option<u64>,
  pub start_brk: Option<u64>,
  pub arg_start: Option<u64>,
  pub arg_end: Option<u64>,
  pub env_start: Option<u64>,
  pub env_end: Option<u64>,
  pub exit_code: Option<i32>
}

const READ_ERROR: &'static str = "Error parsing /proc/../stat file";

macro_rules! stat_parse_num {
  ($item:expr) =>
    (try!(
      $item.ok_or(READ_ERROR.to_owned())
      .and_then(|s|
         s.parse()
         .map_err(err_str)
      )
    ))
}

macro_rules! stat_parse_opt_num {
  ($item:expr) =>
    ($item.and_then(|s|
       s.parse().ok()
     ))
}

impl ProcStat {
  // Generate ProcStat struct given a process directory
  fn new(proc_dir: &str) -> Result<Self, String> {
    let file = try!(
      File::open(Path::new(proc_dir).join("stat"))
        .map_err(err_str)
    );
    let bytes = try!(
      file.bytes().collect::<Result<Vec<_>, _>>()
        .map_err(err_str)
        .and_then(|s|
          String::from_utf8(s)
          .map_err(err_str)
        )
    );
    // /proc/.../stat is "numbers (prog_name) char numbers"
    // prog_name could have arbitrary characters, so we need to parse
    // the file from both ends
    let mut bytes_split = bytes.splitn(2, '(');
    let prefix = try!(bytes_split.next().ok_or(READ_ERROR.to_owned()));
    let mut bytes_split = bytes_split.next().unwrap().rsplitn(2, ')');
    // /proc/.../stat has a newline at the end
    let suffix = try!(bytes_split.next().ok_or(READ_ERROR.to_owned())).trim();
    let prog_name = try!(bytes_split.next().ok_or(READ_ERROR.to_owned()));

    let mut split = suffix.split(' ');

    Ok(ProcStat {
      pid: stat_parse_num!(prefix.split(' ').next()),
      // From here parse from back, since arbitrary data can be in program name
      comm: prog_name.to_owned(),
      state: try!(
        split.next()
          .and_then(|s|
            get_procstate(s)
          ).ok_or(READ_ERROR.to_owned())
      ),
      ppid: stat_parse_num!(split.next()),
      pgrp: stat_parse_num!(split.next()),
      session: stat_parse_num!(split.next()),
      tty_nr: stat_parse_num!(split.next()),
      tpgid: stat_parse_num!(split.next()),
      flags: stat_parse_num!(split.next()),
      minflt: stat_parse_num!(split.next()),
      cminflt: stat_parse_num!(split.next()),
      majflt: stat_parse_num!(split.next()),
      cmajflt: stat_parse_num!(split.next()),
      utime: stat_parse_num!(split.next()),
      stime: stat_parse_num!(split.next()),
      cutime: stat_parse_num!(split.next()),
      cstime: stat_parse_num!(split.next()),
      priority: stat_parse_num!(split.next()),
      nice: stat_parse_num!(split.next()),
      num_threads: stat_parse_num!(split.next()),
      itrealvalue: stat_parse_num!(split.next()),
      starttime: stat_parse_num!(split.next()),
      vsize: stat_parse_num!(split.next()),
      rss: stat_parse_num!(split.next()),
      rsslim: stat_parse_num!(split.next()),
      startcode: stat_parse_num!(split.next()),
      endcode: stat_parse_num!(split.next()),
      startstack: stat_parse_num!(split.next()),
      kstkesp: stat_parse_num!(split.next()),
      kstkeip: stat_parse_num!(split.next()),
      signal: stat_parse_num!(split.next()),
      blocked: stat_parse_num!(split.next()),
      sigignore: stat_parse_num!(split.next()),
      sigcatch: stat_parse_num!(split.next()),
      wchan: stat_parse_num!(split.next()),
      nswap: stat_parse_num!(split.next()),
      cnswap: stat_parse_num!(split.next()),
      exit_signal:
        stat_parse_opt_num!(split.next()),
      processor:
        stat_parse_opt_num!(split.next()),
      rt_priority:
        stat_parse_opt_num!(split.next()),
      policy:
        stat_parse_opt_num!(split.next()),
      delayacct_blkio_ticks:
        stat_parse_opt_num!(split.next()),
      guest_time:
        stat_parse_opt_num!(split.next()),
      cguest_time:
        stat_parse_opt_num!(split.next()),
      start_data:
        stat_parse_opt_num!(split.next()),
      end_data:
        stat_parse_opt_num!(split.next()),
      start_brk:
        stat_parse_opt_num!(split.next()),
      arg_start:
        stat_parse_opt_num!(split.next()),
      arg_end:
        stat_parse_opt_num!(split.next()),
      env_start:
        stat_parse_opt_num!(split.next()),
      env_end:
        stat_parse_opt_num!(split.next()),
      exit_code:
        stat_parse_opt_num!(split.next()),
    })
  }
}

#[derive(Debug)]
pub struct ProcStatus {
  pub pid: TaskId,
  pub ppid: TaskId,
  pub tgid: TaskId,
  pub name: String,
}

macro_rules! extract_key {
  ($map:expr, $key:expr, $func:expr) =>
    (try!(
      $map.remove($key)
        .ok_or(format!("Key '{}' not found", $key))
        .and_then($func)
    ))
}

impl ProcStatus {
  // Generate ProcStatus struct given a process directory
  fn new(proc_dir: &str) -> Result<Self, String> {
    // Try opening file
    let status_file = try!(
      File::open(Path::new(proc_dir).join("status"))
        .map_err(err_str)
    );

    let mut status: HashMap<String, String> =
      BufReader::new(status_file).lines().filter_map(
        |line|
          line
            .map_err(err_str)
            .and_then(|line| {
              let split = line.splitn(2, ':').collect::<Vec<&str>>();

              match (split.get(0), split.get(1)) {
                (Some(key), Some(value)) =>
                  Ok((key.trim().to_owned(),
                     value.trim().to_owned())),
                _ => Err("Error reading status line".to_owned())
              }
             }).ok()
      ).collect();

    let pid = extract_key!(status, "Pid", parse_taskid);
    let ppid = extract_key!(status, "PPid", parse_taskid);
    let tgid = extract_key!(status, "Tgid", parse_taskid);
    let name = extract_key!(status, "Name", |a| Ok(a));

    Ok(ProcStatus{
      pid: pid,
      ppid: ppid,
      tgid: tgid,
      name: name
    })
  }
}

impl PartialEq for Proc {
  fn eq(&self, other: &Self) -> bool {
    self.status.pid.eq(&other.status.pid)
  }
}

impl Eq for Proc {}
impl PartialOrd for Proc {
  fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
    Some(self.cmp(other))
  }
}
impl Ord for Proc {
  fn cmp(&self, other: &Self) -> Ordering {
    self.status.pid.cmp(&other.status.pid)
  }
}

pub struct ProcIter {
  dir_iter: ReadDir,
  query: ProcQuery
}

impl ProcIter {
  pub fn new() -> Result<Self, String> {
    Self::new_query(ProcQuery::NoneQuery)
  }

  pub fn new_query(query: ProcQuery) -> Result<Self, String> {
    let proc_dir = Path::new("/proc");
    let dir_iter = try!(fs::read_dir(proc_dir).map_err(err_str));
    Ok(ProcIter{
      dir_iter: dir_iter,
      query: query
    })
  }

  fn proc_dir_filter(entry_opt: Result<DirEntry, io::Error>, query: &ProcQuery)
    -> Option<Proc> {
    // TODO: This sucks, find a better way
    entry_opt.ok()
      .and_then(|entry|
        entry.file_name().into_string().ok()
      ).and_then(|name|
        name.parse().ok()
      ).and_then(|pid|
        Proc::new(pid).ok()
      ).and_then(|proc_struct|
        if proc_struct.query(query) {
          Some(proc_struct)
        } else {
          None
        }
      )
  }
}

impl Iterator for ProcIter {
  type Item = Proc;

  fn next(&mut self) -> Option<Self::Item> {
    for entry in self.dir_iter.by_ref() {
      match Self::proc_dir_filter(entry, &self.query) {
        Some(p) => return Some(p),
        None => continue
      }
    }
    None
  }

  // Size may be anywhere from 0 to number of dirs
  fn size_hint(&self) -> (usize, Option<usize>) {
    (0, self.dir_iter.size_hint().1)
  }
}

pub type ProcMap = HashMap<TaskId, Proc>;

pub fn get_proc_map() -> Result<ProcMap, String> {
  let iter = try!(ProcIter::new());
  Ok(iter.map(|proc_struct|
    (proc_struct.status.pid, proc_struct)
  ).collect())
}

pub enum ProcQuery {
  PidQuery(TaskId),
  PpidQuery(TaskId),
  NameQuery(String),
  CmdlineQuery(String),
  NoneQuery
}

pub fn create_query(query: &str) -> Result<ProcQuery, String> {
  let splits: Vec<_> = query.splitn(2, '=').collect();

  match splits.len() {
    0 => Ok(ProcQuery::NoneQuery),
    1 => Ok(match query.parse().ok() {
      Some(tid) => ProcQuery::PidQuery(tid),
      None => ProcQuery::NameQuery(query.to_owned())
    }),
    _ => {
      let q_text = splits[1].to_owned();
      let q_tid = q_text.parse();
      match &*splits[0].to_lowercase() {
        "pid" => q_tid.map(|q| ProcQuery::PidQuery(q))
          .or(Err("Query value for type 'pid' not valid".to_owned())),
        "ppid" => q_tid.map(|q| ProcQuery::PpidQuery(q))
          .or(Err("Query value for type 'ppid' not valid".to_owned())),
        "name" => Ok(ProcQuery::NameQuery(q_text)),
        "cmdline" => Ok(ProcQuery::CmdlineQuery(q_text)),
        _ => Err("Invalid query type".to_owned())
      }
    }
  }
}

pub fn taskid_query(tid: TaskId, query: TaskId) -> bool {
  tid == query
}

pub fn string_query(text: &str, query: &str) -> bool {
  text == query
}

pub fn string_s_query(text_vec: &Vec<String>, query: &str) -> bool {
  for text in text_vec {
    if string_query(text, query) {
      return true;
    }
  }
  false
}

<|endoftext|>


<commit_before>use std::rc::Rc;
use std::collections::HashMap;

use regex::{Regex, RegexSet, Captures};

#[doc(hidden)]
pub struct RouteRecognizer<T> {
    prefix: usize,
    patterns: RegexSet,
    routes: Vec<(Pattern, T)>,
}

impl<T> Default for RouteRecognizer<T> {

    fn default() -> Self {
        RouteRecognizer {
            prefix: 0,
            patterns: RegexSet::new([""].iter()).unwrap(),
            routes: Vec::new(),
        }
    }
}

impl<T> RouteRecognizer<T> {

    pub fn new<P: Into<String>, U>(prefix: P, routes: U) -> Self
        where U: IntoIterator<Item=(String, T)>
    {
        let mut paths = Vec::new();
        let mut handlers = Vec::new();
        for item in routes {
            let pat = parse(&item.0);
            handlers.push((Pattern::new(&pat), item.1));
            paths.push(pat);
        };
        let regset = RegexSet::new(&paths);

        RouteRecognizer {
            prefix: prefix.into().len() - 1,
            patterns: regset.unwrap(),
            routes: handlers,
        }
    }

    pub fn set_routes(&mut self, routes: Vec<(&str, T)>) {
        let mut paths = Vec::new();
        let mut handlers = Vec::new();
        for item in routes {
            let pat = parse(item.0);
            handlers.push((Pattern::new(&pat), item.1));
            paths.push(pat);
        };
        self.patterns = RegexSet::new(&paths).unwrap();
        self.routes = handlers;
    }

    pub fn set_prefix<P: Into<String>>(&mut self, prefix: P) {
        let p = prefix.into();
        if p.ends_with('/') {
            self.prefix = p.len() - 1;
        } else {
            self.prefix = p.len();
        }
    }

    pub fn recognize(&self, path: &str) -> Option<(Option<Params>, &T)> {
        let p = &path[self.prefix..];
        if p.is_empty() {
            if let Some(idx) = self.patterns.matches("/").into_iter().next() {
                let (ref pattern, ref route) = self.routes[idx];
                return Some((pattern.match_info(&path[self.prefix..]), route))
            }
        } else if let Some(idx) = self.patterns.matches(p).into_iter().next() {
            let (ref pattern, ref route) = self.routes[idx];
            return Some((pattern.match_info(&path[self.prefix..]), route))
        }
        None
    }
}

struct Pattern {
    re: Regex,
    names: Rc<HashMap<String, usize>>,
}

impl Pattern {
    fn new(pattern: &str) -> Self {
        let re = Regex::new(pattern).unwrap();
        let names = re.capture_names()
            .enumerate()
            .filter_map(|(i, name)| name.map(|name| (name.to_owned(), i)))
            .collect();

        Pattern {
            re,
            names: Rc::new(names),
        }
    }

    fn match_info(&self, text: &str) -> Option<Params> {
        let captures = match self.re.captures(text) {
            Some(captures) => captures,
            None => return None,
        };

        Some(Params::new(Rc::clone(&self.names), text, &captures))
    }
}

pub(crate) fn check_pattern(path: &str) {
    if let Err(err) = Regex::new(&parse(path)) {
        panic!("Wrong path pattern: \"{}\" {}", path, err);
    }
}

fn parse(pattern: &str) -> String {
    const DEFAULT_PATTERN: &str = "[^/]+";

    let mut hard_stop = false;
    let mut re = String::from("^/");
    let mut in_param = false;
    let mut in_param_pattern = false;
    let mut param_name = String::new();
    let mut param_pattern = String::from(DEFAULT_PATTERN);

    for (index, ch) in pattern.chars().enumerate() {
        // All routes must have a leading slash so its optional to have one
        if index == 0 && ch == '/' {
            continue;
        }

        if hard_stop {
            panic!("Tail '*' section has to be last lection of pattern");
        }

        if in_param {
            // In parameter segment: `{....}`
            if ch == '}' {
                if param_pattern == "*" {
                    hard_stop = true;
                    re.push_str(
                        &format!(r"(?P<{}>[%/[:word:][:punct:][:space:]]+)", &param_name));
                } else {
                    re.push_str(&format!(r"(?P<{}>{})", &param_name, &param_pattern));
                }

                param_name.clear();
                param_pattern = String::from(DEFAULT_PATTERN);

                in_param_pattern = false;
                in_param = false;
            } else if ch == ':' {
                // The parameter name has been determined; custom pattern land
                in_param_pattern = true;
                param_pattern.clear();
            } else if in_param_pattern {
                // Ignore leading whitespace for pattern
                if !(ch == ' ' && param_pattern.is_empty()) {
                    param_pattern.push(ch);
                }
            } else {
                param_name.push(ch);
            }
        } else if ch == '{' {
            in_param = true;
        } else {
            re.push(ch);
        }
    }

    re.push('$');
    re
}

/// Route match information
///
/// If resource path contains variable patterns, `Params` stores this variables.
#[derive(Debug)]
pub struct Params {
    text: String,
    matches: Vec<Option<(usize, usize)>>,
    names: Rc<HashMap<String, usize>>,
}

impl Params {
    pub(crate) fn new(names: Rc<HashMap<String, usize>>,
                      text: &str,
                      captures: &Captures) -> Self
    {
        Params {
            names,
            text: text.into(),
            matches: captures
                .iter()
                .map(|capture| capture.map(|m| (m.start(), m.end())))
                .collect(),
        }
    }

    pub(crate) fn empty() -> Self
    {
        Params {
            text: String::new(),
            names: Rc::new(HashMap::new()),
            matches: Vec::new(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.names.is_empty()
    }

    fn by_idx(&self, index: usize) -> Option<&str> {
        self.matches
            .get(index + 1)
            .and_then(|m| m.map(|(start, end)| &self.text[start..end]))
    }

    /// Get matched parameter by name
    pub fn get(&self, key: &str) -> Option<&str> {
        self.names.get(key).and_then(|&i| self.by_idx(i - 1))
    }
}

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

    fn assert_parse(pattern: &str, expected_re: &str) -> Regex {
        let re_str = parse(pattern);
        assert_eq!(&*re_str, expected_re);
        Regex::new(&re_str).unwrap()
    }

    #[test]
    fn test_parse_static() {
        let re = assert_parse("/", r"^/$");
        assert!(re.is_match("/"));
        assert!(!re.is_match("/a"));

        let re = assert_parse("/name", r"^/name$");
        assert!(re.is_match("/name"));
        assert!(!re.is_match("/name1"));
        assert!(!re.is_match("/name/"));
        assert!(!re.is_match("/name~"));

        let re = assert_parse("/name/", r"^/name/$");
        assert!(re.is_match("/name/"));
        assert!(!re.is_match("/name"));
        assert!(!re.is_match("/name/gs"));

        let re = assert_parse("/user/profile", r"^/user/profile$");
        assert!(re.is_match("/user/profile"));
        assert!(!re.is_match("/user/profile/profile"));
    }

    #[test]
    fn test_parse_param() {
        let re = assert_parse("/user/{id}", r"^/user/(?P<id>[^/]+)$");
        assert!(re.is_match("/user/profile"));
        assert!(re.is_match("/user/2345"));
        assert!(!re.is_match("/user/2345/"));
        assert!(!re.is_match("/user/2345/sdg"));

        let captures = re.captures("/user/profile").unwrap();
        assert_eq!(captures.get(1).unwrap().as_str(), "profile");
        assert_eq!(captures.name("id").unwrap().as_str(), "profile");

        let captures = re.captures("/user/1245125").unwrap();
        assert_eq!(captures.get(1).unwrap().as_str(), "1245125");
        assert_eq!(captures.name("id").unwrap().as_str(), "1245125");

        let re = assert_parse(
            "/v{version}/resource/{id}",
            r"^/v(?P<version>[^/]+)/resource/(?P<id>[^/]+)$",
        );
        assert!(re.is_match("/v1/resource/320120"));
        assert!(!re.is_match("/v/resource/1"));
        assert!(!re.is_match("/resource"));

        let captures = re.captures("/v151/resource/adahg32").unwrap();
        assert_eq!(captures.get(1).unwrap().as_str(), "151");
        assert_eq!(captures.name("version").unwrap().as_str(), "151");
        assert_eq!(captures.name("id").unwrap().as_str(), "adahg32");
    }

    #[test]
    fn test_tail_param() {
        let re = assert_parse("/user/{tail:*}",
                              r"^/user/(?P<tail>[%/[:word:][:punct:][:space:]]+)$");
        assert!(re.is_match("/user/profile"));
        assert!(re.is_match("/user/2345"));
        assert!(re.is_match("/user/2345/"));
        assert!(re.is_match("/user/2345/sdg"));
        assert!(re.is_match("/user/2345/sd-_g/"));
        assert!(re.is_match("/user/2345/sdg/asddsasd/index.html"));

        let re = assert_parse("/user/v{tail:*}",
                              r"^/user/v(?P<tail>[%/[:word:][:punct:][:space:]]+)$");
        assert!(!re.is_match("/user/2345/"));
        assert!(re.is_match("/user/vprofile"));
        assert!(re.is_match("/user/v_2345"));
        assert!(re.is_match("/user/v2345/sdg"));
        assert!(re.is_match("/user/v2345/sd-_g/test.html"));
        assert!(re.is_match("/user/v/sdg/asddsasd/index.html"));
    }
}
<commit_msg>recognizer tests<commit_after>use std::rc::Rc;
use std::collections::HashMap;

use regex::{Regex, RegexSet, Captures};

#[doc(hidden)]
pub struct RouteRecognizer<T> {
    prefix: usize,
    patterns: RegexSet,
    routes: Vec<(Pattern, T)>,
}

impl<T> Default for RouteRecognizer<T> {

    fn default() -> Self {
        RouteRecognizer {
            prefix: 0,
            patterns: RegexSet::new([""].iter()).unwrap(),
            routes: Vec::new(),
        }
    }
}

impl<T> RouteRecognizer<T> {

    pub fn new<P: Into<String>, U>(prefix: P, routes: U) -> Self
        where U: IntoIterator<Item=(String, T)>
    {
        let mut paths = Vec::new();
        let mut handlers = Vec::new();
        for item in routes {
            let pat = parse(&item.0);
            handlers.push((Pattern::new(&pat), item.1));
            paths.push(pat);
        };
        let regset = RegexSet::new(&paths);

        RouteRecognizer {
            prefix: prefix.into().len() - 1,
            patterns: regset.unwrap(),
            routes: handlers,
        }
    }

    pub fn set_routes(&mut self, routes: Vec<(&str, T)>) {
        let mut paths = Vec::new();
        let mut handlers = Vec::new();
        for item in routes {
            let pat = parse(item.0);
            handlers.push((Pattern::new(&pat), item.1));
            paths.push(pat);
        };
        self.patterns = RegexSet::new(&paths).unwrap();
        self.routes = handlers;
    }

    pub fn set_prefix<P: Into<String>>(&mut self, prefix: P) {
        let p = prefix.into();
        if p.ends_with('/') {
            self.prefix = p.len() - 1;
        } else {
            self.prefix = p.len();
        }
    }

    pub fn recognize(&self, path: &str) -> Option<(Option<Params>, &T)> {
        let p = &path[self.prefix..];
        if p.is_empty() {
            if let Some(idx) = self.patterns.matches("/").into_iter().next() {
                let (ref pattern, ref route) = self.routes[idx];
                return Some((pattern.match_info(&path[self.prefix..]), route))
            }
        } else if let Some(idx) = self.patterns.matches(p).into_iter().next() {
            let (ref pattern, ref route) = self.routes[idx];
            return Some((pattern.match_info(&path[self.prefix..]), route))
        }
        None
    }
}

struct Pattern {
    re: Regex,
    names: Rc<HashMap<String, usize>>,
}

impl Pattern {
    fn new(pattern: &str) -> Self {
        let re = Regex::new(pattern).unwrap();
        let names = re.capture_names()
            .enumerate()
            .filter_map(|(i, name)| name.map(|name| (name.to_owned(), i)))
            .collect();

        Pattern {
            re,
            names: Rc::new(names),
        }
    }

    fn match_info(&self, text: &str) -> Option<Params> {
        let captures = match self.re.captures(text) {
            Some(captures) => captures,
            None => return None,
        };

        Some(Params::new(Rc::clone(&self.names), text, &captures))
    }
}

pub(crate) fn check_pattern(path: &str) {
    if let Err(err) = Regex::new(&parse(path)) {
        panic!("Wrong path pattern: \"{}\" {}", path, err);
    }
}

fn parse(pattern: &str) -> String {
    const DEFAULT_PATTERN: &str = "[^/]+";

    let mut hard_stop = false;
    let mut re = String::from("^/");
    let mut in_param = false;
    let mut in_param_pattern = false;
    let mut param_name = String::new();
    let mut param_pattern = String::from(DEFAULT_PATTERN);

    for (index, ch) in pattern.chars().enumerate() {
        // All routes must have a leading slash so its optional to have one
        if index == 0 && ch == '/' {
            continue;
        }

        if hard_stop {
            panic!("Tail '*' section has to be last lection of pattern");
        }

        if in_param {
            // In parameter segment: `{....}`
            if ch == '}' {
                if param_pattern == "*" {
                    hard_stop = true;
                    re.push_str(
                        &format!(r"(?P<{}>[%/[:word:][:punct:][:space:]]+)", &param_name));
                } else {
                    re.push_str(&format!(r"(?P<{}>{})", &param_name, &param_pattern));
                }

                param_name.clear();
                param_pattern = String::from(DEFAULT_PATTERN);

                in_param_pattern = false;
                in_param = false;
            } else if ch == ':' {
                // The parameter name has been determined; custom pattern land
                in_param_pattern = true;
                param_pattern.clear();
            } else if in_param_pattern {
                // Ignore leading whitespace for pattern
                if !(ch == ' ' && param_pattern.is_empty()) {
                    param_pattern.push(ch);
                }
            } else {
                param_name.push(ch);
            }
        } else if ch == '{' {
            in_param = true;
        } else {
            re.push(ch);
        }
    }

    re.push('$');
    re
}

/// Route match information
///
/// If resource path contains variable patterns, `Params` stores this variables.
#[derive(Debug)]
pub struct Params {
    text: String,
    matches: Vec<Option<(usize, usize)>>,
    names: Rc<HashMap<String, usize>>,
}

impl Params {
    pub(crate) fn new(names: Rc<HashMap<String, usize>>,
                      text: &str,
                      captures: &Captures) -> Self
    {
        Params {
            names,
            text: text.into(),
            matches: captures
                .iter()
                .map(|capture| capture.map(|m| (m.start(), m.end())))
                .collect(),
        }
    }

    pub(crate) fn empty() -> Self
    {
        Params {
            text: String::new(),
            names: Rc::new(HashMap::new()),
            matches: Vec::new(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.names.is_empty()
    }

    fn by_idx(&self, index: usize) -> Option<&str> {
        self.matches
            .get(index + 1)
            .and_then(|m| m.map(|(start, end)| &self.text[start..end]))
    }

    /// Get matched parameter by name
    pub fn get(&self, key: &str) -> Option<&str> {
        self.names.get(key).and_then(|&i| self.by_idx(i - 1))
    }
}

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

    #[test]
    fn test_recognizer() {
        let mut rec = RouteRecognizer::<usize>::default();

        let routes = vec![
            ("/name", 1),
            ("/name/{val}", 2),
            ("/name/{val}/index.html", 3),
            ("/v{val}/{val2}/index.html", 4),
            ("/v/{tail:*}", 5),
        ];
        rec.set_routes(routes);

        let (params, val) = rec.recognize("/name").unwrap();
        assert_eq!(*val, 1);
        assert!(params.unwrap().is_empty());

        let (params, val) = rec.recognize("/name/value").unwrap();
        assert_eq!(*val, 2);
        assert!(!params.as_ref().unwrap().is_empty());
        assert_eq!(params.as_ref().unwrap().get("val").unwrap(), "value");

        let (params, val) = rec.recognize("/name/value2/index.html").unwrap();
        assert_eq!(*val, 3);
        assert!(!params.as_ref().unwrap().is_empty());
        assert_eq!(params.as_ref().unwrap().get("val").unwrap(), "value2");
        assert_eq!(params.as_ref().unwrap().by_idx(0).unwrap(), "value2");

        let (params, val) = rec.recognize("/vtest/ttt/index.html").unwrap();
        assert_eq!(*val, 4);
        assert!(!params.as_ref().unwrap().is_empty());
        assert_eq!(params.as_ref().unwrap().get("val").unwrap(), "test");
        assert_eq!(params.as_ref().unwrap().get("val2").unwrap(), "ttt");
        assert_eq!(params.as_ref().unwrap().by_idx(0).unwrap(), "test");
        assert_eq!(params.as_ref().unwrap().by_idx(1).unwrap(), "ttt");

        let (params, val) = rec.recognize("/v/blah-blah/index.html").unwrap();
        assert_eq!(*val, 5);
        assert!(!params.as_ref().unwrap().is_empty());
        assert_eq!(params.as_ref().unwrap().get("tail").unwrap(), "blah-blah/index.html");
    }

    fn assert_parse(pattern: &str, expected_re: &str) -> Regex {
        let re_str = parse(pattern);
        assert_eq!(&*re_str, expected_re);
        Regex::new(&re_str).unwrap()
    }

    #[test]
    fn test_parse_static() {
        let re = assert_parse("/", r"^/$");
        assert!(re.is_match("/"));
        assert!(!re.is_match("/a"));

        let re = assert_parse("/name", r"^/name$");
        assert!(re.is_match("/name"));
        assert!(!re.is_match("/name1"));
        assert!(!re.is_match("/name/"));
        assert!(!re.is_match("/name~"));

        let re = assert_parse("/name/", r"^/name/$");
        assert!(re.is_match("/name/"));
        assert!(!re.is_match("/name"));
        assert!(!re.is_match("/name/gs"));

        let re = assert_parse("/user/profile", r"^/user/profile$");
        assert!(re.is_match("/user/profile"));
        assert!(!re.is_match("/user/profile/profile"));
    }

    #[test]
    fn test_parse_param() {
        let re = assert_parse("/user/{id}", r"^/user/(?P<id>[^/]+)$");
        assert!(re.is_match("/user/profile"));
        assert!(re.is_match("/user/2345"));
        assert!(!re.is_match("/user/2345/"));
        assert!(!re.is_match("/user/2345/sdg"));

        let captures = re.captures("/user/profile").unwrap();
        assert_eq!(captures.get(1).unwrap().as_str(), "profile");
        assert_eq!(captures.name("id").unwrap().as_str(), "profile");

        let captures = re.captures("/user/1245125").unwrap();
        assert_eq!(captures.get(1).unwrap().as_str(), "1245125");
        assert_eq!(captures.name("id").unwrap().as_str(), "1245125");

        let re = assert_parse(
            "/v{version}/resource/{id}",
            r"^/v(?P<version>[^/]+)/resource/(?P<id>[^/]+)$",
        );
        assert!(re.is_match("/v1/resource/320120"));
        assert!(!re.is_match("/v/resource/1"));
        assert!(!re.is_match("/resource"));

        let captures = re.captures("/v151/resource/adahg32").unwrap();
        assert_eq!(captures.get(1).unwrap().as_str(), "151");
        assert_eq!(captures.name("version").unwrap().as_str(), "151");
        assert_eq!(captures.name("id").unwrap().as_str(), "adahg32");
    }

    #[test]
    fn test_tail_param() {
        let re = assert_parse("/user/{tail:*}",
                              r"^/user/(?P<tail>[%/[:word:][:punct:][:space:]]+)$");
        assert!(re.is_match("/user/profile"));
        assert!(re.is_match("/user/2345"));
        assert!(re.is_match("/user/2345/"));
        assert!(re.is_match("/user/2345/sdg"));
        assert!(re.is_match("/user/2345/sd-_g/"));
        assert!(re.is_match("/user/2345/sdg/asddsasd/index.html"));

        let re = assert_parse("/user/v{tail:*}",
                              r"^/user/v(?P<tail>[%/[:word:][:punct:][:space:]]+)$");
        assert!(!re.is_match("/user/2345/"));
        assert!(re.is_match("/user/vprofile"));
        assert!(re.is_match("/user/v_2345"));
        assert!(re.is_match("/user/v2345/sdg"));
        assert!(re.is_match("/user/v2345/sd-_g/test.html"));
        assert!(re.is_match("/user/v/sdg/asddsasd/index.html"));
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A small fast RNG

use rand_core::{RngCore, SeedableRng, Error};

#[cfg(all(not(target_os = "emscripten"), target_pointer_width = "64"))]
type Rng = rand_pcg::Pcg64Mcg;
#[cfg(not(all(not(target_os = "emscripten"), target_pointer_width = "64")))]
type Rng = rand_pcg::Pcg32;

/// A small-state, fast non-crypto PRNG
///
/// `SmallRng` may be a good choice when a PRNG with small state, cheap
/// initialization, good statistical quality and good performance are required.
/// It is **not** a good choice when security against prediction or
/// reproducibility are important.
///
/// This PRNG is **feature-gated**: to use, you must enable the crate feature
/// `small_rng`.
///
/// The algorithm is deterministic but should not be considered reproducible
/// due to dependence on platform and possible replacement in future
/// library versions. For a reproducible generator, use a named PRNG from an
/// external crate, e.g. [rand_pcg] or [rand_chacha].
/// Refer also to [The Book](https://rust-random.github.io/book/guide-rngs.html).
///
/// The PRNG algorithm in `SmallRng` is chosen to be
/// efficient on the current platform, without consideration for cryptography
/// or security. The size of its state is much smaller than [`StdRng`].
/// The current algorithm is [`Pcg64Mcg`](rand_pcg::Pcg64Mcg) on 64-bit
/// platforms and [`Pcg32`](rand_pcg::Pcg32) on 32-bit platforms. Both are
/// implemented by the [rand_pcg] crate.
///
/// # Examples
///
/// Initializing `SmallRng` with a random seed can be done using [`SeedableRng::from_entropy`]:
///
/// ```
/// use rand::{Rng, SeedableRng};
/// use rand::rngs::SmallRng;
///
/// // Create small, cheap to initialize and fast RNG with a random seed.
/// // The randomness is supplied by the operating system.
/// let mut small_rng = SmallRng::from_entropy();
/// # let v: u32 = small_rng.gen();
/// ```
///
/// When initializing a lot of `SmallRng`'s, using [`thread_rng`] can be more
/// efficient:
///
/// ```
/// use std::iter;
/// use rand::{SeedableRng, thread_rng};
/// use rand::rngs::SmallRng;
///
/// // Create a big, expensive to initialize and slower, but unpredictable RNG.
/// // This is cached and done only once per thread.
/// let mut thread_rng = thread_rng();
/// // Create small, cheap to initialize and fast RNGs with random seeds.
/// // One can generally assume this won't fail.
/// let rngs: Vec<SmallRng> = iter::repeat(())
///     .map(|()| SmallRng::from_rng(&mut thread_rng).unwrap())
///     .take(10)
///     .collect();
/// ```
///
/// [`StdRng`]: crate::rngs::StdRng
/// [`thread_rng`]: crate::thread_rng
/// [rand_chacha]: https://crates.io/crates/rand_chacha
/// [rand_pcg]: https://crates.io/crates/rand_pcg
#[derive(Clone, Debug)]
pub struct SmallRng(Rng);

impl RngCore for SmallRng {
    #[inline(always)]
    fn next_u32(&mut self) -> u32 {
        self.0.next_u32()
    }

    #[inline(always)]
    fn next_u64(&mut self) -> u64 {
        self.0.next_u64()
    }

    #[inline(always)]
    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.0.fill_bytes(dest);
    }

    #[inline(always)]
    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.0.try_fill_bytes(dest)
    }
}

impl SeedableRng for SmallRng {
    type Seed = <Rng as SeedableRng>::Seed;

    #[inline(always)]
    fn from_seed(seed: Self::Seed) -> Self {
        SmallRng(Rng::from_seed(seed))
    }

    #[inline(always)]
    fn from_rng<R: RngCore>(rng: R) -> Result<Self, Error> {
        Rng::from_rng(rng).map(SmallRng)
    }
}
<commit_msg>Small simplification of doc example<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A small fast RNG

use rand_core::{RngCore, SeedableRng, Error};

#[cfg(all(not(target_os = "emscripten"), target_pointer_width = "64"))]
type Rng = rand_pcg::Pcg64Mcg;
#[cfg(not(all(not(target_os = "emscripten"), target_pointer_width = "64")))]
type Rng = rand_pcg::Pcg32;

/// A small-state, fast non-crypto PRNG
///
/// `SmallRng` may be a good choice when a PRNG with small state, cheap
/// initialization, good statistical quality and good performance are required.
/// It is **not** a good choice when security against prediction or
/// reproducibility are important.
///
/// This PRNG is **feature-gated**: to use, you must enable the crate feature
/// `small_rng`.
///
/// The algorithm is deterministic but should not be considered reproducible
/// due to dependence on platform and possible replacement in future
/// library versions. For a reproducible generator, use a named PRNG from an
/// external crate, e.g. [rand_pcg] or [rand_chacha].
/// Refer also to [The Book](https://rust-random.github.io/book/guide-rngs.html).
///
/// The PRNG algorithm in `SmallRng` is chosen to be
/// efficient on the current platform, without consideration for cryptography
/// or security. The size of its state is much smaller than [`StdRng`].
/// The current algorithm is [`Pcg64Mcg`](rand_pcg::Pcg64Mcg) on 64-bit
/// platforms and [`Pcg32`](rand_pcg::Pcg32) on 32-bit platforms. Both are
/// implemented by the [rand_pcg] crate.
///
/// # Examples
///
/// Initializing `SmallRng` with a random seed can be done using [`SeedableRng::from_entropy`]:
///
/// ```
/// use rand::{Rng, SeedableRng};
/// use rand::rngs::SmallRng;
///
/// // Create small, cheap to initialize and fast RNG with a random seed.
/// // The randomness is supplied by the operating system.
/// let mut small_rng = SmallRng::from_entropy();
/// # let v: u32 = small_rng.gen();
/// ```
///
/// When initializing a lot of `SmallRng`'s, using [`thread_rng`] can be more
/// efficient:
///
/// ```
/// use rand::{SeedableRng, thread_rng};
/// use rand::rngs::SmallRng;
///
/// // Create a big, expensive to initialize and slower, but unpredictable RNG.
/// // This is cached and done only once per thread.
/// let mut thread_rng = thread_rng();
/// // Create small, cheap to initialize and fast RNGs with random seeds.
/// // One can generally assume this won't fail.
/// let rngs: Vec<SmallRng> = (0..10)
///     .map(|_| SmallRng::from_rng(&mut thread_rng).unwrap())
///     .collect();
/// ```
///
/// [`StdRng`]: crate::rngs::StdRng
/// [`thread_rng`]: crate::thread_rng
/// [rand_chacha]: https://crates.io/crates/rand_chacha
/// [rand_pcg]: https://crates.io/crates/rand_pcg
#[derive(Clone, Debug)]
pub struct SmallRng(Rng);

impl RngCore for SmallRng {
    #[inline(always)]
    fn next_u32(&mut self) -> u32 {
        self.0.next_u32()
    }

    #[inline(always)]
    fn next_u64(&mut self) -> u64 {
        self.0.next_u64()
    }

    #[inline(always)]
    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.0.fill_bytes(dest);
    }

    #[inline(always)]
    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.0.try_fill_bytes(dest)
    }
}

impl SeedableRng for SmallRng {
    type Seed = <Rng as SeedableRng>::Seed;

    #[inline(always)]
    fn from_seed(seed: Self::Seed) -> Self {
        SmallRng(Rng::from_seed(seed))
    }

    #[inline(always)]
    fn from_rng<R: RngCore>(rng: R) -> Result<Self, Error> {
        Rng::from_rng(rng).map(SmallRng)
    }
}
<|endoftext|>


<commit_before>//! Implements the searching of the game tree.

pub mod alpha_beta;
pub mod threading;

use std::thread;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver};
use std::time::SystemTime;
use basetypes::*;
use chess_move::*;
use position::Position;
use tt::*;
use self::threading::*;

/// The maximum search depth in half-moves.
pub const MAX_DEPTH: u8 = 63; // Should be less than 127.

/// The half-with of the initial aspiration window in centipawns.
pub const INITIAL_ASPIRATION_WINDOW: Value = 17; // 16;

/// The number of nodes that will be searched without reporting search
/// progress.
///
/// If this value is too small the engine may become slow, if this
/// value is too big the engine may become unresponsive.
pub const NODE_COUNT_REPORT_INTERVAL: NodeCount = 10000;

/// A sequence of moves from some starting position, together with the
/// value assigned to the final position.
pub struct Variation {
    /// A sequence of moves from some starting position.
    pub moves: Vec<Move>,

    /// The value assigned to the final position.
    pub value: Value,

    /// The accuracy of the assigned value.
    pub bound: BoundType,
}

/// Contains information about the current progress of a search.
pub struct SearchStatus {
    started_at: Option<SystemTime>,

    /// `true` if the search is finished, `false` otherwise.
    pub done: bool,

    /// The reached search depth.
    pub depth: u8,

    pub variations: Vec<Variation>,

    /// Number of milliseconds since the beginning of the search.
    pub duration_millis: u64,

    /// Number of analyzed nodes since the beginning of the search.
    pub searched_nodes: NodeCount,

    /// Number of analyzed nodes per second.
    pub nps: NodeCount,
}

/// Executes game tree searches in different positions.
pub struct SearchExecutor {
    tt: Arc<Tt>,
    position: Position,
    status: SearchStatus,

    // A handle to the search thread.
    search_thread: Option<thread::JoinHandle<()>>,

    // A channel for sending commands to the search thread.
    commands: Sender<Command>,

    // A channel for receiving reports from the search thread.
    reports: Receiver<Report>,
}

impl SearchExecutor {
    /// Creates a new instance.
    pub fn new(tt: Arc<Tt>) -> SearchExecutor {

        // Spawn the search thread.
        let (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        let tt_clone = tt.clone();
        let search_thread = thread::spawn(move || {
            serve_deepening(tt_clone, commands_rx, reports_tx);
        });

        SearchExecutor {
            tt: tt,
            search_thread: Some(search_thread),
            position: Position::from_fen("k7/8/8/8/8/8/8/7K w - - 0 1").ok().unwrap(),
            commands: commands_tx,
            reports: reports_rx,
            status: SearchStatus {
                started_at: None,
                done: true,
                depth: 0,
                variations: vec![Variation {
                                     value: -19999,
                                     bound: BOUND_LOWER,
                                     moves: vec![],
                                 }],
                searched_nodes: 0,
                duration_millis: 0,
                nps: 0,
            },
        }
    }

    /// Starts a new search.
    #[allow(unused_variables)]
    pub fn start(&mut self, p: &Position, searchmoves: Option<Vec<String>>, pv_count: usize) {
        // TODO: We ignore the "pv_count" parameter.
        assert_eq!(pv_count, 1);

        // TODO: We ignore the "searchmoves" parameter.

        // TODO: Add `self.legal_moves_count` filed.

        self.stop();
        self.position = p.clone();
        self.status = SearchStatus {
            started_at: Some(SystemTime::now()),
            done: false,
            depth: 0,
            variations: vec![Variation {
                                 value: -19999,
                                 bound: BOUND_LOWER,
                                 moves: vec![],
                             }], // TODO: Set good initial value (vec![best_move]).
            searched_nodes: 0,
            duration_millis: 0,
            ..self.status
        };
        self.commands
            .send(Command::Search {
                search_id: 0,
                position: p.clone(),
                depth: MAX_DEPTH,
                lower_bound: -29999,
                upper_bound: 29999,
            })
            .unwrap();
    }

    /// Stops the current search.
    pub fn stop(&mut self) {
        if !self.status.done {
            self.commands.send(Command::Stop).unwrap();
            loop {
                if let Ok(Report::Done { .. }) = self.reports.recv() {
                    break;
                }
            }
            self.status.done = true;
        }
    }

    /// Stops the current search and joins the search thread.
    ///
    /// After calling `exit`, no other methods on this instance should
    /// be called.
    pub fn exit(&mut self) {
        self.commands.send(Command::Exit).unwrap();
        self.search_thread.take().unwrap().join().unwrap();
    }

    /// Returns current search's status.
    #[inline(always)]
    pub fn status(&self) -> &SearchStatus {
        &self.status
    }

    /// Update current search's status.
    pub fn update_status(&mut self) {
        while let Ok(report) = self.reports.try_recv() {
            match report {
                Report::Progress { searched_depth, searched_nodes, .. } => {
                    self.register_progress(searched_depth, searched_nodes);
                }
                Report::Done { .. } => {
                    self.status.done = true;
                }
            }
        }
    }

    // A helper method. It updates the search status info and makes
    // sure that a new PV is sent to the GUI for each newly reached
    // depth.
    fn register_progress(&mut self, depth: u8, searched_nodes: NodeCount) {
        let duration = self.status.started_at.unwrap().elapsed().unwrap();
        self.status.duration_millis = 1000 * duration.as_secs() +
                                      (duration.subsec_nanos() / 1000000) as u64;
        self.status.searched_nodes = searched_nodes;
        self.status.nps = 1000 * (self.status.nps + self.status.searched_nodes) /
                          (1000 + self.status.duration_millis);
        if self.status.depth < depth {
            self.status.depth = depth;
            self.extract_pv(depth);
        }
    }

    // A helper method. It extracts the primary variation (PV) from
    // the transposition table (TT) and updates `status.current_pv`.
    //
    // **Note:** Because the PV is a moving target (the search
    // continues to run in parallel), imperfections in the reported
    // PVs are unavoidable. To deal with this, we turn a blind eye if
    // the value at the root of the PV differs from the value at the
    // leaf by no more than `EPSILON`.
    fn extract_pv(&mut self, depth: u8) {
        if depth == 0 {
            return;
        }

        // First: Extract the PV, the leaf value, the root value, and
        // the bound type from the TT.
        let mut p = self.position.clone();
        let mut our_turn = true;
        let mut prev_move = None;
        let mut moves = Vec::new();
        let mut leaf_value = -19999;
        let mut root_value = leaf_value;
        let mut bound = BOUND_LOWER;
        while let Some(entry) = self.tt.peek(p.hash()) {
            if entry.bound() != BOUND_NONE {
                // Get the value and the bound type. In half of the
                // cases the value stored in `entry` is from other
                // side's perspective.
                if our_turn {
                    leaf_value = entry.value();
                    bound = entry.bound();
                } else {
                    leaf_value = -entry.value();
                    bound = match entry.bound() {
                        BOUND_UPPER => BOUND_LOWER,
                        BOUND_LOWER => BOUND_UPPER,
                        x => x,
                    };
                }

                // The values under -19999 and over 19999 carry
                // information about in how many moves is the
                // inevitable checkmate. However, do not show this to
                // the user, because it is sometimes incorrect.
                if leaf_value >= 20000 {
                    leaf_value = 19999;
                    if bound == BOUND_LOWER {
                        bound = BOUND_EXACT
                    }
                }
                if leaf_value <= -20000 {
                    leaf_value = -19999;
                    if bound == BOUND_UPPER {
                        bound = BOUND_EXACT
                    }
                }

                if let Some(m) = prev_move {
                    // Extend the PV with the previous move.
                    moves.push(m);
                } else {
                    // We are at the root -- set the root value.
                    root_value = leaf_value;
                }

                if moves.len() < depth as usize && (leaf_value - root_value).abs() <= EPSILON {
                    if let Some(m) = p.try_move_digest(entry.move16()) {
                        if p.do_move(m) {
                            if bound == BOUND_EXACT {
                                // Extend the PV with one more move.
                                prev_move = Some(m);
                                our_turn = !our_turn;
                                continue;
                            } else {
                                // This is the last move in the PV.
                                moves.push(m);
                            }
                        }
                    }
                }
            }
            break;
        }

        // Second: Change the bound type if the leaf value in the PV
        // differs too much from the root value.
        bound = match leaf_value - root_value {
            x if x > EPSILON && bound != BOUND_UPPER => BOUND_LOWER,
            x if x < -EPSILON && bound != BOUND_LOWER => BOUND_UPPER,
            _ => bound,
        };

        // Third: Update `status`.
        self.status.variations = vec![Variation {
                                          value: root_value,
                                          bound: bound,
                                          moves: moves,
                                      }];
    }
}

// A sufficiently small value (in centipawns).
const EPSILON: Value = 8;
<commit_msg>docs + stop() in exit()<commit_after>//! Implements the searching of the game tree.

pub mod alpha_beta;
pub mod threading;

use std::thread;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver};
use std::time::SystemTime;
use basetypes::*;
use chess_move::*;
use position::Position;
use tt::*;
use self::threading::*;

/// The maximum search depth in half-moves.
pub const MAX_DEPTH: u8 = 63; // Should be less than 127.

/// The half-with of the initial aspiration window in centipawns.
pub const INITIAL_ASPIRATION_WINDOW: Value = 17; // 16;

/// The number of nodes that will be searched without reporting search
/// progress.
///
/// If this value is too small the engine may become slow, if this
/// value is too big the engine may become unresponsive.
pub const NODE_COUNT_REPORT_INTERVAL: NodeCount = 10000;

/// A sequence of moves from some starting position, together with the
/// value assigned to the final position.
pub struct Variation {
    /// A sequence of moves from some starting position.
    pub moves: Vec<Move>,

    /// The value assigned to the final position.
    pub value: Value,

    /// The accuracy of the assigned value.
    pub bound: BoundType,
}

/// Contains information about the current progress of a search.
pub struct SearchStatus {
    started_at: Option<SystemTime>,

    /// `true` if the search is finished or has been stopped, `false`
    /// otherwise.
    pub done: bool,

    /// The reached search depth.
    pub depth: u8,

    pub variations: Vec<Variation>,

    /// Number of milliseconds since the beginning of the search.
    pub duration_millis: u64,

    /// Number of analyzed nodes since the beginning of the search.
    pub searched_nodes: NodeCount,

    /// Number of analyzed nodes per second.
    pub nps: NodeCount,
}

/// Executes searches in different positions.
pub struct SearchExecutor {
    tt: Arc<Tt>,
    position: Position,
    status: SearchStatus,

    // A handle to the search thread.
    search_thread: Option<thread::JoinHandle<()>>,

    // A channel for sending commands to the search thread.
    commands: Sender<Command>,

    // A channel for receiving reports from the search thread.
    reports: Receiver<Report>,
}

impl SearchExecutor {
    /// Creates a new instance.
    pub fn new(tt: Arc<Tt>) -> SearchExecutor {

        // Spawn the search thread.
        let (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        let tt_clone = tt.clone();
        let search_thread = thread::spawn(move || {
            serve_deepening(tt_clone, commands_rx, reports_tx);
        });

        SearchExecutor {
            tt: tt,
            search_thread: Some(search_thread),
            position: Position::from_fen("k7/8/8/8/8/8/8/7K w - - 0 1").ok().unwrap(),
            commands: commands_tx,
            reports: reports_rx,
            status: SearchStatus {
                started_at: None,
                done: true,
                depth: 0,
                variations: vec![Variation {
                                     value: -19999,
                                     bound: BOUND_LOWER,
                                     moves: vec![],
                                 }],
                searched_nodes: 0,
                duration_millis: 0,
                nps: 0,
            },
        }
    }

    /// Starts a new search.
    #[allow(unused_variables)]
    pub fn start(&mut self, p: &Position, searchmoves: Option<Vec<String>>, pv_count: usize) {
        // TODO: We ignore the "pv_count" parameter.
        assert_eq!(pv_count, 1);

        // TODO: We ignore the "searchmoves" parameter.

        // TODO: Add `self.legal_moves_count` filed.

        self.stop();
        self.position = p.clone();
        self.status = SearchStatus {
            started_at: Some(SystemTime::now()),
            done: false,
            depth: 0,
            variations: vec![Variation {
                                 value: -19999,
                                 bound: BOUND_LOWER,
                                 moves: vec![],
                             }], // TODO: Set good initial value (vec![best_move]).
            searched_nodes: 0,
            duration_millis: 0,
            ..self.status
        };
        self.commands
            .send(Command::Search {
                search_id: 0,
                position: p.clone(),
                depth: MAX_DEPTH,
                lower_bound: -29999,
                upper_bound: 29999,
            })
            .unwrap();
    }

    /// Stops the current search.
    pub fn stop(&mut self) {
        if !self.status.done {
            self.commands.send(Command::Stop).unwrap();
            loop {
                if let Ok(Report::Done { .. }) = self.reports.recv() {
                    break;
                }
            }
            self.status.done = true;
        }
    }

    /// Stops the current search and retires the current instance.
    ///
    /// After calling `exit`, no other methods on this instance should
    /// be called.
    pub fn exit(&mut self) {
        self.stop();
        self.commands.send(Command::Exit).unwrap();
        self.search_thread.take().unwrap().join().unwrap();
    }

    /// Returns the status of the current search.
    ///
    /// **Important note:** Consecutive calls to this method will
    /// return the same unchanged result. Only after calling
    /// `update_status` or `start`, the result returned by `status`
    /// may change.
    #[inline(always)]
    pub fn status(&self) -> &SearchStatus {
        &self.status
    }

    /// Updates the status of the current search.
    pub fn update_status(&mut self) {
        while let Ok(report) = self.reports.try_recv() {
            match report {
                Report::Progress { searched_depth, searched_nodes, .. } => {
                    self.register_progress(searched_depth, searched_nodes);
                }
                Report::Done { .. } => {
                    self.status.done = true;
                }
            }
        }
    }

    // A helper method. It updates the search status info and makes
    // sure that a new PV is sent to the GUI for each newly reached
    // depth.
    fn register_progress(&mut self, depth: u8, searched_nodes: NodeCount) {
        let duration = self.status.started_at.unwrap().elapsed().unwrap();
        self.status.duration_millis = 1000 * duration.as_secs() +
                                      (duration.subsec_nanos() / 1000000) as u64;
        self.status.searched_nodes = searched_nodes;
        self.status.nps = 1000 * (self.status.nps + self.status.searched_nodes) /
                          (1000 + self.status.duration_millis);
        if self.status.depth < depth {
            self.status.depth = depth;
            self.extract_pv(depth);
        }
    }

    // A helper method. It extracts the primary variation (PV) from
    // the transposition table (TT) and updates `status.current_pv`.
    //
    // **Note:** Because the PV is a moving target (the search
    // continues to run in parallel), imperfections in the reported
    // PVs are unavoidable. To deal with this, we turn a blind eye if
    // the value at the root of the PV differs from the value at the
    // leaf by no more than `EPSILON`.
    fn extract_pv(&mut self, depth: u8) {
        if depth == 0 {
            return;
        }

        // First: Extract the PV, the leaf value, the root value, and
        // the bound type from the TT.
        let mut p = self.position.clone();
        let mut our_turn = true;
        let mut prev_move = None;
        let mut moves = Vec::new();
        let mut leaf_value = -19999;
        let mut root_value = leaf_value;
        let mut bound = BOUND_LOWER;
        while let Some(entry) = self.tt.peek(p.hash()) {
            if entry.bound() != BOUND_NONE {
                // Get the value and the bound type. In half of the
                // cases the value stored in `entry` is from other
                // side's perspective.
                if our_turn {
                    leaf_value = entry.value();
                    bound = entry.bound();
                } else {
                    leaf_value = -entry.value();
                    bound = match entry.bound() {
                        BOUND_UPPER => BOUND_LOWER,
                        BOUND_LOWER => BOUND_UPPER,
                        x => x,
                    };
                }

                // The values under -19999 and over 19999 carry
                // information about in how many moves is the
                // inevitable checkmate. However, do not show this to
                // the user, because it is sometimes incorrect.
                if leaf_value >= 20000 {
                    leaf_value = 19999;
                    if bound == BOUND_LOWER {
                        bound = BOUND_EXACT
                    }
                }
                if leaf_value <= -20000 {
                    leaf_value = -19999;
                    if bound == BOUND_UPPER {
                        bound = BOUND_EXACT
                    }
                }

                if let Some(m) = prev_move {
                    // Extend the PV with the previous move.
                    moves.push(m);
                } else {
                    // We are at the root -- set the root value.
                    root_value = leaf_value;
                }

                if moves.len() < depth as usize && (leaf_value - root_value).abs() <= EPSILON {
                    if let Some(m) = p.try_move_digest(entry.move16()) {
                        if p.do_move(m) {
                            if bound == BOUND_EXACT {
                                // Extend the PV with one more move.
                                prev_move = Some(m);
                                our_turn = !our_turn;
                                continue;
                            } else {
                                // This is the last move in the PV.
                                moves.push(m);
                            }
                        }
                    }
                }
            }
            break;
        }

        // Second: Change the bound type if the leaf value in the PV
        // differs too much from the root value.
        bound = match leaf_value - root_value {
            x if x > EPSILON && bound != BOUND_UPPER => BOUND_LOWER,
            x if x < -EPSILON && bound != BOUND_LOWER => BOUND_UPPER,
            _ => bound,
        };

        // Third: Update `status`.
        self.status.variations = vec![Variation {
                                          value: root_value,
                                          bound: bound,
                                          moves: moves,
                                      }];
    }
}

// A sufficiently small value (in centipawns).
const EPSILON: Value = 8;
<|endoftext|>