<commit_before>use super::Color;
use enum_map::EnumMap;
use toml;

use std::collections::HashMap;
use std::ops::{Index, IndexMut};

/// Color configuration for the application.
///
/// Assign each color role an actual color.
///
/// It implements `Index` and `IndexMut` to access and modify this mapping:
///
/// # Example
///
/// ```rust
/// # use cursive::theme::Palette;
/// use cursive::theme::PaletteColor::*;
/// use cursive::theme::Color::*;
/// use cursive::theme::BaseColor::*;
///
/// let mut palette = Palette::default();
///
/// assert_eq!(palette[Background], Dark(Blue));
/// palette[Shadow] = Light(Red);
/// ```
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct Palette {
    basic: EnumMap<PaletteColor, Color>,
    custom: HashMap<String, PaletteNode>,
}

/// A node in the palette tree.
///
/// This can either be a color, or a nested namespace with its own mapping.
#[derive(PartialEq, Eq, Clone, Debug)]
enum PaletteNode {
    Color(Color),
    Namespace(HashMap<String, PaletteNode>),
}

// Basic usage: only use basic colors
impl Index<PaletteColor> for Palette {
    type Output = Color;

    fn index(&self, palette_color: PaletteColor) -> &Color {
        &self.basic[palette_color]
    }
}

// We can alter existing color if needed (but why?...)
impl IndexMut<PaletteColor> for Palette {
    fn index_mut(&mut self, palette_color: PaletteColor) -> &mut Color {
        &mut self.basic[palette_color]
    }
}

impl Palette {
    /// Returns a custom color from this palette.
    ///
    /// Returns `None` if the given key was not found.
    pub fn custom<'a>(&'a self, key: &str) -> Option<&'a Color> {
        self.custom.get(key).and_then(|node| {
            if let &PaletteNode::Color(ref color) = node {
                Some(color)
            } else {
                None
            }
        })
    }

    /// Returns a new palette where the given namespace has been merged.
    ///
    /// All values in the namespace will override previous values.
    pub fn merge(&self, namespace: &str) -> Palette {
        let mut result = self.clone();

        if let Some(&PaletteNode::Namespace(ref palette)) =
            self.custom.get(namespace)
        {
            // Merge `result` and `palette`
            for (key, value) in palette.iter() {
                match *value {
                    PaletteNode::Color(color) => result.set_color(key, color),
                    PaletteNode::Namespace(ref map) => {
                        result.add_namespace(key, map.clone())
                    }
                }
            }
        }

        result
    }

    /// Sets the color for the given key.
    ///
    /// This will update either the basic palette or the custom values.
    pub fn set_color(&mut self, key: &str, color: Color) {
        use theme::PaletteColor::*;

        match key {
            "background" => self.basic[Background] = color,
            "shadow" => self.basic[Shadow] = color,
            "view" => self.basic[View] = color,
            "primary" => self.basic[Primary] = color,
            "secondary" => self.basic[Secondary] = color,
            "tertiary" => self.basic[Tertiary] = color,
            "title_primary" => self.basic[TitlePrimary] = color,
            "title_secondary" => self.basic[TitleSecondary] = color,
            "highlight" => self.basic[Highlight] = color,
            "highlight_inactive" => self.basic[HighlightInactive] = color,
            other => {
                self.custom
                    .insert(other.to_string(), PaletteNode::Color(color));
            }
        }
    }

    /// Adds a color namespace to this palette.
    fn add_namespace(
        &mut self, key: &str, namespace: HashMap<String, PaletteNode>,
    ) {
        self.custom
            .insert(key.to_string(), PaletteNode::Namespace(namespace));
    }
}

/// Returns the default palette for a cursive application.
///
/// * `Background` => `Dark(Blue)`
/// * `Shadow` => `Dark(Black)`
/// * `View` => `Dark(White)`
/// * `Primary` => `Dark(Black)`
/// * `Secondary` => `Dark(Blue)`
/// * `Tertiary` => `Dark(White)`
/// * `TitlePrimary` => `Dark(Red)`
/// * `TitleSecondary` => `Dark(Yellow)`
/// * `Highlight` => `Dark(Red)`
/// * `HighlightInactive` => `Dark(Blue)`
impl Default for Palette {
    fn default() -> Palette {
        use self::PaletteColor::*;
        use theme::BaseColor::*;
        use theme::Color::*;

        Palette {
            basic: enum_map!{
                Background => Dark(Blue),
                Shadow => Dark(Black),
                View => Dark(White),
                Primary => Dark(Black),
                Secondary => Dark(Blue),
                Tertiary => Dark(White),
                TitlePrimary => Dark(Red),
                TitleSecondary => Dark(Yellow),
                Highlight => Dark(Red),
                HighlightInactive => Dark(Blue),
            },
            custom: HashMap::new(),
        }
    }
}

// Iterate over a toml
fn iterate_toml<'a>(
    table: &'a toml::value::Table,
) -> impl Iterator<Item = (&'a str, PaletteNode)> + 'a {
    table.iter().flat_map(|(key, value)| {
        let node = match value {
            toml::Value::Table(table) => {
                // This should define a new namespace
                // Treat basic colors as simple string.
                // We'll convert them back in the merge method.
                let map = iterate_toml(table)
                    .map(|(key, value)| (key.to_string(), value))
                    .collect();
                // Should we only return something if it's non-empty?
                Some(PaletteNode::Namespace(map))
            }
            toml::Value::Array(colors) => {
                // This should be a list of colors - just pick the first valid one.
                colors
                    .iter()
                    .flat_map(toml::Value::as_str)
                    .flat_map(Color::parse)
                    .map(PaletteNode::Color)
                    .next()
            }
            toml::Value::String(color) => {
                // This describe a new color - easy!
                Color::parse(color).map(PaletteNode::Color)
            }
            other => {
                // Other - error?
                debug!(
                    "Found unexpected value in theme: {} = {:?}",
                    key, other
                );
                None
            }
        };

        node.map(|node| (key.as_str(), node))
    })
}

/// Fills `palette` with the colors from the given `table`.
pub(crate) fn load_toml(palette: &mut Palette, table: &toml::value::Table) {
    // TODO: use serde for that?
    // Problem: toml-rs doesn't do well with Enums...

    for (key, value) in iterate_toml(table) {
        match value {
            PaletteNode::Color(color) => palette.set_color(key, color),
            PaletteNode::Namespace(map) => palette.add_namespace(key, map),
        }
    }
}

/// Color entry in a palette.
///
/// Each `ColorRole` is used for a specific role in a default application.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, EnumMap)]
pub enum PaletteColor {
    /// Color used for the application background.
    Background,
    /// Color used for View shadows.
    Shadow,
    /// Color used for View backgrounds.
    View,
    /// Primary color used for the text.
    Primary,
    /// Secondary color used for the text.
    Secondary,
    /// Tertiary color used for the text.
    Tertiary,
    /// Primary color used for title text.
    TitlePrimary,
    /// Secondary color used for title text.
    TitleSecondary,
    /// Color used for highlighting text.
    Highlight,
    /// Color used for highlighting inactive text.
    HighlightInactive,
}

impl PaletteColor {
    /// Given a palette, resolve `self` to a concrete color.
    pub fn resolve(self, palette: &Palette) -> Color {
        palette[self]
    }
}
<commit_msg>Make Palette::add_namespace public<commit_after>use super::Color;
use enum_map::EnumMap;
use toml;

use std::collections::HashMap;
use std::ops::{Index, IndexMut};

/// Color configuration for the application.
///
/// Assign each color role an actual color.
///
/// It implements `Index` and `IndexMut` to access and modify this mapping:
///
/// # Example
///
/// ```rust
/// # use cursive::theme::Palette;
/// use cursive::theme::PaletteColor::*;
/// use cursive::theme::Color::*;
/// use cursive::theme::BaseColor::*;
///
/// let mut palette = Palette::default();
///
/// assert_eq!(palette[Background], Dark(Blue));
/// palette[Shadow] = Light(Red);
/// ```
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct Palette {
    basic: EnumMap<PaletteColor, Color>,
    custom: HashMap<String, PaletteNode>,
}

/// A node in the palette tree.
///
/// This describes a value attached to a custom keyword in the palette.
///
/// This can either be a color, or a nested namespace with its own mapping.
#[derive(PartialEq, Eq, Clone, Debug)]
pub enum PaletteNode {
    /// A single color.
    Color(Color),
    /// A group of values bundled in the same namespace.
    ///
    /// Namespaces can be merged in the palette with `Palette::merge`.
    Namespace(HashMap<String, PaletteNode>),
}

// Basic usage: only use basic colors
impl Index<PaletteColor> for Palette {
    type Output = Color;

    fn index(&self, palette_color: PaletteColor) -> &Color {
        &self.basic[palette_color]
    }
}

// We can alter existing color if needed (but why?...)
impl IndexMut<PaletteColor> for Palette {
    fn index_mut(&mut self, palette_color: PaletteColor) -> &mut Color {
        &mut self.basic[palette_color]
    }
}

impl Palette {
    /// Returns a custom color from this palette.
    ///
    /// Returns `None` if the given key was not found.
    pub fn custom<'a>(&'a self, key: &str) -> Option<&'a Color> {
        self.custom.get(key).and_then(|node| {
            if let &PaletteNode::Color(ref color) = node {
                Some(color)
            } else {
                None
            }
        })
    }

    /// Returns a new palette where the given namespace has been merged.
    ///
    /// All values in the namespace will override previous values.
    pub fn merge(&self, namespace: &str) -> Palette {
        let mut result = self.clone();

        if let Some(&PaletteNode::Namespace(ref palette)) =
            self.custom.get(namespace)
        {
            // Merge `result` and `palette`
            for (key, value) in palette.iter() {
                match *value {
                    PaletteNode::Color(color) => result.set_color(key, color),
                    PaletteNode::Namespace(ref map) => {
                        result.add_namespace(key, map.clone())
                    }
                }
            }
        }

        result
    }

    /// Sets the color for the given key.
    ///
    /// This will update either the basic palette or the custom values.
    pub fn set_color(&mut self, key: &str, color: Color) {
        use theme::PaletteColor::*;

        match key {
            "background" => self.basic[Background] = color,
            "shadow" => self.basic[Shadow] = color,
            "view" => self.basic[View] = color,
            "primary" => self.basic[Primary] = color,
            "secondary" => self.basic[Secondary] = color,
            "tertiary" => self.basic[Tertiary] = color,
            "title_primary" => self.basic[TitlePrimary] = color,
            "title_secondary" => self.basic[TitleSecondary] = color,
            "highlight" => self.basic[Highlight] = color,
            "highlight_inactive" => self.basic[HighlightInactive] = color,
            other => {
                self.custom
                    .insert(other.to_string(), PaletteNode::Color(color));
            }
        }
    }

    /// Adds a color namespace to this palette.
    pub fn add_namespace(
        &mut self, key: &str, namespace: HashMap<String, PaletteNode>,
    ) {
        self.custom
            .insert(key.to_string(), PaletteNode::Namespace(namespace));
    }
}

