boids_compute.glsl

#[compute]
#version 450

layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;

// Boid data
layout(set = 0, binding = 0, std430) restrict buffer PositionBuffer {
    vec2 data[];
} positions;

layout(set = 0, binding = 1, std430) restrict buffer VelocityBuffer {
    vec2 data[];
} velocities;

// Parameters - exactly 64 bytes
layout(push_constant) uniform Params {
    float boid_count;        // 0-4
    float neighbor_radius;   // 4-8
    float neighbor_radius_sq;// 8-12
    float separation_weight; // 12-16
    float alignment_weight;  // 16-20
    float cohesion_weight;   // 20-24
    float max_speed;         // 24-28
    float max_force;         // 28-32
    float delta_time;        // 32-36
    float box_top_left_x;    // 36-40
    float box_top_left_y;    // 40-44
    float box_size_x;        // 44-48
    float box_size_y;        // 48-52
    float padding1;          // 52-56
    float padding2;          // 56-60
    float padding3;          // 60-64
} params;

// Clamp function for vec2
vec2 clamp_vec2(vec2 v, float min_val, float max_val) {
    return vec2(
        clamp(v.x, min_val, max_val),
        clamp(v.y, min_val, max_val)
    );
}

void main() {
    uint idx = gl_GlobalInvocationID.x;
    uint boid_count = uint(params.boid_count);
    if (idx >= boid_count) return;
    
    vec2 box_top_left = vec2(params.box_top_left_x, params.box_top_left_y);
    vec2 box_size = vec2(params.box_size_x, params.box_size_y);
    
    vec2 pos = positions.data[idx];
    vec2 vel = velocities.data[idx];
    
    // Flocking forces
    vec2 separation = vec2(0.0);
    vec2 alignment = vec2(0.0);
    vec2 cohesion = vec2(0.0);
    uint count = 0u;
    
    // Check neighbors
    for (uint j = 0u; j < boid_count; j++) {
        if (j == idx) continue;
        
        vec2 other_pos = positions.data[j];
        vec2 diff = pos - other_pos;
        float dist_sq = dot(diff, diff);
        
        if (dist_sq < params.neighbor_radius_sq && dist_sq > 0.0) {
            // Separation: inverse square
            separation += diff / max(dist_sq, 0.01);
            
            // Alignment
            alignment += velocities.data[j];
            
            // Cohesion
            cohesion += other_pos;
            count++;
        }
    }
    
    vec2 acceleration = vec2(0.0);
    
    if (count > 0u) {
        float count_f = float(count);
        
        // Separation
        float sep_len = length(separation);
        if (sep_len > 0.0) {
            separation = normalize(separation) * params.max_speed - vel;
            separation = clamp_vec2(separation, -params.max_force, params.max_force);
            acceleration += separation * params.separation_weight;
        }
        
        // Alignment
        vec2 avg_vel = alignment / count_f;
        float avg_vel_len = length(avg_vel);
        if (avg_vel_len > 0.0) {
            alignment = normalize(avg_vel) * params.max_speed - vel;
            alignment = clamp_vec2(alignment, -params.max_force, params.max_force);
            acceleration += alignment * params.alignment_weight;
        }
        
        // Cohesion
        vec2 center = cohesion / count_f;
        vec2 to_center = center - pos;
        float to_center_len = length(to_center);
        if (to_center_len > 0.0) {
            cohesion = normalize(to_center) * params.max_speed - vel;
            cohesion = clamp_vec2(cohesion, -params.max_force, params.max_force);
            acceleration += cohesion * params.cohesion_weight;
        }
    }
    
    // Update velocity
    vel += acceleration * params.delta_time;
    float vel_len = length(vel);
    if (vel_len > params.max_speed) {
        vel = (vel / vel_len) * params.max_speed;
    }
    
    // Update position
    pos += vel * params.delta_time;
    
    // Wrap boundaries
    vec2 box_min = box_top_left;
    vec2 box_max = box_top_left + box_size;
    
    if (pos.x < box_min.x) pos.x = box_max.x;
    else if (pos.x > box_max.x) pos.x = box_min.x;
    
    if (pos.y < box_min.y) pos.y = box_max.y;
    else if (pos.y > box_max.y) pos.y = box_min.y;
    
    // Write back
    positions.data[idx] = pos;
    velocities.data[idx] = vel;
}