1 files changed, 8 insertions, 4 deletions
diff --git a/klm/search/edge_generator.cc b/klm/search/edge_generator.cc
index eacf5de5..dd9d61e4 100644
--- a/klm/search/edge_generator.cc
+++ b/klm/search/edge_generator.cc
@@ -54,20 +54,20 @@ template <class Model> PartialEdge EdgeGenerator::Pop(Context<Model> &context) {
   Arity victim = 0;
   Arity victim_completed;
   Arity incomplete;
+  unsigned char lowest_niceness = 255;
   // 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()) {
         ++completed;
-      } else if (top_nt[i].Length() < lowest_length) {
-        lowest_length = top_nt[i].Length();
+      } else if (top_nt[i].Niceness() < lowest_niceness) {
+        lowest_niceness = top_nt[i].Niceness();
         victim = i;
         victim_completed = completed;
       }
     }
-    if (lowest_length == 255) {
+    if (lowest_niceness == 255) {
       return top;
     }
     incomplete = arity - completed;
@@ -92,10 +92,14 @@ template <class Model> PartialEdge EdgeGenerator::Pop(Context<Model> &context) {
     generate_.push(alternate);
   }
 
+#ifndef NDEBUG  
+  Score before = top.GetScore();
+#endif
   // top is now the continuation.
   FastScore(context, victim, victim - victim_completed, incomplete, old_value, top);
   // TODO: dedupe?  
   generate_.push(top);
+  assert(lowest_niceness != 254 || top.GetScore() == before);
 
   // Invalid indicates no new hypothesis generated.  
   return PartialEdge();