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

//! A small fast RNG

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

#[cfg(target_pointer_width = "64")]
type Rng = super::xoshiro256plusplus::Xoshiro256PlusPlus;
#[cfg(not(target_pointer_width = "64"))]
type Rng = super::xoshiro128plusplus::Xoshiro128PlusPlus;

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

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

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

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

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

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

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

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

//! A small fast RNG

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

#[cfg(target_pointer_width = "64")]
type Rng = super::xoshiro256plusplus::Xoshiro256PlusPlus;
#[cfg(not(target_pointer_width = "64"))]
type Rng = super::xoshiro128plusplus::Xoshiro128PlusPlus;

/// A small-state, fast non-crypto PRNG
///
/// `SmallRng` may be a good choice when a PRNG with small state, cheap
/// initialization, good statistical quality and good performance are required.
/// Note that depending on the application, [`StdRng`] is faster on many modern
/// platforms while providing higher-quality randomness. Furthermore, `SmallRng`
/// is **not** a good choice when:
/// - Security against prediction or reproducibility are important.
///   Use [`StdRng`] instead.
/// - Seeds with many zeros are provided. In such cases, it takes `SmallRng`
///   about 10 samples to produce 0 and 1 bits with equal probability. Either
///   provide seeds with an approximately equal number of 0 and 1 (for example
///   by using [`SeedableRng::from_entropy`] or [`SeedableRng::seed_from_u64`]),
///   or use [`StdRng`] instead.
///
/// The algorithm is deterministic but should not be considered reproducible
/// due to dependence on platform and possible replacement in future
/// library versions. For a reproducible generator, use a named PRNG from an
/// external crate, e.g. [rand_xoshiro] or [rand_chacha].
/// Refer also to [The Book](https://rust-random.github.io/book/guide-rngs.html).
///
/// The PRNG algorithm in `SmallRng` is chosen to be efficient on the current
/// platform, without consideration for cryptography or security. The size of
/// its state is much smaller than [`StdRng`]. The current algorithm is
/// `Xoshiro256PlusPlus` on 64-bit platforms and `Xoshiro128PlusPlus` on 32-bit
/// platforms. Both are also implemented by the [rand_xoshiro] crate.
///
/// # Examples
///
/// Initializing `SmallRng` with a random seed can be done using [`SeedableRng::from_entropy`]:
///
/// ```
/// use rand::{Rng, SeedableRng};
/// use rand::rngs::SmallRng;
///
/// // Create small, cheap to initialize and fast RNG with a random seed.
/// // The randomness is supplied by the operating system.
/// let mut small_rng = SmallRng::from_entropy();
/// # let v: u32 = small_rng.gen();
/// ```
///
/// When initializing a lot of `SmallRng`'s, using [`thread_rng`] can be more
/// efficient:
///
/// ```
/// use rand::{SeedableRng, thread_rng};
/// use rand::rngs::SmallRng;
///
/// // Create a big, expensive to initialize and slower, but unpredictable RNG.
/// // This is cached and done only once per thread.
/// let mut thread_rng = thread_rng();
/// // Create small, cheap to initialize and fast RNGs with random seeds.
/// // One can generally assume this won't fail.
/// let rngs: Vec<SmallRng> = (0..10)
///     .map(|_| SmallRng::from_rng(&mut thread_rng).unwrap())
///     .collect();
/// ```
///
/// [`StdRng`]: crate::rngs::StdRng
/// [`thread_rng`]: crate::thread_rng
/// [rand_chacha]: https://crates.io/crates/rand_chacha
/// [rand_xoshiro]: https://crates.io/crates/rand_xoshiro
#[cfg_attr(doc_cfg, doc(cfg(feature = "small_rng")))]
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SmallRng(Rng);

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

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

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

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

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

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

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


<commit_before>extern crate rand;
extern crate evo;
extern crate petgraph;
extern crate graph_annealing;
extern crate clap;
extern crate crossbeam;
extern crate simple_parallel;
extern crate num_cpus;
extern crate pcg;
extern crate graph_sgf;
extern crate triadic_census;

use std::f32;
use std::fmt::Debug;
use std::str::FromStr;
use clap::{App, Arg};
use rand::{Rng, SeedableRng};
use rand::distributions::{IndependentSample, Range};
use rand::os::OsRng;
use pcg::PcgRng;

use evo::Probability;
use evo::nsga2::{self, FitnessEval, MultiObjective3};
use graph_annealing::repr::edge_ops_genome::{EdgeOpsGenome, Toolbox};
use graph_annealing::helper::draw_graph;
use graph_annealing::goal::Goal;
use simple_parallel::Pool;

use petgraph::{Directed, Graph};
use graph_sgf::{PetgraphReader, Unweighted};
use triadic_census::OptDenseDigraph;

use std::io::BufReader;
use std::fs::File;

// maximum connectivity
const MAX_DEGREE: u32 = 20;

#[derive(Debug, Copy, Clone)]
enum FitnessFunction {
    Null,
    ConnectedComponents,
    StronglyConnectedComponents,
    NeighborMatching,
    TriadicDistance,
}

fn apply_fitness_function<N: Clone + Default, E: Clone + Default>(fitfun: FitnessFunction,
                                                                  goal: &Goal<N, E>,
                                                                  g: &OptDenseDigraph<(), ()>)
                                                                  -> f32 {
    match fitfun {
        FitnessFunction::Null => {
            0.0
        }
        FitnessFunction::ConnectedComponents => {
            goal.connected_components_distance(g) as f32
        }
        FitnessFunction::StronglyConnectedComponents => {
            goal.strongly_connected_components_distance(g) as f32
        }
        FitnessFunction::NeighborMatching => {
            goal.neighbor_matching_score(g) as f32
        }
        FitnessFunction::TriadicDistance => {
            goal.triadic_distance(g) as f32
        }
    }
}

struct MyEval<N, E> {
    goal: Goal<N, E>,
    pool: Pool,
    fitness_functions: (FitnessFunction, FitnessFunction, FitnessFunction),
}

fn fitness<N: Clone + Default, E: Clone + Default>(fitness_functions: (FitnessFunction,
                                                                       FitnessFunction,
                                                                       FitnessFunction),
                                                   goal: &Goal<N, E>,
                                                   ind: &EdgeOpsGenome)
                                                   -> MultiObjective3<f32> {
    let g = ind.to_graph(MAX_DEGREE);
    MultiObjective3::from((apply_fitness_function(fitness_functions.0, goal, &g),
                           apply_fitness_function(fitness_functions.1, goal, &g),
                           apply_fitness_function(fitness_functions.2, goal, &g)))

}

impl<N:Clone+Sync+Default,E:Clone+Sync+Default> FitnessEval<EdgeOpsGenome, MultiObjective3<f32>> for MyEval<N,E> {
    fn fitness(&mut self, pop: &[EdgeOpsGenome]) -> Vec<MultiObjective3<f32>> {
        let pool = &mut self.pool;
        let goal = &self.goal;

        let fitness_functions = self.fitness_functions;

        crossbeam::scope(|scope| {
            pool.map(scope, pop, |ind| fitness(fitness_functions, goal, ind))
                .collect()
        })
    }
}

#[derive(Debug)]
struct Stat<T: Debug> {
    min: T,
    max: T,
    avg: T,
}

