Bug 1488599 - Part 2: Fix will-change budget r=mattwoodrow

Depends on D5245

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

--HG--
extra : moz-landing-system : lando

layout/base/nsLayoutUtils.cpp

@@ -3700,11 +3700,15 @@ nsLayoutUtils::PaintFrame(gfxContext* aRenderingContext, nsIFrame* aFrame,
 
       if (updateState == PartialUpdateResult::Failed) {
         list.DeleteAll(&builder);
+
+        builder.ClearRetainedWindowRegions();
+        builder.ClearWillChangeBudget();
+
         builder.EnterPresShell(aFrame);
         builder.SetDirtyRect(visibleRect);
-        builder.ClearRetainedWindowRegions();
         aFrame->BuildDisplayListForStackingContext(&builder, &list);
-        AddExtraBackgroundItems(builder, list, aFrame, canvasArea, visibleRegion, aBackstop);
+        AddExtraBackgroundItems(
+          builder, list, aFrame, canvasArea, visibleRegion, aBackstop);
 
         builder.LeavePresShell(aFrame, &list);
         updateState = PartialUpdateResult::Updated;

layout/generic/nsFrame.cpp

@@ -986,6 +986,12 @@ nsIFrame::RemoveDisplayItemDataForDeletion()
 
   RetainedDisplayListData* data = GetOrSetRetainedDisplayListData(rootFrame);
 
+  if (MayHaveWillChangeBudget()) {
+    // Keep the frame in list, so it can be removed from the will-change budget.
+    data->Flags(this) = RetainedDisplayListData::FrameFlags::HadWillChange;
+    return;
+  }
+
   if (IsFrameModified() || HasOverrideDirtyRegion()) {
     // Remove deleted frames from RetainedDisplayListData.
     DebugOnly<bool> removed = data->Remove(this);
@@ -3486,7 +3492,7 @@ nsIFrame::BuildDisplayListForChild(nsDisplayListBuilder*   aBuilder,
        NS_FRAME_TOO_DEEP_IN_FRAME_TREE | NS_FRAME_IS_NONDISPLAY))
     return;
 
-  aBuilder->ClearWillChangeBudget(child);
+  aBuilder->RemoveFromWillChangeBudget(child);
 
   const bool shortcutPossible = aBuilder->IsPaintingToWindow() &&
      aBuilder->BuildCompositorHitTestInfo();
@@ -3563,7 +3569,7 @@ nsIFrame::BuildDisplayListForChild(nsDisplayListBuilder*   aBuilder,
     isPlaceholder = true;
     nsPlaceholderFrame* placeholder = static_cast<nsPlaceholderFrame*>(child);
     child = placeholder->GetOutOfFlowFrame();
-    aBuilder->ClearWillChangeBudget(child);
+    aBuilder->RemoveFromWillChangeBudget(child);
     NS_ASSERTION(child, "No out of flow frame?");
     // If 'child' is a pushed float then it's owned by a block that's not an
     // ancestor of the placeholder, and it will be painted by that block and

layout/painting/RetainedDisplayListBuilder.cpp

@@ -660,6 +660,7 @@ RetainedDisplayListBuilder::MergeDisplayLists(
 
 static void
 TakeAndAddModifiedAndFramesWithPropsFromRootFrame(
+  nsDisplayListBuilder* aBuilder,
   nsTArray<nsIFrame*>* aModifiedFrames,
   nsTArray<nsIFrame*>* aFramesWithProps,
   nsIFrame* aRootFrame)
@@ -683,6 +684,10 @@ TakeAndAddModifiedAndFramesWithPropsFromRootFrame(
     if (flags & RetainedDisplayListData::FrameFlags::HasProps) {
       aFramesWithProps->AppendElement(frame);
     }
+
+    if (flags & RetainedDisplayListData::FrameFlags::HadWillChange) {
+      aBuilder->RemoveFromWillChangeBudget(frame);
+    }
   }
 
   data->Clear();
@@ -748,7 +753,7 @@ SubDocEnumCb(nsIDocument* aDocument, void* aData)
   nsIFrame* rootFrame = GetRootFrameForPainting(data->builder, aDocument);
   if (rootFrame) {
     TakeAndAddModifiedAndFramesWithPropsFromRootFrame(
-      data->modifiedFrames, data->framesWithProps, rootFrame);
+      data->builder, data->modifiedFrames, data->framesWithProps, rootFrame);
 
     nsIDocument* innerDoc = rootFrame->PresShell()->GetDocument();
     if (innerDoc) {
@@ -763,14 +768,13 @@ GetModifiedAndFramesWithProps(nsDisplayListBuilder* aBuilder,
                               nsTArray<nsIFrame*>* aOutModifiedFrames,
                               nsTArray<nsIFrame*>* aOutFramesWithProps)
 {
-  MOZ_ASSERT(aBuilder->RootReferenceFrame());
+  nsIFrame* rootFrame = aBuilder->RootReferenceFrame();
+  MOZ_ASSERT(rootFrame);
 
   TakeAndAddModifiedAndFramesWithPropsFromRootFrame(
-    aOutModifiedFrames, aOutFramesWithProps, aBuilder->RootReferenceFrame());
-
-  nsIDocument* rootdoc =
-    aBuilder->RootReferenceFrame()->PresContext()->Document();
+    aBuilder, aOutModifiedFrames, aOutFramesWithProps, rootFrame);
 
+  nsIDocument* rootdoc = rootFrame->PresContext()->Document();
   if (rootdoc) {
     CbData data = { aBuilder, aOutModifiedFrames, aOutFramesWithProps };
 

layout/painting/nsDisplayList.cpp

@@ -1060,6 +1060,7 @@ nsDisplayListBuilder::EndFrame()
   NS_ASSERTION(!mInInvalidSubtree,
                "Someone forgot to cleanup mInInvalidSubtree!");
   mFrameToAnimatedGeometryRootMap.Clear();
+  mAGRBudgetSet.Clear();
   mActiveScrolledRoots.Clear();
   FreeClipChains();
   FreeTemporaryItems();
@@ -2120,29 +2121,20 @@ nsDisplayListBuilder::AddToWillChangeBudget(nsIFrame* aFrame,
     return true; // Already accounted
   }
 
-  nsPresContext* key = aFrame->PresContext();
-  DocumentWillChangeBudget budget;
-  auto willChangeBudgetEntry = mWillChangeBudget.LookupForAdd(key);
-  if (willChangeBudgetEntry) {
-    // We have an existing entry.
-    budget = willChangeBudgetEntry.Data();
-  } else {
-    budget = DocumentWillChangeBudget();
-    willChangeBudgetEntry.OrInsert([&budget]() { return budget; });
-  }
-
-  nsRect area = aFrame->PresContext()->GetVisibleArea();
+  nsPresContext* presContext = aFrame->PresContext();
+  nsRect area = presContext->GetVisibleArea();
   uint32_t budgetLimit = nsPresContext::AppUnitsToIntCSSPixels(area.width) *
                          nsPresContext::AppUnitsToIntCSSPixels(area.height);
-
   uint32_t cost = GetLayerizationCost(aSize);
+
+  DocumentWillChangeBudget& budget = mWillChangeBudget.GetOrInsert(presContext);
+
   bool onBudget =
     (budget.mBudget + cost) / gWillChangeAreaMultiplier < budgetLimit;
 
   if (onBudget) {
     budget.mBudget += cost;
-    willChangeBudgetEntry.Data() = budget;
-    mWillChangeBudgetSet.Put(aFrame, cost);
+    mWillChangeBudgetSet.Put(aFrame, FrameWillChangeBudget(presContext, cost));
     aFrame->SetMayHaveWillChangeBudget(true);
   }
 
@@ -2179,23 +2171,29 @@ nsDisplayListBuilder::IsInWillChangeBudget(nsIFrame* aFrame,
 }
 
 void
-nsDisplayListBuilder::ClearWillChangeBudget(nsIFrame* aFrame)
+nsDisplayListBuilder::RemoveFromWillChangeBudget(nsIFrame* aFrame)
 {
-  if (!aFrame->MayHaveWillChangeBudget()) {
-    return;
-  }
-  aFrame->SetMayHaveWillChangeBudget(false);
+  FrameWillChangeBudget* frameBudget = mWillChangeBudgetSet.GetValue(aFrame);
 
-  uint32_t cost = 0;
-  if (!mWillChangeBudgetSet.Get(aFrame, &cost)) {
+  if (!frameBudget) {
     return;
   }
+
+  DocumentWillChangeBudget* budget =
+    mWillChangeBudget.GetValue(frameBudget->mPresContext);
+
+  if (budget) {
+    budget->mBudget -= frameBudget->mUsage;
+  }
+
   mWillChangeBudgetSet.Remove(aFrame);
+}
 
-  DocumentWillChangeBudget& budget =
-    mWillChangeBudget.GetOrInsert(aFrame->PresContext());
-  MOZ_ASSERT(budget.mBudget >= cost);
-  budget.mBudget -= cost;
+void
+nsDisplayListBuilder::ClearWillChangeBudget()
+{
+  mWillChangeBudgetSet.Clear();
+  mWillChangeBudget.Clear();
 }
 
 #ifdef MOZ_GFX_OPTIMIZE_MOBILE

layout/painting/nsDisplayList.h

@@ -1815,7 +1815,9 @@ public:
    */
   bool IsInWillChangeBudget(nsIFrame* aFrame, const nsSize& aSize);
 
-  void ClearWillChangeBudget(nsIFrame* aFrame);
+  void RemoveFromWillChangeBudget(nsIFrame* aFrame);
+
+  void ClearWillChangeBudget();
 
   void EnterSVGEffectsContents(nsDisplayList* aHoistedItemsStorage);
   void ExitSVGEffectsContents();
@@ -2040,13 +2042,19 @@ private:
 
   struct FrameWillChangeBudget
   {
-    FrameWillChangeBudget(nsIFrame* aFrame, uint32_t aUsage)
-      : mFrame(aFrame)
+    FrameWillChangeBudget()
+      : mPresContext(nullptr)
+      , mUsage(0)
+    {
+    }
+
+    FrameWillChangeBudget(nsPresContext* aPresContext, uint32_t aUsage)
+      : mPresContext(aPresContext)
       , mUsage(aUsage)
     {
     }
 
-    nsIFrame* mFrame;
+    nsPresContext* mPresContext;
     uint32_t mUsage;
   };
 
@@ -2084,7 +2092,8 @@ private:
 
   // Any frame listed in this set is already counted in the budget
   // and thus is in-budget.
-  nsDataHashtable<nsPtrHashKey<nsIFrame>, uint32_t> mWillChangeBudgetSet;
+  nsDataHashtable<nsPtrHashKey<nsIFrame>, FrameWillChangeBudget>
+    mWillChangeBudgetSet;
 
   // Area of animated geometry root budget already allocated
   uint32_t mUsedAGRBudget;