nbody arrayed capture

// LANGUAGE SHOOTOUT: N-body (ARRAYED CAPTURE VERSION)
//
// STATUS: Testing array-indexed captures for velocity delta accumulation.
//
// This version uses [5][3]f64 array instead of 15 scalar fields (dv0x..dv4z).
// Demonstrates the `captured { arr[i]: value }` syntax for clean array updates.
//
// COMPARISON:
//   Scalar version: capture({ dv0x: 0, dv0y: 0, ... dv4z: 0 })  // 15 fields!
//   Array version:  capture({ dv: [5][3]f64{ ... } })           // 1 field!

const std = @import("std");
~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;

const Body = struct { x: f64, y: f64, z: f64, vx: f64, vy: f64, vz: f64, mass: f64 };

const BODIES = [_]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 },
};

// Zero-initialized velocity deltas array
const ZERO_DV = [5][3]f64{
    .{ 0, 0, 0 },
    .{ 0, 0, 0 },
    .{ 0, 0, 0 },
    .{ 0, 0, 0 },
    .{ 0, 0, 0 },
};

// ============================================================================
// Events
// ============================================================================

~event init {}
| bodies [5]Body

~proc init { return .{ .bodies = BODIES }; }

~event offset { bodies: []Body }

~proc offset {
    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;
}

~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);
        var j = i + 1;
        while (j < 5) : (j += 1) {
            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 };
}

~event print_e { e: f64 }

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

~event parse {}
| n u32

~proc parse {
    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 <n>\n", .{}); unreachable; }
    return .{ .n = std.fmt.parseInt(u32, args[1], 10) catch unreachable };
}

// Apply velocity deltas from [5][3] array and advance positions
~event apply { bodies: []Body, dt: f64, dv: [5][3]f64 }

~proc apply {
    for (bodies, 0..) |*b, i| {
        b.vx += dv[i][0];
        b.vy += dv[i][1];
        b.vz += dv[i][2];
        b.x += dt * b.vx;
        b.y += dt * b.vy;
        b.z += dt * b.vz;
    }
}

// Compute force contribution for one body pair
~event pair_force { bi: Body, bj: Body, dt: f64 }
| result { fx: f64, fy: f64, fz: f64, mi: f64, mj: f64 }

~proc pair_force {
    const dx = bi.x - bj.x;
    const dy = bi.y - bj.y;
    const dz = bi.z - bj.z;
    const dsq = dx*dx + dy*dy + dz*dz;
    const mag = dt / (dsq * @sqrt(dsq));
    return .{ .result = .{
        .fx = dx * mag,
        .fy = dy * mag,
        .fz = dz * mag,
        .mi = bi.mass,
        .mj = bj.mass,
    } };
}

// ============================================================================
// MAIN FLOW - Using arrayed capture for velocity deltas!
// ============================================================================

~parse()
| n iters |> init()
    | bodies b[mutable] |> offset(bodies: b[0..])
        |> energy(bodies: b[0..])
            | e e1 |> print_e(e: e1)
                |> for(0..iters)
                    | each |> capture({ dv: ZERO_DV })
                        | as acc |> for(0..5)
                            | each i |> for(i+1..5)
                                | each j |> pair_force(bi: b[i], bj: b[j], dt: DT)
                                    | result f |> captured { dv[i][0]: acc.dv[i][0] - f.fx*f.mj, dv[i][1]: acc.dv[i][1] - f.fy*f.mj, dv[i][2]: acc.dv[i][2] - f.fz*f.mj, dv[j][0]: acc.dv[j][0] + f.fx*f.mi, dv[j][1]: acc.dv[j][1] + f.fy*f.mi, dv[j][2]: acc.dv[j][2] + f.fz*f.mi }
                                | done |> _
                            | done |> _
                        | captured deltas |> apply(bodies: b[0..], dt: DT, dv: deltas.dv)
                    | done |> energy(bodies: b[0..])
                        | e e2 |> print_e(e: e2)