react-native-webgpu 0.10.2 → 0.10.4

package/android/cpp/cpp-adapter.cpp

@@ -6,6 +6,7 @@
 #include <jsi/jsi.h>
 
 #include <ReactCommon/CallInvokerHolder.h>
+#include <android/hardware_buffer_jni.h>
 #include <android/native_window_jni.h>
 #include <webgpu/webgpu_cpp.h>
 
@@ -25,7 +26,8 @@ extern "C" JNIEXPORT void JNICALL Java_com_webgpu_WebGPUModule_initializeNative(
   auto jsCallInvoker{
       facebook::jni::alias_ref<facebook::react::CallInvokerHolder::javaobject>{
           reinterpret_cast<facebook::react::CallInvokerHolder::javaobject>(
-              jsCallInvokerHolder)} -> cthis()->getCallInvoker()};
+              jsCallInvokerHolder)} -> cthis()
+          ->getCallInvoker()};
   auto platformContext =
       std::make_shared<rnwgpu::AndroidPlatformContext>(globalBlobModule);
   manager = std::make_shared<rnwgpu::RNWebGPUManager>(runtime, jsCallInvoker,
@@ -91,3 +93,142 @@ extern "C" JNIEXPORT void JNICALL Java_com_webgpu_WebGPUView_onViewDestroyed(
   }
   registry.removeSurfaceInfo(contextId);
 }
