feat: IR transform passes between parse and render

.github/workflows/ci.yml

@@ -85,7 +85,7 @@ jobs:
           fi
       - name: Run examples
         run: |
-          for ex in basic custom_selectors configuration ir_access; do
+          for ex in basic custom_selectors configuration ir_access transforms; do
             cargo run -q -p markkit --example "$ex"
           done
           cargo run -q -p markkit --features tool --example agent_tool

CHANGELOG.md

@@ -20,6 +20,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
   no longer dropped.
 - `Document.ir_version` (`markkit::ir::IR_VERSION = 2`) is serialized with the IR JSON.
 - Math: `$tex$` and `$$` blocks from MathML (`application/x-tex` annotation), KaTeX and MathJax output.
+- IR transforms: `MarkkitBuilder::transform(..)` runs `Transform` passes after parsing and before
+  rendering, for every format. Built-ins `transform::NormalizeHeadings` and `transform::DropImages`;
+  `transform::from_fn` for closures; `Document::visit` / `visit_mut` to walk the tree.
 - JSON Schema for the IR at `schema/ir.schema.json`; `json-schema` feature adds `markkit::ir::json_schema()`.
 - Per-block provenance: `MarkkitBuilder::track_sources(true)` records a `Source` on each IR block
   (HTML element XPath, DOCX body block index, spreadsheet sheet and A1 range) for citations.
@@ -50,6 +53,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 
 ### Changed
 
+- **Breaking:** `Error` is `#[non_exhaustive]` (new `Error::Transform`); matches need a `_` arm.
 - **Breaking (IR API):** `Node`, `Document`, `ListItem` and the new IR structs are `#[non_exhaustive]`;
   matches need a `_` arm and structs are built with constructors. `Node::Table` cells are now
   `TableCell` (`children`, `colspan`, `rowspan`) with an `alignments` field. Block variants gain a

crates/markkit/Cargo.toml

@@ -98,6 +98,11 @@ name = "agent_tool"
 path = "../../examples/agent_tool.rs"
 required-features = ["html", "tool"]
 
+[[example]]
+name = "transforms"
+path = "../../examples/transforms.rs"
+required-features = ["html"]
+
 [[bench]]
 name = "extraction"
 harness = false

crates/markkit/src/config.rs

@@ -44,6 +44,8 @@ pub struct MarkkitBuilder {
     /// HTML element XPath, DOCX block index, spreadsheet sheet and range.
     /// Default: false (no cost, no `source` in IR JSON).
     pub track_sources: bool,
+    /// IR passes run after parsing, before rendering (see [`crate::transform`]).
+    pub transforms: crate::transform::TransformChain,
 }
 
 impl Default for MarkkitBuilder {
@@ -74,6 +76,7 @@ impl Default for MarkkitBuilder {
             allowed_hosts: Vec::new(),
             max_redirects: 5,
             track_sources: false,
+            transforms: Default::default(),
         }
     }
 }
@@ -105,6 +108,12 @@ impl MarkkitBuilder {
         self
     }
 
+    /// Append an IR transform pass; passes run in registration order.
+    pub fn transform(mut self, t: impl crate::transform::Transform + 'static) -> Self {
+        self.transforms.push(t);
+        self
+    }
+
     pub fn include_selector(mut self, selector: impl Into<String>) -> Self {
         self.include_selectors.push(selector.into());
         self

crates/markkit/src/error.rs

@@ -1,5 +1,6 @@
 /// Errors that can occur during markdown extraction.
 #[derive(Debug, thiserror::Error)]
+#[non_exhaustive]
 pub enum Error {
     #[error("HTML parsing error: {0}")]
     HtmlParse(String),
@@ -30,6 +31,10 @@ pub enum Error {
     #[error("Resource limit exceeded: {0}")]
     ResourceLimit(String),
 
+    /// A [`crate::transform::Transform`] failed.
+    #[error("Transform error: {0}")]
+    Transform(String),
+
     #[error("I/O error: {0}")]
     Io(#[from] std::io::Error),
 }

crates/markkit/src/extract.rs

@@ -167,8 +167,7 @@ impl Markkit {
 
         #[cfg(feature = "html")]
         {
-            let parser = crate::html::HtmlParser::new(&self.config);
-            let doc = parser.parse(html)?;
+            let doc = self.parse_html(html)?;
             let mut renderer = crate::render::MarkdownRenderer::new();
             Ok(renderer.render(&doc))
         }
@@ -189,7 +188,9 @@ impl Markkit {
             )));
         }
         let parser = crate::html::HtmlParser::new(&self.config);
-        parser.parse(html)
+        let mut doc = parser.parse(html)?;
+        self.config.transforms.apply(&mut doc)?;
+        Ok(doc)
     }
 
     /// Extract markdown from raw bytes, auto-detecting the format.
