<commit_before>// Copyright 2015 Nicholas Bishop
//
// 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.

//! Vec2D is a very simple 2D container for storing rectangular data

#![deny(missing_docs)]

/// 2D coordinate
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Coord {
    /// X component
    pub x: usize,

    /// Y component
    pub y: usize
}

/// Rectangle defined by inclusive minimum and maximum coordinates
#[derive(Clone, Copy, Eq, Debug, PartialEq)]
pub struct Rect {
    /// Minimum coordinate (inclusive)
    min_coord: Coord,

    /// Maximum coordinate (inclusive)
    max_coord: Coord
}

/// Rectangle dimensions
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Size {
    /// Width of rectangle
    pub width: usize,
    /// Height of rectangle
    pub height: usize
}

/// Container for 2D data
#[derive(Clone, Debug)]
pub struct Vec2D<T> {
    elems: Vec<T>,
    size: Size
}

/// Mutable iterator over a Vec2D
pub struct RectIterMut<'a, Elem: 'a> {
    grid: std::marker::PhantomData<&'a mut Vec2D<Elem>>,

    rect: Rect,
    cur_elem: *mut Elem,
    cur_coord: Coord,
    stride: isize
}

impl Coord {
    /// Create a coordinate at (x, y)
    pub fn new(x: usize, y: usize) -> Coord {
        Coord {
            x: x,
            y: y
        }
    }
}

impl Rect {
    /// Calculate rectangle width
    pub fn width(&self) -> usize {
        return self.max_coord.x - self.min_coord.x + 1;
    }

    /// Calculate rectangle height
    pub fn height(&self) -> usize {
        return self.max_coord.y - self.min_coord.y + 1;
    }

    /// Calculate rectangle size
    pub fn size(&self) -> Size {
        Size::new(self.width(), self.height())
    }
}

impl Size {
    /// Create a 2D size of (width, height)
    pub fn new(width: usize, height: usize) -> Size {
        Size {
            width: width,
            height: height
        }
    }

    /// width * height
    pub fn area(&self) -> usize {
        self.width * self.height
    }

    /// Return true if the coordinate fits within self's width and
    /// height, false otherwise.
    pub fn contains_coord(&self, coord: Coord) -> bool {
        coord.x < self.width && coord.y < self.height
    }

    /// Create a rectangle starting at (0, 0) with `self`'s size.
    pub fn rect(&self) -> Rect {
        Rect {
            min_coord: Coord::new(0, 0),
            max_coord: Coord::new(self.width - 1, self.height - 1)
        }
    }
}

impl<Elem: Copy> Vec2D<Elem> {
    /// Create a Vec2D with the given `size`. All elements are
    /// initialized as copies of the `example` element.
    pub fn from_example(size: Size, example: &Elem) -> Vec2D<Elem> {
        Vec2D {
            elems: vec![*example; size.area()],
            size: size
        }
    }

    /// Create a Vec2D with the given `size`. The contents are set to
    /// `src`. None is returned if the `size` does not match the
    /// length of `src`.
    pub fn from_vec(size: Size, src: Vec<Elem>) -> Option<Vec2D<Elem>> {
        if size.area() == src.len() {
            Some(Vec2D {
                elems: src,
                size: size
            })
        }
        else {
            None
        }
    }

    /// Create a mutable iterator over a rectangular region of the
    /// Vec2D. None is returned if the given `rect` does not fit
    /// entirely within the Vec2D.
    pub fn rect_iter_mut<'a>(&'a mut self, rect: Rect) -> Option<RectIterMut<'a, Elem>> {
        if self.size.contains_coord(rect.max_coord) {
            Some(RectIterMut {
                grid: std::marker::PhantomData,
                stride: (self.size.width - rect.width() + 1) as isize,
                cur_elem: self.elems.as_mut_ptr(),
                rect: rect,
                cur_coord: rect.min_coord
            })
        }
        else {
            None
        }
    }
}

impl<'a, Elem> Iterator for RectIterMut<'a, Elem> {
    type Item = (Coord, &'a mut Elem);

    fn next(&mut self) -> Option<Self::Item> {
        if self.cur_coord.y <= self.rect.max_coord.y {
            let result = (self.cur_coord, unsafe { &mut *self.cur_elem });

            self.cur_coord.x += 1;
            if self.cur_coord.x <= self.rect.max_coord.x {
                unsafe { self.cur_elem = self.cur_elem.offset(1); }
            }
            else {
                self.cur_coord.x = self.rect.min_coord.x;
                self.cur_coord.y += 1;
                unsafe { self.cur_elem = self.cur_elem.offset(self.stride); }
            }
            Some(result)
        }
        else {
            None
        }
    }
}

impl Rect {
    /// Create a new Rect defined by inclusive minimum and maximum
    /// coordinates. If min_coord is greater than max_coord on either
    /// axis then None is returned.
    pub fn new(min_coord: Coord, max_coord: Coord) -> Option<Rect> {
        if min_coord.x <= max_coord.x && min_coord.y <= max_coord.y {
            Some(Rect {
                min_coord: min_coord,
                max_coord: max_coord
            })
        }
        else {
            None
        }
    }

}

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

    #[test]
    fn test_coord() {
        let coord = Coord::new(1, 2);
        assert_eq!(coord.x, 1);
        assert_eq!(coord.y, 2);
    }

    #[test]
    fn test_rect() {
        let rect = Rect::new(Coord::new(1, 2), Coord::new(5, 3)).unwrap();
        assert_eq!(rect.width(), 5);
        assert_eq!(rect.height(), 2);

        assert_eq!(rect.width(), rect.size().width);
        assert_eq!(rect.height(), rect.size().height);
    }

    #[test]
    fn test_bad_rect() {
        assert_eq!(Rect::new(Coord::new(2, 1), Coord::new(1, 1)).is_none(),
                   true);
        assert_eq!(Rect::new(Coord::new(1, 2), Coord::new(1, 1)).is_none(),
                   true);
    }

    #[test]
    fn test_size() {
        let size = Size::new(3, 2);
        assert_eq!(size.width, 3);
        assert_eq!(size.height, 2);

        assert_eq!(size.area(), 6);

        assert_eq!(size.contains_coord(Coord::new(1, 1)), true);
        assert_eq!(size.contains_coord(Coord::new(4, 1)), false);
        assert_eq!(size.contains_coord(Coord::new(1, 3)), false);

        let rect = size.rect();
        assert_eq!(rect.min_coord, Coord::new(0, 0));
        assert_eq!(rect.max_coord, Coord::new(2, 1));
    }

    #[test]
    fn test_rect_iter_mut() {
        let elems = vec![1, 2, 3, 4];
        let mut grid = Vec2D::from_vec(Size::new(2, 2), elems).unwrap();
        let rect = Rect::new(Coord::new(0, 0), Coord::new(1, 1)).unwrap();

        let mut actual_coords = Vec::new();
        for (coord, elem) in grid.rect_iter_mut(rect).unwrap() {
            *elem = -(*elem);
            actual_coords.push((coord.x, coord.y));
        }
        assert_eq!(actual_coords, [(0, 0), (1, 0), (0, 1), (1, 1)]);
        assert_eq!(grid.elems, [-1, -2, -3, -4]);
    }
}
<commit_msg>Add non-mutable iterator<commit_after>// Copyright 2015 Nicholas Bishop
//
// 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.

//! Vec2D is a very simple 2D container for storing rectangular data

#![deny(missing_docs)]

/// 2D coordinate
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Coord {
    /// X component
    pub x: usize,

    /// Y component
    pub y: usize
}

/// Rectangle defined by inclusive minimum and maximum coordinates
#[derive(Clone, Copy, Eq, Debug, PartialEq)]
pub struct Rect {
    /// Minimum coordinate (inclusive)
    min_coord: Coord,

    /// Maximum coordinate (inclusive)
    max_coord: Coord
}

/// Rectangle dimensions
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Size {
    /// Width of rectangle
    pub width: usize,
    /// Height of rectangle
    pub height: usize
}

/// Container for 2D data
#[derive(Clone, Debug)]
pub struct Vec2D<T> {
    elems: Vec<T>,
    size: Size
}

/// Iterator over a rectangle within a Vec2D
pub struct RectIter<'a, Elem: 'a> {
    grid: std::marker::PhantomData<&'a Vec2D<Elem>>,

    rect: Rect,
    cur_elem: *const Elem,
    cur_coord: Coord,
    stride: isize
}

/// Mutable iterator over a rectangle within a Vec2D
pub struct RectIterMut<'a, Elem: 'a> {
    grid: std::marker::PhantomData<&'a mut Vec2D<Elem>>,

    rect: Rect,
    cur_elem: *mut Elem,
    cur_coord: Coord,
    stride: isize
}

impl Coord {
    /// Create a coordinate at (x, y)
    pub fn new(x: usize, y: usize) -> Coord {
        Coord {
            x: x,
            y: y
        }
    }
}

impl Rect {
    /// Calculate rectangle width
    pub fn width(&self) -> usize {
        return self.max_coord.x - self.min_coord.x + 1;
    }

    /// Calculate rectangle height
    pub fn height(&self) -> usize {
        return self.max_coord.y - self.min_coord.y + 1;
    }

    /// Calculate rectangle size
    pub fn size(&self) -> Size {
        Size::new(self.width(), self.height())
    }
}

impl Size {
    /// Create a 2D size of (width, height)
    pub fn new(width: usize, height: usize) -> Size {
        Size {
            width: width,
            height: height
        }
    }

    /// width * height
    pub fn area(&self) -> usize {
        self.width * self.height
    }

    /// Return true if the coordinate fits within self's width and
    /// height, false otherwise.
    pub fn contains_coord(&self, coord: Coord) -> bool {
        coord.x < self.width && coord.y < self.height
    }

    /// Create a rectangle starting at (0, 0) with `self`'s size.
    pub fn rect(&self) -> Rect {
        Rect {
            min_coord: Coord::new(0, 0),
            max_coord: Coord::new(self.width - 1, self.height - 1)
        }
    }
}

impl<Elem: Copy> Vec2D<Elem> {
    /// Create a Vec2D with the given `size`. All elements are
    /// initialized as copies of the `example` element.
    pub fn from_example(size: Size, example: &Elem) -> Vec2D<Elem> {
        Vec2D {
            elems: vec![*example; size.area()],
            size: size
        }
    }

    /// Create a Vec2D with the given `size`. The contents are set to
    /// `src`. None is returned if the `size` does not match the
    /// length of `src`.
    pub fn from_vec(size: Size, src: Vec<Elem>) -> Option<Vec2D<Elem>> {
        if size.area() == src.len() {
            Some(Vec2D {
                elems: src,
                size: size
            })
        }
        else {
            None
        }
    }

    fn stride(&self, rect: &Rect) -> isize {
        (self.size.width + 1 - rect.width()) as isize
    }

    /// Create an iterator over a rectangular region of the
    /// Vec2D. None is returned if the given `rect` does not fit
    /// entirely within the Vec2D.
    pub fn rect_iter<'a>(&'a self, rect: Rect) -> Option<RectIter<'a, Elem>> {
        if self.size.contains_coord(rect.max_coord) {
            Some(RectIter {
                grid: std::marker::PhantomData,
                stride: self.stride(&rect),
                cur_elem: self.elems.as_ptr(),
                rect: rect,
                cur_coord: rect.min_coord
            })
        }
        else {
            None
        }
    }

    /// Create a mutable iterator over a rectangular region of the
    /// Vec2D. None is returned if the given `rect` does not fit
    /// entirely within the Vec2D.
    pub fn rect_iter_mut<'a>(&'a mut self, rect: Rect) -> Option<RectIterMut<'a, Elem>> {
        if self.size.contains_coord(rect.max_coord) {
            Some(RectIterMut {
                grid: std::marker::PhantomData,
                stride: self.stride(&rect),
                cur_elem: self.elems.as_mut_ptr(),
                rect: rect,
                cur_coord: rect.min_coord
            })
        }
        else {
            None
        }
    }
}

