1 files changed, 19 insertions, 25 deletions

diff --git a/decoder/apply_models.cc b/decoder/apply_models.cc
index 3f3f6a79..278516c1 100644
--- a/decoder/apply_models.cc
+++ b/decoder/apply_models.cc
@@ -233,24 +233,19 @@ public:
   void IncorporateIntoPlusLMForest(size_t head_node_hash, Candidate* item, State2Node* s2n, CandidateList* freelist) {
     Hypergraph::Edge* new_edge = out.AddEdge(item->out_edge_);
     new_edge->edge_prob_ = item->out_edge_.edge_prob_;
-    //start: new code by lijunhui
-    FFState real_state;
-    FFState* real_state_ref;
-    if (models.HaveEraseState()) {
-       models.GetRealFFState(item->state_, real_state);
-       real_state_ref = &real_state;
+
+    Candidate** o_item_ptr = nullptr;
+    if (item->state_.size() && models.NeedsStateErasure()) {
+      // When erasure of certain state bytes is needed, we must make a copy of
+      // the state instead of doing the erasure in-place because future
+      // candidates may require the information in the bytes to be erased.
+      FFState state(item->state_);
+      models.EraseIgnoredBytes(&state);
+      o_item_ptr = &(*s2n)[state];
+    } else {
+      o_item_ptr = &(*s2n)[item->state_];
     }
-    else
-       real_state_ref = &(item->state_);
-    Candidate*& o_item = (*s2n)[(*real_state_ref)];
-    /*FFState real_state;
-    models.GetRealFFState(item->state_, real_state);
-    Candidate*& o_item = (*s2n)[real_state];*/
-    //end: new code by lijunhui
-
-    //start: original code
-    //Candidate*& o_item = (*s2n)[item->state_];
-    //end: original code
+    Candidate*& o_item = *o_item_ptr;
 
     if (!o_item) o_item = item;
 
@@ -272,18 +267,17 @@ public:
     // score is the same for all items with a common residual DP
     // state
     if (item->vit_prob_ > o_item->vit_prob_) {
-      //start: new code by lijunhui
-      if (models.HaveEraseState()) {
-        assert(models.GetRealFFState(o_item->state_) == models.GetRealFFState(item->state_));    // sanity check!
+      if (item->state_.size() && models.NeedsStateErasure()) {
+        // node_states_ should still point to the unerased state.
         node_states_[o_item->node_index_] = item->state_;
+        // sanity check!
+        FFState item_state(item->state_), o_item_state(o_item->state_);
+        models.EraseIgnoredBytes(&item_state);
+        models.EraseIgnoredBytes(&o_item_state);
+        assert(item_state == o_item_state);
       } else {
         assert(o_item->state_ == item->state_);    // sanity check!
       }
-      //end: new code by lijunhui
-
-      //start: original code
-      //assert(o_item->state_ == item->state_);    // sanity check!
-      //end: original code
 
       o_item->est_prob_ = item->est_prob_;
       o_item->vit_prob_ = item->vit_prob_;