nbody kernel pairwise

// LANGUAGE SHOOTOUT: N-body (KERNEL PAIRWISE VERSION)
//
// Full benchmark using kernel:pairwise with noalias optimization.
// Target: match Rust performance (~1.36s on 50M iterations)

const std = @import("std");
~import "$std/kernel"
~import "$std/control"

const PI: f64 = 3.141592653589793;
const SOLAR_MASS: f64 = 4.0 * PI * PI;
const DAYS_PER_YEAR: f64 = 365.24;
const DT: f64 = 0.01;

// Kernel shape
~std.kernel:shape(Body) {
    x: f64,
    y: f64,
    z: f64,
    vx: f64,
    vy: f64,
    vz: f64,
    mass: f64,
}

// Parse command line args
~event parse_args {}
| n u32

~proc parse_args {
    const args = std.process.argsAlloc(std.heap.page_allocator) catch unreachable;
    defer std.process.argsFree(std.heap.page_allocator, args);
    if (args.len < 2) {
        std.debug.print("Usage: nbody <iterations>\n", .{});
        std.process.exit(1);
    }
    return .{ .n = std.fmt.parseInt(u32, args[1], 10) catch unreachable };
}

// Offset momentum so system has zero total momentum
~event offset_momentum { bodies: []Body }

~proc offset_momentum {
    var px: f64 = 0;
    var py: f64 = 0;
    var pz: f64 = 0;
    for (bodies) |b| {
        px += b.vx * b.mass;
        py += b.vy * b.mass;
        pz += b.vz * b.mass;
    }
    bodies[0].vx = -px / SOLAR_MASS;
    bodies[0].vy = -py / SOLAR_MASS;
    bodies[0].vz = -pz / SOLAR_MASS;
}

// Calculate total energy
~event energy { bodies: []const Body }
| e f64

~proc energy {
    var e: f64 = 0;
    for (bodies, 0..) |b, i| {
        e += 0.5 * b.mass * (b.vx * b.vx + b.vy * b.vy + b.vz * b.vz);
        for (i + 1..5) |j| {
            const dx = b.x - bodies[j].x;
            const dy = b.y - bodies[j].y;
            const dz = b.z - bodies[j].z;
            e -= (b.mass * bodies[j].mass) / @sqrt(dx * dx + dy * dy + dz * dz);
        }
    }
    return .{ .e = e };
}

// Print energy value
~event print_energy { e: f64 }

~proc print_energy {
    std.debug.print("{d:.9}\n", .{e});
}

// Update positions after velocity changes
~event advance_positions { bodies: []Body }

~proc advance_positions {
    for (bodies) |*b| {
        b.x += DT * b.vx;
        b.y += DT * b.vy;
        b.z += DT * b.vz;
    }
}

// Main flow
~parse_args()
| n iterations |>
    std.kernel:init(Body) {
        { x: 0, y: 0, z: 0, vx: 0, vy: 0, vz: 0, mass: SOLAR_MASS },
        { x: 4.84143144246472090e+00, y: -1.16032004402742839e+00, z: -1.03622044471123109e-01, vx: 1.66007664274403694e-03 * DAYS_PER_YEAR, vy: 7.69901118419740425e-03 * DAYS_PER_YEAR, vz: -6.90460016972063023e-05 * DAYS_PER_YEAR, mass: 9.54791938424326609e-04 * SOLAR_MASS },
        { x: 8.34336671824457987e+00, y: 4.12479856412430479e+00, z: -4.03523417114321381e-01, vx: -2.76742510726862411e-03 * DAYS_PER_YEAR, vy: 4.99852801234917238e-03 * DAYS_PER_YEAR, vz: 2.30417297573763929e-05 * DAYS_PER_YEAR, mass: 2.85885980666130812e-04 * SOLAR_MASS },
        { x: 1.28943695621391310e+01, y: -1.51111514016986312e+01, z: -2.23307578892655734e-01, vx: 2.96460137564761618e-03 * DAYS_PER_YEAR, vy: 2.37847173959480950e-03 * DAYS_PER_YEAR, vz: -2.96589568540237556e-05 * DAYS_PER_YEAR, mass: 4.36624404335156298e-05 * SOLAR_MASS },
        { x: 1.53796971148509165e+01, y: -2.59193146099879641e+01, z: 1.79258772950371181e-01, vx: 2.68067772490389322e-03 * DAYS_PER_YEAR, vy: 1.62824170038242295e-03 * DAYS_PER_YEAR, vz: -9.51592254519715870e-05 * DAYS_PER_YEAR, mass: 5.15138902046611451e-05 * SOLAR_MASS },
    }
    | kernel k |>
        offset_momentum(bodies: k.ptr[0..k.len])
        |> energy(bodies: k.ptr[0..k.len])
            | e e1 |> print_energy(e: e1)
                |> for(0..iterations)
                    | each _ |>
                        std.kernel:pairwise {
                            const dx = k.x - k.other.x;
                            const dy = k.y - k.other.y;
                            const dz = k.z - k.other.z;
                            const dsq = dx*dx + dy*dy + dz*dz;
                            const mag = DT / (dsq * @sqrt(dsq));
                            k.vx -= dx * k.other.mass * mag;
                            k.vy -= dy * k.other.mass * mag;
                            k.vz -= dz * k.other.mass * mag;
                            k.other.vx += dx * k.mass * mag;
                            k.other.vy += dy * k.mass * mag;
                            k.other.vz += dz * k.mass * mag;
                        }
                        |> advance_positions(bodies: k.ptr[0..k.len])
                    | done |> energy(bodies: k.ptr[0..k.len])
                        | e e2 |> print_energy(e: e2)