#[allow(non_snake_case)]
fn main() {
    let ncpus = num_cpus::get();
    println!("Using {} CPUs", ncpus);
    let matches = App::new("nsga2_edge")
                      .arg(Arg::with_name("NGEN")
                               .long("ngen")
                               .help("Number of generations to run")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("OPS")
                               .long("ops")
                               .help("Edge operation and weight specification, e.g. \
                                      Dup:1,Split:3,Parent:2")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("GRAPH")
                               .long("graph")
                               .help("SGF graph file")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("MU")
                               .long("mu")
                               .help("Population size")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("LAMBDA")
                               .long("lambda")
                               .help("Size of offspring population")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("K")
                               .long("k")
                               .help("Tournament selection (default: 2)")
                               .takes_value(true))
                      .arg(Arg::with_name("SEED")
                               .long("seed")
                               .help("Seed value for Rng")
                               .takes_value(true))
                      .arg(Arg::with_name("PMUT")
                               .long("pmut")
                               .help("Probability for element mutation")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("ILEN")
                               .long("ilen")
                               .help("Initial genome length (random range)")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("OBJECTIVES")
                               .long("objectives")
                               .help("Specify 3 objective functions comma separated (null, cc, \
                                      scc, nm, td), e.g. cc,nm,td")
                               .takes_value(true)
                               .required(true))
                      .get_matches();

    // number of generations
    let NGEN: usize = FromStr::from_str(matches.value_of("NGEN").unwrap()).unwrap();
    println!("NGEN: {}", NGEN);

    // size of population
    let MU: usize = FromStr::from_str(matches.value_of("MU").unwrap()).unwrap();
    println!("MU: {}", MU);

    // size of offspring population
    let LAMBDA: usize = FromStr::from_str(matches.value_of("LAMBDA").unwrap()).unwrap();
    println!("LAMBDA: {}", LAMBDA);

    // tournament selection
    let K: usize = FromStr::from_str(matches.value_of("K").unwrap_or("2")).unwrap();
    assert!(K > 0);
    println!("K: {}", K);

    let seed: Vec<u64>;
    if let Some(seed_str) = matches.value_of("SEED") {
        seed = seed_str.split(",").map(|s| FromStr::from_str(s).unwrap()).collect();
    } else {
        println!("Use OsRng to generate seed..");
        let mut rng = OsRng::new().unwrap();
        seed = (0..2).map(|_| rng.next_u64()).collect();
    }
    println!("SEED: {:?}", seed);

    let PMUT: f32 = FromStr::from_str(matches.value_of("PMUT").unwrap()).unwrap();
    println!("PMUT: {:?}", PMUT);

    let ilen_str = matches.value_of("ILEN").unwrap();
    let ilen: Vec<usize> = ilen_str.split(",").map(|s| FromStr::from_str(s).unwrap()).collect();
    assert!(ilen.len() == 1 || ilen.len() == 2);

    let ilen_from = ilen[0];
    let ilen_to = if ilen.len() == 1 {
        ilen[0]
    } else {
        ilen[1]
    };
    assert!(ilen_from <= ilen_to);

    // read objective functions
    let mut objectives_arr: Vec<FitnessFunction> = Vec::new();
    for s in matches.value_of("OBJECTIVES").unwrap().split(",") {
        objectives_arr.push(match s {
            "null" => {
                FitnessFunction::Null
            }
            "cc" => {
                FitnessFunction::ConnectedComponents
            }
            "scc" => {
                FitnessFunction::StronglyConnectedComponents
            }
            "nm" => {
                FitnessFunction::NeighborMatching
            }
            "td" => {
                FitnessFunction::TriadicDistance
            }
            _ => {
                panic!("Invalid objective function");
            }
        });
    }

    while objectives_arr.len() < 3 {
        objectives_arr.push(FitnessFunction::Null);
    }

    if objectives_arr.len() > 3 {
        panic!("Max 3 objectives allowed");
    }

    println!("objectives={:?}", objectives_arr);

    // read graph
    let graph_file = matches.value_of("GRAPH").unwrap();
    println!("Using graph file: {}", graph_file);
    let graph = {
        let mut f = BufReader::new(File::open(graph_file).unwrap());
        let graph: Graph<Unweighted, Unweighted, Directed> = PetgraphReader::from_sgf(&mut f);
        graph
    };

    // Parse weighted operation choice from command line
    let ops = Toolbox::parse_weighted_op_choice_list(matches.value_of("OPS").unwrap()).unwrap();
    println!("edge ops: {:?}", ops);
    let mut toolbox = Toolbox::new(Probability::new(PMUT), &ops[..]);

    let mut evaluator = MyEval {
        goal: Goal::new(OptDenseDigraph::from(graph)),
        pool: Pool::new(ncpus),
        fitness_functions: (objectives_arr[0], objectives_arr[1], objectives_arr[2]),
    };

    assert!(seed.len() == 2);
    let mut rng: PcgRng = SeedableRng::from_seed([seed[0], seed[1]]);

    // create initial random population
    let initial_population: Vec<EdgeOpsGenome> = (0..MU)
                                                     .map(|_| {
                                                         let len = if ilen_from == ilen_to {
                                                             ilen_from
                                                         } else {
                                                             Range::new(ilen_from, ilen_to)
                                                                 .ind_sample(&mut rng)
                                                         };
                                                         toolbox.random_genome(&mut rng, len)
                                                     })
                                                     .collect();

    // evaluate fitness
    let fitness: Vec<_> = evaluator.fitness(&initial_population[..]);
    assert!(fitness.len() == initial_population.len());

    let mut pop = initial_population;
    let mut fit = fitness;

    for i in 0..NGEN {
        println!("Iteration: {}", i);
        let (new_pop, new_fit) = nsga2::iterate(&mut rng,
                                                pop,
                                                fit,
                                                &mut evaluator,
                                                MU,
                                                LAMBDA,
                                                K,
                                                &mut toolbox);
        pop = new_pop;
        fit = new_fit;

        let stats: Vec<Stat<f32>> = [0, 1, 2]
                                        .iter()
                                        .map(|&i| {
                                            let min = fit.iter().fold(f32::INFINITY, |acc, f| {
                                                let x = f.objectives[i];
                                                if x < acc {
                                                    x
                                                } else {
                                                    acc
                                                }
                                            });
                                            let max = fit.iter().fold(-f32::INFINITY, |acc, f| {
                                                let x = f.objectives[i];
                                                if x > acc {
                                                    x
                                                } else {
                                                    acc
                                                }
                                            });
                                            let sum = fit.iter()
                                                         .fold(0.0, |acc, f| acc + f.objectives[i]);
                                            Stat {
                                                min: min,
                                                max: max,
                                                avg: sum / fit.len() as f32,
                                            }
                                        })
                                        .collect();
        println!("stats: {:?}", stats);
        let mut found_optimum = false;
        for f in fit.iter() {
            if f.objectives[0] < 1.0 && f.objectives[1] <= 1.0 && f.objectives[2] < 0.01 {
                found_optimum = true;
                break;
            }
        }
        if found_optimum {
            println!("Found premature optimum in Iteration {}", i);
            break;
        }
    }
    println!("===========================================================");

    // finally evaluate rank and crowding distance (using select()).
    let rank_dist = nsga2::select(&fit[..], MU);
    assert!(rank_dist.len() == MU);

    let mut j = 0;

    for rd in rank_dist.iter() {
        if rd.rank == 0 {
            println!("-------------------------------------------");
            println!("rd: {:?}", rd);
            println!("fitness: {:?}", fit[rd.idx]);
            println!("genome: {:?}", pop[rd.idx]);

            if fit[rd.idx].objectives[0] < 1.0 {
                draw_graph((&pop[rd.idx].to_graph(MAX_DEGREE)).ref_graph(),
                           // XXX: name
                           &format!("nsga2edgeops_g{}_f{}_i{}.svg",
                                    NGEN,
                                    fit[rd.idx].objectives[1] as usize,
                                    j));
                j += 1;
            }
        }
    }
}
<commit_msg>Read --threshold from commandline<commit_after>extern crate rand;
extern crate evo;
extern crate petgraph;
extern crate graph_annealing;
extern crate clap;
extern crate crossbeam;
extern crate simple_parallel;
extern crate num_cpus;
extern crate pcg;
extern crate graph_sgf;
extern crate triadic_census;

use std::f32;
use std::fmt::Debug;
use std::str::FromStr;
use clap::{App, Arg};
use rand::{Rng, SeedableRng};
use rand::distributions::{IndependentSample, Range};
use rand::os::OsRng;
use pcg::PcgRng;

use evo::Probability;
use evo::nsga2::{self, FitnessEval, MultiObjective3};
use graph_annealing::repr::edge_ops_genome::{EdgeOpsGenome, Toolbox};
use graph_annealing::helper::draw_graph;
use graph_annealing::goal::Goal;
use simple_parallel::Pool;

use petgraph::{Directed, Graph};
use graph_sgf::{PetgraphReader, Unweighted};
use triadic_census::OptDenseDigraph;

use std::io::BufReader;
use std::fs::File;

// maximum connectivity
const MAX_DEGREE: u32 = 20;

#[derive(Debug, Copy, Clone)]
enum FitnessFunction {
    Null,
    ConnectedComponents,
    StronglyConnectedComponents,
    NeighborMatching,
    TriadicDistance,
}

fn apply_fitness_function<N: Clone + Default, E: Clone + Default>(fitfun: FitnessFunction,
                                                                  goal: &Goal<N, E>,
                                                                  g: &OptDenseDigraph<(), ()>)
                                                                  -> f32 {
    match fitfun {
        FitnessFunction::Null => {
            0.0
        }
        FitnessFunction::ConnectedComponents => {
            goal.connected_components_distance(g) as f32
        }
        FitnessFunction::StronglyConnectedComponents => {
            goal.strongly_connected_components_distance(g) as f32
        }
        FitnessFunction::NeighborMatching => {
            goal.neighbor_matching_score(g) as f32
        }
        FitnessFunction::TriadicDistance => {
            goal.triadic_distance(g) as f32
        }
    }
}

struct MyEval<N, E> {
    goal: Goal<N, E>,
    pool: Pool,
    fitness_functions: (FitnessFunction, FitnessFunction, FitnessFunction),
}

fn fitness<N: Clone + Default, E: Clone + Default>(fitness_functions: (FitnessFunction,
                                                                       FitnessFunction,
                                                                       FitnessFunction),
                                                   goal: &Goal<N, E>,
                                                   ind: &EdgeOpsGenome)
                                                   -> MultiObjective3<f32> {
    let g = ind.to_graph(MAX_DEGREE);
    MultiObjective3::from((apply_fitness_function(fitness_functions.0, goal, &g),
                           apply_fitness_function(fitness_functions.1, goal, &g),
                           apply_fitness_function(fitness_functions.2, goal, &g)))

}

impl<N:Clone+Sync+Default,E:Clone+Sync+Default> FitnessEval<EdgeOpsGenome, MultiObjective3<f32>> for MyEval<N,E> {
    fn fitness(&mut self, pop: &[EdgeOpsGenome]) -> Vec<MultiObjective3<f32>> {
        let pool = &mut self.pool;
        let goal = &self.goal;

        let fitness_functions = self.fitness_functions;

        crossbeam::scope(|scope| {
            pool.map(scope, pop, |ind| fitness(fitness_functions, goal, ind))
                .collect()
        })
    }
}

#[derive(Debug)]
struct Stat<T: Debug> {
    min: T,
    max: T,
    avg: T,
}

