1 files changed, 28 insertions, 27 deletions

diff --git a/klm/search/vertex.hh b/klm/search/vertex.hh
index e1a9ad11..52bc1dfe 100644
--- a/klm/search/vertex.hh
+++ b/klm/search/vertex.hh
@@ -18,7 +18,7 @@ class ContextBase;
 
 class VertexNode {
   public:
-    VertexNode() : end_(NULL) {}
+    VertexNode() {}
 
     void InitRoot() {
       extend_.clear();
@@ -26,8 +26,7 @@ class VertexNode {
       state_.left.length = 0;
       state_.right.length = 0;
       right_full_ = false;
-      bound_ = -kScoreInf;
-      end_ = NULL;
+      end_ = Final();
     }
 
     lm::ngram::ChartState &MutableState() { return state_; }
@@ -37,19 +36,20 @@ class VertexNode {
       extend_.push_back(next);
     }
 
-    void SetEnd(Final *end) { end_ = end; }
+    void SetEnd(Final end) {
+      assert(!end_.Valid());
+      end_ = end;
+    }
     
-    Final &MutableEnd() { return *end_; }
-
     void SortAndSet(ContextBase &context, VertexNode **parent_pointer);
 
     // Should only happen to a root node when the entire vertex is empty.   
     bool Empty() const {
-      return !end_ && extend_.empty();
+      return !end_.Valid() && extend_.empty();
     }
 
     bool Complete() const {
-      return end_;
+      return end_.Valid();
     }
 
     const lm::ngram::ChartState &State() const { return state_; }
@@ -63,8 +63,8 @@ class VertexNode {
       return state_.left.length + state_.right.length;
     }
 
-    // May be NULL.
-    const Final *End() const { return end_; }
+    // Will be invalid unless this is a leaf.   
+    const Final End() const { return end_; }
 
     const VertexNode &operator[](size_t index) const {
       return *extend_[index];
@@ -81,7 +81,7 @@ class VertexNode {
     bool right_full_;
 
     Score bound_;
-    Final *end_;
+    Final end_;
 };
 
 class PartialVertex {
@@ -97,7 +97,7 @@ class PartialVertex {
     const lm::ngram::ChartState &State() const { return back_->State(); }
     bool RightFull() const { return back_->RightFull(); }
 
-    Score Bound() const { return Complete() ? back_->End()->Bound() : (*back_)[index_].Bound(); }
+    Score Bound() const { return Complete() ? back_->End().GetScore() : (*back_)[index_].Bound(); }
 
     unsigned char Length() const { return back_->Length(); }
 
@@ -105,20 +105,24 @@ 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 {
-      return *back_->End();
+    const Final End() const {
+      return back_->End();
     }
 
   private:
@@ -126,25 +130,22 @@ class PartialVertex {
     unsigned int index_;
 };
 
-extern PartialVertex kBlankPartialVertex;
-
 class Vertex {
   public:
     Vertex() {}
 
     PartialVertex RootPartial() const { return PartialVertex(root_); }
 
-    const Final *BestChild() const {
+    const Final BestChild() const {
       PartialVertex top(RootPartial());
       if (top.Empty()) {
-        return NULL;
+        return Final();
       } else {
         PartialVertex continuation;
         while (!top.Complete()) {
           top.Split(continuation);
-          top = continuation;
         }
-        return &top.End();
+        return top.End();
       }
     }