Bug 1808565, part 2: Use `SelectionBatcher` when updating Highlights. r=masayuki

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

dom/base/Highlight.cpp

@@ -6,6 +6,7 @@
 
 #include "Highlight.h"
 #include "HighlightRegistry.h"
+#include "mozilla/StaticAnalysisFunctions.h"
 #include "mozilla/dom/HighlightBinding.h"
 
 #include "mozilla/AlreadyAddRefed.h"
@@ -87,7 +88,8 @@ already_AddRefed<Selection> Highlight::CreateHighlightSelection(
   RefPtr<Selection> selection =
       MakeRefPtr<Selection>(SelectionType::eHighlight, aFrameSelection);
   selection->SetHighlightName(aHighlightName);
-  SelectionBatcher selectionBatcher(selection, __FUNCTION__);
+  AutoFrameSelectionBatcher selectionBatcher(__FUNCTION__);
+  selectionBatcher.AddFrameSelection(aFrameSelection);
   // NOLINTNEXTLINE(performance-for-range-copy)
   for (const RefPtr<AbstractRange> range : mRanges) {
     if (range->GetComposedDocOfContainers() ==
@@ -106,8 +108,13 @@ void Highlight::Add(AbstractRange& aRange, ErrorResult& aRv) {
   }
   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);
+
       registry->MaybeAddRangeToHighlightSelection(aRange, *this, aRv);
       if (aRv.Failed()) {
         return;
@@ -120,9 +127,16 @@ void Highlight::Clear(ErrorResult& aRv) {
   Highlight_Binding::SetlikeHelpers::Clear(this, aRv);
   if (!aRv.Failed()) {
     mRanges.Clear();
-    for (const RefPtr<HighlightRegistry> registry :
+    AutoFrameSelectionBatcher selectionBatcher(__FUNCTION__,
+                                               mHighlightRegistries.Count());
+
+    for (const RefPtr<HighlightRegistry>& registry :
          mHighlightRegistries.Keys()) {
-      registry->RemoveHighlightSelection(*this);
+      auto frameSelection = registry->GetFrameSelection();
+      selectionBatcher.AddFrameSelection(frameSelection);
+      // Because of the selection batcher, this call does *not* run script.
+      // MOZ_KnownLive() is needed regardless.
+      MOZ_KnownLive(registry)->RemoveHighlightSelection(*this);
     }
   }
 }
@@ -130,9 +144,17 @@ void Highlight::Clear(ErrorResult& aRv) {
 bool Highlight::Delete(AbstractRange& aRange, ErrorResult& aRv) {
   if (Highlight_Binding::SetlikeHelpers::Delete(this, aRange, aRv)) {
     mRanges.RemoveElement(&aRange);
-    for (const RefPtr<HighlightRegistry> registry :
+    AutoFrameSelectionBatcher selectionBatcher(__FUNCTION__,
+                                               mHighlightRegistries.Count());
+
+    for (const RefPtr<HighlightRegistry>& registry :
          mHighlightRegistries.Keys()) {
-      registry->MaybeRemoveRangeFromHighlightSelection(aRange, *this);
+      auto frameSelection = registry->GetFrameSelection();
+      selectionBatcher.AddFrameSelection(frameSelection);
+      // Because of the selection batcher, this call does *not* run script.
+      // MOZ_KnownLive() is needed regardless.
+      MOZ_KnownLive(registry)->MaybeRemoveRangeFromHighlightSelection(aRange,
+                                                                      *this);
     }
     return true;
   }

dom/base/HighlightRegistry.cpp

@@ -172,14 +172,20 @@ void HighlightRegistry::Clear(ErrorResult& aRv) {
     return;
   }
   auto frameSelection = GetFrameSelection();
+  AutoFrameSelectionBatcher batcher(__FUNCTION__);
+  batcher.AddFrameSelection(frameSelection);
   for (auto const& iter : mHighlightsOrdered) {
-    const RefPtr<const nsAtom> highlightName = iter.first();
+    const RefPtr<const nsAtom>& highlightName = iter.first();
     const RefPtr<Highlight>& highlight = iter.second();
     highlight->RemoveFromHighlightRegistry(*this, *highlightName);
     if (frameSelection) {
-      frameSelection->RemoveHighlightSelection(highlightName);
+      // The selection batcher makes sure that no script is run in this call.
+      // However, `nsFrameSelection::RemoveHighlightSelection` is marked
+      // `MOZ_CAN_RUN_SCRIPT`, therefore `MOZ_KnownLive` is needed regardless.
+      frameSelection->RemoveHighlightSelection(MOZ_KnownLive(highlightName));
     }
   }
+
   mHighlightsOrdered.Clear();
 }
 

dom/base/HighlightRegistry.h

@@ -128,12 +128,12 @@ class HighlightRegistry final : public nsISupports, public nsWrapperCache {
    */
   MOZ_CAN_RUN_SCRIPT bool Delete(const nsAString& aKey, ErrorResult& aRv);
 
- private:
   /**
    * @brief Get the `FrameSelection` object if available. Can return nullptr.
    */
   RefPtr<nsFrameSelection> GetFrameSelection();
 
+ private:
   /**
    * Parent document.
    */

layout/generic/nsFrameSelection.h

@@ -1157,4 +1157,32 @@ class nsFrameSelection final {
   bool mIsDoubleClickSelection{false};
 };
 
+/**
+ * Selection Batcher class that supports multiple FrameSelections.
+ */
+class MOZ_STACK_CLASS AutoFrameSelectionBatcher {
+ public:
+  explicit AutoFrameSelectionBatcher(const char* aFunctionName,
+                                     size_t aEstimatedSize = 1)
+      : mFunctionName(aFunctionName) {
+    mFrameSelections.SetCapacity(aEstimatedSize);
+  }
+  ~AutoFrameSelectionBatcher() {
+    for (const auto& frameSelection : mFrameSelections) {
+      frameSelection->EndBatchChanges(mFunctionName);
+    }
+  }
+  void AddFrameSelection(nsFrameSelection* aFrameSelection) {
+    if (!aFrameSelection) {
+      return;
+    }
+    aFrameSelection->StartBatchChanges(mFunctionName);
+    mFrameSelections.AppendElement(aFrameSelection);
+  }
+
+ private:
+  const char* mFunctionName;
+  AutoTArray<RefPtr<nsFrameSelection>, 1> mFrameSelections;
+};
+
 #endif /* nsFrameSelection_h___ */