nightshade 0.13.3

const PI: f32 = 3.141592653589793;
const NUM_STEPS: i32 = 32;
const ITER_GEOMETRY: i32 = 3;
const ITER_FRAGMENT: i32 = 5;
const MAX_POLYGON_VERTICES: u32 = 16u;

struct Uniforms {
    view: mat4x4<f32>,
    projection: mat4x4<f32>,
    inv_view_projection: mat4x4<f32>,
    camera_position: vec4<f32>,
    resolution: vec2<f32>,
    time: f32,
    z_near: f32,
    z_far: f32,
    _padding0: f32,
    _padding1: f32,
    _padding2: f32,
    _padding3: f32,
    _padding4: f32,
    _padding5: f32,
    _padding6: f32,
}

struct WaterParams {
    base_color: vec4<f32>,
    water_color: vec4<f32>,
    base_height: f32,
    wave_height: f32,
    choppy: f32,
    speed: f32,
    frequency: f32,
    specular_strength: f32,
    fresnel_power: f32,
    edge_feather_distance: f32,
}

struct WaterPolygon {
    vertices: array<vec4<f32>, 16>,
    vertex_count: u32,
    use_bounds: u32,
    _padding0: f32,
    _padding1: f32,
}

struct VertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) uv: vec2<f32>,
}

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

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var<uniform> water_params: WaterParams;
@group(0) @binding(2) var<uniform> water_polygon: WaterPolygon;

fn point_in_polygon(point: vec2<f32>) -> bool {
    let vertex_count = water_polygon.vertex_count;
    if (vertex_count < 3u) {
        return false;
    }

    var inside = false;
    var j = vertex_count - 1u;

    for (var i = 0u; i < vertex_count; i = i + 1u) {
        let vi = water_polygon.vertices[i].xy;
        let vj = water_polygon.vertices[j].xy;

        if ((vi.y > point.y) != (vj.y > point.y)) {
            let x_intersect = (vj.x - vi.x) * (point.y - vi.y) / (vj.y - vi.y) + vi.x;
            if (point.x < x_intersect) {
                inside = !inside;
            }
        }

        j = i;
    }

    return inside;
}

fn distance_to_segment(point: vec2<f32>, segment_a: vec2<f32>, segment_b: vec2<f32>) -> f32 {
    let ab = segment_b - segment_a;
    let ap = point - segment_a;
    let t = clamp(dot(ap, ab) / max(dot(ab, ab), 0.0001), 0.0, 1.0);
    let closest = segment_a + ab * t;
    return length(point - closest);
}

fn distance_to_polygon_edge(point: vec2<f32>) -> f32 {
    let vertex_count = water_polygon.vertex_count;
    if (vertex_count < 3u) {
        return 0.0;
    }

    var min_dist = 1000000.0;
    var j = vertex_count - 1u;

    for (var i = 0u; i < vertex_count; i = i + 1u) {
        let vi = water_polygon.vertices[i].xy;
        let vj = water_polygon.vertices[j].xy;
        let d = distance_to_segment(point, vi, vj);
        min_dist = min(min_dist, d);
        j = i;
    }

    return min_dist;
}

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

fn noise(p: vec2<f32>) -> f32 {
    let i = floor(p);
    let f = fract(p);
    let u = f * f * (3.0 - 2.0 * f);
    return -1.0 + 2.0 * mix(
        mix(hash(i + vec2<f32>(0.0, 0.0)), hash(i + vec2<f32>(1.0, 0.0)), u.x),
        mix(hash(i + vec2<f32>(0.0, 1.0)), hash(i + vec2<f32>(1.0, 1.0)), u.x),
        u.y
    );
}

fn sea_octave(uv_in: vec2<f32>, choppy: f32) -> f32 {
    var uv = uv_in + noise(uv_in);
    var wv = 1.0 - abs(sin(uv));
    let swv = abs(cos(uv));
    wv = mix(wv, swv, wv);
    return pow(1.0 - pow(wv.x * wv.y, 0.65), choppy);
}

fn get_sea_time() -> f32 {
    return 1.0 + uniforms.time * water_params.speed;
}

fn map(p: vec3<f32>, iter: i32) -> f32 {
    var freq = water_params.frequency;
    var amp = water_params.wave_height;
    var choppy = water_params.choppy;
    var uv = p.xz;
    uv.x = uv.x * 0.75;
    let sea_time = get_sea_time();
    let octave_m = mat2x2<f32>(1.6, 1.2, -1.2, 1.6);

    var d: f32;
    var h: f32 = 0.0;
    for (var i: i32 = 0; i < iter; i = i + 1) {
        d = sea_octave((uv + sea_time) * freq, choppy);
        d = d + sea_octave((uv - sea_time) * freq, choppy);
        h = h + d * amp;
        uv = octave_m * uv;
        freq = freq * 1.9;
        amp = amp * 0.22;
        choppy = mix(choppy, 1.0, 0.2);
    }
    return p.y - (water_params.base_height + h);
}

fn map_detailed(p: vec3<f32>) -> f32 {
    return map(p, ITER_FRAGMENT);
}

fn get_sky_color(e: vec3<f32>) -> vec3<f32> {
    let y_val = (max(e.y, 0.0) * 0.8 + 0.2) * 0.8;
    return vec3<f32>(pow(1.0 - y_val, 2.0), 1.0 - y_val, 0.6 + (1.0 - y_val) * 0.4) * 1.1;
}

fn get_normal(p: vec3<f32>, eps: f32) -> vec3<f32> {
    var n: vec3<f32>;
    n.y = map_detailed(p);
    n.x = map_detailed(vec3<f32>(p.x + eps, p.y, p.z)) - n.y;
    n.z = map_detailed(vec3<f32>(p.x, p.y, p.z + eps)) - n.y;
    n.y = eps;
    return normalize(n);
}

