<commit_before>extern mod std;

#[test]
fn test_tokenize() {
    assert tokenize( ~"(1 2 3)" ) == ~[~"(", ~"1", ~"2", ~"3", ~")"];
}

#[test]
fn test_tokenize_empty() {
    assert tokenize( ~"" ) == ~[];
}

#[test]
fn test_tokenize_no_spaces() {
    assert tokenize( ~"one" ) == ~[~"one"];
}

#[test]
fn test_pad_parentheses() {
    assert pad_parentheses( ~"(1 2 3)" ) == ~" ( 1 2 3 ) "
}

#[test]
fn test_pad_parentheses_empty() {
    assert pad_parentheses( ~"" ) == ~""
}

#[test]
fn test_pad_parentheses_one() {
    assert pad_parentheses( ~"(" ) == ~" ( ";
    assert pad_parentheses( ~")" ) == ~" ) ";
}

fn pad_parentheses( input:&str ) -> ~str {
    str::replace(str::replace(input, ")", " ) "), "(", " ( ")
}

fn tokenize( input:&str ) -> ~[~str] {
    str::words(pad_parentheses(input))
}

#[test]
fn test_that_atom_can_read_a_symbol() {
    match atom(~"hello") {
        Symbol(~"hello") => (),
            _ => fail
    }
}

#[test]
fn test_that_atom_can_read_an_int() {
    match atom(~"10") {
        Int(10) => (),
            _ => fail
    }
}

#[test]
fn test_that_atom_can_read_a_float() {
    match atom(~"10.1") {
        Float(10.1) => (),
            _ => fail
    }
}

fn atom( input:~str ) -> Expression {
    match int::from_str(input) {
        Some(number) => Int(number),
        None => match float::from_str(input) {
            Some(number) => Float(number),
            None => Symbol(input)
        }
    }
}

#[test]
fn test_that_read_can_read_an_atom() {
    let atom = read( ~[~"12"] );
    match atom {
        Int(12) => (),
            _ => fail
    }
}

#[test]
fn test_that_read_can_read_a_list() {
    let list = read( ~[ ~"(", ~"1", ~")" ] );
    match list {
        List([Int(1)]) => (),
            _ => fail ~"not a list"
    }
}

#[test]
fn test_that_read_can_read_a_nested_list() {
    let list = read( ~[ ~"(", ~"1", ~"(", ~"2", ~")", ~"3", ~")" ] );
    match list {
        List([Int(1), List([Int(2)]), Int(3)]) => (),
        _ => fail
    }
}

fn parse( program:&str ) -> Expression {
    read( tokenize( program ) )
}

fn read( tokens:~[~str] ) -> Expression {
    fn subexpression( tokens:~[~str] ) -> (Expression, ~[~str]) {
        let mut remainder = copy tokens;
        let token = remainder.remove(0);
        match token {
            ~"(" => {
                let mut accumulator:~[Expression] = ~[];
                while remainder.len() > 0 && remainder[0] != ~")" {
                    let (expr, new_remainder) = subexpression( remainder );
                    accumulator.push(expr);
                    remainder = new_remainder
                }
                // remove the final close paren, this will fail if the parens
                // aren't closed properly due to an assert in remove
                remainder.remove(0);
                (List(accumulator), remainder)
            }
            ~")" => fail,
                _ => (atom(token), remainder)
        }
    }

    let (expression, _remainder) = subexpression( tokens );
    expression
}

enum Expression {
    Int(int),
    Float(float),
    Symbol(~str),
    List(~[Expression])
}

fn stringify_expression( expression:Expression ) -> ~str {
    match expression {
        Int(number) => { fmt!("%d", number) }
        Float(number) => {
            if number == (number as int) as float {
                fmt!("%.1f", number)
            } else {
                fmt!("%f", number)
            }
        }
        Symbol(string) => { string }
        List(expressions) => {
            let strings = expressions.map( | &expr | {stringify_expression(expr)} );
            ~"(" + strings.foldl(~"", |&x, &y| { x + ~" " + y } ).trim() + ~")"
        }
    }
}

#[test]
fn test_eval_returns_number_when_passed_number() {
    let value = eval( Int(1) );
    match value {
        Int(1) => (),
        _ => fail
    }
}

#[test]
fn test_eval_returns_expression_when_passed_quote() {
    let value = eval( List( ~[ Symbol(~"quote"), List( ~[ Symbol(~"a") ] ) ] ));
    io::println( stringify_expression( value ) );
    match value {
        List( [ Symbol( ~"a" ) ] ) => (),
        _ => fail
    }
}

#[test]
fn test_eval_returns_last_expression_when_passed_begin() {
    let expression = parse( ~"(begin 1 2 3)" );
    let value = eval( expression );
    match value {
        Int(3) => (),
        _ => fail
    }
}

#[test]
fn test_if_returns_third_part_when_if_is_true() {
    let expression = parse( ~"(if 1 2 3)" );
    let value = eval( expression );
    match value {
        Int(2) => (),
        _ => fail
    }
}

#[test]
fn test_if_returns_fourth_part_when_if_is_false() {
    let expression = parse( ~"(if 0 2 3)" );
    let value = eval( expression );
    match value {
        Int(3) => (),
        _ => fail
    }
}

fn eval( expression:Expression ) -> Expression {
    fn quote(expressions:~[Expression]) -> Expression {
        match expressions {
            [_, expr] => expr,
            _ => fail ~"Syntax Error: quote must take a single argument"
        }
    }

    fn if_(expressions:~[Expression]) -> Expression {
        match expressions {
            [_, test, true_expr, false_expr] => {
                if is_truthy( test ) { true_expr } else { false_expr }
            }
            _ => fail ~"Syntax Error: if must take three arguments"
        }
    }

    fn begin(expressions:~[Expression]) -> Expression {
        eval( copy expressions[ expressions.len() - 1] )
    }

    match copy expression {
        List( expressions ) => {
            match expressions[0] {
                Symbol(~"quote") => quote(expressions),
                Symbol(~"begin") => begin(expressions),
                Symbol(~"if") => if_(expressions),
                _ => expression
            }
        }
        _ => {
            expression
        }
    }
}

#[test]
fn test_is_truthy_returns_true_for_non_zero_numbers() {
    assert is_truthy( Int(1) );
    assert is_truthy( Int(-1) );
    assert is_truthy( Float(1.0) );
    assert is_truthy( Float(-1.0) );
}

#[test]
fn test_is_truthy_returns_false_for_zero_numbers() {
    assert !is_truthy( Int(0) );
    assert !is_truthy( Float(0.0) );
}

fn is_truthy( expression:Expression ) -> bool {
    match expression {
        Int( number ) => 0 != number,
        Float( number ) => 0.0 != number,
        _ => true
    }
}

#[test]
fn test_environment_accepts_new_value_and_returns_it() {
    let env:Environment = Environment::new_global_environment();
    env.define( ~"monkey", Int(1) );
    env.define( ~"turkey", Int(2) );
    match env.lookup( ~"monkey" ) {
        Some(Int(1)) => (),
        _ => fail ~"Monkey wasn't an Int(1)"
    }
    match env.lookup( ~"turkey" ) {
        Some(Int(2)) => (),
        _ => fail ~"Turkey wasn't an Int(2)"
    }
}

#[test]
fn test_environment_allows_values_to_change() {
    let env:Environment = Environment::new_global_environment();
    env.define( ~"monkey", Int(1) );
    match env.lookup( ~"monkey" ) {
        Some(Int(1)) => (),
        _ => fail ~"Monkey wasn't an Int(1) before mutation"
    }
    env.define( ~"monkey", Int(2) );
    match env.lookup( ~"monkey" ) {
        Some(Int(2)) => (),
        _ => fail ~"Monkey wasn't an Int(2) after mutation"
    }
}

#[test]
fn test_environment_checks_enclosing_environment() {
    let enclosure:~Environment = ~Environment::new_global_environment();
    enclosure.define( ~"monkey", Int(1) );
    let env:Environment = Environment::new(enclosure);
    match env.lookup( ~"monkey" ) {
        Some(Int(1)) => (),
        _ => fail ~"monkey wasn't found in env or enclosure... he's escaped?"
    }
}

struct Environment {
    enclosure:Option<~Environment>,
    mappings:std::treemap::TreeMap<~str,Expression>
}

impl Environment {
    fn define(&self, key:~str, value:Expression) {
        std::treemap::insert(self.mappings, key, value);
    }

    fn check_enclosure(&self, key:~str) -> Option<Expression> {
        match copy self.enclosure {
            Some(environment) => environment.lookup(key),
            _ => None
        }
    }

    fn lookup(&self, key:~str) -> Option<Expression> {
        let local_definition = std::treemap::find(self.mappings, copy key);
        match local_definition {
            None => self.check_enclosure(key),
            _ => local_definition
        }
    }

    static fn new_global_environment() -> Environment {
        let mappings = std::treemap::TreeMap();
        Environment {enclosure:None, mappings:mappings}
    }

    static fn new(enclosure:~Environment) -> Environment {
        let mappings = std::treemap::TreeMap();
        Environment {enclosure:Some(enclosure), mappings:mappings}
    }
}

fn main() {
    io::println(stringify_expression( parse( "(1 2 3 (1 2 3))" ) ));
    io::println(stringify_expression( parse( "((1 2) 3 (1 2 3))" ) ));
    let blah:Expression = List(~[Int(1), List(~[Float(1.0), Symbol(~"xyz")])]);
    io::println(stringify_expression(blah));
    io::println( "(begin 1 2)" )
}
<commit_msg>Add tests on if to confirm that everything is evaluated<commit_after>extern mod std;

#[test]
fn test_tokenize() {
    assert tokenize( ~"(1 2 3)" ) == ~[~"(", ~"1", ~"2", ~"3", ~")"];
}

#[test]
fn test_tokenize_empty() {
    assert tokenize( ~"" ) == ~[];
}

#[test]
fn test_tokenize_no_spaces() {
    assert tokenize( ~"one" ) == ~[~"one"];
}

#[test]
fn test_pad_parentheses() {
    assert pad_parentheses( ~"(1 2 3)" ) == ~" ( 1 2 3 ) "
}

#[test]
fn test_pad_parentheses_empty() {
    assert pad_parentheses( ~"" ) == ~""
}

#[test]
fn test_pad_parentheses_one() {
    assert pad_parentheses( ~"(" ) == ~" ( ";
    assert pad_parentheses( ~")" ) == ~" ) ";
}