+
+// --- WebGPUHardwareBufferView (AHB-pool presentation)
+// ---------------------------------
+
+namespace {
+
+// Weak handle on a WebGPUHardwareBufferView used by the native fence waiter to
+// wake the UI thread when a frame is ready. Holding it weakly means a view that
+// RN has dropped is not kept alive by a SurfaceInfo that outlives it (the
+// registry keeps the SurfaceInfo until onDropViewInstance).
+struct HardwareBufferViewWaker {
+  jweak view = nullptr;
+  jmethodID onFrameReady = nullptr;
+
+  HardwareBufferViewWaker(JNIEnv *env, jobject thiz) {
+    view = env->NewWeakGlobalRef(thiz);
+    jclass cls = env->GetObjectClass(thiz);
+    onFrameReady = env->GetMethodID(cls, "onNativeFrameReady", "()V");
+    env->DeleteLocalRef(cls);
+  }
+
+  ~HardwareBufferViewWaker() {
+    // May run on the waiter or JS thread; ThreadScope attaches if needed.
+    facebook::jni::ThreadScope scope;
+    JNIEnv *env = facebook::jni::Environment::current();
+    if (env != nullptr && view != nullptr) {
+      env->DeleteWeakGlobalRef(view);
+    }
+  }
+
+  // Called on the waiter thread, outside the pool lock.
+  void operator()() const {
+    facebook::jni::ThreadScope scope;
+    JNIEnv *env = facebook::jni::Environment::current();
+    if (env == nullptr || onFrameReady == nullptr) {
+      return;
+    }
+    jobject local = env->NewLocalRef(view);
+    if (local == nullptr) {
+      return; // the view was garbage collected
+    }
+    env->CallVoidMethod(local, onFrameReady);
+    if (env->ExceptionCheck()) {
+      env->ExceptionClear();
+    }
+    env->DeleteLocalRef(local);
+  }
+};
+
+} // namespace
+
+// Turn on pool mode for this context. dpW/dpH is the canvas-client (dp) size;
+// the native pool buffers are sized from the canvas drawing buffer lazily.
+extern "C" JNIEXPORT void JNICALL
+Java_com_webgpu_WebGPUHardwareBufferView_nEnablePool(JNIEnv *env, jobject thiz,
+                                                     jint contextId, jint dpW,
+                                                     jint dpH) {
+  auto &registry = rnwgpu::SurfaceRegistry::getInstance();
+  auto info = registry.getSurfaceInfoOrCreate(
+      contextId, manager->_gpu, static_cast<int>(dpW), static_cast<int>(dpH));
+  auto waker = std::make_shared<HardwareBufferViewWaker>(env, thiz);
+  info->setPoolFrameReadyCallback([waker]() { (*waker)(); });
+  info->enablePool(static_cast<int>(dpW), static_cast<int>(dpH));
+}
+
+// Keep the canvas-client (dp) size in sync on resize.
+extern "C" JNIEXPORT void JNICALL
+Java_com_webgpu_WebGPUHardwareBufferView_nSetClientSize(JNIEnv *env,
+                                                        jobject thiz,
+                                                        jint contextId,
+                                                        jint dpW, jint dpH) {
+  auto &registry = rnwgpu::SurfaceRegistry::getInstance();
+  auto info = registry.getSurfaceInfo(contextId);
+  if (info != nullptr) {
+    info->setPoolClientSize(static_cast<int>(dpW), static_cast<int>(dpH));
+  }
+}
+
+// The HardwareBuffer backing a (generation, slot), for the view to wrap in a
+// Bitmap. Returns null for a retired generation.
+extern "C" JNIEXPORT jobject JNICALL
+Java_com_webgpu_WebGPUHardwareBufferView_nGetHardwareBuffer(
+    JNIEnv *env, jobject thiz, jint contextId, jint generation, jint slot) {
+  auto &registry = rnwgpu::SurfaceRegistry::getInstance();
+  auto info = registry.getSurfaceInfo(contextId);
+  if (info == nullptr) {
+    return nullptr;
+  }
+  void *ahb = info->poolBufferForDisplay(static_cast<uint32_t>(generation),
+                                         static_cast<int>(slot));
+  if (ahb == nullptr) {
+    return nullptr;
+  }
+  // toHardwareBuffer acquires its own ref for the returned jobject; the pool
+  // keeps the underlying buffer alive independently.
+  return AHardwareBuffer_toHardwareBuffer(env,
+                                          static_cast<AHardwareBuffer *>(ahb));
+}
+
+// Latest signaled frame ready to display, encoded (generation << 32 | slot), or
+// -1 when nothing is new. Called on the UI thread after onNativeFrameReady.
+extern "C" JNIEXPORT jlong JNICALL
+Java_com_webgpu_WebGPUHardwareBufferView_nPollReady(JNIEnv *env, jobject thiz,
+                                                    jint contextId) {
+  auto &registry = rnwgpu::SurfaceRegistry::getInstance();
+  auto info = registry.getSurfaceInfo(contextId);
+  if (info == nullptr) {
+    return -1;
+  }
+  return info->poolPollReady();
+}
+
+// The view is done displaying (and holding) a slot; return it to the pool.
+extern "C" JNIEXPORT void JNICALL
+Java_com_webgpu_WebGPUHardwareBufferView_nReleaseSlot(JNIEnv *env, jobject thiz,
+                                                      jint contextId,
+                                                      jint generation,
+                                                      jint slot) {
+  auto &registry = rnwgpu::SurfaceRegistry::getInstance();
+  auto info = registry.getSurfaceInfo(contextId);
+  if (info != nullptr) {
+    info->poolReleaseSlot(static_cast<uint32_t>(generation),
+                          static_cast<int>(slot));
+  }
+}
+
+// View detached / hidden: keep the canvas alive by falling back to an offscreen
+// texture (mirrors switchToOffscreenSurface for the surface path).
+extern "C" JNIEXPORT void JNICALL
+Java_com_webgpu_WebGPUHardwareBufferView_nSwitchToOffscreen(JNIEnv *env,
+                                                            jobject thiz,
+                                                            jint contextId) {
+  auto &registry = rnwgpu::SurfaceRegistry::getInstance();
+  auto info = registry.getSurfaceInfo(contextId);
+  if (info != nullptr) {
+    info->setPoolFrameReadyCallback(nullptr);
+    info->switchToOffscreen();
+  }
+}

