C++ : B-Tree in C++20 (+Iterator support)

Question

I wrote a B-Tree implementation in C++20, based on my previous Red-Black Tree implementation.

Unit Test Demo : https://wandbox.org/permlink/Brw6TgAhdy89OIyj

Any feedback will be welcomed!

#include <algorithm>
#include <cassert>
#include <concepts>
#include <cstddef>
#include <functional>
#include <iostream>
#include <iterator>
#include <memory>
#include <numeric>
#include <random>
#include <ranges>
#include <utility>
#include <vector>

template <std::ranges::random_access_range R>
auto range_from(R&& r, int t) {
    return std::ranges::subrange(std::begin(r) + t, std::end(r));
}

template <typename T>
concept Key = std::regular<T> && std::totally_ordered<T>;

template <Key T, std::size_t t, std::predicate<T, T> Comp = std::less<T>>
class BTree {
    static_assert(t >= 2);
    class Node {
        std::size_t n = 0;
        bool leaf = true;
        Node* parent = nullptr;
        std::size_t index = 0;
        std::vector<T> key;
        std::vector<std::unique_ptr<Node>> child;
    public:
        // modifying n must be done only with setN,
        // for the invariants key.size() == n and child.size() == n + 1 (if exists)
        void setN(std::size_t N) {
            n = N;
            key.resize(n);
            if (!leaf) {
                child.resize(n + 1);
            }
        }

        void validateChild() {
            if (leaf) {
                return;
            }
            for (std::size_t i = 0; i <= n; i++) {
                child[i]->index = i;
                child[i]->parent = this;
            }
        }

        [[nodiscard]] std::size_t getN() const {
            return n;
        }

        [[nodiscard]] bool isFull() const {
            return n == 2 * t - 1;
        }

        [[nodiscard]] bool canTakeKey() const {
            return n > t - 1;
        }

        [[nodiscard]] bool hasMinimalKeys() const {
            return n == t - 1;
        }

        [[nodiscard]] bool isEmpty() const {
            return n == 0;
        }

        friend class BTree;

        Node* RightmostLeaf() {
            Node* curr = this;
            while (!curr->leaf) {
                curr = curr->child.back().get();
            }
            return curr;
        }

        Node* LeftmostLeaf() {
            Node* curr = this;
            while (!curr->leaf) {
                curr = curr->child.front().get();
            }
            return curr;
        }

        // merge child[i + 1] and key[i] into child[i]
        void Merge(std::size_t i) noexcept {
            assert(!leaf && child[i]->hasMinimalKeys() && child[i + 1]->hasMinimalKeys());
            child[i]->setN(2 * t - 1);
            child[i]->key[t - 1] = key[i];
            // bring keys of child[i + 1]
            std::ranges::move(child[i + 1]->key, child[i]->key.begin() + t);
            // bring children of child[i + 1]
            if (!child[i]->leaf) {
                std::ranges::move(child[i + 1]->child, child[i]->child.begin() + t);
            }
            // shift children from i + 1 left by 1 (because child[i + 1] is merged)
            std::shift_left(child.begin() + i + 1, child.end(), 1);
            // shift keys from i left by 1 (because key[i] is merged)
            std::shift_left(key.begin() + i, key.end(), 1);
            setN(n - 1);
            validateChild();
            child[i]->validateChild();
        }

        void LeftRotate() noexcept {
            assert(index + 1 < parent->child.size());
            auto sibling = parent->child[index + 1].get();
            // left rotation
            setN(n + 1);
            key.back() = parent->key[index];
            parent->key[index] = sibling->key.front();
            // shift all keys of right sibling left by 1
            std::shift_left(sibling->key.begin(), sibling->key.end(), 1);
            if (!leaf) {
                child[n] = std::move(sibling->child[0]);
                // shift all children of right sibling left by 1
                std::shift_left(sibling->child.begin(), sibling->child.end(), 1);
            }
            sibling->setN(sibling->getN() - 1);
            validateChild();
            sibling->validateChild();
        }

