onedrop_hlsl/rewrite/

swizzle_assign.rs

//! Pass: lower multi-component swizzle assignments to per-component writes.
//!
//! HLSL accepts `vec.xy = expr` / `vec.xy += expr` / `color.rgb = …`
//! verbatim. WGSL refuses with `invalid left-hand side of assignment` and
//! the hint *"WGSL does not support assignments to swizzles; consider
//! assigning each component individually"*. This is the dominant warp-side
//! failure on the in-the-wild corpus (≈ 100 cases on the 2 000-sample).
//!
//! Rewrite shape: `lhs.SWIZ op= rhs` → `floatN _swztmp_K = (rhs); lhs.c0 op= _swztmp_K.x; lhs.c1 op= _swztmp_K.y; …`
//! where `K` is a per-pass counter so nested or chained rewrites don't
//! collide. We emit HLSL-shaped output (`floatN`) so the regex pipeline
//! can lower it to `vec<f32>` like any other declaration. Single-component
//! swizzle assigns (`uv.x = …`) are left alone — WGSL accepts them.
//!
//! We deliberately skip:
//! - swizzles with repeated components (`uv.xx = …`) — undefined in HLSL
//!   too, surface the WGSL error instead of synthesising a fake last-write.
//! - assignments that sit inside a comma-chain (`a.xy = …, b += 1;`).
//!   The parser splits these into separate `Stmt::Assign`s but the
//!   intervening `,` would dangle if we expand one side into multiple
//!   statements. Conservative: leave the chain alone (rare in the corpus).

use super::*;

pub(crate) fn rewrite_swizzle_assigns(src: &str) -> String {
    let Ok(tu) = parse_hlsl(src) else {
        return src.to_string();
    };
    let mut state = RewriteState {
        src,
        edits: Vec::new(),
        counter: 0,
    };
    if let Some(body) = &tu.shader_body {
        walk_block(body, &mut state);
    }
    for item in &tu.items {
        if let Item::Function(f) = item {
            walk_block(&f.body, &mut state);
        }
    }
    apply_edits(src, &mut state.edits)
}

struct RewriteState<'a> {
    src: &'a str,
    edits: Vec<TextEdit>,
    counter: usize,
}

fn walk_block(b: &Block, st: &mut RewriteState) {
    for s in &b.stmts {
        walk_stmt(s, st);
    }
}

fn walk_stmt(s: &Stmt, st: &mut RewriteState) {
    match s {
        Stmt::Assign(a) => try_emit(a, st),
        Stmt::If(i) => {
            walk_stmt(&i.then_branch, st);
            if let Some(e) = &i.else_branch {
                walk_stmt(e, st);
            }
        }
        Stmt::While(w) => walk_stmt(&w.body, st),
        Stmt::For(f) => {
            if let Some(init) = &f.init {
                walk_stmt(init, st);
            }
            walk_stmt(&f.body, st);
        }
        Stmt::Block(b) => walk_block(b, st),
        Stmt::LocalDecl(_) | Stmt::Expr(_) | Stmt::Return(_) | Stmt::Break | Stmt::Continue => {}
    }
}

fn try_emit(a: &AssignStmt, st: &mut RewriteState) {
    let Expr::Swizzle(sw) = &a.target else {
        return;
    };
    let comps = sw.components.as_str();
    if comps.len() < 2 {
        return; // single-component swizzle assign is valid WGSL
    }
    if !components_unique(comps) {
        return; // duplicate writes — leave WGSL to surface the malformed source
    }
    if followed_by_comma(st.src, a.span.end) {
        return; // part of a comma-chain — rewrite would dangle the comma
    }

    let n = comps.len();
    let temp_ty = match n {
        2 => "float2",
        3 => "float3",
        4 => "float4",
        _ => return,
    };
    let id = st.counter;
    st.counter += 1;
    let tmp = format!("_swztmp_{id}");
    let op_str = assign_op_str(a.op);

    // Slice the original LHS base text + RHS text so any inner edits emitted
    // by other passes (binop_vec, vec_cmp, …) on the same source still apply
    // — we only own the wrapper text around them.
    let base_text = slice(st.src, sw.base.span());
    let rhs_text = slice(st.src, a.value.span());

    let mut out = String::new();
    out.push_str(temp_ty);
    out.push(' ');
    out.push_str(&tmp);
    out.push_str(" = (");
    out.push_str(rhs_text);
    out.push(')');
    for (i, c) in comps.chars().enumerate() {
        let lhs_comp = normalise_component(c);
        let rhs_comp = match i {
            0 => 'x',
            1 => 'y',
            2 => 'z',
            3 => 'w',
            _ => unreachable!(),
        };
        out.push_str("; ");
        out.push_str(base_text);
        out.push('.');
        out.push(lhs_comp);
        out.push_str(op_str);
        out.push_str(&tmp);
        out.push('.');
        out.push(rhs_comp);
    }

    st.edits.push(TextEdit {
        start: a.span.start,
        end: a.span.end,
        replacement: out,
    });
}

fn components_unique(s: &str) -> bool {
    let mut seen = [false; 8]; // indexed by normalised xyzw -> 0..3
    for c in s.chars() {
        let idx = match normalise_component(c) {
            'x' => 0,
            'y' => 1,
            'z' => 2,
            'w' => 3,
            _ => return false,
        };
        if seen[idx] {
            return false;
        }
        seen[idx] = true;
    }
    true
}

fn normalise_component(c: char) -> char {
    match c {
        'r' => 'x',
        'g' => 'y',
        'b' => 'z',
        'a' => 'w',
        c => c,
    }
}

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

fn followed_by_comma(src: &str, end: u32) -> bool {
    let bytes = src.as_bytes();
    let mut i = end as usize;
    while i < bytes.len() {
        match bytes[i] {
            b' ' | b'\t' | b'\r' | b'\n' => i += 1,
            b',' => return true,
            _ => return false,
        }
    }
    false
}

fn slice(src: &str, sp: Span) -> &str {
    &src[sp.start as usize..sp.end as usize]
}