fix(🐛): install the native module before the flag constants are read

apps/example/src/Tests.tsx

@@ -7,7 +7,16 @@ import {
   PixelRatio,
   findNodeHandle,
 } from "react-native";
-import { GPUOffscreenCanvas } from "react-native-webgpu";
+import {
+  GPUOffscreenCanvas,
+  // The flags as the package exports them, not the globals of the same name
+  // (see the `exportedFlags` spec in Constants.spec.ts).
+  GPUBufferUsage as exportedBufferUsage,
+  GPUColorWrite as exportedColorWrite,
+  GPUMapMode as exportedMapMode,
+  GPUShaderStage as exportedShaderStage,
+  GPUTextureUsage as exportedTextureUsage,
+} from "react-native-webgpu";
 import { mat4, vec3, mat3 } from "wgpu-matrix";
 
 import { useClient } from "./useClient";
@@ -123,6 +132,13 @@ export const Tests = ({ assets: { di3D, saturn, moon } }: AssetProps) => {
                 },
                 workletDeviceStress,
                 workletSerialization,
+                exportedFlags: {
+                  GPUBufferUsage: exportedBufferUsage,
+                  GPUColorWrite: exportedColorWrite,
+                  GPUMapMode: exportedMapMode,
+                  GPUShaderStage: exportedShaderStage,
+                  GPUTextureUsage: exportedTextureUsage,
+                },
                 mat4,
                 vec3,
                 mat3,

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

@@ -61,6 +61,38 @@ describe("WebGPUConstants", () => {
       RENDER_ATTACHMENT: 16,
     });
   });
