Terrain shader

terrain.wgsl

struct Uniforms {
    offset: vec2<f32>,
    scale: vec2<f32>,
    tile_size: vec2<u32>,
    time_ms: u32,
}

@group(0) @binding(2)
var<uniform> u: Uniforms;


fn height2depth(height: u32) -> f32 {
    let ieeeOne: u32 = 0x3F800000u; // 1.0 in IEEE binary32
    return bitcast<f32>(height + 1u + ieeeOne) - 1.0;
}

fn top2depth(layer: u32) -> f32 {
    let ieeeOne: u32 = 0x3F800000u; // 1.0 in IEEE binary32
    return bitcast<f32>(0x7FFFFFu - layer + ieeeOne) - 1.0;
}

fn base_depth() -> f32 {
    return height2depth(0u);
}

//#include "depth.wgsl"


// Gold Noise ©2015 [email protected]
// - based on the Golden Ratio
// - uniform normalized distribution
// - fastest static noise generator function (also runs at low precision)
// - use with indicated fractional seeding method. 

let ieeeMantissa: u32 = 0x007FFFFFu; // binary32 mantissa bitmask
let ieeeOne: u32      = 0x3F800000u; // 1.0 in IEEE binary32
let PHI: f32 = 1.61803398874989484820459;  // Φ = Golden Ratio

struct Rng {
    seed: vec3<u32>,
}

fn u32_to_normalized_f32(v: u32) -> f32 {
    return bitcast<f32>(v & ieeeMantissa | ieeeOne) - 1.0;
}

fn gold_noise(xy: vec2<f32>, seed: f32) -> f32 {
    return fract(tan(distance(xy*PHI, xy)*seed)*xy.x);
}

fn xy2seed(v: vec2<u32>) -> Rng {
    var rng = Rng();
    rng.seed = vec3<u32>(v, 1u);
    return rng;
}

fn xy_time2seed(v: vec2<u32>, t: u32) -> Rng {
    var rng = Rng();
    rng.seed = vec3<u32>(v, t);
    return rng;
}

fn next_random_f32(rng: ptr<function, Rng>) -> f32 {
    (*rng).seed.z += 23u;
    (*rng).seed.z *= 2654435769u;
    let xy = vec2<f32>(f32((*rng).seed.x & 0x3FFu) + 17.0, f32((*rng).seed.y & 0x3FFu) + 17.0);
    let s = f32((*rng).seed.z & 0x1Fu) + 11.0;
    let g = tan(distance(xy * PHI, xy) * s) * xy.x;
    return fract(g); //  + f32((*rng).seed.x >> 13u + (*rng).seed.x >> 13u + (*rng).seed.z >> 7u)
}

fn next_random_u32(rng: ptr<function, Rng>) -> u32 {
    let g = next_random_f32(rng);
    let v = reverseBits(bitcast<u32>(g + 1.0) & ieeeMantissa);
    return v;
}

fn uniform01(rng: ptr<function, Rng>) -> f32 {
    return next_random_f32(rng);
}

fn bernoulli(rng: ptr<function, Rng>, p: f32) -> bool {
    return uniform01(rng) < p;
}

fn choose(rng: ptr<function, Rng>, len: u32) -> u32 {
    return u32(next_random_f32(rng) * f32(len));
}

//fn choose(rng: ptr<function, Rng>, len: u32) -> u32 {
//    return next_random_u32(rng) % len;
//}

//#include "rand.wgsl"

struct Pixel {
    color: vec4<f32>,
    normal: vec3<f32>,
    height: u32,
    blend: f32,
}

//#include "pixel.wgsl"

let PI: f32 = 3.14159265358979323846264338327950288;
let TAU: f32 = 6.28318530717958647692528676655900577;

//#include "consts.wgsl"

fn wave(rng: ptr<function, Rng>, time_ms: u32) -> bool {
    return sin((uniform01(rng) + f32(time_ms) / 20000.0) * TAU) <= -0.99;
}

//#include "wave.wgsl"

