fix: Check for context instead of comptime mode in comptime

packages/typegpu/setupVitest.ts

@@ -1,5 +1,9 @@
 import { setup } from '@ark/attest';
 import { type } from 'arktype';
+import { dirname, join } from 'node:path';
+import { fileURLToPath } from 'node:url';
+
+const packageDir = dirname(fileURLToPath(import.meta.url));
 
 const truthyString = type('"0"|"1"').pipe.try((value) => Boolean(Number.parseInt(value)));
 
@@ -12,6 +16,7 @@ const env = ProcessEnvType.assert(process.env);
 export default () =>
   setup({
     formatCmd: 'pnpm fix',
+    tsconfig: join(packageDir, 'tsconfig.json'),
     // Skipping type tests by default
     skipTypes: !env.ENABLE_ATTEST,
   });

packages/typegpu/src/core/slot/accessor.ts

@@ -2,7 +2,7 @@ import { type AnyData, isData } from '../../data/dataTypes.ts';
 import { schemaCallWrapper } from '../../data/schemaCallWrapper.ts';
 import { isSnippet, type ResolvedSnippet, snip } from '../../data/snippet.ts';
 import type { BaseData } from '../../data/wgslTypes.ts';
-import { getResolutionCtx, inCodegenMode } from '../../execMode.ts';
+import { getResolutionCtx } from '../../execMode.ts';
 import { getName, hasTinyestMetadata, setName } from '../../shared/meta.ts';
 import type { InferGPU } from '../../shared/repr.ts';
 import {
@@ -14,6 +14,7 @@ import {
 } from '../../shared/symbols.ts';
 import type { UnwrapRuntimeConstructor } from '../../tgpuBindGroupLayout.ts';
 import {
+  CodegenState,
   getOwnSnippet,
   NormalState,
   type ResolutionCtx,
@@ -174,13 +175,33 @@ export class TgpuAccessorImpl<T extends BaseData>
   }
 
   get $(): InferGPU<T> {
-    if (inCodegenMode()) {
-      return this[$gpuValueOf];
+    const ctx = getResolutionCtx();
+    if (!ctx) {
+      throw new Error(
+        '`tgpu.accessor` relies on GPU resources and cannot be accessed outside of a compute dispatch or draw call',
+      );
     }
 
-    throw new Error(
-      '`tgpu.accessor` relies on GPU resources and cannot be accessed outside of a compute dispatch or draw call',
-    );
+    if (ctx.mode.type !== 'codegen') {
+      const slotValue = ctx.unwrap(this.slot);
+
+      if (
+        typeof slotValue !== 'function' &&
+        !hasTinyestMetadata(slotValue) &&
+        !isTgpuFn(slotValue)
+      ) {
+        return slotValue as unknown as InferGPU<T>;
+      }
+
+      ctx.pushMode(new CodegenState());
+      try {
+        return this[$gpuValueOf];
+      } finally {
+        ctx.popMode('codegen');
+      }
+    }
+
+    return this[$gpuValueOf];
   }
 }
 
@@ -198,12 +219,32 @@ export class TgpuMutableAccessorImpl<T extends BaseData>
   }
 
   get $(): InferGPU<T> {
-    if (inCodegenMode()) {
-      return this[$gpuValueOf];
+    const ctx = getResolutionCtx();
+    if (!ctx) {
+      throw new Error(
+        '`tgpu.mutableAccessor` relies on GPU resources and cannot be accessed outside of a compute dispatch or draw call',
+      );
     }
 
-    throw new Error(
-      '`tgpu.mutableAccessor` relies on GPU resources and cannot be accessed outside of a compute dispatch or draw call',
-    );
+    if (ctx.mode.type !== 'codegen') {
+      const slotValue = ctx.unwrap(this.slot);
+
+      if (
+        typeof slotValue !== 'function' &&
+        !hasTinyestMetadata(slotValue) &&
+        !isTgpuFn(slotValue)
+      ) {
+        return slotValue as unknown as InferGPU<T>;
+      }
+
+      ctx.pushMode(new CodegenState());
+      try {
+        return this[$gpuValueOf];
+      } finally {
+        ctx.popMode('codegen');
+      }
+    }
+
+    return this[$gpuValueOf];
   }
 }

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