fn pad_parentheses( input:&str ) -> ~str {
    str::replace(str::replace(input, ")", " ) "), "(", " ( ")
}

fn tokenize( input:&str ) -> ~[~str] {
    str::words(pad_parentheses(input))
}

#[test]
fn test_that_atom_can_read_a_symbol() {
    match atom(~"hello") {
        Symbol(~"hello") => (),
            _ => fail
    }
}

#[test]
fn test_that_atom_can_read_an_int() {
    match atom(~"10") {
        Int(10) => (),
            _ => fail
    }
}

#[test]
fn test_that_atom_can_read_a_float() {
    match atom(~"10.1") {
        Float(10.1) => (),
            _ => fail
    }
}

fn atom( input:~str ) -> Expression {
    match int::from_str(input) {
        Some(number) => Int(number),
        None => match float::from_str(input) {
            Some(number) => Float(number),
            None => Symbol(input)
        }
    }
}

#[test]
fn test_that_read_can_read_an_atom() {
    let atom = read( ~[~"12"] );
    match atom {
        Int(12) => (),
            _ => fail
    }
}

#[test]
fn test_that_read_can_read_a_list() {
    let list = read( ~[ ~"(", ~"1", ~")" ] );
    match list {
        List([Int(1)]) => (),
            _ => fail ~"not a list"
    }
}

#[test]
fn test_that_read_can_read_a_nested_list() {
    let list = read( ~[ ~"(", ~"1", ~"(", ~"2", ~")", ~"3", ~")" ] );
    match list {
        List([Int(1), List([Int(2)]), Int(3)]) => (),
        _ => fail
    }
}

fn parse( program:&str ) -> Expression {
    read( tokenize( program ) )
}

fn read( tokens:~[~str] ) -> Expression {
    fn subexpression( tokens:~[~str] ) -> (Expression, ~[~str]) {
        let mut remainder = copy tokens;
        let token = remainder.remove(0);
        match token {
            ~"(" => {
                let mut accumulator:~[Expression] = ~[];
                while remainder.len() > 0 && remainder[0] != ~")" {
                    let (expr, new_remainder) = subexpression( remainder );
                    accumulator.push(expr);
                    remainder = new_remainder
                }
                // remove the final close paren, this will fail if the parens
                // aren't closed properly due to an assert in remove
                remainder.remove(0);
                (List(accumulator), remainder)
            }
            ~")" => fail,
                _ => (atom(token), remainder)
        }
    }

    let (expression, _remainder) = subexpression( tokens );
    expression
}

enum Expression {
    Int(int),
    Float(float),
    Symbol(~str),
    List(~[Expression])
}

fn stringify_expression( expression:Expression ) -> ~str {
    match expression {
        Int(number) => { fmt!("%d", number) }
        Float(number) => {
            if number == (number as int) as float {
                fmt!("%.1f", number)
            } else {
                fmt!("%f", number)
            }
        }
        Symbol(string) => { string }
        List(expressions) => {
            let strings = expressions.map( | &expr | {stringify_expression(expr)} );
            ~"(" + strings.foldl(~"", |&x, &y| { x + ~" " + y } ).trim() + ~")"
        }
    }
}

#[test]
fn test_eval_returns_number_when_passed_number() {
    let value = eval( Int(1) );
    match value {
        Int(1) => (),
        _ => fail
    }
}

#[test]
fn test_eval_returns_expression_when_passed_quote() {
    let value = eval( List( ~[ Symbol(~"quote"), List( ~[ Symbol(~"a") ] ) ] ));
    io::println( stringify_expression( value ) );
    match value {
        List( [ Symbol( ~"a" ) ] ) => (),
        _ => fail
    }
}

#[test]
fn test_eval_returns_last_expression_when_passed_begin() {
    let expression = parse( ~"(begin 1 2 3)" );
    let value = eval( expression );
    match value {
        Int(3) => (),
        _ => fail
    }
}

#[test]
fn test_if_returns_third_part_when_if_is_true() {
    let expression = parse( ~"(if 1 2 3)" );
    let value = eval( expression );
    match value {
        Int(2) => (),
        _ => fail
    }
}

#[test]
fn test_if_returns_fourth_part_when_if_is_false() {
    let expression = parse( ~"(if 0 2 3)" );
    let value = eval( expression );
    match value {
        Int(3) => (),
        _ => fail
    }
}

#[test]
fn test_that_if_evaluates_the_then_branch() {
    let expression = parse( ~"(if 1 (begin 1 2) 7)" );
    let value = eval( expression );
    match value {
        Int(2) => (),
        List( [Symbol(~"begin"), Int(1), Int(2)]) => fail ~"If just returned the then branch",
        _ => fail fmt!("If returned something unusual (%s)", stringify_expression(value))
    }
}

#[test]
fn test_that_if_evaluates_the_else_branch() {
    let expression = parse( ~"(if 0 7 (begin 1 2))" );
    let value = eval( expression );
    match value {
        Int(2) => (),
        List( [Symbol(~"begin"), Int(1), Int(2)]) => fail ~"If just returned the else branch",
        _ => fail fmt!("If returned something unusual (%s)", stringify_expression(value))
    }
}

#[test]
fn test_that_if_evaluates_the_test() {
    let expression = parse( ~"(if (begin 1 0) 1 2)" );
    let value = eval( expression );
    match value {
        Int(1) => fail ~"If didn't evaluate the test",
        Int(2) => (),
        _ => fail fmt!("If returned something unusual (%s)", stringify_expression(value))
    }
}

fn eval( expression:Expression ) -> Expression {
    fn quote(expressions:~[Expression]) -> Expression {
        match expressions {
            [_, expr] => expr,
            _ => fail ~"Syntax Error: quote must take a single argument"
        }
    }

    fn if_(expressions:~[Expression]) -> Expression {
        match expressions {
            [_, test, true_expr, false_expr] => {
                eval(if is_truthy( eval(test) ) { true_expr } else { false_expr })
            }
            _ => fail ~"Syntax Error: if must take three arguments"
        }
    }

    fn begin(expressions:~[Expression]) -> Expression {
        eval( copy expressions[ expressions.len() - 1] )
    }

    match copy expression {
        List( expressions ) => {
            match expressions[0] {
                Symbol(~"quote") => quote(expressions),
                Symbol(~"begin") => begin(expressions),
                Symbol(~"if") => if_(expressions),
                _ => expression
            }
        }
        _ => {
            expression
        }
    }
}

#[test]
fn test_is_truthy_returns_true_for_non_zero_numbers() {
    assert is_truthy( Int(1) );
    assert is_truthy( Int(-1) );
    assert is_truthy( Float(1.0) );
    assert is_truthy( Float(-1.0) );
}

#[test]
fn test_is_truthy_returns_false_for_zero_numbers() {
    assert !is_truthy( Int(0) );
    assert !is_truthy( Float(0.0) );
}

fn is_truthy( expression:Expression ) -> bool {
    match expression {
        Int( number ) => 0 != number,
        Float( number ) => 0.0 != number,
        _ => true
    }
}

#[test]
fn test_environment_accepts_new_value_and_returns_it() {
    let env:Environment = Environment::new_global_environment();
    env.define( ~"monkey", Int(1) );
    env.define( ~"turkey", Int(2) );
    match env.lookup( ~"monkey" ) {
        Some(Int(1)) => (),
        _ => fail ~"Monkey wasn't an Int(1)"
    }
    match env.lookup( ~"turkey" ) {
        Some(Int(2)) => (),
        _ => fail ~"Turkey wasn't an Int(2)"
    }
}

#[test]
fn test_environment_allows_values_to_change() {
    let env:Environment = Environment::new_global_environment();
    env.define( ~"monkey", Int(1) );
    match env.lookup( ~"monkey" ) {
        Some(Int(1)) => (),
        _ => fail ~"Monkey wasn't an Int(1) before mutation"
    }
    env.define( ~"monkey", Int(2) );
    match env.lookup( ~"monkey" ) {
        Some(Int(2)) => (),
        _ => fail ~"Monkey wasn't an Int(2) after mutation"
    }
}

#[test]
fn test_environment_checks_enclosing_environment() {
    let enclosure:~Environment = ~Environment::new_global_environment();
    enclosure.define( ~"monkey", Int(1) );
    let env:Environment = Environment::new(enclosure);
    match env.lookup( ~"monkey" ) {
        Some(Int(1)) => (),
        _ => fail ~"monkey wasn't found in env or enclosure... he's escaped?"
    }
}

struct Environment {
    enclosure:Option<~Environment>,
    mappings:std::treemap::TreeMap<~str,Expression>
}

impl Environment {
    fn define(&self, key:~str, value:Expression) {
        std::treemap::insert(self.mappings, key, value);
    }

    fn check_enclosure(&self, key:~str) -> Option<Expression> {
        match copy self.enclosure {
            Some(environment) => environment.lookup(key),
            _ => None
        }
    }

    fn lookup(&self, key:~str) -> Option<Expression> {
        let local_definition = std::treemap::find(self.mappings, copy key);
        match local_definition {
            None => self.check_enclosure(key),
            _ => local_definition
        }
    }

    static fn new_global_environment() -> Environment {
        let mappings = std::treemap::TreeMap();
        Environment {enclosure:None, mappings:mappings}
    }

    static fn new(enclosure:~Environment) -> Environment {
        let mappings = std::treemap::TreeMap();
        Environment {enclosure:Some(enclosure), mappings:mappings}
    }
}

fn main() {
    io::println(stringify_expression( parse( "(1 2 3 (1 2 3))" ) ));
    io::println(stringify_expression( parse( "((1 2) 3 (1 2 3))" ) ));
    let blah:Expression = List(~[Int(1), List(~[Float(1.0), Symbol(~"xyz")])]);
    io::println(stringify_expression(blah));
    io::println( "(begin 1 2)" )
}
<|endoftext|>


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

//! Weighted index sampling
//!
//! This module provides two implementations for sampling indices:
//!
//! *   [`WeightedIndex`] allows `O(log N)` sampling
//! *   [`alias_method::WeightedIndex`] allows `O(1)` sampling, but with
//!      much greater set-up cost
//!
//! [`alias_method::WeightedIndex`]: alias_method/struct.WeightedIndex.html

pub mod alias_method;

use crate::Rng;
use crate::distributions::Distribution;
use crate::distributions::uniform::{UniformSampler, SampleUniform, SampleBorrow};
use core::cmp::PartialOrd;
use core::fmt;