fn grass_pixel(uv: vec2<u32>, time_ms: u32) -> Pixel {
    var rng0 = xy2seed(uv);

    let shift = wave(&rng0, time_ms);

    var rng1 = xy2seed(vec2<u32>(uv.x + u32(shift), uv.y - 1u));
    var rng2 = xy2seed(vec2<u32>(uv.x + u32(!shift), uv.y - 1u));

    var rng3 = xy2seed(vec2<u32>(uv.x, uv.y - 2u));

    var pixel: Pixel;

    let h0 = choose(&rng0, 32u) + 1u;
    let h1 = choose(&rng1, 32u) + 1u;
    let h2 = choose(&rng2, 32u) + 1u;
    let h3 = choose(&rng3, 32u) + 1u;

    let b = choose(&rng0, 32u) + 1u;
    pixel.blend = f32(b >> 4u);

    let grass_color = vec4<f32>(52.0, 140.0, 49.0, 255.0) / 255.0;
    let dark_grass_color = vec4<f32>(31.0, 100.0, 32.0, 255.0) / 255.0;

    if (h3 >> 4u) > 0u {
        if bernoulli(&rng3, 0.2) {
            pixel.color = grass_color;
        } else {
            pixel.color = dark_grass_color;
        }

        pixel.height = 2u;
    } else if (h2 >> 3u) > 0u {
        if bernoulli(&rng2, 0.2) {
            pixel.color = grass_color;
        } else {
            pixel.color = dark_grass_color;
        }

        pixel.height = 2u;
        pixel.blend += 1.0;
    } else if (h1 >> 2u) > 0u {
        if bernoulli(&rng1, 0.2) {
            pixel.color = grass_color;
        } else {
            pixel.color = dark_grass_color;
        }

        pixel.height = 1u;
    } else {
        if bernoulli(&rng0, 0.2) {
            pixel.color = grass_color;
        } else {
            pixel.color = dark_grass_color;
        }

        pixel.height = 0u;
    }

    return pixel;
}

//#include "grass.wgsl"

//#include "rand.wgsl"
//#include "pixel.wgsl"

fn dirt_pixel(uv: vec2<u32>, time_ms: u32) -> Pixel {
    var rng0 = xy2seed(uv);
    var rng1 = xy2seed(uv / 2u);
    var rng2 = xy2seed(uv / 3u);

    var pixel: Pixel;

    pixel.height = 0;
    pixel.blend = 1.0;

    if bernoulli(&rng0, 0.99) && bernoulli(&rng1, 0.995) && bernoulli(&rng2, 0.995) {
        pixel.color = vec4<f32>(0.35, 0.3, 0.22, 1.0);
    } else {
        pixel.color = vec4<f32>(0.3, 0.22, 0.28, 1.0);
    }

    return pixel;
}

//#include "dirt.wgsl"
//#include "rand.wgsl"
//#include "pixel.wgsl"
//#include "wave.wgsl"

fn water_pixel(uv: vec2<u32>, time_ms: u32) -> Pixel {
    var rng_tile = xy2seed(uv / u.tile_size);
    var rng = xy2seed(vec2<u32>(uv.x / 2u, uv.y));

    var pixel: Pixel;

    pixel.height = choose(&rng_tile, 10u);
    pixel.blend = 0.0;

    if wave(&rng, time_ms) {
        pixel.height = 4u;
        pixel.color = vec4<f32>(105.0, 251.0, 242.0, 200.0) / 255.0;
    } else {
        pixel.color = vec4<f32>(16.0, 204.0, 255.0, 100.0) / 255.0;
    }

    return pixel;
}

//#include "water.wgsl"



struct VertexInput {
    @location(0) pos: vec2<f32>,

    // Index of the terrain types
    // For each vertex, 4 terrain types of the quad is specified
    @location(1) terrains: vec4<u32>,
}

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


@vertex
fn vs_main(
    in: VertexInput,
) -> VertexOutput {
    let pos = u.scale * in.pos - u.offset;

    var out: VertexOutput;
    out.pos = vec4<f32>(pos, base_depth(), 1.0);
    out.uv = in.pos;
    out.terrains = in.terrains;
    return out;
}

struct FragInput {
    @builtin(position) pos: vec4<f32>,
    @location(0) @interpolate(flat) terrains: vec4<u32>,
    @location(1) uv: vec2<f32>,
}

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

