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

@@ -7,6 +7,7 @@ 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),
       );
     }
 
@@ -156,6 +168,7 @@ abstract class AccessorBase<
       ctx.resolve(snippet.value, snippet.dataType).value,
       snippet.dataType as T,
       snippet.origin,
+      snippet.possibleSideEffects,
     );
   }
 }

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

@@ -571,7 +571,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()}.$`,
@@ -666,7 +666,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

@@ -20,7 +20,8 @@ export class UnrollableIterable implements SelfResolvable {
   }
 
   [$resolve](_ctx: ResolutionCtx): ResolvedSnippet {
-    return snip(stitch`${this.snippet}`, this.snippet.dataType as BaseData, this.snippet.origin);
+    const { dataType, origin, possibleSideEffects } = this.snippet;
+    return snip(stitch`${this.snippet}`, dataType as BaseData, origin, possibleSideEffects);
   }
 }
 
@@ -95,7 +96,12 @@ export const unroll = (() => {
   impl[$internal] = true;
   impl[$gpuCallable] = {
     call(_ctx, [value]) {
-      return snip(new UnrollableIterable(value), value.dataType, value.origin);
+      return snip(
+        new UnrollableIterable(value),
+        value.dataType,
+        value.origin,
+        value.possibleSideEffects,
+      );
     },
   };
 

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

@@ -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 snip(value.value, explicitFrom(value.dataType), value.origin, false);
       }
 
       /**
@@ -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,14 @@ 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, false);
   }
 }
 
@@ -194,8 +199,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.value.snippet.possibleSideEffects,
+    );
   }
 
-  return snip(stitch`(*${snippet})`, innerType, snippet.origin);
+  return snip(stitch`(*${snippet})`, innerType, snippet.origin, false);
 }

packages/typegpu/src/resolutionCtx.ts

@@ -1006,7 +1006,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
     }
 
     if (typeof item === 'boolean') {
-      return snip(item ? 'true' : 'false', bool, /* origin */ 'constant');
+      return snip(item ? 'true' : 'false', bool, /* origin */ 'constant', false);
     }
 
     if (typeof item === 'string') {
@@ -1052,6 +1052,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
       this.resolve(snippet.value, snippet.dataType).value,
       snippet.dataType,
       snippet.origin,
+      snippet.possibleSideEffects,
     ) as ResolvedSnippet;
   }
 

packages/typegpu/src/std/boolean.ts

@@ -367,7 +367,7 @@ export const isCloseTo = dualImpl({
     return false;
   },
   // GPU implementation
-  codegenImpl: (_ctx, [lhs, rhs, precision = snip(0.01, f32, /* origin */ 'constant')]) => {
+  codegenImpl: (_ctx, [lhs, rhs, precision = snip(0.01, f32, /* origin */ 'constant', false)]) => {
     if (isSnippetNumeric(lhs) && isSnippetNumeric(rhs)) {
       return stitch`(abs(f32(${lhs}) - f32(${rhs})) <= ${precision})`;
     }

packages/typegpu/src/tgsl/conversion.ts

@@ -256,7 +256,12 @@ function applyActionToSnippet(
 
   switch (action.action) {
     case 'ref':
-      return snip(new RefOperator(snippet, targetType as Ptr), targetType, snippet.origin);
+      return snip(
+        new RefOperator(snippet, targetType as Ptr),
+        targetType,
+        snippet.origin,
+        snippet.possibleSideEffects,
+      );
     case 'deref':
       return derefSnippet(snippet);
     case 'cast': {
@@ -368,18 +373,23 @@ export function tryConvertSnippet(
 ): Snippet {
   const targets = Array.isArray(targetDataTypes) ? targetDataTypes : [targetDataTypes];
 
-  const { value, dataType, origin } = snippet;
+  const { value, dataType, origin, possibleSideEffects } = snippet;
 
   if (targets.length === 1) {
     const target = targets[0] as AnyWgslData;
 
     if (target === dataType) {
-      return snip(value, target, origin);
+      return snip(value, target, origin, possibleSideEffects);
     }
 
     if (dataType === UnknownData) {
       // Commit unknown to the expected type.
-      return snip(stitch`${snip(value, target, origin)}`, target, origin);
+      return snip(
+        stitch`${snip(value, target, origin, possibleSideEffects)}`,
+        target,
+        origin,
+        possibleSideEffects,
+      );
     }
   }