nightshade 0.8.0

struct PlacementUniforms {
    patch_origin: vec2<f32>,
    patch_size: f32,
    blades_per_patch: u32,
    patch_index: u32,
    species_count: u32,
    seed: u32,
    _padding: u32,
    bounds_min: vec3<f32>,
    _padding2: f32,
    bounds_max: vec3<f32>,
    _padding3: f32,
}

struct GrassInstance {
    position: vec3<f32>,
    rotation: f32,
    height: f32,
    width: f32,
    species_index: u32,
    lod: u32,
    base_color: vec4<f32>,
    tip_color: vec4<f32>,
    bend: vec2<f32>,
    _padding: vec2<f32>,
}

struct GrassSpeciesData {
    base_color: vec4<f32>,
    tip_color: vec4<f32>,
    sss_color: vec4<f32>,
    sss_intensity: f32,
    specular_power: f32,
    specular_strength: f32,
    blade_curvature: f32,
    blade_width: f32,
    blade_height_min: f32,
    blade_height_max: f32,
    density_scale: f32,
}

struct SpeciesWeights {
    weights0: vec4<f32>,
    weights1: vec4<f32>,
    count: u32,
    _pad0: u32,
    _pad1: u32,
    _pad2: u32,
    _padding: vec4<u32>,
    _padding2: vec4<u32>,
}

@group(0) @binding(0) var<uniform> uniforms: PlacementUniforms;
@group(0) @binding(1) var<storage, read_write> instances: array<GrassInstance>;
@group(0) @binding(2) var<storage, read> species_data: array<GrassSpeciesData>;
@group(0) @binding(3) var<uniform> species_weights: SpeciesWeights;
@group(0) @binding(4) var heightmap: texture_2d<f32>;
@group(0) @binding(5) var heightmap_sampler: sampler;

const PI: f32 = 3.14159265359;

fn pcg_hash(input: u32) -> u32 {
    var state = input * 747796405u + 2891336453u;
    var word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
    return (word >> 22u) ^ word;
}

fn rand_float(seed: ptr<function, u32>) -> f32 {
    *seed = pcg_hash(*seed);
    return f32(*seed) / 4294967295.0;
}

fn rand_range(seed: ptr<function, u32>, min_val: f32, max_val: f32) -> f32 {
    return min_val + rand_float(seed) * (max_val - min_val);
}

fn get_weight(index: u32) -> f32 {
    if index < 4u {
        switch index {
            case 0u: { return species_weights.weights0.x; }
            case 1u: { return species_weights.weights0.y; }
            case 2u: { return species_weights.weights0.z; }
            case 3u: { return species_weights.weights0.w; }
            default: { return 0.0; }
        }
    } else {
        switch index - 4u {
            case 0u: { return species_weights.weights1.x; }
            case 1u: { return species_weights.weights1.y; }
            case 2u: { return species_weights.weights1.z; }
            case 3u: { return species_weights.weights1.w; }
            default: { return 0.0; }
        }
    }
}

fn select_species(random: f32) -> u32 {
    var cumulative = 0.0;
    for (var index = 0u; index < species_weights.count; index++) {
        cumulative += get_weight(index);
        if random < cumulative {
            return index;
        }
    }
    return species_weights.count - 1u;
}

fn sample_terrain_height(position: vec2<f32>) -> f32 {
    let bounds_size = uniforms.bounds_max.xz - uniforms.bounds_min.xz;
    let uv = (position - uniforms.bounds_min.xz) / bounds_size;

    if uv.x < 0.0 || uv.x > 1.0 || uv.y < 0.0 || uv.y > 1.0 {
        return 0.0;
    }

    return textureSampleLevel(heightmap, heightmap_sampler, uv, 0.0).r;
}

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let blade_index = global_id.x;

    if blade_index >= uniforms.blades_per_patch {
        return;
    }

    let instance_index = uniforms.patch_index * uniforms.blades_per_patch + blade_index;

    var seed = uniforms.seed ^ (instance_index * 1973u + blade_index * 9277u + uniforms.patch_index * 26699u);

    let grid_size = u32(ceil(sqrt(f32(uniforms.blades_per_patch))));
    let grid_x = blade_index % grid_size;
    let grid_z = blade_index / grid_size;

    let cell_size = uniforms.patch_size / f32(grid_size);
    let base_x = uniforms.patch_origin.x + f32(grid_x) * cell_size;
    let base_z = uniforms.patch_origin.y + f32(grid_z) * cell_size;

    let jitter_x = rand_range(&seed, 0.0, cell_size);
    let jitter_z = rand_range(&seed, 0.0, cell_size);

    let position_x = base_x + jitter_x;
    let position_z = base_z + jitter_z;

    if position_x < uniforms.bounds_min.x || position_x > uniforms.bounds_max.x ||
       position_z < uniforms.bounds_min.z || position_z > uniforms.bounds_max.z {
        instances[instance_index].height = 0.0;
        return;
    }

    let position_y = sample_terrain_height(vec2<f32>(position_x, position_z));

    let species_random = rand_float(&seed);
    let species_index = select_species(species_random);
    let species = species_data[species_index];

    let density_check = rand_float(&seed);
    if density_check > species.density_scale {
        instances[instance_index].height = 0.0;
        return;
    }

    let height = rand_range(&seed, species.blade_height_min, species.blade_height_max);
    let width = species.blade_width * rand_range(&seed, 0.8, 1.2);
    let rotation = rand_range(&seed, 0.0, 2.0 * PI);

    let color_variation = rand_range(&seed, 0.9, 1.1);
    var base_color = species.base_color;
    base_color.x *= color_variation;
    base_color.y *= color_variation;
    base_color.z *= color_variation;

    var tip_color = species.tip_color;
    tip_color.x *= color_variation;
    tip_color.y *= color_variation;
    tip_color.z *= color_variation;

    instances[instance_index].position = vec3<f32>(position_x, position_y, position_z);
    instances[instance_index].rotation = rotation;
    instances[instance_index].height = height;
    instances[instance_index].width = width;
    instances[instance_index].species_index = species_index;
    instances[instance_index].lod = 0u;
    instances[instance_index].base_color = base_color;
    instances[instance_index].tip_color = tip_color;
    instances[instance_index].bend = vec2<f32>(0.0, 0.0);
    instances[instance_index]._padding = vec2<f32>(0.0, 0.0);
}