<commit_before>#[cfg(feature = "serde")]
mod serde;

mod entry;
mod iter;

use core::iter::{IntoIterator, Iterator};
use iter::*;

pub use entry::*;

#[derive(Clone)]
pub struct IntMap<V> {
    cache: Vec<Vec<(u64, V)>>,
    size: u32,
    mod_mask: u64,
    count: usize,
    load_factor: usize,
}

impl<V> IntMap<V> {
    /// Creates a new IntMap.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// ```
    pub fn new() -> Self {
        IntMap::with_capacity(4)
    }

    /// Creates a new IntMap with at least the given capacity, rounded
    /// to the next power of two.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::with_capacity(20);
    /// ```
    pub fn with_capacity(capacity: usize) -> Self {
        let mut map = IntMap {
            cache: Vec::new(),
            size: 0,
            count: 0,
            mod_mask: 0,
            load_factor: 700, // 70.0%
        };

        map.increase_cache();

        while map.lim() < capacity {
            map.increase_cache();
        }

        map
    }

    /// Sets load rate of IntMap rounded to the first decimal point.
    ///
    /// Values above 1.0 is allowed.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::with_capacity(20);
    /// map.set_load_factor(0.909); // Sets load factor to 90.9%
    /// ```
    pub fn set_load_factor(&mut self, load_factor: f32) {
        self.load_factor = (load_factor * 1000.) as usize;
        self.ensure_load_rate();
    }

    /// Returns current load_factor
    pub fn get_load_factor(&self) -> f32 {
        self.load_factor as f32 / 1000.
    }

    /// Ensures that the IntMap has space for at least `additional` more elements
    pub fn reserve(&mut self, additional: usize) {
        let capacity = (self.count + additional).next_power_of_two();
        while self.lim() < capacity {
            self.increase_cache();
        }
    }

    /// Insert key/value into the IntMap if the key is not yet inserted.
    ///
    /// This function returns true if key/value were inserted and false otherwise.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map = IntMap::new();
    /// assert!(map.insert(21, "Eat my shorts"));
    /// assert!(!map.insert(21, "Ay, caramba"));
    /// assert_eq!(map.get(21), Some(&"Eat my shorts"));
    /// ```
    pub fn insert(&mut self, key: u64, value: V) -> bool {
        let ix = self.calc_index(key);

        let ref mut vals = self.cache[ix];
        if vals.iter().any(|kv| kv.0 == key) {
            return false;
        }

        self.count += 1;
        vals.push((key, value));
        if (self.count & 4) == 4 {
            self.ensure_load_rate();
        }

        true
    }

    /// Get value from the IntMap.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(21, 42);
    /// let val = map.get(21);
    /// assert!(val.is_some());
    /// assert_eq!(*val.unwrap(), 42);
    /// assert!(map.contains_key(21));
    /// ```
    pub fn get(&self, key: u64) -> Option<&V> {
        let ix = self.calc_index(key);

        let ref vals = self.cache[ix];

        if vals.len() > 0 {
            return vals.iter().find_map(|kv| (kv.0 == key).then(|| &kv.1));
        } else {
            return None;
        }
    }

    /// Get mutable value from the IntMap.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(21, 42);
    ///
    /// assert_eq!(*map.get(21).unwrap(), 42);
    /// assert!(map.contains_key(21));
    ///
    /// {
    ///     let mut val = map.get_mut(21).unwrap();
    ///     *val+=1;
    /// }
    ///     assert_eq!(*map.get(21).unwrap(), 43);
    /// ```
    pub fn get_mut(&mut self, key: u64) -> Option<&mut V> {
        let ix = self.calc_index(key);

        let ref mut vals = self.cache[ix];

        if vals.len() > 0 {
            return vals
                .iter_mut()
                .find_map(|kv| (kv.0 == key).then(move || &mut kv.1));
        } else {
            return None;
        }
    }

    /// Remove value from the IntMap.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(21, 42);
    /// let val = map.remove(21);
    /// assert!(val.is_some());
    /// assert_eq!(val.unwrap(), 42);
    /// assert!(!map.contains_key(21));
    /// ```
    pub fn remove(&mut self, key: u64) -> Option<V> {
        let ix = self.calc_index(key);

        let ref mut vals = self.cache[ix];

        if vals.len() > 0 {
            for i in 0..vals.len() {
                let peek = vals[i].0;

                if peek == key {
                    self.count -= 1;
                    let kv = vals.swap_remove(i);
                    return Some(kv.1);
                }
            }

            return None;
        } else {
            return None;
        }
    }

    /// Returns true if key is in map.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(21, 42);
    /// assert!(map.contains_key(21));
    /// ```
    pub fn contains_key(&self, key: u64) -> bool {
        match self.get(key) {
            Some(_) => true,
            None => false,
        }
    }

    /// Removes all elements from map.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(21, 42);
    /// map.clear();
    /// assert_eq!(map.len(), 0);
    /// ```
    pub fn clear(&mut self) {
        for vals in &mut self.cache {
            vals.clear();
        }

        self.count = 0;
    }

    /// Retains only the elements specified by the predicate.
    ///
    /// In other words, remove all elements such that `f(key, &value)` returns false.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(1, 11);
    /// map.insert(2, 12);
    /// map.insert(4, 13);
    ///
    /// // retain only the odd values
    /// map.retain(|k, v| *v % 2 == 1);
    ///
    /// assert_eq!(map.len(), 2);
    /// assert!(map.contains_key(1));
    /// assert!(map.contains_key(4));
    /// ```
    pub fn retain<F>(&mut self, mut f: F)
    where
        F: FnMut(u64, &V) -> bool,
    {
        let mut removed = 0;
        for vals in &mut self.cache {
            vals.retain(|(k, v)| {
                let keep = (f)(*k, v);
                if !keep {
                    removed += 1;
                }
                keep
            });
        }

        self.count -= removed;
    }

    /// Returns true if map is empty
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(21, 42);
    /// assert!(!map.is_empty());
    /// map.remove(21);
    /// assert!(map.is_empty());
    /// ```
    pub fn is_empty(&mut self) -> bool {
        self.count == 0
    }

    //**** Iterators *****

    pub fn iter(&self) -> Iter<u64, V> {
        Iter::new(&self.cache)
    }

    pub fn iter_mut(&mut self) -> IterMut<u64, V> {
        IterMut::new(&mut self.cache)
    }

    pub fn keys(&self) -> Keys<u64, V> {
        Keys { inner: self.iter() }
    }

    pub fn values(&self) -> Values<u64, V> {
        Values { inner: self.iter() }
    }

    pub fn values_mut(&mut self) -> ValuesMut<u64, V> {
        ValuesMut {
            inner: self.iter_mut(),
        }
    }

    pub fn drain(&mut self) -> Drain<u64, V> {
        Drain::new(&mut self.cache, &mut self.count)
    }

    //**** Internal hash stuff *****

    #[inline]
    fn hash_u64(seed: u64) -> u64 {
        let a = 11400714819323198549u64;
        let val = a.wrapping_mul(seed);
        val
    }

    #[inline]
    pub(crate) fn calc_index(&self, key: u64) -> usize {
        let hash = Self::hash_u64(key);
        // Faster modulus
        (hash & self.mod_mask) as usize
    }

    #[inline]
    fn lim(&self) -> usize {
        2u64.pow(self.size) as usize
    }

    fn increase_cache(&mut self) {
        self.size += 1;
        let new_lim = self.lim();
        self.mod_mask = (new_lim as u64) - 1;

        let mut vec: Vec<Vec<(u64, V)>> = Vec::new();

        vec.append(&mut self.cache);

        for _ in 0..new_lim {
            self.cache.push(Vec::with_capacity(0));
        }

        for k in vec.into_iter().flatten() {
            let ix = self.calc_index(k.0);

            let ref mut vals = self.cache[ix];
            vals.push(k);
        }

        debug_assert!(
            self.cache.len() == self.lim(),
            "cache vector the wrong length, lim: {:?} cache: {:?}",
            self.lim(),
            self.cache.len()
        );
    }

    fn ensure_load_rate(&mut self) {
        // Tried using floats here but insert performance tanked.
        while ((self.count * 1000) / self.cache.len()) > self.load_factor {
            self.increase_cache();
        }
    }

    /// Number of elements in map.
    ///
    pub fn len(&self) -> usize {
        self.count as usize
    }

    /// Force count number of slots filled.
    ///
    pub fn load(&self) -> u64 {
        self.cache.iter().filter(|vals| !vals.is_empty()).count() as u64
    }

    pub fn load_rate(&self) -> f64 {
        (self.count as f64) / (self.cache.len() as f64) * 100f64
    }

    /// Total number of slots available.
    ///
    pub fn capacity(&self) -> usize {
        self.cache.len()
    }

    pub fn assert_count(&self) -> bool {
        let count = self.cache.iter().flatten().count();

        self.count == count
    }

    pub fn collisions(&self) -> IntMap<u64> {
        let mut map = IntMap::new();

        for s in self.cache.iter() {
            let key = s.len() as u64;
            if key > 1 {
                if !map.contains_key(key) {
                    map.insert(key, 1);
                } else {
                    let counter = map.get_mut(key).unwrap();
                    *counter += 1;
                }
            }
        }

        // map.sort();

        map
    }

    /// Gets the [`Entry`] that corresponds to the given key.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::{IntMap, Entry};
    ///
    /// let mut counters = IntMap::new();
    ///
    /// for number in [10, 30, 10, 40, 50, 50, 60, 50] {
    ///     let counter = match counters.entry(number) {
    ///         Entry::Occupied(entry) => entry.into_mut(),
    ///         Entry::Vacant(entry) => entry.insert(0),
    ///     };
    ///     *counter += 1;
    /// }
    ///
    /// assert_eq!(counters.get(10), Some(&2));
    /// assert_eq!(counters.get(20), None);
    /// assert_eq!(counters.get(30), Some(&1));
    /// assert_eq!(counters.get(40), Some(&1));
    /// assert_eq!(counters.get(50), Some(&3));
    /// assert_eq!(counters.get(60), Some(&1));
    /// ```
    pub fn entry(&mut self, key: u64) -> Entry<V> {
        Entry::new(key, self)
    }
}

// ***************** Equality *********************

impl<V> PartialEq for IntMap<V>
where
    V: PartialEq,
{
    fn eq(&self, other: &IntMap<V>) -> bool {
        self.iter().all(|(k, a)| other.get(*k) == Some(a))
    }
}
impl<V> Eq for IntMap<V> where V: Eq {}

// ***************** Debug *********************

impl<V> std::fmt::Debug for IntMap<V>
where
    V: std::fmt::Debug,
{
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        fmt.debug_map().entries(self.iter()).finish()
    }
}
<commit_msg>Clean up some branching in get and remove<commit_after>#[cfg(feature = "serde")]
mod serde;

mod entry;
mod iter;

use core::iter::{IntoIterator, Iterator};
use iter::*;

pub use entry::*;

#[derive(Clone)]
pub struct IntMap<V> {
    cache: Vec<Vec<(u64, V)>>,
    size: u32,
    mod_mask: u64,
    count: usize,
    load_factor: usize,
}

impl<V> IntMap<V> {
    /// Creates a new IntMap.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// ```
    pub fn new() -> Self {
        IntMap::with_capacity(4)
    }

    /// Creates a new IntMap with at least the given capacity, rounded
    /// to the next power of two.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::with_capacity(20);
    /// ```
    pub fn with_capacity(capacity: usize) -> Self {
        let mut map = IntMap {
            cache: Vec::new(),
            size: 0,
            count: 0,
            mod_mask: 0,
            load_factor: 700, // 70.0%
        };

        map.increase_cache();

        while map.lim() < capacity {
            map.increase_cache();
        }

        map
    }