        void RightRotate() noexcept {
            assert(index - 1 < parent->child.size());
            auto sibling = parent->child[index - 1].get();
            // right rotation
            setN(n + 1);
            // shift all keys of node right by 1
            std::shift_right(key.begin(), key.end(), 1);
            key.front() = parent->key[index - 1];
            parent->key[index - 1] = sibling->key.back();
            if (!leaf) {
                // shift all children of node right by 1
                std::shift_right(child.begin(), child.end(), 1);
                child[0] = std::move(sibling->child[sibling->getN()]);
            }
            sibling->setN(sibling->getN() - 1);
            validateChild();
            sibling->validateChild();
        }
    };
    
    struct BTreeIterator {
        using difference_type = std::ptrdiff_t;
        using value_type = T;
        using pointer = T*;
        using reference = T&;
        using iterator_category = std::bidirectional_iterator_tag;
        
        Node* node;
        std::vector<T>::iterator it;
        
        void Increment() {
            if (it == node->key.end()) {
                return;
            }
            if (node->leaf) {
                ++it;
                while (node->parent && it == node->key.end()) {
                    it = node->parent->key.begin() + node->index;
                    node = node->parent;
                }
            } else {
                auto i = std::distance(node->key.begin(), it);
                node = node->child[i + 1]->LeftmostLeaf();
                it = node->key.begin();
            }
        }
        
        void Decrement() {
            auto i = std::distance(node->key.begin(), it);
            if (!node->leaf) {
                node = node->child[i]->RightmostLeaf();
                it = node->key.begin() + node->key.size() - 1;
            } else {
                if (i > 0) {
                    --it;
                } else {
                    while (node->parent && node->index == 0) {
                        node = node->parent;
                    }
                    if (node->index > 0) {
                        it = node->parent->key.begin() + node->index - 1;
                        node = node->parent;
                    }
                }
            }
            
        }
        
        BTreeIterator(Node* node, std::size_t i) : node {node} {
            assert(node && i <= node->key.size());
            it = node->key.begin() + i;
        }
        
        reference operator*() const {
            return *it;
        }
        
        pointer operator->() const {
            return it;
        }
        
        BTreeIterator& operator++() {
            Increment();
            return *this;
        }
        
        BTreeIterator operator++(int) {
            BTreeIterator temp = *this;
            Increment();
            return temp;
        }
        
        BTreeIterator& operator--() {
            Decrement();
            return *this;
        }
        
        BTreeIterator operator--(int) {
            BTreeIterator temp = *this;
            Decrement();
            return temp;
        }
        
        friend bool operator==(const BTreeIterator& x, const BTreeIterator& y) {
            return x.node == y.node && x.it == y.it;
        }
        
        friend bool operator!=(const BTreeIterator& x, const BTreeIterator& y) {
            return !(x == y);
        }
    };
    
    struct BTreeConstIterator {
        using difference_type = std::ptrdiff_t;
        using value_type = T;
        using pointer = const T*;
        using reference = const T&;
        using iterator_category = std::bidirectional_iterator_tag;
        
        const Node* node;
        std::vector<T>::const_iterator it;
        
        void Increment() {
            if (it == node->key.cend()) {
                return;
            }
            if (node->leaf) {
                ++it;
                while (node->parent && it == node->key.cend()) {
                    it = node->parent->key.cbegin() + node->index;
                    node = node->parent;
                }
            } else {
                auto i = std::distance(node->key.cbegin(), it);
                node = node->child[i + 1]->LeftmostLeaf();
                it = node->key.cbegin();
            }
        }
        
        void Decrement() {
            auto i = std::distance(node->key.cbegin(), it);
            if (!node->leaf) {
                node = node->child[i]->RightmostLeaf();
                it = node->key.cbegin() + node->key.size() - 1;
            } else {
                if (i > 0) {
                    --it;
                } else {
                    while (node->parent && node->index == 0) {
                        node = node->parent;
                    }
                    if (node->index > 0) {
                        it = node->parent->key.cbegin() + node->index - 1;
                        node = node->parent;
                    }
                }
            }
            
        }
        
