feat: Introduce Index buffers and improve buffer UX

apps/typegpu-docs/src/content/examples/simple/square/index.html

@@ -0,0 +1 @@
+<canvas></canvas>

apps/typegpu-docs/src/content/examples/simple/square/index.ts

@@ -0,0 +1,128 @@
+import tgpu from 'typegpu';
+import * as d from 'typegpu/data';
+
+const presentationFormat = navigator.gpu.getPreferredCanvasFormat();
+const canvas = document.querySelector('canvas') as HTMLCanvasElement;
+const context = canvas.getContext('webgpu') as GPUCanvasContext;
+
+const root = await tgpu.init();
+
+context.configure({
+  device: root.device,
+  format: presentationFormat,
+  alphaMode: 'premultiplied',
+});
+
+const colors = {
+  bottomLeft: d.vec4f(1, 0, 0, 1),
+  bottomRight: d.vec4f(0, 1, 0, 1),
+  topRight: d.vec4f(0, 0, 1, 1),
+  topLeft: d.vec4f(1, 1, 0, 1),
+};
+
+const colorIndices = {
+  bottomLeft: 0,
+  bottomRight: 1,
+  topRight: 2,
+  topLeft: 3,
+} as const;
+
+const colorBuffer = root
+  .createBuffer(d.arrayOf(d.vec4f, 4), Object.values(colors))
+  .$usage('vertex');
+const vertexLayout = tgpu.vertexLayout((n) => d.arrayOf(d.vec4f, n));
+
+const vertex = tgpu['~unstable'].vertexFn({
+  in: {
+    idx: d.builtin.vertexIndex,
+    color: d.vec4f,
+  },
+  out: {
+    color: d.vec4f,
+    pos: d.builtin.position,
+  },
+})(({ idx, color }) => {
+  const vertices = [
+    d.vec2f(-1, -1),
+    d.vec2f(1, -1),
+    d.vec2f(1, 1),
+    d.vec2f(-1, 1),
+  ];
+  return {
+    color,
+    pos: d.vec4f(vertices[idx], 0, 1),
+  };
+});
+
+const mainFragment = tgpu['~unstable'].fragmentFn({
+  in: {
+    color: d.vec4f,
+  },
+  out: d.vec4f,
+})((input) => input.color);
+
+const indexBuffer = root.createBuffer(d.arrayOf(d.u16, 6), [0, 2, 1, 0, 3, 2])
+  .$usage('index');
+
+const pipeline = root['~unstable']
+  .withVertex(vertex, { color: vertexLayout.attrib })
+  .withFragment(mainFragment, { format: presentationFormat })
+  .createPipeline()
+  .withIndexBuffer(indexBuffer);
+
+function render() {
+  pipeline
+    .with(vertexLayout, colorBuffer)
+    .withColorAttachment({
+      view: context.getCurrentTexture().createView(),
+      loadOp: 'clear',
+      storeOp: 'store',
+    })
+    .drawIndexed(6);
+}
+render();
+
+// #region Example controls & Cleanup
+
+function updateColor(
+  color: readonly [number, number, number],
+  position: keyof typeof colors,
+): void {
+  colors[position] = d.vec4f(color[0], color[1], color[2], 1);
+  const idx = colorIndices[position];
+  colorBuffer.writePartial([
+    {
+      idx,
+      value: colors[position],
+    },
+  ]);
+  render();
+}
+
+export const controls = {
+  topLeft: {
+    onColorChange: (value: readonly [number, number, number]) =>
+      updateColor(value, 'topLeft'),
+    initial: [...colors.topLeft.xyz],
+  },
+  topRight: {
+    onColorChange: (value: readonly [number, number, number]) =>
+      updateColor(value, 'topRight'),
+    initial: [...colors.topRight.xyz],
+  },
+  bottomLeft: {
+    onColorChange: (value: readonly [number, number, number]) =>
+      updateColor(value, 'bottomLeft'),
+    initial: [...colors.bottomLeft.xyz],
+  },
+  bottomRight: {
+    onColorChange: (value: readonly [number, number, number]) =>
+      updateColor(value, 'bottomRight'),
+    initial: [...colors.bottomRight.xyz],
+  },
+};
+
+export function onCleanup() {
+  root.destroy();
+}
+// #endregion

apps/typegpu-docs/src/content/examples/simple/square/meta.json

