Bug 2069391 - Make main thread atom caches work for the UTF-8 path as well. r=smaug

Claude wrote the tests but they look sensible to me (basically a copy of
the utf-16 string).

Differential Revision: https://phabricator.services.mozilla.com/D323649

layout/style/GeckoBindings.cpp

@@ -952,7 +952,11 @@ SERVO_IMPL_ELEMENT_ATTR_MATCHING_FUNCTIONS(Gecko_Snapshot,
 #undef SERVO_IMPL_ELEMENT_ATTR_MATCHING_FUNCTIONS
 
 nsAtom* Gecko_Atomize(const char* aString, uint32_t aLength) {
-  return NS_Atomize(nsDependentCSubstring(aString, aLength)).take();
+  nsDependentCSubstring str(aString, aLength);
+  if (NS_IsMainThread()) {
+    return NS_AtomizeMainThread(str).take();
+  }
+  return NS_Atomize(str).take();
 }
 
 nsAtom* Gecko_Atomize16(const nsAString* aString) {

xpcom/ds/nsAtom.h

@@ -273,8 +273,9 @@ already_AddRefed<nsAtom> NS_Atomize(const nsAString& aUTF16String);
 already_AddRefed<nsAtom> NS_Atomize(const nsAString& aUTF16String,
                                     uint32_t aKnownHash);
 
-// An optimized version of the method above for the main thread.
+// Optimized versions of the methods above for the main thread.
 already_AddRefed<nsAtom> NS_AtomizeMainThread(const nsAString& aUTF16String);
+already_AddRefed<nsAtom> NS_AtomizeMainThread(const nsACString& aUTF8String);
 
 // Return a count of the total number of atoms currently alive in the system.
 //

xpcom/ds/nsAtomTable.cpp

@@ -4,6 +4,8 @@
 
 #include "nsAtomTable.h"
 
+#include <type_traits>
+
 #include "PLDHashTable.h"
 #include "mozilla/AppShutdown.h"
 #include "mozilla/Assertions.h"
@@ -147,6 +149,15 @@ struct AtomTableKey {
         mLength(aLength),
         mHash(HashUTF8AsUTF16(aUTF8String, aLength)) {}
 
+  bool Equals(const nsAtom* aAtom) const {
+    if (mUTF8String) {
+      return CompareUTF8toUTF16(
+                 nsDependentCSubstring(mUTF8String, mUTF8String + mLength),
+                 nsDependentAtomString(aAtom)) == 0;
+    }
+    return aAtom->Equals(mUTF16String, mLength);
+  }
+
   const char16_t* mUTF16String;
   const char* mUTF8String;
   uint32_t mLength;
@@ -164,16 +175,7 @@ struct AtomTableEntry : public PLDHashEntryHdr {
   AtomTableEntry(AtomTableEntry&&) = default;
 
   // NOTE: GetKey cannot be implemented.
-  bool KeyEquals(KeyTypePointer aKey) const {
+  bool KeyEquals(KeyTypePointer aKey) const { return aKey->Equals(mAtom); }
-    if (aKey->mUTF8String) {
-      return CompareUTF8toUTF16(
-                 nsDependentCSubstring(aKey->mUTF8String,
-                                       aKey->mUTF8String + aKey->mLength),
-                 nsDependentAtomString(mAtom)) == 0;
-    }
-
-    return mAtom->Equals(aKey->mUTF16String, aKey->mLength);
-  }
 
   static KeyTypePointer KeyToPointer(KeyType aKey) { return &aKey; }
   static PLDHashNumber HashKey(KeyTypePointer aKey) { return aKey->mHash; }
@@ -189,9 +191,7 @@ struct AtomTableEntry : public PLDHashEntryHdr {
 struct AtomCache : public MruCache<AtomTableKey, nsAtom*, AtomCache> {
   static HashNumber Hash(const AtomTableKey& aKey) { return aKey.mHash; }
   static bool Match(const AtomTableKey& aKey, const nsAtom* aVal) {
-    MOZ_ASSERT(aKey.mUTF16String);
+    return aVal->hash() == aKey.mHash && aKey.Equals(aVal);
-    return (aVal->hash() == aKey.mHash) &&
-           aVal->Equals(aKey.mUTF16String, aKey.mLength);
   }
 };
 
@@ -232,17 +232,23 @@ struct ShortAtomCache {
   // Rotates to indicate which way to replace next if all are full in some set.
   uint8_t mNextWay = 0;
 
-  static Signature TryMakeSignature(const char16_t* aStr, size_t aLength) {
+  // The UTF-8 and UTF-16 overloads must produce the same signature for the
+  // same string, so the UTF-8 one only accepts ASCII: any byte > 0x7f starts
+  // a multi-byte sequence that would need decoding first.
+  template <typename CharT>
+  static Signature TryMakeSignature(const CharT* aStr, size_t aLength) {
     static_assert(sizeof(Signature) >= kMaxLength + 1);
+    constexpr uint32_t kMaxChar = sizeof(CharT) == 1 ? 0x7f : 0xff;
     if (aLength == 0 || aLength > kMaxLength) {
       return kInvalidSignature;
     }
     Signature signature = aLength;
     for (size_t i = 0; i < aLength; i++) {
-      if (aStr[i] > 0xff) {
+      const auto c = static_cast<std::make_unsigned_t<CharT>>(aStr[i]);
+      if (c > kMaxChar) {
         return kInvalidSignature;
       }
-      signature = (signature << 8) | static_cast<uint8_t>(aStr[i]);
+      signature = (signature << 8) | static_cast<uint8_t>(c);
     }
     return signature;
   }
@@ -333,8 +339,11 @@ class nsAtomTable {
                                    uint32_t aHash);
   already_AddRefed<nsAtom> Atomize(const nsACString& aUTF8String);
   already_AddRefed<nsAtom> AtomizeMainThread(const nsAString& aUTF16String);
+  already_AddRefed<nsAtom> AtomizeMainThread(const nsACString& aUTF8String);
   already_AddRefed<nsAtom> GetOrInsert(const nsAString& aUTF16String,
                                        AtomTableKey& key);
+  already_AddRefed<nsAtom> GetOrInsert(const nsACString& aUTF8String,
+                                       AtomTableKey& key);
   nsStaticAtom* GetStaticAtom(const nsAString& aUTF16String);
   void RegisterStaticAtoms(const nsStaticAtom* aAtoms, size_t aAtomsLen);
 
@@ -644,16 +653,22 @@ already_AddRefed<nsAtom> NS_Atomize(const char* aUTF8String) {
 
 already_AddRefed<nsAtom> nsAtomTable::Atomize(const nsACString& aUTF8String) {
   AtomTableKey key(aUTF8String.Data(), aUTF8String.Length());
-  nsAtomSubTable& table = SelectSubTable(key);
+  return GetOrInsert(aUTF8String, key);
+}
+
+already_AddRefed<nsAtom> nsAtomTable::GetOrInsert(const nsACString& aUTF8String,
+                                                  AtomTableKey& aKey) {
+  MOZ_ASSERT(aKey.mUTF8String == aUTF8String.Data());
+  nsAtomSubTable& table = SelectSubTable(aKey);
   {
     AutoReadLock lock(table.mLock);
-    if (AtomTableEntry* he = table.Search(key)) {
+    if (AtomTableEntry* he = table.Search(aKey)) {
       return do_AddRef(he->mAtom);
     }
   }
 
   AutoWriteLock lock(table.mLock);
-  AtomTableEntry* he = table.Add(key);
+  AtomTableEntry* he = table.Add(aKey);
 
   if (he->mAtom) {
     return do_AddRef(he->mAtom);
@@ -662,7 +677,7 @@ already_AddRefed<nsAtom> nsAtomTable::Atomize(const nsACString& aUTF8String) {
   nsString str;
   CopyUTF8toUTF16(aUTF8String, str);
   MOZ_ASSERT(str.GetStringBuffer(), "Should create a string buffer");
-  RefPtr<nsAtom> atom = dont_AddRef(nsDynamicAtom::Create(str, key.mHash));
+  RefPtr<nsAtom> atom = dont_AddRef(nsDynamicAtom::Create(str, aKey.mHash));
 
   he->mAtom = atom;
 
@@ -749,6 +764,39 @@ already_AddRefed<nsAtom> NS_AtomizeMainThread(const nsAString& aUTF16String) {
   return gAtomTable->AtomizeMainThread(aUTF16String);
 }
 
+already_AddRefed<nsAtom> nsAtomTable::AtomizeMainThread(
+    const nsACString& aUTF8String) {
+  MOZ_ASSERT(NS_IsMainThread());
+
+  size_t length = aUTF8String.Length();
+  const char* str = aUTF8String.Data();
+
+  if (auto sig = ShortAtomCache::TryMakeSignature(str, length)) {
+    if (nsAtom* cached = sShortAtomCache.Lookup(sig)) {
+      return do_AddRef(cached);
+    }
+    RefPtr<nsAtom> retVal = Atomize(aUTF8String);
+    sShortAtomCache.Put(sig, retVal);
+    return retVal.forget();
+  }
+
+  AtomTableKey key(str, length);
+  RefPtr<nsAtom> retVal;
+  auto p = sRecentlyUsedMainThreadAtoms.Lookup(key);
+  if (p) {
+    retVal = p.Data();
+  } else {
+    retVal = GetOrInsert(aUTF8String, key);
+    p.Set(retVal);
+  }
+  return retVal.forget();
+}
+
+already_AddRefed<nsAtom> NS_AtomizeMainThread(const nsACString& aUTF8String) {
+  MOZ_ASSERT(gAtomTable);
+  return gAtomTable->AtomizeMainThread(aUTF8String);
+}
+
 nsrefcnt NS_GetNumberOfAtoms(void) {
   MOZ_ASSERT(gAtomTable);
   return gAtomTable->RacySlowCount();

xpcom/tests/gtest/TestAtoms.cpp

@@ -86,6 +86,52 @@ TEST(Atoms, ShortStringCacheCorrectness)
   }
 }
 
+static constexpr nsLiteralCString kShortCacheTestUTF8Strings[] = {
+    ""_ns,
+    "abcdefgh"_ns,
+    "a"_ns,
+    "ab"_ns,
+    "\0"_ns,
+    "abcdefg"_ns,
+    "a\0"_ns,
+    "\0a"_ns,
+    "\u00ff"_ns,
+    "\u0100"_ns,
+    "abc\u00ff"_ns,
+    "\u0100abc"_ns,
+    "abcdef\0"_ns,
+    "abcdef\u0100"_ns,
+    "a much longer string"_ns,
+    "a much longer string with \u00ff\u0100\U0001f600"_ns,
+};
+
+TEST(Atoms, ShortStringCacheCorrectnessUTF8)
+{
+  for (const auto& str : kShortCacheTestUTF8Strings) {
+    RefPtr<nsAtom> uncached = NS_Atomize(str);
+    RefPtr<nsAtom> miss = NS_AtomizeMainThread(str);
+    RefPtr<nsAtom> hit = NS_AtomizeMainThread(str);
+
+    EXPECT_EQ(uncached, miss);
+    EXPECT_EQ(miss, hit);
+    EXPECT_TRUE(hit->Equals(NS_ConvertUTF8toUTF16(str)));
+  }
+}
+
+TEST(Atoms, MainThreadCacheUTF8vs16)
+{
+  for (const auto& str : kShortCacheTestUTF8Strings) {
+    NS_ConvertUTF8toUTF16 str16(str);
+    RefPtr<nsAtom> atom16 = NS_AtomizeMainThread(str16);
+    RefPtr<nsAtom> atom8 = NS_AtomizeMainThread(str);
+    EXPECT_EQ(atom16, atom8);
+    RefPtr<nsAtom> atom8Again = NS_AtomizeMainThread(str);
+    EXPECT_EQ(atom8, atom8Again);
+    RefPtr<nsAtom> atom16Again = NS_AtomizeMainThread(str16);
+    EXPECT_EQ(atom8, atom16Again);
+  }
+}
+
 TEST(Atoms, ShortStringCacheEviction)
 {
   const uint64_t kCount = TestGetShortAtomCacheSize() + 10;
@@ -105,6 +151,25 @@ TEST(Atoms, ShortStringCacheEviction)
   }
 }
 
+TEST(Atoms, ShortStringCacheEvictionUTF8)
+{
+  const uint64_t kCount = TestGetShortAtomCacheSize() + 10;
+
+  constexpr size_t len = 2;
+  ASSERT_LE(kCount, uint64_t(1) << (len * 7));
+
+  for (uint64_t pass = 0; pass < 2; ++pass) {
+    for (uint64_t i = 0; i < kCount; ++i) {
+      char buf[len] = {char(i & 0x7f), char((i >> 7) & 0x7f)};
+      nsDependentCSubstring str(buf, len);
+      RefPtr<nsAtom> atom = NS_AtomizeMainThread(str);
+      RefPtr<nsAtom> atom2 = NS_Atomize(str);
+      EXPECT_TRUE(atom->Equals(NS_ConvertUTF8toUTF16(str)));
+      EXPECT_EQ(atom, atom2);
+    }
+  }
+}
+
 TEST(Atoms, Invalid)
 {
   for (unsigned int i = 0; i < std::size(Invalid16Strings); ++i) {