pochoir-parser 0.12.2

#![allow(clippy::panic)]

use pochoir_common::Spanned;
use self_cell::self_cell;
use std::{
    collections::BTreeMap,
    ops::{Add, AddAssign},
    path::{Path, PathBuf},
};

use crate::{Node, ParsedNode};

mod selection;
mod traverse;
mod tree_ref;

pub use selection::build_selector_list;
pub use tree_ref::{TreeRef, TreeRefMut};

self_cell! {
    /// An owned version of a [`Tree`] containing the HTML data used to build it.
    ///
    /// This structure uses [self_cell](https://docs.rs/self_cell) to build a self-referencing structure without a lifetime.
    pub struct OwnedTree {
        owner: String,

        #[covariant]
        dependent: Tree,
    }

    impl { Debug }
}

impl OwnedTree {
    /// Get the inner HTML of the tree.
    pub fn get_data(&self) -> &str {
        self.borrow_owner()
    }

    /// Get an immutable reference to the inner tree.
    pub fn get_tree(&self) -> &Tree<'_> {
        self.borrow_dependent()
    }

    /// Mutate the tree temporarily by executing the given function.
    pub fn mutate<Return>(&mut self, func: impl FnOnce(&mut Tree) -> Return) -> Return {
        self.with_dependent_mut(|_, tree| func(tree))
    }
}

#[derive(Debug, PartialEq, Eq, Clone, Copy, PartialOrd, Ord, Hash)]
pub enum TreeRefId {
    Root,
    Node(usize),
}

impl Add for TreeRefId {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        match (self, rhs) {
            (_, Self::Root) | (Self::Root, _) => Self::Root,
            (Self::Node(a), Self::Node(b)) => Self::Node(a + b),
        }
    }
}

impl AddAssign<Self> for TreeRefId {
    fn add_assign(&mut self, rhs: Self) {
        match (self, rhs) {
            (Self::Root | Self::Node(_), Self::Root) | (Self::Root, Self::Node(_)) => (),
            (Self::Node(a), Self::Node(b)) => *a += b,
        }
    }
}

impl Add<usize> for TreeRefId {
    type Output = Self;

    fn add(self, rhs: usize) -> Self::Output {
        match self {
            Self::Root => Self::Root,
            Self::Node(a) => Self::Node(a + rhs),
        }
    }
}

impl AddAssign<usize> for TreeRefId {
    fn add_assign(&mut self, rhs: usize) {
        match self {
            Self::Root => (),
            Self::Node(a) => *a += rhs,
        }
    }
}

#[derive(Debug, Clone)]
struct TreeNode<'a> {
    data: ParsedNode<'a>,
    parent: TreeRefId,
    children: Vec<TreeRefId>,
}

#[derive(Debug, Clone)]
pub struct Tree<'a> {
    file_path: PathBuf,
    nodes: BTreeMap<TreeRefId, TreeNode<'a>>,
    next_id: TreeRefId,
}

impl<'a> Tree<'a> {
    /// Create an empty [`Tree`].
    pub fn new<P: AsRef<Path>>(file_path: P) -> Self {
        Self {
            file_path: file_path.as_ref().into(),
            nodes: BTreeMap::from_iter([(
                TreeRefId::Root,
                TreeNode {
                    data: Spanned::new(Node::Root),
                    parent: TreeRefId::Root,
                    children: vec![],
                },
            )]),
            next_id: TreeRefId::Node(0),
        }
    }

    pub fn file_path(&self) -> &Path {
        &self.file_path
    }

    pub fn next_id(&self) -> TreeRefId {
        self.next_id
    }

    /// Insert a new node in the tree as a child of `parent`.
    ///
    /// # Panics
    ///
    /// This function panics if the provided parent ID does not exist in the tree. It may happen if
    /// the ID is wrong, does not exist in the tree or if the node having this ID was removed.
    pub fn insert(&mut self, parent: TreeRefId, data: ParsedNode<'a>) -> TreeRefId {
        let id = self.next_id;
        self.nodes.insert(
            id,
            TreeNode {
                data,
                parent,
                children: vec![],
            },
        );

        // Update parent
        self.nodes
            .get_mut(&parent)
            .expect("failed to find parent node in tree")
            .children
            .push(id);

        self.next_id += 1;
        id
    }

    fn insert_with_id(
        &mut self,
        id: TreeRefId,
        parent: TreeRefId,
        children: Vec<TreeRefId>,
        data: ParsedNode<'a>,
    ) {
        self.nodes.insert(
            id,
            TreeNode {
                data,
                parent,
                children,
            },
        );
    }

