[DL]: Handle whitespace reattachment for locked elements

ReattachLayoutTree/RebuildLayoutTree can't be called on elements
with the NeedsStyleRecalc/ChildNeedsStyleRecalc flag set, which might
be the case for locked elements because we block style recalc.
Marking of the NeedsReattachLayoutTree/ChildNeedsReattachLayoutTree
is done inside RecalcStyle and whitespace reattachment marking.
The first case is OK because we already block RecalcStyle, this CL is
fixing the second case.

In this CL, we save a per-DisplayLockContext |whitespace_reattach_set_|
that will save locked elements needing whitespace reattachment. After
we finished style recalc on locked elements (after forced update or
when the budget allows), we will mark those elements as needing layout
tree reattachment, just like the global whitespace reattachment in
StyleEngine.

Bug: 882663, 912949
Change-Id: Ib83845601db24342a30d556689c2a08244491a33
Reviewed-on: https://chromium-review.googlesource.com/c/1460581
Commit-Queue: Rakina Zata Amni <rakina@chromium.org>
Reviewed-by: Rune Lillesveen <futhark@chromium.org>
Reviewed-by: vmpstr <vmpstr@chromium.org>
Cr-Commit-Position: refs/heads/master@{#635430}
diff --git a/third_party/blink/renderer/core/css/style_engine.cc b/third_party/blink/renderer/core/css/style_engine.cc
index ed1b1a4..70b6e34 100644
--- a/third_party/blink/renderer/core/css/style_engine.cc
+++ b/third_party/blink/renderer/core/css/style_engine.cc
@@ -50,6 +50,7 @@
 #include "third_party/blink/renderer/core/css/style_sheet_contents.h"
 #include "third_party/blink/renderer/core/dom/element.h"
 #include "third_party/blink/renderer/core/dom/element_traversal.h"
+#include "third_party/blink/renderer/core/dom/flat_tree_traversal.h"
 #include "third_party/blink/renderer/core/dom/layout_tree_builder_traversal.h"
 #include "third_party/blink/renderer/core/dom/node_computed_style.h"
 #include "third_party/blink/renderer/core/dom/processing_instruction.h"
@@ -1573,6 +1574,30 @@
   for (auto element : whitespace_reattach_set_) {
     if (element->NeedsReattachLayoutTree() || !element->GetLayoutObject())
       continue;
+    if (RuntimeEnabledFeatures::DisplayLockingEnabled() &&
+        GetDocument().LockedDisplayLockCount() > 0) {
+      // This element might be located inside a display locked subtree, so we
+      // might mark it for ReattachLayoutTree later on instead.
+      // TODO(crbug.com/924550): Once we figure out a more efficient way to
+      // determine whether we're inside a locked subtree or not, change this.
+      bool is_in_locked_subtree = false;
+      for (const Node& ancestor :
+           FlatTreeTraversal::InclusiveAncestorsOf(*element)) {
+        if (!ancestor.IsElementNode())
+          continue;
+        if (auto* context = ToElement(ancestor).GetDisplayLockContext()) {
+          if (!context->IsLocked())
+            continue;
+          is_in_locked_subtree = true;
+          context->AddToWhitespaceReattachSet(*element);
+          break;
+        }
+      }
+      if (is_in_locked_subtree)
+        continue;
+      DCHECK(!element->NeedsStyleRecalc());
+      DCHECK(!element->ChildNeedsStyleRecalc());
+    }
     if (Node* first_child = LayoutTreeBuilderTraversal::FirstChild(*element))
       first_child->MarkAncestorsWithChildNeedsReattachLayoutTree();
   }
diff --git a/third_party/blink/renderer/core/display_lock/display_lock_context.cc b/third_party/blink/renderer/core/display_lock/display_lock_context.cc
index 3ad4b4c..4124e3adc 100644
--- a/third_party/blink/renderer/core/display_lock/display_lock_context.cc
+++ b/third_party/blink/renderer/core/display_lock/display_lock_context.cc
@@ -78,6 +78,7 @@
   visitor->Trace(acquire_resolver_);
   visitor->Trace(element_);
   visitor->Trace(document_);
+  visitor->Trace(whitespace_reattach_set_);
   ScriptWrappable::Trace(visitor);
   ActiveScriptWrappable::Trace(visitor);
   ContextLifecycleObserver::Trace(visitor);
@@ -294,6 +295,7 @@
     return;
   }
 
+  MarkElementsForWhitespaceReattachment();
   // We must have "contain: style layout" for display locking.
   // Note that we should also have this containment even if we're forcing
   // this update to happen. Otherwise, proceeding with layout may cause
