node.hpp

#ifndef BV_NODE_HPP
#define BV_NODE_HPP
 
#include <cassert>
#include <cstdint>
#include <cstring>
 
#ifndef CACHE_LINE
// Apparently the most common cache line size is 64.
#define CACHE_LINE 64
#endif
 
#include "branch_selection.hpp"
 
namespace bv {
 
template <class leaf_type, class dtype, uint64_t leaf_size, uint8_t branches>
class node {
   protected:
    typedef branchless_scan<dtype, branches> branching;
    uint8_t meta_data_;
    uint8_t child_count_;  // Bad word alignment. betewen 16 and 48 dead bits...
    branching child_sizes_;
    branching child_sums_;
    void* children_[branches];  
 
    static_assert(leaf_size >= 256, "leaf size needs to be a at least 256");
    static_assert((leaf_size % 128) == 0,
                  "leaf size needs to be divisible by 128");
    static_assert(leaf_size < 0xffffff, "leaf size must fit in 24 bits");
    static_assert((branches == 8) || (branches == 16) || (branches == 32) ||
                      (branches == 64) || (branches == 128),
                  "branching factor needs to be a reasonable power of 2");
 
   public:
    node() {
        meta_data_ = 0;
        child_count_ = 0;
        child_sizes_ = branching();
        child_sums_ = branching();
    }
 
    template <class qds>
    void generate_query_structure(qds* qs) {
        if (has_leaves()) {
            leaf_type** children = reinterpret_cast<leaf_type**>(children_);
            for (size_t i = 0; i < child_count_; i++) {
                children[i]->commit();
                qs->append(children[i]);
            }
        } else {
            node** children = reinterpret_cast<node**>(children_);
            for (size_t i = 0; i < child_count_; i++) {
                children[i]->generate_query_structure(qs);
            }
        }
    }
 
    void has_leaves(bool leaves) {
        meta_data_ = leaves ? meta_data_ | 0b10000000 : meta_data_ & 0b01111111;
    }
 
    bool has_leaves() const { return meta_data_ >> 7; }
 
    bool at(dtype index) const {
        uint8_t child_index = child_sizes_.find(index + 1);
        index -= child_index != 0 ? child_sizes_.get(child_index - 1) : 0;
        if (has_leaves()) {
            [[unlikely]] return reinterpret_cast<leaf_type*>(
                children_[child_index])
                ->at(index);
        } else {
            return reinterpret_cast<node*>(children_[child_index])->at(index);
        }
    }
 
    int set(dtype index, bool v) {
        uint8_t child_index = child_sizes_.find(index + 1);
        index -= child_index != 0 ? child_sizes_.get(child_index - 1) : 0;
        dtype change = 0;
        if (has_leaves()) {
            leaf_type* child =
                reinterpret_cast<leaf_type*>(children_[child_index]);
            [[unlikely]] change = child->set(index, v);
        } else {
            node* child = reinterpret_cast<node*>(children_[child_index]);
            change = child->set(index, v);
        }
        uint8_t c_count = child_count_;
        child_sums_.increment(child_index, c_count, change);
        return change;
    }
 
    dtype rank(dtype index) const {
        uint8_t child_index = child_sizes_.find(index);
        dtype res = 0;
        if (child_index != 0) {
            res = child_sums_.get(child_index - 1);
            [[likely]] index -= child_sizes_.get(child_index - 1);
        }
        if (has_leaves()) {
            leaf_type* child =
                reinterpret_cast<leaf_type*>(children_[child_index]);
            [[unlikely]] return res + child->rank(index);
        } else {
            node* child = reinterpret_cast<node*>(children_[child_index]);
            return res + child->rank(index);
        }
    }
 
    dtype select(dtype count) const {
        uint8_t child_index = child_sums_.find(count);
        dtype res = 0;
        if (child_index != 0) {
            res = child_sizes_.get(child_index - 1);
            [[likely]] count -= child_sums_.get(child_index - 1);
        }
        if (has_leaves()) {
            leaf_type* child =
                reinterpret_cast<leaf_type*>(children_[child_index]);
            [[unlikely]] return res + child->select(count);
        } else {
            node* child = reinterpret_cast<node*>(children_[child_index]);
            return res + child->select(count);
        }
    }
 