impl<'a, Elem> Iterator for RectIter<'a, Elem> {
    type Item = (Coord, &'a Elem);

    fn next(&mut self) -> Option<Self::Item> {
        if self.cur_coord.y <= self.rect.max_coord.y {
            let result = (self.cur_coord, unsafe { & *self.cur_elem });

            self.cur_coord.x += 1;
            if self.cur_coord.x <= self.rect.max_coord.x {
                unsafe { self.cur_elem = self.cur_elem.offset(1); }
            }
            else {
                self.cur_coord.x = self.rect.min_coord.x;
                self.cur_coord.y += 1;
                unsafe { self.cur_elem = self.cur_elem.offset(self.stride); }
            }
            Some(result)
        }
        else {
            None
        }
    }
}

impl<'a, Elem> Iterator for RectIterMut<'a, Elem> {
    type Item = (Coord, &'a mut Elem);

    fn next(&mut self) -> Option<Self::Item> {
        if self.cur_coord.y <= self.rect.max_coord.y {
            let result = (self.cur_coord, unsafe { &mut *self.cur_elem });

            self.cur_coord.x += 1;
            if self.cur_coord.x <= self.rect.max_coord.x {
                unsafe { self.cur_elem = self.cur_elem.offset(1); }
            }
            else {
                self.cur_coord.x = self.rect.min_coord.x;
                self.cur_coord.y += 1;
                unsafe { self.cur_elem = self.cur_elem.offset(self.stride); }
            }
            Some(result)
        }
        else {
            None
        }
    }
}

impl Rect {
    /// Create a new Rect defined by inclusive minimum and maximum
    /// coordinates. If min_coord is greater than max_coord on either
    /// axis then None is returned.
    pub fn new(min_coord: Coord, max_coord: Coord) -> Option<Rect> {
        if min_coord.x <= max_coord.x && min_coord.y <= max_coord.y {
            Some(Rect {
                min_coord: min_coord,
                max_coord: max_coord
            })
        }
        else {
            None
        }
    }

}

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

    #[test]
    fn test_coord() {
        let coord = Coord::new(1, 2);
        assert_eq!(coord.x, 1);
        assert_eq!(coord.y, 2);
    }

    #[test]
    fn test_rect() {
        let rect = Rect::new(Coord::new(1, 2), Coord::new(5, 3)).unwrap();
        assert_eq!(rect.width(), 5);
        assert_eq!(rect.height(), 2);

        assert_eq!(rect.width(), rect.size().width);
        assert_eq!(rect.height(), rect.size().height);
    }

    #[test]
    fn test_bad_rect() {
        assert_eq!(Rect::new(Coord::new(2, 1), Coord::new(1, 1)).is_none(),
                   true);
        assert_eq!(Rect::new(Coord::new(1, 2), Coord::new(1, 1)).is_none(),
                   true);
    }

    #[test]
    fn test_size() {
        let size = Size::new(3, 2);
        assert_eq!(size.width, 3);
        assert_eq!(size.height, 2);

        assert_eq!(size.area(), 6);

        assert_eq!(size.contains_coord(Coord::new(1, 1)), true);
        assert_eq!(size.contains_coord(Coord::new(4, 1)), false);
        assert_eq!(size.contains_coord(Coord::new(1, 3)), false);

        let rect = size.rect();
        assert_eq!(rect.min_coord, Coord::new(0, 0));
        assert_eq!(rect.max_coord, Coord::new(2, 1));
    }

    #[test]
    fn test_rect_iter_mut() {
        let elems = vec![1, 2, 3, 4];
        let mut grid = Vec2D::from_vec(Size::new(2, 2), elems).unwrap();
        let rect = Rect::new(Coord::new(0, 0), Coord::new(1, 1)).unwrap();

        let mut actual_coords = Vec::new();
        for (coord, elem) in grid.rect_iter_mut(rect).unwrap() {
            *elem = -(*elem);
            actual_coords.push((coord.x, coord.y));
        }
        assert_eq!(actual_coords, [(0, 0), (1, 0), (0, 1), (1, 1)]);
        assert_eq!(grid.elems, [-1, -2, -3, -4]);
    }

    #[test]
    fn test_two_iterators() {
        let size = Size::new(2, 1);
        let v = Vec2D::from_vec(size, vec![0, 1]).unwrap();

        let iter1 = v.rect_iter(size.rect()).unwrap();
        let iter2 = v.rect_iter(size.rect()).unwrap();

        for ((coord1, elem1), (coord2, elem2)) in iter1.zip(iter2) {
            assert_eq!(coord1, coord2);
            assert_eq!(elem1, elem2);
        }
    }
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0
// <http://www.apache.org/licenses/LICENSE-2.0>

//! # Rust Random Choice
//! This is an implementation of the stochastic universal sampling algorithm:
//! https://en.wikipedia.org/wiki/Stochastic_universal_sampling
//!
//! ## Advantages
//! - Blazingly fast: O(n) (Roulette wheel selection algorithm: O(n * log n))
//! - Low Memory Usage: O(n); in place variant: O(1)
//! - There is a good diversity for the case, that all weights are equally distributed (in contrast to the roulette wheel selection algorithm which tends to select the same sample n times)
//! - The sum of the weights don't have to be 1.0, but must not overflow
//!
//! ## Applications
//! - **Evolutionary algorithms**: Choose the _n_ fittest populations by their fitness **_fi_**
//! - **Monte Carlo Localization**: Resampling of _n_ particles by their weight **_w_**
//!
//! # Examples
//! ## In Place Variant
//! ```
//! extern crate random_choice;
//! use self::random_choice::random_choice;
//!
//! # fn main() {
//! let mut samples = vec!["hi", "this", "is", "a", "test!"];
//! let weights: Vec<f64> = vec![5.6, 7.8, 9.7, 1.1, 2.0];
//!
//! random_choice().random_choice_in_place_f64(&mut samples, &weights);
//!
//! for sample in samples {
//!     print!("{}, ", sample);
//! }
//! # }
//! ```
//! ## N Selection Variant
//! ```
//! extern crate random_choice;
//! use self::random_choice::random_choice;
//!
//!# fn main() {
//! let capacity: usize = 500;
//! let mut samples: Vec<usize> = Vec::with_capacity(capacity);
//! let mut weights: Vec<f64> = Vec::with_capacity(capacity);
//!
//! for i in 0..capacity {
//!     samples.push(i);
//!     weights.push(i as f64);
//! }
//!
//! let number_choices = 10000;
//! let choices = random_choice().random_choice_f64(&samples, &weights, number_choices);
//!
//! assert!(choices.len() == number_choices);
//!
//! for choice in choices {
//!     print!("{}, ", choice);
//! }
//! # }
//! ```
//! ## With custom seed
//! ```
//! extern crate rand;
//! extern crate random_choice;
//! use random_choice::RandomChoice;
//! use rand::SeedableRng;

//! fn main() {
//!     let mut samples = vec!["hi", "this", "is", "a", "test!"];
//!     let weights: Vec<f64> = vec![5.6, 7.8, 9.7, 1.1, 2.0];
//!
//!     let rng = rand::StdRng::from_seed(&[5000, 44, 55, 199]);
//!
//!     let mut random_choice = RandomChoice::new(rng);
//!     random_choice.random_choice_in_place_f64(&mut samples, &weights);

//!     for sample in samples {
//!         print!("{}, ", sample);
//!     }
//!#}
//! ```

extern crate rand;

use self::rand::{thread_rng, ThreadRng, Rng};

pub struct RandomChoice<RNG: Rng> {
    rng: RNG,
}

/// Creates a new RandomChoice struct using the ThreadRng
pub fn random_choice() -> RandomChoice<ThreadRng> {
    RandomChoice::new(thread_rng())
}

impl<RNG: Rng> RandomChoice<RNG> {
    /// Creates a new RandomChoice struct.
    /// @param rng the random number generator to use with this stuct.
    pub fn new(rng: RNG) -> Self {
        RandomChoice { rng: rng }
    }

    /// Chooses n samples by their weights. The greater their weights the more likely they get chosen.
    ///
    /// @invariant sum of weights must not overflow.
    /// @param samples The to be selected samples
    /// @param weights Weights that get chosen by their weight/probability. One weight can be greater 1.
    /// @param n Number of randomly chosen samples by weight.
    /// @return randomly selected samples by their weights
    pub fn random_choice_f64<'a, T>(&mut self,
                                    samples: &'a [T],
                                    weights: &[f64],
                                    n: usize)
                                    -> Vec<&'a T> {
        if weights.len() == 0 || n == 0 {
            return Vec::new();
        }

        let sum: f64 = weights.iter().fold(0.0, |acc, &i| acc + i);
        let spoke_gap: f64 = sum / n as f64;

        // next_f64() ∈ [0.0, 1.0)
        let spin = self.rng.next_f64() * spoke_gap;

        let mut i: usize = 0;
        let mut accumulated_weights = weights[0];
        let mut choices: Vec<&T> = Vec::with_capacity(n);
        let mut current_spoke: f64 = spin;

        while current_spoke < sum {
            while accumulated_weights < current_spoke {
                i += 1;
                accumulated_weights += weights[i];
            }
            choices.push(&samples[i]);
            current_spoke += spoke_gap;
        }

        // add this condition, because float leads to inaccurate
        // calculations which can miss some samples
        while choices.len() < weights.len() {
            choices.push(&samples[i]);
        }

        choices
    }

    /// Chooses n samples by their weights. The greater their weights the more likely they get chosen.
    /// The result gets saved directly in the samples argument.
    /// @invariant sum of weights must not overflow.
    /// @param samples The to be selected samples
    /// @param weights Weights that get chosen by their weight/probability. One weight can be greater 1.
    pub fn random_choice_in_place_f64<T: Clone>(&mut self, samples: &mut [T], weights: &[f64]) {
        if weights.len() < 2 {
            return;
        }

        let sum: f64 = weights.iter().fold(0.0, |acc, &i| acc + i);
        let n: usize = weights.len();
        let spoke_gap: f64 = sum / n as f64;

        // next_f64() ∈ [0.0, 1.0)
        let spin = self.rng.next_f64() * spoke_gap;

        let mut i: usize = 0;
        let mut j: usize = 0;
        let mut accumulated_weights = weights[0];
        let mut current_spoke: f64 = spin;

        while current_spoke < sum {
            while accumulated_weights < current_spoke {
                j += 1;
                accumulated_weights += weights[j];
            }
            samples[i] = samples[j].clone();
            current_spoke += spoke_gap;
            i += 1
        }

        // add this condition, because float leads to inaccurate
        // calculations which can miss some samples
        while i < weights.len() {
            samples[i] = samples[j].clone();
            i += 1
        }

    }

    pub fn random_choice_f32<'a, T>(&mut self,
                                    samples: &'a [T],
                                    weights: &[f32],
                                    n: usize)
                                    -> Vec<&'a T> {
        if weights.len() == 0 || n == 0 {
            return Vec::new();
        }

        let sum: f64 = weights.iter().fold(0.0, |acc, &i| acc + i as f64);
        let spoke_gap: f64 = sum / n as f64;

        // next_f64() ∈ [0.0, 1.0)
        let spin = self.rng.next_f64() * spoke_gap;

        let mut i: usize = 0;
        let mut accumulated_weights = weights[0] as f64;
        let mut choices: Vec<&T> = Vec::with_capacity(n);
        let mut current_spoke: f64 = spin;

        while current_spoke < sum {
            while accumulated_weights < current_spoke {
                i += 1;
                accumulated_weights += weights[i] as f64;
            }
            choices.push(&samples[i]);
            current_spoke += spoke_gap;
        }

        // add this condition, because float leads to inaccurate
        // calculations which can miss some samples
        while choices.len() < weights.len() {
            choices.push(&samples[i]);
        }

        choices
    }

    pub fn random_choice_in_place_f32<T: Clone>(&mut self, samples: &mut [T], weights: &[f32]) {
        if weights.len() < 2 {
            return;
        }

        let sum: f64 = weights.iter().fold(0.0, |acc, &i| acc + i as f64);
        let n: usize = weights.len();
        let spoke_gap: f64 = sum / n as f64;

        // next_f64() ∈ [0.0, 1.0)
        let spin = self.rng.next_f64() * spoke_gap;

        let mut i: usize = 0;
        let mut j: usize = 0;
        let mut accumulated_weights = weights[0] as f64;
        let mut current_spoke: f64 = spin;

        while current_spoke < sum {
            while accumulated_weights < current_spoke {
                j += 1;
                accumulated_weights += weights[j] as f64;
            }
            samples[i] = samples[j].clone();
            current_spoke += spoke_gap;
            i += 1
        }

        // add this condition, because float leads to inaccurate
        // calculations which can miss some samples
        while i < weights.len() {
            samples[i] = samples[j].clone();
            i += 1
        }
    }
}<commit_msg>Fixed a test failure I only just noticed.<commit_after>// Licensed under the Apache License, Version 2.0
// <http://www.apache.org/licenses/LICENSE-2.0>