@@ -514,6 +516,22 @@
   return nullptr;
 }
 
+void DisplayLockContext::AddToWhitespaceReattachSet(Element& element) {
+  whitespace_reattach_set_.insert(&element);
+}
+
+void DisplayLockContext::MarkElementsForWhitespaceReattachment() {
+  for (auto element : whitespace_reattach_set_) {
+    if (!element || element->NeedsReattachLayoutTree() ||
+        !element->GetLayoutObject())
+      continue;
+
+    if (Node* first_child = LayoutTreeBuilderTraversal::FirstChild(*element))
+      first_child->MarkAncestorsWithChildNeedsReattachLayoutTree();
+  }
+  whitespace_reattach_set_.clear();
+}
+
 bool DisplayLockContext::MarkAncestorsForStyleRecalcIfNeeded() {
   if (IsElementDirtyForStyleRecalc()) {
     element_->MarkAncestorsWithChildNeedsStyleRecalc();
diff --git a/third_party/blink/renderer/core/display_lock/display_lock_context.h b/third_party/blink/renderer/core/display_lock/display_lock_context.h
index 770ebd2a..4020de1 100644
--- a/third_party/blink/renderer/core/display_lock/display_lock_context.h
+++ b/third_party/blink/renderer/core/display_lock/display_lock_context.h
@@ -151,6 +151,8 @@
   // right document's view.
   void DidMoveToNewDocument(Document& old_document);
 
+  void AddToWhitespaceReattachSet(Element& element);
+
   // LifecycleNotificationObserver overrides.
   void WillStartLifecycleUpdate() override;
   void DidFinishLifecycleUpdate() override;
@@ -190,6 +192,10 @@
   // Initiate an update.
   void StartUpdateIfNeeded();
 
+  // Marks ancestors of elements in |whitespace_reattach_set_| with
+  // ChildNeedsReattachLayoutTree and clears the set.
+  void MarkElementsForWhitespaceReattachment();
+
   // The following functions propagate dirty bits from the locked element up to
   // the ancestors in order to be reached. They return true if the element or
   // its subtree were dirty, and false otherwise.
@@ -252,6 +258,15 @@
   WeakMember<Element> element_;
   WeakMember<Document> document_;
 
+  // See StyleEngine's |whitespace_reattach_set_|.
+  // Set of elements that had at least one rendered children removed
+  // since its last lifecycle update. For such elements that are located
+  // in a locked subtree, we save it here instead of the global set in
+  // StyleEngine because we don't want to accidentally mark elements
+  // in a locked subtree for layout tree reattachment before we did
+  // style recalc on them.
+  HeapHashSet<Member<Element>> whitespace_reattach_set_;
+
   StateChangeHelper state_;
   LayoutRect pending_frame_rect_;
   base::Optional<LayoutRect> locked_frame_rect_;
diff --git a/third_party/blink/renderer/core/display_lock/display_lock_context_test.cc b/third_party/blink/renderer/core/display_lock/display_lock_context_test.cc
index 7b1cd5a..c66ae4d 100644
--- a/third_party/blink/renderer/core/display_lock/display_lock_context_test.cc
+++ b/third_party/blink/renderer/core/display_lock/display_lock_context_test.cc
@@ -279,6 +279,89 @@
   client.Reset();
 }
 
+TEST_F(DisplayLockContextTest, CallUpdateStyleAndLayoutAfterChange) {
+  ResizeAndFocus();
+  SetHtmlInnerHTML(R"HTML(
+    <style>
+    #container {
+      width: 100px;
+      height: 100px;
+      contain: content;
+    }
+    </style>
+    <body><div id="container"><b>t</b>esting</div></body>
+  )HTML");
+  auto* element = GetDocument().getElementById("container");
+  auto* script_state = ToScriptStateForMainWorld(GetDocument().GetFrame());
+  {
+    ScriptState::Scope scope(script_state);
+    element->getDisplayLockForBindings()->acquire(script_state, nullptr);
+  }
+  UpdateAllLifecyclePhasesForTest();
+
+  // Sanity checks to ensure the element is locked.
+  EXPECT_FALSE(element->GetDisplayLockContext()->ShouldStyle());
+  EXPECT_FALSE(element->GetDisplayLockContext()->ShouldLayout());
+  EXPECT_FALSE(element->GetDisplayLockContext()->ShouldPaint());
+  EXPECT_EQ(GetDocument().LockedDisplayLockCount(), 1);
+  EXPECT_EQ(GetDocument().ActivationBlockingDisplayLockCount(), 1);
+
+  EXPECT_FALSE(element->NeedsStyleRecalc());
+  EXPECT_FALSE(element->ChildNeedsStyleRecalc());
+  EXPECT_FALSE(element->NeedsReattachLayoutTree());
+  EXPECT_FALSE(element->ChildNeedsReattachLayoutTree());
+
+  // Testing whitespace reattachment, shouldn't mark for reattachment.
+  element->firstChild()->remove();
+
+  EXPECT_FALSE(element->NeedsStyleRecalc());
+  EXPECT_FALSE(element->ChildNeedsStyleRecalc());
+  EXPECT_FALSE(element->NeedsReattachLayoutTree());
+  EXPECT_FALSE(element->ChildNeedsReattachLayoutTree());
+
+  GetDocument().UpdateStyleAndLayoutIgnorePendingStylesheets();
+
+  EXPECT_FALSE(element->NeedsStyleRecalc());
+  EXPECT_FALSE(element->ChildNeedsStyleRecalc());
+  EXPECT_FALSE(element->NeedsReattachLayoutTree());
+  EXPECT_FALSE(element->ChildNeedsReattachLayoutTree());
+
+  // Testing whitespace reattachment + dirty style.
+  element->SetInnerHTMLFromString("<div>something</div>");
+
+  EXPECT_FALSE(element->NeedsStyleRecalc());
+  EXPECT_TRUE(element->ChildNeedsStyleRecalc());
+  EXPECT_FALSE(element->NeedsReattachLayoutTree());
+  EXPECT_FALSE(element->ChildNeedsReattachLayoutTree());
+
+  GetDocument().UpdateStyleAndLayoutIgnorePendingStylesheets();
+
+  EXPECT_FALSE(element->NeedsStyleRecalc());
+  EXPECT_TRUE(element->ChildNeedsStyleRecalc());
+  EXPECT_FALSE(element->NeedsReattachLayoutTree());
+  EXPECT_FALSE(element->ChildNeedsReattachLayoutTree());
+
+  {
+    ScriptState::Scope scope(script_state);
+    element->getDisplayLockForBindings()->commit(script_state);
+  }
+
+  EXPECT_FALSE(element->NeedsStyleRecalc());
+  EXPECT_TRUE(element->ChildNeedsStyleRecalc());
+  EXPECT_FALSE(element->NeedsReattachLayoutTree());
+  EXPECT_FALSE(element->ChildNeedsReattachLayoutTree());
+
+  // Simulating style recalc happening, will mark for reattachment.
+  element->ClearChildNeedsStyleRecalc();
+  element->firstChild()->ClearNeedsStyleRecalc();
+  element->GetDisplayLockContext()->DidStyle();
+
+  EXPECT_FALSE(element->NeedsStyleRecalc());
+  EXPECT_FALSE(element->ChildNeedsStyleRecalc());
+  EXPECT_FALSE(element->NeedsReattachLayoutTree());
+  EXPECT_TRUE(element->ChildNeedsReattachLayoutTree());
+}
+
 TEST_F(DisplayLockContextTest, LockedElementAndDescendantsAreNotFocusable) {
   ResizeAndFocus();
   SetHtmlInnerHTML(R"HTML(
diff --git a/third_party/blink/renderer/core/dom/element.cc b/third_party/blink/renderer/core/dom/element.cc
index 8c090d9b..dd376d1 100644
--- a/third_party/blink/renderer/core/dom/element.cc
+++ b/third_party/blink/renderer/core/dom/element.cc
@@ -2442,6 +2442,7 @@
 void Element::RebuildLayoutTree(WhitespaceAttacher& whitespace_attacher) {
   DCHECK(InActiveDocument());
   DCHECK(parentNode());
+  DCHECK(!StyleRecalcBlockedByDisplayLock());
 
   if (NeedsReattachLayoutTree()) {
     AttachContext reattach_context;
diff --git a/third_party/blink/renderer/core/dom/element.h b/third_party/blink/renderer/core/dom/element.h
index b3861cd..0898c83 100644
--- a/third_party/blink/renderer/core/dom/element.h
+++ b/third_party/blink/renderer/core/dom/element.h
@@ -897,6 +897,7 @@
   DisplayLockContext* GetDisplayLockContext() const;
 
   bool StyleRecalcBlockedByDisplayLock() const;
+
   void ActivateDisplayLockIfNeeded();
 
  protected: