Merge pull request #8724 from processing/fix/update-storage-swizzle

Fix swizzle assignment for array element properties in compute shaders

Author: davepagurek

src/strands/ir_builders.js

@@ -689,7 +689,28 @@ export function createStructArrayElementProxy(strandsContext, bufferNode, indexN
         });
         const id = DAG.getOrCreateNode(dag, nodeData);
         CFG.recordInBasicBlock(cfg, cfg.currentBlock, id);
-        return createStrandsNode(id, field.dim, strandsContext);
+        // When a swizzle assignment fires (e.g. buf[i].vel.y *= -1), onRebind
+        // receives the new vector ID and writes it back to the buffer field,
+        // equivalent to buf[i].vel = newVec.
+        const onRebind = (newFieldID) => {
+          const accessData = DAG.createNodeData({
+            nodeType: NodeType.OPERATION,
+            opCode: OpCode.Binary.ARRAY_ACCESS,
+            dependsOn: [bufferNode.id, index.id],
+            dimension: field.dim,
+            baseType: BaseType.FLOAT,
+            identifier: field.name,
+          });
+          const accessID = DAG.getOrCreateNode(dag, accessData);
+          const assignData = DAG.createNodeData({
+            nodeType: NodeType.ASSIGNMENT,
+            dependsOn: [accessID, newFieldID],
+            phiBlocks: [],
+          });
+          const assignID = DAG.getOrCreateNode(dag, assignData);
+          CFG.recordInBasicBlock(cfg, cfg.currentBlock, assignID);
+        };
+        return createStrandsNode(id, field.dim, strandsContext, onRebind);
       },
       set(val) {
         // Create access node as assignment target (field name in identifier)

test/unit/visual/cases/webgpu.js

@@ -1422,6 +1422,88 @@ visualSuite("WebGPU", function () {
         await screenshot();
       }
     );
+
+    visualTest(
+      'Compute shader assigns to a swizzle of a struct vector field',
+      async function(p5, screenshot) {
+        await p5.createCanvas(50, 50, p5.WEBGPU);
+
+        const particles = p5.createStorage([
+          { position: [15, 10] },
+        ]);
+
+        // Negate position.y via swizzle assignment
+        const computeShader = p5.buildComputeShader(() => {
+          const buf = p5.uniformStorage('buf', particles);
+          const idx = p5.index.x;
+          buf[idx].position.y *= -1;
+        }, { p5, particles });
+        p5.compute(computeShader, 1);
+
+        const sphereShader = p5.baseMaterialShader().modify(() => {
+          const buf = p5.uniformStorage('buf', particles);
+          p5.getWorldInputs((inputs) => {
+            const pos = buf[0].position;
+            inputs.position.x += pos.x;
+            inputs.position.y += pos.y;
+            return inputs;
+          });
+        }, { p5, particles });
+
+        const geo = p5.buildGeometry(() => p5.sphere(5));
+        p5.background(200);
+        p5.noStroke();
+        p5.fill(255, 0, 0);
+        p5.shader(sphereShader);
+        p5.model(geo, 1);
+
+        await screenshot();
+      }
+    );
+
+    visualTest(
+      'Compute shader assigns to a swizzle of a struct vector field inside an if statement',
+      async function(p5, screenshot) {
+        await p5.createCanvas(50, 50, p5.WEBGPU);
+
+        const particles = p5.createStorage([
+          { position: [0, 0], velocity: [5, 5] },
+        ]);
+
+        // Move by velocity, then negate velocity.y if position.y > 0.
+        // After 1st run: position=[5,5], velocity=[5,-5].
+        // After 2nd run: position=[10,0], velocity=[5,-5].
+        const computeShader = p5.buildComputeShader(() => {
+          const buf = p5.uniformStorage('buf', particles);
+          const idx = p5.index.x;
+          buf[idx].position += buf[idx].velocity;
+          if (buf[idx].position.y > 0) {
+            buf[idx].velocity.y *= -1;
+          }
+        }, { p5, particles });
+        p5.compute(computeShader, 1);
+        p5.compute(computeShader, 1);
+
+        const sphereShader = p5.baseMaterialShader().modify(() => {
+          const buf = p5.uniformStorage('buf', particles);
+          p5.getWorldInputs((inputs) => {
+            const pos = buf[0].position;
+            inputs.position.x += pos.x;
+            inputs.position.y += pos.y;
+            return inputs;
+          });
+        }, { p5, particles });
+
+        const geo = p5.buildGeometry(() => p5.sphere(5));
+        p5.background(200);
+        p5.noStroke();
+        p5.fill(255, 0, 0);
+        p5.shader(sphereShader);
+        p5.model(geo, 1);
+
+        await screenshot();
+      }
+    );
   });
 
   visualSuite('Feedback', function() {

test/unit/visual/screenshots/WebGPU/Compute shaders/Compute shader assigns to a swizzle of a struct vector field inside an if statement/000.png

@@ -0,0 +1,3 @@
+{
+  "numScreenshots": 1
+}

test/unit/visual/screenshots/WebGPU/Compute shaders/Compute shader assigns to a swizzle of a struct vector field inside an if statement/metadata.json

@@ -0,0 +1,3 @@
+{
+  "numScreenshots": 1
+}