    /// Sets load rate of IntMap rounded to the first decimal point.
    ///
    /// Values above 1.0 is allowed.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::with_capacity(20);
    /// map.set_load_factor(0.909); // Sets load factor to 90.9%
    /// ```
    pub fn set_load_factor(&mut self, load_factor: f32) {
        self.load_factor = (load_factor * 1000.) as usize;
        self.ensure_load_rate();
    }

    /// Returns current load_factor
    pub fn get_load_factor(&self) -> f32 {
        self.load_factor as f32 / 1000.
    }

    /// Ensures that the IntMap has space for at least `additional` more elements
    pub fn reserve(&mut self, additional: usize) {
        let capacity = (self.count + additional).next_power_of_two();
        while self.lim() < capacity {
            self.increase_cache();
        }
    }

    /// Insert key/value into the IntMap if the key is not yet inserted.
    ///
    /// This function returns true if key/value were inserted and false otherwise.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map = IntMap::new();
    /// assert!(map.insert(21, "Eat my shorts"));
    /// assert!(!map.insert(21, "Ay, caramba"));
    /// assert_eq!(map.get(21), Some(&"Eat my shorts"));
    /// ```
    pub fn insert(&mut self, key: u64, value: V) -> bool {
        let ix = self.calc_index(key);

        let vals = &mut self.cache[ix];
        if vals.iter().any(|kv| kv.0 == key) {
            return false;
        }

        self.count += 1;
        vals.push((key, value));

        if (self.count & 4) == 4 {
            self.ensure_load_rate();
        }

        true
    }

    /// Get value from the IntMap.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(21, 42);
    /// let val = map.get(21);
    /// assert!(val.is_some());
    /// assert_eq!(*val.unwrap(), 42);
    /// assert!(map.contains_key(21));
    /// ```
    pub fn get(&self, key: u64) -> Option<&V> {
        let ix = self.calc_index(key);

        let vals = &self.cache[ix];

        vals.iter().find_map(|kv| (kv.0 == key).then(|| &kv.1))
    }

    /// Get mutable value from the IntMap.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(21, 42);
    ///
    /// assert_eq!(*map.get(21).unwrap(), 42);
    /// assert!(map.contains_key(21));
    ///
    /// {
    ///     let mut val = map.get_mut(21).unwrap();
    ///     *val+=1;
    /// }
    ///     assert_eq!(*map.get(21).unwrap(), 43);
    /// ```
    pub fn get_mut(&mut self, key: u64) -> Option<&mut V> {
        let ix = self.calc_index(key);

        let vals = &mut self.cache[ix];

        return vals
            .iter_mut()
            .find_map(|kv| (kv.0 == key).then(move || &mut kv.1));
    }

    /// Remove value from the IntMap.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(21, 42);
    /// let val = map.remove(21);
    /// assert!(val.is_some());
    /// assert_eq!(val.unwrap(), 42);
    /// assert!(!map.contains_key(21));
    /// ```
    pub fn remove(&mut self, key: u64) -> Option<V> {
        let ix = self.calc_index(key);

        let ref mut vals = self.cache[ix];

        for i in 0..vals.len() {
            let peek = vals[i].0;

            if peek == key {
                self.count -= 1;
                let kv = vals.swap_remove(i);
                return Some(kv.1);
            }
        }

        return None;
    }

    /// Returns true if key is in map.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(21, 42);
    /// assert!(map.contains_key(21));
    /// ```
    pub fn contains_key(&self, key: u64) -> bool {
        self.get(key).is_some()
    }

    /// Removes all elements from map.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(21, 42);
    /// map.clear();
    /// assert_eq!(map.len(), 0);
    /// ```
    pub fn clear(&mut self) {
        for vals in &mut self.cache {
            vals.clear();
        }

        self.count = 0;
    }

    /// Retains only the elements specified by the predicate.
    ///
    /// In other words, remove all elements such that `f(key, &value)` returns false.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(1, 11);
    /// map.insert(2, 12);
    /// map.insert(4, 13);
    ///
    /// // retain only the odd values
    /// map.retain(|k, v| *v % 2 == 1);
    ///
    /// assert_eq!(map.len(), 2);
    /// assert!(map.contains_key(1));
    /// assert!(map.contains_key(4));
    /// ```
    pub fn retain<F>(&mut self, mut f: F)
    where
        F: FnMut(u64, &V) -> bool,
    {
        let mut removed = 0;
        for vals in &mut self.cache {
            vals.retain(|(k, v)| {
                let keep = (f)(*k, v);
                if !keep {
                    removed += 1;
                }
                keep
            });
        }

        self.count -= removed;
    }

    /// Returns true if map is empty
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::IntMap;
    ///
    /// let mut map: IntMap<u64> = IntMap::new();
    /// map.insert(21, 42);
    /// assert!(!map.is_empty());
    /// map.remove(21);
    /// assert!(map.is_empty());
    /// ```
    pub fn is_empty(&mut self) -> bool {
        self.count == 0
    }

    //**** Iterators *****

    pub fn iter(&self) -> Iter<u64, V> {
        Iter::new(&self.cache)
    }

    pub fn iter_mut(&mut self) -> IterMut<u64, V> {
        IterMut::new(&mut self.cache)
    }

    pub fn keys(&self) -> Keys<u64, V> {
        Keys { inner: self.iter() }
    }

    pub fn values(&self) -> Values<u64, V> {
        Values { inner: self.iter() }
    }

    pub fn values_mut(&mut self) -> ValuesMut<u64, V> {
        ValuesMut {
            inner: self.iter_mut(),
        }
    }

    pub fn drain(&mut self) -> Drain<u64, V> {
        Drain::new(&mut self.cache, &mut self.count)
    }

    //**** Internal hash stuff *****

    #[inline]
    fn hash_u64(seed: u64) -> u64 {
        let a = 11400714819323198549u64;
        let val = a.wrapping_mul(seed);
        val
    }

    #[inline]
    pub(crate) fn calc_index(&self, key: u64) -> usize {
        let hash = Self::hash_u64(key);
        // Faster modulus
        (hash & self.mod_mask) as usize
    }

    #[inline]
    fn lim(&self) -> usize {
        2u64.pow(self.size) as usize
    }

    fn increase_cache(&mut self) {
        self.size += 1;
        let new_lim = self.lim();
        self.mod_mask = (new_lim as u64) - 1;

        let mut vec: Vec<Vec<(u64, V)>> = Vec::new();

        vec.append(&mut self.cache);

        for _ in 0..new_lim {
            self.cache.push(Vec::with_capacity(0));
        }

        for k in vec.into_iter().flatten() {
            let ix = self.calc_index(k.0);

            let ref mut vals = self.cache[ix];
            vals.push(k);
        }

        debug_assert!(
            self.cache.len() == self.lim(),
            "cache vector the wrong length, lim: {:?} cache: {:?}",
            self.lim(),
            self.cache.len()
        );
    }

    fn ensure_load_rate(&mut self) {
        // Tried using floats here but insert performance tanked.
        while ((self.count * 1000) / self.cache.len()) > self.load_factor {
            self.increase_cache();
        }
    }

    /// Number of elements in map.
    ///
    pub fn len(&self) -> usize {
        self.count as usize
    }

    /// Force count number of slots filled.
    ///
    pub fn load(&self) -> u64 {
        self.cache.iter().filter(|vals| !vals.is_empty()).count() as u64
    }

    pub fn load_rate(&self) -> f64 {
        (self.count as f64) / (self.cache.len() as f64) * 100f64
    }

    /// Total number of slots available.
    ///
    pub fn capacity(&self) -> usize {
        self.cache.len()
    }

    pub fn assert_count(&self) -> bool {
        let count = self.cache.iter().flatten().count();

        self.count == count
    }

    pub fn collisions(&self) -> IntMap<u64> {
        let mut map = IntMap::new();

        for s in self.cache.iter() {
            let key = s.len() as u64;
            if key > 1 {
                if !map.contains_key(key) {
                    map.insert(key, 1);
                } else {
                    let counter = map.get_mut(key).unwrap();
                    *counter += 1;
                }
            }
        }

        // map.sort();

        map
    }

    /// Gets the [`Entry`] that corresponds to the given key.
    ///
    /// # Examples
    ///
    /// ```
    /// use intmap::{IntMap, Entry};
    ///
    /// let mut counters = IntMap::new();
    ///
    /// for number in [10, 30, 10, 40, 50, 50, 60, 50] {
    ///     let counter = match counters.entry(number) {
    ///         Entry::Occupied(entry) => entry.into_mut(),
    ///         Entry::Vacant(entry) => entry.insert(0),
    ///     };
    ///     *counter += 1;
    /// }
    ///
    /// assert_eq!(counters.get(10), Some(&2));
    /// assert_eq!(counters.get(20), None);
    /// assert_eq!(counters.get(30), Some(&1));
    /// assert_eq!(counters.get(40), Some(&1));
    /// assert_eq!(counters.get(50), Some(&3));
    /// assert_eq!(counters.get(60), Some(&1));
    /// ```
    pub fn entry(&mut self, key: u64) -> Entry<V> {
        Entry::new(key, self)
    }
}

// ***************** Equality *********************

impl<V> PartialEq for IntMap<V>
where
    V: PartialEq,
{
    fn eq(&self, other: &IntMap<V>) -> bool {
        self.iter().all(|(k, a)| other.get(*k) == Some(a))
    }
}
impl<V> Eq for IntMap<V> where V: Eq {}

// ***************** Debug *********************

impl<V> std::fmt::Debug for IntMap<V>
where
    V: std::fmt::Debug,
{
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        fmt.debug_map().entries(self.iter()).finish()
    }
}
<|endoftext|>


<commit_before>// Timexif is copyright © 2016, Dominic van Berkel <dev+timexif@baudvine.net>
// Distributed under the MIT open source license; see the file LICENSE.txt for
// details.

//! The sequencer module turns a collection or glob of paths to images into
//! a sequence of `Image` with known time.

use std::path::{Path, PathBuf};
use rexif::ExifData;
use walkdir::{IterFilterEntry, DirEntry, WalkDir};

use ::Image;

pub fn get_images<P: AsRef<Path>>(dir: P) -> Vec<Image> {
    WalkDir::new(dir).into_iter()
        .filter_map(|r| r.ok())
        .filter(contains_exif)
        .filter(has_timestamp)
        .map(to_image)
        .collect::<Vec<Image>>()
}

fn contains_exif(path: &DirEntry) -> bool {
    unimplemented!();
}

fn has_timestamp(path: &DirEntry) -> bool {
    unimplemented!();
}

fn to_image(path: DirEntry) -> Image {
    unimplemented!();
}
<commit_msg>Sequencer appears to be complete?<commit_after>// Timexif is copyright © 2016, Dominic van Berkel <dev+timexif@baudvine.net>
// Distributed under the MIT open source license; see the file LICENSE.txt for
// details.

//! The sequencer module turns a collection or glob of paths to images into
//! a sequence of `Image` with known time.

use std::path::Path;

use rexif;
use walkdir::WalkDir;

use ::Image;

pub fn get_images<P: AsRef<Path>>(dir: P) -> Vec<Image> {
    WalkDir::new(dir).into_iter()
        // First, all DirEntries that are not errors
        .filter_map(|d| d.ok())
        // Convert to String so we can pass it around more easily
        .filter_map(|d| d.path().to_str().map(String::from))
        .filter_map(to_image)
        .collect::<Vec<Image>>()
}

fn get_exif_datetime(path: &String) -> Option<String> {
    let tagvalue = match rexif::parse_file(path) {
        Ok(record) => record.entries.iter()
            // See if any tags are DateTimeOriginal
            // TODO: Might also want DateTime?
            .find(|&e| e.tag == rexif::ExifTag::DateTimeOriginal)
            .map(|e| e.value.clone()),
        Err(_) => return None
    };

    match tagvalue {
        Some(tagvalue) => match tagvalue {
            rexif::TagValue::Ascii(dt) => Some(dt),
            _ => None
        },
        _ => None
    }
}

