From 2946822248d682c7e244037d3d5d6f849f3657c9 Mon Sep 17 00:00:00 2001
From: Jan-Niklas Jaeschke <jjaschke@mozilla.com>
Date: Tue, 23 Apr 2024 09:42:15 +0000
Subject: [PATCH] Bug 1892589 - Custom Highlight API: Implemented speedup for
 `Highlight::Add()`. r=edgar,dom-core

To preserve uniqueness in the mirrored data structure for the JS-setlike object,
`nsTArray::Contains()` was used. This showed up in profiles due to its `O(n)` time complexity.
Performance could be improved by replacing it with `SetlikeHelpers::Has()`, which uses the more efficient version of the underlying data structure of the setlike.

The same idea has been applied to `HighlightRegistry`, which uses a maplike. However, this is mostly for consistency, since `Highlight::Add()` will likely be called more often and hence the bigger performance bottleneck.

Differential Revision: https://phabricator.services.mozilla.com/D208071
---
 dom/base/Highlight.cpp         | 40 ++++++++++++++++++++++------------
 dom/base/HighlightRegistry.cpp | 22 ++++++++++++++-----
 2 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/dom/base/Highlight.cpp b/dom/base/Highlight.cpp
index cd0efccce58c..0bac8193d72c 100644
--- a/dom/base/Highlight.cpp
+++ b/dom/base/Highlight.cpp
@@ -101,24 +101,36 @@ already_AddRefed<Selection> Highlight::CreateHighlightSelection(
 }
 
 void Highlight::Add(AbstractRange& aRange, ErrorResult& aRv) {
+  // Manually check if the range `aKey` is already present in this highlight,
+  // because `SetlikeHelpers::Add()` doesn't indicate this.
+  // To keep the setlike and the mirrored array in sync, the range must not
+  // be added to `mRanges` if it was already present.
+  // `SetlikeHelpers::Has()` is much faster in checking this than
+  // `nsTArray<>::Contains()`.
+  if (Highlight_Binding::SetlikeHelpers::Has(this, aRange, aRv) ||
+      aRv.Failed()) {
+    return;
+  }
   Highlight_Binding::SetlikeHelpers::Add(this, aRange, aRv);
   if (aRv.Failed()) {
     return;
   }
-  if (!mRanges.Contains(&aRange)) {
-    mRanges.AppendElement(&aRange);
-    AutoFrameSelectionBatcher selectionBatcher(__FUNCTION__,
-                                               mHighlightRegistries.Count());
-    for (const RefPtr<HighlightRegistry>& registry :
-         mHighlightRegistries.Keys()) {
-      auto frameSelection = registry->GetFrameSelection();
-      selectionBatcher.AddFrameSelection(frameSelection);
-      // since this is run in a context guarded by a selection batcher,
-      // no strong reference is needed to keep `registry` alive.
-      MOZ_KnownLive(registry)->MaybeAddRangeToHighlightSelection(aRange, *this);
-      if (aRv.Failed()) {
-        return;
-      }
+
+  MOZ_ASSERT(!mRanges.Contains(&aRange),
+             "setlike and DOM mirror are not in sync");
+
+  mRanges.AppendElement(&aRange);
+  AutoFrameSelectionBatcher selectionBatcher(__FUNCTION__,
+                                             mHighlightRegistries.Count());
+  for (const RefPtr<HighlightRegistry>& registry :
+       mHighlightRegistries.Keys()) {
+    auto frameSelection = registry->GetFrameSelection();
+    selectionBatcher.AddFrameSelection(frameSelection);
+    // since this is run in a context guarded by a selection batcher,
+    // no strong reference is needed to keep `registry` alive.
+    MOZ_KnownLive(registry)->MaybeAddRangeToHighlightSelection(aRange, *this);
+    if (aRv.Failed()) {
+      return;
     }
   }
 }
diff --git a/dom/base/HighlightRegistry.cpp b/dom/base/HighlightRegistry.cpp
index 035aadb78a82..755977286f23 100644
--- a/dom/base/HighlightRegistry.cpp
+++ b/dom/base/HighlightRegistry.cpp
@@ -133,18 +133,28 @@ void HighlightRegistry::AddHighlightSelectionsToFrameSelection() {
 
 void HighlightRegistry::Set(const nsAString& aKey, Highlight& aValue,
                             ErrorResult& aRv) {
+  // manually check if the highlight `aKey` is already registered to be able to
+  // provide a fast path later that avoids calling `std::find_if()`.
+  const bool highlightAlreadyPresent =
+      HighlightRegistry_Binding::MaplikeHelpers::Has(this, aKey, aRv);
+  if (aRv.Failed()) {
+    return;
+  }
   HighlightRegistry_Binding::MaplikeHelpers::Set(this, aKey, aValue, aRv);
   if (aRv.Failed()) {
     return;
   }
   RefPtr<nsFrameSelection> frameSelection = GetFrameSelection();
   RefPtr<nsAtom> highlightNameAtom = NS_AtomizeMainThread(aKey);
-  auto foundIter =
-      std::find_if(mHighlightsOrdered.begin(), mHighlightsOrdered.end(),
-                   [&highlightNameAtom](auto const& aElm) {
-                     return aElm.first() == highlightNameAtom;
-                   });
-  if (foundIter != mHighlightsOrdered.end()) {
+  if (highlightAlreadyPresent) {
+    // If the highlight named `aKey` was present before, replace its value.
+    auto foundIter =
+        std::find_if(mHighlightsOrdered.begin(), mHighlightsOrdered.end(),
+                     [&highlightNameAtom](auto const& aElm) {
+                       return aElm.first() == highlightNameAtom;
+                     });
+    MOZ_ASSERT(foundIter != mHighlightsOrdered.end(),
+               "webIDL maplike and DOM mirror are not in sync");
     foundIter->second()->RemoveFromHighlightRegistry(*this, *highlightNameAtom);
     if (frameSelection) {
       frameSelection->RemoveHighlightSelection(highlightNameAtom);