        BTreeConstIterator(const Node* node, std::size_t i) : node {node} {
            assert(node && i <= node->key.size());
            it = node->key.cbegin() + i;
        }
        
        reference operator*() const {
            return *it;
        }
        
        pointer operator->() const {
            return it;
        }
        
        BTreeConstIterator& operator++() {
            Increment();
            return *this;
        }
        
        BTreeConstIterator operator++(int) {
            BTreeConstIterator temp = *this;
            Increment();
            return temp;
        }
        
        BTreeConstIterator& operator--() {
            Decrement();
            return *this;
        }
        
        BTreeConstIterator operator--(int) {
            BTreeConstIterator temp = *this;
            Decrement();
            return temp;
        }
        
        friend bool operator==(const BTreeConstIterator& x, const BTreeConstIterator& y) {
            return x.node == y.node && x.it == y.it;
        }
        
        friend bool operator!=(const BTreeConstIterator& x, const BTreeConstIterator& y) {
            return !(x == y);
        }
    };

    std::unique_ptr<Node> root;
    
    using iterator = BTreeIterator;
    using const_iterator = BTreeConstIterator;
    using reverse_iterator = std::reverse_iterator<iterator>;
    using const_reverse_iterator = std::reverse_iterator<const_iterator>;
    
    std::pair<const Node*, std::size_t> Search(const Node* x, const T& k) const {
        std::size_t i = std::distance(x->key.begin(), std::ranges::lower_bound(x->key, k, Comp()));
        if (i < x->getN() && k == x->key[i]) { // equal? key found
            return {x, i};
        } else if (x->leaf) { // no child, key is not in the tree
            return {nullptr, 0};
        } else { // search on child between range
            return Search(x->child[i].get(), k);
        }
    }

    void InsertNonFull(Node* x, const T& k) {
        if (x->leaf) { // key should be inserted only at leaf
            InsertToLeaf(x, k);
        } else {
            auto i = std::distance(x->key.begin(), std::ranges::upper_bound(x->key, k, Comp()));
            if (x->child[i]->isFull()) { // is full? then split
                SplitChild(x, i);
                if (Comp()(x->key[i], k)) {
                    i++;
                }
            }
            InsertNonFull(x->child[i].get(), k); // recursively insert
        }
    }

    void InsertToLeaf(Node* node, const T& k) {
        assert(node->leaf);
        auto i = std::distance(node->key.begin(), std::ranges::upper_bound(node->key, k, Comp()));
        node->setN(node->getN() + 1);
        std::shift_right(node->key.begin() + i, node->key.end(), 1);
        node->key[i] = k;
        // n cannot exceeds 2t - 1, because in that case
        // its parent should've called SplitChild on x before
        assert(node->getN() < 2 * t);
    }

    // split x->child[i]
    void SplitChild(Node* x, std::size_t i) noexcept {
        if (!x) {
            return;
        }
        auto y = x->child[i].get();
        if (!y) {
            return;
        }
        // x cannot be full, because in that case its parent should've called SplitChild on x before
        assert(!x->isFull() && y->isFull());
        auto z = std::make_unique<Node>(); // will be y's right sibling
        z->leaf = y->leaf;
        z->setN(t - 1);
        // bring right half keys from y
        std::ranges::move(range_from(y->key, t), z->key.begin());
        if (!y->leaf) {
            // bring right half children from y
            std::ranges::move(range_from(y->child, t), z->child.begin());
            z->validateChild();
        }
        x->setN(x->getN() + 1);
        // shift children of x right by 1 from i + 1
        std::shift_right(x->child.begin() + i + 1, x->child.end(), 1);
        x->child[i + 1] = std::move(z);
        // shift keys of x right by 1 from i
        std::shift_right(x->key.begin() + i, x->key.end(), 1);
        x->key[i] = y->key[t - 1];
        y->setN(t - 1);
        x->validateChild();
        y->validateChild();
    }