@@ -75,4 +75,145 @@ describe('comptime', () => {
       }"
     `);
   });
+
+  it('can read accessors during shader resolution', () => {
+    const value = tgpu.accessor(d.f32, 1);
+    const readValue = tgpu.comptime(() => value.$);
+
+    const myFn = tgpu.fn(
+      [],
+      d.f32,
+    )(() => {
+      return readValue();
+    });
+
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "fn myFn() -> f32 {
+        return 1f;
+      }"
+    `);
+
+    expect(tgpu.resolve([myFn.with(value, 2)])).toMatchInlineSnapshot(`
+      "fn myFn() -> f32 {
+        return 2f;
+      }"
+    `);
+  });
+
+  it('can read and work with accessors', () => {
+    const valueAccess = tgpu.accessor(d.f32, 1);
+    const doubleValue = tgpu.comptime(() => valueAccess.$ * 2);
+
+    const myFn = tgpu.fn(
+      [],
+      d.f32,
+    )(() => {
+      return doubleValue();
+    });
+
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "fn myFn() -> f32 {
+        return 2f;
+      }"
+    `);
+
+    expect(tgpu.resolve([myFn.with(valueAccess, 2)])).toMatchInlineSnapshot(`
+      "fn myFn() -> f32 {
+        return 4f;
+      }"
+    `);
+  });
+
+  it('can read "use gpu" callback accessors', () => {
+    const colorAccess = tgpu.accessor(d.vec3f, () => {
+      'use gpu';
+      return d.vec3f(0, 1, 0);
+    });
+    const readColor = tgpu.comptime(() => colorAccess.$);
+
+    const myFn = tgpu.fn(
+      [],
+      d.vec3f,
+    )(() => {
+      return readColor();
+    });
+
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "fn colorAccess() -> vec3f {
+        return vec3f(0, 1, 0);
+      }
+
+      fn myFn() -> vec3f {
+        return colorAccess();
+      }"
+    `);
+  });
+
+  it('can read GPU-resource accessors', ({ root }) => {
+    const Camera = d.struct({ pos: d.vec3f });
+    const camera = root.createUniform(Camera);
+
+    const posAccess = tgpu.accessor(d.vec3f, () => camera.$.pos);
+    const readPos = tgpu.comptime(() => posAccess.$);
+
+    const myFn = tgpu.fn(
+      [],
+      d.vec3f,
+    )(() => {
+      return readPos();
+    });
+
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "struct Camera {
+        pos: vec3f,
+      }
+
+      @group(0) @binding(0) var<uniform> camera: Camera;
+
+      fn myFn() -> vec3f {
+        return camera.pos;
+      }"
+    `);
+  });
+
+  it('throws when reading "use gpu" callback accessor in js', () => {
+    const colorAccess = tgpu.accessor(d.vec3f, () => {
+      'use gpu';
+      return d.vec3f(0, 1, 0);
+    });
+    const readColor = tgpu.comptime(() => colorAccess.$);
+
+    expect(() => readColor()).toThrowErrorMatchingInlineSnapshot(
+      `[Error: \`tgpu.accessor\` relies on GPU resources and cannot be accessed outside of a compute dispatch or draw call]`,
+    );
+  });
+
+  it('throws when reading GPU-resource accessor in js', ({ root }) => {
+    const Camera = d.struct({ pos: d.vec3f });
+    const camera = root.createUniform(Camera);
+
+    const posAccess = tgpu.accessor(d.vec3f, () => camera.$.pos);
+    const readPos = tgpu.comptime(() => posAccess.$);
+
+    expect(() => readPos()).toThrowErrorMatchingInlineSnapshot(
+      `[Error: \`tgpu.accessor\` relies on GPU resources and cannot be accessed outside of a compute dispatch or draw call]`,
+    );
+  });
+
+  it('throws when a comptime-read accessor has no value', () => {
+    const value = tgpu.accessor(d.f32);
+    const readValue = tgpu.comptime(() => value.$);
+    const myFn = () => {
+      'use gpu';
+      return readValue();
+    };
+
+    expect(() => tgpu.resolve([myFn])).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:myFn
+      - fn*:myFn()
+      - fn:readValue: Missing value for 'slot:value']
+    `);
+  });
 });