From 7a196c73005f5b4f265dc59b1896c5876ef993cb Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Wed, 18 Dec 2024 14:15:59 +0100
Subject: [PATCH 1/9] initial conversion (not working yet)

---
 .../simulation/fluid-with-atomics/index.ts    | 482 ++++++++++--------
 packages/typegpu/src/core/buffer/buffer.ts    |  15 +-
 packages/typegpu/src/data/alignmentOf.ts      |   1 +
 packages/typegpu/src/data/array.ts            |   4 +-
 packages/typegpu/src/data/atomic.ts           |   4 +-
 packages/typegpu/src/data/attributes.ts       |   3 +-
 packages/typegpu/src/data/sizeOf.ts           |   1 +
 packages/typegpu/src/data/struct.ts           |   4 +-
 packages/typegpu/src/data/wgslTypes.ts        |  18 +-
 packages/typegpu/src/shared/repr.ts           |  10 +
 packages/typegpu/tests/data/atomic.test.ts    |   4 +
 11 files changed, 328 insertions(+), 218 deletions(-)

diff --git a/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts b/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
index 22e748e34..35021cc8d 100644
--- a/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
+++ b/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
@@ -1,16 +1,5 @@
-// @ts-nocheck
-// TODO: Reenable type checking when new pipelines are implemented.
-
 import * as d from 'typegpu/data';
-import tgpu, {
-  asMutable,
-  asReadonly,
-  asUniform,
-  asVertex,
-  builtin,
-  wgsl,
-  type TgpuBuffer,
-} from 'typegpu/experimental';
+import tgpu, { asUniform, asMutable, asReadonly } from 'typegpu/experimental';
 
 const root = await tgpu.init();
 
@@ -30,6 +19,10 @@ canvas.addEventListener('contextmenu', (event) => {
   }
 });
 
