From 740aae06bad9e5ff864c914117cab1e74a727614 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fernando=20Jim=C3=A9nez=20Moreno?= <ferjmoreno@gmail.com>
Date: Tue, 5 Mar 2019 10:23:18 +0100
Subject: [PATCH] Register named elements in either the document or shadow tree

---
 components/script/dom/document.rs             | 51 +++-------------
 components/script/dom/documentfragment.rs     | 26 ++++----
 components/script/dom/documentorshadowroot.rs | 60 ++++++++++++++++++-
 components/script/dom/element.rs              | 28 +++++++--
 components/script/dom/shadowroot.rs           | 22 +++++++
 .../tests/mozilla/partial_shadow_dom.html     |  5 +-
 6 files changed, 131 insertions(+), 61 deletions(-)

diff --git a/components/script/dom/document.rs b/components/script/dom/document.rs
index 4ec2e67521c..998055d3547 100644
--- a/components/script/dom/document.rs
+++ b/components/script/dom/document.rs
@@ -74,7 +74,6 @@ use crate::dom::keyboardevent::KeyboardEvent;
 use crate::dom::location::Location;
 use crate::dom::messageevent::MessageEvent;
 use crate::dom::mouseevent::MouseEvent;
-use crate::dom::node::VecPreOrderInsertionHelper;
 use crate::dom::node::{self, document_from_node, window_from_node, CloneChildrenFlag};
 use crate::dom::node::{LayoutNodeHelpers, Node, NodeDamage, NodeFlags, ShadowIncluding};
 use crate::dom::nodeiterator::NodeIterator;
