1 files changed, 66 insertions, 6 deletions
diff --git a/dom/base/nsIdentifierMapEntry.h b/dom/base/nsIdentifierMapEntry.h
index ea60d3322..41b8b4a83 100644
--- a/dom/base/nsIdentifierMapEntry.h
+++ b/dom/base/nsIdentifierMapEntry.h
@@ -40,22 +40,45 @@ class nsIContent;
  * Perhaps the document.all results should have their own hashtable
  * in nsHTMLDocument.
  */
-class nsIdentifierMapEntry : public nsStringHashKey
+class nsIdentifierMapEntry : public PLDHashEntryHdr
 {
 public:
+  struct AtomOrString
+  {
+    MOZ_IMPLICIT AtomOrString(nsIAtom* aAtom) : mAtom(aAtom) {}
+    MOZ_IMPLICIT AtomOrString(const nsAString& aString) : mString(aString) {}
+    AtomOrString(const AtomOrString& aOther)
+      : mAtom(aOther.mAtom)
+      , mString(aOther.mString)
+    {
+    }
+
+    AtomOrString(AtomOrString&& aOther)
+      : mAtom(aOther.mAtom.forget())
+      , mString(aOther.mString)
+    {
+    }
+
+    nsCOMPtr<nsIAtom> mAtom;
+    const nsString mString;
+  };
+
+  typedef const AtomOrString& KeyType;
+  typedef const AtomOrString* KeyTypePointer;
+
   typedef mozilla::dom::Element Element;
   typedef mozilla::net::ReferrerPolicy ReferrerPolicy;
 
-  explicit nsIdentifierMapEntry(const nsAString& aKey) :
-    nsStringHashKey(&aKey), mNameContentList(nullptr)
+  explicit nsIdentifierMapEntry(const AtomOrString& aKey)
+    : mKey(aKey)
   {
   }
-  explicit nsIdentifierMapEntry(const nsAString* aKey) :
-    nsStringHashKey(aKey), mNameContentList(nullptr)
+  explicit nsIdentifierMapEntry(const AtomOrString* aKey)
+    : mKey(aKey ? *aKey : nullptr)
   {
   }
   nsIdentifierMapEntry(nsIdentifierMapEntry&& aOther) :
-    nsStringHashKey(&aOther.GetKey()),
+    mKey(mozilla::Move(aOther.GetKey())),
     mIdContentList(mozilla::Move(aOther.mIdContentList)),
     mNameContentList(aOther.mNameContentList.forget()),
     mChangeCallbacks(aOther.mChangeCallbacks.forget()),
@@ -64,6 +87,42 @@ public:
   }
   ~nsIdentifierMapEntry();
 
+  KeyType GetKey() const { return mKey; }
+
+  nsString GetKeyAsString() const
+  {
+    if (mKey.mAtom) {
+      return nsAtomString(mKey.mAtom);
+    }
+
+    return mKey.mString;
+  }
+
+  bool KeyEquals(const KeyTypePointer aOtherKey) const
+  {
+    if (mKey.mAtom) {
+      if (aOtherKey->mAtom) {
+        return mKey.mAtom == aOtherKey->mAtom;
+      }
+
+      return mKey.mAtom->Equals(aOtherKey->mString);
+    }
+
+    if (aOtherKey->mAtom) {
+      return aOtherKey->mAtom->Equals(mKey.mString);
+    }
+
+    return mKey.mString.Equals(aOtherKey->mString);
+  }
+
+  static KeyTypePointer KeyToPointer(KeyType aKey) { return &aKey; }
+
+  static PLDHashNumber HashKey(const KeyTypePointer aKey)
+  {
+    return aKey->mAtom ?
+      aKey->mAtom->hash() : mozilla::HashString(aKey->mString);
+  }
+
   enum { ALLOW_MEMMOVE = false };
 
   void AddNameElement(nsINode* aDocument, Element* aElement);
@@ -167,6 +226,7 @@ private:
   void FireChangeCallbacks(Element* aOldElement, Element* aNewElement,
                            bool aImageOnly = false);
 
+  AtomOrString mKey;
   // empty if there are no elements with this ID.
   // The elements are stored as weak pointers.
   AutoTArray<Element*, 1> mIdContentList;