+const MAX_WATER_LEVEL_UNPRESSURIZED = '510u';
+const MAX_WATER_LEVEL = '((1u << 24) - 1u)';
+const MAX_PRESSURE = '12u';
+
 const options = {
   size: 64,
   timestep: 25,
@@ -55,157 +48,207 @@ function encodeBrushType(brushType: (typeof BrushTypes)[number]) {
 }
 
 const sizeBuffer = root.createBuffer(d.vec2u).$name('size').$usage('uniform');
+const sizeUniform = asUniform(sizeBuffer);
+
 const viscosityBuffer = root
   .createBuffer(d.u32)
   .$name('viscosity')
   .$usage('uniform');
+const viscosityUniform = asUniform(viscosityBuffer);
 
 const currentStateBuffer = root
   .createBuffer(d.arrayOf(d.u32, 1024 ** 2))
   .$name('current')
   .$usage('storage', 'vertex');
+const currentStateStorage = asReadonly(currentStateBuffer);
 
 const nextStateBuffer = root
   .createBuffer(d.arrayOf(d.atomic(d.u32), 1024 ** 2))
   .$name('next')
   .$usage('storage');
-
-const viscosityData = asUniform(viscosityBuffer);
-const currentStateData = asReadonly(currentStateBuffer);
-const currentStateVertex = asVertex(currentStateBuffer, 'instance');
-const sizeData = asUniform(sizeBuffer);
-const nextStateData = asMutable(nextStateBuffer);
-
-const maxWaterLevelUnpressurized = wgsl.constant(wgsl`510u`);
-const maxWaterLevel = wgsl.constant(wgsl`(1u << 24) - 1u`);
-const maxCompress = wgsl.constant(wgsl`12u`);
+const nextStateStorage = asMutable(nextStateBuffer);
 
 const squareBuffer = root
-  .createBuffer(d.arrayOf(d.vec2u, 4), [
-    d.vec2u(0, 0),
-    d.vec2u(0, 1),
-    d.vec2u(1, 0),
-    d.vec2u(1, 1),
+  .createBuffer(d.arrayOf(d.vec2f, 4), [
+    d.vec2f(0, 0),
+    d.vec2f(0, 1),
+    d.vec2f(1, 0),
+    d.vec2f(1, 1),
   ])
-  .$usage('uniform', 'vertex')
+  .$usage('vertex')
   .$name('square');
 
-const squareBufferData = asVertex(squareBuffer, 'vertex');
-
-const getIndex = wgsl.fn`(x: u32, y: u32) -> u32 {
-  let h = ${sizeData}.y;
-  let w = ${sizeData}.x;
+const getIndex = tgpu
+  .fn([d.u32, d.u32], d.u32)
+  .does(/* wgsl */ `(x: u32, y: u32) -> u32 {
+  let h = sizeData.y;
+  let w = sizeData.x;
   return (y % h) * w + (x % w);
-}`;
-
-const getCell = wgsl.fn`(x: u32, y: u32) -> u32 {
-  return ${currentStateData}[${getIndex}(x, y)];
-}`;
-
-const getCellNext = wgsl.fn`(x: u32, y: u32) -> u32 {
-  return atomicLoad(&${nextStateData}[${getIndex}(x, y)]);
-}`;
-
-const updateCell = wgsl.fn`(x: u32, y: u32, value: u32) {
-  atomicStore(&${nextStateData}[${getIndex}(x, y)], value);
-}`;
-
-const addToCell = wgsl.fn`(x: u32, y: u32, value: u32) {
-  let cell = ${getCellNext}(x, y);
-  let waterLevel = cell & ${maxWaterLevel};
-  let newWaterLevel = min(waterLevel + value, ${maxWaterLevel});
-  atomicAdd(&${nextStateData}[${getIndex}(x, y)], newWaterLevel - waterLevel);
-}`;
-
-const subtractFromCell = wgsl.fn`(x: u32, y: u32, value: u32) {
-  let cell = ${getCellNext}(x, y);
-  let waterLevel = cell & ${maxWaterLevel};
+}`)
+  .$uses({ sizeData: sizeUniform });
+
+const getCell = tgpu
+  .fn([d.u32, d.u32], d.u32)
+  .does(/* wgsl */ `(x: u32, y: u32) -> u32 {
+  return currentStateData[getIndex(x, y)];
+}`)
+  .$uses({ currentStateData: currentStateStorage, getIndex });
+
+const getCellNext = tgpu
+  .fn([d.u32, d.u32], d.u32)
+  .does(/* wgsl */ `(x: u32, y: u32) -> u32 {
+  return atomicLoad(&nextStateData[getIndex(x, y)]);
+}`)
+  .$uses({ nextStateData: nextStateStorage, getIndex });
+
+const updateCell = tgpu
+  .fn([d.u32, d.u32, d.u32])
+  .does(/* wgsl */ `(x: u32, y: u32, value: u32) {
+  atomicStore(&nextStateData[getIndex(x, y)], value);
+}`)
+  .$uses({ nextStateData: nextStateStorage, getIndex });
+
+const addToCell = tgpu
+  .fn([d.u32, d.u32, d.u32])
+  .does(/* wgsl */ `(x: u32, y: u32, value: u32) {
+  let cell = getCellNext(x, y);
+  let waterLevel = cell & ${MAX_WATER_LEVEL};
+  let newWaterLevel = min(waterLevel + value, ${MAX_WATER_LEVEL});
+  atomicAdd(&nextStateData[getIndex(x, y)], newWaterLevel - waterLevel);
+}`)
+  .$uses({ getCellNext, nextStateData: nextStateStorage, getIndex });
+
+const subtractFromCell = tgpu
+  .fn([d.u32, d.u32, d.u32])
+  .does(/* wgsl */ `(x: u32, y: u32, value: u32) {
+  let cell = getCellNext(x, y);
+  let waterLevel = cell & ${MAX_WATER_LEVEL};
   let newWaterLevel = max(waterLevel - min(value, waterLevel), 0u);
-  atomicSub(&${nextStateData}[${getIndex}(x, y)], waterLevel - newWaterLevel);
-}`;
-
-const persistFlags = wgsl.fn`(x: u32, y: u32) {
-  let cell = ${getCell}(x, y);
-  let waterLevel = cell & ${maxWaterLevel};
+  atomicSub(&nextStateData[getIndex(x, y)], waterLevel - newWaterLevel);
+}`)
+  .$uses({ getCellNext, nextStateData: nextStateStorage, getIndex });
+
+const persistFlags = tgpu
+  .fn([d.u32, d.u32])
+  .does(/* wgsl */ `(x: u32, y: u32) {
+  let cell = getCell(x, y);
+  let waterLevel = cell & ${MAX_WATER_LEVEL};
   let flags = cell >> 24;
-  ${updateCell}(x, y, (flags << 24) | waterLevel);
-}`;
+  updateCell(x, y, (flags << 24) | waterLevel);
+}`)
+  .$uses({ getCell, updateCell });
 
-const getStableStateBelow = wgsl.fn`(upper: u32, lower: u32) -> u32 {
+const getStableStateBelow = tgpu.fn([d.u32, d.u32], d.u32).does(/* wgsl */ `(upper: u32, lower: u32) -> u32 {
   let totalMass = upper + lower;
-  if (totalMass <= ${maxWaterLevelUnpressurized}) {
+  if (totalMass <= ${MAX_WATER_LEVEL_UNPRESSURIZED}) {
     return totalMass;
-  } else if (totalMass >= ${maxWaterLevelUnpressurized}*2 && upper > lower) {
-    return totalMass/2 + ${maxCompress};
+  } else if (totalMass >= ${MAX_WATER_LEVEL_UNPRESSURIZED}*2 && upper > lower) {
+    return totalMass/2 + ${MAX_PRESSURE};
   }
-  return ${maxWaterLevelUnpressurized};
-}`;
-
-const isWall = wgsl.fn`(x: u32, y: u32) -> bool {
-  return (${getCell}(x, y) >> 24) == 1u;
-}`;
-
-const isWaterSource = wgsl.fn`(x: u32, y: u32) -> bool {
-  return (${getCell}(x, y) >> 24) == 2u;
-}`;
-
-const isWaterDrain = wgsl.fn`(x: u32, y: u32) -> bool {
-  return (${getCell}(x, y) >> 24) == 3u;
-}`;
-
-const isClearCell = wgsl.fn`(x: u32, y: u32) -> bool {
-  return (${getCell}(x, y) >> 24) == 4u;
-}`;
-
-const getWaterLevel = wgsl.fn`(x: u32, y: u32) -> u32 {
-  return ${getCell}(x, y) & ${maxWaterLevel};
-}`;
-
-const checkForFlagsAndBounds = wgsl.fn`(x: u32, y: u32) -> bool {
-  if (${isClearCell}(x, y)) {
-    ${updateCell}(x, y, 0u);
+  return ${MAX_WATER_LEVEL_UNPRESSURIZED};
+}`);
+
+const isWall = tgpu
+  .fn([d.u32, d.u32], d.bool)
+  .does(/* wgsl */ `(x: u32, y: u32) -> bool {
+  return (getCell(x, y) >> 24) == 1u;
+}`)
+  .$uses({ getCell });
+
+const isWaterSource = tgpu
+  .fn([d.u32, d.u32], d.bool)
+  .does(/* wgsl */ `(x: u32, y: u32) -> bool {
+  return (getCell(x, y) >> 24) == 2u;
+}`)
+  .$uses({ getCell });
+
+const isWaterDrain = tgpu
+  .fn([d.u32, d.u32], d.bool)
+  .does(/* wgsl */ `(x: u32, y: u32) -> bool {
+  return (getCell(x, y) >> 24) == 3u;
+}`)
+  .$uses({ getCell });
+
+const isClearCell = tgpu
+  .fn([d.u32, d.u32], d.bool)
+  .does(/* wgsl */ `(x: u32, y: u32) -> bool {
+  return (getCell(x, y) >> 24) == 4u;
+}`)
+  .$uses({ getCell });
+
+const getWaterLevel = tgpu
+  .fn([d.u32, d.u32], d.u32)
+  .does(/* wgsl */ `(x: u32, y: u32) -> u32 {
+  return getCell(x, y) & ${MAX_WATER_LEVEL};
+}`)
+  .$uses({ getCell });
+
+const checkForFlagsAndBounds = tgpu
+  .fn([d.u32, d.u32], d.bool)
+  .does(/* wgsl */ `(x: u32, y: u32) -> bool {
+  if (isClearCell(x, y)) {
+    updateCell(x, y, 0u);
     return true;
   }
-  if (${isWall}(x, y)) {
-    ${persistFlags}(x, y);
+
+  if (isWall(x, y)) {
+    persistFlags(x, y);
     return true;
   }
-  if (${isWaterSource}(x, y)) {
-    ${persistFlags}(x, y);
-    ${addToCell}(x, y, 10u);
+
+  if (isWaterSource(x, y)) {
+    persistFlags(x, y);
+    addToCell(x, y, 10u);
     return false;
   }
-  if (${isWaterDrain}(x, y)) {
-    ${persistFlags}(x, y);
-    ${updateCell}(x, y, 3u << 24);
+
+  if (isWaterDrain(x, y)) {
+    persistFlags(x, y);
+    updateCell(x, y, 3u << 24);
     return true;
   }
-  if (y == 0 || y == ${sizeData}.y - 1u || x == 0 || x == ${sizeData}.x - 1u) {
-    ${subtractFromCell}(x, y, ${getWaterLevel}(x, y));
+
+  if (y == 0 || y == sizeData.y - 1u || x == 0 || x == sizeData.x - 1u) {
+    subtractFromCell(x, y, getWaterLevel(x, y));
     return true;
   }
+
   return false;
-}`;
+}`)
+  .$uses({
+    isClearCell,
+    updateCell,
+    persistFlags,
+    addToCell,
+    subtractFromCell,
+    isWall,
+    isWaterSource,
+    isWaterDrain,
+    getWaterLevel,
+    sizeData: sizeUniform,
+  });
 
-const decideWaterLevel = wgsl.fn`(x: u32, y: u32) {
-  if (${checkForFlagsAndBounds}(x, y)) {
+const decideWaterLevel = tgpu
+  .fn([d.u32, d.u32])
+  .does(/* wgsl */ `(x: u32, y: u32) {
+  if (checkForFlagsAndBounds(x, y)) {
     return;
   }
 
-  var remainingWater: u32 = ${getWaterLevel}(x, y);
+  var remainingWater: u32 = getWaterLevel(x, y);
 
   if (remainingWater == 0u) {
     return;
   }
 
-  if (!${isWall}(x, y - 1u)) {
-    let waterLevelBelow = ${getWaterLevel}(x, y - 1u);
-    let stable = ${getStableStateBelow}(remainingWater, waterLevelBelow);
+  if (!isWall(x, y - 1u)) {
+    let waterLevelBelow = getWaterLevel(x, y - 1u);
+    let stable = getStableStateBelow(remainingWater, waterLevelBelow);
     if (waterLevelBelow < stable) {
       let change = stable - waterLevelBelow;
-      let flow = min(change, ${viscosityData});
-      ${subtractFromCell}(x, y, flow);
-      ${addToCell}(x, y - 1u, flow);
+      let flow = min(change, viscosityData);
+      subtractFromCell(x, y, flow);
+      addToCell(x, y - 1u, flow);
       remainingWater -= flow;
     }
   }
@@ -215,13 +258,13 @@ const decideWaterLevel = wgsl.fn`(x: u32, y: u32) {
   }
 
   let waterLevelBefore = remainingWater;
-  if (!${isWall}(x - 1u, y)) {
-    let flowRaw = (i32(waterLevelBefore) - i32(${getWaterLevel}(x - 1u, y)));
+  if (!isWall(x - 1u, y)) {
+    let flowRaw = (i32(waterLevelBefore) - i32(getWaterLevel(x - 1u, y)));
     if (flowRaw > 0) {
       let change = max(min(4u, remainingWater), u32(flowRaw)/4);
-      let flow = min(change, ${viscosityData});
-      ${subtractFromCell}(x, y, flow);
-      ${addToCell}(x - 1u, y, flow);
+      let flow = min(change, viscosityData);
+      subtractFromCell(x, y, flow);
+      addToCell(x - 1u, y, flow);
       remainingWater -= flow;
     }
   }
@@ -230,13 +273,13 @@ const decideWaterLevel = wgsl.fn`(x: u32, y: u32) {
     return;
   }
 
-  if (!${isWall}(x + 1u, y)) {
-    let flowRaw = (i32(waterLevelBefore) - i32(${getWaterLevel}(x + 1, y)));
+  if (!isWall(x + 1u, y)) {
+    let flowRaw = (i32(waterLevelBefore) - i32(getWaterLevel(x + 1, y)));
     if (flowRaw > 0) {
       let change = max(min(4u, remainingWater), u32(flowRaw)/4);
-      let flow = min(change, ${viscosityData});
-      ${subtractFromCell}(x, y, flow);
-      ${addToCell}(x + 1u, y, flow);
+      let flow = min(change, viscosityData);
+      subtractFromCell(x, y, flow);
+      addToCell(x + 1u, y, flow);
       remainingWater -= flow;
     }
   }
@@ -245,31 +288,59 @@ const decideWaterLevel = wgsl.fn`(x: u32, y: u32) {
     return;
   }
 
-  if (!${isWall}(x, y + 1u)) {
-    let stable = ${getStableStateBelow}(${getWaterLevel}(x, y + 1u), remainingWater);
+  if (!isWall(x, y + 1u)) {
+    let stable = getStableStateBelow(getWaterLevel(x, y + 1u), remainingWater);
     if (stable < remainingWater) {
-      let flow = min(remainingWater - stable, ${viscosityData});
-      ${subtractFromCell}(x, y, flow);
-      ${addToCell}(x, y + 1u, flow);
+      let flow = min(remainingWater - stable, viscosityData);
+      subtractFromCell(x, y, flow);
+      addToCell(x, y + 1u, flow);
       remainingWater -= flow;
     }
   }
-}`;
-
-const computeWGSL = wgsl`
-  let x = ${builtin.globalInvocationId}.x;
-  let y = ${builtin.globalInvocationId}.y;
-  ${decideWaterLevel}(x, y);
-`;
-
-const vertWGSL = wgsl`
-  let w = ${sizeData}.x;
-  let h = ${sizeData}.y;
-  let x = (f32(${builtin.instanceIndex} % w + ${squareBufferData}.x) / f32(w) - 0.5) * 2. * f32(w) / f32(max(w, h));
-  let y = (f32((${builtin.instanceIndex} - (${builtin.instanceIndex} % w)) / w + ${squareBufferData}.y) / f32(h) - 0.5) * 2. * f32(h) / f32(max(w, h));
-  let cellFlags = ${currentStateVertex} >> 24;
-  let cellVal = f32(${currentStateVertex} & 0xFFFFFF);
-  let pos = vec4<f32>(x, y, 0., 1.);
+}`)
+  .$uses({
+    checkForFlagsAndBounds,
+    getWaterLevel,
+    isWall,
+    getStableStateBelow,
+    subtractFromCell,
+    addToCell,
+    viscosityData: viscosityUniform,
+  });
+
+const compute = tgpu
+  .computeFn([d.builtin.globalInvocationId], {
+    workgroupSize: [options.workgroupSize, options.workgroupSize],
+  })
+  .does(/* wgsl */ `(@builtin(global_invocation_id) gid: vec3u) {
+  let x = gid.x;
+  let y = gid.y;
+  decideWaterLevel(x, y);
+}`)
+  .$uses({ decideWaterLevel });
+
+console.log(
+  tgpu.resolve({
+    input: [compute],
+  }),
+);
+
+const vertex = tgpu
+  .vertexFn(
+    {
+      squareData: d.vec2f,
+      currentStateData: d.u32,
+      idx: d.builtin.instanceIndex,
+    },
+    { pos: d.builtin.position, cell: d.f32 },
+  )
+  .does(/* wgsl */ `(@builtin(instance_index) idx: u32, @location(0) squareData: vec2f, @location(1) currentStateData: u32) -> VertexOut {
+  let w = sizeData.x;
+  let h = sizeData.y;
+  let x = ((f32(idx % w) + squareData.x) / f32(w) - 0.5) * 2. * f32(w) / f32(max(w, h));
+  let y = (f32((idx - (idx % w)) / w + u32(squareData.y)) / f32(h) - 0.5) * 2. * f32(h) / f32(max(w, h));
+  let cellFlags = currentStateData >> 24;
+  let cellVal = f32(currentStateData & 0xFFFFFF);
   var cell: f32;
   cell = cellVal;
   if (cellFlags == 1u) {
@@ -281,17 +352,21 @@ const vertWGSL = wgsl`
   if (cellFlags == 3u) {
     cell = -3.;
   }
-`;
+  return VertexOut(vec4<f32>(x, y, 0., 1.), cell);
+}`)
+  .$uses({
+    sizeData: sizeUniform,
+  });
 
-const fragWGSL = wgsl`
+const fragment = tgpu.fragmentFn({ cell: d.f32 }, d.vec4f).does(/* wgsl */ `(@location(0) cell: f32) -> @location(0) vec4f {
   if (cell == -1.) {
-    return vec4f(0.5, 0.5, 0.5, 1.);
+    return vec4<f32>(0.5, 0.5, 0.5, 1.);
   }
   if (cell == -2.) {
-    return vec4f(0., 1., 0., 1.);
+    return vec4<f32>(0., 1., 0., 1.);
   }
   if (cell == -3.) {
-    return vec4f(1., 0., 0., 1.);
+    return vec4<f32>(1., 0., 0., 1.);
   }
 
   var r = f32((u32(cell) >> 16) & 0xFF)/255.;
@@ -301,8 +376,21 @@ const fragWGSL = wgsl`
   if (g > 0.) { b = 1.;}
   if (b > 0. && b < 0.5) { b = 0.5;}
 
+  if (r == 0. && g == 0. && b == 0.) {
+    return vec4<f32>(0.937, 0.937, 0.976, 1.);
+  }
+
   return vec4f(r, g, b, 1.);
-`;
+}`);
+
+const vertexInstanceLayout = tgpu.vertexLayout(
+  (n: number) => d.arrayOf(d.u32, n),
+  'instance',
+);
+const vertexLayout = tgpu.vertexLayout(
+  (n: number) => d.arrayOf(d.vec2f, n),
+  'vertex',
+);
 
 let drawCanvasData = new Uint32Array(options.size * options.size);
 
@@ -314,60 +402,45 @@ let renderChanges: () => void;
 function resetGameData() {
   drawCanvasData = new Uint32Array(options.size * options.size);
 
-  const computePipeline = root.makeComputePipeline({
-    workgroupSize: [options.workgroupSize, options.workgroupSize],
-    code: computeWGSL,
-  });
-
-  const renderPipeline = root.makeRenderPipeline({
-    vertex: {
-      code: vertWGSL,
-      output: {
-        [builtin.position.s]: 'pos',
-        cell: d.f32,
-      },
-    },
-    fragment: {
-      code: fragWGSL,
-      target: [{ format: presentationFormat }],
-    },
-    primitive: { topology: 'triangle-strip' },
-  });
+  const cpp = root.withCompute(compute).createPipeline();
+  const vp = root
+    .withVertex(vertex, {
+      squareData: vertexLayout.attrib,
+      currentStateData: vertexInstanceLayout.attrib,
+    })
+    .withFragment(fragment, {
+      format: presentationFormat,
+    })
+    .withPrimitive({
+      topology: 'triangle-strip',
+    })
+    .createPipeline();
 
   currentStateBuffer.write(Array.from({ length: 1024 ** 2 }, () => 0));
   nextStateBuffer.write(Array.from({ length: 1024 ** 2 }, () => 0));
   sizeBuffer.write(d.vec2u(options.size, options.size));
 
   render = () => {
-    const view = context.getCurrentTexture().createView();
-
     // compute
-    computePipeline.execute({
-      workgroups: [
-        options.size / options.workgroupSize,
-        options.size / options.workgroupSize,
-      ],
-    });
+    cpp.dispatchWorkgroups(
+      options.size / options.workgroupSize,
+      options.size / options.workgroupSize,
+    );
 
     // render
-    renderPipeline.execute({
-      colorAttachments: [
-        {
-          view,
-          loadOp: 'clear',
-          storeOp: 'store',
-        },
-      ],
-      vertexCount: 4,
-      instanceCount: options.size ** 2,
-    });
+    vp.withColorAttachment({
+      view: context.getCurrentTexture().createView(),
+      clearValue: [0, 0, 0, 0],
+      loadOp: 'clear' as const,
+      storeOp: 'store' as const,
+    })
+      .with(vertexLayout, squareBuffer)
+      .with(vertexInstanceLayout, currentStateBuffer)
+      .draw(4, options.size ** 2);
 
     root.flush();
 
-    currentStateBuffer.write(
-      // The atomic<> prevents this from being a 1-to-1 match.
-      nextStateBuffer as unknown as TgpuBuffer<d.TgpuArray<d.U32>>,
-    );
+    nextStateBuffer.copyFrom(currentStateBuffer);
   };
 
   applyDrawCanvas = () => {
@@ -395,18 +468,15 @@ function resetGameData() {
   };
 
   renderChanges = () => {
-    const view = context.getCurrentTexture().createView();
-    renderPipeline.execute({
-      colorAttachments: [
-        {
-          view,
-          loadOp: 'clear',
-          storeOp: 'store',
-        },
-      ],
-      vertexCount: 4,
-      instanceCount: options.size ** 2,
-    });
+    vp.withColorAttachment({
+      view: context.getCurrentTexture().createView(),
+      clearValue: [0, 0, 0, 0],
+      loadOp: 'clear' as const,
+      storeOp: 'store' as const,
+    })
+      .with(vertexLayout, squareBuffer)
+      .with(vertexInstanceLayout, currentStateBuffer)
+      .draw(4, options.size ** 2);
     root.flush();
   };
 
diff --git a/packages/typegpu/src/core/buffer/buffer.ts b/packages/typegpu/src/core/buffer/buffer.ts
index 76383ce41..6c349086e 100644
--- a/packages/typegpu/src/core/buffer/buffer.ts
+++ b/packages/typegpu/src/core/buffer/buffer.ts
@@ -3,10 +3,11 @@ import { isWgslData } from '../../data';
 import { readData, writeData } from '../../data/dataIO';
 import type { AnyData } from '../../data/dataTypes';
 import { sizeOf } from '../../data/sizeOf';
-import type { WgslTypeLiteral } from '../../data/wgslTypes';
+import type { BaseWgslData, WgslTypeLiteral } from '../../data/wgslTypes';
 import type { Storage } from '../../extension';
 import type { TgpuNamable } from '../../namable';
 import type { Infer } from '../../shared/repr';
+import type { MemIdentity } from '../../shared/repr';
 import type { UnionToIntersection } from '../../shared/utilityTypes';
 import { isGPUBuffer } from '../../types';
 import type { ExperimentalTgpuRoot } from '../root/rootTypes';
@@ -35,7 +36,7 @@ type LiteralToUsageType<T extends 'uniform' | 'storage' | 'vertex'> =
         ? Vertex
         : never;
 
-export interface TgpuBuffer<TData extends AnyData> extends TgpuNamable {
+export interface TgpuBuffer<TData extends BaseWgslData> extends TgpuNamable {
   readonly resourceType: 'buffer';
   readonly dataType: TData;
   readonly initial?: Infer<TData> | undefined;
@@ -50,7 +51,7 @@ export interface TgpuBuffer<TData extends AnyData> extends TgpuNamable {
   $addFlags(flags: GPUBufferUsageFlags): this;
 
   write(data: Infer<TData>): void;
-  copyFrom(srcBuffer: TgpuBuffer<TData>): void;
+  copyFrom(srcBuffer: TgpuBuffer<MemIdentity<TData>>): void;
   read(): Promise<Infer<TData>>;
   destroy(): void;
 }
@@ -91,13 +92,13 @@ export function isUsableAsVertex<T extends TgpuBuffer<AnyData>>(
 // Implementation
 // --------------
 
-type RestrictVertexUsages<TData extends AnyData> = TData extends {
+type RestrictVertexUsages<TData extends BaseWgslData> = TData extends {
   readonly type: WgslTypeLiteral;
 }
   ? ('uniform' | 'storage' | 'vertex')[]
   : 'vertex'[];
 
-class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
+class TgpuBufferImpl<TData extends BaseWgslData> implements TgpuBuffer<TData> {
   public readonly resourceType = 'buffer';
   public flags: GPUBufferUsageFlags =
     GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC;
@@ -168,7 +169,7 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
     return this;
   }
 
-  $usage<T extends RestrictVertexUsages<TData>>(
+  $usage<T extends ('uniform' | 'storage' | 'vertex')[]>(
     ...usages: T
   ): this & UnionToIntersection<LiteralToUsageType<T[number]>> {
     for (const usage of usages) {
@@ -213,7 +214,7 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
     device.queue.writeBuffer(gpuBuffer, 0, hostBuffer, 0, size);
   }
 
-  copyFrom(srcBuffer: TgpuBuffer<TData>): void {
+  copyFrom(srcBuffer: TgpuBuffer<MemIdentity<TData>>): void {
     if (this.buffer.mapState === 'mapped') {
       throw new Error('Cannot copy to a mapped buffer.');
     }
diff --git a/packages/typegpu/src/data/alignmentOf.ts b/packages/typegpu/src/data/alignmentOf.ts
index 79ce0c926..26b3e208a 100644
--- a/packages/typegpu/src/data/alignmentOf.ts
+++ b/packages/typegpu/src/data/alignmentOf.ts
@@ -34,6 +34,7 @@ const knownAlignmentMap: Record<string, number> = {
   mat2x2f: 8,
   mat3x3f: 16,
   mat4x4f: 16,
+  atomic: 4,
 };
 
 function computeAlignment(data: object): number {
diff --git a/packages/typegpu/src/data/array.ts b/packages/typegpu/src/data/array.ts
index 314a475af..0fd903737 100644
--- a/packages/typegpu/src/data/array.ts
+++ b/packages/typegpu/src/data/array.ts
@@ -1,4 +1,4 @@
-import type { Infer } from '../shared/repr';
+import type { Infer, MemIdentity } from '../shared/repr';
 import type { Exotic } from './exotic';
 import { sizeOf } from './sizeOf';
 import type { AnyWgslData, WgslArray } from './wgslTypes';
@@ -48,6 +48,8 @@ class TgpuArrayImpl<TElement extends AnyWgslData>
   public readonly '~repr'!: Infer<TElement>[];
   /** Type-token, not available at runtime */
   public readonly '~exotic'!: WgslArray<Exotic<TElement>>;
+  /** Type-token, not available at runtime */
+  public readonly '~memIdent'!: WgslArray<MemIdentity<TElement>>;
 
   constructor(
     public readonly elementType: TElement,
diff --git a/packages/typegpu/src/data/atomic.ts b/packages/typegpu/src/data/atomic.ts
index fe250c972..7a632e7fc 100644
--- a/packages/typegpu/src/data/atomic.ts
+++ b/packages/typegpu/src/data/atomic.ts
@@ -1,4 +1,4 @@
-import type { Infer } from '../shared/repr';
+import type { Infer, MemIdentity } from '../shared/repr';
 import type { Exotic } from './exotic';
 import type { Atomic, I32, U32 } from './wgslTypes';
 
@@ -29,6 +29,8 @@ class AtomicImpl<TSchema extends U32 | I32> implements Atomic<TSchema> {
   public readonly type = 'atomic';
   /** Type-token, not available at runtime */
   public readonly '~repr'!: Infer<TSchema>;
+  /** Type-token, not available at runtime */
+  public readonly '~memIdent'!: MemIdentity<TSchema>;
 
   constructor(public readonly inner: TSchema) {}
 }
diff --git a/packages/typegpu/src/data/attributes.ts b/packages/typegpu/src/data/attributes.ts
index 657a1c057..1944340f6 100644
--- a/packages/typegpu/src/data/attributes.ts
+++ b/packages/typegpu/src/data/attributes.ts
@@ -1,4 +1,4 @@
-import type { Infer } from '../shared/repr';
+import type { Infer, MemIdentity } from '../shared/repr';
 import { alignmentOf } from './alignmentOf';
 import {
   type AnyData,
@@ -268,6 +268,7 @@ class DecoratedImpl<TInner extends BaseWgslData, TAttribs extends unknown[]>
   implements Decorated<TInner, TAttribs>
 {
   public readonly type = 'decorated';
+  public readonly '~memIdent'!: MemIdentity<TInner>;
 }
 
 class LooseDecoratedImpl<
diff --git a/packages/typegpu/src/data/sizeOf.ts b/packages/typegpu/src/data/sizeOf.ts
index 09687afb4..ff47ab4de 100644
--- a/packages/typegpu/src/data/sizeOf.ts
+++ b/packages/typegpu/src/data/sizeOf.ts
@@ -72,6 +72,7 @@ const knownSizesMap: Record<string, number> = {
   sint32x4: 16,
   'unorm10-10-10-2': 4,
   'unorm8x4-bgra': 4,
+  atomic: 4,
 } satisfies Partial<Record<WgslTypeLiteral | LooseTypeLiteral, number>>;
 
 function sizeOfStruct(struct: WgslStruct) {
diff --git a/packages/typegpu/src/data/struct.ts b/packages/typegpu/src/data/struct.ts
index d06dc9ab5..70eb76d4b 100644
--- a/packages/typegpu/src/data/struct.ts
+++ b/packages/typegpu/src/data/struct.ts
@@ -1,5 +1,5 @@
 import type { TgpuNamable } from '../namable';
-import type { InferRecord } from '../shared/repr';
+import type { InferRecord, MemIdentityRecord } from '../shared/repr';
 import type { Prettify } from '../shared/utilityTypes';
 import type { ExoticRecord } from './exotic';
 import type { AnyWgslData, BaseWgslData, WgslStruct } from './wgslTypes';
@@ -51,6 +51,8 @@ class TgpuStructImpl<TProps extends Record<string, AnyWgslData>>
   public readonly '~repr'!: InferRecord<TProps>;
   /** Type-token, not available at runtime */
   public readonly '~exotic'!: WgslStruct<ExoticRecord<TProps>>;
+  /** Type-token, not available at runtime */
+  public readonly '~memIdent'!: WgslStruct<MemIdentityRecord<TProps>>;
 
   constructor(public readonly propTypes: TProps) {}
 
diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts
index b90288f70..f6dd8e72b 100644
--- a/packages/typegpu/src/data/wgslTypes.ts
+++ b/packages/typegpu/src/data/wgslTypes.ts
@@ -1,4 +1,14 @@
-import type { Infer, InferRecord } from '../shared/repr';
+import type {
+  Infer,
+  InferRecord,
+  MemIdentity,
+  MemIdentityRecord,
+} from '../shared/repr';
+
+type DecoratedLocation<T extends BaseWgslData> = Decorated<
+  T,
+  Location<number>[] | []
+>;
 
 export interface NumberArrayView {
   readonly length: number;
@@ -555,12 +565,14 @@ export interface I32 {
   readonly type: 'i32';
   /** Type-token, not available at runtime */
   readonly '~repr': number;
+  readonly '~memIdent': I32 | Atomic<I32> | DecoratedLocation<I32>;
 }
 
 export interface U32 {
   readonly type: 'u32';
   /** Type-token, not available at runtime */
   readonly '~repr': number;
+  readonly '~memIdent': U32 | Atomic<U32> | DecoratedLocation<U32>;
 }
 
 export interface Vec2f {
@@ -661,6 +673,7 @@ export interface WgslStruct<
   readonly propTypes: TProps;
   /** Type-token, not available at runtime */
   readonly '~repr': InferRecord<TProps>;
+  readonly '~memIdent': WgslStruct<MemIdentityRecord<TProps>>;
 }
 
 export interface WgslArray<TElement = BaseWgslData> {
@@ -669,6 +682,7 @@ export interface WgslArray<TElement = BaseWgslData> {
   readonly elementType: TElement;
   /** Type-token, not available at runtime */
   readonly '~repr': Infer<TElement>[];
+  readonly '~memIdent': WgslArray<MemIdentity<TElement>>;
 }
 
 /**
@@ -679,6 +693,7 @@ export interface Atomic<TInner extends U32 | I32 = U32 | I32> {
   readonly inner: TInner;
   /** Type-token, not available at runtime */
   readonly '~repr': Infer<TInner>;
+  readonly '~memIdent': MemIdentity<TInner>;
 }
 
 export interface Align<T extends number> {
@@ -710,6 +725,7 @@ export interface Decorated<
   readonly attribs: TAttribs;
   /** Type-token, not available at runtime */
   readonly '~repr': Infer<TInner>;
+  readonly '~memIdent': MemIdentity<TInner>;
 }
 
 export const wgslTypeLiterals = [
diff --git a/packages/typegpu/src/shared/repr.ts b/packages/typegpu/src/shared/repr.ts
index a89b73c7c..ffcc517b0 100644
--- a/packages/typegpu/src/shared/repr.ts
+++ b/packages/typegpu/src/shared/repr.ts
@@ -9,3 +9,13 @@ export type Infer<T> = T extends { readonly '~repr': infer TRepr } ? TRepr : T;
 export type InferRecord<T extends Record<string | number | symbol, unknown>> = {
   [Key in keyof T]: Infer<T[Key]>;
 };
+
+export type MemIdentity<T> = T extends { readonly '~memIdent': infer TMemIdent }
+  ? TMemIdent
+  : T;
+
+export type MemIdentityRecord<
+  T extends Record<string | number | symbol, unknown>,
+> = {
+  [Key in keyof T]: MemIdentity<T[Key]>;
+};
diff --git a/packages/typegpu/tests/data/atomic.test.ts b/packages/typegpu/tests/data/atomic.test.ts
index 20cfa59b7..d054834f1 100644
--- a/packages/typegpu/tests/data/atomic.test.ts
+++ b/packages/typegpu/tests/data/atomic.test.ts
@@ -23,6 +23,7 @@ describe('d.atomic', () => {
       type: 'u32' as const,
       '~repr': undefined as unknown as number,
       '~exotic': undefined as unknown as d.U32,
+      '~memIdent': undefined as unknown as d.U32,
     };
 
     const u32Atomic = d.atomic(exoticU32);
@@ -34,6 +35,7 @@ describe('d.atomic', () => {
       type: 'i32' as const,
       '~repr': undefined as unknown as number,
       '~exotic': undefined as unknown as d.I32,
+      '~memIdent': undefined as unknown as d.I32,
     };
 
     const i32Atomic = d.atomic(exoticI32);
@@ -51,11 +53,13 @@ describe('d.isAtomic', () => {
     const atomicU32 = d.atomic({
       type: 'u32',
       '~repr': undefined as unknown as number,
+      '~memIdent': undefined as unknown as d.U32,
     });
 
     const atomicI32 = d.atomic({
       type: 'i32',
       '~repr': undefined as unknown as number,
+      '~memIdent': undefined as unknown as d.I32,
     });
 
     expect(d.isAtomic(atomicU32)).toEqual(true);

From 34ce175933aa71e6c7e59a2236fe5516306ebc2a Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Tue, 7 Jan 2025 12:53:05 +0100
Subject: [PATCH 2/9] another approach

---
 packages/typegpu/src/core/buffer/buffer.ts | 10 ++---
 packages/typegpu/src/data/wgslTypes.ts     |  4 +-
 packages/typegpu/src/shared/repr.ts        |  6 ++-
 packages/typegpu/tests/buffer.test.ts      | 52 ++++++++++++++++++++++
 4 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/packages/typegpu/src/core/buffer/buffer.ts b/packages/typegpu/src/core/buffer/buffer.ts
index 6c349086e..99704c946 100644
--- a/packages/typegpu/src/core/buffer/buffer.ts
+++ b/packages/typegpu/src/core/buffer/buffer.ts
@@ -3,7 +3,7 @@ import { isWgslData } from '../../data';
 import { readData, writeData } from '../../data/dataIO';
 import type { AnyData } from '../../data/dataTypes';
 import { sizeOf } from '../../data/sizeOf';
-import type { BaseWgslData, WgslTypeLiteral } from '../../data/wgslTypes';
+import type { WgslTypeLiteral } from '../../data/wgslTypes';
 import type { Storage } from '../../extension';
 import type { TgpuNamable } from '../../namable';
 import type { Infer } from '../../shared/repr';
@@ -36,7 +36,7 @@ type LiteralToUsageType<T extends 'uniform' | 'storage' | 'vertex'> =
         ? Vertex
         : never;
 
-export interface TgpuBuffer<TData extends BaseWgslData> extends TgpuNamable {
+export interface TgpuBuffer<TData extends AnyData> extends TgpuNamable {
   readonly resourceType: 'buffer';
   readonly dataType: TData;
   readonly initial?: Infer<TData> | undefined;
@@ -92,13 +92,13 @@ export function isUsableAsVertex<T extends TgpuBuffer<AnyData>>(
 // Implementation
 // --------------
 
-type RestrictVertexUsages<TData extends BaseWgslData> = TData extends {
+type RestrictVertexUsages<TData extends AnyData> = TData extends {
   readonly type: WgslTypeLiteral;
 }
   ? ('uniform' | 'storage' | 'vertex')[]
   : 'vertex'[];
 
-class TgpuBufferImpl<TData extends BaseWgslData> implements TgpuBuffer<TData> {
+class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
   public readonly resourceType = 'buffer';
   public flags: GPUBufferUsageFlags =
     GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC;
@@ -169,7 +169,7 @@ class TgpuBufferImpl<TData extends BaseWgslData> implements TgpuBuffer<TData> {
     return this;
   }
 
-  $usage<T extends ('uniform' | 'storage' | 'vertex')[]>(
+  $usage<T extends RestrictVertexUsages<TData>>(
     ...usages: T
   ): this & UnionToIntersection<LiteralToUsageType<T[number]>> {
     for (const usage of usages) {
diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts
index f6dd8e72b..58ce77b27 100644
--- a/packages/typegpu/src/data/wgslTypes.ts
+++ b/packages/typegpu/src/data/wgslTypes.ts
@@ -725,7 +725,9 @@ export interface Decorated<
   readonly attribs: TAttribs;
   /** Type-token, not available at runtime */
   readonly '~repr': Infer<TInner>;
-  readonly '~memIdent': MemIdentity<TInner>;
+  readonly '~memIdent': TAttribs extends Location<number>[]
+    ? MemIdentity<TInner>
+    : Decorated<TInner, TAttribs>;
 }
 
 export const wgslTypeLiterals = [
diff --git a/packages/typegpu/src/shared/repr.ts b/packages/typegpu/src/shared/repr.ts
index ffcc517b0..c9d73aef2 100644
--- a/packages/typegpu/src/shared/repr.ts
+++ b/packages/typegpu/src/shared/repr.ts
@@ -1,3 +1,5 @@
+import type { AnyData } from '../data';
+
 /**
  * Extracts the inferred representation of a resource.
  * @example
@@ -10,7 +12,9 @@ export type InferRecord<T extends Record<string | number | symbol, unknown>> = {
   [Key in keyof T]: Infer<T[Key]>;
 };
 
-export type MemIdentity<T> = T extends { readonly '~memIdent': infer TMemIdent }
+export type MemIdentity<T> = T extends {
+  readonly '~memIdent': infer TMemIdent extends AnyData;
+}
   ? TMemIdent
   : T;
 
diff --git a/packages/typegpu/tests/buffer.test.ts b/packages/typegpu/tests/buffer.test.ts
index 00173744f..0626219b3 100644
--- a/packages/typegpu/tests/buffer.test.ts
+++ b/packages/typegpu/tests/buffer.test.ts
@@ -181,4 +181,56 @@ describe('TgpuBuffer', () => {
         .$usage('storage');
     }).toThrow();
   });
+
+  it('should be able to copy from a buffer identical on the byte level', ({
+    root,
+  }) => {
+    const buffer = root.createBuffer(d.u32);
+    const copy = root.createBuffer(d.atomic(d.u32));
+
+    buffer.copyFrom(copy);
+
+    const buffer2 = root.createBuffer(
+      d.struct({
+        one: d.location(0, d.f32), // does nothing outside an IO struct
+        two: d.atomic(d.u32),
+        three: d.arrayOf(d.u32, 3),
+      }),
+    );
+    const copy2 = root.createBuffer(
+      d.struct({
+        one: d.f32,
+        two: d.u32,
+        three: d.arrayOf(d.atomic(d.u32), 3),
+      }),
+    );
+
+    buffer2.copyFrom(copy2);
+
+    const buffer3 = root.createBuffer(
+      d.struct({
+        one: d.size(16, d.f32),
+        two: d.atomic(d.u32),
+      }),
+    );
+
+    const copy3 = root.createBuffer(
+      d.struct({
+        one: d.f32,
+        two: d.u32,
+      }),
+    );
+
+    const copy31 = root.createBuffer(
+      d.struct({
+        one: d.location(0, d.f32),
+        two: d.u32,
+      }),
+    );
+
+    // @ts-expect-error
+    buffer3.copyFrom(copy3);
+    // @ts-expect-error
+    buffer3.copyFrom(copy31);
+  });
 });

From 784b4dd8ab1540d553393419424585bfea0eb2c7 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Tue, 7 Jan 2025 12:59:01 +0100
Subject: [PATCH 3/9] better mem ident for decorated (i think???)

---
 packages/typegpu/src/data/attributes.ts | 4 +++-
 packages/typegpu/src/data/wgslTypes.ts  | 4 ++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/packages/typegpu/src/data/attributes.ts b/packages/typegpu/src/data/attributes.ts
index 1944340f6..49677d01d 100644
--- a/packages/typegpu/src/data/attributes.ts
+++ b/packages/typegpu/src/data/attributes.ts
@@ -268,7 +268,9 @@ class DecoratedImpl<TInner extends BaseWgslData, TAttribs extends unknown[]>
   implements Decorated<TInner, TAttribs>
 {
   public readonly type = 'decorated';
-  public readonly '~memIdent'!: MemIdentity<TInner>;
+  public readonly '~memIdent'!: TAttribs extends Location<number>[]
+    ? MemIdentity<TInner> | Decorated<MemIdentity<TInner>, TAttribs>
+    : Decorated<TInner, TAttribs>;
 }
 
 class LooseDecoratedImpl<
diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts
index 58ce77b27..94e9d8249 100644
--- a/packages/typegpu/src/data/wgslTypes.ts
+++ b/packages/typegpu/src/data/wgslTypes.ts
@@ -7,7 +7,7 @@ import type {
 
 type DecoratedLocation<T extends BaseWgslData> = Decorated<
   T,
-  Location<number>[] | []
+  Location<number>[]
 >;
 
 export interface NumberArrayView {
@@ -726,7 +726,7 @@ export interface Decorated<
   /** Type-token, not available at runtime */
   readonly '~repr': Infer<TInner>;
   readonly '~memIdent': TAttribs extends Location<number>[]
-    ? MemIdentity<TInner>
+    ? MemIdentity<TInner> | Decorated<MemIdentity<TInner>, TAttribs>
     : Decorated<TInner, TAttribs>;
 }
 

From 5df189f53925bfd1910f36cbdc3f974b5988f376 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Fri, 10 Jan 2025 10:19:14 +0100
Subject: [PATCH 4/9] fix import

---
 packages/typegpu/src/shared/repr.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/typegpu/src/shared/repr.ts b/packages/typegpu/src/shared/repr.ts
index c9d73aef2..0a3bec5ea 100644
--- a/packages/typegpu/src/shared/repr.ts
+++ b/packages/typegpu/src/shared/repr.ts
@@ -1,4 +1,4 @@
-import type { AnyData } from '../data';
+import type { AnyData } from '../data/dataTypes';
 
 /**
  * Extracts the inferred representation of a resource.

From 8722abf9ed35f64b850951fccdcdb34cc4ebd226 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Mon, 20 Jan 2025 09:27:20 +0100
Subject: [PATCH 5/9] a lot of fixes

---
 .../simulation/fluid-with-atomics/index.ts    | 436 +++++++++---------
 1 file changed, 224 insertions(+), 212 deletions(-)

diff --git a/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts b/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
index 294997108..af58b5305 100644
--- a/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
+++ b/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
@@ -1,5 +1,9 @@
+import tgpu, {
+  unstable_asReadonly,
+  unstable_asMutable,
+  unstable_asUniform,
+} from 'typegpu';
 import * as d from 'typegpu/data';
-import tgpu, { asMutable, asReadonly, asUniform } from 'typegpu/experimental';
 
 const root = await tgpu.init();
 
@@ -48,25 +52,25 @@ function encodeBrushType(brushType: (typeof BrushTypes)[number]) {
 }
 
 const sizeBuffer = root.createBuffer(d.vec2u).$name('size').$usage('uniform');
-const sizeUniform = asUniform(sizeBuffer);
+const sizeUniform = unstable_asUniform(sizeBuffer);
 
 const viscosityBuffer = root
   .createBuffer(d.u32)
   .$name('viscosity')
   .$usage('uniform');
-const viscosityUniform = asUniform(viscosityBuffer);
+const viscosityUniform = unstable_asUniform(viscosityBuffer);
 
 const currentStateBuffer = root
   .createBuffer(d.arrayOf(d.u32, 1024 ** 2))
   .$name('current')
   .$usage('storage', 'vertex');
-const currentStateStorage = asReadonly(currentStateBuffer);
+const currentStateStorage = unstable_asReadonly(currentStateBuffer);
 
 const nextStateBuffer = root
   .createBuffer(d.arrayOf(d.atomic(d.u32), 1024 ** 2))
   .$name('next')
   .$usage('storage');
-const nextStateStorage = asMutable(nextStateBuffer);
+const nextStateStorage = unstable_asMutable(nextStateBuffer);
 
 const squareBuffer = root
   .createBuffer(d.arrayOf(d.vec2f, 4), [
@@ -78,143 +82,145 @@ const squareBuffer = root
   .$usage('vertex')
   .$name('square');
 
-const getIndex = tgpu
+const getIndex = tgpu['~unstable']
   .fn([d.u32, d.u32], d.u32)
   .does(/* wgsl */ `(x: u32, y: u32) -> u32 {
-  let h = sizeData.y;
-  let w = sizeData.x;
-  return (y % h) * w + (x % w);
-}`)
+    let h = sizeData.y;
+    let w = sizeData.x;
+    return (y % h) * w + (x % w);
+  }`)
   .$uses({ sizeData: sizeUniform });
 
-const getCell = tgpu
+const getCell = tgpu['~unstable']
   .fn([d.u32, d.u32], d.u32)
   .does(/* wgsl */ `(x: u32, y: u32) -> u32 {
-  return currentStateData[getIndex(x, y)];
-}`)
+    return currentStateData[getIndex(x, y)];
+  }`)
   .$uses({ currentStateData: currentStateStorage, getIndex });
 
-const getCellNext = tgpu
+const getCellNext = tgpu['~unstable']
   .fn([d.u32, d.u32], d.u32)
   .does(/* wgsl */ `(x: u32, y: u32) -> u32 {
-  return atomicLoad(&nextStateData[getIndex(x, y)]);
-}`)
+    return atomicLoad(&nextStateData[getIndex(x, y)]);
+  }`)
   .$uses({ nextStateData: nextStateStorage, getIndex });
 
-const updateCell = tgpu
+const updateCell = tgpu['~unstable']
   .fn([d.u32, d.u32, d.u32])
   .does(/* wgsl */ `(x: u32, y: u32, value: u32) {
-  atomicStore(&nextStateData[getIndex(x, y)], value);
-}`)
+    atomicStore(&nextStateData[getIndex(x, y)], value);
+  }`)
   .$uses({ nextStateData: nextStateStorage, getIndex });
 
-const addToCell = tgpu
+const addToCell = tgpu['~unstable']
   .fn([d.u32, d.u32, d.u32])
   .does(/* wgsl */ `(x: u32, y: u32, value: u32) {
-  let cell = getCellNext(x, y);
-  let waterLevel = cell & ${MAX_WATER_LEVEL};
-  let newWaterLevel = min(waterLevel + value, ${MAX_WATER_LEVEL});
-  atomicAdd(&nextStateData[getIndex(x, y)], newWaterLevel - waterLevel);
-}`)
+    let cell = getCellNext(x, y);
+    let waterLevel = cell & ${MAX_WATER_LEVEL};
+    let newWaterLevel = min(waterLevel + value, ${MAX_WATER_LEVEL});
+    atomicAdd(&nextStateData[getIndex(x, y)], newWaterLevel - waterLevel);
+  }`)
   .$uses({ getCellNext, nextStateData: nextStateStorage, getIndex });
 
-const subtractFromCell = tgpu
+const subtractFromCell = tgpu['~unstable']
   .fn([d.u32, d.u32, d.u32])
   .does(/* wgsl */ `(x: u32, y: u32, value: u32) {
-  let cell = getCellNext(x, y);
-  let waterLevel = cell & ${MAX_WATER_LEVEL};
-  let newWaterLevel = max(waterLevel - min(value, waterLevel), 0u);
-  atomicSub(&nextStateData[getIndex(x, y)], waterLevel - newWaterLevel);
-}`)
+    let cell = getCellNext(x, y);
+    let waterLevel = cell & ${MAX_WATER_LEVEL};
+    let newWaterLevel = max(waterLevel - min(value, waterLevel), 0u);
+    atomicSub(&nextStateData[getIndex(x, y)], waterLevel - newWaterLevel);
+  }`)
   .$uses({ getCellNext, nextStateData: nextStateStorage, getIndex });
 
-const persistFlags = tgpu
+const persistFlags = tgpu['~unstable']
   .fn([d.u32, d.u32])
   .does(/* wgsl */ `(x: u32, y: u32) {
-  let cell = getCell(x, y);
-  let waterLevel = cell & ${MAX_WATER_LEVEL};
-  let flags = cell >> 24;
-  updateCell(x, y, (flags << 24) | waterLevel);
-}`)
+    let cell = getCell(x, y);
+    let waterLevel = cell & ${MAX_WATER_LEVEL};
+    let flags = cell >> 24;
+    updateCell(x, y, (flags << 24) | waterLevel);
+  }`)
   .$uses({ getCell, updateCell });
 
-const getStableStateBelow = tgpu.fn([d.u32, d.u32], d.u32).does(/* wgsl */ `(upper: u32, lower: u32) -> u32 {
-  let totalMass = upper + lower;
-  if (totalMass <= ${MAX_WATER_LEVEL_UNPRESSURIZED}) {
-    return totalMass;
-  } else if (totalMass >= ${MAX_WATER_LEVEL_UNPRESSURIZED}*2 && upper > lower) {
-    return totalMass/2 + ${MAX_PRESSURE};
-  }
-  return ${MAX_WATER_LEVEL_UNPRESSURIZED};
-}`);
+const getStableStateBelow = tgpu['~unstable']
+  .fn([d.u32, d.u32], d.u32)
+  .does(/* wgsl */ `(upper: u32, lower: u32) -> u32 {
+    let totalMass = upper + lower;
+    if (totalMass <= ${MAX_WATER_LEVEL_UNPRESSURIZED}) {
+      return totalMass;
+    } else if (totalMass >= ${MAX_WATER_LEVEL_UNPRESSURIZED}*2 && upper > lower) {
+      return totalMass/2 + ${MAX_PRESSURE};
+    }
+    return ${MAX_WATER_LEVEL_UNPRESSURIZED};
+  }`);
 
-const isWall = tgpu
+const isWall = tgpu['~unstable']
   .fn([d.u32, d.u32], d.bool)
   .does(/* wgsl */ `(x: u32, y: u32) -> bool {
-  return (getCell(x, y) >> 24) == 1u;
-}`)
+    return (getCell(x, y) >> 24) == 1u;
+  }`)
   .$uses({ getCell });
 
-const isWaterSource = tgpu
+const isWaterSource = tgpu['~unstable']
   .fn([d.u32, d.u32], d.bool)
   .does(/* wgsl */ `(x: u32, y: u32) -> bool {
-  return (getCell(x, y) >> 24) == 2u;
-}`)
+    return (getCell(x, y) >> 24) == 2u;
+  }`)
   .$uses({ getCell });
 