//! # Rust Random Choice
//! This is an implementation of the stochastic universal sampling algorithm:
//! https://en.wikipedia.org/wiki/Stochastic_universal_sampling
//!
//! ## Advantages
//! - Blazingly fast: O(n) (Roulette wheel selection algorithm: O(n * log n))
//! - Low Memory Usage: O(n); in place variant: O(1)
//! - There is a good diversity for the case, that all weights are equally distributed (in contrast to the roulette wheel selection algorithm which tends to select the same sample n times)
//! - The sum of the weights don't have to be 1.0, but must not overflow
//!
//! ## Applications
//! - **Evolutionary algorithms**: Choose the _n_ fittest populations by their fitness **_fi_**
//! - **Monte Carlo Localization**: Resampling of _n_ particles by their weight **_w_**
//!
//! # Examples
//! ## In Place Variant
//! ```
//! extern crate random_choice;
//! use self::random_choice::random_choice;
//!
//! # fn main() {
//! let mut samples = vec!["hi", "this", "is", "a", "test!"];
//! let weights: Vec<f64> = vec![5.6, 7.8, 9.7, 1.1, 2.0];
//!
//! random_choice().random_choice_in_place_f64(&mut samples, &weights);
//!
//! for sample in samples {
//!     print!("{}, ", sample);
//! }
//! # }
//! ```
//! ## N Selection Variant
//! ```
//! extern crate random_choice;
//! use self::random_choice::random_choice;
//!
//! # fn main() {
//! let capacity: usize = 500;
//! let mut samples: Vec<usize> = Vec::with_capacity(capacity);
//! let mut weights: Vec<f64> = Vec::with_capacity(capacity);
//!
//! for i in 0..capacity {
//!     samples.push(i);
//!     weights.push(i as f64);
//! }
//!
//! let number_choices = 10000;
//! let choices = random_choice().random_choice_f64(&samples, &weights, number_choices);
//!
//! assert!(choices.len() == number_choices);
//!
//! for choice in choices {
//!     print!("{}, ", choice);
//! }
//! # }
//! ```
//! ## With custom seed
//! ```
//! extern crate rand;
//! extern crate random_choice;
//! use random_choice::RandomChoice;
//! use rand::SeedableRng;
//!
//! fn main() {
//!     let mut samples = vec!["hi", "this", "is", "a", "test!"];
//!     let weights: Vec<f64> = vec![5.6, 7.8, 9.7, 1.1, 2.0];
//!
//!     let rng = rand::StdRng::from_seed(&[5000, 44, 55, 199]);
//!
//!     let mut random_choice = RandomChoice::new(rng);
//!     random_choice.random_choice_in_place_f64(&mut samples, &weights);
//!
//!     for sample in samples {
//!         print!("{}, ", sample);
//!     }
//! }
//! ```

extern crate rand;

use self::rand::{thread_rng, ThreadRng, Rng};

pub struct RandomChoice<RNG: Rng> {
    rng: RNG,
}

/// Creates a new RandomChoice struct using the ThreadRng
pub fn random_choice() -> RandomChoice<ThreadRng> {
    RandomChoice::new(thread_rng())
}

impl<RNG: Rng> RandomChoice<RNG> {
    /// Creates a new RandomChoice struct.
    /// @param rng the random number generator to use with this stuct.
    pub fn new(rng: RNG) -> Self {
        RandomChoice { rng: rng }
    }

    /// Chooses n samples by their weights. The greater their weights the more likely they get chosen.
    ///
    /// @invariant sum of weights must not overflow.
    /// @param samples The to be selected samples
    /// @param weights Weights that get chosen by their weight/probability. One weight can be greater 1.
    /// @param n Number of randomly chosen samples by weight.
    /// @return randomly selected samples by their weights
    pub fn random_choice_f64<'a, T>(&mut self,
                                    samples: &'a [T],
                                    weights: &[f64],
                                    n: usize)
                                    -> Vec<&'a T> {
        if weights.len() == 0 || n == 0 {
            return Vec::new();
        }

        let sum: f64 = weights.iter().fold(0.0, |acc, &i| acc + i);
        let spoke_gap: f64 = sum / n as f64;

        // next_f64() ∈ [0.0, 1.0)
        let spin = self.rng.next_f64() * spoke_gap;

        let mut i: usize = 0;
        let mut accumulated_weights = weights[0];
        let mut choices: Vec<&T> = Vec::with_capacity(n);
        let mut current_spoke: f64 = spin;

        while current_spoke < sum {
            while accumulated_weights < current_spoke {
                i += 1;
                accumulated_weights += weights[i];
            }
            choices.push(&samples[i]);
            current_spoke += spoke_gap;
        }

        // add this condition, because float leads to inaccurate
        // calculations which can miss some samples
        while choices.len() < weights.len() {
            choices.push(&samples[i]);
        }

        choices
    }

    /// Chooses n samples by their weights. The greater their weights the more likely they get chosen.
    /// The result gets saved directly in the samples argument.
    /// @invariant sum of weights must not overflow.
    /// @param samples The to be selected samples
    /// @param weights Weights that get chosen by their weight/probability. One weight can be greater 1.
    pub fn random_choice_in_place_f64<T: Clone>(&mut self, samples: &mut [T], weights: &[f64]) {
        if weights.len() < 2 {
            return;
        }

        let sum: f64 = weights.iter().fold(0.0, |acc, &i| acc + i);
        let n: usize = weights.len();
        let spoke_gap: f64 = sum / n as f64;

        // next_f64() ∈ [0.0, 1.0)
        let spin = self.rng.next_f64() * spoke_gap;

        let mut i: usize = 0;
        let mut j: usize = 0;
        let mut accumulated_weights = weights[0];
        let mut current_spoke: f64 = spin;

        while current_spoke < sum {
            while accumulated_weights < current_spoke {
                j += 1;
                accumulated_weights += weights[j];
            }
            samples[i] = samples[j].clone();
            current_spoke += spoke_gap;
            i += 1
        }

        // add this condition, because float leads to inaccurate
        // calculations which can miss some samples
        while i < weights.len() {
            samples[i] = samples[j].clone();
            i += 1
        }

    }

    pub fn random_choice_f32<'a, T>(&mut self,
                                    samples: &'a [T],
                                    weights: &[f32],
                                    n: usize)
                                    -> Vec<&'a T> {
        if weights.len() == 0 || n == 0 {
            return Vec::new();
        }

        let sum: f64 = weights.iter().fold(0.0, |acc, &i| acc + i as f64);
        let spoke_gap: f64 = sum / n as f64;

        // next_f64() ∈ [0.0, 1.0)
        let spin = self.rng.next_f64() * spoke_gap;

        let mut i: usize = 0;
        let mut accumulated_weights = weights[0] as f64;
        let mut choices: Vec<&T> = Vec::with_capacity(n);
        let mut current_spoke: f64 = spin;

        while current_spoke < sum {
            while accumulated_weights < current_spoke {
                i += 1;
                accumulated_weights += weights[i] as f64;
            }
            choices.push(&samples[i]);
            current_spoke += spoke_gap;
        }

        // add this condition, because float leads to inaccurate
        // calculations which can miss some samples
        while choices.len() < weights.len() {
            choices.push(&samples[i]);
        }

        choices
    }

    pub fn random_choice_in_place_f32<T: Clone>(&mut self, samples: &mut [T], weights: &[f32]) {
        if weights.len() < 2 {
            return;
        }

        let sum: f64 = weights.iter().fold(0.0, |acc, &i| acc + i as f64);
        let n: usize = weights.len();
        let spoke_gap: f64 = sum / n as f64;

        // next_f64() ∈ [0.0, 1.0)
        let spin = self.rng.next_f64() * spoke_gap;

        let mut i: usize = 0;
        let mut j: usize = 0;
        let mut accumulated_weights = weights[0] as f64;
        let mut current_spoke: f64 = spin;

        while current_spoke < sum {
            while accumulated_weights < current_spoke {
                j += 1;
                accumulated_weights += weights[j] as f64;
            }
            samples[i] = samples[j].clone();
            current_spoke += spoke_gap;
            i += 1
        }

        // add this condition, because float leads to inaccurate
        // calculations which can miss some samples
        while i < weights.len() {
            samples[i] = samples[j].clone();
            i += 1
        }
    }
}<|endoftext|>


<commit_before>// `error_chain!` can recurse deeply
#![recursion_limit = "1024"]
#[macro_use]
extern crate error_chain;

// We'll put our errors in an `errors` module, and other modules in
// this crate will `use errors::*;` to get access to everything
// `error_chain!` creates.
pub mod errors {
    // Create the Error, ErrorKind, ResultExt, and Result types
    error_chain!{
        foreign_links {
            Io(::std::io::Error) #[cfg(unix)];
        }
    }
}
use errors::*;

extern crate libinspire;
extern crate libads;

/// Re-export slog
///
/// Users of this library can, but don't have to use slog to build their own loggers
#[macro_use]
pub extern crate slog ;
extern crate slog_stdlog;

use slog::DrainExt;

#[macro_use]
extern crate lazy_static;
extern crate regex;
use regex::Regex;

use std::fs::File;
use std::io::{Read, BufReader};
use std::io::{Write, BufWriter};

pub struct Inspirer {
    logger: slog::Logger,
    inspire: libinspire::Api,
    ads: libads::Api,
}

impl Inspirer {
    /// Initialize 'Inspirer'
    ///
    /// Either provide a custom slog::Logger or default to the standard `log`
    /// crate.
    ///
    /// # Examples
    /// ```
    /// inspirer::Inspirer::init(None);
    /// ```
    pub fn init(logger: Option<slog::Logger>) -> Self {
        let logger = logger.unwrap_or_else(|| slog::Logger::root(slog_stdlog::StdLog.fuse(), o!()));

        Inspirer {
            logger: logger,
            // inspire: libinspire::Api::init(Some(logger)),
            inspire: libinspire::Api::init(None),
            ads: libads::Api::init(None),
        }
    }

    /// Read input from file or stdin
    ///
    /// # Examples
    /// ```
    /// ```
    pub fn get_input(&self, input_source: Option<&str>) -> Result<String> {
        let mut input_data = String::new();

        let mut input_file: File;
        let mut stdin = std::io::stdin();

        let reader: &mut Read = match input_source {
            Some(file_name) => {
                info!(self.logger, "Reading from file";
                      "file_name" => file_name);
                input_file = File::open(file_name)
                    .chain_err(|| format!("File \"{}\" not found", file_name))?;
                &mut input_file
            }
            None => {
                info!(self.logger, "Reading from stdin");
                &mut stdin
            }
        };
        let mut reader = BufReader::new(reader);
        reader
            .read_to_string(&mut input_data)
            .chain_err(|| "Could not read input")?;

        Ok(input_data)
    }

    /// Write output to file or stdout
    pub fn put_output(&self, output_dest: Option<&str>, output: Vec<String>) -> Result<()> {
        let mut stdout = std::io::stdout();
        let mut output_file: std::fs::File;

        let writer: &mut Write = match output_dest {
            Some(file_name) => {
                info!(self.logger, "Writing to file";
                      "file_name" => file_name);
                output_file = std::fs::OpenOptions::new()
                    .append(true)
                    .create(true)
                    .open(file_name)
                    .chain_err(|| format!("Could not open file \"{}\" to write", file_name))?;
                &mut output_file
            }
            None => {
                info!(self.logger, "Writing to stdout");
                // stdout.lock();
                &mut stdout
            }
        };

        let mut writer = BufWriter::new(writer);

        for o in output {
            writer
                .write_all(o.as_bytes())
                .chain_err(|| "Failed to write output")?;
        }

        writer.flush().chain_err(|| "Failed to write output")?;

        Ok(())
    }

    /// The `aux2key` function extracts TeX keys from LaTeX .aux files. These can be for either
    /// BibTeX or BibLaTeX.
    ///
    /// # Examples
    ///
    /// ## bibtex
    ///
    /// Inspire-formatted BibTeX key:
    ///
    /// ```
    /// let inspirer = inspirer::Inspirer::init(None);
    ///
    /// let input =
    /// r"\relax
    /// \citation{Abramovici:1992ah}".to_string();
    ///
    /// assert_eq!(inspirer.aux2key(input), vec!("Abramovici:1992ah"));
    /// ```
    ///
    /// ADS-formatted BibTeX Key:
    ///
    /// ```
    /// let inspirer = inspirer::Inspirer::init(None);
    ///
    /// let input =
    /// r"\relax
    /// \citation{1998PhRvD..58h4020O}".to_string();
    ///
    /// assert_eq!(inspirer.aux2key(input), vec!("1998PhRvD..58h4020O"));
    /// ```
    ///
    /// ## biber
    ///
    /// Inspire-formatted BibLaTeX key:
    ///
    /// ```
    /// let inspirer = inspirer::Inspirer::init(None);
    ///
    /// let input =
    /// r"\relax
    /// \abx@aux@cite{Cutler:1992tc}".to_string();
    ///
    /// assert_eq!(inspirer.aux2key(input), vec!("Cutler:1992tc"));
    /// ```
    pub fn aux2key(&self, input_data: String) -> Vec<String> {

        lazy_static! {
            // TODO: check on the exact characters allowed in keys
            static ref AUX_REGEX: Regex = Regex::new(
                r"(\\citation|\\abx@aux@cite)\{(.+)\}").expect("aux regex compiled during development");
        }

        AUX_REGEX
            .captures_iter(&input_data)
            .map(|c| c.get(2).unwrap().as_str().to_string())
            // TODO just return the iterator: makes more sense with rayon
            .collect()
    }