/// Returns the default palette for a cursive application.
///
/// * `Background` => `Dark(Blue)`
/// * `Shadow` => `Dark(Black)`
/// * `View` => `Dark(White)`
/// * `Primary` => `Dark(Black)`
/// * `Secondary` => `Dark(Blue)`
/// * `Tertiary` => `Dark(White)`
/// * `TitlePrimary` => `Dark(Red)`
/// * `TitleSecondary` => `Dark(Yellow)`
/// * `Highlight` => `Dark(Red)`
/// * `HighlightInactive` => `Dark(Blue)`
impl Default for Palette {
    fn default() -> Palette {
        use self::PaletteColor::*;
        use theme::BaseColor::*;
        use theme::Color::*;

        Palette {
            basic: enum_map!{
                Background => Dark(Blue),
                Shadow => Dark(Black),
                View => Dark(White),
                Primary => Dark(Black),
                Secondary => Dark(Blue),
                Tertiary => Dark(White),
                TitlePrimary => Dark(Red),
                TitleSecondary => Dark(Yellow),
                Highlight => Dark(Red),
                HighlightInactive => Dark(Blue),
            },
            custom: HashMap::new(),
        }
    }
}

// Iterate over a toml
fn iterate_toml<'a>(
    table: &'a toml::value::Table,
) -> impl Iterator<Item = (&'a str, PaletteNode)> + 'a {
    table.iter().flat_map(|(key, value)| {
        let node = match value {
            toml::Value::Table(table) => {
                // This should define a new namespace
                // Treat basic colors as simple string.
                // We'll convert them back in the merge method.
                let map = iterate_toml(table)
                    .map(|(key, value)| (key.to_string(), value))
                    .collect();
                // Should we only return something if it's non-empty?
                Some(PaletteNode::Namespace(map))
            }
            toml::Value::Array(colors) => {
                // This should be a list of colors - just pick the first valid one.
                colors
                    .iter()
                    .flat_map(toml::Value::as_str)
                    .flat_map(Color::parse)
                    .map(PaletteNode::Color)
                    .next()
            }
            toml::Value::String(color) => {
                // This describe a new color - easy!
                Color::parse(color).map(PaletteNode::Color)
            }
            other => {
                // Other - error?
                debug!(
                    "Found unexpected value in theme: {} = {:?}",
                    key, other
                );
                None
            }
        };

        node.map(|node| (key.as_str(), node))
    })
}

/// Fills `palette` with the colors from the given `table`.
pub(crate) fn load_toml(palette: &mut Palette, table: &toml::value::Table) {
    // TODO: use serde for that?
    // Problem: toml-rs doesn't do well with Enums...

    for (key, value) in iterate_toml(table) {
        match value {
            PaletteNode::Color(color) => palette.set_color(key, color),
            PaletteNode::Namespace(map) => palette.add_namespace(key, map),
        }
    }
}

/// Color entry in a palette.
///
/// Each `ColorRole` is used for a specific role in a default application.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, EnumMap)]
pub enum PaletteColor {
    /// Color used for the application background.
    Background,
    /// Color used for View shadows.
    Shadow,
    /// Color used for View backgrounds.
    View,
    /// Primary color used for the text.
    Primary,
    /// Secondary color used for the text.
    Secondary,
    /// Tertiary color used for the text.
    Tertiary,
    /// Primary color used for title text.
    TitlePrimary,
    /// Secondary color used for title text.
    TitleSecondary,
    /// Color used for highlighting text.
    Highlight,
    /// Color used for highlighting inactive text.
    HighlightInactive,
}

impl PaletteColor {
    /// Given a palette, resolve `self` to a concrete color.
    pub fn resolve(self, palette: &Palette) -> Color {
        palette[self]
    }
}
<|endoftext|>


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

use typed_arena::Arena;

// 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, Clone)]
struct Node {
    weight: usize,
    leaf_counts: Vec<usize>,
}

impl Node {
    fn new(weight: usize, initial_count: usize, 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 {
    weight: usize,
    count: usize,
}
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, Clone)]
struct List {
    lookahead0: Node,
    lookahead1: Node,
    next_leaf_index: usize,
}

/// Calculates the bitlengths for the Huffman tree, based on the counts of each
/// symbol.
pub fn length_limited_code_lengths(frequencies: &[usize], max_bits: usize) -> Vec<u32> {
    let num_freqs = frequencies.len();
    let mut bit_lengths = vec![0; num_freqs];

    // Count used symbols and place them in the leaves.
    let mut leaves: Vec<_> = frequencies.iter()
        .enumerate()
        .filter(|&(_, &freq)| freq != 0)
        .map(|(i, &freq)| Leaf { weight: freq, count: i })
        .collect();

    let num_symbols = leaves.len();

    // Short circuit some special cases

    // TODO:
    // if ((1 << maxbits) < numsymbols) {
    //   free(leaves);
    //   return 1;  /* Error, too few maxbits to represent symbols. */
    // }

    if num_symbols == 0 {
        // There are no non-zero frequencies.
        return bit_lengths;
    }
    if num_symbols == 1 {
        bit_lengths[leaves[0].count] = 1;
        return bit_lengths;
    }
    if num_symbols == 2 {
        bit_lengths[leaves[0].count] = 1;
        bit_lengths[leaves[1].count] = 1;
        return bit_lengths;
    }

    // Sort the leaves from least frequent to most frequent.
    leaves.sort();

    let max_bits = if num_symbols - 1 < max_bits {
        num_symbols - 1
    } else {
        max_bits
    };

    let arena_capacity = max_bits * 2 * num_symbols;
    let node_arena: Arena<Node> = Arena::with_capacity(arena_capacity);

    // let max_num_leaves = 2 * num_symbols - 2;
    // let mut lists = vec![
    //     List {
    //         lookahead0: Node::new(leaves[0].weight, 1, max_num_leaves),
    //         lookahead1: Node::new(leaves[1].weight, 2, max_num_leaves),
    //         next_leaf_index: 2,
    //     };
    //     max_bits
    // ];
    //
    // // 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 mut a = lists.pop().unwrap().lookahead1.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 leaf in &leaves[next_count..leaf_count] {
    //         bit_lengths[leaf.count] = bitlength_value;
    //     }
    //     bitlength_value += 1;
    // }

    bit_lengths
}

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 = &mut lists[0];
    let next_leaf = &leaves[current_list.next_leaf_index];
    current_list.lookahead1.weight = next_leaf.weight;

    current_list.lookahead1.leaf_counts[0] = current_list.lookahead0.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 = &mut lists[current_list_index];

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

fn next_tree(weight_sum: usize, 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.lookahead1.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.lookahead1.weight = weight_sum;
        current_list.lookahead1.leaf_counts.clear();

        current_list.lookahead1.leaf_counts.extend_from_slice(previous_list_leaf_counts);
        current_list.lookahead1.leaf_counts.push(*current_list.lookahead0.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;
    let num_symbols = leaves.len();

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

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

    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[current_list_index - 1];
            previous_list.lookahead0.weight + previous_list.lookahead1.weight
        };

        if next_leaf_index < num_symbols && 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 max_bits_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 max_bits_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>Initialize lists with immutable refs to arena Nodes [tests fail]<commit_after>use std::cmp::Ordering;
// use std::{mem};

use typed_arena::Arena;

// 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<'a> {
    weight: usize,
    count: usize,
    tail: Option<&'a mut Node<'a>>,
}

#[derive(Debug)]
struct Leaf {
    weight: usize,
    count: usize,
}
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, Clone)]
struct List<'a> {
    lookahead0: &'a Node<'a>,
    lookahead1: &'a Node<'a>,
}

/// Calculates the bitlengths for the Huffman tree, based on the counts of each
/// symbol.
pub fn length_limited_code_lengths(frequencies: &[usize], max_bits: usize) -> Vec<u32> {
    let num_freqs = frequencies.len();
    let mut bit_lengths = vec![0; num_freqs];

    // Count used symbols and place them in the leaves.
    let mut leaves: Vec<_> = frequencies.iter()
        .enumerate()
        .filter(|&(_, &freq)| freq != 0)
        .map(|(i, &freq)| Leaf { weight: freq, count: i })
        .collect();

    let num_symbols = leaves.len();

    // Short circuit some special cases

    // TODO:
    // if ((1 << maxbits) < numsymbols) {
    //   free(leaves);
    //   return 1;  /* Error, too few maxbits to represent symbols. */
    // }

    if num_symbols == 0 {
        // There are no non-zero frequencies.
        return bit_lengths;
    }
    if num_symbols == 1 {
        bit_lengths[leaves[0].count] = 1;
        return bit_lengths;
    }
    if num_symbols == 2 {
        bit_lengths[leaves[0].count] = 1;
        bit_lengths[leaves[1].count] = 1;
        return bit_lengths;
    }

    // Sort the leaves from least frequent to most frequent.
    leaves.sort();

    let max_bits = if num_symbols - 1 < max_bits {
        num_symbols - 1
    } else {
        max_bits
    };

    let arena_capacity = max_bits * 2 * num_symbols;
    let node_arena: Arena<Node> = Arena::with_capacity(arena_capacity);

    let node0 = node_arena.alloc(Node {
        weight: leaves[0].weight,
        count: 1,
        tail: None,
    });

    let node1 = node_arena.alloc(Node {
        weight: leaves[1].weight,
        count: 2,
        tail: None,
    });

    let mut lists: Vec<List> = vec![
        List {
            lookahead0: node0,
            lookahead1: node1,
        };
        max_bits
    ];

    // let max_num_leaves = 2 * num_symbols - 2;
    // let mut lists = vec![
    //     List {
    //         lookahead0: Node::new(leaves[0].weight, 1, max_num_leaves),
    //         lookahead1: Node::new(leaves[1].weight, 2, max_num_leaves),
    //         next_leaf_index: 2,
    //     };
    //     max_bits
    // ];
    //
    // // 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 mut a = lists.pop().unwrap().lookahead1.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 leaf in &leaves[next_count..leaf_count] {
    //         bit_lengths[leaf.count] = bitlength_value;
    //     }
    //     bitlength_value += 1;
    // }

    bit_lengths
}

