feat(@typegpu/react): Implement useUniformValue hook (#1694)

Author: piaccho

apps/typegpu-docs/src/content/examples/react/triangle/index.html

@@ -1 +1 @@
-<div id='example-app'></div>
+<div id="example-app"></div>

apps/typegpu-docs/src/content/examples/react/triangle/index.tsx

@@ -1,17 +1,23 @@
-import { vec4f } from 'typegpu/data';
-import { useRender } from '@typegpu/react';
+import * as d from 'typegpu/data';
+import { useFrame, useRender, useUniformValue } from '@typegpu/react';
+import { hsvToRgb } from '@typegpu/color';
 
 function App() {
+  const time = useUniformValue(d.f32, 0);
+
+  useFrame(() => {
+    time.value = performance.now() / 1000;
+  });
+
   const { ref } = useRender({
-    fragment: ({ uv }) => {
+    fragment: () => {
       'kernel';
-      return vec4f(uv.x, uv.y, 1, 1);
+      const t = time.$;
+      const rgb = hsvToRgb(d.vec3f(t * 0.5, 1, 1));
+      return d.vec4f(rgb, 1);
     },
   });
 
-  // TODO: Provide a time variable to the shader with useUniformValue
-  // TODO: Make the gradient shift colors over time using hsvToRgb from @typegpu/color
-
   return (
     <main>
       <canvas ref={ref} width='256' height='256' />
@@ -28,7 +34,7 @@ const reactRoot = createRoot(
 reactRoot.render(<App />);
 
 export function onCleanup() {
-  reactRoot.unmount();
+  setTimeout(() => reactRoot.unmount(), 0);
 }
 
 // #endregion

apps/typegpu-docs/vitest.config.mts

@@ -1,5 +1,5 @@
 import { createJiti } from 'jiti';
-import { defineConfig } from 'vitest/config';
+import { defineConfig, type Plugin } from 'vitest/config';
 import { imagetools } from 'vite-imagetools';
 import type TypeGPUPlugin from 'unplugin-typegpu/vite';
 
@@ -13,7 +13,7 @@ export default defineConfig({
   plugins: [
     typegpu({ include: [/\.m?[jt]sx?/] }),
     /** @type {any} */ imagetools(),
-  ],
+  ] as Plugin[],
   server: {
     proxy: {
       '/TypeGPU': {

packages/typegpu-react/src/use-uniform-value.ts

@@ -1,13 +1,62 @@
 import type * as d from 'typegpu/data';
+import { useRoot } from './root-context.tsx';
+import { useEffect, useMemo, useRef, useState } from 'react';
+import type { ValidateUniformSchema } from 'typegpu';
 
 interface UniformValue<TSchema, TValue extends d.Infer<TSchema>> {
   schema: TSchema;
-  value: TValue;
+  value: TValue | undefined;
+  readonly $: d.InferGPU<TSchema>;
 }
 
-export function useUniformValue<TSchema, TValue extends d.Infer<TSchema>>(
-  schema: d.AnyWgslData,
+export function useUniformValue<
+  TSchema extends d.AnyWgslData,
+  TValue extends d.Infer<TSchema>,
+>(
+  schema: ValidateUniformSchema<TSchema>,
   initialValue?: TValue | undefined,
 ): UniformValue<TSchema, TValue> {
-  // TODO: Implement
+  const root = useRoot();
+
+  const [uniformBuffer] = useState(() => {
+    return root.createUniform(
+      schema,
+      initialValue,
+    );
+  });
+
+  const cleanupRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+  useEffect(() => {
+    if (cleanupRef.current) {
+      clearTimeout(cleanupRef.current);
+    }
+
+    return () => {
+      cleanupRef.current = setTimeout(() => {
+        uniformBuffer.buffer.destroy();
+      }, 200);
+    };
+  }, [uniformBuffer]);
+
+  // biome-ignore lint/correctness/useExhaustiveDependencies: This value needs to be stable
+  const uniformValue = useMemo(() => {
+    let currentValue = initialValue;
+    return {
+      schema,
+      get value() {
+        return currentValue;
+      },
+      set value(newValue: TValue | undefined) {
+        currentValue = newValue;
+        if (newValue !== undefined) {
+          uniformBuffer.write(newValue);
+        }
+      },
+      get $() {
+        return uniformBuffer.$;
+      },
+    };
+  }, []);
+
+  return uniformValue as UniformValue<TSchema, TValue>;
 }

packages/typegpu/src/index.ts

@@ -94,6 +94,8 @@ export type {
   Configurable,
   TgpuRoot,
   ValidateBufferSchema,
+  ValidateStorageSchema,
+  ValidateUniformSchema,
   WithBinding,
   WithCompute,
   WithFragment,