zig_particle_SIMD

const std = @import("std");
const print = std.debug.print;
const mem = std.mem;
const tMilli = std.time.milliTimestamp;
const math = std.math;
const rand = std.rand;
const testing = std.testing;

const NUM_PARTICLES = 100000;
const SIMD_WIDTH: usize = 8; // AVX equivalent for 256-bit
const GRAVITY: f32 = -9.81;
const DELTA_TIME: f32 = 0.01;
const FIXED_SEED: u64 = 12345;
const Vector = @Vector(SIMD_WIDTH, f32);

// Structure to hold particle data
pub const ParticleData = struct {
    positions_x: []f32,
    positions_y: []f32,
    positions_z: []f32,
    velocities_x: []f32,
    velocities_y: []f32,
    velocities_z: []f32,
    accelerations_x: []f32,
    accelerations_y: []f32,
    accelerations_z: []f32,
    masses: []f32,
};

pub fn allocateParticleData(allocator: mem.Allocator) !ParticleData {
    const positions_x = try allocator.alloc(f32, NUM_PARTICLES);
    const positions_y = try allocator.alloc(f32, NUM_PARTICLES);
    const positions_z = try allocator.alloc(f32, NUM_PARTICLES);
    const velocities_x = try allocator.alloc(f32, NUM_PARTICLES);
    const velocities_y = try allocator.alloc(f32, NUM_PARTICLES);
    const velocities_z = try allocator.alloc(f32, NUM_PARTICLES);
    const accelerations_x = try allocator.alloc(f32, NUM_PARTICLES);
    const accelerations_y = try allocator.alloc(f32, NUM_PARTICLES);
    const accelerations_z = try allocator.alloc(f32, NUM_PARTICLES);
    const masses = try allocator.alloc(f32, NUM_PARTICLES);

    return ParticleData{
        .positions_x = positions_x,
        .positions_y = positions_y,
        .positions_z = positions_z,
        .velocities_x = velocities_x,
        .velocities_y = velocities_y,
        .velocities_z = velocities_z,
        .accelerations_x = accelerations_x,
        .accelerations_y = accelerations_y,
        .accelerations_z = accelerations_z,
        .masses = masses,
    };
}

pub fn initializeParticles(data: *ParticleData) void {
    // var rng = rand.DefaultPrng.init(@as(u64, @bitCast(tMilli())));
    var rng = rand.DefaultPrng.init(FIXED_SEED);
    for (0..NUM_PARTICLES) |i| {
        data.positions_x[i] = rng.random().float(f32) * 10.0;
        data.positions_y[i] = rng.random().float(f32) * 10.0;
        data.positions_z[i] = rng.random().float(f32) * 10.0;

        data.velocities_x[i] = rng.random().float(f32) * 2.0 - 1.0;
        data.velocities_y[i] = rng.random().float(f32) * 2.0 - 1.0;
        data.velocities_z[i] = rng.random().float(f32) * 2.0 - 1.0;

        data.accelerations_x[i] = 0.0;
        data.accelerations_y[i] = 0.0;
        data.accelerations_z[i] = 0.0;

        data.masses[i] = rng.random().float(f32) * 2.0 + 1.0;
    }
}

fn freeParticleData(allocator: mem.Allocator, data: *ParticleData) void {
  inline for (std.meta.fields(ParticleData)) |field| {
       allocator.free(@field(data, field.name));
    }
}

fn loadVector(ptr: []f32, i: usize) Vector {
    return @as(Vector, ptr[i*SIMD_WIDTH..][0..SIMD_WIDTH].*);
}

fn storeVector(ptr: []f32, i: usize, vec: Vector) void {
    @as(*align(1) @Vector(SIMD_WIDTH, f32), @ptrCast(ptr[i*SIMD_WIDTH..][0..SIMD_WIDTH].ptr)).* = vec;
}

fn updateParticlesSimd(data: *ParticleData) void {
    const gravity_vector: Vector = @splat(GRAVITY);
    const delta_time_vector: Vector = @splat(DELTA_TIME);

    for (0..NUM_PARTICLES / SIMD_WIDTH) |i| {

        // Load data
        const px = loadVector(data.positions_x, i);
        const py = loadVector(data.positions_y, i);
        const pz = loadVector(data.positions_z, i);
        const vx = loadVector(data.velocities_x, i);
        const vy = loadVector(data.velocities_y, i);
        const vz = loadVector(data.velocities_z, i);
        const ax = loadVector(data.accelerations_x, i);
        const ay = loadVector(data.accelerations_y, i);
        var az = loadVector(data.accelerations_z, i);

        // Update accelerations (gravity)
        az = az + gravity_vector;

        // Update velocities
        const nvx: Vector = vx + ax * delta_time_vector;
        const nvy: Vector = vy + ay * delta_time_vector;
        const nvz: Vector = vz + az * delta_time_vector;

        // Update positions
        const npx: Vector = px + nvx * delta_time_vector;
        const npy: Vector = py + nvy * delta_time_vector;
        const npz: Vector = pz + nvz * delta_time_vector;

        // Store results
        storeVector(data.positions_x, i, npx);
        storeVector(data.positions_y, i, npy);
        storeVector(data.positions_z, i, npz);
        storeVector(data.velocities_x, i, nvx);
        storeVector(data.velocities_y, i, nvy);
        storeVector(data.velocities_z, i, nvz);
        storeVector(data.accelerations_x, i, ax);
        storeVector(data.accelerations_y, i, ay);
        storeVector(data.accelerations_z, i, az);
    }
}

pub fn main() !void {

    const allocator = std.heap.page_allocator;

    var particle_data = try allocateParticleData(allocator);
    defer freeParticleData(allocator, &particle_data);
    initializeParticles(&particle_data);

    const num_steps: i32 = 100;

    // Time SIMD version
    print("Starting SIMD test...\n", .{});
    const start_simd = tMilli();
    for (0..num_steps) |_| {
        updateParticlesSimd(&particle_data);
    }
    const end_simd = tMilli();
    const time_simd = @as(f32, @floatFromInt(end_simd - start_simd));
    print("SIMD Time: {d:.2} ms\n", .{time_simd});

    print("First position Zig: x={d:.4}, y={d:.4}, z={d:.4}\n", .{
        particle_data.positions_x[0],
        particle_data.positions_y[0],
        particle_data.positions_z[0],
    });
}