// 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 = &mut lists[0];
//     let next_leaf = &leaves[current_list.next_leaf_index];
//     current_list.lookahead1.weight = next_leaf.weight;
//
//     current_list.lookahead1.leaf_counts[0] = current_list.lookahead0.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 = &mut lists[current_list_index];
//
//     // The next leaf goes next; counting itself makes the leaf_count increase by one.
//     current_list.lookahead1.weight = leaves[current_list.next_leaf_index].weight;
//     current_list.lookahead1.leaf_counts.clear();
//     current_list.lookahead1.leaf_counts.extend_from_slice(&current_list.lookahead0.leaf_counts);
//     let last_index = current_list.lookahead1.leaf_counts.len() - 1;
//     current_list.lookahead1.leaf_counts[last_index] += 1;
//     current_list.next_leaf_index += 1;
// }

// fn next_tree(weight_sum: usize, 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.lookahead1.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.lookahead1.weight = weight_sum;
//         current_list.lookahead1.leaf_counts.clear();
//
//         current_list.lookahead1.leaf_counts.extend_from_slice(previous_list_leaf_counts);
//         current_list.lookahead1.leaf_counts.push(*current_list.lookahead0.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;
//     let num_symbols = leaves.len();
//
//     if current_list_index == 0 && next_leaf_index == num_symbols {
//         // We've added all the leaves to the lowest list, so we're done here
//         return;
//     }
//
//     mem::swap(&mut lists[current_list_index].lookahead0, &mut lists[current_list_index].lookahead1);
//
//     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[current_list_index - 1];
//             previous_list.lookahead0.weight + previous_list.lookahead1.weight
//         };
//
//         if next_leaf_index < num_symbols && 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 max_bits_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 max_bits_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>pub mod test_common;
pub mod distance;

mod partition;
mod bounds;

use self::bounds::*;
use self::distance::*;

use std::cmp;

pub trait KdtreePointTrait: Copy + PartialEq {
    fn dims(&self) -> &[f64];
}

pub struct Kdtree<KdtreePoint> {
    nodes: Vec<KdtreeNode<KdtreePoint>>,

    node_adding_dimension: usize,
    node_depth_during_last_rebuild: usize,
    current_node_depth: usize,
}

impl<KdtreePoint: KdtreePointTrait> Kdtree<KdtreePoint> {
    pub fn new(mut points: &mut [KdtreePoint]) -> Kdtree<KdtreePoint> {
        if points.len() == 0 {
            panic!("empty vector point not allowed");
        }

        let mut tree = Kdtree {
            nodes: vec![],
            node_adding_dimension: 0,
            node_depth_during_last_rebuild: 0,
            current_node_depth: 0,
        };

        tree.rebuild_tree(&mut points);

        tree
    }

    pub fn rebuild_tree(&mut self, points : &mut [KdtreePoint]) {
        self.nodes.clear();

        self.node_depth_during_last_rebuild = 0;
        self.current_node_depth = 0;

        let rect = Bounds::new_from_points(points);
        self.build_tree(points, &rect, 1);
    }

    /// Can be used if you are sure that the tree is degenerated or if you will never again insert the nodes into the tree.
    pub fn gather_points_and_rebuild(&mut self) {
        let mut points : Vec<KdtreePoint> = vec![];
        self.gather_points(0,&mut points);

        self.rebuild_tree(&mut points);
    }

    pub fn nearest_search(&self, node: &KdtreePoint) -> KdtreePoint
    {
        let mut nearest_neighbor = 0usize;
        let mut best_distance = squared_euclidean(node.dims(), &self.nodes[0].point.dims());
        self.nearest_search_impl(node, 0usize, &mut best_distance, &mut nearest_neighbor);

        self.nodes[nearest_neighbor].point
    }

    pub fn has_neighbor_in_range(&self, node: &KdtreePoint, range: f64) -> bool {
        let squared_range = range * range;

        self.distance_squared_to_nearest(node) <= squared_range
    }

    pub fn distance_squared_to_nearest(&self, node: &KdtreePoint) -> f64 {
        squared_euclidean(&self.nearest_search(node).dims(), node.dims())
    }

    pub fn insert_node(&mut self, node_to_add : KdtreePoint) {

        let mut current_index = 0;
        let dimension = self.node_adding_dimension;
        let index_of_new_node = self.add_node(node_to_add,dimension,node_to_add.dims()[dimension]);
        self.node_adding_dimension = ( dimension + 1) % node_to_add.dims().len();
        let mut should_pop_node = false;

        let mut depth = 0;
        loop {

            depth +=1 ;
            let current_node = &mut self.nodes[current_index];

            if node_to_add.dims()[current_node.dimension] <= current_node.split_on {
                if let Some(left_node_index) = current_node.left_node {
                    current_index = left_node_index
                } else {
                    if current_node.point.eq(&node_to_add) {
                        should_pop_node = true;
                    } else {
                        current_node.left_node = Some(index_of_new_node);
                    }
                    break;
                }
            } else {
                if let Some(right_node_index) = current_node.right_node {
                    current_index = right_node_index
                } else {
                    if current_node.point.eq(&node_to_add) {
                        should_pop_node = true;
                    } else {
                        current_node.right_node = Some(index_of_new_node);
                    }
                    break;
                }
            }
        }

        if should_pop_node {
            self.nodes.pop();
        }

        if self.node_depth_during_last_rebuild as f64 * 4.0 < depth as f64  {
            self.gather_points_and_rebuild();
        }
    }

    fn nearest_search_impl(&self, p: &KdtreePoint, searched_index: usize, best_distance_squared: &mut f64, best_leaf_found: &mut usize) {
        let node = &self.nodes[searched_index];

        let splitting_value = node.split_on;
        let point_splitting_dim_value = p.dims()[node.dimension];

        let (closer_node, farther_node) = if point_splitting_dim_value <= splitting_value {
            (node.left_node, node.right_node)
        } else {
            (node.right_node, node.left_node)
        };

        if let Some(closer_node) = closer_node {
            self.nearest_search_impl(p, closer_node, best_distance_squared, best_leaf_found);
        }

        let distance = squared_euclidean(p.dims(), node.point.dims());
        if distance < *best_distance_squared {
            *best_distance_squared = distance;
            *best_leaf_found = searched_index;
        }

        if let Some(farther_node) = farther_node {
            let distance_on_single_dimension = squared_euclidean(&[splitting_value], &[point_splitting_dim_value]);

            if distance_on_single_dimension <= *best_distance_squared {
                self.nearest_search_impl(p, farther_node, best_distance_squared, best_leaf_found);
            }
        }
    }

    fn add_node(&mut self, p: KdtreePoint, dimension: usize, split_on: f64) -> usize {
        let node = KdtreeNode::new(p, dimension, split_on);

        self.nodes.push(node);
        self.nodes.len() - 1
    }

    fn build_tree(&mut self, nodes: &mut [KdtreePoint], bounds: &Bounds, depth : usize) -> usize {
        let splitting_index = partition::partition_sliding_midpoint(nodes, bounds.get_midvalue_of_widest_dim(), bounds.get_widest_dim());
        let pivot_value = nodes[splitting_index].dims()[bounds.get_widest_dim()];

        let node_id = self.add_node(nodes[splitting_index], bounds.get_widest_dim(), pivot_value);
        let nodes_len = nodes.len();

        if splitting_index > 0 {
            let left_rect = bounds.clone_moving_max(pivot_value, bounds.get_widest_dim());
            let left_child_id = self.build_tree(&mut nodes[0..splitting_index], &left_rect, depth+1);
            self.nodes[node_id].left_node = Some(left_child_id);
        }

        if splitting_index < nodes.len() - 1 {
            let right_rect = bounds.clone_moving_min(pivot_value, bounds.get_widest_dim());

            let right_child_id = self.build_tree(&mut nodes[splitting_index + 1..nodes_len], &right_rect, depth+1);
            self.nodes[node_id].right_node = Some(right_child_id);
        }

        self.node_depth_during_last_rebuild =  cmp::max(self.node_depth_during_last_rebuild,depth);

        node_id
    }

    fn gather_points(&self, current_index: usize, points : &mut Vec<KdtreePoint>){
        points.push(self.nodes[current_index].point);
        if let Some(left_index) = self.nodes[current_index].left_node {
            self.gather_points(left_index, points);
        }

        if let Some(right_index) = self.nodes[current_index].right_node {
            self.gather_points(right_index, points);
        }
    }
}

pub struct KdtreeNode<T> {
    left_node: Option<usize>,
    right_node: Option<usize>,

    point: T,
    dimension: usize,
    split_on: f64
}

impl<T: KdtreePointTrait> KdtreeNode<T> {
    fn new(p: T, splitting_dimension: usize, split_on_value: f64) -> KdtreeNode<T> {
        KdtreeNode {
            left_node: None,
            right_node: None,

            point: p,
            dimension: splitting_dimension,
            split_on: split_on_value
        }
    }
}

#[cfg(test)]
mod tests {
    use ::kdtree::test_common::Point2WithId;

    use super::*;

    #[test]
    #[should_panic(expected = "empty vector point not allowed")]
    fn should_panic_given_empty_vector() {
        let mut empty_vec: Vec<Point2WithId> = vec![];

        Kdtree::new(&mut empty_vec);
    }

