webassembly: Reuse PyProxy objects when they are the same Python object.

This commit makes it so that PyProxy objects are reused (on the JavaScript side) when they correspond to an existing Python object that is the same object. For example, proxying the same Python function to JavaScript, the same PyProxy instance is now used. This means that if `foo` is a Python function then accessing it on the JavaScript side such as `api.globals().get("foo")` has the property that: api.globals().get("foo") === api.globals().get("foo") Prior to this commit the above was not true because new PyProxy instances were created each time `foo` was accessed. Signed-off-by: Damien George <[email protected]>
picoscratch · Jul 19, 2024 · 77bd8fe · 77bd8fe
1 parent 5147dc5
commit 77bd8fe
Show file tree

Hide file tree

Showing 5 changed files with 129 additions and 27 deletions.
diff --git a/ports/webassembly/proxy_c.c b/ports/webassembly/proxy_c.c
@@ -49,6 +49,7 @@ enum {
     PROXY_KIND_MP_GENERATOR = 7,
     PROXY_KIND_MP_OBJECT = 8,
     PROXY_KIND_MP_JSPROXY = 9,
+    PROXY_KIND_MP_EXISTING = 10,
 };
 
 enum {
@@ -79,40 +80,76 @@ static size_t proxy_c_ref_next;
 
 void proxy_c_init(void) {
     MP_STATE_PORT(proxy_c_ref) = mp_obj_new_list(0, NULL);
+    MP_STATE_PORT(proxy_c_dict) = mp_obj_new_dict(0);
     mp_obj_list_append(MP_STATE_PORT(proxy_c_ref), MP_OBJ_NULL);
     proxy_c_ref_next = PROXY_C_REF_NUM_STATIC;
 }
 
 MP_REGISTER_ROOT_POINTER(mp_obj_t proxy_c_ref);
+MP_REGISTER_ROOT_POINTER(mp_obj_t proxy_c_dict);
 
 // obj cannot be MP_OBJ_NULL.
 static inline size_t proxy_c_add_obj(mp_obj_t obj) {
     // Search for the first free slot in proxy_c_ref.
+    size_t id = 0;
     mp_obj_list_t *l = (mp_obj_list_t *)MP_OBJ_TO_PTR(MP_STATE_PORT(proxy_c_ref));
     while (proxy_c_ref_next < l->len) {
         if (l->items[proxy_c_ref_next] == MP_OBJ_NULL) {
             // Free slot found, reuse it.
-            size_t id = proxy_c_ref_next;
+            id = proxy_c_ref_next;
             ++proxy_c_ref_next;
             l->items[id] = obj;
-            return id;
+            break;
         }
         ++proxy_c_ref_next;
     }
 
-    // No free slots, so grow proxy_c_ref by one (append at the end of the list).
-    size_t id = l->len;
-    mp_obj_list_append(MP_STATE_PORT(proxy_c_ref), obj);
-    proxy_c_ref_next = l->len;
+    if (id == 0) {
+        // No free slots, so grow proxy_c_ref by one (append at the end of the list).
+        id = l->len;
+        mp_obj_list_append(MP_STATE_PORT(proxy_c_ref), obj);
+        proxy_c_ref_next = l->len;
+    }
+
+    // Add the object to proxy_c_dict, keyed by the object pointer, with value the object id.
+    mp_obj_t obj_key = mp_obj_new_int_from_uint((uintptr_t)obj);
+    mp_map_elem_t *elem = mp_map_lookup(mp_obj_dict_get_map(MP_STATE_PORT(proxy_c_dict)), obj_key, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND);
+    elem->value = mp_obj_new_int_from_uint(id);
+
     return id;
 }
 
+EM_JS(int, js_check_existing, (int c_ref), {
+    return proxy_js_check_existing(c_ref);
+});
+
+// obj cannot be MP_OBJ_NULL.
+static inline int proxy_c_check_existing(mp_obj_t obj) {
+    mp_obj_t obj_key = mp_obj_new_int_from_uint((uintptr_t)obj);
+    mp_map_elem_t *elem = mp_map_lookup(mp_obj_dict_get_map(MP_STATE_PORT(proxy_c_dict)), obj_key, MP_MAP_LOOKUP);
+    if (elem == NULL) {
+        return -1;
+    }
+    uint32_t c_ref = mp_obj_int_get_truncated(elem->value);
+    return js_check_existing(c_ref);
+}
+
 static inline mp_obj_t proxy_c_get_obj(uint32_t c_ref) {
     return ((mp_obj_list_t *)MP_OBJ_TO_PTR(MP_STATE_PORT(proxy_c_ref)))->items[c_ref];
 }
 
 void proxy_c_free_obj(uint32_t c_ref) {
     if (c_ref >= PROXY_C_REF_NUM_STATIC) {
+        // Remove the object from proxy_c_dict if the c_ref in that dict corresponds to this object.
+        // (It may be that this object exists in the dict but with a different c_ref from a more
+        // recent proxy of this object.)
+        mp_obj_t obj_key = mp_obj_new_int_from_uint((uintptr_t)proxy_c_get_obj(c_ref));
+        mp_map_elem_t *elem = mp_map_lookup(mp_obj_dict_get_map(MP_STATE_PORT(proxy_c_dict)), obj_key, MP_MAP_LOOKUP);
+        if (elem != NULL && mp_obj_int_get_truncated(elem->value) == c_ref) {
+            mp_map_lookup(mp_obj_dict_get_map(MP_STATE_PORT(proxy_c_dict)), obj_key, MP_MAP_LOOKUP_REMOVE_IF_FOUND);
+        }
+
+        // Clear the slot in proxy_c_ref used by this object, so the GC can reclaim the object.
         ((mp_obj_list_t *)MP_OBJ_TO_PTR(MP_STATE_PORT(proxy_c_ref)))->items[c_ref] = MP_OBJ_NULL;
         proxy_c_ref_next = MIN(proxy_c_ref_next, c_ref);
     }
@@ -143,6 +180,7 @@ mp_obj_t proxy_convert_js_to_mp_obj_cside(uint32_t *value) {
 
 void proxy_convert_mp_to_js_obj_cside(mp_obj_t obj, uint32_t *out) {
     uint32_t kind;
+    int js_ref;
     if (obj == MP_OBJ_NULL) {
         kind = PROXY_KIND_MP_NULL;
     } else if (obj == mp_const_none) {
@@ -168,6 +206,9 @@ void proxy_convert_mp_to_js_obj_cside(mp_obj_t obj, uint32_t *out) {
     } else if (mp_obj_is_jsproxy(obj)) {
         kind = PROXY_KIND_MP_JSPROXY;
         out[1] = mp_obj_jsproxy_get_ref(obj);
+    } else if ((js_ref = proxy_c_check_existing(obj)) >= 0) {
+        kind = PROXY_KIND_MP_EXISTING;
+        out[1] = js_ref;
     } else if (mp_obj_get_type(obj) == &mp_type_JsException) {
         mp_obj_exception_t *exc = MP_OBJ_TO_PTR(obj);
         if (exc->args->len > 0 && mp_obj_is_jsproxy(exc->args->items[0])) {

diff --git a/ports/webassembly/proxy_js.js b/ports/webassembly/proxy_js.js
@@ -40,6 +40,7 @@ const PROXY_KIND_MP_CALLABLE = 6;
 const PROXY_KIND_MP_GENERATOR = 7;
 const PROXY_KIND_MP_OBJECT = 8;
 const PROXY_KIND_MP_JSPROXY = 9;
+const PROXY_KIND_MP_EXISTING = 10;
 
 const PROXY_KIND_JS_UNDEFINED = 0;
 const PROXY_KIND_JS_NULL = 1;
@@ -61,13 +62,39 @@ class PythonError extends Error {
 function proxy_js_init() {
     globalThis.proxy_js_ref = [globalThis, undefined];
     globalThis.proxy_js_ref_next = PROXY_JS_REF_NUM_STATIC;
+    globalThis.proxy_js_map = new Map();
+    globalThis.proxy_js_existing = [undefined];
     globalThis.pyProxyFinalizationRegistry = new FinalizationRegistry(
         (cRef) => {
+            globalThis.proxy_js_map.delete(cRef);
             Module.ccall("proxy_c_free_obj", "null", ["number"], [cRef]);
         },
     );
 }
 
+// Check if the c_ref (Python proxy index) has a corresponding JavaScript-side PyProxy
+// associated with it.  If so, take a concrete reference to this PyProxy from the WeakRef
+// and put it in proxy_js_existing, to be referenced and reused by PROXY_KIND_MP_EXISTING.
+function proxy_js_check_existing(c_ref) {
+    const existing_obj = globalThis.proxy_js_map.get(c_ref)?.deref();
+    if (existing_obj === undefined) {
+        return -1;
+    }
+
+    // Search for a free slot in proxy_js_existing.
+    for (let i = 0; i < globalThis.proxy_js_existing.length; ++i) {
+        if (globalThis.proxy_js_existing[i] === undefined) {
+            // Free slot found, put existing_obj here and return the index.
+            globalThis.proxy_js_existing[i] = existing_obj;
+            return i;
+        }
+    }
+
+    // No free slot, so append to proxy_js_existing and return the new index.
+    globalThis.proxy_js_existing.push(existing_obj);
+    return globalThis.proxy_js_existing.length - 1;
+}
+
 // js_obj cannot be undefined
 function proxy_js_add_obj(js_obj) {
     // Search for the first free slot in proxy_js_ref.
@@ -241,6 +268,10 @@ function proxy_convert_mp_to_js_obj_jsside(value) {
         // js proxy
         const id = Module.getValue(value + 4, "i32");
         obj = proxy_js_ref[id];
+    } else if (kind === PROXY_KIND_MP_EXISTING) {
+        const id = Module.getValue(value + 4, "i32");
+        obj = globalThis.proxy_js_existing[id];
+        globalThis.proxy_js_existing[id] = undefined;
     } else {
         // obj
         const id = Module.getValue(value + 4, "i32");
@@ -257,6 +288,7 @@ function proxy_convert_mp_to_js_obj_jsside(value) {
             obj = new Proxy(target, py_proxy_handler);
         }
         globalThis.pyProxyFinalizationRegistry.register(obj, id);
+        globalThis.proxy_js_map.set(id, new WeakRef(obj));
     }
     return obj;
 }

diff --git a/tests/ports/webassembly/heap_expand.mjs.exp b/tests/ports/webassembly/heap_expand.mjs.exp
@@ -1,27 +1,27 @@
-135241360
 135241328
 135241296
 135241264
-135241216
-135241168
-135241088
-135240944
-135240640
-135240112
-135239072
-135237008
-135232896
-135224688
-135208288
-135175504
-135109888
-134978800
-134716640
-135216784
-136217152
-138217840
-142219296
-150222224
+135241232
+135241184
+135241136
+135241056
+135240912
+135240608
+135240080
+135239040
+135236976
+135232864
+135224656
+135208256
+135175472
+135109856
+134978768
+134716608
+135216752
+136217120
+138217808
+142219264
+150222192
 1
 2
 4

diff --git a/tests/ports/webassembly/py_proxy_identity.mjs b/tests/ports/webassembly/py_proxy_identity.mjs
@@ -0,0 +1,26 @@
+// Test identity of PyProxy when they are the same Python object.
+
+const mp = await (await import(process.argv[2])).loadMicroPython();
+
+mp.runPython(`
+l = []
+`);
+
+const l1 = mp.globals.get("l");
+const l2 = mp.globals.get("l");
+console.log(l1, l2);
+console.log(l1 === l2);
+
+globalThis.eventTarget = new EventTarget();
+globalThis.event = new Event("event");
+
+mp.runPython(`
+import js
+
+def callback(ev):
+    print("callback", ev)
+js.eventTarget.addEventListener("event", callback)
+js.eventTarget.dispatchEvent(js.event)
+js.eventTarget.removeEventListener("event", callback)
+js.eventTarget.dispatchEvent(js.event)
+`);
diff --git a/tests/ports/webassembly/py_proxy_identity.mjs.exp b/tests/ports/webassembly/py_proxy_identity.mjs.exp
@@ -0,0 +1,3 @@
+PyProxy { _ref: 3 } PyProxy { _ref: 3 }
+true
+callback <JsProxy 7>