From db748e924b629a8f06019f6d0ce7251c6f9410fc Mon Sep 17 00:00:00 2001
From: changq <changq>
Date: Thu, 20 Nov 2014 22:34:28 +0800
Subject: [PATCH] Add support for order=1 to TrieModel

---
 lm/model.cc       |  1 -
 lm/search_trie.cc | 30 +++++++++++++++++++++---------
 lm/trie_sort.hh   |  2 +-
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/lm/model.cc b/lm/model.cc
index a5a16bf8e..feed6a597 100644
--- a/lm/model.cc
+++ b/lm/model.cc
@@ -97,7 +97,6 @@ template <class Search, class VocabularyT> void GenericModel<Search, VocabularyT
     // File counts do not include pruned trigrams that extend to quadgrams etc.   These will be fixed by search_.
     ReadARPACounts(f, counts);
     CheckCounts(counts);
-    if (counts.size() < 2) UTIL_THROW(FormatLoadException, "This ngram implementation assumes at least a bigram model.");
     if (config.probing_multiplier <= 1.0) UTIL_THROW(ConfigException, "probing multiplier must be > 1.0");
 
     std::size_t vocab_size = util::CheckOverflow(VocabularyT::Size(counts[0], config));
diff --git a/lm/search_trie.cc b/lm/search_trie.cc
index 7fc70f4eb..bb4aeb4aa 100644
--- a/lm/search_trie.cc
+++ b/lm/search_trie.cc
@@ -201,9 +201,11 @@ class SRISucks {
       }
       messages_[0].Apply(it_, unigram_file);
       BackoffMessages *messages = messages_ + 1;
-      const RecordReader *end = reader + total_order - 2 /* exclude unigrams and longest order */;
-      for (; reader != end; ++messages, ++reader) {
-        messages->Apply(it_, *reader);
+      if (total_order >= 2) {
+        const RecordReader *end = reader + total_order - 2 /* exclude unigrams and longest order */;
+        for (; reader != end; ++messages, ++reader) {
+          messages->Apply(it_, *reader);
+        }
       }
     }
 
@@ -480,8 +482,10 @@ template <class Quant, class Bhiksha> void BuildTrie(SortedFiles &files, std::ve
     fixed_counts = finder.Counts();
   }
   unigram_file.reset(util::FDOpenOrThrow(unigram_fd));
-  for (const RecordReader *i = inputs; i != inputs + counts.size() - 2; ++i) {
-    if (*i) UTIL_THROW(FormatLoadException, "There's a bug in the trie implementation: the " << (i - inputs + 2) << "-gram table did not complete reading");
+  if (counts.size() > 1) {
+    for (const RecordReader *i = inputs; i != inputs + counts.size() - 2; ++i) {
+      if (*i) UTIL_THROW(FormatLoadException, "There's a bug in the trie implementation: the " << (i - inputs + 2) << "-gram table did not complete reading");
+    }
   }
   SanityCheckCounts(counts, fixed_counts);
   counts = fixed_counts;
@@ -514,7 +518,7 @@ template <class Quant, class Bhiksha> void BuildTrie(SortedFiles &files, std::ve
     inputs[i-2].Rewind();
   }
   // Fill entries except unigram probabilities.
-  {
+  if (counts.size() > 1) {
     WriteEntries<Quant, Bhiksha> writer(contexts, quant, unigrams, out.middle_begin_, out.longest_, counts.size(), sri);
     RecursiveInsert(counts.size(), counts[0], inputs, config.ProgressMessages(), "Writing trie", writer);
     // Write the last unigram entry, which is the end pointer for the bigrams.  
@@ -552,9 +556,17 @@ template <class Quant, class Bhiksha> uint8_t *TrieSearch<Quant, Bhiksha>::Setup
   unigram_.Init(start);
   start += Unigram::Size(counts[0]);
   FreeMiddles();
-  middle_begin_ = static_cast<Middle*>(malloc(sizeof(Middle) * (counts.size() - 2)));
-  middle_end_ = middle_begin_ + (counts.size() - 2);
-  std::vector<uint8_t*> middle_starts(counts.size() - 2);
+  std::vector<uint8_t*> middle_starts;
+  if (counts.size() > 2) {
+      middle_begin_ = static_cast<Middle*>(malloc(sizeof(Middle) * (counts.size() - 2)));
+      middle_end_ = middle_begin_ + (counts.size() - 2);
+      middle_starts.resize(counts.size() - 2);
+  }
+  else {
+      middle_begin_ = NULL;
+      middle_end_ = NULL;
+  }
+
   for (unsigned char i = 2; i < counts.size(); ++i) {
     middle_starts[i-2] = start;
     start += Middle::Size(Quant::MiddleBits(config), counts[i-1], counts[0], counts[i], config);
diff --git a/lm/trie_sort.hh b/lm/trie_sort.hh
index e5406d9b6..65dd7f5e2 100644
--- a/lm/trie_sort.hh
+++ b/lm/trie_sort.hh
@@ -48,7 +48,7 @@ class EntryCompare : public std::binary_function<const void*, const void*, bool>
 
 class RecordReader {
   public:
-    RecordReader() : remains_(true) {}
+    RecordReader() : remains_(false) {}
 
     void Init(FILE *file, std::size_t entry_size);