tx2-iff 0.1.0

// Symplectic warp field computation shader
// Calculates velocity field from vortex primitives

struct Metadata {
    width: u32,
    height: u32,
    levels: u32,
    _padding: u32,
}

struct Vortex {
    x: u32,
    y: u32,
    strength: i32,
    radius: u32,
    decay: u32,
}

@group(0) @binding(0) var<storage, read> vortices: array<Vortex>;
@group(0) @binding(1) var<uniform> metadata: Metadata;
@group(0) @binding(2) var<storage, read_write> output: array<u32>;

// Fixed-point math (16.16 format)
const FIXED_ONE: i32 = 65536;
const FIXED_FRAC_BITS: u32 = 16u;

fn fixed_from_int(x: i32) -> i32 {
    return x << 16;
}

fn fixed_from_f32(x: f32) -> i32 {
    return i32(x * 65536.0);
}

fn fixed_to_f32(x: i32) -> f32 {
    return f32(x) / 65536.0;
}

fn fixed_mul(a: i32, b: i32) -> i32 {
    let product = i64(a) * i64(b);
    return i32(product >> 16);
}

fn fixed_div(a: i32, b: i32) -> i32 {
    if (b == 0) {
        return 0;
    }
    let numerator = i64(a) << 16;
    return i32(numerator / i64(b));
}

// Fast square root using Newton-Raphson
fn fixed_sqrt(x: i32) -> i32 {
    if (x <= 0) {
        return 0;
    }

    var guess = x / 2;
    if (guess == 0) {
        guess = 1;
    }

    for (var i = 0; i < 8; i++) {
        let new_guess = (guess + fixed_div(x, guess)) / 2;
        if (abs(new_guess - guess) < 10) {
            break;
        }
        guess = new_guess;
    }

    return guess;
}

// Fixed-point atan2 approximation
fn atan2_fixed(y: i32, x: i32) -> i32 {
    if (x == 0 && y == 0) {
        return 0;
    }

    let abs_y = abs(y);
    let abs_x = abs(x);

    var angle: i32;

    if (abs_x > abs_y) {
        let ratio = fixed_div(abs_y, abs_x);
        angle = fixed_mul(ratio, FIXED_ONE * 4 / 5); // Approximation
    } else {
        let ratio = fixed_div(abs_x, abs_y);
        angle = FIXED_ONE * 3 / 2 - fixed_mul(ratio, FIXED_ONE * 4 / 5);
    }

    if (x < 0) {
        angle = FIXED_ONE * 3 - angle; // π - angle
    }
    if (y < 0) {
        angle = -angle;
    }

    return angle;
}

// Fixed-point sin approximation (Taylor series)
fn sin_fixed(x: i32) -> i32 {
    // Normalize to [-π, π]
    var angle = x;
    let pi = fixed_from_f32(3.14159265);

    // Series: sin(x) ≈ x - x³/6 + x⁵/120
    let x2 = fixed_mul(angle, angle);
    let x3 = fixed_mul(x2, angle);
    let x5 = fixed_mul(x3, x2);

    var result = angle;
    result -= x3 / 6;
    result += x5 / 120;

    return result;
}

// Fixed-point cos approximation
fn cos_fixed(x: i32) -> i32 {
    let pi_half = fixed_from_f32(1.57079633);
    return sin_fixed(pi_half - x);
}

// Fast exponential decay approximation
fn exp_decay_fixed(x: i32) -> i32 {
    // For negative x, exp(-x) ≈ 1 / (1 + x + x²/2)
    // Simplified for speed
    if (x <= 0) {
        return FIXED_ONE;
    }

    let x2 = fixed_mul(x, x);
    let denom = FIXED_ONE + x + x2 / 2;
    return fixed_div(FIXED_ONE, denom);
}

// Calculate warp field contribution from vortex
fn warp_at_point(px: u32, py: u32, vortex: Vortex) -> vec2<f32> {
    let dx = i32(px) - i32(vortex.x);
    let dy = i32(py) - i32(vortex.y);

    let dx_fixed = fixed_from_int(dx);
    let dy_fixed = fixed_from_int(dy);

    // Distance from vortex center
    let r_sq = dx_fixed * dx_fixed + dy_fixed * dy_fixed;
    let r = fixed_sqrt(r_sq);

    if (r == 0) {
        return vec2<f32>(0.0, 0.0);
    }

    // Angle to point
    let theta = atan2_fixed(dy_fixed, dx_fixed);

    // Falloff: exp(-decay * r / radius)
    let radius_fixed = fixed_from_int(i32(vortex.radius));
    let decay_fixed = fixed_from_int(i32(vortex.decay));

    let exp_arg = fixed_div(fixed_mul(decay_fixed, r), radius_fixed);
    let falloff = exp_decay_fixed(exp_arg);

    // Velocity: strength * (sin(θ), -cos(θ)) * falloff
    let strength_fixed = fixed_from_int(vortex.strength);

    let u = fixed_mul(fixed_mul(strength_fixed, sin_fixed(theta)), falloff);
    let v = -fixed_mul(fixed_mul(strength_fixed, cos_fixed(theta)), falloff);

    return vec2<f32>(fixed_to_f32(u), fixed_to_f32(v));
}

@compute @workgroup_size(8, 8, 1)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let x = global_id.x;
    let y = global_id.y;

    if (x >= metadata.width || y >= metadata.height) {
        return;
    }

    let idx = y * metadata.width + x;

    // Read current pixel
    let current = output[idx];

    // Sum contributions from all vortices
    var total_u = 0.0;
    var total_v = 0.0;

    let vortex_count = arrayLength(&vortices);
    for (var i = 0u; i < vortex_count; i++) {
        let vel = warp_at_point(x, y, vortices[i]);
        total_u += vel.x;
        total_v += vel.y;
    }

    // Backward warp: sample from (x - u, y - v)
    let src_x = f32(x) - total_u;
    let src_y = f32(y) - total_v;

    // Bounds check and simple nearest-neighbor sampling
    let sx = u32(clamp(src_x, 0.0, f32(metadata.width - 1u)));
    let sy = u32(clamp(src_y, 0.0, f32(metadata.height - 1u)));

    let src_idx = sy * metadata.width + sx;

    // Copy warped pixel
    output[idx] = output[src_idx];
}