onedrop_hlsl/rewrite/
for_init.rs

1//! Pass: `for (int i = ...)` init type rewrite (lowering).
2
3use super::*;
4
5// ---------------------------------------------------------------------------
6// Pass: `for (int i = …)` init rewrite
7// ---------------------------------------------------------------------------
8
9/// HLSL allows `for (int i = 0; i < 5; i++)` and the parser surfaces the
10/// init clause as a `Stmt::LocalDecl` with `ty.name == "int"`. The
11/// downstream `replace_types` regex rewrites every `\bint\b` into `f32`,
12/// which turns the declarator into `for (f32 i = 0; …)` — invalid WGSL
13/// syntax (for-init expects a `var`-shaped declaration). We catch the
14/// pattern up front and rewrite it into `for (var i: i32 = 0; …)` so the
15/// loop var is usable as an array index and `replace_types` has nothing
16/// to mangle.
17///
18/// Targets the Stahlregen + suksma `dotes` pair (~2 presets).
19#[cfg(test)]
20pub(crate) fn rewrite_for_int_init(src: &str) -> String {
21    let Ok(tu) = parse_hlsl(src) else {
22        return src.to_string();
23    };
24    let mut edits = Vec::new();
25    if let Some(body) = &tu.shader_body {
26        collect_for_int_edits(body, &mut edits);
27    }
28    for item in &tu.items {
29        if let Item::Function(f) = item {
30            collect_for_int_edits(&f.body, &mut edits);
31        }
32    }
33    apply_edits(src, &mut edits)
34}
35
36pub(crate) fn collect_for_int_edits(block: &Block, edits: &mut Vec<TextEdit>) {
37    for s in &block.stmts {
38        collect_for_int_in_stmt(s, edits);
39    }
40}
41
42fn collect_for_int_in_stmt(s: &Stmt, edits: &mut Vec<TextEdit>) {
43    match s {
44        Stmt::For(f) => {
45            if let Some(init) = &f.init
46                && let Stmt::LocalDecl(d) = &**init
47                && d.ty.name == "int"
48                && d.array_len.is_none()
49            {
50                edits.push(TextEdit {
51                    start: d.ty.span.start,
52                    end: d.ty.span.end,
53                    replacement: "var".to_string(),
54                });
55                edits.push(TextEdit {
56                    start: d.span.end,
57                    end: d.span.end,
58                    replacement: ": i32".to_string(),
59                });
60            }
61            collect_for_int_in_stmt(&f.body, edits);
62        }
63        Stmt::Block(b) => collect_for_int_edits(b, edits),
64        Stmt::If(i) => {
65            collect_for_int_in_stmt(&i.then_branch, edits);
66            if let Some(e) = &i.else_branch {
67                collect_for_int_in_stmt(e, edits);
68            }
69        }
70        Stmt::While(w) => collect_for_int_in_stmt(&w.body, edits),
71        _ => {}
72    }
73}