--- a/layout/base/nsLayoutUtils.cpp
+++ b/layout/base/nsLayoutUtils.cpp
@@ -3695,21 +3695,25 @@ nsLayoutUtils::PaintFrame(gfxContext* aR
             }
           }
         }
         dlStart = TimeStamp::Now();
       }
 
       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;
 
         if (afterMergeChecker) {
           DisplayListChecker nonRetainedChecker(&list, "NR");
           std::stringstream ss;
           ss << "**** Differences between retained-after-merged (AM) and "

--- a/layout/generic/nsFrame.cpp
+++ b/layout/generic/nsFrame.cpp
@@ -981,16 +981,22 @@ nsIFrame::RemoveDisplayItemDataForDeleti
     return;
   }
 
   nsIFrame* rootFrame = PresShell()->GetRootFrame();
   MOZ_ASSERT(rootFrame);
 
   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);
     MOZ_ASSERT(removed,
                "Frame had flags set, but it was not found in DisplayListData!");
   }
 }
 
@@ -3481,17 +3487,17 @@ nsIFrame::BuildDisplayListForChild(nsDis
     }
   }
 
   nsIFrame* child = aChild;
   if (child->HasAnyStateBits(
        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();
 
   const bool doingShortcut = shortcutPossible &&
     (child->GetStateBits() & NS_FRAME_SIMPLE_DISPLAYLIST) &&
     // Animations may change the value of |HasOpacity()|.
     !(child->GetContent() &&
@@ -3558,17 +3564,17 @@ nsIFrame::BuildDisplayListForChild(nsDis
   }
 
   nsDisplayListBuilder::OutOfFlowDisplayData* savedOutOfFlowData = nullptr;
   bool isPlaceholder = false;
   if (child->IsPlaceholderFrame()) {
     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
     // should not be painted through the placeholder.
     if (!child || nsLayoutUtils::IsPopup(child) ||
         (child->GetStateBits() & NS_FRAME_IS_PUSHED_FLOAT))
       return;
     MOZ_ASSERT(child->GetStateBits() & NS_FRAME_OUT_OF_FLOW);

--- a/layout/painting/RetainedDisplayListBuilder.cpp
+++ b/layout/painting/RetainedDisplayListBuilder.cpp
@@ -655,16 +655,17 @@ RetainedDisplayListBuilder::MergeDisplay
 
   *aOutList = merge.Finalize();
   aOutContainerASR = merge.mContainerASR;
   return merge.mResultIsModified;
 }
 
 static void
 TakeAndAddModifiedAndFramesWithPropsFromRootFrame(
+  nsDisplayListBuilder* aBuilder,
   nsTArray<nsIFrame*>* aModifiedFrames,
   nsTArray<nsIFrame*>* aFramesWithProps,
   nsIFrame* aRootFrame)
 {
   MOZ_ASSERT(aRootFrame);
 
   RetainedDisplayListData* data = GetRetainedDisplayListData(aRootFrame);
 
@@ -678,16 +679,20 @@ TakeAndAddModifiedAndFramesWithPropsFrom
 
     if (flags & RetainedDisplayListData::FrameFlags::Modified) {
       aModifiedFrames->AppendElement(frame);
     }
 
     if (flags & RetainedDisplayListData::FrameFlags::HasProps) {
       aFramesWithProps->AppendElement(frame);
     }
+
+    if (flags & RetainedDisplayListData::FrameFlags::HadWillChange) {
+      aBuilder->RemoveFromWillChangeBudget(frame);
+    }
   }
 
   data->Clear();
 }
 
 struct CbData
 {
   nsDisplayListBuilder* builder;
@@ -743,39 +748,38 @@ SubDocEnumCb(nsIDocument* aDocument, voi
   MOZ_ASSERT(aDocument);
   MOZ_ASSERT(aData);
 
   CbData* data = static_cast<CbData*>(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) {
       innerDoc->EnumerateSubDocuments(SubDocEnumCb, aData);
     }
   }
   return true;
 }
 
 static void
 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());
+    aBuilder, aOutModifiedFrames, aOutFramesWithProps, rootFrame);
 
-  nsIDocument* rootdoc =
-    aBuilder->RootReferenceFrame()->PresContext()->Document();
-
+  nsIDocument* rootdoc = rootFrame->PresContext()->Document();
   if (rootdoc) {
     CbData data = { aBuilder, aOutModifiedFrames, aOutFramesWithProps };
 
     rootdoc->EnumerateSubDocuments(SubDocEnumCb, &data);
   }
 }
 
 // ComputeRebuildRegion  debugging