    quickcheck! {
        fn tree_build_creates_tree_with_as_many_leafs_as_there_is_points(xs : Vec<f64>) -> bool {
            if xs.len() == 0 {
                return true;
            }
            let mut vec : Vec<Point2WithId> = vec![];
            for i in 0 .. xs.len() {
                let p = Point2WithId::new(i as i32, xs[i], xs[i]);

                vec.push(p);
            }

            let tree = Kdtree::new(&mut qc_value_vec_to_2d_points_vec(&xs));

            let mut to_iterate : Vec<usize> = vec![];
            to_iterate.push(0);

            while to_iterate.len() > 0 {
                let last_index = to_iterate.last().unwrap().clone();
                let ref x = tree.nodes.get(last_index).unwrap();
                to_iterate.pop();
                if x.left_node.is_some() {
                    to_iterate.push(x.left_node.unwrap());
                }
                if x.right_node.is_some() {
                    to_iterate.push(x.right_node.unwrap());
                }
            }
            xs.len() == tree.nodes.len()
        }
    }

    quickcheck! {
        fn nearest_neighbor_search_using_qc(xs : Vec<f64>) -> bool {
            if xs.len() == 0 {
                return true;
            }

            let point_vec = qc_value_vec_to_2d_points_vec(&xs);
            let tree = Kdtree::new(&mut point_vec.clone());

            for p in &point_vec {
                let found_nn = tree.nearest_search(p);

                assert_eq!(p.id,found_nn.id);
            }

            true
        }
    }

    #[test]
    fn has_neighbor_in_range() {
        let mut vec: Vec<Point2WithId> = vec![Point2WithId::new(0,2.,0.)];

        let tree = Kdtree::new(&mut vec);

        assert_eq!(false,tree.has_neighbor_in_range(&Point2WithId::new(0,0.,0.), 0.));
        assert_eq!(false,tree.has_neighbor_in_range(&Point2WithId::new(0,0.,0.), 1.));
        assert_eq!(true,tree.has_neighbor_in_range(&Point2WithId::new(0,0.,0.), 2.));
        assert_eq!(true,tree.has_neighbor_in_range(&Point2WithId::new(0,0.,0.), 300.));
    }

    #[test]
    fn incremental_add_adds_as_expected() {
        //this test is tricky because it can have problems with the automatic tree rebuild.

        let mut vec = vec![Point2WithId::new(0,0.,0.)];

        let mut tree = Kdtree::new(&mut vec);

        tree.insert_node(Point2WithId::new(0,1.,0.));
        tree.insert_node(Point2WithId::new(0,-1.,0.));

        assert_eq!(tree.nodes.len(), 3);
        assert_eq!(tree.nodes[0].dimension, 0);

        assert_eq!(tree.nodes[0].left_node.is_some(), true);
        assert_eq!(tree.nodes[1].point.dims()[0], 1.);
        assert_eq!(tree.nodes[2].point.dims()[0], -1.);

        assert_eq!(tree.nodes[0].right_node.is_some(), true);
    }

    #[test]
    fn incremental_add_filters_duplicates() {
        let mut vec = vec![Point2WithId::new(0,0.,0.)];

        let mut tree = Kdtree::new(&mut vec);

        let node = Point2WithId::new(0,1.,0.);
        tree.insert_node(node);
        tree.insert_node(node);

        assert_eq!(tree.nodes.len(), 2);
    }

    fn qc_value_vec_to_2d_points_vec(xs: &Vec<f64>) -> Vec<Point2WithId> {
        let mut vec: Vec<Point2WithId> = vec![];
        for i in 0..xs.len() {
            let mut is_duplicated_value = false;
            for j in 0..i {
                if xs[i] == xs[j] {
                    is_duplicated_value = true;
                    break;
                }
            }
            if !is_duplicated_value {
                let p = Point2WithId::new(i as i32, xs[i], xs[i]);
                vec.push(p);
            }
        }

        vec
    }
}<commit_msg>added insert_nodes_and_rebuild<commit_after>pub mod test_common;
pub mod distance;

mod partition;
mod bounds;

use self::bounds::*;
use self::distance::*;

use std::cmp;

pub trait KdtreePointTrait: Copy + PartialEq {
    fn dims(&self) -> &[f64];
}

pub struct Kdtree<KdtreePoint> {
    nodes: Vec<KdtreeNode<KdtreePoint>>,

    node_adding_dimension: usize,
    node_depth_during_last_rebuild: usize,
    current_node_depth: usize,
}

impl<KdtreePoint: KdtreePointTrait> Kdtree<KdtreePoint> {
    pub fn new(mut points: &mut [KdtreePoint]) -> Kdtree<KdtreePoint> {
        if points.len() == 0 {
            panic!("empty vector point not allowed");
        }

        let mut tree = Kdtree {
            nodes: vec![],
            node_adding_dimension: 0,
            node_depth_during_last_rebuild: 0,
            current_node_depth: 0,
        };

        tree.rebuild_tree(&mut points);

        tree
    }

    pub fn rebuild_tree(&mut self, points : &mut [KdtreePoint]) {
        self.nodes.clear();

        self.node_depth_during_last_rebuild = 0;
        self.current_node_depth = 0;

        let rect = Bounds::new_from_points(points);
        self.build_tree(points, &rect, 1);
    }

    /// Can be used if you are sure that the tree is degenerated or if you will never again insert the nodes into the tree.
    pub fn gather_points_and_rebuild(&mut self) {
        let mut points : Vec<KdtreePoint> = vec![];
        self.gather_points(0,&mut points);

        self.rebuild_tree(&mut points);
    }

    pub fn nearest_search(&self, node: &KdtreePoint) -> KdtreePoint
    {
        let mut nearest_neighbor = 0usize;
        let mut best_distance = squared_euclidean(node.dims(), &self.nodes[0].point.dims());
        self.nearest_search_impl(node, 0usize, &mut best_distance, &mut nearest_neighbor);

        self.nodes[nearest_neighbor].point
    }

    pub fn has_neighbor_in_range(&self, node: &KdtreePoint, range: f64) -> bool {
        let squared_range = range * range;

        self.distance_squared_to_nearest(node) <= squared_range
    }

    pub fn distance_squared_to_nearest(&self, node: &KdtreePoint) -> f64 {
        squared_euclidean(&self.nearest_search(node).dims(), node.dims())
    }

    pub fn insert_nodes_and_rebuild(&mut self, nodes_to_add : &mut [KdtreePoint]) {
        let mut pts : Vec<KdtreePoint> = vec![];
        self.gather_points(0, &mut pts);
        pts.extend(nodes_to_add.iter());

        self.rebuild_tree(&mut pts);
    }

    pub fn insert_node(&mut self, node_to_add : KdtreePoint) {

        let mut current_index = 0;
        let dimension = self.node_adding_dimension;
        let index_of_new_node = self.add_node(node_to_add,dimension,node_to_add.dims()[dimension]);
        self.node_adding_dimension = ( dimension + 1) % node_to_add.dims().len();
        let mut should_pop_node = false;

        let mut depth = 0;
        loop {

            depth +=1 ;
            let current_node = &mut self.nodes[current_index];

            if node_to_add.dims()[current_node.dimension] <= current_node.split_on {
                if let Some(left_node_index) = current_node.left_node {
                    current_index = left_node_index
                } else {
                    if current_node.point.eq(&node_to_add) {
                        should_pop_node = true;
                    } else {
                        current_node.left_node = Some(index_of_new_node);
                    }
                    break;
                }
            } else {
                if let Some(right_node_index) = current_node.right_node {
                    current_index = right_node_index
                } else {
                    if current_node.point.eq(&node_to_add) {
                        should_pop_node = true;
                    } else {
                        current_node.right_node = Some(index_of_new_node);
                    }
                    break;
                }
            }
        }

        if should_pop_node {
            self.nodes.pop();
        }

        if self.node_depth_during_last_rebuild as f64 * 4.0 < depth as f64  {
            self.gather_points_and_rebuild();
        }
    }

    fn nearest_search_impl(&self, p: &KdtreePoint, searched_index: usize, best_distance_squared: &mut f64, best_leaf_found: &mut usize) {
        let node = &self.nodes[searched_index];

        let splitting_value = node.split_on;
        let point_splitting_dim_value = p.dims()[node.dimension];

        let (closer_node, farther_node) = if point_splitting_dim_value <= splitting_value {
            (node.left_node, node.right_node)
        } else {
            (node.right_node, node.left_node)
        };

        if let Some(closer_node) = closer_node {
            self.nearest_search_impl(p, closer_node, best_distance_squared, best_leaf_found);
        }

        let distance = squared_euclidean(p.dims(), node.point.dims());
        if distance < *best_distance_squared {
            *best_distance_squared = distance;
            *best_leaf_found = searched_index;
        }

        if let Some(farther_node) = farther_node {
            let distance_on_single_dimension = squared_euclidean(&[splitting_value], &[point_splitting_dim_value]);

            if distance_on_single_dimension <= *best_distance_squared {
                self.nearest_search_impl(p, farther_node, best_distance_squared, best_leaf_found);
            }
        }
    }

    fn add_node(&mut self, p: KdtreePoint, dimension: usize, split_on: f64) -> usize {
        let node = KdtreeNode::new(p, dimension, split_on);

        self.nodes.push(node);
        self.nodes.len() - 1
    }

    fn build_tree(&mut self, nodes: &mut [KdtreePoint], bounds: &Bounds, depth : usize) -> usize {
        let splitting_index = partition::partition_sliding_midpoint(nodes, bounds.get_midvalue_of_widest_dim(), bounds.get_widest_dim());
        let pivot_value = nodes[splitting_index].dims()[bounds.get_widest_dim()];

        let node_id = self.add_node(nodes[splitting_index], bounds.get_widest_dim(), pivot_value);
        let nodes_len = nodes.len();

        if splitting_index > 0 {
            let left_rect = bounds.clone_moving_max(pivot_value, bounds.get_widest_dim());
            let left_child_id = self.build_tree(&mut nodes[0..splitting_index], &left_rect, depth+1);
            self.nodes[node_id].left_node = Some(left_child_id);
        }

        if splitting_index < nodes.len() - 1 {
            let right_rect = bounds.clone_moving_min(pivot_value, bounds.get_widest_dim());

            let right_child_id = self.build_tree(&mut nodes[splitting_index + 1..nodes_len], &right_rect, depth+1);
            self.nodes[node_id].right_node = Some(right_child_id);
        }

        self.node_depth_during_last_rebuild =  cmp::max(self.node_depth_during_last_rebuild,depth);

        node_id
    }

