feat: add support for tagged documents with transparent rendering

src/alloc.rs

@@ -154,6 +154,12 @@ pub trait DocAllocator<'a> {
         DocBuilder::from_utf8_text(self, data.into())
     }
 
+    /// Allocate a document and attach a transparent tag to it.
+    #[inline]
+    fn tagged(&'a self, id: u32, doc: impl Pretty<'a, Self>) -> DocBuilder<'a, Self> {
+        doc.pretty(self).tag(id)
+    }
+
     /// Allocate a document containing the given text, which is assumed to be ASCII.
     ///
     /// This can avoid checking for unicode width.

src/builder.rs

@@ -134,6 +134,16 @@ where
         }
     }
 
+    /// Attach a transparent tag to this document.
+    #[inline]
+    pub fn tag(self, id: u32) -> Self {
+        let Self(allocator, this) = self;
+        match *this {
+            Doc::Nil => Self(allocator, this),
+            _ => Self(allocator, Doc::Tagged(id, allocator.alloc_cow(this)).into()),
+        }
+    }
+
     /// Acts as `self` when laid out on multiple lines and acts as `that` when laid out on a single line.
     ///
     /// ```

src/debug_tags.rs

@@ -0,0 +1,115 @@
+use std::fmt;
+
+use crate::{Doc, DocPtr, text::Text};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct DebugTagRange {
+    pub tag: u32,
+    pub start: usize,
+    pub end: usize,
+}
+
+/// Render debug IR text and collect output ranges for tagged text leaves.
+///
+/// The returned text is identical to `format!("{doc:#?}")`.
+pub fn debug_with_tag_ranges<'a, T>(doc: &Doc<'a, T>) -> (String, Vec<DebugTagRange>)
+where
+    T: DocPtr<'a> + fmt::Debug,
+{
+    let text = format!("{doc:#?}");
+    let mut tagged_tokens = Vec::new();
+    collect_tagged_tokens(doc, None, &mut tagged_tokens);
+
+    let mut cursor = 0;
+    let mut ranges = Vec::new();
+    for (tag, token) in tagged_tokens {
+        if token.is_empty() {
+            continue;
+        }
+        if let Some(found) = text[cursor..].find(&token) {
+            let start = cursor + found;
+            let end = start + token.len();
+            ranges.push(DebugTagRange { tag, start, end });
+            cursor = end;
+        }
+    }
+
+    (text, ranges)
+}
+
+fn collect_tagged_tokens<'a, T>(
+    doc: &Doc<'a, T>,
+    active_tag: Option<u32>,
+    out: &mut Vec<(u32, String)>,
+) where
+    T: DocPtr<'a>,
+{
+    match doc {
+        Doc::Nil | Doc::Fail | Doc::HardLine | Doc::ExpandParent => {}
+        Doc::Text(text) => {
+            if let Some(tag) = active_tag {
+                out.push((tag, format!("{text:?}")));
+            }
+        }
+        Doc::TextWithLen(_, inner) => collect_tagged_tokens(inner, active_tag, out),
+        Doc::Append(left, right) => {
+            collect_tagged_tokens(left, active_tag, out);
+            collect_tagged_tokens(right, active_tag, out);
+        }
+        Doc::LineSuffix(inner)
+        | Doc::Nest(_, inner)
+        | Doc::DedentToRoot(inner)
+        | Doc::Align(inner)
+        | Doc::Flatten(inner) => collect_tagged_tokens(inner, active_tag, out),
+        Doc::Tagged(id, inner) => collect_tagged_tokens(inner, Some(*id), out),
+        Doc::BreakOrFlat(break_doc, flat_doc) => match (&**break_doc, &**flat_doc) {
+            (Doc::HardLine, Doc::Text(Text::Borrowed(" "))) => {}
+            (Doc::HardLine, Doc::Nil) => {}
+            (_, Doc::Nil) => collect_tagged_tokens(break_doc, active_tag, out),
+            (Doc::Nil, _) => collect_tagged_tokens(flat_doc, active_tag, out),
+            _ => {
+                // Debug prints FlatOrBreak(y, x).
+                collect_tagged_tokens(flat_doc, active_tag, out);
+                collect_tagged_tokens(break_doc, active_tag, out);
+            }
+        },
+        Doc::Group(inner) => match &**inner {
+            Doc::BreakOrFlat(break_doc, flat_doc)
+                if matches!(
+                    (&**break_doc, &**flat_doc),
+                    (Doc::HardLine, Doc::Text(Text::Borrowed(" ")))
+                ) => {}
+            Doc::BreakOrFlat(break_doc, flat_doc)
+                if matches!((&**break_doc, &**flat_doc), (Doc::HardLine, Doc::Nil)) => {}
+            _ => collect_tagged_tokens(inner, active_tag, out),
+        },
+        Doc::Union(left, right) | Doc::PartialUnion(left, right) => {
+            collect_tagged_tokens(left, active_tag, out);
+            collect_tagged_tokens(right, active_tag, out);
+        }
+        #[cfg(feature = "contextual")]
+        Doc::OnColumn(_) | Doc::OnNesting(_) => {}
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{Arena, DocAllocator};
+
+    use super::debug_with_tag_ranges;
+
+    #[test]
+    fn keeps_debug_output_stable_and_collects_ranges() {
+        let arena = Arena::new();
+        let doc = (arena.text("alpha").tag(1) + arena.space() + arena.text("beta").tag(2))
+            .group();
+
+        let (text, ranges) = debug_with_tag_ranges(&doc);
+        assert_eq!(text, format!("{doc:#?}"));
+        assert_eq!(ranges.len(), 2);
+        assert_eq!(ranges[0].tag, 1);
+        assert_eq!(ranges[1].tag, 2);
+        assert_eq!(&text[ranges[0].start..ranges[0].end], "\"alpha\"");
+        assert_eq!(&text[ranges[1].start..ranges[1].end], "\"beta\"");
+    }
+}

src/doc.rs

@@ -40,6 +40,7 @@ where
     // Structural
     Append(T, T),  // Sequencing
     LineSuffix(T), // A document that is appended to the end of the current line.
+    Tagged(u32, T),
 
     // Indentation and Alignment
     Nest(isize, T),  // Changes the indentation level
@@ -141,6 +142,8 @@ where
                 f.finish()
             }
             Doc::LineSuffix(ref doc) => write_compact(f, doc, "LineSuffix"),