#[allow(non_snake_case)]
fn main() {
    let ncpus = num_cpus::get();
    println!("Using {} CPUs", ncpus);
    let matches = App::new("nsga2_edge")
                      .arg(Arg::with_name("NGEN")
                               .long("ngen")
                               .help("Number of generations to run")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("OPS")
                               .long("ops")
                               .help("Edge operation and weight specification, e.g. \
                                      Dup:1,Split:3,Parent:2")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("GRAPH")
                               .long("graph")
                               .help("SGF graph file")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("MU")
                               .long("mu")
                               .help("Population size")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("LAMBDA")
                               .long("lambda")
                               .help("Size of offspring population")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("K")
                               .long("k")
                               .help("Tournament selection (default: 2)")
                               .takes_value(true))
                      .arg(Arg::with_name("SEED")
                               .long("seed")
                               .help("Seed value for Rng")
                               .takes_value(true))
                      .arg(Arg::with_name("PMUT")
                               .long("pmut")
                               .help("Probability for element mutation")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("ILEN")
                               .long("ilen")
                               .help("Initial genome length (random range)")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("OBJECTIVES")
                               .long("objectives")
                               .help("Specify 3 objective functions comma separated (null, cc, \
                                      scc, nm, td), e.g. cc,nm,td")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("THRESHOLD")
                               .long("threshold")
                               .help("Abort if for all fitness[i] <= value[i] (default: 0,0,0)")
                               .takes_value(true)
                               .required(false))
                      .get_matches();

    // number of generations
    let NGEN: usize = FromStr::from_str(matches.value_of("NGEN").unwrap()).unwrap();
    println!("NGEN: {}", NGEN);

    // size of population
    let MU: usize = FromStr::from_str(matches.value_of("MU").unwrap()).unwrap();
    println!("MU: {}", MU);

    // size of offspring population
    let LAMBDA: usize = FromStr::from_str(matches.value_of("LAMBDA").unwrap()).unwrap();
    println!("LAMBDA: {}", LAMBDA);

    // tournament selection
    let K: usize = FromStr::from_str(matches.value_of("K").unwrap_or("2")).unwrap();
    assert!(K > 0);
    println!("K: {}", K);

    let seed: Vec<u64>;
    if let Some(seed_str) = matches.value_of("SEED") {
        seed = seed_str.split(",").map(|s| FromStr::from_str(s).unwrap()).collect();
    } else {
        println!("Use OsRng to generate seed..");
        let mut rng = OsRng::new().unwrap();
        seed = (0..2).map(|_| rng.next_u64()).collect();
    }
    println!("SEED: {:?}", seed);

    let PMUT: f32 = FromStr::from_str(matches.value_of("PMUT").unwrap()).unwrap();
    println!("PMUT: {:?}", PMUT);

    let ilen_str = matches.value_of("ILEN").unwrap();
    let ilen: Vec<usize> = ilen_str.split(",").map(|s| FromStr::from_str(s).unwrap()).collect();
    assert!(ilen.len() == 1 || ilen.len() == 2);

    let ilen_from = ilen[0];
    let ilen_to = if ilen.len() == 1 {
        ilen[0]
    } else {
        ilen[1]
    };
    assert!(ilen_from <= ilen_to);

    // read objective functions
    let mut objectives_arr: Vec<FitnessFunction> = Vec::new();
    for s in matches.value_of("OBJECTIVES").unwrap().split(",") {
        objectives_arr.push(match s {
            "null" => {
                FitnessFunction::Null
            }
            "cc" => {
                FitnessFunction::ConnectedComponents
            }
            "scc" => {
                FitnessFunction::StronglyConnectedComponents
            }
            "nm" => {
                FitnessFunction::NeighborMatching
            }
            "td" => {
                FitnessFunction::TriadicDistance
            }
            _ => {
                panic!("Invalid objective function");
            }
        });
    }

    while objectives_arr.len() < 3 {
        objectives_arr.push(FitnessFunction::Null);
    }

    if objectives_arr.len() > 3 {
        panic!("Max 3 objectives allowed");
    }

    println!("objectives={:?}", objectives_arr);

    // read objective functions
    let mut threshold_arr: Vec<f32> = Vec::new();
    for s in matches.value_of("THRESHOLD").unwrap_or("").split(",") {
        let value: f32 = FromStr::from_str(s).unwrap();
        threshold_arr.push(value);
    }

    while threshold_arr.len() < 3 {
        threshold_arr.push(0.0);
    }

    if threshold_arr.len() > 3 {
        panic!("Max 3 threshold values allowed");
    }

    // read graph
    let graph_file = matches.value_of("GRAPH").unwrap();
    println!("Using graph file: {}", graph_file);
    let graph = {
        let mut f = BufReader::new(File::open(graph_file).unwrap());
        let graph: Graph<Unweighted, Unweighted, Directed> = PetgraphReader::from_sgf(&mut f);
        graph
    };

    // Parse weighted operation choice from command line
    let ops = Toolbox::parse_weighted_op_choice_list(matches.value_of("OPS").unwrap()).unwrap();
    println!("edge ops: {:?}", ops);
    let mut toolbox = Toolbox::new(Probability::new(PMUT), &ops[..]);

    let mut evaluator = MyEval {
        goal: Goal::new(OptDenseDigraph::from(graph)),
        pool: Pool::new(ncpus),
        fitness_functions: (objectives_arr[0], objectives_arr[1], objectives_arr[2]),
    };

    assert!(seed.len() == 2);
    let mut rng: PcgRng = SeedableRng::from_seed([seed[0], seed[1]]);

    // create initial random population
    let initial_population: Vec<EdgeOpsGenome> = (0..MU)
                                                     .map(|_| {
                                                         let len = if ilen_from == ilen_to {
                                                             ilen_from
                                                         } else {
                                                             Range::new(ilen_from, ilen_to)
                                                                 .ind_sample(&mut rng)
                                                         };
                                                         toolbox.random_genome(&mut rng, len)
                                                     })
                                                     .collect();

    // evaluate fitness
    let fitness: Vec<_> = evaluator.fitness(&initial_population[..]);
    assert!(fitness.len() == initial_population.len());

    let mut pop = initial_population;
    let mut fit = fitness;

    for i in 0..NGEN {
        println!("Iteration: {}", i);
        let (new_pop, new_fit) = nsga2::iterate(&mut rng,
                                                pop,
                                                fit,
                                                &mut evaluator,
                                                MU,
                                                LAMBDA,
                                                K,
                                                &mut toolbox);
        pop = new_pop;
        fit = new_fit;

        let stats: Vec<Stat<f32>> = [0, 1, 2]
                                        .iter()
                                        .map(|&i| {
                                            let min = fit.iter().fold(f32::INFINITY, |acc, f| {
                                                let x = f.objectives[i];
                                                if x < acc {
                                                    x
                                                } else {
                                                    acc
                                                }
                                            });
                                            let max = fit.iter().fold(-f32::INFINITY, |acc, f| {
                                                let x = f.objectives[i];
                                                if x > acc {
                                                    x
                                                } else {
                                                    acc
                                                }
                                            });
                                            let sum = fit.iter()
                                                         .fold(0.0, |acc, f| acc + f.objectives[i]);
                                            Stat {
                                                min: min,
                                                max: max,
                                                avg: sum / fit.len() as f32,
                                            }
                                        })
                                        .collect();
        println!("stats: {:?}", stats);
        let mut found_optimum = false;
        for f in fit.iter() {
            if f.objectives[0] <= threshold_arr[0] && f.objectives[1] <= threshold_arr[1] &&
               f.objectives[2] <= threshold_arr[2] {
                found_optimum = true;
                break;
            }
        }
        if found_optimum {
            println!("Found premature optimum in Iteration {}", i);
            break;
        }
    }
    println!("===========================================================");

    // finally evaluate rank and crowding distance (using select()).
    let rank_dist = nsga2::select(&fit[..], MU);
    assert!(rank_dist.len() == MU);

    let mut j = 0;

    for rd in rank_dist.iter() {
        if rd.rank == 0 {
            println!("-------------------------------------------");
            println!("rd: {:?}", rd);
            println!("fitness: {:?}", fit[rd.idx]);
            println!("genome: {:?}", pop[rd.idx]);

            // if fit[rd.idx].objectives[0] < 1.0 {
            draw_graph((&pop[rd.idx].to_graph(MAX_DEGREE)).ref_graph(),
                       // XXX: name
                       &format!("nsga2edgeops_g{}_f{}_i{}.svg",
                                NGEN,
                                fit[rd.idx].objectives[1] as usize,
                                j));
            j += 1;
            // }
        }
    }
}
<|endoftext|>


<commit_before>use libc::{uint32_t, c_void};

use sys::timer as ll;

pub fn get_ticks() -> u32 {
    unsafe { ll::SDL_GetTicks() }
}

pub fn get_performance_counter() -> u64 {
    unsafe { ll::SDL_GetPerformanceCounter() }
}

pub fn get_performance_frequency() -> u64 {
    unsafe { ll::SDL_GetPerformanceFrequency() }
}

pub fn delay(ms: u32) {
    unsafe { ll::SDL_Delay(ms) }
}

pub type TimerCallback = extern "C" fn (interval: uint32_t, param: *const c_void) -> u32;

pub struct Timer {
    delay: usize,
    raw: ll::SDL_TimerID,
    callback: TimerCallback,
    param: *const c_void,
    remove_on_drop: bool,
}

impl Timer {
    pub fn new(delay: usize, callback: TimerCallback, param: *const c_void, remove_on_drop: bool) -> Timer {
        Timer { delay: delay, raw: 0, callback: callback, param: param, remove_on_drop: remove_on_drop }
    }

    pub fn start(&mut self) {
        unsafe {
            let timer_id = ll::SDL_AddTimer(self.delay as u32, Some(self.callback), self.param);
            self.raw = timer_id;
        }
    }

    pub fn remove(&mut self) -> bool {
        let ret = unsafe { ll::SDL_RemoveTimer(self.raw) };
        if self.raw != 0 {
            self.raw = 0
        }
        ret == 1
    }
}

#[unsafe_destructor]
impl Drop for Timer {
    fn drop(&mut self) {
        if self.remove_on_drop {
            let ret = unsafe { ll::SDL_RemoveTimer(self.raw) };
            if ret != 1 {
                println!("error dropping timer {}, maybe already removed.", self.raw);
            }
        }
    }
}

#[cfg(test)]
extern "C" fn test_timer_1_callback(_interval: uint32_t, param: *const c_void) -> uint32_t {
    use std::sync::{Arc, Mutex};
    use std::mem;

    let locked_num: &Arc<Mutex<u32>> = unsafe { mem::transmute(param) };
    let mut num = locked_num.lock().unwrap();
    *num = *num + 1;
    10
}

#[test]
fn test_timer_1() {
    use std::sync::{Arc, Mutex};
    use std::mem;

    let local_num: Arc<Mutex<u32>> = Arc::new(Mutex::new(0));
    let timer_num = local_num.clone();
    {
        let param = unsafe { mem::transmute(&timer_num) };
        let mut timer = Timer::new(10, test_timer_1_callback, param, true);
        timer.start();
        delay(100);
        let num = local_num.lock().unwrap();
        assert!(*num == 9);
    }

    // Check that timer has stopped
    delay(100);
    let num = local_num.lock().unwrap();
    assert!(*num == 9);
}