// Note that this whole module is only imported if feature="alloc" is enabled.
#[cfg(not(feature="std"))] use crate::alloc::vec::Vec;

/// A distribution using weighted sampling to pick a discretely selected
/// item.
///
/// Sampling a `WeightedIndex` distribution returns the index of a randomly
/// selected element from the iterator used when the `WeightedIndex` was
/// created. The chance of a given element being picked is proportional to the
/// value of the element. The weights can use any type `X` for which an
/// implementation of [`Uniform<X>`] exists.
///
/// # Performance
///
/// A `WeightedIndex<X>` contains a `Vec<X>` and a [`Uniform<X>`] and so its
/// size is the sum of the size of those objects, possibly plus some alignment.
///
/// Creating a `WeightedIndex<X>` will allocate enough space to hold `N - 1`
/// weights of type `X`, where `N` is the number of weights. However, since
/// `Vec` doesn't guarantee a particular growth strategy, additional memory
/// might be allocated but not used. Since the `WeightedIndex` object also
/// contains, this might cause additional allocations, though for primitive
/// types, ['Uniform<X>`] doesn't allocate any memory.
///
/// Time complexity of sampling from `WeightedIndex` is `O(log N)` where
/// `N` is the number of weights.
///
/// Sampling from `WeightedIndex` will result in a single call to
/// `Uniform<X>::sample` (method of the [`Distribution`] trait), which typically
/// will request a single value from the underlying [`RngCore`], though the
/// exact number depends on the implementaiton of `Uniform<X>::sample`.
///
/// # Example
///
/// ```
/// use rand::prelude::*;
/// use rand::distributions::WeightedIndex;
///
/// let choices = ['a', 'b', 'c'];
/// let weights = [2,   1,   1];
/// let dist = WeightedIndex::new(&weights).unwrap();
/// let mut rng = thread_rng();
/// for _ in 0..100 {
///     // 50% chance to print 'a', 25% chance to print 'b', 25% chance to print 'c'
///     println!("{}", choices[dist.sample(&mut rng)]);
/// }
///
/// let items = [('a', 0), ('b', 3), ('c', 7)];
/// let dist2 = WeightedIndex::new(items.iter().map(|item| item.1)).unwrap();
/// for _ in 0..100 {
///     // 0% chance to print 'a', 30% chance to print 'b', 70% chance to print 'c'
///     println!("{}", items[dist2.sample(&mut rng)].0);
/// }
/// ```
///
/// [`Uniform<X>`]: crate::distributions::uniform::Uniform
/// [`RngCore`]: crate::RngCore
#[derive(Debug, Clone)]
pub struct WeightedIndex<X: SampleUniform + PartialOrd> {
    cumulative_weights: Vec<X>,
    total_weight: X,
    weight_distribution: X::Sampler,
}

impl<X: SampleUniform + PartialOrd> WeightedIndex<X> {
    /// Creates a new a `WeightedIndex` [`Distribution`] using the values
    /// in `weights`. The weights can use any type `X` for which an
    /// implementation of [`Uniform<X>`] exists.
    ///
    /// Returns an error if the iterator is empty, if any weight is `< 0`, or
    /// if its total value is 0.
    ///
    /// [`Uniform<X>`]: crate::distributions::uniform::Uniform
    pub fn new<I>(weights: I) -> Result<WeightedIndex<X>, WeightedError>
        where I: IntoIterator,
              I::Item: SampleBorrow<X>,
              X: for<'a> ::core::ops::AddAssign<&'a X> +
                 Clone +
                 Default {
        let mut iter = weights.into_iter();
        let mut total_weight: X = iter.next()
                                      .ok_or(WeightedError::NoItem)?
                                      .borrow()
                                      .clone();

        let zero = <X as Default>::default();
        if total_weight < zero {
            return Err(WeightedError::InvalidWeight);
        }

        let mut weights = Vec::<X>::with_capacity(iter.size_hint().0);
        for w in iter {
            if *w.borrow() < zero {
                return Err(WeightedError::InvalidWeight);
            }
            weights.push(total_weight.clone());
            total_weight += w.borrow();
        }

        if total_weight == zero {
            return Err(WeightedError::AllWeightsZero);
        }
        let distr = X::Sampler::new(zero, total_weight.clone());

        Ok(WeightedIndex { cumulative_weights: weights, total_weight, weight_distribution: distr })
    }

    /// Update a subset of weights, without changing the number of weights.
    ///
    /// Using this method instead of `new` might be more efficient if only a small number of
    /// weights is modified. For weights that are `Copy`, no allocations are performed.
    ///
    /// In case of error, `self` is not modified.
    pub fn update_weights(&mut self, new_weights: &[(usize, &X)]) -> Result<(), WeightedError>
        where X: for<'a> ::core::ops::AddAssign<&'a X> +
                 for<'a> ::core::ops::SubAssign<&'a X> +
                 Clone +
                 Default {
        let zero = <X as Default>::default();

        let mut total_weight = self.total_weight.clone();

        for &(i, w) in new_weights {
            if *w < zero {
                return Err(WeightedError::InvalidWeight);
            }
            if i >= self.cumulative_weights.len() {
                return Err(WeightedError::TooMany);
            }

            // Unfortunately, we will have to calculate the non-cumulative weight a second time, to
            // avoid producing an invalid state of `self`.
            let mut old_w = self.cumulative_weights[i].clone();
            if i > 0 {
                old_w -= &self.cumulative_weights[i - 1];
            }

            total_weight -= &old_w;
            total_weight += w;
        }
        if total_weight == zero {
            return Err(WeightedError::AllWeightsZero);
        }

        for &(i, w) in new_weights {
            let mut old_w = self.cumulative_weights[i].clone();
            if i > 0 {
                old_w -= &self.cumulative_weights[i - 1];
            }

            for j in i..self.cumulative_weights.len() {
                self.cumulative_weights[j] -= &old_w;
                self.cumulative_weights[j] += w;
            }
        }
        self.total_weight = total_weight;
        self.weight_distribution = X::Sampler::new(zero, self.total_weight.clone());

        Ok(())
    }
}

impl<X> Distribution<usize> for WeightedIndex<X> where
    X: SampleUniform + PartialOrd {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
        use ::core::cmp::Ordering;
        let chosen_weight = self.weight_distribution.sample(rng);
        // Find the first item which has a weight *higher* than the chosen weight.
        self.cumulative_weights.binary_search_by(
            |w| if *w <= chosen_weight { Ordering::Less } else { Ordering::Greater }).unwrap_err()
    }
}

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

    #[test]
    #[cfg(not(miri))] // Miri is too slow
    fn test_weightedindex() {
        let mut r = crate::test::rng(700);
        const N_REPS: u32 = 5000;
        let weights = [1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7];
        let total_weight = weights.iter().sum::<u32>() as f32;

        let verify = |result: [i32; 14]| {
            for (i, count) in result.iter().enumerate() {
                let exp = (weights[i] * N_REPS) as f32 / total_weight;
                let mut err = (*count as f32 - exp).abs();
                if err != 0.0 {
                    err /= exp;
                }
                assert!(err <= 0.25);
            }
        };

        // WeightedIndex from vec
        let mut chosen = [0i32; 14];
        let distr = WeightedIndex::new(weights.to_vec()).unwrap();
        for _ in 0..N_REPS {
            chosen[distr.sample(&mut r)] += 1;
        }
        verify(chosen);

        // WeightedIndex from slice
        chosen = [0i32; 14];
        let distr = WeightedIndex::new(&weights[..]).unwrap();
        for _ in 0..N_REPS {
            chosen[distr.sample(&mut r)] += 1;
        }
        verify(chosen);

        // WeightedIndex from iterator
        chosen = [0i32; 14];
        let distr = WeightedIndex::new(weights.iter()).unwrap();
        for _ in 0..N_REPS {
            chosen[distr.sample(&mut r)] += 1;
        }
        verify(chosen);

        for _ in 0..5 {
            assert_eq!(WeightedIndex::new(&[0, 1]).unwrap().sample(&mut r), 1);
            assert_eq!(WeightedIndex::new(&[1, 0]).unwrap().sample(&mut r), 0);
            assert_eq!(WeightedIndex::new(&[0, 0, 0, 0, 10, 0]).unwrap().sample(&mut r), 4);
        }

        assert_eq!(WeightedIndex::new(&[10][0..0]).unwrap_err(), WeightedError::NoItem);
        assert_eq!(WeightedIndex::new(&[0]).unwrap_err(), WeightedError::AllWeightsZero);
        assert_eq!(WeightedIndex::new(&[10, 20, -1, 30]).unwrap_err(), WeightedError::InvalidWeight);
        assert_eq!(WeightedIndex::new(&[-10, 20, 1, 30]).unwrap_err(), WeightedError::InvalidWeight);
        assert_eq!(WeightedIndex::new(&[-10]).unwrap_err(), WeightedError::InvalidWeight);
    }

    #[test]
    fn test_update_weights() {
        let weights = [1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7];
        let total_weight = weights.iter().sum::<u32>();
        let mut distr = WeightedIndex::new(weights.to_vec()).unwrap();
        assert_eq!(distr.total_weight, total_weight);

        distr.update_weights(&[(2, &4), (5, &1)]).unwrap();
        let expected_weights = [1u32, 2, 4, 0, 5, 1, 7, 1, 2, 3, 4, 5, 6, 7];
        let expected_total_weight = expected_weights.iter().sum::<u32>();
        let expected_distr = WeightedIndex::new(expected_weights.to_vec()).unwrap();
        assert_eq!(distr.total_weight, expected_total_weight);
        assert_eq!(distr.total_weight, expected_distr.total_weight);
        assert_eq!(distr.cumulative_weights, expected_distr.cumulative_weights);
    }
}

/// Error type returned from `WeightedIndex::new`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WeightedError {
    /// The provided weight collection contains no items.
    NoItem,

    /// A weight is either less than zero, greater than the supported maximum or
    /// otherwise invalid.
    InvalidWeight,

    /// All items in the provided weight collection are zero.
    AllWeightsZero,

    /// Too many weights are provided (length greater than `u32::MAX`)
    TooMany,
}

impl WeightedError {
    fn msg(&self) -> &str {
        match *self {
            WeightedError::NoItem => "No weights provided.",
            WeightedError::InvalidWeight => "A weight is invalid.",
            WeightedError::AllWeightsZero => "All weights are zero.",
            WeightedError::TooMany => "Too many weights (hit u32::MAX)",
        }
    }
}

