From c95483844a419864ec0975b89f57cc5645246690 Mon Sep 17 00:00:00 2001
From: s-hadinger <49731213+s-hadinger@users.noreply.github.com>
Date: Fri, 22 Mar 2024 18:51:48 +0100
Subject: [PATCH] Berry fix walrus bug when assigning to self (#21015)

---
 CHANGELOG.md                       |  2 +-
 lib/libesp32/berry/src/be_code.c   | 14 +++-----------
 lib/libesp32/berry/tests/walrus.be | 14 ++++++++++++++
 3 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0d5e8724e..0ce36f348 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,7 +14,7 @@ All notable changes to this project will be documented in this file.
 - LVGL library from v9.0.0 to v9.1.0
 
 ### Fixed
-
+- Berry fix walrus bug when assigning to self
 
 ### Removed
 
diff --git a/lib/libesp32/berry/src/be_code.c b/lib/libesp32/berry/src/be_code.c
index 56cb3b527..d6e1be99e 100644
--- a/lib/libesp32/berry/src/be_code.c
+++ b/lib/libesp32/berry/src/be_code.c
@@ -728,23 +728,15 @@ int be_code_setvar(bfuncinfo *finfo, bexpdesc *e1, bexpdesc *e2, bbool keep_reg)
         setsupvar(finfo, OP_SETUPV, e1, src);
         break;
     case ETMEMBER: /* store to member R(A).RK(B) <- RK(C) */
-        setsfxvar(finfo, OP_SETMBR, e1, src);
-        if (keep_reg && e2->type == ETREG) {
+    case ETINDEX: /* store to member R(A)[RK(B)] <- RK(C) */
+        setsfxvar(finfo, (e1->type == ETMEMBER) ? OP_SETMBR : OP_SETIDX, e1, src);
+        if (keep_reg && e2->type == ETREG && e1->v.ss.obj >= be_list_count(finfo->local)) {
             /* 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 adfef27f2..43ef77cfa 100644
--- a/lib/libesp32/berry/tests/walrus.be
+++ b/lib/libesp32/berry/tests/walrus.be
@@ -56,3 +56,17 @@ import global
 def f() return id(global.l[0] := 42) end
 assert(f() == 42)
 # bug: returns [42, 11]
+
+# bug when using member for self
+class confused_walrus
+    var b
+    def f()
+        var c = 1
+        if self.b := true
+            c = 2
+        end
+        return self
+    end
+end
+var ins = confused_walrus()
+assert(ins.f() == ins)