WebGPU-FFT-Ocean-Demo/shaders/ocean_render.wgsl

@group(0) @binding(0) var<storage, read> heightField : array<f32>;
@group(0) @binding(1) var<storage, read> renderParams : array<f32>;
@group(0) @binding(2) var<storage, read> viewProjection : array<f32>;
@group(0) @binding(3) var<storage, read> cameraPosition : array<f32>;
@group(0) @binding(4) var bottomColorTex  : texture_2d<f32>;
@group(0) @binding(5) var bottomHeightTex : texture_2d<f32>;
@group(0) @binding(6) var bottomSampler   : sampler;

struct VertexOutput {
	@builtin(position) position : vec4<f32>,
	@location(0) worldPosition  : vec3<f32>,
	@location(1) normal         : vec3<f32>,
	@location(2) heightValue    : f32,
};

fn indexFromCoord( x : u32, y : u32, size : u32 ) -> u32 {

	return y * size + x;

}

fn loadViewProjection( ) -> mat4x4<f32> {

	return mat4x4<f32>(
		vec4<f32>( viewProjection[ 0 ],  viewProjection[ 1 ],  viewProjection[ 2 ],  viewProjection[ 3 ] ),
		vec4<f32>( viewProjection[ 4 ],  viewProjection[ 5 ],  viewProjection[ 6 ],  viewProjection[ 7 ] ),
		vec4<f32>( viewProjection[ 8 ],  viewProjection[ 9 ],  viewProjection[ 10 ], viewProjection[ 11 ] ),
		vec4<f32>( viewProjection[ 12 ], viewProjection[ 13 ], viewProjection[ 14 ], viewProjection[ 15 ] )
	);

}

fn sampleHeight( fx : f32, fy : f32, simSize : u32, heightAmp : f32 ) -> f32 {

	let sx : f32 = clamp( fx, 0.0, f32( simSize - 1u ) );
	let sy : f32 = clamp( fy, 0.0, f32( simSize - 1u ) );

	let ix0 : u32 = u32( sx );
	let iy0 : u32 = u32( sy );

	let ix1 : u32 = min( simSize - 1u, ix0 + 1u );
	let iy1 : u32 = min( simSize - 1u, iy0 + 1u );

	let tx  : f32 = sx - f32( ix0 );
	let ty  : f32 = sy - f32( iy0 );

	let idx00 : u32 = indexFromCoord( ix0, iy0, simSize );
	let idx10 : u32 = indexFromCoord( ix1, iy0, simSize );
	let idx01 : u32 = indexFromCoord( ix0, iy1, simSize );
	let idx11 : u32 = indexFromCoord( ix1, iy1, simSize );

	let parity00 : u32 = ( ix0 + iy0 ) & 1u;
	let parity10 : u32 = ( ix1 + iy0 ) & 1u;
	let parity01 : u32 = ( ix0 + iy1 ) & 1u;
	let parity11 : u32 = ( ix1 + iy1 ) & 1u;

	var sign00 : f32 = 1.0;
	var sign10 : f32 = 1.0;
	var sign01 : f32 = 1.0;
	var sign11 : f32 = 1.0;

	if ( parity00 == 1u ) { sign00 = -1.0; }
	if ( parity10 == 1u ) { sign10 = -1.0; }
	if ( parity01 == 1u ) { sign01 = -1.0; }
	if ( parity11 == 1u ) { sign11 = -1.0; }

	let h00 : f32 = heightField[ idx00 ] * heightAmp * sign00;
	let h10 : f32 = heightField[ idx10 ] * heightAmp * sign10;
	let h01 : f32 = heightField[ idx01 ] * heightAmp * sign01;
	let h11 : f32 = heightField[ idx11 ] * heightAmp * sign11;

	let hx0 : f32 = mix( h00, h10, tx );
	let hx1 : f32 = mix( h01, h11, tx );

	return mix( hx0, hx1, ty );

}