    /// The blg2key function extracts missing references from bibtex logs
    ///
    /// # Examples
    ///
    /// ADS-formatted BibTeX key:
    ///
    /// ```
    /// let inspirer = inspirer::Inspirer::init(None);
    ///
    /// let input =
    /// r##"
    /// This is BibTeX, Version 0.99d (TeX Live 2016/Arch Linux)
    /// Capacity: max_strings=35307, hash_size=35307, hash_prime=30011
    /// The top-level auxiliary file: test_bibtex.aux
    /// The style file: unsrt.bst
    /// Database file #1: test_bibtex.bib
    /// Warning--I didn't find a database entry for "2015CQGra..32g4001L"
    /// You've used 0 entries,
    /// ....
    /// "##.to_string();
    ///
    /// assert_eq!(inspirer.blg2key(input), vec!("2015CQGra..32g4001L"));
    /// ```
    pub fn blg2key(&self, input_data: String) -> Vec<String> {

        lazy_static! {
            static ref BLG_REGEX: Regex = Regex::new(
                r#"(Warning--|WARN - )I didn't find a database entry for ["'](.+)["']"#,
            ).expect("blg regex compiled during development");
        }

        BLG_REGEX
            .captures_iter(&input_data)
            .map(|c| c.get(2).unwrap().as_str().to_string())
            // TODO just return the iterator: makes more sense with rayon
            .collect()
    }

