Parser: add explicit body discriminators

crates/kirin-chumsky/src/ast/blocks.rs

@@ -42,14 +42,19 @@ pub struct BlockHeader<'src, TypeOutput> {
 
 /// A basic block containing a label, arguments, and statements.
 ///
-/// Represents syntax like:
+/// Represents the untagged block syntax:
 /// ```ignore
 /// ^bb0(%arg: i32) {
 ///     %x = add %arg, %arg;
 ///     return %x;
 /// }
 /// ```
 ///
+/// This AST node covers both roles the untagged form plays: a CFG's member
+/// block (where the enclosing `cfg { .. }` supplies the discriminator) and a
+/// standalone `Block` body, which the parser and printer tag with the `block`
+/// keyword around this same shape.
+///
 /// Fields are flat to support both full parsing (with block header) and
 /// projection-based parsing (where pieces come from different format positions).
 ///
@@ -69,12 +74,15 @@ pub struct Block<'src, TypeOutput, StmtOutput> {
 ///
 /// Represents syntax like:
 /// ```ignore
-/// {
+/// cfg {
 ///     ^entry(%arg: i32) { ... };
 ///     ^bb1() { ... };
 /// }
 /// ```
 ///
+/// The `cfg` discriminator belongs to the whole container; the member blocks
+/// stay untagged.
+///
 /// The `TypeOutput` parameter is the parsed type representation.
 /// The `StmtOutput` parameter is the parsed statement representation.
 #[derive(Debug, Clone, PartialEq)]

crates/kirin-chumsky/src/function_text/parse_text.rs

@@ -43,7 +43,7 @@
 //!
 //! ```text
 //! stage @A fn @foo(()) -> ();
-//! specialize @A fn @foo(()) -> () { ^0() {} }
+//! specialize @A fn @foo(()) -> () cfg { ^0() {} }
 //! ```
 //!
 //! - Pass 1 creates/finds function `@foo` and staged function `(A, foo)`.
@@ -53,9 +53,9 @@
 //!
 //! ```text
 //! stage @A fn @foo(()) -> ();
-//! specialize @A fn @foo(()) -> () { ^0() {} }
+//! specialize @A fn @foo(()) -> () cfg { ^0() {} }
 //! stage @B fn @bar(i32) -> i32;
-//! specialize @B fn @bar(i32) -> i32 { ^0() {} }
+//! specialize @B fn @bar(i32) -> i32 cfg { ^0() {} }
 //! ```
 //!
 //! - declarations for `@A` are parsed with stage `A`'s dialect;
@@ -64,7 +64,7 @@
 //! Missing header before specialize:
 //!
 //! ```text
-//! specialize @A fn @missing(()) -> () { ^0() {} }
+//! specialize @A fn @missing(()) -> () cfg { ^0() {} }
 //! ```
 //!
 //! - pass 2 cannot find `(A, missing)` in the staged lookup;

crates/kirin-chumsky/src/function_text/syntax.rs

@@ -69,18 +69,35 @@ where
         .labelled("function signature")
 }
 
