onedrop_hlsl/rewrite/

ternary.rs

//! Pass: ternary `cond ? a : b` → WGSL `select(b, a, cond)`.

use super::*;

// ---------------------------------------------------------------------------
// Pass: ternary `?:` → WGSL `select(...)`
// ---------------------------------------------------------------------------

/// Lower every `cond ? a : b` to `select(b, a, cond)`. WGSL has no
/// ternary operator and the regex translator doesn't recognise the
/// `?`/`:` tokens. Many presets hit `expected ';'; found '?'` before this
/// pass.
///
/// Handles nested ternaries by emitting one edit per *top-level* ternary
/// (in source order) and rendering inner ternaries inside the replacement
/// string via a recursive expression emitter. That keeps text edits
/// non-overlapping, which the [`apply_edits`] machinery requires.
pub(crate) fn rewrite_ternary_to_select(src: &str) -> String {
    let Ok(tu) = parse_hlsl(src) else {
        return src.to_string();
    };
    let mut edits = Vec::new();
    if let Some(body) = &tu.shader_body {
        scan_block_for_ternary(body, src, &mut edits);
    }
    for item in &tu.items {
        match item {
            Item::Function(f) => scan_block_for_ternary(&f.body, src, &mut edits),
            Item::GlobalVar(g) => {
                if let Some(init) = &g.init {
                    scan_expr_for_ternary(init, src, &mut edits);
                }
            }
            _ => {}
        }
    }
    apply_edits(src, &mut edits)
}

fn scan_block_for_ternary(b: &Block, src: &str, edits: &mut Vec<TextEdit>) {
    for s in &b.stmts {
        scan_stmt_for_ternary(s, src, edits);
    }
}

fn scan_stmt_for_ternary(s: &Stmt, src: &str, edits: &mut Vec<TextEdit>) {
    match s {
        Stmt::LocalDecl(d) => {
            if let Some(init) = &d.init {
                scan_expr_for_ternary(init, src, edits);
            }
            if let Some(len) = &d.array_len {
                scan_expr_for_ternary(len, src, edits);
            }
        }
        Stmt::Assign(a) => {
            scan_expr_for_ternary(&a.target, src, edits);
            scan_expr_for_ternary(&a.value, src, edits);
        }
        Stmt::Expr(e) => scan_expr_for_ternary(e, src, edits),
        Stmt::If(i) => {
            scan_expr_for_ternary(&i.cond, src, edits);
            scan_stmt_for_ternary(&i.then_branch, src, edits);
            if let Some(b) = &i.else_branch {
                scan_stmt_for_ternary(b, src, edits);
            }
        }
        Stmt::While(w) => {
            scan_expr_for_ternary(&w.cond, src, edits);
            scan_stmt_for_ternary(&w.body, src, edits);
        }
        Stmt::For(f) => {
            if let Some(init) = &f.init {
                scan_stmt_for_ternary(init, src, edits);
            }
            if let Some(c) = &f.cond {
                scan_expr_for_ternary(c, src, edits);
            }
            if let Some(st) = &f.step {
                scan_expr_for_ternary(st, src, edits);
            }
            scan_stmt_for_ternary(&f.body, src, edits);
        }
        Stmt::Return(Some(e)) => scan_expr_for_ternary(e, src, edits),
        Stmt::Block(b) => scan_block_for_ternary(b, src, edits),
        Stmt::Return(None) | Stmt::Break | Stmt::Continue => {}
    }
}

fn scan_expr_for_ternary(e: &Expr, src: &str, edits: &mut Vec<TextEdit>) {
    match e {
        Expr::Ternary(t) => {
            // One edit for the whole ternary span. Recursive emit handles
            // nested ternaries, so we never descend into the children
            // here (doing so would double-emit and overlap).
            let replacement = emit_ternary_aware(e, src);
            edits.push(TextEdit {
                start: t.span.start,
                end: t.span.end,
                replacement,
            });
        }
        Expr::Binary(b) => {
            scan_expr_for_ternary(&b.lhs, src, edits);
            scan_expr_for_ternary(&b.rhs, src, edits);
        }
        Expr::Unary(u) => scan_expr_for_ternary(&u.operand, src, edits),
        Expr::Call(c) => {
            for a in &c.args {
                scan_expr_for_ternary(a, src, edits);
            }
        }
        Expr::Member(m) => scan_expr_for_ternary(&m.base, src, edits),
        Expr::Swizzle(s) => scan_expr_for_ternary(&s.base, src, edits),
        Expr::Index(i) => {
            scan_expr_for_ternary(&i.base, src, edits);
            scan_expr_for_ternary(&i.index, src, edits);
        }
        Expr::InitList(l) => {
            for e in &l.elems {
                scan_expr_for_ternary(e, src, edits);
            }
        }
        Expr::Assign(a) => {
            scan_expr_for_ternary(&a.target, src, edits);
            scan_expr_for_ternary(&a.value, src, edits);
        }
        Expr::Lit(_) | Expr::Ident(_, _) => {}
    }
}

