<commit_before>use std::cmp::Ordering;
use std::{mem};

use libc::{size_t, c_int, c_uint};

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

#[derive(Debug)]
struct Node {
    weight: size_t,
    leaf_counts: Vec<c_int>,
}

impl Node {
    pub fn new(weight: size_t, initial_count: c_int, capacity: usize) -> Node {
        let mut n = Node {
            weight: weight,
            leaf_counts: Vec::with_capacity(capacity),
        };
        n.leaf_counts.push(initial_count);
        n
    }
}

#[derive(Debug)]
struct Leaf {
    pub weight: size_t,
    pub index: size_t,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug)]
struct List {
    lookahead1: Node,
    lookahead2: Node,
    next_leaf_index: size_t,
}

/// Calculates the bitlengths for the Huffman tree, based on the counts of each
/// symbol.
pub fn length_limited_code_lengths(frequencies: &[size_t], maxbits: c_int) -> Vec<c_uint> {
    let mut leaves = vec![];

    // Count used symbols and place them in the leaves.
    for (i, &freq) in frequencies.iter().enumerate() {
        if freq != 0 {
            leaves.push(Leaf { weight: freq, index: i });
        }
    }

    // Short circuit some special cases
    if leaves.is_empty() {
        // There are no non-zero frequencies.
        return vec![0; frequencies.len()];
    }
    if leaves.len() == 1 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        return result;
    }
    if leaves.len() == 2 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        result[leaves[1].index] = 1;
        return result;
    }

    // Sort the leaves from least frequent to most frequent.
    // Add index into the same variable for stable sorting.
    for leaf in &mut leaves {
        leaf.weight = (leaf.weight << 9) | leaf.index;
    }
    leaves.sort();
    for leaf in &mut leaves {
        leaf.weight >>= 9;
    }

    let max_num_leaves = 2 * leaves.len() - 2;

    let mut lists = Vec::with_capacity(maxbits as usize);
    for _ in 0..maxbits {
        lists.push(List {
            lookahead1: Node::new(leaves[0].weight, 1, max_num_leaves),
            lookahead2: Node::new(leaves[1].weight, 2, max_num_leaves),
            next_leaf_index: 2,
        });
    }

    // In the last list, 2 * numsymbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = max_num_leaves - 2;
    for _ in 0..num_boundary_pm_runs {
        boundary_pm_toplevel(&mut lists[..], &leaves);
    }

    let n = frequencies.len();
    let mut result = vec![0; n];

    let mut a = lists.pop().unwrap().lookahead2.leaf_counts.into_iter().rev().peekable();

    let mut bitlength_value = 1;
    while let Some(leaf_count) = a.next() {
        let next_count = *a.peek().unwrap_or(&0);
        for i in next_count..leaf_count {
            result[leaves[i as usize].index] = bitlength_value;
        }
        bitlength_value += 1;
    }

    result
}

fn lowest_list(lists: &mut [List], leaves: &[Leaf]) {
    // We're in the lowest list, just add another leaf to the lookaheads
    // There will always be more leaves to be added on level 0 so this is safe.
    let mut current_list = lists.get_mut(0).unwrap();
    let next_leaf = &leaves[current_list.next_leaf_index];
    current_list.lookahead2.weight = next_leaf.weight;

    current_list.lookahead2.leaf_counts[0] = current_list.lookahead1.leaf_counts.last().unwrap() + 1;
    current_list.next_leaf_index += 1;
}

fn next_leaf(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
    let mut current_list = lists.get_mut(current_list_index).unwrap();

    // The next leaf goes next; counting itself makes the leaf_count increase by one.
    current_list.lookahead2.weight = leaves[current_list.next_leaf_index].weight;
    current_list.lookahead2.leaf_counts.clear();
    current_list.lookahead2.leaf_counts.extend_from_slice(&current_list.lookahead1.leaf_counts);
    let last_index = current_list.lookahead2.leaf_counts.len() - 1;
    current_list.lookahead2.leaf_counts[last_index] += 1;
    current_list.next_leaf_index += 1;
}

fn next_tree(weight_sum: size_t, lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
    let num_leaf_counts = lists[current_list_index - 1].lookahead2.leaf_counts.len();
    let previous_list_leaf_counts = lists[current_list_index - 1].lookahead2.leaf_counts.as_ptr();
    {
        let mut current_list = &mut lists[current_list_index];

        // Make a tree from the lookaheads from the previous list; that goes next.
        // This is not a leaf node, so the leaf count stays the same.
        current_list.lookahead2.weight = weight_sum;
        current_list.lookahead2.leaf_counts.clear();
        for i in 0..num_leaf_counts {
            current_list.lookahead2.leaf_counts.push(unsafe {
                *previous_list_leaf_counts.offset(i as isize)
            });
        }
        current_list.lookahead2.leaf_counts.push(*current_list.lookahead1.leaf_counts.last().unwrap());
    }

    // The previous list needs two new lookahead nodes.
    boundary_pm(lists, leaves, current_list_index - 1);
    boundary_pm(lists, leaves, current_list_index - 1);
}

fn boundary_pm_toplevel(lists: &mut [List], leaves: &[Leaf]) {
    let last_index = lists.len() - 1;
    boundary_pm(lists, leaves, last_index);
}

fn boundary_pm(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
    let next_leaf_index = lists[current_list_index].next_leaf_index;

    if current_list_index == 0 && next_leaf_index == leaves.len() {
        // We've added all the leaves to the lowest list, so we're done here
        return;
    }

    mem::swap(&mut lists[current_list_index].lookahead1, &mut lists[current_list_index].lookahead2);

    if current_list_index == 0 {
        lowest_list(lists, leaves);
    } else {
        // We're at a list other than the lowest list.
        let weight_sum = {
            let previous_list = lists.get(current_list_index - 1).unwrap();
            previous_list.lookahead1.weight + previous_list.lookahead2.weight
        };

        if next_leaf_index < leaves.len() && weight_sum > leaves[next_leaf_index].weight {
            next_leaf(lists, leaves, current_list_index);
        } else {
            next_tree(weight_sum, lists, leaves, current_list_index);
        }
    }
}

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<commit_msg>Remove unsafe indexing into a vector<commit_after>use std::cmp::Ordering;
use std::{mem};

use libc::{size_t, c_int, c_uint};

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

#[derive(Debug)]
struct Node {
    weight: size_t,
    leaf_counts: Vec<c_int>,
}

impl Node {
    pub fn new(weight: size_t, initial_count: c_int, capacity: usize) -> Node {
        let mut n = Node {
            weight: weight,
            leaf_counts: Vec::with_capacity(capacity),
        };
        n.leaf_counts.push(initial_count);
        n
    }
}

#[derive(Debug)]
struct Leaf {
    pub weight: size_t,
    pub index: size_t,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug)]
struct List {
    lookahead1: Node,
    lookahead2: Node,
    next_leaf_index: size_t,
}

/// Calculates the bitlengths for the Huffman tree, based on the counts of each
/// symbol.
pub fn length_limited_code_lengths(frequencies: &[size_t], maxbits: c_int) -> Vec<c_uint> {
    let mut leaves = vec![];

    // Count used symbols and place them in the leaves.
    for (i, &freq) in frequencies.iter().enumerate() {
        if freq != 0 {
            leaves.push(Leaf { weight: freq, index: i });
        }
    }

    // Short circuit some special cases
    if leaves.is_empty() {
        // There are no non-zero frequencies.
        return vec![0; frequencies.len()];
    }
    if leaves.len() == 1 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        return result;
    }
    if leaves.len() == 2 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        result[leaves[1].index] = 1;
        return result;
    }

    // Sort the leaves from least frequent to most frequent.
    // Add index into the same variable for stable sorting.
    for leaf in &mut leaves {
        leaf.weight = (leaf.weight << 9) | leaf.index;
    }
    leaves.sort();
    for leaf in &mut leaves {
        leaf.weight >>= 9;
    }

    let max_num_leaves = 2 * leaves.len() - 2;

    let mut lists = Vec::with_capacity(maxbits as usize);
    for _ in 0..maxbits {
        lists.push(List {
            lookahead1: Node::new(leaves[0].weight, 1, max_num_leaves),
            lookahead2: Node::new(leaves[1].weight, 2, max_num_leaves),
            next_leaf_index: 2,
        });
    }

    // In the last list, 2 * numsymbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = max_num_leaves - 2;
    for _ in 0..num_boundary_pm_runs {
        boundary_pm_toplevel(&mut lists[..], &leaves);
    }

    let n = frequencies.len();
    let mut result = vec![0; n];

    let mut a = lists.pop().unwrap().lookahead2.leaf_counts.into_iter().rev().peekable();

    let mut bitlength_value = 1;
    while let Some(leaf_count) = a.next() {
        let next_count = *a.peek().unwrap_or(&0);
        for i in next_count..leaf_count {
            result[leaves[i as usize].index] = bitlength_value;
        }
        bitlength_value += 1;
    }

    result
}

fn lowest_list(lists: &mut [List], leaves: &[Leaf]) {
    // We're in the lowest list, just add another leaf to the lookaheads
    // There will always be more leaves to be added on level 0 so this is safe.
    let mut current_list = lists.get_mut(0).unwrap();
    let next_leaf = &leaves[current_list.next_leaf_index];
    current_list.lookahead2.weight = next_leaf.weight;

    current_list.lookahead2.leaf_counts[0] = current_list.lookahead1.leaf_counts.last().unwrap() + 1;
    current_list.next_leaf_index += 1;
}

