nbody subflow

? Unknown Status unknown.

Code

// LANGUAGE SHOOTOUT: N-body (SUBFLOW VERSION)
//
// STATUS: Testing terse subflow syntax with arrayed captures.
//
// This version uses ~subflow = branch { field: expr } syntax
// for maximum terseness. No ~proc ceremony!

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

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 },
};

const ZERO_DV = [5][3]f64{
    .{ 0, 0, 0 },
    .{ 0, 0, 0 },
    .{ 0, 0, 0 },
    .{ 0, 0, 0 },
    .{ 0, 0, 0 },
};

// ============================================================================
// Events with SUBFLOW implementations - no ~proc needed!
// ============================================================================

~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

// Subflow with chained captures: kinetic first, then potential
~energy = capture({ e: @as(f64, 0) })
| as ke |> for(bodies)
    | each b |> captured { e: ke.e + 0.5 * b.mass * (b.vx*b.vx + b.vy*b.vy + b.vz*b.vz) }
| captured kinetic |> capture({ e: kinetic.e })
    | as pe |> for(0..5)
        | each i |> for(i+1..5)
            | each j |> captured {
                e: pe.e - (bodies[i].mass * bodies[j].mass) / @sqrt(
                    (bodies[i].x - bodies[j].x)*(bodies[i].x - bodies[j].x) +
                    (bodies[i].y - bodies[j].y)*(bodies[i].y - bodies[j].y) +
                    (bodies[i].z - bodies[j].z)*(bodies[i].z - bodies[j].z)
                )
            }
    | captured result |> e { result.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 };
}

~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;
    }
}

// ============================================================================
// THE MAGIC: Terse subflow chain for force calculation!
// ============================================================================

// Step 1: Compute deltas - use const to compute dx/dy/dz once, then dsq from them
~event pair_delta { bi: Body, bj: Body, dt: f64 }
| delta { dx: f64, dy: f64, dz: f64, dsq: f64, dt: f64, mi: f64, mj: f64 }

~pair_delta = const({ dx: bi.x - bj.x, dy: bi.y - bj.y, dz: bi.z - bj.z })
| as d |> delta {
    d.dx,
    d.dy,
    d.dz,
    dsq: d.dx*d.dx + d.dy*d.dy + d.dz*d.dz,
    dt,
    mi: bi.mass,
    mj: bj.mass
}

// Step 2: Compute force - use const to compute mag once (inv_d^3 pattern)
~event pair_force { dx: f64, dy: f64, dz: f64, dsq: f64, dt: f64, mi: f64, mj: f64 }
| result { fx: f64, fy: f64, fz: f64, mi: f64, mj: f64 }

~pair_force = const({ mag: dt / (dsq * @sqrt(dsq)) })
| as c |> result {
    fx: dx * c.mag,
    fy: dy * c.mag,
    fz: dz * c.mag,
    mi,
    mj
}

// ============================================================================
// MAIN FLOW - Subflows + arrayed captures + shortform params!
// ============================================================================

~parse()
| n iters |> init()
    | bodies b[mutable] |> offset(bodies: b[0..])
        |> energy(bodies: b[0..])
            | e e1 |> std.io:print.ln("{{ e1:d:.9 }}")
                |> for(0..iters)
                    | each |> capture({ dv: ZERO_DV })
                        | as acc |> for(0..5)
                            | each i |> for(i+1..5)
                                | each j |> pair_delta(bi: b[i], bj: b[j], dt: DT)
                                    | delta d |> pair_force(d.dx, d.dy, d.dz, d.dsq, d.dt, d.mi, d.mj)
                                        | 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
                                        }
                        | captured deltas |> apply(bodies: b[0..], dt: DT, deltas.dv)
                    | done |> energy(bodies: b[0..])
                        | e e2 |> std.io:print.ln("{{ e2:d:.9 }}")
input.kz