#[cfg(feature="std")]
impl ::std::error::Error for WeightedError {
    fn description(&self) -> &str {
        self.msg()
    }
    fn cause(&self) -> Option<&dyn (::std::error::Error)> {
        None
    }
}

impl fmt::Display for WeightedError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.msg())
    }
}
<commit_msg>WeightedIndex::update_weights: Correct comment<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.

//! Weighted index sampling
//!
//! This module provides two implementations for sampling indices:
//!
//! *   [`WeightedIndex`] allows `O(log N)` sampling
//! *   [`alias_method::WeightedIndex`] allows `O(1)` sampling, but with
//!      much greater set-up cost
//!
//! [`alias_method::WeightedIndex`]: alias_method/struct.WeightedIndex.html

pub mod alias_method;

use crate::Rng;
use crate::distributions::Distribution;
use crate::distributions::uniform::{UniformSampler, SampleUniform, SampleBorrow};
use core::cmp::PartialOrd;
use core::fmt;

// Note that this whole module is only imported if feature="alloc" is enabled.
#[cfg(not(feature="std"))] use crate::alloc::vec::Vec;

/// A distribution using weighted sampling to pick a discretely selected
/// item.
///
/// Sampling a `WeightedIndex` distribution returns the index of a randomly
/// selected element from the iterator used when the `WeightedIndex` was
/// created. The chance of a given element being picked is proportional to the
/// value of the element. The weights can use any type `X` for which an
/// implementation of [`Uniform<X>`] exists.
///
/// # Performance
///
/// A `WeightedIndex<X>` contains a `Vec<X>` and a [`Uniform<X>`] and so its
/// size is the sum of the size of those objects, possibly plus some alignment.
///
/// Creating a `WeightedIndex<X>` will allocate enough space to hold `N - 1`
/// weights of type `X`, where `N` is the number of weights. However, since
/// `Vec` doesn't guarantee a particular growth strategy, additional memory
/// might be allocated but not used. Since the `WeightedIndex` object also
/// contains, this might cause additional allocations, though for primitive
/// types, ['Uniform<X>`] doesn't allocate any memory.
///
/// Time complexity of sampling from `WeightedIndex` is `O(log N)` where
/// `N` is the number of weights.
///
/// Sampling from `WeightedIndex` will result in a single call to
/// `Uniform<X>::sample` (method of the [`Distribution`] trait), which typically
/// will request a single value from the underlying [`RngCore`], though the
/// exact number depends on the implementaiton of `Uniform<X>::sample`.
///
/// # Example
///
/// ```
/// use rand::prelude::*;
/// use rand::distributions::WeightedIndex;
///
/// let choices = ['a', 'b', 'c'];
/// let weights = [2,   1,   1];
/// let dist = WeightedIndex::new(&weights).unwrap();
/// let mut rng = thread_rng();
/// for _ in 0..100 {
///     // 50% chance to print 'a', 25% chance to print 'b', 25% chance to print 'c'
///     println!("{}", choices[dist.sample(&mut rng)]);
/// }
///
/// let items = [('a', 0), ('b', 3), ('c', 7)];
/// let dist2 = WeightedIndex::new(items.iter().map(|item| item.1)).unwrap();
/// for _ in 0..100 {
///     // 0% chance to print 'a', 30% chance to print 'b', 70% chance to print 'c'
///     println!("{}", items[dist2.sample(&mut rng)].0);
/// }
/// ```
///
/// [`Uniform<X>`]: crate::distributions::uniform::Uniform
/// [`RngCore`]: crate::RngCore
#[derive(Debug, Clone)]
pub struct WeightedIndex<X: SampleUniform + PartialOrd> {
    cumulative_weights: Vec<X>,
    total_weight: X,
    weight_distribution: X::Sampler,
}

impl<X: SampleUniform + PartialOrd> WeightedIndex<X> {
    /// Creates a new a `WeightedIndex` [`Distribution`] using the values
    /// in `weights`. The weights can use any type `X` for which an
    /// implementation of [`Uniform<X>`] exists.
    ///
    /// Returns an error if the iterator is empty, if any weight is `< 0`, or
    /// if its total value is 0.
    ///
    /// [`Uniform<X>`]: crate::distributions::uniform::Uniform
    pub fn new<I>(weights: I) -> Result<WeightedIndex<X>, WeightedError>
        where I: IntoIterator,
              I::Item: SampleBorrow<X>,
              X: for<'a> ::core::ops::AddAssign<&'a X> +
                 Clone +
                 Default {
        let mut iter = weights.into_iter();
        let mut total_weight: X = iter.next()
                                      .ok_or(WeightedError::NoItem)?
                                      .borrow()
                                      .clone();

        let zero = <X as Default>::default();
        if total_weight < zero {
            return Err(WeightedError::InvalidWeight);
        }

        let mut weights = Vec::<X>::with_capacity(iter.size_hint().0);
        for w in iter {
            if *w.borrow() < zero {
                return Err(WeightedError::InvalidWeight);
            }
            weights.push(total_weight.clone());
            total_weight += w.borrow();
        }

        if total_weight == zero {
            return Err(WeightedError::AllWeightsZero);
        }
        let distr = X::Sampler::new(zero, total_weight.clone());

        Ok(WeightedIndex { cumulative_weights: weights, total_weight, weight_distribution: distr })
    }

    /// Update a subset of weights, without changing the number of weights.
    ///
    /// Using this method instead of `new` might be more efficient if only a small number of
    /// weights is modified. No allocations are performed, unless the weight type `X` uses
    /// allocation internally.
    ///
    /// In case of error, `self` is not modified.
    pub fn update_weights(&mut self, new_weights: &[(usize, &X)]) -> Result<(), WeightedError>
        where X: for<'a> ::core::ops::AddAssign<&'a X> +
                 for<'a> ::core::ops::SubAssign<&'a X> +
                 Clone +
                 Default {
        let zero = <X as Default>::default();

        let mut total_weight = self.total_weight.clone();

        for &(i, w) in new_weights {
            if *w < zero {
                return Err(WeightedError::InvalidWeight);
            }
            if i >= self.cumulative_weights.len() {
                return Err(WeightedError::TooMany);
            }

            // Unfortunately, we will have to calculate the non-cumulative weight a second time, to
            // avoid producing an invalid state of `self`.
            let mut old_w = self.cumulative_weights[i].clone();
            if i > 0 {
                old_w -= &self.cumulative_weights[i - 1];
            }

            total_weight -= &old_w;
            total_weight += w;
        }
        if total_weight == zero {
            return Err(WeightedError::AllWeightsZero);
        }

        for &(i, w) in new_weights {
            let mut old_w = self.cumulative_weights[i].clone();
            if i > 0 {
                old_w -= &self.cumulative_weights[i - 1];
            }

            for j in i..self.cumulative_weights.len() {
                self.cumulative_weights[j] -= &old_w;
                self.cumulative_weights[j] += w;
            }
        }
        self.total_weight = total_weight;
        self.weight_distribution = X::Sampler::new(zero, self.total_weight.clone());

        Ok(())
    }
}

impl<X> Distribution<usize> for WeightedIndex<X> where
    X: SampleUniform + PartialOrd {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
        use ::core::cmp::Ordering;
        let chosen_weight = self.weight_distribution.sample(rng);
        // Find the first item which has a weight *higher* than the chosen weight.
        self.cumulative_weights.binary_search_by(
            |w| if *w <= chosen_weight { Ordering::Less } else { Ordering::Greater }).unwrap_err()
    }
}

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

    #[test]
    #[cfg(not(miri))] // Miri is too slow
    fn test_weightedindex() {
        let mut r = crate::test::rng(700);
        const N_REPS: u32 = 5000;
        let weights = [1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7];
        let total_weight = weights.iter().sum::<u32>() as f32;

        let verify = |result: [i32; 14]| {
            for (i, count) in result.iter().enumerate() {
                let exp = (weights[i] * N_REPS) as f32 / total_weight;
                let mut err = (*count as f32 - exp).abs();
                if err != 0.0 {
                    err /= exp;
                }
                assert!(err <= 0.25);
            }
        };

        // WeightedIndex from vec
        let mut chosen = [0i32; 14];
        let distr = WeightedIndex::new(weights.to_vec()).unwrap();
        for _ in 0..N_REPS {
            chosen[distr.sample(&mut r)] += 1;
        }
        verify(chosen);

        // WeightedIndex from slice
        chosen = [0i32; 14];
        let distr = WeightedIndex::new(&weights[..]).unwrap();
        for _ in 0..N_REPS {
            chosen[distr.sample(&mut r)] += 1;
        }
        verify(chosen);

        // WeightedIndex from iterator
        chosen = [0i32; 14];
        let distr = WeightedIndex::new(weights.iter()).unwrap();
        for _ in 0..N_REPS {
            chosen[distr.sample(&mut r)] += 1;
        }
        verify(chosen);

        for _ in 0..5 {
            assert_eq!(WeightedIndex::new(&[0, 1]).unwrap().sample(&mut r), 1);
            assert_eq!(WeightedIndex::new(&[1, 0]).unwrap().sample(&mut r), 0);
            assert_eq!(WeightedIndex::new(&[0, 0, 0, 0, 10, 0]).unwrap().sample(&mut r), 4);
        }

        assert_eq!(WeightedIndex::new(&[10][0..0]).unwrap_err(), WeightedError::NoItem);
        assert_eq!(WeightedIndex::new(&[0]).unwrap_err(), WeightedError::AllWeightsZero);
        assert_eq!(WeightedIndex::new(&[10, 20, -1, 30]).unwrap_err(), WeightedError::InvalidWeight);
        assert_eq!(WeightedIndex::new(&[-10, 20, 1, 30]).unwrap_err(), WeightedError::InvalidWeight);
        assert_eq!(WeightedIndex::new(&[-10]).unwrap_err(), WeightedError::InvalidWeight);
    }

    #[test]
    fn test_update_weights() {
        let weights = [1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7];
        let total_weight = weights.iter().sum::<u32>();
        let mut distr = WeightedIndex::new(weights.to_vec()).unwrap();
        assert_eq!(distr.total_weight, total_weight);

        distr.update_weights(&[(2, &4), (5, &1)]).unwrap();
        let expected_weights = [1u32, 2, 4, 0, 5, 1, 7, 1, 2, 3, 4, 5, 6, 7];
        let expected_total_weight = expected_weights.iter().sum::<u32>();
        let expected_distr = WeightedIndex::new(expected_weights.to_vec()).unwrap();
        assert_eq!(distr.total_weight, expected_total_weight);
        assert_eq!(distr.total_weight, expected_distr.total_weight);
        assert_eq!(distr.cumulative_weights, expected_distr.cumulative_weights);
    }
}