    template <class allocator>
    void deallocate(allocator* alloc) {
        if (has_leaves()) {
            leaf_type** children = reinterpret_cast<leaf_type**>(children_);
            for (uint8_t i = 0; i < child_count_; i++) {
                leaf_type* l = children[i];
                alloc->deallocate_leaf(l);
            }
        } else {
            node** children = reinterpret_cast<node**>(children_);
            for (uint8_t i = 0; i < child_count_; i++) {
                node* n = children[i];
                n->deallocate(alloc);
                alloc->deallocate_node(n);
            }
        }
    }
 
    uint8_t child_count() const { return child_count_; }
 
    void** children() { return children_; }
 
    branching* child_sizes() { return &child_sizes_; }
 
    branching* child_sums() { return &child_sums_; }
 
    dtype size() const {
        return child_count_ > 0 ? child_sizes_.get(child_count_ - 1) : 0;
    }
 
    dtype p_sum() const {
        return child_count_ > 0 ? child_sums_.get(child_count_ - 1) : 0;
    }
 
    template <class child>
    void append_child(child* new_child) {
        child_sizes_.append(child_count_, new_child->size());
        child_sums_.append(child_count_, new_child->p_sum());
        children_[child_count_] = new_child;
        child_count_++;
    }
 
    void* child(uint8_t i) { return children_[i]; }
 
    template <class allocator>
    void insert(dtype index, bool value, allocator* alloc) {
        if (has_leaves()) {
            leaf_insert(index, value, alloc);
        } else {
            [[likely]] node_insert(index, value, alloc);
        }
    }
 
    template <class allocator>
    bool remove(dtype index, allocator* alloc) {
        if (has_leaves()) {
            return leaf_remove(index, alloc);
        } else {
            [[likely]] return node_remove(index, alloc);
        }
    }
 
    void clear_first(uint8_t elems) {
        child_sizes_.clear_first(elems, child_count_);
        child_sums_.clear_first(elems, child_count_);
        for (uint8_t i = 0; i < child_count_ - elems; i++) {
            children_[i] = children_[i + elems];
        }
        child_count_ -= elems;
    }
 
    void transfer_append(node* other, uint8_t elems) {
        void** o_children = other->children();
        branching* o_sizes = other->child_sizes();
        branching* o_sums = other->child_sums();
        uint8_t local_index = child_count_;
        for (uint8_t i = 0; i < elems; i++) {
            children_[local_index] = o_children[i];
            local_index++;
        }
        child_sizes_.append(elems, child_count_, o_sizes);
        child_sums_.append(elems, child_count_, o_sums);
        child_count_ += elems;
        other->clear_first(elems);
    }
 
    void clear_last(uint8_t elems) {
        child_sizes_.clear_last(elems, child_count_);
        child_sums_.clear_last(elems, child_count_);
        child_count_ -= elems;
    }
 
    void transfer_prepend(node* other, uint8_t elems) {
        void** o_children = other->children();
        branching* o_sizes = other->child_sizes();
        branching* o_sums = other->child_sums();
        uint8_t o_size = other->child_count();
        memmove(children_ + elems, children_, child_count_ * sizeof(void*));
        for (uint8_t i = 0; i < elems; i++) {
            children_[i] = o_children[o_size - elems + i];
        }
        child_sizes_.prepend(elems, child_count_, o_size, o_sizes);
        child_sums_.prepend(elems, child_count_, o_size, o_sums);
        child_count_ += elems;
        other->clear_last(elems);
    }
 
    void append_all(node* other) {
        void** o_children = other->children();
        branching* o_sizes = other->child_sizes();
        branching* o_sums = other->child_sums();
        uint8_t o_size = other->child_count();
        for (uint8_t i = 0; i < o_size; i++) {
            children_[child_count_ + i] = o_children[i];
        }
        child_sizes_.append(o_size, child_count_, o_sizes);
        child_sums_.append(o_size, child_count_, o_sums);
        child_count_ += o_size;
    }
 
