fix(🌎): type MakeWebGPUCanvasContext as nullable and keep getNativeSurface() working after unmount

packages/webgpu/src/WebPolyfillGPUModule.ts

@@ -1,11 +1,33 @@
+import type { NativeCanvas } from "./types";
 import { contextIdToId } from "./utils";
 
 const fabric = true;
 
-function getNativeSurface(contextId: number) {
-  const canvas = document.getElementById(
+function getCanvasElement(contextId: number) {
+  return document.getElementById(
     contextIdToId(contextId),
-  ) as HTMLCanvasElement;
+  ) as HTMLCanvasElement | null;
+}
+
+// Native's answer for a context it no longer knows (RNWebGPU::getNativeSurface).
+const makeRetiredSurface = (): NativeCanvas => ({
+  surface: BigInt(0),
+  height: 0,
+  width: 0,
+  clientHeight: 0,
+  clientWidth: 0,
+  addEventListener: () => {},
+  removeEventListener: () => {},
+  dispatchEvent: () => {},
+  setPointerCapture: () => {},
+  releasePointerCapture: () => {},
+});
+
+function getNativeSurface(contextId: number): NativeCanvas {
+  const canvas = getCanvasElement(contextId);
+  if (canvas === null) {
+    return makeRetiredSurface();
+  }
 
   // Layout size, like the native view size in dp: unlike
   // getBoundingClientRect() it ignores CSS transforms.
@@ -31,9 +53,10 @@ function makeWebGPUCanvasContext(
   width: number,
   height: number,
 ) {
-  const canvas = document.getElementById(
-    contextIdToId(contextId),
-  ) as HTMLCanvasElement;
+  const canvas = getCanvasElement(contextId);
+  if (canvas === null) {
+    return null;
+  }
 
   const dpr = window.devicePixelRatio || 1;
   const pixelWidth = (width * dpr).toString();
@@ -58,14 +81,16 @@ function makeWebGPUCanvasContext(
   return Object.assign(context, { present: () => {} });
 }
 
-// @ts-expect-error - polyfill for RNWebGPU native module
-window.RNWebGPU = {
+const polyfill: Pick<
+  typeof RNWebGPU,
+  "getNativeSurface" | "MakeWebGPUCanvasContext" | "destroyContext" | "fabric"
+> = {
   getNativeSurface,
-  // Unlike native, this can return null; Canvas.getContext() is typed
-  // `RNCanvasContext | null` and passes it through.
-  MakeWebGPUCanvasContext:
-    makeWebGPUCanvasContext as typeof RNWebGPU.MakeWebGPUCanvasContext,
+  MakeWebGPUCanvasContext: makeWebGPUCanvasContext,
   // On web the browser owns the canvas/context lifecycle; nothing to retire.
   destroyContext: (_contextId: number) => {},
   fabric,
 };
+
+// The web provides these members of the native module and no others.
+window.RNWebGPU = polyfill as typeof RNWebGPU;

packages/webgpu/src/__tests__/CanvasContext.spec.ts

@@ -8,6 +8,9 @@ describe("Native canvas context", () => {
     const result = await client.eval(({ device, gpu }) => {
       // The harness uses a JS offscreen polyfill; exercise the native context.
       const ctx = RNWebGPU.MakeWebGPUCanvasContext(-1, 16, 16);
+      if (ctx === null) {
+        throw new Error("native returned no canvas context");
+      }
       ctx.configure({ device, format: gpu.getPreferredCanvasFormat() });
       try {
         const texture = ctx.getCurrentTexture();
@@ -37,9 +40,15 @@ describe("Native canvas context", () => {
     const result = await client.eval(({ device, gpu }) => {
       const config = { device, format: gpu.getPreferredCanvasFormat() };
       const kept = RNWebGPU.MakeWebGPUCanvasContext(-2, 16, 16);
+      if (kept === null) {
+        throw new Error("native returned no canvas context");
+      }
       kept.configure(config);
       RNWebGPU.destroyContext(-2);
       const next = RNWebGPU.MakeWebGPUCanvasContext(-2, 16, 16);
+      if (next === null) {
+        throw new Error("native returned no canvas context");
+      }
       next.configure(config);
       try {
         const texture = next.getCurrentTexture();
@@ -63,6 +72,9 @@ describe("Native canvas context", () => {
       const result = await client.eval(
         ({ device, gpu, operation: transition }) => {
           const ctx = RNWebGPU.MakeWebGPUCanvasContext(-1, 16, 16);
+          if (ctx === null) {
+            throw new Error("native returned no canvas context");
+          }
           const config = { device, format: gpu.getPreferredCanvasFormat() };
           ctx.configure(config);
           try {

packages/webgpu/src/__tests__/WebPolyfill.spec.ts

@@ -23,15 +23,12 @@ const makeCanvas = (context: object | null): FakeCanvas => ({
   },
 });
 
-type Polyfill = {
-  MakeWebGPUCanvasContext: (
-    contextId: number,
-    width: number,
-    height: number,
-  ) => (object & { present: () => void }) | null;
-};
+type Polyfill = Pick<
+  typeof RNWebGPU,
+  "getNativeSurface" | "MakeWebGPUCanvasContext"
+>;
 
-const loadPolyfill = (canvas: FakeCanvas): Polyfill => {
+const loadPolyfill = (canvas: FakeCanvas | null): Polyfill => {
   const g = globalThis as Record<string, unknown>;
   g.window = { devicePixelRatio: 2 };
   g.document = { getElementById: () => canvas };
@@ -64,4 +61,32 @@ describe("Web polyfill", () => {
     const RNWebGPU = loadPolyfill(makeCanvas(null));
     expect(RNWebGPU.MakeWebGPUCanvasContext(1, 100, 50)).toBeNull();
   });
+
+  it("reports a 0 x 0 surface once the canvas element is gone, as native does", () => {
+    const surface = loadPolyfill(null).getNativeSurface(1);
+    expect({
+      surface: surface.surface,
+      width: surface.width,
+      height: surface.height,
+      clientWidth: surface.clientWidth,
+      clientHeight: surface.clientHeight,
+    }).toEqual({
+      surface: BigInt(0),
+      width: 0,
+      height: 0,
+      clientWidth: 0,
+      clientHeight: 0,
+    });
+    expect(() => {
+      surface.addEventListener("pointerdown", () => {});
+      surface.removeEventListener("pointerdown", () => {});
+      surface.dispatchEvent(new Event("pointerdown"));
+      surface.setPointerCapture(7);
+      surface.releasePointerCapture(7);
+    }).not.toThrow();
+  });
+
+  it("returns no context once the canvas element is gone", () => {
+    expect(loadPolyfill(null).MakeWebGPUCanvasContext(1, 100, 50)).toBeNull();
+  });
 });

packages/webgpu/src/index.tsx

@@ -71,11 +71,12 @@ declare global {
     gpu: GPU;
     fabric: boolean;
     getNativeSurface: (contextId: number) => NativeCanvas;
+    // null on the web, without WebGPU support or once the canvas is gone.
     MakeWebGPUCanvasContext: (
       contextId: number,
       width: number,
       height: number,
-    ) => RNCanvasContext;
+    ) => RNCanvasContext | null;
     // Retires a canvas context; called by Canvas on unmount (the Canvas owns
     // the native registry entry for its contextId).
     destroyContext: (contextId: number) => void;