@@ -682,55 +681,23 @@ impl Document {
 
     /// Remove any existing association between the provided id and any elements in this document.
     pub fn unregister_named_element(&self, to_unregister: &Element, id: Atom) {
-        debug!(
-            "Removing named element from document {:p}: {:p} id={}",
-            self, to_unregister, id
-        );
-        // Limit the scope of the borrow because id_map might be borrowed again by
-        // GetElementById through the following sequence of calls
-        // reset_form_owner_for_listeners -> reset_form_owner -> GetElementById
-        {
-            let mut id_map = self.id_map.borrow_mut();
-            let is_empty = match id_map.get_mut(&id) {
-                None => false,
-                Some(elements) => {
-                    let position = elements
-                        .iter()
-                        .position(|element| &**element == to_unregister)
-                        .expect("This element should be in registered.");
-                    elements.remove(position);
-                    elements.is_empty()
-                },
-            };
-            if is_empty {
-                id_map.remove(&id);
-            }
-        }
+        self.document_or_shadow_root
+            .unregister_named_element(&self.id_map, to_unregister, &id);
         self.reset_form_owner_for_listeners(&id);
     }
 
     /// Associate an element present in this document with the provided id.
     pub fn register_named_element(&self, element: &Element, id: Atom) {
-        debug!(
-            "Adding named element to document {:p}: {:p} id={}",
-            self, element, id
-        );
-        assert!(element.upcast::<Node>().is_connected());
-        assert!(!id.is_empty());
-
         let root = self.GetDocumentElement().expect(
             "The element is in the document, so there must be a document \
              element.",
         );
-
-        // Limit the scope of the borrow because id_map might be borrowed again by
-        // GetElementById through the following sequence of calls
-        // reset_form_owner_for_listeners -> reset_form_owner -> GetElementById
-        {
-            let mut id_map = self.id_map.borrow_mut();
-            let elements = id_map.entry(id.clone()).or_insert(Vec::new());
-            elements.insert_pre_order(element, root.upcast::<Node>());
-        }
+        self.document_or_shadow_root.register_named_element(
+            &self.id_map,
+            element,
+            &id,
+            DomRoot::from_ref(root.upcast::<Node>()),
+        );
         self.reset_form_owner_for_listeners(&id);
     }
 
@@ -2619,11 +2586,11 @@ impl Document {
             url: DomRefCell::new(url),
             // https://dom.spec.whatwg.org/#concept-document-quirks
             quirks_mode: Cell::new(QuirksMode::NoQuirks),
+            id_map: DomRefCell::new(HashMap::new()),
             // https://dom.spec.whatwg.org/#concept-document-encoding
             encoding: Cell::new(encoding),
             is_html_document: is_html_document == IsHTMLDocument::HTMLDocument,
             activity: Cell::new(activity),
-            id_map: DomRefCell::new(HashMap::new()),
             tag_map: DomRefCell::new(HashMap::new()),
             tagns_map: DomRefCell::new(HashMap::new()),
             classes_map: DomRefCell::new(HashMap::new()),
diff --git a/components/script/dom/documentfragment.rs b/components/script/dom/documentfragment.rs
index 315e04f830d..1561b0d007e 100644
--- a/components/script/dom/documentfragment.rs
+++ b/components/script/dom/documentfragment.rs
@@ -2,27 +2,31 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 
+use crate::dom::bindings::cell::DomRefCell;
 use crate::dom::bindings::codegen::Bindings::DocumentFragmentBinding;
 use crate::dom::bindings::codegen::Bindings::DocumentFragmentBinding::DocumentFragmentMethods;
 use crate::dom::bindings::codegen::Bindings::WindowBinding::WindowMethods;
 use crate::dom::bindings::codegen::UnionTypes::NodeOrString;
 use crate::dom::bindings::error::{ErrorResult, Fallible};
 use crate::dom::bindings::inheritance::Castable;
-use crate::dom::bindings::root::DomRoot;
+use crate::dom::bindings::root::{Dom, DomRoot};
 use crate::dom::bindings::str::DOMString;
 use crate::dom::document::Document;
 use crate::dom::element::Element;
 use crate::dom::htmlcollection::HTMLCollection;
-use crate::dom::node::{window_from_node, Node, ShadowIncluding};
+use crate::dom::node::{window_from_node, Node};
 use crate::dom::nodelist::NodeList;
 use crate::dom::window::Window;
 use dom_struct::dom_struct;
 use servo_atoms::Atom;
+use std::collections::HashMap;
 
 // https://dom.spec.whatwg.org/#documentfragment
 #[dom_struct]
 pub struct DocumentFragment {
     node: Node,
+    /// Caches for the getElement methods
+    id_map: DomRefCell<HashMap<Atom, Vec<Dom<Element>>>>,
 }
 
 impl DocumentFragment {
@@ -30,6 +34,7 @@ impl DocumentFragment {
     pub fn new_inherited(document: &Document) -> DocumentFragment {
         DocumentFragment {
             node: Node::new_inherited(document),
+            id_map: DomRefCell::new(HashMap::new()),
         }
     }
 
@@ -46,6 +51,10 @@ impl DocumentFragment {
 
         Ok(DocumentFragment::new(&document))
     }
+
+    pub fn id_map(&self) -> &DomRefCell<HashMap<Atom, Vec<Dom<Element>>>> {
+        &self.id_map
+    }
 }
 
 impl DocumentFragmentMethods for DocumentFragment {
@@ -57,16 +66,11 @@ impl DocumentFragmentMethods for DocumentFragment {
 
     // https://dom.spec.whatwg.org/#dom-nonelementparentnode-getelementbyid
     fn GetElementById(&self, id: DOMString) -> Option<DomRoot<Element>> {
-        let node = self.upcast::<Node>();
         let id = Atom::from(id);
-        node.traverse_preorder(ShadowIncluding::No)
-            .filter_map(DomRoot::downcast::<Element>)
-            .find(
-                |descendant| match descendant.get_attribute(&ns!(), &local_name!("id")) {
-                    None => false,
-                    Some(attr) => *attr.value().as_atom() == id,
-                },
-            )
+        self.id_map
+            .borrow()
+            .get(&id)
+            .map(|ref elements| DomRoot::from_ref(&*(*elements)[0]))
     }
 
     // https://dom.spec.whatwg.org/#dom-parentnode-firstelementchild
diff --git a/components/script/dom/documentorshadowroot.rs b/components/script/dom/documentorshadowroot.rs
index 0165c0b8384..5c8428242ee 100644
--- a/components/script/dom/documentorshadowroot.rs
+++ b/components/script/dom/documentorshadowroot.rs
@@ -2,6 +2,7 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 
+use crate::dom::bindings::cell::DomRefCell;
 use crate::dom::bindings::codegen::Bindings::NodeBinding::NodeBinding::NodeMethods;
 use crate::dom::bindings::inheritance::Castable;
 use crate::dom::bindings::num::Finite;
@@ -9,13 +10,15 @@ use crate::dom::bindings::root::{Dom, DomRoot};
 use crate::dom::element::Element;
 use crate::dom::htmlelement::HTMLElement;
 use crate::dom::htmlmetaelement::HTMLMetaElement;
-use crate::dom::node;
+use crate::dom::node::{self, Node, VecPreOrderInsertionHelper};
 use crate::dom::window::Window;
 use euclid::Point2D;
 use js::jsapi::JS_GetRuntime;
 use script_layout_interface::message::{NodesFromPointQueryType, QueryMsg};
 use script_traits::UntrustedNodeAddress;
 use servo_arc::Arc;
+use servo_atoms::Atom;
+use std::collections::HashMap;
 use std::fmt;
 use style::context::QuirksMode;
 use style::invalidation::media_queries::{MediaListKey, ToMediaListKey};
@@ -262,4 +265,59 @@ impl DocumentOrShadowRoot {
             },
         }
     }
+
+    /// Remove any existing association between the provided id and any elements in this document.
+    pub fn unregister_named_element(
+        &self,
+        id_map: &DomRefCell<HashMap<Atom, Vec<Dom<Element>>>>,
+        to_unregister: &Element,
+        id: &Atom,
+    ) {
+        debug!(
+            "Removing named element {:p}: {:p} id={}",
+            self, to_unregister, id
+        );
+        // Limit the scope of the borrow because id_map might be borrowed again by
+        // GetElementById through the following sequence of calls
+        // reset_form_owner_for_listeners -> reset_form_owner -> GetElementById
+        {
+            let mut id_map = id_map.borrow_mut();
+            let is_empty = match id_map.get_mut(&id) {
+                None => false,
+                Some(elements) => {
+                    let position = elements
+                        .iter()
+                        .position(|element| &**element == to_unregister)
+                        .expect("This element should be in registered.");
+                    elements.remove(position);
+                    elements.is_empty()
+                },
+            };
+            if is_empty {
+                id_map.remove(&id);
+            }
+        }
+    }
+
+    /// Associate an element present in this document with the provided id.
+    pub fn register_named_element(
+        &self,
+        id_map: &DomRefCell<HashMap<Atom, Vec<Dom<Element>>>>,
+        element: &Element,
+        id: &Atom,
+        root: DomRoot<Node>,
+    ) {
+        debug!("Adding named element {:p}: {:p} id={}", self, element, id);
+        assert!(element.upcast::<Node>().is_connected());
+        assert!(!id.is_empty());
+
+        // Limit the scope of the borrow because id_map might be borrowed again by
+        // GetElementById through the following sequence of calls
+        // reset_form_owner_for_listeners -> reset_form_owner -> GetElementById
+        {
+            let mut id_map = id_map.borrow_mut();
+            let elements = id_map.entry(id.clone()).or_insert(Vec::new());
+            elements.insert_pre_order(element, &root);
+        }
+    }
 }
diff --git a/components/script/dom/element.rs b/components/script/dom/element.rs
index 9d92f196a77..610df336541 100644
--- a/components/script/dom/element.rs
+++ b/components/script/dom/element.rs
@@ -2741,21 +2741,34 @@ impl VirtualMethods for Element {
                         None
                     }
                 });