-const isWaterDrain = tgpu
+const isWaterDrain = tgpu['~unstable']
   .fn([d.u32, d.u32], d.bool)
   .does(/* wgsl */ `(x: u32, y: u32) -> bool {
-  return (getCell(x, y) >> 24) == 3u;
-}`)
+    return (getCell(x, y) >> 24) == 3u;
+  }`)
   .$uses({ getCell });
 
-const isClearCell = tgpu
+const isClearCell = tgpu['~unstable']
   .fn([d.u32, d.u32], d.bool)
   .does(/* wgsl */ `(x: u32, y: u32) -> bool {
-  return (getCell(x, y) >> 24) == 4u;
-}`)
+    return (getCell(x, y) >> 24) == 4u;
+  }`)
   .$uses({ getCell });
 
-const getWaterLevel = tgpu
+const getWaterLevel = tgpu['~unstable']
   .fn([d.u32, d.u32], d.u32)
   .does(/* wgsl */ `(x: u32, y: u32) -> u32 {
-  return getCell(x, y) & ${MAX_WATER_LEVEL};
-}`)
+    return getCell(x, y) & ${MAX_WATER_LEVEL};
+  }`)
   .$uses({ getCell });
 
-const checkForFlagsAndBounds = tgpu
+const checkForFlagsAndBounds = tgpu['~unstable']
   .fn([d.u32, d.u32], d.bool)
   .does(/* wgsl */ `(x: u32, y: u32) -> bool {
-  if (isClearCell(x, y)) {
-    updateCell(x, y, 0u);
-    return true;
-  }
+    if (isClearCell(x, y)) {
+      updateCell(x, y, 0u);
+      return true;
+    }
 
-  if (isWall(x, y)) {
-    persistFlags(x, y);
-    return true;
-  }
+    if (isWall(x, y)) {
+      persistFlags(x, y);
+      return true;
+    }
 
-  if (isWaterSource(x, y)) {
-    persistFlags(x, y);
-    addToCell(x, y, 20u);
-    return false;
-  }
+    if (isWaterSource(x, y)) {
+      persistFlags(x, y);
+      addToCell(x, y, 20u);
+      return false;
+    }
 
-  if (isWaterDrain(x, y)) {
-    persistFlags(x, y);
-    updateCell(x, y, 3u << 24);
-    return true;
-  }
+    if (isWaterDrain(x, y)) {
+      persistFlags(x, y);
+      updateCell(x, y, 3u << 24);
+      return true;
+    }
 
-  if (y == 0 || y == sizeData.y - 1u || x == 0 || x == sizeData.x - 1u) {
-    subtractFromCell(x, y, getWaterLevel(x, y));
-    return true;
-  }
+    if (y == 0 || y == sizeData.y - 1u || x == 0 || x == sizeData.x - 1u) {
+      subtractFromCell(x, y, getWaterLevel(x, y));
+      return true;
+    }
 
-  return false;
-}`)
+    return false;
+  }`)
   .$uses({
     isClearCell,
     updateCell,
@@ -228,76 +234,76 @@ const checkForFlagsAndBounds = tgpu
     sizeData: sizeUniform,
   });
 
