webassembly/proxy_js: Reuse JsProxy ref if object matches.

This reduces memory use by reusing objects, and improves identity/equality
relationships of JavaScript objects on the Python side.

In 77bd8fe5b80b0e7e02cdb6b4272c401ae3dca638 PyProxy's were reused when the
same Python object was proxied across to JavaScript.  This commit does the
same thing but for JsProxy's going from JS to Python.  If an existing
JsProxy reference exists for the JS object about to be proxied across, then
it's reused.

This helps reduce the number of alive objects (memory use), and, more
importantly, improves equality relationships of JavaScript objects on the
Python side.  Eg we now get, on the Python side:

    import js

    print(js.Object == js.Object)

that prints True.  Previously it was False.

Note that this change does not make identity work with `is`, for example
`js.Object is js.Object` is actually False.  With more work that could be
made True but for now we leave that as-is.

The behaviour with this commit matches Pyodide semantics.

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

5 files changed, 32 insertions, 5 deletions

diff --git a/ports/webassembly/objjsproxy.c b/ports/webassembly/objjsproxy.c
index 45f329d7e..a8b21a744 100644
--- a/ports/webassembly/objjsproxy.c
+++ b/ports/webassembly/objjsproxy.c
@@ -285,6 +285,7 @@ static mp_obj_t jsproxy_binary_op(mp_binary_op_t op, mp_obj_t lhs_in, mp_obj_t r
 
 EM_JS(void, proxy_js_free_obj, (int js_ref), {
     if (js_ref >= PROXY_JS_REF_NUM_STATIC) {
+        proxy_js_ref_map.delete(proxy_js_ref[js_ref]);
         proxy_js_ref[js_ref] = undefined;
         if (js_ref < proxy_js_ref_next) {
             proxy_js_ref_next = js_ref;
diff --git a/ports/webassembly/proxy_js.js b/ports/webassembly/proxy_js.js
index 9e7c233e3..cbd6e5b00 100644
--- a/ports/webassembly/proxy_js.js
+++ b/ports/webassembly/proxy_js.js
@@ -62,6 +62,7 @@ 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_ref_map = new Map();
     globalThis.proxy_js_map = new Map();
     globalThis.proxy_js_existing = [undefined];
     globalThis.pyProxyFinalizationRegistry = new FinalizationRegistry(
@@ -95,8 +96,15 @@ function proxy_js_check_existing(c_ref) {
     return globalThis.proxy_js_existing.length - 1;
 }
 
-// js_obj cannot be undefined
+// 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) {
@@ -104,6 +112,7 @@ function proxy_js_add_obj(js_obj) {
             const id = proxy_js_ref_next;
             ++proxy_js_ref_next;
             proxy_js_ref[id] = js_obj;
+            proxy_js_ref_map.set(js_obj, id);
             return id;
         }
         ++proxy_js_ref_next;
@@ -113,6 +122,7 @@ function proxy_js_add_obj(js_obj) {
     const id = proxy_js_ref.length;
     proxy_js_ref[id] = js_obj;
     proxy_js_ref_next = proxy_js_ref.length;
+    proxy_js_ref_map.set(js_obj, id);
     return id;
 }
 
diff --git a/tests/ports/webassembly/py_proxy_identity.mjs b/tests/ports/webassembly/py_proxy_identity.mjs
index d4a720b73..97dab2e78 100644
--- a/tests/ports/webassembly/py_proxy_identity.mjs
+++ b/tests/ports/webassembly/py_proxy_identity.mjs
@@ -23,4 +23,13 @@ js.eventTarget.addEventListener("event", callback)
 js.eventTarget.dispatchEvent(js.event)
 js.eventTarget.removeEventListener("event", callback)
 js.eventTarget.dispatchEvent(js.event)
+
+print("Object equality")
+print(js.Object == js.Object)
+print(js.Object.assign == js.Object.assign)
+
+print("Array equality")
+print(js.Array == js.Array)
+print(js.Array.prototype == js.Array.prototype)
+print(js.Array.prototype.push == js.Array.prototype.push)
 `);
diff --git a/tests/ports/webassembly/py_proxy_identity.mjs.exp b/tests/ports/webassembly/py_proxy_identity.mjs.exp
index 01ccf0d89..344a0a202 100644
--- a/tests/ports/webassembly/py_proxy_identity.mjs.exp
+++ b/tests/ports/webassembly/py_proxy_identity.mjs.exp
@@ -1,3 +1,10 @@
 PyProxy { _ref: 3 } PyProxy { _ref: 3 }
 true
-callback <JsProxy 7>
+callback <JsProxy 5>
+Object equality
+True
+True
+Array equality
+True
+True
+True
diff --git a/tests/ports/webassembly/run_python_async.mjs.exp b/tests/ports/webassembly/run_python_async.mjs.exp
index ad6c49e33..4dff64a60 100644
--- a/tests/ports/webassembly/run_python_async.mjs.exp
+++ b/tests/ports/webassembly/run_python_async.mjs.exp
@@ -2,16 +2,16 @@
 1
 <JsProxy 2>
 py 1
-<JsProxy 5>
+<JsProxy 4>
 py 2
 2
 resolved 123
 3
 = TEST 2 ==========
 1
-<JsProxy 6>
+<JsProxy 5>
 py 1
-<JsProxy 9>
+<JsProxy 6>
 py 2
 2
 setTimeout resolved