Fluid

// fluid · advect — semi-Lagrangian advection of one field through the
// velocity field. Serves two pass instances per tick: velocity advecting
// itself (srcTex = velocityTex = the velocity field) and dye advecting
// through velocity, each with its own uniforms (JS uniform writes are
// frame-global, so per-pass values need per-pass effect instances).

import { advectBacktrace, bilerp2, bilerp4 } from "@vshaders/fluid/advect";

struct Uniforms {
  // Resolution of this pass's render target (the field being advected).
  // Declared by every pass per the standard interface; the backtrace itself
  // works through the velocity grid below.
  resolution: vec2f,
  // Resolution of the velocity grid. Velocity is stored in *its* texels per
  // unit of dt, so displacement in uv converts through the velocity texel —
  // the dye field is finer than the sim grid and must not rescale the flow.
  velocityResolution: vec2f,
  // Timestep in seconds, clamped by the runner (tab restores can't explode).
  dt: f32,
  // Retention factor per 60Hz frame, applied frame-rate independently as
  // pow(dissipation, dt * 60): 1 keeps everything, lower values fade.
  dissipation: f32,
}

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var srcTex: texture_2d<f32>;
@group(0) @binding(2) var velocityTex: texture_2d<f32>;
@group(0) @binding(3) var linearSampler: sampler;

@fragment
fn fs_main(@location(0) uv: vec2f) -> @location(0) vec4f {
  let velTexel = 1.0 / max(uniforms.velocityResolution, vec2f(1.0));
  let velocity = bilerp2(velocityTex, linearSampler, uv);
  let back = advectBacktrace(uv, velocity, max(uniforms.dt, 0.0), velTexel);
  let keep = pow(clamp(uniforms.dissipation, 0.0, 1.0), max(uniforms.dt, 0.0) * 60.0);
  return bilerp4(srcTex, linearSampler, back) * keep;
}

// fluid · splat — injects force or dye along the stroke segment from
// prevCenter to center: the Gaussian falloff is evaluated at the closest
// point on the segment, so fast pointer motion paints a continuous capsule
// instead of a trail of beads. One shader serves every splat instance per
// tick (pointer velocity, pointer dye, idle emitters), each instance with
// its own uniforms: colorMix 0 adds splatValue as-is (velocity, xy into an
// rg target), colorMix 1 adds a cosine-palette color scaled by splatValue.x
// (dye). Works in linear light; the dye field stays linear until display.

import { gaussianSplat } from "@vshaders/fluid/splat";
import { cosinePalette } from "@vshaders/color/palette";

struct Uniforms {
  // Resolution of this pass's render target (the field being splatted).
  resolution: vec2f,
  // Seconds; drives the palette index of dye splats.
  time: f32,
  // Gaussian 1/e radius, in short-side uv units (aspect-corrected below).
  radius: f32,
  // Stroke endpoints in uv of the field; equal for a stationary splat.
  center: vec2f,
  prevCenter: vec2f,
  // Velocity instances: xy = force in velocity-grid texels per unit dt.
  // Dye instances: x = dye amount multiplying the palette color.
  splatValue: vec3f,
  // 0 = raw splatValue (velocity), 1 = palette-colored dye.
  colorMix: f32,
  // Cosine palette shape (linear light) and its time scale/offset.
  paletteBias: vec3f,
  paletteSpeed: f32,
  paletteAmplitude: vec3f,
  paletteShift: f32,
}

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var srcTex: texture_2d<f32>;

@fragment
fn fs_main(@builtin(position) position: vec4f, @location(0) uv: vec2f) -> @location(0) vec4f {
  let res = max(uniforms.resolution, vec2f(1.0));
  // Aspect correction: distances in short-side units on both axes, so the
  // splat is round and `radius` means the same thing at any aspect ratio.
  let aspect = res / min(res.x, res.y);
  let p = uv * aspect;
  let a = uniforms.prevCenter * aspect;
  let b = uniforms.center * aspect;

  // Closest point on the stroke segment a→b; a zero-length segment (guarded
  // against dividing by zero) degrades to the point splat at `center`.
  let ab = b - a;
  let lenSq = dot(ab, ab);
  var nearest = b;
  if (lenSq > 1e-10) {
    let t = clamp(dot(p - a, ab) / lenSq, 0.0, 1.0);
    nearest = a + ab * t;
  }
  // gaussianSplat guards non-positive radius by depositing nothing.
  let falloff = gaussianSplat(p, nearest, uniforms.radius);

  let paletteIndex = uniforms.time * uniforms.paletteSpeed + uniforms.paletteShift;
  let paletteColor = max(
    cosinePalette(
      paletteIndex,
      uniforms.paletteBias,
      uniforms.paletteAmplitude,
      vec3f(1.0),
      vec3f(0.0, 0.33, 0.67),
    ),
    vec3f(0.0),
  );
  let amount = mix(uniforms.splatValue, paletteColor * uniforms.splatValue.x, clamp(uniforms.colorMix, 0.0, 1.0));

  let last = max(vec2i(textureDimensions(srcTex)) - vec2i(1), vec2i(0));
  let base = textureLoad(srcTex, clamp(vec2i(position.xy), vec2i(0), last), 0);
  return vec4f(base.rgb + falloff * amount, 1.0);
}

