webassembly/proxy_js: Allow a Python proxy of a function to be undone.

This optimises the case where a Python function is, for example, stored to
a JavaScript attribute and then later retrieved from Python.  The Python
function no longer needs to be a proxy with double proxying needed for the
call from Python -> JavaScript -> Python.

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

4 files changed, 92 insertions, 7 deletions

diff --git a/ports/webassembly/objjsproxy.c b/ports/webassembly/objjsproxy.c
index 5e2aeb6a3..098f4e75f 100644
--- a/ports/webassembly/objjsproxy.c
+++ b/ports/webassembly/objjsproxy.c
@@ -48,9 +48,17 @@ EM_JS(bool, lookup_attr, (int jsref, const char *str, uint32_t * out), {
     const attr = UTF8ToString(str);
     if (attr in base) {
         let value = base[attr];
-        if (typeof value == "function") {
+        if (typeof value === "function") {
             if (base !== globalThis) {
-                value = value.bind(base);
+                if ("_ref" in value) {
+                    // This is a proxy of a Python function, it doesn't need
+                    // binding.  And not binding it means if it's passed back
+                    // to Python then it can be extracted from the proxy as a
+                    // true Python function.
+                } else {
+                    // A function that is not a Python function.  Bind it.
+                    value = value.bind(base);
+                }
             }
         }
         proxy_convert_js_to_mp_obj_jsside(value, out);
diff --git a/ports/webassembly/proxy_js.js b/ports/webassembly/proxy_js.js
index 1f60bf41e..042deb7b0 100644
--- a/ports/webassembly/proxy_js.js
+++ b/ports/webassembly/proxy_js.js
@@ -135,10 +135,11 @@ function proxy_convert_js_to_mp_obj_jsside(js_obj, out) {
         Module.stringToUTF8(js_obj, buf, len + 1);
         Module.setValue(out + 4, len, "i32");
         Module.setValue(out + 8, buf, "i32");
-    } else if (js_obj instanceof PyProxy) {
-        kind = PROXY_KIND_JS_PYPROXY;
-        Module.setValue(out + 4, js_obj._ref, "i32");
-    } else if (js_obj instanceof PyProxyThenable) {
+    } else if (
+        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 {
@@ -151,7 +152,11 @@ function proxy_convert_js_to_mp_obj_jsside(js_obj, out) {
 }
 
 function proxy_convert_js_to_mp_obj_jsside_force_double_proxy(js_obj, out) {
-    if (js_obj instanceof PyProxy) {
+    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_ref.length;
         proxy_js_ref[id] = js_obj;
@@ -212,6 +217,7 @@ function proxy_convert_mp_to_js_obj_jsside(value) {
             obj = (...args) => {
                 return proxy_call_python(id, args);
             };
+            obj._ref = id;
         } else if (kind === PROXY_KIND_MP_GENERATOR) {
             obj = new PyProxyThenable(id);
         } else {
diff --git a/tests/ports/webassembly/fun_proxy.mjs b/tests/ports/webassembly/fun_proxy.mjs
new file mode 100644
index 000000000..87cce44f2
--- /dev/null
+++ b/tests/ports/webassembly/fun_proxy.mjs
@@ -0,0 +1,52 @@
+// Test proxying of functions between Python and JavaScript.
+
+const mp = await (await import(process.argv[2])).loadMicroPython();
+
+// Create JS functions living on the JS and Py sides.
+globalThis.jsFunAdd = (x, y) => x + y;
+mp.globals.set("js_fun_sub", (x, y) => x - y);
+
+console.log("== JavaScript side ==");
+
+// JS function living on the JS side, should be a function.
+console.log(globalThis.jsFunAdd);
+console.log(globalThis.jsFunAdd(1, 2));
+
+// JS function living on the Py side, should be a function.
+console.log(mp.globals.get("js_fun_sub"));
+console.log(mp.globals.get("js_fun_sub")(1, 2));
+
+mp.runPython(`
+import js
+
+print("== Python side ==")
+
+py_fun_mul = lambda x, y: x * y
+js.pyFunDiv = lambda x, y: x / y
+
+# JS function living on the JS side, should be a JsProxy.
+print(type(js.jsFunAdd))
+print(js.jsFunAdd(1, 2))
+
+# JS function living on the Py side, should be a JsProxy.
+print(type(js_fun_sub))
+print(js_fun_sub(1, 2))
+
+# Py function living on the Py side, should be a function.
+print(type(py_fun_mul))
+print(py_fun_mul(2, 3))
+
+# Py function living on the JS side, should be a function.
+print(type(js.pyFunDiv))
+print(js.pyFunDiv(6, 2))
+`);
+
+console.log("== JavaScript side ==");
+
+// Py function living on the Py side, should be a proxy function.
+console.log(mp.globals.get("py_fun_mul"));
+console.log(mp.globals.get("py_fun_mul")(2, 3));
+
+// Py function living on the JS side, should be a proxy function.
+console.log(globalThis.pyFunDiv);
+console.log(globalThis.pyFunDiv(6, 2));
diff --git a/tests/ports/webassembly/fun_proxy.mjs.exp b/tests/ports/webassembly/fun_proxy.mjs.exp
new file mode 100644
index 000000000..4aeaa47b2
--- /dev/null
+++ b/tests/ports/webassembly/fun_proxy.mjs.exp
@@ -0,0 +1,19 @@
+== JavaScript side ==
+[Function (anonymous)]
+3
+[Function (anonymous)]
+-1
+== Python side ==
+<class 'JsProxy'>
+3
+<class 'JsProxy'>
+-1
+<class 'function'>
+6
+<class 'function'>
+3.0
+== JavaScript side ==
+[Function: obj] { _ref: 4 }
+6
+[Function: obj] { _ref: 3 }
+3