feat: Allow null at comptime

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

@@ -78,13 +78,14 @@ since we treat `tgpu.const` as a way to opt-out of inlining.
 Branch pruning also works for ternary operators when the condition is known at resolution time.
 
 ```ts twoslash
-import { tgpu, d } from 'typegpu';
+import { tgpu, d, std } from 'typegpu';
 const root = await tgpu.init();
-const counterEnabledSlot = tgpu.slot<boolean>(false);
-const counter = root.createMutable(d.u32);
 // ---cut---
+// May be overridden by null.
+const preprocessSlot = tgpu.slot<null | ((a: number) => number)>(std.abs);
+
 const myFunction = tgpu.fn([])(() => {
-  counterEnabledSlot.$ ? counter.$++ : undefined;
+  const n = preprocessSlot.$ !== null ? preprocessSlot.$(-5) : -5;
 });
 ```
 

packages/tinyest-for-wgsl/src/parsers.ts

@@ -134,6 +134,9 @@ const Transpilers: Partial<{
     if (node.bigint) {
       console.warn('BigInt literals are represented as numbers - loss of precision may occur.');
     }
+    if (node.raw === 'null') {
+      return [NODE.nullLiteral];
+    }
     return [NODE.numericLiteral, String(Number(node.value))];
   },
 
@@ -284,6 +287,10 @@ const Transpilers: Partial<{
     return [NODE.break];
   },
 
+  NullLiteral() {
+    return [NODE.nullLiteral];
+  },
+
   TSAsExpression: tsFallthrough,
   TSSatisfiesExpression: tsFallthrough,
   TSNonNullExpression: tsFallthrough,

packages/tinyest-for-wgsl/tests/parsers.test.ts

@@ -31,6 +31,21 @@ describe('transpileFn', () => {
     }),
   );
 
+  it(
+    'parses null',
+    dualTest((p) => {
+      const { params, body, externalNames } = transpileFn(
+        p(`() => {
+          const a = null;
+        }`),
+      );
+
+      expect(params).toStrictEqual([]);
+      expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[[13,"a",[106]]]]"`);
+      expect(externalNames).toMatchInlineSnapshot(`Map {}`);
+    }),
+  );
+
   it(
     'fails when the input is not a function',
     dualTest((p) => {

packages/tinyest/src/nodes.ts

@@ -32,6 +32,7 @@ export const NodeTypeCatalog = {
   stringLiteral: 103,
   objectExpr: 104,
   conditionalExpr: 105,
+  nullLiteral: 106,
 } as const;
 
 export type NodeTypeCatalog = typeof NodeTypeCatalog;
@@ -237,7 +238,10 @@ export type Num = readonly [type: NodeTypeCatalog['numericLiteral'], string];
 /** A string literal */
 export type Str = readonly [type: NodeTypeCatalog['stringLiteral'], string];
 
-export type Literal = Num | Str | boolean;
+/** null literal */
+export type Null = readonly [type: NodeTypeCatalog['nullLiteral']];
+
+export type Literal = Num | Str | boolean | Null;
 
 /** Identifiers are just strings, since string literals are rare in WGSL, and identifiers are everywhere. */
 export type Expression =

packages/typegpu/src/resolutionCtx.ts

@@ -1102,6 +1102,12 @@ export class ResolutionCtxImpl implements ResolutionCtx {
       );
     }
 
+    if (item === null) {
+      throw new WgslTypeError(
+        `'null' is not resolvable. 'null' is only allowed in comptime checks.`,
+      );
+    }
+
     throw new WgslTypeError(
       `Value ${safeStringify(item)} is not resolvable${
         schema && schema !== UnknownData ? ` to type ${safeStringify(schema)}` : ''

packages/typegpu/src/shared/tseynit.ts

@@ -157,6 +157,10 @@ function stringifyExpression(node: tinyest.Expression, ident: string): string {
     return `${wrapIfComplex(node[1], ident)} ? ${wrapIfComplex(node[2], ident)} : ${wrapIfComplex(node[3], ident)}`;
   }
 
+  if (node[0] === NODE.nullLiteral) {
+    return 'null';
+  }
+
   assertExhaustive(node);
 }
 
@@ -181,7 +185,8 @@ function isExpression(node: tinyest.AnyNode): node is tinyest.Expression {
     node[0] === NODE.preUpdate ||
     node[0] === NODE.postUpdate ||
     node[0] === NODE.objectExpr ||
-    node[0] === NODE.conditionalExpr
+    node[0] === NODE.conditionalExpr ||
+    node[0] === NODE.nullLiteral
   ) {
     node satisfies tinyest.Expression;
     return true;

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -1048,6 +1048,10 @@ export class WgslGenerator implements ShaderGenerator {
       throw new Error('Cannot use pre-updates in TypeGPU functions.');
     }
 
+    if (expression[0] === NODE.nullLiteral) {
+      return snip(null, UnknownData, 'constant', false);
+    }
+
     assertExhaustive(expression);
   }
 

packages/typegpu/tests/internal/tseynit.test.ts

@@ -339,5 +339,21 @@ describe('ast to JS transformation', () => {
         }"
       `);
     });