    /// Get the list of all the tree nodes.
    ///
    /// Removed nodes won't be in the list.
    pub fn all_nodes(&self) -> Vec<TreeRefId> {
        self.nodes.iter().map(|n| *n.0).collect()
    }

    /// Get the list of the tree root nodes.
    pub fn root_nodes(&self) -> Vec<TreeRefId> {
        self.nodes
            .iter()
            .filter(|n| *n.0 != TreeRefId::Root && n.1.parent == TreeRefId::Root)
            .map(|n| *n.0)
            .collect()
    }

    /// Get a reference to a tree node or panic if the ID does not exist.
    ///
    /// For a non-panicking version, see [`Tree::try_get`].
    ///
    /// # Panics
    ///
    /// This function panics if the provided ID does not exist in the tree. It may happen if the
    /// node ID is hard-coded and is not yet present in the tree. Keep in mind that removing a node
    /// consumes the reference so this function is very unlikely to panic.
    pub fn get(&self, id: TreeRefId) -> TreeRef<'a, '_> {
        self.try_get(id).unwrap_or_else(|| {
            panic!("failed to get a reference to the tree node with ID {id:?}");
        })
    }

    /// Get a mutable reference to a tree node or panic if the ID does not exist.
    ///
    /// For a non-panicking version, see [`Tree::try_get_mut`].
    ///
    /// # Panics
    ///
    /// This function panics if the provided ID does not exist in the tree. It may happen if the
    /// node ID is hard-coded and is not yet present in the tree. Keep in mind that removing a node
    /// consumes the reference so this function is very unlikely to panic.
    pub fn get_mut(&mut self, id: TreeRefId) -> TreeRefMut<'a, '_> {
        self.try_get_mut(id).unwrap_or_else(|| {
            panic!("failed to get a mutable reference to the tree node with ID {id:?}");
        })
    }

    /// Try to get a tree reference from a node ID.
    ///
    /// Returns `None` if the ID is not present in tree.
    pub fn try_get(&self, id: TreeRefId) -> Option<TreeRef<'a, '_>> {
        if self.nodes.contains_key(&id) {
            Some(TreeRef { id, tree: self })
        } else {
            None
        }
    }

    /// Try to get a mutable tree reference from a node ID.
    ///
    /// Returns `None` if the ID is not present in tree.
    pub fn try_get_mut(&mut self, id: TreeRefId) -> Option<TreeRefMut<'a, '_>> {
        if self.nodes.contains_key(&id) {
            Some(TreeRefMut { id, tree: self })
        } else {
            None
        }
    }

    pub fn traverse_breadth(&self) -> impl Iterator<Item = TreeRef<'a, '_>> {
        let parent = self.get(TreeRefId::Root);

        traverse::TraverseBreadthIter {
            children: parent.children().collect(),
            index: 0,
        }
    }

    pub fn traverse_depth(&self) -> impl Iterator<Item = TreeRef<'a, '_>> {
        let mut queue: Vec<TreeRef> = self.get(TreeRefId::Root).children().collect();
        queue.reverse();

        traverse::TraverseDepthIter { queue }
    }

    /// Search the first node matching the CSS selector in the tree, relatively to this node.
    ///
    /// The traversal order used is the same as in the [official DOM specification](https://dom.spec.whatwg.org/#concept-tree-order)
    /// ie preorder, depth-first.
    ///
    /// It is similar to the `document.querySelector` function in JavaScript.
    ///
    /// # Errors
    ///
    /// Returns [`selectors::parser::SelectorParseError`] if the CSS selector is not a valid CSS selector
    /// or `Ok(None)` if no node matching the CSS selector is found in the tree
    pub fn select<'b>(
        &self,
        selector: &'b str,
    ) -> Result<Option<TreeRefId>, crate::error::SelectorParseError<'b>> {
        let selector_list =
            build_selector_list(selector)?;
        Ok(self.traverse_depth()
            .find(|tree_ref| tree_ref.is_matching(&selector_list))
            .map(|tree_ref| tree_ref.id()))
    }

    /// Select several nodes in the tree using a CSS selector.
    ///
    /// It is similar to the `document.querySelectorAll` function in JavaScript.
    ///
    /// Returns an empty `Vec` if the selector is not a valid CSS selector.
    pub fn select_all(&self, selector: &str) -> Vec<TreeRefId> {
        if let Ok(selector_list) = build_selector_list(selector) {
            self.traverse_depth()
                .filter(|tree_ref| tree_ref.is_matching(&selector_list))
                .map(|tree_ref| tree_ref.id())
                .collect()
        } else {
            vec![]
        }
    }
}

