From a053e639147d97c4a306ab272c12d9520a80e805 Mon Sep 17 00:00:00 2001
From: Damien George <damien@micropython.org>
Date: Mon, 17 Jun 2024 23:32:09 +1000
Subject: [PATCH] webassembly/objjsproxy: Implement proxying of JS iterable
 protocol.

This allows Python to iterate over JavaScript objects that provide
Symbol.iterator.

Signed-off-by: Damien George <damien@micropython.org>
---
 ports/webassembly/objjsproxy.c                | 33 ++++++++++---------
 tests/ports/webassembly/js_proxy_iterator.mjs | 14 ++++++++
 .../webassembly/js_proxy_iterator.mjs.exp     |  4 +++
 3 files changed, 36 insertions(+), 15 deletions(-)
 create mode 100644 tests/ports/webassembly/js_proxy_iterator.mjs
 create mode 100644 tests/ports/webassembly/js_proxy_iterator.mjs.exp

diff --git a/ports/webassembly/objjsproxy.c b/ports/webassembly/objjsproxy.c
index 5c09e003f..cbfe8be49 100644
--- a/ports/webassembly/objjsproxy.c
+++ b/ports/webassembly/objjsproxy.c
@@ -153,14 +153,21 @@ EM_JS(void, js_reflect_construct, (int f_ref, uint32_t n_args, uint32_t * args,
     proxy_convert_js_to_mp_obj_jsside(ret, out);
 });
 
-EM_JS(int, js_get_len, (int f_ref), {
-    return proxy_js_ref[f_ref].length;
+EM_JS(void, js_get_iter, (int f_ref, uint32_t * out), {
+    const f = proxy_js_ref[f_ref];
+    const ret = f[Symbol.iterator]();
+    proxy_convert_js_to_mp_obj_jsside(ret, out);
 });
 
-EM_JS(void, js_subscr_int, (int f_ref, int idx, uint32_t * out), {
+EM_JS(bool, js_iter_next, (int f_ref, uint32_t * out), {
     const f = proxy_js_ref[f_ref];
-    const ret = f[idx];
-    proxy_convert_js_to_mp_obj_jsside(ret, out);
+    const ret = f.next();
+    if (ret.done) {
+        return false;
+    } else {
+        proxy_convert_js_to_mp_obj_jsside(ret.value, out);
+        return true;
+    }
 });
 
 EM_JS(void, js_subscr_load, (int f_ref, uint32_t * index_ref, uint32_t * out), {
@@ -320,17 +327,13 @@ void mp_obj_jsproxy_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
 typedef struct _jsproxy_it_t {
     mp_obj_base_t base;
     mp_fun_1_t iternext;
-    mp_obj_jsproxy_t *obj;
-    uint16_t cur;
-    uint16_t len;
+    mp_obj_jsproxy_t *iter;
 } jsproxy_it_t;
 
 static mp_obj_t jsproxy_it_iternext(mp_obj_t self_in) {
     jsproxy_it_t *self = MP_OBJ_TO_PTR(self_in);
-    if (self->cur < self->len) {
-        uint32_t out[3];
-        js_subscr_int(self->obj->ref, self->cur, out);
-        self->cur += 1;
+    uint32_t out[3];
+    if (js_iter_next(self->iter->ref, out)) {
         return proxy_convert_js_to_mp_obj_cside(out);
     } else {
         return MP_OBJ_STOP_ITERATION;
@@ -343,9 +346,9 @@ static mp_obj_t jsproxy_new_it(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
     jsproxy_it_t *o = (jsproxy_it_t *)iter_buf;
     o->base.type = &mp_type_polymorph_iter;
     o->iternext = jsproxy_it_iternext;
-    o->obj = self;
-    o->cur = 0;
-    o->len = js_get_len(self->ref);
+    uint32_t out[3];
+    js_get_iter(self->ref, out);
+    o->iter = proxy_convert_js_to_mp_obj_cside(out);
     return MP_OBJ_FROM_PTR(o);
 }
 
diff --git a/tests/ports/webassembly/js_proxy_iterator.mjs b/tests/ports/webassembly/js_proxy_iterator.mjs
new file mode 100644
index 000000000..34b54b483
--- /dev/null
+++ b/tests/ports/webassembly/js_proxy_iterator.mjs
@@ -0,0 +1,14 @@
+// Test accessing JavaScript iterables (objects with Symbol.iterator) from Python.
+
+const mp = await (await import(process.argv[2])).loadMicroPython();
+
+mp.runPython(`
+import js
+
+for v in js.Set.new([1, 2]):
+    print(v)
+
+url_search_params = js.URLSearchParams.new("one=1&two=2")
+for key in url_search_params.keys():
+    print(key, list(url_search_params.getAll(key)))
+`);
diff --git a/tests/ports/webassembly/js_proxy_iterator.mjs.exp b/tests/ports/webassembly/js_proxy_iterator.mjs.exp
new file mode 100644
index 000000000..e8d11286c
--- /dev/null
+++ b/tests/ports/webassembly/js_proxy_iterator.mjs.exp
@@ -0,0 +1,4 @@
+1
+2
+one ['1']
+two ['2']