<commit_before>use std::ffi::CStr;
use std::marker::PhantomData;

use llvm_sys::prelude::*;
use llvm_sys::core::*;

use id::IdRef;

use llvm::{Context, BasicBlock, Value};

pub struct Builder<'cid> {
    _context_id: IdRef<'cid>,
    llvm_builder: LLVMBuilderRef
}

impl<'cid> Drop for Builder<'cid> {
    fn drop(&mut self) {
        unsafe {
            LLVMDisposeBuilder(self.llvm_builder);
        }
    }
}

impl<'cid> Builder<'cid> {
    pub fn new(context: &Context<'cid>) -> Builder<'cid> {
        Builder {
            _context_id: IdRef::new(),
            llvm_builder: unsafe { LLVMCreateBuilderInContext(context.as_raw()) }
        }
    }

    pub fn position_at_end<'fid, 'function, 'builder>(&'builder mut self, block: BasicBlock<'cid, 'fid, 'function>) -> PositionedBuilder<'cid, 'fid, 'function, 'builder> {
        unsafe {
            LLVMPositionBuilderAtEnd(self.llvm_builder, block.as_raw());
        }
        PositionedBuilder {
            _context_id: IdRef::new(),
            _function_id: IdRef::new(),
            _function: PhantomData,
            _builder: PhantomData,
            llvm_builder: self.llvm_builder
        }
    }
}

pub struct PositionedBuilder<'cid: 'builder, 'fid, 'function, 'builder> {
    _context_id: IdRef<'cid>,
    _function_id: IdRef<'fid>,
    _function: PhantomData<&'function ()>,
    _builder: PhantomData<&'builder mut Builder<'cid>>,
    llvm_builder: LLVMBuilderRef
}

impl<'cid: 'builder, 'fid, 'function, 'builder> PositionedBuilder<'cid, 'fid, 'function, 'builder> {
    pub fn br(&self, target: BasicBlock<'cid, 'fid, 'function>) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildBr(self.llvm_builder, target.as_raw()))
        }
    }

    pub fn or(&self, lhs: Value<'cid, 'fid, 'function>, rhs: Value<'cid, 'fid, 'function>, name: &CStr) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildOr(self.llvm_builder, lhs.as_raw(), rhs.as_raw(), name.as_ptr()))
        }
    }

    pub fn add(&self, lhs: Value<'cid, 'fid, 'function>, rhs: Value<'cid, 'fid, 'function>, name: &CStr) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildAdd(self.llvm_builder, lhs.as_raw(), rhs.as_raw(), name.as_ptr()))
        }
    }

    pub fn and(&self, lhs: Value<'cid, 'fid, 'function>, rhs: Value<'cid, 'fid, 'function>, name: &CStr) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildAnd(self.llvm_builder, lhs.as_raw(), rhs.as_raw(), name.as_ptr()))
        }
    }

    pub fn mul(&self, lhs: Value<'cid, 'fid, 'function>, rhs: Value<'cid, 'fid, 'function>, name: &CStr) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildMul(self.llvm_builder, lhs.as_raw(), rhs.as_raw(), name.as_ptr()))
        }
    }

    pub fn neg(&self, value: Value<'cid, 'fid, 'function>, name: &CStr) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildNeg(self.llvm_builder, value.as_raw(), name.as_ptr()))
        }
    }

    pub fn not(&self, value: Value<'cid, 'fid, 'function>, name: &CStr) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildNot(self.llvm_builder, value.as_raw(), name.as_ptr()))
        }
    }

    pub fn ret(&self, value: Value<'cid, 'fid, 'function>) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildRet(self.llvm_builder, value.as_raw()))
        }
    }
}
<commit_msg>Continue unboxing LLVM types: Builder and PositionedBuilder<commit_after>use std::ffi::CStr;
use std::marker::PhantomData;

use llvm_sys::prelude::*;
use llvm_sys::core::*;

use id::IdRef;
use owned::{Owned, DropInPlace};

use llvm::{Context, BasicBlock, Value};

pub struct Builder<'cid: 'context, 'context> {
    _context: PhantomData<&'context Context<'cid>>
}

impl<'cid, 'context> DropInPlace for Builder<'cid, 'context> {
    unsafe fn drop_in_place(&mut self) {
        LLVMDisposeBuilder(self.as_raw());
    }
}

impl<'cid, 'context> Builder<'cid, 'context> {
    pub fn new(context: &'context Context<'cid>) -> Owned<Builder<'cid, 'context>> {
        unsafe {
            Owned::from_raw(
                LLVMCreateBuilderInContext(context.as_raw()) as *mut Builder
            )
        }
    }

    pub fn position_at_end<'fid, 'function, 'builder>(&'builder mut self, block: BasicBlock<'cid, 'fid, 'function>) -> &'builder mut PositionedBuilder<'cid, 'context, 'fid, 'function> {
        unsafe {
            LLVMPositionBuilderAtEnd(self.as_raw(), block.as_raw());
            &mut *(self as *mut Builder as *mut PositionedBuilder)
        }
    }

    pub fn as_raw(&self) -> LLVMBuilderRef {
        self as *const Builder as *mut Builder as LLVMBuilderRef
    }
}

pub struct PositionedBuilder<'cid: 'context, 'context, 'fid, 'function> {
    _function_id: IdRef<'fid>,
    _function: PhantomData<&'function ()>,
    _builder: PhantomData<Builder<'cid, 'context>>
}

impl<'cid, 'context, 'fid, 'function> PositionedBuilder<'cid, 'context, 'fid, 'function> {
    pub fn br(&mut self, target: BasicBlock<'cid, 'fid, 'function>) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildBr(self.as_raw(), target.as_raw()))
        }
    }

    pub fn or(&mut self, lhs: Value<'cid, 'fid, 'function>, rhs: Value<'cid, 'fid, 'function>, name: &CStr) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildOr(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()))
        }
    }

    pub fn add(&mut self, lhs: Value<'cid, 'fid, 'function>, rhs: Value<'cid, 'fid, 'function>, name: &CStr) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildAdd(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()))
        }
    }

    pub fn and(&mut self, lhs: Value<'cid, 'fid, 'function>, rhs: Value<'cid, 'fid, 'function>, name: &CStr) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildAnd(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()))
        }
    }

    pub fn mul(&mut self, lhs: Value<'cid, 'fid, 'function>, rhs: Value<'cid, 'fid, 'function>, name: &CStr) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildMul(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()))
        }
    }

    pub fn neg(&mut self, value: Value<'cid, 'fid, 'function>, name: &CStr) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildNeg(self.as_raw(), value.as_raw(), name.as_ptr()))
        }
    }

    pub fn not(&mut self, value: Value<'cid, 'fid, 'function>, name: &CStr) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildNot(self.as_raw(), value.as_raw(), name.as_ptr()))
        }
    }

    pub fn ret(&mut self, value: Value<'cid, 'fid, 'function>) -> Value<'cid, 'fid, 'function> {
        unsafe {
            Value::from_raw(LLVMBuildRet(self.as_raw(), value.as_raw()))
        }
    }

    pub fn as_raw(&self) -> LLVMBuilderRef {
        self as *const PositionedBuilder as *mut PositionedBuilder as LLVMBuilderRef
    }
}
<|endoftext|>


<commit_before>use super::*;

pub struct Iter<'a> {
    pub(super) config: &'a Config,
    pub(super) segment_iter: Box<Iterator<Item = (Lsn, LogID)>>,
    pub(super) segment_base: Option<LogID>,
    pub(super) segment_len: usize,
    pub(super) use_compression: bool,
    pub(super) max_lsn: Lsn,
    pub(super) cur_lsn: Lsn,
    pub(super) trailer: Option<Lsn>,
}

impl<'a> Iterator for Iter<'a> {
    type Item = (Lsn, LogID, Vec<u8>);

    fn next(&mut self) -> Option<Self::Item> {
        // If segment is None, get next on segment_iter, panic
        // if we can't read something we expect to be able to,
        // return None if there are no more remaining segments.
        loop {
            // println!("cur is {}", self.cur_lsn);
            // println!("max is {}", self.max_lsn);
            // println!("base is {:?}", self.segment_base);
            // println!("segment_len {}", self.segment_len);
            let at_end = !valid_entry_offset(self.cur_lsn, self.segment_len);
            if self.trailer.is_none() && at_end {
                // We've read to the end of a torn
                // segment and should stop now.
                return None;
            } else if self.segment_base.is_none() || at_end {
                if let Some((next_lsn, next_lid)) = self.segment_iter.next() {
                    self.read_segment(next_lsn, next_lid).unwrap();
                } else {
                    return None;
                }
            }

            if self.cur_lsn > self.max_lsn {
                // all done
                return None;
            }

            let lid = self.segment_base.unwrap() +
                (self.cur_lsn % self.segment_len as LogID);

            if self.max_lsn <= lid {
                // we've hit the end of the log.
                return None;
            }

            // println!("lid is {}", lid);
            let cached_f = self.config.cached_file();
            let mut f = cached_f.borrow_mut();
            match f.read_message(lid, self.segment_len, self.use_compression) {
                Ok(LogRead::Flush(lsn, buf, on_disk_len)) => {
                    // println!("{}", on_disk_len);
                    // println!("cur_lsn before: {}", self.cur_lsn);
                    self.cur_lsn += (MSG_HEADER_LEN + on_disk_len) as LogID;
                    // println!("cur_lsn after: {}", self.cur_lsn);
                    return Some((lsn, lid, buf));
                }
                Ok(LogRead::Zeroed(on_disk_len)) => {
                    self.cur_lsn += on_disk_len as LogID;
                }
                _ => {
                    if self.trailer.is_none() {
                        // This segment was torn, nothing left to read.
                        return None;
                    }
                    self.segment_base.take();
                    self.trailer.take();
                    continue;
                }
            }
        }
    }
}

impl<'a> Iter<'a> {
    /// read a segment of log messages. Only call after
    /// pausing segment rewriting on the segment accountant!
    fn read_segment(&mut self, lsn: Lsn, offset: LogID) -> std::io::Result<()> {
        // println!("reading segment {}", offset);

        let cached_f = self.config.cached_file();
        let mut f = cached_f.borrow_mut();
        let segment_header = f.read_segment_header(offset)?;
        // println!("read sh {:?} at {}", sh, offset);
        // TODO turn those asserts into returned Errors? Torn pages or invariant?
        assert_eq!(offset % self.segment_len as Lsn, 0);
        assert_eq!(segment_header.lsn % self.segment_len as Lsn, 0);

        // FIXME left: `83886080` right: 0, merge -> write_snapshot -> Iter::next
        assert_eq!(segment_header.lsn, lsn);

        assert!(segment_header.lsn + self.segment_len as LogID >= self.cur_lsn);

        let trailer_offset = offset + self.segment_len as LogID -
            SEG_TRAILER_LEN as LogID;
        let trailer_lsn = segment_header.lsn + self.segment_len as Lsn -
            SEG_TRAILER_LEN as Lsn;

        trace!("trying to read trailer from {}", trailer_offset);
        let segment_trailer = f.read_segment_trailer(trailer_offset);

        trace!("read segment header {:?}", segment_header);
        trace!("read segment trailer {:?}", segment_trailer);

        let trailer_lsn = segment_trailer.ok().and_then(|st| if st.ok &&
            st.lsn == trailer_lsn
        {
            Some(st.lsn)
        } else {
            None
        });

        self.trailer = trailer_lsn;
        self.cur_lsn = segment_header.lsn + SEG_HEADER_LEN as Lsn;
        self.segment_base = Some(offset);

        Ok(())
    }
}