    fn gather_points(&self, current_index: usize, points : &mut Vec<KdtreePoint>){
        points.push(self.nodes[current_index].point);
        if let Some(left_index) = self.nodes[current_index].left_node {
            self.gather_points(left_index, points);
        }

        if let Some(right_index) = self.nodes[current_index].right_node {
            self.gather_points(right_index, points);
        }
    }
}

pub struct KdtreeNode<T> {
    left_node: Option<usize>,
    right_node: Option<usize>,

    point: T,
    dimension: usize,
    split_on: f64
}

impl<T: KdtreePointTrait> KdtreeNode<T> {
    fn new(p: T, splitting_dimension: usize, split_on_value: f64) -> KdtreeNode<T> {
        KdtreeNode {
            left_node: None,
            right_node: None,

            point: p,
            dimension: splitting_dimension,
            split_on: split_on_value
        }
    }
}

#[cfg(test)]
mod tests {
    use ::kdtree::test_common::Point2WithId;

    use super::*;

    #[test]
    #[should_panic(expected = "empty vector point not allowed")]
    fn should_panic_given_empty_vector() {
        let mut empty_vec: Vec<Point2WithId> = vec![];

        Kdtree::new(&mut empty_vec);
    }

    quickcheck! {
        fn tree_build_creates_tree_with_as_many_leafs_as_there_is_points(xs : Vec<f64>) -> bool {
            if xs.len() == 0 {
                return true;
            }
            let mut vec : Vec<Point2WithId> = vec![];
            for i in 0 .. xs.len() {
                let p = Point2WithId::new(i as i32, xs[i], xs[i]);

                vec.push(p);
            }

            let tree = Kdtree::new(&mut qc_value_vec_to_2d_points_vec(&xs));

            let mut to_iterate : Vec<usize> = vec![];
            to_iterate.push(0);

            while to_iterate.len() > 0 {
                let last_index = to_iterate.last().unwrap().clone();
                let ref x = tree.nodes.get(last_index).unwrap();
                to_iterate.pop();
                if x.left_node.is_some() {
                    to_iterate.push(x.left_node.unwrap());
                }
                if x.right_node.is_some() {
                    to_iterate.push(x.right_node.unwrap());
                }
            }
            xs.len() == tree.nodes.len()
        }
    }

    quickcheck! {
        fn nearest_neighbor_search_using_qc(xs : Vec<f64>) -> bool {
            if xs.len() == 0 {
                return true;
            }

            let point_vec = qc_value_vec_to_2d_points_vec(&xs);
            let tree = Kdtree::new(&mut point_vec.clone());

            for p in &point_vec {
                let found_nn = tree.nearest_search(p);

                assert_eq!(p.id,found_nn.id);
            }

            true
        }
    }

    #[test]
    fn has_neighbor_in_range() {
        let mut vec: Vec<Point2WithId> = vec![Point2WithId::new(0,2.,0.)];

        let tree = Kdtree::new(&mut vec);

        assert_eq!(false,tree.has_neighbor_in_range(&Point2WithId::new(0,0.,0.), 0.));
        assert_eq!(false,tree.has_neighbor_in_range(&Point2WithId::new(0,0.,0.), 1.));
        assert_eq!(true,tree.has_neighbor_in_range(&Point2WithId::new(0,0.,0.), 2.));
        assert_eq!(true,tree.has_neighbor_in_range(&Point2WithId::new(0,0.,0.), 300.));
    }

    #[test]
    fn incremental_add_adds_as_expected() {
        //this test is tricky because it can have problems with the automatic tree rebuild.

        let mut vec = vec![Point2WithId::new(0,0.,0.)];

        let mut tree = Kdtree::new(&mut vec);

        tree.insert_node(Point2WithId::new(0,1.,0.));
        tree.insert_node(Point2WithId::new(0,-1.,0.));

        assert_eq!(tree.nodes.len(), 3);
        assert_eq!(tree.nodes[0].dimension, 0);

        assert_eq!(tree.nodes[0].left_node.is_some(), true);
        assert_eq!(tree.nodes[1].point.dims()[0], 1.);
        assert_eq!(tree.nodes[2].point.dims()[0], -1.);

        assert_eq!(tree.nodes[0].right_node.is_some(), true);
    }

    #[test]
    fn incremental_add_filters_duplicates() {
        let mut vec = vec![Point2WithId::new(0,0.,0.)];

        let mut tree = Kdtree::new(&mut vec);

        let node = Point2WithId::new(0,1.,0.);
        tree.insert_node(node);
        tree.insert_node(node);

        assert_eq!(tree.nodes.len(), 2);
    }

    fn qc_value_vec_to_2d_points_vec(xs: &Vec<f64>) -> Vec<Point2WithId> {
        let mut vec: Vec<Point2WithId> = vec![];
        for i in 0..xs.len() {
            let mut is_duplicated_value = false;
            for j in 0..i {
                if xs[i] == xs[j] {
                    is_duplicated_value = true;
                    break;
                }
            }
            if !is_duplicated_value {
                let p = Point2WithId::new(i as i32, xs[i], xs[i]);
                vec.push(p);
            }
        }

        vec
    }
}<|endoftext|>


<commit_before>use std::io;
use std::fmt;
use std::net::SocketAddr;
use std::collections::{VecDeque, HashMap};
use std::sync::mpsc::SendError;
use fibers::{Spawn, BoxSpawn};
use fibers::net::{TcpStream, TcpListener};
use fibers::net::futures::{TcpListenerBind, Connected};
use fibers::net::streams::Incoming;
use fibers::sync::mpsc;
use fibers::sync::oneshot::Link;
use futures::{Future, BoxFuture, Async, Poll, Stream, Sink, AsyncSink, StartSend};
use handy_async::io::{ReadFrom, AsyncWrite};
use handy_async::io::futures::WriteAll;
use handy_async::sync_io::ReadExt;
use handy_async::pattern::{Pattern, Window};
use trackable::error::ErrorKindExt;

use {Result, Error, ErrorKind};
use message::RawMessage;
use super::{MessageStream, MessageSink, MessageSinkItem, Transport};

#[derive(Debug)]
enum OutgoingCommand {
    Send(RawMessage, Option<Link<(), Error, (), ()>>),
}

#[derive(Debug)]
enum IncomingCommand {
    Recv(SocketAddr, Result<RawMessage>),
    Exit(SocketAddr, Result<()>),
}

#[derive(Debug)]
enum Listener {
    Bind(TcpListenerBind),
    Incoming(Incoming),
}
impl Stream for Listener {
    type Item = (Connected, SocketAddr);
    type Error = Error;
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        loop {
            let next = match *self {
                Listener::Bind(ref mut future) => {
                    if let Async::Ready(listener) = track_try!(future.poll()) {
                        Listener::Incoming(listener.incoming())
                    } else {
                        return Ok(Async::NotReady);
                    }
                }
                Listener::Incoming(ref mut stream) => return track_err!(stream.poll()),
            };
            *self = next;
        }
    }
}

/// TCP based server-side implementation of [Transport](trait.Transport.html) trait.
#[derive(Debug)]
pub struct TcpServerTransport {
    spawner: BoxSpawn,
    listener: Listener,
    clients: HashMap<SocketAddr, mpsc::Sender<OutgoingCommand>>,
    incoming_tx: mpsc::Sender<IncomingCommand>,
    incoming_rx: mpsc::Receiver<IncomingCommand>,
}
impl TcpServerTransport {
    /// Makes a new `TcpServerTransport` instance.
    pub fn new<S>(spawner: S, bind_addr: SocketAddr) -> Self
        where S: Spawn + Send + 'static
    {
        let (incoming_tx, incoming_rx) = mpsc::channel();
        TcpServerTransport {
            spawner: spawner.boxed(),
            listener: Listener::Bind(TcpListener::bind(bind_addr)),
            clients: HashMap::new(),
            incoming_tx: incoming_tx,
            incoming_rx: incoming_rx,
        }
    }
    fn handle_new_connection(&mut self, client: SocketAddr, connected: Connected) {
        let incoming_tx0 = self.incoming_tx.clone();
        let incoming_tx1 = self.incoming_tx.clone();
        let (outgoing_tx, outgoing_rx) = mpsc::channel();
        self.spawner
            .spawn(track_err!(connected)
                       .and_then(move |stream| {
                                     TcpHandleClientLoop::new(client,
                                                              stream,
                                                              incoming_tx0,
                                                              outgoing_rx)
                                 })
                       .then(move |result| {
                                 let _ = incoming_tx1.send(IncomingCommand::Exit(client, result));
                                 Ok(())
                             }));
        self.clients.insert(client, outgoing_tx);
    }
    fn handle_incoming_command(&mut self,
                               command: IncomingCommand)
                               -> Option<(SocketAddr, Result<RawMessage>)> {
        match command {
            IncomingCommand::Recv(addr, message) => Some((addr, message)),
            IncomingCommand::Exit(addr, result) => {
                self.clients.remove(&addr);
                if let Err(e) = result {
                    Some((addr, Err(e)))
                } else {
                    None
                }
            }
        }
    }
}
impl Transport for TcpServerTransport {}
impl MessageSink for TcpServerTransport {}
impl MessageStream for TcpServerTransport {}
impl Sink for TcpServerTransport {
    type SinkItem = MessageSinkItem;
    type SinkError = Error;
    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        let (peer, message, link) = item;
        let link = if let Some(outgoing_tx) = self.clients.get(&peer) {
            match outgoing_tx.send(OutgoingCommand::Send(message, link)) {
                Ok(_) => None,
                Err(e) => {
                    let SendError(OutgoingCommand::Send(_, link)) = e;
                    link
                }
            }
        } else {
            link
        };
        if let Some(link) = link {
            let e = ErrorKind::Other.cause(format!("No such client found: {}", peer));
            link.exit(Err(track!(e)));
        }
        Ok(AsyncSink::Ready)
    }
    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        Ok(Async::Ready(()))
    }
}
impl Stream for TcpServerTransport {
    type Item = (SocketAddr, Result<RawMessage>);
    type Error = Error;
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        loop {
            match track_try!(self.listener.poll()) {
                Async::NotReady => {}
                Async::Ready(None) => return Ok(Async::Ready(None)),
                Async::Ready(Some((connected, client))) => {
                    self.handle_new_connection(client, connected);
                    continue;
                }
            }
            match track_try!(self.incoming_rx.poll().map_err(|()| ErrorKind::Other)) {
                Async::NotReady => return Ok(Async::NotReady),
                Async::Ready(None) => unreachable!(),
                Async::Ready(Some(command)) => {
                    if let Some(item) = self.handle_incoming_command(command) {
                        return Ok(Async::Ready(Some(item)));
                    }
                }
            }
        }
    }
}