    /// Fetch BibTeX entries
    pub fn bibtex(&self, key: &str) -> Option<String> {
        let key = Sources::from(key);

        match key {
            Sources::Inspire(k) => {
                debug!(self.logger, "Record type: Inspire"; "key" => k.id);
                self.inspire.fetch_bibtex_with_key(k)
            }
            Sources::Ads(k) => {
                debug!(self.logger, "Record type: ADS"; "key" => k.bibcode);
                self.ads.fetch_bibtex_with_key(k)
            }
            _ => {
                // debug!(self.logger, "Unknown record source"; "key" => key);
                debug!(self.logger, "Record type: unknown");
                None
            }
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum Sources<'a> {
    Inspire(libinspire::RecID<'a>),
    Ads(libads::BibCode<'a>),
    Arxiv,
    None,
}

/// Guess a likely source for a BibTeX key
///
/// Returns `Sources::None` if unable to make a good guess.
///
/// # Examples
/// ```
/// extern crate inspirer;
/// extern crate libinspire;
/// let inspirer = inspirer::Inspirer::init(None);
///
/// assert_eq!(
///     inspirer::Sources::from("Randall:1999ee"),
///     inspirer::Sources::Inspire(libinspire::RecID::new("Randall:1999ee").unwrap())
/// );
/// ```
///
/// ```
/// extern crate inspirer;
/// extern crate libads;
/// let inspirer = inspirer::Inspirer::init(None);
///
/// assert_eq!(
///     inspirer::Sources::from("1999PhRvL..83.3370R"),
///     inspirer::Sources::Ads(libads::BibCode::new("1999PhRvL..83.3370R").unwrap())
/// );
/// ```
impl<'a> From<&'a str> for Sources<'a> {
    fn from(s: &'a str) -> Sources<'a> {
        if libinspire::validate_recid(s) {
            Sources::Inspire(libinspire::RecID::new(s).unwrap())
        } else if libads::validate_bib_code(s) {
            Sources::Ads(libads::BibCode::new(s).unwrap())
        } else {
            Sources::None
        }
    }
}

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

    #[test]
    fn test_aux_bibtex_0_citations() {
        let input = r"\relax ".to_string();

        let output: Vec<String> = Vec::new();
        assert_eq!(Inspirer::init(None).aux2key(input), output);
    }

    #[test]
    fn test_aux_bibtex_1_citation() {
        let input = r"\relax
            \citation{Abramovici:1992ah}"
            .to_string();

        assert_eq!(
            Inspirer::init(None).aux2key(input),
            vec!["Abramovici:1992ah"]
        );
    }

    #[test]
    fn test_aux_bibtex_2_citation() {
        let input = r"\relax
            \citation{Abramovici:1992ah}
            \citation{Thorne:1992sdb}"
            .to_string();

        assert_eq!(
            Inspirer::init(None).aux2key(input),
            vec!["Abramovici:1992ah", "Thorne:1992sdb"]
        );
    }
}
<commit_msg>fix: Items cited together should be split at comma<commit_after>// `error_chain!` can recurse deeply
#![recursion_limit = "1024"]
#[macro_use]
extern crate error_chain;

// We'll put our errors in an `errors` module, and other modules in
// this crate will `use errors::*;` to get access to everything
// `error_chain!` creates.
pub mod errors {
    // Create the Error, ErrorKind, ResultExt, and Result types
    error_chain!{
        foreign_links {
            Io(::std::io::Error) #[cfg(unix)];
        }
    }
}
use errors::*;

extern crate libinspire;
extern crate libads;

/// Re-export slog
///
/// Users of this library can, but don't have to use slog to build their own loggers
#[macro_use]
pub extern crate slog ;
extern crate slog_stdlog;

use slog::DrainExt;

#[macro_use]
extern crate lazy_static;
extern crate regex;
use regex::Regex;

use std::fs::File;
use std::io::{Read, BufReader};
use std::io::{Write, BufWriter};

pub struct Inspirer {
    logger: slog::Logger,
    inspire: libinspire::Api,
    ads: libads::Api,
}

impl Inspirer {
    /// Initialize 'Inspirer'
    ///
    /// Either provide a custom slog::Logger or default to the standard `log`
    /// crate.
    ///
    /// # Examples
    /// ```
    /// inspirer::Inspirer::init(None);
    /// ```
    pub fn init(logger: Option<slog::Logger>) -> Self {
        let logger = logger.unwrap_or_else(|| slog::Logger::root(slog_stdlog::StdLog.fuse(), o!()));

        Inspirer {
            logger: logger,
            // inspire: libinspire::Api::init(Some(logger)),
            inspire: libinspire::Api::init(None),
            ads: libads::Api::init(None),
        }
    }

    /// Read input from file or stdin
    ///
    /// # Examples
    /// ```
    /// ```
    pub fn get_input(&self, input_source: Option<&str>) -> Result<String> {
        let mut input_data = String::new();

        let mut input_file: File;
        let mut stdin = std::io::stdin();

        let reader: &mut Read = match input_source {
            Some(file_name) => {
                info!(self.logger, "Reading from file";
                      "file_name" => file_name);
                input_file = File::open(file_name)
                    .chain_err(|| format!("File \"{}\" not found", file_name))?;
                &mut input_file
            }
            None => {
                info!(self.logger, "Reading from stdin");
                &mut stdin
            }
        };
        let mut reader = BufReader::new(reader);
        reader
            .read_to_string(&mut input_data)
            .chain_err(|| "Could not read input")?;

        Ok(input_data)
    }

    /// Write output to file or stdout
    pub fn put_output(&self, output_dest: Option<&str>, output: Vec<String>) -> Result<()> {
        let mut stdout = std::io::stdout();
        let mut output_file: std::fs::File;

        let writer: &mut Write = match output_dest {
            Some(file_name) => {
                info!(self.logger, "Writing to file";
                      "file_name" => file_name);
                output_file = std::fs::OpenOptions::new()
                    .append(true)
                    .create(true)
                    .open(file_name)
                    .chain_err(|| format!("Could not open file \"{}\" to write", file_name))?;
                &mut output_file
            }
            None => {
                info!(self.logger, "Writing to stdout");
                // stdout.lock();
                &mut stdout
            }
        };

        let mut writer = BufWriter::new(writer);

        for o in output {
            writer
                .write_all(o.as_bytes())
                .chain_err(|| "Failed to write output")?;
        }

        writer.flush().chain_err(|| "Failed to write output")?;

        Ok(())
    }

    /// The `aux2key` function extracts TeX keys from LaTeX .aux files. These can be for either
    /// BibTeX or BibLaTeX.
    ///
    /// # Examples
    ///
    /// ## bibtex
    ///
    /// Inspire-formatted BibTeX key:
    ///
    /// ```
    /// let inspirer = inspirer::Inspirer::init(None);
    ///
    /// let input =
    /// r"\relax
    /// \citation{Abramovici:1992ah}".to_string();
    ///
    /// assert_eq!(inspirer.aux2key(input), vec!("Abramovici:1992ah"));
    /// ```
    ///
    /// ADS-formatted BibTeX Key:
    ///
    /// ```
    /// let inspirer = inspirer::Inspirer::init(None);
    ///
    /// let input =
    /// r"\relax
    /// \citation{1998PhRvD..58h4020O}".to_string();
    ///
    /// assert_eq!(inspirer.aux2key(input), vec!("1998PhRvD..58h4020O"));
    /// ```
    ///
    /// ## biber
    ///
    /// Inspire-formatted BibLaTeX key:
    ///
    /// ```
    /// let inspirer = inspirer::Inspirer::init(None);
    ///
    /// let input =
    /// r"\relax
    /// \abx@aux@cite{Cutler:1992tc}".to_string();
    ///
    /// assert_eq!(inspirer.aux2key(input), vec!("Cutler:1992tc"));
    /// ```
    pub fn aux2key(&self, input_data: String) -> Vec<String> {

        lazy_static! {
            // TODO: check on the exact characters allowed in keys
            // Just find groups of keys which are cited together
            static ref AUX_REGEX: Regex = Regex::new(
                r"(?:\\citation|\\abx@aux@cite)\{(?P<key>.+)\}").expect("aux regex compiled during development");
        }

        lazy_static! {
            // TODO: check on the exact characters allowed in keys
            // Split keys which are cited together
            static ref INNER_REGEX: Regex = Regex::new(
                r"(?P<key>[^,]+),?").expect("aux regex compiled during development");
        }

        AUX_REGEX
            .captures_iter(&input_data)
            .map(|c| c["key"].to_string())
            .collect::<Vec<String>>()
            .iter()
            .flat_map(|s| {
                INNER_REGEX.captures_iter(s).map(|c| c["key"].to_string())
            })
            .collect()
    }

    /// The blg2key function extracts missing references from bibtex logs
    ///
    /// # Examples
    ///
    /// ADS-formatted BibTeX key:
    ///
    /// ```
    /// let inspirer = inspirer::Inspirer::init(None);
    ///
    /// let input =
    /// r##"
    /// This is BibTeX, Version 0.99d (TeX Live 2016/Arch Linux)
    /// Capacity: max_strings=35307, hash_size=35307, hash_prime=30011
    /// The top-level auxiliary file: test_bibtex.aux
    /// The style file: unsrt.bst
    /// Database file #1: test_bibtex.bib
    /// Warning--I didn't find a database entry for "2015CQGra..32g4001L"
    /// You've used 0 entries,
    /// ....
    /// "##.to_string();
    ///
    /// assert_eq!(inspirer.blg2key(input), vec!("2015CQGra..32g4001L"));
    /// ```
    pub fn blg2key(&self, input_data: String) -> Vec<String> {

        lazy_static! {
            static ref BLG_REGEX: Regex = Regex::new(
                r#"(Warning--|WARN - )I didn't find a database entry for ["'](.+)["']"#,
            ).expect("blg regex compiled during development");
        }

        BLG_REGEX
            .captures_iter(&input_data)
            .map(|c| c.get(2).unwrap().as_str().to_string())
            // TODO just return the iterator: makes more sense with rayon
            .collect()
    }

    /// Fetch BibTeX entries
    pub fn bibtex(&self, key: &str) -> Option<String> {
        let key = Sources::from(key);

        match key {
            Sources::Inspire(k) => {
                debug!(self.logger, "Record type: Inspire"; "key" => k.id);
                self.inspire.fetch_bibtex_with_key(k)
            }
            Sources::Ads(k) => {
                debug!(self.logger, "Record type: ADS"; "key" => k.bibcode);
                self.ads.fetch_bibtex_with_key(k)
            }
            _ => {
                // debug!(self.logger, "Unknown record source"; "key" => key);
                debug!(self.logger, "Record type: unknown");
                None
            }
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum Sources<'a> {
    Inspire(libinspire::RecID<'a>),
    Ads(libads::BibCode<'a>),
    Arxiv,
    None,
}

/// Guess a likely source for a BibTeX key
///
/// Returns `Sources::None` if unable to make a good guess.
///
/// # Examples
/// ```
/// extern crate inspirer;
/// extern crate libinspire;
/// let inspirer = inspirer::Inspirer::init(None);
///
/// assert_eq!(
///     inspirer::Sources::from("Randall:1999ee"),
///     inspirer::Sources::Inspire(libinspire::RecID::new("Randall:1999ee").unwrap())
/// );
/// ```
///
/// ```
/// extern crate inspirer;
/// extern crate libads;
/// let inspirer = inspirer::Inspirer::init(None);
///
/// assert_eq!(
///     inspirer::Sources::from("1999PhRvL..83.3370R"),
///     inspirer::Sources::Ads(libads::BibCode::new("1999PhRvL..83.3370R").unwrap())
/// );
/// ```
impl<'a> From<&'a str> for Sources<'a> {
    fn from(s: &'a str) -> Sources<'a> {
        if libinspire::validate_recid(s) {
            Sources::Inspire(libinspire::RecID::new(s).unwrap())
        } else if libads::validate_bib_code(s) {
            Sources::Ads(libads::BibCode::new(s).unwrap())
        } else {
            Sources::None
        }
    }
}

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

    #[test]
    fn test_aux_bibtex_0_citations() {
        let input = r"\relax ".to_string();

        let output: Vec<String> = Vec::new();
        assert_eq!(Inspirer::init(None).aux2key(input), output);
    }

    #[test]
    fn test_aux_bibtex_1_citation() {
        let input = r"\relax
            \citation{Abramovici:1992ah}"
            .to_string();

        assert_eq!(
            Inspirer::init(None).aux2key(input),
            vec!["Abramovici:1992ah"]
        );
    }

    #[test]
    fn test_aux_bibtex_2_citation() {
        let input = r"\relax
            \citation{Abramovici:1992ah}
            \citation{Thorne:1992sdb}"
            .to_string();

        assert_eq!(
            Inspirer::init(None).aux2key(input),
            vec!["Abramovici:1992ah", "Thorne:1992sdb"]
        );
    }

    #[test]
    fn test_aux_bibtex_2_citation_together() {
        let input = r"\relax
            \citation{Abramovici:1992ah,Thorne:1992sdb}"
            .to_string();

        assert_eq!(
            Inspirer::init(None).aux2key(input),
            vec!["Abramovici:1992ah", "Thorne:1992sdb"]
        );
    }

    #[test]
    fn test_aux_bibtex_3_citation() {
        let input = r"\relax
            \citation{Abramovici:1992ah}
            \citation{Thorne:1992sdb}
            \citation{Thorne:1992sdb}"
            .to_string();

        assert_eq!(
            Inspirer::init(None).aux2key(input),
            vec!["Abramovici:1992ah", "Thorne:1992sdb", "Thorne:1992sdb"]
        );
    }

    #[test]
    fn test_aux_bibtex_3_citation_together() {
        let input = r"\relax
            \citation{Abramovici:1992ah,Thorne:1992sdb,Bildsten:1992my}"
            .to_string();

        assert_eq!(
            Inspirer::init(None).aux2key(input),
            vec!["Abramovici:1992ah", "Thorne:1992sdb", "Bildsten:1992my"]
        );
    }
}
<|endoftext|>


<commit_before>// (c) 2015 Joost Yervante Damad <joost@productize.be>

// loosely based on https://github.com/cgaebel/sexp
// latest version can be found at https://github.com/andete/rust_sexp

use std::fmt;
use std::str::FromStr;
use std::f64;

use std::fs::File;
use std::io::prelude::*;

pub enum Atom {
  S(String),
  Q(String),
  I(i64),
  F(f64),
}

impl Atom {
    pub fn f(&self) -> Result<f64,String> {
        match *self {
            Atom::F(f) => Ok(f),
            Atom::I(i) => Ok(i as f64),
            ref x => Err(format!("not a float: {}", x))
        }
    }
    pub fn i(&self) -> Result<i64,String> {
        match *self {
            Atom::I(i) => Ok(i),
            ref x => Err(format!("not an int: {}", x))
        }
    }
    pub fn string(&self) -> Result<String,String> {
        match *self {
            Atom::S(ref s) => Ok(s.clone()),
            Atom::Q(ref s) => Ok(s.clone()),
            ref x => Err(format!("not a string: {}", x))
        }
    }
    pub fn as_string(&self) -> Result<String,String> {
        match *self {
            Atom::S(ref s) => Ok(s.clone()),
            Atom::Q(ref s) => Ok(s.clone()),
            Atom::F(ref s) => Ok(format!("{}", s)),
            Atom::I(ref s) => Ok(format!("{}", s)),
        }
    }
}

pub enum Sexp {
  Atom(Atom),
  Empty,
  List(Vec<Sexp>),
}

impl Sexp {
    pub fn atom(&self) -> Result<&Atom,String> {
        match *self {
            Sexp::Atom(ref f) => Ok(f),
            ref x => Err(format!("not an atom: {}", x))
        }
    }
    pub fn list(&self) -> Result<&Vec<Sexp>,String> {
        match *self {
            Sexp::List(ref v) => Ok(v),
            ref x => Err(format!("not a list: {}", x))
        }
    }
}

pub struct ParseError {
    msg: &'static str,
    line: usize,
    lpos: usize,
}

pub enum Error {
    ParseError(ParseError),
    IOError(std::io::Error),
}

pub type ERes<T> = Result<T, Error>;

fn err_parse<T>(msg: &'static str, state: &ParseState) -> ERes<T> {
    Err(Error::ParseError(ParseError { msg: msg, line: state.line, lpos: state.lpos}))
}

fn err_io<T>(err: std::io::Error) -> ERes<T> {
    Err(Error::IOError(err))
}

struct ParseState {
    pos: usize,
    line: usize,
    lpos: usize,
    vec: Vec<char>,
    len: usize,
}

impl ParseState {
    fn peek(&self) -> ERes<char> {
        if self.pos >= self.len {
            return err_parse("end of document reached", self)
        }
        Ok(self.vec[self.pos])
    }
    fn peek_option(&self) -> Option<char> {
        if self.pos >= self.len {
            return None
        }
        Some(self.vec[self.pos])
    }
    fn next(&mut self) -> ERes<()> {
        if self.pos >= self.len {
            return err_parse("end of document reached", self)
        }
        let c = self.vec[self.pos];
        self.pos += 1;
        match c {
            '\r' | '\n' => {
                self.lpos = 0;
                self.line += 1;
            }
            _ => {
                self.lpos += 1;
            }
        }
        Ok(())
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Error::ParseError(ref p) =>
                write!(f, "Parse Error {}:{}: {}", p.line, p.lpos, p.msg),
            Error::IOError(ref i) =>
                write!(f, "{}", i)
        }
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Error::ParseError(ref p) =>
                write!(f, "Parse Error {}:{}: {}", p.line, p.lpos, p.msg),
            Error::IOError(ref i) =>
                write!(f, "{}", i)
        }
    }
}

impl fmt::Display for Atom {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    match *self {
        Atom::S(ref s) => write!(f, "{}", s),
        Atom::Q(ref s) => write!(f, "\"{}\"", s),
        Atom::I(i)     => write!(f, "{}", i),
        Atom::F(d)     => {
            let z = d.floor();
            if d - z < f64::EPSILON {
                write!(f, "{}.0", z)
            } else {
                write!(f, "{}", d)
            }
        }
    }
  }
}

impl fmt::Display for Sexp {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    match *self {
      Sexp::Atom(ref a) => write!(f, "{}", a),
      Sexp::List(ref xs) => {
        try!(write!(f, "("));
        for (i, x) in xs.iter().enumerate() {
          let s = if i == 0 { "" } else { " " };
          try!(write!(f, "{}{}", s, x));
        }
        write!(f, ")")
      },
      Sexp::Empty => write!(f, ""),
    }
  }
}

fn parse_list(state: &mut ParseState) -> ERes<Sexp> {
    //println!("list");
    try!(state.next()); // skip (
    let mut l: Vec<Sexp> = Vec::new();
    loop {
        match try!(state.peek()) {
            ')' => {
                try!(state.next());
                break;
            }
            _ => {
                l.push(try!(parse_sexp(state)));
            }
        }
    }
    //println!("Found list");
    Ok(Sexp::List(l))
}

fn parse_quoted_string(state: &mut ParseState) -> ERes<Atom> {
    //println!("qstring");
    try!(state.next()); // skip "
    let mut s = String::new();
    loop {
        match try!(state.peek()) {
            '"' => {
                try!(state.next());
                break
            }
            x @ '\r' | x @ '\n' => {
                s.push(x);
                try!(state.next());
                }
            x => {
                s.push(x);
                try!(state.next());
            }
        }
    }
    //println!("Found quoted string {}", s);
    Ok(Atom::Q(s))
}

fn parse_string(state: &mut ParseState) -> ERes<Atom> {
    //println!("string");
    let mut s = String::new();
    loop {
        match state.peek_option() {
            Some(x) => {
                match x {
                    ' ' | '\t' | '\r' | '\n' | ')' => break,
                    '"' => return err_parse("unexpected \" in string", state),
                    x => s.push(x),
                }
            }
            None => {
                break;
            }
        }
        try!(state.next())
    }
    //println!("Found string {}", s);
    Ok(Atom::S(s))
}

fn parse_number(state: &mut ParseState) -> ERes<Atom> {
    //println!("number");
    let mut s = String::new();
    loop {
        match state.peek_option() {
            Some(x) => {
                match x {
                    ' ' | '\r' | '\n' | '\t' | ')' => {
                        break
                    },
                    '0' ... '9' | '.' | '-' => {
                        s.push(state.vec[state.pos])
                    },
                    _ => {
                        return err_parse("unexpected char in number", state)
                    },
                }
            }
            None => {
                break
            }
        }
        try!(state.next())
    }
    //println!("Found number {}", s);
    let s2: &str = &s[..];
    if s.contains('.') {
        Ok(Atom::F(f64::from_str(s2).unwrap()))
    } else {
        Ok(Atom::I(i64::from_str(s2).unwrap()))
    }
}

fn parse_atom(state: &mut ParseState) -> ERes<Sexp> {
    //println!("atom");
    let a = match try!(state.peek()) {
        '"' => {
            try!(parse_quoted_string(state))
        }
        '0' ... '9' | '.' | '-' => {
            try!(parse_number(state))
        }
        _ => {
            try!(parse_string(state))
        }
    };
    Ok(Sexp::Atom(a))
}

fn parse_sexp(state: &mut ParseState) -> ERes<Sexp> {
    loop {
        match try!(state.peek()) {
            '(' => {
                return parse_list(state)
            }
            ' ' | '\t' | '\r' | '\n' => {
                try!(state.next());
            }
            ')' => {
                return err_parse("unmatched )", state)
            }
            _ => {
                return parse_atom(state)
            }
        }
    }
}

fn parse(data: &str) -> ERes<Sexp> {
    if data.len() == 0 {
        Ok(Sexp::Empty)
    } else {
        let vec: Vec<char> = data.chars().collect();
        let len = vec.len();
        let state = &mut ParseState { pos: 0, line: 1, lpos: 0, vec: vec, len: len };
        parse_sexp(state)
    }
}

fn read_file(name: &str) -> Result<String, std::io::Error> {
    let mut f = try!(File::open(name));
    let mut s = String::new();
    try!(f.read_to_string(&mut s));
    Ok(s)
}

pub fn parse_str(s: &str) -> ERes<Sexp> {
    parse(s)
}

pub fn parse_file(name: &str) -> ERes<Sexp> {
    let s = try!(match read_file(name) {
        Ok(s) => Ok(s),
        Err(x) => err_io(x),
    });
    parse(&s[..])
}

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

    #[allow(dead_code)]
    fn check_parse(s: &str) {
        let e = parse_str(s).unwrap();
        let t = format!("{}", e);
        assert_eq!(s, t);
    }

    #[allow(dead_code)]
    fn parse_fail(s: &str) {
        parse_str(s).unwrap();
    }

    #[test]
    fn test_empty() { check_parse("") }

    #[test]
    fn test_minimal() { check_parse("()") }

    #[test]
    fn test_string() { check_parse("hello") }

    #[test]
    fn test_qstring() { check_parse("\"hello\"") }

    #[test]
    fn test_number() { check_parse("1.3") }

    #[test]
    fn test_float_vs_int() { check_parse("2.0") }

    #[test]
    fn test_double() { check_parse("(())") }

    #[test]
    fn test_br_string() { check_parse("(world)") }

    #[test]
    fn test_br_qstring() { check_parse("(\"world\")") }

    #[test]
    fn test_br_int() { check_parse("(42)") }

    #[test]
    fn test_br_float() { check_parse("(12.7)") }

    #[test]
    fn test_br_qbrstring() { check_parse("(\"(()\")") }

    #[test]
    #[should_panic(expected="Parse Error 1:1: end of document reached")]
    fn test_invalid1() { parse_fail("(") }

    #[test]
    #[should_panic(expected="Parse Error 1:0: unmatched )")]
    fn test_invalid2() { parse_fail(")") }

    #[test]
    #[should_panic(expected="Parse Error 1:6: end of document reached")]
    fn test_invalid3() { parse_fail("\"hello") }

    #[test]
    fn test_complex() { check_parse("(module SWITCH_3W_SIDE_MMP221-R (layer F.Cu) (descr \"\") (pad 1 thru_hole rect (size 1.2 1.2) (at -2.5 -1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 2 thru_hole rect (size 1.2 1.2) (at 0.0 -1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 3 thru_hole rect (size 1.2 1.2) (at 2.5 -1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 5 thru_hole rect (size 1.2 1.2) (at 0.0 1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 6 thru_hole rect (size 1.2 1.2) (at -2.5 1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 4 thru_hole rect (size 1.2 1.2) (at 2.5 1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (fp_line (start -4.5 -1.75) (end 4.5 -1.75) (layer F.SilkS) (width 0.127)) (fp_line (start 4.5 -1.75) (end 4.5 1.75) (layer F.SilkS) (width 0.127)) (fp_line (start 4.5 1.75) (end -4.5 1.75) (layer F.SilkS) (width 0.127)) (fp_line (start -4.5 1.75) (end -4.5 -1.75) (layer F.SilkS) (width 0.127)))") }
}
<commit_msg>fallback on string when parsing as a number fails!<commit_after>// (c) 2015 Joost Yervante Damad <joost@productize.be>

// loosely based on https://github.com/cgaebel/sexp
// latest version can be found at https://github.com/andete/rust_sexp

use std::fmt;
use std::str::FromStr;
use std::f64;

use std::fs::File;
use std::io::prelude::*;

pub enum Atom {
  S(String),
  Q(String),
  I(i64),
  F(f64),
}

impl Atom {
    pub fn f(&self) -> Result<f64,String> {
        match *self {
            Atom::F(f) => Ok(f),
            Atom::I(i) => Ok(i as f64),
            ref x => Err(format!("not a float: {}", x))
        }
    }
    pub fn i(&self) -> Result<i64,String> {
        match *self {
            Atom::I(i) => Ok(i),
            ref x => Err(format!("not an int: {}", x))
        }
    }
    pub fn string(&self) -> Result<String,String> {
        match *self {
            Atom::S(ref s) => Ok(s.clone()),
            Atom::Q(ref s) => Ok(s.clone()),
            ref x => Err(format!("not a string: {}", x))
        }
    }
    pub fn as_string(&self) -> Result<String,String> {
        match *self {
            Atom::S(ref s) => Ok(s.clone()),
            Atom::Q(ref s) => Ok(s.clone()),
            Atom::F(ref s) => Ok(format!("{}", s)),
            Atom::I(ref s) => Ok(format!("{}", s)),
        }
    }
}

pub enum Sexp {
  Atom(Atom),
  Empty,
  List(Vec<Sexp>),
}

impl Sexp {
    pub fn atom(&self) -> Result<&Atom,String> {
        match *self {
            Sexp::Atom(ref f) => Ok(f),
            ref x => Err(format!("not an atom: {}", x))
        }
    }
    pub fn list(&self) -> Result<&Vec<Sexp>,String> {
        match *self {
            Sexp::List(ref v) => Ok(v),
            ref x => Err(format!("not a list: {}", x))
        }
    }
}

pub struct ParseError {
    msg: &'static str,
    line: usize,
    lpos: usize,
}

pub enum Error {
    ParseError(ParseError),
    IOError(std::io::Error),
}

pub type ERes<T> = Result<T, Error>;

fn err_parse<T>(msg: &'static str, state: &ParseState) -> ERes<T> {
    Err(Error::ParseError(ParseError { msg: msg, line: state.line, lpos: state.lpos}))
}

fn err_io<T>(err: std::io::Error) -> ERes<T> {
    Err(Error::IOError(err))
}

struct ParseState {
    pos: usize,
    line: usize,
    lpos: usize,
    vec: Vec<char>,
    len: usize,
    store_pos: usize,
    store_line: usize,
    store_lpos: usize,
}

impl ParseState {
    fn peek(&self) -> ERes<char> {
        if self.pos >= self.len {
            return err_parse("end of document reached", self)
        }
        Ok(self.vec[self.pos])
    }
    fn peek_option(&self) -> Option<char> {
        if self.pos >= self.len {
            return None
        }
        Some(self.vec[self.pos])
    }
    fn next(&mut self) -> ERes<()> {
        if self.pos >= self.len {
            return err_parse("end of document reached", self)
        }
        let c = self.vec[self.pos];
        self.pos += 1;
        match c {
            '\r' | '\n' => {
                self.lpos = 0;
                self.line += 1;
            }
            _ => {
                self.lpos += 1;
            }
        }
        Ok(())
    }

