B Tree

btree.cpp

#include <iostream>
#include <vector>
#include <algorithm>
#include <unordered_set>

#include "btree.h"

auto Node::FindExactKeyIdx(KeyType key, int *idx) -> bool {
    int i = 0;
    while (i < num_keys_) {
        if (keys_[i] == key) {
            *idx = i;
            return true;
        }
        i++;
    }
    return false;
}

auto Node::FindKeyIdx(KeyType key) -> int {
    int i = 0;
    while (i < num_keys_ && key > keys_[i]) {
        i++;
    }
    return i;
}

void Node::RemoveKVAt(int idx) {
    // 将 idx 之后的 Key 和 Value 往前挪一格，覆盖掉原有的
    for (int i = idx; i < GetNumKeys() - 1; ++i) {
        SetKey(i, GetKey(i + 1));
        SetValue(i, GetValue(i + 1));
    }
    DecrementNumKeys();
}

auto BTree::Search(const KeyType key, ValueType *value) -> bool {
    if (root_page_id_ == INVALID_PAGE_ID) {
        return false;
    }

    auto current_page = disk_manager_.GetPage(root_page_id_);
    while (true) {
        auto current_node = current_page->GetAs();

        // 是否为叶子节点
        if (current_node->IsLeaf()) {
            int idx;
            auto result = current_node->FindExactKeyIdx(key, &idx);
            if (result) {
                *value = current_node->GetValue(idx);
            }
            return result;
        }

        // 是内部节点
        auto idx = current_node->FindKeyIdx(key);

        // 检查 idx 是否刚好就是 key
        if (idx < current_node->GetNumKeys() && current_node->GetKey(idx) == key) {
            *value = current_node->GetValue(idx);
            return true;
        }

        // 不是 key，则获取子节点
        auto child_page_id = current_node->GetChild(idx);
        current_page = disk_manager_.GetPage(child_page_id);
    }
}

auto Node::InsertKeyValue(KeyType key, ValueType value, page_id_t child_page_id) -> bool {
    auto idx = FindKeyIdx(key);

    // 移动 child
    for (int i = GetNumKeys(); i > idx; --i) {
        SetChild(i + 1, GetChild(i));
    }

    // 移动 kv
    for (int i = GetNumKeys() - 1; i >= idx; --i) {
        SetKey(i + 1, GetKey(i));
        SetValue(i + 1, GetValue(i));
    }

    SetKey(idx, key);
    SetValue(idx, value);
    if (!IsLeaf() && child_page_id != INVALID_PAGE_ID) {
        SetChild(idx + 1, child_page_id); // 应该是插入到右侧子指针
    }
    IncrementNumKeys();
    return true;
}

