Berry various fixes for Walrus Operator (#18982)

CHANGELOG.md

@@ -11,6 +11,7 @@ All notable changes to this project will be documented in this file.
 ### Changed
 
 ### Fixed
+- Berry various fixes for Walrus Operator
 
 ### Removed
 

lib/libesp32/berry/src/be_code.c

@@ -78,20 +78,11 @@ static int codeABx(bfuncinfo *finfo, bopcode op, int a, int bx)
     return codeinst(finfo, ISET_OP(op) | ISET_RA(a) | ISET_Bx(bx));
 }
 
-/* Move value from register b to register a */
-static void code_move_nooptim(bfuncinfo *finfo, int a, int b)
-{
-    if (isK(b)) {
-        codeABx(finfo, OP_LDCONST, a, b & 0xFF);
-    } else {
-        codeABC(finfo, OP_MOVE, a, b, 0);
-    }
-}
-
 /* Move value from register b to register a */
 /* Check the previous instruction to compact both instruction as one if possible */
 /* If b is a constant, add LDCONST or add MOVE otherwise */
-static void code_move(bfuncinfo *finfo, int a, int b)
+/* returns false if the move operation happened, or true if there was a register optimization and `b` should be replaced by `a` */
+static bbool code_move(bfuncinfo *finfo, int a, int b)
 {
     if (finfo->pc) {  /* If not the first instruction of the function */
         binstruction *i = be_vector_end(&finfo->code);  /* get the last instruction */
@@ -101,11 +92,16 @@ static void code_move(bfuncinfo *finfo, int a, int b)
             int x = IGET_RA(*i), y = IGET_RKB(*i), z = IGET_RKC(*i);
             if (b == x && (a == y || (op < OP_NEG && a == z))) {
                 *i = (*i & ~IRA_MASK) | ISET_RA(a);
-                return;
+                return btrue;
             }
         }
     }
-    code_move_nooptim(finfo, a, b);
+    if (isK(b)) {
+        codeABx(finfo, OP_LDCONST, a, b & 0xFF);
+    } else {
+        codeABC(finfo, OP_MOVE, a, b, 0);
+    }
+    return bfalse;
 }
 
 /* Free register at top (checks that it´s a register) */
@@ -113,7 +109,7 @@ static void code_move(bfuncinfo *finfo, int a, int b)
 static void free_expreg(bfuncinfo *finfo, bexpdesc *e)
 {
     /* release temporary register */
-    if (e && e->type == ETREG) {
+    if (e && e->type == ETREG && e->v.idx == finfo->freereg - 1) {      /* free ETREG only if it's top of stack */
         be_code_freeregs(finfo, 1);
     }
 }
@@ -690,10 +686,10 @@ int be_code_setvar(bfuncinfo *finfo, bexpdesc *e1, bexpdesc *e2, bbool keep_reg)
     switch (e1->type) {
     case ETLOCAL: /* It can't be ETREG. */
         if (e1->v.idx != src) {
-            if (keep_reg) {
-                code_move_nooptim(finfo, e1->v.idx, src); /* always do explicit move */
-            } else {
-                code_move(finfo, e1->v.idx, src); /* do explicit move only if needed */
+            bbool reg_optimized = code_move(finfo, e1->v.idx, src); /* do explicit move only if needed */
+            if (reg_optimized) {
+                free_expreg(finfo, e2); /* free source (checks only ETREG) */
+                *e2 = *e1;      /* now e2 is e1 ETLOCAL */
             }
         }
         break;
@@ -725,7 +721,7 @@ int be_code_nextreg(bfuncinfo *finfo, bexpdesc *e)
 {
     int dst = finfo->freereg;
     int src = exp2anyreg(finfo, e); /* get variable register index */
-    if (e->type != ETREG) { /* move local and const to new register */
+    if ((e->type != ETREG) || (src < dst - 1)) { /* move local and const to new register, don't move if already top of stack */
         code_move(finfo, dst, src);
         be_code_allocregs(finfo, 1);
     } else {

lib/libesp32/berry/src/be_lexer.c

@@ -41,7 +41,7 @@ static const char* const kwords_tab[] = {
     "for", "def", "end", "class", "break", "continue",
     "return", "true", "false", "nil", "var", "do",
     "import", "as", "try", "except", "raise", "static",
-    // ".f"
+    ":=",
 };
 
 void be_lexerror(blexer *lexer, const char *msg)

lib/libesp32/berry/src/be_parser.c

@@ -1131,9 +1131,11 @@ static void walrus_expr(bparser *parser, bexpdesc *e)
     sub_expr(parser, e, ASSIGN_OP_PRIO);    /* left expression */
     btokentype op = next_type(parser);
     if (op == OptWalrus) {
+        check_symbol(parser, e);
         bexpdesc e1 = *e;           /* copy var to e1, e will get the result of expression */
         scan_next_token(parser);    /* skip ':=' */
         expr(parser, e);
+        check_var(parser, e);
         if (check_newvar(parser, &e1)) { /* new variable */
             new_var(parser, e1.v.s, e);
         }