ai-app/client-core/src/markdown_blocks.rs

//! Split a markdown message into its top-level **blocks** -- one
//! paragraph, heading, fenced code block, list, table or quote each, as a
//! byte slice of the original source.
//!
//! This exists for streaming. A transcript row used to be one text widget
//! holding the whole message, so a single streamed delta re-shaped every
//! paragraph of it through the text engine again; the phone's bench v2 put
//! the stream phase at p50 18.2ms against Compose's 13.4ms for exactly
//! that reason (docs/IRIS_TODO.md). A row is a column of one widget per
//! block now, and a delta that lands in the last block leaves every
//! earlier block's layout alone. `docs/DECISIONS.md`'s 2026-09-06 entry has
//! what that rejected and why the split lives here rather than in the UI
//! crate: `docs/CLIENT_CORE.md` already wanted a block model for P1, and
//! keeping it here means iris stays a text renderer that knows nothing
//! about markdown.
//!
//! **Blocks only.** Inline styling (bold, links, inline code) is still the
//! renderer's own job, per block -- this deliberately does not build a
//! full AST, because nothing needs one yet.
//!
//! ## Appending is not guaranteed to leave earlier blocks alone
//!
//! It nearly always does, which is what makes the fast path worth having,
//! but markdown has no such rule: appending a "```" line can turn text
//! that was three paragraphs into one fenced block, and appending "---"
//! under a paragraph turns that paragraph into a heading. So a caller
//! taking the O(last block) path **must compare the prefix it is about to
//! keep** rather than assume it. [`common_prefix`] is that comparison, and
//! it is cheap next to laying the text out again.

use pulldown_cmark::{Event, Options, Parser, Tag};

/// What a block is, for a renderer that wants to style or space blocks
/// differently. `Other` is deliberately present rather than a panic or a
/// silent fallback to `Paragraph`: markdown has more block kinds than this
/// list and more get added, and a renderer treating an unknown one as
/// prose is right, but it should be able to *tell* that is what it is
/// doing.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BlockKind {
    Paragraph,
    Heading,
    /// A fenced or indented code block.
    Code,
    List,
    Table,
    Quote,
    /// A thematic break, raw HTML, a footnote -- anything with no
    /// distinguished treatment here.
    Other,
}

/// One top-level block: its kind and the exact source that produced it.
/// `source` is a slice of the input with trailing whitespace removed, so
/// two splits of the same prefix compare equal even when one of them had a
/// delta arriving after it.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Block {
    pub kind: BlockKind,
    pub source: String,
}

fn kind_of(tag: &Tag) -> BlockKind {
    match tag {
        Tag::Paragraph => BlockKind::Paragraph,
        Tag::Heading { .. } => BlockKind::Heading,
        Tag::CodeBlock(_) => BlockKind::Code,
        Tag::List(_) => BlockKind::List,
        Tag::Table(_) => BlockKind::Table,
        Tag::BlockQuote(_) => BlockKind::Quote,
        _ => BlockKind::Other,
    }
}

fn options() -> Options {
    // The same set `transcript-ui`'s renderer parses with, so a block
    // boundary here and the styling there cannot disagree about what the
    // source means.
    Options::ENABLE_STRIKETHROUGH | Options::ENABLE_TABLES | Options::ENABLE_TASKLISTS
}

/// Split `src` into its top-level blocks, in source order. An empty or
/// whitespace-only input gives no blocks; text the parser does not put
/// inside any block (a stray fence marker mid-stream) still comes back,
/// as `Other`, rather than being dropped.
pub fn split_blocks(src: &str) -> Vec<Block> {
    let mut out: Vec<Block> = Vec::new();
    let mut depth = 0usize;
    let mut kind = BlockKind::Other;
    for (event, range) in Parser::new_ext(src, options()).into_offset_iter() {
        match event {
            Event::Start(tag) => {
                if depth == 0 {
                    kind = kind_of(&tag);
                }
                depth += 1;
            }
            Event::End(_) => {
                depth -= 1;
                if depth == 0 {
                    push(&mut out, kind, &src[range]);
                }
            }
            // A top-level event that is not part of any block -- a
            // thematic break, a block of raw HTML. Inside one, it is the
            // enclosing block's business and this does nothing.
            _ => {
                if depth == 0 {
                    push(&mut out, BlockKind::Other, &src[range]);
                }
            }
        }
    }
    out
}