#[cfg(test)]
extern "C" fn test_timer_2_callback(_interval: uint32_t, _param: *const c_void) -> uint32_t {
    0
}

#[test]
fn test_timer_2() {
    use std::ptr;

    let _ = {
        let mut timer = Timer::new(1000, test_timer_2_callback, ptr::null(), true);
        timer.start();
        timer
    };
    delay(200);
    delay(200);
    delay(200);
    delay(200);
    delay(200);
    delay(200);
}
<commit_msg>Adjust test_timer_1 to be more robust<commit_after>use libc::{uint32_t, c_void};

use sys::timer as ll;

pub fn get_ticks() -> u32 {
    unsafe { ll::SDL_GetTicks() }
}

pub fn get_performance_counter() -> u64 {
    unsafe { ll::SDL_GetPerformanceCounter() }
}

pub fn get_performance_frequency() -> u64 {
    unsafe { ll::SDL_GetPerformanceFrequency() }
}

pub fn delay(ms: u32) {
    unsafe { ll::SDL_Delay(ms) }
}

pub type TimerCallback = extern "C" fn (interval: uint32_t, param: *const c_void) -> u32;

pub struct Timer {
    delay: usize,
    raw: ll::SDL_TimerID,
    callback: TimerCallback,
    param: *const c_void,
    remove_on_drop: bool,
}

impl Timer {
    pub fn new(delay: usize, callback: TimerCallback, param: *const c_void, remove_on_drop: bool) -> Timer {
        Timer { delay: delay, raw: 0, callback: callback, param: param, remove_on_drop: remove_on_drop }
    }

    pub fn start(&mut self) {
        unsafe {
            let timer_id = ll::SDL_AddTimer(self.delay as u32, Some(self.callback), self.param);
            self.raw = timer_id;
        }
    }

    pub fn remove(&mut self) -> bool {
        let ret = unsafe { ll::SDL_RemoveTimer(self.raw) };
        if self.raw != 0 {
            self.raw = 0
        }
        ret == 1
    }
}

#[unsafe_destructor]
impl Drop for Timer {
    fn drop(&mut self) {
        if self.remove_on_drop {
            let ret = unsafe { ll::SDL_RemoveTimer(self.raw) };
            if ret != 1 {
                println!("error dropping timer {}, maybe already removed.", self.raw);
            }
        }
    }
}

#[cfg(test)]
extern "C" fn test_timer_1_callback(_interval: uint32_t, param: *const c_void) -> uint32_t {
    use std::sync::{Arc, Mutex};
    use std::mem;

    let locked_num: &Arc<Mutex<u32>> = unsafe { mem::transmute(param) };
    let mut num = locked_num.lock().unwrap();
    *num = *num + 1;
    10
}

#[test]
fn test_timer_1() {
    use std::sync::{Arc, Mutex};
    use std::mem;

    let local_num: Arc<Mutex<u32>> = Arc::new(Mutex::new(0));
    let timer_num = local_num.clone();

    let final_num = {
        let param = unsafe { mem::transmute(&timer_num) };
        let mut timer = Timer::new(10, test_timer_1_callback, param, true);
        timer.start();
        delay(500);
        let num = local_num.lock().unwrap();
        assert!(*num > 0);

        *num
    };

    // Check that timer has stopped
    delay(500);
    let num = local_num.lock().unwrap();
    assert!(*num == final_num);
}

#[cfg(test)]
extern "C" fn test_timer_2_callback(_interval: uint32_t, _param: *const c_void) -> uint32_t {
    0
}

#[test]
fn test_timer_2() {
    use std::ptr;

    let _ = {
        let mut timer = Timer::new(1000, test_timer_2_callback, ptr::null(), true);
        timer.start();
        timer
    };
    delay(200);
    delay(200);
    delay(200);
    delay(200);
    delay(200);
    delay(200);
}
<|endoftext|>


<commit_before>use std::mem;
use libc::{uint32_t, c_void};
use std::raw;
use std::kinds::marker::ContravariantLifetime;

#[allow(non_camel_case_types)]
pub mod ll {
    use libc::{uint32_t, uint64_t, c_void, c_int};

    //SDL_timer.h
    pub type SDL_TimerCallback =
        ::std::option::Option<extern "C" fn(arg1: uint32_t, arg2: *const c_void)
                                            -> uint32_t>;
    pub type SDL_TimerID = c_int;
    extern "C" {
        pub fn SDL_GetTicks() -> uint32_t;
        pub fn SDL_GetPerformanceCounter() -> uint64_t;
        pub fn SDL_GetPerformanceFrequency() -> uint64_t;
        pub fn SDL_Delay(ms: uint32_t);

        pub fn SDL_AddTimer(interval: uint32_t, callback: SDL_TimerCallback,
                            param: *const c_void) -> SDL_TimerID;
        pub fn SDL_RemoveTimer(id: SDL_TimerID) -> c_int;
    }
}

pub fn get_ticks() -> uint {
    unsafe { ll::SDL_GetTicks() as uint }
}

pub fn get_performance_counter() -> u64 {
    unsafe { ll::SDL_GetPerformanceCounter() }
}

pub fn get_performance_frequency() -> u64 {
    unsafe { ll::SDL_GetPerformanceFrequency() }
}

pub fn delay(ms: uint) {
    unsafe { ll::SDL_Delay(ms as u32) }
}

pub struct Timer<'a> {
    delay: uint,
    raw: ll::SDL_TimerID,
    closure: raw::Closure,
    remove_on_drop: bool,
    lifetime: ContravariantLifetime<'a>,
}

impl<'a> Timer<'a> {
    pub fn new<'b>(delay: uint, callback: ||: 'b -> uint, remove_on_drop: bool) -> Timer<'b> {
        unsafe {
            let c_param = mem::transmute::<_, raw::Closure>(callback);
            Timer { delay: delay, raw: 0, closure: c_param, remove_on_drop: remove_on_drop, lifetime: ContravariantLifetime }
        }
    }

    pub fn start(&mut self) {
        unsafe {
            let timer_id = ll::SDL_AddTimer(self.delay as u32, Some(c_timer_callback), mem::transmute(&self.closure));
            self.raw = timer_id;
        }
    }

    pub fn remove(&mut self) -> bool {
        let ret = unsafe { ll::SDL_RemoveTimer(self.raw) };
        if self.raw != 0 {
            self.raw = 0
        }
        ret == 1
    }
}

#[unsafe_destructor]
impl<'a> Drop for Timer<'a> {
    fn drop(&mut self) {
        if self.remove_on_drop {
            let ret = unsafe { ll::SDL_RemoveTimer(self.raw) };
            if ret != 1 {
                println!("error dropping timer {}, maybe already removed.", self.raw);
            }
        }
    }
}

extern "C" fn c_timer_callback(_interval: uint32_t, param: *const c_void) -> uint32_t {
    let f : &mut || -> uint = unsafe { mem::transmute(param) };
    (*f)() as uint32_t
}
<commit_msg>Use already present lifetime<commit_after>use std::mem;
use libc::{uint32_t, c_void};
use std::raw;
use std::kinds::marker::ContravariantLifetime;

#[allow(non_camel_case_types)]
pub mod ll {
    use libc::{uint32_t, uint64_t, c_void, c_int};

    //SDL_timer.h
    pub type SDL_TimerCallback =
        ::std::option::Option<extern "C" fn(arg1: uint32_t, arg2: *const c_void)
                                            -> uint32_t>;
    pub type SDL_TimerID = c_int;
    extern "C" {
        pub fn SDL_GetTicks() -> uint32_t;
        pub fn SDL_GetPerformanceCounter() -> uint64_t;
        pub fn SDL_GetPerformanceFrequency() -> uint64_t;
        pub fn SDL_Delay(ms: uint32_t);

        pub fn SDL_AddTimer(interval: uint32_t, callback: SDL_TimerCallback,
                            param: *const c_void) -> SDL_TimerID;
        pub fn SDL_RemoveTimer(id: SDL_TimerID) -> c_int;
    }
}

pub fn get_ticks() -> uint {
    unsafe { ll::SDL_GetTicks() as uint }
}

pub fn get_performance_counter() -> u64 {
    unsafe { ll::SDL_GetPerformanceCounter() }
}

pub fn get_performance_frequency() -> u64 {
    unsafe { ll::SDL_GetPerformanceFrequency() }
}

pub fn delay(ms: uint) {
    unsafe { ll::SDL_Delay(ms as u32) }
}

pub struct Timer<'a> {
    delay: uint,
    raw: ll::SDL_TimerID,
    closure: raw::Closure,
    remove_on_drop: bool,
    lifetime: ContravariantLifetime<'a>,
}

