feat: isBeingTranspiled + getTargetShaderLanguage

packages/typegpu/src/std/environment.ts

@@ -0,0 +1,62 @@
+import { comptime } from '../core/function/comptime.ts';
+import { bool } from '../data/numeric.ts';
+import { snip } from '../data/snippet.ts';
+import { getExecMode, getResolutionCtx } from '../execMode.ts';
+import { $gpuCallable } from '../shared/symbols.ts';
+import type { DualFn } from '../types.ts';
+
+const impl = (() => false) as DualFn<() => boolean>;
+impl.toString = () => 'isBeingTranspiled';
+impl[$gpuCallable] = {
+  call(_ctx, _args) {
+    return snip(true, bool, 'constant');
+  },
+};
+
+/**
+ * Returns `true` when the direct callee is being transpiled for GPU, otherwise `false`.
+ *
+ * @example
+ * const f = () => {
+ *   'use gpu';
+ *   return isBeingTranspiled() ? 1 : 0;
+ * };
+ *
+ * f() // returns 0, but resolved WGSL looks like this:
+ *
+ * fn f() -> i32 {
+ *   return 1;
+ * }
+ *
+ * @note
+ * Inside `comptime`, `lazy` or `simulate`, it always returns `false`.
+ */
+export const isBeingTranspiled = impl;
+
+/**
+ * Returns `wgsl` if invoked during the resolution process; otherwise, returns `undefined`.
+ *
+ * @example
+ * const f = () => {
+ *   'use gpu';
+ *   return getTargetShaderLanguage() === 'wgsl';
+ * };
+ *
+ * f() // returns false, but resolved WGSL looks like this:
+ *
+ * fn f() -> bool {
+ *   return true;
+ * }
+ *
+ * @note
+ * Inside `lazy`, it always returns `wgsl`.
+ * Inside `simulate`, it always returns `undefined`.
+ * Inside `comptime`, it returns `wgsl` if called during the resolution process; otherwise, `undefined`.
+ */
+export const getTargetShaderLanguage = comptime((() => {
+  const ctx = getResolutionCtx();
+  if (!ctx) {
+    return undefined;
+  }
+  return getExecMode().type !== 'simulate' ? 'wgsl' : undefined;
+}) as () => string | undefined);

packages/typegpu/src/std/index.ts

@@ -189,3 +189,5 @@ export { extensionEnabled } from './extensions.ts';
 export { bitcastU32toF32, bitcastU32toI32 } from './bitcast.ts';
 
 export { range } from './range.ts';
+
+export { isBeingTranspiled, getTargetShaderLanguage } from './environment.ts';

packages/typegpu/tests/std/environment.test.ts

@@ -0,0 +1,172 @@
+import { it } from 'typegpu-testing-utility';
+import { expect, describe } from 'vitest';
+
+import tgpu, { d, std } from '../../src/index.js';
+
+describe('isBeingTranspiled', () => {
+  it('returns false top level', () => {
+    expect(std.isBeingTranspiled()).toBe(false);
+  });
+
+  it('returns true during function resolution', () => {
+    const f = () => {
+      'use gpu';
+      if (std.isBeingTranspiled()) {
+        return 7;
+      } else {
+        return -7;
+      }
+    };
+
+    expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
+      "fn f() -> i32 {
+        {
+          return 7;
+        }
+      }"
+    `);
+  });
+
+  it('returns false inside comptime', () => {
+    const checkTranspilation = tgpu.comptime(std.isBeingTranspiled);
+    expect(checkTranspilation()).toBe(false);
+
+    const f = () => {
+      'use gpu';
+      const _transpilation = checkTranspilation();
+    };
+
+    expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
+      "fn f() {
+        const _transpilation = false;
+      }"
+    `);
+
+    expect(checkTranspilation()).toBe(false);
+  });
+
+  it('returns false inside lazy', () => {
+    const checkTranspilation = tgpu.lazy(std.isBeingTranspiled);
+
+    const f = () => {
+      'use gpu';
+      const _transpilation = checkTranspilation.$;
+    };
+
+    expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
+      "fn f() {
+        const _transpilation = false;
+      }"
+    `);
+  });
+
+  it('returns false inside simulate', () => {
+    const counter = tgpu.privateVar(d.u32, 0);
+
+    const result = tgpu['~unstable'].simulate(() => {
+      if (!std.isBeingTranspiled()) {
+        counter.$ += 1;
+      }
+      return counter.$;
+    });
+
+    expect(result.value).toBe(1);
+  });
+
+  it('correctly branches during js execution', () => {
+    const f = () => {
+      'use gpu';
+      if (std.isBeingTranspiled()) {
+        return 7;
+      } else {
+        return -7;
+      }
+    };
+
+    expect(f()).toBe(-7);
+  });
+});
+
+describe('getTargetShaderLanguage', () => {
+  it('returns undefined top level', () => {
+    expect(std.getTargetShaderLanguage()).toBe(undefined);
+  });
+
+  it('returns `wgsl` during function resolution', () => {
+    const f = () => {
+      'use gpu';
+      if (std.getTargetShaderLanguage() === 'wgsl') {
+        return 7;
+      } else {
+        return -7;
+      }
+    };
+
+    expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
+      "fn f() -> i32 {
+        {
+          return 7;
+        }
+      }"
+    `);
+  });
+
+  it('returns undefined inside comptime outside of resolution and `wgsl` during function resolution', () => {
+    const checkTranspilation = tgpu.comptime(std.getTargetShaderLanguage);
+    expect(checkTranspilation()).toBe(undefined);
+
+    const f = () => {
+      'use gpu';
+      const _transpilation = checkTranspilation() === 'wgsl';
+    };
+
+    expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
+      "fn f() {
+        const _transpilation = true;
+      }"
+    `);
+
+    expect(checkTranspilation()).toBe(undefined);
+  });
+
+  it('returns `wgsl` inside lazy', () => {
+    const checkTranspilation = tgpu.lazy(std.getTargetShaderLanguage);
+
+    const f = () => {
+      'use gpu';
+      const _transpilation = checkTranspilation.$ === 'wgsl';
+    };
+
+    expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
+      "fn f() {
+        const _transpilation = true;
+      }"
+    `);
+  });
+
+  it('returns undefined inside simulate', () => {
+    const counter = tgpu.privateVar(d.u32, 0);
+
+    const result = tgpu['~unstable'].simulate(() => {
+      if (std.getTargetShaderLanguage() !== 'wgsl') {
+        counter.$ += 1;
+      }
+      return counter.$;
+    });
+
+    expect(result.value).toBe(1);
+  });
+
+  it('correctly branches during js execution', () => {
+    const f = () => {
+      'use gpu';
+      if (std.getTargetShaderLanguage() === 'wgsl') {
+        return 7;
+      } else {
+        return -7;
+      }
+    };
+
+    expect(f()).toBe(-7);
+  });
+});