+  // The package exports the flags as importable values (src/constants.ts),
+  // read from the globals the native module installs. They used to be read
+  // before the install ran: every exported flag was undefined. Native app only:
+  // the web and node runners read the flags of the browser and of the `webgpu`
+  // package.
+  it("exports the flags installed by the native module", async () => {
+    if (client.OS !== "ios" && client.OS !== "android") {
+      return;
+    }
+    const result = await client.eval(({ exportedFlags }) => {
+      if (!exportedFlags) {
+        return null;
+      }
+      const globals = globalThis as unknown as Record<string, unknown>;
+      return {
+        GPUBufferUsage: exportedFlags.GPUBufferUsage === globals.GPUBufferUsage,
+        GPUColorWrite: exportedFlags.GPUColorWrite === globals.GPUColorWrite,
+        GPUMapMode: exportedFlags.GPUMapMode === globals.GPUMapMode,
+        GPUShaderStage: exportedFlags.GPUShaderStage === globals.GPUShaderStage,
+        GPUTextureUsage:
+          exportedFlags.GPUTextureUsage === globals.GPUTextureUsage &&
+          exportedFlags.GPUTextureUsage.RENDER_ATTACHMENT === 16,
+      };
+    });
+    expect(result).toEqual({
+      GPUBufferUsage: true,
+      GPUColorWrite: true,
+      GPUMapMode: true,
+      GPUShaderStage: true,
+      GPUTextureUsage: true,
+    });
+  });
   it("GPUAdapter", async () => {
     const result = await client.eval(({ device, gpu }) => {
       return [

packages/webgpu/src/__tests__/setup.ts

@@ -70,6 +70,15 @@ interface GPUTestingContext {
       value: object,
     ) => Promise<{ threw: boolean; message: string; typeOnWorklet: string }>;
   };
+  // Native app harness only: the flags as the package exports them
+  // (src/constants.ts), next to the globals of the same name.
+  exportedFlags?: {
+    GPUBufferUsage: typeof GPUBufferUsage;
+    GPUColorWrite: typeof GPUColorWrite;
+    GPUMapMode: typeof GPUMapMode;
+    GPUShaderStage: typeof GPUShaderStage;
+    GPUTextureUsage: typeof GPUTextureUsage;
+  };
   // Native app harness only: a mounted probe view for the
   // drawElementImageToTexture spec (its React tag, its size in points and the
   // device pixel ratio). See apps/example/src/Tests.tsx.

packages/webgpu/src/constants.ts

@@ -22,9 +22,12 @@
 //     device.createBuffer({ usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ });
 //   };
 //
-// These are read at module evaluation time. The package entry (`index.tsx`)
-// re-exports `./main` before `./constants`, and `./main` installs the native
-// module synchronously, so the globals always exist by the time this runs.
+// These are read at module evaluation time, so the native module must be
+// installed first: the package entry (`./index.tsx`) starts with a bare
+// `import "./main/installNative"` (so does `./main`), which runs the install
+// before this file is reached, directly or through `./install`. An import is
+// evaluated before the body of the module that declares it, so installing in
+// the body of `./main` left every flag undefined.
 
 export const { GPUBufferUsage } = globalThis;
 

packages/webgpu/src/index.tsx

@@ -1,4 +1,13 @@
 /// <reference types="@webgpu/types" />
+// Must stay the FIRST import, and a bare one: it installs the native module
+// (`./main/installNative.ts`; the web variant loads the RNWebGPU polyfill),
+// and `./constants` reads the globals it installs while being evaluated.
+// Putting it here, rather than relying on `./main` being exported before
+// `./constants` below, makes the order of the `export *` lines irrelevant. A
+// bare import is never deferred by `inlineRequires`, and `import/order --fix`
+// never moves an import across one.
+import "./main/installNative";
+
 import type {
   GPUDrawElementImageDestination,
   GPUDrawElementImageSource,

packages/webgpu/src/install.ts

@@ -19,23 +19,15 @@ const constants = {
   GPUMapMode,
 };
 
-// The GPU instance destined for `navigator.gpu` on other runtimes, wrapped in
-// a holder object. It cannot be read at module evaluation: this module can be
-// evaluated before the native install has populated `RNWebGPU.gpu`, so
-// `main/index.tsx` fills the holder right after installing. The holder itself
-// is captured into the `installWebGPU` worklet closure; when a worklet is
-// serialized (always after startup, so after the holder is filled), the
-// Worklets custom serializer boxes the GPU object inside it, and unboxing on
-// the target runtime installs the GPU prototype there.
-const holder: { gpu?: GPU } = {};
-
-/**
- * @internal Called once by `main/index.tsx` after the native install, so
- * `installWebGPU()` can put `navigator.gpu` on other runtimes.
- */
-export const provideGPUForInstall = (gpu: GPU) => {
-  holder.gpu = gpu;
-};
+// The GPU instance destined for `navigator.gpu` on other runtimes, read from
+// the native install like the flags above (the package entry runs the install
+// before this module is evaluated). Undefined on web, where the RNWebGPU
+// polyfill has no gpu and the browser already provides navigator.gpu. Captured
+// into the `installWebGPU` worklet closure: when a worklet is serialized, the
+// Worklets custom serializer boxes the GPU object, and unboxing on the target
+// runtime installs the GPU prototype there.
+const gpu: GPU | undefined =
+  typeof RNWebGPU === "undefined" ? undefined : RNWebGPU.gpu;
 
 /**
  * Install WebGPU on the runtime that calls it.
@@ -88,7 +80,6 @@ export const installWebGPU = () => {
       g[key] = value;
     }
   }
-  const { gpu } = holder;
   if (gpu !== undefined) {
     const nav = g.navigator as { gpu?: GPU; userAgent?: string } | undefined;
     if (nav === undefined) {

packages/webgpu/src/main/index.tsx

@@ -1,8 +1,13 @@
+// Must stay the FIRST import, and a bare one: it installs the native module
+// (the RNWebGPU global and the WebGPU globals) before anything else here is
+// evaluated; the package entry imports it first as well, see ./installNative.ts.
+// An import with bindings could be deferred to its first use by
+// `inlineRequires`; a bare import is not.
+import "./installNative";
+
 import { findNodeHandle } from "react-native";
 
 import { registerWebGPUForReanimated } from "../external";
-import { provideGPUForInstall } from "../install";
-import WebGPUModule from "../NativeWebGPUModule";
 import type { GPUDrawElementImageSourceView } from "../types";
 
 // Resolve the `source` of drawElementImageToTexture to a React tag on the
@@ -61,8 +66,6 @@ export * from "../formats";
 
 export { default as WebGPUModule } from "../NativeWebGPUModule";
 
-const _installOk = WebGPUModule.install();
-
 registerWebGPUForReanimated();
 
 if (typeof RNWebGPU !== "undefined") {
@@ -87,10 +90,4 @@ if (typeof RNWebGPU !== "undefined") {
   global.createImageBitmap =
     global.createImageBitmap ?? RNWebGPU.createImageBitmap.bind(RNWebGPU);
   installDrawElementImageToTexture();
-  // Let installWebGPU() put navigator.gpu on other runtimes (see install.ts).
-  provideGPUForInstall(RNWebGPU.gpu);
-} else {
-  console.warn(
-    `[react-native-webgpu] install() returned ${_installOk} but RNWebGPU global is not available`,
-  );
 }

packages/webgpu/src/main/index.web.tsx

@@ -1,4 +1,4 @@
-import "../WebPolyfillGPUModule";
+import "./installNative";
 
 export * from "../Canvas";
 export * from "../Offscreen";

packages/webgpu/src/main/installNative.ts

@@ -0,0 +1,20 @@
+import WebGPUModule from "../NativeWebGPUModule";
+
+// Installs the native module: the RNWebGPU global, and the WebGPU globals of
+// the main runtime (GPUBufferUsage, GPUTextureUsage, GPUShaderStage...).
+//
+// A module of its own, imported FIRST (as a bare import) by the package entry
+// `../index.tsx` and by `./index.tsx`, because `../constants` reads those
+// globals while it is evaluated, and every import of a module is evaluated
+// before that module's body. Installing in the body of `./index.tsx` came too
+// late: the flags exported by the package were all undefined.
+const installOk = WebGPUModule.install();
+
+if (typeof RNWebGPU === "undefined") {
+  // install() logs the reason natively (e.g. a null jsi::Runtime on iOS).
+  // Nothing works from here on (navigator.gpu, the flag constants, Canvas),
+  // so fail at import time rather than leave every export undefined.
+  throw new Error(
+    `[react-native-webgpu] the native module failed to install (install() returned ${installOk}): the RNWebGPU global is not available. Make sure the native dependencies are installed and rebuild the app.`,
+  );
+}

packages/webgpu/src/main/installNative.web.ts

@@ -0,0 +1,5 @@
+// Web counterpart of ./installNative.ts, so the package entry can import
+// `./main/installNative` on every platform without reaching the TurboModule
+// here. The browser provides the WebGPU globals (navigator.gpu, the flag
+// constants); the polyfill provides the RNWebGPU global that Canvas relies on.
+import "../WebPolyfillGPUModule";

packages/webgpu/src/mock.ts

@@ -16,9 +16,68 @@ import type { CanvasProps, NativeCanvas, RNCanvasContext } from "./Canvas";
 // missing native install (see `importDevice`/`adoptTexture` in
 // `./importDevice.ts`).
 
-export * from "./constants";
 export * from "./formats";
 
+// The flag constants. The real package re-exports the globals the native
+// module installs (`./constants.ts`); nothing installs them under Jest, so the
+// mock defines the spec values itself and, like the native install, also puts
+// them on globalThis for code that references the bare globals.
+export const GPUBufferUsage: GPUBufferUsage = {
+  MAP_READ: 0x0001,
+  MAP_WRITE: 0x0002,
+  COPY_SRC: 0x0004,
+  COPY_DST: 0x0008,
+  INDEX: 0x0010,
+  VERTEX: 0x0020,
+  UNIFORM: 0x0040,
+  STORAGE: 0x0080,
+  INDIRECT: 0x0100,
+  QUERY_RESOLVE: 0x0200,
+};
+
+export const GPUTextureUsage: GPUTextureUsage = {
+  COPY_SRC: 0x01,
+  COPY_DST: 0x02,
+  TEXTURE_BINDING: 0x04,
+  STORAGE_BINDING: 0x08,
+  RENDER_ATTACHMENT: 0x10,
+  TRANSIENT_ATTACHMENT: 0x20,
+};
+
+export const GPUShaderStage: GPUShaderStage = {
+  VERTEX: 0x1,
+  FRAGMENT: 0x2,
+  COMPUTE: 0x4,
+};
+
+export const GPUColorWrite: GPUColorWrite = {
+  RED: 0x1,
+  GREEN: 0x2,
+  BLUE: 0x4,
+  ALPHA: 0x8,
+  ALL: 0xf,
+};
+
+export const GPUMapMode: GPUMapMode = {
+  READ: 0x1,
+  WRITE: 0x2,
+};
+
+{
+  const g = globalThis as unknown as Record<string, unknown>;
+  for (const [key, value] of Object.entries({
+    GPUBufferUsage,
+    GPUTextureUsage,
+    GPUShaderStage,
+    GPUColorWrite,
+    GPUMapMode,
+  })) {
+    if (g[key] === undefined) {
+      g[key] = value;
+    }
+  }
+}
+
 const NOT_AVAILABLE_UNDER_JEST =
   "is not available under Jest - react-native-webgpu has no native implementation in tests (see react-native-webgpu/jestSetup.js)";