nightshade 0.13.3

const PI: f32 = 3.141592653589793;
const MAX_STEPS: i32 = 64;

const SHAPE_BOX: u32 = 0u;
const SHAPE_CYLINDER: u32 = 1u;
const SHAPE_SPHERE: u32 = 2u;

const FLOW_WATERFALL: u32 = 0u;
const FLOW_MIST: u32 = 1u;
const FLOW_CASCADE: u32 = 2u;

struct Uniforms {
    view: mat4x4<f32>,
    projection: mat4x4<f32>,
    view_projection: mat4x4<f32>,
    camera_position: vec4<f32>,
    time: f32,
    _pad0: f32,
    _pad1: f32,
    _pad2: f32,
}

struct WaterMaterial {
    base_color: vec4<f32>,
    water_color: vec4<f32>,
    wave_height: f32,
    choppy: f32,
    speed: f32,
    freq: f32,
    specular_strength: f32,
    fresnel_power: f32,
    volume_shape: u32,
    volume_flow_type: u32,
    volume_size: vec3<f32>,
    is_volumetric: u32,
    flow_direction: vec2<f32>,
    flow_strength: f32,
    _flow_padding: f32,
}

struct WaterInstance {
    model_0: vec4<f32>,
    model_1: vec4<f32>,
    model_2: vec4<f32>,
    model_3: vec4<f32>,
    material_index: u32,
    _pad0: u32,
    _pad1: u32,
    _pad2: u32,
}

struct VertexInput {
    @location(0) position: vec3<f32>,
    @location(1) normal: vec3<f32>,
    @location(2) tex_coords: vec2<f32>,
}

struct VertexOutput {
    @builtin(position) clip_position: vec4<f32>,
    @location(0) world_pos: vec3<f32>,
    @location(1) local_pos: vec3<f32>,
    @location(2) @interpolate(flat) material_index: u32,
    @location(3) @interpolate(flat) instance_index: u32,
}

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var<storage, read> materials: array<WaterMaterial>;
@group(0) @binding(2) var<storage, read> instances: array<WaterInstance>;
@group(0) @binding(3) var<storage, read> visible_indices: array<u32>;

fn hash3(p: vec3<f32>) -> f32 {
    var p3 = fract(p * 0.1031);
    p3 = p3 + dot(p3, p3.yzx + 33.33);
    return fract((p3.x + p3.y) * p3.z);
}

fn noise3d(p: vec3<f32>) -> f32 {
    let i = floor(p);
    let f = fract(p);
    let u = f * f * (3.0 - 2.0 * f);

    return mix(
        mix(
            mix(hash3(i + vec3<f32>(0.0, 0.0, 0.0)), hash3(i + vec3<f32>(1.0, 0.0, 0.0)), u.x),
            mix(hash3(i + vec3<f32>(0.0, 1.0, 0.0)), hash3(i + vec3<f32>(1.0, 1.0, 0.0)), u.x),
            u.y
        ),
        mix(
            mix(hash3(i + vec3<f32>(0.0, 0.0, 1.0)), hash3(i + vec3<f32>(1.0, 0.0, 1.0)), u.x),
            mix(hash3(i + vec3<f32>(0.0, 1.0, 1.0)), hash3(i + vec3<f32>(1.0, 1.0, 1.0)), u.x),
            u.y
        ),
        u.z
    );
}

fn fbm3d(p: vec3<f32>, octaves: i32) -> f32 {
    var value = 0.0;
    var amplitude = 0.5;
    var frequency = 1.0;

    for (var index = 0; index < octaves; index = index + 1) {
        value = value + amplitude * noise3d(p * frequency);
        amplitude = amplitude * 0.5;
        frequency = frequency * 2.0;
    }
    return value;
}

fn sdf_box(p: vec3<f32>, half_size: vec3<f32>) -> f32 {
    let d = abs(p) - half_size;
    return length(max(d, vec3<f32>(0.0))) + min(max(d.x, max(d.y, d.z)), 0.0);
}

fn sdf_cylinder(p: vec3<f32>, radius: f32, half_height: f32) -> f32 {
    let d = vec2<f32>(length(p.xz) - radius, abs(p.y) - half_height);
    return min(max(d.x, d.y), 0.0) + length(max(d, vec2<f32>(0.0)));
}

fn sdf_ellipsoid(p: vec3<f32>, radii: vec3<f32>) -> f32 {
    let k0 = length(p / radii);
    let k1 = length(p / (radii * radii));
    return k0 * (k0 - 1.0) / k1;
}

fn sdf_volume(p: vec3<f32>, shape: u32, size: vec3<f32>) -> f32 {
    let half_size = size * 0.5;
    if shape == SHAPE_CYLINDER {
        return sdf_cylinder(p, half_size.x, half_size.y);
    } else if shape == SHAPE_SPHERE {
        return sdf_ellipsoid(p, half_size);
    } else {
        return sdf_box(p, half_size);
    }
}

