onedrop_hlsl/rewrite/

array_lower.rs

//! Pass: HLSL array → WGSL `array<T, N>` lowering (globals + locals).

use super::*;

// ---------------------------------------------------------------------------
// Pass 2: array global lowering
// ---------------------------------------------------------------------------

/// Lower MD2-style global arrays (`const float4 samples[5] = {…};`) to a
/// WGSL-friendly comment-block form that the regex pipeline can survive.
/// The original `const` line confuses `replace_types` (which doesn't know
/// about `[N]`) and `rewrite_local_declarations` (which doesn't accept
/// init lists). We replace it with a flat `var<private>` decl in
/// HLSL-shaped sugar plus a strip marker; the regex pass then leaves it
/// alone because the syntactic shape is already WGSL-friendly.
///
/// Targets the `const float4 samples[5] = {…};` idiom (cluster K + H).
/// Falls back to passing the source through when parsing fails or no
/// matching globals are found.
#[cfg(test)]
pub(crate) fn lower_array_globals(src: &str) -> String {
    let Ok(tu) = parse_hlsl(src) else {
        return src.to_string();
    };
    let mut edits = Vec::new();
    collect_global_array_edits(&tu, src, &mut edits);
    apply_edits(src, &mut edits)
}

/// Walk every top-level `Item::GlobalVar` with an array length + InitList
/// and emit one text edit per match. Shared between [`lower_array_globals`]
/// (the standalone pass) and [`apply_lowerings`] (the combined pass used
/// by [`apply_all`]).
pub(crate) fn collect_global_array_edits(
    tu: &TranslationUnit,
    src: &str,
    edits: &mut Vec<TextEdit>,
) {
    for item in &tu.items {
        let Item::GlobalVar(g) = item else { continue };
        if g.array_len.is_none() {
            continue;
        }
        let Some(init) = &g.init else {
            continue;
        };
        let Expr::InitList(list) = init else {
            continue;
        };
        let comp_ty = type_from_typeref(&g.ty);
        let comp_size = vec_size(comp_ty);
        if comp_size == 0 {
            continue;
        }
        let Some(array_len_expr) = &g.array_len else {
            continue;
        };
        let Expr::Lit(Lit {
            value: LitValue::Int(n),
            ..
        }) = array_len_expr
        else {
            continue;
        };
        let n = *n as usize;
        let comp_wgsl = comp_ty.wgsl_name();
        let mut elems: Vec<String> = Vec::new();
        // The HLSL flat init `{ 0.0, 1.0, ... }` (component-by-component)
        // needs grouping into vecN. We also accept the already-vecN form
        // `{ float4(...), ... }`.
        if list.elems.len() == n {
            // Each element is one vector — emit as-is in WGSL constructor
            // shape.
            for e in &list.elems {
                elems.push(emit_expr_as_wgsl(e, src, comp_ty));
            }
        } else if list.elems.len() == n * comp_size {
            // Flat list: pack `comp_size` scalars per element.
            for chunk in list.elems.chunks(comp_size) {
                let parts: Vec<String> = chunk
                    .iter()
                    .map(|e| emit_expr_as_wgsl(e, src, WgslType::F32))
                    .collect();
                elems.push(format!("{comp_wgsl}({})", parts.join(", ")));
            }
        } else {
            continue; // shape mismatch, leave it alone
        }
        let body = elems.join(", ");
        // Emit as WGSL `var`-initialised array. `const` would be cleaner
        // but requires every element to be const-evaluable; `var` survives
        // arithmetic like `11.0/3.0` and matches the address-space-free
        // form that's legal at both module scope and inside `fs_main`
        // (the existing pipeline drops the const-array body into the
        // function body, so we can't assume module scope).
        let replacement = format!(
            "var {name}: array<{comp_wgsl}, {n}> = array<{comp_wgsl}, {n}>({body});",
            name = g.name,
        );
        edits.push(TextEdit {
            start: g.span.start,
            end: span_includes_semi(src, g.span).end,
            replacement,
        });
    }
}

/// Lower local array declarations (`float3 m[3];`, `float arr[5] = {…};`)
/// to WGSL-valid `var` shape. The downstream regex `rewrite_local_declarations`
/// doesn't understand the `[N]` suffix or init-lists, so the lowered form
/// drops through to naga as-is. Targets the ougiel local-array failures
/// and any preset with bare-typed local arrays.
#[cfg(test)]
pub(crate) fn lower_local_arrays(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 {
        collect_local_array_edits(body, src, &mut edits);
    }
    for item in &tu.items {
        if let Item::Function(f) = item {
            collect_local_array_edits(&f.body, src, &mut edits);
        }
    }
    apply_edits(src, &mut edits)
}