fn to_image(path: String) -> Option<Image> {
    get_exif_datetime(&path)
        .and_then(|dt| Image::from_exif_datetime(&path, &dt).ok())
}
<|endoftext|>


<commit_before>#![allow(unstable)]

#[derive(Show, PartialEq)]
pub struct Error;

pub type ParseResult<O, I> = Result<(O, I), Error>;
pub type StrResult<'a, O> = ParseResult<O, &'a str>;

pub trait Parser {
    type Input: Clone;
    type Output;
    fn parse(&mut self, input: <Self as Parser>::Input) ->
        Result<(<Self as Parser>::Output, <Self as Parser>::Input), Error>;
}
impl <'a, I, O, P> Parser for &'a mut P
    where I: Clone, P: Parser<Input=I, Output=O> {
    type Input = I;
    type Output = O;
    fn parse(&mut self, input: I) -> Result<(O, I), Error> {
        (*self).parse(input)
    }
}

pub fn char<'a>(input: &'a str) -> ParseResult<char, &'a str> {
    match input.slice_shift_char() {
        Some(x) => Ok(x),
        None => Err(Error)
    }
}

pub struct Many
 {
    parser: P
}
impl
 Parser for Many
 {
    type Input =
::Input;
    type Output = Vec<
::Output>;
    fn parse(&mut self, mut input:
::Input) -> Result<(Vec<
::Output>,
::Input), Error> {
        let mut result = Vec::new();
        loop {
            match self.parser.parse(input.clone()) {
                Ok((x, rest)) => {
                    result.push(x);
                    input = rest;
                }
                Err(_) => break
            }
        }
        Ok((result, input))
    }
}
pub fn many
(p: P) -> Many
 {
    Many { parser: p }
}

pub fn many1<'a, P: Clone + 'a>(mut p: P) -> FnParser<'a,
::Input, Vec<
::Output>>
    where P: Parser {
    Box::new(move |&mut:input| {
        let (first, input) = try!(p.parse(input));
        let (mut rest, input) = try!(many(&mut p).parse(input));
        rest.insert(0, first);
        Ok((rest, input))
    })
}

pub struct SepBy
 {
    parser: P,
    separator: S
}
impl
 Parser for SepBy

    where P: Parser, S: Parser<Input=
::Input> {

    type Input =
::Input;
    type Output = Vec<
::Output>;
    fn parse(&mut self, mut input:
::Input) -> Result<(Vec<
::Output>,
::Input), Error> {
        let mut result = Vec::new();
        match self.parser.parse(input.clone()) {
            Ok((x, rest)) => {
                result.push(x);
                input = rest;
            }
            Err(_) => return Ok((result, input))
        }
        loop {
            match self.separator.parse(input.clone()) {
                Ok((_, rest)) => input = rest,
                Err(e) => break
            }
            match self.parser.parse(input.clone()) {
                Ok((x, rest)) => {
                    result.push(x);
                    input = rest;
                }
                Err(e) => return Err(e)
            }
        }
        Ok((result, input))
    }
}
pub fn sep_by
(parser: P, separator: S) -> SepBy
 {
    SepBy { parser: parser, separator: separator }
}

pub type FnParser<'a, I, O> = Box<FnMut(I) -> Result<(O, I), Error> + 'a>;

impl <'a, I: Clone, O> Parser for FnParser<'a, I, O> {
    type Input = I;
    type Output = O;
    fn parse(&mut self, input: I) -> Result<(O, I), Error> {
        self(input)
    }
}

impl <'a, I: Clone, O> Parser for fn (I) -> ParseResult<O, I> {
    type Input = I;
    type Output = O;
    fn parse(&mut self, input: I) -> ParseResult<O, I> {
        self(input)
    }
}

pub struct Satisfy<Pred> { pred: Pred }

impl <'a, Pred> Parser for Satisfy<Pred>
    where Pred: FnMut(char) -> bool {

    type Input = &'a str;
    type Output = char;
    fn parse(&mut self, input: &'a str) -> Result<(char, &'a str), Error> {
        match input.slice_shift_char() {
            Some((c, s)) => {
                if (self.pred)(c) { Ok((c, s)) }
                else { Err(Error) }
            }
            None => Err(Error)
        }
    }
}

pub fn satisfy<Pred>(pred: Pred) -> Satisfy<Pred>
    where Pred: FnMut(char) -> bool {
    Satisfy { pred: pred }
}

pub fn space() -> Satisfy<fn (char) -> bool> {
    satisfy(CharExt::is_whitespace as fn (char) -> bool)
}

pub struct StringP<'a> { s: &'a str }
impl <'a, 'b> Parser for StringP<'b> {
    type Input = &'a str;
    type Output = &'a str;
    fn parse(&mut self, input: &'a str) -> Result<(&'a str, &'a str), Error> {
        if input.starts_with(self.s) {
            let l = self.s.len();
            Ok((&input[..l], &input[l..]))
        }
        else {
            Err(Error)
        }
    }
}

pub fn string(s: &str) -> StringP {
    StringP { s: s }
}

pub struct AndThen<P1, P2>(P1, P2);
impl  Parser for AndThen<P1, P2>
    where P1: Parser<Input=I, Output=A>, P2: Parser<Input=I, Output=B> {

    type Input = I;
    type Output = (A, B);
    fn parse(&mut self, input: I) -> Result<((A, B), I), Error> {
        let (a, rest) = try!(self.0.parse(input));
        let (b, rest) = try!(self.1.parse(rest));
        Ok(((a, b), rest))
    }
}
pub fn and_then<P1, P2>(p1: P1, p2: P2) -> AndThen<P1, P2>
    where P1: Parser, P2: Parser {
    AndThen(p1, p2)
}

pub struct Optional
(P);
impl
 Parser for Optional

    where P: Parser {
    type Input =
::Input;
    type Output = Option<
::Output>;
    fn parse(&mut self, input:
::Input) -> Result<(Option<
::Output>,
::Input), Error> {
        match self.0.parse(input.clone()) {
            Ok((x, rest)) => Ok((Some(x), rest)),
            Err(_) => Ok((None, input))
        }
    }
}
pub fn optional
(parser: P) -> Optional
 {
    Optional(parser)
}

pub struct Env {
    input: I
}

impl  Env {
    pub fn new(input: I) -> Env {
        Env { input: input }
    }

    pub fn with
(&mut self, mut parser: P) -> Result<O, Error>
        where P: Parser<Input=I, Output=O> {
        let (o, rest) = try!(parser.parse(self.input.clone()));
        self.input = rest;
        Ok(o)
    }

    pub fn result<O>(self, output: O) -> Result<(O, I), Error> {
        Ok((output, self.input))
    }
}

pub fn digit<'a>(input: &'a str) -> ParseResult<char, &'a str> {
    match input.slice_shift_char() {
        Some(x) if x.0.is_digit(10) => Ok(x),
        Some((c, _)) => Err(Error),
        None => Err(Error)
    }
}

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

    fn integer<'a>(input: &'a str) -> Result<(i64, &'a str), Error> {
        let mut env = Env::new(input);
        let chars = try!(env.with(many(digit as fn(_) -> _)));
        let mut n = 0;
        for &c in chars.iter() {
            n = n * 10 + (c as i64 - '0' as i64);
        }
        env.result(n)
    }

    #[test]
    fn test_integer() {
        assert_eq!((integer as fn(_) -> _).parse("123"), Ok((123i64, "")));
    }
    #[test]
    fn list() {
        let mut p = sep_by(integer as fn(_) -> _, satisfy(|c| c == ','));
        assert_eq!(p.parse("123,4,56"), Ok((vec![123, 4, 56], "")));
    }
}

/*

impl Parser for  {
    type Input = ;
    type Output = ;
    fn parse(&mut self, input: ) -> Result<(, ), Error>;
}
 */
<commit_msg>Added FnParser and ManyAppend<commit_after>#![allow(unstable)]

#[derive(Show, PartialEq)]
pub struct Error;

pub type ParseResult<O, I> = Result<(O, I), Error>;
pub type StrResult<'a, O> = ParseResult<O, &'a str>;

pub trait Parser {
    type Input: Clone;
    type Output;
    fn parse(&mut self, input: <Self as Parser>::Input) ->
        Result<(<Self as Parser>::Output, <Self as Parser>::Input), Error>;
}
impl <'a, I, O, P> Parser for &'a mut P
    where I: Clone, P: Parser<Input=I, Output=O> {
    type Input = I;
    type Output = O;
    fn parse(&mut self, input: I) -> Result<(O, I), Error> {
        (*self).parse(input)
    }
}

pub fn char<'a>(input: &'a str) -> ParseResult<char, &'a str> {
    match input.slice_shift_char() {
        Some(x) => Ok(x),
        None => Err(Error)
    }
}

pub struct ManyAppend<'a, O: 'a, P: Parser<Output=O> + 'a> {
    parser: P,
    vec: &'a mut Vec<O>
}
impl <'a, O, P: Parser<Output=O> + 'a> Parser for ManyAppend<'a, O, P> {
    type Input =
::Input;
    type Output = ();
    fn parse(&mut self, mut input:
::Input) -> Result<((),
::Input), Error> {
        loop {
            match self.parser.parse(input.clone()) {
                Ok((x, rest)) => {
                    self.vec.push(x);
                    input = rest;
                }
                Err(_) => break
            }
        }
        Ok(((), input))
    }
}

pub fn many_append<'a, O, P: Parser<Output=O>>(parser: P, vec: &'a mut Vec<O>) -> ManyAppend<'a, O, P> {
    ManyAppend { parser: parser, vec: vec }
}

pub struct Many
 {
    parser: P
}
impl
 Parser for Many
 {
    type Input =
::Input;
    type Output = Vec<
::Output>;
    fn parse(&mut self, input:
::Input) -> Result<(Vec<
::Output>,
::Input), Error> {
        let mut result = Vec::new();
        let ((), input) = try!(many_append(&mut self.parser, &mut result).parse(input));
        Ok((result, input))
    }
}
pub fn many
(p: P) -> Many
 {
    Many { parser: p }
}

pub fn many1<'a, P: Clone + 'a>(mut p: P) -> Box<Parser<Input=
::Input, Output=Vec<
::Output>> + 'a>
    where P: Parser {
    Box::new(FnParser(move |&mut:input| {
        let (first, input) = try!(p.parse(input));
        let mut result = vec![first];
        let ((), input) = try!(many_append(&mut p, &mut result).parse(input));
        Ok((result, input))
    }))
}

pub struct SepBy
 {
    parser: P,
    separator: S
}
impl
 Parser for SepBy

    where P: Parser, S: Parser<Input=
::Input> {

    type Input =
::Input;
    type Output = Vec<
::Output>;
    fn parse(&mut self, mut input:
::Input) -> Result<(Vec<
::Output>,
::Input), Error> {
        let mut result = Vec::new();
        match self.parser.parse(input.clone()) {
            Ok((x, rest)) => {
                result.push(x);
                input = rest;
            }
            Err(_) => return Ok((result, input))
        }
        let rest = FnParser(|input| {
            let mut env = Env::new(input);
            try!(env.with(&mut self.separator));
            let v = try!(env.with(&mut self.parser));
            env.result(v)
        });
        let ((), input) = try!(many_append(rest, &mut result).parse(input));
        Ok((result, input))
    }
}
pub fn sep_by
(parser: P, separator: S) -> SepBy
 {
    SepBy { parser: parser, separator: separator }
}