-const decideWaterLevel = tgpu
+const decideWaterLevel = tgpu['~unstable']
   .fn([d.u32, d.u32])
   .does(/* wgsl */ `(x: u32, y: u32) {
-  if (checkForFlagsAndBounds(x, y)) {
-    return;
-  }
+    if (checkForFlagsAndBounds(x, y)) {
+      return;
+    }
 
-  var remainingWater = getWaterLevel(x, y);
+    var remainingWater = getWaterLevel(x, y);
 
-  if (remainingWater == 0u) {
-    return;
-  }
+    if (remainingWater == 0u) {
+      return;
+    }
 
-  if (!isWall(x, y - 1u)) {
-    let waterLevelBelow = getWaterLevel(x, y - 1u);
-    let stable = getStableStateBelow(remainingWater, waterLevelBelow);
-    if (waterLevelBelow < stable) {
-      let change = stable - waterLevelBelow;
-      let flow = min(change, viscosityData);
-      subtractFromCell(x, y, flow);
-      addToCell(x, y - 1u, flow);
-      remainingWater -= flow;
+    if (!isWall(x, y - 1u)) {
+      let waterLevelBelow = getWaterLevel(x, y - 1u);
+      let stable = getStableStateBelow(remainingWater, waterLevelBelow);
+      if (waterLevelBelow < stable) {
+        let change = stable - waterLevelBelow;
+        let flow = min(change, viscosityData);
+        subtractFromCell(x, y, flow);
+        addToCell(x, y - 1u, flow);
+        remainingWater -= flow;
+      }
     }
-  }
 
-  if (remainingWater == 0u) {
-    return;
-  }
+    if (remainingWater == 0u) {
+      return;
+    }
 
-  let waterLevelBefore = remainingWater;
-  if (!isWall(x - 1u, y)) {
-    let flowRaw = (i32(waterLevelBefore) - i32(getWaterLevel(x - 1u, y)));
-    if (flowRaw > 0) {
-      let change = max(min(4u, remainingWater), u32(flowRaw)/4);
-      let flow = min(change, viscosityData);
-      subtractFromCell(x, y, flow);
-      addToCell(x - 1u, y, flow);
-      remainingWater -= flow;
+    let waterLevelBefore = remainingWater;
+    if (!isWall(x - 1u, y)) {
+      let flowRaw = (i32(waterLevelBefore) - i32(getWaterLevel(x - 1u, y)));
+      if (flowRaw > 0) {
+        let change = max(min(4u, remainingWater), u32(flowRaw)/4);
+        let flow = min(change, viscosityData);
+        subtractFromCell(x, y, flow);
+        addToCell(x - 1u, y, flow);
+        remainingWater -= flow;
+      }
     }
-  }
 
-  if (remainingWater == 0u) {
-    return;
-  }
+    if (remainingWater == 0u) {
+      return;
+    }
 
-  if (!isWall(x + 1u, y)) {
-    let flowRaw = (i32(waterLevelBefore) - i32(getWaterLevel(x + 1, y)));
-    if (flowRaw > 0) {
-      let change = max(min(4u, remainingWater), u32(flowRaw)/4);
-      let flow = min(change, viscosityData);
-      subtractFromCell(x, y, flow);
-      addToCell(x + 1u, y, flow);
-      remainingWater -= flow;
+    if (!isWall(x + 1u, y)) {
+      let flowRaw = (i32(waterLevelBefore) - i32(getWaterLevel(x + 1, y)));
+      if (flowRaw > 0) {
+        let change = max(min(4u, remainingWater), u32(flowRaw)/4);
+        let flow = min(change, viscosityData);
+        subtractFromCell(x, y, flow);
+        addToCell(x + 1u, y, flow);
+        remainingWater -= flow;
+      }
     }
-  }
 
-  if (remainingWater == 0u) {
-    return;
-  }
+    if (remainingWater == 0u) {
+      return;
+    }
 
-  if (!isWall(x, y + 1u)) {
-    let stable = getStableStateBelow(getWaterLevel(x, y + 1u), remainingWater);
-    if (stable < remainingWater) {
-      let flow = min(remainingWater - stable, viscosityData);
-      subtractFromCell(x, y, flow);
-      addToCell(x, y + 1u, flow);
-      remainingWater -= flow;
+    if (!isWall(x, y + 1u)) {
+      let stable = getStableStateBelow(getWaterLevel(x, y + 1u), remainingWater);
+      if (stable < remainingWater) {
+        let flow = min(remainingWater - stable, viscosityData);
+        subtractFromCell(x, y, flow);
+        addToCell(x, y + 1u, flow);
+        remainingWater -= flow;
+      }
     }
-  }
-}`)
+  }`)
   .$uses({
     checkForFlagsAndBounds,
     getWaterLevel,
@@ -309,7 +315,7 @@ const decideWaterLevel = tgpu
     viscosityData: viscosityUniform,
   });
 
-const vertex = tgpu
+const vertex = tgpu['~unstable']
   .vertexFn(
     {
       squareData: d.vec2f,
@@ -319,54 +325,55 @@ const vertex = tgpu
     { pos: d.builtin.position, cell: d.f32 },
   )
   .does(/* wgsl */ `(@builtin(instance_index) idx: u32, @location(0) squareData: vec2f, @location(1) currentStateData: u32) -> VertexOut {
-  let w = sizeData.x;
-  let h = sizeData.y;
-  let x = ((f32(idx % w) + squareData.x) / f32(w) - 0.5) * 2. * f32(w) / f32(max(w, h));
-  let y = (f32((idx - (idx % w)) / w + u32(squareData.y)) / f32(h) - 0.5) * 2. * f32(h) / f32(max(w, h));
-  let cellFlags = currentStateData >> 24;
-  var cell = f32(currentStateData & 0xFFFFFF);
-  if (cellFlags == 1u) {
-    cell = -1.;
-  }
-  if (cellFlags == 2u) {
-    cell = -2.;
-  }
-  if (cellFlags == 3u) {
-    cell = -3.;
-  }
-  return VertexOut(vec4f(x, y, 0., 1.), cell);
-}`)
+    let w = sizeData.x;
+    let h = sizeData.y;
+    let x = ((f32(idx % w) + squareData.x) / f32(w) - 0.5) * 2. * f32(w) / f32(max(w, h));
+    let y = (f32((idx - (idx % w)) / w + u32(squareData.y)) / f32(h) - 0.5) * 2. * f32(h) / f32(max(w, h));
+    let cellFlags = currentStateData >> 24;
+    var cell = f32(currentStateData & 0xFFFFFF);
+    if (cellFlags == 1u) {
+      cell = -1.;
+    }
+    if (cellFlags == 2u) {
+      cell = -2.;
+    }
+    if (cellFlags == 3u) {
+      cell = -3.;
+    }
+    return VertexOut(vec4f(x, y, 0., 1.), cell);
+  }`)
   .$uses({
     sizeData: sizeUniform,
   });
 