fn next_leaf(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
    let mut current_list = lists.get_mut(current_list_index).unwrap();

    // The next leaf goes next; counting itself makes the leaf_count increase by one.
    current_list.lookahead2.weight = leaves[current_list.next_leaf_index].weight;
    current_list.lookahead2.leaf_counts.clear();
    current_list.lookahead2.leaf_counts.extend_from_slice(&current_list.lookahead1.leaf_counts);
    let last_index = current_list.lookahead2.leaf_counts.len() - 1;
    current_list.lookahead2.leaf_counts[last_index] += 1;
    current_list.next_leaf_index += 1;
}

fn next_tree(weight_sum: size_t, lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
    {
        let (head, tail) = lists.split_at_mut(current_list_index);
        let prev_list = head.last_mut().unwrap();
        let current_list = tail.first_mut().unwrap();

        let previous_list_leaf_counts = &prev_list.lookahead2.leaf_counts;

        // Make a tree from the lookaheads from the previous list; that goes next.
        // This is not a leaf node, so the leaf count stays the same.
        current_list.lookahead2.weight = weight_sum;
        current_list.lookahead2.leaf_counts.clear();

        current_list.lookahead2.leaf_counts.extend(previous_list_leaf_counts);
        current_list.lookahead2.leaf_counts.push(*current_list.lookahead1.leaf_counts.last().unwrap());
    }

    // The previous list needs two new lookahead nodes.
    boundary_pm(lists, leaves, current_list_index - 1);
    boundary_pm(lists, leaves, current_list_index - 1);
}

fn boundary_pm_toplevel(lists: &mut [List], leaves: &[Leaf]) {
    let last_index = lists.len() - 1;
    boundary_pm(lists, leaves, last_index);
}

fn boundary_pm(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
    let next_leaf_index = lists[current_list_index].next_leaf_index;

    if current_list_index == 0 && next_leaf_index == leaves.len() {
        // We've added all the leaves to the lowest list, so we're done here
        return;
    }

    mem::swap(&mut lists[current_list_index].lookahead1, &mut lists[current_list_index].lookahead2);

    if current_list_index == 0 {
        lowest_list(lists, leaves);
    } else {
        // We're at a list other than the lowest list.
        let weight_sum = {
            let previous_list = lists.get(current_list_index - 1).unwrap();
            previous_list.lookahead1.weight + previous_list.lookahead2.weight
        };

        if next_leaf_index < leaves.len() && weight_sum > leaves[next_leaf_index].weight {
            next_leaf(lists, leaves, current_list_index);
        } else {
            next_tree(weight_sum, lists, leaves, current_list_index);
        }
    }
}

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<|endoftext|>


<commit_before>// Lumol, an extensible molecular simulation engine
// Copyright (C) 2015-2016 G. Fraux — BSD license

//! Caching energy components for speeding up the energy computation in
//! Monte-Carlo simulations.
//!
//! In most of Monte-Carlo moves, only a very small subset of the system changes.
//! We can use that property to remove the need of recomputing most of the
//! energy components, by storing them and providing update callbacks.
use std::mem;

use super::System;
use types::{Vector3D, Array2};

/// Callback for updating a cache. It also take an `&mut System` argument for
/// updating the cache inside the global potentials.
type UpdateCallback = Box<Fn(&mut EnergyCache, &mut System)>;

/// This is a cache for energy computation.
///
/// Cache integrity is left up to the user of this structure: any function with
/// `update_` prefix must be called as needed to ensure cache consistency.
pub struct EnergyCache {
    /// 2-D array containing the pairs interactions between particles `i` and
    /// `j` at index `i, j`
    pairs_cache: Array2<f64>,
    /// Energy of all the pairs in the system
    pairs: f64,
    /// Energy of all the bonds in the system
    bonds: f64,
    /// Energy of all the angles in the system
    angles: f64,
    /// Energy of all the dihedrals angles in the system
    dihedrals: f64,
    /// Energy of coulombic interactions
    coulomb: f64,
    /// Energy of global interactions
    global: f64,
    /// Callback to be called to update the cache if the system is modified
    updater: Option<UpdateCallback>
}

impl Default for EnergyCache {
    fn default() -> EnergyCache {
        EnergyCache::new()
    }
}

impl EnergyCache {
    /// Create a new empty energy cache.
    pub fn new() -> EnergyCache {
        EnergyCache {
            pairs_cache: Array2::zeros((0, 0)),
            pairs: 0.0,
            bonds: 0.0,
            angles: 0.0,
            dihedrals: 0.0,
            coulomb: 0.0,
            global: 0.0,
            updater: None,
        }
    }

    /// Clear all values in the cache by setting them to 0
    fn clear(&mut self) {
        self.pairs_cache.assign(0.0);
        self.pairs = 0.0;
        self.bonds = 0.0;
        self.angles = 0.0;
        self.dihedrals = 0.0;
        self.coulomb = 0.0;
        self.global = 0.0;
    }

    /// Initialize the cache to be used with `system`. After a call to this
    /// function, the cache is only usable with the same system. To change
    /// the associated system, one must call this function again.
    pub fn init(&mut self, system: &System) {
        self.clear();
        self.pairs_cache.resize_if_different((system.size(), system.size()));

        let evaluator = system.energy_evaluator();

        for i in 0..system.size() {
            for j in (i + 1)..system.size() {
                let r = system.nearest_image(i, j).norm();
                let energy = evaluator.pair(r, i, j);
                self.pairs_cache[(i, j)] = energy;
                self.pairs_cache[(j, i)] = energy;
                self.pairs += energy;
            }
        }

        self.bonds = evaluator.bonds();
        self.angles = evaluator.angles();
        self.dihedrals = evaluator.dihedrals();
        self.coulomb = evaluator.coulomb();
        self.global = evaluator.global();
    }

    /// Get the cached energy
    pub fn energy(&self) -> f64 {
        let mut energy = 0.0;
        energy += self.pairs;

        energy += self.bonds;
        energy += self.angles;
        energy += self.dihedrals;

        energy += self.coulomb;
        energy += self.global;
        return energy;
    }

    /// Update the cache after a call to a `EnergyCache::*_cost` function or
    /// `EnergyCache::unused`.
    pub fn update(&mut self, system: &mut System) {
        let updater = mem::replace(&mut self.updater, None);
        if let Some(updater) = updater {
            updater(self, system);
        } else {
            fatal_error!("Error in `EnergyCache::update`: \
                         This function MUST be called after a call to a `*_cost` function");
        }
    }

    /// This function should be called whenever the cache is not used, but one
    /// still want it to be updated. Future call to `EnergyCache::update` will
    /// recompute the full cache.
    pub fn unused(&mut self) {
        self.updater = Some(Box::new(|cache, system| {
            cache.init(system);
        }))
    }
}

impl EnergyCache {
    /// Get the cost of moving the set of particles with indexes in `idxes` to
    /// `newpos`. This function DO NOT update the cache, the
    /// `update_particles_moved` function MUST be called if the particles are
    /// effectively moved.
    pub fn move_particles_cost(&mut self, system: &System, idxes: Vec<usize>, newpos: &[Vector3D]) -> f64 {
        let evaluator = system.energy_evaluator();

        // First, go for pair interactions
        let mut new_pairs = Array2::<f64>::zeros((system.size(), system.size()));
        let mut pairs_delta = 0.0;
        // Interactions with the sub-system not being moved
        for (i, &part_i) in idxes.iter().enumerate() {
            for part_j in 0..system.size() {
                // Exclude interactions inside the sub-system.
                if idxes.contains(&part_j) {continue}

                let r = system.cell().distance(&system[part_j].position, &newpos[i]);
                let energy = evaluator.pair(r, part_i, part_j);

                pairs_delta += energy;
                new_pairs[(part_i, part_j)] += energy;
                new_pairs[(part_j, part_i)] += energy;

                pairs_delta -= self.pairs_cache[(part_i, part_j)];
            }
        }

        // Interactions withing the sub-system being moved
        for (i, &part_i) in idxes.iter().enumerate() {
            for (j, &part_j) in idxes.iter().enumerate().skip(i + 1) {
                let r = system.cell().distance(&newpos[i], &newpos[j]);
                let energy = evaluator.pair(r, part_i, part_j);

                pairs_delta += energy;
                new_pairs[(part_i, part_j)] += energy;
                new_pairs[(part_j, part_i)] += energy;

                pairs_delta -= self.pairs_cache[(part_i, part_j)];
            }
        }

        // Recompute everything here. One day we could store some terms and do
        // the same thing than for pair interactions.
        let mut bonds = 0.0;
        let mut angles = 0.0;
        let mut dihedrals = 0.0;
        for molecule in system.molecules() {
            for bond in molecule.bonds() {
                let (i, j) = (bond.i(), bond.j());
                let ri = new_position(system, i, &idxes, newpos);
                let rj = new_position(system, j, &idxes, newpos);
                let r = system.cell().distance(ri, rj);
                bonds += evaluator.bond(r, i, j);
            }

            for angle in molecule.angles() {
                let (i, j, k) = (angle.i(), angle.j(), angle.k());
                let ri = new_position(system, i, &idxes, newpos);
                let rj = new_position(system, j, &idxes, newpos);
                let rk = new_position(system, k, &idxes, newpos);
                let theta = system.cell().angle(ri, rj, rk);
                angles += evaluator.angle(theta, i, j, k);
            }

            for dihedral in molecule.dihedrals() {
                let (i, j, k, m) = (dihedral.i(), dihedral.j(), dihedral.k(), dihedral.m());
                let ri = new_position(system, i, &idxes, newpos);
                let rj = new_position(system, j, &idxes, newpos);
                let rk = new_position(system, k, &idxes, newpos);
                let rm = new_position(system, m, &idxes, newpos);
                let phi = system.cell().dihedral(ri, rj, rk, rm);
                dihedrals += evaluator.dihedral(phi, i, j, k, m);
            }
        }

        let coulomb_delta = if let Some(coulomb) = system.interactions().coulomb() {
            coulomb.borrow_mut().move_particles_cost(system, &idxes, newpos)
        } else {
            0.0
        };

        let mut global_delta = 0.0;
        for potential in system.interactions().globals() {
            global_delta += potential.borrow_mut().move_particles_cost(system, &idxes, newpos);
        }

        let cost = pairs_delta + (bonds - self.bonds)
                               + (angles - self.angles)
                               + (dihedrals - self.dihedrals)
                               + coulomb_delta + global_delta;

        self.updater = Some(Box::new(move |cache, system| {
            cache.bonds = bonds;
            cache.angles = angles;
            cache.dihedrals = dihedrals;

            cache.pairs += pairs_delta;
            cache.coulomb += coulomb_delta;
            cache.global += global_delta;

            let (n, m) = new_pairs.shape();
            debug_assert!(n == m);
            debug_assert!((n, m) == cache.pairs_cache.shape());
            // only loop over the indices that actually changed
            for &i in idxes.iter() {
                for j in 0..n {
                    if idxes.contains(&j) {continue}
                    cache.pairs_cache[(i, j)] = new_pairs[(i, j)];
                    cache.pairs_cache[(j, i)] = new_pairs[(i, j)];
                }
                for &k in idxes.iter().skip(i + 1) {
                    cache.pairs_cache[(i, k)] = new_pairs[(i, k)];
                    cache.pairs_cache[(k, i)] = new_pairs[(i, k)];
                }
            }

            // Update the cache for the global potentials
            if let Some(coulomb) = system.interactions().coulomb() {
                coulomb.borrow_mut().update();
            }
            for potential in system.interactions().globals() {
                potential.borrow_mut().update();
            }
        }));
        return cost;
    }
}

