huffman.hpp

// Copyright Take Vos 2020.
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at https://www.boost.org/LICENSE_1_0.txt)
 
#include "required.hpp"
#include "bits.hpp"
#include "cast.hpp"
#include <span>
#include <vector>
 
namespace tt {
 
template<typename T>
class huffman_tree {
    static_assert(std::is_integral_v<T> && std::is_signed_v<T>);
 
    using state_t = T const *;
 
    std::vector<T> tree;
 
public:
    huffman_tree() {
        tree.push_back(0);
        tree.push_back(0);
    }
 
    void add(int symbol, int code, int code_length) noexcept {
        tt_axiom(code_length >= 1);
 
        int offset = 0;
        while (--code_length > 0) {
            int select = (code >> code_length) & 1;
            offset += select;
          
            int value = tree[offset];
 
            // value may not be a leaf.
            tt_axiom(value <= 0); 
 
            if (value == 0) {
                // Unused node entry. Point to the first of two new entries.
                value = -(static_cast<int>(std::ssize(tree)) - offset);
 
                tree[offset] = narrow_cast<T>(value);
                tree.push_back(0);
                tree.push_back(0);
            }
 
            // Go to the next entry.
            offset -= value;
        }
 
        // place the symbol as a leaf.
        int select = code & 1;
        offset += select;
 
        tt_axiom(tree[offset] == 0); 
        tree[offset] = narrow_cast<T>(symbol + 1);
    }
 
    [[nodiscard]] state_t start() const noexcept {
        return tree.data();
    }
 
    [[nodiscard]] int get(bool code_bit, state_t &state) const {
        state += static_cast<ptrdiff_t>(code_bit);
 
        if (*state == 0) {
            throw parse_error("Code not in huffman tree.");
        }
 
        auto value = *state;
        state -= static_cast<ptrdiff_t>(value);
        return value - 1;
    }
 
    [[nodiscard]] int get_symbol(std::span<std::byte const> bytes, ssize_t &bit_offset) const noexcept {
        auto state = start();
        while (true) {
            int symbol;
            if ((symbol = get(get_bit(bytes, bit_offset), state)) >= 0) {
                return symbol;
            }
        }
    }
    
    [[nodiscard]] static huffman_tree from_lengths(int const *lengths, ssize_t nr_symbols) {
        struct symbol_length_t {
            int symbol;
            int length;
 
            symbol_length_t(int symbol, int length) : symbol(symbol), length(length) {}
        };
 
        std::vector<symbol_length_t> symbol_lengths;
        symbol_lengths.reserve(nr_symbols);
 
        for (int symbol = 0; symbol != nr_symbols; ++symbol) {
            symbol_lengths.emplace_back(symbol, lengths[symbol]);
        }
 
        // Sort the table based on the length of the code, followed by symbol
        std::sort(symbol_lengths.begin(), symbol_lengths.end(), [](ttlet &a, ttlet &b) {
            if (a.length == b.length) {
                return a.symbol < b.symbol;
            } else {
                return a.length < b.length;
            }
        });
 
        auto r = huffman_tree{};
 
        int code = 0;
        int prev_length = 0;
        for (auto &&entry: symbol_lengths) {
            if (entry.length != 0) {
                code <<= (entry.length - prev_length);
 
                r.add(entry.symbol, code, entry.length);
                ++code;
            }
 
            prev_length = entry.length;
        }
 
        return r;
    }
 
    [[nodiscard]] static huffman_tree from_lengths(std::vector<int> const &lengths) {
        return from_lengths(lengths.data(), std::ssize(lengths));
    }
};
 
 
 
}