feat/unreleased: annotate viz nodes using //# not # (#2610)

Author: sxlijin

engine/baml-lib/ast/src/parser/datamodel.pest

@@ -1,5 +1,5 @@
 schema = {
-    SOI ~ (mdx_header | expr_fn | top_level_assignment | value_expression_block | type_expression_block | template_declaration | type_alias | comment_block | raw_string_literal | empty_lines | CATCH_ALL)* ~ EOI
+    SOI ~ (expr_fn | top_level_assignment | value_expression_block | type_expression_block | template_declaration | type_alias | comment_block | raw_string_literal | empty_lines | CATCH_ALL)* ~ EOI
 }
 
 // ######################################
@@ -27,7 +27,7 @@ field_type_with_attr = { field_type ~ (NEWLINE? ~ (field_attribute | trailing_co
 value_expression_keyword  = { FUNCTION_KEYWORD | TEST_KEYWORD | CLIENT_KEYWORD | RETRY_POLICY_KEYWORD | GENERATOR_KEYWORD }
 value_expression_block    = { value_expression_keyword ~ identifier ~ named_argument_list? ~ ARROW? ~ field_type_chain? ~ SPACER_TEXT ~ BLOCK_OPEN ~ value_expression_contents ~ BLOCK_CLOSE }
 value_expression_contents = {
-    (mdx_header | stmt | type_builder_block | value_expression | comment_block | block_attribute | empty_lines | BLOCK_LEVEL_CATCH_ALL)*
+    (stmt | type_builder_block | value_expression | comment_block | block_attribute | empty_lines | BLOCK_LEVEL_CATCH_ALL)*
 }
 value_expression          = { identifier ~ config_expression? ~ (NEWLINE? ~ field_attribute)* ~ trailing_comment? }
 
@@ -339,7 +339,7 @@ top_level_stmt = {
 expr_fn = { "function" ~ identifier ~ named_argument_list ~ ARROW? ~ field_type_chain? ~ expr_block }
 
 // Body of a function (including curly brackets).
-expr_block = { BLOCK_OPEN ~ NEWLINE? ~ (mdx_header | expr_body_stmt | stmt | comment_block | empty_lines)* ~ expression? ~ (comment_block | empty_lines)* ~ BLOCK_CLOSE }
+expr_block = { BLOCK_OPEN ~ NEWLINE? ~ (expr_body_stmt | stmt | comment_block | empty_lines)* ~ expression? ~ (comment_block | empty_lines)* ~ BLOCK_CLOSE }
 
 // More forgiving statement rule for function bodies - only for statements that commonly miss semicolons
 expr_body_stmt = {
@@ -356,8 +356,6 @@ expr_body_stmt = {
     ~ NEWLINE?
 }
 
-// Headers can include any characters except a newline; do not use unquoted_string_literal restrictions
-mdx_header = { WHITESPACE* ~ "#"+ ~ WHITESPACE* ~ (!NEWLINE ~ ANY)* ~ NEWLINE? }
 
 // Statement.
 stmt = {

engine/baml-lib/ast/src/parser/parse.rs

@@ -7,7 +7,7 @@ use pest::Parser;
 
 use super::{
     parse_assignment::parse_assignment,
-    parse_expr::{parse_expr_fn, parse_header, parse_top_level_assignment},
+    parse_expr::{parse_comment_header_pair, parse_expr_fn, parse_top_level_assignment},
     parse_expression::parse_expression,
     parse_template_string::parse_template_string,
     parse_type_expression_block::parse_type_expression_block,
@@ -167,12 +167,13 @@ pub fn parse(root_path: &Path, source: &SourceFile) -> Result<(Ast, Diagnostics)
                     ));
                         break;
                     }
-                    Rule::mdx_header => {
-                        if let Some(header) = parse_header(current, &mut diagnostics) {
-                            pending_headers.push(header);
-                        }
-                    }
                     Rule::comment_block => {
+                        let headers =
+                            headers_from_comment_block_top_level(current.clone(), &mut diagnostics);
+                        if !headers.is_empty() {
+                            pending_headers.extend(headers);
+                            continue;
+                        }
                         match pairs.peek().map(|b| b.as_rule()) {
                             Some(Rule::empty_lines) => {
                                 // free floating
@@ -226,6 +227,25 @@ pub fn parse(root_path: &Path, source: &SourceFile) -> Result<(Ast, Diagnostics)
     }
 }
 
+fn headers_from_comment_block_top_level(
+    token: pest::iterators::Pair<'_, Rule>,
+    diagnostics: &mut Diagnostics,
+) -> Vec<Header> {
+    if token.as_rule() != Rule::comment_block {
+        return Vec::new();
+    }
+
+    let mut headers = Vec::new();
+    for current in token.into_inner() {
+        if current.as_rule() == Rule::comment {
+            if let Some(header) = parse_comment_header_pair(&current, diagnostics) {
+                headers.push(header);
+            }
+        }
+    }
+    headers
+}
+
 fn get_expected_from_error(positives: &[Rule]) -> String {
     use std::fmt::Write as _;
     let mut out = String::with_capacity(positives.len() * 6);

engine/baml-lib/ast/src/parser/parse_expr.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use internal_baml_diagnostics::{DatamodelError, Diagnostics};
 
 use super::{
@@ -765,34 +767,56 @@ pub fn parse_expr_block(token: Pair<'_>, diagnostics: &mut Diagnostics) -> Optio
     let mut expr = None;
     let _open_bracket = tokens.next()?;
 
-    // Collect all items first to process headers together
+    // Collect all items first so we can gather every header before we bind them
+    // to statements. We need two passes: the first pass collects and normalizes
+    // the headers (including establishing their relative levels), the second
+    // pass walks the statements in source order and attaches those normalized
+    // headers. If we tried to attach while parsing in a single pass, headers
+    // appearing inside comment blocks would be seen after their statements and
+    // could not participate in markdown hierarchy normalization.
     let mut items: Vec<Pair<'_>> = Vec::new();
     for item in tokens {
         items.push(item);
     }
 
     // Track headers with their hierarchy
+    // NB(sam): I don't entirely understand why we need to wrap Headers in Arc<>,
+    // but here are the notes from codex:
+    // <codex>
+    // Most AST nodes are owned outright—each node sits in exactly one place in
+    // the tree—so ordinary struct fields work fine. Header annotations are the
+    // odd case: the parser needs to attach the same logical header instance to
+    // multiple spots (statements, trailing expressions, top‑level block etc.)
+    // while also normalizing them later. To avoid copying or moving those
+    // structs repeatedly, the parser promotes headers into shared references
+    // (Arc<Header>). That lets the first pass create and normalize a header
+    // once, stash it in the lookup map, and then hand out clones of the pointer
+    // wherever the header appears, without duplication or life‑time juggling.
+    // Functionally, Arc is central here because headers get reused across many
+    // nodes, not because other AST structures require special thread‑safety
+    // treatment.
+    // </codex>
     let mut all_headers_in_block: Vec<std::sync::Arc<Header>> = Vec::new();
 
     // First pass: collect all headers
     for item in &items {
-        if item.as_rule() == Rule::mdx_header {
-            let header = parse_header(item.clone(), diagnostics);
-            if let Some(header) = header {
-                let header_arc = std::sync::Arc::new(header);
-                all_headers_in_block.push(header_arc.clone());
+        if item.as_rule() == Rule::comment_block {
+            let headers = headers_from_comment_block(item.clone(), diagnostics);
+            if !headers.is_empty() {
+                all_headers_in_block.extend(headers);
             }
         }
     }
 
-    // Normalize all headers in the block together
+    // normalize_headers adjusts header levels so the shallowest header in the
+    // scope becomes an h1
     normalize_headers(&mut all_headers_in_block);
 
-    // Debug: Print normalized headers (disabled)
-    // println!("PARSER: Normalized headers in block:");
-    // for (i, header) in all_headers_in_block.iter().enumerate() {
-    //     println!("  [{}] '{}' (Level: {})", i, header.title, header.level);
-    // }
+    // Lookup by span so we can reuse normalized headers later.
+    let mut header_lookup: HashMap<(usize, usize), std::sync::Arc<Header>> = HashMap::new();
+    for header in &all_headers_in_block {
+        header_lookup.insert((header.span.start, header.span.end), header.clone());
+    }
 
     // Second pass: process statements and expressions with normalized headers
     let mut current_headers: Vec<std::sync::Arc<Header>> = Vec::new();
@@ -826,26 +850,6 @@ pub fn parse_expr_block(token: Pair<'_>, diagnostics: &mut Diagnostics) -> Optio
                     continue;
                 }
             }
-            Rule::mdx_header => {
-                // Headers are already processed, just update current headers
-                let header = parse_header(item, diagnostics);
-                if let Some(header) = header {
-                    let header_arc = std::sync::Arc::new(header);
-
-                    // Find the corresponding normalized header
-                    if let Some(normalized_header) = all_headers_in_block
-                        .iter()
-                        .find(|h| h.title == header_arc.title)
-                    {
-                        // Implement header hierarchy logic
-                        filter_headers_by_hierarchy(&mut current_headers, normalized_header);
-
-                        // Add to current headers and headers since last statement
-                        current_headers.push(normalized_header.clone());
-                        headers_since_last_stmt.push(normalized_header.clone());
-                    }
-                }
-            }
             Rule::BLOCK_CLOSE => {
                 // Commentend out because we can't have blocks without return
                 // expressions otherwise. Plus we need functions with no return
@@ -863,8 +867,18 @@ pub fn parse_expr_block(token: Pair<'_>, diagnostics: &mut Diagnostics) -> Optio
                 continue;
             }
             Rule::comment_block => {
-                // Skip comments in function bodies
-                continue;
+                let headers = headers_from_comment_block(item, diagnostics);
+                if headers.is_empty() {
+                    continue;
+                }
+                for header in headers {
+                    attach_header_if_known(
+                        &header,
+                        &header_lookup,
+                        &mut current_headers,
+                        &mut headers_since_last_stmt,
+                    );
+                }
             }
             Rule::empty_lines => {
                 // Skip empty lines in function bodies
@@ -942,44 +956,75 @@ pub fn parse_expr_block(token: Pair<'_>, diagnostics: &mut Diagnostics) -> Optio
     })
 }
 
-/// Parse a single header from an MDX header token
-pub fn parse_header(token: Pair<'_>, diagnostics: &mut Diagnostics) -> Option<Header> {
-    let full_text = token.as_str();
-    let header_span = diagnostics.span(token.as_span());
-
-    // Find the start of the hash sequence
-    let hash_start = full_text.find('#')?;
-    let after_whitespace = full_text[hash_start..].trim_start();
-
-    // Count consecutive hash characters
-    let hash_count = after_whitespace.chars().take_while(|&c| c == '#').count();
+fn headers_from_comment_block(
+    token: Pair<'_>,
+    diagnostics: &mut Diagnostics,
+) -> Vec<std::sync::Arc<Header>> {
+    if token.as_rule() != Rule::comment_block {
+        return Vec::new();
+    }
 
-    // Extract the title after the hash sequence and whitespace
-    let after_hashes = &after_whitespace[hash_count..];
-    let title_text = after_hashes.trim().to_string();
+    let mut headers = Vec::new();
+    for current in token.into_inner() {
+        if current.as_rule() == Rule::comment {
+            if let Some(header) = parse_comment_header_pair(&current, diagnostics) {
+                headers.push(std::sync::Arc::new(header));
+            }
+        }
+    }
+    headers
+}
 
-    // Remove trailing newline if present
-    let title_text = title_text
-        .trim_end_matches('\n')
-        .trim_end_matches('\r')
-        .to_string();
+pub(crate) fn parse_comment_header_pair(
+    comment: &Pair<'_>,
+    diagnostics: &mut Diagnostics,
+) -> Option<Header> {
+    let span = diagnostics.span(comment.as_span());
+    let mut text = comment.as_str().trim_start();
+    if !text.starts_with("//") {
+        return None;
+    }
+    text = &text[2..];
+    let text = text.trim_start();
+    if !text.starts_with('#') {
+        return None;
+    }
 
-    let level = hash_count as u8;
+    let mut level = 0usize;
+    for ch in text.chars() {
+        if ch == '#' {
+            level += 1;
+        } else {
+            break;
+        }
+    }
+    if level == 0 {
+        return None;
+    }
 
-    // Print debug information about the header (disabled)
-    // let indent = " ".repeat(level as usize);
-    // println!(
-    //     "{}└ HEADER Level {}: '{}' (hash count: {})",
-    //     indent, level, title_text, level
-    // );
+    let rest = text[level..].trim().to_string();
 
     Some(Header {
-        level,
-        title: title_text,
-        span: header_span,
+        level: level as u8,
+        title: rest,
+        span,
     })
 }
 
+fn attach_header_if_known(
+    header: &std::sync::Arc<Header>,
+    lookup: &HashMap<(usize, usize), std::sync::Arc<Header>>,
+    current_headers: &mut Vec<std::sync::Arc<Header>>,
+    headers_since_last_stmt: &mut Vec<std::sync::Arc<Header>>,
+) {
+    let key = (header.span.start, header.span.end);
+    if let Some(normalized_header) = lookup.get(&key) {
+        filter_headers_by_hierarchy(current_headers, normalized_header);
+        current_headers.push(normalized_header.clone());
+        headers_since_last_stmt.push(normalized_header.clone());
+    }
+}
+
 /// Filter headers based on hierarchy rules (markdown-style nesting)
 fn filter_headers_by_hierarchy(
     pending_headers: &mut Vec<std::sync::Arc<Header>>,

engine/baml-lib/ast/src/parser/parse_expression.rs

@@ -711,19 +711,19 @@ mod tests {
     }
 
     #[test]
-    fn test_mdx_header_parsing() {
-        println!("\n=== Testing MDX Header Parsing ===");
+    fn test_comment_header_parsing() {
+        println!("\n=== Testing Comment Header Parsing ===");
 
         let input = r#"{
-            # Level 1 Header
+            //# Level 1 Header
             let x = "hello";
 
-            ## Level 2 Header
+            //## Level 2 Header
             let y = "world";
 
-            ########### Level 11 Header
+            //########### Level 11 Header
 
-            ### Level 3 Headers
+            //### Level 3 Headers
             x + y
         }"#;
 
@@ -732,7 +732,7 @@ mod tests {
         let mut diagnostics = Diagnostics::new(root_path.into());
         diagnostics.set_source(&source);
 
-        println!("Parsing expression block with mdx headers...");
+        println!("Parsing expression block with comment headers...");
 
         let pair_result = BAMLParser::parse(Rule::expr_block, input);
         match pair_result {
@@ -762,21 +762,21 @@ mod tests {
     fn test_complex_header_hierarchy() {
         println!("\n=== Testing Complex Header Hierarchy ===");
 
-        let input = r#"# Loop Processing
+        let input = r#"//# Loop Processing
 fn ForLoopWithHeaders() -> int {
     let items = [1, 2, 3, 4, 5];
     let result = 0;
 
-    ## Main Loop
+    //## Main Loop
     for (item in items) {
-        ### Item Processing
+        //### Item Processing
         let processed = item * 2;
 
-        #### Accumulation
+        //#### Accumulation
         result = result + processed;
     }
 
-    ## Final Result
+    //## Final Result
     result
 }"#;
 

engine/baml-lib/ast/src/parser/parse_value_expression_block.rs

@@ -127,7 +127,7 @@ pub(crate) fn parse_value_expression_block(
                             }
                         }
                         Rule::empty_lines => {}
-                        Rule::BLOCK_LEVEL_CATCH_ALL | Rule::mdx_header => {
+                        Rule::BLOCK_LEVEL_CATCH_ALL => {
                             diagnostics.push_error(DatamodelError::new_validation_error(
                                 "This line is not a valid field or attribute definition. A valid property may look like: 'myProperty \"some value\"' for example, with no colons.",
                                 diagnostics.span(item.as_span()),