@takram/three-geospatial 0.8.0 → 0.9.1

package/src/webgpu/FnLayout.ts

@@ -1,17 +1,17 @@
 import type {
   ProxiedTuple,
   ShaderCallNodeInternal,
-  ShaderNodeFn,
-  Struct
+  ShaderNodeFn
 } from 'three/src/nodes/TSL.js'
 import { Fn } from 'three/tsl'
-import type { NodeBuilder, Texture3DNode, TextureNode } from 'three/webgpu'
+import { StructTypeNode, type NodeBuilder } from 'three/webgpu'
 
 import type { Node, NodeType } from './node'
 
-// Note that "texture" and "texture3D" are just placeholders until TSL supports
-// texture types.
-type FnLayoutType = NodeType | Struct | 'texture' | 'texture3D'
+type FnLayoutType =
+  | NodeType
+  | (new (...args: any[]) => any)
+  | ((...args: any[]) => any)
 
 export interface FnLayoutInput<T extends FnLayoutType = FnLayoutType> {
   name: string
@@ -30,13 +30,11 @@ export interface FnLayout<
 
 type InferNodeObject<T extends FnLayoutType> = T extends NodeType
   ? Node<T>
-  : T extends Struct
-    ? ReturnType<T>
-    : T extends 'texture'
-      ? TextureNode
-      : T extends 'texture3D'
-        ? Texture3DNode
-        : never
+  : T extends new (...args: any[]) => any
+    ? InstanceType<T>
+    : T extends (...args: any[]) => any
+      ? ReturnType<T>
+      : never
 
 type InferNodeObjects<Inputs extends readonly FnLayoutInput[]> = {
   [K in keyof Inputs]: Inputs[K] extends FnLayoutInput<infer T>
@@ -58,10 +56,13 @@ function transformType(type: FnLayoutType): string {
   if (typeof type === 'string') {
     return type
   }
-  if (type.layout.name == null) {
-    throw new Error('Struct name is required.')
+  if ('layout' in type && type.layout instanceof StructTypeNode) {
+    if (type.layout.name == null) {
+      throw new Error('Struct name is required.')
+    }
+    return type.layout.name
   }
-  return type.layout.name
+  throw new Error(`Unsupported layout type: ${type}`)
 }
 
 export function FnLayout<

package/src/webgpu/HighpVelocityNode.ts

@@ -20,10 +20,10 @@ export class HighpVelocityNode extends TempNode {
 
   projectionMatrix?: Matrix4 | null
 
-  private readonly currentProjectionMatrix = uniform(new Matrix4())
+  private readonly currentProjectionMatrix = uniform('mat4')
   private readonly previousProjectionMatrix = uniform('mat4')
 
-  private readonly currentModelViewMatrix = uniform(new Matrix4())
+  private readonly currentModelViewMatrix = uniform('mat4')
   private readonly previousModelViewMatrix = uniform('mat4')
   private readonly objectModelViewMatrices = new WeakMap<Object3D, Matrix4>()
 

package/src/webgpu/LensFlareNode.ts

@@ -45,7 +45,7 @@ export class LensFlareNode extends TempNode {
     this.glareNode = new LensGlareNode()
 
     this.featuresNode = rtt(add(this.ghostNode, this.haloNode))
-    this.featuresNode.value.name = 'LensFlareNode.Features'
+    this.featuresNode.value.name = 'LensFlare_features'
     this.featuresNode.pixelRatio = 0.5
 
     // Use the full resolution because the thresholdNode already downsamples the

package/src/webgpu/LensGlareNode.ts

@@ -29,9 +29,9 @@ import {
   vec4
 } from 'three/tsl'
 import {
+  IndirectStorageBufferAttribute,
   MeshBasicNodeMaterial,
   RendererUtils,
-  StorageBufferAttribute,
   type ComputeNode,
   type NodeBuilder,
   type NodeFrame,
@@ -92,7 +92,12 @@ export class LensGlareNode extends FilterNode {
 
   private computeNode?: ComputeNode
 
-  private readonly counterBuffer = new StorageBufferAttribute(1, 1)
+  // drawIndexedIndirect format:
+  // [indexCount, instanceCount, firstIndex, baseVertex, firstInstance]
+  private readonly indirectBuffer = new IndirectStorageBufferAttribute(
+    new Uint32Array([6, 0, 0, 0, 0]),
+    1
+  )
   private instanceBuffer = instancedArray(1, instanceStruct)
 
   private readonly renderTarget = this.createRenderTarget()
@@ -106,16 +111,19 @@ export class LensGlareNode extends FilterNode {
   private readonly camera = new PerspectiveCamera()
   private rendererState?: RendererUtils.RendererState
 
-  private readonly inputTexelSize = uniform(new Vector2())
-  private readonly outputTexelSize = uniform(new Vector2())
-  private readonly geometryRatio = uniform(new Vector2())
+  private readonly inputTexelSize = uniform('vec2')
+  private readonly outputTexelSize = uniform('vec2')
+  private readonly geometryRatio = uniform('vec2')
 
   constructor(inputNode?: TextureNode | null) {
     super(inputNode)
+    this.material.name = 'LensGlare'
+
     this.inputNode = inputNode
     this.resolutionScale = 0.5
 
     this.outputTexture = this.renderTarget.texture
+    this.mesh.geometry.indirect = this.indirectBuffer
   }
 
   override customCacheKey(): number {
@@ -151,7 +159,7 @@ export class LensGlareNode extends FilterNode {
     const { width: inputWidth, height: inputHeight } = inputNode.value
     this.setSize(inputWidth, inputHeight) // Compute node is initialized here.
 
-    const { computeNode, counterBuffer, renderTarget } = this
+    const { computeNode, indirectBuffer, renderTarget } = this
     invariant(computeNode != null)
 
     this.inputTexelSize.value.set(1 / inputWidth, 1 / inputHeight)
@@ -165,24 +173,12 @@ export class LensGlareNode extends FilterNode {
     const { width: outputWidth, height: outputHeight } = renderTarget
     this.outputTexelSize.value.set(1 / outputWidth, 1 / outputHeight)
 
-    // Reset the counter:
-    counterBuffer.array[0] = 0
-    counterBuffer.needsUpdate = true
+    // Reset instanceCount in the indirect buffer:
+    indirectBuffer.array[1] = 0
+    indirectBuffer.needsUpdate = true
 
     void renderer.compute(computeNode)
 
-    renderer
-      .getArrayBufferAsync(counterBuffer)
-      .then(arrayBuffer => {
-        // TODO: This is indeed a couple of frames behind, thus the number of
-        // computed instances above and the number of instances to be drawn by
-        // the mesh differ.
-        this.mesh.count = new Uint32Array(arrayBuffer)[0]
-      })
-      .catch((error: unknown) => {
-        console.error(error)
-      })
-
     this.rendererState = resetRendererState(renderer, this.rendererState)
 
     renderer.setRenderTarget(renderTarget)
@@ -195,16 +191,16 @@ export class LensGlareNode extends FilterNode {
     const {
       spikePairCount,
       inputNode,
-      counterBuffer,
+      indirectBuffer,
       instanceBuffer,
       outputTexelSize
     } = this
     invariant(inputNode != null)
 
-    const counterStorage = storage(
-      counterBuffer,
+    const indirectStorage = storage(
+      indirectBuffer,
       'uint',
-      counterBuffer.count
+      indirectBuffer.count
     ).toAtomic()
 
     this.computeNode = Fn(() => {
@@ -218,7 +214,11 @@ export class LensGlareNode extends FilterNode {
       const inputLuminance = inputColor.a // Alpha channel stores luminance
 
       If(inputLuminance.greaterThan(0.1), () => {
-        const countBefore = atomicAdd(counterStorage.element(0), spikePairCount)
+        // The first element is instanceCount in the drawIndexedIndirect buffer.
+        const countBefore = atomicAdd(
+          indirectStorage.element(1),
+          spikePairCount
+        )
         for (let i = 0; i < spikePairCount; ++i) {
           const instance = instanceBuffer.element(countBefore.add(i))
           instance.get('color').assign(inputColor.rgb)

package/src/webgpu/LensHaloNode.ts

@@ -38,8 +38,9 @@ export class LensHaloNode extends TempNode {
 
   constructor(inputNode?: TextureNode | null) {
     super('vec3')
-    this.inputNode = inputNode
     this.updateBeforeType = NodeUpdateType.FRAME
+
+    this.inputNode = inputNode
   }
 
   override updateBefore({ renderer }: NodeFrame): void {

package/src/webgpu/OutputTexture3DNode.ts

@@ -31,7 +31,16 @@ export class OutputTexture3DNode extends Texture3DNode {
 
   override clone(): this {
     // @ts-expect-error Ignore
-    return new this.constructor(this.owner, this.value)
+    const copy = new this.constructor(this.owner, this.value)
+    copy.uvNode = this.uvNode
+    copy.levelNode = this.levelNode
+    copy.biasNode = this.biasNode
+    copy.sampler = this.sampler
+    copy.depthNode = this.depthNode
+    copy.compareNode = this.compareNode
+    copy.gradNode = this.gradNode
+    copy.offsetNode = this.offsetNode
+    return copy
   }
 }
 

package/src/webgpu/OutputTextureNode.ts

@@ -31,7 +31,16 @@ export class OutputTextureNode extends TextureNode {
 
   override clone(): this {
     // @ts-expect-error Ignore
-    return new this.constructor(this.owner, this.value)
+    const copy = new this.constructor(this.owner, this.value)
+    copy.uvNode = this.uvNode
+    copy.levelNode = this.levelNode
+    copy.biasNode = this.biasNode
+    copy.sampler = this.sampler
+    copy.depthNode = this.depthNode
+    copy.compareNode = this.compareNode
+    copy.gradNode = this.gradNode
+    copy.offsetNode = this.offsetNode
+    return copy
   }
 }
 

package/src/webgpu/STBNTextureNode.ts

@@ -0,0 +1,72 @@
+import { hashString } from 'three/src/nodes/core/NodeUtils.js'
+import { Fn, frameId, nodeImmutable, screenCoordinate, vec3 } from 'three/tsl'
+import {
+  Data3DTexture,
+  NearestFilter,
+  RedFormat,
+  RepeatWrapping,
+  Texture3DNode,
+  type NodeBuilder
+} from 'three/webgpu'
+
+import { DEFAULT_STBN_URL } from '../constants'
+import { STBNLoader } from '../STBNLoader'
+
+const emptyTexture3D = /*#__PURE__*/ (() => {
+  const texture = new Data3DTexture(new Uint8Array(1))
+  texture.format = RedFormat
+  // BUG: TextureNode doesn't update these when the texture is swapped.
+  texture.minFilter = NearestFilter
+  texture.magFilter = NearestFilter
+  texture.wrapS = RepeatWrapping
+  texture.wrapT = RepeatWrapping
+  texture.wrapR = RepeatWrapping
+  texture.needsUpdate = true
+  return texture
+})()
+
+export class STBNTextureNode extends Texture3DNode {
+  url = DEFAULT_STBN_URL
+
+  constructor() {
+    super(emptyTexture3D)
+  }
+
+  override customCacheKey(): number {
+    return hashString(this.url)
+  }
+
+  override setup(builder: NodeBuilder): unknown {
+    new STBNLoader()
+      .loadAsync(this.url)
+      .then(texture => {
+        this.value = texture
+      })
+      .catch((error: unknown) => {
+        console.error(error)
+      })
+    return super.setup(builder)
+  }
+
+  override clone(): this {
+    // @ts-expect-error Ignore
+    const copy = new this.constructor()
+    copy.uvNode = this.uvNode
+    copy.levelNode = this.levelNode
+    copy.biasNode = this.biasNode
+    copy.sampler = this.sampler
+    copy.depthNode = this.depthNode
+    copy.compareNode = this.compareNode
+    copy.gradNode = this.gradNode
+    copy.offsetNode = this.offsetNode
+    return copy
+  }
+}
+
+export const stbnTexture = /*#__PURE__*/ nodeImmutable(STBNTextureNode)
+
+export const stbn = /*#__PURE__*/ Fn(() => {
+  return stbnTexture
+    .sample(vec3(screenCoordinate.xy, frameId.mod(64)).div(vec3(128, 128, 64)))
+    .r.toConst('stbn')
+}).once()()

package/src/webgpu/ScreenSpaceShadowNode.ts

@@ -14,6 +14,8 @@
  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  * License for the specific language governing permissions and limitations under
  * the License.
+ *
+ * Modified from the original source code.
  */
 
 import {
@@ -38,7 +40,6 @@ import {
   ivec2,
   min,
   mix,
-  textureSize,
   textureStore,
   uniform,
   vec2,
@@ -56,7 +57,6 @@ import {
   type NodeFrame,
   type TextureNode
 } from 'three/webgpu'
-import invariant from 'tiny-invariant'
 
 import { cameraFar, cameraNear } from './accessors'
 import type { Node } from './node'
@@ -107,14 +107,12 @@ export class ScreenSpaceShadowNode extends TempNode {
   shadowContrast = uniform(4)
   shadowIntensity = uniform(1)
   bilinearThreshold = uniform(0.02)
-  nearDepth = uniform(0)
-  farDepth = uniform(1)
 
   // xy: Screen coordinate
   // z: Normalized Z
   // w: Direction sign
-  private readonly lightCoordinate = uniform(new Vector4())
-  private readonly dispatchOffset = uniform(new Vector2(), 'ivec2')
+  private readonly lightCoordinate = uniform('vec4')
+  private readonly dispatchOffset = uniform('ivec2')
   private readonly dispatchIndex = uniform(0)
 
   private readonly dispatches: readonly Dispatch[] = Array.from(
@@ -124,7 +122,7 @@ export class ScreenSpaceShadowNode extends TempNode {
   )
   private dispatchCount = 0
 
-  private computeNode?: ComputeNode
+  private readonly computeNode: ComputeNode
 
   constructor(
     depthNode: TextureNode,
@@ -132,6 +130,8 @@ export class ScreenSpaceShadowNode extends TempNode {
     mainLight: DirectionalLight
   ) {
     super('float')
+    this.updateBeforeType = NodeUpdateType.FRAME
+
     this.depthNode = depthNode
     this.camera = camera
     this.mainLight = mainLight
@@ -141,12 +141,11 @@ export class ScreenSpaceShadowNode extends TempNode {
     texture.minFilter = LinearFilter
     texture.magFilter = LinearFilter
     texture.generateMipmaps = false
-    texture.name = 'ScreenSpaceShadowNode'
+    texture.name = 'ScreenSpaceShadow'
 
     this.outputTexture = texture
     this.textureNode = outputTexture(this, texture)
-
-    this.updateBeforeType = NodeUpdateType.FRAME
+    this.computeNode = this.createComputeNode()
   }
 
   override customCacheKey(): number {
@@ -198,7 +197,6 @@ export class ScreenSpaceShadowNode extends TempNode {
 
     this.updateDispatchList(lightProjection, size)
 
-    invariant(this.computeNode != null)
     for (let index = 0; index < this.dispatchCount; ++index) {
       const dispatch = this.dispatches[index]
       this.dispatchOffset.value.set(dispatch.offset.x, dispatch.offset.y)
@@ -341,7 +339,7 @@ export class ScreenSpaceShadowNode extends TempNode {
   }
 
   // See bend_sss_gpu.h
-  private setupCompute(builder: NodeBuilder): void {
+  private createComputeNode(): ComputeNode {
     const {
       depthNode,
       camera,
@@ -353,8 +351,6 @@ export class ScreenSpaceShadowNode extends TempNode {
       shadowContrast,
       shadowIntensity,
       bilinearThreshold,
-      farDepth,
-      nearDepth,
       lightCoordinate,
       dispatchOffset
     } = this
@@ -442,27 +438,30 @@ export class ScreenSpaceShadowNode extends TempNode {
       return { pixelXY, pixelDistance, xyDelta, xAxisMajor }
     }
 
-    const loadDepth = (coord: Node<'ivec2'>): Node<'float'> => {
-      const depth = depthNode.load(coord).toVar()
-      if (builder.renderer.logarithmicDepthBuffer) {
-        depth.assign(
-          logarithmicToPerspectiveDepth(
+    return Fn(builder => {
+      const [nearDepth, farDepth] = builder.renderer.reversedDepthBuffer
+        ? [float(1), float(0)]
+        : [float(0), float(1)]
+
+      const loadDepth = (coord: Node<'ivec2'>): Node<'float'> => {
+        let depth: Node = depthNode.load(coord)
+        if (builder.renderer.logarithmicDepthBuffer) {
+          depth = logarithmicToPerspectiveDepth(
             depth,
             cameraNear(camera),
             cameraFar(camera)
           )
-        )
+        }
+        depth = depth.toConst()
+
+        // Emulate a point sampler in bend_sss_gpu.h, with Wrap Mode set to
+        // Clamp-To-Border-Color, and Border Color set to farDepth.
+        return and(
+          coord.greaterThanEqual(0).all(),
+          coord.lessThan(depthNode.size()).all()
+        ).select(depth, farDepth)
       }
 
-      // Emulate a point sampler in bend_sss_gpu.h, with Wrap Mode set to
-      // Clamp-To-Border-Color, and Border Color set to farDepth.
-      return and(
-        coord.greaterThanEqual(0).all(),
-        coord.lessThan(textureSize(depthNode)).all()
-      ).select(depth, farDepth)
-    }
-
-    this.computeNode = Fn(() => {
       const { pixelXY, xyDelta, pixelDistance, xAxisMajor } =
         getWorkgroupExtents()
 
@@ -656,7 +655,7 @@ export class ScreenSpaceShadowNode extends TempNode {
         shadowValue.mul(shadowContrast).add(contrastOffset).saturate()
       )
 
-      const result = float().toVar()
+      const result = float(0).toVar()
 
       // Take the average of 4 samples, this is useful to reduce aliasing noise
       // in the source depth, especially with long shadows.
@@ -669,14 +668,10 @@ export class ScreenSpaceShadowNode extends TempNode {
       // Asking the GPU to write scattered single-byte pixels isn't great,
       // But thankfully the latency is hidden by all the work we're doing...
       textureStore(outputTexture, writeXY, mix(1, result, shadowIntensity))
-    })().compute(
-      0, // Determine this later
-      [GROUP_SIZE, 1, 1]
-    )
+    })().computeKernel([GROUP_SIZE, 1, 1])
   }
 
   override setup(builder: NodeBuilder): unknown {
-    this.setupCompute(builder)
     return this.textureNode
   }
 

package/src/webgpu/SeparableFilterNode.ts

@@ -1,4 +1,4 @@
-import { Vector2, type RenderTarget } from 'three'
+import type { RenderTarget } from 'three'
 import { uniform } from 'three/tsl'
 import {
   NodeMaterial,
@@ -24,14 +24,17 @@ export abstract class SeparableFilterNode extends FilterNode {
   private readonly mesh = new QuadMesh(this.material)
   private rendererState?: RendererUtils.RendererState
 
-  protected readonly inputTexelSize = uniform(new Vector2())
-  protected readonly direction = uniform(new Vector2())
+  protected readonly inputTexelSize = uniform('vec2')
+  protected readonly direction = uniform('vec2')
 
   constructor(inputNode?: TextureNode | null) {
     super(inputNode)
+    const typeName = (this.constructor as typeof Node).type.replace(/Node$/, '')
+    this.material.name = typeName
+    this.mesh.name = typeName
 
-    this.horizontalRT = this.createRenderTarget('Horizontal')
-    this.verticalRT = this.createRenderTarget('Vertical')
+    this.horizontalRT = this.createRenderTarget('horizontal')
+    this.verticalRT = this.createRenderTarget('vertical')
     this.outputTexture = this.verticalRT.texture
   }
 

package/src/webgpu/SingleFilterNode.ts

@@ -1,4 +1,4 @@
-import { Vector2, type RenderTarget } from 'three'
+import type { RenderTarget } from 'three'
 import { uniform } from 'three/tsl'
 import {
   NodeMaterial,
@@ -21,10 +21,13 @@ export abstract class SingleFilterNode extends FilterNode {
   private readonly mesh = new QuadMesh(this.material)
   private rendererState?: RendererUtils.RendererState
 
-  protected readonly inputTexelSize = uniform(new Vector2())
+  protected readonly inputTexelSize = uniform('vec2')
 
   constructor(inputNode?: TextureNode | null) {
     super(inputNode)
+    const typeName = (this.constructor as typeof Node).type.replace(/Node$/, '')
+    this.material.name = typeName
+    this.mesh.name = typeName
 
     this.renderTarget = this.createRenderTarget()
     this.outputTexture = this.renderTarget.texture

package/src/webgpu/StorageTexture3DNode.ts

@@ -0,0 +1,30 @@
+/* eslint-disable @typescript-eslint/class-methods-use-this */
+
+import { vec3 } from 'three/tsl'
+import { StorageTextureNode, type Node, type NodeBuilder } from 'three/webgpu'
+
+// WORKAROUND: StorageTextureNode on Storage3DTexture breaks UV.
+// TODO: File a PR in the upstream.
+export class StorageTexture3DNode extends StorageTextureNode {
+  static override get type(): string {
+    return 'StorageTexture3DNode'
+  }
+
+  override getDefaultUV(): Node {
+    return vec3(0.5, 0.5, 0.5)
+  }
+
+  setUpdateMatrix(_value: boolean): void {}
+
+  generateUV(builder: NodeBuilder, uvNode: Node): string {
+    return uvNode.build(builder, this.sampler ? 'vec3' : 'ivec3') as string
+  }
+
+  generateOffset(builder: NodeBuilder, offsetNode: Node): string {
+    return offsetNode.build(builder, 'ivec3') as string
+  }
+}
+
+export const storageTexture3D = (
+  ...args: ConstructorParameters<typeof StorageTexture3DNode>
+): StorageTexture3DNode => new StorageTexture3DNode(...args)