1 files changed, 24 insertions, 7 deletions

diff --git a/runtime/safe_map.h b/runtime/safe_map.h
index bf3a15eec..941fd0e2d 100644
--- a/runtime/safe_map.h
+++ b/runtime/safe_map.h
@@ -34,10 +34,12 @@ class SafeMap {
 
  public:
   typedef typename ::std::map<K, V, Comparator, Allocator>::key_compare key_compare;
+  typedef typename ::std::map<K, V, Comparator, Allocator>::value_compare value_compare;
   typedef typename ::std::map<K, V, Comparator, Allocator>::allocator_type allocator_type;
   typedef typename ::std::map<K, V, Comparator, Allocator>::iterator iterator;
   typedef typename ::std::map<K, V, Comparator, Allocator>::const_iterator const_iterator;
   typedef typename ::std::map<K, V, Comparator, Allocator>::size_type size_type;
+  typedef typename ::std::map<K, V, Comparator, Allocator>::key_type key_type;
   typedef typename ::std::map<K, V, Comparator, Allocator>::value_type value_type;
 
   SafeMap() = default;
@@ -50,6 +52,9 @@ class SafeMap {
     return *this;
   }
 
+  key_compare key_comp() const { return map_.key_comp(); }
+  value_compare value_comp() const { return map_.value_comp(); }
+
   iterator begin() { return map_.begin(); }
   const_iterator begin() const { return map_.begin(); }
   iterator end() { return map_.end(); }
@@ -58,8 +63,9 @@ class SafeMap {
   bool empty() const { return map_.empty(); }
   size_type size() const { return map_.size(); }
 
+  void swap(Self& other) { map_.swap(other.map_); }
   void clear() { map_.clear(); }
-  void erase(iterator it) { map_.erase(it); }
+  iterator erase(iterator it) { return map_.erase(it); }
   size_type erase(const K& k) { return map_.erase(k); }
 
   iterator find(const K& k) { return map_.find(k); }
@@ -78,9 +84,18 @@ class SafeMap {
   }
 
   // Used to insert a new mapping.
-  void Put(const K& k, const V& v) {
-    std::pair<iterator, bool> result = map_.insert(std::make_pair(k, v));
+  iterator Put(const K& k, const V& v) {
+    std::pair<iterator, bool> result = map_.emplace(k, v);
     DCHECK(result.second);  // Check we didn't accidentally overwrite an existing value.
+    return result.first;
+  }
+
+  // Used to insert a new mapping at a known position for better performance.
+  iterator PutBefore(iterator pos, const K& k, const V& v) {
+    // Check that we're using the correct position and the key is not in the map.
+    DCHECK(pos == map_.end() || map_.key_comp()(k, pos->first));
+    DCHECK(pos == map_.begin() || map_.key_comp()((--iterator(pos))->first, k));
+    return map_.emplace_hint(pos, k, v);
   }
 
   // Used to insert a new mapping or overwrite an existing mapping. Note that if the value type
@@ -102,13 +117,15 @@ class SafeMap {
   ::std::map<K, V, Comparator, Allocator> map_;
 };
 
-template <typename K, typename V, typename Comparator>
-bool operator==(const SafeMap<K, V, Comparator>& lhs, const SafeMap<K, V, Comparator>& rhs) {
+template <typename K, typename V, typename Comparator, typename Allocator>
+bool operator==(const SafeMap<K, V, Comparator, Allocator>& lhs,
+                const SafeMap<K, V, Comparator, Allocator>& rhs) {
   return lhs.Equals(rhs);
 }
 
-template <typename K, typename V, typename Comparator>
-bool operator!=(const SafeMap<K, V, Comparator>& lhs, const SafeMap<K, V, Comparator>& rhs) {
+template <typename K, typename V, typename Comparator, typename Allocator>
+bool operator!=(const SafeMap<K, V, Comparator, Allocator>& lhs,
+                const SafeMap<K, V, Comparator, Allocator>& rhs) {
   return !(lhs == rhs);
 }