fn density_waterfall(normalized_pos: vec3<f32>, time: f32, mat: WaterMaterial) -> f32 {
    let flow_speed = mat.speed * 2.0;
    let flow_offset = vec3<f32>(0.0, time * flow_speed, 0.0);

    let stretched_pos = vec3<f32>(
        normalized_pos.x * 3.0,
        normalized_pos.y * 0.5,
        normalized_pos.z * 3.0
    );

    let noise_pos = stretched_pos * mat.freq * 2.0 + flow_offset;
    var density = fbm3d(noise_pos, 4);

    let height_factor = smoothstep(-1.0, -0.6, normalized_pos.y) * (1.0 - smoothstep(0.6, 1.0, normalized_pos.y));
    density = density * height_factor;

    let turbulence = fbm3d(noise_pos * 2.0 + vec3<f32>(17.3, 0.0, 23.7), 3) * mat.choppy * 0.1;
    density = density + turbulence;

    return density;
}

fn density_mist(normalized_pos: vec3<f32>, time: f32, mat: WaterMaterial) -> f32 {
    let slow_time = time * mat.speed * 0.2;

    let rise_offset = vec3<f32>(0.0, -slow_time * 2.0, 0.0);
    let drift_offset = vec3<f32>(slow_time * 0.3, 0.0, slow_time * 0.2);

    let noise_pos1 = normalized_pos * mat.freq * 1.5 + rise_offset + drift_offset;
    let noise_pos2 = normalized_pos * mat.freq * 3.0 + rise_offset * 0.7 - drift_offset;

    var density = fbm3d(noise_pos1, 3) * 0.6;
    density = density + fbm3d(noise_pos2, 2) * 0.4;

    density = smoothstep(0.3, 0.7, density);

    let dist_from_center = length(normalized_pos.xz);
    let center_falloff = 1.0 - smoothstep(0.0, 1.0, dist_from_center);
    density = density * center_falloff * center_falloff;

    let height_falloff = 1.0 - smoothstep(-0.5, 0.8, normalized_pos.y);
    let bottom_falloff = smoothstep(-1.0, -0.5, normalized_pos.y);
    density = density * height_falloff * bottom_falloff;

    let wisps = noise3d(normalized_pos * 6.0 + rise_offset * 3.0);
    density = density * smoothstep(-0.2, 0.3, wisps);

    return density * 0.25;
}

fn density_cascade(normalized_pos: vec3<f32>, time: f32, mat: WaterMaterial) -> f32 {
    let flow_speed = mat.speed * 1.5;
    let flow_offset = vec3<f32>(0.0, time * flow_speed, 0.0);

    let cascade_pos = normalized_pos * mat.freq * 4.0 + flow_offset;

    var density = 0.0;

    for (var index = 0; index < 3; index = index + 1) {
        let offset = vec3<f32>(f32(index) * 7.3, f32(index) * 3.7, f32(index) * 5.1);
        let cascade_noise = fbm3d(cascade_pos + offset, 3);
        let cascade = smoothstep(0.4, 0.6, cascade_noise);
        density = density + cascade * 0.4;
    }

    let height_factor = 1.0 - smoothstep(0.5, 1.0, normalized_pos.y);
    density = density * height_factor;

    let dist_from_center = length(normalized_pos.xz);
    let center_factor = 1.0 - smoothstep(0.0, 0.8, dist_from_center);
    density = density * center_factor;

    return density;
}

fn volume_density(local_pos: vec3<f32>, time: f32, mat: WaterMaterial) -> f32 {
    let half_size = mat.volume_size * 0.5;

    let sdf_dist = sdf_volume(local_pos, mat.volume_shape, mat.volume_size);
    if sdf_dist > 0.0 {
        return 0.0;
    }

    let normalized_pos = local_pos / half_size;

    var density: f32;
    if mat.volume_flow_type == FLOW_MIST {
        density = density_mist(normalized_pos, time, mat);
    } else if mat.volume_flow_type == FLOW_CASCADE {
        density = density_cascade(normalized_pos, time, mat);
    } else {
        density = density_waterfall(normalized_pos, time, mat);
    }

    let edge_factor = 1.0 - smoothstep(-0.15, 0.0, sdf_dist);
    density = density * edge_factor;

    return clamp(density * 1.5, 0.0, 1.0);
}

fn ray_volume_intersection(ray_origin: vec3<f32>, ray_dir: vec3<f32>, shape: u32, size: vec3<f32>) -> vec2<f32> {
    let half_size = size * 0.5;
    let inv_dir = 1.0 / ray_dir;
    let t1 = (-half_size - ray_origin) * inv_dir;
    let t2 = (half_size - ray_origin) * inv_dir;
    let t_min_vec = min(t1, t2);
    let t_max_vec = max(t1, t2);
    let t_near = max(max(t_min_vec.x, t_min_vec.y), t_min_vec.z);
    let t_far = min(min(t_max_vec.x, t_max_vec.y), t_max_vec.z);
    return vec2<f32>(max(t_near, 0.0), t_far);
}