fn valid_entry_offset(lid: LogID, segment_len: usize) -> bool {
    let seg_start = lid / segment_len as LogID * segment_len as LogID;

    let max_lid = seg_start + segment_len as LogID -
        SEG_TRAILER_LEN as LogID - MSG_HEADER_LEN as LogID;

    let min_lid = seg_start + SEG_HEADER_LEN as LogID;

    lid >= min_lid && lid <= max_lid
}
<commit_msg>small refactor of log iterator<commit_after>use super::*;

pub struct Iter<'a> {
    pub(super) config: &'a Config,
    pub(super) segment_iter: Box<Iterator<Item = (Lsn, LogID)>>,
    pub(super) segment_base: Option<LogID>,
    pub(super) segment_len: usize,
    pub(super) use_compression: bool,
    pub(super) max_lsn: Lsn,
    pub(super) cur_lsn: Lsn,
    pub(super) trailer: Option<Lsn>,
}

impl<'a> Iterator for Iter<'a> {
    type Item = (Lsn, LogID, Vec<u8>);

    fn next(&mut self) -> Option<Self::Item> {
        // If segment is None, get next on segment_iter, panic
        // if we can't read something we expect to be able to,
        // return None if there are no more remaining segments.
        loop {
            let at_end = !valid_entry_offset(self.cur_lsn, self.segment_len);
            if self.trailer.is_none() && at_end {
                // We've read to the end of a torn
                // segment and should stop now.
                return None;
            } else if self.segment_base.is_none() || at_end {
                if let Some((next_lsn, next_lid)) = self.segment_iter.next() {
                    self.read_segment(next_lsn, next_lid).unwrap();
                } else {
                    return None;
                }
            }

            if self.cur_lsn > self.max_lsn {
                // all done
                return None;
            }

            let lid = self.segment_base.unwrap() +
                (self.cur_lsn % self.segment_len as LogID);

            if self.max_lsn <= lid {
                // we've hit the end of the log.
                return None;
            }

            let cached_f = self.config.cached_file();
            let mut f = cached_f.borrow_mut();
            match f.read_message(lid, self.segment_len, self.use_compression) {
                Ok(LogRead::Flush(lsn, buf, on_disk_len)) => {
                    self.cur_lsn += (MSG_HEADER_LEN + on_disk_len) as LogID;
                    return Some((lsn, lid, buf));
                }
                Ok(LogRead::Zeroed(on_disk_len)) => {
                    self.cur_lsn += on_disk_len as LogID;
                }
                _ => {
                    if self.trailer.is_none() {
                        // This segment was torn, nothing left to read.
                        return None;
                    }
                    self.segment_base.take();
                    self.trailer.take();
                    continue;
                }
            }
        }
    }
}

impl<'a> Iter<'a> {
    /// read a segment of log messages. Only call after
    /// pausing segment rewriting on the segment accountant!
    fn read_segment(&mut self, lsn: Lsn, offset: LogID) -> std::io::Result<()> {
        let cached_f = self.config.cached_file();
        let mut f = cached_f.borrow_mut();
        let segment_header = f.read_segment_header(offset)?;
        // TODO turn those asserts into returned Errors? Torn pages or invariant?
        assert_eq!(offset % self.segment_len as Lsn, 0);
        assert_eq!(segment_header.lsn % self.segment_len as Lsn, 0);

        // FIXME left: `83886080` right: 0, merge -> write_snapshot -> Iter::next
        assert_eq!(segment_header.lsn, lsn);

        assert!(segment_header.lsn + self.segment_len as LogID >= self.cur_lsn);

        let trailer_offset = offset + self.segment_len as LogID -
            SEG_TRAILER_LEN as LogID;
        let trailer_lsn = segment_header.lsn + self.segment_len as Lsn -
            SEG_TRAILER_LEN as Lsn;

        trace!("trying to read trailer from {}", trailer_offset);
        let segment_trailer = f.read_segment_trailer(trailer_offset);

        trace!("read segment header {:?}", segment_header);
        trace!("read segment trailer {:?}", segment_trailer);

        let trailer_lsn = segment_trailer.ok().and_then(|st| if st.ok &&
            st.lsn == trailer_lsn
        {
            Some(st.lsn)
        } else {
            None
        });

        self.trailer = trailer_lsn;
        self.cur_lsn = segment_header.lsn + SEG_HEADER_LEN as Lsn;
        self.segment_base = Some(offset);

        Ok(())
    }
}

fn valid_entry_offset(lid: LogID, segment_len: usize) -> bool {
    let seg_start = lid / segment_len as LogID * segment_len as LogID;

    let max_lid = seg_start + segment_len as LogID -
        SEG_TRAILER_LEN as LogID - MSG_HEADER_LEN as LogID;

    let min_lid = seg_start + SEG_HEADER_LEN as LogID;

    lid >= min_lid && lid <= max_lid
}
<|endoftext|>


<commit_before>/// A genetic algorithm in Rust
/// Copyright (C) 2015  Andrew Schwartzmeyer

use Parameters;
use Problem;
use individual::Individual;
use rand::{Rng, thread_rng};
use std::thread;
use time::precise_time_s;

pub struct Results {
    pub problem: Problem,
    pub parameters: Parameters,
    pub individual: Individual,
    pub iterations: usize,
    pub duration: f64
}

/// A genetic algorithm that searches for convergence to the given
/// tolerance for the problem across the n-dimensional hypercube,
/// using a population of individuals, up to a maximum iterations
/// number of generations.
pub fn search(problem: Problem, parameters: Parameters) -> Results {
    // get thread local random number generator
    let mut rng = thread_rng();

    // initialize population of individuals
    let mut population: Vec<_> = (0..parameters.population).map(|_| {
        Individual::new(problem, parameters.dimension, &mut rng)
    }).collect();

    // start timing the search
    let start_time = precise_time_s();

    // search iterations number of generations
    for i in 0..parameters.iterations {
        // select, mutate, and recombine individuals for next generation
        let mut offspring: Vec<Individual> = Vec::with_capacity(population.len());
        for _ in 0..population.len()/2 {
            let (mut x, mut y) = (select(&population, &mut rng),
                                  select(&population, &mut rng));
            x.mutate(&mut rng);
            y.mutate(&mut rng);
            Individual::combine(&mut x, &mut y, &mut rng);
            offspring.push(x);
            offspring.push(y);
        }
        assert!(offspring.len() == population.len());

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

        // replace population with next generation
        population = offspring;

        // examine best individual for convergence
        if let Some(x) = population.iter().min() {
            if x.fitness < parameters.tolerance {
                return Results {
                    problem: problem, parameters: parameters,
                    individual: x.clone(), iterations: i,
                    duration: precise_time_s() - start_time
                };
            }
            // print verbose information
            if parameters.verbosity > 0 && i % 10 == 0 {
                let fitness = x.fitness;
                let solution = x.solution.clone();
                thread::spawn(move || {
                    if parameters.verbosity >= 1 {
                        println!("{}th fitness {}", i, fitness);
                    }
                    if parameters.verbosity >= 2 {
                        println!{"{:?}", solution};
                    }
                });
            }
        }
    }
    if let Some(x) = population.iter().min() {
        Results { problem: problem, parameters: parameters,
                  individual: x.clone(), iterations: parameters.iterations,
                  duration: precise_time_s() - start_time }
    } else {
        unimplemented!();
    }
}

/// Tournament selection from 4 random individuals
fn select<R: Rng>(population: &[Individual], rng: &mut R) -> Individual {
    if let Some(selected) = (0..4).map(|_| rng.choose(population)).min() {
        selected.unwrap().clone()
    } else {
        unimplemented!();
    }
}
<commit_msg>Shorten parameters to params in algorithm<commit_after>/// A genetic algorithm in Rust
/// Copyright (C) 2015  Andrew Schwartzmeyer

use Parameters;
use Problem;
use individual::Individual;
use rand::{Rng, thread_rng};
use std::thread;
use time::precise_time_s;

pub struct Results {
    pub problem: Problem,
    pub parameters: Parameters,
    pub individual: Individual,
    pub iterations: usize,
    pub duration: f64
}

/// A genetic algorithm that searches for convergence to the given
/// tolerance for the problem across the n-dimensional hypercube,
/// using a population of individuals, up to a maximum iterations
/// number of generations.
pub fn search(problem: Problem, params: Parameters) -> Results {
    // get thread local random number generator
    let mut rng = thread_rng();

    // initialize population of individuals
    let mut population: Vec<_> = (0..params.population).map(|_| {
        Individual::new(problem, params.dimension, &mut rng)
    }).collect();

    // start timing the search
    let start_time = precise_time_s();

    // search iterations number of generations
    for i in 0..params.iterations {
        // select, mutate, and crossover individuals for next generation
        let mut offspring: Vec<Individual> = Vec::with_capacity(population.len());
        for _ in 0..population.len()/2 {
            let (mut x, mut y) = (select(&population, &mut rng),
                                  select(&population, &mut rng));
            x.mutate(&mut rng);
            y.mutate(&mut rng);
            Individual::combine(&mut x, &mut y, &mut rng);
            offspring.push(x);
            offspring.push(y);
        }
        assert!(offspring.len() == population.len());

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

        // replace population with next generation
        population = offspring;

        // examine best individual for convergence
        if let Some(x) = population.iter().min() {
            if x.fitness < params.tolerance {
                return Results {
                    problem: problem, parameters: params,
                    individual: x.clone(), iterations: i,
                    duration: precise_time_s() - start_time
                };
            }
            // print verbose information
            if params.verbosity > 0 && i % 10 == 0 {
                let fitness = x.fitness;
                let solution = x.solution.clone();
                thread::spawn(move || {
                    if params.verbosity >= 1 {
                        println!("{}th fitness {}", i, fitness);
                    }
                    if params.verbosity >= 2 {
                        println!{"{:?}", solution};
                    }
                });
            }
        }
    }
    if let Some(x) = population.iter().min() {
        Results { problem: problem, parameters: params,
                  individual: x.clone(), iterations: params.iterations,
                  duration: precise_time_s() - start_time }
    } else {
        unimplemented!();
    }
}

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


