Project: finder, src: Finder.cc

#include "Finder.h"
#include "docid.h"
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <cstddef>
#include <cstdlib>
#include <vector>
#include <algorithm>

/**
 * コンストラクタ
 */
Finder::Finder(const std::string &docfile, const std::string &filename, bool bTrunc) : bptree(filename, bTrunc), analyzeCount(0), bCreate(true) {
  int flags = O_RDWR | O_CREAT;
  if (bTrunc) {
    flags |= O_TRUNC;
  }
  docfd = open(docfile.c_str(), flags, 0644);
  fd = bptree.getFileDescriptor();
}

/**
 * ドキュメントを解析して保存する
 */
void Finder::add(unsigned long long id) {
  // DOC アーカイブから文書を取得する
  int len8;
  char *utf8 = getDocFromId(docfd, id, &len8);
  if (utf8 == 0) return; // 文書がなければ終了

  // 解析結果をローカルに保存
  work.analyze(id, utf8, len8);
  if (++analyzeCount > MAX_DOCUMENT_LIST) {
    // 一定数たまったので二次媒体に移す
    flush();
  }
  free(utf8);
}

/**
 * 解析結果の hash table を B++ Tree に保存する
 */
void Finder::flush() {
  work.save2BPTree(bptree);
  work.clear();
  analyzeCount = 0;
}

/**
 * 入力文字列が文書集合に含まれているかどうか調査し、その結果を返す
 */
std::list<Document> Finder::find(const char *utf8, int len8) {
  std::list<Document> result;

  // UCS4 文字列に変換
  int len4;
  UCS4 *ucs4 = ucs4FromUtf8(utf8, len8, &len4);
  if (ucs4 == 0) return result;

  // 長さが足りない場合は空を返す
  if (len4 < NGRAM_SIZE) {
    return result;
  }

  // 検索文字列を含む「列」を作成する
  std::vector<BPTree::ObjectStream> objStreamArray;
  for (int i = 0; i <= len4 - NGRAM_SIZE; ++i) {
    NGram key(ucs4 + i);
    objStreamArray.push_back(bptree.find(key));
  }

  std::vector<DocumentList> doclistArray;
  std::vector<BPTree::ObjectStream>::iterator itr = objStreamArray.begin();
  for ( ; itr != objStreamArray.end(); ++itr) {
    // この段階で Document がない場合は空を返す
    if (itr->isEOF()) return result;
    DocumentList doclist;
    itr->load(doclist);
    doclistArray.push_back(doclist);
  }

  std::vector<DocumentList::iterator> docItrArray;
  std::vector<DocumentList>::iterator j = doclistArray.begin();
  for ( ; j != doclistArray.end(); ++j) {
    docItrArray.push_back((*j).begin());
  }

  Document doc;
  unsigned long long nextId;
  if (docItrArray[0] == doclistArray[0].end()) return result;
  doc = *(docItrArray[0]);

  while (find(&nextId, result, doc, objStreamArray, doclistArray, docItrArray)) {
    // nextId まで読み飛ばす
    while (doc.getDocid() < nextId) {
      if (!incrementDocItr(objStreamArray[0], doclistArray[0], docItrArray[0])) return result;
      doc = *(docItrArray[0]);
    }
  }
  return result;
}

/**
 * ObjectStream 上、「次」の Document を探索する
 */
bool Finder::incrementDocItr(BPTree::ObjectStream &os, DocumentList &doclist, DocumentList::iterator &itr) {
  if (itr != doclist.end()) {
    // まだ iterator が残っている
    ++itr;
    if (itr != doclist.end()) return true;
  }

  // ObjectStream から取得する
  if (os.isEOF()) return false; // もうなかった
  if (!os.load(doclist)) return false; // load に失敗した場合
  itr = doclist.begin();

  return (itr != doclist.end());
}

/**
 * 各「列」を比較して並行成分を除いて同じ index があったら result に追記する
 */
bool Finder::find(unsigned long long *nextId, std::list<Document> &result, Document &doc, std::vector<BPTree::ObjectStream> &objStreamArray, std::vector<DocumentList> &doclistArray, std::vector<DocumentList::iterator> &docItrArray) {
  std::size_t sz = objStreamArray.size();
  std::set<int> bookmarks(doc.getBookmarks());

  for (int m = 1; m < sz; ++m) {
    Document mdoc = *(docItrArray[m]);
    while(true) {
      if (mdoc.getDocid() >= doc.getDocid()) break;
      if (!incrementDocItr(objStreamArray[m], doclistArray[m], docItrArray[m])) {
        // 各「列」にこれ以上文書がない場合は中断する
        return false;
      }
      mdoc = *(docItrArray[m]);
    }
    if (mdoc.getDocid() > doc.getDocid()) {
      // 文書 ID を追い越した場合は探索元の ID をそこまで上げる
      // 探索自体は継続する
      *nextId = mdoc.getDocid();
      return true;
    }
    const std::set<int> &mdocbm = mdoc.getBookmarks();
    std::set<int> tmpset;
    for (std::set<int>::const_iterator j = mdocbm.begin(); j != mdocbm.end(); ++j) {
      int index = (*j) - m;
      if (bookmarks.find(index) != bookmarks.end()) {
        // 両方の集合に index が含まれていれば検索先候補とする
        tmpset.insert(index);
      }
    }
    bookmarks = tmpset;
  }

  // この時点で bookmarks に含まれている index は
  // すべての「列」に含まれている成分
  // 文書 ID と bookmarks から結果を構成する
  if (!bookmarks.empty()) {
    result.push_back(Document(doc.getDocid(), bookmarks.begin(), bookmarks.end()));
  }

  // 次の ID を探索する準備をする
  if (!incrementDocItr(objStreamArray[0], doclistArray[0], docItrArray[0])) {
    // これ以上はない
    return false;
  } else {
    // nextId に次の ID を設定する
    doc = *(docItrArray[0]);
    *nextId = doc.getDocid();
    return true;
  }
}