    void Delete(Node* x, const T& k) noexcept {
        std::size_t i = std::distance(x->key.begin(), std::ranges::lower_bound(x->key, k, Comp()));
        if (i < x->getN() && k == x->key[i]) { // equal? key found
            Delete(x, k, i);
        } else if (x->leaf) { // no child, key is not in the tree
            return;
        } else { // search on child between range
            Node* next = x->child[i].get();
            if (x->child[i]->hasMinimalKeys()) {
                if (i + 1 < x->child.size() && x->child[i + 1]->canTakeKey()) {
                    x->child[i]->LeftRotate();
                } else if (i - 1 < x->child.size() && x->child[i - 1]->canTakeKey()) {
                    x->child[i]->RightRotate();
                } else if (i + 1 < x->child.size()) {
                    x->Merge(i);
                    next = x->child[i].get();
                    if (x == root.get() && x->isEmpty()) {
                        // shrink tree in height, merged child should be a new root
                        root = std::move(x->child[i]);
                        root->parent = nullptr;
                    }
                } else if (i - 1 < x->child.size()) {
                    x->Merge(i - 1);
                    next = x->child[i - 1].get();
                    if (x == root.get() && x->isEmpty()) {
                        // shrink tree in height, merged child should be a new root
                        root = std::move(x->child[i - 1]);
                        root->parent = nullptr;
                    }
                }
            }
            Delete(next, k);
        }
    }

    void Delete(Node* x, const T& k, std::size_t i) noexcept {
        assert(x->key[i] == k);
        if (x->leaf) {
            // directly erase from leaf
            DeleteToLeaf(x, i);
        } else if (x->child[i]->canTakeKey()) {
            // find predecessor and swap keys
            auto predLeaf = x->child[i]->RightmostLeaf();
            std::swap(predLeaf->key.back(), x->key[i]);
            // now k is in left child, search there
            Delete(x->child[i].get(), k);
        } else if (x->child[i + 1]->canTakeKey()) {
            // find successor and swap keys
            auto succLeaf = x->child[i + 1]->LeftmostLeaf();
            std::swap(succLeaf->key.front(), x->key[i]);
            // now k is in right child, search there
            Delete(x->child[i + 1].get(), k);
        } else {
            // merge two children
            x->Merge(i);
            Node* next = x->child[i].get();
            if (x == root.get() && x->getN() == 0) {
                // shrink tree in height, merged child should be a new root
                root = std::move(x->child[i]);
                root->parent = nullptr;
            }
            Delete(next, k);
        }
    }

    void DeleteToLeaf(Node* node, std::size_t i) {
        assert(node->leaf);
        std::shift_left(node->key.begin() + i, node->key.end(), 1);
        node->setN(node->getN() - 1);
        assert(node == root.get() || node->getN() >= t - 1);
    }
    
    void ValidateIterators() {
        begin_ = iterator(root->LeftmostLeaf(), 0);
        cbegin_ = const_iterator(root->LeftmostLeaf(), 0);
        end_ = iterator(root.get(), root->getN());
        cend_ = const_iterator(root.get(), root->getN());
    }
    
    iterator begin_;
    const_iterator cbegin_;
    iterator end_;
    const_iterator cend_;

public:
    BTree() : root {std::make_unique<Node>()},
        begin_ {root.get(), 0}, cbegin_ {root.get(), 0}, end_ {root.get(), 0}, cend_ {root.get(), 0} {
    }

    [[nodiscard]] std::pair<const Node*, std::size_t> Search(const T& k) const {
        return Search(root.get(), k);
    }

    void Insert(const T& k) noexcept {
        if (root->isFull()) { // if root is full then make it as a child of new root - and split
            auto s = std::make_unique<Node>();
            s->leaf = false;
            s->setN(0);
            s->child[0] = std::move(root);
            root = std::move(s);
            SplitChild(root.get(), 0);
            InsertNonFull(root.get(), k);
        } else {
            InsertNonFull(root.get(), k);
        }
        ValidateIterators();
    }

