Refactor SystemDictionary::LookupPredictive using new LoudsTrie::Node APIs

In the previous implementation, LookupPredictive was implemented by DFS and some short key list management to emulate BFS because it was difficult to implement BFS due to LOUDS trie implementation.  This CL implements BFS using a new LOUDS trie APIs.  The resulting code is simpler and more efficient.

Benchmark shows significant performance improvement for suggestion and prediction.  For example, observed 1.2x and 1.4x speed up for mobile one character prediction and suggestion, respectively.  Note also that this is just a code refactoring, no user-visible behavior change is intended.

BUG=none
TEST=unittest,benchmark

git-svn-id: https://mozc.googlecode.com/svn/trunk@526 a6090854-d499-a067-5803-1114d4e51264
diff --git a/src/dictionary/system/system_dictionary.cc b/src/dictionary/system/system_dictionary.cc
index 311bfd1..54ae426 100644
--- a/src/dictionary/system/system_dictionary.cc
+++ b/src/dictionary/system/system_dictionary.cc
@@ -50,6 +50,7 @@
 #include <cstring>
 #include <limits>
 #include <map>
+#include <queue>
 #include <string>
 #include <utility>
 #include <vector>
@@ -716,141 +717,73 @@
   return false;
 }
 
-namespace {
+void SystemDictionary::CollectPredictiveNodesInBfsOrder(
+    StringPiece encoded_key,
+    const KeyExpansionTable &table,
+    size_t limit,
+    vector<PredictiveLookupSearchState> *result) const {
+  queue<PredictiveLookupSearchState> queue;
+  queue.push(PredictiveLookupSearchState(LoudsTrie::Node(), 0, false));
+  do {
+    PredictiveLookupSearchState state = queue.front();
+    queue.pop();
 
-// Collects short keys preferentially.
-class ShortKeyCollector : public LoudsTrie::Callback {
- public:
-  // Holds a lookup result from trie.
-  struct Entry {
-    string encoded_key;  // Encoded lookup key
-    string encoded_actual_key;  // Encoded actual key in trie (expanded key)
-    size_t actual_key_len;  // Decoded actual key length
-    int key_id;  // Key ID in trie
-  };
+    // Update traversal state for |encoded_key| and its expanded keys.
+    if (state.key_pos < encoded_key.size()) {
+      const char target_char = encoded_key[state.key_pos];
+      const ExpandedKey &chars = table.ExpandKey(target_char);
 
-  ShortKeyCollector(const SystemDictionaryCodecInterface *codec,
-                    StringPiece original_encoded_key,
-                    size_t min_key_len,
-                    size_t limit)
-      : codec_(codec),
-        original_encoded_key_(original_encoded_key),
-        min_key_len_(min_key_len),
-        limit_(limit),
-        current_max_key_len_(0),
-        num_max_key_length_entries_(0) {
-    entry_list_.reserve(limit);
-  }
-
-  virtual ResultType Run(const char *trie_key,
-                         size_t trie_key_len, int key_id) {
-    const StringPiece encoded_actual_key(trie_key, trie_key_len);
-
-    // First calculate the length of decoded key.
-    // Note: In the current kana modifier insensitive lookup mechanism, the
-    // lengths of the original lookup key and its expanded key are equal, so we
-    // can omit the construction of lookup key by calculating the length of
-    // decoded actual key. Just DCHECK here.
-    const size_t key_len = codec_->GetDecodedKeyLength(encoded_actual_key);
-    DCHECK_EQ(key_len, codec_->GetDecodedKeyLength(
-        original_encoded_key_.as_string() +
-        string(trie_key + original_encoded_key_.size(),
-               trie_key_len - original_encoded_key_.size())));
-    // Uninterested in too short key.
-    if (key_len < min_key_len_) {
-      return SEARCH_CONTINUE;
+      for (key_trie_->MoveToFirstChild(&state.node);
+           key_trie_->IsValidNode(state.node);
+           key_trie_->MoveToNextSibling(&state.node)) {
+        const char c = key_trie_->GetEdgeLabelToParentNode(state.node);
+        if (!chars.IsHit(c)) {
+          continue;
+        }
+        const bool is_expanded = state.is_expanded || c != target_char;
+        queue.push(PredictiveLookupSearchState(state.node,
+                                               state.key_pos + 1,
+                                               is_expanded));
+      }
+      continue;
     }
 
-    // Check the key length after decoding and update the internal state. As
-    // explained above, the length of actual key (expanded key) is equal to that
-    // of key.
-    const size_t actual_key_len = key_len;
-    if (actual_key_len > current_max_key_len_) {
-      if (entry_list_.size() > limit_) {
-        return SEARCH_CULL;
-      }
-      current_max_key_len_ = actual_key_len;
-      num_max_key_length_entries_ = 1;
-    } else if (actual_key_len == current_max_key_len_) {
-      ++num_max_key_length_entries_;
-    } else {
-      if (entry_list_.size() - num_max_key_length_entries_ + 1 >= limit_) {
-        RemoveAllMaxKeyLengthEntries();
-        UpdateMaxKeyLengthInternal();
-      }
+    // Collect prediction keys (state.key_pos >= encoded_key.size()).
+    if (key_trie_->IsTerminalNode(state.node)) {
+      result->push_back(state);
     }
 
-    // Keep this entry at the back.
-    entry_list_.push_back(Entry());
-    entry_list_.back().encoded_key.reserve(trie_key_len);
-    original_encoded_key_.CopyToString(&entry_list_.back().encoded_key);
-    entry_list_.back().encoded_key.append(
-        trie_key + original_encoded_key_.size(),
-        trie_key_len - original_encoded_key_.size());
-    entry_list_.back().encoded_actual_key.assign(trie_key, trie_key_len);
-    entry_list_.back().actual_key_len = actual_key_len;
-    entry_list_.back().key_id = key_id;
-
-    return SEARCH_CONTINUE;
-  }
-
-  const vector<Entry> &entry_list() const {
-    return entry_list_;
-  }
-
- private:
-  void RemoveAllMaxKeyLengthEntries() {
-    for (size_t i = 0; i < entry_list_.size(); ) {
-      Entry *result = &entry_list_[i];
-      if (result->actual_key_len < current_max_key_len_) {
-        // This result should still be kept.
-        ++i;
-        continue;
+    // Collected enough entries.  Collect all the remaining keys that have the
+    // same length as the longest key.
+    if (result->size() > limit) {
+      // The current key is the longest because of BFS property.
+      const size_t max_key_len = state.key_pos;
+      while (!queue.empty()) {
+        state = queue.front();
+        queue.pop();
+        if (state.key_pos > max_key_len) {
+          // Key length in the queue is monotonically increasing because of BFS
+          // property.  So we don't need to check all the elements in the queue.
+          break;
+        }
+        DCHECK_EQ(state.key_pos, max_key_len);
+        if (key_trie_->IsTerminalNode(state.node)) {
+          result->push_back(state);
+        }
       }
-
-      // Remove the i-th result by swapping it for the last one.
-      Entry *last_result = &entry_list_.back();
-      swap(result->encoded_key, last_result->encoded_key);
-      swap(result->encoded_actual_key, last_result->encoded_actual_key);
-      result->actual_key_len = last_result->actual_key_len;
-      result->key_id = last_result->key_id;
-      entry_list_.pop_back();
-
-      // Do not update the |i| here.
+      break;
     }
-  }
 
-  void UpdateMaxKeyLengthInternal() {
-    current_max_key_len_ = 0;
-    num_max_key_length_entries_ = 0;
-    for (size_t i = 0; i < entry_list_.size(); ++i) {
-      if (entry_list_[i].actual_key_len > current_max_key_len_) {
-        current_max_key_len_ = entry_list_[i].actual_key_len;
-        num_max_key_length_entries_ = 1;
-      } else if (entry_list_[i].actual_key_len == current_max_key_len_) {
-        ++num_max_key_length_entries_;
-      }
+    // Update traversal state for children.
+    for (key_trie_->MoveToFirstChild(&state.node);
+         key_trie_->IsValidNode(state.node);
+         key_trie_->MoveToNextSibling(&state.node)) {
+      queue.push(PredictiveLookupSearchState(state.node,
+                                             state.key_pos + 1,
+                                             state.is_expanded));
     }
-  }
-
-  const SystemDictionaryCodecInterface *codec_;
-  const StringPiece original_encoded_key_;
-
-  // Filter conditions.
-  const size_t min_key_len_;
-  const size_t limit_;
-
-  // Internal state for tracking current maximum key length.
-  size_t current_max_key_len_;
-  size_t num_max_key_length_entries_;
-
-  // Contains lookup results.
-  vector<Entry> entry_list_;
-
-  DISALLOW_COPY_AND_ASSIGN(ShortKeyCollector);
-};
-
-}  // namespace
+  } while (!queue.empty());
+}
 
 void SystemDictionary::LookupPredictive(
     StringPiece key, bool use_kana_modifier_insensitive_lookup,
@@ -861,31 +794,52 @@
     return;
   }
 
-  string lookup_key_str;
-  codec_->EncodeKey(key, &lookup_key_str);
-  if (lookup_key_str.size() > LoudsTrie::kMaxDepth) {
+  string encoded_key;
+  codec_->EncodeKey(key, &encoded_key);
+  if (encoded_key.size() > LoudsTrie::kMaxDepth) {
     return;
   }
 
-  // First, collect up to 64 keys so that results are as short as possible,
-  // which emulates BFS over trie.
-  ShortKeyCollector collector(codec_, lookup_key_str, 0, 64);
-  const KeyExpansionTable &table =
-      use_kana_modifier_insensitive_lookup ?
-      hiragana_expansion_table_ : KeyExpansionTable::GetDefaultInstance();
-  key_trie_->PredictiveSearchWithKeyExpansion(lookup_key_str, table,
-                                              &collector);
+  const KeyExpansionTable &table = use_kana_modifier_insensitive_lookup
+      ? hiragana_expansion_table_
+      : KeyExpansionTable::GetDefaultInstance();
 
-  string decoded_key, actual_key;
-  for (size_t i = 0; i < collector.entry_list().size(); ++i) {
-    const ShortKeyCollector::Entry &entry = collector.entry_list()[i];
+  // TODO(noriyukit): Lookup limit should be implemented at caller side by using
+  // callback mechanism.  This hard-coding limits the capability and generality
+  // of dictionary module.  CollectPredictiveNodesInBfsOrder() and the following
+  // loop for callback should be integrated for this purpose.
+  const size_t kLookupLimit = 64;
+  vector<PredictiveLookupSearchState> result;
+  result.reserve(kLookupLimit);
+  CollectPredictiveNodesInBfsOrder(encoded_key, table, kLookupLimit, &result);
 
+  // Reused buffer and instances inside the following loop.
+  char encoded_actual_key_buffer[LoudsTrie::kMaxDepth + 1];
+  string decoded_key, actual_key_str;
+  decoded_key.reserve(key.size() * 2);
+  actual_key_str.reserve(key.size() * 2);
+  for (size_t i = 0; i < result.size(); ++i) {
+    const PredictiveLookupSearchState &state = result[i];
+
+    // Computes the actual key.  For example:
+    // key = "くー"
+    // encoded_actual_key = encode("ぐーぐる")  [expanded]
+    // encoded_actual_key_prediction_suffix = encode("ぐる")
+    const StringPiece encoded_actual_key =
+        key_trie_->RestoreKeyString(state.node, encoded_actual_key_buffer);
+    const StringPiece encoded_actual_key_prediction_suffix(
+        encoded_actual_key,
+        encoded_key.size(),
+        encoded_actual_key.size() - encoded_key.size());
+
+    // decoded_key = "くーぐる" (= key + prediction suffix)
     decoded_key.clear();
-    codec_->DecodeKey(entry.encoded_key, &decoded_key);
+    key.CopyToString(&decoded_key);
+    codec_->DecodeKey(encoded_actual_key_prediction_suffix, &decoded_key);
     switch (callback->OnKey(decoded_key)) {
-      case DictionaryInterface::Callback::TRAVERSE_DONE:
+      case Callback::TRAVERSE_DONE:
         return;
-      case DictionaryInterface::Callback::TRAVERSE_NEXT_KEY:
+      case Callback::TRAVERSE_NEXT_KEY:
         continue;
       case DictionaryInterface::Callback::TRAVERSE_CULL:
         LOG(FATAL) << "Culling is not implemented.";
@@ -894,38 +848,41 @@
         break;
     }
 
-    actual_key.clear();
-    codec_->DecodeKey(entry.encoded_actual_key, &actual_key);
-    const bool is_expanded = entry.encoded_key != entry.encoded_actual_key;
-    switch (callback->OnActualKey(decoded_key, actual_key, is_expanded)) {
-      case DictionaryInterface::Callback::TRAVERSE_DONE:
+    StringPiece actual_key;
+    if (state.is_expanded) {
+      actual_key_str.clear();
+      codec_->DecodeKey(encoded_actual_key, &actual_key_str);
+      actual_key = actual_key_str;
+    } else {
+      actual_key = decoded_key;
+    }
+    switch (callback->OnActualKey(decoded_key, actual_key, state.is_expanded)) {
+      case Callback::TRAVERSE_DONE:
         return;
-      case DictionaryInterface::Callback::TRAVERSE_NEXT_KEY:
+      case Callback::TRAVERSE_NEXT_KEY:
         continue;
-      case DictionaryInterface::Callback::TRAVERSE_CULL:
+      case Callback::TRAVERSE_CULL:
         LOG(FATAL) << "Culling is not implemented.";
         continue;
       default:
         break;
     }
 
-    const uint8 *encoded_tokens_ptr =
-        GetTokenArrayPtr(*token_array_, entry.key_id);
+    const int key_id = key_trie_->GetKeyIdOfTerminalNode(state.node);
     for (TokenDecodeIterator iter(codec_, value_trie_.get(),
                                   frequent_pos_, actual_key,
-                                  encoded_tokens_ptr);
+                                  GetTokenArrayPtr(*token_array_, key_id));
          !iter.Done(); iter.Next()) {
       const TokenInfo &token_info = iter.Get();
-      const DictionaryInterface::Callback::ResultType result =
+      const Callback::ResultType result =
           callback->OnToken(decoded_key, actual_key, *token_info.token);
-      if (result == DictionaryInterface::Callback::TRAVERSE_DONE) {
+      if (result == Callback::TRAVERSE_DONE) {
         return;
       }
-      if (result == DictionaryInterface::Callback::TRAVERSE_NEXT_KEY) {
+      if (result == Callback::TRAVERSE_NEXT_KEY) {
         break;
       }
-      DCHECK_NE(DictionaryInterface::Callback::TRAVERSE_CULL, result)
-          << "Culling is not implemented.";
+      DCHECK_NE(Callback::TRAVERSE_CULL, result) << "Not implemented";
     }
   }
 }
diff --git a/src/dictionary/system/system_dictionary.h b/src/dictionary/system/system_dictionary.h
index 98dea2c..54de3ee 100644
--- a/src/dictionary/system/system_dictionary.h
+++ b/src/dictionary/system/system_dictionary.h
@@ -44,6 +44,7 @@
 #include "dictionary/system/codec_interface.h"
 #include "dictionary/system/words_info.h"
 #include "storage/louds/bit_vector_based_array.h"
+#include "storage/louds/key_expansion_table.h"
 #include "storage/louds/louds_trie.h"
 // for FRIEND_TEST
 #include "testing/base/public/gunit_prod.h"
@@ -227,6 +228,23 @@
       char *actual_key_buffer,
       string *actual_prefix) const;
 
+  struct PredictiveLookupSearchState {
+    PredictiveLookupSearchState() : key_pos(0), is_expanded(false) {}
+    PredictiveLookupSearchState(const storage::louds::LoudsTrie::Node &n,
+                                size_t pos, bool expanded)
+        : node(n), key_pos(pos), is_expanded(expanded) {}
+
+    storage::louds::LoudsTrie::Node node;
+    size_t key_pos;
+    bool is_expanded;
+  };
+
+  void CollectPredictiveNodesInBfsOrder(
+      StringPiece encoded_key,
+      const storage::louds::KeyExpansionTable &table,
+      size_t limit,
+      vector<PredictiveLookupSearchState> *result) const;
+
   scoped_ptr<storage::louds::LoudsTrie> key_trie_;
   scoped_ptr<storage::louds::LoudsTrie> value_trie_;
   scoped_ptr<storage::louds::BitVectorBasedArray> token_array_;
diff --git a/src/dictionary/system/value_dictionary.cc b/src/dictionary/system/value_dictionary.cc
index 03b59ce..b9ea3ab 100644
--- a/src/dictionary/system/value_dictionary.cc
+++ b/src/dictionary/system/value_dictionary.cc
@@ -30,6 +30,7 @@
 #include "dictionary/system/value_dictionary.h"
 
 #include <limits>
+#include <queue>
 #include <string>
 
 #include "base/logging.h"
@@ -135,56 +136,30 @@
   token->attributes = Token::NONE;
 }
 
-// Converts a value of SystemDictionary::Callback::ResultType to the
-// corresponding value of LoudsTrie::Callback::ResultType.
-inline LoudsTrie::Callback::ResultType ConvertResultType(
-    const DictionaryInterface::Callback::ResultType result) {
-  switch (result) {
-    case DictionaryInterface::Callback::TRAVERSE_DONE:
-      return LoudsTrie::Callback::SEARCH_DONE;
-    case DictionaryInterface::Callback::TRAVERSE_NEXT_KEY:
-    case DictionaryInterface::Callback::TRAVERSE_CONTINUE:
-      return LoudsTrie::Callback::SEARCH_CONTINUE;
-    case DictionaryInterface::Callback::TRAVERSE_CULL:
-      return LoudsTrie::Callback::SEARCH_CULL;
-    default:
-      LOG(DFATAL) << "Enum value " << result << " cannot be converted";
-      return LoudsTrie::Callback::SEARCH_DONE;  // dummy
+inline DictionaryInterface::Callback::ResultType HandleTerminalNode(
+    const LoudsTrie &value_trie,
+    const SystemDictionaryCodecInterface &codec,
+    const uint16 suggestion_only_word_id,
+    const LoudsTrie::Node &node,
+    DictionaryInterface::Callback *callback,
+    char *encoded_value_buffer,
+    string *value,
+    Token *token) {
+  const StringPiece encoded_value =
+      value_trie.RestoreKeyString(node, encoded_value_buffer);
+
+  value->clear();
+  codec.DecodeValue(encoded_value, value);
+  const DictionaryInterface::Callback::ResultType result =
+      callback->OnKey(*value);
+  if (result != DictionaryInterface::Callback::TRAVERSE_CONTINUE) {
+    return result;
   }
+
+  FillToken(suggestion_only_word_id, *value, token);
+  return callback->OnToken(*value, *value, *token);
 }
 
-class PredictiveTraverser : public LoudsTrie::Callback {
- public:
-  PredictiveTraverser(const SystemDictionaryCodecInterface *codec,
-                      const uint16 suggestion_only_word_id,
-                      DictionaryInterface::Callback *callback)
-      : codec_(codec),
-        suggestion_only_word_id_(suggestion_only_word_id),
-        callback_(callback) {}
-  virtual ~PredictiveTraverser() {}
-
-  virtual ResultType Run(const char *key_begin, size_t len, int key_id) {
-    // The decoded key of value trie corresponds to value (surface form).
-    string value;
-    codec_->DecodeValue(StringPiece(key_begin, len), &value);
-    DictionaryInterface::Callback::ResultType result = callback_->OnKey(value);
-    if (result != DictionaryInterface::Callback::TRAVERSE_CONTINUE) {
-      return ConvertResultType(result);
-    }
-    FillToken(suggestion_only_word_id_, value, &token_);
-    result = callback_->OnToken(value, value, token_);
-    return ConvertResultType(result);
-  }
-
- private:
-  const SystemDictionaryCodecInterface *codec_;
-  const uint16 suggestion_only_word_id_;
-  DictionaryInterface::Callback *callback_;
-  Token token_;
-
-  DISALLOW_COPY_AND_ASSIGN(PredictiveTraverser);
-};
-
 }  // namespace
 
 void ValueDictionary::LookupPredictive(
@@ -196,11 +171,46 @@
   if (key.empty()) {
     return;
   }
-  string lookup_key_str;
-  codec_->EncodeValue(key, &lookup_key_str);
-  DCHECK(value_trie_.get() != NULL);
-  PredictiveTraverser traverser(codec_, suggestion_only_word_id_, callback);
-  value_trie_->PredictiveSearch(lookup_key_str.c_str(), &traverser);
+  string encoded_key;
+  codec_->EncodeValue(key, &encoded_key);
+
+  LoudsTrie::Node node;
+  if (!value_trie_->Traverse(encoded_key, &node)) {
+    return;
+  }
+
+  char encoded_value_buffer[LoudsTrie::kMaxDepth + 1];
+  string value;
+  value.reserve(key.size() * 2);
+  Token token;
+
+  // Traverse subtree rooted at |node|.
+  queue<LoudsTrie::Node> queue;
+  queue.push(node);
+  do {
+    node = queue.front();
+    queue.pop();
+
+    if (value_trie_->IsTerminalNode(node)) {
+      switch (HandleTerminalNode(*value_trie_, *codec_,
+                                 suggestion_only_word_id_,
+                                 node, callback, encoded_value_buffer,
+                                 &value, &token)) {
+        case Callback::TRAVERSE_DONE:
+          return;
+        case Callback::TRAVERSE_CULL:
+          continue;
+        default:
+          break;
+      }
+    }
+
+    for (value_trie_->MoveToFirstChild(&node);
+         value_trie_->IsValidNode(node);
+         value_trie_->MoveToNextSibling(&node)) {
+      queue.push(node);
+    }
+  } while (!queue.empty());
 }
 
 void ValueDictionary::LookupPrefix(
diff --git a/src/mozc_version_template.txt b/src/mozc_version_template.txt
index 2d27f45..6e9fe1b 100644
--- a/src/mozc_version_template.txt
+++ b/src/mozc_version_template.txt
@@ -1,6 +1,6 @@
 MAJOR=2
 MINOR=16
-BUILD=2042
+BUILD=2043
 REVISION=102
 # NACL_DICTIONARY_VERSION is the target version of the system dictionary to be
 # downloaded by NaCl Mozc.
diff --git a/src/storage/louds/louds_trie.cc b/src/storage/louds/louds_trie.cc
index ea9c939..e953a76 100644
--- a/src/storage/louds/louds_trie.cc
+++ b/src/storage/louds/louds_trie.cc
@@ -120,100 +120,13 @@
   return -1;
 }
 
-namespace {
-
-// Traverses the subtree rooted at |node| in DFS order and runs |callback| at
-// each terminal node.  Returns true if traversal should be stopped.
-bool TraverseSubTree(
-    const LoudsTrie &trie,
-    const KeyExpansionTable &key_expansion_table,
-    LoudsTrie::Callback *callback,
-    LoudsTrie::Node node,
-    char *key_buffer,
-    StringPiece::size_type key_len) {
-  if (trie.IsTerminalNode(node)) {
-    const LoudsTrie::Callback::ResultType result =
-        callback->Run(key_buffer, key_len, trie.GetKeyIdOfTerminalNode(node));
-    switch (result) {
-      case LoudsTrie::Callback::SEARCH_DONE:
-        return true;
-      case LoudsTrie::Callback::SEARCH_CULL:
-        return false;
-      default:
-        break;
-    }
-  }
-
-  for (trie.MoveToFirstChild(&node); trie.IsValidNode(node);
-       trie.MoveToNextSibling(&node)) {
-    key_buffer[key_len] = trie.GetEdgeLabelToParentNode(node);
-    if (TraverseSubTree(trie, key_expansion_table, callback, node,
-                        key_buffer, key_len + 1)) {
-      return true;
-    }
-  }
-  return false;
-}
-
-// Recursively traverses the trie in DFS order until terminal nodes for each
-// expanded key is found.  Then TraverseSubTree() is called at each terminal
-// node.  Returns true if traversal should be stopped.
-bool PredictiveSearchWithKeyExpansionImpl(
-    const LoudsTrie &trie,
-    StringPiece key,
-    const KeyExpansionTable &key_expansion_table,
-    LoudsTrie::Callback *callback,
-    LoudsTrie::Node node,
-    char *key_buffer,
-    StringPiece::size_type key_len) {
-  if (key_len == key.size()) {
-    return TraverseSubTree(trie, key_expansion_table, callback, node,
-                           key_buffer, key_len);
-  }
-
-  const char target_char = key[key_len];
-  const ExpandedKey &chars = key_expansion_table.ExpandKey(target_char);
-  for (trie.MoveToFirstChild(&node); trie.IsValidNode(node);
-       trie.MoveToNextSibling(&node)) {
-    const char c = trie.GetEdgeLabelToParentNode(node);
-    if (chars.IsHit(c)) {
-      key_buffer[key_len] = c;
-      if (PredictiveSearchWithKeyExpansionImpl(trie, key, key_expansion_table,
-                                               callback, node,
-                                               key_buffer, key_len + 1)) {
-        return true;
-      }
-    }
-  }
-  return false;
-}
-
-}  // namespace
-
-void LoudsTrie::PredictiveSearchWithKeyExpansion(
-    StringPiece key, const KeyExpansionTable &key_expansion_table,
-    Callback *callback) const {
-  char key_buffer[kMaxDepth + 1];
-  PredictiveSearchWithKeyExpansionImpl(
-      *this, key, key_expansion_table, callback,
-      Node(),  // Root
-      key_buffer, 0);
-}
-
-StringPiece LoudsTrie::RestoreKeyString(int key_id, char *buf) const {
-  // TODO(noriyukit): Check if it's really necessary to handle this case.
-  if (key_id < 0) {
-    return StringPiece();
-  }
-
+StringPiece LoudsTrie::RestoreKeyString(Node node, char *buf) const {
   // Ensure the returned StringPiece is null-terminated.
   char *const buf_end = buf + kMaxDepth;
   *buf_end = '\0';
 
   // Climb up the trie to the root and fill |buf| backward.
   char *ptr = buf_end;
-  Node node;
-  GetTerminalNodeFromKeyId(key_id, &node);
   for (; !louds_.IsRoot(node); louds_.MoveToParent(&node)) {
     *--ptr = GetEdgeLabelToParentNode(node);
   }
diff --git a/src/storage/louds/louds_trie.h b/src/storage/louds/louds_trie.h
index 5d8551b..d7044b9 100644
--- a/src/storage/louds/louds_trie.h
+++ b/src/storage/louds/louds_trie.h
@@ -32,7 +32,6 @@
 
 #include "base/port.h"
 #include "base/string_piece.h"
-#include "storage/louds/key_expansion_table.h"
 #include "storage/louds/louds.h"
 #include "storage/louds/simple_succinct_bit_vector_index.h"
 
@@ -48,32 +47,6 @@
   // This class stores a traversal state.
   typedef Louds::Node Node;
 
-  // Interface which is called back when the result is found.
-  class Callback {
-   public:
-    enum ResultType {
-      // Finishes the current (prefix or predictive) search.
-      SEARCH_DONE,
-
-      // Continues the current (prefix or predictive) search.
-      SEARCH_CONTINUE,
-
-      // Finished the (prefix or predictive) search of the current edge,
-      // but still continues to search for other edges.
-      SEARCH_CULL,
-    };
-
-    virtual ~Callback() {
-    }
-
-    // The callback will be invoked when a target word is found.
-    // "s" may not be null-terminated.
-    virtual ResultType Run(const char *s, size_t len, int key_id) = 0;
-
-   protected:
-    Callback() {}
-  };
-
   LoudsTrie() : edge_character_(nullptr) {}
   ~LoudsTrie() {}
 
@@ -122,12 +95,24 @@
     return node;
   }
 
+  // Restores the key string that reaches to |node|.  The caller is
+  // responsible for allocating a buffer for the result StringPiece, which needs
+  // to be passed in |buf|.  The returned StringPiece points to a piece of
+  // |buf|.
+  // REQUIRES: |buf| is longer than kMaxDepth + 1.
+  StringPiece RestoreKeyString(Node node, char *buf) const;
+
   // Restores the key string corresponding to |key_id|.  The caller is
   // responsible for allocating a buffer for the result StringPiece, which needs
   // to be passed in |buf|.  The returned StringPiece points to a piece of
   // |buf|.
   // REQUIRES: |buf| is longer than kMaxDepth + 1.
-  StringPiece RestoreKeyString(int key_id, char *buf) const;
+  StringPiece RestoreKeyString(int key_id, char *buf) const {
+    // TODO(noriyukit): Check if it's necessary to handle negative IDs.
+    return key_id < 0
+        ? StringPiece()
+        : RestoreKeyString(GetTerminalNodeFromKeyId(key_id), buf);
+  }
 
   // Methods for moving node exported from Louds class; see louds.h.
   void MoveToFirstChild(Node *node) const {
@@ -195,25 +180,6 @@
     }
   }
 
-  // TODO(noriyukit): The following search methods rely on Callback.  However,
-  // this results in the nested callback to implement SystemDictionary's lookup
-  // methods (i.e., inside implementations of DictionaryInterface::Callback,
-  // LoudsTrie::Callback needs to be called; see system_dictionary.cc.  This is
-  // somewhat inefficient because both requires virtual method calls.
-  // Therefore, it'd be better to implement the following search methods
-  // directly in system_dictionary.cc using more generic APIs defined above.
-
-  // Searches the trie structure, and invokes callback->Run when for each word
-  // which begins with key is found.
-  void PredictiveSearch(StringPiece key, Callback *callback) const {
-    PredictiveSearchWithKeyExpansion(
-        key, KeyExpansionTable::GetDefaultInstance(), callback);
-  }
-
-  void PredictiveSearchWithKeyExpansion(
-      StringPiece key, const KeyExpansionTable &key_expansion_table,
-      Callback *callback) const;
-
  private:
   Louds louds_;  // Tree structure representation by LOUDS.
 
diff --git a/src/storage/louds/louds_trie_test.cc b/src/storage/louds/louds_trie_test.cc
index 213e61c..2427faf 100644
--- a/src/storage/louds/louds_trie_test.cc
+++ b/src/storage/louds/louds_trie_test.cc
@@ -48,55 +48,6 @@
 namespace louds {
 namespace {
 
-class LoudsTrieTest : public ::testing::Test {
-};
-
-class TestCallback : public LoudsTrie::Callback {
- public:
-  TestCallback() : current_index_(0), limit_(numeric_limits<size_t>::max()) {
-  }
-
-  void AddExpectation(
-      const string &s, size_t len, int id, ResultType result_type) {
-    expectation_list_.push_back(Expectation());
-    Expectation &expectation = expectation_list_.back();
-    expectation.s = s;
-    expectation.len = len;
-    expectation.id = id;
-    expectation.result_type = result_type;
-  }
-
-  virtual ResultType Run(const char *s, size_t len, int id) {
-    if (current_index_ >= expectation_list_.size()) {
-      ADD_FAILURE() << "Too many invocations: " << expectation_list_.size()
-                    << " vs " << current_index_;
-      // Quit the callback.
-      return SEARCH_DONE;
-    }
-    const Expectation &expectation = expectation_list_[current_index_];
-    EXPECT_EQ(expectation.s, string(s, len)) << current_index_;
-    EXPECT_EQ(expectation.len, len) << current_index_;
-    EXPECT_EQ(expectation.id, id) << current_index_;
-    ++current_index_;
-    return expectation.result_type;
-  }
-
-  size_t num_invoked() const { return current_index_; }
-
- private:
-  struct Expectation {
-    string s;
-    size_t len;
-    int id;
-    ResultType result_type;
-  };
-  vector<Expectation> expectation_list_;
-  size_t current_index_;
-  size_t limit_;
-
-  DISALLOW_COPY_AND_ASSIGN(TestCallback);
-};
-
 class RecordCallbackArgs {
  public:
   struct CallbackArgs {
@@ -128,7 +79,7 @@
   return node;
 }
 
-TEST_F(LoudsTrieTest, NodeBasedApis) {
+TEST(LoudsTrieTest, NodeBasedApis) {
   // Create the following trie (* stands for non-terminal nodes):
   //
   //        *          Key   ID
@@ -351,7 +302,7 @@
   }
 }
 
-TEST_F(LoudsTrieTest, HasKey) {
+TEST(LoudsTrieTest, HasKey) {
   LoudsTrieBuilder builder;
   builder.Add("a");
   builder.Add("abc");
@@ -384,7 +335,7 @@
   EXPECT_FALSE(trie.HasKey("bcxyz"));
 }
 
-TEST_F(LoudsTrieTest, ExactSearch) {
+TEST(LoudsTrieTest, ExactSearch) {
   LoudsTrieBuilder builder;
   builder.Add("a");
   builder.Add("abc");
@@ -418,7 +369,7 @@
   trie.Close();
 }
 
-TEST_F(LoudsTrieTest, PrefixSearch) {
+TEST(LoudsTrieTest, PrefixSearch) {
   LoudsTrieBuilder builder;
   builder.Add("aa");
   builder.Add("ab");
@@ -492,229 +443,7 @@
   }
 }
 
-TEST_F(LoudsTrieTest, PredictiveSearch) {
-  LoudsTrieBuilder builder;
-  builder.Add("aa");
-  builder.Add("ab");
-  builder.Add("abc");
-  builder.Add("abcd");
-  builder.Add("abcde");
-  builder.Add("abcdef");
-  builder.Add("abcea");
-  builder.Add("abcef");
-  builder.Add("abd");
-  builder.Add("ebd");
-  builder.Add("\x01\xFF");
-  builder.Add("\x01\xFF\xFF");
-
-  builder.Build();
-  LoudsTrie trie;
-  trie.Open(reinterpret_cast<const uint8 *>(builder.image().data()));
-
-  {
-    TestCallback callback;
-    callback.AddExpectation("abc", 3, builder.GetId("abc"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abcd", 4, builder.GetId("abcd"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abcde", 5, builder.GetId("abcde"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abcdef", 6, builder.GetId("abcdef"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abcea", 5, builder.GetId("abcea"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abcef", 5, builder.GetId("abcef"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-
-    trie.PredictiveSearch("abc", &callback);
-    EXPECT_EQ(6, callback.num_invoked());
-  }
-  {
-    TestCallback callback;
-    callback.AddExpectation("aa", 2, builder.GetId("aa"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("ab", 2, builder.GetId("ab"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abc", 3, builder.GetId("abc"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abcd", 4, builder.GetId("abcd"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abcde", 5, builder.GetId("abcde"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abcdef", 6, builder.GetId("abcdef"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abcea", 5, builder.GetId("abcea"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abcef", 5, builder.GetId("abcef"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abd", 3, builder.GetId("abd"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-
-    trie.PredictiveSearch("a", &callback);
-    EXPECT_EQ(9, callback.num_invoked());
-  }
-  {
-    // Make sure non-ascii characters can be found.
-    TestCallback callback;
-    callback.AddExpectation("\x01\xFF", 2, builder.GetId("\x01\xFF"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("\x01\xFF\xFF", 3, builder.GetId("\x01\xFF\xFF"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-
-    trie.PredictiveSearch("\x01", &callback);
-    EXPECT_EQ(2, callback.num_invoked());
-  }
-
-  trie.Close();
-}
-
-TEST_F(LoudsTrieTest, PredictiveSearchWithLimit) {
-  LoudsTrieBuilder builder;
-  builder.Add("aa");
-  builder.Add("ab");
-  builder.Add("abc");
-  builder.Add("abcd");
-  builder.Add("abcde");
-  builder.Add("abcdef");
-  builder.Add("abcea");
-  builder.Add("abcef");
-  builder.Add("abd");
-  builder.Add("ebd");
-
-  builder.Build();
-  LoudsTrie trie;
-  trie.Open(reinterpret_cast<const uint8 *>(builder.image().data()));
-
-  {
-    TestCallback callback;
-    callback.AddExpectation("abc", 3, builder.GetId("abc"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abcd", 4, builder.GetId("abcd"),
-                            LoudsTrie::Callback::SEARCH_DONE);
-
-    trie.PredictiveSearch("abc", &callback);
-    EXPECT_EQ(2, callback.num_invoked());
-  }
-  {
-    TestCallback callback;
-    callback.AddExpectation("aa", 2, builder.GetId("aa"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("ab", 2, builder.GetId("ab"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abc", 3, builder.GetId("abc"),
-                            LoudsTrie::Callback::SEARCH_DONE);
-
-    trie.PredictiveSearch("a", &callback);
-    EXPECT_EQ(3, callback.num_invoked());
-  }
-
-  trie.Close();
-}
-
-TEST_F(LoudsTrieTest, PredictiveSearchWithKeyExpansion) {
-  LoudsTrieBuilder builder;
-  builder.Add("abc");
-  builder.Add("adc");
-  builder.Add("cbc");
-  builder.Add("ddc");
-
-  builder.Build();
-  LoudsTrie trie;
-  trie.Open(reinterpret_cast<const uint8 *>(builder.image().data()));
-
-  KeyExpansionTable key_expansion_table;
-  key_expansion_table.Add('b', "d");
-
-  {
-    TestCallback callback;
-    callback.AddExpectation("abc", 3, builder.GetId("abc"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("adc", 3, builder.GetId("adc"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-
-    trie.PredictiveSearchWithKeyExpansion(
-        "ab", key_expansion_table, &callback);
-    EXPECT_EQ(2, callback.num_invoked());
-  }
-
-  {
-    TestCallback callback;
-    callback.AddExpectation("adc", 3, builder.GetId("adc"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    trie.PredictiveSearchWithKeyExpansion(
-        "ad", key_expansion_table, &callback);
-    EXPECT_EQ(1, callback.num_invoked());
-  }
-
-  trie.Close();
-}
-
-TEST_F(LoudsTrieTest, PredictiveSearchCulling) {
-  LoudsTrieBuilder builder;
-  builder.Add("a");
-  builder.Add("ab");
-  builder.Add("abc");
-  builder.Add("abcd");
-  builder.Add("ae");
-  builder.Add("aec");
-  builder.Add("aecd");
-
-  builder.Build();
-  LoudsTrie trie;
-  trie.Open(reinterpret_cast<const uint8 *>(builder.image().data()));
-
-  {
-    TestCallback callback;
-    callback.AddExpectation("a", 1, builder.GetId("a"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("ab", 2, builder.GetId("ab"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abc", 3, builder.GetId("abc"),
-                            LoudsTrie::Callback::SEARCH_CULL);
-    // No callback for abcd.
-    callback.AddExpectation("ae", 2, builder.GetId("ae"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("aec", 3, builder.GetId("aec"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("aecd", 4, builder.GetId("aecd"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-
-    trie.PredictiveSearch("a", &callback);
-    EXPECT_EQ(6, callback.num_invoked());
-  }
-
-  trie.Close();
-}
-
-TEST_F(LoudsTrieTest, PredictiveSearchCulling2) {
-  LoudsTrieBuilder builder;
-  builder.Add("abc");
-  builder.Add("abcd");
-  builder.Add("abcde");
-  builder.Add("abxy");
-  builder.Add("abxyz");
-
-  builder.Build();
-  LoudsTrie trie;
-  trie.Open(reinterpret_cast<const uint8 *>(builder.image().data()));
-
-  {
-    TestCallback callback;
-    callback.AddExpectation("abc", 3, builder.GetId("abc"),
-                            LoudsTrie::Callback::SEARCH_CULL);
-    callback.AddExpectation("abxy", 4, builder.GetId("abxy"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-    callback.AddExpectation("abxyz", 5, builder.GetId("abxyz"),
-                            LoudsTrie::Callback::SEARCH_CONTINUE);
-
-    trie.PredictiveSearch("a", &callback);
-    EXPECT_EQ(3, callback.num_invoked());
-  }
-
-  trie.Close();
-}
-
-TEST_F(LoudsTrieTest, RestoreKeyString) {
+TEST(LoudsTrieTest, RestoreKeyString) {
   LoudsTrieBuilder builder;
   builder.Add("aa");
   builder.Add("ab");