Fix Berry parsing from upstream (#25064)

lib/libesp32/berry/berry_port/be_code.py

@@ -1619,13 +1619,20 @@ if BE_USE_SCRIPT_COMPILER:
     # static void package_suffix(bfuncinfo *finfo, bexpdesc *c, bexpdesc *k)
     # {
     #     int key;
-    #     if (hasjump(k)) {
-    #         /* The key still carries open short-circuit jumps. Close them before
-    #          * materializing the object, otherwise the code that loads the object
-    #          * lands between those jumps and is skipped whenever the short circuit
-    #          * is taken. Materializing the object first is what 2785c08 changed in
-    #          * order to save a register in `self.a[128]`, so keep that order for
-    #          * every key that has no pending jumps. */
+    #     if (hasjump(k) || k->type == ETINDEX || k->type == ETMEMBER) {
+    #         /* Materialize the key before the object in two cases:
+    #          * - The key still carries open short-circuit jumps. Loading the object
+    #          *   first would place that load between those jumps, where it is
+    #          *   skipped whenever the short circuit is taken.
+    #          * - The key is itself an index or member expression that is not
+    #          *   materialized yet, as in `a[b[i+1]]` or `a[o.(n)]`. Loading the
+    #          *   object first puts it in a register above the key's own object and
+    #          *   index, and suffix_destreg() then releases it together with them
+    #          *   although it is still in use; an assignment then leaves freereg
+    #          *   below the number of active locals.
+    #          * Materializing the object first is what 2785c08 changed in order to
+    #          * save a register in `self.a[128]`, so keep that order for every
+    #          * other key. */
     #         key = exp2anyreg(finfo, k);
     #         c->v.ss.obj = exp2anyreg(finfo, c);
     #     } else {
@@ -1637,13 +1644,20 @@ if BE_USE_SCRIPT_COMPILER:
     # }
     def package_suffix(finfo, c, k):
         """Package a suffix object from c with key k."""
-        if hasjump(k):
-            # The key still carries open short-circuit jumps. Close them before
-            # materializing the object, otherwise the code that loads the object
-            # lands between those jumps and is skipped whenever the short circuit
-            # is taken. Materializing the object first is what 2785c08 changed in
-            # order to save a register in `self.a[128]`, so keep that order for
-            # every key that has no pending jumps.
+        if hasjump(k) or k.type == ETINDEX or k.type == ETMEMBER:
+            # Materialize the key before the object in two cases:
+            # - The key still carries open short-circuit jumps. Loading the object
+            #   first would place that load between those jumps, where it is
+            #   skipped whenever the short circuit is taken.
+            # - The key is itself an index or member expression that is not
+            #   materialized yet, as in `a[b[i+1]]` or `a[o.(n)]`. Loading the
+            #   object first puts it in a register above the key's own object and
+            #   index, and suffix_destreg() then releases it together with them
+            #   although it is still in use; an assignment then leaves freereg
+            #   below the number of active locals.
+            # Materializing the object first is what 2785c08 changed in order to
+            # save a register in `self.a[128]`, so keep that order for every
+            # other key.
             key = exp2anyreg(finfo, k)
             c.v.ss.obj = exp2anyreg(finfo, c)
         else:

lib/libesp32/berry/src/be_code.c

@@ -884,15 +884,22 @@ void be_code_ret(bfuncinfo *finfo, bexpdesc *e)
 static void package_suffix(bfuncinfo *finfo, bexpdesc *c, bexpdesc *k)
 {
     int key;
-    if (hasjump(k)) {
-        /* The key still carries open short-circuit jumps. Close them before
-         * materializing the object, otherwise the code that loads the object
-         * lands between those jumps and is skipped whenever the short circuit
-         * is taken. Materializing the object first is what 2785c08 changed in
-         * order to save a register in `self.a[128]`, so keep that order for
-         * every key that has no pending jumps. */
+    if (hasjump(k) || k->type == ETINDEX || k->type == ETMEMBER) {
+        /* Materialize the key before the object in two cases:
+         * - The key still carries open short-circuit jumps. Loading the object
+         *   first would place that load between those jumps, where it is
+         *   skipped whenever the short circuit is taken.
+         * - The key is itself an index or member expression that is not
+         *   materialized yet, as in `a[b[i+1]]` or `a[o.(n)]`. Loading the
+         *   object first puts it in a register above the key's own object and
+         *   index, and suffix_destreg() then releases it together with them
+         *   although it is still in use; an assignment then leaves freereg
+         *   below the number of active locals.
+         * Materializing the object first is what 2785c08 changed in order to
+         * save a register in `self.a[128]`, so keep that order for every
+         * other key. */
         key = exp2anyreg(finfo, k);
-    c->v.ss.obj = exp2anyreg(finfo, c);
+        c->v.ss.obj = exp2anyreg(finfo, c);
     } else {
         c->v.ss.obj = exp2anyreg(finfo, c);
         key = exp2anyreg(finfo, k);

lib/libesp32/berry/tests/compiler.be

@@ -78,3 +78,79 @@ def test()
   for n : 1..2 end
 end
 test()
+
+# Index assignment whose key is itself an index or member expression.
+# The object of the assignment was loaded into a register above the key's
+# own object and index, suffix_destreg() released it together with them, and
+# freereg ended up below the number of active locals. The statement itself
+# was right, but the next one used the register of the last local as a
+# scratch register:
+#
+# ```berry
+# def f()
+#   var a = 10
+#   var b = 20
+#   m[m[4]] = 1
+#   var c = m[0] + 1    # overwrote `b`: returned [10, 2, 2]
+#   return [a, b, c]
+# end
+# ```
+#
+# With no local in scope, compilation failed with
+# `register overflow (more than 255)` instead.
+nk_m = [0, 0, 0, 0, 0]
+def nk_f()
+  var a = 10
+  var b = 20
+  nk_m[nk_m[4]] = 1
+  var c = nk_m[0] + 1
+  return [a, b, c]
+end
+assert(nk_f() == [10, 20, 2])
+
+# computed inner index and compound assignments
+nk_l = [0, 1, 2, 3, 4, 5, 6, 7]
+def nk_g(i)
+  var a = 10
+  var b = 20
+  nk_l[nk_l[i + 1]] = 9     # nk_l[2] = 9
+  nk_l[nk_l[5]] <<= 1       # nk_l[5] = 10
+  nk_l[nk_l[7]] += 1        # nk_l[7] = 8
+  var c = nk_l[0] + 1
+  return [a, b, c]
+end
+assert(nk_g(1) == [10, 20, 1])
+assert(nk_l == [0, 1, 9, 3, 4, 10, 6, 8])
+
+# the same through instance members
+class nk_C
+  var map, keys
+  def init() self.map = {} self.keys = ['x', 'y'] end
+  def set(i, v)
+    var n = 42
+    self.map[self.keys[i + 1]] = v
+    var s = self.keys[0]
+    return [n, s]
+  end
+end
+nk_c = nk_C()
+assert(nk_c.set(0, 1) == [42, 'x'])
+assert(nk_c.map['y'] == 1)
+
+# a member with a computed name as the key
+class nk_D var x end
+nk_d = nk_D()
+nk_d.x = 1
+def nk_h()
+  var a = 10
+  var n = 'x'
+  nk_m[nk_d.(n + '')] = 5
+  var s = nk_m[0] + 1
+  return [a, n, s]
+end
+assert(nk_h() == [10, 'x', 2])
+
+# no local in scope: used to fail to compile
+nk_n = nil
+compile("nk_n = [0, 0, 0, 0, 0]  nk_n[nk_n[4]] = 1  nk_n[1] = nk_n[0] + 1")()
+assert(nk_n == [1, 2, 0, 0, 0])