auto BTree::Insert(KeyType key, ValueType value) -> bool {
    if (root_page_id_ == INVALID_PAGE_ID) {
        auto new_root_page_id = disk_manager_.GetNewPage(true);
        root_page_id_ = new_root_page_id;
        auto root_page = disk_manager_.GetPage(root_page_id_);
        auto root_node = root_page->GetMutAs();
        root_node->InsertKeyValue(key, value);
        return true;
    }

    auto current_page = disk_manager_.GetPage(root_page_id_);
    Node *parent_node = nullptr;
    while (true) {
        auto current_node = current_page->GetAs();

        if (current_node->IsFull()) {
            if (current_node->IsLeaf()) {
                auto new_page_id = disk_manager_.GetNewPage(true);
                auto new_page = disk_manager_.GetPage(new_page_id);
                auto new_node = new_page->GetMutAs();
                // split
                std::vector<std::pair<KeyType, ValueType> > tmp_kv;

                auto count = current_node->GetNumKeys() - MIN_KEY_NUM;
                while (count--) {
                    tmp_kv.push_back({
                        current_node->GetKey(current_node->GetNumKeys() - 1),
                        current_node->GetValue(current_node->GetNumKeys() - 1)
                    });
                    current_node->DecrementNumKeys();
                }
                // 反转 tmp_kv
                std::reverse(tmp_kv.begin(), tmp_kv.end());

                // 第一个 kv 作为新页的分隔 kv
                auto [separator_key, separator_value] = tmp_kv[0];
                for (auto i = 1; i < tmp_kv.size(); ++i) {
                    auto [k, v] = tmp_kv[i];
                    new_node->SetKey(new_node->GetNumKeys(), k);
                    new_node->SetValue(new_node->GetNumKeys(), v);
                    new_node->IncrementNumKeys();
                }

                if (parent_node == nullptr) {
                    // 是根节点，创建新内部节点，高度 +1
                    auto new_root_page_id = disk_manager_.GetNewPage(false);
                    auto new_root_page = disk_manager_.GetPage(new_root_page_id);
                    auto new_root_node = new_root_page->GetMutAs();
                    new_root_node->SetKey(0, separator_key);
                    new_root_node->SetValue(0, separator_value);
                    new_root_node->SetChild(0, current_page->GetPageId());
                    new_root_node->SetChild(1, new_page_id);
                    new_root_node->IncrementNumKeys();
                    root_page_id_ = new_root_page_id;

                    parent_node = new_root_node;
                } else {
                    parent_node->InsertKeyValue(separator_key, separator_value, new_page_id);
                }

                if (separator_key == key) {
                    int idx;
                    parent_node->FindExactKeyIdx(key, &idx);
                    parent_node->SetValue(idx, value);
                    return true;
                }

                if (separator_key < key) {
                    current_node = new_node;
                }
            } else {
                auto new_page_id = disk_manager_.GetNewPage(false);
                auto new_page = disk_manager_.GetPage(new_page_id);
                auto new_node = new_page->GetMutAs();

                // split
                std::vector<std::pair<KeyType, ValueType> > tmp_kv;
                std::vector<page_id_t> tmp_child_page_ids;

                auto count = current_node->GetNumKeys() - MIN_KEY_NUM;
                while (count--) {
                    auto num_keys = current_node->GetNumKeys();
                    tmp_kv.push_back({current_node->GetKey(num_keys - 1), current_node->GetValue(num_keys - 1)});
                    tmp_child_page_ids.push_back(current_node->GetChild(num_keys));
                    current_node->DecrementNumKeys();
                }

                // 不需要额外取，因为 tmp_kv[0] 是放入 parent_node 的
                // // 额外取一个 child
                // tmp_child_page_ids.push_back(current_node->GetChild(current_node->GetNumKeys()));

                std::reverse(tmp_kv.begin(), tmp_kv.end());
                std::reverse(tmp_child_page_ids.begin(), tmp_child_page_ids.end());

                auto [separator_key, separator_value] = tmp_kv[0];
                for (auto i = 1; i < tmp_kv.size(); ++i) {
                    auto [k, v] = tmp_kv[i];
                    new_node->SetKey(new_node->GetNumKeys(), k);
                    new_node->SetValue(new_node->GetNumKeys(), v);
                    new_node->IncrementNumKeys();
                }

                // 放入 child
                for (auto i = 0; i < tmp_child_page_ids.size(); ++i) {
                    new_node->SetChild(i, tmp_child_page_ids[i]);
                }

                if (parent_node == nullptr) {
                    // 是根节点，创建新内部节点，高度 +1
                    auto new_root_page_id = disk_manager_.GetNewPage(false);
                    auto new_root_page = disk_manager_.GetPage(new_root_page_id);
                    auto new_root_node = new_root_page->GetMutAs();
                    new_root_node->SetKey(0, separator_key);
                    new_root_node->SetValue(0, separator_value);
                    new_root_node->SetChild(0, current_page->GetPageId());
                    new_root_node->SetChild(1, new_page_id);
                    new_root_node->IncrementNumKeys();
                    root_page_id_ = new_root_page_id;

                    parent_node = new_root_node;
                } else {
                    parent_node->InsertKeyValue(separator_key, separator_value, new_page_id);
                }

                if (separator_key == key) {
                    int idx;
                    parent_node->FindExactKeyIdx(key, &idx);
                    parent_node->SetValue(idx, value);
                    return true;
                }

                if (separator_key < key) {
                    current_node = new_node;
                }
            }
        }

        // 是否为叶子节点
        if (current_node->IsLeaf()) {
            int idx;
            auto result = current_node->FindExactKeyIdx(key, &idx);
            if (result) {
                current_node->SetValue(idx, value);
            } else {
                // 不存在 key，InsertKeyValue
                current_node->InsertKeyValue(key, value);
            }
            return true;
        }

        // 是内部节点
        auto idx = current_node->FindKeyIdx(key);

        // 检查 idx 是否刚好就是 key
        if (idx < current_node->GetNumKeys() && current_node->GetKey(idx) == key) {
            current_node->SetValue(idx, value);
            return true;
        }

        // 不是 key，则获取子节点
        auto child_page_id = current_node->GetChild(idx);
        parent_node = current_node;
        current_page = disk_manager_.GetPage(child_page_id);
    }
}

