rust-sponges/shaders/ray_marching.wgsl

struct FractalParameters
{
    translation: vec3<f32>,
    rotation: vec3<f32>,
    scale: f32,
    iterations: u32
}

struct RayMarchingPushConstants
{
    inverse_proj: mat4x4<f32>,
    view_matrix: mat4x4<f32>,
    camera_pos: vec3<f32>,
    fractal_parameters: FractalParameters
}


var<push_constant> constants: RayMarchingPushConstants; 

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

@vertex
fn vertex(@builtin(vertex_index) i: u32) -> Vertex
{
    let vertices = array<vec4<f32>, 4>(
        vec4<f32>(-1., -1., 0., 1.),
        vec4<f32>(1., -1., 0., 1.),
        vec4<f32>(-1., 1., 0., 1.),
        vec4<f32>(1., 1., 0., 1.),
    );

    let uvs = array<vec2<f32>, 4>(
        vec2<f32>(0., 0.),
        vec2<f32>(1., 0.),
        vec2<f32>(0., 1.),
        vec2<f32>(1., 1.),
    );

    let indices = array<u32, 6>(
        0, 1, 2,
        2, 1, 3
    );

    var v: Vertex;
    v.pos = vertices[indices[i]];
    v.uv = uvs[indices[i]];
    return v;
}

@fragment
fn fragment(in: Vertex) -> @location(0) vec4<f32>
{
    // Produce input ray
    let ndc_ray = vec4<f32>(in.uv.xy * 2. - vec2<f32>(1.), -1.0, 1.0);
    var unproj =constants.inverse_proj * ndc_ray; 
    unproj.w = 0.;


    // Reproject into frustum
    var ray = (constants.view_matrix * unproj).xyz;
    ray /= length(ray);

    //return vec4<f32>(ray, 1.);

    // Ray march
    var t = 0.;
    var pos = constants.camera_pos;

    loop
    {
        // Sample sdf
        let sdf = sdf(pos);

        if sdf < 0.00001
        {
            let grad = sdf_gradient(pos);
            return vec4<f32>((dot(grad, normalize(vec3<f32>(1.))) + 1.) * 0.5);
        }
        if sdf > 1000
        {
            break;
        }

        t += sdf;
        pos += ray * sdf;
    }

    // Ray escaped, get skybox
    return skybox(ray);
}

fn sdf(pos: vec3<f32>) -> f32
{
    var x = pos;
    var t = 1.;

    for(var i = 0u; i < constants.fractal_parameters.iterations; i++)
    {
        x *= constants.fractal_parameters.scale;
        t *= constants.fractal_parameters.scale;
        x = abs(x);
        x += constants.fractal_parameters.translation;
        x *= rot();
    }

    return sdf_octahedron(x, 1.) / t;
}

fn rot() -> mat3x3<f32>
{
    let rx = constants.fractal_parameters.rotation.x;
    let ry = constants.fractal_parameters.rotation.y;
    let rz = constants.fractal_parameters.rotation.z;
    return mat3x3<f32>(
        cos(rx), sin(rx), 0.,
        -sin(rx), cos(rx), 0.,
        0., 0., 1.
    )
    *
    mat3x3<f32>(
         cos(ry), 0., sin(ry),
        0., 1., 0.,
         -sin(ry), 0., cos(ry)
    )
*
    mat3x3<f32>(
        1., 0., 0.,
        0., cos(rz), sin(rz),
        0, -sin(rz), cos(rz)
    );

}

fn sdf_box(pos: vec3<f32>, b: vec3<f32>) -> f32
{
    let q = abs(pos) - b;
    return length(max(q, vec3<f32>(0.))) + min(max(q.x, max(q.y, q.z)), 0.);
}

fn sdf_octahedron(pos: vec3<f32>, s: f32) -> f32
{
    let p = abs(pos);
    return (p.x + p.y + p.z - s) *0.57735;
}

fn sdf_gradient(p: vec3<f32>) -> vec3<f32>
{
    let eps = 0.0001;
    let h = vec2<f32>(eps, 0);

    return normalize(
        vec3<f32>(
            sdf(p + h.xyy) - sdf(p-h.xyy),
            sdf(p + h.yxy) - sdf(p-h.yxy),
            sdf(p + h.yyx) - sdf(p-h.yyx)
        )
    );
}


fn skybox(dir: vec3<f32>) -> vec4<f32>
{
    let sun_dir = normalize(vec3<f32>(1., 1., 1.));

    let gnd_under = vec4<f32>(0.423, 0.450, 0.448, 1.0);
    let gnd_top = vec4<f32>(0.323, 0.350, 0.348, 1.0);
    let gnd = interpolate(gnd_top, gnd_under, map(dir.y, -0.7, 0, 0, 1));
    
    let b = vec4<f32>(0.545, 0.874, 0.940, 1.0);
    let top: vec4<f32> = vec4<f32>(0.0891, 0.464, 0.990, 1.0);
    let sky: vec4<f32> = interpolate(b, top, dir.y);


    let height = map(dir.y, -0.01, 0.01, 0.0, 1.0);
    var res = interpolate(gnd, sky, height);

    var dt = map( dot(dir, normalize(sun_dir)), 0.999, 1, 0, 1);
    if(dt < 0) {dt = 0;}
    res += dt * vec4(0.990, 0.973, 0.782, 1.0);

    return res;
}


fn interpolate(a: vec4<f32>, b: vec4<f32>, x: f32) -> vec4<f32>
{
    var t = x;
    if(t > 1.) {t = 1.;};
    if(t < 0.) {t = 0.;};
    let at = 1. - t;
    return vec4<f32>
    (
        a.x * at + b.x * t,
        a.y * at + b.y * t,
        a.z * at + b.z * t,
        a.w * at + b.w * t
    );
}

fn map(x: f32, xmin: f32, xmax: f32, ymin: f32, ymax: f32) -> f32
{
    return ((x - xmin) / (xmax - xmin))  * (ymax - ymin) + ymin;
}