bplus_tree.h

/* Copyright (c) 2021 Xie Meiyi(xiemeiyi@hust.edu.cn) and OceanBase and/or its affiliates. All rights reserved.
miniob is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:
         http://license.coscl.org.cn/MulanPSL2
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
See the Mulan PSL v2 for more details. */
 
//
//
// Created by Xie Meiyi
// Rewritten by Longda & Wangyunlai
//
//
 
#pragma once
 
#include <string.h>
 
#include "common/lang/comparator.h"
#include "common/lang/memory.h"
#include "common/lang/sstream.h"
#include "common/lang/functional.h"
#include "common/log/log.h"
#include "sql/parser/parse_defs.h"
#include "storage/buffer/disk_buffer_pool.h"
#include "storage/record/record_manager.h"
#include "storage/index/latch_memo.h"
#include "storage/index/bplus_tree_log.h"
 
class BplusTreeHandler;
class BplusTreeMiniTransaction;
 
enum class BplusTreeOperationType
{
  READ,
  INSERT,
  DELETE,
};
 
class AttrComparator
{
public:
  void init(AttrType type, int length)
  {
    attr_type_   = type;
    attr_length_ = length;
  }
 
  int attr_length() const { return attr_length_; }
 
  int operator()(const char *v1, const char *v2) const
  {
    // TODO: optimized the comparison
    Value left;
    left.set_type(attr_type_);
    left.set_data(v1, attr_length_);
    Value right;
    right.set_type(attr_type_);
    right.set_data(v2, attr_length_);
    return DataType::type_instance(attr_type_)->compare(left, right);
  }
 
private:
  AttrType attr_type_;
  int      attr_length_;
};
 
class KeyComparator
{
public:
  void init(AttrType type, int length) { attr_comparator_.init(type, length); }
 
  const AttrComparator &attr_comparator() const { return attr_comparator_; }
 
  int operator()(const char *v1, const char *v2) const
  {
    int result = attr_comparator_(v1, v2);
    if (result != 0) {
      return result;
    }
 
    const RID *rid1 = (const RID *)(v1 + attr_comparator_.attr_length());
    const RID *rid2 = (const RID *)(v2 + attr_comparator_.attr_length());
    return RID::compare(rid1, rid2);
  }
 
private:
  AttrComparator attr_comparator_;
};
 
class AttrPrinter
{
public:
  void init(AttrType type, int length)
  {
    attr_type_   = type;
    attr_length_ = length;
  }
 
  int attr_length() const { return attr_length_; }
 
  string operator()(const char *v) const
  {
    Value value(attr_type_, const_cast<char *>(v), attr_length_);
    return value.to_string();
  }
 
private:
  AttrType attr_type_;
  int      attr_length_;
};
 
class KeyPrinter
{
public:
  void init(AttrType type, int length) { attr_printer_.init(type, length); }
 
  const AttrPrinter &attr_printer() const { return attr_printer_; }
 
  string operator()(const char *v) const
  {
    stringstream ss;
    ss << "{key:" << attr_printer_(v) << ",";
 
    const RID *rid = (const RID *)(v + attr_printer_.attr_length());
    ss << "rid:{" << rid->to_string() << "}}";
    return ss.str();
  }
 
private:
  AttrPrinter attr_printer_;
};
 
struct IndexFileHeader
{
  IndexFileHeader()
  {
    memset(this, 0, sizeof(IndexFileHeader));
    root_page = BP_INVALID_PAGE_NUM;
  }
  PageNum  root_page;          
  int32_t  internal_max_size;  
  int32_t  leaf_max_size;      
  int32_t  attr_length;        
  int32_t  key_length;         
  AttrType attr_type;          
 
  const string to_string() const
  {
    stringstream ss;
 
    ss << "attr_length:" << attr_length << ","
       << "key_length:" << key_length << ","
       << "attr_type:" << attr_type_to_string(attr_type) << ","
       << "root_page:" << root_page << ","
       << "internal_max_size:" << internal_max_size << ","
       << "leaf_max_size:" << leaf_max_size << ";";
 
    return ss.str();
  }
};
 
struct IndexNode
{
  static constexpr int HEADER_SIZE = 12;
 
  bool    is_leaf;  
  int     key_num;  
  PageNum parent;   
};
 
struct LeafIndexNode : public IndexNode
{
  static constexpr int HEADER_SIZE = IndexNode::HEADER_SIZE + 4;
 
  PageNum next_brother;
  char array[0];
};
 
struct InternalIndexNode : public IndexNode
{
  static constexpr int HEADER_SIZE = IndexNode::HEADER_SIZE;
 
  char array[0];
};
 
class IndexNodeHandler
{
public:
  IndexNodeHandler(BplusTreeMiniTransaction &mtr, const IndexFileHeader &header, Frame *frame);
  virtual ~IndexNodeHandler() = default;
 
