Skip to main content

blitz_dom/
query_selector.rs

1use blitz_traits::node_id::NodeId;
2use selectors::SelectorList;
3use smallvec::SmallVec;
4use style::dom::{TDocument, TNode};
5use style::dom_apis::{
6    MayUseInvalidation, QueryAll, QueryFirst, element_closest, element_matches, query_selector,
7};
8use style::selector_parser::{SelectorImpl, SelectorParser};
9use style_traits::ParseError;
10
11use crate::{BaseDocument, Node};
12
13impl BaseDocument {
14    /// Find the node with the specified id attribute (if one exists).
15    /// If multiple nodes have the same id, the first in tree order is returned.
16    pub fn get_element_by_id(&self, id: &str) -> Option<NodeId> {
17        match self.nodes_to_id.get(id)?.as_slice() {
18            [] => None,
19            [node_id] => Some(*node_id),
20            candidates => self.first_in_tree_order(candidates),
21        }
22    }
23
24    /// Find the first of `candidates` in tree order
25    fn first_in_tree_order(&self, candidates: &[NodeId]) -> Option<NodeId> {
26        let mut stack = vec![self.root_node_id];
27        while let Some(node_id) = stack.pop() {
28            if candidates.contains(&node_id) {
29                return Some(node_id);
30            }
31            stack.extend(self.nodes[node_id].children.iter().rev().copied());
32        }
33        None
34    }
35
36    /// Add a node to the id-to-node map
37    pub(crate) fn add_to_id_map(&mut self, id: &str, node_id: NodeId) {
38        if id.is_empty() {
39            return;
40        }
41        let node_ids = self.nodes_to_id.entry(id.to_string()).or_default();
42        if !node_ids.contains(&node_id) {
43            node_ids.push(node_id);
44        }
45    }
46
47    /// Remove a node from the id-to-node map
48    pub(crate) fn remove_from_id_map(&mut self, id: &str, node_id: NodeId) {
49        if let Some(node_ids) = self.nodes_to_id.get_mut(id) {
50            node_ids.retain(|nid| *nid != node_id);
51            if node_ids.is_empty() {
52                self.nodes_to_id.remove(id);
53            }
54        }
55    }
56
57    /// Find the first node that matches the selector specified as a string
58    /// Returns:
59    ///   - Err(_) if parsing the selector fails
60    ///   - Ok(None) if nothing matches
61    ///   - Ok(Some(node_id)) with the first node ID that matches if one is found
62    pub fn query_selector<'input>(
63        &self,
64        selector: &'input str,
65    ) -> Result<Option<NodeId>, ParseError<'input>> {
66        self.query_selector_in(self.root_node_id, selector)
67    }
68
69    /// Find the first descendant of `scope` that matches the selector specified
70    /// as a string.
71    ///
72    /// The scope node itself is never matched. Selector parts may match nodes
73    /// outside the scope while evaluating relationships between descendants.
74    ///
75    /// Returns:
76    ///   - `Err(_)` if parsing the selector fails
77    ///   - `Ok(None)` if nothing matches
78    ///   - `Ok(Some(node_id))` with the first matching descendant ID otherwise
79    pub fn query_selector_in<'input>(
80        &self,
81        scope: NodeId,
82        selector: &'input str,
83    ) -> Result<Option<NodeId>, ParseError<'input>> {
84        let selector_list = self.try_parse_selector_list(selector)?;
85        Ok(self.query_selector_in_raw(scope, &selector_list))
86    }
87
88    /// Find the first descendant of the document root that matches the
89    /// selector(s) specified in `selector_list`.
90    ///
91    /// The document root itself is never matched. Selector parts may match
92    /// nodes outside the scope while evaluating relationships between
93    /// descendants.
94    pub fn query_selector_raw(&self, selector_list: &SelectorList<SelectorImpl>) -> Option<NodeId> {
95        self.query_selector_in_raw(self.root_node_id, selector_list)
96    }
97
98    /// Find the first descendant of `scope` that matches the selector(s)
99    /// specified in `selector_list`.
100    ///
101    /// The scope node itself is never matched. Selector parts may match nodes
102    /// outside the scope while evaluating relationships between descendants.
103    pub fn query_selector_in_raw(
104        &self,
105        scope: NodeId,
106        selector_list: &SelectorList<SelectorImpl>,
107    ) -> Option<NodeId> {
108        let root_node = &self.nodes[scope];
109        let mut result = None;
110        query_selector::<&Node, QueryFirst>(
111            root_node,
112            selector_list,
113            &mut result,
114            self.may_use_invalidation_for(scope),
115        );
116
117        result.map(|node| node.id)
118    }
119
120    /// Find all nodes that match the selector specified as a string
121    /// Returns:
122    ///   - `Err(_)` if parsing the selector fails
123    ///   - `Ok(SmallVec<usize>)` with all matching nodes otherwise
124    pub fn query_selector_all<'input>(
125        &self,
126        selector: &'input str,
127    ) -> Result<SmallVec<[NodeId; 32]>, ParseError<'input>> {
128        self.query_selector_all_in(self.root_node_id, selector)
129    }
130
131    /// Find all descendants of `scope` that match the selector specified as a
132    /// string, in tree order.
133    ///
134    /// The scope node itself is never matched. Selector parts may match nodes
135    /// outside the scope while evaluating relationships between descendants.
136    ///
137    /// Returns:
138    ///   - `Err(_)` if parsing the selector fails
139    ///   - `Ok(_)` with all matching descendant IDs otherwise
140    pub fn query_selector_all_in<'input>(
141        &self,
142        scope: NodeId,
143        selector: &'input str,
144    ) -> Result<SmallVec<[NodeId; 32]>, ParseError<'input>> {
145        let selector_list = self.try_parse_selector_list(selector)?;
146        Ok(self.query_selector_all_in_raw(scope, &selector_list))
147    }
148
149    /// Find all descendants of the document root that match the selector(s)
150    /// specified in `selector_list`, in tree order.
151    ///
152    /// The document root itself is never matched. Selector parts may match
153    /// nodes outside the scope while evaluating relationships between
154    /// descendants.
155    pub fn query_selector_all_raw(
156        &self,
157        selector_list: &SelectorList<SelectorImpl>,
158    ) -> SmallVec<[NodeId; 32]> {
159        self.query_selector_all_in_raw(self.root_node_id, selector_list)
160    }
161
162    /// Find all descendants of `scope` that match the selector(s) specified in
163    /// `selector_list`, in tree order.
164    ///
165    /// The scope node itself is never matched. Selector parts may match nodes
166    /// outside the scope while evaluating relationships between descendants.
167    pub fn query_selector_all_in_raw(
168        &self,
169        scope: NodeId,
170        selector_list: &SelectorList<SelectorImpl>,
171    ) -> SmallVec<[NodeId; 32]> {
172        let root_node = &self.nodes[scope];
173        let mut results = SmallVec::new();
174        query_selector::<&Node, QueryAll>(
175            root_node,
176            selector_list,
177            &mut results,
178            self.may_use_invalidation_for(scope),
179        );
180
181        results.iter().map(|node| node.id).collect()
182    }
183
184    fn may_use_invalidation_for(&self, scope: NodeId) -> MayUseInvalidation {
185        if scope == self.root_node_id {
186            MayUseInvalidation::Yes
187        } else {
188            MayUseInvalidation::No
189        }
190    }
191
192    /// Test whether the node identified by `node_id` matches the selector
193    /// specified as a string.
194    ///
195    /// Non-element nodes never match.
196    pub fn matches_selector<'input>(
197        &self,
198        node_id: NodeId,
199        selector: &'input str,
200    ) -> Result<bool, ParseError<'input>> {
201        let selector_list = self.try_parse_selector_list(selector)?;
202        Ok(self.nodes[node_id].matches_selector_raw(&selector_list))
203    }
204
205    /// Find the closest matching element at or above the node identified by
206    /// `node_id`.
207    ///
208    /// Non-element nodes never match and return `None`.
209    pub fn closest<'input>(
210        &self,
211        node_id: NodeId,
212        selector: &'input str,
213    ) -> Result<Option<NodeId>, ParseError<'input>> {
214        let selector_list = self.try_parse_selector_list(selector)?;
215        Ok(self.nodes[node_id].closest_raw(&selector_list))
216    }
217
218    pub fn try_parse_selector_list<'input>(
219        &self,
220        input: &'input str,
221    ) -> Result<SelectorList<SelectorImpl>, ParseError<'input>> {
222        let url_extra_data = self.url.url_extra_data();
223        SelectorParser::parse_author_origin_no_namespace(input, &url_extra_data)
224    }
225}
226
227impl Node {
228    /// Find the first descendant of this node that matches the selector(s)
229    /// specified in `selector_list`.
230    ///
231    /// The scope node itself is never matched. Selector parts may match nodes
232    /// outside the scope while evaluating relationships between descendants.
233    ///
234    /// Text and comment scope nodes return no matches.
235    pub fn query_selector_raw(&self, selector_list: &SelectorList<SelectorImpl>) -> Option<NodeId> {
236        let mut result = None;
237        query_selector::<&Node, QueryFirst>(
238            self,
239            selector_list,
240            &mut result,
241            MayUseInvalidation::No,
242        );
243        result.map(|node| node.id)
244    }
245
246    /// Find all descendants of this node that match the selector(s) specified
247    /// in `selector_list`, in tree order.
248    ///
249    /// The scope node itself is never matched. Selector parts may match nodes
250    /// outside the scope while evaluating relationships between descendants.
251    ///
252    /// Text and comment scope nodes return no matches.
253    pub fn query_selector_all_raw(
254        &self,
255        selector_list: &SelectorList<SelectorImpl>,
256    ) -> SmallVec<[NodeId; 32]> {
257        let mut results = SmallVec::new();
258        query_selector::<&Node, QueryAll>(
259            self,
260            selector_list,
261            &mut results,
262            MayUseInvalidation::No,
263        );
264        results.iter().map(|node| node.id).collect()
265    }
266
267    /// Test whether this element matches the selector(s) specified in
268    /// `selector_list`.
269    ///
270    /// Non-element nodes never match.
271    pub fn matches_selector_raw(&self, selector_list: &SelectorList<SelectorImpl>) -> bool {
272        if !self.is_element() {
273            return false;
274        }
275
276        element_matches(&self, selector_list, self.owner_doc().quirks_mode())
277    }
278
279    /// Find the closest matching element at or above this element.
280    ///
281    /// Non-element nodes never match and return `None`.
282    pub fn closest_raw(&self, selector_list: &SelectorList<SelectorImpl>) -> Option<NodeId> {
283        if !self.is_element() {
284            return None;
285        }
286
287        element_closest(self, selector_list, self.owner_doc().quirks_mode()).map(|node| node.id)
288    }
289}