    uint64_t bits_size() const {
        uint64_t ret = sizeof(node) * 8;
        if (has_leaves()) {
            leaf_type* const* children =
                reinterpret_cast<leaf_type* const*>(children_);
            for (uint8_t i = 0; i < child_count_; i++) {
                ret += children[i]->bits_size();
            }
        } else {
            node* const* children = reinterpret_cast<node* const*>(children_);
            for (uint8_t i = 0; i < child_count_; i++) {
                ret += children[i]->bits_size();
            }
        }
        return ret;
    }
 
    void flush() {
        if (has_leaves()) {
            leaf_type** children = reinterpret_cast<leaf_type**>(children_);
            for (uint8_t i = 0; i < child_count_; i++) {
                children[i]->commit();
            }
        } else {
            node** children = reinterpret_cast<node**>(children_);
            for (uint8_t i = 0; i < child_count_; i++) {
                children[i]->flush();
            }
        }
    }
 
    uint64_t validate() const {
        uint64_t ret = 1;
        uint64_t child_s_sum = 0;
        uint64_t child_p_sum = 0;
        if (has_leaves()) {
            leaf_type* const* children =
                reinterpret_cast<leaf_type* const*>(children_);
            for (uint8_t i = 0; i < child_count_; i++) {
                uint64_t child_size = children[i]->size();
                assert(child_size >= leaf_size / 3);
                child_s_sum += child_size;
                assert(child_sizes_.get(i) == child_s_sum);
                child_p_sum += children[i]->p_sum();
                assert(child_sums_.get(i) == child_p_sum);
                assert(children[i]->capacity() * WORD_BITS <= leaf_size);
                assert(children[i]->size() <= leaf_size);
                ret += children[i]->validate();
            }
        } else {
            node* const* children = reinterpret_cast<node* const*>(children_);
            for (uint8_t i = 0; i < child_count_; i++) {
                uint64_t child_size = children[i]->size();
                assert(child_size >= branches / 3);
                child_s_sum += child_size;
                assert(child_sizes_.get(i) == child_s_sum);
                child_p_sum += children[i]->p_sum();
                assert(child_sums_.get(i) == child_p_sum);
                ret += children[i]->validate();
            }
        }
        return ret;
    }
 
    void print(bool internal_only) const {
        std::cout << "{\n\"type\": \"node\",\n"
                  << "\"has_leaves\": " << has_leaves() << ",\n"
                  << "\"child_count\": " << int(child_count_) << ",\n"
                  << "\"size\": " << size() << ",\n"
                  << "\"child_sizes\": [";
        for (uint8_t i = 0; i < branches; i++) {
            std::cout << child_sizes_.get(i);
            if (i != branches - 1) {
                std::cout << ", ";
            }
        }
        std::cout << "],\n"
                  << "\"p_sum\": " << p_sum() << ",\n"
                  << "\"child_sums\": [";
        for (uint8_t i = 0; i < branches; i++) {
            std::cout << child_sums_.get(i);
            if (i != branches - 1) {
                std::cout << ", ";
            }
        }
        std::cout << "],\n"
                  << "\"children\": [\n";
        if (has_leaves()) {
            leaf_type* const* children =
                reinterpret_cast<leaf_type* const*>(children_);
            for (uint8_t i = 0; i < child_count_; i++) {
                children[i]->print(internal_only);
                if (i != child_count_ - 1) {
                    std::cout << ",";
                }
                std::cout << "\n";
            }
        } else {
            node* const* children = reinterpret_cast<node* const*>(children_);
            for (uint8_t i = 0; i < child_count_; i++) {
                children[i]->print(internal_only);
                if (i != child_count_ - 1) {
                    std::cout << ",";
                }
                std::cout << "\n";
            }
        }
        std::cout << "]}";
    }
 
