webassembly: Improve identity and fix bug with lost JsProxy refs.

Commit ffa98cb0143c43af9f4c61142784a08a19f660c5 improved equality for
`JsProxy` objects so that, eg, `js.Object == js.Object` is true.

As mentioned in #17758, a further optimisation is to make identity work in
that case, eg `js.Object is js.Object` should be true (on the Python side).

This commit implements that, by keeping track of all `JsProxy` Python
objects and reusing them where possible: where the underlying JS ref is
equal, ie they point to the same JS object.  That reduces memory churn and
gives better identity behaviour of JS objects proxied over to Python.

As part of this, a bug is fixed where JS objects can be freed while there's
still a `JsProxy` referring to that JS object.  A test is added for that
exact scenario, and the test now passes.

Signed-off-by: Damien George <damien@micropython.org>

9 files changed, 155 insertions, 52 deletions

diff --git a/ports/webassembly/objjsproxy.c b/ports/webassembly/objjsproxy.c
index a8b21a744..2d46702ff 100644
--- a/ports/webassembly/objjsproxy.c
+++ b/ports/webassembly/objjsproxy.c
@@ -32,6 +32,9 @@
 #include "py/runtime.h"
 #include "proxy_c.h"
 
+static mp_obj_t *jsproxy_table = NULL;
+static size_t jsproxy_table_len = 0;
+
 EM_JS(bool, has_attr, (int jsref, const char *str), {
     const base = proxy_js_ref[jsref];
     const attr = UTF8ToString(str);
@@ -295,6 +298,7 @@ EM_JS(void, proxy_js_free_obj, (int js_ref), {
 
 static mp_obj_t jsproxy___del__(mp_obj_t self_in) {
     mp_obj_jsproxy_t *self = MP_OBJ_TO_PTR(self_in);
+    jsproxy_table[self->ref] = MP_OBJ_NULL;
     proxy_js_free_obj(self->ref);
     return mp_const_none;
 }
@@ -590,13 +594,40 @@ MP_DEFINE_CONST_OBJ_TYPE(
     iter, jsproxy_getiter
     );
 
+void mp_obj_jsproxy_init(void) {
+    jsproxy_table = NULL;
+    jsproxy_table_len = 0;
+    MP_STATE_PORT(jsproxy_global_this) = mp_obj_new_jsproxy(MP_OBJ_JSPROXY_REF_GLOBAL_THIS);
+}
+
+MP_REGISTER_ROOT_POINTER(mp_obj_t jsproxy_global_this);
+
 mp_obj_t mp_obj_new_jsproxy(int ref) {
+    // The proxy for this ref should not exist.
+    assert(ref >= jsproxy_table_len || jsproxy_table[ref] == MP_OBJ_NULL);
+
     mp_obj_jsproxy_t *o = mp_obj_malloc_with_finaliser(mp_obj_jsproxy_t, &mp_type_jsproxy);
     o->ref = ref;
     o->bind_to_self = false;
+    if (ref >= jsproxy_table_len) {
+        size_t new_len = MAX(16, ref * 2);
+        jsproxy_table = realloc(jsproxy_table, new_len * sizeof(mp_obj_t));
+        for (size_t i = jsproxy_table_len; i < new_len; ++i) {
+            jsproxy_table[i] = MP_OBJ_NULL;
+        }
+        jsproxy_table_len = new_len;
+    }
+    jsproxy_table[ref] = MP_OBJ_FROM_PTR(o);
     return MP_OBJ_FROM_PTR(o);
 }
 
+mp_obj_t mp_obj_get_jsproxy(int ref) {
+    // The proxy for this ref should exist.
+    assert(ref < jsproxy_table_len && jsproxy_table[ref] != MP_OBJ_NULL);
+
+    return jsproxy_table[ref];
+}
+
 // Load/delete/store an attribute from/to the JavaScript globalThis entity.
 void mp_obj_jsproxy_global_this_attr(qstr attr, mp_obj_t *dest) {
     if (dest[0] == MP_OBJ_NULL) {
diff --git a/ports/webassembly/proxy_c.c b/ports/webassembly/proxy_c.c
index 459f6a6da..790ad90ef 100644
--- a/ports/webassembly/proxy_c.c
+++ b/ports/webassembly/proxy_c.c
@@ -59,8 +59,9 @@ enum {
     PROXY_KIND_JS_INTEGER = 3,
     PROXY_KIND_JS_DOUBLE = 4,
     PROXY_KIND_JS_STRING = 5,
-    PROXY_KIND_JS_OBJECT = 6,
-    PROXY_KIND_JS_PYPROXY = 7,
+    PROXY_KIND_JS_OBJECT_EXISTING = 6,
+    PROXY_KIND_JS_OBJECT = 7,
+    PROXY_KIND_JS_PYPROXY = 8,
 };
 
 MP_DEFINE_CONST_OBJ_TYPE(
@@ -83,6 +84,9 @@ void proxy_c_init(void) {
     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;
+
+    void mp_obj_jsproxy_init(void);
+    mp_obj_jsproxy_init();
 }
 
 MP_REGISTER_ROOT_POINTER(mp_obj_t proxy_c_ref);
@@ -172,6 +176,8 @@ mp_obj_t proxy_convert_js_to_mp_obj_cside(uint32_t *value) {
         return s;
     } else if (value[0] == PROXY_KIND_JS_PYPROXY) {
         return proxy_c_get_obj(value[1]);
+    } else if (value[0] == PROXY_KIND_JS_OBJECT_EXISTING) {
+        return mp_obj_get_jsproxy(value[1]);
     } else {
         // PROXY_KIND_JS_OBJECT
         return mp_obj_new_jsproxy(value[1]);
diff --git a/ports/webassembly/proxy_c.h b/ports/webassembly/proxy_c.h
index bac0a90bd..2f37aedbc 100644
--- a/ports/webassembly/proxy_c.h
+++ b/ports/webassembly/proxy_c.h
@@ -53,6 +53,7 @@ void proxy_convert_mp_to_js_obj_cside(mp_obj_t obj, uint32_t *out);
 void proxy_convert_mp_to_js_exc_cside(void *exc, uint32_t *out);
 
 mp_obj_t mp_obj_new_jsproxy(int ref);
+mp_obj_t mp_obj_get_jsproxy(int ref);
 void mp_obj_jsproxy_global_this_attr(qstr attr, mp_obj_t *dest);
 
 static inline bool mp_obj_is_jsproxy(mp_obj_t o) {
diff --git a/ports/webassembly/proxy_js.js b/ports/webassembly/proxy_js.js
index cbd6e5b00..60b832d81 100644
--- a/ports/webassembly/proxy_js.js
+++ b/ports/webassembly/proxy_js.js
@@ -48,8 +48,9 @@ const PROXY_KIND_JS_BOOLEAN = 2;
 const PROXY_KIND_JS_INTEGER = 3;
 const PROXY_KIND_JS_DOUBLE = 4;
 const PROXY_KIND_JS_STRING = 5;
-const PROXY_KIND_JS_OBJECT = 6;
-const PROXY_KIND_JS_PYPROXY = 7;
+const PROXY_KIND_JS_OBJECT_EXISTING = 6;
+const PROXY_KIND_JS_OBJECT = 7;
+const PROXY_KIND_JS_PYPROXY = 8;
 
 class PythonError extends Error {
     constructor(exc_type, exc_details) {
@@ -63,6 +64,7 @@ function proxy_js_init() {
     globalThis.proxy_js_ref = [globalThis, undefined];
     globalThis.proxy_js_ref_next = PROXY_JS_REF_NUM_STATIC;
     globalThis.proxy_js_ref_map = new Map();
+    globalThis.proxy_js_ref_map.set(globalThis, 0);
     globalThis.proxy_js_map = new Map();
     globalThis.proxy_js_existing = [undefined];
     globalThis.pyProxyFinalizationRegistry = new FinalizationRegistry(
@@ -99,12 +101,6 @@ function proxy_js_check_existing(c_ref) {
 // The `js_obj` argument cannot be `undefined`.
 // Returns an integer reference to the given `js_obj`.
 function proxy_js_add_obj(js_obj) {
-    // See if there is an existing JsProxy reference, and use that if there is.
-    const existing_ref = proxy_js_ref_map.get(js_obj);
-    if (existing_ref !== undefined) {
-        return existing_ref;
-    }
-
     // Search for the first free slot in proxy_js_ref.
     while (proxy_js_ref_next < proxy_js_ref.length) {
         if (proxy_js_ref[proxy_js_ref_next] === undefined) {
@@ -175,7 +171,7 @@ function proxy_call_python(target, argumentsList) {
     return ret;
 }
 
-function proxy_convert_js_to_mp_obj_jsside(js_obj, out) {
+function proxy_convert_js_to_mp_obj_jsside_helper(js_obj, out, allow_pyproxy) {
     let kind;
     if (js_obj === undefined) {
         kind = PROXY_KIND_JS_UNDEFINED;
@@ -206,33 +202,35 @@ function proxy_convert_js_to_mp_obj_jsside(js_obj, out) {
         Module.setValue(out + 4, len, "i32");
         Module.setValue(out + 8, buf, "i32");
     } else if (
-        js_obj instanceof PyProxy ||
-        (typeof js_obj === "function" && "_ref" in js_obj) ||
-        js_obj instanceof PyProxyThenable
+        allow_pyproxy &&
+        (js_obj instanceof PyProxy ||
+            (typeof js_obj === "function" && "_ref" in js_obj) ||
+            js_obj instanceof PyProxyThenable)
     ) {
         kind = PROXY_KIND_JS_PYPROXY;
         Module.setValue(out + 4, js_obj._ref, "i32");
     } else {
-        kind = PROXY_KIND_JS_OBJECT;
-        const id = proxy_js_add_obj(js_obj);
+        let id;
+        // See if there is an existing JsProxy reference, and use that if there is.
+        const existing_ref = proxy_js_ref_map.get(js_obj);
+        if (existing_ref !== undefined) {
+            kind = PROXY_KIND_JS_OBJECT_EXISTING;
+            id = existing_ref;
+        } else {
+            kind = PROXY_KIND_JS_OBJECT;
+            id = proxy_js_add_obj(js_obj);
+        }
         Module.setValue(out + 4, id, "i32");
     }
     Module.setValue(out + 0, kind, "i32");
 }
 
+function proxy_convert_js_to_mp_obj_jsside(js_obj, out) {
+    proxy_convert_js_to_mp_obj_jsside_helper(js_obj, out, true);
+}
+
 function proxy_convert_js_to_mp_obj_jsside_force_double_proxy(js_obj, out) {
-    if (
-        js_obj instanceof PyProxy ||
-        (typeof js_obj === "function" && "_ref" in js_obj) ||
-        js_obj instanceof PyProxyThenable
-    ) {
-        const kind = PROXY_KIND_JS_OBJECT;
-        const id = proxy_js_add_obj(js_obj);
-        Module.setValue(out + 4, id, "i32");
-        Module.setValue(out + 0, kind, "i32");
-    } else {
-        proxy_convert_js_to_mp_obj_jsside(js_obj, out);
-    }
+    proxy_convert_js_to_mp_obj_jsside_helper(js_obj, out, false);
 }
 
 function proxy_convert_mp_to_js_obj_jsside(value) {
diff --git a/tests/ports/webassembly/heap_expand.mjs.exp b/tests/ports/webassembly/heap_expand.mjs.exp
index 563413514..67ebe98e7 100644
--- a/tests/ports/webassembly/heap_expand.mjs.exp
+++ b/tests/ports/webassembly/heap_expand.mjs.exp
@@ -1,27 +1,27 @@
-135241328
-135241296
-135241264
-135241232
-135241184
-135241136
-135241056
-135240912
-135240608
-135240080
-135239040
-135236976
-135232864
-135224656
-135208256
-135175472
-135109856
-134978768
-134716608
-135216752
-136217120
-138217808
-142219264
-150222192
+135241312
+135241280
+135241248
+135241216
+135241168
+135241120
+135241040
+135240896
+135240592
+135240064
+135239024
+135236960
+135232848
+135224640
+135208240
+135175456
+135109840
+134978752
+134716592
+135216800
+136217168
+138217984
+142219568
+150222816
 1
 2
 4
diff --git a/tests/ports/webassembly/js_proxy_identity.mjs b/tests/ports/webassembly/js_proxy_identity.mjs
index e279d219d..ca2f3980a 100644
--- a/tests/ports/webassembly/js_proxy_identity.mjs
+++ b/tests/ports/webassembly/js_proxy_identity.mjs
@@ -11,8 +11,15 @@ print("Object equality")
 print(js.Object == js.Object)
 print(js.Object.assign == js.Object.assign)
 
+print("Object identity")
+print(js.Object is js.Object)
+
 print("Array equality")
 print(js.Array == js.Array)
 print(js.Array.prototype == js.Array.prototype)
 print(js.Array.prototype.push == js.Array.prototype.push)
+
+print("Array identity")
+print(js.Array is js.Array)
+print(js.Array.prototype is js.Array.prototype)
 `);
diff --git a/tests/ports/webassembly/js_proxy_identity.mjs.exp b/tests/ports/webassembly/js_proxy_identity.mjs.exp
index 5791d911b..d8f1ae891 100644
--- a/tests/ports/webassembly/js_proxy_identity.mjs.exp
+++ b/tests/ports/webassembly/js_proxy_identity.mjs.exp
@@ -2,7 +2,12 @@
 Object equality
 True
 True
+Object identity
+True
 Array equality
 True
 True
 True
+Array identity
+True
+True
diff --git a/tests/ports/webassembly/js_proxy_reuse_free.mjs b/tests/ports/webassembly/js_proxy_reuse_free.mjs
new file mode 100644
index 000000000..ebca86f0b
--- /dev/null
+++ b/tests/ports/webassembly/js_proxy_reuse_free.mjs
@@ -0,0 +1,48 @@
+// Test reuse of JsProxy references and freeing of JsProxy objects.
+// This ensures that a Python-side JsProxy that refers to a JavaScript object retains
+// the correct JavaScript object in the case that another JsProxy that refers to the
+// same JavaScript object is freed.
+
+const mp = await (await import(process.argv[2])).loadMicroPython();
+
+globalThis.obj = [1, 2];
+globalThis.obj2 = [3, 4];
+
+console.log("JS obj:", globalThis.obj);
+
+mp.runPython(`
+import gc
+import js
+
+# Create 2 proxies of the same JS object.
+# They should refer to the same underlying JS-side reference.
+obj = js.obj
+obj_copy = js.obj
+print(obj, obj_copy, obj == obj_copy)
+
+# Print out the object.
+js.console.log("Py obj:", obj)
+
+# Forget obj_copy and trigger a GC when the Python code finishes.
+obj_copy = None
+gc.collect()
+`);
+
+console.log("JS obj:", globalThis.obj);
+
+mp.runPython(`
+# Create a new proxy of a different object.
+# It should not clobber the existing obj proxy reference.
+obj2 = js.obj2
+
+# Create a copy of the existing obj proxy.
+obj_copy = js.obj
+
+# Print the JS proxy, it should be the same reference as before.
+print(obj, obj_copy, obj == obj_copy)
+
+# Print out the object.
+js.console.log("Py obj:", obj)
+`);
+
+console.log("JS obj:", globalThis.obj);
diff --git a/tests/ports/webassembly/js_proxy_reuse_free.mjs.exp b/tests/ports/webassembly/js_proxy_reuse_free.mjs.exp
new file mode 100644
index 000000000..c74e4f49e
--- /dev/null
+++ b/tests/ports/webassembly/js_proxy_reuse_free.mjs.exp
@@ -0,0 +1,7 @@
+JS obj: [ 1, 2 ]
+<JsProxy 2> <JsProxy 2> True
+Py obj: [ 1, 2 ]
+JS obj: [ 1, 2 ]
+<JsProxy 2> <JsProxy 2> True
+Py obj: [ 1, 2 ]
+JS obj: [ 1, 2 ]