Tweak list of shadow roots attached to doc

2 files changed, 10 insertions, 9 deletions

diff --git a/components/script/dom/document.rs b/components/script/dom/document.rs
index 998055d3547..697d2938115 100644
--- a/components/script/dom/document.rs
+++ b/components/script/dom/document.rs
@@ -377,7 +377,7 @@ pub struct Document {
     /// https://html.spec.whatwg.org/multipage/#completely-loaded
     completely_loaded: Cell<bool>,
     /// List of shadow roots bound to the document tree.
-    shadow_roots: DomRefCell<Vec<Dom<ShadowRoot>>>,
+    shadow_roots: DomRefCell<HashSet<Dom<ShadowRoot>>>,
     /// Whether any of the shadow roots need the stylesheets flushed.
     shadow_roots_styles_changed: Cell<bool>,
 }
@@ -2674,7 +2674,7 @@ impl Document {
             completely_loaded: Cell::new(false),
             script_and_layout_blockers: Cell::new(0),
             delayed_tasks: Default::default(),
-            shadow_roots: DomRefCell::new(Vec::new()),
+            shadow_roots: DomRefCell::new(HashSet::new()),
             shadow_roots_styles_changed: Cell::new(false),
         }
     }
@@ -3132,16 +3132,13 @@ impl Document {
     pub fn register_shadow_root(&self, shadow_root: &ShadowRoot) {
         self.shadow_roots
             .borrow_mut()
-            .push(Dom::from_ref(shadow_root));
+            .insert(Dom::from_ref(shadow_root));
         self.invalidate_shadow_roots_stylesheets();
     }
 
     pub fn unregister_shadow_root(&self, shadow_root: &ShadowRoot) {
         let mut shadow_roots = self.shadow_roots.borrow_mut();
-        let position = shadow_roots.iter().position(|sr| **sr == *shadow_root);
-        if let Some(index) = position {
-            shadow_roots.remove(index);
-        }
+        shadow_roots.remove(&Dom::from_ref(shadow_root));
     }
 
     pub fn invalidate_shadow_roots_stylesheets(&self) {
diff --git a/components/script/dom/element.rs b/components/script/dom/element.rs
index 34a4ec41cfa..e0659729673 100644
--- a/components/script/dom/element.rs
+++ b/components/script/dom/element.rs
@@ -489,7 +489,9 @@ impl Element {
             .shadow_root
             .or_init(|| ShadowRoot::new(self, &*self.node.owner_doc()));
 
-        self.node.owner_doc().register_shadow_root(&*shadow_root);
+        if self.is_connected() {
+            self.node.owner_doc().register_shadow_root(&*shadow_root);
+        }
 
         Ok(shadow_root)
     }
@@ -2802,7 +2804,10 @@ impl VirtualMethods for Element {
             f.bind_form_control_to_tree();
         }
 
+        let doc = document_from_node(self);
+
         if let Some(shadow_root) = self.rare_data.shadow_root.get() {
+            doc.register_shadow_root(&shadow_root);
             let shadow_root = shadow_root.upcast::<Node>();
             shadow_root.set_flag(NodeFlags::IS_CONNECTED, context.tree_connected);
             for node in shadow_root.children() {
@@ -2815,7 +2820,6 @@ impl VirtualMethods for Element {
             return;
         }
 
-        let doc = document_from_node(self);
         if let Some(ref value) = *self.id_attribute.borrow() {
             if let Some(shadow_root) = self.upcast::<Node>().owner_shadow_root() {
                 shadow_root.register_named_element(self, value.clone());