/// Return either the new position of a particle (from `newpos`) if its index
/// is in `idxes`, or its old position in the system.
fn new_position<'a>(system: &'a System, i: usize, idxes: &[usize], newpos: &'a[Vector3D]) -> &'a Vector3D {
    for (idx, &particle) in idxes.iter().enumerate() {
        if particle == i {
            return &newpos[idx];
        }
    }
    return &system[i].position;
}

#[cfg(test)]
mod tests {
    use super::*;
    use sys::{System, UnitCell};
    use energy::*;
    use utils::{system_from_xyz, unit_from};
    use types::Vector3D;

    fn testing_system() -> System {
        let mut system = system_from_xyz("4
        bonds
        O     0.000000     0.000000     0.000000
        O     0.000000     0.000000     1.480000
        H     0.895669     0.000000    -0.316667
        H    -0.895669     0.000000     1.796667");
        system.set_cell(UnitCell::cubic(10.0));
        assert!(system.molecules().len() == 1);

        system.interactions_mut().add_pair("H", "H", PairInteraction::new(
            Box::new(LennardJones{sigma: 3.0, epsilon: unit_from(0.5, "kJ/mol")}), 3.0
        ));

        system.interactions_mut().add_pair("O", "O", PairInteraction::new(
            Box::new(NullPotential), 3.0
        ));

        system.interactions_mut().add_pair("O", "H", PairInteraction::new(
            Box::new(LennardJones{sigma: 1.0, epsilon: unit_from(0.3, "kJ/mol")}), 3.0
        ));

        system.interactions_mut().add_bond("O", "H",
            Box::new(Harmonic{x0: 3.4, k: unit_from(522.0, "kJ/mol/A^2")})
        );

        system.interactions_mut().add_angle("O", "O", "H",
            Box::new(Harmonic{x0: f64::to_radians(120.0), k: unit_from(150.0, "kJ/mol/deg^2")})
        );

        system.interactions_mut().add_dihedral("H", "O", "O", "H",
            Box::new(Harmonic{x0: f64::to_radians(180.0), k: unit_from(800.0, "kJ/mol/deg^2")})
        );

        system.interactions_mut().set_coulomb(Box::new(Wolf::new(8.0)));

        for atom in &mut system {
            if atom.name() == "O" {
                atom.charge = -0.5;
            } else if atom.name() == "H" {
                atom.charge = 0.5;
            }
        }

        return system;
    }

    #[test]
    fn cache_energy() {
        let system = testing_system();
        let mut cache = EnergyCache::new();
        cache.init(&system);

        assert_approx_eq!(cache.energy(), system.potential_energy());
    }

    #[test]
    fn move_atoms() {
        let mut system = testing_system();
        let mut cache = EnergyCache::new();
        let old_e = system.potential_energy();
        cache.init(&system);

        let idxes = vec![0, 3];
        let newpos = &[Vector3D::new(0.0, 0.0, 0.5), Vector3D::new(-0.7, 0.2, 1.5)];

        let cost = cache.move_particles_cost(&system, idxes, newpos);

        system[0].position = newpos[0];
        system[3].position = newpos[1];
        let new_e = system.potential_energy();
        assert_approx_eq!(cost, new_e - old_e, 1e-14);

        cache.update(&mut system);
        assert_approx_eq!(cache.energy(), new_e, 1e-14);

        // Check that the cache is really updated
        let old_e = new_e;
        let idxes = vec![2, 3];
        let newpos = &[Vector3D::new(0.9, 0.2, -0.4), Vector3D::new(-0.9, 0.0, 1.8)];
        let cost = cache.move_particles_cost(&system, idxes, newpos);
        system[2].position = newpos[0];
        system[3].position = newpos[1];
        let new_e = system.potential_energy();
        assert_approx_eq!(cost, new_e - old_e, 1e-14);
    }
}
<commit_msg>Added tail corrections to cache<commit_after>// Lumol, an extensible molecular simulation engine
// Copyright (C) 2015-2016 G. Fraux — BSD license

//! Caching energy components for speeding up the energy computation in
//! Monte-Carlo simulations.
//!
//! In most of Monte-Carlo moves, only a very small subset of the system changes.
//! We can use that property to remove the need of recomputing most of the
//! energy components, by storing them and providing update callbacks.
use std::mem;

use super::System;
use types::{Vector3D, Array2};

/// Callback for updating a cache. It also take an `&mut System` argument for
/// updating the cache inside the global potentials.
type UpdateCallback = Box<Fn(&mut EnergyCache, &mut System)>;

/// This is a cache for energy computation.
///
/// Cache integrity is left up to the user of this structure: any function with
/// `update_` prefix must be called as needed to ensure cache consistency.
pub struct EnergyCache {
    /// 2-D array containing the pairs interactions between particles `i` and
    /// `j` at index `i, j`
    pairs_cache: Array2<f64>,
    /// Energy of all the pairs in the system
    pairs: f64,
    /// Contribution of long range corrections
    pairs_tail: f64,
    /// Energy of all the bonds in the system
    bonds: f64,
    /// Energy of all the angles in the system
    angles: f64,
    /// Energy of all the dihedrals angles in the system
    dihedrals: f64,
    /// Energy of coulombic interactions
    coulomb: f64,
    /// Energy of global interactions
    global: f64,
    /// Callback to be called to update the cache if the system is modified
    updater: Option<UpdateCallback>
}

impl Default for EnergyCache {
    fn default() -> EnergyCache {
        EnergyCache::new()
    }
}

impl EnergyCache {
    /// Create a new empty energy cache.
    pub fn new() -> EnergyCache {
        EnergyCache {
            pairs_cache: Array2::zeros((0, 0)),
            pairs: 0.0,
            pairs_tail: 0.0,
            bonds: 0.0,
            angles: 0.0,
            dihedrals: 0.0,
            coulomb: 0.0,
            global: 0.0,
            updater: None,
        }
    }

    /// Clear all values in the cache by setting them to 0
    fn clear(&mut self) {
        self.pairs_cache.assign(0.0);
        self.pairs = 0.0;
        self.pairs_tail = 0.0;
        self.bonds = 0.0;
        self.angles = 0.0;
        self.dihedrals = 0.0;
        self.coulomb = 0.0;
        self.global = 0.0;
    }

    /// Initialize the cache to be used with `system`. After a call to this
    /// function, the cache is only usable with the same system. To change
    /// the associated system, one must call this function again.
    pub fn init(&mut self, system: &System) {
        self.clear();
        self.pairs_cache.resize_if_different((system.size(), system.size()));

        let evaluator = system.energy_evaluator();

        for i in 0..system.size() {
            for j in (i + 1)..system.size() {
                let r = system.nearest_image(i, j).norm();
                let energy = evaluator.pair(r, i, j);
                self.pairs_cache[(i, j)] = energy;
                self.pairs_cache[(j, i)] = energy;
                self.pairs += energy;
            }
        }

        self.pairs_tail = evaluator.pairs_tail();
        self.bonds = evaluator.bonds();
        self.angles = evaluator.angles();
        self.dihedrals = evaluator.dihedrals();
        self.coulomb = evaluator.coulomb();
        self.global = evaluator.global();
    }

    /// Get the cached energy
    pub fn energy(&self) -> f64 {
        let mut energy = 0.0;
        energy += self.pairs;
        energy += self.pairs_tail;

        energy += self.bonds;
        energy += self.angles;
        energy += self.dihedrals;

        energy += self.coulomb;
        energy += self.global;
        return energy;
    }

    /// Update the cache after a call to a `EnergyCache::*_cost` function or
    /// `EnergyCache::unused`.
    pub fn update(&mut self, system: &mut System) {
        let updater = mem::replace(&mut self.updater, None);
        if let Some(updater) = updater {
            updater(self, system);
        } else {
            fatal_error!("Error in `EnergyCache::update`: \
                         This function MUST be called after a call to a `*_cost` function");
        }
    }

    /// This function should be called whenever the cache is not used, but one
    /// still want it to be updated. Future call to `EnergyCache::update` will
    /// recompute the full cache.
    pub fn unused(&mut self) {
        self.updater = Some(Box::new(|cache, system| {
            cache.init(system);
        }))
    }
}

impl EnergyCache {
    /// Get the cost of moving the set of particles with indexes in `idxes` to
    /// `newpos`. This function DO NOT update the cache, the
    /// `update_particles_moved` function MUST be called if the particles are
    /// effectively moved.
    pub fn move_particles_cost(&mut self, system: &System, idxes: Vec<usize>, newpos: &[Vector3D]) -> f64 {
        let evaluator = system.energy_evaluator();

        // First, go for pair interactions
        let mut new_pairs = Array2::<f64>::zeros((system.size(), system.size()));
        let mut pairs_delta = 0.0;
        // Interactions with the sub-system not being moved
        for (i, &part_i) in idxes.iter().enumerate() {
            for part_j in 0..system.size() {
                // Exclude interactions inside the sub-system.
                if idxes.contains(&part_j) {continue}

                let r = system.cell().distance(&system[part_j].position, &newpos[i]);
                let energy = evaluator.pair(r, part_i, part_j);

                pairs_delta += energy;
                new_pairs[(part_i, part_j)] += energy;
                new_pairs[(part_j, part_i)] += energy;

                pairs_delta -= self.pairs_cache[(part_i, part_j)];
            }
        }

        // Interactions withing the sub-system being moved
        for (i, &part_i) in idxes.iter().enumerate() {
            for (j, &part_j) in idxes.iter().enumerate().skip(i + 1) {
                let r = system.cell().distance(&newpos[i], &newpos[j]);
                let energy = evaluator.pair(r, part_i, part_j);

                pairs_delta += energy;
                new_pairs[(part_i, part_j)] += energy;
                new_pairs[(part_j, part_i)] += energy;

                pairs_delta -= self.pairs_cache[(part_i, part_j)];
            }
        }

        // Recompute everything here. One day we could store some terms and do
        // the same thing than for pair interactions.
        let mut bonds = 0.0;
        let mut angles = 0.0;
        let mut dihedrals = 0.0;
        for molecule in system.molecules() {
            for bond in molecule.bonds() {
                let (i, j) = (bond.i(), bond.j());
                let ri = new_position(system, i, &idxes, newpos);
                let rj = new_position(system, j, &idxes, newpos);
                let r = system.cell().distance(ri, rj);
                bonds += evaluator.bond(r, i, j);
            }

            for angle in molecule.angles() {
                let (i, j, k) = (angle.i(), angle.j(), angle.k());
                let ri = new_position(system, i, &idxes, newpos);
                let rj = new_position(system, j, &idxes, newpos);
                let rk = new_position(system, k, &idxes, newpos);
                let theta = system.cell().angle(ri, rj, rk);
                angles += evaluator.angle(theta, i, j, k);
            }

            for dihedral in molecule.dihedrals() {
                let (i, j, k, m) = (dihedral.i(), dihedral.j(), dihedral.k(), dihedral.m());
                let ri = new_position(system, i, &idxes, newpos);
                let rj = new_position(system, j, &idxes, newpos);
                let rk = new_position(system, k, &idxes, newpos);
                let rm = new_position(system, m, &idxes, newpos);
                let phi = system.cell().dihedral(ri, rj, rk, rm);
                dihedrals += evaluator.dihedral(phi, i, j, k, m);
            }
        }

        let coulomb_delta = if let Some(coulomb) = system.interactions().coulomb() {
            coulomb.borrow_mut().move_particles_cost(system, &idxes, newpos)
        } else {
            0.0
        };

        let mut global_delta = 0.0;
        for potential in system.interactions().globals() {
            global_delta += potential.borrow_mut().move_particles_cost(system, &idxes, newpos);
        }

        let pairs_tail = evaluator.pairs_tail();

        let cost = pairs_delta + (pairs_tail - self.pairs_tail)
                               + (bonds - self.bonds)
                               + (angles - self.angles)
                               + (dihedrals - self.dihedrals)
                               + coulomb_delta + global_delta;

        self.updater = Some(Box::new(move |cache, system| {
            cache.pairs_tail = pairs_tail;
            cache.bonds = bonds;
            cache.angles = angles;
            cache.dihedrals = dihedrals;

            cache.pairs += pairs_delta;
            cache.coulomb += coulomb_delta;
            cache.global += global_delta;

            let (n, m) = new_pairs.shape();
            debug_assert!(n == m);
            debug_assert!((n, m) == cache.pairs_cache.shape());
            // only loop over the indices that actually changed
            for &i in idxes.iter() {
                for j in 0..n {
                    if idxes.contains(&j) {continue}
                    cache.pairs_cache[(i, j)] = new_pairs[(i, j)];
                    cache.pairs_cache[(j, i)] = new_pairs[(i, j)];
                }
                for &k in idxes.iter().skip(i + 1) {
                    cache.pairs_cache[(i, k)] = new_pairs[(i, k)];
                    cache.pairs_cache[(k, i)] = new_pairs[(i, k)];
                }
            }

            // Update the cache for the global potentials
            if let Some(coulomb) = system.interactions().coulomb() {
                coulomb.borrow_mut().update();
            }
            for potential in system.interactions().globals() {
                potential.borrow_mut().update();
            }
        }));
        return cost;
    }
}

/// Return either the new position of a particle (from `newpos`) if its index
/// is in `idxes`, or its old position in the system.
fn new_position<'a>(system: &'a System, i: usize, idxes: &[usize], newpos: &'a[Vector3D]) -> &'a Vector3D {
    for (idx, &particle) in idxes.iter().enumerate() {
        if particle == i {
            return &newpos[idx];
        }
    }
    return &system[i].position;
}

#[cfg(test)]
mod tests {
    use super::*;
    use sys::{System, UnitCell};
    use energy::*;
    use utils::{system_from_xyz, unit_from};
    use types::Vector3D;

    fn testing_system() -> System {
        let mut system = system_from_xyz("4
        bonds
        O     0.000000     0.000000     0.000000
        O     0.000000     0.000000     1.480000
        H     0.895669     0.000000    -0.316667
        H    -0.895669     0.000000     1.796667");
        system.set_cell(UnitCell::cubic(10.0));
        assert!(system.molecules().len() == 1);

        system.interactions_mut().add_pair("H", "H", PairInteraction::new(
            Box::new(LennardJones{sigma: 3.0, epsilon: unit_from(0.5, "kJ/mol")}), 3.0
        ));

        system.interactions_mut().add_pair("O", "O", PairInteraction::new(
            Box::new(NullPotential), 3.0
        ));

        system.interactions_mut().add_pair("O", "H", PairInteraction::new(
            Box::new(LennardJones{sigma: 1.0, epsilon: unit_from(0.3, "kJ/mol")}), 3.0
        ));

        system.interactions_mut().add_bond("O", "H",
            Box::new(Harmonic{x0: 3.4, k: unit_from(522.0, "kJ/mol/A^2")})
        );

        system.interactions_mut().add_angle("O", "O", "H",
            Box::new(Harmonic{x0: f64::to_radians(120.0), k: unit_from(150.0, "kJ/mol/deg^2")})
        );

        system.interactions_mut().add_dihedral("H", "O", "O", "H",
            Box::new(Harmonic{x0: f64::to_radians(180.0), k: unit_from(800.0, "kJ/mol/deg^2")})
        );

        system.interactions_mut().set_coulomb(Box::new(Wolf::new(8.0)));

        for atom in &mut system {
            if atom.name() == "O" {
                atom.charge = -0.5;
            } else if atom.name() == "H" {
                atom.charge = 0.5;
            }
        }

        return system;
    }

    #[test]
    fn cache_energy() {
        let system = testing_system();
        let mut cache = EnergyCache::new();
        cache.init(&system);

        assert_approx_eq!(cache.energy(), system.potential_energy());
    }

    #[test]
    fn move_atoms() {
        let mut system = testing_system();
        let mut cache = EnergyCache::new();
        let old_e = system.potential_energy();
        cache.init(&system);

        let idxes = vec![0, 3];
        let newpos = &[Vector3D::new(0.0, 0.0, 0.5), Vector3D::new(-0.7, 0.2, 1.5)];

        let cost = cache.move_particles_cost(&system, idxes, newpos);

        system[0].position = newpos[0];
        system[3].position = newpos[1];
        let new_e = system.potential_energy();
        assert_approx_eq!(cost, new_e - old_e, 1e-14);

        cache.update(&mut system);
        assert_approx_eq!(cache.energy(), new_e, 1e-14);

        // Check that the cache is really updated
        let old_e = new_e;
        let idxes = vec![2, 3];
        let newpos = &[Vector3D::new(0.9, 0.2, -0.4), Vector3D::new(-0.9, 0.0, 1.8)];
        let cost = cache.move_particles_cost(&system, idxes, newpos);
        system[2].position = newpos[0];
        system[3].position = newpos[1];
        let new_e = system.potential_energy();
        assert_approx_eq!(cost, new_e - old_e, 1e-14);
    }
}
<|endoftext|>


<commit_before>/*
Module: map

A hashmap
*/

/* Section: Types */

/*
Type: hashfn

A function that returns a hash of a value.
The hash should concentrate entropy in the
lower bits.
*/
type hashfn<K> = fn(K) -> uint;

/*
Type: eqfn

Equality
*/
type eqfn<K> = fn(K, K) -> bool;

/*
Type: hashset

A convenience type to treat a hashmap as a set
*/
type hashset<K> = hashmap<K, ()>;

/*
Obj: hashmap
*/
type hashmap<K, V> = obj {
    /*
    Method: size

    Return the number of elements in the map
    */
    fn size() -> uint;
    /*
    Method: insert

    Add a value to the map. If the map already contains a value for
    the specified key then the original value is replaced.

    Returns:

    True if the key did not already exist in the map
    */
    fn insert(K, V) -> bool;
    /*
    Method: contains_key

    Returns true if the map contains a value for the specified key
    */
    fn contains_key(K) -> bool;
    /*
    Method: get

    Get the value for the specified key

    Failure:

    If the key does not exist in the map
    */
    fn get(K) -> V;
    /*
    Method: find

    Get the value for the specified key. If the key does not exist
    in the map then returns none.
    */
    fn find(K) -> option::t<V>;
    /*
    Method: remove

    Remove and return a value from the map. If the key does not exist
    in the map then returns none.
    */
    fn remove(K) -> option::t<V>;
    /*
    Method: rehash

    Force map growth and rehashing
    */
    fn rehash();
    /*
    Method: items

    Iterate over all the key/value pairs in the map
    */
    fn items(block(K, V));
    /*
    Method: keys

    Iterate over all the keys in the map
    */
    fn keys(block(K));
    /*
    Iterate over all the values in the map
    */
    fn values(block(V));
};

/* Section: Operations */

mod chained {
    type entry<copy K, copy V> = {
        hash: uint,
        key: K,
        mutable value: V,
        mutable next: chain<K, V>
    };

