Project: finder, src: BPTreeCore.cc


B+ Tree 基本コード
#define _LARGEFILE64_SOURCE
#include "BPTreeCore.h"
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <stack>
#include <assert.h>

using namespace std;

BPTreeCore::BPTreeCore(const std::string &filename, bool bTruncate) : fd(0) {
  int flags = O_RDWR | O_CREAT;
  if (bTruncate) {
    flags |= O_TRUNC;
  }
  fd = open(filename.c_str(), flags, 0644);
  bCreated = true;
  initialFill();
}

BPTreeCore::BPTreeCore(int fd) : fd(fd) {
  bCreated = false;
  initialFill();
}

BPTreeCore::~BPTreeCore() {
  if (bCreated && fd >= 0) {
    close(fd);
    fd = -1;
  }
}

// -------------------------- インタフェース ---------------------

void BPTreeCore::append(const NGram &key, off_t value) {
  Node node;
  std::stack<Node> nodeStack;
  int index;

  // root node を最初に読み込む
  lseek(fd, 0, SEEK_SET);
  node.readObject(fd);

  while (node.level > 0) {
    // 中間ノードの場合はたどる
    off_t offset;
    index = bsearch(node.elements, key);
    nodeStack.push(node);
    if (index >= 0) {
      // キーは範囲内
      offset = node.elements[index].offset;
    } else if (key < node.elements[0].key) {
      // キーは B++ Tree に登録されているキーよりも小さい
      offset = node.elements[0].offset;
    } else {
      // キーは B++ Tree に登録されているキーよりも大きい
      offset = node.elements[node.last - 1].offset;
    }
    lseek(fd, offset, SEEK_SET);
    node.readObject(fd);
    node.offset = offset;
  }

  // 葉ノード
  if (node.last == BPTREE_ELEMENT_SIZE) {
    // 空きがない場合は split 操作を行う
    splitNode(nodeStack, node, key);
  }

  Element e(key, value);
  if (node.last == 0) {
    // 新規 element
    node.elements.push_back(e);
  } else {
    index = bsearch(node.elements, key);
    if (index >= 0) {
      if (node.elements[index].key == key) {
        // 完全一致なので上書きする
        node.elements[index].offset = value;
      } else if (index == node.last) {
        node.elements.push_back(e);
      } else {
        // 新規 element
        node.elements.insert(node.elements.begin() + index + 1, e);
      }
    } else if (key < node.elements[0].key) {
      node.elements.insert(node.elements.begin(), e);
    } else {
      node.elements.push_back(e);
    }
  }
  node.last++;

  // file descriptor に書き込む
  lseek(fd, node.offset, SEEK_SET);
  node.writeObject(fd);

  // stack 上の最小要素よりも小さい要素を追加した場合、ノードを修正する
  while (!nodeStack.empty()) {
    node = nodeStack.top();
    if (node.last > 0 && key < node.elements[0].key) {
      node.elements[0].key = key;
      lseek(fd, node.offset, SEEK_SET);
      node.writeObject(fd);
    }
    nodeStack.pop();
  }
}

bool BPTreeCore::find(const NGram &key, off_t *pval) const {
  Node node;
  lseek(fd, 0, SEEK_SET);
  node.readObject(fd);

  // last == 0 の場合、B++ Tree は空
  if (node.last == 0) return false;

  // root node を検索する
  int index;
  if (node.elements[node.last - 1].key < key) {
    index = node.last - 1;
  } else {
    index = bsearch(node.elements, key);
  }

  // 範囲内に key はなかった
  if (index < 0) return false;

  while(node.level > 0) {
    // 中間ノードの場合はツリーをたどっていく
    lseek(fd, node.elements[index].offset, SEEK_SET);
    node.readObject(fd);
    if (node.elements[node.last - 1].key < key) {
      index = node.last - 1;
    } else {
      index = bsearch(node.elements, key);
    }
    if (index < 0) return false;
  }

  // node が葉ノードの場合は厳密に一致することを要求する
  if (!(node.elements[index].key == key)) return false;

  // 一致
  if (pval) *pval = node.elements[index].offset;
  return true;
}

// -------------------------- 実装 -------------------------------
/**
 * Element オブジェクトの書き込み
 */
void BPTreeCore::Element::writeObject(int fd) const {
  key.writeObject(fd);
  write(fd, &offset, sizeof(offset));
}

