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

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// <NAME>, <NAME>, <NAME>".

#[derive(Debug)]
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,
    index: 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)]
struct List {
    lookahead1: Node,
    lookahead2: 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], maxbits: i32) -> Vec<u32> {
    let mut leaves = vec![];

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

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

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

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

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

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

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

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

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

    result
}

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

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

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

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

fn next_tree(weight_sum: 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.lookahead2.leaf_counts;

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

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

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

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

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

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

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

    if current_list_index == 0 {
        lowest_list(lists, leaves);
    } else {
        // We're at a list other than the lowest list.
        let weight_sum = {
            let previous_list = lists.get(current_list_index - 1).unwrap();
            previous_list.lookahead1.weight + previous_list.lookahead2.weight
        };<commit_msg>Get rid of get().unwrap()s<commit_after>use std::cmp::Ordering;
use std::{mem};

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// <NAME>, <NAME>, <NAME>".

#[derive(Debug)]
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,
    index: 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)]
struct List {
    lookahead1: Node,
    lookahead2: 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], maxbits: i32) -> Vec<u32> {
    let mut leaves = vec![];

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

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

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

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

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

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

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

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

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

    result
}

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

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

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

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

fn next_tree(weight_sum: 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.lookahead2.leaf_counts;

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

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

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

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

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

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

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

    if current_list_index == 0 {
        lowest_list(lists, leaves);
    } else {
        // We're at a list other than the lowest list.
        let weight_sum = {
            let previous_list = &lists[current_list_index - 1];
            previous_list.lookahead1.weight + previous_list.lookahead2.weight
        };
{'fraction_non_alphanumeric': 0.12154952447228021, 'fraction_numerical': 0.011134307585247043, 'mean_word_length': 2.852844801906464, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '64819', 'n_tokens_mistral': 4343, 'n_tokens_neox': 4220, 'n_words': 1087}