   protected:
    template <class allocator>
    void rebalance_leaf(uint8_t index, leaf_type* leaf, allocator* alloc) {
        // Number of leaves that can fit in the "left" sibling (with potential
        // reallocation).
        uint32_t l_cap = 0;
        if (index > 0) {
            l_cap = child_sizes_.get(index - 1);
            l_cap -= index > 1 ? child_sizes_.get(index - 2) : 0;
            [[likely]] l_cap = leaf_size - l_cap;
        }
        // Number of leaves that can fir in the "right" sibling (with potential
        // reallocation).
        uint32_t r_cap = 0;
        if (index < child_count_ - 1) {
            r_cap = child_sizes_.get(index + 1);
            r_cap -= child_sizes_.get(index);
            [[likely]] r_cap = leaf_size - r_cap;
        }
        if (l_cap < 2 * leaf_size / 9 && r_cap < 2 * leaf_size / 9) {
            // Rebalancing without creating a new leaf is impossible
            // (impracitcal).
            leaf_type* a_child;
            leaf_type* b_child;
            leaf_type* new_child;
            if (index == 0) {
                // If the full leaf is the first child, a new leaf is created
                // between indexes 0 and 1.
                dtype n_elem = child_sizes_.get(1) / 3;
                dtype n_cap = (n_elem + 2 * WORD_BITS) / WORD_BITS;
                n_cap += n_cap % 2;
                n_cap = n_cap * WORD_BITS > leaf_size ? leaf_size / WORD_BITS
                                                      : n_cap;
                a_child = reinterpret_cast<leaf_type*>(children_[0]);
                new_child = alloc->template allocate_leaf<leaf_type>(n_cap);
                b_child = reinterpret_cast<leaf_type*>(children_[1]);
                new_child->transfer_append(b_child, b_child->size() - n_elem);
                new_child->transfer_prepend(a_child, a_child->size() - n_elem);
                [[unlikely]] index++;
            } else {
                // If the full leaf is not the first child, a new leaf is
                // created to the left of the full leaf.
                a_child = reinterpret_cast<leaf_type*>(children_[index - 1]);
                b_child = reinterpret_cast<leaf_type*>(children_[index]);
                dtype n_elem = (a_child->size() + b_child->size()) / 3;
                dtype n_cap = (n_elem + 2 * WORD_BITS) / WORD_BITS;
                n_cap += n_cap % 2;
                n_cap = n_cap * WORD_BITS > leaf_size ? leaf_size / WORD_BITS
                                                      : n_cap;
                new_child = alloc->template allocate_leaf<leaf_type>(n_cap);
                new_child->transfer_append(b_child, b_child->size() - n_elem);
                new_child->transfer_prepend(a_child, a_child->size() - n_elem);
            }
            // Update cumulative sizes and sums.
            for (uint8_t i = child_count_; i > index; i--) {
                children_[i] = children_[i - 1];
            }
            children_[index] = new_child;
            child_sizes_.insert(index, child_count_, a_child->size(),
                                new_child->size());
            child_sums_.insert(index, child_count_, a_child->p_sum(),
                               new_child->p_sum());
            [[unlikely]] child_count_++;
        } else if (r_cap > l_cap) {
            // Right sibling has more space than the left sibling. Move elements
            // to right sibling
            leaf_type* sibling =
                reinterpret_cast<leaf_type*>(children_[index + 1]);
            uint32_t n_size = sibling->size() + r_cap / 2;
            if (sibling->capacity() * WORD_BITS < n_size) {
                n_size = n_size / WORD_BITS + 1;
                n_size += n_size % 2;
                n_size = n_size * WORD_BITS <= leaf_size
                             ? n_size
                             : leaf_size / WORD_BITS;
                children_[index + 1] = alloc->reallocate_leaf(
                    sibling, sibling->capacity(), n_size);
                sibling = reinterpret_cast<leaf_type*>(children_[index + 1]);
            }
            sibling->transfer_prepend(leaf, r_cap / 2);
            child_sizes_.set(
                index, index != 0 ? child_sizes_.get(index - 1) + leaf->size()
                                  : leaf->size());
            child_sums_.set(
                index, index != 0 ? child_sums_.get(index - 1) + leaf->p_sum()
                                  : leaf->p_sum());
        } else {
            // Left sibling has more space than the right sibling. Move elements
            // to the left sibling.
            leaf_type* sibling =
                reinterpret_cast<leaf_type*>(children_[index - 1]);
            uint32_t n_size = sibling->size() + l_cap / 2;
            if (sibling->capacity() * WORD_BITS < n_size) {
                n_size = n_size / WORD_BITS + 1;
                n_size += n_size % 2;
                n_size = n_size * WORD_BITS <= leaf_size
                             ? n_size
                             : leaf_size / WORD_BITS;
                children_[index - 1] = alloc->reallocate_leaf(
                    sibling, sibling->capacity(), n_size);
                sibling = reinterpret_cast<leaf_type*>(children_[index - 1]);
            }
            sibling->transfer_append(leaf, l_cap / 2);
            child_sizes_.set(index - 1,
                             index > 1
                                 ? child_sizes_.get(index - 2) + sibling->size()
                                 : sibling->size());
            child_sums_.set(index - 1, index > 1 ? child_sums_.get(index - 2) +
                                                       sibling->p_sum()
                                                 : sibling->p_sum());
            [[likely]] (void(0));
        }
    }
 
