From eea863c791a454296b5cda4e99c7caf373d66661 Mon Sep 17 00:00:00 2001
From: Brad Werth <bwerth@mozilla.com>
Date: Wed, 1 Aug 2018 15:33:06 -0700
Subject: [PATCH] Bug 1479508 Part 2: Change GetElementsWithGrid to use a more
 conservative traversal that skips subtrees without frames. r=bz

MozReview-Commit-ID: 1hTQD8LoLQw

--HG--
extra : rebase_source : e8651e1a78818a25c98cc34db2ac217756596666
---
 dom/base/Element.cpp | 43 +++++++++++++++++++------------------------
 dom/base/Element.h   | 14 +-------------
 2 files changed, 20 insertions(+), 37 deletions(-)

diff --git a/dom/base/Element.cpp b/dom/base/Element.cpp
index 0b24fbf503a5..66481d0224b7 100644
--- a/dom/base/Element.cpp
+++ b/dom/base/Element.cpp
@@ -1555,36 +1555,31 @@ Element::GetElementsByClassName(const nsAString& aClassNames)
 void
 Element::GetElementsWithGrid(nsTArray<RefPtr<Element>>& aElements)
 {
-  // This helper function is passed to GetElementsByMatching()
-  // to identify elements with styling which will cause them to
-  // generate a nsGridContainerFrame during layout.
-  auto IsDisplayGrid = [](Element* aElement) -> bool
-  {
-    RefPtr<ComputedStyle> computedStyle =
-      nsComputedDOMStyle::GetComputedStyle(aElement, nullptr);
-    if (computedStyle) {
-      const nsStyleDisplay* display = computedStyle->StyleDisplay();
-      return (display->mDisplay == StyleDisplay::Grid ||
-              display->mDisplay == StyleDisplay::InlineGrid);
-    }
-    return false;
-  };
+  nsINode* cur = this;
+  while (cur) {
+    if (cur->IsElement()) {
+      Element* elem = cur->AsElement();
 
-  GetElementsByMatching(IsDisplayGrid, aElements);
-}
+      if (elem->GetPrimaryFrame()) {
+        // See if this has a GridContainerFrame. Use the same method that
+        // nsGridContainerFrame uses, which deals with some edge cases.
+        if (nsGridContainerFrame::GetGridContainerFrame(elem->GetPrimaryFrame())) {
+          aElements.AppendElement(elem);
+        }
 
-void
-Element::GetElementsByMatching(nsElementMatchFunc aFunc,
-                               nsTArray<RefPtr<Element>>& aElements)
-{
-  for (nsINode* cur = this; cur; cur = cur->GetNextNode(this)) {
-    if (cur->IsElement() && aFunc(cur->AsElement())) {
-      aElements.AppendElement(cur->AsElement());
+        // This element has a frame, so allow the traversal to go through
+        // the children.
+        cur = cur->GetNextNode(this);
+        continue;
+      }
     }
+
+    // Either this isn't an element, or it has no frame. Continue with the
+    // traversal but ignore all the children.
+    cur = cur->GetNextNonChildNode(this);
   }
 }
 
-
 /**
  * Returns the count of descendants (inclusive of aContent) in
  * the uncomposed document that are explicitly set as editable.
diff --git a/dom/base/Element.h b/dom/base/Element.h
index 52630edd5770..1a9175497de0 100644
--- a/dom/base/Element.h
+++ b/dom/base/Element.h
@@ -1167,24 +1167,12 @@ public:
 
   /**
    * Return an array of all elements in the subtree rooted at this
-   * element that are styled as grid containers. This includes
-   * elements that don't actually generate any frames (by virtue of
-   * being in a 'display:none' subtree), but this does not include
+   * element that have grid container frames. This does not include
    * pseudo-elements.
    */
   void GetElementsWithGrid(nsTArray<RefPtr<Element>>& aElements);
 
 private:
-  /**
-   * Define a general matching function that can be passed to
-   * GetElementsByMatching(). Each Element being considered is
-   * passed in.
-   */
-  typedef bool (*nsElementMatchFunc)(Element* aElement);
-
-  void GetElementsByMatching(nsElementMatchFunc aFunc,
-                             nsTArray<RefPtr<Element>>& aElements);
-
   /**
    * Implement the algorithm specified at
    * https://dom.spec.whatwg.org/#insert-adjacent for both