    tag chain<copy K, copy V> {
        present(@entry<K, V>);
        absent;
    }

    type t<copy K, copy V> = {
        mutable size: uint,
        mutable chains: [mutable chain<K,V>],
        hasher: hashfn<K>,
        eqer: eqfn<K>
    };

    tag search_result<copy K, copy V> {
        not_found;
        found_first(uint, @entry<K,V>);
        found_after(@entry<K,V>, @entry<K,V>);
    }

    fn search_rem<copy K, copy V>(tbl: t<K,V>,
                                  k: K,
                                  h: uint,
                                  idx: uint,
                                  e_root: @entry<K,V>) -> search_result<K,V> {
        let e0 = e_root;
        let comp = 1u;   // for logging
        while true {
            alt e0.next {
              absent. {
                log("search_tbl", "absent", "comparisons", comp,
                    "hash", h, "idx", idx);

                ret not_found;
              }
              present(e1) {
                comp += 1u;
                let e1_key = e1.key; // Satisfy alias checker.
                if e1.hash == h && tbl.eqer(e1_key, k) {
                    log("search_tbl", "present", "comparisons", comp,
                        "hash", h, "idx", idx);
                    ret found_after(e0, e1);
                } else {
                    e0 = e1;
                }
              }
            }
        }
        util::unreachable();
    }

    fn search_tbl<copy K, copy V>(
        tbl: t<K,V>, k: K, h: uint) -> search_result<K,V> {
        let idx = h % vec::len(tbl.chains);
        alt tbl.chains[idx] {
          absent. {
            log("search_tbl", "absent", "comparisons", 0u,
                "hash", h, "idx", idx);
            ret not_found;
          }
          present(e) {
            let e_key = e.key; // Satisfy alias checker.
            if e.hash == h && tbl.eqer(e_key, k) {
                log("search_tbl", "present", "comparisons", 1u,
                    "hash", h, "idx", idx);
                ret found_first(idx, e);
            } else {
                ret search_rem(tbl, k, h, idx, e);
            }
          }
        }
    }

    fn insert<copy K, copy V>(tbl: t<K,V>, k: K, v: V) -> bool {
        let hash = tbl.hasher(k);
        alt search_tbl(tbl, k, hash) {
          not_found. {
            tbl.size += 1u;
            let idx = hash % vec::len(tbl.chains);
            let old_chain = tbl.chains[idx];
            tbl.chains[idx] = present(@{
                hash: hash,
                key: k,
                mutable value: v,
                mutable next: old_chain});
            ret true;
          }
          found_first(_, entry) {
            entry.value = v;
            ret false;
          }
          found_after(_, entry) {
            entry.value = v;
            ret false
          }
        }
    }

    fn get<copy K, copy V>(tbl: t<K,V>, k: K) -> option::t<V> {
        alt search_tbl(tbl, k, tbl.hasher(k)) {
          not_found. {
            ret option::none;
          }

          found_first(_, entry) {
            ret option::some(entry.value);
          }

          found_after(_, entry) {
            ret option::some(entry.value);
          }
        }
    }

    fn remove<copy K, copy V>(tbl: t<K,V>, k: K) -> option::t<V> {
        alt search_tbl(tbl, k, tbl.hasher(k)) {
          not_found. {
            ret option::none;
          }

          found_first(idx, entry) {
            tbl.size -= 1u;
            tbl.chains[idx] = entry.next;
            ret option::some(entry.value);
          }

          found_after(eprev, entry) {
            tbl.size -= 1u;
            eprev.next = entry.next;
            ret option::some(entry.value);
          }
        }
    }

    fn chains<copy K, copy V>(nchains: uint) -> [mutable chain<K,V>] {
        ret vec::init_elt_mut(absent, nchains);
    }

    fn foreach_entry<copy K, copy V>(chain0: chain<K,V>,
                                     blk: block(@entry<K,V>)) {
        let chain = chain0;
        while true {
            alt chain {
              absent. { ret; }
              present(entry) {
                let next = entry.next;
                blk(entry); // may modify entry.next!
                chain = next;
              }
            }
        }
    }

    fn foreach_chain<copy K, copy V>(chains: [const chain<K,V>],
                                     blk: block(@entry<K,V>)) {
        let i = 0u, n = vec::len(chains);
        while i < n {
            foreach_entry(chains[i], blk);
            i += 1u;
        }
    }

    fn rehash<copy K, copy V>(tbl: t<K,V>) {
        let old_chains = tbl.chains;
        let n_old_chains = vec::len(old_chains);
        let n_new_chains: uint = uint::next_power_of_two(n_old_chains + 1u);
        tbl.chains = chains(n_new_chains);
        foreach_chain(old_chains) { |entry|
            let idx = entry.hash % n_new_chains;
            entry.next = tbl.chains[idx];
            tbl.chains[idx] = present(entry);
        }
    }

    fn items<copy K, copy V>(tbl: t<K,V>, blk: block(K,V)) {
        let tbl_chains = tbl.chains;  // Satisfy alias checker.
        foreach_chain(tbl_chains) { |entry|
            let key = entry.key;
            let value = entry.value;
            blk(key, value);
        }
    }

    obj o<copy K, copy V>(tbl: @t<K,V>,
                          lf: float) {
        fn size() -> uint {
            ret tbl.size;
        }

        fn insert(k: K, v: V) -> bool {
            let nchains = vec::len(tbl.chains);
            let load = (tbl.size + 1u as float) / (nchains as float);
            if load > lf {
                rehash(*tbl);
            }
            ret insert(*tbl, k, v);
        }

        fn contains_key(k: K) -> bool {
            ret option::is_some(get(*tbl, k));
        }

        fn get(k: K) -> V {
            ret option::get(get(*tbl, k));
        }

        fn find(k: K) -> option::t<V> {
            ret get(*tbl, k);
        }

        fn remove(k: K) -> option::t<V> {
            ret remove(*tbl, k);
        }

        fn rehash() {
            rehash(*tbl);
        }

        fn items(blk: block(K, V)) {
            items(*tbl, blk);
        }

        fn keys(blk: block(K)) {
            items(*tbl) { |k, _v| blk(k) }
        }

        fn values(blk: block(V)) {
            items(*tbl) { |_k, v| blk(v) }
        }
    }

