feat: Expose schema read-write APIs (#2411)

Author: aleksanderkatan

apps/typegpu-docs/src/content/docs/apis/buffers.mdx

@@ -236,6 +236,11 @@ If you pass an unmapped buffer, the data will be written to the buffer using `GP
 If you passed your own buffer to the `root.createBuffer` function, you need to ensure it has the `GPUBufferUsage.COPY_DST` usage flag if you want to write to it using the `write` method.
 :::
 
+:::tip
+To write to, patch, or read from an `ArrayBuffer` directly, TypeGPU provides [a dedicated API](/TypeGPU/apis/data-schemas/#arraybuffer-io) for that purpose. 
+The same API is used internally for buffer writes.
+:::
+
 ### Permissive write inputs
 
 `.write()` accepts several equivalent forms — you don't need to construct typed instances:

apps/typegpu-docs/src/content/docs/apis/data-schemas.mdx

@@ -577,3 +577,49 @@ fn myFn(){
   s2.n = 1; // s1 is not modified
 }
 ```
+
+## ArrayBuffer IO
+
+TypeGPU exports functions for reading and writing to a raw `ArrayBuffer`.
+Signatures are analogous to `buffer.write`, `buffer.patch` and `buffer.read`,
+although each one additionally requires a data schema.
+
+```ts twoslash
+import { d, patchArrayBuffer, readFromArrayBuffer, writeToArrayBuffer } from 'typegpu';
+// ---cut---
+const buffer = new ArrayBuffer(16);
+
+// update entire buffer
+writeToArrayBuffer(buffer, d.vec4u, d.vec4u(1, 2, 3, 4));
+```
+
+```ts twoslash
+import { d, patchArrayBuffer, readFromArrayBuffer, writeToArrayBuffer } from 'typegpu';
+// ---cut---
+const buffer = new ArrayBuffer(64);
+const Numbers = d.arrayOf(d.u32, 16);
+
+// update a slice
+const layout = d.memoryLayoutOf(Numbers, (a) => a[4]);
+writeToArrayBuffer(buffer, Numbers, [4, 5, 6, 7], { startOffset: layout.offset });
+```
+
+```ts twoslash
+import { d, patchArrayBuffer, readFromArrayBuffer, writeToArrayBuffer } from 'typegpu';
+// ---cut---
+const buffer = new ArrayBuffer(64);
+const Boid = d.struct({ pos: d.vec2u, id: d.u32 });
+const Boids = d.arrayOf(Boid, 4);
+
+// patch
+patchArrayBuffer(buffer, Boids, { 2: { pos: d.vec2u() } });
+```
+
+```ts twoslash
+import { d, patchArrayBuffer, readFromArrayBuffer, writeToArrayBuffer } from 'typegpu';
+// ---cut---
+const buffer = new ArrayBuffer(64);
+
+// read
+const mat = readFromArrayBuffer(buffer, d.mat4x4f);
+```

packages/typegpu/src/core/buffer/buffer.ts

@@ -1,15 +1,9 @@
-import { BufferReader, BufferWriter, getSystemEndianness } from 'typed-binary';
 import { getCompiledWriter } from '../../data/compiledIO.ts';
-import { readData, writeData } from '../../data/dataIO.ts';
 import type { AnyData } from '../../data/dataTypes.ts';
-import {
-  type WriteInstruction,
-  convertPartialToPatch,
-  getPatchInstructions,
-} from '../../data/partialIO.ts';
+import { convertPartialToPatch, getPatchInstructions } from '../../data/partialIO.ts';
 import { sizeOf } from '../../data/sizeOf.ts';
 import type { BaseData } from '../../data/wgslTypes.ts';
-import { isWgslArray, isWgslData } from '../../data/wgslTypes.ts';
+import { isWgslData } from '../../data/wgslTypes.ts';
 import type { StorageFlag } from '../../extension.ts';
 import type { TgpuNamable } from '../../shared/meta.ts';
 import { getName, setName } from '../../shared/meta.ts';
@@ -37,8 +31,8 @@ import {
   type TgpuFixedBufferUsage,
   uniform,
 } from './bufferUsage.ts';
-import { alignmentOf } from '../../data/alignmentOf.ts';
-import { roundUp } from '../../mathUtils.ts';
+import { calculateOffsets, readFromArrayBuffer, writeToArrayBuffer } from '../../data/dataIO.ts';
+import { patchArrayBuffer } from '../../data/partialIO.ts';
 
 // ----------
 // Public API
@@ -193,8 +187,6 @@ export function isUsableAsIndex<T extends TgpuBuffer<BaseData>>(
 // --------------
 // Implementation
 // --------------
-const endianness = getSystemEndianness();
-
 class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
   readonly [$internal] = true;
   readonly resourceType = 'buffer';
@@ -261,7 +253,7 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
         if (this.#initialCallback) {
           this.#initialCallback(this);
         } else if (this.initial) {
-          this.#writeToTarget(this.#getMappedRange(), this.initial);
+          writeToArrayBuffer(this.#getMappedRange(), this.dataType, this.initial);
         }
         this.#unmapBuffer();
       }
@@ -354,77 +346,10 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
     getCompiledWriter(this.dataType);
   }
 
-  #writeToTarget(
-    target: ArrayBuffer,
-    data: InferInput<TData> | ArrayBuffer,
-    options?: BufferWriteOptions,
-  ): void {
-    const startOffset = options?.startOffset ?? 0;
-    const endOffset = options?.endOffset ?? target.byteLength;
-
-    // Fast path: raw byte copy, user guarantees the padded layout
-    if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) {
-      const src =
-        data instanceof ArrayBuffer
-          ? new Uint8Array(data)
-          : new Uint8Array(data.buffer, data.byteOffset, data.byteLength);
-      const regionSize = endOffset - startOffset;
-      if (src.byteLength !== regionSize) {
-        console.warn(
-          `Buffer size mismatch: expected ${regionSize} bytes, got ${src.byteLength}. ` +
-            (src.byteLength < regionSize ? 'Data truncated.' : 'Excess ignored.'),
-        );
-      }
-      const copyLen = Math.min(src.byteLength, regionSize);
-      new Uint8Array(target).set(src.subarray(0, copyLen), startOffset);
-      return;
-    }
-
-    const dataView = new DataView(target);
-    const isLittleEndian = endianness === 'little';
-
-    const compiledWriter = getCompiledWriter(this.dataType);
-
-    if (compiledWriter) {
-      try {
-        compiledWriter(dataView, startOffset, data, isLittleEndian, endOffset);
-        return;
-      } catch (error) {
-        console.error(
-          `Error when using compiled writer for buffer ${
-            getName(this) ?? '<unnamed>'
-          } - this is likely a bug, please submit an issue at https://github.com/software-mansion/TypeGPU/issues\nUsing fallback writer instead.`,
-          error,
-        );
-      }
-    }
-
-    const writer = new BufferWriter(target);
-    writer.seekTo(startOffset);
-    writeData(writer, this.dataType, data as Infer<TData>);
-  }
-
   write(data: InferInput<TData>, options?: BufferWriteOptions): void;
   write(data: ArrayBuffer, options?: BufferWriteOptions): void;
   write(data: InferInput<TData> | ArrayBuffer, options?: BufferWriteOptions): void {
     const gpuBuffer = this.buffer;
-    const bufferSize = sizeOf(this.dataType);
-    const startOffset = options?.startOffset ?? 0;
-
-    let naturalSize: number | undefined = undefined;
-    if (isWgslArray(this.dataType) && Array.isArray(data)) {
-      const arrayData = data as unknown[];
-      naturalSize =
-        arrayData.length *
-        roundUp(sizeOf(this.dataType.elementType), alignmentOf(this.dataType.elementType));
-    } else if (ArrayBuffer.isView(data) || data instanceof ArrayBuffer) {
-      naturalSize = data.byteLength;
-    }
-    const naturalEndOffset =
-      naturalSize !== undefined ? Math.min(startOffset + naturalSize, bufferSize) : undefined;
-
-    const endOffset = options?.endOffset ?? naturalEndOffset ?? bufferSize;
-    const size = endOffset - startOffset;
 
     if (gpuBuffer.mapState === 'mapped') {
       const mapped = this.#getMappedRange();
@@ -433,43 +358,34 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
         // via arrayBuffer. Nothing to do here
         return;
       }
-      this.#writeToTarget(mapped, data, options);
+      writeToArrayBuffer(mapped, this.dataType, data, options);
       return;
     }
 
     // If the caller already wrote directly into #hostBuffer via
     // arrayBuffer, skip the redundant copy, the data is already in place.
     if (!(data instanceof ArrayBuffer && data === this.#hostBuffer)) {
-      this.#writeToTarget(this.#hostBuffer, data, options);
+      writeToArrayBuffer(this.#hostBuffer, this.dataType, data, options);
     }
+
+    const { startOffset, endOffset } = calculateOffsets(options, this.dataType, data);
+    const size = endOffset - startOffset;
+
     this.#device.queue.writeBuffer(gpuBuffer, startOffset, this.#hostBuffer, startOffset, size);
   }
 
   /** @deprecated Use {@link patch} instead. */
   public writePartial(data: InferPartial<TData>): void {
-    this.#applyInstructions(
-      getPatchInstructions(
-        this.dataType,
-        convertPartialToPatch(this.dataType, data),
-        this.#hostBuffer,
-      ),
-    );
+    this.patch(convertPartialToPatch(this.dataType, data) as InferPatch<TData>);
   }
 
   public patch(data: InferPatch<TData>): void {
-    this.#applyInstructions(getPatchInstructions(this.dataType, data, this.#hostBuffer));
-  }
-
-  #applyInstructions(instructions: WriteInstruction[]): void {
     const gpuBuffer = this.buffer;
 
     if (gpuBuffer.mapState === 'mapped') {
-      const mappedRange = this.#getMappedRange();
-      const mappedView = new Uint8Array(mappedRange);
-      for (const { data, gpuOffset } of instructions) {
-        mappedView.set(data, gpuOffset);
-      }
+      patchArrayBuffer(this.#getMappedRange(), this.dataType, data);
     } else {
+      const instructions = getPatchInstructions(this.dataType, data, this.#hostBuffer);
       for (const { data, gpuOffset } of instructions) {
         this.#device.queue.writeBuffer(gpuBuffer, gpuOffset, data);
       }
@@ -504,14 +420,12 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
     const gpuBuffer = this.buffer;
 
     if (gpuBuffer.mapState === 'mapped') {
-      const mapped = this.#getMappedRange();
-      return readData(new BufferReader(mapped), this.dataType);
+      return readFromArrayBuffer(this.#getMappedRange(), this.dataType);
     }
 
     if (gpuBuffer.usage & GPUBufferUsage.MAP_READ) {
       await gpuBuffer.mapAsync(GPUMapMode.READ);
-      const mapped = this.#getMappedRange();
-      const res = readData(new BufferReader(mapped), this.dataType);
+      const res = readFromArrayBuffer(this.#getMappedRange(), this.dataType);
       this.#unmapBuffer();
       return res;
     }
@@ -527,7 +441,7 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
     this.#device.queue.submit([commandEncoder.finish()]);
     await stagingBuffer.mapAsync(GPUMapMode.READ, 0, sizeOf(this.dataType));
 
-    const res = readData(new BufferReader(stagingBuffer.getMappedRange()), this.dataType);
+    const res = readFromArrayBuffer(stagingBuffer.getMappedRange(), this.dataType);
 
     stagingBuffer.unmap();
     stagingBuffer.destroy();

packages/typegpu/src/data/dataIO.ts

@@ -1,5 +1,11 @@
-import type { ISerialInput, ISerialOutput } from 'typed-binary';
-import type { Infer, InferRecord } from '../shared/repr.ts';
+import {
+  BufferReader,
+  BufferWriter,
+  getSystemEndianness,
+  type ISerialInput,
+  type ISerialOutput,
+} from 'typed-binary';
+import type { Infer, InferInput, InferRecord } from '../shared/repr.ts';
 import alignIO from './alignIO.ts';
 import { alignmentOf, customAlignmentOf } from './alignmentOf.ts';
 import type { AnyConcreteData, AnyData, Disarray, LooseDecorated, Unstruct } from './dataTypes.ts';
@@ -20,6 +26,11 @@ import {
   vec4u,
 } from './vector.ts';
 import type * as wgsl from './wgslTypes.ts';
+import { isWgslArray, type BaseData } from './wgslTypes.ts';
+import type { BufferWriteOptions } from '../core/buffer/buffer.ts';
+import { getCompiledWriter } from './compiledIO.ts';
+import { getName } from '../shared/meta.ts';
+import { roundUp } from '../mathUtils.ts';
 
 type DataWriter<TSchema extends wgsl.BaseData> = (
   output: ISerialOutput,
@@ -788,3 +799,82 @@ export function readData<TData extends wgsl.BaseData>(
 
   return reader(input, schema);
 }
+
+const endianness = getSystemEndianness();
+
+export function calculateOffsets<T extends BaseData>(
+  options: BufferWriteOptions | undefined,
+  schema: T,
+  data: InferInput<T> | ArrayBuffer,
+): { startOffset: number; endOffset: number } {
+  const bufferSize = sizeOf(schema);
+  const startOffset = options?.startOffset ?? 0;
+  let naturalSize: number | undefined = undefined;
+  if (isWgslArray(schema) && Array.isArray(data)) {
+    const arrayData = data as unknown[];
+    naturalSize =
+      arrayData.length * roundUp(sizeOf(schema.elementType), alignmentOf(schema.elementType));
+  } else if (ArrayBuffer.isView(data) || data instanceof ArrayBuffer) {
+    naturalSize = data.byteLength;
+  }
+  const naturalEndOffset =
+    naturalSize !== undefined ? Math.min(startOffset + naturalSize, bufferSize) : undefined;
+
+  const endOffset = options?.endOffset ?? naturalEndOffset ?? bufferSize;
+
+  return { startOffset, endOffset };
+}
+
+export function writeToArrayBuffer<T extends BaseData>(
+  buffer: ArrayBuffer,
+  schema: T,
+  data: InferInput<T> | ArrayBuffer,
+  options?: BufferWriteOptions,
+) {
+  const { startOffset, endOffset } = calculateOffsets(options, schema, data);
+
+  // Fast path: raw byte copy, user guarantees the padded layout
+  if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) {
+    const src =
+      data instanceof ArrayBuffer
+        ? new Uint8Array(data)
+        : new Uint8Array(data.buffer, data.byteOffset, data.byteLength);
+    const regionSize = endOffset - startOffset;
+    if (src.byteLength !== regionSize) {
+      console.warn(
+        `Buffer size mismatch: expected ${regionSize} bytes, got ${src.byteLength}. ` +
+          (src.byteLength < regionSize ? 'Data truncated.' : 'Excess ignored.'),
+      );
+    }
+    const copyLen = Math.min(src.byteLength, regionSize);
+    new Uint8Array(buffer).set(src.subarray(0, copyLen), startOffset);
+    return;
+  }
+
+  const dataView = new DataView(buffer);
+  const isLittleEndian = endianness === 'little';
+
+  const compiledWriter = getCompiledWriter(schema);
+
+  if (compiledWriter) {
+    try {
+      compiledWriter(dataView, startOffset, data, isLittleEndian, endOffset);
+      return;
+    } catch (error) {
+      console.error(
+        `Error when using compiled writer for data type '${
+          schema.type
+        }' (${getName(schema) ?? 'unnamed'}) - this is likely a bug, please submit an issue at https://github.com/software-mansion/TypeGPU/issues\nUsing fallback writer instead.`,
+        error,
+      );
+    }
+  }
+
+  const writer = new BufferWriter(buffer);
+  writer.seekTo(startOffset);
+  writeData(writer, schema, data as Infer<T>);
+}
+
+export function readFromArrayBuffer<T extends BaseData>(buffer: ArrayBuffer, schema: T): Infer<T> {
+  return readData(new BufferReader(buffer), schema);
+}