    template <class allocator>
    void leaf_insert(dtype index, bool value, allocator* alloc) {
        uint8_t child_index = child_sizes_.find(index);
        leaf_type* child = reinterpret_cast<leaf_type*>(children_[child_index]);
#ifdef DEBUG
        if (child->size() > leaf_size) {
            std::cerr << child->size() << " > " << leaf_size << "\n";
            std::cerr << "invalid leaf size" << std::endl;
            assert(child->size() <= leaf_size);
        }
#endif
        if (child->need_realloc()) {
            if (child->size() >= leaf_size) {
                rebalance_leaf(child_index, child, alloc);
            } else {
                dtype cap = child->capacity();
                children_[child_index] =
                    alloc->reallocate_leaf(child, cap, cap + 2);
            }
            return leaf_insert(index, value, alloc);
        }
        if (child_index != 0) {
            [[likely]] index -= child_sizes_.get(child_index - 1);
        }
        child_sizes_.increment(child_index, child_count_, 1u);
        child_sums_.increment(child_index, child_count_, value);
        child->insert(index, value);
    }
 
    template <class allocator>
    void rebalance_node(uint8_t index, allocator* alloc) {
        // Number of elements that can be added to the left sibling.
        uint32_t l_cap = 0;
        if (index > 0) {
            [[likely]] l_cap =
                branches -
                reinterpret_cast<node*>(children_[index - 1])->child_count();
        }
        // Number of elements that can be aded to the right sibling.
        uint32_t r_cap = 0;
        if (index < child_count_ - 1) {
            [[likely]] r_cap =
                branches -
                reinterpret_cast<node*>(children_[index + 1])->child_count();
        }
        node* a_node;
        node* b_node;
        if (l_cap <= 1 && r_cap <= 1) {
            // There is no room in either sibling.
            if (index == 0) {
                a_node = reinterpret_cast<node*>(children_[0]);
                b_node = reinterpret_cast<node*>(children_[1]);
                [[unlikely]] index++;
            } else {
                a_node = reinterpret_cast<node*>(children_[index - 1]);
                b_node = reinterpret_cast<node*>(children_[index]);
            }
            node* new_child = alloc->template allocate_node<node>();
            new_child->has_leaves(a_node->has_leaves());
            new_child->transfer_append(b_node, branches / 3);
            new_child->transfer_prepend(a_node, branches / 3);
            for (size_t i = child_count_; i > index; i--) {
                children_[i] = children_[i - 1];
            }
            child_sizes_.insert(index, child_count_, a_node->size(),
                                new_child->size());
            child_sums_.insert(index, child_count_, a_node->p_sum(),
                               new_child->p_sum());
            children_[index] = new_child;
            child_count_++;
            [[unlikely]] return;
        } else if (l_cap > r_cap) {
            // There is more room in the left sibling.
            a_node = reinterpret_cast<node*>(children_[index - 1]);
            b_node = reinterpret_cast<node*>(children_[index]);
            a_node->transfer_append(b_node, l_cap / 2);
            index--;
        } else {
            // There is more room in the right sibling.
            a_node = reinterpret_cast<node*>(children_[index]);
            b_node = reinterpret_cast<node*>(children_[index + 1]);
            b_node->transfer_prepend(a_node, r_cap / 2);
        }
        // Fix cumulative sums and sizes.
        if (index == 0) {
            child_sizes_.set(0, a_node->size());
            [[unlikely]] child_sums_.set(0, a_node->p_sum());
        } else {
            child_sizes_.set(index,
                             child_sizes_.get(index - 1) + a_node->size());
            child_sums_.set(index,
                            child_sums_.get(index - 1) + a_node->p_sum());
        }
    }
 
