1 files changed, 18 insertions, 15 deletions

diff --git a/src/include/fst/extensions/pdt/collection.h b/src/include/fst/extensions/pdt/collection.h
index 26be504..24a443f 100644
--- a/src/include/fst/extensions/pdt/collection.h
+++ b/src/include/fst/extensions/pdt/collection.h
@@ -16,7 +16,7 @@
 // Author: riley@google.com (Michael Riley)
 //
 // \file
-// Class to store a collection of sets with elements of type T.
+// Class to store a collection of ordered (multi-)sets with elements of type T.
 
 #ifndef FST_EXTENSIONS_PDT_COLLECTION_H__
 #define FST_EXTENSIONS_PDT_COLLECTION_H__
@@ -29,11 +29,11 @@ using std::vector;
 
 namespace fst {
 
-// Stores a collection of non-empty sets with elements of type T. A
-// default constructor, equality ==, a total order <, and an STL-style
-// hash class must be defined on the elements. Provides signed
-// integer ID (of type I) of each unique set. The IDs are allocated
-// starting from 0 in order.
+// Stores a collection of non-empty, ordered (multi-)sets with elements
+// of type T. A default constructor, equality ==, and an STL-style
+// hash class must be defined on the elements. Provides signed integer
+// ID (of type I) of each unique set. The IDs are allocated starting
+// from 0 in order.
 template <class I, class T>
 class Collection {
  public:
@@ -80,31 +80,34 @@ class Collection {
 
   Collection() {}
 
-  // Lookups integer ID from set. If it doesn't exist, then adds it.
-  // Set elements should be in strict order (and therefore unique).
-  I FindId(const vector<T> &set) {
+  // Lookups integer ID from ordered multi-set. If it doesn't exist
+  // and 'insert' is true, then adds it. Otherwise returns -1.
+  I FindId(const vector<T> &set, bool insert = true) {
     I node_id = kNoNodeId;
     for (ssize_t i = set.size() - 1; i >= 0; --i) {
       Node node(node_id, set[i]);
-      node_id = node_table_.FindId(node);
+      node_id = node_table_.FindId(node, insert);
+      if (node_id == -1) break;
     }
     return node_id;
   }
 
-  // Finds set given integer ID. Returns true if ID corresponds
-  // to set. Use iterators below to traverse result.
+  // Finds ordered (multi-)set given integer ID. Returns set iterator
+  // to traverse result.
   SetIterator FindSet(I id) {
-    if (id < 0 && id >= node_table_.Size()) {
+    if (id < 0 || id >= node_table_.Size()) {
       return SetIterator(kNoNodeId, Node(kNoNodeId, T()), &node_table_);
     } else {
       return SetIterator(id, node_table_.FindEntry(id), &node_table_);
     }
   }
 
+  I Size() const { return node_table_.Size(); }
+
  private:
   static const I kNoNodeId;
   static const size_t kPrime;
-  static std::tr1::hash<T> hash_;
+  static std::hash<T> hash_;
 
   NodeTable node_table_;
 
@@ -115,7 +118,7 @@ template<class I, class T> const I Collection<I, T>::kNoNodeId = -1;
 
 template <class I, class T> const size_t Collection<I, T>::kPrime = 7853;
 
-template <class I, class T> std::tr1::hash<T> Collection<I, T>::hash_;
+template <class I, class T> std::hash<T> Collection<I, T>::hash_;
 
 }  // namespace fst