/// Render an expression to text, replacing any ternaries in its subtree
/// with `select(...)`. Non-ternary nodes that contain ternaries get
/// rebuilt from the AST so the ternary substitution can flow up; pure
/// non-ternary subtrees fall back to the original source slice (which
/// preserves whitespace, comments and HLSL idioms that downstream regex
/// passes still need to see verbatim).
fn emit_ternary_aware(e: &Expr, src: &str) -> String {
    if !subtree_has_ternary(e) {
        return src[e.span().start as usize..e.span().end as usize].to_string();
    }
    match e {
        Expr::Ternary(t) => format!(
            "select(({}), ({}), ({}))",
            emit_ternary_aware(&t.else_expr, src),
            emit_ternary_aware(&t.then_expr, src),
            emit_ternary_aware(&t.cond, src),
        ),
        Expr::Binary(b) => format!(
            "({} {} {})",
            emit_ternary_aware(&b.lhs, src),
            binop_text(b.op),
            emit_ternary_aware(&b.rhs, src),
        ),
        Expr::Unary(u) => format!("{}{}", unop_text(u.op), emit_ternary_aware(&u.operand, src)),
        Expr::Call(c) => {
            let args: Vec<String> = c.args.iter().map(|a| emit_ternary_aware(a, src)).collect();
            format!("{}({})", c.callee, args.join(", "))
        }
        Expr::Member(m) => format!("{}.{}", emit_ternary_aware(&m.base, src), m.member),
        Expr::Swizzle(s) => format!("{}.{}", emit_ternary_aware(&s.base, src), s.components),
        Expr::Index(i) => format!(
            "{}[{}]",
            emit_ternary_aware(&i.base, src),
            emit_ternary_aware(&i.index, src)
        ),
        Expr::Assign(a) => format!(
            "{} {} {}",
            emit_ternary_aware(&a.target, src),
            assign_op_text(a.op),
            emit_ternary_aware(&a.value, src)
        ),
        // Lit/Ident/InitList don't (or rarely) contain ternaries; the
        // `subtree_has_ternary` short-circuit above catches the common
        // cases, but InitList could still hit here for nested inits.
        Expr::InitList(l) => {
            let elems: Vec<String> = l.elems.iter().map(|e| emit_ternary_aware(e, src)).collect();
            format!("{{ {} }}", elems.join(", "))
        }
        Expr::Lit(_) | Expr::Ident(_, _) => {
            src[e.span().start as usize..e.span().end as usize].to_string()
        }
    }
}

fn subtree_has_ternary(e: &Expr) -> bool {
    match e {
        Expr::Ternary(_) => true,
        Expr::Binary(b) => subtree_has_ternary(&b.lhs) || subtree_has_ternary(&b.rhs),
        Expr::Unary(u) => subtree_has_ternary(&u.operand),
        Expr::Call(c) => c.args.iter().any(subtree_has_ternary),
        Expr::Member(m) => subtree_has_ternary(&m.base),
        Expr::Swizzle(s) => subtree_has_ternary(&s.base),
        Expr::Index(i) => subtree_has_ternary(&i.base) || subtree_has_ternary(&i.index),
        Expr::InitList(l) => l.elems.iter().any(subtree_has_ternary),
        Expr::Assign(a) => subtree_has_ternary(&a.target) || subtree_has_ternary(&a.value),
        Expr::Lit(_) | Expr::Ident(_, _) => false,
    }
}

fn binop_text(op: BinaryOp) -> &'static str {
    match op {
        BinaryOp::Add => "+",
        BinaryOp::Sub => "-",
        BinaryOp::Mul => "*",
        BinaryOp::Div => "/",
        BinaryOp::Rem => "%",
        BinaryOp::Eq => "==",
        BinaryOp::Ne => "!=",
        BinaryOp::Lt => "<",
        BinaryOp::Le => "<=",
        BinaryOp::Gt => ">",
        BinaryOp::Ge => ">=",
        BinaryOp::And => "&&",
        BinaryOp::Or => "||",
        BinaryOp::BitAnd => "&",
        BinaryOp::BitOr => "|",
        BinaryOp::BitXor => "^",
        BinaryOp::Shl => "<<",
        BinaryOp::Shr => ">>",
    }
}

fn unop_text(op: UnaryOp) -> &'static str {
    match op {
        UnaryOp::Neg => "-",
        UnaryOp::Pos => "+",
        UnaryOp::Not => "!",
        UnaryOp::BitNot => "~",
    }
}

fn assign_op_text(op: AssignOp) -> &'static str {
    match op {
        AssignOp::Set => "=",
        AssignOp::Add => "+=",
        AssignOp::Sub => "-=",
        AssignOp::Mul => "*=",
        AssignOp::Div => "/=",
        AssignOp::Rem => "%=",
    }
}