onedrop_hlsl/

parse.rs

//! HLSL recursive-descent parser.
//!
//! Consumes the [`Token`] stream from [`crate::lex::tokenize`] and produces a
//! [`TranslationUnit`] AST. Pratt-style operator precedence for expressions,
//! straightforward recursive descent for statements and top-level items.
//!
//! The parser is pragmatic, not exhaustive: it handles the HLSL subset that
//! MilkDrop 2 user comp shaders actually use. Constructs outside that subset
//! (templates, attributes, geometry-shader semantics, etc.) produce a
//! [`ParseError`] — the caller falls back to the existing regex pipeline.
//!
//! This parser is not yet wired into [`crate::translate_shader`]. The AST
//! sits available for a future WGSL emitter, while the existing regex passes
//! continue to drive the current comp-shader pass-rate on
//! `test-presets-200/`.

use crate::ast::*;
use crate::lex::{Keyword, LexError, Span, Token, TokenKind, tokenize};
use std::fmt;

/// Parse error: human-readable message plus the offending source span.
#[derive(Debug, Clone, PartialEq)]
pub struct ParseError {
    pub message: String,
    pub span: Span,
}

impl fmt::Display for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "parse error at {}:{}: {}",
            self.span.line, self.span.col, self.message
        )
    }
}

impl std::error::Error for ParseError {}

impl From<LexError> for ParseError {
    fn from(e: LexError) -> Self {
        Self {
            message: e.message,
            span: Span {
                start: e.offset,
                end: e.offset,
                line: e.line,
                col: e.col,
            },
        }
    }
}

/// Top-level entry point: lex then parse a full HLSL source string.
///
/// Accepts either a bare statement block (no `shader_body` wrapper) or a
/// full MD2 comp shader (`<global decls> shader_body { ... }`). Returns the
/// translation unit AST or a [`ParseError`] at the first unrecoverable
/// position.
pub fn parse_hlsl(src: &str) -> Result<TranslationUnit, ParseError> {
    let tokens = tokenize(src)?;
    let mut p = Parser::new(&tokens, src);
    p.parse_translation_unit()
}

/// Hand-written recursive-descent parser. Holds a borrow on the token slice
/// and source string; produces AST nodes that own their data (so the
/// parser's lifetime doesn't leak into downstream consumers).
struct Parser<'a> {
    tokens: &'a [Token],
    src: &'a str,
    pos: usize,
}

impl<'a> Parser<'a> {
    fn new(tokens: &'a [Token], src: &'a str) -> Self {
        Self {
            tokens,
            src,
            pos: 0,
        }
    }

    // -------- token cursor helpers --------

    fn peek(&self) -> &Token {
        &self.tokens[self.pos]
    }

    fn peek_at(&self, offset: usize) -> Option<&Token> {
        self.tokens.get(self.pos + offset)
    }

    fn bump(&mut self) -> &Token {
        let t = &self.tokens[self.pos];
        if !matches!(t.kind, TokenKind::Eof) {
            self.pos += 1;
        }
        t
    }

    fn at(&self, kind: &TokenKind) -> bool {
        std::mem::discriminant(&self.peek().kind) == std::mem::discriminant(kind)
    }

    fn eat(&mut self, kind: &TokenKind) -> bool {
        if self.at(kind) {
            self.bump();
            true
        } else {
            false
        }
    }

    fn expect(&mut self, kind: TokenKind, ctx: &str) -> Result<(), ParseError> {
        if std::mem::discriminant(&self.peek().kind) == std::mem::discriminant(&kind) {
            self.bump();
            Ok(())
        } else {
            Err(ParseError {
                message: format!("expected {ctx}, found {:?}", self.peek().kind),
                span: self.peek().span,
            })
        }
    }