/// Error type returned from `WeightedIndex::new`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WeightedError {
    /// The provided weight collection contains no items.
    NoItem,

    /// A weight is either less than zero, greater than the supported maximum or
    /// otherwise invalid.
    InvalidWeight,

    /// All items in the provided weight collection are zero.
    AllWeightsZero,

    /// Too many weights are provided (length greater than `u32::MAX`)
    TooMany,
}

impl WeightedError {
    fn msg(&self) -> &str {
        match *self {
            WeightedError::NoItem => "No weights provided.",
            WeightedError::InvalidWeight => "A weight is invalid.",
            WeightedError::AllWeightsZero => "All weights are zero.",
            WeightedError::TooMany => "Too many weights (hit u32::MAX)",
        }
    }
}

#[cfg(feature="std")]
impl ::std::error::Error for WeightedError {
    fn description(&self) -> &str {
        self.msg()
    }
    fn cause(&self) -> Option<&dyn (::std::error::Error)> {
        None
    }
}

impl fmt::Display for WeightedError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.msg())
    }
}
<|endoftext|>


<commit_before>use crate::parser::trivia::from_utf8_unchecked;
use combine::parser::byte::{byte, bytes, digit, hex_digit, oct_digit};
use combine::parser::range::{range, recognize};
use combine::stream::RangeStream;
use combine::*;

// ;; Boolean

// boolean = true / false
parse!(boolean() -> bool, {
    choice((
        (byte(b't'), range(&b"rue"[..]),),
        (byte(b'f'), range(&b"alse"[..]),),
    )).map(|p| p.0 == b't')
});

// ;; Integer

// integer = dec-int / hex-int / oct-int / bin-int
parse!(integer() -> i64, {
    choice!(
        attempt(hex_int()),
        attempt(oct_int()),
        attempt(bin_int()),
        dec_int()
            .and_then(|s| s.replace("_", "").parse())
            .message("While parsing an Integer")
    )
});

// dec-int = [ minus / plus ] unsigned-dec-int
// unsigned-dec-int = DIGIT / digit1-9 1*( DIGIT / underscore DIGIT )
parse!(dec_int() -> &'a str, {
    recognize((
        optional(choice([byte(b'-'), byte(b'+')])),
        choice((
            byte(b'0'),
            (
                satisfy(|c| (b'1'..=b'9').contains(&c)),
                skip_many((
                    optional(byte(b'_')),
                    skip_many1(digit()),
                )),
            ).map(|t| t.0),
        )),
    )).map(|b: &[u8]| {
        unsafe { from_utf8_unchecked(b, "`digit` and `_` filter out npn-ASCII") }
    })
});

// hex-prefix = %x30.78               ; 0x
// hex-int = hex-prefix HEXDIG *( HEXDIG / underscore HEXDIG )
parse!(hex_int() -> i64, {
    bytes(b"0x").with(
        recognize((
            hex_digit(),
            skip_many((
                optional(byte(b'_')),
                skip_many1(hex_digit()),
            )),
        ).map(|t| t.0)
    )).and_then(|b: &[u8]| {
        let s = unsafe { from_utf8_unchecked(b, "`hex_digit` and `_` filter out npn-ASCII") };
        i64::from_str_radix(&s.replace("_", ""), 16)
    }).message("While parsing a hexadecimal Integer")
});

// oct-prefix = %x30.6F               ; 0o
// oct-int = oct-prefix digit0-7 *( digit0-7 / underscore digit0-7 )
parse!(oct_int() -> i64, {
    bytes(b"0o").with(
        recognize((
            oct_digit(),
            skip_many((
                optional(byte(b'_')),
                skip_many1(oct_digit()),
            )),
        ).map(|t| t.0)
    )).and_then(|b: &[u8]| {
        let s = unsafe { from_utf8_unchecked(b, "`oct_digit` and `_` filter out npn-ASCII") };
        i64::from_str_radix(&s.replace("_", ""), 8)
    }).message("While parsing a octal Integer")
});

// bin-prefix = %x30.62               ; 0b
// bin-int = bin-prefix digit0-1 *( digit0-1 / underscore digit0-1 )
parse!(bin_int() -> i64, {
    bytes(b"0b").with(
        recognize((
            satisfy(|c: u8| c == b'0' || c == b'1'),
            skip_many((
                optional(byte(b'_')),
                skip_many1(satisfy(|c: u8| c == b'0' || c == b'1')),
            )),
        ).map(|t| t.0)
    )).and_then(|b: &[u8]| {
        let s = unsafe { from_utf8_unchecked(b, "`is_digit` and `_` filter out npn-ASCII") };
        i64::from_str_radix(&s.replace("_", ""), 2)
    }).message("While parsing a binary Integer")
});

// ;; Float

// float = float-int-part ( exp / frac [ exp ] )
// float =/ special-float
// float-int-part = dec-int
parse!(float() -> f64, {
    choice((
        parse_float()
            .and_then(|s| s.replace("_", "").parse()),
        special_float(),
    )).message("While parsing a Float")
});

parse!(parse_float() -> &'a str, {
    recognize((
        attempt((dec_int(), look_ahead(one_of([b'e', b'E', b'.'])))),
        choice((
            exp(),
            (
                frac(),
                optional(exp()),
            ).map(|_| "")
        )),
    )).map(|b: &[u8]| {
        unsafe { from_utf8_unchecked(b, "`dec_int`, `one_of`, `exp`, and `frac` filter out npn-ASCII") }
    })
});

// frac = decimal-point zero-prefixable-int
// decimal-point = %x2E               ; .
parse!(frac() -> &'a str, {
    recognize((
        byte(b'.'),
        parse_zero_prefixable_int(),
    )).map(|b: &[u8]| {
        unsafe { from_utf8_unchecked(b, "`.` and `parse_zero_prefixable_int` filter out npn-ASCII") }
    })
});

// zero-prefixable-int = DIGIT *( DIGIT / underscore DIGIT )
parse!(parse_zero_prefixable_int() -> &'a str, {
    recognize((
        skip_many1(digit()),
        skip_many((
            optional(byte(b'_')),
            skip_many1(digit()),
        )),
    )).map(|b: &[u8]| {
        unsafe { from_utf8_unchecked(b, "`digit` and `_` filter out npn-ASCII") }
    })
});

// exp = "e" float-exp-part
// float-exp-part = [ minus / plus ] zero-prefixable-int
parse!(exp() -> &'a str, {
    recognize((
        one_of([b'e', b'E']),
        optional(one_of([b'+', b'-'])),
        parse_zero_prefixable_int(),
    )).map(|b: &[u8]| {
        unsafe { from_utf8_unchecked(b, "`one_of` and `parse_zero_prefixable_int` filter out npn-ASCII") }
    })
});

// special-float = [ minus / plus ] ( inf / nan )
parse!(special_float() -> f64, {
    attempt(optional(one_of([b'+', b'-'])).and(choice((inf(), nan()))).map(|(s, f)| {
        match s {
            Some(b'+') | None => f,
            Some(b'-') => -f,
            _ => unreachable!("one_of should prevent this"),
        }
    }))
});

// inf = %x69.6e.66  ; inf
parse!(inf() -> f64, {
    range(&b"inf"[..]).map(|_| f64::INFINITY)
});

// nan = %x6e.61.6e  ; nan
parse!(nan() -> f64, {
    range(&b"nan"[..]).map(|_| f64::NAN)
});
<commit_msg>refactor(parser): Pull out constants<commit_after>use crate::parser::trivia::from_utf8_unchecked;
use combine::parser::byte::{byte, bytes, digit, hex_digit, oct_digit};
use combine::parser::range::{range, recognize};
use combine::stream::RangeStream;
use combine::*;

// ;; Boolean

// boolean = true / false
pub(crate) const TRUE: &[u8] = b"true";
pub(crate) const FALSE: &[u8] = b"false";
parse!(boolean() -> bool, {
    choice((
        (byte(TRUE[0]), range(&TRUE[1..]),),
        (byte(FALSE[0]), range(&FALSE[1..]),),
    )).map(|p| p.0 == b't')
});

// ;; Integer

// integer = dec-int / hex-int / oct-int / bin-int
parse!(integer() -> i64, {
    choice!(
        attempt(hex_int()),
        attempt(oct_int()),
        attempt(bin_int()),
        dec_int()
            .and_then(|s| s.replace("_", "").parse())
            .message("While parsing an Integer")
    )
});

// dec-int = [ minus / plus ] unsigned-dec-int
// unsigned-dec-int = DIGIT / digit1-9 1*( DIGIT / underscore DIGIT )
parse!(dec_int() -> &'a str, {
    recognize((
        optional(choice([byte(b'-'), byte(b'+')])),
        choice((
            byte(b'0'),
            (
                satisfy(|c| (b'1'..=b'9').contains(&c)),
                skip_many((
                    optional(byte(b'_')),
                    skip_many1(digit()),
                )),
            ).map(|t| t.0),
        )),
    )).map(|b: &[u8]| {
        unsafe { from_utf8_unchecked(b, "`digit` and `_` filter out npn-ASCII") }
    })
});

// hex-prefix = %x30.78               ; 0x
// hex-int = hex-prefix HEXDIG *( HEXDIG / underscore HEXDIG )
parse!(hex_int() -> i64, {
    bytes(b"0x").with(
        recognize((
            hex_digit(),
            skip_many((
                optional(byte(b'_')),
                skip_many1(hex_digit()),
            )),
        ).map(|t| t.0)
    )).and_then(|b: &[u8]| {
        let s = unsafe { from_utf8_unchecked(b, "`hex_digit` and `_` filter out npn-ASCII") };
        i64::from_str_radix(&s.replace("_", ""), 16)
    }).message("While parsing a hexadecimal Integer")
});

// oct-prefix = %x30.6F               ; 0o
// oct-int = oct-prefix digit0-7 *( digit0-7 / underscore digit0-7 )
parse!(oct_int() -> i64, {
    bytes(b"0o").with(
        recognize((
            oct_digit(),
            skip_many((
                optional(byte(b'_')),
                skip_many1(oct_digit()),
            )),
        ).map(|t| t.0)
    )).and_then(|b: &[u8]| {
        let s = unsafe { from_utf8_unchecked(b, "`oct_digit` and `_` filter out npn-ASCII") };
        i64::from_str_radix(&s.replace("_", ""), 8)
    }).message("While parsing a octal Integer")
});