#[cfg(test)]
#[allow(clippy::similar_names)]
mod tests {
    use pochoir_template_engine::TemplateBlock;
    use std::borrow::Cow;

    use crate::Attrs;

    use super::*;

    #[test]
    #[allow(clippy::similar_names)]
    fn make_tree() {
        let mut tree = Tree::new("index.html");

        let container_node = Node::Element(Cow::Borrowed("div"), Attrs::new());
        let text_node = Node::TemplateBlock(TemplateBlock::RawText(Cow::Borrowed("Hello world!")));
        let text2_node =
            Node::TemplateBlock(TemplateBlock::RawText(Cow::Borrowed("Hello world 2!")));
        let link_node = Node::Element(
            Cow::Borrowed("a"),
            Attrs::from_iter([(
                Cow::Borrowed("href"),
                vec![Spanned::new(TemplateBlock::text("https://example.com"))],
            )]),
        );

        let container_id = tree.insert(TreeRefId::Root, Spanned::new(container_node.clone()));
        let text_id = tree.insert(container_id, Spanned::new(text_node.clone()));
        let text2_id = tree.insert(container_id, Spanned::new(text2_node.clone()));
        let link_id = tree.insert(container_id, Spanned::new(link_node.clone()));

        let container = tree.get(container_id);
        let text = tree.get(text_id); // Test relations
        let text2 = tree.get(text2_id);
        let link = tree.get(link_id);

        // Test data
        assert_eq!(*container.data(), container_node);
        assert_eq!(*text.data(), text_node);
        assert_eq!(*text2.data(), text2_node);
        assert_eq!(*link.data(), link_node);

        // Test relations
        assert_eq!(container.parent(), tree.get(TreeRefId::Root));
        assert_eq!(text.parent(), tree.get(container_id));
        assert_eq!(text2.parent(), tree.get(container_id));
        assert_eq!(link.parent(), tree.get(container_id));

        assert_eq!(container.prev_sibling(), None);
        assert_eq!(text.prev_sibling(), None);
        assert_eq!(text2.prev_sibling(), Some(tree.get(text_id)));
        assert_eq!(link.prev_sibling(), Some(tree.get(text2_id)));

        assert_eq!(container.next_sibling(), None);
        assert_eq!(text.next_sibling(), Some(tree.get(text2_id)));
        assert_eq!(text2.next_sibling(), Some(tree.get(link_id)));
        assert_eq!(link.next_sibling(), None);

        assert_eq!(container.children().collect::<Vec<TreeRef>>(), vec![text, text2, link]);
        assert_eq!(text.children().collect::<Vec<TreeRef>>(), vec![]);
        assert_eq!(text2.children().collect::<Vec<TreeRef>>(), vec![]);
        assert_eq!(link.children().collect::<Vec<TreeRef>>(), vec![]);
    }

    #[test]
    fn traverse_tree() {
        let mut tree = Tree::new("index.html");

        let container_node = Node::Element(Cow::Borrowed("div"), Attrs::new());
        let text_node = Node::TemplateBlock(TemplateBlock::RawText(Cow::Borrowed("Hello world!")));
        let text2_node =
            Node::TemplateBlock(TemplateBlock::RawText(Cow::Borrowed("Hello world 2!")));
        let link_node = Node::Element(
            Cow::Borrowed("a"),
            Attrs::from_iter([(
                Cow::Borrowed("href"),
                vec![Spanned::new(TemplateBlock::text("https://example.com"))],
            )]),
        );

        let container_id = tree.insert(TreeRefId::Root, Spanned::new(container_node.clone()));
        let link_id = tree.insert(container_id, Spanned::new(link_node.clone()));
        let text_id = tree.insert(container_id, Spanned::new(text_node.clone()));
        let text2_id = tree.insert(link_id, Spanned::new(text2_node.clone()));

        let ids: Vec<TreeRefId> = tree
            .traverse_breadth()
            .map(|tree_ref| tree_ref.id())
            .collect();
        assert_eq!(ids, vec![container_id, link_id, text_id, text2_id]);

        let ids: Vec<TreeRefId> = tree
            .traverse_depth()
            .map(|tree_ref| tree_ref.id())
            .collect();
        assert_eq!(ids, vec![container_id, link_id, text2_id, text_id]);
    }