pub(crate) fn collect_local_array_edits(block: &Block, src: &str, edits: &mut Vec<TextEdit>) {
    for s in &block.stmts {
        match s {
            Stmt::LocalDecl(d) if d.array_len.is_some() => {
                let Some(len_expr) = &d.array_len else {
                    continue;
                };
                let Expr::Lit(Lit {
                    value: LitValue::Int(n),
                    ..
                }) = len_expr
                else {
                    continue;
                };
                let comp_ty = type_from_typeref(&d.ty);
                if comp_ty == WgslType::Unknown {
                    continue;
                }
                let comp_wgsl = comp_ty.wgsl_name();
                let replacement = match &d.init {
                    None => format!("var {name}: array<{comp_wgsl}, {n}>;", name = d.name),
                    Some(Expr::InitList(l)) => {
                        // Scalar element types (`float arr[N]`) treat as
                        // 1 component per element so the flat-init branch
                        // works without a special case.
                        let comp_size = if comp_ty.is_scalar() {
                            1
                        } else {
                            vec_size(comp_ty)
                        };
                        if comp_size == 0 {
                            continue;
                        }
                        let n_usize = *n as usize;
                        let elems: Vec<String> = if comp_ty.is_scalar() && l.elems.len() == n_usize
                        {
                            l.elems
                                .iter()
                                .map(|e| emit_expr_as_wgsl(e, src, comp_ty))
                                .collect()
                        } else if l.elems.len() == n_usize {
                            // One vector per array element.
                            l.elems
                                .iter()
                                .map(|e| emit_expr_as_wgsl(e, src, comp_ty))
                                .collect()
                        } else if l.elems.len() == n_usize * comp_size {
                            l.elems
                                .chunks(comp_size)
                                .map(|chunk| {
                                    let parts: Vec<String> = chunk
                                        .iter()
                                        .map(|e| emit_expr_as_wgsl(e, src, WgslType::F32))
                                        .collect();
                                    format!("{comp_wgsl}({})", parts.join(", "))
                                })
                                .collect()
                        } else {
                            continue;
                        };
                        format!(
                            "var {name}: array<{comp_wgsl}, {n}> = array<{comp_wgsl}, {n}>({});",
                            elems.join(", "),
                            name = d.name,
                        )
                    }
                    Some(_) => continue, // unusual init shape, leave alone
                };
                edits.push(TextEdit {
                    start: d.span.start,
                    end: span_includes_semi(src, d.span).end,
                    replacement,
                });
            }
            Stmt::If(i) => {
                if let Stmt::Block(b) = &*i.then_branch {
                    collect_local_array_edits(b, src, edits);
                }
                if let Some(eb) = &i.else_branch
                    && let Stmt::Block(b) = &**eb
                {
                    collect_local_array_edits(b, src, edits);
                }
            }
            Stmt::While(w) => {
                if let Stmt::Block(b) = &*w.body {
                    collect_local_array_edits(b, src, edits);
                }
            }
            Stmt::For(f) => {
                if let Stmt::Block(b) = &*f.body {
                    collect_local_array_edits(b, src, edits);
                }
            }
            Stmt::Block(b) => collect_local_array_edits(b, src, edits),
            _ => {}
        }
    }
}

/// Extend a span forward to include the trailing `;` (if present). The
/// AST's GlobalVar span ends at the name, not the semicolon — for textual
/// replacement we want to swallow the whole declaration.
fn span_includes_semi(src: &str, span: Span) -> Span {
    let bytes = src.as_bytes();
    let mut i = span.end as usize;
    while i < bytes.len() && bytes[i] != b';' {
        i += 1;
    }
    if i < bytes.len() {
        i += 1; // include the `;`
    }
    Span {
        end: i as u32,
        ..span
    }
}

/// Emit an HLSL expression as WGSL text — narrow subset, enough for the
/// array-global lowering case where elements are literals or simple
/// arithmetic. Falls back to the raw original text if the expression is
/// outside the supported shape.
fn emit_expr_as_wgsl(e: &Expr, src: &str, expected: WgslType) -> String {
    let raw = || -> &str { &src[e.span().start as usize..e.span().end as usize] };
    match e {
        Expr::Lit(Lit {
            value: LitValue::Int(v),
            ..
        }) => {
            if matches!(expected, WgslType::F32) {
                format!("{v}.0")
            } else {
                v.to_string()
            }
        }
        Expr::Lit(Lit {
            value: LitValue::Float(v),
            ..
        }) => {
            // Force `.0` suffix when stringification would otherwise omit
            // the dot, so the literal stays a float in WGSL.
            let s = format!("{v}");
            if s.contains('.') || s.contains('e') || s.contains('E') {
                s
            } else {
                format!("{s}.0")
            }
        }
        Expr::Unary(u) if matches!(u.op, UnaryOp::Neg) => {
            format!("-{}", emit_expr_as_wgsl(&u.operand, src, expected))
        }
        Expr::Binary(b) => {
            let op = match b.op {
                BinaryOp::Add => " + ",
                BinaryOp::Sub => " - ",
                BinaryOp::Mul => " * ",
                BinaryOp::Div => " / ",
                _ => return raw().to_string(),
            };
            format!(
                "({}){}({})",
                emit_expr_as_wgsl(&b.lhs, src, expected),
                op,
                emit_expr_as_wgsl(&b.rhs, src, expected)
            )
        }
        Expr::Call(c) => {
            // Constructor calls map to their WGSL spellings.
            let head = match c.callee.as_str() {
                "float2" | "vec2" => "vec2<f32>",
                "float3" | "vec3" => "vec3<f32>",
                "float4" | "vec4" => "vec4<f32>",
                other => other,
            };
            let args: Vec<String> = c
                .args
                .iter()
                .map(|a| emit_expr_as_wgsl(a, src, WgslType::F32))
                .collect();
            format!("{head}({})", args.join(", "))
        }
        _ => raw().to_string(),
    }
}