feat(@typegpu/react): Add useMirroredUniform hook

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

@@ -1,10 +1,43 @@
 import * as d from 'typegpu/data';
-import { useFrame, useRender, useUniformValue } from '@typegpu/react';
+import {
+  useFrame,
+  useMirroredUniform,
+  useRender,
+  useUniformValue,
+} from '@typegpu/react';
 import { hsvToRgb } from '@typegpu/color';
 
+// TODO: We can come up with a more sophisticated example later
+function ColorBox(props: { color: d.v3f }) {
+  const color = useMirroredUniform(d.vec3f, props.color);
+
+  const { ref } = useRender({
+    fragment: () => {
+      'kernel';
+      return d.vec4f(color.$, 1);
+    },
+  });
+
+  return (
+    <canvas
+      ref={ref}
+      width='32'
+      height='32'
+      style={{ border: '2px black solid' }}
+    />
+  );
+}
+
+let randomizeColor: () => void;
+
 function App() {
+  const [currentColor, setCurrentColor] = useState(d.vec3f(1, 0, 0));
   const time = useUniformValue(d.f32, 0);
 
+  randomizeColor = () => {
+    setCurrentColor(d.vec3f(Math.random(), Math.random(), Math.random()));
+  };
+
   useFrame(() => {
     time.value = performance.now() / 1000;
   });
@@ -14,25 +47,36 @@ function App() {
       'kernel';
       const t = time.$;
       const rgb = hsvToRgb(d.vec3f(t * 0.5, 1, 1));
+
       return d.vec4f(rgb, 1);
     },
   });
 
   return (
-    <main>
+    <main style={{ display: 'flex', alignItems: 'center', gap: 32 }}>
       <canvas ref={ref} width='256' height='256' />
+      <ColorBox color={currentColor} />
     </main>
   );
 }
 
 // #region Example controls and cleanup
 
 import { createRoot } from 'react-dom/client';
+import { useState } from 'react';
 const reactRoot = createRoot(
   document.getElementById('example-app') as HTMLDivElement,
 );
 reactRoot.render(<App />);
 
+export const controls = {
+  'Randomize box color': {
+    onButtonClick: () => {
+      randomizeColor();
+    },
+  },
+};
+
 export function onCleanup() {
   setTimeout(() => reactRoot.unmount(), 0);
 }

packages/typegpu-react/package.json

@@ -30,13 +30,16 @@
   "keywords": [],
   "license": "MIT",
   "peerDependencies": {
-    "typegpu": "^0.7.0",
-    "react": "^19.0.0"
+    "react": "^19.0.0",
+    "typegpu": "^0.7.0"
   },
   "devDependencies": {
+    "@testing-library/dom": "^10.4.1",
+    "@testing-library/react": "^16.3.0",
     "@typegpu/tgpu-dev-cli": "workspace:*",
+    "@types/react": "^19.1.8",
+    "@types/react-dom": "^19.1.6",
     "@webgpu/types": "catalog:types",
-    "@types/react": "^19.0.0",
     "tsdown": "catalog:build",
     "typegpu": "workspace:*",
     "typescript": "catalog:types",

packages/typegpu-react/src/index.ts

@@ -1,3 +1,4 @@
 export { useFrame } from './use-frame.ts';
 export { useRender } from './use-render.ts';
 export { useUniformValue } from './use-uniform-value.ts';
+export { useMirroredUniform } from './use-mirrored-uniform.ts';

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

@@ -0,0 +1,69 @@
+import * as d from 'typegpu/data';
+import { useRoot } from './root-context.tsx';
+import { useEffect, useMemo, useRef, useState } from 'react';
+import type { ValidateUniformSchema } from 'typegpu';
+
+interface MirroredValue<TSchema> {
+  schema: TSchema;
+  readonly $: d.InferGPU<TSchema>;
+}
+
+export function useMirroredUniform<
+  TSchema extends d.AnyWgslData,
+  TValue extends d.Infer<TSchema>,
+>(
+  schema: ValidateUniformSchema<TSchema>,
+  value: TValue,
+): MirroredValue<TSchema> {
+  const root = useRoot();
+  const [uniformBuffer, setUniformBuffer] = useState(() => {
+    return root.createUniform(schema, value);
+  });
+  const prevSchemaRef = useRef(schema);
+  const currentSchemaRef = useRef(schema);
+  const cleanupRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+
+  useEffect(() => {
+    if (cleanupRef.current) {
+      clearTimeout(cleanupRef.current);
+    }
+
+    return () => {
+      cleanupRef.current = setTimeout(() => {
+        uniformBuffer.buffer.destroy();
+      }, 200);
+    };
+  }, [uniformBuffer]);
+
+  useEffect(() => {
+    if (!d.deepEqual(prevSchemaRef.current as d.AnyData, schema as d.AnyData)) {
+      uniformBuffer.buffer.destroy();
+      setUniformBuffer(root.createUniform(schema, value));
+      prevSchemaRef.current = schema;
+    } else {
+      uniformBuffer.write(value);
+    }
+  }, [schema, value, root, uniformBuffer]);
+
+  if (
+    !d.deepEqual(currentSchemaRef.current as d.AnyData, schema as d.AnyData)
+  ) {
+    currentSchemaRef.current = schema;
+  }
+
+  // Using current schema ref instead of schema directly
+  // to prevent unnecessary re-memoization when schema object
+  // reference changes but content is structurally equivalent.
+  // biome-ignore lint/correctness/useExhaustiveDependencies: This value needs to be stable
+  const mirroredValue = useMemo(
+    () => ({
+      schema,
+      get $() {
+        return uniformBuffer.$;
+      },
+    }),
+    [currentSchemaRef.current, uniformBuffer],
+  );
+
+  return mirroredValue as MirroredValue<TSchema>;
+}