grid.slang

#include "shared.h"
#include "renderer_types.h"

// https://bgolus.medium.com/the-best-darn-grid-shader-yet-727f9278b9d8
f32 PristineGrid(vec2 uv, vec2 lineWidth) {
    lineWidth = saturate(lineWidth);
    vec4 uvDDXY = vec4(ddx(uv), ddy(uv));
    vec2 uvDeriv = vec2(length(uvDDXY.xz), length(uvDDXY.yw));
    uvDeriv = clamp(uvDeriv, 1e-5, 5.0);
    bvec2 invertLine = lineWidth > 0.5;
    vec2 targetWidth = select(invertLine, 1.0 - lineWidth, lineWidth);
    vec2 drawWidth = clamp(targetWidth, uvDeriv, 0.5);
    vec2 lineAA = max(uvDeriv, 0.000001) * 1.5;
    vec2 gridUV = abs(frac(uv) * 2.0 - 1.0);
    gridUV = select(invertLine, gridUV, 1.0 - gridUV);
    vec2 grid2 = smoothstep(drawWidth + lineAA, drawWidth - lineAA, gridUV);
    grid2 *= saturate(targetWidth / drawWidth);
    grid2 = lerp(grid2, targetWidth, saturate(uvDeriv * 2.0 - 1.0));
    grid2 = select(invertLine, 1.0 - grid2, grid2);
    return lerp(grid2.x, 1.0, grid2.y);
}

f32 ScreenSpaceLine(f32 value, f32 pixelWidth) {
    f32 deriv = fwidth(value);
    f32 width = deriv * pixelWidth;
    return 1.0 - smoothstep(width, width + deriv, abs(value));
}

[vk::push_constant]
GridPushConstants pc;

static vec2 GRID_SCALE = vec2(1.0, 1.0);
static f32 THIN_LINE_WIDTH = 0.1;
static vec4 THIN_COLOR = vec4(0.25, 0.25, 0.25, 1.0f);

static f32 MAJOR_DIVISION = 10.0;
static f32 THICK_LINE_WIDTH = 0.01; // gets lerped with thin so has to be lower than it
static vec4 THICK_COLOR = vec4(0.55, 0.55, 0.55, 1.0f);

static f32 FADE_DISTANCE = 250.f;

static f32 AXIS_WIDTH_PX = 0.01;
static vec4 X_AXIS_COLOR = vec4(1.0, 0.0, 0.0, 1.0);
static vec4 Y_AXIS_COLOR = vec4(0.0, 1.0, 0.0, 1.0);
static vec4 Z_AXIS_COLOR = vec4(0.0, 0.5, 1.0, 1.0);

static const vec2 fullscreen_quad_verts[3] = { vec2(-1.0, -1.0), vec2(3.0, -1.0), vec2(-1.0, 3.0) };

struct VSOutput {
    vec4 position : SV_Position;
    vec3 ray_dir : TEXCOORD0;
    vec3 ray_org : TEXCOORD1;
}

[shader("vertex")]
VSOutput vsMain(u32 vertexID: SV_VertexID) {
    VSOutput o;

    FrameUBO frame = *pc.frame;

    vec2 ndc = fullscreen_quad_verts[vertexID];
    vec4 clip = vec4(ndc, 1.0, 1.0);
    vec4 world = clip * frame.inv_view_proj;

    world.xyz /= world.w;
    o.ray_org = frame.position_ws.xyz;
    o.ray_dir = world.xyz - frame.position_ws.xyz;
    o.position = vec4(ndc, 0.0, 1.0);

    return o;
}

[shader("fragment")]
vec4 fsMain(VSOutput input) : SV_Target {
    FrameUBO frame = *pc.frame;

    vec3 rayOrigin = input.ray_org;
    vec3 rayDir = normalize(input.ray_dir);

    if (abs(rayDir.z) < 1e-4) {
        discard;
    }

    f32 t = -rayOrigin.z / rayDir.z;

    if (t <= 0.0) {
        discard;
    }

    vec3 worldPos = rayOrigin + rayDir * t;
    vec2 cameraGridOrigin = floor(frame.position_ws.xy / MAJOR_DIVISION) * MAJOR_DIVISION;
    vec2 uv = (worldPos.xy - cameraGridOrigin) * GRID_SCALE;

    f32 minor = PristineGrid(uv, vec2(THIN_LINE_WIDTH));
    f32 major = PristineGrid(uv / MAJOR_DIVISION, vec2(THICK_LINE_WIDTH));

    vec4 color = THIN_COLOR * minor;
    color = lerp(color, THICK_COLOR, major);

    // override if  axis
    f32 xAxis = ScreenSpaceLine(worldPos.y, AXIS_WIDTH_PX);
    f32 yAxis = ScreenSpaceLine(worldPos.x, AXIS_WIDTH_PX);
    color = lerp(color, X_AXIS_COLOR, xAxis);
    color = lerp(color, Y_AXIS_COLOR, yAxis);

    f32 dist = distance(frame.position_ws.xy, worldPos.xy);
    f32 fade = 1.0 - saturate(dist / FADE_DISTANCE);
    color.a *= fade;

    return color;
}