    fn store_pos(&mut self) {
        self.store_pos = self.pos;
        self.store_line = self.line;
        self.store_lpos = self.lpos;

    }

    fn restore_pos(&mut self) {
        self.pos = self.store_pos;
        self.line = self.store_line;
        self.lpos = self.store_lpos;
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Error::ParseError(ref p) =>
                write!(f, "Parse Error {}:{}: {}", p.line, p.lpos, p.msg),
            Error::IOError(ref i) =>
                write!(f, "{}", i)
        }
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Error::ParseError(ref p) =>
                write!(f, "Parse Error {}:{}: {}", p.line, p.lpos, p.msg),
            Error::IOError(ref i) =>
                write!(f, "{}", i)
        }
    }
}

impl fmt::Display for Atom {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    match *self {
        Atom::S(ref s) => write!(f, "{}", s),
        Atom::Q(ref s) => write!(f, "\"{}\"", s),
        Atom::I(i)     => write!(f, "{}", i),
        Atom::F(d)     => {
            let z = d.floor();
            if d - z < f64::EPSILON {
                write!(f, "{}.0", z)
            } else {
                write!(f, "{}", d)
            }
        }
    }
  }
}

impl fmt::Display for Sexp {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    match *self {
      Sexp::Atom(ref a) => write!(f, "{}", a),
      Sexp::List(ref xs) => {
        try!(write!(f, "("));
        for (i, x) in xs.iter().enumerate() {
          let s = if i == 0 { "" } else { " " };
          try!(write!(f, "{}{}", s, x));
        }
        write!(f, ")")
      },
      Sexp::Empty => write!(f, ""),
    }
  }
}

fn parse_list(state: &mut ParseState) -> ERes<Sexp> {
    //println!("list");
    try!(state.next()); // skip (
    let mut first = true;
    let mut l: Vec<Sexp> = Vec::new();
    loop {
        match try!(state.peek()) {
            ')' => {
                try!(state.next());
                break;
            }
            ' ' | '\t' | '\r' | '\n' => {
                try!(state.next());
            }
            _ => {
                let res = try!(parse_sexp(state));
                if first {
                    println!("List: {}", &res);
                    first = false;
                }
                l.push(res)
            }
        }
    }
    let l = Sexp::List(l);
    println!("Found list {}", &l);
    Ok(l)
}

fn parse_quoted_string(state: &mut ParseState) -> ERes<Atom> {
    //println!("qstring");
    try!(state.next()); // skip "
    let mut s = String::new();
    loop {
        match try!(state.peek()) {
            '"' => {
                try!(state.next());
                break
            }
            x @ '\r' | x @ '\n' => {
                s.push(x);
                try!(state.next());
                }
            x => {
                s.push(x);
                try!(state.next());
            }
        }
    }
    //println!("Found quoted string {}", s);
    Ok(Atom::Q(s))
}

fn parse_string(state: &mut ParseState) -> ERes<Atom> {
    //println!("string");
    let mut s = String::new();
    loop {
        match state.peek_option() {
            Some(x) => {
                match x {
                    ' ' | '\t' | '\r' | '\n' | ')' => break,
                    '"' => return err_parse("unexpected \" in string", state),
                    x => s.push(x),
                }
            }
            None => {
                break;
            }
        }
        try!(state.next())
    }
    //println!("Found string {}", s);
    Ok(Atom::S(s))
}

fn parse_number(state: &mut ParseState) -> ERes<Atom> {
    //println!("number");
    let mut s = String::new();
    loop {
        match state.peek_option() {
            Some(x) => {
                match x {
                    ' ' | '\r' | '\n' | '\t' | ')' => {
                        break
                    },
                    '0' ... '9' | '.' | '-' => {
                        s.push(state.vec[state.pos])
                    },
                    _ => {
                        return err_parse("unexpected char in number", &state)
                    },
                }
            }
            None => {
                break
            }
        }
        try!(state.next())
    }
    //println!("Found number {}", s);
    let s2: &str = &s[..];
    if s.contains('.') {
        Ok(Atom::F(f64::from_str(s2).unwrap()))
    } else {
        Ok(Atom::I(i64::from_str(s2).unwrap()))
    }
}

fn parse_atom(state: &mut ParseState) -> ERes<Sexp> {
    //println!("atom");
    let a = match try!(state.peek()) {
        '"' => {
            try!(parse_quoted_string(state))
        },
        '0' ... '9' | '.' | '-' => {
            state.store_pos();
            match parse_number(state) {
                Ok(a) => {
                    a
                },
                Err(_) => {
                    state.restore_pos();
                    try!(parse_string(state))
                },
            }
        }
        _ => {
            try!(parse_string(state))
        }
    };
    Ok(Sexp::Atom(a))
}

fn parse_sexp(state: &mut ParseState) -> ERes<Sexp> {
    loop {
        match try!(state.peek()) {
            '(' => {
                return parse_list(state)
            }
            ' ' | '\t' | '\r' | '\n' => {
                try!(state.next());
            }
            ')' => {
                return err_parse("unmatched )", state)
            }
            _ => {
                return parse_atom(state)
            }
        }
    }
}

fn parse(data: &str) -> ERes<Sexp> {
    if data.len() == 0 {
        Ok(Sexp::Empty)
    } else {
        let vec: Vec<char> = data.chars().collect();
        let len = vec.len();
        let state = &mut ParseState { pos: 0, line: 1, lpos: 0, vec: vec, len: len, store_pos: 0, store_lpos: 0, store_line: 0, };
        parse_sexp(state)
    }
}

fn read_file(name: &str) -> Result<String, std::io::Error> {
    let mut f = try!(File::open(name));
    let mut s = String::new();
    try!(f.read_to_string(&mut s));
    Ok(s)
}

pub fn parse_str(s: &str) -> ERes<Sexp> {
    parse(s)
}

pub fn parse_file(name: &str) -> ERes<Sexp> {
    let s = try!(match read_file(name) {
        Ok(s) => Ok(s),
        Err(x) => err_io(x),
    });
    parse(&s[..])
}

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

    #[allow(dead_code)]
    fn check_parse(s: &str) {
        let e = parse_str(s).unwrap();
        let t = format!("{}", e);
        assert_eq!(s, t);
    }

    #[allow(dead_code)]
    fn parse_fail(s: &str) {
        parse_str(s).unwrap();
    }

    #[test]
    fn test_empty() { check_parse("") }

    #[test]
    fn test_minimal() { check_parse("()") }

    #[test]
    fn test_string() { check_parse("hello") }

    #[test]
    fn test_qstring() { check_parse("\"hello\"") }

    #[test]
    fn test_number() { check_parse("1.3") }

    #[test]
    fn test_float_vs_int() { check_parse("2.0") }

    #[test]
    fn test_double() { check_parse("(())") }

    #[test]
    fn test_br_string() { check_parse("(world)") }

    #[test]
    fn test_br_qstring() { check_parse("(\"world\")") }

    #[test]
    fn test_br_int() { check_parse("(42)") }

    #[test]
    fn test_br_float() { check_parse("(12.7)") }

    #[test]
    fn test_br_qbrstring() { check_parse("(\"(()\")") }

    #[test]
    #[should_panic(expected="Parse Error 1:1: end of document reached")]
    fn test_invalid1() { parse_fail("(") }

    #[test]
    #[should_panic(expected="Parse Error 1:0: unmatched )")]
    fn test_invalid2() { parse_fail(")") }

    #[test]
    #[should_panic(expected="Parse Error 1:6: end of document reached")]
    fn test_invalid3() { parse_fail("\"hello") }

    #[test]
    fn test_complex() { check_parse("(module SWITCH_3W_SIDE_MMP221-R (layer F.Cu) (descr \"\") (pad 1 thru_hole rect (size 1.2 1.2) (at -2.5 -1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 2 thru_hole rect (size 1.2 1.2) (at 0.0 -1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 3 thru_hole rect (size 1.2 1.2) (at 2.5 -1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 5 thru_hole rect (size 1.2 1.2) (at 0.0 1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 6 thru_hole rect (size 1.2 1.2) (at -2.5 1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (pad 4 thru_hole rect (size 1.2 1.2) (at 2.5 1.6 0) (layers *.Cu *.Mask) (drill 0.8)) (fp_line (start -4.5 -1.75) (end 4.5 -1.75) (layer F.SilkS) (width 0.127)) (fp_line (start 4.5 -1.75) (end 4.5 1.75) (layer F.SilkS) (width 0.127)) (fp_line (start 4.5 1.75) (end -4.5 1.75) (layer F.SilkS) (width 0.127)) (fp_line (start -4.5 1.75) (end -4.5 -1.75) (layer F.SilkS) (width 0.127)))") }

    #[test]
    fn test_temp() {
        let s = parse_file("/tmp/usbser.kicad_pcb").unwrap();
        println!("{}", s);
    }
}

