src/lib/components/3d/GridCustom.svelte

<!-- Credits to Fyrestar for the https://github.com/Fyrestar/THREE.InfiniteGridHelper  -->
<script lang="ts">
	import { T, useTask, useThrelte } from "@threlte/core";
	import { Color, DoubleSide, Plane, Vector3, Mesh } from "three";
	import * as THREE from "three";

	// Grid shader code with improved precision and stability
	import { revision } from "@threlte/core";

	// Props
	let {
		cellColor = "#71717A",
		sectionColor = "#707070",
		cellSize = 1,
		backgroundColor = "#000000",
		backgroundOpacity = 0,
		sectionSize = 10,
		plane = "xz",
		gridSize = [20, 20],
		followCamera = false,
		infiniteGrid = false,
		fadeDistance = 100,
		fadeStrength = 1,
		fadeOrigin = undefined,
		cellThickness = 1,
		sectionThickness = 2,
		side = DoubleSide,
		type = "grid",
		axis = "x",
		maxRadius = 0,
		cellDividers = 6,
		sectionDividers = 2,
		floorColor = "#2a2a2a",
		floorOpacity = 0.3,
		ref = $bindable(),
		children = undefined,
		...props
	} = $props();

	// Shared fade calculation function for both shaders
	const fadeCalculation = /*glsl*/ `
		float calculateFade(vec3 worldPos, float viewZ, vec3 fadeOrigin, float fadeDistance, float fadeStrength, float cameraNear, float cameraFar) {
			float dist = distance(fadeOrigin, worldPos);
			float fadeFactor = 1.0 - clamp(dist / fadeDistance, 0.0, 1.0);
			fadeFactor = pow(fadeFactor, fadeStrength);
			
			float viewDepthFade = 1.0 - clamp((viewZ - cameraNear) / (cameraFar - cameraNear), 0.0, 1.0);
			viewDepthFade = smoothstep(0.0, 0.3, viewDepthFade);
			
			return min(fadeFactor, viewDepthFade);
		}
	`;

	const vertexShader = /*glsl*/ `
		varying vec3 localPosition;
		varying vec4 worldPosition;
		varying float vViewZ;

		uniform vec3 worldCamProjPosition;
		uniform vec3 worldPlanePosition;
		uniform float fadeDistance;
		uniform bool infiniteGrid;
		uniform bool followCamera;
		uniform int coord0, coord1, coord2;

		void main() {
			localPosition = vec3(position[coord0], position[coord1], position[coord2]);
			if (infiniteGrid) localPosition *= 1.0 + fadeDistance;

			worldPosition = modelMatrix * vec4(localPosition, 1.0);
			if (followCamera) {
				worldPosition.xyz += (worldCamProjPosition - worldPlanePosition);
				localPosition = (inverse(modelMatrix) * worldPosition).xyz;
			}

			vec4 mvPosition = viewMatrix * worldPosition;
			vViewZ = -mvPosition.z;
			gl_Position = projectionMatrix * mvPosition;
		}
	`;

	const fragmentShader = /*glsl*/ `
		#define PI 3.141592653589793
		
		varying vec3 localPosition;
		varying vec4 worldPosition;
		varying float vViewZ;

		uniform float cellSize, sectionSize, cellThickness, sectionThickness;
		uniform vec3 cellColor, sectionColor, backgroundColor, fadeOrigin;
		uniform float backgroundOpacity, fadeDistance, fadeStrength, cameraNear, cameraFar;
		uniform bool infiniteGrid;
		uniform int coord0, coord1, coord2, gridType, lineGridCoord;
		uniform float circleGridMaxRadius, polarCellDividers, polarSectionDividers;

		${fadeCalculation}

		float getSquareGrid(float size, float thickness, vec3 localPos) {
			vec2 coord = localPos.xy / size;
			vec2 derivative = fwidth(coord);
			vec2 grid = abs(fract(coord - 0.5) - 0.5) / derivative;
			float line = min(grid.x, grid.y) + 1.0 - thickness;
			return clamp(1.0 - line, 0.0, 1.0);
		}

		float getLinesGrid(float size, float thickness, vec3 localPos) {
			float coord = localPos[lineGridCoord] / size;
			float derivative = fwidth(coord);
			float line = abs(fract(coord - 0.5) - 0.5) / derivative - thickness * 0.5;
			return clamp(1.0 - line, 0.0, 1.0);
		}

		float getCirclesGrid(float size, float thickness, vec3 localPos) {
			float coord = length(localPos.xy) / size;
			float derivative = fwidth(coord);
			float line = abs(fract(coord - 0.5) - 0.5) / derivative - thickness * 0.5;
			if (!infiniteGrid && circleGridMaxRadius > 0.0 && coord > circleGridMaxRadius + thickness * 0.1) discard;
			return clamp(1.0 - line, 0.0, 1.0);
		}

		float getPolarGrid(float size, float thickness, float polarDividers, vec3 localPos) {
			float rad = length(localPos.xy) / size;
			vec2 coord = vec2(rad, atan(localPos.x, localPos.y) * polarDividers / PI);
			vec2 derivative = fwidth(coord);
			vec2 grid = abs(fract(coord - 0.5) - 0.5) / derivative;
			float line = min(grid.x, grid.y) + 1.0 - thickness;
			if (!infiniteGrid && circleGridMaxRadius > 0.0 && rad > circleGridMaxRadius + thickness * 0.1) discard;
			return clamp(1.0 - line, 0.0, 1.0);
		}

		void main() {
			float g1 = 0.0, g2 = 0.0;
			vec3 localPos = vec3(localPosition[coord0], localPosition[coord1], localPosition[coord2]);

			if (gridType == 0) {
				g1 = getSquareGrid(cellSize, cellThickness, localPos);
				g2 = getSquareGrid(sectionSize, sectionThickness, localPos);
			} else if (gridType == 1) {
				g1 = getLinesGrid(cellSize, cellThickness, localPos);
				g2 = getLinesGrid(sectionSize, sectionThickness, localPos);
			} else if (gridType == 2) {
				g1 = getCirclesGrid(cellSize, cellThickness, localPos);
				g2 = getCirclesGrid(sectionSize, sectionThickness, localPos);
			} else if (gridType == 3) {
				g1 = getPolarGrid(cellSize, cellThickness, polarCellDividers, localPos);
				g2 = getPolarGrid(sectionSize, sectionThickness, polarSectionDividers, localPos);
			}

			float fadeFactor = calculateFade(worldPosition.xyz, vViewZ, fadeOrigin, fadeDistance, fadeStrength, cameraNear, cameraFar);
			vec3 color = mix(cellColor, sectionColor, clamp(sectionThickness * g2, 0.0, 1.0));
			float gridAlpha = clamp((g1 + g2) * fadeFactor, 0.0, 1.0);

			if (backgroundOpacity > 0.0) {
				vec3 finalColor = mix(backgroundColor, color, gridAlpha);
				float blendedAlpha = clamp(max(gridAlpha, backgroundOpacity * fadeFactor), 0.0, 1.0);
				gl_FragColor = vec4(finalColor, blendedAlpha);
			} else {
				gl_FragColor = vec4(color, gridAlpha);
			}

			if (gl_FragColor.a < 0.05) discard;

			#include <tonemapping_fragment>
			#include <${revision < 154 ? "encodings_fragment" : "colorspace_fragment"}>
		}
	`;

	// Simple floor shader
	const floorVertexShader = /*glsl*/ `
		varying vec3 vWorldPosition;
		varying float vViewZ;
		
		void main() {
			vec4 worldPosition = modelMatrix * vec4(position, 1.0);
			vWorldPosition = worldPosition.xyz;
			vec4 mvPosition = viewMatrix * worldPosition;
			vViewZ = -mvPosition.z;
			gl_Position = projectionMatrix * mvPosition;
		}
	`;

	const floorFragmentShader = /*glsl*/ `
		uniform vec3 floorColor, fadeOrigin;
		uniform float floorOpacity, fadeDistance, fadeStrength, cameraNear, cameraFar;
		varying vec3 vWorldPosition;
		varying float vViewZ;
		
		${fadeCalculation}
		
		void main() {
			float fadeFactor = calculateFade(vWorldPosition, vViewZ, fadeOrigin, fadeDistance, fadeStrength, cameraNear, cameraFar);
			float finalOpacity = floorOpacity * fadeFactor;
			gl_FragColor = vec4(floorColor, finalOpacity);
			if (gl_FragColor.a < 0.01) discard;
		}
	`;

	const mesh = new Mesh();
	const { invalidate, camera } = useThrelte();
	const gridPlane = new Plane();
	const upVector = new Vector3(0, 1, 0);
	const zeroVector = new Vector3(0, 0, 0);

	const axisToInt: Record<string, number> = { x: 0, y: 1, z: 2 };
	const planeToAxes: Record<string, string> = { xz: "xzy", xy: "xyz", zy: "zyx" };
	const gridType = { grid: 0, lines: 1, circular: 2, polar: 3 };

	// Shared uniforms (used by both grid and floor)
	const sharedUniforms = {
		fadeOrigin: { value: new Vector3() },
		fadeDistance: { value: fadeDistance },
		fadeStrength: { value: fadeStrength },
		cameraNear: { value: 0.1 },
		cameraFar: { value: 1000 }
	};

	// Grid uniforms
	const uniforms = {
		...sharedUniforms,
		cellSize: { value: cellSize },
		sectionSize: { value: sectionSize },
		cellColor: { value: new Color(cellColor) },
		sectionColor: { value: new Color(sectionColor) },
		backgroundColor: { value: new Color(backgroundColor) },
		backgroundOpacity: { value: backgroundOpacity },
		cellThickness: { value: cellThickness },
		sectionThickness: { value: sectionThickness },
		infiniteGrid: { value: infiniteGrid },
		followCamera: { value: followCamera },
		coord0: { value: 0 },
		coord1: { value: 2 },
		coord2: { value: 1 },
		gridType: { value: gridType.grid },
		lineGridCoord: { value: axisToInt[axis as keyof typeof axisToInt] || 0 },
		circleGridMaxRadius: { value: maxRadius },
		polarCellDividers: { value: cellDividers },
		polarSectionDividers: { value: sectionDividers },
		worldCamProjPosition: { value: new Vector3() },
		worldPlanePosition: { value: new Vector3() }
	};

	// Floor uniforms (simpler, reusing shared uniforms)
	const floorUniforms = {
		...sharedUniforms,
		floorColor: { value: new Color(floorColor) },
		floorOpacity: { value: floorOpacity }
	};

	// Single update effect for all uniforms
	$effect.pre(() => {
		const axes = planeToAxes[plane] || "xzy";
		const [c0, c1, c2] = [axes.charAt(0), axes.charAt(1), axes.charAt(2)].map(
			(c) => axisToInt[c as keyof typeof axisToInt]
		);

		// Update grid uniforms
		Object.assign(uniforms, {
			coord0: { value: c0 },
			coord1: { value: c1 },
			coord2: { value: c2 },
			cellSize: { value: cellSize },
			sectionSize: { value: sectionSize },
			cellThickness: { value: cellThickness },
			sectionThickness: { value: sectionThickness },
			backgroundOpacity: { value: backgroundOpacity },
			infiniteGrid: { value: infiniteGrid },
			followCamera: { value: followCamera }
		});

		uniforms.cellColor.value.set(cellColor);
		uniforms.sectionColor.value.set(sectionColor);
		uniforms.backgroundColor.value.set(backgroundColor);

		// Update shared uniforms (affects both grid and floor)
		sharedUniforms.fadeDistance.value = fadeDistance;
		sharedUniforms.fadeStrength.value = fadeStrength;
		floorUniforms.floorColor.value.set(floorColor);
		floorUniforms.floorOpacity.value = floorOpacity;

		// Update camera uniforms
		if (camera.current && "near" in camera.current && "far" in camera.current) {
			const cam = camera.current as THREE.PerspectiveCamera;
			sharedUniforms.cameraNear.value = cam.near;
			sharedUniforms.cameraFar.value = cam.far;
		}

		// Update grid type
		const typeMap = { grid: 0, lines: 1, circular: 2, polar: 3 };
		uniforms.gridType.value = typeMap[type as keyof typeof typeMap] || 0;
		if (type === "lines")
			uniforms.lineGridCoord.value = axisToInt[axis as keyof typeof axisToInt] || 0;
		if (type === "circular" || type === "polar") {
			uniforms.circleGridMaxRadius.value = maxRadius;
			if (type === "polar") {
				uniforms.polarCellDividers.value = cellDividers;
				uniforms.polarSectionDividers.value = sectionDividers;
			}
		}

		invalidate();
	});

	// Single task for both grid and floor fade origins
	useTask(
		() => {
			gridPlane.setFromNormalAndCoplanarPoint(upVector, zeroVector).applyMatrix4(mesh.matrixWorld);
			const material = mesh.material as THREE.ShaderMaterial;
			if (material?.uniforms) {
				const { worldCamProjPosition, worldPlanePosition } = material.uniforms;
				const projectedPoint = gridPlane.projectPoint(
					camera.current.position,
					worldCamProjPosition.value
				);
				if (!fadeOrigin) sharedUniforms.fadeOrigin.value = projectedPoint;
				worldPlanePosition.value.set(0, 0, 0).applyMatrix4(mesh.matrixWorld);
			}
		},
		{ autoInvalidate: false }
	);
