markdown_ppp/ast_transform/

query.rs

//! Query API for finding elements in AST
//!
//! This module provides query methods for finding elements in the AST that match certain conditions.
//!
//! # Example
//!
//! ```rust
//! use markdown_ppp::ast::*;
//! use markdown_ppp::ast_transform::Query;
//!
//! let doc = Document {
//!     blocks: vec![
//!         Block::Paragraph(vec![
//!             Inline::Text("hello".to_string()),
//!             Inline::Autolink("https://example.com".to_string()),
//!         ]),
//!     ],
//! };
//!
//! // Find all autolinks
//! let autolinks = doc.find_all_inlines(|inline| {
//!     matches!(inline, Inline::Autolink(_))
//! });
//! assert_eq!(autolinks.len(), 1);
//!
//! // Count text nodes
//! let text_count = doc.count_inlines(|inline| {
//!     matches!(inline, Inline::Text(_))
//! });
//! assert_eq!(text_count, 1);
//! ```

use crate::ast::*;

/// Query trait for finding elements in AST structures
pub trait Query {
    /// Find all inline elements matching a predicate
    fn find_all_inlines<F>(&self, predicate: F) -> Vec<&Inline>
    where
        F: Fn(&Inline) -> bool;

    /// Find all block elements matching a predicate
    fn find_all_blocks<F>(&self, predicate: F) -> Vec<&Block>
    where
        F: Fn(&Block) -> bool;

    /// Find the first inline element matching a predicate
    fn find_first_inline<F>(&self, predicate: F) -> Option<&Inline>
    where
        F: Fn(&Inline) -> bool;

    /// Find the first block element matching a predicate
    fn find_first_block<F>(&self, predicate: F) -> Option<&Block>
    where
        F: Fn(&Block) -> bool;

    /// Count inline elements matching a predicate
    fn count_inlines<F>(&self, predicate: F) -> usize
    where
        F: Fn(&Inline) -> bool,
    {
        self.find_all_inlines(predicate).len()
    }

    /// Count block elements matching a predicate
    fn count_blocks<F>(&self, predicate: F) -> usize
    where
        F: Fn(&Block) -> bool,
    {
        self.find_all_blocks(predicate).len()
    }

    /// Check if any inline element matches a predicate
    fn any_inline<F>(&self, predicate: F) -> bool
    where
        F: Fn(&Inline) -> bool,
    {
        self.find_first_inline(predicate).is_some()
    }

    /// Check if any block element matches a predicate
    fn any_block<F>(&self, predicate: F) -> bool
    where
        F: Fn(&Block) -> bool,
    {
        self.find_first_block(predicate).is_some()
    }

    /// Find all links in the document
    fn find_all_links(&self) -> Vec<&Link> {
        self.find_all_inlines(|inline| matches!(inline, Inline::Link(_)))
            .into_iter()
            .filter_map(|inline| match inline {
                Inline::Link(link) => Some(link),
                _ => None,
            })
            .collect()
    }

    /// Find all images in the document
    fn find_all_images(&self) -> Vec<&Image> {
        self.find_all_inlines(|inline| matches!(inline, Inline::Image(_)))
            .into_iter()
            .filter_map(|inline| match inline {
                Inline::Image(image) => Some(image),
                _ => None,
            })
            .collect()
    }

    /// Find all headings in the document
    fn find_all_headings(&self) -> Vec<&Heading> {
        self.find_all_blocks(|block| matches!(block, Block::Heading(_)))
            .into_iter()
            .filter_map(|block| match block {
                Block::Heading(heading) => Some(heading),
                _ => None,
            })
            .collect()
    }

    /// Find all autolinks in the document
    fn find_all_autolinks(&self) -> Vec<&str> {
        self.find_all_inlines(|inline| matches!(inline, Inline::Autolink(_)))
            .into_iter()
            .filter_map(|inline| match inline {
                Inline::Autolink(url) => Some(url.as_str()),
                _ => None,
            })
            .collect()
    }