// bin-prefix = %x30.62               ; 0b
// bin-int = bin-prefix digit0-1 *( digit0-1 / underscore digit0-1 )
parse!(bin_int() -> i64, {
    bytes(b"0b").with(
        recognize((
            satisfy(|c: u8| c == b'0' || c == b'1'),
            skip_many((
                optional(byte(b'_')),
                skip_many1(satisfy(|c: u8| c == b'0' || c == b'1')),
            )),
        ).map(|t| t.0)
    )).and_then(|b: &[u8]| {
        let s = unsafe { from_utf8_unchecked(b, "`is_digit` and `_` filter out npn-ASCII") };
        i64::from_str_radix(&s.replace("_", ""), 2)
    }).message("While parsing a binary Integer")
});

// ;; Float

// float = float-int-part ( exp / frac [ exp ] )
// float =/ special-float
// float-int-part = dec-int
parse!(float() -> f64, {
    choice((
        parse_float()
            .and_then(|s| s.replace("_", "").parse()),
        special_float(),
    )).message("While parsing a Float")
});

parse!(parse_float() -> &'a str, {
    recognize((
        attempt((dec_int(), look_ahead(one_of([b'e', b'E', b'.'])))),
        choice((
            exp(),
            (
                frac(),
                optional(exp()),
            ).map(|_| "")
        )),
    )).map(|b: &[u8]| {
        unsafe { from_utf8_unchecked(b, "`dec_int`, `one_of`, `exp`, and `frac` filter out npn-ASCII") }
    })
});

// frac = decimal-point zero-prefixable-int
// decimal-point = %x2E               ; .
parse!(frac() -> &'a str, {
    recognize((
        byte(b'.'),
        parse_zero_prefixable_int(),
    )).map(|b: &[u8]| {
        unsafe { from_utf8_unchecked(b, "`.` and `parse_zero_prefixable_int` filter out npn-ASCII") }
    })
});

// zero-prefixable-int = DIGIT *( DIGIT / underscore DIGIT )
parse!(parse_zero_prefixable_int() -> &'a str, {
    recognize((
        skip_many1(digit()),
        skip_many((
            optional(byte(b'_')),
            skip_many1(digit()),
        )),
    )).map(|b: &[u8]| {
        unsafe { from_utf8_unchecked(b, "`digit` and `_` filter out npn-ASCII") }
    })
});

// exp = "e" float-exp-part
// float-exp-part = [ minus / plus ] zero-prefixable-int
parse!(exp() -> &'a str, {
    recognize((
        one_of([b'e', b'E']),
        optional(one_of([b'+', b'-'])),
        parse_zero_prefixable_int(),
    )).map(|b: &[u8]| {
        unsafe { from_utf8_unchecked(b, "`one_of` and `parse_zero_prefixable_int` filter out npn-ASCII") }
    })
});

// special-float = [ minus / plus ] ( inf / nan )
parse!(special_float() -> f64, {
    attempt(optional(one_of([b'+', b'-'])).and(choice((inf(), nan()))).map(|(s, f)| {
        match s {
            Some(b'+') | None => f,
            Some(b'-') => -f,
            _ => unreachable!("one_of should prevent this"),
        }
    }))
});

// inf = %x69.6e.66  ; inf
pub(crate) const INF: &[u8] = b"inf";
parse!(inf() -> f64, {
    range(INF).map(|_| f64::INFINITY)
});

// nan = %x6e.61.6e  ; nan
pub(crate) const NAN: &[u8] = b"nan";
parse!(nan() -> f64, {
    range(NAN).map(|_| f64::NAN)
});
<|endoftext|>


<commit_before>//! VFS stands for Virtual File System.
//!
//! When doing analysis, we don't want to do any IO, we want to keep all source
//! code in memory. However, the actual source code is stored on disk, so you
//! need to get it into the memory in the first place somehow. VFS is the
//! component which does this.
//!
//! It is also responsible for watching the disk for changes, and for merging
//! editor state (modified, unsaved files) with disk state.
//! TODO: Some LSP clients support watching the disk, so this crate should
//! to support custom watcher events (related to https://github.com/rust-analyzer/rust-analyzer/issues/131)
//!
//! VFS is based on a concept of roots: a set of directories on the file system
//! whihc are watched for changes. Typically, there will be a root for each
//! Cargo package.
mod arena;
mod io;

use std::{
    fmt,
    mem,
    thread,
    cmp::Reverse,
    path::{Path, PathBuf},
    ffi::OsStr,
    sync::Arc,
    fs,
};

use rustc_hash::{FxHashMap, FxHashSet};
use relative_path::RelativePathBuf;
use crossbeam_channel::Receiver;
use walkdir::DirEntry;
use thread_worker::{WorkerHandle};

use crate::{
    arena::{ArenaId, Arena},
};

pub use crate::io::TaskResult as VfsTask;

/// `RootFilter` is a predicate that checks if a file can belong to a root. If
/// several filters match a file (nested dirs), the most nested one wins.
struct RootFilter {
    root: PathBuf,
    file_filter: fn(&Path) -> bool,
}

impl RootFilter {
    fn new(root: PathBuf) -> RootFilter {
        RootFilter {
            root,
            file_filter: has_rs_extension,
        }
    }
    /// Check if this root can contain `path`. NB: even if this returns
    /// true, the `path` might actually be conained in some nested root.
    fn can_contain(&self, path: &Path) -> Option<RelativePathBuf> {
        if !(self.file_filter)(path) {
            return None;
        }
        if !(path.starts_with(&self.root)) {
            return None;
        }
        let path = path.strip_prefix(&self.root).unwrap();
        let path = RelativePathBuf::from_path(path).unwrap();
        Some(path)
    }
}

fn has_rs_extension(p: &Path) -> bool {
    p.extension() == Some(OsStr::new("rs"))
}

#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
pub struct VfsRoot(pub u32);

impl ArenaId for VfsRoot {
    fn from_u32(idx: u32) -> VfsRoot {
        VfsRoot(idx)
    }
    fn to_u32(self) -> u32 {
        self.0
    }
}

#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
pub struct VfsFile(pub u32);

impl ArenaId for VfsFile {
    fn from_u32(idx: u32) -> VfsFile {
        VfsFile(idx)
    }
    fn to_u32(self) -> u32 {
        self.0
    }
}

struct VfsFileData {
    root: VfsRoot,
    path: RelativePathBuf,
    text: Arc<String>,
}

pub struct Vfs {
    roots: Arena<VfsRoot, RootFilter>,
    files: Arena<VfsFile, VfsFileData>,
    root2files: FxHashMap<VfsRoot, FxHashSet<VfsFile>>,
    pending_changes: Vec<VfsChange>,
    worker: io::Worker,
    worker_handle: WorkerHandle,
}

impl fmt::Debug for Vfs {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("Vfs { ... }")
    }
}

impl Vfs {
    pub fn new(mut roots: Vec<PathBuf>) -> (Vfs, Vec<VfsRoot>) {
        let (worker, worker_handle) = io::start();

        let mut res = Vfs {
            roots: Arena::default(),
            files: Arena::default(),
            root2files: FxHashMap::default(),
            worker,
            worker_handle,
            pending_changes: Vec::new(),
        };

        // A hack to make nesting work.
        roots.sort_by_key(|it| Reverse(it.as_os_str().len()));
        for (i, path) in roots.iter().enumerate() {
            let root = res.roots.alloc(RootFilter::new(path.clone()));
            res.root2files.insert(root, Default::default());
            let nested = roots[..i]
                .iter()
                .filter(|it| it.starts_with(path))
                .map(|it| it.clone())
                .collect::<Vec<_>>();
            let filter = move |entry: &DirEntry| {
                if entry.file_type().is_file() {
                    has_rs_extension(entry.path())
                } else {
                    nested.iter().all(|it| it != entry.path())
                }
            };
            let task = io::Task {
                root,
                path: path.clone(),
                filter: Box::new(filter),
            };
            res.worker.inp.send(task);
        }
        let roots = res.roots.iter().map(|(id, _)| id).collect();
        (res, roots)
    }

    pub fn root2path(&self, root: VfsRoot) -> PathBuf {
        self.roots[root].root.clone()
    }

    pub fn path2file(&self, path: &Path) -> Option<VfsFile> {
        if let Some((_root, _path, Some(file))) = self.find_root(path) {
            return Some(file);
        }
        None
    }

    pub fn file2path(&self, file: VfsFile) -> PathBuf {
        let rel_path = &self.files[file].path;
        let root_path = &self.roots[self.files[file].root].root;
        rel_path.to_path(root_path)
    }

    pub fn file_for_path(&self, path: &Path) -> Option<VfsFile> {
        if let Some((_root, _path, Some(file))) = self.find_root(path) {
            return Some(file);
        }
        None
    }

    pub fn load(&mut self, path: &Path) -> Option<VfsFile> {
        if let Some((root, rel_path, file)) = self.find_root(path) {
            return if let Some(file) = file {
                Some(file)
            } else {
                let text = fs::read_to_string(path).unwrap_or_default();
                let text = Arc::new(text);
                let file = self.add_file(root, rel_path.clone(), Arc::clone(&text));
                let change = VfsChange::AddFile {
                    file,
                    text,
                    root,
                    path: rel_path,
                };
                self.pending_changes.push(change);
                Some(file)
            };
        }
        None
    }

    pub fn task_receiver(&self) -> &Receiver<io::TaskResult> {
        &self.worker.out
    }

    pub fn handle_task(&mut self, task: io::TaskResult) {
        let mut files = Vec::new();
        // While we were scanning the root in the backgound, a file might have
        // been open in the editor, so we need to account for that.
        let exising = self.root2files[&task.root]
            .iter()
            .map(|&file| (self.files[file].path.clone(), file))
            .collect::<FxHashMap<_, _>>();
        for (path, text) in task.files {
            if let Some(&file) = exising.get(&path) {
                let text = Arc::clone(&self.files[file].text);
                files.push((file, path, text));
                continue;
            }
            let text = Arc::new(text);
            let file = self.add_file(task.root, path.clone(), Arc::clone(&text));
            files.push((file, path, text));
        }

        let change = VfsChange::AddRoot {
            root: task.root,
            files,
        };
        self.pending_changes.push(change);
    }