</script>

<!-- Shadow-receiving floor underneath -->
<T.Mesh rotation={[-Math.PI / 2, 0, 0]} receiveShadow position.y={0} {...props}>
	<T.PlaneGeometry
		args={infiniteGrid
			? [1000, 1000]
			: typeof gridSize == "number"
				? [gridSize, gridSize]
				: gridSize}
	/>
	<T.ShadowMaterial
		transparent={true}
		opacity={0.3}
		polygonOffset={true}
		polygonOffsetFactor={1}
		polygonOffsetUnits={1}
	/>
</T.Mesh>

<!-- Fading floor -->
<T.Mesh rotation={[-Math.PI / 2, 0, 0]} position.y={0} {...props}>
	<T.PlaneGeometry
		args={infiniteGrid
			? [1000, 1000]
			: typeof gridSize == "number"
				? [gridSize, gridSize]
				: gridSize}
	/>
	<T.ShaderMaterial
		vertexShader={floorVertexShader}
		fragmentShader={floorFragmentShader}
		uniforms={floorUniforms}
		transparent={true}
		side={THREE.DoubleSide}
		depthTest={true}
		depthWrite={false}
		polygonOffset={true}
		polygonOffsetFactor={-1}
		polygonOffsetUnits={-1}
	/>
</T.Mesh>

<!-- Grid lines -->
<T is={mesh} bind:ref frustumCulled={false} position.y={0.005} {...props}>
	<T.ShaderMaterial
		{fragmentShader}
		{vertexShader}
		{uniforms}
		transparent
		{side}
		depthTest={true}
		depthWrite={false}
	/>
	{#if children}
		{@render children({ ref: mesh })}
	{:else}
		<T.PlaneGeometry args={typeof gridSize == "number" ? [gridSize, gridSize] : gridSize} />
	{/if}
</T>