1 files changed, 22 insertions, 1 deletions
diff --git a/klm/util/probing_hash_table.hh b/klm/util/probing_hash_table.hh
index 57866ff9..51a2944d 100644
--- a/klm/util/probing_hash_table.hh
+++ b/klm/util/probing_hash_table.hh
@@ -109,9 +109,20 @@ template <class EntryT, class HashT, class EqualT = std::equal_to<typename Entry
         if (equal_(got, key)) { out = i; return true; }
         if (equal_(got, invalid_)) return false;
         if (++i == end_) i = begin_;
-      }   
+      }
+    }
+
+    // Like UnsafeMutableFind, but the key must be there.
+    template <class Key> MutableIterator UnsafeMutableMustFind(const Key key) {
+       for (MutableIterator i(begin_ + (hash_(key) % buckets_));;) {
+        Key got(i->GetKey());
+        if (equal_(got, key)) { return i; }
+        assert(!equal_(got, invalid_));
+        if (++i == end_) i = begin_;
+      }
     }
 
+
     template <class Key> bool Find(const Key key, ConstIterator &out) const {
 #ifdef DEBUG
       assert(initialized_);
@@ -124,6 +135,16 @@ template <class EntryT, class HashT, class EqualT = std::equal_to<typename Entry
       }    
     }
 
+    // Like Find but we're sure it must be there.
+    template <class Key> ConstIterator MustFind(const Key key) const {
+      for (ConstIterator i(begin_ + (hash_(key) % buckets_));;) {
+        Key got(i->GetKey());
+        if (equal_(got, key)) { return i; }
+        assert(!equal_(got, invalid_));
+        if (++i == end_) i = begin_;
+      }
+    }
+
     void Clear() {
       Entry invalid;
       invalid.SetKey(invalid_);