@vertex
fn vs_main(
    vertex: VertexInput,
    @builtin(instance_index) instance_index: u32
) -> VertexOutput {
    var out: VertexOutput;

    let visible_idx = visible_indices[instance_index];
    let instance = instances[visible_idx];

    let model = mat4x4<f32>(
        instance.model_0,
        instance.model_1,
        instance.model_2,
        instance.model_3
    );

    let mat = materials[instance.material_index];
    let scaled_pos = vertex.position * mat.volume_size;
    let world_pos = (model * vec4<f32>(scaled_pos, 1.0)).xyz;

    out.world_pos = world_pos;
    out.local_pos = scaled_pos;
    out.clip_position = uniforms.view_projection * vec4<f32>(world_pos, 1.0);
    out.material_index = instance.material_index;
    out.instance_index = visible_idx;

    return out;
}

struct FragmentOutput {
    @location(0) color: vec4<f32>,
    @builtin(frag_depth) depth: f32,
}

@fragment
fn fs_main(in: VertexOutput) -> FragmentOutput {
    let mat = materials[in.material_index];
    let camera_pos = uniforms.camera_position.xyz;

    let ray_origin = camera_pos;
    let ray_dir = normalize(in.world_pos - camera_pos);

    let instance = instances[in.instance_index];
    let model = mat4x4<f32>(
        instance.model_0,
        instance.model_1,
        instance.model_2,
        instance.model_3
    );
    let model_inv = mat4x4<f32>(
        vec4<f32>(1.0, 0.0, 0.0, 0.0),
        vec4<f32>(0.0, 1.0, 0.0, 0.0),
        vec4<f32>(0.0, 0.0, 1.0, 0.0),
        vec4<f32>(-model[3].x, -model[3].y, -model[3].z, 1.0)
    );

    let local_ray_origin = (model_inv * vec4<f32>(ray_origin, 1.0)).xyz;
    let local_ray_dir = normalize((model_inv * vec4<f32>(ray_dir, 0.0)).xyz);

    let t_bounds = ray_volume_intersection(local_ray_origin, local_ray_dir, mat.volume_shape, mat.volume_size);

    if t_bounds.x > t_bounds.y {
        discard;
    }

    let step_size = 1.0 / f32(MAX_STEPS);
    let ray_length = t_bounds.y - t_bounds.x;
    let actual_step = ray_length * step_size;

    var accumulated_color = vec3<f32>(0.0);
    var accumulated_alpha = 0.0;

    let sun_direction = normalize(vec3<f32>(0.5, 0.8, 0.3));
    let sun_color = vec3<f32>(1.0, 0.95, 0.85);
    let ambient = vec3<f32>(0.3, 0.4, 0.5);

    for (var index = 0; index < MAX_STEPS; index = index + 1) {
        if accumulated_alpha > 0.95 {
            break;
        }

        let t = t_bounds.x + (f32(index) + 0.5) * actual_step;
        let sample_pos = local_ray_origin + local_ray_dir * t;

        let density = volume_density(sample_pos, uniforms.time, mat);

        if density > 0.01 {
            let foam_amount = smoothstep(0.3, 0.7, density);
            let base_water = mat.base_color.rgb;
            let foam_white = vec3<f32>(0.9, 0.95, 1.0);
            var sample_color = mix(base_water, foam_white, foam_amount);

            let light_density = volume_density(sample_pos + sun_direction * 0.1, uniforms.time, mat);
            let shadow = exp(-light_density * 2.0);
            let lighting = ambient + sun_color * shadow * 0.7;
            sample_color = sample_color * lighting;

            let sample_alpha = density * actual_step * 4.0;

            accumulated_color = accumulated_color + sample_color * sample_alpha * (1.0 - accumulated_alpha);
            accumulated_alpha = accumulated_alpha + sample_alpha * (1.0 - accumulated_alpha);
        }
    }

    if accumulated_alpha < 0.01 {
        discard;
    }

    let view_dir = -ray_dir;
    let fresnel = pow(1.0 - max(dot(view_dir, vec3<f32>(0.0, 0.0, 1.0)), 0.0), mat.fresnel_power);
    accumulated_color = accumulated_color + vec3<f32>(0.1, 0.15, 0.2) * fresnel * accumulated_alpha;

    var out: FragmentOutput;
    out.color = vec4<f32>(accumulated_color, accumulated_alpha);
    if t_bounds.x > 0.001 {
        let entry_world = ray_origin + ray_dir * t_bounds.x;
        let entry_clip = uniforms.view_projection * vec4<f32>(entry_world, 1.0);
        out.depth = entry_clip.z / entry_clip.w;
    } else {
        out.depth = in.clip_position.z;
    }
    return out;
}