Adding color to terrain

src/components/editor/index.tsx

@@ -25,7 +25,7 @@ const initialNodes: (Node & nodeTypes.All)[] = [
     id: 'terrain-out',
     position: { x: 250, y: 0 },
     type: 'terrain',
-    data: {},
+    data: { biome: 'Grassland' },
   },
 ];
 

src/components/nodes/terrain-node.tsx

@@ -1,14 +1,46 @@
-import type { NodeProps } from 'reactflow';
+import { useEffect } from 'react';
+import { useReactFlow, type NodeProps } from 'reactflow';
+
+import * as helpers from './helpers';
 
 import * as TerrainGenNode from '@/components/common/terraingen-node';
+import { useGraphGlobals } from '@/hooks/use-graph-globals';
 import * as nodeTypes from '@/lib/graph/node-types';
 
+type TerrainNodeData = nodeTypes.Terrain['data'];
 const HANDLES = nodeTypes.HANDLES.terrain;
 
-function TerrainNode({ ...props }: NodeProps) {
+function TerrainNode({ id, data, ...props }: NodeProps<TerrainNodeData>) {
+  const { setNodes } = useReactFlow();
+  const { triggerNodePipelineUpdate } = useGraphGlobals();
+
+  const onBiomeChange = (biome: TerrainNodeData['biome']) => {
+    helpers.updateNodeData<TerrainNodeData>({ id, setNodes, newData: { biome } });
+  };
+
+  useEffect(() => {
+    triggerNodePipelineUpdate(id);
+  }, [data.biome, id, triggerNodePipelineUpdate]);
+
   return (
     <TerrainGenNode.Root title="Terrain (Output)" {...props}>
       <TerrainGenNode.HandleInput label="Height" handleId={HANDLES.in.height} valueType="f32" />
+      <TerrainGenNode.HandleInput
+        label="Water Height"
+        handleId={HANDLES.in.waterHeight}
+        valueType="f32"
+      />
+      <TerrainGenNode.SelectInput
+        label="Biome"
+        value={data.biome}
+        onChange={onBiomeChange}
+        options={[
+          { label: 'Grassland', value: 'Grassland' },
+          { label: 'Desert', value: 'Desert' },
+          { label: 'Mountain', value: 'Mountain' },
+          { label: 'Tundra', value: 'Tundra' },
+        ]}
+      />
     </TerrainGenNode.Root>
   );
 }

src/lib/graph/index.ts

@@ -76,15 +76,49 @@ function generatePipelines(
     const heightEdgeSourceNode = orderedDependencyNodes.find(
       (node) => node.id === heightEdge?.source,
     );
-    const outputs: scene.DisplacePipeline['outputs'] = {
-      height: nodeMapping.getHandleKey({
-        // TODO: wow these type assertions are awesome (evil as fuck)
-        sourceNode: heightEdgeSourceNode!,
-        outgoingHandleId: heightEdge!.sourceHandle!,
-      }),
-    };
 
-    displacePipeline = { instructionSet, uniforms, outputs };
+    // Add this check - if no height edge, skip building the pipeline
+    if (!heightEdge || !heightEdgeSourceNode) {
+      console.warn('Terrain node requires a height input connection');
+      displacePipeline = undefined;
+    } else {
+      const waterHeightEdge = edges.find(
+        (edge) =>
+          edge.target === terrainNode.id &&
+          edge.targetHandle === nodeTypes.HANDLES.terrain.in.waterHeight,
+      );
+      const waterHeightSourceNode = waterHeightEdge
+        ? orderedDependencyNodes.find((node) => node.id === waterHeightEdge.source)
+        : undefined;
+
+      const outputs: scene.DisplacePipeline['outputs'] = {
+        height: nodeMapping.getHandleKey({
+          sourceNode: heightEdgeSourceNode,
+          outgoingHandleId: heightEdge.sourceHandle!,
+        }),
+
+        waterHeight:
+          waterHeightSourceNode && waterHeightEdge
+            ? nodeMapping.getHandleKey({
+                sourceNode: waterHeightSourceNode,
+                outgoingHandleId: waterHeightEdge.sourceHandle!,
+              })
+            : undefined,
+
+        biome: (() => {
+          const biome = (terrainNode as nodeTypes.Terrain & { id: string }).data.biome;
+          const biomeMap: Record<string, number> = {
+            Grassland: 0,
+            Desert: 1,
+            Mountain: 2,
+            Tundra: 3,
+          };
+          return biomeMap[biome] ?? 0;
+        })(),
+      };
+
+      displacePipeline = { instructionSet, uniforms, outputs };
+    }
   }
 
   // find water pipeline
@@ -111,6 +145,7 @@ function generatePipelines(
         sourceNode: heightEdgeSourceNode!,
         outgoingHandleId: heightEdge!.sourceHandle!,
       }),