    fn mk<copy K, copy V>(hasher: hashfn<K>, eqer: eqfn<K>) -> hashmap<K,V> {
        let initial_capacity: uint = 32u; // 2^5
        let t = @{mutable size: 0u,
                  mutable chains: chains(initial_capacity),
                  hasher: hasher,
                  eqer: eqer};
        ret o(t, 0.75);
    }
}

/*
Function: mk_flat_hashmap

Construct a "flat" hashmap, meaning that there are
not chains per buckets, but rather we search a sequence
of buckets for each key.

Warning: it is unclear to me that this code is correct
on 32-bit processors.  Check out the 'hash-tearing' code
in hash() and the comment surrounding it. - Niko

Parameters:

hasher - The hash function for key type K
eqer - The equality function for key type K
*/
fn mk_flat_hashmap<copy K, copy V>(hasher: hashfn<K>, eqer: eqfn<K>)
    -> hashmap<K, V> {
    let initial_capacity: uint = 32u; // 2^5

    let load_factor: util::rational = {num: 3, den: 4};
    tag bucket<copy K, copy V> { nil; deleted; some(K, V); }
    fn make_buckets<copy K, copy V>(nbkts: uint) -> [mutable bucket<K, V>] {
        ret vec::init_elt_mut::<bucket<K, V>>(nil::<K, V>, nbkts);
    }
    // Derive two hash functions from the one given by taking the upper
    // half and lower half of the uint bits.  Our bucket probing
    // sequence is then defined by
    //
    //   hash(key, i) := hashl(key) * i + hashr(key)   for i = 0, 1, 2, ...
    //
    // Tearing the hash function apart this way is kosher in practice
    // as, assuming 32-bit uints, the table would have to be at 2^32
    // buckets before the resulting pair of hash functions no longer
    // probes all buckets for a fixed key.  Note that hashl is made to
    // output odd numbers (hence coprime to the number of nbkts, which
    // is always a power? of 2), so that all buckets are probed for a
    // fixed key.

    fn hashl(n: uint) -> uint { ret (n >>> 16u) * 2u + 1u; }
    fn hashr(n: uint) -> uint { ret 0x0000_ffff_u & n; }
    fn hash(h: uint, nbkts: uint, i: uint) -> uint {
        ret (hashl(h) * i + hashr(h)) % nbkts;
    }

    /**
     * We attempt to never call this with a full table.  If we do, it
     * will fail.
     */
    fn insert_common<copy K, copy V>(hasher: hashfn<K>, eqer: eqfn<K>,
                                     bkts: [mutable bucket<K, V>],
                                     nbkts: uint, key: K, val: V) -> bool {
        let i: uint = 0u;
        let h = hasher(key);
        while i < nbkts {
            let j: uint = hash(h, nbkts, i);
            alt bkts[j] {
              some(k, _) {
                // Copy key to please alias analysis.

                let k_ = k;
                if eqer(key, k_) {
                    log("map updated", "i", i, "h", h, "nbkts", nbkts);
                    bkts[j] = some(k_, val);
                    ret false;
                }
                i += 1u;
              }
              _ {
                log("map inserted", "i", i, "h", h, "nbkts", nbkts);
                bkts[j] = some(key, val);
                ret true;
              }
            }
        }
        fail; // full table
    }
    fn find_common<copy K, copy V>(hasher: hashfn<K>, eqer: eqfn<K>,
                                   bkts: [mutable bucket<K, V>],
                                   nbkts: uint, key: K) -> option::t<V> {
        let i: uint = 0u;
        let h = hasher(key);
        while i < nbkts {
            let j: uint = hash(h, nbkts, i);
            alt bkts[j] {
              some(k, v) {
                // Copy to please alias analysis.
                let k_ = k;
                let v_ = v;
                if eqer(key, k_) {
                    log("map present", "i", i, "h", h, "nbkts", nbkts);
                    ret option::some(v_);
                }
              }
              nil. {
                log("map absent", "i", i, "h", h, "nbkts", nbkts);
                ret option::none;
              }
              deleted. { }
            }
            i += 1u;
        }
        ret option::none;
    }
    fn rehash<copy K, copy V>(hasher: hashfn<K>, eqer: eqfn<K>,
                              oldbkts: [mutable bucket<K, V>],
                              _noldbkts: uint,
                              newbkts: [mutable bucket<K, V>],
                              nnewbkts: uint) {
        for b: bucket<K, V> in oldbkts {
            alt b {
              some(k_, v_) {
                let k = k_;
                let v = v_;
                insert_common(hasher, eqer, newbkts, nnewbkts, k, v);
              }
              _ { }
            }
        }
    }
    obj hashmap<copy K, copy V>(hasher: hashfn<K>,
                                eqer: eqfn<K>,
                                mutable bkts: [mutable bucket<K, V>],
                                mutable nbkts: uint,
                                mutable nelts: uint,
                                lf: util::rational) {
        fn size() -> uint { ret nelts; }
        fn insert(key: K, val: V) -> bool {
            let load: util::rational =
                {num: nelts + 1u as int, den: nbkts as int};
            if !util::rational_leq(load, lf) {
                let nnewbkts: uint = uint::next_power_of_two(nbkts + 1u);
                let newbkts = make_buckets(nnewbkts);
                rehash(hasher, eqer, bkts, nbkts, newbkts, nnewbkts);
                bkts = newbkts;
                nbkts = nnewbkts;
            }
            if insert_common(hasher, eqer, bkts, nbkts, key, val) {
                nelts += 1u;
                ret true;
            }
            ret false;
        }
        fn contains_key(key: K) -> bool {
            ret alt find_common(hasher, eqer, bkts, nbkts, key) {
                  option::some(_) { true }
                  _ { false }
                };
        }
        fn get(key: K) -> V {
            ret alt find_common(hasher, eqer, bkts, nbkts, key) {
                  option::some(val) { val }
                  _ { fail }
                };
        }
        fn find(key: K) -> option::t<V> {
            be find_common(hasher, eqer, bkts, nbkts, key);
        }
        fn remove(key: K) -> option::t<V> {
            let i: uint = 0u;
            let h = hasher(key);
            while i < nbkts {
                let j: uint = hash(h, nbkts, i);
                alt bkts[j] {
                  some(k, v) {
                    let k_ = k;
                    let vo = option::some(v);
                    if eqer(key, k_) {
                        bkts[j] = deleted;
                        nelts -= 1u;
                        ret vo;
                    }
                  }
                  deleted. { }
                  nil. { ret option::none; }
                }
                i += 1u;
            }
            ret option::none;
        }
        fn rehash() {
            let newbkts = make_buckets(nbkts);
            rehash(hasher, eqer, bkts, nbkts, newbkts, nbkts);
            bkts = newbkts;
        }
        fn items(it: block(K, V)) {
            for b in bkts {
                alt b { some(k, v) { it(copy k, copy v); } _ { } }
            }
        }
        fn keys(it: block(K)) {
            for b in bkts {
                alt b { some(k, _) { it(copy k); } _ { } }
            }
        }
        fn values(it: block(V)) {
            for b in bkts {
                alt b { some(_, v) { it(copy v); } _ { } }
            }
        }
    }
    let bkts = make_buckets(initial_capacity);
    ret hashmap(hasher, eqer, bkts, initial_capacity, 0u, load_factor);
}

/*
Function: mk_hashmap

Construct a hashmap.

Parameters:

hasher - The hash function for key type K
eqer - The equality function for key type K
*/
fn mk_hashmap<copy K, copy V>(hasher: hashfn<K>, eqer: eqfn<K>)
    -> hashmap<K, V> {
    ret chained::mk(hasher, eqer);
}

/*
Function: new_str_hash

Construct a hashmap for string keys
*/
fn new_str_hash<copy V>() -> hashmap<str, V> {
    ret mk_hashmap(str::hash, str::eq);
}

/*
Function: new_int_hash

Construct a hashmap for int keys
*/
fn new_int_hash<copy V>() -> hashmap<int, V> {
    fn hash_int(&&x: int) -> uint { ret x as uint; }
    fn eq_int(&&a: int, &&b: int) -> bool { ret a == b; }
    ret mk_hashmap(hash_int, eq_int);
}

/*
Function: new_uint_hash

Construct a hashmap for uint keys
*/
fn new_uint_hash<copy V>() -> hashmap<uint, V> {
    fn hash_uint(&&x: uint) -> uint { ret x; }
    fn eq_uint(&&a: uint, &&b: uint) -> bool { ret a == b; }
    ret mk_hashmap(hash_uint, eq_uint);
}

/*
Function: set_add

Convenience function for adding keys to a hashmap with nil type keys
*/
fn set_add<K>(set: hashset<K>, key: K) -> bool { ret set.insert(key, ()); }

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// End:
<commit_msg>Purge old map impl: more collisions, not clear it is safe on 64bit<commit_after>/*
Module: map

A hashmap
*/

/* Section: Types */

/*
Type: hashfn

A function that returns a hash of a value.
The hash should concentrate entropy in the
lower bits.
*/
type hashfn<K> = fn(K) -> uint;

/*
Type: eqfn

Equality
*/
type eqfn<K> = fn(K, K) -> bool;

/*
Type: hashset

A convenience type to treat a hashmap as a set
*/
type hashset<K> = hashmap<K, ()>;

/*
Obj: hashmap
*/
type hashmap<K, V> = obj {
    /*
    Method: size

    Return the number of elements in the map
    */
    fn size() -> uint;
    /*
    Method: insert

    Add a value to the map. If the map already contains a value for
    the specified key then the original value is replaced.

    Returns:

    True if the key did not already exist in the map
    */
    fn insert(K, V) -> bool;
    /*
    Method: contains_key

    Returns true if the map contains a value for the specified key
    */
    fn contains_key(K) -> bool;
    /*
    Method: get

    Get the value for the specified key

    Failure:

    If the key does not exist in the map
    */
    fn get(K) -> V;
    /*
    Method: find

    Get the value for the specified key. If the key does not exist
    in the map then returns none.
    */
    fn find(K) -> option::t<V>;
    /*
    Method: remove

    Remove and return a value from the map. If the key does not exist
    in the map then returns none.
    */
    fn remove(K) -> option::t<V>;
    /*
    Method: rehash

    Force map growth and rehashing
    */
    fn rehash();
    /*
    Method: items

    Iterate over all the key/value pairs in the map
    */
    fn items(block(K, V));
    /*
    Method: keys

    Iterate over all the keys in the map
    */
    fn keys(block(K));
    /*
    Iterate over all the values in the map
    */
    fn values(block(V));
};

/* Section: Operations */

mod chained {
    type entry<copy K, copy V> = {
        hash: uint,
        key: K,
        mutable value: V,
        mutable next: chain<K, V>
    };

    tag chain<copy K, copy V> {
        present(@entry<K, V>);
        absent;
    }

    type t<copy K, copy V> = {
        mutable size: uint,
        mutable chains: [mutable chain<K,V>],
        hasher: hashfn<K>,
        eqer: eqfn<K>
    };

    tag search_result<copy K, copy V> {
        not_found;
        found_first(uint, @entry<K,V>);
        found_after(@entry<K,V>, @entry<K,V>);
    }