  void init_empty(bool leaf);
 
  bool is_leaf() const;
 
  virtual int key_size() const;
  virtual int value_size() const;
  virtual int item_size() const;
 
  void    increase_size(int n);
  int     size() const;
  int     max_size() const;
  int     min_size() const;
  RC      set_parent_page_num(PageNum page_num);
  PageNum parent_page_num() const;
  PageNum page_num() const;
 
  bool is_safe(BplusTreeOperationType op, bool is_root_node);
 
  bool validate() const;
 
  Frame *frame() const { return frame_; }
 
  friend string to_string(const IndexNodeHandler &handler);
 
  RC recover_insert_items(int index, const char *items, int num);
  RC recover_remove_items(int index, int num);
 
protected:
  virtual char *__item_at(int index) const { return nullptr; }
  char         *__key_at(int index) const { return __item_at(index); }
  char         *__value_at(int index) const { return __item_at(index) + key_size(); };
 
protected:
  BplusTreeMiniTransaction &mtr_;
  const IndexFileHeader    &header_;
  Frame                    *frame_ = nullptr;
  IndexNode                *node_  = nullptr;
};
 
class LeafIndexNodeHandler final : public IndexNodeHandler
{
public:
  LeafIndexNodeHandler(BplusTreeMiniTransaction &mtr, const IndexFileHeader &header, Frame *frame);
  virtual ~LeafIndexNodeHandler() = default;
 
  RC      init_empty();
  RC      set_next_page(PageNum page_num);
  PageNum next_page() const;
 
  char *key_at(int index);
  char *value_at(int index);
 
  int lookup(const KeyComparator &comparator, const char *key, bool *found = nullptr) const;
 
  RC  insert(int index, const char *key, const char *value);
  RC  remove(int index);
  int remove(const char *key, const KeyComparator &comparator);
  RC  move_half_to(LeafIndexNodeHandler &other);
  RC  move_first_to_end(LeafIndexNodeHandler &other);
  RC  move_last_to_front(LeafIndexNodeHandler &other);
  RC move_to(LeafIndexNodeHandler &other);
 
  bool validate(const KeyComparator &comparator, DiskBufferPool *bp) const;
 
  friend string to_string(const LeafIndexNodeHandler &handler, const KeyPrinter &printer);
 
protected:
  char *__item_at(int index) const override;
 
  RC append(const char *items, int num);
  RC append(const char *item);
  RC preappend(const char *item);
 
private:
  LeafIndexNode *leaf_node_ = nullptr;
};
 
class InternalIndexNodeHandler final : public IndexNodeHandler
{
public:
  InternalIndexNodeHandler(BplusTreeMiniTransaction &mtr, const IndexFileHeader &header, Frame *frame);
  virtual ~InternalIndexNodeHandler() = default;
 
  RC init_empty();
  RC create_new_root(PageNum first_page_num, const char *key, PageNum page_num);
 
  RC      insert(const char *key, PageNum page_num, const KeyComparator &comparator);
  char   *key_at(int index);
  PageNum value_at(int index);
 
  int  value_index(PageNum page_num);
  void set_key_at(int index, const char *key);
  void remove(int index);
 
  int lookup(
      const KeyComparator &comparator, const char *key, bool *found = nullptr, int *insert_position = nullptr) const;
 
  RC move_to(InternalIndexNodeHandler &other);
  RC move_first_to_end(InternalIndexNodeHandler &other);
  RC move_last_to_front(InternalIndexNodeHandler &other);
  RC move_half_to(InternalIndexNodeHandler &other);
 
  bool validate(const KeyComparator &comparator, DiskBufferPool *bp) const;
 
  friend string to_string(const InternalIndexNodeHandler &handler, const KeyPrinter &printer);
 
private:
  RC insert_items(int index, const char *items, int num);
  RC append(const char *items, int num);
  RC append(const char *item);
  RC preappend(const char *item);
 
private:
  char *__item_at(int index) const override;
 
  int value_size() const override;
  int item_size() const override;
 
private:
  InternalIndexNode *internal_node_ = nullptr;
};
 
class BplusTreeHandler
{
public:
  RC create(LogHandler &log_handler, BufferPoolManager &bpm, const char *file_name, AttrType attr_type, int attr_length,
      int internal_max_size = -1, int leaf_max_size = -1);
  RC create(LogHandler &log_handler, DiskBufferPool &buffer_pool, AttrType attr_type, int attr_length,
      int internal_max_size = -1, int leaf_max_size = -1);
 
  RC open(LogHandler &log_handler, BufferPoolManager &bpm, const char *file_name);
  RC open(LogHandler &log_handler, DiskBufferPool &buffer_pool);
 
  RC close();
 
  RC insert_entry(const char *user_key, const RID *rid);
 