// fluid · curl — scalar (out-of-plane) curl of the velocity field, written
// once per tick into an r16float target for the vorticity-confinement pass.
// The point curl is tent-filtered over the 3x3 neighborhood before it is
// stored: the confinement force normalizes the gradient of |curl|, and raw
// per-texel curl makes that direction flip texel to texel, injecting
// texel-scale divergence the collocated projection mathematically cannot
// remove (its wide Laplacian is blind at the Nyquist frequency). Filtered
// curl keeps the vortex structure and moves the force onto scales the
// Jacobi solve actually reaches.
// The kernel reads textureLoad neighborhoods clamped to the extent, so no
// sampler is needed and the boundary behaves as clamp-to-edge.

import { curlAt } from "@vshaders/fluid/vorticity";

struct Uniforms {
  // Resolution of this pass's render target (the sim grid). Declared by
  // every pass per the standard interface; the kernel clamps its own reads.
  resolution: vec2f,
}

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var velocityTex: texture_2d<f32>;

@fragment
fn fs_main(@builtin(position) position: vec4f) -> @location(0) vec4f {
  let p = vec2i(position.xy);
  // 3x3 tent filter (4-2-1)/16 over the point curl.
  var filtered = 4.0 * curlAt(velocityTex, p);
  filtered += 2.0 * curlAt(velocityTex, p + vec2i(1, 0));
  filtered += 2.0 * curlAt(velocityTex, p - vec2i(1, 0));
  filtered += 2.0 * curlAt(velocityTex, p + vec2i(0, 1));
  filtered += 2.0 * curlAt(velocityTex, p - vec2i(0, 1));
  filtered += curlAt(velocityTex, p + vec2i(1, 1));
  filtered += curlAt(velocityTex, p + vec2i(1, -1));
  filtered += curlAt(velocityTex, p + vec2i(-1, 1));
  filtered += curlAt(velocityTex, p + vec2i(-1, -1));
  return vec4f(filtered / 16.0, 0.0, 0.0, 1.0);
}

// fluid · vorticity — vorticity confinement (Fedkiw, Stam, Jensen 2001):
// adds a force that steers flow back around the vortices grid advection
// smears out, keeping the motion swirly instead of syrupy. Reads the curl
// field written by curl.wgsl and the current velocity; writes velocity.

import { vorticityForceAt } from "@vshaders/fluid/vorticity";

struct Uniforms {
  // Resolution of this pass's render target (the sim grid). Declared by
  // every pass per the standard interface; the kernel clamps its own reads.
  resolution: vec2f,
  // Timestep in seconds; the confinement force integrates as force * dt.
  dt: f32,
  // Confinement strength; 5–30 is typical on a 128 grid, more is stormier.
  // The kernel returns zero force for non-positive strength.
  curlStrength: f32,
}

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var velocityTex: texture_2d<f32>;
@group(0) @binding(2) var curlTex: texture_2d<f32>;

@fragment
fn fs_main(@builtin(position) position: vec4f) -> @location(0) vec4f {
  let p = vec2i(position.xy);
  let last = max(vec2i(textureDimensions(velocityTex)) - vec2i(1), vec2i(0));
  let velocity = textureLoad(velocityTex, clamp(p, vec2i(0), last), 0).xy;
  let force = vorticityForceAt(curlTex, p, uniforms.curlStrength);
  return vec4f(velocity + force * max(uniforms.dt, 0.0), 0.0, 1.0);
}

// fluid · divergence — central-difference divergence of the velocity field,
// written once per tick into an r16float target. It is the right-hand side
// of the pressure Poisson equation the Jacobi passes relax.
//
// `contained` picks the boundary condition (uniform, so control flow stays
// uniform): 1 treats the domain edge as a reflective container wall (the
// off-grid velocity neighbor is the negation of the center's, so flow into
// a wall registers as compression); 0 keeps the free clamped-edge reads,
// under which fluid drifts out at the edges.

import { divergenceAt, divergenceBoundedAt } from "@vshaders/fluid/project";

struct Uniforms {
  // Resolution of this pass's render target (the sim grid). Declared by
  // every pass per the standard interface; the kernel clamps its own reads.
  resolution: vec2f,
  // 1 = contained (reflective walls), 0 = free. Set by the runner.
  contained: f32,
}

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var velocityTex: texture_2d<f32>;

