Files
mandelbrot/src/shaders/colorize.wgsl
T

198 lines
7.7 KiB
WebGPU Shading Language

// Colourise pass: map the iteration pass's per-pixel escape data (from
// `mandelbrot.wgsl`'s `fs_data`) through the palette. This is the only
// color-dependent step, so changing the palette / colour scale / offset (e.g.
// colour cycling) re-runs just this cheap pass — the expensive perturbation
// iteration in the data texture is reused untouched.
//
// The data texture holds, per texel: R = ci (palette parameter), G = DE
// darkening factor, B = interior fraction (for boundary anti-aliasing). It is
// the same resolution as this pass's target, so we read it with `textureLoad`
// at the fragment's integer pixel coordinate (nearest — iteration data must not
// be linearly filtered across escape boundaries).
@group(0) @binding(0) var<uniform> u: Uniforms;
@group(0) @binding(1) var data_tex: texture_2d<f32>;
@group(0) @binding(2) var<uniform> lights: array<Light, 16>;
@vertex
fn vs_main(@builtin(vertex_index) idx: u32) -> @builtin(position) vec4<f32> {
return vec4<f32>(fullscreen_triangle_pos(idx), 0.0, 1.0);
}
fn shadow_fragment(pos: vec2<f32>) -> vec4<f32> {
let x = i32(pos.x);
let y = i32(pos.y);
let size = textureDimensions(data_tex);
let here = textureLoad(data_tex, vec2<i32>(x, y), 0);
if here.b != 0. {
return vec4<f32>(shadow_interior_color(), 1.0);
}
// Forward differences, except on the last column/row where x+1 / y+1
// is off the texture: fall back to a backward difference, mirrored
// (h0 + (h0 - h[-1])) so the slope keeps the sign normal_from_heights
// expects — plugging h[-1] in directly would flip the normal there.
let h0 = here.g;
var h1: f32;
if x + 1 < i32(size.x) {
h1 = textureLoad(data_tex, vec2<i32>(x + 1, y), 0).g;
} else {
h1 = 2.0 * h0 - textureLoad(data_tex, vec2<i32>(x - 1, y), 0).g;
}
var h2: f32;
if y + 1 < i32(size.y) {
h2 = textureLoad(data_tex, vec2<i32>(x, y + 1), 0).g;
} else {
h2 = 2.0 * h0 - textureLoad(data_tex, vec2<i32>(x, y - 1), 0).g;
}
let normal = normal_from_heights(h0, h1, h2);
return vec4<f32>(shadow_color(normal, here.r), 1.0);
}
@fragment
fn fs_main(@builtin(position) pos: vec4<f32>) -> @location(0) vec4<f32> {
if u.shadow == 2u {
return ray_marching(pos);
} else if u.shadow == 1u {
return shadow_fragment(pos.xy);
} else {
let d = textureLoad(data_tex, vec2<i32>(i32(pos.x), i32(pos.y)), 0);
let ci = d.r;
let de = d.g;
let interior_frac = d.b;
var col = classic_color(ci, de);
// Anti-alias the set boundary: fade toward black by the fraction of the
// pixel's sub-samples that landed in the interior.
col = col * (1.0 - interior_frac);
return vec4<f32>(col, 1.0);
}
}
// Colour of rays that miss the fractal's footprint.
const RAY_MISS: vec4<f32> = vec4<f32>(1.0, 0.0, 0.0, 1.0);
// Per-frame constants of the raymarch, computed once per pixel in
// `ray_marching` rather than on each of the up-to-100 `sdf` steps.
struct MarchConsts {
size: vec2<f32>,
// (size.x / aspect_ratio, size.y): world xy -> texel scale.
to_texel: vec2<f32>,
size_i: vec2<i32>,
inv_size_y: f32,
};
fn sdf(pos: vec3<f32>, k: MarchConsts) -> f32 {
let texture_pos_f32 = pos.xy * k.to_texel;
let texture_pos = clamp(vec2<i32>(texture_pos_f32), vec2<i32>(0, 0), k.size_i - vec2<i32>(1, 1));
let to_texture = max(-min(texture_pos_f32, vec2(0.)), max(texture_pos_f32 - k.size, vec2(0.)));
let dist_to_texture = length(to_texture) * k.inv_size_y;
let px = textureLoad(data_tex, texture_pos, 0);
let de = (px.g * k.inv_size_y) * 0.5;
// Height is measured toward -z, the side the camera sits on (it looks
// along +z), so the terrain is solid on +z: interior plateau at z = 0,
// exterior sloping away from the camera as `de` grows.
let signed_z = -pos.z;
let z = max(signed_z, 0.);
var d: f32;
if px.b != 0. {
d = z;
} else {
d = min(sqrt(z * z + de * de), signed_z + 1. - exp(-de * 5.));
}
// Outside the texture footprint, `d` is the distance from the clamped
// point q on the footprint's edge. The terrain lies over the (convex)
// footprint, so |p - x|² ≥ |q - x|² + |p - q|² for every terrain point x:
// combine in quadrature (not by adding, which overshoots). p can't be in
// the solid out here, so a negative `d` counts as 0.
if dist_to_texture > 0. {
let d_pos = max(d, 0.);
return sqrt(d_pos * d_pos + dist_to_texture * dist_to_texture);
}
return d;
}
fn ray_marching(pos: vec4<f32>) -> vec4<f32> {
let size_i = vec2<i32>(textureDimensions(data_tex));
let size = vec2<f32>(size_i);
let aspect_ratio = u.screen_dim.x / u.screen_dim.y;
let k = MarchConsts(size, vec2<f32>(size.x / aspect_ratio, size.y), size_i, 1.0 / size.y);
let in_texture = vec2<f32>(
(pos.x / size.x) * 2. - 1.,
(pos.y / size.y) * 2. - 1.,
);
var world_pos = u.camera_inv_proj * vec4<f32>(in_texture, 0., 1.0);
let ray_origin = world_pos.xyz;
let ray_dir = u.camera_direction;
// Start where the ray crosses z = 0, the topmost possible surface (the
// camera pitch is clamped short of ±90°, so ray_dir.z > 0).
let start = ray_origin - ray_dir * (ray_origin.z / ray_dir.z);
// The terrain only exists over the footprint x in [0, aspect],
// y in [0, 1]: clip the ray's xy to it up front, so rays that miss it cost
// nothing and the rest start marching at its edge. A huge finite 1/d on
// an axis the ray doesn't move along (top-down, during the 2D <-> 3D
// transition) keeps the slab maths finite.
let inv = select(1.0 / ray_dir.xy, vec2<f32>(1e30), abs(ray_dir.xy) < vec2<f32>(1e-20));
let ta = -start.xy * inv;
let tb = (vec2<f32>(aspect_ratio, 1.0) - start.xy) * inv;
let t_leave = min(max(ta.x, tb.x), max(ta.y, tb.y));
var t = max(max(min(ta.x, tb.x), min(ta.y, tb.y)), 0.0);
if t >= t_leave {
return RAY_MISS;
}
// About 1/50 of a texel at typical sizes: tighter only adds steps
// without visibly moving the hit.
let dist_threshold = 0.00001;
// Rays grazing the exponential slope see a tiny `dist` for many steps in
// a row and would crawl along it until the step budget runs out. Force a
// step of at least half a texel (the height field is nearest-sampled, so
// nothing finer exists), and bisect back if that lands inside the solid.
let min_step = 0.5 * k.inv_size_y;
var hit = false;
var t_prev = t;
for (var i = 0u; i < 100u; i++) {
let dist = sdf(start + t * ray_dir, k);
if dist < dist_threshold {
hit = true;
if dist < 0. {
// Overshot: t_prev is outside, t inside. Refine the crossing.
var lo = t_prev;
var hi = t;
for (var j = 0u; j < 8u; j++) {
let mid = 0.5 * (lo + hi);
if sdf(start + mid * ray_dir, k) < dist_threshold {
hi = mid;
} else {
lo = mid;
}
}
t = hi;
}
break;
}
t_prev = t;
t += max(dist, min_step);
// Past the footprint's far edge: nothing left to hit.
if t >= t_leave {
break;
}
}
// Out of steps while still over the footprint: the ray is skimming the
// surface, so shade where it got to rather than reporting a miss.
if !hit && t < t_leave {
hit = true;
}
// Shade outside the loop, so its registers don't weigh on the march.
if !hit {
return RAY_MISS;
}
return shadow_fragment((start.xy + t * ray_dir.xy) * k.to_texel);
}