+      waterHeight: undefined,
     };
 
     waterPipeline = { instructionSet, uniforms, outputs };

src/lib/graph/node-types.ts

@@ -21,7 +21,10 @@ export type Mix = Node<'mix', { nodeType: 'Vec3' | 'Float' }>;
 export type SmoothstepFloat = Node<'smoothstepFloat'>;
 export type SmoothstepVec3 = Node<'smoothstepVec3'>;
 export type VertexData = Node<'vertexData'>;
-export type Terrain = Node<'terrain'>;
+export type Terrain = Node<
+  'terrain',
+  { biome: 'Grassland' | 'Desert' | 'Mountain' | 'Tundra' }
+>;
 export type Separate = Node<'separate'>;
 export type Combine = Node<'combine'>;
 export type Float = Node<'float', { value: number }>;
@@ -114,7 +117,10 @@ export const HANDLES = {
     out: { position: 'vec3-pos-out' },
   },
   terrain: {
-    in: { height: 'float-trans-in' },
+    in: {
+      height: 'float-trans-in',
+      waterHeight: 'float-waterHeight-in',
+    },
     out: {},
   },
   separate: {
@@ -174,7 +180,10 @@ export const NODE_PREFABS: { [nodeType in All['type']]: All & { type: nodeType }
   },
   mix: { type: 'mix', data: { nodeType: 'Float' } },
   vertexData: { type: 'vertexData', data: {} },
-  terrain: { type: 'terrain', data: {} },
+  terrain: {
+    type: 'terrain',
+    data: { biome: 'Grassland' },
+  },
   vector: { type: 'vector', data: { x: 0, y: 0, z: 0 } },
   mathFloat: {
     type: 'mathFloat',

src/lib/renderers/terrain-renderer.ts

@@ -78,11 +78,19 @@ export class TerrainRenderer implements IRenderer {
   private customNodeGraphUniformsBindGroupLayout: GPUBindGroupLayout;
   private customNodeGraphUniformsBindGroup: GPUBindGroup;
 
+  private biomeUniformBuffer: GPUBuffer;
+  private biomeBindGroupLayout: GPUBindGroupLayout;
+  private biomeBindGroup: GPUBindGroup;
+
   // custom uniform buffers
   private nodeGraphUniformBuffer!: GPUBuffer;
   private nodeGraphUniformLayout!: Map<string, number> | undefined;
   private nodeGraphUniformConfig!: scene.DisplacePipeline['uniforms'] | undefined;
 
+  private waterHeightUniformBuffer: GPUBuffer;
+  private waterHeightBindGroupLayout: GPUBindGroupLayout;
+  private waterHeightBindGroup: GPUBindGroup;
+
   /** Custom displace pipeline, gets reconfigured whenever node structure is changed */
   private customDisplacePipeline: GPUComputePipeline;
 
@@ -194,6 +202,68 @@ export class TerrainRenderer implements IRenderer {
       ],
     });
 
+    this.biomeUniformBuffer = this.device.createBuffer({
+      label: 'biome uniform',
+      size: 16,
+      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
+    });
+
+    const defaultBiome = new Uint32Array([0]);
+    this.device.queue.writeBuffer(this.biomeUniformBuffer, 0, defaultBiome);
+
+    this.biomeBindGroupLayout = this.device.createBindGroupLayout({
+      label: 'biome bind group layout',
+      entries: [
+        {
+          binding: 0,
+          visibility: GPUShaderStage.FRAGMENT,
+          buffer: { type: 'uniform' },
+        },
+      ],
+    });
+
+    this.biomeBindGroup = this.device.createBindGroup({
+      label: 'biome bind group',
+      layout: this.biomeBindGroupLayout,
+      entries: [
+        {
+          binding: 0,
+          resource: { buffer: this.biomeUniformBuffer },
+        },
+      ],
+    });
+
+    this.waterHeightUniformBuffer = this.device.createBuffer({
+      label: 'water height uniform',
+      size: 16, // f32 with padding
+      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
+    });
+
+    // Initialize with 0
+    this.device.queue.writeBuffer(this.waterHeightUniformBuffer, 0, new Float32Array([0.0]));
+
+    this.waterHeightBindGroupLayout = this.device.createBindGroupLayout({
+      label: 'water height bind group layout',
+      entries: [
+        {
+          binding: 0,
+          visibility: GPUShaderStage.FRAGMENT,
+          buffer: { type: 'uniform' },
+        },
+      ],
+    });
+
+    this.waterHeightBindGroup = this.device.createBindGroup({
+      label: 'water height bind group',
+      layout: this.waterHeightBindGroupLayout,
+      entries: [
+        {
+          binding: 0,
+          resource: { buffer: this.waterHeightUniformBuffer },
+        },
+      ],
+    });
+
     // initialize depth texture and depth texture view
     this.depthTexture = this.createDepthTexture({
       width: this.webGPU.canvas.width * window.devicePixelRatio,
@@ -621,7 +691,11 @@ export class TerrainRenderer implements IRenderer {
     return this.device.createRenderPipeline({
       layout: this.device.createPipelineLayout({
         label: 'naive pipeline layout',
-        bindGroupLayouts: [this.sceneUniformsBindGroupLayout],
+        bindGroupLayouts: [
+          this.sceneUniformsBindGroupLayout,
+          this.waterHeightBindGroupLayout,
+          this.biomeBindGroupLayout,
+        ],
       }),
       depthStencil: {
         depthWriteEnabled: true,
@@ -882,6 +956,8 @@ export class TerrainRenderer implements IRenderer {
 
     renderPass.setPipeline(this.pipeline);
     renderPass.setBindGroup(0, this.sceneUniformsBindGroup);
+    renderPass.setBindGroup(1, this.waterHeightBindGroup);
+    renderPass.setBindGroup(2, this.biomeBindGroup);
     renderPass.setVertexBuffer(0, this.stage.groundPlane.vertexBuffer);
     renderPass.setIndexBuffer(this.stage.groundPlane.indexBuffer!, 'uint32');
     renderPass.drawIndexedIndirect(this.stage.groundPlane.indirectBuffer!, 0);
@@ -928,6 +1004,26 @@ export class TerrainRenderer implements IRenderer {
 
     // also TODO: reuse code between this and our constructor
 
+    if (config.outputs.waterHeight) {
+      console.log('Water height output key:', config.outputs.waterHeight);
+
+      const waterHeightKey = config.outputs.waterHeight.replace('nodeGraphUniforms.', '');
+      const waterHeightUniform = config.uniforms.find((u) => u.key === waterHeightKey);
+
+      if (
+        waterHeightUniform &&
+        waterHeightUniform.type === 'f32' &&
+        waterHeightUniform.initialValue !== null
+      ) {
+        this.setWaterHeightForTerrain(waterHeightUniform.initialValue);
+      }
+    }
+
+    if (config.outputs.biome !== undefined) {
+      const biomeData = new Uint32Array([config.outputs.biome]);
+      this.device.queue.writeBuffer(this.biomeUniformBuffer, 0, biomeData);
+    }
+
     const customComputeShader = jit.generateDisplaceShaderCode(
       config,
       displaceComputeShaderTemplate,
@@ -1176,6 +1272,7 @@ export class TerrainRenderer implements IRenderer {
 
   setMeshUniforms(size: number, resolution: number) {
     this.stage.groundPlane.updateUniforms(this.device, size, resolution);
+    this.stage.waterPlane.updateUniforms(this.device, size, resolution);
 
     const encoder = this.device.createCommandEncoder();
     this.runComputes(encoder);
@@ -1216,4 +1313,9 @@ export class TerrainRenderer implements IRenderer {
     this.device.queue.submit([encoder.finish()]);
     console.log('rerunning compute after uniform update');
   }
+
+  setWaterHeightForTerrain(height: number) {
+    this.device.queue.writeBuffer(this.waterHeightUniformBuffer, 0, new Float32Array([height]));
+    console.log('Updated water height for terrain fragment shader:', height);
+  }
 }

src/lib/shaders/jit/types/shaders.ts

@@ -13,6 +13,8 @@ type ShaderConfigBase = {
 export type DisplaceShaderConfig = ShaderConfigBase & {
   outputs: {
     height: util.ReferenceKey;
+    waterHeight?: util.ReferenceKey;
+    biome?: number;
   };
 };