-const fragment = tgpu.fragmentFn({ cell: d.f32 }, d.vec4f).does(/* wgsl */ `(@location(0) cell: f32) -> @location(0) vec4f {
-  if (cell == -1.) {
-    return vec4f(0.5, 0.5, 0.5, 1.);
-  }
-  if (cell == -2.) {
-    return vec4f(0., 1., 0., 1.);
-  }
-  if (cell == -3.) {
-    return vec4f(1., 0., 0., 1.);
-  }
-
-  let normalized = min(cell / f32(0xFF), 1.);
+const fragment = tgpu['~unstable']
+  .fragmentFn({ cell: d.f32 }, d.vec4f)
+  .does(/* wgsl */ `(@location(0) cell: f32) -> @location(0) vec4f {
+    if (cell == -1.) {
+      return vec4f(0.5, 0.5, 0.5, 1.);
+    }
+    if (cell == -2.) {
+      return vec4f(0., 1., 0., 1.);
+    }
+    if (cell == -3.) {
+      return vec4f(1., 0., 0., 1.);
+    }
 
-  if (normalized == 0.) {
-    return vec4f(0.937, 0.937, 0.976, 1.);
-  }
+    let normalized = min(cell / f32(0xFF), 1.);
 
-  let res = 1. / (1. + exp(-(normalized - 0.2) * 10.));
-  return vec4f(0, 0, max(0.5, res), res);
+    if (normalized == 0.) {
+      return vec4f();
+    }
 
-}`);
+    let res = 1. / (1. + exp(-(normalized - 0.2) * 10.));
+    return vec4f(0, 0, max(0.5, res), res);
+  }`);
 