    /// Find all text nodes in the document
    fn find_all_text(&self) -> Vec<&str> {
        self.find_all_inlines(|inline| matches!(inline, Inline::Text(_)))
            .into_iter()
            .filter_map(|inline| match inline {
                Inline::Text(text) => Some(text.as_str()),
                _ => None,
            })
            .collect()
    }

    /// Find all code spans in the document
    fn find_all_code_spans(&self) -> Vec<&str> {
        self.find_all_inlines(|inline| matches!(inline, Inline::Code(_)))
            .into_iter()
            .filter_map(|inline| match inline {
                Inline::Code(code) => Some(code.as_str()),
                _ => None,
            })
            .collect()
    }

    /// Find all code blocks in the document
    fn find_all_code_blocks(&self) -> Vec<&CodeBlock> {
        self.find_all_blocks(|block| matches!(block, Block::CodeBlock(_)))
            .into_iter()
            .filter_map(|block| match block {
                Block::CodeBlock(code_block) => Some(code_block),
                _ => None,
            })
            .collect()
    }

    /// Find all tables in the document
    fn find_all_tables(&self) -> Vec<&Table> {
        self.find_all_blocks(|block| matches!(block, Block::Table(_)))
            .into_iter()
            .filter_map(|block| match block {
                Block::Table(table) => Some(table),
                _ => None,
            })
            .collect()
    }

    /// Find all lists in the document
    fn find_all_lists(&self) -> Vec<&List> {
        self.find_all_blocks(|block| matches!(block, Block::List(_)))
            .into_iter()
            .filter_map(|block| match block {
                Block::List(list) => Some(list),
                _ => None,
            })
            .collect()
    }
}

impl Query for Document {
    fn find_all_inlines<F>(&self, predicate: F) -> Vec<&Inline>
    where
        F: Fn(&Inline) -> bool,
    {
        let mut results = Vec::new();
        for block in &self.blocks {
            results.extend(block.find_all_inlines(&predicate));
        }
        results
    }

    fn find_all_blocks<F>(&self, predicate: F) -> Vec<&Block>
    where
        F: Fn(&Block) -> bool,
    {
        let mut results = Vec::new();
        for block in &self.blocks {
            results.extend(block.find_all_blocks(&predicate));
        }
        results
    }

    fn find_first_inline<F>(&self, predicate: F) -> Option<&Inline>
    where
        F: Fn(&Inline) -> bool,
    {
        for block in &self.blocks {
            if let Some(inline) = block.find_first_inline(&predicate) {
                return Some(inline);
            }
        }
        None
    }

    fn find_first_block<F>(&self, predicate: F) -> Option<&Block>
    where
        F: Fn(&Block) -> bool,
    {
        for block in &self.blocks {
            if let Some(found) = block.find_first_block(&predicate) {
                return Some(found);
            }
        }
        None
    }
}

impl Query for Block {
    fn find_all_inlines<F>(&self, predicate: F) -> Vec<&Inline>
    where
        F: Fn(&Inline) -> bool,
    {
        let mut results = Vec::new();
        collect_inlines_from_block(self, &predicate, &mut results);
        results
    }

    fn find_all_blocks<F>(&self, predicate: F) -> Vec<&Block>
    where
        F: Fn(&Block) -> bool,
    {
        let mut results = Vec::new();
        collect_blocks_from_block(self, &predicate, &mut results);
        results
    }

    fn find_first_inline<F>(&self, predicate: F) -> Option<&Inline>
    where
        F: Fn(&Inline) -> bool,
    {
        find_first_inline_in_block(self, &predicate)
    }

    fn find_first_block<F>(&self, predicate: F) -> Option<&Block>
    where
        F: Fn(&Block) -> bool,
    {
        find_first_block_in_block(self, &predicate)
    }
}

impl Query for Vec<Inline> {
    fn find_all_inlines<F>(&self, predicate: F) -> Vec<&Inline>
    where
        F: Fn(&Inline) -> bool,
    {
        let mut results = Vec::new();
        for inline in self {
            collect_inlines_from_inline(inline, &predicate, &mut results);
        }
        results
    }