#[derive(Debug)]
struct TcpHandleClientLoop {
    peer: SocketAddr,
    transport: TcpClientTransport,
    incoming_tx: mpsc::Sender<IncomingCommand>,
    outgoing_rx: mpsc::Receiver<OutgoingCommand>,
}
impl TcpHandleClientLoop {
    fn new(peer: SocketAddr,
           stream: TcpStream,
           incoming_tx: mpsc::Sender<IncomingCommand>,
           outgoing_rx: mpsc::Receiver<OutgoingCommand>)
           -> Self {
        TcpHandleClientLoop {
            peer: peer,
            transport: TcpClientTransport::new(peer, stream),
            incoming_tx: incoming_tx,
            outgoing_rx: outgoing_rx,
        }
    }
}
impl Future for TcpHandleClientLoop {
    type Item = ();
    type Error = Error;
    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            track_try!(self.transport.poll_complete());
            match track_try!(self.outgoing_rx.poll().map_err(|()| ErrorKind::Other)) {
                Async::NotReady => {}
                Async::Ready(None) => return Ok(Async::Ready(())),
                Async::Ready(Some(command)) => {
                    let OutgoingCommand::Send(message, link) = command;
                    track_try!(self.transport.start_send((self.peer, message, link)));
                    continue;
                }
            }
            match track_try!(self.transport.poll()) {
                Async::NotReady => return Ok(Async::NotReady),
                Async::Ready(None) => return Ok(Async::Ready(())),
                Async::Ready(Some((peer, message))) => {
                    assert_eq!(peer, self.peer);
                    let command = IncomingCommand::Recv(peer, message);
                    let _ = self.incoming_tx.send(command);
                }
            }
        }
    }
}

/// TCP based client-side implementation of [Transport](trait.Transport.html) trait.
///
/// This can only communicate with pre-connected peer (server).
#[derive(Debug)]
pub struct TcpClientTransport {
    sink: TcpMessageSink,
    stream: TcpMessageStream,
}
impl TcpClientTransport {
    /// Makes a new `TcpClientTransport` instance.
    pub fn new(peer: SocketAddr, stream: TcpStream) -> Self {
        TcpClientTransport {
            sink: TcpMessageSink::new(peer, stream.clone()),
            stream: TcpMessageStream::new(peer, stream),
        }
    }
}
impl Transport for TcpClientTransport {}
impl MessageSink for TcpClientTransport {}
impl MessageStream for TcpClientTransport {}
impl Sink for TcpClientTransport {
    type SinkItem = MessageSinkItem;
    type SinkError = Error;
    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        self.sink.start_send(item)
    }
    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        self.sink.poll_complete()
    }
}
impl Stream for TcpClientTransport {
    type Item = (SocketAddr, Result<RawMessage>);
    type Error = Error;
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.stream.poll()
    }
}

type MessageSinkState = Option<::std::result::Result<TcpStream, WriteAll<TcpStream, Vec<u8>>>>;

#[derive(Debug)]
struct TcpMessageSink {
    peer: SocketAddr,
    state: MessageSinkState,
    queue: VecDeque<(RawMessage, Option<Link<(), Error, (), ()>>)>,
}
impl TcpMessageSink {
    fn new(peer: SocketAddr, stream: TcpStream) -> Self {
        TcpMessageSink {
            peer: peer,
            state: Some(Ok(stream)),
            queue: VecDeque::new(),
        }
    }
    fn poll_complete_impl(&mut self) -> Poll<(), Error> {
        let mut state = self.state.take().expect("unreachable");
        loop {
            match state {
                Err(mut future) => {
                    let polled = track_try!(future.poll().map_err(|e| e.into_error()));
                    if let Async::Ready((socket, _)) = polled {
                        let (_, link) = self.queue.pop_front().unwrap();
                        if let Some(link) = link {
                            link.exit(Ok(()));
                        }
                        state = Ok(socket);
                    } else {
                        self.state = Some(Err(future));
                        return Ok(Async::NotReady);
                    }
                }
                Ok(stream) => {
                    if let Some(&(ref message, _)) = self.queue.front() {
                        state = Err(stream.async_write_all(message.to_bytes()));
                    } else {
                        self.state = Some(Ok(stream));
                        return Ok(Async::Ready(()));
                    }
                }
            }
        }
    }
}
impl Sink for TcpMessageSink {
    type SinkItem = MessageSinkItem;
    type SinkError = Error;
    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        let (addr, message, link) = item;
        track_assert_eq!(addr, self.peer, ErrorKind::Invalid);
        self.queue.push_back((message, link));
        Ok(AsyncSink::Ready)
    }
    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        match track_err!(self.poll_complete_impl()) {
            Err(e) => {
                for link in self.queue.drain(..).filter_map(|(_, link)| link) {
                    link.exit(Err(e.clone()));
                }
                Err(e)
            }
            Ok(v) => Ok(v),

        }
    }
}
impl MessageSink for TcpMessageSink {}

type RecvMessageBytes = BoxFuture<(TcpStream, Vec<u8>), io::Error>;

struct TcpMessageStream {
    peer: SocketAddr,
    future: RecvMessageBytes,
}
impl TcpMessageStream {
    fn new(peer: SocketAddr, stream: TcpStream) -> Self {
        TcpMessageStream {
            peer: peer,
            future: Self::recv_message_bytes(stream),
        }
    }
    fn recv_message_bytes(stream: TcpStream) -> RecvMessageBytes {
        let pattern = vec![0; 20]
            .and_then(|mut buf| {
                          let message_len = (&mut &buf[2..4]).read_u16be().unwrap();
                          buf.resize(20 + message_len as usize, 0);
                          Window::new(buf).skip(20)
                      })
            .and_then(|buf| buf.into_inner());
        pattern.read_from(stream).map_err(|e| e.into_error()).boxed()
    }
}
impl Stream for TcpMessageStream {
    type Item = (SocketAddr, Result<RawMessage>);
    type Error = Error;
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        let polled = track_try!(match self.future.poll() {
                                    Err(e) => {
            if e.kind() == io::ErrorKind::UnexpectedEof {
                return Ok(Async::Ready(None));
            }
            Err(e)
        }
                                    Ok(v) => Ok(v),
                                });
        if let Async::Ready((stream, bytes)) = polled {
            let message = track_err!(RawMessage::read_from(&mut &bytes[..]));
            self.future = Self::recv_message_bytes(stream);
            Ok(Async::Ready(Some((self.peer, message))))
        } else {
            Ok(Async::NotReady)
        }
    }
}
impl MessageStream for TcpMessageStream {}
impl fmt::Debug for TcpMessageStream {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "TcpMessageStream {{ peer: {:?}, future: _ }}", self.peer)
    }
}
<commit_msg>Fix bug in reading STUN message using TCP<commit_after>use std::io;
use std::fmt;
use std::net::SocketAddr;
use std::collections::{VecDeque, HashMap};
use std::sync::mpsc::SendError;
use fibers::{Spawn, BoxSpawn};
use fibers::net::{TcpStream, TcpListener};
use fibers::net::futures::{TcpListenerBind, Connected};
use fibers::net::streams::Incoming;
use fibers::sync::mpsc;
use fibers::sync::oneshot::Link;
use futures::{Future, BoxFuture, Async, Poll, Stream, Sink, AsyncSink, StartSend};
use handy_async::io::{ReadFrom, AsyncWrite};
use handy_async::io::futures::WriteAll;
use handy_async::sync_io::ReadExt;
use handy_async::pattern::{Pattern, Window};
use trackable::error::ErrorKindExt;

use {Result, Error, ErrorKind};
use message::RawMessage;
use super::{MessageStream, MessageSink, MessageSinkItem, Transport};

#[derive(Debug)]
enum OutgoingCommand {
    Send(RawMessage, Option<Link<(), Error, (), ()>>),
}

#[derive(Debug)]
enum IncomingCommand {
    Recv(SocketAddr, Result<RawMessage>),
    Exit(SocketAddr, Result<()>),
}

#[derive(Debug)]
enum Listener {
    Bind(TcpListenerBind),
    Incoming(Incoming),
}
impl Stream for Listener {
    type Item = (Connected, SocketAddr);
    type Error = Error;
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        loop {
            let next = match *self {
                Listener::Bind(ref mut future) => {
                    if let Async::Ready(listener) = track_try!(future.poll()) {
                        Listener::Incoming(listener.incoming())
                    } else {
                        return Ok(Async::NotReady);
                    }
                }
                Listener::Incoming(ref mut stream) => return track_err!(stream.poll()),
            };
            *self = next;
        }
    }
}

