Corner case left state was not being fully minimized

1 files changed, 24 insertions, 15 deletions

diff --git a/kenlm/lm/left.hh b/kenlm/lm/left.hh
index 15464c82c..41f71f849 100644
--- a/kenlm/lm/left.hh
+++ b/kenlm/lm/left.hh
@@ -175,24 +175,14 @@ template <class M> class RuleScore {
 
       float backoffs[kMaxOrder - 1], backoffs2[kMaxOrder - 1];
       float *back = backoffs, *back2 = backoffs2;
-      unsigned char next_use;
+      unsigned char next_use = out_.right.length;
 
       // First word
-      ProcessRet(model_.ExtendLeft(out_.right.words, out_.right.words + out_.right.length, out_.right.backoff, in.left.pointers[0], 1, back, next_use));
-      if (!next_use) {
-        left_done_ = true;
-        out_.right = in.right;
-        return;
-      }
+      if (ExtendLeft(in, next_use, 1, out_.right.backoff, back)) return;
+
       // Words after the first, so extending a bigram to begin with
-      unsigned char extend_length = 2;
-      for (const uint64_t *i = in.left.pointers + 1; i < in.left.pointers + in.left.length; ++i, ++extend_length) {
-        ProcessRet(model_.ExtendLeft(out_.right.words, out_.right.words + next_use, back, *i, extend_length, back2, next_use));
-        if (!next_use) {
-          left_done_ = true;
-          out_.right = in.right;
-          return;
-        }
+      for (unsigned char extend_length = 2; extend_length <= in.left.length; ++extend_length) {
+        if (ExtendLeft(in, next_use, extend_length, back, back2)) return;
         std::swap(back, back2);
       }
 
@@ -228,6 +218,25 @@ template <class M> class RuleScore {
     }
 
   private:
+    bool ExtendLeft(const ChartState &in, unsigned char &next_use, unsigned char extend_length, const float *back_in, float *back_out) {
+      ProcessRet(model_.ExtendLeft(
+            out_.right.words, out_.right.words + next_use, // Words to extend into
+            back_in, // Backoffs to use
+            in.left.pointers[extend_length - 1], extend_length, // Words to be extended
+            back_out, // Backoffs for the next score
+            next_use)); // Length of n-gram to use in next scoring.  
+      if (next_use != out_.right.length) {
+        left_done_ = true;
+        if (!next_use) {
+          out_.right = in.right;
+          // Early exit.  
+          return true;
+        }
+      }
+      // Continue scoring.  
+      return false;
+    }
+
     void ProcessRet(const FullScoreReturn &ret) {
       prob_ += ret.prob;
       if (left_done_) return;