auto BTree::Remove(KeyType key) -> bool {
    if (root_page_id_ == INVALID_PAGE_ID) {
        return false; // 树为空
    }

    bool result = RemoveInternal(root_page_id_, key);

    // 根节点为空，唯一子节点成为根节点
    auto root_page = disk_manager_.GetPage(root_page_id_);
    auto root_node = root_page->GetMutAs();

    if (result && root_node->GetNumKeys() == 0) {
        if (root_node->IsLeaf()) {
            // 树全空
            root_page_id_ = INVALID_PAGE_ID;
        } else {
            root_page_id_ = root_node->GetChild(0);
        }
    }
    return result;
}

auto BTree::RemoveInternal(page_id_t page_id, KeyType key) -> bool {
    auto page = disk_manager_.GetPage(page_id);
    auto node = page->GetMutAs();
    int idx = node->FindKeyIdx(key);

    // Key 存在于当前节点 node 中
    if (idx < node->GetNumKeys() && node->GetKey(idx) == key) {
        // [1]：如果是叶子节点，直接删除
        if (node->IsLeaf()) {
            node->RemoveKVAt(idx);
            return true;
        }

        // [2]：如果是内部节点
        page_id_t left_child_id = node->GetChild(idx);
        page_id_t right_child_id = node->GetChild(idx + 1);
        auto left_child = disk_manager_.GetPage(left_child_id)->GetMutAs();
        auto right_child = disk_manager_.GetPage(right_child_id)->GetMutAs();

        // [2a]：左子节点有多余 key，可以是借用
        if (left_child->GetNumKeys() >= MIN_DEGREE) {
            auto pred = GetPredecessor(left_child_id); // 找前驱
            node->SetKey(idx, pred.first);
            node->SetValue(idx, pred.second);
            return RemoveInternal(left_child_id, pred.first); //
        }

        // [2b]：右子节点有多余 key，借用
        if (right_child->GetNumKeys() >= MIN_DEGREE) {
            auto succ = GetSuccessor(right_child_id); // 找后继
            node->SetKey(idx, succ.first);
            node->SetValue(idx, succ.second);
            return RemoveInternal(right_child_id, succ.first);
        }

        // [2c]：左右孩子都只有 t-1 个 key，执行合并
        Merge(node, idx); // 将 key 和 right_child 合并进 left_child
        return RemoveInternal(left_child_id, key); // 在合并后的左孩子中继续递归删除 key
    }

    // Key 当前不在内部节点 node 中
    if (node->IsLeaf()) {
        return false; // 不存在
    }

    // 找到包含 key 的子节点
    bool is_last_child = (idx == node->GetNumKeys());
    page_id_t child_id = node->GetChild(idx);
    auto child = disk_manager_.GetPage(child_id)->GetMutAs();

    // 子节点只有 t-1 个 key，执行 3a 或 3b 填充它
    if (child->GetNumKeys() < MIN_DEGREE) {
        FillChild(node, idx);
    }

    // 如何合并后 idx 变了
    if (is_last_child && idx > node->GetNumKeys()) {
        return RemoveInternal(node->GetChild(idx - 1), key);
    }

    return RemoveInternal(node->GetChild(idx), key);
}

