Bug 1505489 - Keep track of elements with part attributes in a shadow tree. r=bzbarsky

I need this to make style invalidation work with Shadow Parts in a way that's
fast. If something in the ancestor shadow root or any of its ancestor changes,
that makes a ::part() selector change, I don't want to walk the whole shadow
tree over and over in order to find the parts that I need to restyle.

Unfortunately we cannot just use the mutation observer setup from ShadowRoot
since, unlike for slotted elements, there's no restriction of where a part can
appear in the shadow tree.

That means that the regular nsIMutationObserver notifications don't quite cut
it, since we'd get notified only of subtree roots and we'd need to tree-walk
around in order to figure out if we have any new part.

I thought that I was going to be able to share more code with other bits if I
moved them away from nsIMutationObserver, thus bug 1554498, but in the end it
was not the case, so here's the "without bug 1554498" version of the patch.

The patch on top of that bug (that as I mention in the commit message I'm
ambivalent about whether we should land or not) should be pretty similar either
way.

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

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

dom/base/Element.cpp

@@ -1717,10 +1717,16 @@ nsresult Element::BindToTree(BindContext& aContext, nsINode& aParent) {
     nsNodeUtils::NativeAnonymousChildListChange(this, false);
   }
 
-  // Ensure we only add to the table once, in the case we move the ShadowRoot
-  // around.
-  if (HasID() && aContext.SubtreeRootChanges()) {
-    AddToIdTable(DoGetID());
+  // Ensure we only run this once, in the case we move the ShadowRoot around.
+  if (aContext.SubtreeRootChanges()) {
+    if (HasPartAttribute()) {
+      if (ShadowRoot* shadow = GetContainingShadow()) {
+        shadow->PartAdded(*this);
+      }
+    }
+    if (HasID()) {
+      AddToIdTable(DoGetID());
+    }
   }
 
   if (MayHaveStyle() && !IsXULElement()) {
@@ -1798,26 +1804,29 @@ RemoveFromBindingManagerRunnable::Run() {
   return NS_OK;
 }
 
-static bool ShouldRemoveFromIdTableOnUnbind(const Element& aElement,
-                                            bool aNullParent) {
-  if (aElement.IsInUncomposedDoc()) {
-    return true;
-  }
-
-  if (!aElement.IsInShadowTree()) {
-    return false;
-  }
-
-  return aNullParent || !aElement.GetParent()->IsInShadowTree();
+bool WillDetachFromShadowOnUnbind(const Element& aElement, bool aNullParent) {
+  // If our parent still is in a shadow tree by now, and we're not removing
+  // ourselves from it, then we're still going to be in a shadow tree after
+  // this.
+  return aElement.IsInShadowTree() &&
+         (aNullParent || !aElement.GetParent()->IsInShadowTree());
 }
 
 void Element::UnbindFromTree(bool aNullParent) {
+  const bool detachingFromShadow =
+      WillDetachFromShadowOnUnbind(*this, aNullParent);
   // Make sure to only remove from the ID table if our subtree root is actually
   // changing.
-  if (ShouldRemoveFromIdTableOnUnbind(*this, aNullParent)) {
+  if (IsInUncomposedDoc() || detachingFromShadow) {
     RemoveFromIdTable();
   }
 
+  if (detachingFromShadow && HasPartAttribute()) {
+    if (ShadowRoot* shadow = GetContainingShadow()) {
+      shadow->PartRemoved(*this);
+    }
+  }
+
   // Make sure to unbind this node before doing the kids
   Document* document = GetComposedDoc();
 
@@ -2605,6 +2614,28 @@ nsresult Element::BeforeSetAttr(int32_t aNamespaceID, nsAtom* aName,
   return NS_OK;
 }
 
+nsresult Element::AfterSetAttr(int32_t aNamespaceID, nsAtom* aName,
+                               const nsAttrValue* aValue,
+                               const nsAttrValue* aOldValue,
+                               nsIPrincipal* aMaybeScriptedPrincipal,
+                               bool aNotify) {
+  if (aNamespaceID == kNameSpaceID_None && aName == nsGkAtoms::part) {
+    bool isPart = !!aValue;
+    if (HasPartAttribute() != isPart) {
+      SetHasPartAttribute(isPart);
+      if (ShadowRoot* shadow = GetContainingShadow()) {
+        if (isPart) {
+          shadow->PartAdded(*this);
+        } else {
+          shadow->PartRemoved(*this);
+        }
+      }
+    }
+    MOZ_ASSERT(HasPartAttribute() == isPart);
+  }
+  return NS_OK;
+}
+
 void Element::PreIdMaybeChange(int32_t aNamespaceID, nsAtom* aName,
                                const nsAttrValueOrString* aValue) {
   if (aNamespaceID != kNameSpaceID_None || aName != nsGkAtoms::id) {

dom/base/Element.h

@@ -1778,15 +1778,11 @@ class Element : public FragmentOrElement {
    *        principal is directly responsible for the attribute change.
    * @param aNotify Whether we plan to notify document observers.
    */
-  // Note that this is inlined so that when subclasses call it it gets
-  // inlined.  Those calls don't go through a vtable.
   virtual nsresult AfterSetAttr(int32_t aNamespaceID, nsAtom* aName,
                                 const nsAttrValue* aValue,
                                 const nsAttrValue* aOldValue,
                                 nsIPrincipal* aMaybeScriptedPrincipal,
-                                bool aNotify) {
-    return NS_OK;
-  }
+                                bool aNotify);
 
   /**
    * This function shall be called just before the id attribute changes. It will

dom/base/ShadowRoot.cpp

@@ -183,6 +183,18 @@ void ShadowRoot::InvalidateStyleAndLayoutOnSubtree(Element* aElement) {
   presShell->DestroyFramesForAndRestyle(aElement);
 }
 
+void ShadowRoot::PartAdded(const Element& aPart) {
+  MOZ_ASSERT(aPart.HasPartAttribute());
+  MOZ_ASSERT(!mParts.Contains(&aPart));
+  mParts.AppendElement(&aPart);
+}
+
+void ShadowRoot::PartRemoved(const Element& aPart) {
+  MOZ_ASSERT(mParts.Contains(&aPart));
+  mParts.RemoveElement(&aPart);
+  MOZ_ASSERT(!mParts.Contains(&aPart));
+}
+
 void ShadowRoot::AddSlot(HTMLSlotElement* aSlot) {
   MOZ_ASSERT(aSlot);
 

dom/base/ShadowRoot.h

@@ -154,6 +154,13 @@ class ShadowRoot final : public DocumentFragment,
   void RemoveSlot(HTMLSlotElement* aSlot);
   bool HasSlots() const { return !mSlotMap.IsEmpty(); };
 
+  void PartAdded(const Element&);
+  void PartRemoved(const Element&);
+
+  const nsTArray<const Element*>& Parts() const {
+    return mParts;
+  }
+
   const RawServoAuthorStyles* GetServoStyles() const {
     return mServoStyles.get();
   }
@@ -259,6 +266,10 @@ class ShadowRoot final : public DocumentFragment,
   // are in the shadow tree and should be kept alive by its parent.
   nsClassHashtable<nsStringHashKey, SlotArray> mSlotMap;
 
+  // Unordered array of all elements that have a part attribute in this shadow
+  // tree.
+  nsTArray<const Element*> mParts;
+
   bool mIsUAWidget;
 
   nsresult Clone(dom::NodeInfo*, nsINode** aResult) const override;

dom/base/nsINode.h

@@ -1405,6 +1405,8 @@ class nsINode : public mozilla::dom::EventTarget {
     ElementMayHaveStyle,
     // Set if the element has a name attribute set.
     ElementHasName,
+    // Set if the element has a part attribute set.
+    ElementHasPart,
     // Set if the element might have a contenteditable attribute set.
     ElementMayHaveContentEditableAttr,
     // Set if the node is the common ancestor of the start/end nodes of a Range
@@ -1498,6 +1500,7 @@ class nsINode : public mozilla::dom::EventTarget {
   void SetMayHaveClass() { SetBoolFlag(ElementMayHaveClass); }
   bool MayHaveStyle() const { return GetBoolFlag(ElementMayHaveStyle); }
   bool HasName() const { return GetBoolFlag(ElementHasName); }
+  bool HasPartAttribute() const { return GetBoolFlag(ElementHasPart); }
   bool MayHaveContentEditableAttr() const {
     return GetBoolFlag(ElementMayHaveContentEditableAttr);
   }
@@ -1605,6 +1608,7 @@ class nsINode : public mozilla::dom::EventTarget {
   void SetMayHaveStyle() { SetBoolFlag(ElementMayHaveStyle); }
   void SetHasName() { SetBoolFlag(ElementHasName); }
   void ClearHasName() { ClearBoolFlag(ElementHasName); }
+  void SetHasPartAttribute(bool aPart) { SetBoolFlag(ElementHasPart, aPart); }
   void SetMayHaveContentEditableAttr() {
     SetBoolFlag(ElementMayHaveContentEditableAttr);
   }