    pub fn add_file_overlay(&mut self, path: &Path, text: String) -> Option<VfsFile> {
        let mut res = None;
        if let Some((root, path, file)) = self.find_root(path) {
            let text = Arc::new(text);
            let change = if let Some(file) = file {
                res = Some(file);
                self.change_file(file, Arc::clone(&text));
                VfsChange::ChangeFile { file, text }
            } else {
                let file = self.add_file(root, path.clone(), Arc::clone(&text));
                res = Some(file);
                VfsChange::AddFile {
                    file,
                    text,
                    root,
                    path,
                }
            };
            self.pending_changes.push(change);
        }
        res
    }

    pub fn change_file_overlay(&mut self, path: &Path, new_text: String) {
        if let Some((_root, _path, file)) = self.find_root(path) {
            let file = file.expect("can't change a file which wasn't added");
            let text = Arc::new(new_text);
            self.change_file(file, Arc::clone(&text));
            let change = VfsChange::ChangeFile { file, text };
            self.pending_changes.push(change);
        }
    }

    pub fn remove_file_overlay(&mut self, path: &Path) -> Option<VfsFile> {
        let mut res = None;
        if let Some((root, path, file)) = self.find_root(path) {
            let file = file.expect("can't remove a file which wasn't added");
            res = Some(file);
            let full_path = path.to_path(&self.roots[root].root);
            let change = if let Ok(text) = fs::read_to_string(&full_path) {
                let text = Arc::new(text);
                self.change_file(file, Arc::clone(&text));
                VfsChange::ChangeFile { file, text }
            } else {
                self.remove_file(file);
                VfsChange::RemoveFile { root, file, path }
            };
            self.pending_changes.push(change);
        }
        res
    }

    pub fn commit_changes(&mut self) -> Vec<VfsChange> {
        mem::replace(&mut self.pending_changes, Vec::new())
    }

    /// Sutdown the VFS and terminate the background watching thread.
    pub fn shutdown(self) -> thread::Result<()> {
        let _ = self.worker.shutdown();
        self.worker_handle.shutdown()
    }

    fn add_file(&mut self, root: VfsRoot, path: RelativePathBuf, text: Arc<String>) -> VfsFile {
        let data = VfsFileData { root, path, text };
        let file = self.files.alloc(data);
        self.root2files.get_mut(&root).unwrap().insert(file);
        file
    }

    fn change_file(&mut self, file: VfsFile, new_text: Arc<String>) {
        self.files[file].text = new_text;
    }

    fn remove_file(&mut self, file: VfsFile) {
        //FIXME: use arena with removal
        self.files[file].text = Default::default();
        self.files[file].path = Default::default();
        let root = self.files[file].root;
        let removed = self.root2files.get_mut(&root).unwrap().remove(&file);
        assert!(removed);
    }

    fn find_root(&self, path: &Path) -> Option<(VfsRoot, RelativePathBuf, Option<VfsFile>)> {
        let (root, path) = self
            .roots
            .iter()
            .find_map(|(root, data)| data.can_contain(path).map(|it| (root, it)))?;
        let file = self.root2files[&root]
            .iter()
            .map(|&it| it)
            .find(|&file| self.files[file].path == path);
        Some((root, path, file))
    }
}

#[derive(Debug, Clone)]
pub enum VfsChange {
    AddRoot {
        root: VfsRoot,
        files: Vec<(VfsFile, RelativePathBuf, Arc<String>)>,
    },
    AddFile {
        root: VfsRoot,
        file: VfsFile,
        path: RelativePathBuf,
        text: Arc<String>,
    },
    RemoveFile {
        root: VfsRoot,
        file: VfsFile,
        path: RelativePathBuf,
    },
    ChangeFile {
        file: VfsFile,
        text: Arc<String>,
    },
}
<commit_msg>Fix instance of uneeded brackets in use_statement (thanks to #333)<commit_after>//! VFS stands for Virtual File System.
//!
//! When doing analysis, we don't want to do any IO, we want to keep all source
//! code in memory. However, the actual source code is stored on disk, so you
//! need to get it into the memory in the first place somehow. VFS is the
//! component which does this.
//!
//! It is also responsible for watching the disk for changes, and for merging
//! editor state (modified, unsaved files) with disk state.
//! TODO: Some LSP clients support watching the disk, so this crate should
//! to support custom watcher events (related to https://github.com/rust-analyzer/rust-analyzer/issues/131)
//!
//! VFS is based on a concept of roots: a set of directories on the file system
//! whihc are watched for changes. Typically, there will be a root for each
//! Cargo package.
mod arena;
mod io;

use std::{
    fmt,
    mem,
    thread,
    cmp::Reverse,
    path::{Path, PathBuf},
    ffi::OsStr,
    sync::Arc,
    fs,
};

use rustc_hash::{FxHashMap, FxHashSet};
use relative_path::RelativePathBuf;
use crossbeam_channel::Receiver;
use walkdir::DirEntry;
use thread_worker::WorkerHandle;

use crate::{
    arena::{ArenaId, Arena},
};

pub use crate::io::TaskResult as VfsTask;

/// `RootFilter` is a predicate that checks if a file can belong to a root. If
/// several filters match a file (nested dirs), the most nested one wins.
struct RootFilter {
    root: PathBuf,
    file_filter: fn(&Path) -> bool,
}

impl RootFilter {
    fn new(root: PathBuf) -> RootFilter {
        RootFilter {
            root,
            file_filter: has_rs_extension,
        }
    }
    /// Check if this root can contain `path`. NB: even if this returns
    /// true, the `path` might actually be conained in some nested root.
    fn can_contain(&self, path: &Path) -> Option<RelativePathBuf> {
        if !(self.file_filter)(path) {
            return None;
        }
        if !(path.starts_with(&self.root)) {
            return None;
        }
        let path = path.strip_prefix(&self.root).unwrap();
        let path = RelativePathBuf::from_path(path).unwrap();
        Some(path)
    }
}

fn has_rs_extension(p: &Path) -> bool {
    p.extension() == Some(OsStr::new("rs"))
}

#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
pub struct VfsRoot(pub u32);

impl ArenaId for VfsRoot {
    fn from_u32(idx: u32) -> VfsRoot {
        VfsRoot(idx)
    }
    fn to_u32(self) -> u32 {
        self.0
    }
}

#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
pub struct VfsFile(pub u32);

impl ArenaId for VfsFile {
    fn from_u32(idx: u32) -> VfsFile {
        VfsFile(idx)
    }
    fn to_u32(self) -> u32 {
        self.0
    }
}

struct VfsFileData {
    root: VfsRoot,
    path: RelativePathBuf,
    text: Arc<String>,
}

pub struct Vfs {
    roots: Arena<VfsRoot, RootFilter>,
    files: Arena<VfsFile, VfsFileData>,
    root2files: FxHashMap<VfsRoot, FxHashSet<VfsFile>>,
    pending_changes: Vec<VfsChange>,
    worker: io::Worker,
    worker_handle: WorkerHandle,
}

impl fmt::Debug for Vfs {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("Vfs { ... }")
    }
}

impl Vfs {
    pub fn new(mut roots: Vec<PathBuf>) -> (Vfs, Vec<VfsRoot>) {
        let (worker, worker_handle) = io::start();

        let mut res = Vfs {
            roots: Arena::default(),
            files: Arena::default(),
            root2files: FxHashMap::default(),
            worker,
            worker_handle,
            pending_changes: Vec::new(),
        };

        // A hack to make nesting work.
        roots.sort_by_key(|it| Reverse(it.as_os_str().len()));
        for (i, path) in roots.iter().enumerate() {
            let root = res.roots.alloc(RootFilter::new(path.clone()));
            res.root2files.insert(root, Default::default());
            let nested = roots[..i]
                .iter()
                .filter(|it| it.starts_with(path))
                .map(|it| it.clone())
                .collect::<Vec<_>>();
            let filter = move |entry: &DirEntry| {
                if entry.file_type().is_file() {
                    has_rs_extension(entry.path())
                } else {
                    nested.iter().all(|it| it != entry.path())
                }
            };
            let task = io::Task {
                root,
                path: path.clone(),
                filter: Box::new(filter),
            };
            res.worker.inp.send(task);
        }
        let roots = res.roots.iter().map(|(id, _)| id).collect();
        (res, roots)
    }

    pub fn root2path(&self, root: VfsRoot) -> PathBuf {
        self.roots[root].root.clone()
    }

    pub fn path2file(&self, path: &Path) -> Option<VfsFile> {
        if let Some((_root, _path, Some(file))) = self.find_root(path) {
            return Some(file);
        }
        None
    }

    pub fn file2path(&self, file: VfsFile) -> PathBuf {
        let rel_path = &self.files[file].path;
        let root_path = &self.roots[self.files[file].root].root;
        rel_path.to_path(root_path)
    }

    pub fn file_for_path(&self, path: &Path) -> Option<VfsFile> {
        if let Some((_root, _path, Some(file))) = self.find_root(path) {
            return Some(file);
        }
        None
    }

    pub fn load(&mut self, path: &Path) -> Option<VfsFile> {
        if let Some((root, rel_path, file)) = self.find_root(path) {
            return if let Some(file) = file {
                Some(file)
            } else {
                let text = fs::read_to_string(path).unwrap_or_default();
                let text = Arc::new(text);
                let file = self.add_file(root, rel_path.clone(), Arc::clone(&text));
                let change = VfsChange::AddFile {
                    file,
                    text,
                    root,
                    path: rel_path,
                };
                self.pending_changes.push(change);
                Some(file)
            };
        }
        None
    }

    pub fn task_receiver(&self) -> &Receiver<io::TaskResult> {
        &self.worker.out
    }

    pub fn handle_task(&mut self, task: io::TaskResult) {
        let mut files = Vec::new();
        // While we were scanning the root in the backgound, a file might have
        // been open in the editor, so we need to account for that.
        let exising = self.root2files[&task.root]
            .iter()
            .map(|&file| (self.files[file].path.clone(), file))
            .collect::<FxHashMap<_, _>>();
        for (path, text) in task.files {
            if let Some(&file) = exising.get(&path) {
                let text = Arc::clone(&self.files[file].text);
                files.push((file, path, text));
                continue;
            }
            let text = Arc::new(text);
            let file = self.add_file(task.root, path.clone(), Arc::clone(&text));
            files.push((file, path, text));
        }

        let change = VfsChange::AddRoot {
            root: task.root,
            files,
        };
        self.pending_changes.push(change);
    }