    fn search_rem<copy K, copy V>(tbl: t<K,V>,
                                  k: K,
                                  h: uint,
                                  idx: uint,
                                  e_root: @entry<K,V>) -> search_result<K,V> {
        let e0 = e_root;
        let comp = 1u;   // for logging
        while true {
            alt e0.next {
              absent. {
                log("search_tbl", "absent", "comparisons", comp,
                    "hash", h, "idx", idx);

                ret not_found;
              }
              present(e1) {
                comp += 1u;
                let e1_key = e1.key; // Satisfy alias checker.
                if e1.hash == h && tbl.eqer(e1_key, k) {
                    log("search_tbl", "present", "comparisons", comp,
                        "hash", h, "idx", idx);
                    ret found_after(e0, e1);
                } else {
                    e0 = e1;
                }
              }
            }
        }
        util::unreachable();
    }

    fn search_tbl<copy K, copy V>(
        tbl: t<K,V>, k: K, h: uint) -> search_result<K,V> {
        let idx = h % vec::len(tbl.chains);
        alt tbl.chains[idx] {
          absent. {
            log("search_tbl", "absent", "comparisons", 0u,
                "hash", h, "idx", idx);
            ret not_found;
          }
          present(e) {
            let e_key = e.key; // Satisfy alias checker.
            if e.hash == h && tbl.eqer(e_key, k) {
                log("search_tbl", "present", "comparisons", 1u,
                    "hash", h, "idx", idx);
                ret found_first(idx, e);
            } else {
                ret search_rem(tbl, k, h, idx, e);
            }
          }
        }
    }

    fn insert<copy K, copy V>(tbl: t<K,V>, k: K, v: V) -> bool {
        let hash = tbl.hasher(k);
        alt search_tbl(tbl, k, hash) {
          not_found. {
            tbl.size += 1u;
            let idx = hash % vec::len(tbl.chains);
            let old_chain = tbl.chains[idx];
            tbl.chains[idx] = present(@{
                hash: hash,
                key: k,
                mutable value: v,
                mutable next: old_chain});
            ret true;
          }
          found_first(_, entry) {
            entry.value = v;
            ret false;
          }
          found_after(_, entry) {
            entry.value = v;
            ret false
          }
        }
    }

    fn get<copy K, copy V>(tbl: t<K,V>, k: K) -> option::t<V> {
        alt search_tbl(tbl, k, tbl.hasher(k)) {
          not_found. {
            ret option::none;
          }

          found_first(_, entry) {
            ret option::some(entry.value);
          }

          found_after(_, entry) {
            ret option::some(entry.value);
          }
        }
    }

    fn remove<copy K, copy V>(tbl: t<K,V>, k: K) -> option::t<V> {
        alt search_tbl(tbl, k, tbl.hasher(k)) {
          not_found. {
            ret option::none;
          }

          found_first(idx, entry) {
            tbl.size -= 1u;
            tbl.chains[idx] = entry.next;
            ret option::some(entry.value);
          }

          found_after(eprev, entry) {
            tbl.size -= 1u;
            eprev.next = entry.next;
            ret option::some(entry.value);
          }
        }
    }

    fn chains<copy K, copy V>(nchains: uint) -> [mutable chain<K,V>] {
        ret vec::init_elt_mut(absent, nchains);
    }

    fn foreach_entry<copy K, copy V>(chain0: chain<K,V>,
                                     blk: block(@entry<K,V>)) {
        let chain = chain0;
        while true {
            alt chain {
              absent. { ret; }
              present(entry) {
                let next = entry.next;
                blk(entry); // may modify entry.next!
                chain = next;
              }
            }
        }
    }

    fn foreach_chain<copy K, copy V>(chains: [const chain<K,V>],
                                     blk: block(@entry<K,V>)) {
        let i = 0u, n = vec::len(chains);
        while i < n {
            foreach_entry(chains[i], blk);
            i += 1u;
        }
    }

    fn rehash<copy K, copy V>(tbl: t<K,V>) {
        let old_chains = tbl.chains;
        let n_old_chains = vec::len(old_chains);
        let n_new_chains: uint = uint::next_power_of_two(n_old_chains + 1u);
        tbl.chains = chains(n_new_chains);
        foreach_chain(old_chains) { |entry|
            let idx = entry.hash % n_new_chains;
            entry.next = tbl.chains[idx];
            tbl.chains[idx] = present(entry);
        }
    }

    fn items<copy K, copy V>(tbl: t<K,V>, blk: block(K,V)) {
        let tbl_chains = tbl.chains;  // Satisfy alias checker.
        foreach_chain(tbl_chains) { |entry|
            let key = entry.key;
            let value = entry.value;
            blk(key, value);
        }
    }

    obj o<copy K, copy V>(tbl: @t<K,V>,
                          lf: float) {
        fn size() -> uint {
            ret tbl.size;
        }

        fn insert(k: K, v: V) -> bool {
            let nchains = vec::len(tbl.chains);
            let load = (tbl.size + 1u as float) / (nchains as float);
            if load > lf {
                rehash(*tbl);
            }
            ret insert(*tbl, k, v);
        }

        fn contains_key(k: K) -> bool {
            ret option::is_some(get(*tbl, k));
        }

        fn get(k: K) -> V {
            ret option::get(get(*tbl, k));
        }

        fn find(k: K) -> option::t<V> {
            ret get(*tbl, k);
        }

        fn remove(k: K) -> option::t<V> {
            ret remove(*tbl, k);
        }

        fn rehash() {
            rehash(*tbl);
        }

        fn items(blk: block(K, V)) {
            items(*tbl, blk);
        }

        fn keys(blk: block(K)) {
            items(*tbl) { |k, _v| blk(k) }
        }

        fn values(blk: block(V)) {
            items(*tbl) { |_k, v| blk(v) }
        }
    }

    fn mk<copy K, copy V>(hasher: hashfn<K>, eqer: eqfn<K>) -> hashmap<K,V> {
        let initial_capacity: uint = 32u; // 2^5
        let t = @{mutable size: 0u,
                  mutable chains: chains(initial_capacity),
                  hasher: hasher,
                  eqer: eqer};
        ret o(t, 0.75);
    }
}

/*
Function: mk_hashmap

Construct a hashmap.

Parameters:

hasher - The hash function for key type K
eqer - The equality function for key type K
*/
fn mk_hashmap<copy K, copy V>(hasher: hashfn<K>, eqer: eqfn<K>)
    -> hashmap<K, V> {
    ret chained::mk(hasher, eqer);
}

/*
Function: new_str_hash

Construct a hashmap for string keys
*/
fn new_str_hash<copy V>() -> hashmap<str, V> {
    ret mk_hashmap(str::hash, str::eq);
}

/*
Function: new_int_hash

Construct a hashmap for int keys
*/
fn new_int_hash<copy V>() -> hashmap<int, V> {
    fn hash_int(&&x: int) -> uint { ret x as uint; }
    fn eq_int(&&a: int, &&b: int) -> bool { ret a == b; }
    ret mk_hashmap(hash_int, eq_int);
}

/*
Function: new_uint_hash

Construct a hashmap for uint keys
*/
fn new_uint_hash<copy V>() -> hashmap<uint, V> {
    fn hash_uint(&&x: uint) -> uint { ret x; }
    fn eq_uint(&&a: uint, &&b: uint) -> bool { ret a == b; }
    ret mk_hashmap(hash_uint, eq_uint);
}

/*
Function: set_add

Convenience function for adding keys to a hashmap with nil type keys
*/
fn set_add<K>(set: hashset<K>, key: K) -> bool { ret set.insert(key, ()); }

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// End:
<|endoftext|>


<commit_before>use std::iter::FromIterator;
use std::mem;
use {Union, UnionFind, UnionResult};

#[derive(Copy, Clone, Debug)]
struct Payload<V> {
    data: V,
    link_last_child: usize,
}

/// Union-Find implementation with quick find operation.
#[derive(Debug)]
pub struct QuickFindUf<V> {
    link_root: Vec<usize>,
    link_sibling: Vec<usize>,
    payload: Vec<Option<Payload<V>>>,
}

impl<V> Clone for QuickFindUf<V>
    where V: Clone + Union
{
    #[inline]
    fn clone(&self) -> QuickFindUf<V> {
        QuickFindUf {
            link_root: self.link_root.clone(),
            link_sibling: self.link_sibling.clone(),
            payload: self.payload.clone(),
        }
    }

    #[inline]
    fn clone_from(&mut self, other: &QuickFindUf<V>) {
        self.link_root.clone_from(&other.link_root);
        self.link_sibling.clone_from(&other.link_sibling);
        self.payload.clone_from(&other.payload);
    }
}

impl<V: Union> UnionFind<V> for QuickFindUf<V> {
    #[inline]
    fn size(&self) -> usize {
        self.payload.len()
    }

    #[inline]
    fn insert(&mut self, data: V) -> usize {
        let key = self.payload.len();
        self.link_root.push(key);
        self.link_sibling.push(key);
        self.payload.push(Some(Payload {
            data: data,
            link_last_child: key,
        }));
        key
    }