fn heightmap_tracing(ori: vec3<f32>, dir: vec3<f32>) -> vec3<f32> {
    var tm: f32 = 0.0;
    var tx: f32 = uniforms.z_far;
    var hx = map(ori + dir * tx, ITER_GEOMETRY);
    if (hx > 0.0) {
        return ori + dir * tx;
    }
    var hm = map(ori, ITER_GEOMETRY);
    var tmid: f32 = 0.0;
    for (var i: i32 = 0; i < NUM_STEPS; i = i + 1) {
        let denom = hm - hx;
        let frac = hm / max(abs(denom), 0.0001) * sign(denom);
        tmid = mix(tm, tx, clamp(frac, 0.0, 1.0));
        let p = ori + dir * tmid;
        let hmid = map(p, ITER_GEOMETRY);
        if (hmid < 0.0) {
            tx = tmid;
            hx = hmid;
        } else {
            tm = tmid;
            hm = hmid;
        }
    }
    return ori + dir * tmid;
}

fn diffuse_factor(n: vec3<f32>, l: vec3<f32>, p: f32) -> f32 {
    return pow(dot(n, l) * 0.4 + 0.6, p);
}

fn specular_factor(n: vec3<f32>, l: vec3<f32>, e: vec3<f32>, s: f32) -> f32 {
    let nrm = (s + 8.0) / (PI * 8.0);
    return pow(max(dot(reflect(e, n), l), 0.0), s) * nrm;
}

fn get_sea_color(p: vec3<f32>, n: vec3<f32>, l: vec3<f32>, eye: vec3<f32>, dist: vec3<f32>) -> vec3<f32> {
    let sea_base = water_params.base_color.xyz;
    let sea_water_color = water_params.water_color.xyz;

    var fresnel = clamp(1.0 - dot(n, -eye), 0.0, 1.0);
    fresnel = pow(fresnel, water_params.fresnel_power) * 0.65;

    let reflected = get_sky_color(reflect(eye, n));
    let refracted = sea_base + diffuse_factor(n, l, 80.0) * sea_water_color * 0.12;
    var color = mix(refracted, reflected, fresnel);

    let atten = max(1.0 - dot(dist, dist) * 0.001, 0.0);
    color = color + sea_water_color * (p.y - water_params.base_height - water_params.wave_height) * 0.18 * atten;
    color = color + vec3<f32>(specular_factor(n, l, eye, 60.0)) * water_params.specular_strength;

    return color;
}

@vertex
fn vertex_main(@builtin(vertex_index) vertex_index: u32) -> VertexOutput {
    var out: VertexOutput;
    var positions = array<vec2<f32>, 3>(
        vec2<f32>(-1.0, -1.0),
        vec2<f32>(3.0, -1.0),
        vec2<f32>(-1.0, 3.0)
    );
    let pos = positions[vertex_index];
    out.position = vec4<f32>(pos, 0.0, 1.0);
    out.uv = pos * 0.5 + 0.5;
    return out;
}

@fragment
fn fragment_main(in: VertexOutput) -> FragmentOutput {
    var out: FragmentOutput;

    let screen_uv = in.uv;
    var ndc = screen_uv * 2.0 - 1.0;

    let clip_far = vec4<f32>(ndc.x, ndc.y, 0.0, 1.0);
    let world_pos_h = uniforms.inv_view_projection * clip_far;

    let ori = uniforms.camera_position.xyz;
    var dir: vec3<f32>;
    if (abs(world_pos_h.w) < 1e-5) {
        dir = normalize(world_pos_h.xyz);
    } else {
        let world_pos = world_pos_h.xyz / world_pos_h.w;
        dir = normalize(world_pos - ori);
    }

    if (dir.y >= 0.0) {
        discard;
    }

    let p = heightmap_tracing(ori, dir);
    let point_xz = vec2<f32>(p.x, p.z);

    var edge_alpha = 1.0;

    if (water_polygon.use_bounds == 1u) {
        if (!point_in_polygon(point_xz)) {
            discard;
        }

        let edge_dist = distance_to_polygon_edge(point_xz);
        let feather_dist = water_params.edge_feather_distance;
        if (feather_dist > 0.0) {
            edge_alpha = smoothstep(0.0, feather_dist, edge_dist);
        }
    }

    let dist = p - ori;
    let distance_to_water = length(dist);
    let n = get_normal(p, dot(dist, dist) * (0.1 / uniforms.resolution.x));
    let light = normalize(vec3<f32>(0.0, 1.0, 0.8));

    let sea_color = get_sea_color(p, n, light, dir, dist);
    let sky_color = get_sky_color(dir);

    let horizon_blend = smoothstep(-0.02, 0.0, dir.y);
    let distance_fog = smoothstep(uniforms.z_far * 0.4, uniforms.z_far * 0.9, distance_to_water);
    let fog_blend = max(horizon_blend, distance_fog);

    if (fog_blend > 0.995) {
        discard;
    }

    let color = mix(sea_color, sky_color, fog_blend);

    let water_clip = uniforms.projection * uniforms.view * vec4<f32>(p, 1.0);
    let water_ndc_depth = water_clip.z / water_clip.w;

    let final_alpha = edge_alpha * (1.0 - fog_blend);
    if (final_alpha < 0.01) {
        discard;
    }

    out.color = vec4<f32>(pow(color, vec3<f32>(0.65)), final_alpha);
    out.depth = water_ndc_depth;
    return out;
}