@@ -283,7 +284,8 @@ impl Markkit {
     fn parse_html_doc(&self, bytes: &[u8]) -> crate::Result<ParsedDocument> {
         let text = String::from_utf8_lossy(bytes);
         let parser = crate::html::HtmlParser::new(&self.config);
-        let (doc, page) = parser.parse_page(&text)?;
+        let (mut doc, page) = parser.parse_page(&text)?;
+        self.config.transforms.apply(&mut doc)?;
         let mut p = ParsedDocument::new(doc, Format::Html);
         p.title = page.title;
         p.metadata = page.metadata;
@@ -302,7 +304,9 @@ impl Markkit {
         let out = guarded("DOCX", Error::DocxExtract, || {
             crate::docx::parse(bytes, &self.config)
         })?;
-        let mut p = ParsedDocument::new(out.doc, Format::Docx);
+        let mut doc = out.doc;
+        self.config.transforms.apply(&mut doc)?;
+        let mut p = ParsedDocument::new(doc, Format::Docx);
         p.title = out.title;
         p.stats.paragraphs = Some(out.paragraphs);
         p.stats.tables = Some(out.tables);
@@ -321,7 +325,9 @@ impl Markkit {
         let out = guarded("spreadsheet", Error::SpreadsheetExtract, || {
             crate::spreadsheet::parse(bytes, &self.config)
         })?;
-        let mut p = ParsedDocument::new(out.doc, out.format);
+        let mut doc = out.doc;
+        self.config.transforms.apply(&mut doc)?;
+        let mut p = ParsedDocument::new(doc, out.format);
         p.warnings = out.warnings;
         p.stats.sheets = Some(out.sheets);
         p.stats.rows = Some(out.rows);

crates/markkit/src/ir.rs

@@ -50,6 +50,29 @@ impl Document {
     pub fn is_empty(&self) -> bool {
         self.nodes.is_empty()
     }
+
+    /// Visit every node, pre-order (parents before children).
+    pub fn visit(&self, f: &mut dyn FnMut(&Node)) {
+        fn walk(nodes: &[Node], f: &mut dyn FnMut(&Node)) {
+            for n in nodes {
+                f(n);
+                n.for_each_child_list_ref(&mut |list| walk(list, f));
+            }
+        }
+        walk(&self.nodes, f);
+    }
+
+    /// Visit every node mutably, pre-order. Children are visited after `f`
+    /// ran on their parent, so `f` may rewrite a subtree before it is walked.
+    pub fn visit_mut(&mut self, f: &mut dyn FnMut(&mut Node)) {
+        fn walk(nodes: &mut [Node], f: &mut dyn FnMut(&mut Node)) {
+            for n in nodes {
+                f(n);
+                n.for_each_child_list(&mut |list| walk(list, f));
+            }
+        }
+        walk(&mut self.nodes, f);
+    }
 }
 
 /// A block-level or inline node in the document tree.
@@ -196,6 +219,90 @@ impl Node {
         }
     }
 
+    /// Call `f` on each child list of this node (block or inline children,
+    /// list-item children, table cells, definition terms and bodies,
+    /// figure captions). Leaves have none.
+    pub fn for_each_child_list(&mut self, f: &mut dyn FnMut(&mut Vec<Node>)) {
+        match self {
+            Node::Heading { children, .. }
+            | Node::Paragraph { children, .. }
+            | Node::BlockQuote { children, .. }
+            | Node::FootnoteDefinition { children, .. }
+            | Node::Bold { children }
+            | Node::Italic { children }
+            | Node::Strikethrough { children }
+            | Node::Link { children, .. } => f(children),
+            Node::Figure {
+                children, caption, ..
+            } => {
+                f(children);
+                if let Some(c) = caption {
+                    f(c);
+                }
+            }
+            Node::List { items, .. } => {
+                for item in items {
+                    f(&mut item.children);
+                }
+            }
+            Node::Table { headers, rows, .. } => {
+                for cell in headers.iter_mut().chain(rows.iter_mut().flatten()) {
+                    f(&mut cell.children);
+                }
+            }
+            Node::DefinitionList { items, .. } => {
+                for item in items {
+                    f(&mut item.term);
+                    for d in &mut item.definitions {
+                        f(d);
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+
+    /// Read-only [`Node::for_each_child_list`].
+    pub fn for_each_child_list_ref(&self, f: &mut dyn FnMut(&Vec<Node>)) {
+        match self {
+            Node::Heading { children, .. }
+            | Node::Paragraph { children, .. }
+            | Node::BlockQuote { children, .. }
+            | Node::FootnoteDefinition { children, .. }
+            | Node::Bold { children }
+            | Node::Italic { children }
+            | Node::Strikethrough { children }
+            | Node::Link { children, .. } => f(children),
+            Node::Figure {
+                children, caption, ..
+            } => {
+                f(children);
+                if let Some(c) = caption {
+                    f(c);
+                }
+            }
+            Node::List { items, .. } => {
+                for item in items {
+                    f(&item.children);
+                }
+            }
+            Node::Table { headers, rows, .. } => {
+                for cell in headers.iter().chain(rows.iter().flatten()) {
+                    f(&cell.children);
+                }
+            }
+            Node::DefinitionList { items, .. } => {
+                for item in items {
+                    f(&item.term);
+                    for d in &item.definitions {
+                        f(d);
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+
     /// Where this block came from, when the parser recorded it.
     pub fn source(&self) -> Option<&Source> {
         match self {

crates/markkit/src/lib.rs

@@ -61,6 +61,7 @@ pub mod render;
 mod spreadsheet;
 #[cfg(feature = "tool")]
 pub mod tool;
+pub mod transform;
 #[cfg(feature = "docx")]
 mod xml;