<commit_before>#![crate_name = "ansi_term"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![desc = "A rust library for ANSI terminal colours and styles (bold, underline)"]

//! This is a library for controlling colours and formatting, such as
//! red bold text or blue underlined text, on ANSI terminals.
//!
//! ```rust
//! extern crate ansi_term;
//! use ansi_term::Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, Fixed};
//! use ansi_term::Paint;
//! ```
//!
//! Simple Colours
//! --------------
//!
//! You can format strings by calling the `paint` method on a Colour
//! or a Style object, passing in the string you want to format. For
//! example, to get some red text, call the `paint` method on `Red`:
//!
//! ```rust
//! Red.paint("Red!")
//! ```
//!
//! You can also call the `paint` method on a Style. To turn a Colour
//! into a Style, use the `normal` method, though this does the exact
//! some thing as above:
//!
//! ```rust
//! Red.normal().paint("Still red!")
//! ```
//!
//! Bold, Underline, and Background
//! -------------------------------
//!
//! To do anything more complex than just foreground colours, you need
//! to use Style objects. Calling the `bold` or `underline` method on
//! a Colour returns a Style that has the appropriate property set on
//! it:
//!
//! ```rust
//! Blue.bold().paint("Blue bold!")
//! Yellow.underline().paint("Yellow underline!")
//! ```
//!
//! Again, you can call these methods on Styles as well as just Colours:
//!
//! ```rust
//! Blue.normal().bold().paint("Still blue bold!")
//! ```
//!
//! Finally, you can set the background colour of a Style by using the
//! `on` method:
//!
//! ```rust
//! Blue.on(Yellow).paint("Blue on yellow!")
//! ```
//!
//! Extended Colours
//! ----------------
//!
//! You can access the extended range of 256 colours by using the
//! Fixed constructor, which takes an argument of the colour number to
//! use. This can be used wherever you would use a Colour:
//!
//! ```rust
//! Fixed(134).paint("A sort of light purple.")
//! ```
//!
//! This even works for background colours:
//!
//! ```rust
//! Fixed(221).on(Fixed(124)).paint("Mustard in the ketchup.")
//! ```
//!
//! No Formatting
//! -------------
//!
//! Finally, for the sake of completeness, the Plain style provides
//! neither colours nor formatting.
//!
//! ```rust
//! Plain.paint("No colours here.")
//! ```

#![feature(phase)] extern crate regex;
#[phase(plugin)] extern crate regex_macros;
use Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed};
use Style::{Plain, Foreground, Styled};

pub trait Paint {
    /// Paints the given text with this colour.
    fn paint(&self, input: &str) -> String;

    /// Returns a Style with the underline property set.
    fn underline(&self) -> Style;

    /// Return a Style with the bold property set.
    fn bold(&self) -> Style;

    /// Return a Style with the background colour set.
    fn on(&self, background: Colour) -> Style;
}

/// A colour is one specific type of ANSI escape code, and can refer
/// to either the foreground or background colour.
pub enum Colour {
    Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed(u8),
}

// These are the standard numeric sequences.
// See http://invisible-island.net/xterm/ctlseqs/ctlseqs.html

impl Colour {
    fn foreground_code(&self) -> String {
        match *self {
            Black => "30".to_string(),
            Red => "31".to_string(),
            Green => "32".to_string(),
            Yellow => "33".to_string(),
            Blue => "34".to_string(),
            Purple => "35".to_string(),
            Cyan => "36".to_string(),
            White => "37".to_string(),
            Fixed(num) => format!("38;5;{}", num),
        }
    }

    fn background_code(&self) -> String {
        match *self {
            Black => "40".to_string(),
            Red => "41".to_string(),
            Green => "42".to_string(),
            Yellow => "43".to_string(),
            Blue => "44".to_string(),
            Purple => "45".to_string(),
            Cyan => "46".to_string(),
            White => "47".to_string(),
            Fixed(num) => format!("48;5;{}", num),
        }
    }

    /// Return a Style with the foreground colour set to this colour.
    pub fn normal(&self) -> Style {
        Styled { foreground: *self, background: None, bold: false, underline: false }
    }
}

/// The Paint trait represents a style or colour that can be applied
/// to a piece of text.
impl Paint for Colour {
    /// This is a short-cut so you don't have to use Blue.normal() just
    /// to get blue text.
    fn paint(&self, input: &str) -> String {
        let re = format!("\x1B[{}m{}\x1B[0m", self.foreground_code(), input);
        return re.to_string();
    }

    fn underline(&self) -> Style {
        Styled { foreground: *self, background: None, bold: false, underline: true }
    }

    fn bold(&self) -> Style {
        Styled { foreground: *self, background: None, bold: true, underline: false }
    }

    fn on(&self, background: Colour) -> Style {
        Styled { foreground: *self, background: Some(background), bold: false, underline: false }
    }
}

/// A style is a collection of properties that can format a string
/// using ANSI escape codes.
pub enum Style {

    /// The Plain style provides no formatting.
    Plain,

    /// The Foreground style provides just a foreground colour.
    Foreground(Colour),

    ///
    /// The Styled style is a catch-all for anything more complicated
    /// than that. It's technically possible for there to be other
    /// cases, such as "bold foreground", but probably isn't worth it.
    Styled { foreground: Colour, background: Option<Colour>, bold: bool, underline: bool, },
}

impl Paint for Style {
    fn paint(&self, input: &str) -> String {
        match *self {
            Plain => input.to_string(),
            Foreground(c) => c.paint(input),
            Styled(foreground, background, bold, underline) => {
                let bg = match background {
                    Some(c) => format!("{};", c.background_code()),
                    None => "".to_string()
                };
                let bo = if bold { "1;" } else { "" };
                let un = if underline { "4;" } else { "" };
                let painted = format!("\x1B[{}{}{}{}m{}\x1B[0m", bo, un, bg, foreground.foreground_code(), input.to_string());
                return painted.to_string();
            }
        }
    }

    fn bold(&self) -> Style {
        match *self {
            Plain =>                Styled { foreground: White,  background: None,  bold: true,  underline: false },
            Foreground(c) =>        Styled { foreground: c,      background: None,  bold: true,  underline: false },
            Styled(fg, bg, _, u) => Styled { foreground: fg,     background: bg,    bold: true,  underline: u },
        }
    }

    fn underline(&self) -> Style {
        match *self {
            Plain =>                Styled { foreground: White,  background: None,  bold: false,  underline: true },
            Foreground(c) =>        Styled { foreground: c,      background: None,  bold: false,  underline: true },
            Styled(fg, bg, b, _) => Styled { foreground: fg,     background: bg,    bold: b,      underline: true },
        }
    }

    fn on(&self, background: Colour) -> Style {
        match *self {
            Plain =>               Styled { foreground: White,  background: Some(background),  bold: false,  underline: false },
            Foreground(c) =>       Styled { foreground: c,      background: Some(background),  bold: false,  underline: false },
            Styled(fg, _, b, u) => Styled { foreground: fg,     background: Some(background),  bold: b,      underline: u },
        }
    }
}

/// Return a String with all ANSI escape codes removed. For example:
///
/// ```rust
/// strip_formatting(Blue.paint("hello!")) == "hello!"
/// ```
pub fn strip_formatting(input: String) -> String {
    // What's blue and smells like red paint? Blue paint.
    let re = regex!("\x1B\\[.+?m");
    re.replace_all(input.as_slice(), "").to_string()
}

