use html5ever::{LocalName, Namespace, QualName};
use std::cell::RefCell;
use std::collections::HashMap;
use std::fmt;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct NodeId(pub(crate) u32);
impl NodeId {
pub fn new(val: u32) -> Self {
NodeId(val)
}
pub fn index(self) -> usize {
self.0 as usize
}
pub fn raw(self) -> u32 {
self.0
}
}
impl fmt::Display for NodeId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "NodeId({})", self.0)
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Attribute {
pub name: QualName,
pub value: String,
}
#[derive(Clone, Debug)]
pub enum NodeData {
Document,
Doctype {
name: String,
public_id: String,
system_id: String,
},
Element {
name: QualName,
attrs: Vec<Attribute>,
template_contents: Option<NodeId>,
mathml_annotation_xml_integration_point: bool,
},
Text {
contents: String,
},
Comment {
contents: String,
},
ProcessingInstruction {
target: String,
data: String,
},
}
#[derive(Clone, Debug)]
pub struct Node {
pub id: NodeId,
pub parent: Option<NodeId>,
pub first_child: Option<NodeId>,
pub last_child: Option<NodeId>,
pub prev_sibling: Option<NodeId>,
pub next_sibling: Option<NodeId>,
pub data: NodeData,
}
impl Node {
pub fn is_document(&self) -> bool {
matches!(self.data, NodeData::Document)
}
pub fn is_element(&self) -> bool {
matches!(self.data, NodeData::Element { .. })
}
pub fn is_text(&self) -> bool {
matches!(self.data, NodeData::Text { .. })
}
pub fn as_element(&self) -> Option<&QualName> {
match &self.data {
NodeData::Element { name, .. } => Some(name),
_ => None,
}
}
pub fn attrs(&self) -> Option<&[Attribute]> {
match &self.data {
NodeData::Element { attrs, .. } => Some(attrs),
_ => None,
}
}
pub fn attrs_mut(&mut self) -> Option<&mut Vec<Attribute>> {
match &mut self.data {
NodeData::Element { attrs, .. } => Some(attrs),
_ => None,
}
}
pub fn get_attribute(&self, name: &str) -> Option<&str> {
self.attrs()?.iter().find_map(|a| {
if a.name.local.as_ref() == name {
Some(a.value.as_str())
} else {
None
}
})
}
pub fn set_attribute(&mut self, name: &str, value: String) {
if let NodeData::Element { attrs, .. } = &mut self.data {
if let Some(attr) = attrs.iter_mut().find(|a| a.name.local.as_ref() == name) {
attr.value = value;
} else {
attrs.push(Attribute {
name: QualName::new(None, Namespace::default(), LocalName::from(name)),
value,
});
}
}
}
pub fn text_content_of_text_node(&self) -> Option<&str> {
match &self.data {
NodeData::Text { contents } => Some(contents),
_ => None,
}
}
}
pub struct DomTree {
inner: RefCell<DomTreeInner>,
}
pub(crate) struct DomTreeInner {
pub(crate) nodes: Vec<Option<Node>>,
pub(crate) free_list: Vec<u32>,
pub(crate) document: NodeId,
pub(crate) id_index: HashMap<String, NodeId>,
}
impl DomTree {
pub fn new() -> Self {
let doc_node = Node {
id: NodeId(0),
parent: None,
first_child: None,
last_child: None,
prev_sibling: None,
next_sibling: None,
data: NodeData::Document,
};
DomTree {
inner: RefCell::new(DomTreeInner {
nodes: vec![Some(doc_node)],
free_list: Vec::new(),
document: NodeId(0),
id_index: HashMap::new(),
}),
}
}
pub fn document(&self) -> NodeId {
self.inner.borrow().document
}
pub(crate) fn borrow_inner(&self) -> std::cell::Ref<'_, DomTreeInner> {
self.inner.borrow()
}
pub fn new_node(&self, data: NodeData) -> NodeId {
let mut inner = self.inner.borrow_mut();
let id = if let Some(slot) = inner.free_list.pop() {
NodeId(slot)
} else {
let idx = inner.nodes.len() as u32;
inner.nodes.push(None);
NodeId(idx)
};
if let NodeData::Element { ref attrs, .. } = data
&& let Some(id_attr) = attrs.iter().find(|a| a.name.local.as_ref() == "id")
{
inner.id_index.insert(id_attr.value.clone(), id);
}
inner.nodes[id.index()] = Some(Node {
id,
parent: None,
first_child: None,
last_child: None,
prev_sibling: None,
next_sibling: None,
data,
});
id
}
pub fn get_node(&self, id: NodeId) -> Option<Node> {
self.inner.borrow().nodes.get(id.index())?.clone()
}
pub fn with_node<F, R>(&self, id: NodeId, f: F) -> Option<R>
where
F: FnOnce(&Node) -> R,
{
let inner = self.inner.borrow();
inner.nodes.get(id.index())?.as_ref().map(f)
}
pub fn with_node_mut<F, R>(&self, id: NodeId, f: F) -> Option<R>
where
F: FnOnce(&mut Node) -> R,
{
let mut inner = self.inner.borrow_mut();
inner.nodes.get_mut(id.index())?.as_mut().map(f)
}
pub fn append_child(&self, parent_id: NodeId, child_id: NodeId) {
self.detach(child_id);
let mut inner = self.inner.borrow_mut();
let old_last = inner
.nodes
.get(parent_id.index())
.and_then(|n| n.as_ref())
.and_then(|n| n.last_child);
if let Some(Some(child)) = inner.nodes.get_mut(child_id.index()) {
child.parent = Some(parent_id);
child.prev_sibling = old_last;
child.next_sibling = None;
}
if let Some(old_last_id) = old_last
&& let Some(Some(old_last_node)) = inner.nodes.get_mut(old_last_id.index())
{
old_last_node.next_sibling = Some(child_id);
}
if let Some(Some(parent)) = inner.nodes.get_mut(parent_id.index()) {
if parent.first_child.is_none() {
parent.first_child = Some(child_id);
}
parent.last_child = Some(child_id);
}
}
pub fn insert_before(&self, existing_id: NodeId, new_sibling_id: NodeId) {
let (parent_id, prev_id) = {
let inner = self.inner.borrow();
let node = match inner.nodes.get(existing_id.index()).and_then(|n| n.as_ref()) {
Some(n) => n,
None => return,
};
match node.parent {
Some(p) => (p, node.prev_sibling),
None => return,
}
};
self.detach(new_sibling_id);
let mut inner = self.inner.borrow_mut();
if let Some(Some(node)) = inner.nodes.get_mut(new_sibling_id.index()) {
node.parent = Some(parent_id);
node.prev_sibling = prev_id;
node.next_sibling = Some(existing_id);
}
if let Some(Some(node)) = inner.nodes.get_mut(existing_id.index()) {
node.prev_sibling = Some(new_sibling_id);
}
if let Some(prev) = prev_id {
if let Some(Some(node)) = inner.nodes.get_mut(prev.index()) {
node.next_sibling = Some(new_sibling_id);
}
} else if let Some(Some(parent)) = inner.nodes.get_mut(parent_id.index()) {
parent.first_child = Some(new_sibling_id);
}
}
pub fn detach(&self, node_id: NodeId) {
let mut inner = self.inner.borrow_mut();
let (parent_id, prev_id, next_id) = match inner.nodes.get(node_id.index()).and_then(|n| n.as_ref()) {
Some(node) => (node.parent, node.prev_sibling, node.next_sibling),
None => return,
};
if let Some(prev) = prev_id
&& let Some(Some(node)) = inner.nodes.get_mut(prev.index())
{
node.next_sibling = next_id;
} else if let Some(parent_id) = parent_id
&& let Some(Some(parent)) = inner.nodes.get_mut(parent_id.index())
{
parent.first_child = next_id;
}
if let Some(next) = next_id
&& let Some(Some(node)) = inner.nodes.get_mut(next.index())
{
node.prev_sibling = prev_id;
} else if let Some(parent_id) = parent_id
&& let Some(Some(parent)) = inner.nodes.get_mut(parent_id.index())
{
parent.last_child = prev_id;
}
if let Some(Some(node)) = inner.nodes.get_mut(node_id.index()) {
node.parent = None;
node.prev_sibling = None;
node.next_sibling = None;
}
}
pub fn remove(&self, node_id: NodeId) {
self.detach(node_id);
let descendants = self.descendants(node_id);
let mut inner = self.inner.borrow_mut();
let mut ids_to_remove = Vec::new();
for &desc_id in &descendants {
if let Some(Some(node)) = inner.nodes.get(desc_id.index())
&& let Some(id_val) = node.get_attribute("id")
{
ids_to_remove.push(id_val.to_string());
}
}
if let Some(Some(node)) = inner.nodes.get(node_id.index())
&& let Some(id_val) = node.get_attribute("id")
{
ids_to_remove.push(id_val.to_string());
}
for id_str in ids_to_remove {
inner.id_index.remove(&id_str);
}
for desc_id in descendants {
inner.nodes[desc_id.index()] = None;
inner.free_list.push(desc_id.0);
}
inner.nodes[node_id.index()] = None;
inner.free_list.push(node_id.0);
}
pub fn children(&self, node_id: NodeId) -> Vec<NodeId> {
let inner = self.inner.borrow();
let mut result = Vec::new();
let mut current = inner
.nodes
.get(node_id.index())
.and_then(|n| n.as_ref())
.and_then(|n| n.first_child);
while let Some(child_id) = current {
result.push(child_id);
current = inner
.nodes
.get(child_id.index())
.and_then(|n| n.as_ref())
.and_then(|n| n.next_sibling);
}
result
}
pub fn descendants(&self, node_id: NodeId) -> Vec<NodeId> {
let inner = self.inner.borrow();
let mut result = Vec::new();
let mut stack = Vec::new();
let mut first = inner
.nodes
.get(node_id.index())
.and_then(|n| n.as_ref())
.and_then(|n| n.first_child);
let mut children_to_push = Vec::new();
while let Some(child_id) = first {
children_to_push.push(child_id);
first = inner
.nodes
.get(child_id.index())
.and_then(|n| n.as_ref())
.and_then(|n| n.next_sibling);
}
for child_id in children_to_push.into_iter().rev() {
stack.push(child_id);
}
while let Some(current) = stack.pop() {
result.push(current);
let mut child = inner
.nodes
.get(current.index())
.and_then(|n| n.as_ref())
.and_then(|n| n.first_child);
let mut children_to_push = Vec::new();
while let Some(child_id) = child {
children_to_push.push(child_id);
child = inner
.nodes
.get(child_id.index())
.and_then(|n| n.as_ref())
.and_then(|n| n.next_sibling);
}
for child_id in children_to_push.into_iter().rev() {
stack.push(child_id);
}
}
result
}
pub fn ancestors(&self, node_id: NodeId) -> Vec<NodeId> {
let inner = self.inner.borrow();
let mut result = Vec::new();
let mut current = inner
.nodes
.get(node_id.index())
.and_then(|n| n.as_ref())
.and_then(|n| n.parent);
while let Some(parent_id) = current {
result.push(parent_id);
current = inner
.nodes
.get(parent_id.index())
.and_then(|n| n.as_ref())
.and_then(|n| n.parent);
}
result
}
pub fn get_element_by_id(&self, id: &str) -> Option<NodeId> {
self.inner.borrow().id_index.get(id).copied()
}
pub fn text_content(&self, node_id: NodeId) -> String {
let inner = self.inner.borrow();
let mut result = String::new();
collect_text_inner(&inner, node_id, &mut result);
result
}
pub fn append_text(&self, parent_id: NodeId, text: &str) {
let last_child_is_text = {
let inner = self.inner.borrow();
inner
.nodes
.get(parent_id.index())
.and_then(|n| n.as_ref())
.and_then(|n| n.last_child)
.and_then(|lc| inner.nodes.get(lc.index()))
.and_then(|n| n.as_ref())
.map(|n| n.is_text())
.unwrap_or(false)
};
if last_child_is_text {
let last_child_id = {
let inner = self.inner.borrow();
inner
.nodes
.get(parent_id.index())
.and_then(|n| n.as_ref())
.and_then(|n| n.last_child)
.unwrap()
};
let mut inner = self.inner.borrow_mut();
if let Some(Some(node)) = inner.nodes.get_mut(last_child_id.index())
&& let NodeData::Text { contents } = &mut node.data
{
contents.push_str(text);
return;
}
}
let text_id = self.new_node(NodeData::Text {
contents: text.to_string(),
});
self.append_child(parent_id, text_id);
}
pub fn find_body_or_root(&self) -> NodeId {
let doc = self.document();
for child in self.children(doc) {
if let Some(n) = self.get_node(child)
&& n.as_element()
.map(|name| name.local.as_ref() == "html")
.unwrap_or(false)
{
for html_child in self.children(child) {
if let Some(hc) = self.get_node(html_child)
&& hc
.as_element()
.map(|name| name.local.as_ref() == "body")
.unwrap_or(false)
{
return html_child;
}
}
return child;
}
}
doc
}
pub fn import_children_from(&self, parent_id: NodeId, source: &DomTree, source_node: NodeId) {
let source_children = source.children(source_node);
for source_child_id in source_children {
self.import_node_from(parent_id, source, source_child_id);
}
}
fn import_node_from(&self, parent_id: NodeId, source: &DomTree, source_node_id: NodeId) {
let node_data = {
let source_inner = source.inner.borrow();
match source_inner.nodes.get(source_node_id.index()) {
Some(Some(node)) => node.data.clone(),
_ => return,
}
};
let new_id = self.new_node(node_data);
self.append_child(parent_id, new_id);
let children = source.children(source_node_id);
for child_id in children {
self.import_node_from(new_id, source, child_id);
}
}
pub fn len(&self) -> usize {
self.inner.borrow().nodes.iter().filter(|n| n.is_some()).count()
}
pub fn is_empty(&self) -> bool {
self.len() <= 1
}
pub fn update_id_index(&self, node_id: NodeId, old_id: Option<&str>, new_id: Option<&str>) {
let mut inner = self.inner.borrow_mut();
if let Some(old) = old_id {
inner.id_index.remove(old);
}
if let Some(new) = new_id {
inner.id_index.insert(new.to_string(), node_id);
}
}
}
fn collect_text_inner(inner: &DomTreeInner, node_id: NodeId, buf: &mut String) {
if let Some(Some(node)) = inner.nodes.get(node_id.index()) {
match &node.data {
NodeData::Text { contents } => buf.push_str(contents),
_ => {
let mut child = node.first_child;
while let Some(child_id) = child {
collect_text_inner(inner, child_id, buf);
child = inner
.nodes
.get(child_id.index())
.and_then(|n| n.as_ref())
.and_then(|n| n.next_sibling);
}
}
}
}
}
impl Default for DomTree {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new_tree_has_document() {
let tree = DomTree::new();
assert_eq!(tree.len(), 1);
let node = tree.get_node(tree.document()).unwrap();
assert!(node.is_document());
}
#[test]
fn test_append_child() {
let tree = DomTree::new();
let child = tree.new_node(NodeData::Text {
contents: "hello".into(),
});
let doc = tree.document();
tree.append_child(doc, child);
assert_eq!(tree.len(), 2);
let doc_node = tree.get_node(doc).unwrap();
assert_eq!(doc_node.first_child, Some(child));
assert_eq!(doc_node.last_child, Some(child));
let child_node = tree.get_node(child).unwrap();
assert_eq!(child_node.parent, Some(doc));
}
#[test]
fn test_multiple_children() {
let tree = DomTree::new();
let doc = tree.document();
let c1 = tree.new_node(NodeData::Text { contents: "a".into() });
let c2 = tree.new_node(NodeData::Text { contents: "b".into() });
let c3 = tree.new_node(NodeData::Text { contents: "c".into() });
tree.append_child(doc, c1);
tree.append_child(doc, c2);
tree.append_child(doc, c3);
assert_eq!(tree.children(doc), vec![c1, c2, c3]);
}
#[test]
fn test_detach() {
let tree = DomTree::new();
let doc = tree.document();
let c1 = tree.new_node(NodeData::Text { contents: "a".into() });
let c2 = tree.new_node(NodeData::Text { contents: "b".into() });
tree.append_child(doc, c1);
tree.append_child(doc, c2);
tree.detach(c1);
assert_eq!(tree.children(doc), vec![c2]);
}
#[test]
fn test_insert_before() {
let tree = DomTree::new();
let doc = tree.document();
let c1 = tree.new_node(NodeData::Text { contents: "a".into() });
let c2 = tree.new_node(NodeData::Text { contents: "b".into() });
let c3 = tree.new_node(NodeData::Text { contents: "c".into() });
tree.append_child(doc, c1);
tree.append_child(doc, c3);
tree.insert_before(c3, c2);
assert_eq!(tree.children(doc), vec![c1, c2, c3]);
}
#[test]
fn test_text_content() {
let tree = DomTree::new();
let doc = tree.document();
let div = tree.new_node(NodeData::Element {
name: QualName::new(None, ns!(html), local_name!("div")),
attrs: vec![],
template_contents: None,
mathml_annotation_xml_integration_point: false,
});
tree.append_child(doc, div);
let t1 = tree.new_node(NodeData::Text {
contents: "Hello ".into(),
});
let t2 = tree.new_node(NodeData::Text {
contents: "World".into(),
});
tree.append_child(div, t1);
tree.append_child(div, t2);
assert_eq!(tree.text_content(div), "Hello World");
}
#[test]
fn test_get_element_by_id() {
let tree = DomTree::new();
let doc = tree.document();
let div = tree.new_node(NodeData::Element {
name: QualName::new(None, ns!(html), local_name!("div")),
attrs: vec![Attribute {
name: QualName::new(None, Namespace::default(), LocalName::from("id")),
value: "main".into(),
}],
template_contents: None,
mathml_annotation_xml_integration_point: false,
});
tree.append_child(doc, div);
assert_eq!(tree.get_element_by_id("main"), Some(div));
assert_eq!(tree.get_element_by_id("nonexistent"), None);
}
#[test]
fn test_append_text_merges() {
let tree = DomTree::new();
let doc = tree.document();
tree.append_text(doc, "Hello ");
tree.append_text(doc, "World");
assert_eq!(tree.children(doc).len(), 1);
assert_eq!(tree.text_content(doc), "Hello World");
}
#[test]
fn test_remove_subtree() {
let tree = DomTree::new();
let doc = tree.document();
let div = tree.new_node(NodeData::Element {
name: QualName::new(None, ns!(html), local_name!("div")),
attrs: vec![],
template_contents: None,
mathml_annotation_xml_integration_point: false,
});
tree.append_child(doc, div);
let text = tree.new_node(NodeData::Text { contents: "hi".into() });
tree.append_child(div, text);
assert_eq!(tree.len(), 3);
tree.remove(div);
assert_eq!(tree.len(), 1);
}
}