perf: Use element indexes for faster selector lookups in some cases

Signed-off-by: Dmitry Dygalo <dmitry@dygalo.dev>

Author: Stranger6667

CHANGELOG.md

@@ -12,6 +12,10 @@
 
 - CLI tests compatibility with future Cargo versions. [#602](https://github.com/Stranger6667/css-inline/issues/602)
 
+### Performance
+
+- Use element indexes for faster selector lookups on large documents.
+
 ## [0.18.0] - 2025-11-01
 
 ### Added

bindings/c/CHANGELOG.md

@@ -8,6 +8,10 @@
 - Update `html5ever` to `0.36`.
 - Update `selectors` to `0.33`.
 
+### Performance
+
+- Use element indexes for faster selector lookups on large documents.
+
 ## [0.18.0] - 2025-11-01
 
 ### Added

bindings/java/CHANGELOG.md

@@ -8,6 +8,10 @@
 - Update `html5ever` to `0.36`.
 - Update `selectors` to `0.33`.
 
+### Performance
+
+- Use element indexes for faster selector lookups on large documents.
+
 ## [0.18.0] - 2025-11-01
 
 ### Added

bindings/javascript/CHANGELOG.md

@@ -8,6 +8,10 @@
 - Update `html5ever` to `0.36`.
 - Update `selectors` to `0.33`.
 
+### Performance
+
+- Use element indexes for faster selector lookups on large documents.
+
 ## [0.18.0] - 2025-11-01
 
 ### Added

bindings/python/CHANGELOG.md

@@ -8,6 +8,10 @@
 - Update `html5ever` to `0.36`.
 - Update `selectors` to `0.33`.
 
+### Performance
+
+- Use element indexes for faster selector lookups on large documents.
+
 ## [0.18.0] - 2025-11-01
 
 ### Added

bindings/ruby/CHANGELOG.md

@@ -8,6 +8,10 @@
 - Update `html5ever` to `0.36`.
 - Update `selectors` to `0.33`.
 
+### Performance
+
+- Use element indexes for faster selector lookups on large documents.
+
 ## [0.18.0] - 2025-11-01
 
 ### Added

css-inline/src/html/document.rs

@@ -8,10 +8,15 @@ use super::{
     InliningMode,
 };
 use crate::{html::DocumentStyleMap, InlineError};
-use html5ever::local_name;
+use html5ever::{local_name, tendril::StrTendril, LocalName};
+use rustc_hash::FxHashMap;
 use selectors::context::SelectorCaches;
 use std::{fmt, fmt::Formatter, io::Write, iter::successors};
 
+/// Minimum input size (bytes) to build element indexes.
+/// For smaller inputs, linear scan is faster than hash map operations.
+const INDEX_SIZE_THRESHOLD: usize = 1024;
+
 /// HTML document representation.
 ///
 /// A `Document` holds a collection of nodes, with each node representing an HTML element.
@@ -46,6 +51,14 @@ pub(crate) struct Document {
     /// Ids of `link` nodes, specifically those with the `rel` attribute value set as `stylesheet`.
     /// They represent the locations (URLs) of all linked stylesheet resources in the document.
     linked_stylesheets: Vec<NodeId>,
+    /// Whether element indexes are enabled for this document.
+    use_indexes: bool,
+    /// Index: ID attribute value -> `NodeId` (IDs should be unique per document).
+    by_id: FxHashMap<LocalName, NodeId>,
+    /// Index: class name -> list of `NodeId`s with that class.
+    by_class: FxHashMap<LocalName, Vec<NodeId>>,
+    /// Index: tag name -> list of `NodeId`s with that tag.
+    by_tag: FxHashMap<LocalName, Vec<NodeId>>,
 }
 
 impl fmt::Debug for Document {
@@ -68,16 +81,22 @@ impl Document {
         parser::parse_with_options(bytes, preallocate_node_capacity, mode)
     }
 
-    pub(super) fn with_capacity(capacity: usize) -> Self {
+    pub(super) fn with_capacity(capacity: usize, input_size: usize) -> Self {
         // Dummy node at index 0 so that other indices fit in NonZero
         let mut nodes = vec![Node::new(NodeData::Document), Node::new(NodeData::Document)];
         // Usually there are a lot of nodes, hence, reserve some space for them
         nodes.reserve(capacity);
+        // Only build indexes for larger documents where the lookup benefit outweighs overhead
+        let use_indexes = input_size >= INDEX_SIZE_THRESHOLD;
         Document {
             nodes,
             elements: Vec::with_capacity(capacity),
             styles: Vec::new(),
             linked_stylesheets: Vec::new(),
+            use_indexes,
+            by_id: FxHashMap::default(),
+            by_class: FxHashMap::default(),
+            by_tag: FxHashMap::default(),
         }
     }
 
@@ -130,6 +149,58 @@ impl Document {
         self.elements.push(node);
     }
 
+    /// Whether element indexes are enabled for this document.
+    #[inline]
+    pub(crate) fn has_indexes(&self) -> bool {
+        self.use_indexes
+    }
+
+    /// Index an element by its tag name.
+    #[inline]
+    pub(super) fn index_by_tag(&mut self, tag: LocalName, node_id: NodeId) {
+        if self.use_indexes {
+            self.by_tag.entry(tag).or_default().push(node_id);
+        }
+    }
+
+    /// Index an element by its ID attribute value.
+    #[inline]
+    pub(super) fn index_by_id(&mut self, id: &StrTendril, node_id: NodeId) {
+        if self.use_indexes && !id.is_empty() {
+            let id_atom = LocalName::from(&**id);
+            // IDs should be unique; later declarations override earlier ones
+            self.by_id.insert(id_atom, node_id);
+        }
+    }
+
+    /// Index an element by its class names (space-separated).
+    pub(super) fn index_by_classes(&mut self, classes: &StrTendril, node_id: NodeId) {
+        if self.use_indexes {
+            for class in classes.split_whitespace() {
+                let class_atom = LocalName::from(class);
+                self.by_class.entry(class_atom).or_default().push(node_id);
+            }
+        }
+    }
+
+    /// Lookup an element by its ID attribute value.
+    #[inline]
+    pub(crate) fn get_by_id(&self, id: &LocalName) -> Option<NodeId> {
+        self.by_id.get(id).copied()
+    }
+
+    /// Lookup elements by class name.
+    #[inline]
+    pub(crate) fn get_by_class(&self, class: &LocalName) -> &[NodeId] {
+        self.by_class.get(class).map_or(&[], Vec::as_slice)
+    }
+
+    /// Lookup elements by tag name.
+    #[inline]
+    pub(crate) fn get_by_tag(&self, tag: &LocalName) -> &[NodeId] {
+        self.by_tag.get(tag).map_or(&[], Vec::as_slice)
+    }
+
     /// Detach a node from its siblings and its parent.
     ///
     /// Before:
@@ -402,7 +473,7 @@ mod tests {
 
     #[test]
     fn test_insert_before() {
-        let mut doc = Document::with_capacity(0);
+        let mut doc = Document::with_capacity(0, 0);
 
         let node1_id = doc.push_node(new_element());
         let node2_id = doc.push_node(new_element());
@@ -420,7 +491,7 @@ mod tests {
 
     #[test]
     fn test_append() {
-        let mut doc = Document::with_capacity(0);
+        let mut doc = Document::with_capacity(0, 0);
 
         let node1_id = doc.push_node(new_element());
         let node2_id = doc.push_node(new_element());

css-inline/src/html/iter.rs

@@ -2,31 +2,67 @@ use super::{
     document::Document,
     element::Element,
     node::{NodeData, NodeId},
-    selectors::{ParseError, Selectors},
+    selectors::{ParseError, SelectorAnchor, Selectors},
     Specificity,
 };
 use selectors::context::SelectorCaches;
 
+/// Source of elements to iterate over.
+/// Allows using different indexes for faster lookup.
+enum ElementSource<'a> {
+    /// Iterate over a slice of elements (all elements or indexed subset).
+    Slice(std::slice::Iter<'a, NodeId>),
+    /// Single element from ID lookup.
+    Single(Option<NodeId>),
+}
+
+impl ElementSource<'_> {
+    #[inline]
+    fn next(&mut self) -> Option<NodeId> {
+        match self {
+            ElementSource::Slice(iter) => iter.next().copied(),
+            ElementSource::Single(opt) => opt.take(),
+        }
+    }
+}
+
 /// Compile selectors from a string and create an element iterator that yields elements matching these selectors.
 #[inline]
 pub(crate) fn select<'a, 'b, 'c>(
     document: &'a Document,
     selectors: &'b str,
     caches: &'c mut SelectorCaches,
 ) -> Result<Select<'a, 'c>, ParseError<'b>> {
-    Selectors::compile(selectors).map(|selectors| Select {
-        document,
-        caches,
-        iter: document.elements.iter(),
-        selectors,
+    Selectors::compile(selectors).map(|selectors| {
+        // Only use indexes if they were built during parsing
+        let source = if document.has_indexes() {
+            match selectors.anchor() {
+                SelectorAnchor::Id(id) => ElementSource::Single(document.get_by_id(id.as_inner())),
+                SelectorAnchor::Class(class) => {
+                    ElementSource::Slice(document.get_by_class(class.as_inner()).iter())
+                }
+                SelectorAnchor::Tag(tag) => {
+                    ElementSource::Slice(document.get_by_tag(tag.as_inner()).iter())
+                }
+                SelectorAnchor::None => ElementSource::Slice(document.elements.iter()),
+            }
+        } else {
+            ElementSource::Slice(document.elements.iter())
+        };
+        Select {
+            document,
+            caches,
+            source,
+            selectors,
+        }
     })
 }
 
 /// An element iterator adaptor that yields elements matching given selectors.
 pub(crate) struct Select<'a, 'c> {
     document: &'a Document,
     caches: &'c mut SelectorCaches,
-    iter: std::slice::Iter<'a, NodeId>,
+    source: ElementSource<'a>,
     /// The selectors to be matched.
     selectors: Selectors,
 }
@@ -44,11 +80,11 @@ impl<'a> Iterator for Select<'a, '_> {
 
     #[inline]
     fn next(&mut self) -> Option<Element<'a>> {
-        for element_id in self.iter.by_ref() {
-            let NodeData::Element { element, .. } = &self.document[*element_id].data else {
+        while let Some(element_id) = self.source.next() {
+            let NodeData::Element { element, .. } = &self.document[element_id].data else {
                 unreachable!("Element ids always point to element nodes")
             };
-            let element = Element::new(self.document, *element_id, element);
+            let element = Element::new(self.document, element_id, element);
             for selector in self.selectors.iter() {
                 if element.matches(selector, self.caches) {
                     return Some(element);

css-inline/src/html/parser.rs

@@ -29,7 +29,10 @@ pub(crate) fn parse_with_options(
     mode: InliningMode,
 ) -> Document {
     let sink = Sink {
-        document: RefCell::new(Document::with_capacity(preallocate_node_capacity)),
+        document: RefCell::new(Document::with_capacity(
+            preallocate_node_capacity,
+            bytes.len(),
+        )),
     };
     let options = html5ever::ParseOpts::default();
     match mode {
@@ -80,11 +83,32 @@ impl Sink {
         attributes: Vec<Attribute>,
         inlining_ignored: bool,
     ) -> NodeId {
+        // Extract ID and class values before moving attributes
+        let mut id_value = None;
+        let mut class_value = None;
+        for attr in &attributes {
+            if attr.name.local == local_name!("id") {
+                id_value = Some(attr.value.clone());
+            } else if attr.name.local == local_name!("class") {
+                class_value = Some(attr.value.clone());
+            }
+        }
+        let tag_name = name.local.clone();
         let node_id = self.push_node(NodeData::Element {
             element: ElementData::new(name, attributes),
             inlining_ignored,
         });
-        self.document.borrow_mut().push_element_id(node_id);
+        let mut document = self.document.borrow_mut();
+        document.push_element_id(node_id);
+        // Index by tag name
+        document.index_by_tag(tag_name, node_id);
+        // Index by ID and class attributes
+        if let Some(id) = &id_value {
+            document.index_by_id(id, node_id);
+        }
+        if let Some(class) = &class_value {
+            document.index_by_classes(class, node_id);
+        }
         node_id
     }
 

css-inline/src/html/selectors/local_name.rs

@@ -12,6 +12,11 @@ impl LocalName {
     pub(crate) fn into_inner(self) -> html5ever::LocalName {
         self.0
     }
+    /// Returns a reference to the inner `html5ever::LocalName`.
+    #[inline]
+    pub(crate) fn as_inner(&self) -> &html5ever::LocalName {
+        &self.0
+    }
     pub(crate) fn as_bytes(&self) -> &[u8] {
         self.0.as_bytes()
     }