@@ -0,0 +1,5 @@
+{
+  "title": "Square",
+  "category": "simple",
+  "tags": ["experimental", "index buffer"]
+}

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

@@ -4,7 +4,6 @@ import {
   getCompiledWriterForSchema,
 } from '../../data/compiledIO.ts';
 import { readData, writeData } from '../../data/dataIO.ts';
-import type { AnyData } from '../../data/dataTypes.ts';
 import { getWriteInstructions } from '../../data/partialIO.ts';
 import { sizeOf } from '../../data/sizeOf.ts';
 import type { BaseData, WgslTypeLiteral } from '../../data/wgslTypes.ts';
@@ -25,6 +24,7 @@ import {
   type TgpuBufferUniform,
   type TgpuFixedBufferUsage,
 } from './bufferUsage.ts';
+import type { AnyData, UnwrapDecorated } from '../../data/dataTypes.ts';
 
 // ----------
 // Public API
@@ -43,16 +43,21 @@ export interface VertexFlag {
   usableAsVertex: true;
 }
 
+export interface IndexFlag {
+  usableAsIndex: true;
+}
+
 /**
  * @deprecated Use VertexFlag instead.
  */
 export type Vertex = VertexFlag;
 
-type LiteralToUsageType<T extends 'uniform' | 'storage' | 'vertex'> = T extends
-  'uniform' ? UniformFlag
-  : T extends 'storage' ? StorageFlag
-  : T extends 'vertex' ? VertexFlag
-  : never;
+type LiteralToUsageType<T extends 'uniform' | 'storage' | 'vertex' | 'index'> =
+  T extends 'uniform' ? UniformFlag
+    : T extends 'storage' ? StorageFlag
+    : T extends 'vertex' ? VertexFlag
+    : T extends 'index' ? IndexFlag
+    : never;
 
 type ViewUsages<TBuffer extends TgpuBuffer<BaseData>> =
   | (boolean extends TBuffer['usableAsUniform'] ? never : 'uniform')
@@ -71,6 +76,17 @@ const usageToUsageConstructor = {
   readonly: asReadonly,
 };
 
+type IsIndexCompatible<TData extends BaseData> = UnwrapDecorated<TData> extends
+  {
+    readonly type: 'array';
+    readonly elementType: infer TElement;
+  }
+  ? TElement extends BaseData
+    ? UnwrapDecorated<TElement> extends { readonly type: 'u32' | 'u16' } ? true
+    : false
+  : false
+  : false;
+
 export interface TgpuBuffer<TData extends BaseData> extends TgpuNamable {
   readonly resourceType: 'buffer';
   readonly dataType: TData;
@@ -82,8 +98,9 @@ export interface TgpuBuffer<TData extends BaseData> extends TgpuNamable {
   usableAsUniform: boolean;
   usableAsStorage: boolean;
   usableAsVertex: boolean;
+  usableAsIndex: boolean;
 
-  $usage<T extends RestrictVertexUsages<TData>>(
+  $usage<T extends RestrictUsages<TData>>(
     ...usages: T
   ): this & UnionToIntersection<LiteralToUsageType<T[number]>>;
   $addFlags(flags: GPUBufferUsageFlags): this;
@@ -109,6 +126,7 @@ export function INTERNAL_createBuffer<TData extends AnyData>(
       'uniform',
     ]);
   }
+
   return new TgpuBufferImpl(group, typeSchema, initialOrBuffer);
 }
 
@@ -124,15 +142,39 @@ export function isUsableAsVertex<T extends TgpuBuffer<AnyData>>(
   return !!(buffer as unknown as VertexFlag).usableAsVertex;
 }
 
+export function isUsableAsIndex<T extends TgpuBuffer<AnyData>>(
+  buffer: T,
+): buffer is T & IndexFlag {
+  return !!(buffer as unknown as IndexFlag).usableAsIndex;
+}
+
 // --------------
 // Implementation
 // --------------
 const endianness = getSystemEndianness();
 