#[cfg(test)]
mod tests {
    use super::strip_formatting;
    use super::Style::Plain;
    use super::Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed};
    use super::Paint;

    #[test]
    fn test_plain() {
        let hi = Plain.paint("text/plain");
        assert!(hi == "text/plain".to_string());
    }

    #[test]
    fn test_red() {
        let hi = Red.paint("hi");
        assert!(hi == "\x1B[31mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_black() {
        let hi = Black.normal().paint("hi");
        assert!(hi == "\x1B[30mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_yellow_bold() {
        let hi = Yellow.bold().paint("hi");
        assert!(hi == "\x1B[1;33mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_yellow_bold_2() {
        let hi = Yellow.normal().bold().paint("hi");
        assert!(hi == "\x1B[1;33mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_blue_underline() {
        let hi = Blue.underline().paint("hi");
        assert!(hi == "\x1B[4;34mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_green_bold_underline() {
        let hi = Green.bold().underline().paint("hi");
        assert!(hi == "\x1B[1;4;32mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_green_bold_underline_2() {
        let hi = Green.underline().bold().paint("hi");
        assert!(hi == "\x1B[1;4;32mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_purple_on_white() {
        let hi = Purple.on(White).paint("hi");
        assert!(hi == "\x1B[47;35mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_purple_on_white_2() {
        let hi = Purple.normal().on(White).paint("hi");
        assert!(hi == "\x1B[47;35mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_cyan_bold_on_white() {
        let hi = Cyan.bold().on(White).paint("hi");
        assert!(hi == "\x1B[1;47;36mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_cyan_underline_on_white() {
        let hi = Cyan.underline().on(White).paint("hi");
        assert!(hi == "\x1B[4;47;36mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_cyan_bold_underline_on_white() {
        let hi = Cyan.bold().underline().on(White).paint("hi");
        assert!(hi == "\x1B[1;4;47;36mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_cyan_underline_bold_on_white() {
        let hi = Cyan.underline().bold().on(White).paint("hi");
        assert!(hi == "\x1B[1;4;47;36mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_fixed() {
        let hi = Fixed(100).paint("hi");
        assert!(hi == "\x1B[38;5;100mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_fixed_on_purple() {
        let hi = Fixed(100).on(Purple).paint("hi");
        assert!(hi == "\x1B[45;38;5;100mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_fixed_on_fixed() {
        let hi = Fixed(100).on(Fixed(200)).paint("hi");
        assert!(hi == "\x1B[48;5;200;38;5;100mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_strip_formatting() {
        let hi = strip_formatting(Blue.paint("hi"));
        assert!(hi == "hi".to_string());
    }

    #[test]
    fn test_strip_formatting_2() {
        let hi = strip_formatting(Blue.on(Fixed(230)).bold().paint("hi"));
        assert!(hi == "hi".to_string());
    }
}
<commit_msg>Replace tuples with structs in patterns<commit_after>#![crate_name = "ansi_term"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![desc = "A rust library for ANSI terminal colours and styles (bold, underline)"]

//! This is a library for controlling colours and formatting, such as
//! red bold text or blue underlined text, on ANSI terminals.
//!
//! ```rust
//! extern crate ansi_term;
//! use ansi_term::Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, Fixed};
//! use ansi_term::Paint;
//! ```
//!
//! Simple Colours
//! --------------
//!
//! You can format strings by calling the `paint` method on a Colour
//! or a Style object, passing in the string you want to format. For
//! example, to get some red text, call the `paint` method on `Red`:
//!
//! ```rust
//! Red.paint("Red!")
//! ```
//!
//! You can also call the `paint` method on a Style. To turn a Colour
//! into a Style, use the `normal` method, though this does the exact
//! some thing as above:
//!
//! ```rust
//! Red.normal().paint("Still red!")
//! ```
//!
//! Bold, Underline, and Background
//! -------------------------------
//!
//! To do anything more complex than just foreground colours, you need
//! to use Style objects. Calling the `bold` or `underline` method on
//! a Colour returns a Style that has the appropriate property set on
//! it:
//!
//! ```rust
//! Blue.bold().paint("Blue bold!")
//! Yellow.underline().paint("Yellow underline!")
//! ```
//!
//! Again, you can call these methods on Styles as well as just Colours:
//!
//! ```rust
//! Blue.normal().bold().paint("Still blue bold!")
//! ```
//!
//! Finally, you can set the background colour of a Style by using the
//! `on` method:
//!
//! ```rust
//! Blue.on(Yellow).paint("Blue on yellow!")
//! ```
//!
//! Extended Colours
//! ----------------
//!
//! You can access the extended range of 256 colours by using the
//! Fixed constructor, which takes an argument of the colour number to
//! use. This can be used wherever you would use a Colour:
//!
//! ```rust
//! Fixed(134).paint("A sort of light purple.")
//! ```
//!
//! This even works for background colours:
//!
//! ```rust
//! Fixed(221).on(Fixed(124)).paint("Mustard in the ketchup.")
//! ```
//!
//! No Formatting
//! -------------
//!
//! Finally, for the sake of completeness, the Plain style provides
//! neither colours nor formatting.
//!
//! ```rust
//! Plain.paint("No colours here.")
//! ```

#![feature(phase)] extern crate regex;
#[phase(plugin)] extern crate regex_macros;
use Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed};
use Style::{Plain, Foreground, Styled};

pub trait Paint {
    /// Paints the given text with this colour.
    fn paint(&self, input: &str) -> String;

    /// Returns a Style with the underline property set.
    fn underline(&self) -> Style;

    /// Return a Style with the bold property set.
    fn bold(&self) -> Style;

    /// Return a Style with the background colour set.
    fn on(&self, background: Colour) -> Style;
}

/// A colour is one specific type of ANSI escape code, and can refer
/// to either the foreground or background colour.
pub enum Colour {
    Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed(u8),
}

// These are the standard numeric sequences.
// See http://invisible-island.net/xterm/ctlseqs/ctlseqs.html

impl Colour {
    fn foreground_code(&self) -> String {
        match *self {
            Black => "30".to_string(),
            Red => "31".to_string(),
            Green => "32".to_string(),
            Yellow => "33".to_string(),
            Blue => "34".to_string(),
            Purple => "35".to_string(),
            Cyan => "36".to_string(),
            White => "37".to_string(),
            Fixed(num) => format!("38;5;{}", num),
        }
    }

    fn background_code(&self) -> String {
        match *self {
            Black => "40".to_string(),
            Red => "41".to_string(),
            Green => "42".to_string(),
            Yellow => "43".to_string(),
            Blue => "44".to_string(),
            Purple => "45".to_string(),
            Cyan => "46".to_string(),
            White => "47".to_string(),
            Fixed(num) => format!("48;5;{}", num),
        }
    }

    /// Return a Style with the foreground colour set to this colour.
    pub fn normal(&self) -> Style {
        Styled { foreground: *self, background: None, bold: false, underline: false }
    }
}

/// The Paint trait represents a style or colour that can be applied
/// to a piece of text.
impl Paint for Colour {
    /// This is a short-cut so you don't have to use Blue.normal() just
    /// to get blue text.
    fn paint(&self, input: &str) -> String {
        let re = format!("\x1B[{}m{}\x1B[0m", self.foreground_code(), input);
        return re.to_string();
    }

    fn underline(&self) -> Style {
        Styled { foreground: *self, background: None, bold: false, underline: true }
    }

    fn bold(&self) -> Style {
        Styled { foreground: *self, background: None, bold: true, underline: false }
    }

    fn on(&self, background: Colour) -> Style {
        Styled { foreground: *self, background: Some(background), bold: false, underline: false }
    }
}

/// A style is a collection of properties that can format a string
/// using ANSI escape codes.
pub enum Style {

    /// The Plain style provides no formatting.
    Plain,

    /// The Foreground style provides just a foreground colour.
    Foreground(Colour),

    ///
    /// The Styled style is a catch-all for anything more complicated
    /// than that. It's technically possible for there to be other
    /// cases, such as "bold foreground", but probably isn't worth it.
    Styled { foreground: Colour, background: Option<Colour>, bold: bool, underline: bool, },
}

impl Paint for Style {
    fn paint(&self, input: &str) -> String {
        match *self {
            Plain => input.to_string(),
            Foreground(c) => c.paint(input),
            Styled { foreground, background, bold, underline } => {
                let bg = match background {
                    Some(c) => format!("{};", c.background_code()),
                    None => "".to_string()
                };
                let bo = if bold { "1;" } else { "" };
                let un = if underline { "4;" } else { "" };
                let painted = format!("\x1B[{}{}{}{}m{}\x1B[0m", bo, un, bg, foreground.foreground_code(), input.to_string());
                return painted.to_string();
            }
        }
    }

    fn bold(&self) -> Style {
        match *self {
            Plain => Styled { foreground: White, background: None, bold: true, underline: false },
            Foreground(c) => Styled { foreground: c, background: None, bold: true, underline: false },
            Styled { foreground, background, bold: _, underline } => Styled { foreground: foreground, background: background, bold: true, underline: underline },
        }
    }

    fn underline(&self) -> Style {
        match *self {
            Plain => Styled { foreground: White, background: None, bold: false, underline: true },
            Foreground(c) => Styled { foreground: c, background: None, bold: false, underline: true },
            Styled { foreground, background, bold, underline: _ } => Styled { foreground: foreground, background: background, bold: bold, underline: true },
        }
    }

    fn on(&self, background: Colour) -> Style {
        match *self {
            Plain => Styled { foreground: White,background: Some(background), bold: false, underline: false },
            Foreground(c) => Styled { foreground: c, background: Some(background), bold: false, underline: false },
            Styled { foreground, background: _, bold, underline } => Styled { foreground: foreground, background: Some(background), bold: bold, underline: underline },
        }
    }
}

/// Return a String with all ANSI escape codes removed. For example:
///
/// ```rust
/// strip_formatting(Blue.paint("hello!")) == "hello!"
/// ```
pub fn strip_formatting(input: String) -> String {
    // What's blue and smells like red paint? Blue paint.
    let re = regex!("\x1B\\[.+?m");
    re.replace_all(input.as_slice(), "").to_string()
}

#[cfg(test)]
mod tests {
    use super::strip_formatting;
    use super::Style::Plain;
    use super::Colour::{Black, Red, Green, Yellow, Blue, Purple, Cyan, White, Fixed};
    use super::Paint;

    #[test]
    fn test_plain() {
        let hi = Plain.paint("text/plain");
        assert!(hi == "text/plain".to_string());
    }

    #[test]
    fn test_red() {
        let hi = Red.paint("hi");
        assert!(hi == "\x1B[31mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_black() {
        let hi = Black.normal().paint("hi");
        assert!(hi == "\x1B[30mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_yellow_bold() {
        let hi = Yellow.bold().paint("hi");
        assert!(hi == "\x1B[1;33mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_yellow_bold_2() {
        let hi = Yellow.normal().bold().paint("hi");
        assert!(hi == "\x1B[1;33mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_blue_underline() {
        let hi = Blue.underline().paint("hi");
        assert!(hi == "\x1B[4;34mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_green_bold_underline() {
        let hi = Green.bold().underline().paint("hi");
        assert!(hi == "\x1B[1;4;32mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_green_bold_underline_2() {
        let hi = Green.underline().bold().paint("hi");
        assert!(hi == "\x1B[1;4;32mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_purple_on_white() {
        let hi = Purple.on(White).paint("hi");
        assert!(hi == "\x1B[47;35mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_purple_on_white_2() {
        let hi = Purple.normal().on(White).paint("hi");
        assert!(hi == "\x1B[47;35mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_cyan_bold_on_white() {
        let hi = Cyan.bold().on(White).paint("hi");
        assert!(hi == "\x1B[1;47;36mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_cyan_underline_on_white() {
        let hi = Cyan.underline().on(White).paint("hi");
        assert!(hi == "\x1B[4;47;36mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_cyan_bold_underline_on_white() {
        let hi = Cyan.bold().underline().on(White).paint("hi");
        assert!(hi == "\x1B[1;4;47;36mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_cyan_underline_bold_on_white() {
        let hi = Cyan.underline().bold().on(White).paint("hi");
        assert!(hi == "\x1B[1;4;47;36mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_fixed() {
        let hi = Fixed(100).paint("hi");
        assert!(hi == "\x1B[38;5;100mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_fixed_on_purple() {
        let hi = Fixed(100).on(Purple).paint("hi");
        assert!(hi == "\x1B[45;38;5;100mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_fixed_on_fixed() {
        let hi = Fixed(100).on(Fixed(200)).paint("hi");
        assert!(hi == "\x1B[48;5;200;38;5;100mhi\x1B[0m".to_string());
    }

    #[test]
    fn test_strip_formatting() {
        let hi = strip_formatting(Blue.paint("hi"));
        assert!(hi == "hi".to_string());
    }

    #[test]
    fn test_strip_formatting_2() {
        let hi = strip_formatting(Blue.on(Fixed(230)).bold().paint("hi"));
        assert!(hi == "hi".to_string());
    }
}
<|endoftext|>


<commit_before>use self::StartsWithStateInternal::*;

/// Automaton describes types that behave as a finite automaton.
///
/// All implementors of this trait are represented by *byte based* automata.
/// Stated differently, all transitions in the automata correspond to a single
/// byte in the input.
///
/// This implementation choice is important for a couple reasons:
///
/// 1. The set of possible transitions in each node is small, which may make
///    efficient memory usage easier.
/// 2. The finite state transducers in this crate are all byte based, so any
///    automata used on them must also be byte based.
///
/// In practice, this does present somewhat of a problem, for example, if
/// you're storing UTF-8 encoded strings in a finite state transducer. Consider
/// using a `Levenshtein` automaton, which accepts a query string and an edit
/// distance. The edit distance should apply to some notion of *character*,
/// which could be represented by at least 1-4 bytes in a UTF-8 encoding (for
/// some definition of "character"). Therefore, the automaton must have UTF-8
/// decoding built into it. This can be tricky to implement, so you may find
/// the [`utf8-ranges`](https://crates.io/crates/utf8-ranges) crate useful.
pub trait Automaton {
    /// The type of the state used in the automaton.
    type State;

    /// Returns a single start state for this automaton.
    ///
    /// This method should always return the same value for each
    /// implementation.
    fn start(&self) -> Self::State;

    /// Returns true if and only if `state` is a match state.
    fn is_match(&self, state: &Self::State) -> bool;

    /// Returns true if and only if `state` can lead to a match in zero or more
    /// steps.
    ///
    /// If this returns `false`, then no sequence of inputs from this state
    /// should ever produce a match. If this does not follow, then those match
    /// states may never be reached. In other words, behavior may be incorrect.
    ///
    /// If this returns `true` even when no match is possible, then behavior
    /// will be correct, but callers may be forced to do additional work.
    fn can_match(&self, _state: &Self::State) -> bool {
        true
    }

    /// Returns true if and only if `state` matches and must match no matter what
    /// steps are taken.
    ///
    /// If this returns `true`, then every sequence of inputs from this state
    /// produces a match. If this does not follow, then those match states may
    /// never be reached. In other words, behavior may be incorrect.
    ///
    /// If this returns `false` even when every sequence of inputs will lead to
    /// a match, then behavior will be correct, but callers may be forced to do
    /// additional work.
    fn will_always_match(&self, _state: &Self::State) -> bool {
        false
    }

    /// Return the next state given `state` and an input.
    fn accept(&self, state: &Self::State, byte: u8) -> Self::State;

    /// Returns an automaton that matches the strings that start with something
    /// this automaton matches.
    fn starts_with(self) -> StartsWith<Self> where Self: Sized {
        StartsWith(self)
    }

    /// Returns an automaton that matches the strings matched by either this or
    /// the other automaton.
    fn union<Rhs: Automaton>(
        self,
        rhs: Rhs,
    ) -> Union<Self, Rhs> where Self: Sized {
        Union(self, rhs)
    }

    /// Returns an automaton that matches the strings matched by both this and
    /// the other automaton.
    fn intersection<Rhs: Automaton>(
        self,
        rhs: Rhs,
    ) -> Intersection<Self, Rhs> where Self: Sized {
        Intersection(self, rhs)
    }

    /// Returns an automaton that matches the strings not matched by this
    /// automaton.
    fn complement(self) -> Complement<Self> where Self: Sized {
        Complement(self)
    }
}

impl<'a, T: Automaton> Automaton for &'a T {
    type State = T::State;

    fn start(&self) -> Self::State {
        (*self).start()
    }

    fn is_match(&self, state: &Self::State) -> bool {
        (*self).is_match(state)
    }

    fn can_match(&self, state: &Self::State) -> bool {
        (*self).can_match(state)
    }

    fn will_always_match(&self, state: &Self::State) -> bool {
        (*self).will_always_match(state)
    }

    fn accept(&self, state: &Self::State, byte: u8) -> Self::State {
        (*self).accept(state, byte)
    }
}

/// An automaton that always matches.
///
/// This is useful in a generic context as a way to express that no automaton
/// should be used.
pub struct AlwaysMatch;

impl Automaton for AlwaysMatch {
    type State = ();

    fn start(&self) -> () { () }
    fn is_match(&self, _: &()) -> bool { true }
    fn can_match(&self, _: &()) -> bool { true }
    fn will_always_match(&self, _: &()) -> bool { true }
    fn accept(&self, _: &(), _: u8) -> () { () }
}

/// An automaton that matches and string that begins with something that the
/// wrapped automaton matches.
pub struct StartsWith<A>(A);

/// The `Automaton` state for `StartsWith<A>`.
pub struct StartsWithState<A: Automaton>(StartsWithStateInternal<A>);

enum StartsWithStateInternal<A: Automaton> {
    Done,
    Running(A::State)
}

impl<A: Automaton> Automaton for StartsWith<A> {
    type State = StartsWithState<A>;
    fn start(&self) -> StartsWithState<A> {
        StartsWithState({
            let inner = self.0.start();
            if self.0.is_match(&inner) {
                Done
            } else {
                Running(inner)
            }
        })
    }

    fn is_match(&self, state: &StartsWithState<A>) -> bool {
        match state.0 {
            Done => true,
            Running(_) => false
        }
    }

    fn can_match(&self, state: &StartsWithState<A>) -> bool {
        match state.0 {
            Done => true,
            Running(ref inner) => self.0.can_match(inner)
        }
    }

    fn will_always_match(&self, state: &StartsWithState<A>) -> bool {
        match state.0 {
            Done => true,
            Running(_) => false
        }
    }

    fn accept(&self, state: &StartsWithState<A>, byte: u8) -> StartsWithState<A> {
        StartsWithState(
            match state.0 {
                Done => Done,
                Running(ref inner) => {
                    let next_inner = self.0.accept(inner, byte);
                    if self.0.is_match(&next_inner) {
                        Done
                    } else {
                        Running(next_inner)
                    }
                }
            }
        )
    }
}

/// An automaton that matches when one of its component automata match.
pub struct Union<A, B>(A, B);

/// The `Automaton` state for `Union<A, B>`.
pub struct UnionState<A: Automaton, B: Automaton>(A::State, B::State);

impl<A: Automaton, B: Automaton> Automaton for Union<A, B> {
    type State = UnionState<A, B>;

    fn start(&self) -> UnionState<A, B> {
        UnionState(
            self.0.start(),
            self.1.start()
        )
    }

    fn is_match(&self, state: &UnionState<A, B>) -> bool {
        self.0.is_match(&state.0) || self.1.is_match(&state.1)
    }

    fn can_match(&self, state: &UnionState<A, B>) -> bool {
        self.0.can_match(&state.0) || self.1.can_match(&state.1)
    }

    fn will_always_match(&self, state: &UnionState<A, B>) -> bool {
        self.0.will_always_match(&state.0) || self.1.will_always_match(&state.1)
    }

    fn accept(&self, state: &UnionState<A, B>, byte: u8) -> UnionState<A, B> {
        UnionState(
            self.0.accept(&state.0, byte),
            self.1.accept(&state.1, byte)
        )
    }
}

/// An automaton that matches when both of its component automata match.
pub struct Intersection<A, B>(A, B);

/// The `Automaton` state for `Intersection<A, B>`.
pub struct IntersectionState<A: Automaton, B: Automaton>(A::State, B::State);

impl<A: Automaton, B: Automaton> Automaton for Intersection<A, B> {
    type State = IntersectionState<A, B>;

    fn start(&self) -> IntersectionState<A, B> {
        IntersectionState(
            self.0.start(),
            self.1.start()
        )
    }

    fn is_match(&self, state: &IntersectionState<A, B>) -> bool {
        self.0.is_match(&state.0) && self.1.is_match(&state.1)
    }

    fn can_match(&self, state: &IntersectionState<A, B>) -> bool {
        self.0.can_match(&state.0) && self.1.can_match(&state.1)
    }

    fn will_always_match(&self, state: &IntersectionState<A, B>) -> bool {
        self.0.will_always_match(&state.0) && self.1.will_always_match(&state.1)
    }

    fn accept(&self, state: &IntersectionState<A, B>, byte: u8) -> IntersectionState<A, B> {
        IntersectionState(
            self.0.accept(&state.0, byte),
            self.1.accept(&state.1, byte)
        )
    }
}

/// An automaton that matches exactly when the automaton it wraps does not.
pub struct Complement<A>(A);

/// The `Automaton` state for `Complement<A>`.
pub struct ComplementState<A: Automaton>(A::State);

impl<A: Automaton> Automaton for Complement<A> {
    type State = ComplementState<A>;

    fn start(&self) -> ComplementState<A> {
        ComplementState(self.0.start())
    }

    fn is_match(&self, state: &ComplementState<A>) -> bool {
        !self.0.is_match(&state.0)
    }

    fn can_match(&self, state: &ComplementState<A>) -> bool {
        !self.0.will_always_match(&state.0)
    }

    fn will_always_match(&self, state: &ComplementState<A>) -> bool {
        !self.0.can_match(&state.0)
    }

    fn accept(&self, state: &ComplementState<A>, byte: u8) -> ComplementState<A> {
        ComplementState(self.0.accept(&state.0, byte))
    }
}

#[cfg(test)]
mod test {
    use ::{IntoStreamer, Streamer, Levenshtein, Regex, Set, Automaton};
    use set::OpBuilder;

    static WORDS: &'static str = include_str!("../../data/words-10000");

    fn get_set() -> Set {
        Set::from_iter(WORDS.lines()).unwrap()
    }

    #[test]
    fn complement_small() {
        let keys = vec!["fa", "fo", "fob", "focus", "foo", "food", "foul"];
        let set = Set::from_iter(keys).unwrap();
        let lev = Levenshtein::new("foo", 1).unwrap();
        let stream = set.search(lev.complement()).into_stream();

        let keys = stream.into_strs().unwrap();
        assert_eq!(keys, vec!["fa", "focus", "foul"]);
    }

    #[test]
    fn startswith_small() {
        let keys = vec!["", "cooing", "fa", "fo", "fob", "focus", "foo", "food", "foul", "fritter", "frothing"];
        let set = Set::from_iter(keys).unwrap();
        let lev = Levenshtein::new("foo", 1).unwrap();
        let stream = set.search(lev.starts_with()).into_stream();

        let keys = stream.into_strs().unwrap();
        assert_eq!(keys, vec!["cooing", "fo", "fob", "focus", "foo", "food", "foul", "frothing"]);
    }

    #[test]
    fn intersection_small() {
        let keys = vec!["fa", "fo", "fob", "focus", "foo", "food", "foul"];
        let set = Set::from_iter(keys).unwrap();
        let lev = Levenshtein::new("foo", 1).unwrap();
        let reg = Regex::new("(..)*").unwrap();
        let stream = set.search(lev.intersection(reg)).into_stream();

        let keys = stream.into_strs().unwrap();
        assert_eq!(keys, vec!["fo", "food"]);
    }

    #[test]
    fn union_small() {
        let keys = vec!["fa", "fo", "fob", "focus", "foo", "food", "foul"];
        let set = Set::from_iter(keys).unwrap();
        let lev = Levenshtein::new("foo", 1).unwrap();
        let reg = Regex::new("(..)*").unwrap();
        let stream = set.search(lev.union(reg)).into_stream();

        let keys = stream.into_strs().unwrap();
        assert_eq!(keys, vec!["fa", "fo", "fob", "foo", "food", "foul"]);
    }

    #[test]
    fn intersection_large() {
        let set = get_set();
        let lev = Levenshtein::new("foo", 3).unwrap();
        let reg = Regex::new("(..)*").unwrap();
        let mut stream1 = set.search((&lev).intersection(&reg)).into_stream();
        let mut stream2 = OpBuilder::new().add(set.search(&lev)).add(set.search(&reg)).intersection();
        while let Some(key1) = stream1.next() {
            assert_eq!(stream2.next(), Some(key1));
        }
        assert_eq!(stream2.next(), None);
    }

    #[test]
    fn union_large() {
        let set = get_set();
        let lev = Levenshtein::new("foo", 3).unwrap();
        let reg = Regex::new("(..)*").unwrap();
        let mut stream1 = set.search((&lev).union(&reg)).into_stream();
        let mut stream2 = OpBuilder::new().add(set.search(&lev)).add(set.search(&reg)).union();
        while let Some(key1) = stream1.next() {
            assert_eq!(stream2.next(), Some(key1));
        }
        assert_eq!(stream2.next(), None);
    }
}
<commit_msg>comment fix<commit_after>use self::StartsWithStateInternal::*;

/// Automaton describes types that behave as a finite automaton.
///
/// All implementors of this trait are represented by *byte based* automata.
/// Stated differently, all transitions in the automata correspond to a single
/// byte in the input.
///
/// This implementation choice is important for a couple reasons:
///
/// 1. The set of possible transitions in each node is small, which may make
///    efficient memory usage easier.
/// 2. The finite state transducers in this crate are all byte based, so any
///    automata used on them must also be byte based.
///
/// In practice, this does present somewhat of a problem, for example, if
/// you're storing UTF-8 encoded strings in a finite state transducer. Consider
/// using a `Levenshtein` automaton, which accepts a query string and an edit
/// distance. The edit distance should apply to some notion of *character*,
/// which could be represented by at least 1-4 bytes in a UTF-8 encoding (for
/// some definition of "character"). Therefore, the automaton must have UTF-8
/// decoding built into it. This can be tricky to implement, so you may find
/// the [`utf8-ranges`](https://crates.io/crates/utf8-ranges) crate useful.
pub trait Automaton {
    /// The type of the state used in the automaton.
    type State;

    /// Returns a single start state for this automaton.
    ///
    /// This method should always return the same value for each
    /// implementation.
    fn start(&self) -> Self::State;

    /// Returns true if and only if `state` is a match state.
    fn is_match(&self, state: &Self::State) -> bool;

    /// Returns true if and only if `state` can lead to a match in zero or more
    /// steps.
    ///
    /// If this returns `false`, then no sequence of inputs from this state
    /// should ever produce a match. If this does not follow, then those match
    /// states may never be reached. In other words, behavior may be incorrect.
    ///
    /// If this returns `true` even when no match is possible, then behavior
    /// will be correct, but callers may be forced to do additional work.
    fn can_match(&self, _state: &Self::State) -> bool {
        true
    }

    /// Returns true if and only if `state` matches and must match no matter what
    /// steps are taken.
    ///
    /// If this returns `true`, then every sequence of inputs from this state
    /// produces a match. If this does not follow, then those match states may
    /// never be reached. In other words, behavior may be incorrect.
    ///
    /// If this returns `false` even when every sequence of inputs will lead to
    /// a match, then behavior will be correct, but callers may be forced to do
    /// additional work.
    fn will_always_match(&self, _state: &Self::State) -> bool {
        false
    }

    /// Return the next state given `state` and an input.
    fn accept(&self, state: &Self::State, byte: u8) -> Self::State;

    /// Returns an automaton that matches the strings that start with something
    /// this automaton matches.
    fn starts_with(self) -> StartsWith<Self> where Self: Sized {
        StartsWith(self)
    }

    /// Returns an automaton that matches the strings matched by either this or
    /// the other automaton.
    fn union<Rhs: Automaton>(
        self,
        rhs: Rhs,
    ) -> Union<Self, Rhs> where Self: Sized {
        Union(self, rhs)
    }

    /// Returns an automaton that matches the strings matched by both this and
    /// the other automaton.
    fn intersection<Rhs: Automaton>(
        self,
        rhs: Rhs,
    ) -> Intersection<Self, Rhs> where Self: Sized {
        Intersection(self, rhs)
    }

    /// Returns an automaton that matches the strings not matched by this
    /// automaton.
    fn complement(self) -> Complement<Self> where Self: Sized {
        Complement(self)
    }
}

impl<'a, T: Automaton> Automaton for &'a T {
    type State = T::State;

    fn start(&self) -> Self::State {
        (*self).start()
    }

    fn is_match(&self, state: &Self::State) -> bool {
        (*self).is_match(state)
    }

    fn can_match(&self, state: &Self::State) -> bool {
        (*self).can_match(state)
    }

    fn will_always_match(&self, state: &Self::State) -> bool {
        (*self).will_always_match(state)
    }

    fn accept(&self, state: &Self::State, byte: u8) -> Self::State {
        (*self).accept(state, byte)
    }
}

/// An automaton that always matches.
///
/// This is useful in a generic context as a way to express that no automaton
/// should be used.
pub struct AlwaysMatch;

impl Automaton for AlwaysMatch {
    type State = ();

    fn start(&self) -> () { () }
    fn is_match(&self, _: &()) -> bool { true }
    fn can_match(&self, _: &()) -> bool { true }
    fn will_always_match(&self, _: &()) -> bool { true }
    fn accept(&self, _: &(), _: u8) -> () { () }
}

/// An automaton that matches a string that begins with something that the
/// wrapped automaton matches.
pub struct StartsWith<A>(A);

/// The `Automaton` state for `StartsWith<A>`.
pub struct StartsWithState<A: Automaton>(StartsWithStateInternal<A>);

enum StartsWithStateInternal<A: Automaton> {
    Done,
    Running(A::State)
}

impl<A: Automaton> Automaton for StartsWith<A> {
    type State = StartsWithState<A>;
    fn start(&self) -> StartsWithState<A> {
        StartsWithState({
            let inner = self.0.start();
            if self.0.is_match(&inner) {
                Done
            } else {
                Running(inner)
            }
        })
    }

    fn is_match(&self, state: &StartsWithState<A>) -> bool {
        match state.0 {
            Done => true,
            Running(_) => false
        }
    }

    fn can_match(&self, state: &StartsWithState<A>) -> bool {
        match state.0 {
            Done => true,
            Running(ref inner) => self.0.can_match(inner)
        }
    }

    fn will_always_match(&self, state: &StartsWithState<A>) -> bool {
        match state.0 {
            Done => true,
            Running(_) => false
        }
    }

    fn accept(&self, state: &StartsWithState<A>, byte: u8) -> StartsWithState<A> {
        StartsWithState(
            match state.0 {
                Done => Done,
                Running(ref inner) => {
                    let next_inner = self.0.accept(inner, byte);
                    if self.0.is_match(&next_inner) {
                        Done
                    } else {
                        Running(next_inner)
                    }
                }
            }
        )
    }
}

/// An automaton that matches when one of its component automata match.
pub struct Union<A, B>(A, B);

/// The `Automaton` state for `Union<A, B>`.
pub struct UnionState<A: Automaton, B: Automaton>(A::State, B::State);

impl<A: Automaton, B: Automaton> Automaton for Union<A, B> {
    type State = UnionState<A, B>;

    fn start(&self) -> UnionState<A, B> {
        UnionState(
            self.0.start(),
            self.1.start()
        )
    }

    fn is_match(&self, state: &UnionState<A, B>) -> bool {
        self.0.is_match(&state.0) || self.1.is_match(&state.1)
    }

    fn can_match(&self, state: &UnionState<A, B>) -> bool {
        self.0.can_match(&state.0) || self.1.can_match(&state.1)
    }

    fn will_always_match(&self, state: &UnionState<A, B>) -> bool {
        self.0.will_always_match(&state.0) || self.1.will_always_match(&state.1)
    }

    fn accept(&self, state: &UnionState<A, B>, byte: u8) -> UnionState<A, B> {
        UnionState(
            self.0.accept(&state.0, byte),
            self.1.accept(&state.1, byte)
        )
    }
}

/// An automaton that matches when both of its component automata match.
pub struct Intersection<A, B>(A, B);

/// The `Automaton` state for `Intersection<A, B>`.
pub struct IntersectionState<A: Automaton, B: Automaton>(A::State, B::State);

impl<A: Automaton, B: Automaton> Automaton for Intersection<A, B> {
    type State = IntersectionState<A, B>;

    fn start(&self) -> IntersectionState<A, B> {
        IntersectionState(
            self.0.start(),
            self.1.start()
        )
    }

    fn is_match(&self, state: &IntersectionState<A, B>) -> bool {
        self.0.is_match(&state.0) && self.1.is_match(&state.1)
    }

    fn can_match(&self, state: &IntersectionState<A, B>) -> bool {
        self.0.can_match(&state.0) && self.1.can_match(&state.1)
    }

    fn will_always_match(&self, state: &IntersectionState<A, B>) -> bool {
        self.0.will_always_match(&state.0) && self.1.will_always_match(&state.1)
    }

    fn accept(&self, state: &IntersectionState<A, B>, byte: u8) -> IntersectionState<A, B> {
        IntersectionState(
            self.0.accept(&state.0, byte),
            self.1.accept(&state.1, byte)
        )
    }
}

/// An automaton that matches exactly when the automaton it wraps does not.
pub struct Complement<A>(A);

/// The `Automaton` state for `Complement<A>`.
pub struct ComplementState<A: Automaton>(A::State);

impl<A: Automaton> Automaton for Complement<A> {
    type State = ComplementState<A>;

    fn start(&self) -> ComplementState<A> {
        ComplementState(self.0.start())
    }

    fn is_match(&self, state: &ComplementState<A>) -> bool {
        !self.0.is_match(&state.0)
    }

    fn can_match(&self, state: &ComplementState<A>) -> bool {
        !self.0.will_always_match(&state.0)
    }

    fn will_always_match(&self, state: &ComplementState<A>) -> bool {
        !self.0.can_match(&state.0)
    }

    fn accept(&self, state: &ComplementState<A>, byte: u8) -> ComplementState<A> {
        ComplementState(self.0.accept(&state.0, byte))
    }
}

#[cfg(test)]
mod test {
    use ::{IntoStreamer, Streamer, Levenshtein, Regex, Set, Automaton};
    use set::OpBuilder;

    static WORDS: &'static str = include_str!("../../data/words-10000");

    fn get_set() -> Set {
        Set::from_iter(WORDS.lines()).unwrap()
    }

    #[test]
    fn complement_small() {
        let keys = vec!["fa", "fo", "fob", "focus", "foo", "food", "foul"];
        let set = Set::from_iter(keys).unwrap();
        let lev = Levenshtein::new("foo", 1).unwrap();
        let stream = set.search(lev.complement()).into_stream();

        let keys = stream.into_strs().unwrap();
        assert_eq!(keys, vec!["fa", "focus", "foul"]);
    }

    #[test]
    fn startswith_small() {
        let keys = vec!["", "cooing", "fa", "fo", "fob", "focus", "foo", "food", "foul", "fritter", "frothing"];
        let set = Set::from_iter(keys).unwrap();
        let lev = Levenshtein::new("foo", 1).unwrap();
        let stream = set.search(lev.starts_with()).into_stream();

        let keys = stream.into_strs().unwrap();
        assert_eq!(keys, vec!["cooing", "fo", "fob", "focus", "foo", "food", "foul", "frothing"]);
    }

    #[test]
    fn intersection_small() {
        let keys = vec!["fa", "fo", "fob", "focus", "foo", "food", "foul"];
        let set = Set::from_iter(keys).unwrap();
        let lev = Levenshtein::new("foo", 1).unwrap();
        let reg = Regex::new("(..)*").unwrap();
        let stream = set.search(lev.intersection(reg)).into_stream();

        let keys = stream.into_strs().unwrap();
        assert_eq!(keys, vec!["fo", "food"]);
    }

    #[test]
    fn union_small() {
        let keys = vec!["fa", "fo", "fob", "focus", "foo", "food", "foul"];
        let set = Set::from_iter(keys).unwrap();
        let lev = Levenshtein::new("foo", 1).unwrap();
        let reg = Regex::new("(..)*").unwrap();
        let stream = set.search(lev.union(reg)).into_stream();

        let keys = stream.into_strs().unwrap();
        assert_eq!(keys, vec!["fa", "fo", "fob", "foo", "food", "foul"]);
    }

    #[test]
    fn intersection_large() {
        let set = get_set();
        let lev = Levenshtein::new("foo", 3).unwrap();
        let reg = Regex::new("(..)*").unwrap();
        let mut stream1 = set.search((&lev).intersection(&reg)).into_stream();
        let mut stream2 = OpBuilder::new().add(set.search(&lev)).add(set.search(&reg)).intersection();
        while let Some(key1) = stream1.next() {
            assert_eq!(stream2.next(), Some(key1));
        }
        assert_eq!(stream2.next(), None);
    }

    #[test]
    fn union_large() {
        let set = get_set();
        let lev = Levenshtein::new("foo", 3).unwrap();
        let reg = Regex::new("(..)*").unwrap();
        let mut stream1 = set.search((&lev).union(&reg)).into_stream();
        let mut stream2 = OpBuilder::new().add(set.search(&lev)).add(set.search(&reg)).union();
        while let Some(key1) = stream1.next() {
            assert_eq!(stream2.next(), Some(key1));
        }
        assert_eq!(stream2.next(), None);
    }
}
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

use failure::format_err;
use futures::stream::StreamExt;
use slog_scope::{info, warn};

use mimir2::domain::ports::primary::list_documents::ListDocuments;
use mimir2::domain::{
    model::index::IndexVisibility, ports::primary::generate_index::GenerateIndex,
};
use mimir2::{
    adapters::secondary::elasticsearch::{self, ES_DEFAULT_TIMEOUT, ES_DEFAULT_VERSION_REQ},
    domain::ports::secondary::remote::Remote,
};
use mimirsbrunn::admin_geofinder::AdminGeoFinder;
use mimirsbrunn::osm_reader::make_osm_reader;
use mimirsbrunn::osm_reader::poi::{add_address, compute_weight, pois, PoiConfig};
use mimirsbrunn::osm_reader::street::{compute_street_weight, streets};
use mimirsbrunn::settings::osm2mimir::{Args, Settings};
use places::poi::Poi;

const POI_REVERSE_GEOCODING_CONCURRENCY: usize = 8;

async fn run(args: Args) -> Result<(), mimirsbrunn::Error> {
    let input = args.input.clone(); // we save the input, because args will be consumed by settings.
    validate_args(&args)?;
    let settings = &Settings::new(args.clone())?;

    let mut osm_reader = make_osm_reader(&input)?;

    let import_streets_enabled = settings
        .street
        .as_ref()
        .map(|street| street.import)
        .unwrap_or_else(|| false);

    let import_poi_enabled = settings
        .poi
        .as_ref()
        .map(|poi| poi.import)
        .unwrap_or_else(|| false);

    if !import_streets_enabled && !import_poi_enabled {
        return Err(format_err!(
            "Neither streets nor POIs import is enabled. Nothing to do.\n\
             Use --import-way=true or --import-poi=true"
        ));
    }

    let pool =
        elasticsearch::remote::connection_pool_url(&settings.elasticsearch.connection_string)
            .await
            .map_err(|err| {
                format_err!(
                    "could not create elasticsearch connection pool: {}",
                    err.to_string()
                )
            })?;

    let client = pool
        .conn(ES_DEFAULT_TIMEOUT, ES_DEFAULT_VERSION_REQ)
        .await
        .map_err(|err| format_err!("could not connect elasticsearch pool: {}", err.to_string()))?;

    let admins_geofinder: AdminGeoFinder = match client.list_documents().await {
        Ok(stream) => {
            stream
                .map(|admin| admin.expect("could not parse admin"))
                .collect()
                .await
        }
        Err(err) => {
            warn!("administratives regions not found in es db. {:?}", err);
            std::iter::empty().collect()
        }
    };

    if import_streets_enabled {
        info!("Extracting streets from osm");
        let mut streets = streets(&mut osm_reader, &admins_geofinder, &settings)?;

        info!("computing street weight");
        compute_street_weight(&mut streets);

        let streets = futures::stream::iter(streets);
        let street_config = args.get_street_config()?;

        client
            .generate_index(street_config, streets, IndexVisibility::Public)
            .await
            .map_err(|err| format_err!("could not generate index: {}", err.to_string()))?;
    }

    if import_poi_enabled {
        let config = settings
            .poi
            .as_ref()
            .and_then(|poi| poi.config.clone())
            .unwrap_or_else(PoiConfig::default);

        // Ideally, this pois function would create a stream, which would then map and do other
        // stuff, and then be indexed
        info!("Extracting pois from osm");
        let pois = pois(&mut osm_reader, &config, &admins_geofinder);

        let pois: Vec<Poi> = futures::stream::iter(pois)
            .map(compute_weight)
            .map(|poi| add_address(&client, poi))
            .buffer_unordered(POI_REVERSE_GEOCODING_CONCURRENCY)
            .collect()
            .await;

        let poi_config = args.get_poi_config()?;

        client
            .generate_index(
                poi_config,
                futures::stream::iter(pois),
                IndexVisibility::Public,
            )
            .await
            .map_err(|err| format_err!("could not generate index: {}", err.to_string()))?;
    }
    Ok(())
}

// We need to allow for unused variables, because currently all the checks on
// args require the db-storage feature. If this feature is not used, then there
// is a warning
#[allow(unused_variables)]
fn validate_args(args: &Args) -> Result<(), mimirsbrunn::Error> {
    #[cfg(feature = "db-storage")]
    if args.db_file.is_some() {
        // If the user specified db_file, he must also specify db_buffer_size, or else!
        if args.db_buffer_size.is_none() {
            return Err(failure::format_err!("You need to specify database buffer size if you want to use database storage. Use --db-buffer-size"));
        }
    }
    #[cfg(feature = "db-storage")]
    if args.db_buffer_size.is_some() {
        // If the user specified db_buffer_size, he must also specify db_file, or else!
        if args.db_file.is_none() {
            return Err(failure::format_err!("You need to specify database file if you want to use database storage. Use --db-file"));
        }
    }
    Ok(())
}

#[tokio::main]
async fn main() {
    mimirsbrunn::utils::launch_async(Box::new(run)).await;
}
<commit_msg>Fix clippy warning<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

use failure::format_err;
use futures::stream::StreamExt;
use slog_scope::{info, warn};

use mimir2::domain::ports::primary::list_documents::ListDocuments;
use mimir2::domain::{
    model::index::IndexVisibility, ports::primary::generate_index::GenerateIndex,
};
use mimir2::{
    adapters::secondary::elasticsearch::{self, ES_DEFAULT_TIMEOUT, ES_DEFAULT_VERSION_REQ},
    domain::ports::secondary::remote::Remote,
};
use mimirsbrunn::admin_geofinder::AdminGeoFinder;
use mimirsbrunn::osm_reader::make_osm_reader;
use mimirsbrunn::osm_reader::poi::{add_address, compute_weight, pois, PoiConfig};
use mimirsbrunn::osm_reader::street::{compute_street_weight, streets};
use mimirsbrunn::settings::osm2mimir::{Args, Settings};
use places::poi::Poi;

const POI_REVERSE_GEOCODING_CONCURRENCY: usize = 8;

async fn run(args: Args) -> Result<(), mimirsbrunn::Error> {
    let input = args.input.clone(); // we save the input, because args will be consumed by settings.
    validate_args(&args)?;
    let settings = &Settings::new(args.clone())?;

    let mut osm_reader = make_osm_reader(&input)?;

    let import_streets_enabled = settings
        .street
        .as_ref()
        .map(|street| street.import)
        .unwrap_or_else(|| false);

    let import_poi_enabled = settings
        .poi
        .as_ref()
        .map(|poi| poi.import)
        .unwrap_or_else(|| false);

    if !import_streets_enabled && !import_poi_enabled {
        return Err(format_err!(
            "Neither streets nor POIs import is enabled. Nothing to do.\n\
             Use --import-way=true or --import-poi=true"
        ));
    }

    let pool =
        elasticsearch::remote::connection_pool_url(&settings.elasticsearch.connection_string)
            .await
            .map_err(|err| {
                format_err!(
                    "could not create elasticsearch connection pool: {}",
                    err.to_string()
                )
            })?;

    let client = pool
        .conn(ES_DEFAULT_TIMEOUT, ES_DEFAULT_VERSION_REQ)
        .await
        .map_err(|err| format_err!("could not connect elasticsearch pool: {}", err.to_string()))?;

    let admins_geofinder: AdminGeoFinder = match client.list_documents().await {
        Ok(stream) => {
            stream
                .map(|admin| admin.expect("could not parse admin"))
                .collect()
                .await
        }
        Err(err) => {
            warn!("administratives regions not found in es db. {:?}", err);
            std::iter::empty().collect()
        }
    };

    if import_streets_enabled {
        info!("Extracting streets from osm");
        let mut streets = streets(&mut osm_reader, &admins_geofinder, settings)?;

        info!("computing street weight");
        compute_street_weight(&mut streets);

        let streets = futures::stream::iter(streets);
        let street_config = args.get_street_config()?;

        client
            .generate_index(street_config, streets, IndexVisibility::Public)
            .await
            .map_err(|err| format_err!("could not generate index: {}", err.to_string()))?;
    }

    if import_poi_enabled {
        let config = settings
            .poi
            .as_ref()
            .and_then(|poi| poi.config.clone())
            .unwrap_or_else(PoiConfig::default);

        // Ideally, this pois function would create a stream, which would then map and do other
        // stuff, and then be indexed
        info!("Extracting pois from osm");
        let pois = pois(&mut osm_reader, &config, &admins_geofinder);

        let pois: Vec<Poi> = futures::stream::iter(pois)
            .map(compute_weight)
            .map(|poi| add_address(&client, poi))
            .buffer_unordered(POI_REVERSE_GEOCODING_CONCURRENCY)
            .collect()
            .await;

        let poi_config = args.get_poi_config()?;

        client
            .generate_index(
                poi_config,
                futures::stream::iter(pois),
                IndexVisibility::Public,
            )
            .await
            .map_err(|err| format_err!("could not generate index: {}", err.to_string()))?;
    }
    Ok(())
}

// We need to allow for unused variables, because currently all the checks on
// args require the db-storage feature. If this feature is not used, then there
// is a warning
#[allow(unused_variables)]
fn validate_args(args: &Args) -> Result<(), mimirsbrunn::Error> {
    #[cfg(feature = "db-storage")]
    if args.db_file.is_some() {
        // If the user specified db_file, he must also specify db_buffer_size, or else!
        if args.db_buffer_size.is_none() {
            return Err(failure::format_err!("You need to specify database buffer size if you want to use database storage. Use --db-buffer-size"));
        }
    }
    #[cfg(feature = "db-storage")]
    if args.db_buffer_size.is_some() {
        // If the user specified db_buffer_size, he must also specify db_file, or else!
        if args.db_file.is_none() {
            return Err(failure::format_err!("You need to specify database file if you want to use database storage. Use --db-file"));
        }
    }
    Ok(())
}

#[tokio::main]
async fn main() {
    mimirsbrunn::utils::launch_async(Box::new(run)).await;
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate log;
extern crate env_logger;
extern crate osmpbfreader;
extern crate rustc_serialize;
extern crate docopt;
extern crate mimirsbrunn;

use std::collections::HashSet;
use std::collections::BTreeMap;
use osmpbfreader::{OsmObj, OsmId};

pub type AdminsMap = BTreeMap<OsmId, mimirsbrunn::Admin>;

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_input: String,
    flag_level: Vec<u32>,
}

static USAGE: &'static str = "
Usage:
    osm2mimir --input=<pbf-file> --level=<admin-level>...

Options:
    -h, --help                               Show this message.
    -i <pbf-file>, --input=<pbf-file>        OSM PBF file.
    -l <admin-level>, --level=<admin-level>  Admin levels to keep.
";

fn get_osm_id(obj: &OsmObj) -> OsmId {
    match *obj {
        OsmObj::Node(ref node) => OsmId::Node(node.id),
        OsmObj::Way(ref way) => OsmId::Way(way.id),
        OsmObj::Relation(ref rel) => OsmId::Relation(rel.id),
    }
}

fn update_coordinates(filename: &String, admins: &mut AdminsMap) {
    if admins.is_empty() {
        return;
    }
    // load coord for administratives regions
    let path = std::path::Path::new(&filename);
    let r = std::fs::File::open(&path).unwrap();
    let mut pbf = osmpbfreader::OsmPbfReader::new(r);
    for obj in pbf.iter() {
        if let osmpbfreader::OsmObj::Node(ref node) = obj {
            let mut adm = match admins.get_mut(&get_osm_id(&obj)) {
                Some(val) => val,
                None => continue,
            };
            adm.coord.lat = node.lat;
            adm.coord.lon = node.lon;
        }
    }
}

fn administartive_regions(filename: &String, levels: &HashSet<u32>) -> AdminsMap {
    let mut administrative_regions = AdminsMap::new();
    let path = std::path::Path::new(&filename);
    let r = std::fs::File::open(&path).unwrap();
    let mut pbf = osmpbfreader::OsmPbfReader::new(r);
    // load administratives regions
    for obj in pbf.iter() {
        if let osmpbfreader::OsmObj::Relation(relation) = obj {
            // not administartive region
            if !relation.tags
                        .get("boundary")
                        .map(|s| s == "administrative")
                        .unwrap_or(false) {
                continue;
            }
            // admininstrative region without coordinates
            let admin_centre = match relation.refs.iter().find(|rf| rf.role == "admin_centre") {
                Some(val) => val.member,
                None => {
                    info!("adminstrative regione without coordinates for relation {}.",
                          relation.id);
                    continue;
                }
            };
            let level = relation.tags
                                .get("admin_level")
                                .and_then(|s| s.parse().ok());
            let level = match level {
                None => {
                    info!("invalid admin_level for relation {}: admin_level {:?}",
                          relation.id,
                          relation.tags.get("admin_level"));
                    continue;
                }
                Some(ref l) if !levels.contains(&l) => continue,
                Some(l) => l,
            };
            // administrative region with name ?
            let name = match relation.tags.get("name") {
                Some(val) => val,
                None => {
                    info!("adminstrative region without name for relation {}:  admin_level {} \
                           ignored.",
                          relation.id,
                          level);
                    continue;
                }
            };
            let admin_id = match relation.tags.get("ref:INSEE") {
                Some(val) => format!("admin:fr:{}", val),
                None => format!("admin:osm:{}", relation.id),
            };
            let zip_code = match relation.tags.get("addr:postcode") {
                Some(val) => &val[..],
                None => "",
            };
            let admin = mimirsbrunn::Admin {
                id: admin_id,
                level: level,
                name: name.to_string(),
                zip_code: zip_code.to_string(),
                // TODO weight value ?
                weight: 1,
                coord: mimirsbrunn::Coord {
                    lat: 0.0,
                    lon: 0.0,
                },
            };
            administrative_regions.insert(admin_centre, admin);
        }
    }
    return administrative_regions;
}

fn main() {
    env_logger::init().unwrap();
    debug!("importing adminstrative region into Mimir");
    let args: Args = docopt::Docopt::new(USAGE)
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());
    let levels = args.flag_level.iter().cloned().collect();
    let mut res = administartive_regions(&args.flag_input, &levels);
    update_coordinates(&args.flag_input, &mut res);
}
<commit_msg>mimirsbrunn: url elasticsearch added in Args<commit_after>#[macro_use]
extern crate log;
extern crate env_logger;
extern crate osmpbfreader;
extern crate rustc_serialize;
extern crate docopt;
extern crate mimirsbrunn;

use std::collections::HashSet;
use std::collections::BTreeMap;
use osmpbfreader::{OsmObj, OsmId};

pub type AdminsMap = BTreeMap<OsmId, mimirsbrunn::Admin>;

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_input: String,
    flag_level: Vec<u32>,
    flag_connection_string: String,
}

static USAGE: &'static str = "
Usage:
    osm2mimir --help
    osm2mimir --input=<file> [--connection-string=<connection-string>] --level=<level>...

Options:
    -h, --help            Show this message.
    -i, --input=<file>    OSM PBF file.
    -l, --level=<level>   Admin levels to keep.
    -c, --connection-string=<connection-string>
                          Elasticsearch parameters, [default: http://localhost:9200/munin]
";

fn get_osm_id(obj: &OsmObj) -> OsmId {
    match *obj {
        OsmObj::Node(ref node) => OsmId::Node(node.id),
        OsmObj::Way(ref way) => OsmId::Way(way.id),
        OsmObj::Relation(ref rel) => OsmId::Relation(rel.id),
    }
}

fn update_coordinates(filename: &String, admins: &mut AdminsMap) {
    if admins.is_empty() {
        return;
    }
    // load coord for administratives regions
    let path = std::path::Path::new(&filename);
    let r = std::fs::File::open(&path).unwrap();
    let mut pbf = osmpbfreader::OsmPbfReader::new(r);
    for obj in pbf.iter() {
        if let osmpbfreader::OsmObj::Node(ref node) = obj {
            let mut adm = match admins.get_mut(&get_osm_id(&obj)) {
                Some(val) => val,
                None => continue,
            };
            adm.coord.lat = node.lat;
            adm.coord.lon = node.lon;
        }
    }
}

fn administartive_regions(filename: &String, levels: &HashSet<u32>) -> AdminsMap {
    let mut administrative_regions = AdminsMap::new();
    let path = std::path::Path::new(&filename);
    let r = std::fs::File::open(&path).unwrap();
    let mut pbf = osmpbfreader::OsmPbfReader::new(r);
    // load administratives regions
    for obj in pbf.iter() {
        if let osmpbfreader::OsmObj::Relation(relation) = obj {
            // not administartive region
            if !relation.tags
                        .get("boundary")
                        .map(|s| s == "administrative")
                        .unwrap_or(false) {
                continue;
            }
            // admininstrative region without coordinates
            let admin_centre = match relation.refs.iter().find(|rf| rf.role == "admin_centre") {
                Some(val) => val.member,
                None => {
                    info!("adminstrative regione without coordinates for relation {}.",
                          relation.id);
                    continue;
                }
            };
            let level = relation.tags
                                .get("admin_level")
                                .and_then(|s| s.parse().ok());
            let level = match level {
                None => {
                    info!("invalid admin_level for relation {}: admin_level {:?}",
                          relation.id,
                          relation.tags.get("admin_level"));
                    continue;
                }
                Some(ref l) if !levels.contains(&l) => continue,
                Some(l) => l,
            };
            // administrative region with name ?
            let name = match relation.tags.get("name") {
                Some(val) => val,
                None => {
                    info!("adminstrative region without name for relation {}:  admin_level {} \
                           ignored.",
                          relation.id,
                          level);
                    continue;
                }
            };
            let admin_id = match relation.tags.get("ref:INSEE") {
                Some(val) => format!("admin:fr:{}", val),
                None => format!("admin:osm:{}", relation.id),
            };
            let zip_code = match relation.tags.get("addr:postcode") {
                Some(val) => &val[..],
                None => "",
            };
            let admin = mimirsbrunn::Admin {
                id: admin_id,
                level: level,
                name: name.to_string(),
                zip_code: zip_code.to_string(),
                // TODO weight value ?
                weight: 1,
                coord: mimirsbrunn::Coord {
                    lat: 0.0,
                    lon: 0.0,
                },
            };
            administrative_regions.insert(admin_centre, admin);
        }
    }
    return administrative_regions;
}

fn main() {
    env_logger::init().unwrap();
    debug!("importing adminstrative region into Mimir");
    let args: Args = docopt::Docopt::new(USAGE)
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());
    let levels = args.flag_level.iter().cloned().collect();
    let mut res = administartive_regions(&args.flag_input, &levels);
    update_coordinates(&args.flag_input, &mut res);
}
<|endoftext|>