/// TCP based server-side implementation of [Transport](trait.Transport.html) trait.
#[derive(Debug)]
pub struct TcpServerTransport {
    spawner: BoxSpawn,
    listener: Listener,
    clients: HashMap<SocketAddr, mpsc::Sender<OutgoingCommand>>,
    incoming_tx: mpsc::Sender<IncomingCommand>,
    incoming_rx: mpsc::Receiver<IncomingCommand>,
}
impl TcpServerTransport {
    /// Makes a new `TcpServerTransport` instance.
    pub fn new<S>(spawner: S, bind_addr: SocketAddr) -> Self
        where S: Spawn + Send + 'static
    {
        let (incoming_tx, incoming_rx) = mpsc::channel();
        TcpServerTransport {
            spawner: spawner.boxed(),
            listener: Listener::Bind(TcpListener::bind(bind_addr)),
            clients: HashMap::new(),
            incoming_tx: incoming_tx,
            incoming_rx: incoming_rx,
        }
    }
    fn handle_new_connection(&mut self, client: SocketAddr, connected: Connected) {
        let incoming_tx0 = self.incoming_tx.clone();
        let incoming_tx1 = self.incoming_tx.clone();
        let (outgoing_tx, outgoing_rx) = mpsc::channel();
        self.spawner
            .spawn(track_err!(connected)
                       .and_then(move |stream| {
                                     TcpHandleClientLoop::new(client,
                                                              stream,
                                                              incoming_tx0,
                                                              outgoing_rx)
                                 })
                       .then(move |result| {
                                 let _ = incoming_tx1.send(IncomingCommand::Exit(client, result));
                                 Ok(())
                             }));
        self.clients.insert(client, outgoing_tx);
    }
    fn handle_incoming_command(&mut self,
                               command: IncomingCommand)
                               -> Option<(SocketAddr, Result<RawMessage>)> {
        match command {
            IncomingCommand::Recv(addr, message) => Some((addr, message)),
            IncomingCommand::Exit(addr, result) => {
                self.clients.remove(&addr);
                if let Err(e) = result {
                    Some((addr, Err(e)))
                } else {
                    None
                }
            }
        }
    }
}
impl Transport for TcpServerTransport {}
impl MessageSink for TcpServerTransport {}
impl MessageStream for TcpServerTransport {}
impl Sink for TcpServerTransport {
    type SinkItem = MessageSinkItem;
    type SinkError = Error;
    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        let (peer, message, link) = item;
        let link = if let Some(outgoing_tx) = self.clients.get(&peer) {
            match outgoing_tx.send(OutgoingCommand::Send(message, link)) {
                Ok(_) => None,
                Err(e) => {
                    let SendError(OutgoingCommand::Send(_, link)) = e;
                    link
                }
            }
        } else {
            link
        };
        if let Some(link) = link {
            let e = ErrorKind::Other.cause(format!("No such client found: {}", peer));
            link.exit(Err(track!(e)));
        }
        Ok(AsyncSink::Ready)
    }
    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        Ok(Async::Ready(()))
    }
}
impl Stream for TcpServerTransport {
    type Item = (SocketAddr, Result<RawMessage>);
    type Error = Error;
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        loop {
            match track_try!(self.listener.poll()) {
                Async::NotReady => {}
                Async::Ready(None) => return Ok(Async::Ready(None)),
                Async::Ready(Some((connected, client))) => {
                    self.handle_new_connection(client, connected);
                    continue;
                }
            }
            match track_try!(self.incoming_rx.poll().map_err(|()| ErrorKind::Other)) {
                Async::NotReady => return Ok(Async::NotReady),
                Async::Ready(None) => unreachable!(),
                Async::Ready(Some(command)) => {
                    if let Some(item) = self.handle_incoming_command(command) {
                        return Ok(Async::Ready(Some(item)));
                    }
                }
            }
        }
    }
}

#[derive(Debug)]
struct TcpHandleClientLoop {
    peer: SocketAddr,
    transport: TcpClientTransport,
    incoming_tx: mpsc::Sender<IncomingCommand>,
    outgoing_rx: mpsc::Receiver<OutgoingCommand>,
}
impl TcpHandleClientLoop {
    fn new(peer: SocketAddr,
           stream: TcpStream,
           incoming_tx: mpsc::Sender<IncomingCommand>,
           outgoing_rx: mpsc::Receiver<OutgoingCommand>)
           -> Self {
        TcpHandleClientLoop {
            peer: peer,
            transport: TcpClientTransport::new(peer, stream),
            incoming_tx: incoming_tx,
            outgoing_rx: outgoing_rx,
        }
    }
}
impl Future for TcpHandleClientLoop {
    type Item = ();
    type Error = Error;
    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            track_try!(self.transport.poll_complete());
            match track_try!(self.outgoing_rx.poll().map_err(|()| ErrorKind::Other)) {
                Async::NotReady => {}
                Async::Ready(None) => return Ok(Async::Ready(())),
                Async::Ready(Some(command)) => {
                    let OutgoingCommand::Send(message, link) = command;
                    track_try!(self.transport.start_send((self.peer, message, link)));
                    continue;
                }
            }
            match track_try!(self.transport.poll()) {
                Async::NotReady => return Ok(Async::NotReady),
                Async::Ready(None) => return Ok(Async::Ready(())),
                Async::Ready(Some((peer, message))) => {
                    assert_eq!(peer, self.peer);
                    let command = IncomingCommand::Recv(peer, message);
                    let _ = self.incoming_tx.send(command);
                }
            }
        }
    }
}

/// TCP based client-side implementation of [Transport](trait.Transport.html) trait.
///
/// This can only communicate with pre-connected peer (server).
#[derive(Debug)]
pub struct TcpClientTransport {
    sink: TcpMessageSink,
    stream: TcpMessageStream,
}
impl TcpClientTransport {
    /// Makes a new `TcpClientTransport` instance.
    pub fn new(peer: SocketAddr, stream: TcpStream) -> Self {
        TcpClientTransport {
            sink: TcpMessageSink::new(peer, stream.clone()),
            stream: TcpMessageStream::new(peer, stream),
        }
    }
}
impl Transport for TcpClientTransport {}
impl MessageSink for TcpClientTransport {}
impl MessageStream for TcpClientTransport {}
impl Sink for TcpClientTransport {
    type SinkItem = MessageSinkItem;
    type SinkError = Error;
    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        self.sink.start_send(item)
    }
    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        self.sink.poll_complete()
    }
}
impl Stream for TcpClientTransport {
    type Item = (SocketAddr, Result<RawMessage>);
    type Error = Error;
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.stream.poll()
    }
}

type MessageSinkState = Option<::std::result::Result<TcpStream, WriteAll<TcpStream, Vec<u8>>>>;

#[derive(Debug)]
struct TcpMessageSink {
    peer: SocketAddr,
    state: MessageSinkState,
    queue: VecDeque<(RawMessage, Option<Link<(), Error, (), ()>>)>,
}
impl TcpMessageSink {
    fn new(peer: SocketAddr, stream: TcpStream) -> Self {
        TcpMessageSink {
            peer: peer,
            state: Some(Ok(stream)),
            queue: VecDeque::new(),
        }
    }
    fn poll_complete_impl(&mut self) -> Poll<(), Error> {
        let mut state = self.state.take().expect("unreachable");
        loop {
            match state {
                Err(mut future) => {
                    let polled = track_try!(future.poll().map_err(|e| e.into_error()));
                    if let Async::Ready((socket, _)) = polled {
                        let (_, link) = self.queue.pop_front().unwrap();
                        if let Some(link) = link {
                            link.exit(Ok(()));
                        }
                        state = Ok(socket);
                    } else {
                        self.state = Some(Err(future));
                        return Ok(Async::NotReady);
                    }
                }
                Ok(stream) => {
                    if let Some(&(ref message, _)) = self.queue.front() {
                        state = Err(stream.async_write_all(message.to_bytes()));
                    } else {
                        self.state = Some(Ok(stream));
                        return Ok(Async::Ready(()));
                    }
                }
            }
        }
    }
}
impl Sink for TcpMessageSink {
    type SinkItem = MessageSinkItem;
    type SinkError = Error;
    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        let (addr, message, link) = item;
        track_assert_eq!(addr, self.peer, ErrorKind::Invalid);
        self.queue.push_back((message, link));
        Ok(AsyncSink::Ready)
    }
    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        match track_err!(self.poll_complete_impl()) {
            Err(e) => {
                for link in self.queue.drain(..).filter_map(|(_, link)| link) {
                    link.exit(Err(e.clone()));
                }
                Err(e)
            }
            Ok(v) => Ok(v),

        }
    }
}
impl MessageSink for TcpMessageSink {}

type RecvMessageBytes = BoxFuture<(TcpStream, Vec<u8>), io::Error>;

struct TcpMessageStream {
    peer: SocketAddr,
    future: RecvMessageBytes,
}
impl TcpMessageStream {
    fn new(peer: SocketAddr, stream: TcpStream) -> Self {
        TcpMessageStream {
            peer: peer,
            future: Self::recv_message_bytes(stream),
        }
    }
    fn recv_message_bytes(stream: TcpStream) -> RecvMessageBytes {
        let pattern = vec![0; 20]
            .and_then(|mut buf| {
                          let message_len = (&mut &buf[2..4]).read_u16be().unwrap();
                          buf.resize(20 + message_len as usize, 0);
                          Window::new(buf).skip(20)
                      })
            .and_then(|buf| Ok(buf.into_inner()));
        pattern.read_from(stream).map_err(|e| e.into_error()).boxed()
    }
}
impl Stream for TcpMessageStream {
    type Item = (SocketAddr, Result<RawMessage>);
    type Error = Error;
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        let result = match self.future.poll() {
            Err(e) => {
                if e.kind() == io::ErrorKind::UnexpectedEof {
                    return Ok(Async::Ready(None));
                }
                Err(e)
            }
            Ok(v) => Ok(v),
        };
        let polled = track_try!(result);
        if let Async::Ready((stream, bytes)) = polled {
            let message = track_err!(RawMessage::read_from(&mut &bytes[..]));
            self.future = Self::recv_message_bytes(stream);
            Ok(Async::Ready(Some((self.peer, message))))
        } else {
            Ok(Async::NotReady)
        }
    }
}
impl MessageStream for TcpMessageStream {}
impl fmt::Debug for TcpMessageStream {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "TcpMessageStream {{ peer: {:?}, future: _ }}", self.peer)
    }
}
<|endoftext|>


<commit_before>//! Simple daemon
//!
//! This daemon will start web server, track new packages and build them

