Adapt reflection and screenshot components to WebGPURenderer

docs/components/screenshot.md

@@ -43,6 +43,10 @@ To take a screenshot programmatically and get a canvas, call `getCanvas()`:
 document.querySelector('a-scene').components.screenshot.getCanvas('equirectangular');
 ```
 
+> **NOTE:** When using `THREE.WebGPURenderer`, pixels are read back from the GPU
+> asynchronously, so `getCanvas()` returns a `Promise` resolving to the canvas
+> instead of the canvas itself.
+
 To take a screenshot programmatically and automatically save the file, call `capture()`:
 
 ```js

src/components/scene/reflection.js

@@ -2,6 +2,11 @@
 import * as THREE from 'three';
 import { registerComponent as register } from '../../core/component.js';
 
+// WebGLRenderer uses WebGLCubeRenderTarget, WebGPURenderer uses CubeRenderTarget.
+// Written so that Webpack can't statically determine the export used;
+// only one of the two exists depending on the three.js build.
+var cubeRenderTargetImpl = ['WebGLCubeRenderTarget', 'CubeRenderTarget'].find(function (x) { return THREE[x]; });
+
 // source: view-source:https://storage.googleapis.com/chromium-webxr-test/r886480/proposals/lighting-estimation.html
 function updateLights (estimate, probeLight, directionalLight, directionalLightPosition) {
   var intensityScalar =
@@ -30,9 +35,9 @@ export var Component = register('reflection', {
   sceneOnly: true,
   init: function () {
     var self = this;
-    this.cubeRenderTarget = new THREE.WebGLCubeRenderTarget(16);
+    this.cubeRenderTarget = new THREE[cubeRenderTargetImpl](16);
     this.cubeCamera = new THREE.CubeCamera(0.1, 1000, this.cubeRenderTarget);
-    this.lightingEstimationTexture = (new THREE.WebGLCubeRenderTarget(16)).texture;
+    this.lightingEstimationTexture = (new THREE[cubeRenderTargetImpl](16)).texture;
     this.needsVREnvironmentUpdate = true;
 
     // Update WebXR to support light-estimation

src/components/scene/screenshot.js

@@ -2,6 +2,15 @@
 import { registerComponent } from '../../core/component.js';
 import * as THREE from 'three';
 
+// WebGLRenderer uses WebGLCubeRenderTarget, WebGPURenderer uses CubeRenderTarget.
+// Written so that Webpack can't statically determine the export used;
+// only one of the two exists depending on the three.js build.
+var cubeRenderTargetImpl = ['WebGLCubeRenderTarget', 'CubeRenderTarget'].find(function (x) { return THREE[x]; });
+// TSL and MeshBasicNodeMaterial only exist in the three.js build with
+// WebGPURenderer, same Webpack trick as above.
+var TSL = ['TSL'].map(function (x) { return THREE[x]; })[0];
+var MeshBasicNodeMaterial = ['MeshBasicNodeMaterial'].map(function (x) { return THREE[x]; })[0];
+
 var VERTEX_SHADER = [
   'attribute vec3 position;',
   'attribute vec2 uv;',
@@ -55,15 +64,27 @@ export var Component = registerComponent('screenshot', {
   setup: function () {
     var el = this.el;
     if (this.canvas) { return; }
-    var gl = el.renderer.getContext();
-    if (!gl) { return; }
-    this.cubeMapSize = gl.getParameter(gl.MAX_CUBE_MAP_TEXTURE_SIZE);
-    this.material = new THREE.RawShaderMaterial({
-      uniforms: {map: {type: 't', value: null}},
-      vertexShader: VERTEX_SHADER,
-      fragmentShader: FRAGMENT_SHADER,
-      side: THREE.DoubleSide
-    });
+    var ctx = el.renderer.getContext();
+    if (!ctx) { return; }
+    if (ctx.getParameter) {
+      this.cubeMapSize = ctx.getParameter(ctx.MAX_CUBE_MAP_TEXTURE_SIZE);
+    } else {
+      // ctx is a GPUCanvasContext; cube map faces are limited by maxTextureDimension2D.
+      var device = el.renderer.backend.device;
+      this.cubeMapSize = device ? device.limits.maxTextureDimension2D : 2048;
+    }
+    // WebGPURenderer does not support RawShaderMaterial, use a node material with TSL.
+    if (TSL) {
+      this.material = this.createNodeMaterial();
+    } else {
+      this.material = new THREE.RawShaderMaterial({
+        uniforms: {map: {type: 't', value: null}},
+        vertexShader: VERTEX_SHADER,
+        fragmentShader: FRAGMENT_SHADER,
+        side: THREE.DoubleSide
+      });
+      this.cubeTextureUniform = this.material.uniforms.map;
+    }
     this.quad = new THREE.Mesh(
       new THREE.PlaneGeometry(1, 1),
       this.material
@@ -76,6 +97,25 @@ export var Component = registerComponent('screenshot', {
     this.onKeyDown = this.onKeyDown.bind(this);
   },
 
+  /**
+   * TSL (node material) equivalent of the raw GLSL shaders above, used with
+   * WebGPURenderer which does not support RawShaderMaterial.
+   */
+  createNodeMaterial: function () {
+    var material = new MeshBasicNodeMaterial({side: THREE.DoubleSide});
+    var uv = TSL.uv();
+    var longitude = TSL.float(1).sub(uv.x).mul(2 * Math.PI).sub(Math.PI).add(Math.PI / 2);
+    var latitude = uv.y.mul(Math.PI);
+    var dir = TSL.vec3(
+      TSL.sin(longitude).mul(TSL.sin(latitude)).negate(),
+      TSL.cos(latitude),
+      TSL.cos(longitude).mul(TSL.sin(latitude)).negate()
+    );
+    this.cubeTextureUniform = TSL.cubeTexture(new THREE.CubeTexture(), dir);
+    material.colorNode = TSL.vec4(this.cubeTextureUniform.rgb, 1);
+    return material;
+  },
+
   getRenderTarget: function (width, height) {
     return new THREE.WebGLRenderTarget(width, height, {
       colorSpace: this.el.sceneEl.renderer.outputColorSpace,
@@ -140,10 +180,10 @@ export var Component = registerComponent('screenshot', {
     } else {
       // Use ortho camera.
       camera = this.camera;
-      cubeRenderTarget = new THREE.WebGLCubeRenderTarget(
+      cubeRenderTarget = new THREE[cubeRenderTargetImpl](
         Math.min(this.cubeMapSize, 2048),
         {
-          format: THREE.RGBFormat,
+          format: THREE.RGBAFormat,
           generateMipmaps: true,
           minFilter: THREE.LinearMipmapLinearFilter,
           colorSpace: THREE.SRGBColorSpace
@@ -155,7 +195,7 @@ export var Component = registerComponent('screenshot', {
       el.camera.getWorldQuaternion(cubeCamera.quaternion);
       // Render scene with cube camera.
       cubeCamera.update(el.renderer, el.object3D);
-      this.quad.material.uniforms.map.value = cubeCamera.renderTarget.texture;
+      this.cubeTextureUniform.value = cubeCamera.renderTarget.texture;
       size = {width: this.data.width, height: this.data.height};
       // Use quad to project image taken by the cube camera.
       this.quad.visible = true;
@@ -174,19 +214,25 @@ export var Component = registerComponent('screenshot', {
     var isVREnabled = this.el.renderer.xr.enabled;
     var renderer = this.el.renderer;
     var params;
+    var self = this;
     this.setup();
     // Disable VR.
     renderer.xr.enabled = false;
     params = this.setCapture(projection);
-    this.renderCapture(params.camera, params.size, params.projection);
-    // Trigger file download.
-    this.saveCapture();
+    this.renderCapture(params.camera, params.size, params.projection)
+      .then(function () {
+        // Trigger file download.
+        self.saveCapture();
+      });
     // Restore VR.
     renderer.xr.enabled = isVREnabled;
   },
 
   /**
    * Return canvas instead of triggering download (e.g., for uploading blob to server).
+   * With WebGLRenderer the canvas is returned synchronously.
+   * With WebGPURenderer pixels are read back asynchronously, so a Promise
+   * resolving to the canvas is returned instead.
    */
   getCanvas: function (projection) {
     var isVREnabled = this.el.renderer.xr.enabled;
@@ -195,22 +241,22 @@ export var Component = registerComponent('screenshot', {
     // Disable VR.
     var params = this.setCapture(projection);
     renderer.xr.enabled = false;
-    this.renderCapture(params.camera, params.size, params.projection);
+    var promise = this.renderCapture(params.camera, params.size, params.projection);
     // Restore VR.
     renderer.xr.enabled = isVREnabled;
-    return this.canvas;
+    if (renderer.readRenderTargetPixels) { return this.canvas; }
+    return promise;
   },
 
   renderCapture: function (camera, size, projection) {
     var autoClear = this.el.renderer.autoClear;
     var el = this.el;
-    var imageData;
     var output;
     var pixels;
     var renderer = el.renderer;
+    var self = this;
     // Create rendering target and buffer to store the read pixels.
     output = this.getRenderTarget(size.width, size.height);
-    pixels = new Uint8Array(4 * size.width * size.height);
     // Resize quad, camera, and canvas.
     this.resize(size.width, size.height);
     // Render scene to render target.
@@ -220,9 +266,30 @@ export var Component = registerComponent('screenshot', {
     renderer.render(el.object3D, camera);
     renderer.autoClear = autoClear;
     // Read image pixels back.
-    renderer.readRenderTargetPixels(output, 0, 0, size.width, size.height, pixels);
-    renderer.setRenderTarget(null);
-    if (projection === 'perspective') {
+    if (renderer.readRenderTargetPixels) {
+      pixels = new Uint8Array(4 * size.width * size.height);
+      renderer.readRenderTargetPixels(output, 0, 0, size.width, size.height, pixels);
+      renderer.setRenderTarget(null);
+      this.copyCapture(pixels, size, projection);
+      return Promise.resolve(this.canvas);
+    }
+    // WebGPURenderer only supports reading pixels back asynchronously.
+    return renderer.readRenderTargetPixelsAsync(output, 0, 0, size.width, size.height)
+      .then(function (pixels) {
+        renderer.setRenderTarget(null);
+        self.copyCapture(pixels, size, projection);
+        return self.canvas;
+      });
+  },
+
+  copyCapture: function (pixels, size, projection) {
+    var imageData;
+    // Pixels read back from WebGL are bottom-up, they are top-down with the
+    // WebGPU backend of WebGPURenderer. The equirectangular projection quad
+    // renders vertically inverted, so it needs the opposite flip of the
+    // perspective projection.
+    var topDown = this.el.renderer.coordinateSystem === THREE.WebGPUCoordinateSystem;
+    if ((projection === 'perspective') !== topDown) {
       pixels = this.flipPixelsVertically(pixels, size.width, size.height);
     }
     imageData = new ImageData(new Uint8ClampedArray(pixels), size.width, size.height);