From ad99bd5b210a31bba1707d73fb546a2939b0b128 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveTH@foxmail.com>
Date: Sun, 19 Feb 2023 23:26:57 +0800
Subject: [PATCH] some optimize

---
 src/compiler.h | 9 ++++++++-
 src/pocketpy.h | 4 +++-
 src/vm.h       | 2 +-
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/compiler.h b/src/compiler.h
index d5b167e8..73bbb5ad 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -401,11 +401,18 @@ private:
     }
 
     void exprAssign() {
+        int lhs = co()->codes.empty() ? -1 : co()->codes.size() - 1;
         co()->_rvalue = true;
         TokenIndex op = parser->prev.type;
         if(op == TK("=")) {     // a = (expr)
             EXPR_TUPLE();
-            emit(OP_STORE_REF);
+            if(lhs!=-1 && co()->codes[lhs].op == OP_LOAD_NAME_REF){
+                emit(OP_STORE_NAME, co()->codes[lhs].arg);
+                co()->codes[lhs].op = OP_NO_OP;
+                co()->codes[lhs].arg = -1;
+            }else{
+                emit(OP_STORE_REF);
+            }
         }else{                  // a += (expr) -> a = a + (expr)
             EXPR();
             switch (op) {
diff --git a/src/pocketpy.h b/src/pocketpy.h
index bdd10c15..72918acb 100644
--- a/src/pocketpy.h
+++ b/src/pocketpy.h
@@ -591,7 +591,9 @@ void add_module_math(VM* vm){
 void add_module_dis(VM* vm){
     PyVar mod = vm->new_module("dis");
     vm->bind_func<1>(mod, "dis", [](VM* vm, pkpy::Args& args) {
-        CodeObject_ code = vm->PyFunction_AS_C(args[0]).code;
+        PyVar f = args[0];
+        if(f->is_type(vm->tp_bound_method)) f = vm->PyBoundMethod_AS_C(args[0]).method;
+        CodeObject_ code = vm->PyFunction_AS_C(f).code;
         (*vm->_stdout) << vm->disassemble(code);
         return vm->None;
     });
diff --git a/src/vm.h b/src/vm.h
index 841397af..d13e46a1 100644
--- a/src/vm.h
+++ b/src/vm.h
@@ -483,7 +483,7 @@ public:
             if(byte.op == OP_LOAD_CONST){
                 argStr += " (" + PyStr_AS_C(asRepr(co->consts[byte.arg])) + ")";
             }
-            if(byte.op == OP_LOAD_NAME_REF || byte.op == OP_LOAD_NAME || byte.op == OP_RAISE){
+            if(byte.op == OP_LOAD_NAME_REF || byte.op == OP_LOAD_NAME || byte.op == OP_RAISE || byte.op == OP_STORE_NAME){
                 argStr += " (" + co->names[byte.arg].first.escape(true) + ")";
             }
             if(byte.op == OP_FAST_INDEX || byte.op == OP_FAST_INDEX_REF){