auto BTree::GetPredecessor(page_id_t page_id) -> std::pair<KeyType, ValueType> {
    auto curr = disk_manager_.GetPage(page_id)->GetAs();
    // 一直往最右侧子节点走，直到叶子
    while (!curr->IsLeaf()) {
        curr = disk_manager_.GetPage(curr->GetChild(curr->GetNumKeys()))->GetAs();
    }
    return {curr->GetKey(curr->GetNumKeys() - 1), curr->GetValue(curr->GetNumKeys() - 1)};
}

auto BTree::GetSuccessor(page_id_t page_id) -> std::pair<KeyType, ValueType> {
    auto curr = disk_manager_.GetPage(page_id)->GetAs();
    // 一直往最左侧子节点走，直到叶子
    while (!curr->IsLeaf()) {
        curr = disk_manager_.GetPage(curr->GetChild(0))->GetAs();
    }
    return {curr->GetKey(0), curr->GetValue(0)};
}

// 确保子节点至少有 t 个 key
void BTree::FillChild(Node *parent, int child_idx) {
    page_id_t prev_id = (child_idx > 0) ? parent->GetChild(child_idx - 1) : INVALID_PAGE_ID;
    page_id_t next_id = (child_idx < parent->GetNumKeys()) ? parent->GetChild(child_idx + 1) : INVALID_PAGE_ID;

    Node *prev_child = (prev_id != INVALID_PAGE_ID) ? disk_manager_.GetPage(prev_id)->GetMutAs() : nullptr;
    Node *next_child = (next_id != INVALID_PAGE_ID) ? disk_manager_.GetPage(next_id)->GetMutAs() : nullptr;

    // [3a]：左兄弟至少有 t 个，向左借
    if (prev_child && prev_child->GetNumKeys() >= MIN_DEGREE) {
        BorrowFromPrev(parent, child_idx);
    }
    // [3a]：右兄弟至少有 t 个，向右借
    else if (next_child && next_child->GetNumKeys() >= MIN_DEGREE) {
        BorrowFromNext(parent, child_idx);
    }
    // [3b]：都是 t - 1 则左或者右合并
    else {
        if (child_idx != parent->GetNumKeys()) {
            Merge(parent, child_idx); // 如果不是最右孩子，与后一个合并
        } else {
            Merge(parent, child_idx - 1); // 如果是最右孩子，与前一个合并
        }
    }
}

// 从左兄弟借用一个
void BTree::BorrowFromPrev(Node *parent, int child_idx) {
    auto child = disk_manager_.GetPage(parent->GetChild(child_idx))->GetMutAs();
    auto sibling = disk_manager_.GetPage(parent->GetChild(child_idx - 1))->GetMutAs();

    // 将 child 中的所有元素往后挪一格，在头部腾出空间
    for (int i = child->GetNumKeys(); i > 0; --i) {
        child->SetKey(i, child->GetKey(i - 1));
        child->SetValue(i, child->GetValue(i - 1));
    }
    if (!child->IsLeaf()) {
        for (int i = child->GetNumKeys() + 1; i > 0; --i) {
            child->SetChild(i, child->GetChild(i - 1));
        }
    }

    // 将父节点的 key 放入 child 的头部
    child->SetKey(0, parent->GetKey(child_idx - 1));
    child->SetValue(0, parent->GetValue(child_idx - 1));
    if (!child->IsLeaf()) {
        child->SetChild(0, sibling->GetChild(sibling->GetNumKeys()));
    }
    child->IncrementNumKeys();

    // 将左兄弟的最后一个 key 放入父节点，填补之前下移的空位
    parent->SetKey(child_idx - 1, sibling->GetKey(sibling->GetNumKeys() - 1));
    parent->SetValue(child_idx - 1, sibling->GetValue(sibling->GetNumKeys() - 1));

    sibling->DecrementNumKeys();
}