impl <'a, I: Clone, O> Parser for Box<FnMut(I) -> Result<(O, I), Error> + 'a> {
    type Input = I;
    type Output = O;
    fn parse(&mut self, input: I) -> Result<(O, I), Error> {
        self(input)
    }
}
struct FnParser<'a, I: Clone, O, F: FnMut(I) -> ParseResult<O, I>>(F);

impl <'a, I, O, F> Parser for FnParser<'a, I, O, F>
    where I: Clone, F: FnMut(I) -> ParseResult<O, I> {
    type Input = I;
    type Output = O;
    fn parse(&mut self, input: I) -> ParseResult<O, I> {
        (self.0)(input)
    }
}

impl <'a, I: Clone, O> Parser for fn (I) -> ParseResult<O, I> {
    type Input = I;
    type Output = O;
    fn parse(&mut self, input: I) -> ParseResult<O, I> {
        self(input)
    }
}

pub struct Satisfy<Pred> { pred: Pred }

impl <'a, Pred> Parser for Satisfy<Pred>
    where Pred: FnMut(char) -> bool {

    type Input = &'a str;
    type Output = char;
    fn parse(&mut self, input: &'a str) -> Result<(char, &'a str), Error> {
        match input.slice_shift_char() {
            Some((c, s)) => {
                if (self.pred)(c) { Ok((c, s)) }
                else { Err(Error) }
            }
            None => Err(Error)
        }
    }
}

pub fn satisfy<Pred>(pred: Pred) -> Satisfy<Pred>
    where Pred: FnMut(char) -> bool {
    Satisfy { pred: pred }
}

pub fn space() -> Satisfy<fn (char) -> bool> {
    satisfy(CharExt::is_whitespace as fn (char) -> bool)
}

pub struct StringP<'a> { s: &'a str }
impl <'a, 'b> Parser for StringP<'b> {
    type Input = &'a str;
    type Output = &'a str;
    fn parse(&mut self, input: &'a str) -> Result<(&'a str, &'a str), Error> {
        if input.starts_with(self.s) {
            let l = self.s.len();
            Ok((&input[..l], &input[l..]))
        }
        else {
            Err(Error)
        }
    }
}

pub fn string(s: &str) -> StringP {
    StringP { s: s }
}

pub struct AndThen<P1, P2>(P1, P2);
impl  Parser for AndThen<P1, P2>
    where P1: Parser<Input=I, Output=A>, P2: Parser<Input=I, Output=B> {

    type Input = I;
    type Output = (A, B);
    fn parse(&mut self, input: I) -> Result<((A, B), I), Error> {
        let (a, rest) = try!(self.0.parse(input));
        let (b, rest) = try!(self.1.parse(rest));
        Ok(((a, b), rest))
    }
}
pub fn and_then<P1, P2>(p1: P1, p2: P2) -> AndThen<P1, P2>
    where P1: Parser, P2: Parser {
    AndThen(p1, p2)
}

pub struct Optional
(P);
impl
 Parser for Optional

    where P: Parser {
    type Input =
::Input;
    type Output = Option<
::Output>;
    fn parse(&mut self, input:
::Input) -> Result<(Option<
::Output>,
::Input), Error> {
        match self.0.parse(input.clone()) {
            Ok((x, rest)) => Ok((Some(x), rest)),
            Err(_) => Ok((None, input))
        }
    }
}
pub fn optional
(parser: P) -> Optional
 {
    Optional(parser)
}

pub struct Env {
    input: I
}

impl  Env {
    pub fn new(input: I) -> Env {
        Env { input: input }
    }

    pub fn with
(&mut self, mut parser: P) -> Result<O, Error>
        where P: Parser<Input=I, Output=O> {
        let (o, rest) = try!(parser.parse(self.input.clone()));
        self.input = rest;
        Ok(o)
    }

    pub fn result<O>(self, output: O) -> Result<(O, I), Error> {
        Ok((output, self.input))
    }
}

pub fn digit<'a>(input: &'a str) -> ParseResult<char, &'a str> {
    match input.slice_shift_char() {
        Some(x) if x.0.is_digit(10) => Ok(x),
        Some((c, _)) => Err(Error),
        None => Err(Error)
    }
}

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

    fn integer<'a>(input: &'a str) -> Result<(i64, &'a str), Error> {
        let mut env = Env::new(input);
        let chars = try!(env.with(many(digit as fn(_) -> _)));
        let mut n = 0;
        for &c in chars.iter() {
            n = n * 10 + (c as i64 - '0' as i64);
        }
        env.result(n)
    }

    #[test]
    fn test_integer() {
        assert_eq!((integer as fn(_) -> _).parse("123"), Ok((123i64, "")));
    }
    #[test]
    fn list() {
        let mut p = sep_by(integer as fn(_) -> _, satisfy(|c| c == ','));
        assert_eq!(p.parse("123,4,56"), Ok((vec![123, 4, 56], "")));
    }
}

/*

impl Parser for  {
    type Input = ;
    type Output = ;
    fn parse(&mut self, input: ) -> Result<(, ), Error>;
}
 */
<|endoftext|>


<commit_before>#![feature(test)]

#![feature(plugin)]
#![plugin(clippy)]

extern crate test;
extern crate rand;

use rand::distributions::{IndependentSample, Range};

pub struct Settings<F, R>
    where F: Fn(&[f32]) -> f32,
          R: rand::Rng
{
    pub min_max_pos: Vec<(f32, f32)>,

    pub cr_min: f32,
    pub cr_max: f32,
    pub cr_change_probability: f32,

    pub f_min: f32,
    pub f_max: f32,
    pub f_change_probability: f32,

    pub pop_size: usize,
    pub rng: R,

    pub cost_function: F,
}

impl<F, R> Settings<F, R>
    where F: Fn(&[f32]) -> f32,
          R: rand::Rng
{
    pub fn min_max_rng(min_max_pos: Vec<(f32, f32)>, cost_function: F, rng: R) -> Settings<F, R> {
        // create settings for the algorithm
        Settings {
            min_max_pos: min_max_pos,

            cr_min: 0.0,
            cr_max: 1.0,
            cr_change_probability: 0.1,

            f_min: 0.1,
            f_max: 1.0,
            f_change_probability: 0.1,

            pop_size: 100,
            rng: rng,

            cost_function: cost_function,
        }
    }
}

impl<F> Settings<F, rand::XorShiftRng>
    where F: Fn(&[f32]) -> f32
{
    pub fn new(min_max_pos: Vec<(f32, f32)>, cost_function: F) -> Settings<F, rand::XorShiftRng> {
        Settings::min_max_rng(min_max_pos, cost_function, rand::weak_rng())
    }
}

#[derive(Clone,Debug)]
struct Individual {
    pos: Vec<f32>,
    // the lower, the better.
    cost: Option<f32>,

    // control parameters
    cr: f32,
    f: f32,
}

pub struct Population<F, R>
    where F: Fn(&[f32]) -> f32,
          R: rand::Rng
{
    // TODO use a single vector for curr and best and controlparameters?
    curr: Vec<Individual>,
    best: Vec<Individual>,

    settings: Settings<F, R>,

    // index of global best individual. Might be in best or in curr.
    best_idx: Option<usize>,

    // cost value of the global best individual, for quick access
    best_cost_cache: Option<f32>,
    dim: usize,
    between_popsize: Range<usize>,
    between_dim: Range<usize>,
    between_cr: Range<f32>,
    between_f: Range<f32>,

    pop_countdown: usize,
}

impl<F, R> Iterator for Population<F, R>
    where F: Fn(&[f32]) -> f32,
          R: rand::Rng
{
    type Item = f32;

    /// Returns the cost value of the current best
    fn next(&mut self) -> Option<f32> {
        if 0 == self.pop_countdown {
            // if the whole pop has been evaluated, evolve it to update positions.
            // this also copies curr to best, if better.
            self.update_best();
            self.update_positions();
            self.pop_countdown = self.curr.len();
        }

        // perform a single fitness evaluation
        self.pop_countdown -= 1;
        let curr = &mut self.curr[self.pop_countdown];

        let cost = (self.settings.cost_function)(&curr.pos);
        curr.cost = Some(cost);

        // see if we have improved the global best
        if self.best_cost_cache.is_none() || cost < self.best_cost_cache.unwrap() {
            self.best_cost_cache = Some(cost);
            self.best_idx = Some(self.pop_countdown);
        }

        self.best_cost_cache
    }
}

impl<F> Population<F, rand::XorShiftRng>
    where F: Fn(&[f32]) -> f32
{
    pub fn new(min_max_pos: Vec<(f32, f32)>, cost_function: F) -> Population<F, rand::XorShiftRng> {
        Population::from_settings(Settings::new(min_max_pos, cost_function))
    }
}