--- a/layout/painting/nsDisplayList.cpp
+++ b/layout/painting/nsDisplayList.cpp
@@ -1055,16 +1055,17 @@ nsDisplayListBuilder::BeginFrame()
 }
 
 void
 nsDisplayListBuilder::EndFrame()
 {
   NS_ASSERTION(!mInInvalidSubtree,
                "Someone forgot to cleanup mInInvalidSubtree!");
   mFrameToAnimatedGeometryRootMap.Clear();
+  mAGRBudgetSet.Clear();
   mActiveScrolledRoots.Clear();
   FreeClipChains();
   FreeTemporaryItems();
   nsCSSRendering::EndFrameTreesLocked();
 
   MOZ_ASSERT(!mCompositorHitTestInfo);
 }
 
@@ -2115,39 +2116,30 @@ GetLayerizationCost(const nsSize& aSize)
 bool
 nsDisplayListBuilder::AddToWillChangeBudget(nsIFrame* aFrame,
                                             const nsSize& aSize)
 {
   if (mWillChangeBudgetSet.Get(aFrame, nullptr)) {
     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);
   }
 
   return onBudget;
 }
 
 bool
 nsDisplayListBuilder::IsInWillChangeBudget(nsIFrame* aFrame,
@@ -2174,33 +2166,39 @@ nsDisplayListBuilder::IsInWillChangeBudg
       false,
       params,
       ArrayLength(params));
   }
   return onBudget;
 }
 
 void
-nsDisplayListBuilder::ClearWillChangeBudget(nsIFrame* aFrame)
-{
-  if (!aFrame->MayHaveWillChangeBudget()) {
+nsDisplayListBuilder::RemoveFromWillChangeBudget(nsIFrame* aFrame)
+{
+  FrameWillChangeBudget* frameBudget = mWillChangeBudgetSet.GetValue(aFrame);
+
+  if (!frameBudget) {
     return;
   }
-  aFrame->SetMayHaveWillChangeBudget(false);
-
-  uint32_t cost = 0;
-  if (!mWillChangeBudgetSet.Get(aFrame, &cost)) {
-    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
 const float gAGRBudgetAreaMultiplier = 0.3;
 #else
 const float gAGRBudgetAreaMultiplier = 3.0;
 #endif
 

--- a/layout/painting/nsDisplayList.h
+++ b/layout/painting/nsDisplayList.h
@@ -1810,17 +1810,19 @@ public:
   /**
    * This will add the current frame to the will-change budget the first
    * time it is seen. On subsequent calls this will return the same
    * answer. This effectively implements a first-come, first-served
    * allocation of the will-change budget.
    */
   bool IsInWillChangeBudget(nsIFrame* aFrame, const nsSize& aSize);
 
-  void ClearWillChangeBudget(nsIFrame* aFrame);
+  void RemoveFromWillChangeBudget(nsIFrame* aFrame);
+
+  void ClearWillChangeBudget();
 
   void EnterSVGEffectsContents(nsDisplayList* aHoistedItemsStorage);
   void ExitSVGEffectsContents();
 
   /**
    * Note: if changing the conditions under which scroll info layers
    * are created, make a corresponding change to
    * ScrollFrameWillBuildScrollInfoLayer() in nsSliderFrame.cpp.
@@ -2035,23 +2037,29 @@ private:
     {
     }
 
     uint32_t mBudget;
   };
 
   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;
   };
 
   nsIFrame* const mReferenceFrame;
   nsIFrame* mIgnoreScrollFrame;
   nsDisplayCompositorHitTestInfo* mCompositorHitTestInfo;
 
   nsPresArena mPool;
@@ -2079,17 +2087,18 @@ private:
   RefPtr<AnimatedGeometryRoot> mCurrentAGR;
 
   // will-change budget tracker
   nsDataHashtable<nsPtrHashKey<nsPresContext>, DocumentWillChangeBudget>
     mWillChangeBudget;
 
   // 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;
   // Set of frames already counted in budget
   nsTHashtable<nsPtrHashKey<nsIFrame>> mAGRBudgetSet;
 
   nsTArray<nsIFrame*> mModifiedFramesDuringBuilding;