main.zig

const std = @import("std");

const Vec2 = struct {
    x: f32,
    y: f32,
};

const Body = struct {
    position: Vec2,
    inv_mass: f32,
};

fn getScreenWidthF32() f32 {
    return 100.0;
}

fn getScreenHeightF32() f32 {
    return 100.0;
}

fn getWorldToScreen2D(pos: Vec2) Vec2 {
    return pos; // identity to keep it simple
}

fn shouldRemoveBody(body: Body) bool {
    const pos = getWorldToScreen2D(body.position);
    return body.inv_mass != 0 and
        (pos.x < 0 or pos.x > getScreenWidthF32() or
            pos.y < 0 or pos.y > getScreenHeightF32());
}

fn shouldKeepBody(body: Body) bool {
    return !shouldRemoveBody(body);
}

fn original(bodies: *std.ArrayList(Body)) void {
    var items_to_remove: u32 = 0;
    for (0..bodies.items.len) |idx| {
        const pos = getWorldToScreen2D(bodies.items[idx].position);
        if (bodies.items[idx].inv_mass != 0) {
            if (pos.x < 0 or pos.x > getScreenWidthF32() or
                pos.y < 0 or pos.y > getScreenHeightF32())
            {
                bodies.items[idx] = bodies.items[bodies.items.len - 1 - items_to_remove];
                items_to_remove += 1;
            }
        }
    }
    bodies.items.len -= items_to_remove;
}

fn fixed(bodies: *std.ArrayList(Body)) void {
    var write_idx: usize = 0;
    for (bodies.items) |body| {
        if (shouldKeepBody(body)) {
            bodies.items[write_idx] = body;
            write_idx += 1;
        }
    }
    bodies.items.len = write_idx;
}

fn partition_approach(bodies: *std.ArrayList(Body)) void {
    var front: usize = 0;
    var back: usize = bodies.items.len;

    while (front < back) {
        // Front pointer: find next remove item
        while (front < back and shouldKeepBody(bodies.items[front])) {
            front += 1;
        }

        // Back pointer: find next keep item
        while (front < back and shouldRemoveBody(bodies.items[back - 1])) {
            back -= 1;
        }

        // If both pointers found their targets, swap
        if (front < back) {
            back -= 1;
            std.mem.swap(Body, &bodies.items[front], &bodies.items[back]);
            front += 1;
        }
    }

    bodies.items.len = front; // Everything before 'front' is keep
}

fn make_breaking_test(allocator: std.mem.Allocator) !std.ArrayList(Body) {
    var bodies = std.ArrayList(Body).init(allocator);

    // remove, keep, remove, remove
    try bodies.append(.{ .position = .{ .x = 200, .y = 50 }, .inv_mass = 1 });
    try bodies.append(.{ .position = .{ .x = 50, .y = 50 }, .inv_mass = 1 });
    try bodies.append(.{ .position = .{ .x = 300, .y = 50 }, .inv_mass = 1 });
    try bodies.append(.{ .position = .{ .x = 400, .y = 50 }, .inv_mass = 1 });

    return bodies;
}

test "original function bug test" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    var bodies_original = try make_breaking_test(allocator);
    original(&bodies_original);
    // original have 1 item, but the wrong item
    try std.testing.expectEqual(1, bodies_original.items.len);
    // (x=400 instead of x=50)
    //try std.testing.expectEqual(50.0, bodies_original.items[0].position.x);

    // should have exactly 1 item (x=50)
    var bodies_correct = try make_breaking_test(allocator);
    fixed(&bodies_correct);
    try std.testing.expectEqual(1, bodies_correct.items.len);
    try std.testing.expectEqual(50.0, bodies_correct.items[0].position.x);

    var bodies_partition = try make_breaking_test(allocator);
    partition_approach(&bodies_partition);
    try std.testing.expectEqual(1, bodies_partition.items.len);
    try std.testing.expectEqual(50.0, bodies_partition.items[0].position.x);
}