    fn ident_text(&self, span: Span) -> &'a str {
        &self.src[span.start as usize..span.end as usize]
    }

    /// Consume the current token if it's an `Ident` and return its text.
    /// Errors with `ctx` in the message otherwise.
    fn expect_ident(&mut self, ctx: &str) -> Result<(String, Span), ParseError> {
        let t = self.peek();
        if matches!(t.kind, TokenKind::Ident) {
            let span = t.span;
            let s = self.ident_text(span).to_string();
            self.bump();
            Ok((s, span))
        } else {
            Err(ParseError {
                message: format!("expected {ctx}, found {:?}", t.kind),
                span: t.span,
            })
        }
    }

    // -------- top-level --------

    fn parse_translation_unit(&mut self) -> Result<TranslationUnit, ParseError> {
        let start_span = self.peek().span;
        let mut items = Vec::new();
        let mut shader_body = None;

        while !matches!(self.peek().kind, TokenKind::Eof) {
            if matches!(self.peek().kind, TokenKind::Keyword(Keyword::ShaderBody)) {
                let kw_span = self.peek().span;
                self.bump();
                self.expect(TokenKind::LBrace, "`{` after `shader_body`")?;
                shader_body = Some(self.parse_block_after_brace(kw_span)?);
                break;
            }
            // Sampler decl: `sampler[2D|3D] <name>;`
            if matches!(self.peek().kind, TokenKind::Keyword(Keyword::Sampler)) {
                items.push(self.parse_sampler_decl()?);
                continue;
            }
            // `[static] [const] <type> <name>[, <name>]... ;` global var.
            // Detected by leading qualifier OR by `<type-ident> <name-ident>`
            // shape (when the third token is not `(`, which would mean a
            // function definition).
            if matches!(
                self.peek().kind,
                TokenKind::Keyword(Keyword::Static) | TokenKind::Keyword(Keyword::Const)
            ) {
                self.parse_global_var_decl(&mut items)?;
                continue;
            }
            // Function vs. global-var disambiguation: both start with
            // `<Ident> <Ident>`. The third token tells us which:
            // - `(`     → function
            // - any other → global variable
            if self.is_function_signature() {
                items.push(Item::Function(self.parse_function_def()?));
                continue;
            }
            if self.is_global_var_decl() {
                self.parse_global_var_decl(&mut items)?;
                continue;
            }
            // Anything else at top level is an error.
            return Err(ParseError {
                message: format!("unexpected top-level token {:?}", self.peek().kind),
                span: self.peek().span,
            });
        }

        let end_span = self.peek().span;
        Ok(TranslationUnit {
            items,
            shader_body,
            span: start_span.merge(end_span),
        })
    }

    /// Peek 3 tokens to decide if we're looking at a function definition:
    /// `<type-ident> <name-ident> (`. Doesn't consume.
    fn is_function_signature(&self) -> bool {
        let t0 = matches!(
            self.peek().kind,
            TokenKind::Ident | TokenKind::Keyword(Keyword::Void)
        );
        let t1 = matches!(self.peek_at(1).map(|t| &t.kind), Some(TokenKind::Ident));
        let t2 = matches!(self.peek_at(2).map(|t| &t.kind), Some(TokenKind::LParen));
        t0 && t1 && t2
    }

    /// Top-level global var decl detection: `<type-ident> <name-ident>`
    /// where the third token is *not* `(`. Doesn't consume.
    fn is_global_var_decl(&self) -> bool {
        let t0 = matches!(self.peek().kind, TokenKind::Ident);
        let t1 = matches!(self.peek_at(1).map(|t| &t.kind), Some(TokenKind::Ident));
        let t2_not_paren = !matches!(self.peek_at(2).map(|t| &t.kind), Some(TokenKind::LParen));
        t0 && t1 && t2_not_paren
    }

    /// Parse `[static] [const] <type> <name>[<array>] [= <init>] [, <name> ...] ;`
    /// and flatten the comma list into one [`Item::GlobalVar`] per name.
    /// The qualifier flags propagate to every flattened variable.
    fn parse_global_var_decl(&mut self, items: &mut Vec<Item>) -> Result<(), ParseError> {
        let start = self.peek().span;
        let mut is_static = false;
        let mut is_const = false;
        loop {
            match self.peek().kind {
                TokenKind::Keyword(Keyword::Static) => {
                    is_static = true;
                    self.bump();
                }
                TokenKind::Keyword(Keyword::Const) => {
                    is_const = true;
                    self.bump();
                }
                _ => break,
            }
        }
        let ty = self.parse_type()?;
        loop {
            let (name, name_span) = self.expect_ident("global variable name")?;
            let array_len = if self.eat(&TokenKind::LBracket) {
                let len = self.parse_expr()?;
                self.expect(TokenKind::RBracket, "`]` after array length")?;
                Some(len)
            } else {
                None
            };
            let init = if self.eat(&TokenKind::Eq) {
                Some(self.parse_expr()?)
            } else {
                None
            };
            items.push(Item::GlobalVar(GlobalVar {
                is_static,
                is_const,
                ty: ty.clone(),
                name,
                array_len,
                init,
                span: start.merge(name_span),
            }));
            if !self.eat(&TokenKind::Comma) {
                break;
            }
        }
        self.expect(TokenKind::Semi, "`;` at end of global declaration")?;
        Ok(())
    }

    fn parse_sampler_decl(&mut self) -> Result<Item, ParseError> {
        let kw_tok = self.bump();
        let kw_span = kw_tok.span;
        let kw_text = self.ident_text(kw_span);
        let tag = match kw_text {
            "sampler2D" => SamplerTag::Sampler2D,
            "sampler3D" => SamplerTag::Sampler3D,
            _ => SamplerTag::Sampler,
        };
        let (name, name_span) = self.expect_ident("sampler name")?;
        // Optional `: register(s0)` or `= sampler_state { ... };` — both are
        // legal HLSL. We swallow up to the next top-level `;` so the parser
        // doesn't need to model the state block.
        while !matches!(self.peek().kind, TokenKind::Semi | TokenKind::Eof) {
            // Skip a brace-balanced block if we see one.
            if matches!(self.peek().kind, TokenKind::LBrace) {
                self.skip_braced_block()?;
                continue;
            }
            self.bump();
        }
        self.expect(TokenKind::Semi, "`;` after sampler declaration")?;
        Ok(Item::SamplerDecl(SamplerDecl {
            tag,
            name,
            span: kw_span.merge(name_span),
        }))
    }

    /// `{ ... }` skipper for opaque blocks (sampler_state, struct bodies in
    /// the rare cases MD2 uses them). Caller positions us at `{`; we end
    /// after the matching `}`.
    fn skip_braced_block(&mut self) -> Result<(), ParseError> {
        let open = self.peek().span;
        self.expect(TokenKind::LBrace, "`{`")?;
        let mut depth = 1i32;
        while depth > 0 {
            match self.peek().kind {
                TokenKind::LBrace => {
                    depth += 1;
                    self.bump();
                }
                TokenKind::RBrace => {
                    depth -= 1;
                    self.bump();
                }
                TokenKind::Eof => {
                    return Err(ParseError {
                        message: "unterminated `{` block".to_string(),
                        span: open,
                    });
                }
                _ => {
                    self.bump();
                }
            }
        }
        Ok(())
    }

    fn parse_function_def(&mut self) -> Result<FunctionDef, ParseError> {
        let start = self.peek().span;
        let return_type = self.parse_type()?;
        let (name, _) = self.expect_ident("function name")?;
        self.expect(TokenKind::LParen, "`(` after function name")?;
        let mut params = Vec::new();
        if !matches!(self.peek().kind, TokenKind::RParen) {
            loop {
                params.push(self.parse_param()?);
                if !self.eat(&TokenKind::Comma) {
                    break;
                }
            }
        }
        self.expect(TokenKind::RParen, "`)` after parameter list")?;
        let body = self.parse_block()?;
        let end = body.span;
        Ok(FunctionDef {
            return_type,
            name,
            params,
            body,
            span: start.merge(end),
        })
    }

    fn parse_param(&mut self) -> Result<Param, ParseError> {
        let start = self.peek().span;
        let qualifier = match self.peek().kind {
            TokenKind::Keyword(Keyword::In) => {
                self.bump();
                Some(ParamQualifier::In)
            }
            TokenKind::Keyword(Keyword::Out) => {
                self.bump();
                Some(ParamQualifier::Out)
            }
            TokenKind::Keyword(Keyword::InOut) => {
                self.bump();
                Some(ParamQualifier::InOut)
            }
            _ => None,
        };
        let ty = self.parse_type()?;
        let (name, name_span) = self.expect_ident("parameter name")?;
        // HLSL allows `: SEMANTIC` after a parameter — skip silently.
        if self.eat(&TokenKind::Colon) {
            // Eat one identifier (the semantic).
            if matches!(self.peek().kind, TokenKind::Ident) {
                self.bump();
            }
        }
        Ok(Param {
            qualifier,
            ty,
            name,
            span: start.merge(name_span),
        })
    }

    /// `<type-ident>` — one identifier consumed. We do not look ahead for
    /// matrix bracket syntax; HLSL spells matrices as compound idents
    /// (`float2x2`, `mat3x3`).
    fn parse_type(&mut self) -> Result<TypeRef, ParseError> {
        if matches!(self.peek().kind, TokenKind::Keyword(Keyword::Void)) {
            let span = self.peek().span;
            self.bump();
            return Ok(TypeRef::new("void", span));
        }
        if matches!(self.peek().kind, TokenKind::Ident) {
            let span = self.peek().span;
            let name = self.ident_text(span).to_string();
            self.bump();
            return Ok(TypeRef::new(name, span));
        }
        Err(ParseError {
            message: format!("expected type name, found {:?}", self.peek().kind),
            span: self.peek().span,
        })
    }

    // -------- statements --------

    fn parse_block(&mut self) -> Result<Block, ParseError> {
        let open_span = self.peek().span;
        self.expect(TokenKind::LBrace, "`{` at start of block")?;
        self.parse_block_after_brace(open_span)
    }

    fn parse_block_after_brace(&mut self, open_span: Span) -> Result<Block, ParseError> {
        let mut stmts = Vec::new();
        while !matches!(self.peek().kind, TokenKind::RBrace | TokenKind::Eof) {
            self.parse_stmt_unit(&mut stmts)?;
        }
        let close_span = self.peek().span;
        self.expect(TokenKind::RBrace, "`}` at end of block")?;
        Ok(Block {
            stmts,
            span: open_span.merge(close_span),
        })
    }

    /// Parse one *syntactic* statement unit and push the resulting Stmt(s)
    /// onto `out`. Most kinds push exactly one statement, but two HLSL
    /// idioms produce multiple:
    ///
    /// - Multi-name declarations (`float a, b, c;`) — one Stmt per name.
    /// - Comma-separated statements (`a += 1, b += 2;`) — one Stmt per
    ///   comma-separated assignment/expression.
    ///
    /// The block parser calls this in a loop so the caller-side AST sees a
    /// flat statement list — convenient for downstream emitters that don't
    /// care about source-line grouping.
    fn parse_stmt_unit(&mut self, out: &mut Vec<Stmt>) -> Result<(), ParseError> {
        let single = match self.peek().kind {
            TokenKind::LBrace => Some(Stmt::Block(self.parse_block()?)),
            TokenKind::Keyword(Keyword::If) => Some(self.parse_if()?),
            TokenKind::Keyword(Keyword::While) => Some(self.parse_while()?),
            TokenKind::Keyword(Keyword::Do) => Some(self.parse_do_while()?),
            TokenKind::Keyword(Keyword::For) => Some(self.parse_for()?),
            TokenKind::Keyword(Keyword::Return) => Some(self.parse_return()?),
            TokenKind::Keyword(Keyword::Break) => {
                self.bump();
                self.expect(TokenKind::Semi, "`;` after `break`")?;
                Some(Stmt::Break)
            }
            TokenKind::Keyword(Keyword::Continue) => {
                self.bump();
                self.expect(TokenKind::Semi, "`;` after `continue`")?;
                Some(Stmt::Continue)
            }
            TokenKind::Semi => {
                self.bump();
                None
            }
            _ => {
                self.parse_decl_or_expr_stmts(out)?;
                return Ok(());
            }
        };
        if let Some(s) = single {
            out.push(s);
        }
        Ok(())
    }

    /// Compatibility shim — many internal callsites (for-loop init, etc.)
    /// still want a single Stmt. Calls [`parse_stmt_unit`] and pulls the
    /// first emitted statement; subsequent statements (rare in those
    /// contexts) are concatenated as a synthetic [`Stmt::Block`] so no
    /// declared variables go missing.
    fn parse_stmt(&mut self) -> Result<Stmt, ParseError> {
        let mut buf = Vec::with_capacity(1);
        self.parse_stmt_unit(&mut buf)?;
        match buf.len() {
            0 => Ok(Stmt::Block(Block {
                stmts: Vec::new(),
                span: self.peek().span,
            })),
            1 => Ok(buf.into_iter().next().unwrap()),
            _ => {
                let span = buf
                    .first()
                    .map(stmt_span)
                    .unwrap_or(Span::dummy())
                    .merge(buf.last().map(stmt_span).unwrap_or(Span::dummy()));
                Ok(Stmt::Block(Block { stmts: buf, span }))
            }
        }
    }

    /// The tricky dispatch: a statement starts with an `Ident`. Decide
    /// between local decl (`<type> <name>`) and expression statement
    /// (`<expr>;` or `<lhs> [op]= <rhs>;`).
    ///
    /// Heuristic: if the lookahead is `Ident Ident`, it's a decl. We also
    /// allow the leading qualifier `const` for completeness even though MD2
    /// rarely uses it locally.
    fn parse_decl_or_expr_stmts(&mut self, out: &mut Vec<Stmt>) -> Result<(), ParseError> {
        // Strip leading `const` if present — local `const` decls are legal
        // in HLSL even though MD2 rarely uses them.
        let saved_pos = self.pos;
        let mut had_const = false;
        if matches!(self.peek().kind, TokenKind::Keyword(Keyword::Const)) {
            self.bump();
            had_const = true;
        }

        if matches!(self.peek().kind, TokenKind::Ident)
            && matches!(self.peek_at(1).map(|t| &t.kind), Some(TokenKind::Ident))
        {
            return self.parse_local_decl_tail(out);
        }

        // Not a decl — rewind the `const` skip and parse as expression
        // statement.
        if had_const {
            self.pos = saved_pos;
        }
        self.parse_expr_or_assign_stmts(out)
    }

    /// Parse the rest of a local decl after the type+name lookahead
    /// matched. Eats: `<type> <name>[<dim>] [= <init>] [, <name>...] ;`.
    /// Pushes one [`Stmt::LocalDecl`] per name onto `out`.
    fn parse_local_decl_tail(&mut self, out: &mut Vec<Stmt>) -> Result<(), ParseError> {
        let start = self.peek().span;
        let ty = self.parse_type()?;
        loop {
            let (name, name_span) = self.expect_ident("variable name")?;
            let mut array_len = None;
            if self.eat(&TokenKind::LBracket) {
                array_len = Some(self.parse_expr()?);
                self.expect(TokenKind::RBracket, "`]` after array length")?;
            }
            let init = if self.eat(&TokenKind::Eq) {
                // Chained-assign init — HLSL accepts `T x = y = z = expr;`,
                // evaluating right-to-left. parse_expr stops at `=` (not a
                // binary op), so each `=` we still see after the init expr
                // is the next link in the chain. We collect the chain and
                // emit each intermediate `lhs = rhs` as its own assignment
                // statement *before* the LocalDecl; the leftmost LHS then
                // initialises the new variable. WGSL has no assignment
                // expression, so lowering at parse time keeps the rest of
                // the pipeline free of `var x = (y = ...)` shapes.
                let mut chain = vec![self.parse_expr()?];
                while self.eat(&TokenKind::Eq) {
                    chain.push(self.parse_expr()?);
                }
                for i in (0..chain.len().saturating_sub(1)).rev() {
                    let target = chain[i].clone();
                    let value = chain[i + 1].clone();
                    let span = target.span().merge(value.span());
                    out.push(Stmt::Assign(AssignStmt {
                        target,
                        op: AssignOp::Set,
                        value,
                        span,
                    }));
                }
                Some(chain.into_iter().next().unwrap())
            } else {
                None
            };
            out.push(Stmt::LocalDecl(LocalDecl {
                ty: ty.clone(),
                name,
                array_len,
                init,
                span: start.merge(name_span),
            }));
            if !self.eat(&TokenKind::Comma) {
                break;
            }
        }
        self.expect(TokenKind::Semi, "`;` at end of declaration")?;
        Ok(())
    }

    /// Parse one or more comma-separated `<expr>` or `<lhs> [op]= <rhs>`
    /// statements, terminated by `;`. MD2 user shaders use the comma
    /// form (`a += 1, b += 2;`) freely; we flatten into multiple
    /// statements at parse time so downstream emitters don't need to
    /// understand the comma operator.
    fn parse_expr_or_assign_stmts(&mut self, out: &mut Vec<Stmt>) -> Result<(), ParseError> {
        loop {
            let start = self.peek().span;
            let lhs = self.parse_expr()?;
            if let Some(op) = self.peek_assign_op() {
                self.bump();
                let rhs = self.parse_expr()?;
                let span = start.merge(rhs.span());
                out.push(Stmt::Assign(AssignStmt {
                    target: lhs,
                    op,
                    value: rhs,
                    span,
                }));
            } else {
                out.push(Stmt::Expr(lhs));
            }
            if !self.eat(&TokenKind::Comma) {
                break;
            }
        }
        self.expect(TokenKind::Semi, "`;` after statement")?;
        Ok(())
    }

    fn peek_assign_op(&self) -> Option<AssignOp> {
        Some(match self.peek().kind {
            TokenKind::Eq => AssignOp::Set,
            TokenKind::PlusEq => AssignOp::Add,
            TokenKind::MinusEq => AssignOp::Sub,
            TokenKind::StarEq => AssignOp::Mul,
            TokenKind::SlashEq => AssignOp::Div,
            TokenKind::PercentEq => AssignOp::Rem,
            _ => return None,
        })
    }

    fn parse_if(&mut self) -> Result<Stmt, ParseError> {
        let start = self.peek().span;
        self.bump(); // if
        self.expect(TokenKind::LParen, "`(` after `if`")?;
        let cond = self.parse_expr()?;
        self.expect(TokenKind::RParen, "`)` after `if` condition")?;
        let then_branch = Box::new(self.parse_stmt()?);
        let else_branch = if matches!(self.peek().kind, TokenKind::Keyword(Keyword::Else)) {
            self.bump();
            Some(Box::new(self.parse_stmt()?))
        } else {
            None
        };
        let end_span = else_branch
            .as_ref()
            .map(|b| stmt_span(b))
            .unwrap_or(stmt_span(&then_branch));
        Ok(Stmt::If(IfStmt {
            cond,
            then_branch,
            else_branch,
            span: start.merge(end_span),
        }))
    }

    fn parse_while(&mut self) -> Result<Stmt, ParseError> {
        let start = self.peek().span;
        self.bump(); // while
        self.expect(TokenKind::LParen, "`(` after `while`")?;
        let cond = self.parse_expr()?;
        self.expect(TokenKind::RParen, "`)` after `while` condition")?;
        let body = Box::new(self.parse_stmt()?);
        let end = stmt_span(&body);
        Ok(Stmt::While(WhileStmt {
            cond,
            body,
            do_while: false,
            span: start.merge(end),
        }))
    }

    fn parse_do_while(&mut self) -> Result<Stmt, ParseError> {
        let start = self.peek().span;
        self.bump(); // do
        let body = Box::new(self.parse_stmt()?);
        self.expect(
            TokenKind::Keyword(Keyword::While),
            "`while` after `do` body",
        )?;
        self.expect(TokenKind::LParen, "`(` after `while`")?;
        let cond = self.parse_expr()?;
        self.expect(TokenKind::RParen, "`)` after `while` condition")?;
        let end = self.peek().span;
        self.expect(TokenKind::Semi, "`;` after do-while")?;
        Ok(Stmt::While(WhileStmt {
            cond,
            body,
            do_while: true,
            span: start.merge(end),
        }))
    }

    fn parse_for(&mut self) -> Result<Stmt, ParseError> {
        let start = self.peek().span;
        self.bump(); // for
        self.expect(TokenKind::LParen, "`(` after `for`")?;
        let init = if matches!(self.peek().kind, TokenKind::Semi) {
            self.bump();
            None
        } else {
            // Init position: a decl or an expression statement; both end
            // with `;`. parse_stmt handles either branch.
            Some(Box::new(self.parse_stmt()?))
        };
        let cond = if matches!(self.peek().kind, TokenKind::Semi) {
            None
        } else {
            Some(self.parse_expr()?)
        };
        self.expect(TokenKind::Semi, "`;` between `for` cond and step")?;
        let step = if matches!(self.peek().kind, TokenKind::RParen) {
            None
        } else {
            // `for` step accepts an assignment-shaped expression
            // (`i = i + 1`) — captured as an [`Expr::Assign`] node so the
            // AST keeps a typed handle on the LHS, the operator, and the
            // RHS. Plain expressions (`i++`) parse through `parse_expr`
            // without an assignment tail.
            let base = self.parse_expr()?;
            if let Some(op) = self.peek_assign_op() {
                self.bump();
                let rhs = self.parse_expr()?;
                let span = base.span().merge(rhs.span());
                Some(Expr::Assign(AssignExpr {
                    target: Box::new(base),
                    op,
                    value: Box::new(rhs),
                    span,
                }))
            } else {
                Some(base)
            }
        };
        self.expect(TokenKind::RParen, "`)` to close `for` header")?;
        let body = Box::new(self.parse_stmt()?);
        let end = stmt_span(&body);
        Ok(Stmt::For(ForStmt {
            init,
            cond,
            step,
            body,
            span: start.merge(end),
        }))
    }

    fn parse_return(&mut self) -> Result<Stmt, ParseError> {
        self.bump(); // return
        let val = if matches!(self.peek().kind, TokenKind::Semi) {
            None
        } else {
            Some(self.parse_expr()?)
        };
        self.expect(TokenKind::Semi, "`;` after `return`")?;
        Ok(Stmt::Return(val))
    }

    // -------- expressions (Pratt) --------

    fn parse_expr(&mut self) -> Result<Expr, ParseError> {
        self.parse_ternary()
    }

    fn parse_ternary(&mut self) -> Result<Expr, ParseError> {
        let cond = self.parse_logical_or()?;
        if self.eat(&TokenKind::Question) {
            let then_expr = self.parse_expr()?;
            self.expect(TokenKind::Colon, "`:` in ternary expression")?;
            let else_expr = self.parse_expr()?;
            let span = cond.span().merge(else_expr.span());
            return Ok(Expr::Ternary(TernaryExpr {
                cond: Box::new(cond),
                then_expr: Box::new(then_expr),
                else_expr: Box::new(else_expr),
                span,
            }));
        }
        Ok(cond)
    }

    fn parse_binary_left<F>(
        &mut self,
        mut next: F,
        accept: &[(TokenKind, BinaryOp)],
    ) -> Result<Expr, ParseError>
    where
        F: FnMut(&mut Self) -> Result<Expr, ParseError>,
    {
        let mut lhs = next(self)?;
        while let Some(op) = accept
            .iter()
            .find(|(k, _)| std::mem::discriminant(&self.peek().kind) == std::mem::discriminant(k))
            .map(|(_, op)| *op)
        {
            self.bump();
            let rhs = next(self)?;
            let span = lhs.span().merge(rhs.span());
            lhs = Expr::Binary(BinaryExpr {
                op,
                lhs: Box::new(lhs),
                rhs: Box::new(rhs),
                span,
            });
        }
        Ok(lhs)
    }

    fn parse_logical_or(&mut self) -> Result<Expr, ParseError> {
        self.parse_binary_left(
            |s| s.parse_logical_and(),
            &[(TokenKind::BarBar, BinaryOp::Or)],
        )
    }

    fn parse_logical_and(&mut self) -> Result<Expr, ParseError> {
        self.parse_binary_left(|s| s.parse_bit_or(), &[(TokenKind::AmpAmp, BinaryOp::And)])
    }

    fn parse_bit_or(&mut self) -> Result<Expr, ParseError> {
        self.parse_binary_left(|s| s.parse_bit_xor(), &[(TokenKind::Bar, BinaryOp::BitOr)])
    }

    fn parse_bit_xor(&mut self) -> Result<Expr, ParseError> {
        self.parse_binary_left(
            |s| s.parse_bit_and(),
            &[(TokenKind::Caret, BinaryOp::BitXor)],
        )
    }

    fn parse_bit_and(&mut self) -> Result<Expr, ParseError> {
        self.parse_binary_left(
            |s| s.parse_equality(),
            &[(TokenKind::Amp, BinaryOp::BitAnd)],
        )
    }

    fn parse_equality(&mut self) -> Result<Expr, ParseError> {
        self.parse_binary_left(
            |s| s.parse_relational(),
            &[
                (TokenKind::EqEq, BinaryOp::Eq),
                (TokenKind::BangEq, BinaryOp::Ne),
            ],
        )
    }

    fn parse_relational(&mut self) -> Result<Expr, ParseError> {
        self.parse_binary_left(
            |s| s.parse_shift(),
            &[
                (TokenKind::Lt, BinaryOp::Lt),
                (TokenKind::LtEq, BinaryOp::Le),
                (TokenKind::Gt, BinaryOp::Gt),
                (TokenKind::GtEq, BinaryOp::Ge),
            ],
        )
    }

    fn parse_shift(&mut self) -> Result<Expr, ParseError> {
        self.parse_binary_left(
            |s| s.parse_additive(),
            &[
                (TokenKind::Shl, BinaryOp::Shl),
                (TokenKind::Shr, BinaryOp::Shr),
            ],
        )
    }

    fn parse_additive(&mut self) -> Result<Expr, ParseError> {
        self.parse_binary_left(
            |s| s.parse_multiplicative(),
            &[
                (TokenKind::Plus, BinaryOp::Add),
                (TokenKind::Minus, BinaryOp::Sub),
            ],
        )
    }

    fn parse_multiplicative(&mut self) -> Result<Expr, ParseError> {
        self.parse_binary_left(
            |s| s.parse_unary(),
            &[
                (TokenKind::Star, BinaryOp::Mul),
                (TokenKind::Slash, BinaryOp::Div),
                (TokenKind::Percent, BinaryOp::Rem),
            ],
        )
    }

    fn parse_unary(&mut self) -> Result<Expr, ParseError> {
        let start = self.peek().span;
        let op = match self.peek().kind {
            TokenKind::Minus => Some(UnaryOp::Neg),
            TokenKind::Plus => Some(UnaryOp::Pos),
            TokenKind::Bang => Some(UnaryOp::Not),
            TokenKind::Tilde => Some(UnaryOp::BitNot),
            _ => None,
        };
        if let Some(op) = op {
            self.bump();
            let operand = self.parse_unary()?;
            let span = start.merge(operand.span());
            return Ok(Expr::Unary(UnaryExpr {
                op,
                operand: Box::new(operand),
                span,
            }));
        }
        self.parse_postfix()
    }

    fn parse_postfix(&mut self) -> Result<Expr, ParseError> {
        let mut e = self.parse_primary()?;
        loop {
            match self.peek().kind {
                TokenKind::Dot => {
                    self.bump();
                    let (member, m_span) = self.expect_ident("member name after `.`")?;
                    let span = e.span().merge(m_span);
                    if is_swizzle(&member) {
                        e = Expr::Swizzle(SwizzleExpr {
                            base: Box::new(e),
                            components: member,
                            span,
                        });
                    } else {
                        e = Expr::Member(MemberExpr {
                            base: Box::new(e),
                            member,
                            span,
                        });
                    }
                }
                TokenKind::LBracket => {
                    self.bump();
                    let idx = self.parse_expr()?;
                    let close = self.peek().span;
                    self.expect(TokenKind::RBracket, "`]` after index")?;
                    let span = e.span().merge(close);
                    e = Expr::Index(IndexExpr {
                        base: Box::new(e),
                        index: Box::new(idx),
                        span,
                    });
                }
                TokenKind::PlusPlus | TokenKind::MinusMinus => {
                    // Postfix increment is a statement in HLSL practice;
                    // we model it as `x = x + 1` at emit time. For now we
                    // treat `x++` as just `x` so AST consumers can spot it
                    // and rewrite — though no emitter does so yet. We bump
                    // the token to avoid an infinite loop, and emit a
                    // synthetic binary op to make the value semantics
                    // close-to-right (`x++` evaluates to `x` pre-increment,
                    // which we approximate by leaving `e` unchanged).
                    self.bump();
                }
                _ => break,
            }
        }
        Ok(e)
    }

    fn parse_primary(&mut self) -> Result<Expr, ParseError> {
        let t = self.peek().clone();
        match t.kind {
            TokenKind::IntLit(v) => {
                self.bump();
                Ok(Expr::Lit(Lit {
                    value: LitValue::Int(v),
                    span: t.span,
                }))
            }
            TokenKind::FloatLit(v) => {
                self.bump();
                Ok(Expr::Lit(Lit {
                    value: LitValue::Float(v),
                    span: t.span,
                }))
            }
            TokenKind::BoolLit(v) => {
                self.bump();
                Ok(Expr::Lit(Lit {
                    value: LitValue::Bool(v),
                    span: t.span,
                }))
            }
            TokenKind::Ident => {
                self.bump();
                let name = self.ident_text(t.span).to_string();
                // Call or constructor: `name(...)`.
                if matches!(self.peek().kind, TokenKind::LParen) {
                    self.bump();
                    let mut args = Vec::new();
                    if !matches!(self.peek().kind, TokenKind::RParen) {
                        loop {
                            let base = self.parse_expr()?;
                            // HLSL accepts an assignment expression as a
                            // call argument: `lerp(a, tmp = GetBlur1(uv), b)`
                            // sets `tmp` *and* uses the new value as the
                            // arg. WGSL has no assignment expression, so a
                            // separate AST pass will lift the side-effect
                            // out; here we just capture the shape as
                            // [`Expr::Assign`] so it round-trips.
                            let arg = if let Some(op) = self.peek_assign_op() {
                                self.bump();
                                let rhs = self.parse_expr()?;
                                let span = base.span().merge(rhs.span());
                                Expr::Assign(AssignExpr {
                                    target: Box::new(base),
                                    op,
                                    value: Box::new(rhs),
                                    span,
                                })
                            } else {
                                base
                            };
                            args.push(arg);
                            if !self.eat(&TokenKind::Comma) {
                                break;
                            }
                        }
                    }
                    let close = self.peek().span;
                    self.expect(TokenKind::RParen, "`)` after call arguments")?;
                    return Ok(Expr::Call(CallExpr {
                        callee: name,
                        args,
                        span: t.span.merge(close),
                    }));
                }
                Ok(Expr::Ident(name, t.span))
            }
            TokenKind::LParen => {
                self.bump();
                let mut inner = self.parse_expr()?;
                // C-style comma operator inside parens: `(a, b, c)` evaluates
                // all in order and yields `c`. MD2 occasionally writes
                // `texsize.zx*(q3, q3)` (a degenerate case). We honour the
                // syntax by keeping only the rightmost expression; the
                // earlier ones have no side effects in any preset we've
                // surveyed.
                while self.eat(&TokenKind::Comma) {
                    inner = self.parse_expr()?;
                }
                let close = self.peek().span;
                self.expect(TokenKind::RParen, "`)` to close parenthesised expression")?;
                Ok(self.with_span(inner, t.span.merge(close)))
            }
            TokenKind::LBrace => {
                // `{ a, b, c }` initialiser list.
                self.bump();
                let mut elems = Vec::new();
                if !matches!(self.peek().kind, TokenKind::RBrace) {
                    loop {
                        elems.push(self.parse_expr()?);
                        if !self.eat(&TokenKind::Comma) {
                            break;
                        }
                        // Trailing comma OK.
                        if matches!(self.peek().kind, TokenKind::RBrace) {
                            break;
                        }
                    }
                }
                let close = self.peek().span;
                self.expect(TokenKind::RBrace, "`}` to close initializer list")?;
                Ok(Expr::InitList(InitListExpr {
                    elems,
                    span: t.span.merge(close),
                }))
            }
            _ => Err(ParseError {
                message: format!("expected expression, found {:?}", t.kind),
                span: t.span,
            }),
        }
    }

    /// Replace the outer span on an expression; used after consuming `(`
    /// `expr` `)` so the parenthesised whole gets the wider span.
    fn with_span(&self, e: Expr, span: Span) -> Expr {
        match e {
            Expr::Lit(mut l) => {
                l.span = span;
                Expr::Lit(l)
            }
            Expr::Ident(n, _) => Expr::Ident(n, span),
            Expr::Binary(mut b) => {
                b.span = span;
                Expr::Binary(b)
            }
            Expr::Unary(mut u) => {
                u.span = span;
                Expr::Unary(u)
            }
            Expr::Ternary(mut t) => {
                t.span = span;
                Expr::Ternary(t)
            }
            Expr::Call(mut c) => {
                c.span = span;
                Expr::Call(c)
            }
            Expr::Member(mut m) => {
                m.span = span;
                Expr::Member(m)
            }
            Expr::Swizzle(mut s) => {
                s.span = span;
                Expr::Swizzle(s)
            }
            Expr::Index(mut i) => {
                i.span = span;
                Expr::Index(i)
            }
            Expr::InitList(mut l) => {
                l.span = span;
                Expr::InitList(l)
            }
            Expr::Assign(mut a) => {
                a.span = span;
                Expr::Assign(a)
            }
        }
    }
}