fn terrain_pixel(idx: u32, xy: vec2<u32>) -> Pixel {
    var pixel: Pixel;
    switch idx {
        case 1u: {
            pixel = grass_pixel(xy, u.time_ms);
        }
        case 2u: {
            pixel = dirt_pixel(xy, u.time_ms);
        }
        case 3u: {
            pixel = water_pixel(xy, u.time_ms);
        }
        default: {
            pixel.color = vec4<f32>(0.7, 0.5, 0.6, 1.0);
            pixel.height = 0u;
        }
    }
    return pixel;
}

fn uv2xy(uv: vec2<f32>) -> vec2<u32> {
    let uv = vec2<i32>(floor(uv));
    return vec2<u32>(uv + 2147483647);
}

@fragment
fn fs_main(in: FragInput) -> FragOutput {
    let tile_size = vec2<f32>(u.tile_size);

    let i = floor(fract(in.uv) * tile_size) / tile_size;
    var weights = pow(vec4<f32>((1.0 - i.x) * (1.0 - i.y), i.x * (1.0 - i.y), (1.0 - i.x) * i.y, i.x * i.y), vec4<f32>(2.0, 2.0, 2.0, 2.0));

    let xy_pixel = uv2xy(in.uv * tile_size);
    let xy_tile = uv2xy(in.uv);

    let pixel_x = terrain_pixel(in.terrains.x, xy_pixel);//, terrain_payloads.x);
    let pixel_y = terrain_pixel(in.terrains.y, xy_pixel);//, terrain_payloads.y);
    let pixel_z = terrain_pixel(in.terrains.z, xy_pixel);//, terrain_payloads.z);
    let pixel_w = terrain_pixel(in.terrains.w, xy_pixel);//, terrain_payloads.w);

    var weights2 = weights;

    if in.terrains.x == in.terrains.y {
        weights2.x += weights.y;
        weights2.y += weights.x;
    }

    if in.terrains.z == in.terrains.w {
        weights2.z += weights.w;
        weights2.w += weights.z;
    }

    if in.terrains.x == in.terrains.z {
        weights2.x += weights.z;
        weights2.z += weights.x;
    }

    if in.terrains.y == in.terrains.w {
        weights2.y += weights.w;
        weights2.w += weights.y;
    }

    var rng = xy2seed(xy_tile);

    if bernoulli(&rng, 0.5) {
        if in.terrains.x == in.terrains.w {
            weights2.x += weights.w;
            weights2.w += weights.x;
        } else if in.terrains.y == in.terrains.z {
            weights2.y += weights.z;
            weights2.z += weights.y;
        }
    } else {
        if in.terrains.y == in.terrains.z {
            weights2.y += weights.z;
            weights2.z += weights.y;
        } else if in.terrains.x == in.terrains.w {
            weights2.x += weights.w;
            weights2.w += weights.x;
        }
    }

    let c = pow(i.x * i.y * (1.0 - i.x) * (1.0 - i.y), 0.2);

    let xh = weights2.x * (1.0 + c * uniform01(&rng)) + pixel_x.blend / tile_size.x;
    let yh = weights2.y * (1.0 + c * uniform01(&rng)) + pixel_y.blend / tile_size.x;
    let zh = weights2.z * (1.0 + c * uniform01(&rng)) + pixel_z.blend / tile_size.x;
    let wh = weights2.w * (1.0 + c * uniform01(&rng)) + pixel_w.blend / tile_size.x;

    let m = max(xh, max(yh, max(zh, wh)));

    var pixel: Pixel;
    if m <= xh {
        pixel = pixel_x;
    }
    if m <= yh {
        pixel = pixel_y;
    }
    if m <= zh {
        pixel = pixel_z;
    }
    if m <= wh {
        pixel = pixel_w;
    }

    var out: FragOutput;
    out.color = pixel.color;
    out.depth = height2depth(pixel.height);
    if distance(in.uv, round(in.uv)) < 0.05 {
        out.color = vec4<f32>(0.0, 0.0, 1.0, 1.0);
    }

    return out;
}