From 0518bd6c64b5516e805d3e352f31d4bd10fd3058 Mon Sep 17 00:00:00 2001
From: s-hadinger <49731213+s-hadinger@users.noreply.github.com>
Date: Tue, 12 Mar 2024 23:00:52 +0100
Subject: [PATCH] Berry fix walrus with member or index (#20939)

---
 CHANGELOG.md                       |  1 +
 lib/libesp32/berry/src/be_code.c   | 12 ++++++++++++
 lib/libesp32/berry/tests/walrus.be | 23 +++++++++++++++++++++++
 3 files changed, 36 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a428dfe5e..3f7f5eadc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -27,6 +27,7 @@ All notable changes to this project will be documented in this file.
 - Berry bug when parsing ternary operator (#20839)
 - HASPmota widgets line, btnmatrix, qrcode, bar, checkbox (#20881)
 - Filesystem save of JSON settings data
+- Berry fix walrus with member or index
 
 ### Removed
 
diff --git a/lib/libesp32/berry/src/be_code.c b/lib/libesp32/berry/src/be_code.c
index 103dbb31a..56cb3b527 100644
--- a/lib/libesp32/berry/src/be_code.c
+++ b/lib/libesp32/berry/src/be_code.c
@@ -729,9 +729,21 @@ int be_code_setvar(bfuncinfo *finfo, bexpdesc *e1, bexpdesc *e2, bbool keep_reg)
         break;
     case ETMEMBER: /* store to member R(A).RK(B) <- RK(C) */
         setsfxvar(finfo, OP_SETMBR, e1, src);
+        if (keep_reg && e2->type == ETREG) {
+            /* special case of walrus assignemnt when we need to recreate an ETREG */
+            code_move(finfo, e1->v.ss.obj, src);    /* move from ETREG to MEMBER instance*/
+            free_expreg(finfo, e2); /* free source (checks only ETREG) */
+            e2->v.idx = e1->v.ss.obj; /* update to new register */
+        }
         break;
     case ETINDEX: /* store to member R(A)[RK(B)] <- RK(C) */
         setsfxvar(finfo, OP_SETIDX, e1, src);
+        if (keep_reg && e2->type == ETREG) {
+            /* special case of walrus assignemnt when we need to recreate an ETREG */
+            code_move(finfo, e1->v.ss.obj, src);
+            free_expreg(finfo, e2); /* free source (checks only ETREG) */
+            e2->v.idx = e1->v.ss.obj;
+        }
         break;
     default:
         return 1;
diff --git a/lib/libesp32/berry/tests/walrus.be b/lib/libesp32/berry/tests/walrus.be
index faa92fa15..adfef27f2 100644
--- a/lib/libesp32/berry/tests/walrus.be
+++ b/lib/libesp32/berry/tests/walrus.be
@@ -33,3 +33,26 @@ assert_attribute_error("var a,b def f() a end")
 
 # while the following does have side effect
 def f() a := b end
+
+# bug when using walrus with member
+def id(x) return x end
+var a = 1
+import global
+def f() return id(global.a := 42) end
+assert(f() == 42)
+# bug: returns <module: global>
+
+def concat(x, y, z) return str(x)+str(y)+str(z) end
+var a = 1
+import global
+def f() return concat(global.a := 1, global.a := 42, global.a := 0) end
+assert(f() == "1420")
+# bug: returns '1<module: global>42'
+
+# same bug when using index
+def id(x) return x end
+l = [10,11]
+import global
+def f() return id(global.l[0] := 42) end
+assert(f() == 42)
+# bug: returns [42, 11]