From 8cae20af1be7e81857cee8d5679a0ba8dc220d10 Mon Sep 17 00:00:00 2001
From: Aliaksandr Kalenik <kalenik.aliaksandr@gmail.com>
Date: Tue, 25 Mar 2025 17:30:52 +0000
Subject: [PATCH] LibWeb: Maintain a mapping for fast lookup in
 getElementById()

With this change we maintain a data structure that maps ids to
corresponding elements. This allows us to avoid tree traversal in
getElementById() in all cases except ones when lookup happens for
unconnected elements.
---
 Libraries/LibWeb/CMakeLists.txt             |  1 +
 Libraries/LibWeb/DOM/Document.cpp           | 35 +++++++++++++++++-
 Libraries/LibWeb/DOM/Document.h             |  7 ++--
 Libraries/LibWeb/DOM/DocumentFragment.h     |  4 +--
 Libraries/LibWeb/DOM/Element.cpp            | 16 +++++++--
 Libraries/LibWeb/DOM/Element.h              |  1 +
 Libraries/LibWeb/DOM/ElementByIdMap.cpp     | 39 +++++++++++++++++++++
 Libraries/LibWeb/DOM/ElementByIdMap.h       | 24 +++++++++++++
 Libraries/LibWeb/DOM/NonElementParentNode.h | 38 --------------------
 Libraries/LibWeb/DOM/ParentNode.cpp         | 23 ++++++++++++
 Libraries/LibWeb/DOM/ParentNode.h           |  2 ++
 Libraries/LibWeb/DOM/ShadowRoot.cpp         |  7 ++++
 Libraries/LibWeb/DOM/ShadowRoot.h           |  5 +++
 Libraries/LibWeb/Forward.h                  |  1 +
 Libraries/LibWeb/SVG/SVGSVGElement.h        |  5 +--
 15 files changed, 157 insertions(+), 51 deletions(-)
 create mode 100644 Libraries/LibWeb/DOM/ElementByIdMap.cpp
 create mode 100644 Libraries/LibWeb/DOM/ElementByIdMap.h
 delete mode 100644 Libraries/LibWeb/DOM/NonElementParentNode.h