@vertex
fn v_main(
	@location(0) position : vec3<f32>,
	@location(1) uv       : vec2<f32>,
	@builtin(instance_index) instanceIndex : u32
) -> VertexOutput {

	let simSize   : u32 = u32( renderParams[ 0 ] );
	let meshSize  : f32 = renderParams[ 1 ];
	let heightAmp : f32 = renderParams[ 2 ];
	let meshRes   : u32 = u32( renderParams[ 4 ] );

	// Tiling amount (slider 1..N). We clamp to at least 1.
	let tilingF   : f32 = max( 1.0, renderParams[ 5 ] );
	let tiling    : u32 = max( 1u, u32( tilingF + 0.5 ) );

	let tileRange : u32 = tiling - 1u;
	let tilesPerRow : u32 = tileRange * 2u + 1u;

	let tileIndex : u32 = instanceIndex;

	let tileXIndex : u32 = tileIndex % tilesPerRow;
	let tileZIndex : u32 = tileIndex / tilesPerRow;

	let offsetGridX : i32 = i32( tileXIndex ) - i32( tileRange );
	let offsetGridZ : i32 = i32( tileZIndex ) - i32( tileRange );

	let offsetX  : f32 = f32( offsetGridX ) * meshSize;
	let offsetZ  : f32 = f32( offsetGridZ ) * meshSize;

	let denom    : f32 = max( 1.0, f32( meshRes - 1u ) );
	let fxNorm   : f32 = clamp( uv.x / denom, 0.0, 1.0 );
	let fyNorm   : f32 = clamp( uv.y / denom, 0.0, 1.0 );

	let fx       : f32 = fxNorm * f32( simSize - 1u );
	let fy       : f32 = fyNorm * f32( simSize - 1u );

	let hCenter  : f32 = sampleHeight( fx, fy, simSize, heightAmp );

	let worldX   : f32 = ( fxNorm - 0.5 ) * meshSize + offsetX;
	let worldZ   : f32 = ( fyNorm - 0.5 ) * meshSize + offsetZ;

	// Derivatives via central differences on the sampled height field
	let fxStep   : f32 = 1.0;
	let fyStep   : f32 = 1.0;

	let hXp      : f32 = sampleHeight( fx + fxStep, fy, simSize, heightAmp );
	let hXm      : f32 = sampleHeight( fx - fxStep, fy, simSize, heightAmp );
	let hYp      : f32 = sampleHeight( fx, fy + fyStep, simSize, heightAmp );
	let hYm      : f32 = sampleHeight( fx, fy - fyStep, simSize, heightAmp );

	let h        : f32 = hCenter;

	let dx       : f32 = hXp - hXm;
	let dz       : f32 = hYp - hYm;

	let normal : vec3<f32> = normalize( vec3<f32>( -dx, 2.0, -dz ) );

	let vp : mat4x4<f32> = loadViewProjection();

	var output : VertexOutput;
	output.position      = vp * vec4<f32>( worldX, h, worldZ, 1.0 );
	output.worldPosition = vec3<f32>( worldX, h, worldZ );
	output.normal        = normal;
	output.heightValue   = h;

	return output;

}