-type RestrictVertexUsages<TData extends BaseData> = TData extends {
+type IsArrayOfU32<TData extends BaseData> = UnwrapDecorated<TData> extends {
+  readonly type: 'array';
+  readonly elementType: infer TElement;
+}
+  ? TElement extends BaseData
+    ? UnwrapDecorated<TElement> extends { readonly type: 'u32' } ? true
+    : false
+  : false
+  : false;
+
+type IsWgslLiteral<TData extends BaseData> = TData extends {
   readonly type: WgslTypeLiteral;
-} ? ('uniform' | 'storage' | 'vertex')[]
-  : 'vertex'[];
+} ? true
+  : false;
+
+type RestrictUsages<TData extends BaseData> = IsIndexCompatible<TData> extends
+  true
+  ? IsArrayOfU32<TData> extends true
+    ? ('uniform' | 'storage' | 'vertex' | 'index')[]
+  : ['index']
+  : IsWgslLiteral<TData> extends true ? ('uniform' | 'storage' | 'vertex')[]
+  : ['vertex'];
 
 class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
   public readonly resourceType = 'buffer';
@@ -148,12 +190,14 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
   usableAsUniform = false;
   usableAsStorage = false;
   usableAsVertex = false;
+  usableAsIndex = false;
 
   constructor(
     private readonly _group: ExperimentalTgpuRoot,
     public readonly dataType: TData,
     public readonly initialOrBuffer?: Infer<TData> | GPUBuffer | undefined,
-    private readonly _disallowedUsages?: ('uniform' | 'storage' | 'vertex')[],
+    private readonly _disallowedUsages?:
+      ('uniform' | 'storage' | 'vertex' | 'index')[],
   ) {
     if (isGPUBuffer(initialOrBuffer)) {
       this._ownBuffer = false;
@@ -201,7 +245,7 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
     return this;
   }
 
-  $usage<T extends RestrictVertexUsages<TData>>(
+  $usage<T extends ('uniform' | 'storage' | 'vertex' | 'index')[]>(
     ...usages: T
   ): this & UnionToIntersection<LiteralToUsageType<T[number]>> {
     for (const usage of usages) {
@@ -210,12 +254,15 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
           `Buffer of type ${this.dataType.type} cannot be used as ${usage}`,
         );
       }
+
       this.flags |= usage === 'uniform' ? GPUBufferUsage.UNIFORM : 0;
       this.flags |= usage === 'storage' ? GPUBufferUsage.STORAGE : 0;
       this.flags |= usage === 'vertex' ? GPUBufferUsage.VERTEX : 0;
+      this.flags |= usage === 'index' ? GPUBufferUsage.INDEX : 0;
       this.usableAsUniform = this.usableAsUniform || usage === 'uniform';
       this.usableAsStorage = this.usableAsStorage || usage === 'storage';
       this.usableAsVertex = this.usableAsVertex || usage === 'vertex';
+      this.usableAsIndex = this.usableAsIndex || usage === 'index';
     }
     return this as this & UnionToIntersection<LiteralToUsageType<T[number]>>;
   }

packages/typegpu/src/core/pipeline/computePipeline.ts

@@ -29,8 +29,7 @@ interface ComputePipelineInternals {
 // Public API
 // ----------
 
-export interface TgpuComputePipeline
-  extends TgpuNamable, Timeable<TgpuComputePipeline> {
+export interface TgpuComputePipeline extends TgpuNamable, Timeable {
   readonly [$internal]: ComputePipelineInternals;
   readonly resourceType: 'compute-pipeline';
 
@@ -117,26 +116,26 @@ class TgpuComputePipelineImpl implements TgpuComputePipeline {
 
   withPerformanceCallback(
     callback: (start: bigint, end: bigint) => void | Promise<void>,
-  ): TgpuComputePipeline {
+  ): this {
     const newPriors = createWithPerformanceCallback(
       this._priors,
       callback,
       this._core.branch,
     );
-    return new TgpuComputePipelineImpl(this._core, newPriors);
+    return new TgpuComputePipelineImpl(this._core, newPriors) as this;
   }
 
   withTimestampWrites(options: {
     querySet: TgpuQuerySet<'timestamp'> | GPUQuerySet;
     beginningOfPassWriteIndex?: number;
     endOfPassWriteIndex?: number;
-  }): TgpuComputePipeline {
+  }): this {
     const newPriors = createWithTimestampWrites(
       this._priors,
       options,
       this._core.branch,
     );
-    return new TgpuComputePipelineImpl(this._core, newPriors);
+    return new TgpuComputePipelineImpl(this._core, newPriors) as this;
   }
 
   dispatchWorkgroups(