    void Delete(const T& k) {
        Delete(root.get(), k);
        ValidateIterators();
    }
    
    iterator begin() {
        return begin_;
    }
    
    const_iterator begin() const {
        return cbegin_;
    }
    
    const_iterator cbegin() const {
        return cbegin_;
    }
    
    iterator end() {
        return end_;
    }
    
    const_iterator end() const {
        return cend_;
    }
    
    const_iterator cend() const {
        return cend_;
    }
    
    reverse_iterator rbegin() {
        return reverse_iterator(end_);
    }
    
    const_reverse_iterator rbegin() const {
        return const_reverse_iterator(cend_);
    }
    
    const_reverse_iterator crbegin() const {
        return const_reverse_iterator(cend_);
    }
    
    reverse_iterator rend() {
        return reverse_iterator(begin_);
    }
    
    const_reverse_iterator rend() const {
        return const_reverse_iterator(cbegin_);
    }
    
    const_reverse_iterator crend() const {
        return const_reverse_iterator(cbegin_);
    }
    

};

int main() {
    BTree<int, 2> tree;

    constexpr std::size_t N = 100;

    std::vector<int> v (N);
    std::iota(v.begin(), v.end(), 1);
    std::mt19937 gen(std::random_device{}());
    std::ranges::shuffle(v, gen);

    for (auto n : v) {
        tree.Insert(n);
        for (const auto& key : tree) {
            std::cout << key << ' ';
        }
        std::cout << '\n';
    }

    assert(std::ranges::all_of(v, [&tree](auto n){return tree.Search(n).first != nullptr;}));
   

    // should output 1 2 3 ... N
    for (const auto& key : tree) {
        std::cout << key << ' ';
    }
    std::cout << '\n';

    std::ranges::shuffle(v, gen);

    for (auto n : v) {
        tree.Delete(n);
        for (const auto& key : tree) {
            std::cout << key << ' ';
        }
        std::cout << '\n';
    }
}

Answer 1

Don't restrict the type of key

You'll notice that ordered STL containers such as std::map put no constraints on the type of the key, neither using SFINAE, concepts nor a mention in the description. The reason is that you are allowed to specify your own compare function, which you can use to provide a way to order keys even if the keys themselves do not have a total order.

As an example: suppose I want to store two-dimensional coordinates in a B-tree, and order them based on their x coordinate. Two-dimensional coordinates themselves do not have a well-defined ordering, so your concept Key would prevent it from being usable with BTree, even though it is trivial to write a comparison function that just compares the x coordinate of two keys.

Naming things

Avoid very short names for variables and types, unless it is a commonly used name. For example, T is fine for the template type of a key or value, n is a common abbreviation for a "number" of things. But t is a bad choice for the number of children of a node. I suggest you replace it with fanout.

I also recommend you use the plural for names of containers that can hold multiple elements. So keys instead of key, children instead of child.

Be consistent when naming things. I see both camelCase and PascalCase used for member functions.

Add more assert()s where appropriate

You already use assert() in a few cases, but it can be done in a lot more places. For example, in BTree::Node::setN, you can add:

assert(N <= 2 * t - 1);

The iterator operators could also use some assert() statements to check that you don't try to iterate past the beginning or end of the tree, and so on.

Optimize your iterators

There's a little bit of redundancy in your iterators:

Node* node;
std::vector<T>::iterator it;

Here, it also contains a pointer to node->key. It's better to just store the integer index into node->key. That way, you also don't have to jump through hoops to get an index into node->child as well:

Node *node;
size_t index;

Optimize your Nodes

You also store some redundant information in each Node. Consider that key.size() is equal to n, and child.size() is equal to n + 1 for interior nodes, and child.empty() == true for leaf nodes. So n and leaf store redundant information. Removing those two variables gets rid of 16 bytes on 64-bit architectures, and there's less state that needs to be kept in sync.

In principle you could also remove parent and index. This makes the iterators and the rotation operations more complex though, so I would probably keep it as it is now.