impl<F, R> Population<F, R>
    where F: Fn(&[f32]) -> f32,
          R: rand::Rng
{
    // Creates a new population based on the given settings.
    pub fn from_settings(s: Settings<F, R>) -> Population<F, R> {
        assert!(s.min_max_pos.len() >= 1,
                "need at least one element to optimize");

        // create a vector of randomly initialized individuals for current.
        let dim = s.min_max_pos.len();

        // Empty individual, with no cost value (yet)
        let dummy_individual = Individual {
            pos: vec![0.0; dim],
            cost: None,
            cr: 0.0,
            f: 0.0,
        };

        // creates all the empty individuals
        let mut pop = Population {
            curr: vec![dummy_individual.clone(); s.pop_size],
            best: vec![dummy_individual; s.pop_size],
            best_idx: None,
            best_cost_cache: None,
            dim: dim,
            pop_countdown: s.pop_size,
            between_popsize: Range::new(0, s.pop_size),
            between_dim: Range::new(0, dim),
            between_cr: Range::new(s.cr_min, s.cr_max),
            between_f: Range::new(s.f_min, s.f_max),
            settings: s,
        };

        for ind in &mut pop.curr {
            // init control parameters
            ind.cr = pop.between_cr.ind_sample(&mut pop.settings.rng);
            ind.f = pop.between_f.ind_sample(&mut pop.settings.rng);

            // random range for each dimension
            for d in 0..dim {
                let between_min_max = Range::new(pop.settings.min_max_pos[d].0,
                                                 pop.settings.min_max_pos[d].1);
                ind.pos[d] = between_min_max.ind_sample(&mut pop.settings.rng);
            }
        }

        pop
    }

    fn update_best(&mut self) {
        for i in 0..self.curr.len() {
            let curr = &mut self.curr[i];
            let best = &mut self.best[i];

            // we use <= here, so that the individual moves even if the cost
            // stays the same.
            if best.cost.is_none() || curr.cost.unwrap() <= best.cost.unwrap() {
                // replace individual's best. swap is *much* faster than clone.
                std::mem::swap(curr, best);
            }
        }
    }

    fn update_positions(&mut self) {
        let rng = &mut self.settings.rng;
        for i in 0..self.curr.len() {
            // sample 3 different individuals
            let id1 = self.between_popsize.ind_sample(rng);

            let mut id2 = self.between_popsize.ind_sample(rng);
            while id2 == id1 {
                id2 = self.between_popsize.ind_sample(rng);
            }

            let mut id3 = self.between_popsize.ind_sample(rng);
            while id3 == id1 || id3 == id2 {
                id3 = self.between_popsize.ind_sample(rng);
            }

            let curr = &mut self.curr[i];
            let best = &self.best[i];

            // see "Self-Adapting Control Parameters in Differential Evolution:
            // A Comparative Study on Numerical Benchmark Problems"
            if rng.gen::<f32>() < self.settings.cr_change_probability {
                curr.cr = self.between_cr.ind_sample(rng);
            } else {
                curr.cr = best.cr;
            }
            if rng.gen::<f32>() < self.settings.f_change_probability {
                curr.f = self.between_f.ind_sample(rng);
            } else {
                curr.f = best.f;
            }

            let curr_pos = &mut curr.pos;
            let best_pos = &best.pos;
            let best1_pos = &self.best[id1].pos;
            let best2_pos = &self.best[id2].pos;
            let best3_pos = &self.best[id3].pos;

            let forced_mutation_dim = self.between_dim.ind_sample(rng);

            // This implements the DE/rand/1/bin, the most widely used algorithm.
            for d in 0..self.dim {
                if d == forced_mutation_dim || rng.gen::<f32>() < curr.cr {
                    curr_pos[d] = best3_pos[d] + curr.f * (best1_pos[d] - best2_pos[d]);
                } else {
                    curr_pos[d] = best_pos[d];
                }
            }

            // reset cost, has to be updated by the user.
            curr.cost = None;
        }
    }

    pub fn best(&self) -> Option<(f32, &[f32])> {
        if let Some(bi) = self.best_idx {
            let curr = &self.curr[bi];
            let best = &self.best[bi];

            if curr.cost.is_none() {
                return Some((best.cost.unwrap(), &best.pos));
            }
            if best.cost.is_none() {
                return Some((curr.cost.unwrap(), &curr.pos));
            }
            if curr.cost.unwrap() < best.cost.unwrap() {
                return Some((curr.cost.unwrap(), &curr.pos));
            }
            return Some((best.cost.unwrap(), &best.pos));
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate rand;
    use super::*;
    use test::Bencher;
    use rand::{XorShiftRng, StdRng, IsaacRng, Isaac64Rng, Rng, ChaChaRng};
    use rand::{OsRng, weak_rng};

    fn setup<F: Fn(&[f32]) -> f32, R: rand::Rng>(dim: usize,
                                                 cost_fn: F,
                                                 rng: R)
                                                 -> Population<F, R> {
        let s = Settings::min_max_rng(vec![(-100.0, 100.0); dim], cost_fn, rng);
        Population::from_settings(s)
    }

    #[bench]
    fn bench_square_fitness_opt(b: &mut Bencher) {
        let rng: XorShiftRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_thread_rng(b: &mut Bencher) {
        let mut pop = setup(5, |_| 1.234, rand::thread_rng());
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_xor_shift(b: &mut Bencher) {
        let rng: XorShiftRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_chacha(b: &mut Bencher) {
        let rng: ChaChaRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_isaac(b: &mut Bencher) {
        let rng: IsaacRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_isaac64(b: &mut Bencher) {
        let rng: Isaac64Rng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_std(b: &mut Bencher) {
        let rng: StdRng = StdRng::new().unwrap();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_osrng(b: &mut Bencher) {
        let rng: OsRng = OsRng::new().unwrap();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_weak_rng(b: &mut Bencher) {
        let rng = weak_rng();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }
}
<commit_msg>comment<commit_after>#![feature(test)]

#![feature(plugin)]
#![plugin(clippy)]

extern crate test;
extern crate rand;

use rand::distributions::{IndependentSample, Range};

pub struct Settings<F, R>
    where F: Fn(&[f32]) -> f32,
          R: rand::Rng
{
    pub min_max_pos: Vec<(f32, f32)>,

    pub cr_min: f32,
    pub cr_max: f32,
    pub cr_change_probability: f32,

    pub f_min: f32,
    pub f_max: f32,
    pub f_change_probability: f32,

    pub pop_size: usize,
    pub rng: R,

    pub cost_function: F,
}

impl<F, R> Settings<F, R>
    where F: Fn(&[f32]) -> f32,
          R: rand::Rng
{
    pub fn min_max_rng(min_max_pos: Vec<(f32, f32)>, cost_function: F, rng: R) -> Settings<F, R> {
        // create settings for the algorithm
        Settings {
            min_max_pos: min_max_pos,

            cr_min: 0.0,
            cr_max: 1.0,
            cr_change_probability: 0.1,

            f_min: 0.1,
            f_max: 1.0,
            f_change_probability: 0.1,

            pop_size: 100,
            rng: rng,

            cost_function: cost_function,
        }
    }
}

impl<F> Settings<F, rand::XorShiftRng>
    where F: Fn(&[f32]) -> f32
{
    pub fn new(min_max_pos: Vec<(f32, f32)>, cost_function: F) -> Settings<F, rand::XorShiftRng> {
        Settings::min_max_rng(min_max_pos, cost_function, rand::weak_rng())
    }
}

#[derive(Clone,Debug)]
struct Individual {
    pos: Vec<f32>,
    // the lower, the better.
    cost: Option<f32>,

    // control parameters
    cr: f32,
    f: f32,
}

pub struct Population<F, R>
    where F: Fn(&[f32]) -> f32,
          R: rand::Rng
{
    // TODO use a single vector for curr and best and controlparameters?
    curr: Vec<Individual>,
    best: Vec<Individual>,

    settings: Settings<F, R>,

    // index of global best individual. Might be in best or in curr.
    best_idx: Option<usize>,

    // cost value of the global best individual, for quick access
    best_cost_cache: Option<f32>,
    dim: usize,
    between_popsize: Range<usize>,
    between_dim: Range<usize>,
    between_cr: Range<f32>,
    between_f: Range<f32>,

    pop_countdown: usize,
}

impl<F, R> Iterator for Population<F, R>
    where F: Fn(&[f32]) -> f32,
          R: rand::Rng
{
    type Item = f32;

    /// Returns the cost value of the current best
    fn next(&mut self) -> Option<f32> {
        if 0 == self.pop_countdown {
            // if the whole pop has been evaluated, evolve it to update positions.
            // this also copies curr to best, if better.
            self.update_best();
            self.update_positions();
            self.pop_countdown = self.curr.len();
        }

        // perform a single fitness evaluation
        self.pop_countdown -= 1;
        let curr = &mut self.curr[self.pop_countdown];

        let cost = (self.settings.cost_function)(&curr.pos);
        curr.cost = Some(cost);

        // see if we have improved the global best
        if self.best_cost_cache.is_none() || cost < self.best_cost_cache.unwrap() {
            self.best_cost_cache = Some(cost);
            self.best_idx = Some(self.pop_countdown);
        }

        self.best_cost_cache
    }
}

impl<F> Population<F, rand::XorShiftRng>
    where F: Fn(&[f32]) -> f32
{
    pub fn new(min_max_pos: Vec<(f32, f32)>, cost_function: F) -> Population<F, rand::XorShiftRng> {
        Population::from_settings(Settings::new(min_max_pos, cost_function))
    }
}

impl<F, R> Population<F, R>
    where F: Fn(&[f32]) -> f32,
          R: rand::Rng
{
    // Creates a new population based on the given settings.
    pub fn from_settings(s: Settings<F, R>) -> Population<F, R> {
        assert!(s.min_max_pos.len() >= 1,
                "need at least one element to optimize");

        // create a vector of randomly initialized individuals for current.
        let dim = s.min_max_pos.len();

        // Empty individual, with no cost value (yet)
        let dummy_individual = Individual {
            pos: vec![0.0; dim],
            cost: None,
            cr: 0.0,
            f: 0.0,
        };

        // creates all the empty individuals
        let mut pop = Population {
            curr: vec![dummy_individual.clone(); s.pop_size],
            best: vec![dummy_individual; s.pop_size],
            best_idx: None,
            best_cost_cache: None,
            dim: dim,
            pop_countdown: s.pop_size,
            between_popsize: Range::new(0, s.pop_size),
            between_dim: Range::new(0, dim),
            between_cr: Range::new(s.cr_min, s.cr_max),
            between_f: Range::new(s.f_min, s.f_max),
            settings: s,
        };

        for ind in &mut pop.curr {
            // init control parameters
            ind.cr = pop.between_cr.ind_sample(&mut pop.settings.rng);
            ind.f = pop.between_f.ind_sample(&mut pop.settings.rng);

            // random range for each dimension
            for d in 0..dim {
                let between_min_max = Range::new(pop.settings.min_max_pos[d].0,
                                                 pop.settings.min_max_pos[d].1);
                ind.pos[d] = between_min_max.ind_sample(&mut pop.settings.rng);
            }
        }

        pop
    }

    fn update_best(&mut self) {
        for i in 0..self.curr.len() {
            let curr = &mut self.curr[i];
            let best = &mut self.best[i];

            // we use <= here, so that the individual moves even if the cost
            // stays the same.
            if best.cost.is_none() || curr.cost.unwrap() <= best.cost.unwrap() {
                // replace individual's best. swap is *much* faster than clone.
                std::mem::swap(curr, best);
            }
        }
    }

    // Modifies all the curr positions. This needs a lot of random numbers, so
    // for a fast cost function it is important to use a fast random number
    // generator.
    fn update_positions(&mut self) {
        let rng = &mut self.settings.rng;
        for i in 0..self.curr.len() {
            // sample 3 different individuals
            let id1 = self.between_popsize.ind_sample(rng);

            let mut id2 = self.between_popsize.ind_sample(rng);
            while id2 == id1 {
                id2 = self.between_popsize.ind_sample(rng);
            }

            let mut id3 = self.between_popsize.ind_sample(rng);
            while id3 == id1 || id3 == id2 {
                id3 = self.between_popsize.ind_sample(rng);
            }

            let curr = &mut self.curr[i];
            let best = &self.best[i];

            // see "Self-Adapting Control Parameters in Differential Evolution:
            // A Comparative Study on Numerical Benchmark Problems"
            if rng.gen::<f32>() < self.settings.cr_change_probability {
                curr.cr = self.between_cr.ind_sample(rng);
            } else {
                curr.cr = best.cr;
            }
            if rng.gen::<f32>() < self.settings.f_change_probability {
                curr.f = self.between_f.ind_sample(rng);
            } else {
                curr.f = best.f;
            }

            let curr_pos = &mut curr.pos;
            let best_pos = &best.pos;
            let best1_pos = &self.best[id1].pos;
            let best2_pos = &self.best[id2].pos;
            let best3_pos = &self.best[id3].pos;

            let forced_mutation_dim = self.between_dim.ind_sample(rng);

            // This implements the DE/rand/1/bin, the most widely used algorithm.
            // See "A Comparative Study of Differential Evolution Variants for
            // Global Optimization (2006)".
            for d in 0..self.dim {
                if d == forced_mutation_dim || rng.gen::<f32>() < curr.cr {
                    curr_pos[d] = best3_pos[d] + curr.f * (best1_pos[d] - best2_pos[d]);
                } else {
                    curr_pos[d] = best_pos[d];
                }
            }

            // reset cost, has to be updated by the user.
            curr.cost = None;
        }
    }

    pub fn best(&self) -> Option<(f32, &[f32])> {
        if let Some(bi) = self.best_idx {
            let curr = &self.curr[bi];
            let best = &self.best[bi];

            if curr.cost.is_none() {
                return Some((best.cost.unwrap(), &best.pos));
            }
            if best.cost.is_none() {
                return Some((curr.cost.unwrap(), &curr.pos));
            }
            if curr.cost.unwrap() < best.cost.unwrap() {
                return Some((curr.cost.unwrap(), &curr.pos));
            }
            return Some((best.cost.unwrap(), &best.pos));
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate rand;
    use super::*;
    use test::Bencher;
    use rand::{XorShiftRng, StdRng, IsaacRng, Isaac64Rng, Rng, ChaChaRng};
    use rand::{OsRng, weak_rng};

    fn setup<F: Fn(&[f32]) -> f32, R: rand::Rng>(dim: usize,
                                                 cost_fn: F,
                                                 rng: R)
                                                 -> Population<F, R> {
        let s = Settings::min_max_rng(vec![(-100.0, 100.0); dim], cost_fn, rng);
        Population::from_settings(s)
    }

    #[bench]
    fn bench_square_fitness_opt(b: &mut Bencher) {
        let rng: XorShiftRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_thread_rng(b: &mut Bencher) {
        let mut pop = setup(5, |_| 1.234, rand::thread_rng());
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_xor_shift(b: &mut Bencher) {
        let rng: XorShiftRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_chacha(b: &mut Bencher) {
        let rng: ChaChaRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_isaac(b: &mut Bencher) {
        let rng: IsaacRng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_isaac64(b: &mut Bencher) {
        let rng: Isaac64Rng = OsRng::new().unwrap().gen();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_std(b: &mut Bencher) {
        let rng: StdRng = StdRng::new().unwrap();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_osrng(b: &mut Bencher) {
        let rng: OsRng = OsRng::new().unwrap();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }

    #[bench]
    fn rand_weak_rng(b: &mut Bencher) {
        let rng = weak_rng();
        let mut pop = setup(5, |_| 1.234, rng);
        b.iter(|| pop.next());
    }
}
<|endoftext|>


<commit_before>#![no_std]
//!
//! You can populate [`Petnames`] with your own word lists, but the word lists
//! from upstream [petname](https://github.com/dustinkirkland/petname) are
//! included with the `default_dictionary` feature (enabled by default). See
//! [`Petnames::small`], [`Petnames::medium`], and [`Petnames::large`] to select
//! a particular built-in word list, or use the [`Default`] implementation.
//!
//! The other thing you need is a random number generator from [rand][]:
//!
//! ```rust
//! let mut rng = rand::thread_rng();
//! let pname = petname::Petnames::default().generate(&mut rng, 7, ":");
//! ```
//!
//! It may be more convenient to use the default random number generator:
//!
//! ```rust
//! let pname = petname::Petnames::default().generate_one(7, ":");
//! ```
//!
//! There's a [convenience function][petname] that'll do all of this:
//!
//! ```rust
//! let pname = petname::petname(7, ":");
//! ```
//!
//! But the most flexible approach is to create an [`Iterator`] with
//! [`iter`][`Petnames::iter`]:
//!
//! ```rust
//! let mut rng = rand::thread_rng();
//! let petnames = petname::Petnames::default();
//! let ten_thousand_names: Vec<String> =
//!   petnames.iter(&mut rng, 3, "_").take(10000).collect();
//! ```
//!
//! You can modify the word lists to, for example, only use words beginning with
//! the letter "b":
//!
//! ```rust
//! let mut petnames = petname::Petnames::default();
//! petnames.retain(|s| s.starts_with("b"));
//! petnames.generate_one(3, ".");
//! ```
//!

extern crate alloc;

use alloc::{
    string::{String, ToString},
    vec::Vec,
};

use itertools::Itertools;
use rand::seq::SliceRandom;

/// Convenience function to generate a new petname from default word lists.
#[allow(dead_code)]
#[cfg(feature = "std_rng")]
#[cfg(feature = "default_dictionary")]
pub fn petname(words: u8, separator: &str) -> String {
    Petnames::new().generate_one(words, separator)
}

/// A word list.
pub type Words<'a> = Vec<&'a str>;

/// Word lists and the logic to combine them into _petnames_.
///
/// A _petname_ with `n` words will contain, in order:
///
///   * `n - 2` adverbs when `n >= 2`, otherwise 0 adverbs.
///   * 1 adjective when `n >= 2`, otherwise 0 adjectives.
///   * 1 name / noun when `n >= 1`, otherwise 0 names.
///
#[derive(Clone, Debug, PartialEq)]
pub struct Petnames<'a> {
    pub adjectives: Words<'a>,
    pub adverbs: Words<'a>,
    pub names: Words<'a>,
}

impl<'a> Petnames<'a> {
    /// Constructs a new `Petnames` from the default (small) word lists.
    #[cfg(feature = "default_dictionary")]
    pub fn new() -> Self {
        Self::default()
    }

    /// Constructs a new `Petnames` from the small word lists.
    #[cfg(feature = "default_dictionary")]
    pub fn small() -> Self {
        Self::init(
            include_str!("../words/small/adjectives.txt"),
            include_str!("../words/small/adverbs.txt"),
            include_str!("../words/small/names.txt"),
        )
    }

    /// Constructs a new `Petnames` from the medium word lists.
    #[cfg(feature = "default_dictionary")]
    pub fn medium() -> Self {
        Self::init(
            include_str!("../words/medium/adjectives.txt"),
            include_str!("../words/medium/adverbs.txt"),
            include_str!("../words/medium/names.txt"),
        )
    }

    /// Constructs a new `Petnames` from the large word lists.
    #[cfg(feature = "default_dictionary")]
    pub fn large() -> Self {
        Self::init(
            include_str!("../words/large/adjectives.txt"),
            include_str!("../words/large/adverbs.txt"),
            include_str!("../words/large/names.txt"),
        )
    }

    /// Constructs a new `Petnames` from the given word lists.
    ///
    /// The words are extracted from the given strings by splitting on whitespace.
    pub fn init(adjectives: &'a str, adverbs: &'a str, names: &'a str) -> Self {
        Self {
            adjectives: adjectives.split_whitespace().collect(),
            adverbs: adverbs.split_whitespace().collect(),
            names: names.split_whitespace().collect(),
        }
    }

    /// Keep words matching a predicate.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # #[cfg(feature = "default_dictionary")]
    /// let mut petnames = petname::Petnames::default();
    /// # #[cfg(feature = "default_dictionary")]
    /// petnames.retain(|s| s.starts_with("b"));
    /// # #[cfg(feature = "default_dictionary")]
    /// # #[cfg(feature = "std_rng")]
    /// petnames.generate_one(2, ".");
    /// ```
    ///
    /// This is merely a convenience wrapper that applies the same predicate to
    /// the adjectives, adverbs, and names lists.
    ///
    pub fn retain<F>(&mut self, mut predicate: F)
    where
        F: FnMut(&str) -> bool,
    {
        self.adjectives.retain(|word| predicate(word));
        self.adverbs.retain(|word| predicate(word));
        self.names.retain(|word| predicate(word));
    }

    /// Calculate the cardinality of this `Petnames`.
    ///
    /// If this is low, names may be repeated by the generator with a higher
    /// frequency than your use-case may allow. If it is 0 (zero) the generator
    /// will panic (unless `words` is also zero).
    ///
    /// This can saturate. If the total possible combinations of words exceeds
    /// `u128::MAX` then this will return `u128::MAX`.
    pub fn cardinality(&self, words: u8) -> u128 {
        Lists(self, words)
            .map(|list| list.len() as u128)
            .fold1(u128::saturating_mul)
            .unwrap_or(0u128)
    }

    /// Generate a new petname.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let mut rng = rand::thread_rng();
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// petname::Petnames::default().generate(&mut rng, 7, ":");
    /// ```
    ///
    /// # Notes
    ///
    /// This may return fewer words than you request if one or more of the word
    /// lists are empty. For example, if there are no adverbs, requesting 3 or
    /// more words may still yield only "doubtful-salmon".
    ///
    pub fn generate<RNG>(&self, rng: &mut RNG, words: u8, separator: &str) -> String
    where
        RNG: rand::Rng,
    {
        itertools::Itertools::intersperse(
            Lists(self, words)
                .filter_map(|list| list.choose(rng))
                .cloned(),
            separator,
        )
        .collect::<String>()
    }

    /// Generate a single new petname.
    ///
    /// This is like `generate` but uses `rand::thread_rng` as the random
    /// source. For efficiency use `generate` when creating multiple names, or
    /// when you want to use a custom source of randomness.
    #[cfg(feature = "std_rng")]
    pub fn generate_one(&self, words: u8, separator: &str) -> String {
        self.generate(&mut rand::thread_rng(), words, separator)
    }

    /// Iterator yielding petnames.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let mut rng = rand::thread_rng();
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let petnames = petname::Petnames::default();
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let mut iter = petnames.iter(&mut rng, 4, "_");
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// println!("name: {}", iter.next().unwrap());
    /// ```
    ///
    pub fn iter<RNG>(&self, rng: &'a mut RNG, words: u8, separator: &str) -> Names<RNG>
    where
        RNG: rand::Rng,
    {
        Names {
            petnames: self,
            rng,
            words,
            separator: separator.to_string(),
        }
    }

    /// Iterator yielding unique – i.e. non-repeating – petnames.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let mut rng = rand::thread_rng();
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let petnames = petname::Petnames::default();
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let mut iter = petnames.iter_unique(&mut rng, 4, "_");
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// println!("name: {}", iter.next().unwrap());
    /// ```
    ///
    pub fn iter_unique<RNG>(
        &'a self,
        rng: &'a mut RNG,
        words: u8,
        separator: &str,
    ) -> NamesProduct<'a, core::iter::Cycle<alloc::vec::IntoIter<&'a str>>>
    where
        RNG: rand::Rng,
    {
        let lists: Vec<Words<'a>> = Lists(self, words).cloned().collect();
        NamesProduct::new(&lists, rng, separator)
    }
}

#[cfg(feature = "default_dictionary")]
impl<'a> Default for Petnames<'a> {
    fn default() -> Self {
        Self::small()
    }
}

/// Iterator over a `Petnames`' word lists.
///
/// This yields the appropriate lists from which to select a word when
/// constructing a petname of `n` words. For example, if you want 3 words in
/// your petname, this will first yield the adverbs word list, then adjectives,
/// then names.
struct Lists<'a>(&'a Petnames<'a>, u8);

impl<'a> Iterator for Lists<'a> {
    type Item = &'a Words<'a>;

    fn size_hint(&self) -> (usize, Option<usize>) {
        (0, Some(self.1 as usize))
    }

    fn next(&mut self) -> Option<Self::Item> {
        let Self(petnames, ref mut word) = self;
        match word {
            0 => None,
            1 => {
                *word -= 1;
                Some(&petnames.names)
            }
            2 => {
                *word -= 1;
                Some(&petnames.adjectives)
            }
            _ => {
                *word -= 1;
                Some(&petnames.adverbs)
            }
        }
    }
}

/// Iterator yielding petnames.
pub struct Names<'a, RNG>
where
    RNG: rand::Rng,
{
    petnames: &'a Petnames<'a>,
    rng: &'a mut RNG,
    words: u8,
    separator: String,
}

impl<'a, RNG> Names<'a, RNG>
where
    RNG: rand::Rng,
{
    /// Calculate the cardinality of this iterator; see `Petnames::cardinality`.
    #[allow(dead_code)]
    pub fn cardinality(&self) -> u128 {
        self.petnames.cardinality(self.words)
    }
}

impl<'a, RNG> Iterator for Names<'a, RNG>
where
    RNG: rand::Rng,
{
    type Item = String;

    fn next(&mut self) -> Option<Self::Item> {
        Some(
            self.petnames
                .generate(self.rng, self.words, &self.separator),
        )
    }
}

/// Iterator yielding unique petnames.
pub struct NamesProduct<'a, ITERATOR>
where
    ITERATOR: Iterator<Item = &'a str>,
{
    iters: Vec<(ITERATOR, Option<ITERATOR::Item>)>,
    separator: String,
}

impl<'a> NamesProduct<'a, core::iter::Cycle<alloc::vec::IntoIter<&'a str>>> {
    pub fn new<RNG>(lists: &[Words<'a>], rng: &'a mut RNG, separator: &str) -> Self
    where
        RNG: rand::Rng,
    {
        NamesProduct {
            iters: lists
                .iter()
                .cloned()
                .map(|mut list| {
                    list.shuffle(rng); // Could be expensive.
                    list.push(""); // Cycle marker.
                    (list.into_iter().cycle(), None)
                })
                .collect(),
            separator: separator.to_string(),
        }
    }
}

impl<'a, ITERATOR> Iterator for NamesProduct<'a, ITERATOR>
where
    ITERATOR: Iterator<Item = &'a str>,
{
    type Item = String;

    fn next(&mut self) -> Option<Self::Item> {
        let mut bump = true; // Request advance of next iterator.
        for (iter, word) in self.iters.iter_mut() {
            if bump || word.is_none() {
                match iter.next() {
                    None => {
                        // This shouldn't happen because we expect the iterators
                        // to cycle. However, if it does, we're definitely done.
                        return None;
                    }
                    Some("") => {
                        // This is the cycle end marker. We want to get another
                        // new word from this iterator, and advance the *next*
                        // iterator too.
                        match iter.next() {
                            None => return None,
                            Some("") => return None,
                            Some(s) => *word = Some(s),
                        }
                        bump = true
                    }
                    Some(s) => {
                        // We have a new word from this iterator, so we do not
                        // yet need to advance the next iterator.
                        *word = Some(s);
                        bump = false
                    }
                }
            }
        }
        if bump {
            // We reached the end of the last iterator, hence we're done.
            None
        } else {
            // We can construct a word!
            Some(
                itertools::Itertools::intersperse(
                    self.iters.iter().filter_map(|(_, w)| *w),
                    &self.separator,
                )
                .collect(),
            )
        }
    }
}
<commit_msg>Make NamesProduct::new non-shuffling, but add ::shuffled alternative<commit_after>#![no_std]
//!
//! You can populate [`Petnames`] with your own word lists, but the word lists
//! from upstream [petname](https://github.com/dustinkirkland/petname) are
//! included with the `default_dictionary` feature (enabled by default). See
//! [`Petnames::small`], [`Petnames::medium`], and [`Petnames::large`] to select
//! a particular built-in word list, or use the [`Default`] implementation.
//!
//! The other thing you need is a random number generator from [rand][]:
//!
//! ```rust
//! let mut rng = rand::thread_rng();
//! let pname = petname::Petnames::default().generate(&mut rng, 7, ":");
//! ```
//!
//! It may be more convenient to use the default random number generator:
//!
//! ```rust
//! let pname = petname::Petnames::default().generate_one(7, ":");
//! ```
//!
//! There's a [convenience function][petname] that'll do all of this:
//!
//! ```rust
//! let pname = petname::petname(7, ":");
//! ```
//!
//! But the most flexible approach is to create an [`Iterator`] with
//! [`iter`][`Petnames::iter`]:
//!
//! ```rust
//! let mut rng = rand::thread_rng();
//! let petnames = petname::Petnames::default();
//! let ten_thousand_names: Vec<String> =
//!   petnames.iter(&mut rng, 3, "_").take(10000).collect();
//! ```
//!
//! You can modify the word lists to, for example, only use words beginning with
//! the letter "b":
//!
//! ```rust
//! let mut petnames = petname::Petnames::default();
//! petnames.retain(|s| s.starts_with("b"));
//! petnames.generate_one(3, ".");
//! ```
//!

extern crate alloc;

use alloc::{
    string::{String, ToString},
    vec::Vec,
};

use itertools::Itertools;
use rand::seq::SliceRandom;

/// Convenience function to generate a new petname from default word lists.
#[allow(dead_code)]
#[cfg(feature = "std_rng")]
#[cfg(feature = "default_dictionary")]
pub fn petname(words: u8, separator: &str) -> String {
    Petnames::new().generate_one(words, separator)
}

/// A word list.
pub type Words<'a> = Vec<&'a str>;

/// Word lists and the logic to combine them into _petnames_.
///
/// A _petname_ with `n` words will contain, in order:
///
///   * `n - 2` adverbs when `n >= 2`, otherwise 0 adverbs.
///   * 1 adjective when `n >= 2`, otherwise 0 adjectives.
///   * 1 name / noun when `n >= 1`, otherwise 0 names.
///
#[derive(Clone, Debug, PartialEq)]
pub struct Petnames<'a> {
    pub adjectives: Words<'a>,
    pub adverbs: Words<'a>,
    pub names: Words<'a>,
}

impl<'a> Petnames<'a> {
    /// Constructs a new `Petnames` from the default (small) word lists.
    #[cfg(feature = "default_dictionary")]
    pub fn new() -> Self {
        Self::default()
    }

    /// Constructs a new `Petnames` from the small word lists.
    #[cfg(feature = "default_dictionary")]
    pub fn small() -> Self {
        Self::init(
            include_str!("../words/small/adjectives.txt"),
            include_str!("../words/small/adverbs.txt"),
            include_str!("../words/small/names.txt"),
        )
    }

    /// Constructs a new `Petnames` from the medium word lists.
    #[cfg(feature = "default_dictionary")]
    pub fn medium() -> Self {
        Self::init(
            include_str!("../words/medium/adjectives.txt"),
            include_str!("../words/medium/adverbs.txt"),
            include_str!("../words/medium/names.txt"),
        )
    }

    /// Constructs a new `Petnames` from the large word lists.
    #[cfg(feature = "default_dictionary")]
    pub fn large() -> Self {
        Self::init(
            include_str!("../words/large/adjectives.txt"),
            include_str!("../words/large/adverbs.txt"),
            include_str!("../words/large/names.txt"),
        )
    }

    /// Constructs a new `Petnames` from the given word lists.
    ///
    /// The words are extracted from the given strings by splitting on whitespace.
    pub fn init(adjectives: &'a str, adverbs: &'a str, names: &'a str) -> Self {
        Self {
            adjectives: adjectives.split_whitespace().collect(),
            adverbs: adverbs.split_whitespace().collect(),
            names: names.split_whitespace().collect(),
        }
    }

    /// Keep words matching a predicate.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # #[cfg(feature = "default_dictionary")]
    /// let mut petnames = petname::Petnames::default();
    /// # #[cfg(feature = "default_dictionary")]
    /// petnames.retain(|s| s.starts_with("b"));
    /// # #[cfg(feature = "default_dictionary")]
    /// # #[cfg(feature = "std_rng")]
    /// petnames.generate_one(2, ".");
    /// ```
    ///
    /// This is merely a convenience wrapper that applies the same predicate to
    /// the adjectives, adverbs, and names lists.
    ///
    pub fn retain<F>(&mut self, mut predicate: F)
    where
        F: FnMut(&str) -> bool,
    {
        self.adjectives.retain(|word| predicate(word));
        self.adverbs.retain(|word| predicate(word));
        self.names.retain(|word| predicate(word));
    }

    /// Calculate the cardinality of this `Petnames`.
    ///
    /// If this is low, names may be repeated by the generator with a higher
    /// frequency than your use-case may allow. If it is 0 (zero) the generator
    /// will panic (unless `words` is also zero).
    ///
    /// This can saturate. If the total possible combinations of words exceeds
    /// `u128::MAX` then this will return `u128::MAX`.
    pub fn cardinality(&self, words: u8) -> u128 {
        Lists(self, words)
            .map(|list| list.len() as u128)
            .fold1(u128::saturating_mul)
            .unwrap_or(0u128)
    }

    /// Generate a new petname.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let mut rng = rand::thread_rng();
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// petname::Petnames::default().generate(&mut rng, 7, ":");
    /// ```
    ///
    /// # Notes
    ///
    /// This may return fewer words than you request if one or more of the word
    /// lists are empty. For example, if there are no adverbs, requesting 3 or
    /// more words may still yield only "doubtful-salmon".
    ///
    pub fn generate<RNG>(&self, rng: &mut RNG, words: u8, separator: &str) -> String
    where
        RNG: rand::Rng,
    {
        itertools::Itertools::intersperse(
            Lists(self, words)
                .filter_map(|list| list.choose(rng))
                .cloned(),
            separator,
        )
        .collect::<String>()
    }

    /// Generate a single new petname.
    ///
    /// This is like `generate` but uses `rand::thread_rng` as the random
    /// source. For efficiency use `generate` when creating multiple names, or
    /// when you want to use a custom source of randomness.
    #[cfg(feature = "std_rng")]
    pub fn generate_one(&self, words: u8, separator: &str) -> String {
        self.generate(&mut rand::thread_rng(), words, separator)
    }

    /// Iterator yielding petnames.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let mut rng = rand::thread_rng();
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let petnames = petname::Petnames::default();
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let mut iter = petnames.iter(&mut rng, 4, "_");
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// println!("name: {}", iter.next().unwrap());
    /// ```
    ///
    pub fn iter<RNG>(&self, rng: &'a mut RNG, words: u8, separator: &str) -> Names<RNG>
    where
        RNG: rand::Rng,
    {
        Names {
            petnames: self,
            rng,
            words,
            separator: separator.to_string(),
        }
    }

    /// Iterator yielding unique – i.e. non-repeating – petnames.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let mut rng = rand::thread_rng();
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let petnames = petname::Petnames::default();
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// let mut iter = petnames.iter_unique(&mut rng, 4, "_");
    /// # #[cfg(all(feature = "std_rng", feature = "default_dictionary"))]
    /// println!("name: {}", iter.next().unwrap());
    /// ```
    ///
    pub fn iter_unique<RNG>(
        &'a self,
        rng: &'a mut RNG,
        words: u8,
        separator: &str,
    ) -> NamesProduct<'a, core::iter::Cycle<alloc::vec::IntoIter<&'a str>>>
    where
        RNG: rand::Rng,
    {
        let lists: Vec<Words<'a>> = Lists(self, words).cloned().collect();
        NamesProduct::shuffled(&lists, rng, separator)
    }
}

#[cfg(feature = "default_dictionary")]
impl<'a> Default for Petnames<'a> {
    fn default() -> Self {
        Self::small()
    }
}

/// Iterator over a `Petnames`' word lists.
///
/// This yields the appropriate lists from which to select a word when
/// constructing a petname of `n` words. For example, if you want 3 words in
/// your petname, this will first yield the adverbs word list, then adjectives,
/// then names.
struct Lists<'a>(&'a Petnames<'a>, u8);

impl<'a> Iterator for Lists<'a> {
    type Item = &'a Words<'a>;

    fn size_hint(&self) -> (usize, Option<usize>) {
        (0, Some(self.1 as usize))
    }

    fn next(&mut self) -> Option<Self::Item> {
        let Self(petnames, ref mut word) = self;
        match word {
            0 => None,
            1 => {
                *word -= 1;
                Some(&petnames.names)
            }
            2 => {
                *word -= 1;
                Some(&petnames.adjectives)
            }
            _ => {
                *word -= 1;
                Some(&petnames.adverbs)
            }
        }
    }
}

/// Iterator yielding petnames.
pub struct Names<'a, RNG>
where
    RNG: rand::Rng,
{
    petnames: &'a Petnames<'a>,
    rng: &'a mut RNG,
    words: u8,
    separator: String,
}

impl<'a, RNG> Names<'a, RNG>
where
    RNG: rand::Rng,
{
    /// Calculate the cardinality of this iterator; see `Petnames::cardinality`.
    #[allow(dead_code)]
    pub fn cardinality(&self) -> u128 {
        self.petnames.cardinality(self.words)
    }
}

impl<'a, RNG> Iterator for Names<'a, RNG>
where
    RNG: rand::Rng,
{
    type Item = String;

    fn next(&mut self) -> Option<Self::Item> {
        Some(
            self.petnames
                .generate(self.rng, self.words, &self.separator),
        )
    }
}

/// Iterator yielding petnames from the product of given word lists.
///
/// This can be used to ensure that only unique names are produced.
pub struct NamesProduct<'a, ITERATOR>
where
    ITERATOR: Iterator<Item = &'a str>,
{
    iters: Vec<(ITERATOR, Option<ITERATOR::Item>)>,
    separator: String,
}

impl<'a> NamesProduct<'a, core::iter::Cycle<alloc::vec::IntoIter<&'a str>>> {
    /// Cycles through the product of the `lists`, joining with `separator`. The
    /// leftmost list will cycle most rapidly.
    pub fn new(lists: &[Words<'a>], separator: &str) -> Self {
        NamesProduct {
            iters: lists
                .iter()
                .cloned()
                .map(|mut list| {
                    list.push(""); // Cycle marker.
                    (list.into_iter().cycle(), None)
                })
                .collect(),
            separator: separator.to_string(),
        }
    }

    /// Shuffles each of the given `lists` with `rng`, then cycles through the
    /// product of the lists, joining with `separator`. The leftmost list will
    /// cycle most rapidly.
    pub fn shuffled<RNG>(lists: &[Words<'a>], rng: &'a mut RNG, separator: &str) -> Self
    where
        RNG: rand::Rng,
    {
        NamesProduct {
            iters: lists
                .iter()
                .cloned()
                .map(|mut list| {
                    list.shuffle(rng); // Could be expensive.
                    list.push(""); // Cycle marker.
                    (list.into_iter().cycle(), None)
                })
                .collect(),
            separator: separator.to_string(),
        }
    }
}

impl<'a, ITERATOR> Iterator for NamesProduct<'a, ITERATOR>
where
    ITERATOR: Iterator<Item = &'a str>,
{
    type Item = String;

    fn next(&mut self) -> Option<Self::Item> {
        let mut bump = true; // Request advance of next iterator.
        for (iter, word) in self.iters.iter_mut() {
            if bump || word.is_none() {
                match iter.next() {
                    None => {
                        // This shouldn't happen because we expect the iterators
                        // to cycle. However, if it does, we're definitely done.
                        return None;
                    }
                    Some("") => {
                        // This is the cycle end marker. We want to get another
                        // new word from this iterator, and advance the *next*
                        // iterator too.
                        match iter.next() {
                            None => return None,
                            Some("") => return None,
                            Some(s) => *word = Some(s),
                        }
                        bump = true
                    }
                    Some(s) => {
                        // We have a new word from this iterator, so we do not
                        // yet need to advance the next iterator.
                        *word = Some(s);
                        bump = false
                    }
                }
            }
        }
        if bump {
            // We reached the end of the last iterator, hence we're done.
            None
        } else {
            // We can construct a word!
            Some(
                itertools::Itertools::intersperse(
                    self.iters.iter().filter_map(|(_, w)| *w),
                    &self.separator,
                )
                .collect(),
            )
        }
    }
}
<|endoftext|>


<commit_before>//! Jamal provides a bi-directional interface for transformations between json and yaml documents.
extern crate serde_json;
extern crate serde_yaml;

use serde_json::Value as JsonValue;
use serde_yaml::Value as YamlValue;
use std::error::Error as StdError;
use std::num::ParseFloatError;
use std::collections::BTreeMap;

/// represents potential errors that can happen during the transformation process
#[derive(Debug)]
pub enum Error {
    /// Occurs when an attempt to parse a `serde_yaml::Real` as an `f64` fails
    ParseFloat(ParseFloatError),
    /// Occurs when a transformation is not supported
    InvalidValue,
}

impl StdError for Error {
    fn description(&self) -> &str {
        match *self {
            Error::ParseFloat(ref e) => e.description(),
            Error::InvalidValue => "invalid value"
        }
    }

    fn cause(&self) -> Option<&StdError> {
        match *self {
            Error::ParseFloat(ref e) => Some(e),
            _ => None,
        }
    }
}

impl ::std::fmt::Display for Error {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "{:?}", self)
    }
}

/// A result type with a fixed type for `jamal::Errors`
pub type Result<T> = std::result::Result<T, Error>;

impl From<ParseFloatError> for Error {
    fn from(error: ParseFloatError) -> Error {
        Error::ParseFloat(error)
    }
}

/// converts a `serde_json::Value` into a `serde_yaml::Value`
pub fn to_yaml(json: &JsonValue) -> Result<YamlValue> {
    match json {
        &JsonValue::Null => Ok(YamlValue::Null),
        &JsonValue::Bool(ref value) => Ok(YamlValue::Boolean(value.clone())),
        &JsonValue::I64(ref value) => Ok(YamlValue::Integer(value.clone())),
        &JsonValue::U64(ref value) => Ok(YamlValue::Integer(value.clone() as i64)),
        &JsonValue::F64(ref value) => Ok(YamlValue::Real(value.to_string())),
        &JsonValue::String(ref value) => Ok(YamlValue::String(value.clone())),
        &JsonValue::Array(ref values) => {
            let mut yaml_values = vec![];
            for value in values {
                yaml_values.push(try!(to_yaml(&value)));
            }
            Ok(YamlValue::Array(yaml_values))
        }
        &JsonValue::Object(ref value) => {
            let mut yaml_value = BTreeMap::new();
            for (k, v) in value {
                let yaml = try!(to_yaml(&v));
                yaml_value.insert(YamlValue::String(k.clone()), yaml);
            }
            Ok(YamlValue::Hash(yaml_value))
        }
    }
}

/// converts a `serde_yaml::Value` into a `serde_json::Value`
pub fn to_json(yaml: &YamlValue) -> Result<JsonValue> {
    match yaml {
        &YamlValue::Real(ref value) => Ok(JsonValue::F64(try!(value.parse::<f64>()))),
        &YamlValue::Integer(ref value) => Ok(JsonValue::I64(value.clone())),
        &YamlValue::String(ref value) => Ok(JsonValue::String(value.clone())),
        &YamlValue::Boolean(ref value) => Ok(JsonValue::Bool(value.clone())),
        &YamlValue::Array(ref values) => {
            let mut json_values = vec![];
            for value in values {
                let json = try!(to_json(&value));
                json_values.push(json);
            }
            Ok(JsonValue::Array(json_values))
        }
        &YamlValue::Hash(ref value) => {
            let mut json_value = BTreeMap::new();
            for (k, v) in value {
                match k {
                    &YamlValue::String(ref key) => {
                        let json = try!(to_json(&v));
                        json_value.insert(key.clone(), json);
                    }
                    _ => return Err(Error::InvalidValue),
                }
            }
            Ok(JsonValue::Object(json_value))
        }
        &YamlValue::Alias(_) => Err(Error::InvalidValue), // not supported yet
        &YamlValue::Null => Ok(JsonValue::Null),
        &YamlValue::BadValue => Err(Error::InvalidValue),
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;
    use super::serde_json;
    use super::serde_yaml;
    use super::{to_json, to_yaml};

    #[test]
    fn json_to_yaml() {
        let input = serde_json::to_value(&vec!["foo"]);
        let output = serde_yaml::Value::Array(vec![serde_yaml::Value::String("foo".to_owned())]);
        assert_eq!(to_yaml(&input).unwrap(), output);
    }

    #[test]
    fn yaml_to_json() {
        let input = serde_yaml::to_value(&vec!["foo"]);
        let output = serde_json::Value::Array(vec![serde_json::Value::String("foo".to_owned())]);
        assert_eq!(to_json(&input).unwrap(), output);
    }
}
<commit_msg>document errors<commit_after>//! Jamal provides a bi-directional interface for transformations between json and yaml documents.
extern crate serde_json;
extern crate serde_yaml;

use serde_json::Value as JsonValue;
use serde_yaml::Value as YamlValue;
use std::error::Error as StdError;
use std::num::ParseFloatError;
use std::collections::BTreeMap;

/// represents potential errors that can happen during the transformation process
#[derive(Debug)]
pub enum Error {
    /// Occurs when an attempt to parse a `serde_yaml::Real` as an `f64` fails
    ParseFloat(ParseFloatError),
    /// Occurs when a transformation is requested for an invalid value
    InvalidValue,
    /// Occurs when an unsupported transformation is requested
    UnsupportedValue
}

impl StdError for Error {
    fn description(&self) -> &str {
        match *self {
            Error::ParseFloat(ref e) => e.description(),
            Error::InvalidValue => "invalid value",
            Error::UnsupportedValue => "supported value",
        }
    }

    fn cause(&self) -> Option<&StdError> {
        match *self {
            Error::ParseFloat(ref e) => Some(e),
            _ => None,
        }
    }
}

impl ::std::fmt::Display for Error {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "{:?}", self)
    }
}

/// A result type with a fixed type for `jamal::Errors`
pub type Result<T> = std::result::Result<T, Error>;

impl From<ParseFloatError> for Error {
    fn from(error: ParseFloatError) -> Error {
        Error::ParseFloat(error)
    }
}

/// converts a `serde_json::Value` into a `serde_yaml::Value`
pub fn to_yaml(json: &JsonValue) -> Result<YamlValue> {
    match json {
        &JsonValue::Null => Ok(YamlValue::Null),
        &JsonValue::Bool(ref value) => Ok(YamlValue::Boolean(value.clone())),
        &JsonValue::I64(ref value) => Ok(YamlValue::Integer(value.clone())),
        &JsonValue::U64(ref value) => Ok(YamlValue::Integer(value.clone() as i64)),
        &JsonValue::F64(ref value) => Ok(YamlValue::Real(value.to_string())),
        &JsonValue::String(ref value) => Ok(YamlValue::String(value.clone())),
        &JsonValue::Array(ref values) => {
            let mut yaml_values = vec![];
            for value in values {
                yaml_values.push(try!(to_yaml(&value)));
            }
            Ok(YamlValue::Array(yaml_values))
        }
        &JsonValue::Object(ref value) => {
            let mut yaml_value = BTreeMap::new();
            for (k, v) in value {
                let yaml = try!(to_yaml(&v));
                yaml_value.insert(YamlValue::String(k.clone()), yaml);
            }
            Ok(YamlValue::Hash(yaml_value))
        }
    }
}

/// converts a `serde_yaml::Value` into a `serde_json::Value`
pub fn to_json(yaml: &YamlValue) -> Result<JsonValue> {
    match yaml {
        &YamlValue::Real(ref value) => Ok(JsonValue::F64(try!(value.parse::<f64>()))),
        &YamlValue::Integer(ref value) => Ok(JsonValue::I64(value.clone())),
        &YamlValue::String(ref value) => Ok(JsonValue::String(value.clone())),
        &YamlValue::Boolean(ref value) => Ok(JsonValue::Bool(value.clone())),
        &YamlValue::Array(ref values) => {
            let mut json_values = vec![];
            for value in values {
                let json = try!(to_json(&value));
                json_values.push(json);
            }
            Ok(JsonValue::Array(json_values))
        }
        &YamlValue::Hash(ref value) => {
            let mut json_value = BTreeMap::new();
            for (k, v) in value {
                match k {
                    &YamlValue::String(ref key) => {
                        let json = try!(to_json(&v));
                        json_value.insert(key.clone(), json);
                    }
                    _ => return Err(Error::InvalidValue),
                }
            }
            Ok(JsonValue::Object(json_value))
        }
        &YamlValue::Alias(_) => Err(Error::InvalidValue), // not supported yet
        &YamlValue::Null => Ok(JsonValue::Null),
        &YamlValue::BadValue => Err(Error::InvalidValue),
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;
    use super::serde_json;
    use super::serde_yaml;
    use super::{to_json, to_yaml};

    #[test]
    fn json_to_yaml() {
        let input = serde_json::to_value(&vec!["foo"]);
        let output = serde_yaml::Value::Array(vec![serde_yaml::Value::String("foo".to_owned())]);
        assert_eq!(to_yaml(&input).unwrap(), output);
    }

    #[test]
    fn yaml_to_json() {
        let input = serde_yaml::to_value(&vec!["foo"]);
        let output = serde_json::Value::Array(vec![serde_json::Value::String("foo".to_owned())]);
        assert_eq!(to_json(&input).unwrap(), output);
    }
}
<|endoftext|>