feat: More comprehensive tracking of possible side effects

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

@@ -88,7 +88,8 @@ const myFunction = tgpu.fn([])(() => {
 
 :::caution
 Ternary operator's condition must be a comptime-known value.
-This restriction is due to WGSL having no ternary operator equivalent. 
+This restriction is due to WGSL having no ternary operator equivalent.
+However, if both branches are known to have no side-effects, then the condition can be a runtime value.
 :::
 
 ## *console.log*
@@ -517,7 +518,7 @@ import { tgpu, d } from 'typegpu';
 // final shader bundle, but we cannot
 // refer to it in any other way.
 const existingGlobal = tgpu['~unstable']
-  .rawCodeSnippet('EXISTING_GLOBAL', d.f32, 'constant');
+  .rawCodeSnippet('EXISTING_GLOBAL', d.f32, 'constant', false);
 
 const foo = tgpu.fn([], d.f32)(() => {
   'use gpu';
@@ -546,3 +547,13 @@ optimizations.
 If what the expression is a direct reference to an existing value (e.g. a uniform, a
 storage binding, ...), then choose from `'uniform'`, `'mutable'`, `'readonly'`, `'workgroup'`,
 `'private'` or `'handle'` depending on the address space of the referred value.
+
+### `possibleSideEffects`
+The fourth optional parameter `possibleSideEffects` indicates, whether generating this snippet may produce a WGSL expression with
+observable side-effects (e.g. calling a barrier, discarding a fragment,
+or writing to memory).
+
+Snippets with `possibleSideEffects: true` cannot appear in ternary
+branches that get compiled to `select()`, because `select()` evaluates
+both branches unconditionally - a side-effect meant to be conditional
+would execute regardless of the condition.

packages/typegpu/src/core/function/dualImpl.ts

@@ -35,10 +35,8 @@ interface DualImplOptions<T extends AnyFn> {
    *
    * - `discard` -> `true` - it discards the fragment.
    * - `workgroupBarrier()` -> `true` - the barrier synchronizes threads.
-   * - `atomicLoad(p)` -> `true` - atomic operations may synchronize threads
-   *   through memory ordering.
-   * - `sin(x)`, `abs(x)` -> `false` - these are purely value-producing; the call
-   *   itself has no observable effect beyond the returned value.
+   * - `sin(x)`, `abs(x)` -> `false` - these are purely value-producing; the
+   *   call itself has no observable effect beyond the returned value.
    *
    * When `false`, the result inherits side-effects from its arguments: it
    * only has `possibleSideEffects: true` if at least one argument does.
@@ -103,6 +101,7 @@ export function dualImpl<T extends AnyFn>(options: DualImplOptions<T>): DualFn<T
             returnType,
             // Functions give up ownership of their return value
             /* origin */ 'constant',
+            options.sideEffects,
           );
         } catch (e) {
           // cpuImpl may in some cases be present but implemented only partially.

packages/typegpu/src/core/rawCodeSnippet/tgpuRawCodeSnippet.ts

@@ -41,6 +41,7 @@ export type RawCodeSnippetOrigin = Exclude<
  * @param expression The code snippet that will be injected in place of `foo.$`
  * @param type The type of the expression
  * @param [origin='runtime'] Where the value originates from.
+ * @param [possibleSideEffects=true] Whether generating this snippet may produce a WGSL expression with observable side-effects (e.g. calling a barrier, discarding a fragment, or writing to memory).
  *
  * **-- Which origin to choose?**
  *
@@ -59,7 +60,7 @@ export type RawCodeSnippetOrigin = Exclude<
  * // final shader bundle, but we cannot
  * // refer to it in any other way.
  * const existingGlobal = tgpu['~unstable']
- *   .rawCodeSnippet('EXISTING_GLOBAL', d.f32, 'constant');
+ *   .rawCodeSnippet('EXISTING_GLOBAL', d.f32, 'constant', false);
  *
  * const foo = () => {
  *   'use gpu';
@@ -76,8 +77,9 @@ export function rawCodeSnippet<TDataType extends AnyData>(
   expression: string,
   type: TDataType,
   origin: RawCodeSnippetOrigin | undefined = 'runtime',
+  possibleSideEffects: boolean | undefined = true,
 ): TgpuRawCodeSnippet<TDataType> {
-  return new TgpuRawCodeSnippetImpl(expression, type, origin);
+  return new TgpuRawCodeSnippetImpl(expression, type, origin, possibleSideEffects);
 }
 
 // --------------
@@ -90,14 +92,21 @@ class TgpuRawCodeSnippetImpl<TDataType extends BaseData>
   readonly [$internal]: true;
   readonly dataType: TDataType;
   readonly origin: RawCodeSnippetOrigin;
+  readonly possibleSideEffects: boolean;
 
   #expression: string;
   #externals: ExternalMap | undefined;
 
-  constructor(expression: string, type: TDataType, origin: RawCodeSnippetOrigin) {
+  constructor(
+    expression: string,
+    type: TDataType,
+    origin: RawCodeSnippetOrigin,
+    possibleSideEffects: boolean,
+  ) {
     this[$internal] = true;
     this.dataType = type;
     this.origin = origin;
+    this.possibleSideEffects = possibleSideEffects;
 
     this.#expression = expression;
   }
@@ -115,7 +124,7 @@ class TgpuRawCodeSnippetImpl<TDataType extends BaseData>
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
     const replacedExpression = replaceExternalsInWgsl(ctx, this.#externals ?? {}, this.#expression);
 
-    return snip(replacedExpression, this.dataType, this.origin);
+    return snip(replacedExpression, this.dataType, this.origin, this.possibleSideEffects);
   }
 
   toString() {
@@ -125,12 +134,13 @@ class TgpuRawCodeSnippetImpl<TDataType extends BaseData>
   get [$gpuValueOf](): InferGPU<TDataType> {
     const dataType = this.dataType;
     const origin = this.origin;
+    const possibleSideEffects = this.possibleSideEffects;
 
     return new Proxy(
       {
         [$internal]: true,
         get [$ownSnippet]() {
-          return snip(this, dataType, origin);
+          return snip(this, dataType, origin, possibleSideEffects);
         },
         [$resolve]: (ctx) => ctx.resolve(this),
         toString: () => `raw(${String(this.dataType)}): "${this.#expression}".$`,

packages/typegpu/src/core/sampler/sampler.ts

@@ -121,7 +121,7 @@ export class TgpuLaidOutSamplerImpl<
       {
         [$internal]: true,
         get [$ownSnippet]() {
-          return snip(this, schema, /* origin */ 'handle');
+          return snip(this, schema, /* origin */ 'handle', false);
         },
         [$resolve]: (ctx) => ctx.resolve(this),
         toString: () => `${this.toString()}.$`,
@@ -212,7 +212,7 @@ class TgpuFixedSamplerImpl<T extends WgslSampler | WgslComparisonSampler>
       {
         [$internal]: true,
         get [$ownSnippet]() {
-          return snip(this, schema, /* origin */ 'handle');
+          return snip(this, schema, /* origin */ 'handle', false);
         },
         [$resolve]: (ctx) => ctx.resolve(this),
         toString: () => `${this.toString()}.$`,

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

@@ -1,12 +1,13 @@
 import { type AnyData, isData } from '../../data/dataTypes.ts';
 import { schemaCallWrapper } from '../../data/schemaCallWrapper.ts';
-import { isSnippet, type ResolvedSnippet, snip } from '../../data/snippet.ts';
+import { isSnippet, type ResolvedSnippet, snip, withValue } from '../../data/snippet.ts';
 import type { BaseData } from '../../data/wgslTypes.ts';
 import { getResolutionCtx, inCodegenMode } from '../../execMode.ts';
 import { getName, hasTinyestMetadata, setName } from '../../shared/meta.ts';
 import type { InferGPU } from '../../shared/repr.ts';
 import {
   $getNameForward,
+  $gpuCallable,
   $gpuValueOf,
   $internal,
   $ownSnippet,
@@ -15,6 +16,7 @@ import {
 import type { UnwrapRuntimeConstructor } from '../../tgpuBindGroupLayout.ts';
 import {
   getOwnSnippet,
+  isGPUCallable,
   NormalState,
   type ResolutionCtx,
   type SelfResolvable,
@@ -96,7 +98,12 @@ abstract class AccessorBase<
     const ctx = getResolutionCtx()!;
     let value = getGpuValueRecursively(ctx.unwrap(this.slot));
 
-    while (typeof value === 'function' && !isTgpuFn(value) && !hasTinyestMetadata(value)) {
+    while (
+      typeof value === 'function' &&
+      !isTgpuFn(value) &&
+      !isGPUCallable(value) &&
+      !hasTinyestMetadata(value)
+    ) {
       // Not a GPU function, so has to be a resource accessor (ran in codegen mode) or comptime
       value = value();
       if (isSnippet(value)) {
@@ -109,9 +116,14 @@ abstract class AccessorBase<
       return ownSnippet;
     }
 
+    if (isGPUCallable(value)) {
+      return value[$gpuCallable].call(ctx, []);
+    }
+
     if (isTgpuFn(value) || hasTinyestMetadata(value)) {
+      const fn = ctx.resolve(value);
       return ctx.withResetIndentLevel(() =>
-        snip(`${ctx.resolve(value).value}()`, this.schema, /* origin */ 'runtime'),
+        snip(`${fn.value}()`, this.schema, /* origin */ 'runtime', fn.possibleSideEffects),
       );
     }
 
@@ -152,11 +164,7 @@ abstract class AccessorBase<
 
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
     const snippet = this.#createSnippet();
-    return snip(
-      ctx.resolve(snippet.value, snippet.dataType).value,
-      snippet.dataType as T,
-      snippet.origin,
-    );
+    return withValue(ctx.resolve(snippet.value, snippet.dataType).value, snippet);
   }
 }
 

packages/typegpu/src/core/texture/externalTexture.ts

@@ -56,7 +56,7 @@ export class TgpuExternalTextureImpl implements TgpuExternalTexture, SelfResolva
       {
         [$internal]: true,
         get [$ownSnippet]() {
-          return snip(this, schema, 'handle');
+          return snip(this, schema, 'handle', false);
         },
         [$resolve]: (ctx) => ctx.resolve(this),
         toString: () => `textureExternal:${getName(this) ?? '<unnamed>'}.$`,

packages/typegpu/src/core/texture/texture.ts

@@ -610,7 +610,7 @@ class TgpuFixedTextureViewImpl<T extends WgslTexture | WgslStorageTexture>
       {
         [$internal]: true,
         get [$ownSnippet]() {
-          return snip(this, schema, /* origin */ 'handle');
+          return snip(this, schema, /* origin */ 'handle', false);
         },
         [$resolve]: (ctx) => ctx.resolve(this),
         toString: () => `${this.toString()}.$`,
@@ -705,7 +705,7 @@ export class TgpuLaidOutTextureViewImpl<T extends WgslTexture | WgslStorageTextu
       {
         [$internal]: true,
         get [$ownSnippet]() {
-          return snip(this, schema, /* origin */ 'handle');
+          return snip(this, schema, /* origin */ 'handle', false);
         },
         [$resolve]: (ctx) => ctx.resolve(this),
         toString: () => `${this.toString()}.$`,

packages/typegpu/src/core/unroll/tgpuUnroll.ts

@@ -1,10 +1,8 @@
-import { stitch } from '../resolve/stitch.ts';
 import { $gpuCallable, $internal, $resolve } from '../../shared/symbols.ts';
 import { setName } from '../../shared/meta.ts';
 import type { DualFn } from '../../types.ts';
-import { type ResolvedSnippet, snip, type Snippet } from '../../data/snippet.ts';
+import { type ResolvedSnippet, type Snippet, withValue } from '../../data/snippet.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
-import type { BaseData } from '../../data/wgslTypes.ts';
 
 /**
  * The result of calling `tgpu.unroll(...)`. The code responsible for
@@ -19,8 +17,8 @@ export class UnrollableIterable implements SelfResolvable {
     this.snippet = snippet;
   }
 
-  [$resolve](_ctx: ResolutionCtx): ResolvedSnippet {
-    return snip(stitch`${this.snippet}`, this.snippet.dataType as BaseData, this.snippet.origin);
+  [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
+    return ctx.resolveSnippet(this.snippet);
   }
 }
 
@@ -95,7 +93,7 @@ export const unroll = (() => {
   impl[$internal] = true;
   impl[$gpuCallable] = {
     call(_ctx, [value]) {
-      return snip(new UnrollableIterable(value), value.dataType, value.origin);
+      return withValue(new UnrollableIterable(value), value);
     },
   };
 

packages/typegpu/src/core/variable/tgpuVariable.ts

@@ -112,7 +112,7 @@ class TgpuVarImpl<TScope extends VariableScope, TDataType extends BaseData>
       {
         [$internal]: true,
         get [$ownSnippet]() {
-          return snip(this, dataType, origin);
+          return snip(this, dataType, origin, false);
         },
         [$resolve]: (ctx) => ctx.resolve(this),
         toString: () => `var:${getName(this) ?? '<unnamed>'}.$`,

packages/typegpu/src/data/ref.ts

@@ -5,7 +5,7 @@ import { $gpuCallable, $internal, $ownSnippet, $resolve } from '../shared/symbol
 import type { DualFn, SelfResolvable } from '../types.ts';
 import { UnknownData } from './dataTypes.ts';
 import { createPtrFromOrigin, explicitFrom } from './ptr.ts';
-import { isAlias, type ResolvedSnippet, snip, type Snippet } from './snippet.ts';
+import { isAlias, type ResolvedSnippet, snip, type Snippet, withDataType } from './snippet.ts';
 import { isNaturallyEphemeral, isPtr, type Ptr, type StorableData } from './wgslTypes.ts';
 
 // ----------
@@ -77,7 +77,7 @@ export const _ref = (() => {
       if (isPtr(value.dataType)) {
         // This can happen if we take a reference of an *implicit* pointer, one
         // made by assigning a reference to a `const`.
-        return snip(value.value, explicitFrom(value.dataType), value.origin);
+        return withDataType(explicitFrom(value.dataType), value);
       }
 
       /**
@@ -90,7 +90,12 @@ export const _ref = (() => {
        * ```
        */
       const ptrType = createPtrFromOrigin(value.origin, value.dataType as StorableData);
-      return snip(new RefOperator(value, ptrType), ptrType ?? UnknownData, /* origin */ 'runtime');
+      return snip(
+        new RefOperator(value, ptrType),
+        ptrType ?? UnknownData,
+        /* origin */ 'runtime',
+        value.possibleSideEffects,
+      );
     },
   };
 
@@ -175,14 +180,19 @@ export class RefOperator implements SelfResolvable {
     if (!this.#ptrType) {
       throw new Error(stitch`Cannot take a reference of ${this.snippet}`);
     }
-    return snip(this, this.#ptrType, this.snippet.origin);
+    return snip(this, this.#ptrType, this.snippet.origin, this.snippet.possibleSideEffects);
   }
 
   [$resolve](): ResolvedSnippet {
     if (!this.#ptrType) {
       throw new Error(stitch`Cannot take a reference of ${this.snippet}`);
     }
-    return snip(stitch`(&${this.snippet})`, this.#ptrType, this.snippet.origin);
+    return snip(
+      stitch`(&${this.snippet})`,
+      this.#ptrType,
+      this.snippet.origin,
+      this.snippet.possibleSideEffects,
+    );
   }
 }
 
@@ -194,8 +204,13 @@ export function derefSnippet(snippet: Snippet): Snippet {
   const innerType = snippet.dataType.inner;
 
   if (snippet.value instanceof RefOperator) {
-    return snip(stitch`${snippet.value.snippet}`, innerType, snippet.origin);
+    return snip(
+      stitch`${snippet.value.snippet}`,
+      innerType,
+      snippet.origin,
+      snippet.possibleSideEffects,
+    );
   }
 
-  return snip(stitch`(*${snippet})`, innerType, snippet.origin);
+  return snip(stitch`(*${snippet})`, innerType, snippet.origin, false);
 }

packages/typegpu/src/data/snippet.ts

@@ -180,6 +180,12 @@ export function withDataType(dataType: BaseData | UnknownData, snippet: Snippet)
   return new SnippetImpl(snippet.value, dataType, snippet.origin, snippet.possibleSideEffects);
 }
 
+export function withValue(value: string, snippet: Snippet): ResolvedSnippet;
+export function withValue(value: unknown, snippet: Snippet): Snippet;
+export function withValue(value: unknown, snippet: Snippet): Snippet {
+  return new SnippetImpl(value, snippet.dataType, snippet.origin, snippet.possibleSideEffects);
+}
+
 export function withSideEffects(
   possibleSideEffects: boolean,
   snippet: ResolvedSnippet,