    fn find_all_blocks<F>(&self, _predicate: F) -> Vec<&Block>
    where
        F: Fn(&Block) -> bool,
    {
        Vec::new() // Inline elements don't contain blocks
    }

    fn find_first_inline<F>(&self, predicate: F) -> Option<&Inline>
    where
        F: Fn(&Inline) -> bool,
    {
        for inline in self {
            if let Some(found) = find_first_inline_in_inline(inline, &predicate) {
                return Some(found);
            }
        }
        None
    }

    fn find_first_block<F>(&self, _predicate: F) -> Option<&Block>
    where
        F: Fn(&Block) -> bool,
    {
        None // Inline elements don't contain blocks
    }
}

// Helper functions for recursive collection

fn collect_inlines_from_block<'a, F>(block: &'a Block, predicate: &F, results: &mut Vec<&'a Inline>)
where
    F: Fn(&Inline) -> bool,
{
    match block {
        Block::Paragraph(inlines) => {
            for inline in inlines {
                collect_inlines_from_inline(inline, predicate, results);
            }
        }
        Block::Heading(heading) => {
            for inline in &heading.content {
                collect_inlines_from_inline(inline, predicate, results);
            }
        }
        Block::BlockQuote(blocks) => {
            for block in blocks {
                collect_inlines_from_block(block, predicate, results);
            }
        }
        Block::List(list) => {
            for item in &list.items {
                for block in &item.blocks {
                    collect_inlines_from_block(block, predicate, results);
                }
            }
        }
        Block::Table(table) => {
            for row in &table.rows {
                for cell in row {
                    for inline in cell {
                        collect_inlines_from_inline(inline, predicate, results);
                    }
                }
            }
        }
        Block::FootnoteDefinition(footnote) => {
            for block in &footnote.blocks {
                collect_inlines_from_block(block, predicate, results);
            }
        }
        Block::GitHubAlert(alert) => {
            for block in &alert.blocks {
                collect_inlines_from_block(block, predicate, results);
            }
        }
        Block::Definition(def) => {
            for inline in &def.label {
                collect_inlines_from_inline(inline, predicate, results);
            }
        }
        _ => {} // Terminal blocks
    }
}

fn collect_inlines_from_inline<'a, F>(
    inline: &'a Inline,
    predicate: &F,
    results: &mut Vec<&'a Inline>,
) where
    F: Fn(&Inline) -> bool,
{
    if predicate(inline) {
        results.push(inline);
    }

    match inline {
        Inline::Emphasis(inlines) | Inline::Strong(inlines) | Inline::Strikethrough(inlines) => {
            for inline in inlines {
                collect_inlines_from_inline(inline, predicate, results);
            }
        }
        Inline::Link(link) => {
            for inline in &link.children {
                collect_inlines_from_inline(inline, predicate, results);
            }
        }
        Inline::LinkReference(link_ref) => {
            for inline in &link_ref.label {
                collect_inlines_from_inline(inline, predicate, results);
            }
            for inline in &link_ref.text {
                collect_inlines_from_inline(inline, predicate, results);
            }
        }
        _ => {} // Terminal inlines
    }
}

fn collect_blocks_from_block<'a, F>(block: &'a Block, predicate: &F, results: &mut Vec<&'a Block>)
where
    F: Fn(&Block) -> bool,
{
    if predicate(block) {
        results.push(block);
    }

    match block {
        Block::BlockQuote(blocks) => {
            for block in blocks {
                collect_blocks_from_block(block, predicate, results);
            }
        }
        Block::List(list) => {
            for item in &list.items {
                for block in &item.blocks {
                    collect_blocks_from_block(block, predicate, results);
                }
            }
        }
        Block::FootnoteDefinition(footnote) => {
            for block in &footnote.blocks {
                collect_blocks_from_block(block, predicate, results);
            }
        }
        Block::GitHubAlert(alert) => {
            for block in &alert.blocks {
                collect_blocks_from_block(block, predicate, results);
            }
        }
        _ => {} // Terminal or inline-containing blocks
    }
}