<|endoftext|>


<commit_before>#![feature(macro_rules)]
use std::collections::{HashMap, TreeMap, SmallIntMap, TrieMap};
use std::hash::Hash;

/// Conversion from an `Iterator` of pairs.
pub trait Groupable<K, V> {
    /// Loops through the entire iterator, grouping all keys into a container
    /// implementing `FromKeyedIterator` with a container of values per key.
    /// The values will be aggregated per key into a container implementing
    /// `Extendable` for the value type.
    ///
    /// # Example
    ///
    /// ```rust
    /// use std::collections::HashMap;
    /// use groupable::Groupable;
    ///
    /// let evens = range(0u, 10).map(|i| (i % 2 == 0, i))
    ///                         .group::<HashMap<bool, Vec<uint>>>();
    ///
    /// assert_eq!(evens[true].as_slice(), [0, 2, 4, 6, 8].as_slice());
    /// assert_eq!(evens[false].as_slice(), [1, 3, 5, 7, 9].as_slice());
    /// ```
    fn group<B: FromKeyedIterator<K, V>>(&mut self) -> B;
}

impl<K, V, I: Iterator<(K, V)>> Groupable<K, V> for I {
    fn group<B: FromKeyedIterator<K, V>>(&mut self) -> B {
        FromKeyedIterator::from_keyed_iter(self.by_ref())
    }
}

/// Conversion from an `Iterator` of key-value pairs.
trait FromKeyedIterator<K, V> {
    /// Build a container with groups of elements from an external iterator.
    fn from_keyed_iter<I: Iterator<(K, V)>>(I) -> Self;
}

#[inline(always)]
fn group_into<K,
              V,
              U: Extendable<V>,
              I: Iterator<(K, V)>,
              M: MutableMap<K, U>>(mut iter: I, map: &mut M) {
    for (key, val) in iter {
        let val_iter = Some(val).move_iter();
        match map.find_mut(&key) {
            Some(collection) => {
                collection.extend(val_iter);
                continue
            },
            None => {} // insert below
        }

        map.insert(key, FromIterator::from_iter(val_iter));
    }
}

macro_rules! impl_keyed_iter (
    ($name:ident: $($bounds:ident),+) => (
        impl <K: $($bounds+)+, V, U: Extendable<V>> FromKeyedIterator<K, V> for $name<K, U> {
            fn from_keyed_iter<T: Iterator<(K, V)>>(iter: T) -> $name<K, U> {
                let mut map = $name::<K, U>::new();
                group_into(iter, &mut map);
                map
            }
        }
    )
)

macro_rules! impl_uint_keyed_iter (
    ($name:ident) => (
        impl <V, U: Extendable<V>> FromKeyedIterator<uint, V> for $name<U> {
            fn from_keyed_iter<T: Iterator<(uint, V)>>(iter: T) -> $name<U> {
                let mut map = $name::<U>::new();
                group_into(iter, &mut map);
                map
            }
        }
    )
)

impl_keyed_iter!(HashMap: Ord, Hash)
impl_keyed_iter!(TreeMap: Ord)
impl_uint_keyed_iter!(SmallIntMap)
impl_uint_keyed_iter!(TrieMap)
<commit_msg>Fix private type exports.<commit_after>#![feature(macro_rules)]
use std::collections::{HashMap, TreeMap, SmallIntMap, TrieMap};
use std::hash::Hash;

/// Conversion from an `Iterator` of pairs.
pub trait Groupable<K, V> {
    /// Loops through the entire iterator, grouping all keys into a container
    /// implementing `FromKeyedIterator` with a container of values per key.
    /// The values will be aggregated per key into a container implementing
    /// `Extendable` for the value type.
    ///
    /// # Example
    ///
    /// ```rust
    /// use std::collections::HashMap;
    /// use groupable::Groupable;
    ///
    /// let evens = range(0u, 10).map(|i| (i % 2 == 0, i))
    ///                         .group::<HashMap<bool, Vec<uint>>>();
    ///
    /// assert_eq!(evens[true].as_slice(), [0, 2, 4, 6, 8].as_slice());
    /// assert_eq!(evens[false].as_slice(), [1, 3, 5, 7, 9].as_slice());
    /// ```
    fn group<B: FromKeyedIterator<K, V>>(&mut self) -> B;
}

impl<K, V, I: Iterator<(K, V)>> Groupable<K, V> for I {
    fn group<B: FromKeyedIterator<K, V>>(&mut self) -> B {
        FromKeyedIterator::from_keyed_iter(self.by_ref())
    }
}

/// Conversion from an `Iterator` of key-value pairs.
pub trait FromKeyedIterator<K, V> {
    /// Build a container with groups of elements from an external iterator.
    fn from_keyed_iter<I: Iterator<(K, V)>>(I) -> Self;
}

#[inline(always)]
fn group_into<K,
              V,
              U: Extendable<V>,
              I: Iterator<(K, V)>,
              M: MutableMap<K, U>>(mut iter: I, map: &mut M) {
    for (key, val) in iter {
        let val_iter = Some(val).into_iter();
        match map.find_mut(&key) {
            Some(collection) => {
                collection.extend(val_iter);
                continue
            },
            None => {} // insert below
        }

        map.insert(key, FromIterator::from_iter(val_iter));
    }
}

macro_rules! impl_keyed_iter (
    ($name:ident: $($bounds:ident),+) => (
        impl <K: $($bounds+)+, V, U: Extendable<V>> FromKeyedIterator<K, V> for $name<K, U> {
            fn from_keyed_iter<T: Iterator<(K, V)>>(iter: T) -> $name<K, U> {
                let mut map = $name::<K, U>::new();
                group_into(iter, &mut map);
                map
            }
        }
    )
)

macro_rules! impl_uint_keyed_iter (
    ($name:ident) => (
        impl <V, U: Extendable<V>> FromKeyedIterator<uint, V> for $name<U> {
            fn from_keyed_iter<T: Iterator<(uint, V)>>(iter: T) -> $name<U> {
                let mut map = $name::<U>::new();
                group_into(iter, &mut map);
                map
            }
        }
    )
)

impl_keyed_iter!(HashMap: Ord, Hash)
impl_keyed_iter!(TreeMap: Ord)
impl_uint_keyed_iter!(SmallIntMap)
impl_uint_keyed_iter!(TrieMap)
<|endoftext|>


<commit_before>/*!
Rust bindings to the [sodium library](https://github.com/jedisct1/libsodium).

Sodium is a portable implementation of Dan Bernsteins [NaCl: Networking and
Cryptography library](http://nacl.cr.yp.to)

For most users, if you want public-key (asymmetric) cryptography you should use
the functions in `crypto::asymmetricbox` for encryption/decryption.

If you want secret-key (symmetric) cryptography you should be using the
functions in `crypto::secretbox` for encryption/decryption.

For public-key signatures you should use the functions in `crypto::sign` for
signature creation and verification.

Unless you know what you're doing you most certainly don't want to use the
functions in `crypto::scalarmult`, `crypto::stream`, `crypto::auth` and
`crypto::onetimeauth`.

## Thread Safety
All functions in this library are thread-safe provided that the `init()`
function has been called during program execution.

If `init()` hasn't been called then all functions except the random-number
generation functions and the key-generation functions are thread-safe.

# Public-key cryptography
 `crypto::asymmetricbox`

 `crypto::sign`

# Secret-key cryptography
 `crypto::secretbox`

 `crypto::stream`

 `crypto::auth`

 `crypto::onetimeauth`

# Low-level functions
 `crypto::hash`

 `crypto::verify`

 `crypto::shorthash`
 */
#![crate_name = "sodiumoxide"]
#![crate_type = "lib"]
#![warn(missing_docs)]
#![warn(non_upper_case_globals)]
#![warn(non_camel_case_types)]
#![warn(unused_qualifications)]
#![feature(libc, collections)]

/* workaround: the rust compiler doesn't recognize
   the features path, test and io yet, still it warns
   about using them */
#![allow(unused_features)]
#![feature(path, test, io)]

extern crate "libsodium-sys" as ffi;
extern crate libc;

/**
 * `init()` initializes the sodium library and chooses faster versions of
 * the primitives if possible. `init()` also makes the random number generation
 * functions (`gen_key`, `gen_keypair`, `gen_nonce`, `randombytes`, `randombytes_into`)
 * thread-safe
 */
pub fn init() -> bool {
    unsafe {
        ffi::sodium_init() == 0
    }
}

#[macro_use]
mod utils;

pub mod randombytes;

/**
 * Cryptographic functions
 */
pub mod crypto {
    pub mod asymmetricbox;
    pub mod sign;
    pub mod scalarmult;
    pub mod auth;
    pub mod hash;
    pub mod secretbox;
    pub mod onetimeauth;
    pub mod stream;
    pub mod shorthash;
    pub mod verify;
}

<commit_msg>Added #![feature(core)] to silence waring<commit_after>/*!
Rust bindings to the [sodium library](https://github.com/jedisct1/libsodium).

Sodium is a portable implementation of Dan Bernsteins [NaCl: Networking and
Cryptography library](http://nacl.cr.yp.to)

For most users, if you want public-key (asymmetric) cryptography you should use
the functions in `crypto::asymmetricbox` for encryption/decryption.

If you want secret-key (symmetric) cryptography you should be using the
functions in `crypto::secretbox` for encryption/decryption.

For public-key signatures you should use the functions in `crypto::sign` for
signature creation and verification.

Unless you know what you're doing you most certainly don't want to use the
functions in `crypto::scalarmult`, `crypto::stream`, `crypto::auth` and
`crypto::onetimeauth`.

## Thread Safety
All functions in this library are thread-safe provided that the `init()`
function has been called during program execution.

If `init()` hasn't been called then all functions except the random-number
generation functions and the key-generation functions are thread-safe.

# Public-key cryptography
 `crypto::asymmetricbox`

 `crypto::sign`

# Secret-key cryptography
 `crypto::secretbox`

 `crypto::stream`

 `crypto::auth`

 `crypto::onetimeauth`

# Low-level functions
 `crypto::hash`

 `crypto::verify`

 `crypto::shorthash`
 */
#![crate_name = "sodiumoxide"]
#![crate_type = "lib"]
#![warn(missing_docs)]
#![warn(non_upper_case_globals)]
#![warn(non_camel_case_types)]
#![warn(unused_qualifications)]
#![feature(libc, collections, core)]

/* workaround: the rust compiler doesn't recognize
   the features path, test and io yet, still it warns
   about using them */
#![allow(unused_features)]
#![feature(path, test, io)]

extern crate "libsodium-sys" as ffi;
extern crate libc;

/**
 * `init()` initializes the sodium library and chooses faster versions of
 * the primitives if possible. `init()` also makes the random number generation
 * functions (`gen_key`, `gen_keypair`, `gen_nonce`, `randombytes`, `randombytes_into`)
 * thread-safe
 */
pub fn init() -> bool {
    unsafe {
        ffi::sodium_init() == 0
    }
}

#[macro_use]
mod utils;

pub mod randombytes;

/**
 * Cryptographic functions
 */
pub mod crypto {
    pub mod asymmetricbox;
    pub mod sign;
    pub mod scalarmult;
    pub mod auth;
    pub mod hash;
    pub mod secretbox;
    pub mod onetimeauth;
    pub mod stream;
    pub mod shorthash;
    pub mod verify;
}

<|endoftext|>


<commit_before>//! Replaces the deprecated functionality of std::os::num_cpus.
#![cfg_attr(test, deny(warnings))]
#![deny(missing_docs)]
#![allow(non_snake_case)]

#[cfg(not(windows))]
extern crate libc;

/// Returns the number of CPUs of the current machine.
#[inline]
pub fn get() -> usize {
    get_num_cpus()
}

/// Returns the number of physical cores of the current machine.
/// If not possible on the particular architecture returns same as get() which
/// is the logical CPUs.
#[inline]
pub fn get_physical() -> usize {
    get_num_physical_cpus()
}

#[cfg(not(target_os = "linux"))]
#[inline]
fn get_num_physical_cpus() -> usize {
    // Not implemented, fallback
    get_num_cpus()
}

#[cfg(target_os = "linux")]
fn get_num_physical_cpus() -> usize {
    use std::io::BufReader;
    use std::io::BufRead;
    use std::fs::File;
    use std::collections::HashSet;

    let file = match File::open("/proc/cpuinfo") {
        Ok(val) => val,
        Err(_) => {return get_num_cpus()},
    };
    let reader = BufReader::new(file);
    let mut set = HashSet::new();
    let mut coreid: u32 = 0;
    let mut physid: u32 = 0;
    let mut chgcount = 0;
    for line in reader.lines().filter_map(|result| result.ok()) {
        let parts: Vec<&str> = line.split(':').map(|s| s.trim()).collect();
        if parts.len() != 2 {
            continue
        }
        if parts[0] == "core id" || parts[0] == "physical id" {
            let value = match parts[1].trim().parse() {
              Ok(val) => val,
              Err(_) => break,
            };
            match parts[0] {
                "core id"     => coreid = value,
                "physical id" => physid = value,
                _ => {},
            }
            chgcount += 1;
        }
        if chgcount == 2 {
            set.insert((physid, coreid));
            chgcount = 0;
        }
    }
    let count = set.len();
    if count == 0 { get_num_cpus() } else { count }
}

#[cfg(windows)]
fn get_num_cpus() -> usize {
    #[repr(C)]
    struct SYSTEM_INFO {
        wProcessorArchitecture: u16,
        wReserved: u16,
        dwPageSize: u32,
        lpMinimumApplicationAddress: *mut u8,
        lpMaximumApplicationAddress: *mut u8,
        dwActiveProcessorMask: *mut u8,
        dwNumberOfProcessors: u32,
        dwProcessorType: u32,
        dwAllocationGranularity: u32,
        wProcessorLevel: u16,
        wProcessorRevision: u16,
    }

    extern "system" {
        fn GetSystemInfo(lpSystemInfo: *mut SYSTEM_INFO);
    }

    unsafe {
        let mut sysinfo: SYSTEM_INFO = std::mem::uninitialized();
        GetSystemInfo(&mut sysinfo);
        sysinfo.dwNumberOfProcessors as usize
    }
}

#[cfg(any(target_os = "freebsd",
          target_os = "dragonfly",
          target_os = "bitrig",
          target_os = "netbsd"))]