    #[test]
    fn select() {
        let mut tree = Tree::new("index.html");

        let container_node = Node::Element(Cow::Borrowed("div"), Attrs::new());
        let link_node = Node::Element(
            Cow::Borrowed("a"),
            Attrs::from_iter([
                (
                    Cow::Borrowed("href"),
                    vec![Spanned::new(TemplateBlock::text("https://example.com"))],
                ),
                (
                    Cow::Borrowed("class"),
                    vec![Spanned::new(TemplateBlock::text("link"))],
                ),
            ]),
        );

        let container_id = tree.insert(TreeRefId::Root, Spanned::new(container_node.clone()));
        let link_id = tree.insert(container_id, Spanned::new(link_node.clone()));
        let container_cloned_id = tree.insert(container_id, Spanned::new(container_node.clone()));
        let link_cloned_id = tree.insert(container_cloned_id, Spanned::new(link_node.clone()));

        let container_cloned = tree.get(container_cloned_id);
        let link_cloned = tree.get(link_cloned_id);

        assert_eq!(tree.select("a.link").unwrap(), Some(link_id));
        assert_eq!(tree.select("div").unwrap(), Some(container_id));
        assert_eq!(tree.select("div > div").unwrap(), Some(container_cloned_id));
        assert_eq!(container_cloned.select(".link").unwrap(), Some(link_cloned));
        assert_eq!(
            tree.get(tree.select("div > div").unwrap().unwrap()).select(".link").unwrap(),
            Some(link_cloned)
        );
        assert_eq!(tree.select_all(".link"), vec![link_id, link_cloned_id]);
    }

    #[test]
    fn update_tree() {
        let mut tree = Tree::new("index.html");

        let container_node = Node::Element(Cow::Borrowed("div"), Attrs::new());
        let link_node = Node::Element(
            Cow::Borrowed("a"),
            Attrs::from_iter([
                (
                    Cow::Borrowed("href"),
                    vec![Spanned::new(TemplateBlock::text("https://example.com"))],
                ),
                (
                    Cow::Borrowed("class"),
                    vec![Spanned::new(TemplateBlock::text("link"))],
                ),
            ]),
        );

        let container_id = tree.insert(TreeRefId::Root, Spanned::new(container_node.clone()));
        let link_id = tree.insert(container_id, Spanned::new(link_node.clone()));
        let container_cloned_id = tree.insert(container_id, Spanned::new(container_node.clone()));
        let link_cloned_id = tree.insert(container_cloned_id, Spanned::new(link_node.clone()));

        tree.get_mut(container_cloned_id).remove();

        assert_eq!(tree.try_get(container_cloned_id), None);
        assert_eq!(tree.try_get(link_cloned_id), None);
        assert_eq!(tree.get(container_id).children().collect::<Vec<TreeRef>>(), vec![tree.get(link_id)]);

        let container2_id = tree.insert(container_id, Spanned::new(container_node));
        let _link2_id = tree.insert(container2_id, Spanned::new(link_node));

        let container = tree.get(container_id);

        let removed_nodes: Vec<TreeRefId> = container
            .traverse_depth()
            .filter(|tree_ref| tree_ref.attr("class") == Ok(Some("link".to_string())))
            .map(|tree_ref| tree_ref.id())
            .collect();

        for id in removed_nodes {
            tree.get_mut(id).remove();
        }

        assert_eq!(tree.get(container_id).traverse_breadth().filter(|tree_ref| matches!(&tree_ref.data(), Node::Element(_, attrs) if attrs.get("class") == Some(&vec![Spanned::new(TemplateBlock::text("link"))]))).count(), 0);
    }

    #[test]
    fn sub_tree() {
        let mut tree = Tree::new("index.html");

        let container_node = Node::Element(Cow::Borrowed("div"), Attrs::new());
        let text_node = Node::TemplateBlock(TemplateBlock::RawText(Cow::Borrowed("Hello world!")));
        let text2_node =
            Node::TemplateBlock(TemplateBlock::RawText(Cow::Borrowed("Hello world 2!")));
        let link_node = Node::Element(
            Cow::Borrowed("a"),
            Attrs::from_iter([(
                Cow::Borrowed("href"),
                vec![Spanned::new(TemplateBlock::text("https://example.com"))],
            )]),
        );

        let container_id = tree.insert(TreeRefId::Root, Spanned::new(container_node.clone()));
        let text_id = tree.insert(container_id, Spanned::new(text_node.clone()));
        let link_id = tree.insert(container_id, Spanned::new(link_node.clone()));
        let text2_id = tree.insert(link_id, Spanned::new(text2_node.clone()));

        let link = tree.get(link_id);
        let sub_tree = link.sub_tree();
        assert_eq!(sub_tree.try_get(container_id), None);
        assert_eq!(sub_tree.try_get(text_id), None);
        assert_eq!(sub_tree.get(text2_id).data(), &text2_node);
    }
}