package/android/src/main/java/com/webgpu/WebGPUAPI.java

@@ -22,4 +22,6 @@ public interface WebGPUAPI {
   );
 
   void surfaceOffscreen();
+
+  int getContextId();
 }

package/android/src/main/java/com/webgpu/WebGPUHardwareBufferView.java

@@ -0,0 +1,278 @@
+package com.webgpu;
+
+import android.annotation.SuppressLint;
+import android.content.Context;
+import android.graphics.Bitmap;
+import android.graphics.Canvas;
+import android.graphics.Paint;
+import android.graphics.Rect;
+import android.hardware.HardwareBuffer;
+import android.os.Build;
+import android.os.Handler;
+import android.os.Looper;
+import android.view.View;
+
+import androidx.annotation.Keep;
+import androidx.annotation.NonNull;
+import androidx.annotation.RequiresApi;
+
+import java.util.ArrayDeque;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * A "normal RN view" backend for the WebGPU canvas.
+ *
+ * <p>WebGPU renders into a native pool of AHardwareBuffers (sized from the canvas drawing buffer,
+ * like the swapchain) that Dawn imports as SharedTextureMemory; see SurfaceRegistry.h. Each finished
+ * buffer is drawn inline here via {@link Bitmap#wrapHardwareBuffer} + {@link Canvas#drawBitmap}, so
+ * this is a plain {@link View}: parent transforms, clipping, alpha, z-order and animations all
+ * apply, with no GL interop and no extra copy.
+ *
+ * <p>This view is a pure consumer: it enables pool mode, and whenever the native fence waiter
+ * signals that a frame is ready it asks native for the latest ready (gen, slot), wraps that buffer
+ * in a cached Bitmap, and draws it scaled to the current bounds. Acquire is rigorous (native blocks
+ * on Dawn's render-complete fence before a frame is ready); release is heuristic via a held-ring of
+ * {@link #HELD_MAX} frames.
+ */
+@RequiresApi(api = Build.VERSION_CODES.Q)
+@SuppressLint("ViewConstructor")
+public class WebGPUHardwareBufferView extends View {
+
+  private static final int HELD_MAX = 2; // displayed frames kept before release (release safety)
+
+  private final WebGPUAPI mApi;
+  private final Paint mPaint;
+  private final Rect mSrc = new Rect();
+  private final Rect mDst = new Rect();
+
+  // token (generation << 32 | slot) -> wrapped hardware Bitmap (one per slot per generation).
+  private final Map<Long, Bitmap> mBitmaps = new HashMap<>();
+  // Held-ring of displayed frame tokens (newest last). The newest is what onDraw shows.
+  private final ArrayDeque<Long> mHeld = new ArrayDeque<>();
+
+  // Native -> UI thread wake-up. The fence waiter calls onNativeFrameReady() off the UI thread;
+  // we coalesce those into at most one pending consume() on the main looper.
+  private final Handler mMainHandler = new Handler(Looper.getMainLooper());
+  private final AtomicBoolean mConsumePosted = new AtomicBoolean(false);
+  private final Runnable mConsume = this::consume;
+
+  private Bitmap mDisplayed;
+  private int mDisplayedW;
+  private int mDisplayedH;
+  private int mCurrentGen;
+
+  private boolean mAttached;
+  private boolean mEnabled;
+
+  public WebGPUHardwareBufferView(Context context, WebGPUAPI api) {
+    super(context);
+    mApi = api;
+    mPaint = new Paint(Paint.FILTER_BITMAP_FLAG);
+    setWillNotDraw(false);
+  }
+
+  private int contextId() {
+    return mApi.getContextId();
+  }
+
+  private static int genOf(long token) {
+    return (int) (token >>> 32);
+  }
+
+  private static int slotOf(long token) {
+    return (int) (token & 0xffffffffL);
+  }
+
+  // --- Sizing ----------------------------------------------------------------
+
+  private void enablePool(int w, int h) {
+    final float density = getResources().getDisplayMetrics().density;
+    nEnablePool(contextId(), Math.max(1, Math.round(w / density)),
+        Math.max(1, Math.round(h / density)));
+    mEnabled = true;
+  }
+
+  @Override
+  protected void onSizeChanged(int w, int h, int oldw, int oldh) {
+    super.onSizeChanged(w, h, oldw, oldh);
+    if (w <= 0 || h <= 0) {
+      return;
+    }
+    if (!mEnabled) {
+      enablePool(w, h);
+    } else {
+      // Native reallocates the pool from the canvas drawing buffer; we only keep the dp
+      // canvas-client size in sync so JS computes the right canvas.width/height.
+      final float density = getResources().getDisplayMetrics().density;
+      nSetClientSize(contextId(), Math.max(1, Math.round(w / density)),
+          Math.max(1, Math.round(h / density)));
+    }
+  }
+
+  // --- Consume (UI thread, woken by the native fence waiter) -----------------
+
+  /** Called from the native fence-waiter thread (see cpp-adapter.cpp) when a frame is ready. */
+  @Keep
+  private void onNativeFrameReady() {
+    if (mConsumePosted.compareAndSet(false, true)) {
+      mMainHandler.post(mConsume);
+    }
+  }
+
+  private void consume() {
+    mConsumePosted.set(false);
+    if (!mAttached) {
+      return;
+    }
+    long token = nPollReady(contextId());
+    if (token >= 0) {
+      onFrameReady(token);
+    }
+  }
+
+  private void onFrameReady(long token) {
+    final int gen = genOf(token);
+    final int slot = slotOf(token);
+
+    if (gen > mCurrentGen) {
+      mCurrentGen = gen;
+      recycleStaleGenerations();
+    }
+
+    Bitmap bmp = mBitmaps.get(token);
+    if (bmp == null) {
+      HardwareBuffer hb = nGetHardwareBuffer(contextId(), gen, slot);
+      if (hb == null) {
+        return; // generation already retired; drop this stale frame
+      }
+      // wrapHardwareBuffer takes its own ref, so the wrapper can be closed immediately.
+      bmp = Bitmap.wrapHardwareBuffer(hb, null);
+      hb.close();
+      if (bmp == null) {
+        // Cannot display this frame; hand the slot straight back so the pool
+        // does not drain (native marked it Displayed when we polled it).
+        nReleaseSlot(contextId(), gen, slot);
+        return;
+      }
+      mBitmaps.put(token, bmp);
+    }
+
+    mDisplayed = bmp;
+    mDisplayedW = bmp.getWidth();
+    mDisplayedH = bmp.getHeight();
+
+    mHeld.addLast(token);
+    while (mHeld.size() > HELD_MAX) {
+      long old = mHeld.pollFirst();
+      if (genOf(old) == mCurrentGen) {
+        // Live generation reuses slots: hand it back once HWUI has had time to sample it.
+        nReleaseSlot(contextId(), mCurrentGen, slotOf(old));
+      } else {
+        // Old generation: its slot is never re-rendered, so just drop the Bitmap.
+        recycleToken(old);
+      }
+    }
+    invalidate();
+  }
+
+  private void recycleStaleGenerations() {
+    // Keep the current and previous generation (cross-fade); recycle older cached Bitmaps that are
+    // no longer held.
+    Iterator<Map.Entry<Long, Bitmap>> it = mBitmaps.entrySet().iterator();
+    while (it.hasNext()) {
+      Map.Entry<Long, Bitmap> e = it.next();
+      if (genOf(e.getKey()) < mCurrentGen - 1 && !mHeld.contains(e.getKey())) {
+        Bitmap b = e.getValue();
+        if (b != mDisplayed && !b.isRecycled()) {
+          b.recycle();
+        }
+        it.remove();
+      }
+    }
+  }
+
+  private void recycleToken(long token) {
+    Bitmap b = mBitmaps.remove(token);
+    if (b != null && b != mDisplayed && !b.isRecycled()) {
+      b.recycle();
+    }
+  }
+
+  @Override
+  protected void onDraw(Canvas canvas) {
+    super.onDraw(canvas);
+    Bitmap bmp = mDisplayed;
+    if (bmp == null || bmp.isRecycled()) {
+      return;
+    }
+    // Always scale the last good frame to the current bounds, so a buffer that lags the view size
+    // (mid-resize) is shown stretched rather than blank.
+    mSrc.set(0, 0, mDisplayedW, mDisplayedH);
+    mDst.set(0, 0, getWidth(), getHeight());
+    canvas.drawBitmap(bmp, mSrc, mDst, mPaint);
+  }
+
+  // --- Lifecycle -------------------------------------------------------------
+
+  @Override
+  protected void onAttachedToWindow() {
+    super.onAttachedToWindow();
+    mAttached = true;
+    if (!mEnabled && getWidth() > 0 && getHeight() > 0) {
+      enablePool(getWidth(), getHeight());
+    }
+    // Pick up anything that became ready while we were detached.
+    onNativeFrameReady();
+  }
+
+  @Override
+  protected void onDetachedFromWindow() {
+    super.onDetachedFromWindow();
+    mAttached = false;
+    mEnabled = false;
+    mMainHandler.removeCallbacks(mConsume);
+    mConsumePosted.set(false);
+
+    // Keep the canvas alive offscreen, then drop all GPU resources. The
+    // SurfaceInfo stays registered so a transient detach/re-attach keeps
+    // working; WebGPUViewManager.onDropViewInstance removes it when RN drops
+    // the view for good.
+    nSwitchToOffscreen(contextId());
+
+    mDisplayed = null;
+    mHeld.clear();
+    for (Bitmap b : mBitmaps.values()) {
+      if (!b.isRecycled()) {
+        b.recycle();
+      }
+    }
+    mBitmaps.clear();
+    mCurrentGen = 0;
+  }
+
+  @Override
+  protected void onVisibilityChanged(@NonNull View changedView, int visibility) {
+    super.onVisibilityChanged(changedView, visibility);
+    if (visibility == VISIBLE && mAttached && !mEnabled
+        && getWidth() > 0 && getHeight() > 0) {
+      enablePool(getWidth(), getHeight());
+    }
+  }
+
+  // --- Native (cpp-adapter.cpp) ----------------------------------------------
+
+  private native void nEnablePool(int contextId, int dpW, int dpH);
+
+  private native void nSetClientSize(int contextId, int dpW, int dpH);
+
+  private native long nPollReady(int contextId);
+
+  private native HardwareBuffer nGetHardwareBuffer(int contextId, int generation, int slot);
+
+  private native void nReleaseSlot(int contextId, int generation, int slot);
+
+  private native void nSwitchToOffscreen(int contextId);
+}

package/android/src/main/java/com/webgpu/WebGPUView.java

@@ -1,6 +1,7 @@
 package com.webgpu;
 
 import android.content.Context;
+import android.os.Build;
 import android.view.Surface;
 import android.view.View;
 
@@ -10,12 +11,16 @@ import com.facebook.react.views.view.ReactViewGroup;
 
 public class WebGPUView extends ReactViewGroup implements WebGPUAPI {
 
+  // Backing view kinds, see updateView().
+  private static final int KIND_SURFACE_VIEW = 0;
+  private static final int KIND_TEXTURE_VIEW = 1;
+  private static final int KIND_HARDWARE_BUFFER_VIEW = 2;
 
   private int mContextId;
   private boolean mOpaque = true;
   private String mSurfaceType = "auto";
   private boolean mZOrderOnTop = false;
-  private boolean mAppliedTextureView;
+  private int mAppliedKind = -1;
   private boolean mAppliedZOrderOnTop;
   private WebGPUModule mModule;
   private View mView = null;
@@ -46,35 +51,62 @@ public class WebGPUView extends ReactViewGroup implements WebGPUAPI {
     mZOrderOnTop = value;
   }
 
+  // Resolve the backing view from the props. "auto" picks SurfaceView for an
+  // opaque canvas and TextureView for a non-opaque one. WebGPUHardwareBufferView
+  // (a plain View drawing AHardwareBuffers inline, API 29+) is opt-in via
+  // surfaceType until it gets a runtime fallback to TextureView when the
+  // AHardwareBuffer import fails on a device; below API 29 the request itself
+  // falls back to TextureView.
+  private int resolveKind() {
+    if ("SurfaceView".equals(mSurfaceType)) {
+      return KIND_SURFACE_VIEW;
+    }
+    if ("TextureView".equals(mSurfaceType)) {
+      return KIND_TEXTURE_VIEW;
+    }
+    if ("HardwareBufferView".equals(mSurfaceType)) {
+      boolean hardwareBufferSupported = Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q;
+      return hardwareBufferSupported ? KIND_HARDWARE_BUFFER_VIEW : KIND_TEXTURE_VIEW;
+    }
+    return mOpaque ? KIND_SURFACE_VIEW : KIND_TEXTURE_VIEW;
+  }
+
   // Apply the complete prop transaction once, after contextId and all rendering
   // options have arrived. Only a change of backing view (or of zOrderOnTop,
   // which a SurfaceView must know before it attaches) replaces the child;
   // opacity is applied in place.
   public void updateView() {
-    boolean textureView = "TextureView".equals(mSurfaceType)
-      || (!"SurfaceView".equals(mSurfaceType) && !mOpaque);
-    boolean zOrderOnTop = !textureView && mZOrderOnTop;
-    if (mView == null || textureView != mAppliedTextureView
-        || zOrderOnTop != mAppliedZOrderOnTop) {
+    int kind = resolveKind();
+    boolean zOrderOnTop = kind == KIND_SURFACE_VIEW && mZOrderOnTop;
+    if (mView == null || kind != mAppliedKind || zOrderOnTop != mAppliedZOrderOnTop) {
       if (mView != null) {
         removeView(mView);
       }
-      mAppliedTextureView = textureView;
+      mAppliedKind = kind;
       mAppliedZOrderOnTop = zOrderOnTop;
       Context ctx = getContext();
-      mView = textureView
-        ? new WebGPUTextureView(ctx, this, mOpaque)
-        : new WebGPUSurfaceView(ctx, this, zOrderOnTop, mOpaque);
+      switch (kind) {
+        case KIND_HARDWARE_BUFFER_VIEW:
+          mView = new WebGPUHardwareBufferView(ctx, this);
+          break;
+        case KIND_TEXTURE_VIEW:
+          mView = new WebGPUTextureView(ctx, this, mOpaque);
+          break;
+        default:
+          mView = new WebGPUSurfaceView(ctx, this, zOrderOnTop, mOpaque);
+          break;
+      }
       addView(mView);
       // ReactViewGroup.requestLayout() is a deliberate no-op, so addView outside
       // the layout pass never lays the child out; do it by hand or the new view
       // stays 0x0 and never gets a surface.
       layoutChild();
-    } else if (textureView) {
+    } else if (kind == KIND_TEXTURE_VIEW) {
       ((WebGPUTextureView) mView).setOpaque(mOpaque);
-    } else {
+    } else if (kind == KIND_SURFACE_VIEW) {
       ((WebGPUSurfaceView) mView).setOpaque(mOpaque);
     }
+    // WebGPUHardwareBufferView always composites with alpha; opacity is a no-op.
   }
 
   private void layoutChild() {
@@ -110,6 +142,11 @@ public class WebGPUView extends ReactViewGroup implements WebGPUAPI {
     switchToOffscreenSurface(mContextId);
   }
 
+  @Override
+  public int getContextId() {
+    return mContextId;
+  }
+
   /**
    * Called from WebGPUViewManager.onDropViewInstance when React removes this
    * view: the view dies with its Canvas, so it retires the registry entry.

package/cpp/jsi/JSICache.cpp

@@ -1,5 +1,6 @@
 #include "JSICache.h"
 
+#include <algorithm>
 #include <iterator>
 #include <stdexcept>
 #include <string>
@@ -94,4 +95,26 @@ jsi::Object &JSICache::setPrototype(PrototypeKey key, jsi::Object prototype) {
   return it->second;
 }
 
+jsi::Value JSICache::lockWrapper(jsi::Runtime &runtime,
+                                 const void *nativeObject) {
+  auto it = _wrappers.find(nativeObject);
+  if (it == _wrappers.end()) {
+    return jsi::Value::undefined();
+  }
+  return it->second.lock(runtime);
+}
+
+void JSICache::setWrapper(jsi::Runtime &runtime, const void *nativeObject,
+                          const jsi::Object &wrapper) {
+  _wrappers.insert_or_assign(nativeObject, jsi::WeakObject(runtime, wrapper));
+  if (_wrappers.size() < _pruneWrappersAt) {
+    return;
+  }
+  for (auto it = _wrappers.begin(); it != _wrappers.end();) {
+    it = it->second.lock(runtime).isUndefined() ? _wrappers.erase(it)
+                                                : std::next(it);
+  }
+  _pruneWrappersAt = std::max(kMinPruneWrappersAt, _wrappers.size() * 2);
+}
+
 } // namespace rnwgpu

package/cpp/jsi/JSICache.h

@@ -13,7 +13,8 @@ namespace jsi = facebook::jsi;
 
 /**
  * Per-runtime store for JSI values that native code wants to keep across
- * calls (currently: the prototype object of every NativeObject class).
+ * calls: the prototype object of every NativeObject class, and the JS
+ * wrapper of every native object that has one on the runtime.
  *
  * Ownership model (borrowed from react-native-nitro-modules' JSICache):
  *
@@ -72,6 +73,26 @@ public:
    */
   jsi::Object &setPrototype(PrototypeKey key, jsi::Object prototype);
 
+  /**
+   * Returns the JS wrapper stored for `nativeObject` on this runtime, or
+   * undefined if none was stored or the wrapper has been collected.
+   *
+   * The key is the native object's address. A live wrapper keeps its native
+   * object alive (the object is the wrapper's NativeState), so as long as
+   * the entry locks it still belongs to that object. Once the wrapper has
+   * been collected the entry is stale and is skipped or replaced, so an
+   * address reused by a new object can never resolve to the old wrapper.
+   */
+  jsi::Value lockWrapper(jsi::Runtime &runtime, const void *nativeObject);
+
+  /**
+   * Stores `wrapper` as the JS wrapper of `nativeObject` on this runtime.
+   * Only a weak reference is kept: the wrapper stays collectable, and the
+   * native object is released when it goes away.
+   */
+  void setWrapper(jsi::Runtime &runtime, const void *nativeObject,
+                  const jsi::Object &wrapper);
+
   JSICache() = default;
   ~JSICache() override = default;
   JSICache(const JSICache &) = delete;
@@ -81,6 +102,12 @@ private:
   static std::shared_ptr<JSICache> getOrCreateOnGlobal(jsi::Runtime &runtime);
 
   std::unordered_map<PrototypeKey, jsi::Object> _prototypes;
+  // Native object -> its wrapper on this runtime. Entries of collected
+  // wrappers are pruned in setWrapper once the map has doubled since the
+  // last pruning, so it stays proportional to the number of live wrappers.
+  std::unordered_map<const void *, jsi::WeakObject> _wrappers;
+  size_t _pruneWrappersAt = kMinPruneWrappersAt;
+  static constexpr size_t kMinPruneWrappersAt = 64;
 
   // Runtime -> cache. Weak on purpose: the runtime's global owns the cache.
   static std::mutex _registryMutex;