    #[inline]
    fn union(&mut self, key0: usize, key1: usize) -> bool {
        let k0 = self.find(key0);
        let k1 = self.find(key1);
        if k0 == k1 {
            return false;
        }

        // Temporary replace with dummy to move out the elements of the vector.
        let Payload { data: d0, link_last_child: c0 } = mem::replace(&mut self.payload[k0], None)
                                                            .unwrap();
        let Payload { data: d1, link_last_child: c1 } = mem::replace(&mut self.payload[k1], None)
                                                            .unwrap();

        let (root, child_root, val, last) = match Union::union(d0, d1) {
            UnionResult::Left(val) => (k0, k1, val, c0),
            UnionResult::Right(val) => (k1, k0, val, c1),
        };

        self.link_sibling[last] = child_root;

        let mut elem = child_root;
        while self.link_sibling[elem] != elem {
            debug_assert_eq!(self.link_root[elem], child_root);
            self.link_root[elem] = root;
            elem = self.link_sibling[elem];
        }
        debug_assert_eq!(self.link_root[elem], child_root);
        self.link_root[elem] = root;

        self.payload[root] = Some(Payload {
            data: val,
            link_last_child: elem,
        });

        true
    }

    #[inline]
    fn find(&mut self, key: usize) -> usize {
        self.link_root[key]
    }

    #[inline]
    fn get(&mut self, key: usize) -> &V {
        let root_key = self.find(key);
        &self.payload[root_key].as_ref().unwrap().data
    }

    #[inline]
    fn get_mut(&mut self, key: usize) -> &mut V {
        let root_key = self.find(key);
        &mut self.payload[root_key].as_mut().unwrap().data
    }
}

impl<A: Union> FromIterator<A> for QuickFindUf<A> {
    #[inline]
    fn from_iter<T: IntoIterator<Item = A>>(iterator: T) -> QuickFindUf<A> {
        let mut uf = QuickFindUf {
            link_root: vec![],
            link_sibling: vec![],
            payload: vec![],
        };
        uf.extend(iterator);
        uf
    }
}

impl<A> Extend<A> for QuickFindUf<A> {
    #[inline]
    fn extend<T>(&mut self, iterable: T)
        where T: IntoIterator<Item = A>
    {
        let len = self.payload.len();
        let payload = iterable.into_iter()
                              .zip(len..)
                              .map(|(data, link)| {
                                  Payload {
                                      data: data,
                                      link_last_child: link,
                                  }
                              })
                              .map(Some);
        self.payload.extend(payload);

        let new_len = self.payload.len();
        self.link_root.extend(len..new_len);
        self.link_sibling.extend(len..new_len);
    }
}
<commit_msg>run rustfmt<commit_after>use std::iter::FromIterator;
use std::mem;
use {Union, UnionFind, UnionResult};

#[derive(Copy, Clone, Debug)]
struct Payload<V> {
    data: V,
    link_last_child: usize,
}

/// Union-Find implementation with quick find operation.
#[derive(Debug)]
pub struct QuickFindUf<V> {
    link_root: Vec<usize>,
    link_sibling: Vec<usize>,
    payload: Vec<Option<Payload<V>>>,
}

impl<V> Clone for QuickFindUf<V> where V: Clone + Union
{
    #[inline]
    fn clone(&self) -> QuickFindUf<V> {
        QuickFindUf {
            link_root: self.link_root.clone(),
            link_sibling: self.link_sibling.clone(),
            payload: self.payload.clone(),
        }
    }

    #[inline]
    fn clone_from(&mut self, other: &QuickFindUf<V>) {
        self.link_root.clone_from(&other.link_root);
        self.link_sibling.clone_from(&other.link_sibling);
        self.payload.clone_from(&other.payload);
    }
}

impl<V: Union> UnionFind<V> for QuickFindUf<V> {
    #[inline]
    fn size(&self) -> usize {
        self.payload.len()
    }

    #[inline]
    fn insert(&mut self, data: V) -> usize {
        let key = self.payload.len();
        self.link_root.push(key);
        self.link_sibling.push(key);
        self.payload.push(Some(Payload {
            data: data,
            link_last_child: key,
        }));
        key
    }

    #[inline]
    fn union(&mut self, key0: usize, key1: usize) -> bool {
        let k0 = self.find(key0);
        let k1 = self.find(key1);
        if k0 == k1 {
            return false;
        }

        // Temporary replace with dummy to move out the elements of the vector.
        let Payload { data: d0, link_last_child: c0 } = mem::replace(&mut self.payload[k0], None)
                                                            .unwrap();
        let Payload { data: d1, link_last_child: c1 } = mem::replace(&mut self.payload[k1], None)
                                                            .unwrap();

        let (root, child_root, val, last) = match Union::union(d0, d1) {
            UnionResult::Left(val) => (k0, k1, val, c0),
            UnionResult::Right(val) => (k1, k0, val, c1),
        };

        self.link_sibling[last] = child_root;

        let mut elem = child_root;
        while self.link_sibling[elem] != elem {
            debug_assert_eq!(self.link_root[elem], child_root);
            self.link_root[elem] = root;
            elem = self.link_sibling[elem];
        }
        debug_assert_eq!(self.link_root[elem], child_root);
        self.link_root[elem] = root;

        self.payload[root] = Some(Payload {
            data: val,
            link_last_child: elem,
        });

        true
    }

    #[inline]
    fn find(&mut self, key: usize) -> usize {
        self.link_root[key]
    }

    #[inline]
    fn get(&mut self, key: usize) -> &V {
        let root_key = self.find(key);
        &self.payload[root_key].as_ref().unwrap().data
    }

    #[inline]
    fn get_mut(&mut self, key: usize) -> &mut V {
        let root_key = self.find(key);
        &mut self.payload[root_key].as_mut().unwrap().data
    }
}

impl<A: Union> FromIterator<A> for QuickFindUf<A> {
    #[inline]
    fn from_iter<T: IntoIterator<Item = A>>(iterator: T) -> QuickFindUf<A> {
        let mut uf = QuickFindUf {
            link_root: vec![],
            link_sibling: vec![],
            payload: vec![],
        };
        uf.extend(iterator);
        uf
    }
}

impl<A> Extend<A> for QuickFindUf<A> {
    #[inline]
    fn extend<T>(&mut self, iterable: T)
        where T: IntoIterator<Item = A>
    {
        let len = self.payload.len();
        let payload = iterable.into_iter()
                              .zip(len..)
                              .map(|(data, link)| {
                                  Payload {
                                      data: data,
                                      link_last_child: link,
                                  }
                              })
                              .map(Some);
        self.payload.extend(payload);

        let new_len = self.payload.len();
        self.link_root.extend(len..new_len);
        self.link_sibling.extend(len..new_len);
    }
}
<|endoftext|>


<commit_before>//! Retry sending records configuration.

/// You can calculate retrying interval as the following equation:
///
/// `retry_interval = exp ** (multiplier + retry_counts)`
///
/// see: https://github.com/jimmycuadra/retry/blob/v0.4.0/src/lib.rs#L142-L143
///
/// You can estimate to caluculate concrete values like:
///
/// * retry_counts: 10, e ^ (5 + 10)/1000.0/60.0/60.0 = 0.908060381242253, about 0.9 hour
///
/// * retry_counts: 11, e ^ (5 + 11)/1000.0/60.0/60.0 = 2.4683640334744092, about 2.5 hours
///
/// * retry_counts: 12, e ^ (5 + 12)/1000.0/60.0/60.0 = 6.709709098215361, about 6.7 hours
///
/// where multiplier = 5,
/// e is [exponential function](https://doc.rust-lang.org/std/primitive.f64.html#method.exp).
///
/// ## Default values
///
/// * multiplier: 5_f64
/// * max: 10
#[derive(Debug, Clone, PartialEq)]
pub struct RetryConf {
    max: u64,
    multiplier: f64,
}

impl Default for RetryConf {
    fn default() -> RetryConf {
        RetryConf {
            max: 10,
            multiplier: 5_f64,
        }
    }
}

impl RetryConf {
    pub fn new() -> RetryConf {
        Default::default()
    }

    pub fn max(mut self, max: u64) -> RetryConf {
        self.max = max;
        self
    }

    pub fn multiplier(mut self, multiplier: f64) -> RetryConf {
        self.multiplier = multiplier;
        self
    }

    pub fn build(self) -> (u64, f64) {
        (self.max, self.multiplier)
    }
}
<commit_msg>retry_conf: Fix English<commit_after>//! Retry sending records configuration.

/// You can calculate retrying interval as the following equation:
///
/// `retry_interval = exp ** (multiplier + retry_counts)`
///
/// see: https://github.com/jimmycuadra/retry/blob/v0.4.0/src/lib.rs#L142-L143
///
/// You can estimate to caluculate with concrete values like:
///
/// * retry_counts: 10, e ^ (5 + 10)/1000.0/60.0/60.0 = 0.908060381242253, about 0.9 hour
///
/// * retry_counts: 11, e ^ (5 + 11)/1000.0/60.0/60.0 = 2.4683640334744092, about 2.5 hours
///
/// * retry_counts: 12, e ^ (5 + 12)/1000.0/60.0/60.0 = 6.709709098215361, about 6.7 hours
///
/// where multiplier = 5,
/// e is [exponential function](https://doc.rust-lang.org/std/primitive.f64.html#method.exp).
///
/// ## Default values
///
/// * multiplier: 5_f64
/// * max: 10
#[derive(Debug, Clone, PartialEq)]
pub struct RetryConf {
    max: u64,
    multiplier: f64,
}

impl Default for RetryConf {
    fn default() -> RetryConf {
        RetryConf {
            max: 10,
            multiplier: 5_f64,
        }
    }
}

impl RetryConf {
    pub fn new() -> RetryConf {
        Default::default()
    }

    pub fn max(mut self, max: u64) -> RetryConf {
        self.max = max;
        self
    }

    pub fn multiplier(mut self, multiplier: f64) -> RetryConf {
        self.multiplier = multiplier;
        self
    }

    pub fn build(self) -> (u64, f64) {
        (self.max, self.multiplier)
    }
}
<|endoftext|>