Allocate less, switch to existing being continuation

3 files changed, 43 insertions, 41 deletions

diff --git a/search/edge.hh b/search/edge.hh
index 2cdc05f37..3b096c1ae 100644
--- a/search/edge.hh
+++ b/search/edge.hh
@@ -62,12 +62,16 @@ class PartialEdge {
 
 class PartialEdgePool {
   public:
-    PartialEdge Allocate(Arity arity) {
+    PartialEdge Allocate(Arity arity, Arity chart_states) {
       return PartialEdge(
-          pool_.Allocate(sizeof(Score) + sizeof(Arity) + arity * sizeof(PartialVertex) + (arity + 1) * sizeof(lm::ngram::ChartState)),
+          pool_.Allocate(sizeof(Score) + sizeof(Arity) + arity * sizeof(PartialVertex) + chart_states * sizeof(lm::ngram::ChartState)),
           arity);
     }
 
+    PartialEdge Allocate(Arity arity) {
+      return Allocate(arity, arity + 1);
+    }
+
   private:
     util::Pool pool_;
 };
diff --git a/search/edge_generator.cc b/search/edge_generator.cc
index 43e0188fa..17caab866 100644
--- a/search/edge_generator.cc
+++ b/search/edge_generator.cc
@@ -16,14 +16,9 @@ EdgeGenerator::EdgeGenerator(PartialEdge root, Note note) : top_score_(root.GetS
 
 namespace {
 
-template <class Model> void FastScore(const Context<Model> &context, Arity victim, Arity before_idx, Arity incomplete, const PartialEdge previous, PartialEdge update) {
+template <class Model> void FastScore(const Context<Model> &context, Arity victim, Arity before_idx, Arity incomplete, const PartialVertex &previous_vertex, PartialEdge update) {
   lm::ngram::ChartState *between = update.Between();
-  const lm::ngram::ChartState *previous_between = previous.Between();
-  const search::PartialVertex &previous_vertex = previous.NT()[victim];
-
   lm::ngram::ChartState *before = &between[before_idx], *after = &between[before_idx + 1];
-  // copy [0, after] 
-  memcpy(between, previous_between, sizeof(lm::ngram::ChartState) * (before_idx + 2));
 
   float adjustment = 0.0;
   const lm::ngram::ChartState &previous_reveal = previous_vertex.State();
@@ -43,13 +38,12 @@ template <class Model> void FastScore(const Context<Model> &context, Arity victi
       adjustment += lm::ngram::Subsume(context.LanguageModel(), before->left, before->right, after->left, after->right, update_reveal.left.length);
     }
     before->right = after->right;
-    // Copy the others shifted one down, covering after.  
-    memcpy(after, previous_between + before_idx + 2, sizeof(lm::ngram::ChartState) * (incomplete + 1 - before_idx - 2));
-  } else {
-    // Copy [after + 1, incomplete]
-    memcpy(after + 1, previous_between + before_idx + 2, sizeof(lm::ngram::ChartState) * (incomplete + 1 - before_idx - 2));
+    // Shift the others shifted one down, covering after.  
+    for (lm::ngram::ChartState *cover = after; cover < between + incomplete; ++cover) {
+      *cover = *(cover + 1);
+    }
   }
-  update.SetScore(previous.GetScore() + adjustment * context.GetWeights().LM());
+  update.SetScore(update.GetScore() + adjustment * context.GetWeights().LM());
 }
 
 } // namespace
@@ -62,9 +56,10 @@ template <class Model> PartialEdge EdgeGenerator::Pop(Context<Model> &context, P
 
   Arity victim = 0;
   Arity victim_completed;
-  Arity completed = 0;
+  Arity incomplete;
   // Select victim or return if complete.   
   {
+    Arity completed = 0;
     unsigned char lowest_length = 255;
     for (Arity i = 0; i != arity_; ++i) {
       if (top_nt[i].Complete()) {
@@ -80,31 +75,31 @@ template <class Model> PartialEdge EdgeGenerator::Pop(Context<Model> &context, P
       top_score_ = generate_.empty() ? -kScoreInf : generate_.top().GetScore();
       return top;
     }
+    incomplete = arity_ - completed;
   }
 
-  float old_bound = top_nt[victim].Bound();
-  PartialEdge continuation = partial_edge_pool.Allocate(arity_);
-  PartialVertex *continuation_nt = continuation.NT();
-  // The alternate's score will change because the nt changes.
-  bool split = top_nt[victim].Split(continuation_nt[victim]);
-  // top is now the alternate.  
+  PartialVertex old_value(top_nt[victim]);
+  PartialVertex alternate_changed;
+  if (top_nt[victim].Split(alternate_changed)) {
+    PartialEdge alternate = partial_edge_pool.Allocate(arity_, incomplete + 1);
+    alternate.SetScore(top.GetScore() + alternate_changed.Bound() - old_value.Bound());
 
-  for (Arity i = 0; i < victim; ++i) continuation_nt[i] = top_nt[i];
-  for (Arity i = victim + 1; i < arity_; ++i) continuation_nt[i] = top_nt[i];
-  FastScore(context, victim, victim - victim_completed, arity_ - completed, top, continuation);
-  // TODO: dedupe?  
-  generate_.push(continuation);
+    PartialVertex *alternate_nt = alternate.NT();
+    for (Arity i = 0; i < victim; ++i) alternate_nt[i] = top_nt[i];
+    alternate_nt[victim] = alternate_changed;
+    for (Arity i = victim + 1; i < arity_; ++i) alternate_nt[i] = top_nt[i];
+
+    memcpy(alternate.Between(), top.Between(), sizeof(lm::ngram::ChartState) * (incomplete + 1));
 
-  if (split) {
-    // We have an alternate.  
-    top.SetScore(top.GetScore() + top_nt[victim].Bound() - old_bound);
     // TODO: dedupe?  
-    generate_.push(top);
-  } else {
-    // TODO should free top here. 
-    // Better would be changing Split.  
+    generate_.push(alternate);
   }
 
+  // top is now the continuation.
+  FastScore(context, victim, victim - victim_completed, incomplete, old_value, top);
+  // TODO: dedupe?  
+  generate_.push(top);
+
   top_score_ = generate_.top().GetScore();
   // Invalid indicates no new hypothesis generated.  
   return PartialEdge();
diff --git a/search/vertex.hh b/search/vertex.hh
index 4ccb24713..1cdcae30d 100644
--- a/search/vertex.hh
+++ b/search/vertex.hh
@@ -105,16 +105,20 @@ class PartialVertex {
       return index_ + 1 < back_->Size();
     }
 
-    // Split into continuation and alternative, rendering this the alternative.
-    bool Split(PartialVertex &continuation) {
+    // Split into continuation and alternative, rendering this the continuation.
+    bool Split(PartialVertex &alternative) {
       assert(!Complete());
-      continuation.back_ = &((*back_)[index_]);
-      continuation.index_ = 0;
+      bool ret;
       if (index_ + 1 < back_->Size()) {
-        ++index_;
-        return true;
+        alternative.index_ = index_ + 1;
+        alternative.back_ = back_;
+        ret = true;
+      } else {
+        ret = false;
       }
-      return false;
+      back_ = &((*back_)[index_]);
+      index_ = 0;
+      return ret;
     }
 
     const Final &End() const {
@@ -140,7 +144,6 @@ class Vertex {
         PartialVertex continuation;
         while (!top.Complete()) {
           top.Split(continuation);
-          top = continuation;
         }
         return &top.End();
       }