-const vertexInstanceLayout = tgpu.vertexLayout(
+const vertexInstanceLayout = tgpu['~unstable'].vertexLayout(
   (n: number) => d.arrayOf(d.u32, n),
   'instance',
 );
-const vertexLayout = tgpu.vertexLayout(
+const vertexLayout = tgpu['~unstable'].vertexLayout(
   (n: number) => d.arrayOf(d.vec2f, n),
   'vertex',
 );
@@ -381,17 +388,19 @@ let renderChanges: () => void;
 function resetGameData() {
   drawCanvasData = new Uint32Array(options.size * options.size);
 
-  const compute = tgpu
+  const compute = tgpu['~unstable']
     .computeFn([d.builtin.globalInvocationId], {
       workgroupSize: [options.workgroupSize, options.workgroupSize],
     })
     .does(/* wgsl */ `(@builtin(global_invocation_id) gid: vec3u) {
-    decideWaterLevel(gid.x, gid.y);
-  }`)
+      decideWaterLevel(gid.x, gid.y);
+    }`)
     .$uses({ decideWaterLevel });
 
-  const cpp = root.withCompute(compute).createPipeline();
-  const vp = root
+  const computePipeline = root['~unstable']
+    .withCompute(compute)
+    .createPipeline();
+  const renderPipeline = root['~unstable']
     .withVertex(vertex, {
       squareData: vertexLayout.attrib,
       currentStateData: vertexInstanceLayout.attrib,
@@ -412,20 +421,22 @@ function resetGameData() {
 
   render = () => {
     // compute
-    cpp.dispatchWorkgroups(
+    computePipeline.dispatchWorkgroups(
       options.size / options.workgroupSize,
       options.size / options.workgroupSize,
     );
 
     // render
-    vp.withColorAttachment({
-      view: context.getCurrentTexture().createView(),
-      clearValue: [0, 0, 0, 0],
-      loadOp: 'clear' as const,
-      storeOp: 'store' as const,
-    }).draw(4, options.size ** 2);
+    renderPipeline
+      .withColorAttachment({
+        view: context.getCurrentTexture().createView(),
+        clearValue: [0, 0, 0, 0],
+        loadOp: 'clear' as const,
+        storeOp: 'store' as const,
+      })
+      .draw(4, options.size ** 2);
 
-    root.flush();
+    root['~unstable'].flush();
 
     currentStateBuffer.copyFrom(nextStateBuffer);
   };
@@ -441,7 +452,7 @@ function resetGameData() {
 
         const index = j * options.size + i;
         root.device.queue.writeBuffer(
-          currentStateBuffer.buffer,
+          nextStateBuffer.buffer,
           index * Uint32Array.BYTES_PER_ELEMENT,
           drawCanvasData,
           index,
@@ -455,16 +466,17 @@ function resetGameData() {
   };
 
   renderChanges = () => {
-    vp.withColorAttachment({
-      view: context.getCurrentTexture().createView(),
-      clearValue: [0, 0, 0, 0],
-      loadOp: 'clear' as const,
-      storeOp: 'store' as const,
-    })
+    renderPipeline
+      .withColorAttachment({
+        view: context.getCurrentTexture().createView(),
+        clearValue: [0, 0, 0, 0],
+        loadOp: 'clear' as const,
+        storeOp: 'store' as const,
+      })
       .with(vertexLayout, squareBuffer)
       .with(vertexInstanceLayout, currentStateBuffer)
       .draw(4, options.size ** 2);
-    root.flush();
+    root['~unstable'].flush();
   };
 
   createSampleScene();
@@ -496,7 +508,7 @@ canvas.onmousemove = (event) => {
     options.size -
     Math.floor((event.offsetY * window.devicePixelRatio) / cellSize) -
     1;
-  const allAffectedCells = [];
+  const allAffectedCells = new Set<{ x: number; y: number }>();
   for (let i = -options.brushSize; i <= options.brushSize; i++) {
     for (let j = -options.brushSize; j <= options.brushSize; j++) {
       if (
@@ -506,7 +518,7 @@ canvas.onmousemove = (event) => {
         y + j >= 0 &&
         y + j < options.size
       ) {
-        allAffectedCells.push({ x: x + i, y: y + j });
+        allAffectedCells.add({ x: x + i, y: y + j });
       }
     }
   }
@@ -632,23 +644,23 @@ export const controls = {
   },
 
   viscosity: {
-    initial: 1000,
-    min: 10,
-    max: 1000,
-    step: 1,
+    initial: 0,
+    min: 0,
+    max: 1,
+    step: 0.01,
     onSliderChange: (value: number) => {
-      options.viscosity = value;
-      viscosityBuffer.write(value);
+      options.viscosity = 1000 - value * 990;
+      viscosityBuffer.write(options.viscosity);
     },
   },
 
   'brush size': {
-    initial: 0,
-    min: 0,
-    max: 10,
+    initial: 1,
+    min: 1,
+    max: 20,
     step: 1,
     onSliderChange: (value: number) => {
-      options.brushSize = value;
+      options.brushSize = value - 1;
     },
   },
 

From 84ee1f4d07115edc7ba93a74fe374c809a2ae176 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Mon, 20 Jan 2025 09:28:12 +0100
Subject: [PATCH 6/9] smol fix

---
 .../src/content/examples/simulation/fluid-with-atomics/index.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts b/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
index af58b5305..653751a1f 100644
--- a/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
+++ b/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
@@ -148,7 +148,7 @@ const getStableStateBelow = tgpu['~unstable']
     let totalMass = upper + lower;
     if (totalMass <= ${MAX_WATER_LEVEL_UNPRESSURIZED}) {
       return totalMass;
-    } else if (totalMass >= ${MAX_WATER_LEVEL_UNPRESSURIZED}*2 && upper > lower) {
+    } else if (totalMass >= ${MAX_WATER_LEVEL_UNPRESSURIZED} * 2 && upper > lower) {
       return totalMass/2 + ${MAX_PRESSURE};
     }
     return ${MAX_WATER_LEVEL_UNPRESSURIZED};

From 9d3bf983bcc8b7925943cc22e99154df5cdb8a90 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Mon, 20 Jan 2025 09:34:07 +0100
Subject: [PATCH 7/9] use tgpu.const

---
 .../simulation/fluid-with-atomics/index.ts    | 45 ++++++++++++-------
 1 file changed, 28 insertions(+), 17 deletions(-)

diff --git a/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts b/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
index 653751a1f..723cd963d 100644
--- a/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
+++ b/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
@@ -23,9 +23,9 @@ canvas.addEventListener('contextmenu', (event) => {
   }
 });
 
-const MAX_WATER_LEVEL_UNPRESSURIZED = 'u32(0xFF)';
-const MAX_WATER_LEVEL = '((1u << 24) - 1u)';
-const MAX_PRESSURE = '12u';
+const MAX_WATER_LEVEL_UNPRESSURIZED = tgpu['~unstable'].const(d.u32, 0xff);
+const MAX_WATER_LEVEL = tgpu['~unstable'].const(d.u32, (1 << 24) - 1);
+const MAX_PRESSURE = tgpu['~unstable'].const(d.u32, 12);
 
 const options = {
   size: 32,
@@ -116,43 +116,54 @@ const addToCell = tgpu['~unstable']
   .fn([d.u32, d.u32, d.u32])
   .does(/* wgsl */ `(x: u32, y: u32, value: u32) {
     let cell = getCellNext(x, y);
-    let waterLevel = cell & ${MAX_WATER_LEVEL};
-    let newWaterLevel = min(waterLevel + value, ${MAX_WATER_LEVEL});
+    let waterLevel = cell & MAX_WATER_LEVEL;
+    let newWaterLevel = min(waterLevel + value, MAX_WATER_LEVEL);
     atomicAdd(&nextStateData[getIndex(x, y)], newWaterLevel - waterLevel);
   }`)
-  .$uses({ getCellNext, nextStateData: nextStateStorage, getIndex });
+  .$uses({
+    getCellNext,
+    nextStateData: nextStateStorage,
+    getIndex,
+    MAX_WATER_LEVEL,
+  });
 
 const subtractFromCell = tgpu['~unstable']
   .fn([d.u32, d.u32, d.u32])
   .does(/* wgsl */ `(x: u32, y: u32, value: u32) {
     let cell = getCellNext(x, y);
-    let waterLevel = cell & ${MAX_WATER_LEVEL};
+    let waterLevel = cell & MAX_WATER_LEVEL;
     let newWaterLevel = max(waterLevel - min(value, waterLevel), 0u);
     atomicSub(&nextStateData[getIndex(x, y)], waterLevel - newWaterLevel);
   }`)
-  .$uses({ getCellNext, nextStateData: nextStateStorage, getIndex });
+  .$uses({
+    getCellNext,
+    nextStateData: nextStateStorage,
+    getIndex,
+    MAX_WATER_LEVEL,
+  });
 
 const persistFlags = tgpu['~unstable']
   .fn([d.u32, d.u32])
   .does(/* wgsl */ `(x: u32, y: u32) {
     let cell = getCell(x, y);
-    let waterLevel = cell & ${MAX_WATER_LEVEL};
+    let waterLevel = cell & MAX_WATER_LEVEL;
     let flags = cell >> 24;
     updateCell(x, y, (flags << 24) | waterLevel);
   }`)
-  .$uses({ getCell, updateCell });
+  .$uses({ getCell, updateCell, MAX_WATER_LEVEL });
 
 const getStableStateBelow = tgpu['~unstable']
   .fn([d.u32, d.u32], d.u32)
   .does(/* wgsl */ `(upper: u32, lower: u32) -> u32 {
     let totalMass = upper + lower;
-    if (totalMass <= ${MAX_WATER_LEVEL_UNPRESSURIZED}) {
+    if (totalMass <= MAX_WATER_LEVEL_UNPRESSURIZED) {
       return totalMass;
-    } else if (totalMass >= ${MAX_WATER_LEVEL_UNPRESSURIZED} * 2 && upper > lower) {
-      return totalMass/2 + ${MAX_PRESSURE};
+    } else if (totalMass >= MAX_WATER_LEVEL_UNPRESSURIZED * 2 && upper > lower) {
+      return totalMass/2 + MAX_PRESSURE;
     }
-    return ${MAX_WATER_LEVEL_UNPRESSURIZED};
-  }`);
+    return MAX_WATER_LEVEL_UNPRESSURIZED;
+  }`)
+  .$uses({ MAX_PRESSURE, MAX_WATER_LEVEL_UNPRESSURIZED });
 
 const isWall = tgpu['~unstable']
   .fn([d.u32, d.u32], d.bool)
@@ -185,9 +196,9 @@ const isClearCell = tgpu['~unstable']
 const getWaterLevel = tgpu['~unstable']
   .fn([d.u32, d.u32], d.u32)
   .does(/* wgsl */ `(x: u32, y: u32) -> u32 {
-    return getCell(x, y) & ${MAX_WATER_LEVEL};
+    return getCell(x, y) &MAX_WATER_LEVEL;
   }`)
-  .$uses({ getCell });
+  .$uses({ getCell, MAX_WATER_LEVEL });
 
 const checkForFlagsAndBounds = tgpu['~unstable']
   .fn([d.u32, d.u32], d.bool)

From 94c9d501c3b815dc406db5075f27f77152d1e2a7 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Mon, 20 Jan 2025 10:52:37 +0100
Subject: [PATCH 8/9] remove redundant command encoder (don't know why it took
 me so long to notice this)

---
 .../content/examples/simulation/fluid-with-atomics/index.ts    | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts b/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
index 723cd963d..e69a3107e 100644
--- a/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
+++ b/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
@@ -453,8 +453,6 @@ function resetGameData() {
   };
 
   applyDrawCanvas = () => {
-    const commandEncoder = root.device.createCommandEncoder();
-
     for (let i = 0; i < options.size; i++) {
       for (let j = 0; j < options.size; j++) {
         if (drawCanvasData[j * options.size + i] === 0) {
@@ -472,7 +470,6 @@ function resetGameData() {
       }
     }
 
-    root.device.queue.submit([commandEncoder.finish()]);
     drawCanvasData.fill(0);
   };
 

From 62239e3d0c41aea4285b80db9b7b0e40ec9786db Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Mon, 27 Jan 2025 09:38:14 +0100
Subject: [PATCH 9/9] handle mobile + extra test case

---
 .../simulation/fluid-with-atomics/index.ts    | 84 +++++++++++++++----
 packages/typegpu/src/data/attributes.ts       |  2 +-
 packages/typegpu/src/data/wgslTypes.ts        |  2 +-
 packages/typegpu/tests/buffer.test.ts         |  9 ++
 4 files changed, 81 insertions(+), 16 deletions(-)

diff --git a/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts b/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
index e69a3107e..5a87a2103 100644
--- a/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
+++ b/apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts
@@ -494,28 +494,24 @@ function resetGameData() {
 
 let isDrawing = false;
 let isErasing = false;
+let longTouchTimeout: number | null = null;
+let touchMoved = false;
 
-canvas.onmousedown = (event) => {
+const startDrawing = (erase: boolean) => {
   isDrawing = true;
-  isErasing = event.button === 2;
+  isErasing = erase;
 };
 
-canvas.onmouseup = () => {
+const stopDrawing = () => {
   isDrawing = false;
   renderChanges();
-};
-
-canvas.onmousemove = (event) => {
-  if (!isDrawing) {
-    return;
+  if (longTouchTimeout) {
+    clearTimeout(longTouchTimeout);
+    longTouchTimeout = null;
   }
+};
 
-  const cellSize = canvas.width / options.size;
-  const x = Math.floor((event.offsetX * window.devicePixelRatio) / cellSize);
-  const y =
-    options.size -
-    Math.floor((event.offsetY * window.devicePixelRatio) / cellSize) -
-    1;
+const handleDrawing = (x: number, y: number) => {
   const allAffectedCells = new Set<{ x: number; y: number }>();
   for (let i = -options.brushSize; i <= options.brushSize; i++) {
     for (let j = -options.brushSize; j <= options.brushSize; j++) {
@@ -547,6 +543,66 @@ canvas.onmousemove = (event) => {
   renderChanges();
 };
 
+canvas.onmousedown = (event) => {
+  startDrawing(event.button === 2);
+};
+
+canvas.onmouseup = stopDrawing;
+
+canvas.onmousemove = (event) => {
+  if (!isDrawing) {
+    return;
+  }
+
+  const cellSize = canvas.width / options.size;
+  const x = Math.floor((event.offsetX * window.devicePixelRatio) / cellSize);
+  const y =
+    options.size -
+    Math.floor((event.offsetY * window.devicePixelRatio) / cellSize) -
+    1;
+
+  handleDrawing(x, y);
+};
+
+canvas.ontouchstart = (event) => {
+  event.preventDefault();
+  touchMoved = false;
+  longTouchTimeout = window.setTimeout(() => {
+    if (!touchMoved) {
+      startDrawing(true);
+    }
+  }, 500);
+  startDrawing(false);
+};
+
+canvas.ontouchend = (event) => {
+  event.preventDefault();
+  stopDrawing();
+};
+
+canvas.ontouchmove = (event) => {
+  event.preventDefault();
+  touchMoved = true;
+  if (!isDrawing) {
+    return;
+  }
+
+  const touch = event.touches[0];
+  const cellSize = canvas.width / options.size;
+  const canvasPos = canvas.getBoundingClientRect();
+  const x = Math.floor(
+    ((touch.clientX - canvasPos.left) * window.devicePixelRatio) / cellSize,
+  );
+  const y =
+    options.size -
+    Math.floor(
+      ((touch.clientY - canvasPos.top) * window.devicePixelRatio) / cellSize,
+    ) -
+    1;
+
+  handleDrawing(x, y);
+};
+
 const createSampleScene = () => {
   const middlePoint = Math.floor(options.size / 2);
   const radius = Math.floor(options.size / 8);
diff --git a/packages/typegpu/src/data/attributes.ts b/packages/typegpu/src/data/attributes.ts
index f545f8c11..428693e41 100644
--- a/packages/typegpu/src/data/attributes.ts
+++ b/packages/typegpu/src/data/attributes.ts
@@ -339,7 +339,7 @@ class DecoratedImpl<TInner extends BaseWgslData, TAttribs extends unknown[]>
   public readonly type = 'decorated';
   public readonly '~memIdent'!: TAttribs extends Location<number>[]
     ? MemIdentity<TInner> | Decorated<MemIdentity<TInner>, TAttribs>
-    : Decorated<TInner, TAttribs>;
+    : Decorated<MemIdentity<TInner>, TAttribs>;
 }
 
 class LooseDecoratedImpl<
diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts
index 3c6b9c8a7..e2266ea12 100644
--- a/packages/typegpu/src/data/wgslTypes.ts
+++ b/packages/typegpu/src/data/wgslTypes.ts
@@ -758,7 +758,7 @@ export interface Decorated<
   readonly '~repr': Infer<TInner>;
   readonly '~memIdent': TAttribs extends Location<number>[]
     ? MemIdentity<TInner> | Decorated<MemIdentity<TInner>, TAttribs>
-    : Decorated<TInner, TAttribs>;
+    : Decorated<MemIdentity<TInner>, TAttribs>;
 }
 
 export const wgslTypeLiterals = [
diff --git a/packages/typegpu/tests/buffer.test.ts b/packages/typegpu/tests/buffer.test.ts
index 5fda1c9af..500ffdf87 100644
--- a/packages/typegpu/tests/buffer.test.ts
+++ b/packages/typegpu/tests/buffer.test.ts
@@ -228,9 +228,18 @@ describe('TgpuBuffer', () => {
       }),
     );
 
+    const copy32 = root.createBuffer(
+      d.struct({
+        one: d.size(12, d.f32),
+        two: d.u32,
+      }),
+    );
+
     // @ts-expect-error
     buffer3.copyFrom(copy3);
     // @ts-expect-error
     buffer3.copyFrom(copy31);
+    // @ts-expect-error
+    buffer3.copyFrom(copy32);
   });
 });