impl<'a> Timer<'a> {
    pub fn new(delay: uint, callback: ||: 'a -> uint, remove_on_drop: bool) -> Timer<'a> {
        unsafe {
            let c_param = mem::transmute::<_, raw::Closure>(callback);
            Timer { delay: delay, raw: 0, closure: c_param, remove_on_drop: remove_on_drop, lifetime: ContravariantLifetime }
        }
    }

    pub fn start(&mut self) {
        unsafe {
            let timer_id = ll::SDL_AddTimer(self.delay as u32, Some(c_timer_callback), mem::transmute(&self.closure));
            self.raw = timer_id;
        }
    }

    pub fn remove(&mut self) -> bool {
        let ret = unsafe { ll::SDL_RemoveTimer(self.raw) };
        if self.raw != 0 {
            self.raw = 0
        }
        ret == 1
    }
}

#[unsafe_destructor]
impl<'a> Drop for Timer<'a> {
    fn drop(&mut self) {
        if self.remove_on_drop {
            let ret = unsafe { ll::SDL_RemoveTimer(self.raw) };
            if ret != 1 {
                println!("error dropping timer {}, maybe already removed.", self.raw);
            }
        }
    }
}

extern "C" fn c_timer_callback(_interval: uint32_t, param: *const c_void) -> uint32_t {
    let f : &mut || -> uint = unsafe { mem::transmute(param) };
    (*f)() as uint32_t
}
<|endoftext|>


<commit_before>use std::fmt;
use std::io;
use std::net::{SocketAddr, TcpListener as StdTcpListener};
use std::time::{Duration, Instant};

use futures::{Async, Future, Poll, Stream};
use tokio_reactor::Handle;
use tokio_tcp::TcpListener;
use tokio_timer::Delay;

use self::addr_stream::AddrStream;

/// A stream of connections from binding to an address.
#[must_use = "streams do nothing unless polled"]
pub struct AddrIncoming {
    addr: SocketAddr,
    listener: TcpListener,
    sleep_on_errors: bool,
    tcp_keepalive_timeout: Option<Duration>,
    tcp_nodelay: bool,
    timeout: Option<Delay>,
}

impl AddrIncoming {
    pub(super) fn new(addr: &SocketAddr, handle: Option<&Handle>) -> ::Result<AddrIncoming> {
        let listener = if let Some(handle) = handle {
            let std_listener = StdTcpListener::bind(addr)
                .map_err(::Error::new_listen)?;
            TcpListener::from_std(std_listener, handle)
                .map_err(::Error::new_listen)?
        } else {
            TcpListener::bind(addr).map_err(::Error::new_listen)?
        };

        let addr = listener.local_addr().map_err(::Error::new_listen)?;

        Ok(AddrIncoming {
            addr: addr,
            listener: listener,
            sleep_on_errors: true,
            tcp_keepalive_timeout: None,
            tcp_nodelay: false,
            timeout: None,
        })
    }

    /// Get the local address bound to this listener.
    pub fn local_addr(&self) -> SocketAddr {
        self.addr
    }

    /// Set whether TCP keepalive messages are enabled on accepted connections.
    ///
    /// If `None` is specified, keepalive is disabled, otherwise the duration
    /// specified will be the time to remain idle before sending TCP keepalive
    /// probes.
    pub fn set_keepalive(&mut self, keepalive: Option<Duration>) -> &mut Self {
        self.tcp_keepalive_timeout = keepalive;
        self
    }

    /// Set the value of `TCP_NODELAY` option for accepted connections.
    pub fn set_nodelay(&mut self, enabled: bool) -> &mut Self {
        self.tcp_nodelay = enabled;
        self
    }

    /// Set whether to sleep on accept errors.
    ///
    /// A possible scenario is that the process has hit the max open files
    /// allowed, and so trying to accept a new connection will fail with
    /// `EMFILE`. In some cases, it's preferable to just wait for some time, if
    /// the application will likely close some files (or connections), and try
    /// to accept the connection again. If this option is `true`, the error
    /// will be logged at the `error` level, since it is still a big deal,
    /// and then the listener will sleep for 1 second.
    ///
    /// In other cases, hitting the max open files should be treat similarly
    /// to being out-of-memory, and simply error (and shutdown). Setting
    /// this option to `false` will allow that.
    ///
    /// Default is `true`.
    pub fn set_sleep_on_errors(&mut self, val: bool) {
        self.sleep_on_errors = val;
    }
}

impl Stream for AddrIncoming {
    // currently unnameable...
    type Item = AddrStream;
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        // Check if a previous timeout is active that was set by IO errors.
        if let Some(ref mut to) = self.timeout {
            match to.poll() {
                Ok(Async::Ready(())) => {}
                Ok(Async::NotReady) => return Ok(Async::NotReady),
                Err(err) => {
                    error!("sleep timer error: {}", err);
                }
            }
        }
        self.timeout = None;
        loop {
            match self.listener.poll_accept() {
                Ok(Async::Ready((socket, addr))) => {
                    if let Some(dur) = self.tcp_keepalive_timeout {
                        if let Err(e) = socket.set_keepalive(Some(dur)) {
                            trace!("error trying to set TCP keepalive: {}", e);
                        }
                    }
                    if let Err(e) = socket.set_nodelay(self.tcp_nodelay) {
                        trace!("error trying to set TCP nodelay: {}", e);
                    }
                    return Ok(Async::Ready(Some(AddrStream::new(socket, addr))));
                },
                Ok(Async::NotReady) => return Ok(Async::NotReady),
                Err(e) => {
                    if self.sleep_on_errors {
                        // Connection errors can be ignored directly, continue by
                        // accepting the next request.
                        if is_connection_error(&e) {
                            debug!("accepted connection already errored: {}", e);
                            continue;
                        }
                        // Sleep 1s.
                        let delay = Instant::now() + Duration::from_secs(1);
                        let mut timeout = Delay::new(delay);

                        match timeout.poll() {
                            Ok(Async::Ready(())) => {
                                // Wow, it's been a second already? Ok then...
                                error!("accept error: {}", e);
                                continue
                            },
                            Ok(Async::NotReady) => {
                                error!("accept error: {}", e);
                                self.timeout = Some(timeout);
                                return Ok(Async::NotReady);
                            },
                            Err(timer_err) => {
                                error!("couldn't sleep on error, timer error: {}", timer_err);
                                return Err(e);
                            }
                        }
                    } else {
                        return Err(e);
                    }
                },
            }
        }
    }
}

/// This function defines errors that are per-connection. Which basically
/// means that if we get this error from `accept()` system call it means
/// next connection might be ready to be accepted.
///
/// All other errors will incur a timeout before next `accept()` is performed.
/// The timeout is useful to handle resource exhaustion errors like ENFILE
/// and EMFILE. Otherwise, could enter into tight loop.
fn is_connection_error(e: &io::Error) -> bool {
    e.kind() == io::ErrorKind::ConnectionRefused ||
    e.kind() == io::ErrorKind::ConnectionAborted ||
    e.kind() == io::ErrorKind::ConnectionReset
}

impl fmt::Debug for AddrIncoming {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("AddrIncoming")
            .field("addr", &self.addr)
            .field("sleep_on_errors", &self.sleep_on_errors)
            .field("tcp_keepalive_timeout", &self.tcp_keepalive_timeout)
            .field("tcp_nodelay", &self.tcp_nodelay)
            .finish()
    }
}

mod addr_stream {
    use std::io::{self, Read, Write};
    use std::net::SocketAddr;
    use bytes::{Buf, BufMut};
    use futures::Poll;
    use tokio_tcp::TcpStream;
    use tokio_io::{AsyncRead, AsyncWrite};

    #[derive(Debug)]
    pub struct AddrStream {
        inner: TcpStream,
        pub(super) remote_addr: SocketAddr,
    }

    impl AddrStream {
        pub(super) fn new(tcp: TcpStream, addr: SocketAddr) -> AddrStream {
            AddrStream {
                inner: tcp,
                remote_addr: addr,
            }
        }
    }

    impl Read for AddrStream {
        #[inline]
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            self.inner.read(buf)
        }
    }

    impl Write for AddrStream {
        #[inline]
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            self.inner.write(buf)
        }

        #[inline]
        fn flush(&mut self ) -> io::Result<()> {
            self.inner.flush()
        }
    }

    impl AsyncRead for AddrStream {
        #[inline]
        unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
            self.inner.prepare_uninitialized_buffer(buf)
        }

        #[inline]
        fn read_buf<B: BufMut>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
            self.inner.read_buf(buf)
        }
    }

    impl AsyncWrite for AddrStream {
        #[inline]
        fn shutdown(&mut self) -> Poll<(), io::Error> {
            AsyncWrite::shutdown(&mut self.inner)
        }

        #[inline]
        fn write_buf<B: Buf>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
            self.inner.write_buf(buf)
        }
    }
}
<commit_msg>perf(server): skip noop AddrStream inner flush<commit_after>use std::fmt;
use std::io;
use std::net::{SocketAddr, TcpListener as StdTcpListener};
use std::time::{Duration, Instant};

use futures::{Async, Future, Poll, Stream};
use tokio_reactor::Handle;
use tokio_tcp::TcpListener;
use tokio_timer::Delay;

use self::addr_stream::AddrStream;

/// A stream of connections from binding to an address.
#[must_use = "streams do nothing unless polled"]
pub struct AddrIncoming {
    addr: SocketAddr,
    listener: TcpListener,
    sleep_on_errors: bool,
    tcp_keepalive_timeout: Option<Duration>,
    tcp_nodelay: bool,
    timeout: Option<Delay>,
}

impl AddrIncoming {
    pub(super) fn new(addr: &SocketAddr, handle: Option<&Handle>) -> ::Result<AddrIncoming> {
        let listener = if let Some(handle) = handle {
            let std_listener = StdTcpListener::bind(addr)
                .map_err(::Error::new_listen)?;
            TcpListener::from_std(std_listener, handle)
                .map_err(::Error::new_listen)?
        } else {
            TcpListener::bind(addr).map_err(::Error::new_listen)?
        };

        let addr = listener.local_addr().map_err(::Error::new_listen)?;

        Ok(AddrIncoming {
            addr: addr,
            listener: listener,
            sleep_on_errors: true,
            tcp_keepalive_timeout: None,
            tcp_nodelay: false,
            timeout: None,
        })
    }

    /// Get the local address bound to this listener.
    pub fn local_addr(&self) -> SocketAddr {
        self.addr
    }

    /// Set whether TCP keepalive messages are enabled on accepted connections.
    ///
    /// If `None` is specified, keepalive is disabled, otherwise the duration
    /// specified will be the time to remain idle before sending TCP keepalive
    /// probes.
    pub fn set_keepalive(&mut self, keepalive: Option<Duration>) -> &mut Self {
        self.tcp_keepalive_timeout = keepalive;
        self
    }

    /// Set the value of `TCP_NODELAY` option for accepted connections.
    pub fn set_nodelay(&mut self, enabled: bool) -> &mut Self {
        self.tcp_nodelay = enabled;
        self
    }

