fix inplace op

include/pocketpy/expr.h

@@ -65,6 +65,7 @@ struct Expr{
     virtual bool is_literal() const { return false; }
     virtual bool is_json_object() const { return false; }
     virtual bool is_attrib() const { return false; }
+    virtual bool is_subscr() const { return false; }
     virtual bool is_compare() const { return false; }
     virtual int star_level() const { return 0; }
     virtual bool is_tuple() const { return false; }
@@ -80,6 +81,14 @@ struct Expr{
     [[nodiscard]] virtual bool emit_store(CodeEmitContext* ctx) {
         return false;
     }
+
+    virtual void emit_inplace(CodeEmitContext* ctx) {
+        emit_(ctx);
+    }
+
+    [[nodiscard]] virtual bool emit_store_inplace(CodeEmitContext* ctx) {
+        return emit_store(ctx);
+    }
 };
 
 struct CodeEmitContext{
@@ -316,9 +325,13 @@ struct FStringExpr: Expr{
 struct SubscrExpr: Expr{
     Expr_ a;
     Expr_ b;
+    bool is_subscr() const override { return true; }
     void emit_(CodeEmitContext* ctx) override;
     bool emit_del(CodeEmitContext* ctx) override;
     bool emit_store(CodeEmitContext* ctx) override;
+
+    void emit_inplace(CodeEmitContext* ctx) override;
+    bool emit_store_inplace(CodeEmitContext* ctx) override;
 };
 
 struct AttribExpr: Expr{
@@ -330,7 +343,10 @@ struct AttribExpr: Expr{
     bool emit_del(CodeEmitContext* ctx) override;
     bool emit_store(CodeEmitContext* ctx) override;
     void emit_method(CodeEmitContext* ctx);
+
     bool is_attrib() const override { return true; }
+    void emit_inplace(CodeEmitContext* ctx) override;
+    bool emit_store_inplace(CodeEmitContext* ctx) override;
 };
 
 struct CallExpr: Expr{
@@ -362,6 +378,9 @@ struct BinaryExpr: Expr{
     TokenIndex op;
     Expr_ lhs;
     Expr_ rhs;
+    bool inplace;
+
+    BinaryExpr(bool inplace=false): inplace(inplace) {}
     bool is_compare() const override;
     void _emit_compare(CodeEmitContext*, small_vector_2<int, 6>&);
     void emit_(CodeEmitContext* ctx) override;

include/pocketpy/opcodes.h

@@ -5,6 +5,7 @@ OPCODE(NO_OP)
 /**************************/
 OPCODE(POP_TOP)
 OPCODE(DUP_TOP)
+OPCODE(DUP_TOP_TWO)
 OPCODE(ROT_TWO)
 OPCODE(ROT_THREE)
 OPCODE(PRINT_EXPR)

src/ceval.cpp

@@ -137,8 +137,14 @@ __NEXT_STEP:
     /*****************************************/
     case OP_POP_TOP: POP(); DISPATCH()
     case OP_DUP_TOP: PUSH(TOP()); DISPATCH()
+    case OP_DUP_TOP_TWO:
+        // [a, b]
+        PUSH(SECOND()); // [a, b, a]
+        PUSH(SECOND()); // [a, b, a, b]
+        DISPATCH()
     case OP_ROT_TWO: std::swap(TOP(), SECOND()); DISPATCH()
     case OP_ROT_THREE:{
+        // [a, b, c] -> [c, a, b]
         PyVar _0 = TOP();
         TOP() = SECOND();
         SECOND() = THIRD();

src/compiler.cpp

@@ -790,14 +790,16 @@ __EAT_DOTS_END:
                 if(lhs_p->is_starred()) SyntaxError();
                 if(ctx()->is_compiling_class) SyntaxError("can't use inplace operator in class definition");
                 advance();
-                auto e = make_expr<BinaryExpr>();
+                // a[x] += 1;   a and x should be evaluated only once
+                // a.x += 1;    a should be evaluated only once
+                auto e = make_expr<BinaryExpr>(true);  // inplace=true
                 e->op = prev().type - 1; // -1 to remove =
                 e->lhs = ctx()->s_expr.popx();
                 EXPR_TUPLE();
                 e->rhs = ctx()->s_expr.popx();
-                if(e->is_starred()) SyntaxError();
+                if(e->rhs->is_starred()) SyntaxError();
                 e->emit_(ctx());
-                bool ok = lhs_p->emit_store(ctx());
+                bool ok = lhs_p->emit_store_inplace(ctx());
                 if(!ok) SyntaxError();
             } return true;
             case TK("="): {

src/expr.cpp

@@ -569,6 +569,20 @@ namespace pkpy{
         return true;
     }
 
+    void SubscrExpr::emit_inplace(CodeEmitContext* ctx){
+        a->emit_(ctx);
+        b->emit_(ctx);
+        ctx->emit_(OP_DUP_TOP_TWO, BC_NOARG, line);
+        ctx->emit_(OP_LOAD_SUBSCR, BC_NOARG, line);
+    }
+
+    bool SubscrExpr::emit_store_inplace(CodeEmitContext* ctx){
+        // [a, b, val] -> [val, a, b]
+        ctx->emit_(OP_ROT_THREE, BC_NOARG, line);
+        ctx->emit_(OP_STORE_SUBSCR, BC_NOARG, line);
+        return true;
+    }
+
     bool SubscrExpr::emit_del(CodeEmitContext* ctx){
         a->emit_(ctx);
         b->emit_(ctx);
@@ -598,6 +612,19 @@ namespace pkpy{
         ctx->emit_(OP_LOAD_METHOD, b.index, line);
     }
 
+    void AttribExpr::emit_inplace(CodeEmitContext* ctx) {
+        a->emit_(ctx);
+        ctx->emit_(OP_DUP_TOP, BC_NOARG, line);
+        ctx->emit_(OP_LOAD_ATTR, b.index, line);
+    }
+
+    bool AttribExpr::emit_store_inplace(CodeEmitContext* ctx) {
+        // [a, val] -> [val, a]
+        ctx->emit_(OP_ROT_TWO, BC_NOARG, line);
+        ctx->emit_(OP_STORE_ATTR, b.index, line);
+        return true;
+    }
+
     void CallExpr::emit_(CodeEmitContext* ctx) {
         bool vargs = false;
         bool vkwargs = false;
@@ -689,7 +716,11 @@ namespace pkpy{
             // [b, RES]
         }else{
             // (1 + 2) < c
-            lhs->emit_(ctx);
+            if(inplace){
+                lhs->emit_inplace(ctx);
+            }else{
+                lhs->emit_(ctx);
+            }
         }
 
         rhs->emit_(ctx);

src/vm.cpp

@@ -834,7 +834,7 @@ void VM::__log_s_data(const char* title) {
     }
     output.push_back(']');
     Bytecode byte = *frame->_ip;
-    std::cout << output << " " << OP_NAMES[byte.op] << " " << _opcode_argstr(nullptr, byte, frame->co) << std::endl;
+    std::cout << output << " " << OP_NAMES[byte.op] << " " << _opcode_argstr(nullptr, frame->ip(), byte, frame->co) << std::endl;
 }
 #endif
 

tests/99_bugs.py

@@ -104,3 +104,14 @@ assert g(**g(**ret)) == ret
 
 # other known issues:
 # 1. d.extend(d) if d is deque
+
+g = 0
+def test():
+    global g
+    g += 1
+    return g
+
+a = [1, 10, 3]
+a[test()] += 1
+assert (a == [1, 11, 3]), a
+assert (g == 1), g