fn push(out: &mut Vec<Block>, kind: BlockKind, source: &str) {
    let source = source.trim_end();
    if source.is_empty() {
        return;
    }
    out.push(Block {
        kind,
        source: source.to_string(),
    });
}

/// How many leading blocks of `old` and `new` are identical -- what a
/// caller may keep the laid-out widgets for. See the module doc for why
/// this is a comparison rather than an assumption.
pub fn common_prefix(old: &[Block], new: &[Block]) -> usize {
    old.iter().zip(new).take_while(|(a, b)| a == b).count()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn kinds(src: &str) -> Vec<BlockKind> {
        split_blocks(src).into_iter().map(|b| b.kind).collect()
    }

    #[test]
    fn a_message_splits_into_its_top_level_blocks() {
        let src = "# Title\n\nFirst para.\n\n```rust\nfn main() {}\n```\n\n- a\n- b\n";
        assert_eq!(
            kinds(src),
            vec![
                BlockKind::Heading,
                BlockKind::Paragraph,
                BlockKind::Code,
                BlockKind::List
            ]
        );
        let blocks = split_blocks(src);
        assert_eq!(blocks[1].source, "First para.");
        assert_eq!(blocks[2].source, "```rust\nfn main() {}\n```");
    }

    #[test]
    fn blank_input_has_no_blocks() {
        assert!(split_blocks("").is_empty());
        assert!(split_blocks("   \n\n  ").is_empty());
    }

    /// The property the streaming fast path rests on, in its ordinary
    /// shape: a delta landing in the last paragraph must leave every
    /// earlier block byte-identical.
    #[test]
    fn a_delta_into_the_last_paragraph_leaves_earlier_blocks_untouched() {
        let before = split_blocks("# Title\n\nFirst para.\n\nSecond par");
        let after = split_blocks("# Title\n\nFirst para.\n\nSecond paragraph now.");
        assert_eq!(common_prefix(&before, &after), 2);
        assert_eq!(before.len(), 3);
        assert_eq!(after.len(), 3);
        assert_ne!(before[2], after[2]);
    }

    /// A delta that starts a *new* block keeps every old block, including
    /// the one that was last -- so the fast path appends rather than
    /// replacing.
    #[test]
    fn a_delta_that_starts_a_new_block_keeps_every_old_one() {
        let before = split_blocks("First para.\n\nSecond para.");
        let after = split_blocks("First para.\n\nSecond para.\n\nThird");
        assert_eq!(common_prefix(&before, &after), 2);
        assert_eq!(after.len(), 3);
    }

    /// A code fence arrives one delta at a time and is unterminated for
    /// most of its life. It must still be *one* block the whole way, or
    /// every delta would re-split the message into a different number of
    /// pieces.
    #[test]
    fn an_unterminated_fence_is_one_block_while_it_streams() {
        for src in [
            "Here:\n\n```rust\n",
            "Here:\n\n```rust\nfn main() {\n",
            "Here:\n\n```rust\nfn main() {\n    println!(\"hi\");\n",
        ] {
            assert_eq!(
                kinds(src),
                vec![BlockKind::Paragraph, BlockKind::Code],
                "{src:?}"
            );
        }
    }

    /// The half the fast path had no reason to touch, and the reason
    /// `common_prefix` is a comparison rather than an assumption:
    /// appending can rewrite what came before. `---` under a paragraph
    /// turns that paragraph into a setext heading, so the block that was
    /// already laid out is not the block it is now.
    #[test]
    fn appending_can_rewrite_an_earlier_block_and_the_prefix_says_so() {
        let before = split_blocks("Not a heading\n\nsecond");
        let after = split_blocks("Not a heading\n\nsecond\n---");
        assert_eq!(before[1].kind, BlockKind::Paragraph);
        assert_eq!(after[1].kind, BlockKind::Heading);
        assert_eq!(
            common_prefix(&before, &after),
            1,
            "the rewritten block must not be reported as keepable"
        );
    }

    #[test]
    fn a_thematic_break_is_its_own_block() {
        assert_eq!(
            kinds("one\n\n---\n\ntwo"),
            vec![BlockKind::Paragraph, BlockKind::Other, BlockKind::Paragraph]
        );
    }

    /// The shapes a real transcript actually contains, each checked for
    /// the one property the streaming fast path needs: the *number* of
    /// blocks and every earlier block's source stay put while the message
    /// grows. A fence's own blank lines, a `---` inside one, a nested
    /// list and a table are all places where a naive line-based split
    /// would break the message into more pieces than there are blocks.
    #[test]
    fn the_transcripts_own_block_shapes_survive_a_split() {
        let fence_with_blanks = "Intro.\n\n```rust\nfn a() {}\n\nfn b() {}\n```\n\nAfter.";
        assert_eq!(
            kinds(fence_with_blanks),
            vec![BlockKind::Paragraph, BlockKind::Code, BlockKind::Paragraph],
            "a blank line inside a fence is not a block boundary"
        );
        assert_eq!(
            kinds("```\n---\n```"),
            vec![BlockKind::Code],
            "a thematic break inside a fence is code, not a break"
        );
        assert_eq!(
            kinds("- a\n  - a1\n  - a2\n- b"),
            vec![BlockKind::List],
            "a nested list is one top-level block"
        );
        assert_eq!(
            kinds("## Heading\n```sh\nls\n```"),
            vec![BlockKind::Heading, BlockKind::Code],
            "a fence directly under a heading, with no blank line"
        );
        assert_eq!(
            kinds("| a | b |\n|---|---|\n| 1 | 2 |"),
            vec![BlockKind::Table]
        );
        assert_eq!(
            kinds("> quoted\n> more\n\nplain"),
            vec![BlockKind::Quote, BlockKind::Paragraph]
        );
    }

    /// `apply_delta`'s precondition, stated as the property rather than
    /// the arithmetic: for every prefix of a realistic streamed message,
    /// the blocks before the last one must be exactly the blocks the
    /// previous prefix had. Where markdown breaks that (the `---` case
    /// above), `common_prefix` has to *say* so -- which is what the
    /// `>= len - 1` assertion below checks: the split may rewrite the
    /// last block, never an earlier one, or `RowBlocks::apply_delta`
    /// would keep a widget whose text is no longer what it holds.
    #[test]
    fn every_prefix_of_a_streamed_message_keeps_all_but_its_last_block() {
        let full = "# Report\n\nFirst finding, at some length.\n\n```rust\nfn main() {\n\n                        println!(\"hi\");\n}\n```\n\n- one\n  - nested\n- two\n\n| a | b |\n                    |---|---|\n| 1 | 2 |\n\n> and a closing quote.";
        // Every character boundary, so a delta landing mid-word and one
        // landing exactly on a fence's closing backtick are both covered.
        let mut prev = Vec::new();
        for end in full.char_indices().map(|(i, _)| i).chain([full.len()]) {
            let now = split_blocks(&full[..end]);
            let common = common_prefix(&prev, &now);
            assert!(
                prev.is_empty() || common + 1 >= prev.len(),
                "at {end} bytes the split rewrote block {common} of {}, not just the last one:\n                 before={prev:#?}\nafter={now:#?}",
                prev.len()
            );
            prev = now;
        }
    }

    /// The half a growing message cannot show: a fence that never closes.
    /// The stream ends there and the block must still be the code block
    /// it has been all along, not re-split into paragraphs.
    #[test]
    fn a_stream_that_ends_inside_a_fence_still_ends_with_one_code_block() {
        let src = "Here is the patch:\n\n```diff\n- old line\n+ new line";
        let blocks = split_blocks(src);
        assert_eq!(
            blocks.iter().map(|b| b.kind).collect::<Vec<_>>(),
            vec![BlockKind::Paragraph, BlockKind::Code]
        );
        assert_eq!(blocks[1].source, "```diff\n- old line\n+ new line");
    }

    /// A delta that closes a fence changes the *last* block only, so the
    /// fast path takes it -- the case the module doc says is the reason
    /// `common_prefix` is a comparison.
    #[test]
    fn the_delta_that_closes_a_fence_changes_only_the_last_block() {
        let before = split_blocks("Text.\n\n```\ncode\n");
        let after = split_blocks("Text.\n\n```\ncode\n```");
        assert_eq!(before.len(), after.len());
        assert_eq!(common_prefix(&before, &after), 1);
        assert_ne!(before[1], after[1]);
    }
}