    /// Set whether to sleep on accept errors.
    ///
    /// A possible scenario is that the process has hit the max open files
    /// allowed, and so trying to accept a new connection will fail with
    /// `EMFILE`. In some cases, it's preferable to just wait for some time, if
    /// the application will likely close some files (or connections), and try
    /// to accept the connection again. If this option is `true`, the error
    /// will be logged at the `error` level, since it is still a big deal,
    /// and then the listener will sleep for 1 second.
    ///
    /// In other cases, hitting the max open files should be treat similarly
    /// to being out-of-memory, and simply error (and shutdown). Setting
    /// this option to `false` will allow that.
    ///
    /// Default is `true`.
    pub fn set_sleep_on_errors(&mut self, val: bool) {
        self.sleep_on_errors = val;
    }
}

impl Stream for AddrIncoming {
    // currently unnameable...
    type Item = AddrStream;
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        // Check if a previous timeout is active that was set by IO errors.
        if let Some(ref mut to) = self.timeout {
            match to.poll() {
                Ok(Async::Ready(())) => {}
                Ok(Async::NotReady) => return Ok(Async::NotReady),
                Err(err) => {
                    error!("sleep timer error: {}", err);
                }
            }
        }
        self.timeout = None;
        loop {
            match self.listener.poll_accept() {
                Ok(Async::Ready((socket, addr))) => {
                    if let Some(dur) = self.tcp_keepalive_timeout {
                        if let Err(e) = socket.set_keepalive(Some(dur)) {
                            trace!("error trying to set TCP keepalive: {}", e);
                        }
                    }
                    if let Err(e) = socket.set_nodelay(self.tcp_nodelay) {
                        trace!("error trying to set TCP nodelay: {}", e);
                    }
                    return Ok(Async::Ready(Some(AddrStream::new(socket, addr))));
                },
                Ok(Async::NotReady) => return Ok(Async::NotReady),
                Err(e) => {
                    if self.sleep_on_errors {
                        // Connection errors can be ignored directly, continue by
                        // accepting the next request.
                        if is_connection_error(&e) {
                            debug!("accepted connection already errored: {}", e);
                            continue;
                        }
                        // Sleep 1s.
                        let delay = Instant::now() + Duration::from_secs(1);
                        let mut timeout = Delay::new(delay);

                        match timeout.poll() {
                            Ok(Async::Ready(())) => {
                                // Wow, it's been a second already? Ok then...
                                error!("accept error: {}", e);
                                continue
                            },
                            Ok(Async::NotReady) => {
                                error!("accept error: {}", e);
                                self.timeout = Some(timeout);
                                return Ok(Async::NotReady);
                            },
                            Err(timer_err) => {
                                error!("couldn't sleep on error, timer error: {}", timer_err);
                                return Err(e);
                            }
                        }
                    } else {
                        return Err(e);
                    }
                },
            }
        }
    }
}

/// This function defines errors that are per-connection. Which basically
/// means that if we get this error from `accept()` system call it means
/// next connection might be ready to be accepted.
///
/// All other errors will incur a timeout before next `accept()` is performed.
/// The timeout is useful to handle resource exhaustion errors like ENFILE
/// and EMFILE. Otherwise, could enter into tight loop.
fn is_connection_error(e: &io::Error) -> bool {
    e.kind() == io::ErrorKind::ConnectionRefused ||
    e.kind() == io::ErrorKind::ConnectionAborted ||
    e.kind() == io::ErrorKind::ConnectionReset
}

impl fmt::Debug for AddrIncoming {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("AddrIncoming")
            .field("addr", &self.addr)
            .field("sleep_on_errors", &self.sleep_on_errors)
            .field("tcp_keepalive_timeout", &self.tcp_keepalive_timeout)
            .field("tcp_nodelay", &self.tcp_nodelay)
            .finish()
    }
}

mod addr_stream {
    use std::io::{self, Read, Write};
    use std::net::SocketAddr;
    use bytes::{Buf, BufMut};
    use futures::Poll;
    use tokio_tcp::TcpStream;
    use tokio_io::{AsyncRead, AsyncWrite};

    #[derive(Debug)]
    pub struct AddrStream {
        inner: TcpStream,
        pub(super) remote_addr: SocketAddr,
    }

    impl AddrStream {
        pub(super) fn new(tcp: TcpStream, addr: SocketAddr) -> AddrStream {
            AddrStream {
                inner: tcp,
                remote_addr: addr,
            }
        }
    }

    impl Read for AddrStream {
        #[inline]
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            self.inner.read(buf)
        }
    }

    impl Write for AddrStream {
        #[inline]
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            self.inner.write(buf)
        }

        #[inline]
        fn flush(&mut self) -> io::Result<()> {
            // TcpStream::flush is a noop, so skip calling it...
            Ok(())
        }
    }

    impl AsyncRead for AddrStream {
        #[inline]
        unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
            self.inner.prepare_uninitialized_buffer(buf)
        }

        #[inline]
        fn read_buf<B: BufMut>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
            self.inner.read_buf(buf)
        }
    }

    impl AsyncWrite for AddrStream {
        #[inline]
        fn shutdown(&mut self) -> Poll<(), io::Error> {
            AsyncWrite::shutdown(&mut self.inner)
        }

        #[inline]
        fn write_buf<B: Buf>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
            self.inner.write_buf(buf)
        }
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::path::PathBuf;

use quote::Tokens;
use walkdir::WalkDir;
use cfg::CargonautsConfig;

use ast::Routes;

const FLAG: &'static str = "used_cargonauts_asset_pipeline";