@fragment
fn f_main( input : VertexOutput ) -> @location(0) vec4<f32> {

	let mode      : u32       = u32( renderParams[ 3 ] );

	let N          : vec3<f32> = normalize( input.normal );

	// Debug: visualize normals
	if ( mode == 1u ) {
		let normalColor : vec3<f32> = 0.5 * ( N + vec3<f32>( 1.0, 1.0, 1.0 ) );
		return vec4<f32>( normalColor, 1.0 );
	}

	// Debug: true solid color (no height variation)
	if ( mode == 2u ) {
		let flatColor : vec3<f32> = vec3<f32>( 0.05, 0.4, 0.8 );
		return vec4<f32>( flatColor, 1.0 );
	}

	// Debug: visualize height field as grayscale
	if ( mode == 3u ) {
		// Scale and bias height into a visible range
		let h     : f32 = input.heightValue;
		let hNorm : f32 = clamp( h * 0.05 + 0.5, 0.0, 1.0 );
		let c     : vec3<f32> = vec3<f32>( hNorm, hNorm, hNorm );
		return vec4<f32>( c, 1.0 );
	}

	// Sun / light setup
	let lightDir   : vec3<f32> = normalize( vec3<f32>( 0.2, 0.85, 0.35 ) );
	let sunColor   : vec3<f32> = vec3<f32>( 1.0, 0.97, 0.9 );

	// Camera + view
	let camPos     : vec3<f32> = vec3<f32>( cameraPosition[ 0 ], cameraPosition[ 1 ], cameraPosition[ 2 ] );
	let viewDir    : vec3<f32> = normalize( camPos - input.worldPosition );

	// Normal and basic terms
	let NdotL      : f32       = max( dot( N, lightDir ), 0.0 );
	let NdotV      : f32       = max( dot( N, viewDir ), 0.0 );

	// Water body base color
	let waterBase  : vec3<f32> = vec3<f32>( 0.01, 0.18, 0.55 );

	// Realistic mode: reuse simple lighting but with sand/depth-based underwater color
	if ( mode == 4u ) {

		let meshSize    : f32       = renderParams[ 1 ];
		let uvBase      : vec2<f32> = input.worldPosition.xz / meshSize + vec2<f32>( 0.5, 0.5 );
		let noise       : vec2<f32> = vec2<f32>(
			sin( input.worldPosition.x * 0.05 + input.worldPosition.z * 0.12 ),
			cos( input.worldPosition.x * 0.04 - input.worldPosition.z * 0.09 )
		) * 0.03;
		let uvMain      : vec2<f32> = uvBase + noise;
		let uvAlt       : vec2<f32> = uvBase + vec2<f32>( 0.25, 0.43 ) + noise * 0.8;
		let colorSample : vec4<f32> = textureSample( bottomColorTex, bottomSampler, uvMain );
		let colorAlt    : vec4<f32> = textureSample( bottomColorTex, bottomSampler, uvAlt );
		let finalColor  : vec3<f32> = mix( colorSample.rgb, colorAlt.rgb, 0.38 );
		let heightSample: vec4<f32> = textureSample( bottomHeightTex, bottomSampler, uvMain );

		let heightAlt   : vec4<f32> = textureSample( bottomHeightTex, bottomSampler, uvAlt );

		let heightVal   : f32       = mix( heightSample.r, heightAlt.r, 0.4 );

		let baseColor   : vec3<f32> = finalColor;

		let poolY       : f32       = -6.0;
		let depthGeom   : f32       = clamp( poolY - input.worldPosition.y, 0.0, 20.0 );

		let shallowT    : f32       = clamp( ( heightVal - 0.4 ) / 0.6, 0.0, 1.0 );

		let depthT      : f32       = clamp( depthGeom / 8.0, 0.0, 1.0 );

		let sandWeight  : f32       = shallowT * ( 1.0 - depthT );
		let waterWeight : f32       = 1.0 - sandWeight;

		let sandColor   : vec3<f32> = baseColor;
		let wetColor    : vec3<f32> = mix( baseColor, waterBase, 0.6 );

		let bottomColor : vec3<f32> = sandColor * sandWeight + wetColor * waterWeight;

		let absorbCoeff : vec3<f32> = vec3<f32>( 0.04, 0.09, 0.16 );
		let absorb      : vec3<f32> = exp( -absorbCoeff * depthGeom );

		let bodyColor   : vec3<f32> = bottomColor * absorb;

		let up          : vec3<f32> = vec3<f32>( 0.0, 1.0, 0.0 );
		let reflDirEnv  : vec3<f32> = normalize( reflect( -viewDir, N ) );
		let horizonT    : f32       = clamp( pow( 1.0 - max( dot( reflDirEnv, up ), 0.0 ), 1.5 ), 0.0, 1.0 );
		let skyZenith   : vec3<f32> = vec3<f32>( 0.02, 0.12, 0.28 );
		let skyHorizon  : vec3<f32> = vec3<f32>( 0.30, 0.45, 0.70 );
		let skyAnalytic : vec3<f32> = mix( skyZenith, skyHorizon, horizonT );
		let skyMask     : vec3<f32> = vec3<f32>( 0.08, 0.18, 0.32 );
		let skyColor    : vec3<f32> = mix( skyAnalytic, skyMask, 0.4 );

		let halfDir    : vec3<f32> = normalize( lightDir + viewDir );
		let NdotH      : f32       = max( dot( N, halfDir ), 0.0 );
		let spec       : f32       = pow( NdotH, 140.0 ) * NdotL;
		let specTerm   : vec3<f32> = spec * vec3<f32>( 1.0, 1.0, 1.0 ) * 1.3;

		let F0         : vec3<f32> = vec3<f32>( 0.02, 0.025, 0.03 );
		let oneMinus   : f32       = 1.0 - NdotV;
		let fresnel    : vec3<f32> = F0 + ( vec3<f32>( 1.0, 1.0, 1.0 ) - F0 ) * pow( oneMinus, 5.0 );

		let diffuse    : vec3<f32> = bodyColor * ( 0.15 + 0.85 * NdotL ) * sunColor;

		let envBlend   : vec3<f32> = mix( skyColor, vec3<f32>( 0.08, 0.25, 0.55 ), 0.35 );
		let reflection : vec3<f32> = envBlend + specTerm;

		let color      : vec3<f32> = diffuse * ( vec3<f32>( 1.0, 1.0, 1.0 ) - fresnel ) + reflection * fresnel;

		let colorGamma : vec3<f32> = pow( clamp( color, vec3<f32>( 0.0, 0.0, 0.0 ), vec3<f32>( 1.0, 1.0, 1.0 ) ), vec3<f32>( 1.0 / 2.2, 1.0 / 2.2, 1.0 / 2.2 ) );

		return vec4<f32>( colorGamma, 1.0 );

	}

	// Simple sky / environment tint (towards horizon)
	let skyZenith  : vec3<f32> = vec3<f32>( 0.02, 0.12, 0.28 );
	let skyHorizon : vec3<f32> = vec3<f32>( 0.30, 0.45, 0.70 );

	// Reflection direction for environment lookup
	let reflDirEnv : vec3<f32> = normalize( reflect( -viewDir, N ) );
	let up         : vec3<f32> = vec3<f32>( 0.0, 1.0, 0.0 );
	let horizonT   : f32       = clamp( pow( 1.0 - max( dot( reflDirEnv, up ), 0.0 ), 1.5 ), 0.0, 1.0 );
	let skyAnalytic : vec3<f32> = mix( skyZenith, skyHorizon, horizonT );
	let skyMask     : vec3<f32> = vec3<f32>( 0.08, 0.18, 0.32 );
	let skyColor    : vec3<f32> = mix( skyAnalytic, skyMask, 0.4 );

	// Specular highlight using Blinn-Phong
	let halfDir    : vec3<f32> = normalize( lightDir + viewDir );
	let NdotH      : f32       = max( dot( N, halfDir ), 0.0 );
	let spec       : f32       = pow( NdotH, 140.0 ) * NdotL;
	let specTerm   : vec3<f32> = spec * vec3<f32>( 1.0, 1.0, 1.0 ) * 1.3;

	// Fresnel using Schlick's approximation (F0 ~ 0.02 for water)
	let F0         : vec3<f32> = vec3<f32>( 0.02, 0.025, 0.03 );
	let oneMinus   : f32       = 1.0 - NdotV;
	let fresnel    : vec3<f32> = F0 + ( vec3<f32>( 1.0, 1.0, 1.0 ) - F0 ) * pow( oneMinus, 5.0 );

	// Diffuse-ish underwater term (absorbed light)
	let diffuse    : vec3<f32> = waterBase * ( 0.15 + 0.85 * NdotL ) * sunColor;

	// Reflected environment (sky + sun highlight)
	let envBlend   : vec3<f32> = mix( skyColor, vec3<f32>( 0.08, 0.25, 0.55 ), 0.35 );
	let reflection : vec3<f32> = envBlend + specTerm;

	// Mix refraction (underwater) and reflection via Fresnel
	let color      : vec3<f32> = diffuse * ( vec3<f32>( 1.0, 1.0, 1.0 ) - fresnel ) + reflection * fresnel;

	// Mild contrast to keep it punchy
	let colorGamma : vec3<f32> = pow( clamp( color, vec3<f32>( 0.0, 0.0, 0.0 ), vec3<f32>( 1.0, 1.0, 1.0 ) ), vec3<f32>( 1.0 / 2.2, 1.0 / 2.2, 1.0 / 2.2 ) );

	return vec4<f32>( colorGamma, 1.0 );

}