+
+    it('handles null', () => {
+      const slot = tgpu.slot<number | null>(null);
+      const fn = () => {
+        'use gpu';
+        if (slot.$ !== null) {
+        }
+      };
+      expect(stringifyNode(getBodyAst(fn))).toMatchInlineSnapshot(`
+        "{
+          if (slot.$ !== (null)) {
+
+          }
+        }"
+      `);
+    });
   });
 });

packages/typegpu/tests/slot.test.ts

@@ -440,4 +440,27 @@ describe('tgpu.slot', () => {
       }"
     `);
   });
+
+  it('allows null', () => {
+    const stepsSlot = tgpu.slot<number | null>(null);
+
+    const getSteps = () => {
+      'use gpu';
+      let steps = 0;
+      if (stepsSlot.$ !== null) {
+        steps = stepsSlot.$;
+      } else {
+        steps = 5;
+      }
+      return steps;
+    };
+
+    expect(tgpu.resolve([getSteps])).toMatchInlineSnapshot(`
+      "fn getSteps() -> i32 {
+        var steps = 0;
+        steps = 5i;
+        return steps;
+      }"
+    `);
+  });
 });

packages/typegpu/tests/tgsl/comptime.test.ts

@@ -75,4 +75,18 @@ describe('comptime', () => {
       }"
     `);
   });
+
+  it('can return null', () => {
+    const comptime = tgpu.comptime(() => null);
+    const myFn = () => {
+      'use gpu';
+      return comptime() !== null ? 0 : 1;
+    };
+
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "fn myFn() -> i32 {
+        return 1;
+      }"
+    `);
+  });
 });

packages/typegpu/tests/tgslFn.test.ts

@@ -1310,3 +1310,79 @@ describe('string injection', () => {
     `);
   });
 });
+
+describe('nulls in TGSL', () => {
+  it('throws when assigning to a variable', () => {
+    const myFn = () => {
+      'use gpu';
+      const a = null;
+    };
+
+    expect(() => tgpu.resolve([myFn])).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:myFn
+      - fn*:myFn(): 'const a = null' is invalid, cannot determine WGSL type of 'null'
+      -----
+      - Try using or defining a schema that matches your desired value the most, and wrap the value with it: 'const a = Schema(null)'
+      -----]
+    `);
+  });
+
+  it('allows comptime usage', () => {
+    let externalNum: number | null;
+    const myFn = () => {
+      'use gpu';
+      if (externalNum !== null) {
+        return externalNum;
+      } else {
+        return 1;
+      }
+    };
+
+    externalNum = 0;
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "fn myFn() -> i32 {
+        return 0;
+      }"
+    `);
+    externalNum = null;
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "fn myFn() -> i32 {
+        return 1;
+      }"
+    `);
+  });
+
+  it('allows comptime usage in ternary checks', () => {
+    let externalNum: number | null = 1;
+    const myFn = () => {
+      'use gpu';
+      return externalNum !== null ? externalNum : 1;
+    };
+
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "fn myFn() -> i32 {
+        return 1;
+      }"
+    `);
+  });
+
+  it('is not coerced to false', () => {
+    const myFn = () => {
+      'use gpu';
+      // @ts-ignore
+      if (null) {
+        return false;
+      }
+      return true;
+    };
+
+    expect(() => tgpu.resolve([myFn])).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:myFn
+      - fn*:myFn(): 'null' is not resolvable. 'null' is only allowed in comptime checks.]
+    `);
+  });
+});

packages/unplugin-typegpu/src/core/obfuscate.ts

@@ -172,6 +172,9 @@ const visitors = {
   conditionalExpr(ctx: Context, node: tinyest.ConditionalExpression) {
     return [NODE.conditionalExpr, obf(ctx, node[1]), obf(ctx, node[2]), obf(ctx, node[3])];
   },
+  nullLiteral(_: Context, node: tinyest.Null) {
+    return node;
+  },
 } as const satisfies {
   [N in keyof typeof NODE]: (
     ctx: Context,

packages/unplugin-typegpu/test/obfuscation.test.ts

@@ -269,6 +269,23 @@ describe('obfuscate', () => {
     `);
   });
 
+  it('handles null', () => {
+    const code = `() => {
+      const variable = null;
+    }`;
+    const transpiled = transpileFn(parse(code));
+
+    const { params, body, externalNames } = obfuscate(transpiled);
+
+    expect(params).toStrictEqual([]);
+    expect(stringifyNode(body)).toMatchInlineSnapshot(`
+      "{
+        const a = null;
+      }"
+    `);
+    expect(externalNames).toMatchInlineSnapshot(`Map {}`);
+  });
+
   it('obfuscates parameters', () => {
     const code = `(param1, param2) => { return param2 + param1; }`;
     const transpiled = transpileFn(parse(code));