// 从右兄弟借用一个
void BTree::BorrowFromNext(Node *parent, int child_idx) {
    auto child = disk_manager_.GetPage(parent->GetChild(child_idx))->GetMutAs();
    auto right_sibling = disk_manager_.GetPage(parent->GetChild(child_idx + 1))->GetMutAs();

    // 将父节点的 key 放入 child 的尾部
    child->SetKey(child->GetNumKeys(), parent->GetKey(child_idx));
    child->SetValue(child->GetNumKeys(), parent->GetValue(child_idx));
    if (!child->IsLeaf()) {
        child->SetChild(child->GetNumKeys() + 1, right_sibling->GetChild(0));
    }
    child->IncrementNumKeys();

    // 将右兄弟的第一个 key 放回父节点
    parent->SetKey(child_idx, right_sibling->GetKey(0));
    parent->SetValue(child_idx, right_sibling->GetValue(0));

    // 右兄弟全部前移一个索引，填补回去
    for (int i = 0; i < right_sibling->GetNumKeys() - 1; ++i) {
        right_sibling->SetKey(i, right_sibling->GetKey(i + 1));
        right_sibling->SetValue(i, right_sibling->GetValue(i + 1));
    }
    if (!right_sibling->IsLeaf()) {
        for (int i = 0; i < right_sibling->GetNumKeys(); ++i) {
            // child 比 key 数组长 1
            right_sibling->SetChild(i, right_sibling->GetChild(i + 1));
        }
    }
    right_sibling->DecrementNumKeys();
}

void BTree::Merge(Node *parent, int child_idx) {
    auto left_child = disk_manager_.GetPage(parent->GetChild(child_idx))->GetMutAs();
    auto right_child = disk_manager_.GetPage(parent->GetChild(child_idx + 1))->GetMutAs();

    // 将 parent 里的下放 key 放到 left_child
    left_child->SetKey(MIN_DEGREE - 1, parent->GetKey(child_idx));
    left_child->SetValue(MIN_DEGREE - 1, parent->GetValue(child_idx));

    // 把 right_child 里的所有 key 和 child 移植到 left_child 的后半截
    for (int i = 0; i < right_child->GetNumKeys(); ++i) {
        left_child->SetKey(i + MIN_DEGREE, right_child->GetKey(i));
        left_child->SetValue(i + MIN_DEGREE, right_child->GetValue(i));
    }
    if (!left_child->IsLeaf()) {
        for (int i = 0; i <= right_child->GetNumKeys(); ++i) {
            left_child->SetChild(i + MIN_DEGREE, right_child->GetChild(i));
        }
    }

    // left_child 现在直接满员
    left_child->SetNumKeys(MAX_KEY_NUM);

    // 删掉 parent 中的分隔 key
    for (int i = child_idx; i < parent->GetNumKeys() - 1; ++i) {
        parent->SetKey(i, parent->GetKey(i + 1));
        parent->SetValue(i, parent->GetValue(i + 1));
    }
    for (int i = child_idx + 1; i < parent->GetNumKeys(); ++i) {
        parent->SetChild(i, parent->GetChild(i + 1));
    }
    parent->DecrementNumKeys();
}


int main() {
    BTree btree;
    btree.Insert(1, 1);
    ValueType val;
    btree.Search(1, &val);
    btree.Remove(1);
}

btree.h

#include <cstdint>
#include <iostream>
#include <utility>
#include <vector>
#include <memory>


#define INVALID_PAGE_ID (-1)

#define MIN_DEGREE 2

#define MAX_KEY_NUM (MIN_DEGREE * 2 - 1)
#define MIN_KEY_NUM (MIN_DEGREE - 1)

#define MAX_CHILD_NUM (MIN_DEGREE * 2)
#define MIN_CHILD_NUM MIN_DEGREE

using page_id_t = uint32_t;
using KeyType = uint32_t;
using ValueType = uint32_t;

class Node {
public:
    explicit Node(const bool is_leaf) : isLeaf_(is_leaf), num_keys_(0) {
        keys_.resize(MAX_KEY_NUM + 1);
        values_.resize(MAX_KEY_NUM + 1);
        children_page_id_.resize(MAX_CHILD_NUM + 1);
    }

    auto IsLeaf() const -> bool { return isLeaf_; }
    void SetIsLeaf(const bool is_leaf) { isLeaf_ = is_leaf; }

    auto IsFull() const -> bool { return num_keys_ >= MAX_KEY_NUM; }