-/// Body span scanner. Matches an optional keyword prefix (e.g. `digraph`,
-/// `ungraph`) followed by a brace-balanced `{ ... }` CFG. Returns the
-/// span covering everything from the first non-brace token (or the opening
-/// brace) through the matching closing brace. Does not parse body contents.
+/// Body span scanner. Skips the body's discriminator and header — whatever the
+/// dialect's format string puts before the first `{` — then matches a
+/// brace-balanced `{ ... }`. Returns the span covering everything from the
+/// first token through the matching closing brace. Does not parse body
+/// contents; that is the dialect statement parser's job, which is what keeps
+/// dialect-level validation intact.
+///
+/// All four body kinds carry an explicit textual discriminator, and each is
+/// scanned by the same rule:
+///
+/// ```text
+/// fn @f(..) -> T cfg { ^entry(..) { .. } }      // keyword, then the CFG's braces
+/// fn @f(..) -> T block ^body(..) { .. }         // keyword + header, then braces
+/// fn @f(..) -> T digraph ^g0(..) { .. }         // keyword + header, then braces
+/// fn @f(..) -> T ungraph ^u0(..) { .. }         // keyword + header, then braces
+/// ```
+///
+/// Projected formats (`fn @f(..) -> T (%x: T) { .. }`) work the same way: the
+/// scanner does not care what the prefix tokens are, only where the first `{`
+/// is.
 fn body_span<'src, I>() -> impl Parser<'src, I, SimpleSpan, ParserError<'src>>
 where
     I: TokenInput<'src>,
 {
     chumsky::primitive::custom(|input: &mut chumsky::input::InputRef<'src, '_, I, _>| {
         let start = input.cursor();
-        // Skip tokens until we find the opening brace. This allows keyword
-        // prefixes like `digraph ^name(ports...) {` or `ungraph ^name(...) {`.
+        // Skip tokens until we find the opening brace. This is what lets the
+        // discriminator and header through: `cfg {`, `block ^name(args...) {`,
+        // `digraph ^name(ports...) {`, `ungraph ^name(...) {`.
         loop {
             match input.next() {
                 Some(Token::LBrace) => break,

crates/kirin-chumsky/src/function_text/tests.rs

@@ -131,7 +131,7 @@ enum MixedStage {
 // Helpers
 // ---------------------------------------------------------------------------
 
-const BODY: &str = "{ ^0() {} }";
+const BODY: &str = "cfg { ^0() {} }";
 
 fn unit_sig() -> Signature<UnitType> {
     Signature::new(vec![UnitType], UnitType, ())
@@ -445,7 +445,7 @@ fn test_invalid_body_parse_has_source() {
     let mut pipeline: Pipeline<StageInfo<FunctionBody>> = Pipeline::new();
     // Valid header but invalid body tokens
     let err = pipeline
-        .parse("stage @A fn @foo(()) -> (); specialize @A fn @foo(()) -> () { invalid }")
+        .parse("stage @A fn @foo(()) -> (); specialize @A fn @foo(()) -> () cfg { invalid }")
         .unwrap_err();
     // The body parse failure should chain a source error
     // (may or may not have source depending on where it fails)
@@ -478,3 +478,98 @@ fn test_invalid_declaration_keyword() {
     let err = pipeline.parse("define @A fn @foo(()) -> ();").unwrap_err();
     assert_eq!(err.kind, crate::FunctionParseErrorKind::InvalidHeader);
 }
+
+// ---------------------------------------------------------------------------
+// Body-span scanner
+//
+// The framework only strips `specialize @stage`; the rest of the declaration —
+// discriminator, header, and brace-balanced body — is handed to the dialect
+// statement parser as one span. The scanner skips tokens until the first `{`,
+// so it is agnostic to *which* discriminator precedes the body, and stays
+// correct for all four body kinds without knowing any of them.
+// ---------------------------------------------------------------------------
+
+/// Scan one `specialize` declaration and return the exact source slice handed
+/// to the dialect statement parser.
+///
+/// `LowerBody`'s `I32Type` only affects the *header* type list; the scanner
+/// never parses body contents, so any body kind can be scanned through it.
+fn scanned_body_span(src: &str) -> String {
+    let tokens = super::syntax::tokenize(src);
+    let (declaration, _) = super::syntax::parse_one_declaration::<LowerBody>(&tokens)
+        .expect("declaration should scan");
+    match declaration {
+        super::syntax::Declaration::Specialize { body_span, .. } => {
+            src[body_span.start..body_span.end].to_string()
+        }
+        super::syntax::Declaration::Stage(_) => panic!("expected a specialize declaration"),
+    }
+}
+
+#[test]
+fn test_body_span_scans_tagged_cfg() {
+    let body = "fn @f(i32) -> i32 cfg { ^entry(%x: i32) { %r = add %x, %x; ret %r; } }";
+    assert_eq!(scanned_body_span(&format!("specialize @A {body}")), body);
+}
+
+#[test]
+fn test_body_span_scans_tagged_block() {
+    let body = "fn @f(i32) -> i32 block ^body(%x: i32) { %r = add %x, %x; ret %r; }";
+    assert_eq!(scanned_body_span(&format!("specialize @A {body}")), body);
+}
+
+#[test]
+fn test_body_span_scans_digraph() {
+    let body = "fn @f(i32) -> i32 digraph ^g0(%x: i32) { %r = add %x, %x; yield %r; }";
+    assert_eq!(scanned_body_span(&format!("specialize @A {body}")), body);
+}
+
+#[test]
+fn test_body_span_scans_ungraph() {
+    let body = "fn @f(i32) -> i32 ungraph ^u0(%x: i32) { edge %w = wire; node(%x, %w); }";
+    assert_eq!(scanned_body_span(&format!("specialize @A {body}")), body);
+}
+
+#[test]
+fn test_body_span_scans_a_projected_body_without_a_discriminator() {
+    // Projections stay raw, so a dialect can spell its own wrapper. The
+    // scanner does not require a keyword — only a brace-balanced body.
+    let body = "fn @f(i32) -> i32 (%x: i32) { %r = add %x, %x; ret %r; }";
+    assert_eq!(scanned_body_span(&format!("specialize @A {body}")), body);
+}
+
+#[test]
+fn test_body_span_stops_at_the_matching_brace() {
+    // Two declarations: the first span must end at *its* closing brace, not
+    // run on into the second.
+    let first = "fn @f(i32) -> i32 cfg { ^entry(%x: i32) { ret %x; } }";
+    let second = "fn @g(i32) -> i32 cfg { ^e { } }";
+    let src = format!("specialize @A {first} specialize @A {second}");
+    assert_eq!(scanned_body_span(&src), first);
+}
+
+#[test]
+fn test_body_span_requires_an_opening_brace() {
+    let tokens = super::syntax::tokenize("specialize @A fn @f(i32) -> i32 cfg;");
+    let errors = super::syntax::parse_one_declaration::<LowerBody>(&tokens)
+        .expect_err("a body with no `{` should not scan");
+    assert!(
+        errors
+            .iter()
+            .any(|e| format!("{e}").contains("expected '{'")),
+        "expected a missing-brace diagnostic, got: {errors:?}"
+    );
+}
+
+#[test]
+fn test_body_span_rejects_an_unclosed_brace() {
+    let tokens = super::syntax::tokenize("specialize @A fn @f(i32) -> i32 cfg { ^entry {");
+    let errors = super::syntax::parse_one_declaration::<LowerBody>(&tokens)
+        .expect_err("an unbalanced body should not scan");
+    assert!(
+        errors
+            .iter()
+            .any(|e| format!("{e}").contains("unclosed '{'")),
+        "expected an unclosed-brace diagnostic, got: {errors:?}"
+    );
+}

crates/kirin-chumsky/src/parsers/blocks.rs

@@ -120,14 +120,25 @@ where
         .labelled("block header")
 }
 
-/// Parses a complete block with header and statements.
+/// Parses a block with header and statements, **without** the `block`
+/// discriminator keyword.
+///
+/// Matches: `^bb0(%arg: i32) { ... }`
+///
+/// This is the shared block grammar. It is internal on purpose: the two public
+/// entry points differ only in what wraps it.
+///
+/// - [`block`] prefixes the `block` keyword — a standalone `Block` body is
+///   always tagged.
+/// - [`cfg()`] and [`cfg_body()`] use it directly — a CFG's member blocks stay
+///   untagged, because the enclosing `cfg { ... }` already names the body kind.
 ///
 /// Requires a parser for the language/dialect statements.
 ///
 /// The type parameter `T` specifies the type annotation type (typically the TypeLattice).
 /// The type parameter `S` is the statement AST type produced by the language parser.
 /// The parser produces `Block<'t, <T as HasParser>::Output, S>`.
-pub fn block<'t, I, T, S>(
+fn untagged_block<'t, I, T, S>(
     language: RecursiveParser<'t, I, S>,
 ) -> impl Parser<'t, I, Spanned<Block<'t, <T as HasParser<'t>>::Output, S>>, ParserError<'t>>
 where
@@ -162,18 +173,56 @@ where
     })
 }
 
-/// Parses a CFG containing multiple blocks.
+/// Parses a complete standalone `Block` body, including its `block` discriminator.
 ///
 /// Matches:
 /// ```text
-/// {
+/// block ^bb0(%arg: i32) {
+///     %x = add %arg, %arg;
+///     ret %x;
+/// }
+/// ```
+///
+/// This is the parser behind the default `{body}` interpolation of a `Block`
+/// field. The `block` keyword is **required** — the untagged form
+/// `^bb0(...) { ... }` is only valid for a CFG's member blocks, where
+/// [`cfg()`] supplies the discriminator once for the whole body.
+///
+/// Requires a parser for the language/dialect statements.
+///
+/// The type parameter `T` specifies the type annotation type (typically the TypeLattice).
+/// The type parameter `S` is the statement AST type produced by the language parser.
+/// The parser produces `Block<'t, <T as HasParser>::Output, S>`.
+pub fn block<'t, I, T, S>(
+    language: RecursiveParser<'t, I, S>,
+) -> impl Parser<'t, I, Spanned<Block<'t, <T as HasParser<'t>>::Output, S>>, ParserError<'t>>
+where
+    I: TokenInput<'t>,
+    T: HasParser<'t>,
+    S: Clone,
+{
+    just(Token::Identifier("block"))
+        .ignore_then(untagged_block::<_, T, S>(language))
+        .labelled("block")
+}
+
+/// Parses a CFG containing multiple blocks, including its `cfg` discriminator.
+///
+/// Matches:
+/// ```text
+/// cfg {
 ///     ^bb0(%arg: i32) {
 ///         %x = add %arg, %arg;
-///         return %x;
+///         ret %x;
 ///     }
 /// }
 /// ```
 ///
+/// This is the parser behind the default `{body}` interpolation of a `CFG`
+/// field. The `cfg` keyword is **required**; the member blocks are untagged
+/// (`^bb0 { ... }`, never `block ^bb0 { ... }`) because `cfg` already names
+/// the body kind for the whole container.
+///
 /// The type parameter `T` specifies the type annotation type (typically the TypeLattice).
 /// The type parameter `S` is the statement AST type produced by the language parser.
 /// The parser produces `CFG<'t, <T as HasParser>::Output, S>`.
@@ -185,20 +234,24 @@ where
     T: HasParser<'t>,
     S: Clone,
 {
-    block::<_, T, S>(language)
-        .then_ignore(just(Token::Semicolon).or_not())
-        .repeated()
-        .collect::<Vec<_>>()
-        .delimited_by(just(Token::LBrace), just(Token::RBrace))
+    just(Token::Identifier("cfg"))
+        .ignore_then(
+            untagged_block::<_, T, S>(language)
+                .then_ignore(just(Token::Semicolon).or_not())
+                .repeated()
+                .collect::<Vec<_>>()
+                .delimited_by(just(Token::LBrace), just(Token::RBrace)),
+        )
         .map(|blocks| CFG { blocks })
         .labelled("cfg")
 }
 
 /// Parses block body statements (without header, without braces).
 ///
 /// Matches a sequence of `statement ;` pairs. This is the inner content of
-/// a block body, used for `:body` projections on Block fields where the
-/// caller provides surrounding syntax via the format string.
+/// a block body, used for `{field:body}` projections on Block fields where the
+/// caller provides surrounding syntax via the format string. It stays raw: no
+/// `block` keyword and no braces are injected.
 pub fn block_body_statements<'t, I, S>(
     language: RecursiveParser<'t, I, S>,
 ) -> impl Parser<'t, I, Vec<Spanned<S>>, ParserError<'t>>
@@ -217,11 +270,13 @@ where
         .labelled("block body statements")
 }
 
-/// Parses CFG body (blocks without outer braces).
+/// Parses CFG body (untagged blocks, without the `cfg` keyword or outer braces).
 ///
-/// Matches a sequence of blocks, each optionally terminated by a semicolon.
-/// This is the inner content of a CFG, used for `:body` projections on
-/// CFG fields where the caller provides surrounding syntax via the format string.
+/// Matches a sequence of untagged blocks, each optionally terminated by a
+/// semicolon. This is the inner content of a CFG, used for `{field:body}`
+/// projections on CFG fields where the caller provides surrounding syntax via
+/// the format string. It stays raw: no `cfg` keyword is injected, so a dialect
+/// author can spell their own wrapper.
 pub fn cfg_body<'t, I, T, S>(
     language: RecursiveParser<'t, I, S>,
 ) -> impl Parser<'t, I, Vec<Spanned<Block<'t, <T as HasParser<'t>>::Output, S>>>, ParserError<'t>>
@@ -230,7 +285,7 @@ where
     T: HasParser<'t>,
     S: Clone,
 {
-    block::<_, T, S>(language)
+    untagged_block::<_, T, S>(language)
         .then_ignore(just(Token::Semicolon).or_not())
         .repeated()
         .collect::<Vec<_>>()