Add an insert() method to MapVector. Adds the first MapVector unit test.

[oota-llvm.git] / include / llvm / ADT / MapVector.h

diff --git a/include/llvm/ADT/MapVector.h b/include/llvm/ADT/MapVector.h
index ffe1cd3eef3a6cb8b9a71555bdad822296828850..f29681f644e8d6c39639352b7ba1395ae563e631 100644 (file)
--- a/include/llvm/ADT/MapVector.h
+++ b/include/llvm/ADT/MapVector.h
@@ -19,6 +19,7 @@
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
 #include <vector>
 
 namespace llvm {
@@ -28,7 +29,7 @@ namespace llvm {
 /// mapping is done with DenseMap from Keys to indexes in that vector.
 template<typename KeyT, typename ValueT,
          typename MapType = llvm::DenseMap<KeyT, unsigned>,
-         typename VectorType = std::vector<std::pair<KeyT, ValueT> >>
+         typename VectorType = std::vector<std::pair<KeyT, ValueT> > >
 class MapVector {
   typedef typename VectorType::size_type SizeType;
 
@@ -78,6 +79,40 @@ public:
     }
     return Vector[I].second;
   }
+
+  ValueT lookup(const KeyT &Key) const {
+    typename MapType::const_iterator Pos = Map.find(Key);
+    return Pos == Map.end()? ValueT() : Vector[Pos->second].second;
+  }
+
+  std::pair<iterator, bool> insert(const std::pair<KeyT, ValueT> &KV) {
+    std::pair<KeyT, unsigned> Pair = std::make_pair(KV.first, 0);
+    std::pair<typename MapType::iterator, bool> Result = Map.insert(Pair);
+    unsigned &I = Result.first->second;
+    if (Result.second) {
+      Vector.push_back(std::make_pair(KV.first, KV.second));
+      I = Vector.size() - 1;
+      return std::make_pair(llvm::prior(end()), true);
+    }
+    return std::make_pair(begin() + I, false);
+  }
+
+  unsigned count(const KeyT &Key) const {
+    typename MapType::const_iterator Pos = Map.find(Key);
+    return Pos == Map.end()? 0 : 1;
+  }
+
+  iterator find(const KeyT &Key) {
+    typename MapType::const_iterator Pos = Map.find(Key);
+    return Pos == Map.end()? Vector.end() :
+                            (Vector.begin() + Pos->second);
+  }
+
+  const_iterator find(const KeyT &Key) const {
+    typename MapType::const_iterator Pos = Map.find(Key);
+    return Pos == Map.end()? Vector.end() :
+                            (Vector.begin() + Pos->second);
+  }
 };
 
 }