fn get_num_cpus() -> usize {
    let mut cpus: libc::c_uint = 0;
    let mut cpus_size = std::mem::size_of_val(&cpus);

    unsafe {
        cpus = libc::sysconf(libc::_SC_NPROCESSORS_ONLN) as libc::c_uint;
    }
    if cpus < 1 {
        let mut mib = [libc::CTL_HW, libc::HW_NCPU, 0, 0];
        unsafe {
            libc::sysctl(mib.as_mut_ptr(),
                         2,
                         &mut cpus as *mut _ as *mut _,
                         &mut cpus_size as *mut _ as *mut _,
                         0 as *mut _,
                         0);
        }
        if cpus < 1 {
            cpus = 1;
        }
    }
    cpus as usize
}

#[cfg(target_os = "openbsd")]
fn get_num_cpus() -> usize {
    let mut cpus: libc::c_uint = 0;
    let mut cpus_size = std::mem::size_of_val(&cpus);
    let mut mib = [libc::CTL_HW, libc::HW_NCPU, 0, 0];

    unsafe {
        libc::sysctl(mib.as_mut_ptr(),
                     2,
                     &mut cpus as *mut _ as *mut _,
                     &mut cpus_size as *mut _ as *mut _,
                     0 as *mut _,
                     0);
    }
    if cpus < 1 {
        cpus = 1;
    }
    cpus as usize
}

#[cfg(feature = "nightly")]
#[feature(asm, naked_functions)]
#[naked]
#[inline(always)]
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
fn popcnt(bits: u64) -> u64 {
    let ret: u64;
    unsafe {
        asm!("popcntq %rdi, %rax"
             : "={rax}"(ret)
             : "{rdi}"(bits)
        );
    };
    ret
}

#[cfg(not(feature = "nightly"))]
#[inline(always)]
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
fn popcnt(n: u64) -> u64 {
    let mut c: u64 = (n & 0x5555555555555555) + ((n >> 1) & 0x5555555555555555);
    c = (c & 0x3333333333333333) + ((c >> 2) & 0x3333333333333333);
    c = (c & 0x0f0f0f0f0f0f0f0f) + ((c >> 4) & 0x0f0f0f0f0f0f0f0f);
    c = (c & 0x00ff00ff00ff00ff) + ((c >> 8) & 0x00ff00ff00ff00ff);
    c = (c & 0x0000ffff0000ffff) + ((c >> 16) & 0x0000ffff0000ffff);
    c = (c & 0x00000000ffffffff) + ((c >> 32) & 0x00000000ffffffff);
    c
}

#[inline(always)]
#[cfg(all(target_os = "linux", not(target_arch = "x86_64")))]
fn popcnt(n: u64) -> u64 {
    let mut c: u64 = (n & 0x5555555555555555) + ((n >> 1) & 0x5555555555555555);
    c = (c & 0x3333333333333333) + ((c >> 2) & 0x3333333333333333);
    c = (c & 0x0f0f0f0f0f0f0f0f) + ((c >> 4) & 0x0f0f0f0f0f0f0f0f);
    c = (c & 0x00ff00ff00ff00ff) + ((c >> 8) & 0x00ff00ff00ff00ff);
    c = (c & 0x0000ffff0000ffff) + ((c >> 16) & 0x0000ffff0000ffff);
    c = (c & 0x00000000ffffffff) + ((c >> 32) & 0x00000000ffffffff);
    c
}

#[cfg(target_os = "linux")]
fn get_num_cpus() -> usize {
    let mut set:  libc::cpu_set_t = unsafe { std::mem::zeroed() };
    if unsafe { libc::sched_getaffinity(0, std::mem::size_of::<libc::cpu_set_t>(), &mut set) } == 0 {
        let ptr = unsafe { std::mem::transmute::<&libc::cpu_set_t, &u64>(&set) };
        popcnt(*ptr) as usize
    } else {
        unsafe {
            libc::sysconf(libc::_SC_NPROCESSORS_ONLN) as usize
        }
    }
}

#[cfg(any(
    target_os = "nacl",
    target_os = "macos",
    target_os = "ios",
    target_os = "android",
    target_os = "solaris",
    target_os = "fuchsia")
)]
fn get_num_cpus() -> usize {
    unsafe {
        libc::sysconf(libc::_SC_NPROCESSORS_ONLN) as usize
    }
}
#[cfg(any(target_os = "emscripten", target_os = "redox", target_os = "haiku"))]
fn get_num_cpus() -> usize {
    1
}

#[test]
fn lower_bound() {
    assert!(get() > 0);
    assert!(get_physical() > 0);
}

#[test]
fn upper_bound() {
    assert!(get() < 236_451);
    assert!(get_physical() < 236_451);
}
<commit_msg>Use  instead handcrafted one<commit_after>//! Replaces the deprecated functionality of std::os::num_cpus.
#![cfg_attr(test, deny(warnings))]
#![deny(missing_docs)]
#![allow(non_snake_case)]

#[cfg(not(windows))]
extern crate libc;

/// Returns the number of CPUs of the current machine.
#[inline]
pub fn get() -> usize {
    get_num_cpus()
}

/// Returns the number of physical cores of the current machine.
/// If not possible on the particular architecture returns same as get() which
/// is the logical CPUs.
#[inline]
pub fn get_physical() -> usize {
    get_num_physical_cpus()
}

#[cfg(not(target_os = "linux"))]
#[inline]
fn get_num_physical_cpus() -> usize {
    // Not implemented, fallback
    get_num_cpus()
}

#[cfg(target_os = "linux")]
fn get_num_physical_cpus() -> usize {
    use std::io::BufReader;
    use std::io::BufRead;
    use std::fs::File;
    use std::collections::HashSet;

    let file = match File::open("/proc/cpuinfo") {
        Ok(val) => val,
        Err(_) => {return get_num_cpus()},
    };
    let reader = BufReader::new(file);
    let mut set = HashSet::new();
    let mut coreid: u32 = 0;
    let mut physid: u32 = 0;
    let mut chgcount = 0;
    for line in reader.lines().filter_map(|result| result.ok()) {
        let parts: Vec<&str> = line.split(':').map(|s| s.trim()).collect();
        if parts.len() != 2 {
            continue
        }
        if parts[0] == "core id" || parts[0] == "physical id" {
            let value = match parts[1].trim().parse() {
              Ok(val) => val,
              Err(_) => break,
            };
            match parts[0] {
                "core id"     => coreid = value,
                "physical id" => physid = value,
                _ => {},
            }
            chgcount += 1;
        }
        if chgcount == 2 {
            set.insert((physid, coreid));
            chgcount = 0;
        }
    }
    let count = set.len();
    if count == 0 { get_num_cpus() } else { count }
}

#[cfg(windows)]
fn get_num_cpus() -> usize {
    #[repr(C)]
    struct SYSTEM_INFO {
        wProcessorArchitecture: u16,
        wReserved: u16,
        dwPageSize: u32,
        lpMinimumApplicationAddress: *mut u8,
        lpMaximumApplicationAddress: *mut u8,
        dwActiveProcessorMask: *mut u8,
        dwNumberOfProcessors: u32,
        dwProcessorType: u32,
        dwAllocationGranularity: u32,
        wProcessorLevel: u16,
        wProcessorRevision: u16,
    }

    extern "system" {
        fn GetSystemInfo(lpSystemInfo: *mut SYSTEM_INFO);
    }

    unsafe {
        let mut sysinfo: SYSTEM_INFO = std::mem::uninitialized();
        GetSystemInfo(&mut sysinfo);
        sysinfo.dwNumberOfProcessors as usize
    }
}

#[cfg(any(target_os = "freebsd",
          target_os = "dragonfly",
          target_os = "bitrig",
          target_os = "netbsd"))]
fn get_num_cpus() -> usize {
    let mut cpus: libc::c_uint = 0;
    let mut cpus_size = std::mem::size_of_val(&cpus);

    unsafe {
        cpus = libc::sysconf(libc::_SC_NPROCESSORS_ONLN) as libc::c_uint;
    }
    if cpus < 1 {
        let mut mib = [libc::CTL_HW, libc::HW_NCPU, 0, 0];
        unsafe {
            libc::sysctl(mib.as_mut_ptr(),
                         2,
                         &mut cpus as *mut _ as *mut _,
                         &mut cpus_size as *mut _ as *mut _,
                         0 as *mut _,
                         0);
        }
        if cpus < 1 {
            cpus = 1;
        }
    }
    cpus as usize
}

#[cfg(target_os = "openbsd")]
fn get_num_cpus() -> usize {
    let mut cpus: libc::c_uint = 0;
    let mut cpus_size = std::mem::size_of_val(&cpus);
    let mut mib = [libc::CTL_HW, libc::HW_NCPU, 0, 0];

    unsafe {
        libc::sysctl(mib.as_mut_ptr(),
                     2,
                     &mut cpus as *mut _ as *mut _,
                     &mut cpus_size as *mut _ as *mut _,
                     0 as *mut _,
                     0);
    }
    if cpus < 1 {
        cpus = 1;
    }
    cpus as usize
}

#[cfg(target_os = "linux")]
fn get_num_cpus() -> usize {
    let mut set:  libc::cpu_set_t = unsafe { std::mem::zeroed() };
    if unsafe { libc::sched_getaffinity(0, std::mem::size_of::<libc::cpu_set_t>(), &mut set) } == 0 {
        let ptr = unsafe { std::mem::transmute::<&libc::cpu_set_t, &u64>(&set) };
        (*ptr).count_ones() as usize
    } else {
        unsafe {
            libc::sysconf(libc::_SC_NPROCESSORS_ONLN) as usize
        }
    }
}

#[cfg(any(
    target_os = "nacl",
    target_os = "macos",
    target_os = "ios",
    target_os = "android",
    target_os = "solaris",
    target_os = "fuchsia")
)]
fn get_num_cpus() -> usize {
    unsafe {
        libc::sysconf(libc::_SC_NPROCESSORS_ONLN) as usize
    }
}
#[cfg(any(target_os = "emscripten", target_os = "redox", target_os = "haiku"))]
fn get_num_cpus() -> usize {
    1
}

#[test]
fn lower_bound() {
    assert!(get() > 0);
    assert!(get_physical() > 0);
}

#[test]
fn upper_bound() {
    assert!(get() < 236_451);
    assert!(get_physical() < 236_451);
}
<|endoftext|>