1 files changed, 69 insertions, 17 deletions

diff --git a/dom/media/platforms/SimpleMap.h b/dom/media/platforms/SimpleMap.h
index c26bff1e9a..635ba6f085 100644
--- a/dom/media/platforms/SimpleMap.h
+++ b/dom/media/platforms/SimpleMap.h
@@ -5,49 +5,101 @@
 #ifndef mozilla_SimpleMap_h
 #define mozilla_SimpleMap_h
 
+#include <utility>
+
+#include "mozilla/Maybe.h"
 #include "mozilla/Mutex.h"
 #include "nsTArray.h"
 
-#include <utility>
-
 namespace mozilla {
 
-template <typename T>
+struct ThreadSafePolicy {
+  struct PolicyLock {
+    explicit PolicyLock(const char* aName) : mMutex(aName) {}
+    Mutex mMutex MOZ_UNANNOTATED;
+  };
+  PolicyLock& mPolicyLock;
+  explicit ThreadSafePolicy(PolicyLock& aPolicyLock)
+      : mPolicyLock(aPolicyLock) {
+    mPolicyLock.mMutex.Lock();
+  }
+  ~ThreadSafePolicy() { mPolicyLock.mMutex.Unlock(); }
+};
+
+struct NoOpPolicy {
+  struct PolicyLock {
+    explicit PolicyLock(const char*) {}
+  };
+  explicit NoOpPolicy(PolicyLock&) {}
+  ~NoOpPolicy() = default;
+};
+
+// An map employing an array instead of a hash table to optimize performance,
+// particularly beneficial when the number of expected items in the map is
+// small.
+template <typename K, typename V, typename Policy = NoOpPolicy>
 class SimpleMap {
- public:
-  typedef std::pair<int64_t, T> Element;
+  using ElementType = std::pair<K, V>;
+  using MapType = AutoTArray<ElementType, 16>;
 
-  SimpleMap() : mMutex("SimpleMap") {}
+ public:
+  SimpleMap() : mLock("SimpleMap"){};
 
+  // Check if aKey is in the map.
+  bool Contains(const K& aKey) {
+    struct Comparator {
+      bool Equals(const ElementType& aElement, const K& aKey) const {
+        return aElement.first == aKey;
+      }
+    };
+    Policy guard(mLock);
+    return mMap.Contains(aKey, Comparator());
+  }
   // Insert Key and Value pair at the end of our map.
-  void Insert(int64_t aKey, const T& aValue) {
-    MutexAutoLock lock(mMutex);
+  void Insert(const K& aKey, const V& aValue) {
+    Policy guard(mLock);
     mMap.AppendElement(std::make_pair(aKey, aValue));
   }
   // Sets aValue matching aKey and remove it from the map if found.
   // The element returned is the first one found.
   // Returns true if found, false otherwise.
-  bool Find(int64_t aKey, T& aValue) {
-    MutexAutoLock lock(mMutex);
+  bool Find(const K& aKey, V& aValue) {
+    if (Maybe<V> v = Take(aKey)) {
+      aValue = v.extract();
+      return true;
+    }
+    return false;
+  }
+  // Take the value matching aKey and remove it from the map if found.
+  Maybe<V> Take(const K& aKey) {
+    Policy guard(mLock);
     for (uint32_t i = 0; i < mMap.Length(); i++) {
-      Element& element = mMap[i];
+      ElementType& element = mMap[i];
       if (element.first == aKey) {
-        aValue = element.second;
+        Maybe<V> value = Some(element.second);
         mMap.RemoveElementAt(i);
-        return true;
+        return value;
       }
     }
-    return false;
+    return Nothing();
   }
   // Remove all elements of the map.
   void Clear() {
-    MutexAutoLock lock(mMutex);
+    Policy guard(mLock);
     mMap.Clear();
   }
+  // Iterate through all elements of the map and call the function F.
+  template <typename F>
+  void ForEach(F&& aCallback) {
+    Policy guard(mLock);
+    for (const auto& element : mMap) {
+      aCallback(element.first, element.second);
+    }
+  }
 
  private:
-  Mutex mMutex MOZ_UNANNOTATED;  // To protect mMap.
-  AutoTArray<Element, 16> mMap;
+  typename Policy::PolicyLock mLock;
+  MapType mMap;
 };
 
 }  // namespace mozilla