Merge pull request #8726 from aashu2006/feat/storagebuffer-read

feat(webgpu): add read() to p5.StorageBuffer with tests

Author: davepagurek

src/webgpu/p5.RendererWebGPU.js

@@ -173,6 +173,78 @@ function rendererWebGPU(p5, fn) {
         device.queue.writeBuffer(this.buffer, 0, floatData);
       }
     }
+
+    /**
+     * Reads data from a storage buffer back into JavaScript.
+     *
+     * Copies data from the GPU to the CPU using a temporary buffer,
+     * so it must be awaited. Returns a `Float32Array` for number
+     * buffers, or an array of plain objects for struct buffers.
+     * 
+     * Note: This is a GPU -> CPU read, so calling it often (like every frame)
+     * can be slow.
+     *
+     * ```js example
+     * let data;
+     * let computeShader;
+     *
+     * async function setup() {
+     *   await createCanvas(100, 100, WEBGPU);
+     *
+     *   data = createStorage(new Float32Array([1, 2, 3, 4]));
+     *   computeShader = buildComputeShader(doubleValues);
+     *   compute(computeShader, 4);
+     *
+     *   let result = await data.read();
+     *   // result is Float32Array [2, 4, 6, 8]
+     *   for (let i = 0; i < result.length; i++) {
+     *     print(result[i]);
+     *   }
+     *   describe('Prints the values 2, 4, 6, 8 to the console.');
+     * }
+     *
+     * function doubleValues() {
+     *   let d = uniformStorage(data);
+     *   let idx = index.x;
+     *   d[idx] = d[idx] * 2;
+     * }
+     * ```
+     *
+     * @method read
+     * @for p5.StorageBuffer
+     * @beta
+     * @webgpu
+     * @webgpuOnly
+     * @returns {Promise<Float32Array|Object[]>}
+     */
+    async read() {
+      const device = this._renderer.device;
+      this._renderer.flushDraw();
+
+      const stagingBuffer = device.createBuffer({
+        size: this.size,
+        usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
+      });
+
+      const commandEncoder = device.createCommandEncoder();
+      commandEncoder.copyBufferToBuffer(this.buffer, 0, stagingBuffer, 0, this.size);
+      device.queue.submit([commandEncoder.finish()]);
+
+      await stagingBuffer.mapAsync(GPUMapMode.READ, 0, this.size);
+      const mappedRange = stagingBuffer.getMappedRange(0, this.size);
+
+      // Copy before unmapping because mapped memory becomes invalid after unmap
+      const rawCopy = new Float32Array(mappedRange.byteLength / 4);
+      rawCopy.set(new Float32Array(mappedRange));
+
+      stagingBuffer.unmap();
+      stagingBuffer.destroy();
+
+      if (this._schema !== null) {
+        return this._renderer._unpackStructArray(rawCopy, this._schema);
+      }
+      return rawCopy;
+    }
   }
 
   /**
@@ -3243,19 +3315,24 @@ ${hookUniformFields}}
 
       let maxEnd = 0;
       let maxAlign = 1;
-      const fields = entries.map(([name]) => {
+      const fields = entries.map(([name, value]) => {
         const el = elements[name];
         maxEnd = Math.max(maxEnd, el.offsetEnd);
         // Alignment for scalars/vectors: <=4 -> 4, <=8 -> 8, else 16
         const align = el.size <= 4 ? 4 : el.size <= 8 ? 8 : 16;
         maxAlign = Math.max(maxAlign, align);
+        // Track original JS type for reconstruction during readback
+        const kind = value?.isVector ? 'vector'
+          : value?.isColor ? 'color'
+          : undefined;
         return {
           name,
           baseType: el.baseType,
           size: el.size,
           offset: el.offset,
           packInPlace: el.packInPlace ?? false,
           dim: el.size / 4,
+          kind,
         };
       });
 
@@ -3282,6 +3359,65 @@ ${hookUniformFields}}
       return floatView;
     }
 
+    // Inverse of _packStructArray reads packed buffer back into plain JS objects
+    // using the same schema layout - fields, stride and offsets
+    _unpackStructArray(floatView, schema) {
+      const { fields, stride } = schema;
+      const dataView = new DataView(floatView.buffer);
+      const count = Math.floor(floatView.byteLength / stride);
+      const result = [];
+
+      for (let i = 0; i < count; i++) {
+        const item = {};
+        const baseOffset = i * stride;
+        for (const field of fields) {
+          const byteOffset = baseOffset + field.offset;
+          const n = field.size / 4;
+
+          if (field.baseType === 'u32') {
+            if (n === 1) {
+              item[field.name] = dataView.getUint32(byteOffset, true);
+            } else {
+              item[field.name] = Array.from({ length: n }, (_, j) =>
+                dataView.getUint32(byteOffset + j * 4, true)
+              );
+            }
+          } else if (field.baseType === 'i32') {
+            if (n === 1) {
+              item[field.name] = dataView.getInt32(byteOffset, true);
+            } else {
+              item[field.name] = Array.from({ length: n }, (_, j) =>
+                dataView.getInt32(byteOffset + j * 4, true)
+              );
+            }
+          } else {
+            const idx = byteOffset / 4;
+            if (n === 1) {
+              item[field.name] = floatView[idx];
+            } else {
+              const values = Array.from(floatView.slice(idx, idx + n));
+              if (field.kind === 'vector') {
+                item[field.name] = this._pInst.createVector(...values);
+              } else if (field.kind === 'color') {
+                // Color was packed as normalized RGBA [0-1] via _getRGBA([1,1,1,1])
+                // Scale back to the current colorMode range
+                const maxes = this.states.colorMaxes[this.states.colorMode];
+                item[field.name] = this._pInst.color(
+                  values[0] * maxes[0], values[1] * maxes[1],
+                  values[2] * maxes[2], values[3] * maxes[3]
+                );
+              } else {
+                item[field.name] = values;
+              }
+            }
+          }
+        }
+        result.push(item);
+      }
+
+      return result;
+    }
+
     createStorage(dataOrCount) {
       const device = this.device;
 

test/unit/webgpu/p5.RendererWebGPU.js

@@ -167,4 +167,91 @@ suite('WebGPU p5.RendererWebGPU', function() {
       expect(myp5._renderer).to.exist;
     });
   });
+
+  suite('StorageBuffer.read()', function() {
+    test('reads back float array data', async function() {
+      const input = new Float32Array([1, 2, 3, 4]);
+      const buf = myp5.createStorage(input);
+
+      const result = await buf.read();
+
+      expect(result).to.be.instanceOf(Float32Array);
+      expect(result.length).to.equal(input.length);
+      for (let i = 0; i < input.length; i++) {
+        expect(result[i]).to.be.closeTo(input[i], 0.001);
+      }
+    });
+
+    test('reads back struct array data', async function() {
+      const input = [
+        { x: 1.0, y: 2.0 },
+        { x: 3.0, y: 4.0 },
+      ];
+      const buf = myp5.createStorage(input);
+
+      const result = await buf.read();
+
+      expect(result).to.be.an('array');
+      expect(result.length).to.equal(input.length);
+      for (let i = 0; i < input.length; i++) {
+        expect(result[i].x).to.be.closeTo(input[i].x, 0.001);
+        expect(result[i].y).to.be.closeTo(input[i].y, 0.001);
+      }
+    });
+
+    test('read after update returns new data', async function() {
+      const buf = myp5.createStorage(new Float32Array([10, 20, 30]));
+      const updated = new Float32Array([100, 200, 300]);
+      buf.update(updated);
+
+      const result = await buf.read();
+
+      for (let i = 0; i < updated.length; i++) {
+        expect(result[i]).to.be.closeTo(updated[i], 0.001);
+      }
+    });
+
+    test('reads back struct with vector fields as p5.Vector', async function() {
+      const input = [
+        { position: myp5.createVector(1, 2), speed: 5.0 },
+        { position: myp5.createVector(3, 4), speed: 10.0 },
+      ];
+      const buf = myp5.createStorage(input);
+
+      const result = await buf.read();
+
+      expect(result).to.be.an('array');
+      expect(result.length).to.equal(2);
+      // Vector fields come back as p5.Vector
+      expect(result[0].position.isVector).to.be.true;
+      expect(result[0].position.x).to.be.closeTo(1, 0.001);
+      expect(result[0].position.y).to.be.closeTo(2, 0.001);
+      expect(result[0].speed).to.be.closeTo(5.0, 0.001);
+      expect(result[1].position.isVector).to.be.true;
+      expect(result[1].position.x).to.be.closeTo(3, 0.001);
+      expect(result[1].position.y).to.be.closeTo(4, 0.001);
+      expect(result[1].speed).to.be.closeTo(10.0, 0.001);
+    });
+
+    test('reads back data modified by a compute shader', async function() {
+      const input = new Float32Array([1, 2, 3, 4]);
+      const buf = myp5.createStorage(input);
+
+      const computeShader = myp5.buildComputeShader(() => {
+        const d = myp5.uniformStorage();
+        const idx = myp5.index.x;
+        d[idx] = d[idx] * 2;
+      }, { myp5 });
+
+      computeShader.setUniform('d', buf);
+      myp5.compute(computeShader, 4);
+
+      const result = await buf.read();
+
+      expect(result).to.be.instanceOf(Float32Array);
+      for (let i = 0; i < input.length; i++) {
+        expect(result[i]).to.be.closeTo(input[i] * 2, 0.001);
+      }
+    });
+  });
 });