@fragment
fn fs_main(@builtin(position) position: vec4f) -> @location(0) vec4f {
  let p = vec2i(position.xy);
  let divergence = select(
    divergenceAt(velocityTex, p),
    divergenceBoundedAt(velocityTex, p),
    uniforms.contained > 0.5,
  );
  return vec4f(divergence, 0.0, 0.0, 1.0);
}

// fluid · clear pressure — scales last tick's relaxed pressure down and
// hands it to the Jacobi passes as their starting guess (a warm start:
// pressure fields are strongly correlated frame to frame, so this buys the
// iterations back). A dedicated pass because JS uniform writes are
// frame-global: the decay cannot be a first-iteration uniform on the
// Jacobi effect.

struct Uniforms {
  // Resolution of this pass's render target (the sim grid). Declared by
  // every pass per the standard interface.
  resolution: vec2f,
  // Fraction of the previous pressure kept as the warm start; 0 restarts
  // the relaxation from scratch every tick.
  pressureDecay: f32,
}

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var pressureTex: texture_2d<f32>;

@fragment
fn fs_main(@builtin(position) position: vec4f) -> @location(0) vec4f {
  let last = max(vec2i(textureDimensions(pressureTex)) - vec2i(1), vec2i(0));
  let pressure = textureLoad(pressureTex, clamp(vec2i(position.xy), vec2i(0), last), 0).x;
  return vec4f(pressure * clamp(uniforms.pressureDecay, 0.0, 1.0), 0.0, 0.0, 1.0);
}

// fluid · jacobi pressure — one Jacobi relaxation step of the pressure
// Poisson equation laplacian(p) = divergence. The runner encodes this pass
// N times per tick over a ping-pong r16float pair, rebinding pressureTex to
// the read half before each pass (texture rebinds take effect per pass;
// N is the `iterations` runner param). Jacobi converges fast on the
// high-frequency divergence pointer splats create and slowly on
// domain-scale modes — which is why N is tunable.
//
// `contained` picks the boundary condition (uniform, so control flow stays
// uniform): 1 uses the bounded kernel's explicit pure-Neumann walls (the
// off-grid pressure neighbor equals the center), matching the contained
// divergence and gradient passes; 0 keeps the free clamped-edge reads.

import { jacobiPressureAt, jacobiPressureBoundedAt } from "@vshaders/fluid/project";

struct Uniforms {
  // Resolution of this pass's render target (the sim grid). Declared by
  // every pass per the standard interface; the kernel clamps its own reads.
  resolution: vec2f,
  // 1 = contained (reflective walls), 0 = free. Set by the runner.
  contained: f32,
}

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var pressureTex: texture_2d<f32>;
@group(0) @binding(2) var divergenceTex: texture_2d<f32>;

@fragment
fn fs_main(@builtin(position) position: vec4f) -> @location(0) vec4f {
  let p = vec2i(position.xy);
  let pressure = select(
    jacobiPressureAt(pressureTex, divergenceTex, p),
    jacobiPressureBoundedAt(pressureTex, divergenceTex, p),
    uniforms.contained > 0.5,
  );
  return vec4f(pressure, 0.0, 0.0, 1.0);
}

// fluid · gradient subtract — the projection step that makes the flow
// incompressible: subtracts the gradient of the relaxed pressure from the
// velocity field, leaving it approximately divergence-free. Writes the
// projected velocity back into the velocity ping-pong pair.
//
// `contained` picks the boundary condition (uniform, so control flow stays
// uniform): 1 uses the bounded kernel's pure-Neumann pressure walls, the
// gradient half of the reflective-wall projection (pair with the contained
// divergence and Jacobi passes); 0 keeps the free clamped-edge reads.

import { subtractPressureGradientAt, subtractPressureGradientBoundedAt } from "@vshaders/fluid/project";

struct Uniforms {
  // Resolution of this pass's render target (the sim grid). Declared by
  // every pass per the standard interface; the kernel clamps its own reads.
  resolution: vec2f,
  // 1 = contained (reflective walls), 0 = free. Set by the runner.
  contained: f32,
}

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var velocityTex: texture_2d<f32>;
@group(0) @binding(2) var pressureTex: texture_2d<f32>;

@fragment
fn fs_main(@builtin(position) position: vec4f) -> @location(0) vec4f {
  let p = vec2i(position.xy);
  let projected = select(
    subtractPressureGradientAt(velocityTex, pressureTex, p),
    subtractPressureGradientBoundedAt(velocityTex, pressureTex, p),
    uniforms.contained > 0.5,
  );
  return vec4f(projected, 0.0, 1.0);
}