    auto FindExactKeyIdx(KeyType key, int *idx) -> bool;

    auto FindKeyIdx(KeyType key) -> int;

    auto InsertKeyValue(KeyType key, ValueType value, page_id_t child_page_id = INVALID_PAGE_ID) -> bool;

    auto GetKey(const int key_idx) const -> KeyType { return keys_[key_idx]; }
    void SetKey(const int key_idx, const KeyType key) { keys_[key_idx] = key; }

    auto GetValue(const int idx) const -> ValueType { return values_[idx]; }
    void SetValue(const int idx, const ValueType value) { values_[idx] = value; }

    auto GetChild(const int idx) const -> page_id_t { return children_page_id_[idx]; }
    void SetChild(const int idx, const page_id_t child_page_id) { children_page_id_[idx] = child_page_id; }

    auto GetNumKeys() const -> int { return num_keys_; }
    void SetNumKeys(const int num_keys) { num_keys_ = num_keys; }

    void IncrementNumKeys() { num_keys_++; }
    void DecrementNumKeys() { num_keys_--; }

    void AddKeyValue(KeyType key, ValueType value) {
        keys_.push_back(key);
        values_.push_back(value);
        num_keys_++;
    }

    void AddKey(KeyType key) {
        keys_.push_back(key);
        num_keys_++;
    }

    void AddChild(page_id_t child_page_id) {
        children_page_id_.push_back(child_page_id);
    }

    void ResizeKeys(int size) {
        keys_.resize(size);
    }

    void ResizeValues(int size) {
        values_.resize(size);
    }

    void ResizeChildren(int size) {
        children_page_id_.resize(size);
    }

    void RemoveKVAt(int idx);




private:
    bool isLeaf_;

    int num_keys_;

    std::vector<KeyType> keys_;
    std::vector<page_id_t> children_page_id_;
    std::vector<ValueType> values_;
};

class Page {
public:
    Page(const page_id_t page_id, const bool isLeaf) : page_id_(page_id), node_(isLeaf) {
    }

    auto GetAs() -> Node * {
        return &node_;
    }

    auto GetMutAs() -> Node * {
        return &node_;
    }

    auto GetPageId() const -> page_id_t {return page_id_;}

private:
    page_id_t page_id_;
    Node node_;
};

class DiskManager {
public:
    DiskManager() : num_pages_(0) {
    }

    page_id_t GetNewPage(bool isLeaf) {
        auto new_page_id = static_cast<page_id_t>(num_pages_++);
        pages_.push_back(std::make_unique<Page>(new_page_id, isLeaf));
        return new_page_id;
    }

    auto GetPage(const page_id_t page_id) -> Page * {
        if (page_id >= 0 && page_id < static_cast<page_id_t>(pages_.size())) {
            return pages_[page_id].get();
        }
        return nullptr;
    }

private:
    int num_pages_;
    std::vector<std::unique_ptr<Page>> pages_;
};

class BTree {
public:
    BTree() : root_page_id_(INVALID_PAGE_ID) { disk_manager_ = DiskManager(); }

    BTree(page_id_t root_page_id, DiskManager&& disk_manager)
        : root_page_id_(root_page_id), disk_manager_(std::move(disk_manager)) {}

    auto Search(KeyType key, ValueType *value) -> bool;

    auto Insert(KeyType key, ValueType value) -> bool;

    auto Remove(KeyType key) -> bool;

private:
    auto RemoveInternal(page_id_t page_id, KeyType key) -> bool;

    // 获取前驱和后继
    auto GetPredecessor(page_id_t page_id) -> std::pair<KeyType, ValueType>;
    auto GetSuccessor(page_id_t page_id) -> std::pair<KeyType, ValueType>;

    void FillChild(Node *parent, int child_idx);

    // 3a：向左/右兄弟借 kv
    void BorrowFromPrev(Node *parent, int child_idx);
    void BorrowFromNext(Node *parent, int child_idx);

    // 2c / 3b：合并兄弟节点
    void Merge(Node *parent, int child_idx);

    page_id_t root_page_id_;

    DiskManager disk_manager_;
};