use std::{env, thread};
use std::process::exit;
use std::fs::File;
use std::io::Write;
use std::time::Duration;
use std::path::PathBuf;
use libc::fork;
use time;
use DocBuilderOptions;
use DocBuilder;
use utils::{update_sources, update_release_activity};
use db::{connect_db, update_search_index};

const DAEMON_PID_FILE_PATH: &'static str = "/var/run/cratesfyi.pid";

pub fn start_daemon() {
    // first check required environment variables
    for v in ["CRATESFYI_PREFIX",
              "CRATESFYI_PREFIX",
              "CRATESFYI_GITHUB_USERNAME",
              "CRATESFYI_GITHUB_ACCESSTOKEN"]
        .iter() {
        env::var(v).expect("Environment variable not found");
    }

    // check paths once
    opts().check_paths().unwrap();

    // fork the process
    let pid = unsafe { fork() };
    if pid > 0 {
        let mut file = File::create(DAEMON_PID_FILE_PATH).expect("Failed to create pid file");
        writeln!(&mut file, "{}", pid).expect("Failed to write pid");

        info!("cratesfyi {} daemon started on: {}", ::BUILD_VERSION, pid);
        exit(0);
    }

    // check new crates every 5 minutes
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(300));
            debug!("Checking new crates");
            let doc_builder = DocBuilder::new(opts());
            if let Err(e) = doc_builder.get_new_crates() {
                error!("Failed to get new crates: {}", e);
            }
        }
    });

    // build new crates every 3 minutes
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(180));

            let mut opts = opts();
            opts.skip_if_exists = true;

            // check lock file
            if opts.prefix.join("cratesfyi.lock").exists() {
                warn!("Lock file exits, skipping building new crates");
                continue;
            }

            // update index
            if let Err(e) = update_sources() {
                error!("Failed to update sources: {}", e);
                continue;
            }

            let mut doc_builder = DocBuilder::new(opts);
            if let Err(e) = doc_builder.load_cache() {
                error!("Failed to load cache: {}", e);
                continue;
            }

            debug!("Building new crates");
            if let Err(e) = doc_builder.build_packages_queue() {
                error!("Failed build new crates: {}", e);
            }

            if let Err(e) = doc_builder.save_cache() {
                error!("Failed to save cache: {}", e);
            }
        }
    });

    // update release activity everyday at 02:00
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(60));
            let now = time::now();
            if now.tm_hour == 2 && now.tm_min == 0 {
                info!("Updating release activity");
                if let Err(e) = update_release_activity() {
                    error!("Failed to update release activity: {}", e);
                }
            }
        }
    });

    // update search index every 3 hours
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(60*60*3));
            let conn = connect_db().expect("Failed to connect database");
            if let Err(e) = update_search_index(&conn) {
                error!("Failed to update search index: {}", e);
            }
        }
    });

    // TODO: update ssl certificate every 3 months

    // at least start web server
    info!("Starting web server");
    ::start_web_server(None);
}

fn opts() -> DocBuilderOptions {
    let prefix = PathBuf::from(env::var("CRATESFYI_PREFIX")
                               .expect("CRATESFYI_PREFIX environment variable not found"));
    DocBuilderOptions::from_prefix(prefix)
}
<commit_msg>Save pid file into prefix<commit_after>//! Simple daemon
//!
//! This daemon will start web server, track new packages and build them

use std::{env, thread};
use std::process::exit;
use std::fs::File;
use std::io::Write;
use std::time::Duration;
use std::path::PathBuf;
use libc::fork;
use time;
use DocBuilderOptions;
use DocBuilder;
use utils::{update_sources, update_release_activity};
use db::{connect_db, update_search_index};

pub fn start_daemon() {
    // first check required environment variables
    for v in ["CRATESFYI_PREFIX",
              "CRATESFYI_PREFIX",
              "CRATESFYI_GITHUB_USERNAME",
              "CRATESFYI_GITHUB_ACCESSTOKEN"]
        .iter() {
        env::var(v).expect("Environment variable not found");
    }

    let dbopts = opts();

    // check paths once
    dbopts.check_paths().unwrap();

    // fork the process
    let pid = unsafe { fork() };
    if pid > 0 {
        let mut file = File::create(dbopts.prefix.join("cratesfyi.pid"))
            .expect("Failed to create pid file");
        writeln!(&mut file, "{}", pid).expect("Failed to write pid");

        info!("cratesfyi {} daemon started on: {}", ::BUILD_VERSION, pid);
        exit(0);
    }

    // check new crates every 5 minutes
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(300));
            debug!("Checking new crates");
            let doc_builder = DocBuilder::new(opts());
            if let Err(e) = doc_builder.get_new_crates() {
                error!("Failed to get new crates: {}", e);
            }
        }
    });

    // build new crates every 3 minutes
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(180));

            let mut opts = opts();
            opts.skip_if_exists = true;

            // check lock file
            if opts.prefix.join("cratesfyi.lock").exists() {
                warn!("Lock file exits, skipping building new crates");
                continue;
            }

            // update index
            if let Err(e) = update_sources() {
                error!("Failed to update sources: {}", e);
                continue;
            }

            let mut doc_builder = DocBuilder::new(opts);
            if let Err(e) = doc_builder.load_cache() {
                error!("Failed to load cache: {}", e);
                continue;
            }

            debug!("Building new crates");
            if let Err(e) = doc_builder.build_packages_queue() {
                error!("Failed build new crates: {}", e);
            }

            if let Err(e) = doc_builder.save_cache() {
                error!("Failed to save cache: {}", e);
            }
        }
    });

    // update release activity everyday at 02:00
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(60));
            let now = time::now();
            if now.tm_hour == 2 && now.tm_min == 0 {
                info!("Updating release activity");
                if let Err(e) = update_release_activity() {
                    error!("Failed to update release activity: {}", e);
                }
            }
        }
    });

    // update search index every 3 hours
    thread::spawn(move || {
        loop {
            thread::sleep(Duration::from_secs(60 * 60 * 3));
            let conn = connect_db().expect("Failed to connect database");
            if let Err(e) = update_search_index(&conn) {
                error!("Failed to update search index: {}", e);
            }
        }
    });

    // TODO: update ssl certificate every 3 months

    // at least start web server
    info!("Starting web server");
    ::start_web_server(None);
}

fn opts() -> DocBuilderOptions {
    let prefix = PathBuf::from(env::var("CRATESFYI_PREFIX")
        .expect("CRATESFYI_PREFIX environment variable not found"));
    DocBuilderOptions::from_prefix(prefix)
}
<|endoftext|>


<commit_before>use term_painter::ToStyle;
use term_painter::Attr::*;
use std::fmt;
use std::io;
use std::process;

#[derive(Clone)]
pub struct SearchResult {
    pub title: String,
    pub score: String,
    pub number: String,
    pub genres: Vec<String>,
    pub year: String,
}

impl fmt::Display for SearchResult {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        try!(write!(f,
                    "{}",
                    Bold.paint(format!("{} ({})\n", &self.title, &self.year))));
        let mut stars = String::new();
        let score: f32 = self.score.parse::<f32>().unwrap();
        for i in 0..10 {
            if i < (score + 0.5) as i32 {
                stars.push_str("\u{2605}");
            } else {
                stars.push_str("\u{2606}");
            }
        }
        try!(write!(f, "{}", &self.genres[0]));
        for genre in self.genres.clone().into_iter().skip(1) {
            try!(write!(f, ", {}", genre));
        }
        write!(f,
               "\n{}, ({} bei {} Bewertungen)",
               stars,
               &self.score,
               &self.number)
    }
}

impl fmt::Debug for SearchResult {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "{:?}\n{:?}\n{:?}\n{:?}\n{:?}\n******",
               self.title,
               self.year,
               self.genres,
               self.score,
               self.number)
    }
}

pub fn output_result(results: Vec<SearchResult>) {
    let further = results.clone();
    for res in results.into_iter().take(3) {
        println!("{}", res);
    }
    for output in further.into_iter().skip(3) {
        let mut buffer = String::new();
        if io::stdin().read_line(&mut buffer).is_ok() {
            if buffer != "q\n" {
                println!("{}", output);
            } else {
                process::exit(0);
            }
        }
    }
}

pub fn cancel_request() {
    println!("Sorry, no movie matches to your given Params!");
    process::exit(0);
}
<commit_msg>fixed issues<commit_after>use term_painter::ToStyle;
use term_painter::Attr::*;
use std::fmt;
use std::io;
use std::process;

#[derive(Clone)]
pub struct SearchResult {
    pub title: String,
    pub score: String,
    pub number: String,
    pub genres: Vec<String>,
    pub year: String,
}

impl fmt::Display for SearchResult {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        try!(write!(f,
                    "{}",
                    Bold.paint(format!("{} ({})\n", &self.title, &self.year))));
        let mut stars = String::new();
        let score: f32 = self.score.parse::<f32>().unwrap();
        for i in 0..10 {
            if i < (score + 0.5) as i32 {
                stars.push_str("\u{2605}");
            } else {
                stars.push_str("\u{2606}");
            }
        }
        try!(write!(f, "{}", &self.genres[0]));
        for genre in self.genres.clone().into_iter().skip(1) {
            try!(write!(f, ", {}", genre));
        }
        write!(f,
               "\n{}, ({} at {} ratings)",
               stars,
               &self.score,
               &self.number)
    }
}

impl fmt::Debug for SearchResult {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
               "{:?}\n{:?}\n{:?}\n{:?}\n{:?}\n******",
               self.title,
               self.year,
               self.genres,
               self.score,
               self.number)
    }
}

pub fn output_result(results: Vec<SearchResult>) {
    if results.is_empty() {
        cancel_request();
    }
    let further = results.clone();
    for res in results.into_iter().take(3) {
        println!("\n{}", res);
    }
    for output in further.into_iter().skip(3) {
        let mut buffer = String::new();
        if io::stdin().read_line(&mut buffer).is_ok() {
            if buffer != "q\n" {
                println!("{}", output);
            } else {
                process::exit(0);
            }
        }
    }
}

pub fn cancel_request() {
    println!("Sorry, no movie matches to your given Params!");
    process::exit(0);
}
<|endoftext|>