    template <class allocator>
    void node_insert(dtype index, bool value, allocator* alloc) {
        uint8_t child_index = child_sizes_.find(index);
        node* child = reinterpret_cast<node*>(children_[child_index]);
#ifdef DEBUG
        if (child_index >= child_count_) {
            std::cout << int(child_index) << " >= " << int(child_count_)
                      << std::endl;
            assert(child_index < child_count_);
        }
#endif
        if (child->child_count() == branches) {
            rebalance_node(child_index, alloc);
            child_index = child_sizes_.find(index);
            [[unlikely]] child =
                reinterpret_cast<node*>(children_[child_index]);
        }
        if (child_index != 0) {
            [[likely]] index -= child_sizes_.get(child_index - 1);
        }
        child_sizes_.increment(child_index, child_count_, 1u);
        child_sums_.increment(child_index, child_count_, value);
        child->insert(index, value, alloc);
    }
 
    template <class allocator>
    void rebalance_leaves_right(leaf_type* a, leaf_type* b, allocator* alloc) {
        dtype a_cap = a->capacity();
        dtype addition = (b->size() - leaf_size / 3) / 2;
        if (a_cap * WORD_BITS < a->size() + addition) {
            dtype n_cap = 1 + (a->size() + addition) / WORD_BITS;
            n_cap += n_cap % 2;
            n_cap =
                n_cap * WORD_BITS <= leaf_size ? n_cap : leaf_size / WORD_BITS;
            a = alloc->reallocate_leaf(a, a_cap, n_cap);
            children_[0] = a;
        }
        a->transfer_append(b, addition);
        child_sizes_.set(0, a->size());
        child_sums_.set(0, a->p_sum());
    }
 
    template <class allocator>
    void rebalance_leaves_left(leaf_type* a, leaf_type* b, uint8_t idx,
                               allocator* alloc) {
        dtype b_cap = b->capacity();
        dtype addition = (a->size() - leaf_size / 3) / 2;
        if (b_cap * WORD_BITS < b->size() + addition) {
            dtype n_cap = 1 + (b->size() + addition) / WORD_BITS;
            n_cap += n_cap % 2;
            n_cap =
                n_cap * WORD_BITS <= leaf_size ? n_cap : leaf_size / WORD_BITS;
            b = alloc->reallocate_leaf(b, b_cap, n_cap);
            children_[idx + 1] = b;
        }
        b->transfer_prepend(a, addition);
        if (idx == 0) {
            child_sizes_.set(0, a->size());
            [[unlikely]] child_sums_.set(0, a->p_sum());
        } else {
            child_sizes_.set(idx, child_sizes_.get(idx - 1) + a->size());
            child_sums_.set(idx, child_sums_.get(idx - 1) + a->p_sum());
        }
    }
 
    template <class allocator>
    void merge_leaves(leaf_type* a, leaf_type* b, uint8_t idx,
                      allocator* alloc) {
        dtype a_cap = a->capacity();
        if (a_cap * WORD_BITS < a->size() + b->size()) {
            dtype n_cap = 1 + (a->size() + b->size()) / WORD_BITS;
            n_cap += n_cap % 2;
            n_cap =
                n_cap * WORD_BITS <= leaf_size ? n_cap : leaf_size / WORD_BITS;
            a = alloc->reallocate_leaf(a, a_cap, n_cap);
        }
        a->append_all(b);
        alloc->deallocate_leaf(b);
        for (uint8_t i = idx; i < child_count_ - 1; i++) {
            children_[i] = children_[i + 1];
        }
        children_[idx] = a;
        child_sizes_.remove(idx, child_count_);
        child_sums_.remove(idx, child_count_);
        child_count_--;
    }
 