fn find_first_inline_in_block<'a, F>(block: &'a Block, predicate: &F) -> Option<&'a Inline>
where
    F: Fn(&Inline) -> bool,
{
    match block {
        Block::Paragraph(inlines) => {
            for inline in inlines {
                if let Some(found) = find_first_inline_in_inline(inline, predicate) {
                    return Some(found);
                }
            }
        }
        Block::Heading(heading) => {
            for inline in &heading.content {
                if let Some(found) = find_first_inline_in_inline(inline, predicate) {
                    return Some(found);
                }
            }
        }
        Block::BlockQuote(blocks) => {
            for block in blocks {
                if let Some(found) = find_first_inline_in_block(block, predicate) {
                    return Some(found);
                }
            }
        }
        Block::List(list) => {
            for item in &list.items {
                for block in &item.blocks {
                    if let Some(found) = find_first_inline_in_block(block, predicate) {
                        return Some(found);
                    }
                }
            }
        }
        Block::Table(table) => {
            for row in &table.rows {
                for cell in row {
                    for inline in cell {
                        if let Some(found) = find_first_inline_in_inline(inline, predicate) {
                            return Some(found);
                        }
                    }
                }
            }
        }
        Block::FootnoteDefinition(footnote) => {
            for block in &footnote.blocks {
                if let Some(found) = find_first_inline_in_block(block, predicate) {
                    return Some(found);
                }
            }
        }
        Block::GitHubAlert(alert) => {
            for block in &alert.blocks {
                if let Some(found) = find_first_inline_in_block(block, predicate) {
                    return Some(found);
                }
            }
        }
        Block::Definition(def) => {
            for inline in &def.label {
                if let Some(found) = find_first_inline_in_inline(inline, predicate) {
                    return Some(found);
                }
            }
        }
        _ => {} // Terminal blocks
    }
    None
}

fn find_first_inline_in_inline<'a, F>(inline: &'a Inline, predicate: &F) -> Option<&'a Inline>
where
    F: Fn(&Inline) -> bool,
{
    if predicate(inline) {
        return Some(inline);
    }

    match inline {
        Inline::Emphasis(inlines) | Inline::Strong(inlines) | Inline::Strikethrough(inlines) => {
            for inline in inlines {
                if let Some(found) = find_first_inline_in_inline(inline, predicate) {
                    return Some(found);
                }
            }
        }
        Inline::Link(link) => {
            for inline in &link.children {
                if let Some(found) = find_first_inline_in_inline(inline, predicate) {
                    return Some(found);
                }
            }
        }
        Inline::LinkReference(link_ref) => {
            for inline in &link_ref.label {
                if let Some(found) = find_first_inline_in_inline(inline, predicate) {
                    return Some(found);
                }
            }
            for inline in &link_ref.text {
                if let Some(found) = find_first_inline_in_inline(inline, predicate) {
                    return Some(found);
                }
            }
        }
        _ => {} // Terminal inlines
    }
    None
}

fn find_first_block_in_block<'a, F>(block: &'a Block, predicate: &F) -> Option<&'a Block>
where
    F: Fn(&Block) -> bool,
{
    if predicate(block) {
        return Some(block);
    }

    match block {
        Block::BlockQuote(blocks) => {
            for block in blocks {
                if let Some(found) = find_first_block_in_block(block, predicate) {
                    return Some(found);
                }
            }
        }
        Block::List(list) => {
            for item in &list.items {
                for block in &item.blocks {
                    if let Some(found) = find_first_block_in_block(block, predicate) {
                        return Some(found);
                    }
                }
            }
        }
        Block::FootnoteDefinition(footnote) => {
            for block in &footnote.blocks {
                if let Some(found) = find_first_block_in_block(block, predicate) {
                    return Some(found);
                }
            }
        }
        Block::GitHubAlert(alert) => {
            for block in &alert.blocks {
                if let Some(found) = find_first_block_in_block(block, predicate) {
                    return Some(found);
                }
            }
        }
        _ => {} // Terminal or inline-containing blocks
    }
    None
}