pub fn assets(cfg: &CargonautsConfig, routes: &Routes) -> Tokens {
    let handler = asset_handler(routes);

    let dir = {
        // First, check if they used an asset pipeline
        let used_asset_pipeline = env::args().find(|arg| arg.starts_with(FLAG));
        if let Some(arg) = used_asset_pipeline {
            // Check if they specified an unusual build directory
            let arg =  arg.trim_left_matches(FLAG);
            if arg.len() > 2 && arg.starts_with('[') && arg.ends_with(']') {
                let range = 1..(arg.len() - 1);
                cfg.project_root().join(&arg[range])
            } else {
                // If they didn't, assets are in OUT_DIR/cargonauts/assets
                let out_dir: PathBuf = env::var("OUT_DIR").expect("OUT_DIR not set").into();
                out_dir.join("cargonauts/assets")
            }
        } else {
            // If they didn't use a pipeline, find their normal assets folder
            cfg.assets()
        }
    };

    let assets = WalkDir::new(&dir).into_iter().filter_map(|entry| entry.ok().and_then(|entry| {
        if entry.file_type().is_file() {
            let path = entry.path();
            let file_path = path.to_string_lossy();
            let url_path = if entry.file_name().to_string_lossy().starts_with("_index") {
                path.strip_prefix(&dir).unwrap().parent().unwrap().to_string_lossy()
            } else {
                path.strip_prefix(&dir).unwrap().to_string_lossy()
            };
            Some(quote!({
                let service = ::cargonauts::routing::AssetHandler {
                    handler: #handler,
                    data: include_bytes!(#file_path),
                };
                routes.add(::cargonauts::server::Method::Get, String::from(#url_path), Box::new(service) as ::cargonauts::routing::Handler);
            }))
        } else { None }
    })).collect::<Vec<_>>();

    quote! {
        #(#assets)*
    }
}

fn asset_handler(routes: &Routes) -> Tokens {
    if let Some(handler) = routes.asset_handler.as_ref() {
        quote!(#handler)
    } else {
        quote!(::cargonauts::routing::default_asset_handler)
    }
}
<commit_msg>Use `index` as index name by default.<commit_after>use std::env;
use std::path::PathBuf;

use quote::Tokens;
use walkdir::WalkDir;
use cfg::CargonautsConfig;

use ast::Routes;

const FLAG: &'static str = "used_cargonauts_asset_pipeline";

pub fn assets(cfg: &CargonautsConfig, routes: &Routes) -> Tokens {
    let handler = asset_handler(routes);

    let dir = {
        // First, check if they used an asset pipeline
        let used_asset_pipeline = env::args().find(|arg| arg.starts_with(FLAG));
        if let Some(arg) = used_asset_pipeline {
            // Check if they specified an unusual build directory
            let arg =  arg.trim_left_matches(FLAG);
            if arg.len() > 2 && arg.starts_with('[') && arg.ends_with(']') {
                let range = 1..(arg.len() - 1);
                cfg.project_root().join(&arg[range])
            } else {
                // If they didn't, assets are in OUT_DIR/cargonauts/assets
                let out_dir: PathBuf = env::var("OUT_DIR").expect("OUT_DIR not set").into();
                out_dir.join("cargonauts/assets")
            }
        } else {
            // If they didn't use a pipeline, find their normal assets folder
            cfg.assets()
        }
    };

    let assets = WalkDir::new(&dir).into_iter().filter_map(|entry| entry.ok().and_then(|entry| {
        if entry.file_type().is_file() {
            let path = entry.path();
            let file_path = path.to_string_lossy();
            let url_path = if entry.file_name().to_string_lossy().starts_with("index") {
                path.strip_prefix(&dir).unwrap().parent().unwrap().to_string_lossy()
            } else {
                path.strip_prefix(&dir).unwrap().to_string_lossy()
            };
            Some(quote!({
                let service = ::cargonauts::routing::AssetHandler {
                    handler: #handler,
                    data: include_bytes!(#file_path),
                };
                routes.add(::cargonauts::server::Method::Get, String::from(#url_path), Box::new(service) as ::cargonauts::routing::Handler);
            }))
        } else { None }
    })).collect::<Vec<_>>();

    quote! {
        #(#assets)*
    }
}

fn asset_handler(routes: &Routes) -> Tokens {
    if let Some(handler) = routes.asset_handler.as_ref() {
        quote!(#handler)
    } else {
        quote!(::cargonauts::routing::default_asset_handler)
    }
}
<|endoftext|>


<commit_before>//
// Copyright:: Copyright (c) 2015 Chef Software, Inc.
// License:: Apache License, Version 2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

pub use errors;
use errors::{DeliveryError, Kind};
use std::fs::File;
use std::default::Default;
use utils::say::{say, sayln};
use rustc_serialize::Encodable;
use std::path::PathBuf;
use toml;
use utils::mkdir_recursive;
use std::io::prelude::*;
use utils::path_join_many::PathJoinMany;

#[derive(RustcEncodable, Clone)]
pub struct Config {
    pub server: Option<String>,
    pub api_port: Option<String>,
    pub user: Option<String>,
    pub enterprise: Option<String>,
    pub organization: Option<String>,
    pub project: Option<String>,
    pub git_port: Option<String>,
    pub pipeline: Option<String>
}

impl Default for Config {
    fn default() -> Config {
        Config{
            server: None,
            api_port: None,
            enterprise: None,
            organization: None,
            project: None,
            user: None,
            git_port: Some(String::from("8989")),
            pipeline: Some(String::from("master"))
        }
    }
}

macro_rules! config_accessor_for {
    ($name:ident, $set_name:ident, $err_msg:expr) => (
        impl Config {
            pub fn $name(&self) -> Result<String, DeliveryError> {
                match self.$name {
                    Some(ref v) => Ok(v.clone()),
                    None => Err(DeliveryError{ kind: Kind::MissingConfig, detail: Some(String::from($err_msg)) })
                }
            }

            pub fn $set_name(mut self, $name: &str) -> Config {
                if !$name.is_empty() {
                    self.$name = Some(String::from($name));
                }
                self
            }
        }
    )
}

config_accessor_for!(server, set_server, "Server not set; try --server or set it in your .toml config file");
config_accessor_for!(api_port, set_api_port, "API port not set; try --api-port or set it in your .toml config file");
config_accessor_for!(user, set_user, "User not set; try --user or set it in your .toml config file");
config_accessor_for!(enterprise, set_enterprise, "Enterprise not set; try --ent or set it in your .toml config file");
config_accessor_for!(organization, set_organization, "Organization not set; try --org or set it in your .toml config file");
config_accessor_for!(project, set_project, "Project not set; try --project or set it in your .toml config file");
config_accessor_for!(git_port, set_git_port, "Git Port not set; please set it in your .toml config file");
config_accessor_for!(pipeline, set_pipeline, "Pipeline not set; try --for or set it in your .toml config file");

impl Config {

    /// Return the host and port at which we can access the Delivery
    /// API. By default, we assume the use of HTTPS on the standard
    /// port `443`. Unless a port is specified in the configuration,
    /// we'll just return the server name; otherwise we append the
    /// port.
    pub fn api_host_and_port(&self) -> Result<String, DeliveryError> {
        let s = try!(self.server());
        return Ok(match self.api_port {
            Some(ref p) => format!("{}:{}", s, p),
            None    => s
        });
    }

    /// Returns the SSH URL to talk to Delivery's Git
    pub fn delivery_git_ssh_url(&self) -> Result<String, DeliveryError> {
        let s = try!(self.server());
        let host_and_port = match self.git_port {
            Some(ref p) => format!("{}:{}", s, p),
            None    => s // TODO: Currently we *always* have a git port
        };
        let u = try!(self.user());
        let e = try!(self.enterprise());
        let o = try!(self.organization());
        let p = try!(self.project());
        Ok(format!("ssh://{}@{}@{}/{}/{}/{}", u, e, host_and_port, e, o, p))
    }

    pub fn load_config(cwd: &PathBuf) -> Result<Config, DeliveryError> {
        let have_config = Config::have_dot_delivery_cli(cwd);
        match have_config.as_ref() {
            Some(path) => {
                let toml = try!(Config::read_file(path));
                match Config::parse_config(&toml) {
                    Ok(c) => return Ok(c),
                    Err(_) => return Ok(Default::default())
                }
            },
            None => return Ok(Default::default())
        }
    }

    pub fn write_file(&self, path: &PathBuf) -> Result<(), DeliveryError> {
        let write_dir = path.join_many(&[".delivery"]);
        if !write_dir.is_dir() {
            try!(mkdir_recursive(&write_dir));
        }
        let write_path = path.join_many(&[".delivery", "cli.toml"]);
        say("white", "Writing configuration to ");
        sayln("yellow", &format!("{}", write_path.display()));
        let mut f = try!(File::create(&write_path));
        let toml_string = toml::encode_str(self);
        sayln("magenta", "New configuration");
        sayln("magenta", "-----------------");
        say("white", &toml_string);
        try!(f.write_all(toml_string.as_bytes()));
        Ok(())
    }

    pub fn parse_config(toml: &str) -> Result<Config, DeliveryError> {
        let mut parser = toml::Parser::new(toml);
        match parser.parse() {
            Some(value) => { return Config::set_values_from_toml_table(value); },
            None => {
                return Err(DeliveryError{
                    kind: Kind::ConfigParse,
                    detail: Some(format!("Parse errors: {:?}", parser.errors))
                });
            }
        }
    }

    fn set_values_from_toml_table(table: toml::Table) -> Result<Config, DeliveryError> {
        let mut config: Config = Default::default();
        config.server = Config::stringify_values(table.get("server"));
        config.api_port = Config::stringify_values(table.get("api_port"));
        config.project = Config::stringify_values(table.get("project"));
        config.enterprise = Config::stringify_values(table.get("enterprise"));
        config.organization = Config::stringify_values(table.get("organization"));
        config.user = Config::stringify_values(table.get("user"));
        config.git_port = Config::stringify_values(table.get("git_port"));
        return Ok(config);
    }

    fn read_file(path: &PathBuf) -> Result<String, DeliveryError>  {
        let mut toml_file = try!(File::open(path));
        let mut toml = String::new();
        try!(toml_file.read_to_string(&mut toml));
        Ok(toml)
    }

    fn stringify_values(toml_value: Option<&toml::Value>) -> Option<String> {
        match toml_value {
            Some(value) => {
                let is_string = value.as_str();
                match is_string {
                    Some(vstr) => return Some(String::from(vstr)),
                    None => return None
                }
            },
            None => {
                return None;
            }
        }
    }

    fn check_dot_delivery_cli(path: PathBuf) -> Option<PathBuf> {
        let dot_git = path.join_many(&[".delivery", "cli.toml"]);
        debug!("Checking {}", dot_git.display());
        let is_file: Option<PathBuf> = if dot_git.is_file() {
            Some(dot_git)
        } else {
            None
        };
        is_file
    }

    fn have_dot_delivery_cli(orig_path: &PathBuf) -> Option<PathBuf> {
        let mut path = orig_path.clone();
        loop {
            let check_result: Option<PathBuf> = Config::check_dot_delivery_cli(path.clone());
            match check_result.as_ref() {
                Some(_) => { return check_result.clone() }
                None => {
                    if path.pop() { } else { return check_result.clone() }
                }
            }
        };
    }
}

#[cfg(test)]
mod tests {
    use super::Config;
    use std::default::Default;

    #[test]
    fn parse_config() {
        let toml = r#"
            server = "127.0.0.1"
            enterprise = "chef"
            organization = "chef"
            user = "adam"
"#;
        let config_result = Config::parse_config(toml);
        match config_result {
            Ok(config) => {
                assert_eq!(Some(String::from("127.0.0.1")), config.server);
                assert_eq!(None, config.git_port);
            },
            Err(e) => {
                panic!("Failed to parse: {:?}", e.detail)
            }
        }
    }

    #[test]
    fn test_api_url_with_port() {
        let mut conf  = Config::default();
        conf.server   = Some("127.0.0.1".to_string());
        conf.api_port = Some("2112".to_string());
        assert_eq!("127.0.0.1:2112".to_string(),
                   conf.api_host_and_port().unwrap());
    }

    #[test]
    fn test_api_url_without_port() {
        let mut conf = Config::default();
        conf.server  = Some("127.0.0.1".to_string());
        assert!(conf.api_port.is_none());
        assert_eq!("127.0.0.1".to_string(),
                   conf.api_host_and_port().unwrap());
    }

    #[test]
    fn test_api_url_without_server() {
        let conf = Config::default();
        assert!(conf.server.is_none());
        assert!(conf.api_host_and_port().is_err());
    }

    #[test]
    fn test_git_url_with_default_port() {
        let mut conf      = Config::default();
        conf.server       = Some("127.0.0.1".to_string());
        conf.user         = Some("user".to_string());
        conf.enterprise   = Some("ent".to_string());
        conf.organization = Some("org".to_string());
        conf.project      = Some("proj".to_string());
        assert_eq!("ssh://user@ent@127.0.0.1:8989/ent/org/proj".to_string(),
                   conf.delivery_git_ssh_url().unwrap());
    }

    #[test]
    fn test_git_url_with_port() {
        let mut conf      = Config::default();
        conf.server       = Some("127.0.0.1".to_string());
        conf.user         = Some("user".to_string());
        conf.enterprise   = Some("ent".to_string());
        conf.organization = Some("org".to_string());
        conf.project      = Some("proj".to_string());
        conf.git_port     = Some("2112".to_string());
        assert_eq!("ssh://user@ent@127.0.0.1:2112/ent/org/proj".to_string(),
                   conf.delivery_git_ssh_url().unwrap());
    }

    #[test]
    fn test_git_url_without_server() {
        let conf = Config::default();
        assert!(conf.server.is_none());
        assert!(conf.delivery_git_ssh_url().is_err());
    }
}
<commit_msg>Fix config handling to use defaults<commit_after>//
// Copyright:: Copyright (c) 2015 Chef Software, Inc.
// License:: Apache License, Version 2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

pub use errors;
use errors::{DeliveryError, Kind};
use std::fs::File;
use std::default::Default;
use utils::say::{say, sayln};
use rustc_serialize::Encodable;
use std::path::PathBuf;
use toml;
use utils::mkdir_recursive;
use std::io::prelude::*;
use utils::path_join_many::PathJoinMany;

#[derive(RustcEncodable, Clone)]
pub struct Config {
    pub server: Option<String>,
    pub api_port: Option<String>,
    pub user: Option<String>,
    pub enterprise: Option<String>,
    pub organization: Option<String>,
    pub project: Option<String>,
    pub git_port: Option<String>,
    pub pipeline: Option<String>
}

impl Default for Config {
    fn default() -> Config {
        Config{
            server: None,
            api_port: None,
            enterprise: None,
            organization: None,
            project: None,
            user: None,
            git_port: Some(String::from("8989")),
            pipeline: Some(String::from("master"))
        }
    }
}

macro_rules! config_accessor_for {
    ($name:ident, $set_name:ident, $err_msg:expr) => (
        impl Config {
            pub fn $name(&self) -> Result<String, DeliveryError> {
                match self.$name {
                    Some(ref v) => Ok(v.clone()),
                    None => Err(DeliveryError{ kind: Kind::MissingConfig, detail: Some(String::from($err_msg)) })
                }
            }

            pub fn $set_name(mut self, $name: &str) -> Config {
                if !$name.is_empty() {
                    self.$name = Some(String::from($name));
                }
                self
            }
        }
    )
}

config_accessor_for!(server, set_server, "Server not set; try --server or set it in your .toml config file");
config_accessor_for!(api_port, set_api_port, "API port not set; try --api-port or set it in your .toml config file");
config_accessor_for!(user, set_user, "User not set; try --user or set it in your .toml config file");
config_accessor_for!(enterprise, set_enterprise, "Enterprise not set; try --ent or set it in your .toml config file");
config_accessor_for!(organization, set_organization, "Organization not set; try --org or set it in your .toml config file");
config_accessor_for!(project, set_project, "Project not set; try --project or set it in your .toml config file");
config_accessor_for!(git_port, set_git_port, "Git Port not set; please set it in your .toml config file");
config_accessor_for!(pipeline, set_pipeline, "Pipeline not set; try --for or set it in your .toml config file");

impl Config {

    /// Return the host and port at which we can access the Delivery
    /// API. By default, we assume the use of HTTPS on the standard
    /// port `443`. Unless a port is specified in the configuration,
    /// we'll just return the server name; otherwise we append the
    /// port.
    pub fn api_host_and_port(&self) -> Result<String, DeliveryError> {
        let s = try!(self.server());
        return Ok(match self.api_port {
            Some(ref p) => format!("{}:{}", s, p),
            None    => s
        });
    }

    /// Returns the SSH URL to talk to Delivery's Git
    pub fn delivery_git_ssh_url(&self) -> Result<String, DeliveryError> {
        let s = try!(self.server());
        let host_and_port = match self.git_port {
            Some(ref p) => format!("{}:{}", s, p),
            None    => s // TODO: Currently we *always* have a git port
        };
        let u = try!(self.user());
        let e = try!(self.enterprise());
        let o = try!(self.organization());
        let p = try!(self.project());
        Ok(format!("ssh://{}@{}@{}/{}/{}/{}", u, e, host_and_port, e, o, p))
    }

    pub fn load_config(cwd: &PathBuf) -> Result<Config, DeliveryError> {
        let have_config = Config::have_dot_delivery_cli(cwd);
        match have_config.as_ref() {
            Some(path) => {
                let toml = try!(Config::read_file(path));
                match Config::parse_config(&toml) {
                    Ok(c) => return Ok(c),
                    Err(_) => return Ok(Default::default())
                }
            },
            None => return Ok(Default::default())
        }
    }

    pub fn write_file(&self, path: &PathBuf) -> Result<(), DeliveryError> {
        let write_dir = path.join_many(&[".delivery"]);
        if !write_dir.is_dir() {
            try!(mkdir_recursive(&write_dir));
        }
        let write_path = path.join_many(&[".delivery", "cli.toml"]);
        say("white", "Writing configuration to ");
        sayln("yellow", &format!("{}", write_path.display()));
        let mut f = try!(File::create(&write_path));
        let toml_string = toml::encode_str(self);
        sayln("magenta", "New configuration");
        sayln("magenta", "-----------------");
        say("white", &toml_string);
        try!(f.write_all(toml_string.as_bytes()));
        Ok(())
    }

    pub fn parse_config(toml: &str) -> Result<Config, DeliveryError> {
        let mut parser = toml::Parser::new(toml);
        match parser.parse() {
            Some(value) => { return Config::set_values_from_toml_table(value); },
            None => {
                return Err(DeliveryError{
                    kind: Kind::ConfigParse,
                    detail: Some(format!("Parse errors: {:?}", parser.errors))
                });
            }
        }
    }

    fn set_values_from_toml_table(table: toml::Table) -> Result<Config, DeliveryError> {
        let mut config: Config = Default::default();
        config.server = stringify_or("server", &table, config.server);
        config.api_port = stringify_or("api_port", &table, config.api_port);
        config.pipeline = stringify_or("pipeline", &table, config.pipeline);
        config.project = stringify_or("project", &table, config.project);
        config.enterprise = stringify_or("enterprise", &table,
                                         config.enterprise);
        config.organization = stringify_or("organization", &table,
                                           config.organization);
        config.user = stringify_or("user", &table, config.user);
        config.git_port = stringify_or("git_port", &table, config.git_port);
        return Ok(config);
    }

    fn read_file(path: &PathBuf) -> Result<String, DeliveryError>  {
        let mut toml_file = try!(File::open(path));
        let mut toml = String::new();
        try!(toml_file.read_to_string(&mut toml));
        Ok(toml)
    }

    fn stringify_values(toml_value: Option<&toml::Value>) -> Option<String> {
        match toml_value {
            Some(value) => {
                let is_string = value.as_str();
                match is_string {
                    Some(vstr) => return Some(String::from(vstr)),
                    None => return None
                }
            },
            None => {
                return None;
            }
        }
    }

    fn check_dot_delivery_cli(path: PathBuf) -> Option<PathBuf> {
        let dot_git = path.join_many(&[".delivery", "cli.toml"]);
        debug!("Checking {}", dot_git.display());
        let is_file: Option<PathBuf> = if dot_git.is_file() {
            Some(dot_git)
        } else {
            None
        };
        is_file
    }

    fn have_dot_delivery_cli(orig_path: &PathBuf) -> Option<PathBuf> {
        let mut path = orig_path.clone();
        loop {
            let check_result: Option<PathBuf> = Config::check_dot_delivery_cli(path.clone());
            match check_result.as_ref() {
                Some(_) => { return check_result.clone() }
                None => {
                    if path.pop() { } else { return check_result.clone() }
                }
            }
        };
    }
}

fn stringify_or(key: &str, table: &toml::Table, default: Option<String>) -> Option<String> {
    Config::stringify_values(table.get(key)).or(default)
}

#[cfg(test)]
mod tests {
    use super::Config;
    use std::default::Default;

    #[test]
    fn parse_config_with_defaults() {
        let toml = r#"
            server = "127.0.0.1"
            enterprise = "chef"
            user = "adam"
"#;
        let config_result = Config::parse_config(toml);
        match config_result {
            Ok(config) => {
                assert_eq!(Some(String::from("127.0.0.1")), config.server);
                assert_eq!(Some("8989".to_string()), config.git_port);
                assert_eq!(Some("master".to_string()), config.pipeline);
                assert_eq!(None, config.organization);
            },
            Err(e) => {
                panic!("Failed to parse: {:?}", e.detail)
            }
        }
    }

    #[test]
    fn parse_config_override_defaults() {
        let toml = r#"
            server = "127.0.0.1"
            enterprise = "chef"
            user = "adam"
            git_port = "4151"
            pipeline = "dev"
"#;
        let config_result = Config::parse_config(toml);
        match config_result {
            Ok(config) => {
                assert_eq!(Some("4151".to_string()), config.git_port);
                assert_eq!(Some("dev".to_string()), config.pipeline);
                assert_eq!(Some(String::from("127.0.0.1")), config.server);
                assert_eq!(None, config.organization);
            },
            Err(e) => {
                panic!("Failed to parse: {:?}", e.detail)
            }
        }
    }

    #[test]
    fn test_api_url_with_port() {
        let mut conf  = Config::default();
        conf.server   = Some("127.0.0.1".to_string());
        conf.api_port = Some("2112".to_string());
        assert_eq!("127.0.0.1:2112".to_string(),
                   conf.api_host_and_port().unwrap());
    }

    #[test]
    fn test_api_url_without_port() {
        let mut conf = Config::default();
        conf.server  = Some("127.0.0.1".to_string());
        assert!(conf.api_port.is_none());
        assert_eq!("127.0.0.1".to_string(),
                   conf.api_host_and_port().unwrap());
    }

    #[test]
    fn test_api_url_without_server() {
        let conf = Config::default();
        assert!(conf.server.is_none());
        assert!(conf.api_host_and_port().is_err());
    }

    #[test]
    fn test_git_url_with_default_port() {
        let mut conf      = Config::default();
        conf.server       = Some("127.0.0.1".to_string());
        conf.user         = Some("user".to_string());
        conf.enterprise   = Some("ent".to_string());
        conf.organization = Some("org".to_string());
        conf.project      = Some("proj".to_string());
        assert_eq!("ssh://user@ent@127.0.0.1:8989/ent/org/proj".to_string(),
                   conf.delivery_git_ssh_url().unwrap());
    }

    #[test]
    fn test_git_url_with_port() {
        let mut conf      = Config::default();
        conf.server       = Some("127.0.0.1".to_string());
        conf.user         = Some("user".to_string());
        conf.enterprise   = Some("ent".to_string());
        conf.organization = Some("org".to_string());
        conf.project      = Some("proj".to_string());
        conf.git_port     = Some("2112".to_string());
        assert_eq!("ssh://user@ent@127.0.0.1:2112/ent/org/proj".to_string(),
                   conf.delivery_git_ssh_url().unwrap());
    }

    #[test]
    fn test_git_url_without_server() {
        let conf = Config::default();
        assert!(conf.server.is_none());
        assert!(conf.delivery_git_ssh_url().is_err());
    }
}
<|endoftext|>