/**
 * Element オブジェクトの読み込み
 */
void BPTreeCore::Element::readObject(int fd) {
  key.readObject(fd);
  read(fd, &offset, sizeof(offset));
}

/**
 * Node オブジェクトの書き込み
 */
void BPTreeCore::Node::writeObject(int fd) const {
  NGram key0;
  off_t oft0 = 0;
  Element dummy(key0, oft0);
  write(fd, &level, sizeof(level));
  write(fd, &last, sizeof(last));
  for (int i = 0; i < BPTREE_ELEMENT_SIZE; ++i) {
    if (i < last) {
      elements[i].writeObject(fd);
    } else {
      // 保存領域だけ確保しておく
      dummy.writeObject(fd);
    }
  }
}

/**
 * Node オブジェクトの読み込み
 */
void BPTreeCore::Node::readObject(int fd) {
  read(fd, &level, sizeof(level));
  read(fd, &last, sizeof(last));
  elements.clear();
  for (int i = 0; i < last; ++i) {
    NGram key0;
    off_t oft0 = 0;
    Element e(key0, oft0);
    e.readObject(fd);
    elements.push_back(e);
  }
}

/**
 * Element ベクターの二分検索
 */
int BPTreeCore::bsearch(const std::vector<BPTreeCore::Element> &elements, const NGram &key) {
  int minIndex = 0, maxIndex = elements.size();
  int cur = (minIndex + maxIndex) / 2;

  // 範囲内にあることを調べる
  if (minIndex > maxIndex) return -1;
  if (key < elements[minIndex].key) return -1;
  if (elements[maxIndex - 1].key < key) return -1;

  while (minIndex + 1 < maxIndex) {
    if (elements[cur].key == key) {
      break;
    } else if (elements[cur].key < key) {
      minIndex = cur;
    } else {
      maxIndex = cur;
    }
    cur = (minIndex + maxIndex) / 2;
  }
  return cur;
}

/**
 * Node を分割する
 */
void BPTreeCore::splitNode(std::stack<BPTreeCore::Node> &stack, BPTreeCore::Node &node, const NGram &key) {
  // 分割する必要がない場合は分割しない
  if (node.last < BPTREE_ELEMENT_SIZE) return;

  int k, sp = node.last / 2;
  Node node1, node2;
  std::vector<Element>::iterator i = node.elements.begin();

  // 前半
  node1.elements.insert(node1.elements.begin(), i, i + sp);
  node1.last = sp;
  node1.level = node.level;
  if (node.offset == 0) {
    node1.offset = lseek(fd, 0, SEEK_END);
  } else {
    node1.offset = node.offset;
    lseek(fd, node1.offset, SEEK_SET);
  }
  node1.writeObject(fd);

  // 後半
  node2.elements.insert(node2.elements.begin(), i + sp, node.elements.end());
  node2.last = sp;
  node2.level = node.level;
  node2.offset = lseek(fd, 0, SEEK_END);
  node2.writeObject(fd);

  Node parent;
  Element child1(node1.elements[0].key, node1.offset);
  Element child2(node2.elements[0].key, node2.offset);
  if (stack.empty()) {
    // node は root node
    parent.last = 2;
    parent.level = 1;
    parent.offset = 0;
    parent.elements.push_back(child1);
    parent.elements.push_back(child2);
    lseek(fd, 0, SEEK_SET);
  } else {
    // 親ノードに要素を追加する
    parent = stack.top();
    stack.pop();
    splitNode(stack, parent, key);
    int index = bsearch(parent.elements, child1.key);
    assert (index != -1);
    if (index == BPTREE_ELEMENT_SIZE - 1) {
      parent.elements.push_back(child2);
    } else {
      parent.elements.insert(parent.elements.begin() + (index + 1), child2);
    }
    ++parent.last;
    lseek(fd, parent.offset, SEEK_SET);
  }

  parent.writeObject(fd);
  stack.push(parent);
  if (node2.elements[0].key < key) {
    node = node2;
  } else {
    node = node1;
  }
}

/**
 * root node の初期化
 */
void BPTreeCore::initialFill() {
  off_t last = lseek(fd, 0, SEEK_END);
  if (last == (off_t)0) {
    // 初期ノードを挿入
    Node node;
    node.writeObject(fd);
    lseek(fd, 0, SEEK_SET);
  }
}