    template <class allocator>
    bool leaf_remove(dtype index, allocator* alloc) {
        uint8_t child_index = child_sizes_.find(index + 1);
        leaf_type* child = reinterpret_cast<leaf_type*>(children_[child_index]);
        if (child->size() <= leaf_size / 3) {
            if (child_index == 0) {
                leaf_type* sibling = reinterpret_cast<leaf_type*>(children_[1]);
                if (sibling->size() > leaf_size * 5 / 9) {
                    rebalance_leaves_right(child, sibling, alloc);
                } else {
                    merge_leaves(child, sibling, 0, alloc);
                }
                [[unlikely]] ((void)0);
            } else {
                leaf_type* sibling =
                    reinterpret_cast<leaf_type*>(children_[child_index - 1]);
                if (sibling->size() > leaf_size * 5 / 9) {
                    rebalance_leaves_left(sibling, child, child_index - 1,
                                          alloc);
                } else {
                    merge_leaves(sibling, child, child_index - 1, alloc);
                }
            }
            child_index = child_sizes_.find(index);
            [[unlikely]] child =
                reinterpret_cast<leaf_type*>(children_[child_index]);
        }
        if (child_index != 0) {
            [[likely]] index -= child_sizes_.get(child_index - 1);
        }
        bool value = child->remove(index);
        child_sizes_.increment(child_index, child_count_, -1);
        child_sums_.increment(child_index, child_count_, -int(value));
        return value;
    }
 
    void rebalance_nodes_right(node* a, node* b, uint8_t idx) {
        a->transfer_append(b, (b->child_count() - branches / 3) / 2);
        child_sizes_.set(idx, a->size());
        [[unlikely]] child_sums_.set(idx, a->p_sum());
    }
 
    void rebalance_nodes_left(node* a, node* b, uint8_t idx) {
        b->transfer_prepend(a, (a->child_count() - branches / 3) / 2);
        if (idx == 0) {
            child_sizes_.set(0, a->size());
            [[unlikely]] child_sums_.set(0, a->p_sum());
        } else {
            child_sizes_.set(idx, child_sizes_.get(idx - 1) + a->size());
            child_sums_.set(idx, child_sums_.get(idx - 1) + a->p_sum());
        }
    }
 
    template <class allocator>
    void merge_nodes(node* a, node* b, uint8_t idx, allocator* alloc) {
        a->append_all(b);
        alloc->deallocate_node(b);
        for (uint8_t i = idx; i < child_count_ - 1; i++) {
            children_[i] = children_[i + 1];
        }
        children_[idx] = a;
        child_sizes_.remove(idx, child_count_);
        child_sums_.remove(idx, child_count_);
        child_count_--;
    }
 
    template <class allocator>
    bool node_remove(dtype index, allocator* alloc) {
        uint8_t child_index = child_sizes_.find(index + 1);
        node* child = reinterpret_cast<node*>(children_[child_index]);
        if (child->child_count_ <= branches / 3) {
            if (child_index == 0) {
                node* sibling = reinterpret_cast<node*>(children_[1]);
                if (sibling->child_count_ > branches * 5 / 9) {
                    rebalance_nodes_right(child, sibling, 0);
                } else {
                    merge_nodes(child, sibling, 0, alloc);
                }
                [[unlikely]] ((void)0);
            } else {
                node* sibling =
                    reinterpret_cast<node*>(children_[child_index - 1]);
                if (sibling->child_count_ > branches * 5 / 9) {
                    rebalance_nodes_left(sibling, child, child_index - 1);
                } else {
                    merge_nodes(sibling, child, child_index - 1, alloc);
                }
            }
            child_index = child_sizes_.find(index + 1);
            [[unlikely]] child =
                reinterpret_cast<node*>(children_[child_index]);
        }
        if (child_index != 0) {
            [[likely]] index -= child_sizes_.get(child_index - 1);
        }
        bool value = child->remove(index, alloc);
        child_sizes_.increment(child_index, child_count_, -1);
        child_sums_.increment(child_index, child_count_, -int(value));
        return value;
    }
};
}  // namespace bv
#endif