use std::collections::{BTreeSet, HashMap};
use crate::dom::Dom;
use crate::node::NodeData;
use crate::node_id::NodeId;
pub(crate) type Bucket = BTreeSet<NodeId>;
#[derive(Debug, Default, Clone)]
pub(crate) struct Indexes {
pub(crate) by_id: HashMap<String, Bucket>,
pub(crate) by_tag: HashMap<String, Bucket>,
pub(crate) by_class: HashMap<String, Bucket>,
}
impl Indexes {
fn push_unique(bucket: &mut Bucket, id: NodeId) {
bucket.insert(id);
}
fn remove_from(map: &mut HashMap<String, Bucket>, key: &str, id: NodeId) {
if let Some(bucket) = map.get_mut(key) {
bucket.remove(&id);
if bucket.is_empty() {
map.remove(key);
}
}
}
pub(crate) fn register_id(&mut self, id: NodeId, id_value: &str) {
if id_value.is_empty() {
return;
}
Self::push_unique(self.by_id.entry(id_value.to_string()).or_default(), id);
}
pub(crate) fn unregister_id(&mut self, id: NodeId, id_value: &str) {
if id_value.is_empty() {
return;
}
Self::remove_from(&mut self.by_id, id_value, id);
}
pub(crate) fn register_tag(&mut self, id: NodeId, tag: &str) {
Self::push_unique(self.by_tag.entry(tag.to_string()).or_default(), id);
}
pub(crate) fn unregister_tag(&mut self, id: NodeId, tag: &str) {
Self::remove_from(&mut self.by_tag, tag, id);
}
pub(crate) fn register_class(&mut self, id: NodeId, class: &str) {
Self::push_unique(self.by_class.entry(class.to_string()).or_default(), id);
}
pub(crate) fn unregister_class(&mut self, id: NodeId, class: &str) {
Self::remove_from(&mut self.by_class, class, id);
}
}
impl<Ext> Dom<Ext> {
pub(crate) fn hook_register(&mut self, id: NodeId) {
let Some(node) = self.get_node(id) else {
return;
};
let (tag, id_attr, classes) = match &node.data {
NodeData::Element {
tag,
attrs,
classes,
..
} => {
let tag = tag.clone();
let id_attr = attrs.get("id").cloned();
let classes: Vec<String> = classes.iter().cloned().collect();
(tag, id_attr, classes)
}
_ => return,
};
self.indexes.register_tag(id, &tag);
if let Some(v) = id_attr {
self.indexes.register_id(id, &v);
}
for c in classes {
self.indexes.register_class(id, &c);
}
}
pub(crate) fn hook_unregister(&mut self, id: NodeId) {
let Some(node) = self.get_node(id) else {
return;
};
let (tag, id_attr, classes) = match &node.data {
NodeData::Element {
tag,
attrs,
classes,
..
} => {
let tag = tag.clone();
let id_attr = attrs.get("id").cloned();
let classes: Vec<String> = classes.iter().cloned().collect();
(tag, id_attr, classes)
}
_ => return,
};
self.indexes.unregister_tag(id, &tag);
if let Some(v) = id_attr {
self.indexes.unregister_id(id, &v);
}
for c in classes {
self.indexes.unregister_class(id, &c);
}
}
pub fn get_element_by_id(&self, id_value: &str) -> Option<NodeId> {
let bucket = self.indexes.by_id.get(id_value)?;
if bucket.len() == 1 {
return bucket.iter().next().copied();
}
use crate::position::DocumentPosition;
let root = self.root();
let mut best: Option<NodeId> = None;
for &candidate in bucket {
let connected = self.ancestor_path(candidate).first() == Some(&root);
if !connected {
continue;
}
best = Some(match best {
None => candidate,
Some(b)
if self
.compare_document_position(b, candidate)
.contains(DocumentPosition::PRECEDING) =>
{
candidate
}
Some(b) => b,
});
}
best.or_else(|| bucket.iter().next().copied())
}
pub fn get_elements_by_tag_name_all(&self, tag: &str) -> Vec<NodeId> {
if tag == "*" {
let mut out: Vec<NodeId> = self
.indexes
.by_tag
.values()
.flat_map(|b| b.iter().copied())
.collect();
out.sort_unstable();
out
} else {
self.indexes
.by_tag
.get(tag)
.map(|b| b.iter().copied().collect())
.unwrap_or_default()
}
}
pub fn get_elements_by_class_name_all(&self, names: &str) -> Vec<NodeId> {
let wanted: Vec<&str> = names.split_ascii_whitespace().collect();
if wanted.is_empty() {
return self.get_elements_by_tag_name_all("*");
}
let mut buckets: Vec<&Bucket> = wanted
.iter()
.filter_map(|w| self.indexes.by_class.get(*w))
.collect();
if buckets.len() != wanted.len() {
return Vec::new(); }
buckets.sort_by_key(|b| b.len());
let smallest = buckets[0];
smallest
.iter()
.copied()
.filter(|id| buckets[1..].iter().all(|b| b.contains(id)))
.collect()
}
}
#[cfg(test)]
mod tests {
use crate::Dom;
#[test]
fn duplicate_ids_resolve_in_document_order() {
let mut dom: Dom = Dom::new();
let root = dom.root();
let later = dom.create_element("p");
dom.set_attribute(later, "id", "dup").unwrap();
let earlier = dom.create_element("p");
dom.set_attribute(earlier, "id", "dup").unwrap();
dom.append_child(root, later).unwrap();
dom.insert_before(root, earlier, Some(later)).unwrap();
assert_eq!(dom.get_element_by_id("dup"), Some(earlier));
dom.remove_child_dropping(root, earlier).unwrap();
assert_eq!(dom.get_element_by_id("dup"), Some(later));
}
#[test]
fn tag_index_stays_exact_across_many_nodes() {
let mut dom: Dom = Dom::new();
let root = dom.root();
let ids: Vec<_> = (0..20_000)
.map(|_| {
let d = dom.create_element("div");
dom.append_child(root, d).unwrap();
d
})
.collect();
assert_eq!(dom.get_elements_by_tag_name_all("div").len(), 20_000);
for id in ids {
dom.remove_child_dropping(root, id).unwrap();
}
assert!(dom.get_elements_by_tag_name_all("div").is_empty());
assert!(
!dom.indexes.by_tag.contains_key("div"),
"empty bucket is removed"
);
assert!(dom.validate().is_empty());
}
#[test]
fn id_index_populated_on_set_attribute() {
let mut dom: Dom = Dom::new();
let el = dom.create_element("div");
dom.set_attribute(el, "id", "main").unwrap();
assert_eq!(dom.get_element_by_id("main"), Some(el));
}
#[test]
fn id_index_unregisters_on_removal() {
let mut dom: Dom = Dom::new();
let el = dom.create_element("div");
dom.set_attribute(el, "id", "main").unwrap();
dom.remove_attribute(el, "id").unwrap();
assert_eq!(dom.get_element_by_id("main"), None);
}
#[test]
fn id_index_updates_on_reassignment() {
let mut dom: Dom = Dom::new();
let el = dom.create_element("div");
dom.set_attribute(el, "id", "old").unwrap();
dom.set_attribute(el, "id", "new").unwrap();
assert_eq!(dom.get_element_by_id("old"), None);
assert_eq!(dom.get_element_by_id("new"), Some(el));
}
#[test]
fn id_index_survives_node_drop() {
let mut dom: Dom = Dom::new();
let el = dom.create_element("div");
dom.set_attribute(el, "id", "main").unwrap();
let root = dom.root();
dom.append_child(root, el).unwrap();
dom.drop_subtree(el).unwrap();
assert_eq!(dom.get_element_by_id("main"), None);
}
#[test]
fn tag_index_finds_elements() {
let mut dom: Dom = Dom::new();
let a = dom.create_element("div");
let b = dom.create_element("div");
let c = dom.create_element("span");
let divs = dom.get_elements_by_tag_name_all("div");
assert!(divs.contains(&a));
assert!(divs.contains(&b));
assert!(!divs.contains(&c));
}
#[test]
fn tag_index_wildcard_returns_all() {
let mut dom: Dom = Dom::new();
let _ = dom.create_element("a");
let _ = dom.create_element("b");
let all = dom.get_elements_by_tag_name_all("*");
assert_eq!(all.len(), 2);
}
#[test]
fn tag_index_clears_on_free() {
let mut dom: Dom = Dom::new();
let a = dom.create_element("a");
assert_eq!(dom.get_elements_by_tag_name_all("a"), vec![a]);
let root = dom.root();
dom.append_child(root, a).unwrap();
dom.drop_subtree(a).unwrap();
assert!(dom.get_elements_by_tag_name_all("a").is_empty());
}
#[test]
fn class_index_basic() {
let mut dom: Dom = Dom::new();
let el = dom.create_element("div");
dom.add_class(el, "foo").unwrap();
assert_eq!(dom.get_elements_by_class_name_all("foo"), vec![el]);
}
#[test]
fn class_index_intersection() {
let mut dom: Dom = Dom::new();
let a = dom.create_element("div");
dom.add_class(a, "x").unwrap();
dom.add_class(a, "y").unwrap();
let b = dom.create_element("div");
dom.add_class(b, "x").unwrap(); let c = dom.create_element("div");
dom.add_class(c, "y").unwrap();
assert_eq!(dom.get_elements_by_class_name_all("x y"), vec![a]);
let xs = dom.get_elements_by_class_name_all("x");
assert!(xs.contains(&a) && xs.contains(&b));
}
#[test]
fn class_index_handles_toggle_and_replace() {
let mut dom: Dom = Dom::new();
let el = dom.create_element("div");
dom.add_class(el, "old").unwrap();
assert_eq!(dom.get_elements_by_class_name_all("old"), vec![el]);
dom.replace_class(el, "old", "new").unwrap();
assert!(dom.get_elements_by_class_name_all("old").is_empty());
assert_eq!(dom.get_elements_by_class_name_all("new"), vec![el]);
dom.toggle_class(el, "new").unwrap(); assert!(dom.get_elements_by_class_name_all("new").is_empty());
}
#[test]
fn id_attribute_via_set_id_sugar_indexed() {
let mut dom: Dom = Dom::new();
let el = dom.create_element("div");
dom.set_id(el, "hero").unwrap();
assert_eq!(dom.get_element_by_id("hero"), Some(el));
}
#[test]
fn freed_slot_reuse_does_not_leak_old_index_entries() {
let mut dom: Dom = Dom::new();
let a = dom.create_element("div");
dom.set_attribute(a, "id", "x").unwrap();
dom.add_class(a, "c").unwrap();
dom.free(a);
let b = dom.create_element("span");
assert_eq!(dom.get_element_by_id("x"), None);
assert!(dom.get_elements_by_class_name_all("c").is_empty());
assert_eq!(dom.get_elements_by_tag_name_all("span"), vec![b]);
assert!(dom.get_elements_by_tag_name_all("div").is_empty());
}
#[test]
fn duplicate_ids_first_wins() {
let mut dom: Dom = Dom::new();
let a = dom.create_element("div");
dom.set_attribute(a, "id", "dup").unwrap();
let b = dom.create_element("span");
dom.set_attribute(b, "id", "dup").unwrap();
assert_eq!(dom.get_element_by_id("dup"), Some(a));
dom.remove_attribute(a, "id").unwrap();
assert_eq!(dom.get_element_by_id("dup"), Some(b));
}
}