fn stmt_span(s: &Stmt) -> Span {
    match s {
        Stmt::LocalDecl(d) => d.span,
        Stmt::Assign(a) => a.span,
        Stmt::Expr(e) => e.span(),
        Stmt::If(i) => i.span,
        Stmt::While(w) => w.span,
        Stmt::For(f) => f.span,
        Stmt::Return(Some(e)) => e.span(),
        Stmt::Return(None) => Span::dummy(),
        Stmt::Break | Stmt::Continue => Span::dummy(),
        Stmt::Block(b) => b.span,
    }
}

/// A `.<chars>` member access is a swizzle iff every char is one of the
/// component letters (xyzw / rgba), with length 1–4. HLSL doesn't allow
/// mixed xyzw+rgba in one swizzle, but we accept it here — the emitter is
/// stricter than the parser.
fn is_swizzle(s: &str) -> bool {
    if s.is_empty() || s.len() > 4 {
        return false;
    }
    s.bytes()
        .all(|b| matches!(b, b'x' | b'y' | b'z' | b'w' | b'r' | b'g' | b'b' | b'a'))
}

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

    fn parse(src: &str) -> Result<TranslationUnit, ParseError> {
        parse_hlsl(src)
    }

    #[test]
    fn empty_input_is_empty_unit() {
        let tu = parse("").unwrap();
        assert!(tu.items.is_empty());
        assert!(tu.shader_body.is_none());
    }

    #[test]
    fn shader_body_only() {
        let tu = parse("shader_body { float a = 1.0; }").unwrap();
        assert!(tu.items.is_empty());
        let body = tu.shader_body.expect("shader_body present");
        assert_eq!(body.stmts.len(), 1);
        let Stmt::LocalDecl(d) = &body.stmts[0] else {
            panic!("expected local decl, got {:?}", body.stmts[0]);
        };
        assert_eq!(d.name, "a");
        assert_eq!(d.ty.name, "float");
        assert!(matches!(
            d.init.as_ref().unwrap(),
            Expr::Lit(l) if matches!(l.value, LitValue::Float(v) if (v - 1.0).abs() < 1e-9)
        ));
    }

    #[test]
    fn binary_operator_precedence() {
        // `a + b * c` should parse as `a + (b * c)`.
        let tu = parse("shader_body { float r = a + b * c; }").unwrap();
        let body = tu.shader_body.unwrap();
        let Stmt::LocalDecl(d) = &body.stmts[0] else {
            panic!("expected decl");
        };
        let init = d.init.as_ref().unwrap();
        let Expr::Binary(top) = init else {
            panic!("expected binary at top");
        };
        assert_eq!(top.op, BinaryOp::Add);
        // RHS should be a Mul, LHS should be plain ident.
        assert!(matches!(*top.lhs, Expr::Ident(ref n, _) if n == "a"));
        let Expr::Binary(rhs) = &*top.rhs else {
            panic!("expected mul on rhs");
        };
        assert_eq!(rhs.op, BinaryOp::Mul);
    }

    #[test]
    fn ternary_parses_right_associatively() {
        let tu = parse("shader_body { float r = a ? b : c ? d : e; }").unwrap();
        let body = tu.shader_body.unwrap();
        let Stmt::LocalDecl(d) = &body.stmts[0] else {
            panic!();
        };
        let init = d.init.as_ref().unwrap();
        let Expr::Ternary(t) = init else {
            panic!("expected ternary");
        };
        // Right-associative: else_expr is itself a ternary.
        assert!(matches!(*t.else_expr, Expr::Ternary(_)));
    }

    #[test]
    fn function_call_parses() {
        let tu = parse("shader_body { float l = length(uv2); }").unwrap();
        let body = tu.shader_body.unwrap();
        let Stmt::LocalDecl(d) = &body.stmts[0] else {
            panic!();
        };
        let Some(Expr::Call(c)) = d.init.as_ref() else {
            panic!("expected call init");
        };
        assert_eq!(c.callee, "length");
        assert_eq!(c.args.len(), 1);
    }

    #[test]
    fn swizzle_vs_member_distinguished() {
        let tu = parse("shader_body { float a = uv.x; float2 b = vec.xy; float c = obj.member; }")
            .unwrap();
        let body = tu.shader_body.unwrap();
        let Stmt::LocalDecl(d0) = &body.stmts[0] else {
            panic!();
        };
        assert!(matches!(d0.init.as_ref().unwrap(), Expr::Swizzle(_)));
        let Stmt::LocalDecl(d1) = &body.stmts[1] else {
            panic!();
        };
        assert!(matches!(d1.init.as_ref().unwrap(), Expr::Swizzle(_)));
        let Stmt::LocalDecl(d2) = &body.stmts[2] else {
            panic!();
        };
        assert!(matches!(d2.init.as_ref().unwrap(), Expr::Member(_)));
    }

    #[test]
    fn assignment_statement_with_compound_op() {
        let tu = parse("shader_body { ret += tex2D(sampler_main, uv) * 0.5; }").unwrap();
        let body = tu.shader_body.unwrap();
        let Stmt::Assign(a) = &body.stmts[0] else {
            panic!("expected assignment, got {:?}", body.stmts[0]);
        };
        assert_eq!(a.op, AssignOp::Add);
        assert!(matches!(a.target, Expr::Ident(_, _)));
    }

    #[test]
    fn chained_assign_init_lowered_to_pre_decl_statements() {
        // HLSL: `float2 ruv = uv = 0.5;`
        // Expect: `uv = 0.5;` emitted first, then `float2 ruv = uv;`.
        let tu = parse("shader_body { float2 ruv = uv = 0.5; }").unwrap();
        let body = tu.shader_body.unwrap();
        assert_eq!(body.stmts.len(), 2);
        let Stmt::Assign(a) = &body.stmts[0] else {
            panic!("expected assign first, got {:?}", body.stmts[0]);
        };
        assert_eq!(a.op, AssignOp::Set);
        assert!(matches!(&a.target, Expr::Ident(n, _) if n == "uv"));
        let Stmt::LocalDecl(d) = &body.stmts[1] else {
            panic!("expected decl second, got {:?}", body.stmts[1]);
        };
        assert_eq!(d.name, "ruv");
        assert!(matches!(d.init.as_ref().unwrap(), Expr::Ident(n, _) if n == "uv"));
    }

    #[test]
    fn chained_assign_init_three_deep() {
        // `float2 ruv = uv = uv2 = expr;` should emit two assigns then a decl.
        let tu = parse("shader_body { float2 ruv = uv = uv2 = 0.5; }").unwrap();
        let body = tu.shader_body.unwrap();
        assert_eq!(body.stmts.len(), 3);
        // Innermost first: uv2 = 0.5
        let Stmt::Assign(a0) = &body.stmts[0] else {
            panic!("expected assign, got {:?}", body.stmts[0]);
        };
        assert!(matches!(&a0.target, Expr::Ident(n, _) if n == "uv2"));
        // Next: uv = uv2
        let Stmt::Assign(a1) = &body.stmts[1] else {
            panic!("expected assign, got {:?}", body.stmts[1]);
        };
        assert!(matches!(&a1.target, Expr::Ident(n, _) if n == "uv"));
        assert!(matches!(&a1.value, Expr::Ident(n, _) if n == "uv2"));
        // Then the decl, init = uv
        let Stmt::LocalDecl(d) = &body.stmts[2] else {
            panic!("expected decl, got {:?}", body.stmts[2]);
        };
        assert!(matches!(d.init.as_ref().unwrap(), Expr::Ident(n, _) if n == "uv"));
    }

    #[test]
    fn array_initialiser_local_decl() {
        let tu = parse("shader_body { float arr[3] = {1.0, 2.0, 3.0}; }").unwrap();
        let body = tu.shader_body.unwrap();
        let Stmt::LocalDecl(d) = &body.stmts[0] else {
            panic!();
        };
        assert!(d.array_len.is_some());
        let Some(Expr::InitList(l)) = d.init.as_ref() else {
            panic!("expected init list");
        };
        assert_eq!(l.elems.len(), 3);
    }

    #[test]
    fn if_else_chain() {
        let tu =
            parse("shader_body { if (a > 0) b = 1; else if (a < 0) b = -1; else b = 0; }").unwrap();
        let body = tu.shader_body.unwrap();
        let Stmt::If(top) = &body.stmts[0] else {
            panic!("expected if, got {:?}", body.stmts[0]);
        };
        assert!(top.else_branch.is_some());
        // The nested else should itself be an If.
        let Some(else_b) = &top.else_branch else {
            panic!()
        };
        assert!(matches!(else_b.as_ref(), Stmt::If(_)));
    }

    #[test]
    fn while_loop() {
        let tu = parse("shader_body { while (n < 4) { ret += 1.0; n = n + 1; } }").unwrap();
        let body = tu.shader_body.unwrap();
        assert!(matches!(body.stmts[0], Stmt::While(_)));
    }

    #[test]
    fn for_loop_with_init_and_step() {
        let tu =
            parse("shader_body { for (int i = 0; i < 5; i = i + 1) { ret += 1.0; } }").unwrap();
        let body = tu.shader_body.unwrap();
        let Stmt::For(f) = &body.stmts[0] else {
            panic!("expected for, got {:?}", body.stmts[0]);
        };
        assert!(f.init.is_some());
        assert!(f.cond.is_some());
        assert!(f.step.is_some());
    }

    #[test]
    fn function_definition() {
        let src = "float square(float x) { return x * x; }";
        let tu = parse(src).unwrap();
        assert_eq!(tu.items.len(), 1);
        let Item::Function(f) = &tu.items[0] else {
            panic!("expected function, got {:?}", tu.items[0]);
        };
        assert_eq!(f.name, "square");
        assert_eq!(f.params.len(), 1);
        assert_eq!(f.params[0].name, "x");
        assert_eq!(f.return_type.name, "float");
    }

    #[test]
    fn sampler_declaration_at_top_level() {
        let tu = parse("sampler sampler_fw_sky; shader_body { ret = 0; }").unwrap();
        assert_eq!(tu.items.len(), 1);
        let Item::SamplerDecl(d) = &tu.items[0] else {
            panic!("expected sampler decl");
        };
        assert_eq!(d.name, "sampler_fw_sky");
    }

    #[test]
    fn nested_constructor_call() {
        let tu = parse("shader_body { float2 v = float2(1.0, 2.0); }").unwrap();
        let body = tu.shader_body.unwrap();
        let Stmt::LocalDecl(d) = &body.stmts[0] else {
            panic!();
        };
        let Some(Expr::Call(c)) = d.init.as_ref() else {
            panic!("expected call init");
        };
        assert_eq!(c.callee, "float2");
        assert_eq!(c.args.len(), 2);
    }

    #[test]
    fn unary_negate() {
        let tu = parse("shader_body { float a = -b; }").unwrap();
        let body = tu.shader_body.unwrap();
        let Stmt::LocalDecl(d) = &body.stmts[0] else {
            panic!();
        };
        let Some(Expr::Unary(u)) = d.init.as_ref() else {
            panic!("expected unary, got {:?}", d.init);
        };
        assert_eq!(u.op, UnaryOp::Neg);
    }

    #[test]
    fn realistic_comp_shader_body() {
        let src = r#"
            sampler sampler_fw_sky;
            shader_body {
                float2 uvm = frac(uv + time * 0.0);
                ret = tex2D(sampler_main, uvm);
                float diff = 1 - length(ret.xy - GetBlur1(uvm).xy) * 3.5;
                ret.xy *= diff;
            }
        "#;
        let tu = parse(src).unwrap();
        assert_eq!(tu.items.len(), 1);
        let body = tu.shader_body.expect("body");
        assert_eq!(body.stmts.len(), 4);
        // Last statement is an `*=` assignment to a swizzle target.
        let Stmt::Assign(a) = &body.stmts[3] else {
            panic!("expected assign, got {:?}", body.stmts[3]);
        };
        assert_eq!(a.op, AssignOp::Mul);
        assert!(matches!(a.target, Expr::Swizzle(_)));
    }

    #[test]
    fn parse_error_carries_position() {
        // Missing `;` after assignment — should fail at the `}` token.
        let err = parse("shader_body { a = 1 }").unwrap_err();
        assert!(err.message.contains("`;`"));
        // Position is at the offending token, not byte 0.
        assert!(err.span.start > 0);
    }

    #[test]
    fn realistic_warp_shader_body() {
        // MD2 warp shaders use `ret`/`tex2D`/`q1`/local floats heavily.
        let src = r#"
            shader_body {
                float2 uv2 = uv - 0.5;
                uv2 *= aspect.xy;
                float r = length(uv2);
                float ang = atan2(uv2.y, uv2.x);
                uv2 = float2(r * cos(ang), r * sin(ang));
                ret = tex2D(sampler_main, uv2 + 0.5);
            }
        "#;
        let tu = parse(src).unwrap();
        let body = tu.shader_body.unwrap();
        // 6 statements expected.
        assert_eq!(body.stmts.len(), 6);
    }

    #[test]
    fn lifted_helper_function_before_body() {
        let src = r#"
            float square(float x) { return x * x; }
            float2 normish(float2 v) { return v * (1.0 / length(v)); }
            shader_body {
                float2 d = normish(uv - 0.5);
                ret = float4(square(d.x), square(d.y), 0, 1);
            }
        "#;
        let tu = parse(src).unwrap();
        assert_eq!(tu.items.len(), 2);
        for item in &tu.items {
            assert!(matches!(item, Item::Function(_)));
        }
        let body = tu.shader_body.unwrap();
        assert_eq!(body.stmts.len(), 2);
    }

    #[test]
    fn global_multi_name_decl_flattens_to_one_item_per_name() {
        // MD2 idiom: `float3 ret1, neu, blur;` before `shader_body`.
        // Parser flattens to three GlobalVar items so downstream emitters
        // see one decl per name.
        let tu = parse("float3 ret1, neu, blur; shader_body { ret = 0; }").unwrap();
        assert_eq!(tu.items.len(), 3);
        let names: Vec<&str> = tu
            .items
            .iter()
            .map(|i| match i {
                Item::GlobalVar(g) => g.name.as_str(),
                _ => panic!("expected GlobalVar"),
            })
            .collect();
        assert_eq!(names, vec!["ret1", "neu", "blur"]);
    }

    #[test]
    fn local_multi_name_decl_flattens_to_one_stmt_per_name() {
        let tu = parse("shader_body { float a, b = 2.0, c; }").unwrap();
        let body = tu.shader_body.unwrap();
        assert_eq!(body.stmts.len(), 3);
        for s in &body.stmts {
            assert!(matches!(s, Stmt::LocalDecl(_)));
        }
        let Stmt::LocalDecl(b) = &body.stmts[1] else {
            panic!();
        };
        assert_eq!(b.name, "b");
        assert!(b.init.is_some());
    }

    #[test]
    fn comma_separated_assignments_flatten_into_multiple_stmts() {
        // `a += 1, b += 2;` parses as two Stmt::Assign nodes.
        let tu = parse("shader_body { a += 1, b += 2; }").unwrap();
        let body = tu.shader_body.unwrap();
        assert_eq!(body.stmts.len(), 2);
        for s in &body.stmts {
            assert!(matches!(s, Stmt::Assign(_)));
        }
    }

    #[test]
    fn global_const_array_initialiser() {
        let src = r#"
            const float4 samples[5] = { float4(1,0,0,0), float4(0,1,0,0),
                                        float4(0,0,1,0), float4(0,0,0,1),
                                        float4(1,1,1,1) };
            shader_body { ret = samples[0]; }
        "#;
        let tu = parse(src).unwrap();
        assert_eq!(tu.items.len(), 1);
        let Item::GlobalVar(g) = &tu.items[0] else {
            panic!("expected GlobalVar");
        };
        assert!(g.is_const);
        assert_eq!(g.name, "samples");
        assert!(g.array_len.is_some());
        let init = g.init.as_ref().expect("init list present");
        assert!(matches!(init, Expr::InitList(l) if l.elems.len() == 5));
    }

    #[test]
    fn comma_operator_inside_parens_keeps_rightmost() {
        // `(a, b)` evaluates to `b`. Real preset (`orb - inferno`) writes
        // `(q3, q3)` to confuse the parser; we accept it and keep the
        // tail expression.
        let tu = parse("shader_body { float r = (1.0, 2.0); }").unwrap();
        let body = tu.shader_body.unwrap();
        let Stmt::LocalDecl(d) = &body.stmts[0] else {
            panic!();
        };
        let Some(Expr::Lit(l)) = d.init.as_ref() else {
            panic!("expected float literal init");
        };
        assert!(matches!(l.value, LitValue::Float(v) if (v - 2.0).abs() < 1e-9));
    }

    #[test]
    fn float_array_initialiser_with_int_literals() {
        // `float arr[5] = {1,2,3,4,5};` was a known blocker in the residual
        // failure list. The parser handles it; the emitter will catch up.
        let tu = parse("shader_body { float arr[5] = {1, 2, 3, 4, 5}; }").unwrap();
        let body = tu.shader_body.unwrap();
        let Stmt::LocalDecl(d) = &body.stmts[0] else {
            panic!();
        };
        let Some(Expr::InitList(l)) = d.init.as_ref() else {
            panic!("expected init list");
        };
        assert_eq!(l.elems.len(), 5);
        // All elements are int literals — emitter will need to coerce to
        // float at WGSL emit time.
        for e in &l.elems {
            assert!(matches!(
                e,
                Expr::Lit(Lit {
                    value: LitValue::Int(_),
                    ..
                })
            ));
        }
    }
}