diff --git a/Libraries/LibWeb/CMakeLists.txt b/Libraries/LibWeb/CMakeLists.txt
index 09009f99d40..31417efeebc 100644
--- a/Libraries/LibWeb/CMakeLists.txt
+++ b/Libraries/LibWeb/CMakeLists.txt
@@ -214,6 +214,7 @@ set(SOURCES
     DOM/DocumentType.cpp
     DOM/EditingHostManager.cpp
     DOM/Element.cpp
+    DOM/ElementByIdMap.cpp
     DOM/ElementFactory.cpp
     DOM/Event.cpp
     DOM/EventDispatcher.cpp
diff --git a/Libraries/LibWeb/DOM/Document.cpp b/Libraries/LibWeb/DOM/Document.cpp
index eb6a4bd8e36..9a573e720f6 100644
--- a/Libraries/LibWeb/DOM/Document.cpp
+++ b/Libraries/LibWeb/DOM/Document.cpp
@@ -63,6 +63,7 @@
 #include <LibWeb/DOM/DocumentType.h>
 #include <LibWeb/DOM/EditingHostManager.h>
 #include <LibWeb/DOM/Element.h>
+#include <LibWeb/DOM/ElementByIdMap.h>
 #include <LibWeb/DOM/ElementFactory.h>
 #include <LibWeb/DOM/Event.h>
 #include <LibWeb/DOM/HTMLCollection.h>
@@ -5345,7 +5346,7 @@ static void insert_in_tree_order(Vector<GC::Ref<DOM::Element>>& elements, DOM::E
         elements.append(element);
 }
 
-void Document::element_id_changed(Badge<DOM::Element>, GC::Ref<DOM::Element> element)
+void Document::element_id_changed(Badge<DOM::Element>, GC::Ref<DOM::Element> element, Optional<FlyString> old_id)
 {
     for (auto* form_associated_element : m_form_associated_elements_with_form_attribute)
         form_associated_element->element_id_changed({});
@@ -5354,6 +5355,14 @@ void Document::element_id_changed(Badge<DOM::Element>, GC::Ref<DOM::Element> ele
         insert_in_tree_order(m_potentially_named_elements, element);
     else
         (void)m_potentially_named_elements.remove_first_matching([element](auto& e) { return e == element; });
+
+    auto new_id = element->id();
+    if (old_id.has_value()) {
+        element->document_or_shadow_root_element_by_id_map().remove(old_id.value(), element);
+    }
+    if (new_id.has_value()) {
+        element->document_or_shadow_root_element_by_id_map().add(new_id.value(), element);
+    }
 }
 
 void Document::element_with_id_was_added(Badge<DOM::Element>, GC::Ref<DOM::Element> element)
@@ -5363,6 +5372,10 @@ void Document::element_with_id_was_added(Badge<DOM::Element>, GC::Ref<DOM::Eleme
 
     if (is_potentially_named_element_by_id(*element))
         insert_in_tree_order(m_potentially_named_elements, element);
+
+    if (auto id = element->id(); id.has_value()) {
+        element->document_or_shadow_root_element_by_id_map().add(id.value(), element);
+    }
 }
 
 void Document::element_with_id_was_removed(Badge<DOM::Element>, GC::Ref<DOM::Element> element)
@@ -5372,6 +5385,10 @@ void Document::element_with_id_was_removed(Badge<DOM::Element>, GC::Ref<DOM::Ele
 
     if (is_potentially_named_element_by_id(*element))
         (void)m_potentially_named_elements.remove_first_matching([element](auto& e) { return e == element; });
+
+    if (auto id = element->id(); id.has_value()) {
+        element->document_or_shadow_root_element_by_id_map().remove(id.value(), element);
+    }
 }
 
 void Document::element_name_changed(Badge<DOM::Element>, GC::Ref<DOM::Element> element)
@@ -6441,6 +6458,22 @@ void Document::set_onvisibilitychange(WebIDL::CallbackType* value)
     set_event_handler_attribute(HTML::EventNames::visibilitychange, value);
 }
 
+ElementByIdMap& Document::element_by_id() const
+{
+    if (!m_element_by_id)
+        m_element_by_id = make<ElementByIdMap>();
+    return *m_element_by_id;
+}
+
+GC::Ptr<Element> ElementByIdMap::get(FlyString const& element_id) const
+{
+    if (auto elements = m_map.get(element_id); elements.has_value() && !elements->is_empty()) {
+        if (auto element = elements->first(); element.has_value())
+            return *element;
+    }
+    return {};
+}
+
 StringView to_string(SetNeedsLayoutReason reason)
 {
     switch (reason) {
diff --git a/Libraries/LibWeb/DOM/Document.h b/Libraries/LibWeb/DOM/Document.h
index 935c9bec40f..edd6c9c863a 100644
--- a/Libraries/LibWeb/DOM/Document.h
+++ b/Libraries/LibWeb/DOM/Document.h
@@ -25,7 +25,6 @@
 #include <LibWeb/CSS/CSSStyleSheet.h>
 #include <LibWeb/CSS/StyleSheetList.h>
 #include <LibWeb/Cookie/Cookie.h>
-#include <LibWeb/DOM/NonElementParentNode.h>
 #include <LibWeb/DOM/ParentNode.h>
 #include <LibWeb/HTML/BrowsingContext.h>
 #include <LibWeb/HTML/CrossOrigin/OpenerPolicy.h>
@@ -158,7 +157,6 @@ enum class PolicyControlledFeature : u8 {
 
 class Document
     : public ParentNode
-    , public NonElementParentNode<Document>
     , public HTML::GlobalEventHandlers {
     WEB_PLATFORM_OBJECT(Document, ParentNode);
     GC_DECLARE_ALLOCATOR(Document);
@@ -742,7 +740,7 @@ public:
     GC::Ptr<HTML::SessionHistoryEntry> latest_entry() const { return m_latest_entry; }
     void set_latest_entry(GC::Ptr<HTML::SessionHistoryEntry> e) { m_latest_entry = e; }
 
-    void element_id_changed(Badge<DOM::Element>, GC::Ref<DOM::Element> element);
+    void element_id_changed(Badge<DOM::Element>, GC::Ref<DOM::Element> element, Optional<FlyString> old_id);
     void element_with_id_was_added(Badge<DOM::Element>, GC::Ref<DOM::Element> element);
     void element_with_id_was_removed(Badge<DOM::Element>, GC::Ref<DOM::Element> element);
     void element_name_changed(Badge<DOM::Element>, GC::Ref<DOM::Element> element);
@@ -895,6 +893,8 @@ public:
         m_pending_nodes_for_style_invalidation_due_to_presence_of_has.set(node.make_weak_ptr<Node>());
     }
 
+    ElementByIdMap& element_by_id() const;
+
 protected:
     virtual void initialize(JS::Realm&) override;
     virtual void visit_edges(Cell::Visitor&) override;
@@ -952,6 +952,7 @@ private:
     GC::Ptr<Node> m_active_favicon;
     WeakPtr<HTML::BrowsingContext> m_browsing_context;
     URL::URL m_url;
+    mutable OwnPtr<ElementByIdMap> m_element_by_id;
 
     GC::Ptr<HTML::Window> m_window;
 
diff --git a/Libraries/LibWeb/DOM/DocumentFragment.h b/Libraries/LibWeb/DOM/DocumentFragment.h
index 5a3fe876d5a..4ba3f269cfd 100644
--- a/Libraries/LibWeb/DOM/DocumentFragment.h
+++ b/Libraries/LibWeb/DOM/DocumentFragment.h
@@ -7,14 +7,12 @@
 #pragma once
 
 #include <LibWeb/DOM/Element.h>
-#include <LibWeb/DOM/NonElementParentNode.h>
 #include <LibWeb/DOM/ParentNode.h>
 
 namespace Web::DOM {
 
 class DocumentFragment
-    : public ParentNode
-    , public NonElementParentNode<DocumentFragment> {
+    : public ParentNode {
     WEB_PLATFORM_OBJECT(DocumentFragment, ParentNode);
     GC_DECLARE_ALLOCATOR(DocumentFragment);
 
diff --git a/Libraries/LibWeb/DOM/Element.cpp b/Libraries/LibWeb/DOM/Element.cpp
index cf830055967..028e123f545 100644
--- a/Libraries/LibWeb/DOM/Element.cpp
+++ b/Libraries/LibWeb/DOM/Element.cpp
@@ -3520,8 +3520,12 @@ void Element::attribute_changed(FlyString const& local_name, Optional<String> co
         else
             m_id = value_or_empty;
 
-        if (is_connected())
-            document().element_id_changed({}, *this);
+        if (is_connected()) {
+            Optional<FlyString> old_value_fly_string;
+            if (old_value.has_value())
+                old_value_fly_string = *old_value;
+            document().element_id_changed({}, *this, old_value_fly_string);
+        }
     } else if (local_name == HTML::AttributeNames::name) {
         if (value_or_empty.is_empty())
             m_name = {};
@@ -3591,6 +3595,14 @@ CSS::StyleSheetList& Element::document_or_shadow_root_style_sheets()
     return document().style_sheets();
 }
 
+ElementByIdMap& Element::document_or_shadow_root_element_by_id_map()
+{
+    auto& root_node = root();
+    if (is<ShadowRoot>(root_node))
+        return static_cast<ShadowRoot&>(root_node).element_by_id();
+    return document().element_by_id();
+}
+
 // https://html.spec.whatwg.org/multipage/dynamic-markup-insertion.html#dom-element-gethtml
 WebIDL::ExceptionOr<String> Element::get_html(GetHTMLOptions const& options) const
 {
diff --git a/Libraries/LibWeb/DOM/Element.h b/Libraries/LibWeb/DOM/Element.h
index 3bb562ed201..e62f86a5080 100644
--- a/Libraries/LibWeb/DOM/Element.h
+++ b/Libraries/LibWeb/DOM/Element.h
@@ -218,6 +218,7 @@ public:
     GC::Ref<CSS::CSSStyleProperties> style_for_bindings();
 
     CSS::StyleSheetList& document_or_shadow_root_style_sheets();
+    ElementByIdMap& document_or_shadow_root_element_by_id_map();
 
     WebIDL::ExceptionOr<GC::Ref<DOM::DocumentFragment>> parse_fragment(StringView markup);
 
diff --git a/Libraries/LibWeb/DOM/ElementByIdMap.cpp b/Libraries/LibWeb/DOM/ElementByIdMap.cpp
new file mode 100644
index 00000000000..5031434109e
--- /dev/null
+++ b/Libraries/LibWeb/DOM/ElementByIdMap.cpp
@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) 2025, Aliaksandr Kalenik <kalenik.aliaksandr@gmail.com>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#include <LibWeb/DOM/ElementByIdMap.h>
+
+namespace Web::DOM {
+
+void ElementByIdMap::add(FlyString const& element_id, Element& element)
+{
+    auto& elements_with_id = m_map.ensure(element_id, [] { return Vector<WeakPtr<Element>> {}; });
+
+    // Remove all elements that were deallocated.
+    elements_with_id.remove_all_matching([](WeakPtr<Element>& element) {
+        return !element.has_value();
+    });
+
+    VERIFY(!elements_with_id.contains_slow(element));
+    elements_with_id.insert_before_matching(element, [&](auto& another_element) {
+        return element.is_before(*another_element);
+    });
+}
+
+void ElementByIdMap::remove(FlyString const& element_id, Element& element)
+{
+    auto maybe_elements_with_id = m_map.get(element_id);
+    if (!maybe_elements_with_id.has_value())
+        return;
+    auto& elements_with_id = *maybe_elements_with_id;
+    elements_with_id.remove_all_matching([&](auto& another_element) {
+        if (!another_element.has_value())
+            return true;
+        return &element == another_element.ptr();
+    });
+}
+
+}
diff --git a/Libraries/LibWeb/DOM/ElementByIdMap.h b/Libraries/LibWeb/DOM/ElementByIdMap.h
new file mode 100644
index 00000000000..57b12678342
--- /dev/null
+++ b/Libraries/LibWeb/DOM/ElementByIdMap.h
@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) 2025, Aliaksandr Kalenik <kalenik.aliaksandr@gmail.com>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+#include <LibWeb/DOM/Element.h>
+#include <LibWeb/Forward.h>
+
+namespace Web::DOM {
+
+class ElementByIdMap {
+public:
+    void add(FlyString const& element_id, Element&);
+    void remove(FlyString const& element_id, Element&);
+    GC::Ptr<Element> get(FlyString const& element_id) const;
+
+private:
+    HashMap<FlyString, Vector<WeakPtr<Element>>> m_map;
+};
+
+}
diff --git a/Libraries/LibWeb/DOM/NonElementParentNode.h b/Libraries/LibWeb/DOM/NonElementParentNode.h
deleted file mode 100644
index 57268356128..00000000000
--- a/Libraries/LibWeb/DOM/NonElementParentNode.h
+++ /dev/null
@@ -1,38 +0,0 @@
-/*
- * Copyright (c) 2020, Andreas Kling <andreas@ladybird.org>
- *
- * SPDX-License-Identifier: BSD-2-Clause
- */
-
-#pragma once
-
-#include <AK/FlyString.h>
-#include <AK/Forward.h>
-#include <LibGC/Ptr.h>
-#include <LibWeb/Forward.h>
-#include <LibWeb/HTML/AttributeNames.h>
-#include <LibWeb/TreeNode.h>
-
-namespace Web::DOM {
-
-template<typename NodeType>
-class NonElementParentNode {
-public:
-    GC::Ptr<Element> get_element_by_id(FlyString const& id) const
-    {
-        GC::Ptr<Element> found_element;
-        const_cast<NodeType*>(static_cast<NodeType const*>(this))->template for_each_in_inclusive_subtree_of_type<Element>([&](auto& element) {
-            if (element.id() == id) {
-                found_element = &element;
-                return TraversalDecision::Break;
-            }
-            return TraversalDecision::Continue;
-        });
-        return found_element;
-    }
-
-protected:
-    NonElementParentNode() = default;
-};
-
-}
diff --git a/Libraries/LibWeb/DOM/ParentNode.cpp b/Libraries/LibWeb/DOM/ParentNode.cpp
index 2a44bc40290..2f2b4b64761 100644
--- a/Libraries/LibWeb/DOM/ParentNode.cpp
+++ b/Libraries/LibWeb/DOM/ParentNode.cpp
@@ -260,4 +260,27 @@ GC::Ref<HTMLCollection> ParentNode::get_elements_by_class_name(StringView class_
     });
 }
 
+GC::Ptr<Element> ParentNode::get_element_by_id(FlyString const& id) const
+{
+    // For document and shadow root we have a cache that allows fast lookup.
+    if (is_document()) {
+        auto const& document = static_cast<Document const&>(*this);
+        return document.element_by_id().get(id);
+    }
+    if (is_shadow_root()) {
+        auto const& shadow_root = static_cast<ShadowRoot const&>(*this);
+        return shadow_root.element_by_id().get(id);
+    }
+
+    GC::Ptr<Element> found_element;
+    const_cast<ParentNode&>(*this).for_each_in_inclusive_subtree_of_type<Element>([&](Element& element) {
+        if (element.id() == id) {
+            found_element = &element;
+            return TraversalDecision::Break;
+        }
+        return TraversalDecision::Continue;
+    });
+    return found_element;
+}
+
 }
diff --git a/Libraries/LibWeb/DOM/ParentNode.h b/Libraries/LibWeb/DOM/ParentNode.h
index 6b1e9b70d45..74c4a938436 100644
--- a/Libraries/LibWeb/DOM/ParentNode.h
+++ b/Libraries/LibWeb/DOM/ParentNode.h
@@ -38,6 +38,8 @@ public:
 
     GC::Ref<HTMLCollection> get_elements_by_class_name(StringView);
 
+    GC::Ptr<Element> get_element_by_id(FlyString const& id) const;
+
 protected:
     ParentNode(JS::Realm& realm, Document& document, NodeType type)
         : Node(realm, document, type)
diff --git a/Libraries/LibWeb/DOM/ShadowRoot.cpp b/Libraries/LibWeb/DOM/ShadowRoot.cpp
index d84a92ec958..20d699abe9b 100644
--- a/Libraries/LibWeb/DOM/ShadowRoot.cpp
+++ b/Libraries/LibWeb/DOM/ShadowRoot.cpp
@@ -184,4 +184,11 @@ WebIDL::ExceptionOr<Vector<GC::Ref<Animations::Animation>>> ShadowRoot::get_anim
     return relevant_animations;
 }
 
+ElementByIdMap& ShadowRoot::element_by_id() const
+{
+    if (!m_element_by_id)
+        m_element_by_id = make<ElementByIdMap>();
+    return *m_element_by_id;
+}
+
 }
diff --git a/Libraries/LibWeb/DOM/ShadowRoot.h b/Libraries/LibWeb/DOM/ShadowRoot.h
index ef18117b56a..b55b483fd2d 100644
--- a/Libraries/LibWeb/DOM/ShadowRoot.h
+++ b/Libraries/LibWeb/DOM/ShadowRoot.h
@@ -8,6 +8,7 @@
 
 #include <LibWeb/Bindings/ShadowRootPrototype.h>
 #include <LibWeb/DOM/DocumentFragment.h>
+#include <LibWeb/DOM/ElementByIdMap.h>
 #include <LibWeb/WebIDL/ObservableArray.h>
 
 namespace Web::DOM {
@@ -62,6 +63,8 @@ public:
 
     WebIDL::ExceptionOr<Vector<GC::Ref<Animations::Animation>>> get_animations();
 
+    ElementByIdMap& element_by_id() const;
+
     virtual void finalize() override;
 
 protected:
@@ -90,6 +93,8 @@ private:
     // https://dom.spec.whatwg.org/#shadowroot-serializable
     bool m_serializable { false };
 
+    mutable OwnPtr<ElementByIdMap> m_element_by_id;
+
     GC::Ptr<CSS::StyleSheetList> m_style_sheets;
     mutable GC::Ptr<WebIDL::ObservableArray> m_adopted_style_sheets;
 };
diff --git a/Libraries/LibWeb/Forward.h b/Libraries/LibWeb/Forward.h
index d85324a960c..32c4a802ad2 100644
--- a/Libraries/LibWeb/Forward.h
+++ b/Libraries/LibWeb/Forward.h
@@ -311,6 +311,7 @@ class DOMImplementation;
 class DOMTokenList;
 class EditingHostManager;
 class Element;
+class ElementByIdMap;
 class Event;
 class EventHandler;
 class EventTarget;
diff --git a/Libraries/LibWeb/SVG/SVGSVGElement.h b/Libraries/LibWeb/SVG/SVGSVGElement.h
index 05ca00f8057..f5b14ab457b 100644
--- a/Libraries/LibWeb/SVG/SVGSVGElement.h
+++ b/Libraries/LibWeb/SVG/SVGSVGElement.h
@@ -7,7 +7,6 @@
 #pragma once
 
 #include <LibGfx/Bitmap.h>
-#include <LibWeb/DOM/NonElementParentNode.h>
 #include <LibWeb/Geometry/DOMMatrix.h>
 #include <LibWeb/Geometry/DOMPoint.h>
 #include <LibWeb/SVG/AttributeParser.h>
@@ -22,9 +21,7 @@
 namespace Web::SVG {
 
 class SVGSVGElement final : public SVGGraphicsElement
-    , public SVGViewport
-    // SVGSVGElement is not strictly a NonElementParentNode, but it implements the same get_element_by_id() method.
-    , public DOM::NonElementParentNode<SVGSVGElement> {
+    , public SVGViewport {
     WEB_PLATFORM_OBJECT(SVGSVGElement, SVGGraphicsElement);
     GC_DECLARE_ALLOCATOR(SVGSVGElement);