  RC delete_entry(const char *user_key, const RID *rid);
 
  bool is_empty() const;
 
  RC get_entry(const char *user_key, int key_len, list<RID> &rids);
 
  RC sync();
 
  bool validate_tree();
 
public:
  const IndexFileHeader &file_header() const { return file_header_; }
  DiskBufferPool        &buffer_pool() const { return *disk_buffer_pool_; }
  LogHandler            &log_handler() const { return *log_handler_; }
 
public:
  RC recover_update_root_page(BplusTreeMiniTransaction &mtr, PageNum root_page_num);
  RC recover_init_header_page(BplusTreeMiniTransaction &mtr, Frame *frame, const IndexFileHeader &header);
 
public:
  RC print_tree();
  RC print_leafs();
 
private:
  RC print_leaf(Frame *frame);
  RC print_internal_node_recursive(Frame *frame);
 
  bool validate_leaf_link(BplusTreeMiniTransaction &mtr);
  bool validate_node_recursive(BplusTreeMiniTransaction &mtr, Frame *frame);
 
protected:
  RC find_leaf(BplusTreeMiniTransaction &mtr, BplusTreeOperationType op, const char *key, Frame *&frame);
 
  RC left_most_page(BplusTreeMiniTransaction &mtr, Frame *&frame);
 
  RC find_leaf_internal(BplusTreeMiniTransaction &mtr, BplusTreeOperationType op,
      const function<PageNum(InternalIndexNodeHandler &)> &child_page_getter, Frame *&frame);
 
  RC crabing_protocal_fetch_page(
      BplusTreeMiniTransaction &mtr, BplusTreeOperationType op, PageNum page_num, bool is_root_page, Frame *&frame);
 
  RC delete_entry_internal(BplusTreeMiniTransaction &mtr, Frame *leaf_frame, const char *key);
 
  template <typename IndexNodeHandlerType>
  RC split(BplusTreeMiniTransaction &mtr, Frame *frame, Frame *&new_frame);
 
  template <typename IndexNodeHandlerType>
  RC coalesce_or_redistribute(BplusTreeMiniTransaction &mtr, Frame *frame);
 
  template <typename IndexNodeHandlerType>
  RC coalesce(BplusTreeMiniTransaction &mtr, Frame *neighbor_frame, Frame *frame, Frame *parent_frame, int index);
 
  template <typename IndexNodeHandlerType>
  RC redistribute(BplusTreeMiniTransaction &mtr, Frame *neighbor_frame, Frame *frame, Frame *parent_frame, int index);
 
  RC insert_entry_into_parent(BplusTreeMiniTransaction &mtr, Frame *frame, Frame *new_frame, const char *key);
 
  RC insert_entry_into_leaf_node(BplusTreeMiniTransaction &mtr, Frame *frame, const char *pkey, const RID *rid);
 
  RC create_new_tree(BplusTreeMiniTransaction &mtr, const char *key, const RID *rid);
 
  void update_root_page_num_locked(BplusTreeMiniTransaction &mtr, PageNum root_page_num);
 
  RC adjust_root(BplusTreeMiniTransaction &mtr, Frame *root_frame);
 
private:
  common::MemPoolItem::item_unique_ptr make_key(const char *user_key, const RID &rid);
 
protected:
  LogHandler     *log_handler_      = nullptr;  
  DiskBufferPool *disk_buffer_pool_ = nullptr;  
  bool            header_dirty_     = false;    
  IndexFileHeader file_header_;
 
  // 在调整根节点时，需要加上这个锁。
  // 这个锁可以使用递归读写锁，但是这里偷懒先不改
  common::SharedMutex root_lock_;
 
  KeyComparator key_comparator_;
  KeyPrinter    key_printer_;
 
  unique_ptr<common::MemPoolItem> mem_pool_item_;
 
private:
  friend class BplusTreeScanner;
  friend class BplusTreeTester;
};
 
class BplusTreeScanner
{
public:
  BplusTreeScanner(BplusTreeHandler &tree_handler);
  ~BplusTreeScanner();
 
  RC open(const char *left_user_key, int left_len, bool left_inclusive, const char *right_user_key, int right_len,
      bool right_inclusive);
 
  RC next_entry(RID &rid);
 
  RC close();
 
private:
  RC fix_user_key(const char *user_key, int key_len, bool want_greater, char **fixed_key, bool *should_inclusive);
 
  void fetch_item(RID &rid);
 
  bool touch_end();
 
private:
  bool                     inited_ = false;
  BplusTreeHandler        &tree_handler_;
  BplusTreeMiniTransaction mtr_;
 
  Frame *current_frame_ = nullptr;
 
  common::MemPoolItem::item_unique_ptr right_key_;
  int                                  iter_index_    = -1;
  bool                                 first_emitted_ = false;
};