chore: Shaderbits

Author: iwoplaza

packages/typegpu/src/data/snippet.ts

@@ -2,6 +2,7 @@ import { undecorate } from './dataTypes.ts';
 import type { UnknownData } from './dataTypes.ts';
 import { DEV } from '../shared/env.ts';
 import { type BaseData, isNumericSchema } from './wgslTypes.ts';
+import { reduceIfCompound, type TypeBit } from '../shaderbit/typeBits.ts';
 
 export type Origin =
   | 'uniform'
@@ -88,10 +89,10 @@ export type MapValueToSnippet<T> = { [K in keyof T]: Snippet };
 
 class SnippetImpl implements Snippet {
   readonly value: unknown;
-  readonly dataType: BaseData | UnknownData;
+  readonly dataType: TypeBit | UnknownData;
   readonly origin: Origin;
 
-  constructor(value: unknown, dataType: BaseData | UnknownData, origin: Origin) {
+  constructor(value: unknown, dataType: TypeBit | UnknownData, origin: Origin) {
     this.value = value;
     this.dataType = dataType;
     this.origin = origin;
@@ -121,7 +122,7 @@ export function snip(
   return new SnippetImpl(
     value,
     // We don't care about attributes in snippet land, so we discard that information.
-    undecorate(dataType as BaseData),
+    reduceIfCompound(dataType as BaseData),
     origin,
   );
 }

packages/typegpu/src/data/wgslTypes.ts

@@ -39,10 +39,11 @@ import type { WgslExternalTexture, WgslStorageTexture, WgslTexture } from './tex
 import type { WgslComparisonSampler, WgslSampler } from './sampler.ts';
 import type { _ref as ref } from './ref.ts';
 import type { DualFn } from '../types.ts';
+import type { CompoundBit, TypeBit } from '../shaderbit/typeBits.ts';
 
 type DecoratedLocation<T extends BaseData> = Decorated<T, Location[]>;
 
-export interface BaseData {
+export interface BaseData extends CompoundBit<TypeBit> {
   readonly [$internal]: Record<string, unknown>;
   readonly type: string;
   readonly [$repr]: unknown;

packages/typegpu/src/shaderbit/typeBits.ts

@@ -0,0 +1,67 @@
+export interface CompoundBit<T> {
+  '~shaderbit:wgsl': {
+    reduce(): T;
+  };
+}
+
+export interface AttributeBit {
+  name: string;
+  params: (string | number)[][];
+}
+
+export interface StructPropertyBit<
+  TType extends TypeBit = TypeBit,
+  TAttribs extends AttributeBit[] = AttributeBit[],
+> {
+  type: TType;
+  attribs: TAttribs;
+}
+
+export interface StructBit {
+  kind: 'struct';
+  fields: Record<string, StructPropertyBit>;
+}
+
+export interface ArrayBit {
+  kind: 'array';
+  elem: TypeBit;
+  count: number;
+}
+
+export type TypeBit =
+  | 'bool'
+  | 'f32'
+  | 'f16'
+  | 'i32'
+  | 'u32'
+  | 'vec2f'
+  | 'vec2h'
+  | 'vec2i'
+  | 'vec2u'
+  | 'vec2b'
+  | 'vec3f'
+  | 'vec3h'
+  | 'vec3i'
+  | 'vec3u'
+  | 'vec3b'
+  | 'vec4f'
+  | 'vec4h'
+  | 'vec4i'
+  | 'vec4u'
+  | 'vec4b'
+  | 'mat2x2f'
+  | 'mat3x3f'
+  | 'mat4x4f'
+  | StructBit
+  | ArrayBit;
+
+export function isCompoundBit(bit: unknown): bit is CompoundBit<unknown> {
+  return !!(bit as CompoundBit<unknown>)['~shaderbit:wgsl'];
+}
+
+export function reduceIfCompound<T>(bit: CompoundBit<T> | T): T {
+  if (isCompoundBit(bit)) {
+    return bit['~shaderbit:wgsl'].reduce();
+  }
+  return bit as T;
+}