    pub fn add_file_overlay(&mut self, path: &Path, text: String) -> Option<VfsFile> {
        let mut res = None;
        if let Some((root, path, file)) = self.find_root(path) {
            let text = Arc::new(text);
            let change = if let Some(file) = file {
                res = Some(file);
                self.change_file(file, Arc::clone(&text));
                VfsChange::ChangeFile { file, text }
            } else {
                let file = self.add_file(root, path.clone(), Arc::clone(&text));
                res = Some(file);
                VfsChange::AddFile {
                    file,
                    text,
                    root,
                    path,
                }
            };
            self.pending_changes.push(change);
        }
        res
    }

    pub fn change_file_overlay(&mut self, path: &Path, new_text: String) {
        if let Some((_root, _path, file)) = self.find_root(path) {
            let file = file.expect("can't change a file which wasn't added");
            let text = Arc::new(new_text);
            self.change_file(file, Arc::clone(&text));
            let change = VfsChange::ChangeFile { file, text };
            self.pending_changes.push(change);
        }
    }

    pub fn remove_file_overlay(&mut self, path: &Path) -> Option<VfsFile> {
        let mut res = None;
        if let Some((root, path, file)) = self.find_root(path) {
            let file = file.expect("can't remove a file which wasn't added");
            res = Some(file);
            let full_path = path.to_path(&self.roots[root].root);
            let change = if let Ok(text) = fs::read_to_string(&full_path) {
                let text = Arc::new(text);
                self.change_file(file, Arc::clone(&text));
                VfsChange::ChangeFile { file, text }
            } else {
                self.remove_file(file);
                VfsChange::RemoveFile { root, file, path }
            };
            self.pending_changes.push(change);
        }
        res
    }

    pub fn commit_changes(&mut self) -> Vec<VfsChange> {
        mem::replace(&mut self.pending_changes, Vec::new())
    }

    /// Sutdown the VFS and terminate the background watching thread.
    pub fn shutdown(self) -> thread::Result<()> {
        let _ = self.worker.shutdown();
        self.worker_handle.shutdown()
    }

    fn add_file(&mut self, root: VfsRoot, path: RelativePathBuf, text: Arc<String>) -> VfsFile {
        let data = VfsFileData { root, path, text };
        let file = self.files.alloc(data);
        self.root2files.get_mut(&root).unwrap().insert(file);
        file
    }

    fn change_file(&mut self, file: VfsFile, new_text: Arc<String>) {
        self.files[file].text = new_text;
    }

    fn remove_file(&mut self, file: VfsFile) {
        //FIXME: use arena with removal
        self.files[file].text = Default::default();
        self.files[file].path = Default::default();
        let root = self.files[file].root;
        let removed = self.root2files.get_mut(&root).unwrap().remove(&file);
        assert!(removed);
    }

    fn find_root(&self, path: &Path) -> Option<(VfsRoot, RelativePathBuf, Option<VfsFile>)> {
        let (root, path) = self
            .roots
            .iter()
            .find_map(|(root, data)| data.can_contain(path).map(|it| (root, it)))?;
        let file = self.root2files[&root]
            .iter()
            .map(|&it| it)
            .find(|&file| self.files[file].path == path);
        Some((root, path, file))
    }
}

#[derive(Debug, Clone)]
pub enum VfsChange {
    AddRoot {
        root: VfsRoot,
        files: Vec<(VfsFile, RelativePathBuf, Arc<String>)>,
    },
    AddFile {
        root: VfsRoot,
        file: VfsFile,
        path: RelativePathBuf,
        text: Arc<String>,
    },
    RemoveFile {
        root: VfsRoot,
        file: VfsFile,
        path: RelativePathBuf,
    },
    ChangeFile {
        file: VfsFile,
        text: Arc<String>,
    },
}
<|endoftext|>


<commit_before>use mime::Mime;

header! {
    /// `Content-Type` header, defined in
    /// [RFC7231](http://tools.ietf.org/html/rfc7231#section-3.1.1.5)
    ///
    /// The `Content-Type` header field indicates the media type of the
    /// associated representation: either the representation enclosed in the
    /// message payload or the selected representation, as determined by the
    /// message semantics.  The indicated media type defines both the data
    /// format and how that data is intended to be processed by a recipient,
    /// within the scope of the received message semantics, after any content
    /// codings indicated by Content-Encoding are decoded.
    ///
    /// # ABNF
    /// ```plain
    /// Content-Type = media-type
    /// ```
    ///
    /// # Example values
    /// * `text/html; charset=ISO-8859-4`
    ///
    /// # Examples
    /// ```
    /// use hyper::header::{Headers, ContentType};
    /// use hyper::mime::{Mime, TopLevel, SubLevel};
    ///
    /// let mut headers = Headers::new();
    ///
    /// headers.set(
    ///     ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![]))
    /// );
    /// ```
    /// ```
    /// use hyper::header::{Headers, ContentType};
    /// use hyper::mime::{Mime, TopLevel, SubLevel, Attr, Value};
    ///
    /// let mut headers = Headers::new();
    ///
    /// headers.set(
    ///     ContentType(Mime(TopLevel::Application, SubLevel::Json,
    ///                      vec![(Attr::Charset, Value::Utf8)]))
    /// );
    /// ```
    (ContentType, "Content-Type") => [Mime]

    test_content_type {
        test_header!(
            test1,
            // FIXME: Should be b"text/html; charset=ISO-8859-4" but mime crate lowercases
            // the whole value so parsing and formatting the value gives a different result
            vec![b"text/html; charset=iso-8859-4"],
            Some(HeaderField(Mime(
                TopLevel::Text,
                SubLevel::Html,
                vec![(Attr::Charset, Value::Ext("iso-8859-4".to_owned()))]))));
    }
}

impl ContentType {
    /// A constructor  to easily create a `Content-Type: application/json` header.
    #[inline]
    pub fn json() -> ContentType {
        ContentType(mime!(Application/Json))
    }

    /// A constructor  to easily create a `Content-Type: text/plain; charset=utf-8` header.
    #[inline]
    pub fn plaintext() -> ContentType {
        ContentType(mime!(Text/Plain; Charset=Utf8))
    }

    /// A constructor  to easily create a `Content-Type: text/html; charset=utf-8` header.
    #[inline]
    pub fn html() -> ContentType {
        ContentType(mime!(Text/Html; Charset=Utf8))
    }

    /// A constructor  to easily create a `Content-Type: application/www-form-url-encoded` header.
    #[inline]
    pub fn form_url_encoded() -> ContentType {
        ContentType(mime!(Application/WwwFormUrlEncoded))
    }
    /// A constructor  to easily create a `Content-Type: image/jpeg` header.
    #[inline]
    pub fn jpeg() -> ContentType {
        ContentType(mime!(Image/Jpeg))
    }

    /// A constructor  to easily create a `Content-Type: image/png` header.
    #[inline]
    pub fn png() -> ContentType {
        ContentType(mime!(Image/Png))
    }
}

bench_header!(bench, ContentType, { vec![b"application/json; charset=utf-8".to_vec()] });
<commit_msg>docs(headers): Explain content_type assigned type<commit_after>use mime::Mime;

header! {
    /// `Content-Type` header, defined in
    /// [RFC7231](http://tools.ietf.org/html/rfc7231#section-3.1.1.5)
    ///
    /// The `Content-Type` header field indicates the media type of the
    /// associated representation: either the representation enclosed in the
    /// message payload or the selected representation, as determined by the
    /// message semantics.  The indicated media type defines both the data
    /// format and how that data is intended to be processed by a recipient,
    /// within the scope of the received message semantics, after any content
    /// codings indicated by Content-Encoding are decoded.
    ///
    /// Although the `mime` crate allows the mime options to be any slice, this crate
    /// forces the use of Vec. This is to make sure the same header can't have more than 1 type. If
    /// this is an issue, it's possible to implement `Header` on a custom struct.
    ///
    /// # ABNF
    /// ```plain
    /// Content-Type = media-type
    /// ```
    ///
    /// # Example values
    /// * `text/html; charset=ISO-8859-4`
    ///
    /// # Examples
    /// ```
    /// use hyper::header::{Headers, ContentType};
    /// use hyper::mime::{Mime, TopLevel, SubLevel};
    ///
    /// let mut headers = Headers::new();
    ///
    /// headers.set(
    ///     ContentType(Mime(TopLevel::Text, SubLevel::Html, vec![]))
    /// );
    /// ```
    /// ```
    /// use hyper::header::{Headers, ContentType};
    /// use hyper::mime::{Mime, TopLevel, SubLevel, Attr, Value};
    ///
    /// let mut headers = Headers::new();
    ///
    /// headers.set(
    ///     ContentType(Mime(TopLevel::Application, SubLevel::Json,
    ///                      vec![(Attr::Charset, Value::Utf8)]))
    /// );
    /// ```
    (ContentType, "Content-Type") => [Mime]

    test_content_type {
        test_header!(
            test1,
            // FIXME: Should be b"text/html; charset=ISO-8859-4" but mime crate lowercases
            // the whole value so parsing and formatting the value gives a different result
            vec![b"text/html; charset=iso-8859-4"],
            Some(HeaderField(Mime(
                TopLevel::Text,
                SubLevel::Html,
                vec![(Attr::Charset, Value::Ext("iso-8859-4".to_owned()))]))));
    }
}

impl ContentType {
    /// A constructor  to easily create a `Content-Type: application/json` header.
    #[inline]
    pub fn json() -> ContentType {
        ContentType(mime!(Application/Json))
    }

    /// A constructor  to easily create a `Content-Type: text/plain; charset=utf-8` header.
    #[inline]
    pub fn plaintext() -> ContentType {
        ContentType(mime!(Text/Plain; Charset=Utf8))
    }

    /// A constructor  to easily create a `Content-Type: text/html; charset=utf-8` header.
    #[inline]
    pub fn html() -> ContentType {
        ContentType(mime!(Text/Html; Charset=Utf8))
    }

    /// A constructor  to easily create a `Content-Type: application/www-form-url-encoded` header.
    #[inline]
    pub fn form_url_encoded() -> ContentType {
        ContentType(mime!(Application/WwwFormUrlEncoded))
    }
    /// A constructor  to easily create a `Content-Type: image/jpeg` header.
    #[inline]
    pub fn jpeg() -> ContentType {
        ContentType(mime!(Image/Jpeg))
    }

    /// A constructor  to easily create a `Content-Type: image/png` header.
    #[inline]
    pub fn png() -> ContentType {
        ContentType(mime!(Image/Png))
    }
}

bench_header!(bench, ContentType, { vec![b"application/json; charset=utf-8".to_vec()] });
<|endoftext|>