+            // Keep tags transparent in debug output.
+            Doc::Tagged(_, ref doc) => doc.fmt(f),
 
             Doc::Nest(off, ref doc) => {
                 write!(f, "Nest({off}, ",)?;

src/lib.rs

@@ -144,13 +144,15 @@
 
 mod alloc;
 mod builder;
+mod debug_tags;
 mod doc;
 mod render;
 pub mod text;
 pub(crate) mod visitor;
 
 pub use alloc::{Arena, BoxAllocator, DocAllocator, RcAllocator};
 pub use builder::DocBuilder;
+pub use debug_tags::{DebugTagRange, debug_with_tag_ranges};
 pub use doc::{BoxDoc, BuildDoc, Doc, DocPtr, RcDoc, RefDoc};
 pub use render::{FmtWrite, IoWrite, Render};
 

src/render.rs

@@ -85,5 +85,17 @@ pub trait Render {
         Ok(())
     }
 
+    /// Called when rendering enters a tagged document.
+    #[inline]
+    fn on_tag_enter(&mut self, _id: u32) -> Result<(), Self::Error> {
+        Ok(())
+    }
+
+    /// Called when rendering exits a tagged document.
+    #[inline]
+    fn on_tag_exit(&mut self, _id: u32) -> Result<(), Self::Error> {
+        Ok(())
+    }
+
     fn fail_doc(&self) -> Self::Error;
 }