+                let owner_shadow_root = self.upcast::<Node>().owner_shadow_root();
                 if node.is_connected() {
                     let value = attr.value().as_atom().clone();
                     match mutation {
                         AttributeMutation::Set(old_value) => {
                             if let Some(old_value) = old_value {
                                 let old_value = old_value.as_atom().clone();
-                                doc.unregister_named_element(self, old_value);
+                                if let Some(ref shadow_root) = owner_shadow_root {
+                                    shadow_root.unregister_named_element(self, old_value);
+                                } else {
+                                    doc.unregister_named_element(self, old_value);
+                                }
                             }
                             if value != atom!("") {
-                                doc.register_named_element(self, value);
+                                if let Some(ref shadow_root) = owner_shadow_root {
+                                    shadow_root.register_named_element(self, value);
+                                } else {
+                                    doc.register_named_element(self, value);
+                                }
                             }
                         },
                         AttributeMutation::Removed => {
                             if value != atom!("") {
-                                doc.unregister_named_element(self, value);
+                                if let Some(ref shadow_root) = owner_shadow_root {
+                                    shadow_root.unregister_named_element(self, value);
+                                } else {
+                                    doc.unregister_named_element(self, value);
+                                }
                             }
                         },
                     }
@@ -2800,8 +2813,6 @@ impl VirtualMethods for Element {
             return;
         }
 
-        let doc = document_from_node(self);
-
         if let Some(shadow_root) = self.upcast::<Node>().owner_shadow_root() {
             let shadow_root = shadow_root.upcast::<Node>();
             shadow_root.set_flag(NodeFlags::IS_CONNECTED, tree_connected);
@@ -2811,8 +2822,13 @@ impl VirtualMethods for Element {
             }
         }
 
+        let doc = document_from_node(self);
         if let Some(ref value) = *self.id_attribute.borrow() {
-            doc.register_named_element(self, value.clone());
+            if let Some(shadow_root) = self.upcast::<Node>().owner_shadow_root() {
+                shadow_root.register_named_element(self, value.clone());
+            } else {
+                doc.register_named_element(self, value.clone());
+            }
         }
         // This is used for layout optimization.
         doc.increment_dom_count();
diff --git a/components/script/dom/shadowroot.rs b/components/script/dom/shadowroot.rs
index 078c937d701..55535ad9a9d 100644
--- a/components/script/dom/shadowroot.rs
+++ b/components/script/dom/shadowroot.rs
@@ -20,6 +20,7 @@ use crate::dom::window::Window;
 use dom_struct::dom_struct;
 use selectors::context::QuirksMode;
 use servo_arc::Arc;
+use servo_atoms::Atom;
 use style::author_styles::AuthorStyles;
 use style::dom::TElement;
 use style::media_queries::Device;
@@ -126,6 +127,27 @@ impl ShadowRoot {
             .upcast::<Node>()
             .dirty(NodeDamage::NodeStyleDamaged);
     }
+
+    /// Remove any existing association between the provided id and any elements
+    /// in this shadow tree.
+    pub fn unregister_named_element(&self, to_unregister: &Element, id: Atom) {
+        self.document_or_shadow_root.unregister_named_element(
+            self.document_fragment.id_map(),
+            to_unregister,
+            &id,
+        );
+    }
+
+    /// Associate an element present in this shadow tree with the provided id.
+    pub fn register_named_element(&self, element: &Element, id: Atom) {
+        let root = self.upcast::<Node>().inclusive_ancestors().last().unwrap();
+        self.document_or_shadow_root.register_named_element(
+            self.document_fragment.id_map(),
+            element,
+            &id,
+            root,
+        );
+    }
 }
 
 impl ShadowRootMethods for ShadowRoot {
diff --git a/tests/wpt/mozilla/tests/mozilla/partial_shadow_dom.html b/tests/wpt/mozilla/tests/mozilla/partial_shadow_dom.html
index d97f1422d20..74e308f9403 100644
--- a/tests/wpt/mozilla/tests/mozilla/partial_shadow_dom.html
+++ b/tests/wpt/mozilla/tests/mozilla/partial_shadow_dom.html
@@ -6,7 +6,7 @@
   <script src="/resources/testharnessreport.js"></script>
 </head>
 <body>
-  <h1 class="header">Not in the shadows</h1>
+  <h1 id="header" class="header">Not in the shadows</h1>
   <div id="host">
   </div>
   <script>
@@ -20,12 +20,14 @@
 
       assert_equals(document.body.getElementsByTagName('h1').length, 1);
       assert_equals(document.body.getElementsByClassName('header').length, 1);
+      assert_equals(document.getElementById('header').textContent, "Not in the shadows");
       assert_equals(document.querySelectorAll('h1').length, 1);
       assert_equals(document.body.childNodes.length, 6);
 
       // Append child to shadow tree and check that its content is encapsulated.
       var shadowChild = document.createElement('h1');
       shadowChild.classList.add('header');
+      shadowChild.setAttribute('id', 'header');
       shadowChild.textContent = 'In the shadows';
       shadowRoot.appendChild(shadowChild);
       assert_equals(document.body.getElementsByTagName('h1').length, 1);
@@ -33,6 +35,7 @@
       assert_equals(document.querySelectorAll('h1').length, 1);
       assert_equals(document.body.childNodes.length, 6);
       assert_equals(shadowRoot.querySelectorAll('h1').length, 1);
+      assert_equals(shadowRoot.getElementById('header').textContent, "In the shadows");
     });
   </script>
 </body>