From bb651906c3c5f91fffc781d06188da9a4585eef4 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@hclcat.games>
Date: Tue, 8 Nov 2022 20:13:52 +0800
Subject: [PATCH] add listcomp support

---
 src/codeobject.h | 20 +++++++++++++++-----
 src/compiler.h   | 44 ++++++++++++++++++++++++++++++++++----------
 src/opcodes.h    |  3 +++
 src/vm.h         |  6 ++++++
 4 files changed, 58 insertions(+), 15 deletions(-)

diff --git a/src/codeobject.h b/src/codeobject.h
index 51e67fad..3f3258ba 100644
--- a/src/codeobject.h
+++ b/src/codeobject.h
@@ -55,6 +55,12 @@ struct CodeObject {
         return co_consts.size() - 1;
     }
 
+    void __copyToEnd(int start, int end){
+        auto _start = co_code.begin() + start;
+        auto _end = co_code.begin() + end;
+        co_code.insert(co_code.end(), _start, _end);
+    }
+
     _Str toString(){
         _StrStream ss;
         int prev_line = -1;
@@ -95,7 +101,7 @@ struct CodeObject {
 
 class Frame {
 private:
-    std::stack<PyVar> s_data;
+    std::vector<PyVar> s_data;
     int ip = 0;
 public:
     StlDict* f_globals;
@@ -124,8 +130,8 @@ public:
     }
 
     inline PyVar __pop(){
-        PyVar v = s_data.top();
-        s_data.pop();
+        PyVar v = s_data.back();
+        s_data.pop_back();
         return v;
     }
 
@@ -136,11 +142,15 @@ public:
     }
 
     inline PyVar topValue(VM* vm){
-        return __deref_pointer(vm, s_data.top());
+        return __deref_pointer(vm, s_data.back());
+    }
+
+    inline PyVar topNValue(VM* vm, int n=-1){
+        return __deref_pointer(vm, s_data[s_data.size() + n]);
     }
 
     inline void push(PyVar v){
-        s_data.push(v);
+        s_data.push_back(v);
     }
 
     inline void jumpTo(int i){
diff --git a/src/compiler.h b/src/compiler.h
index 46f78182..91843acf 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -11,7 +11,7 @@
 class Compiler;
 
 typedef void (Compiler::*GrammarFn)();
-typedef void (Compiler::*CompilerAction)();
+typedef std::function<void(Compiler*)> CompilerAction;
 
 struct GrammarRule{
     GrammarFn prefix;
@@ -430,8 +430,33 @@ public:
     }
 
     void exprList() {
-        ExprCommaSplitArgs("]");
+        int _patch = emitCode(OP_NO_OP);
+        int _body_start = getCode()->co_code.size();
+        int ARGC = 0;
+        do {
+            matchNewLines();
+            if (peek() == TK("]")) break;
+            EXPR(); ARGC++;
+            matchNewLines();
+            if(ARGC == 1 && match(TK("for"))) goto __LISTCOMP;
+        } while (match(TK(",")));
+        matchNewLines();
+        consume(TK("]"));
         emitCode(OP_BUILD_LIST, ARGC);
+        return;
+
+__LISTCOMP:
+        int _body_end = getCode()->co_code.size();
+        getCode()->co_code[_patch].op = OP_JUMP_ABSOLUTE;
+        getCode()->co_code[_patch].arg = _body_end;
+        emitCode(OP_BUILD_LIST, 0);
+        __compileForLoop([=](Compiler* compiler){
+            // [list, iter]
+            getCode()->__copyToEnd(_body_start, _body_end);
+            // [list, iter, value]
+            emitCode(OP_LIST_APPEND);
+        });
+        consume(TK("]"));
     }
 
     void exprMap() {
@@ -538,7 +563,7 @@ public:
         }
         consume(TK("@indent"));
         while (peek() != TK("@dedent")) {
-            (this->*action)();
+            action(this);
             matchNewLines();
         }
         consume(TK("@dedent"));
@@ -617,7 +642,7 @@ public:
         loops.pop();
     }
 
-    void compileWhileStatement() {
+    void compileWhileLoop() {
         Loop& loop = enterLoop(false);
         EXPR_TUPLE();
         int patch = emitCode(OP_POP_JUMP_IF_FALSE);
@@ -627,12 +652,11 @@ public:
         exitLoop();
     }
 
-    void compileForStatement() {
+    void __compileForLoop(CompilerAction action) {
         int size = 0;
         do {
             consume(TK("@id"));
-            exprName();     // push a name ptr into stack
-            size++;
+            exprName(); size++;
         } while (match(TK(",")));
         if(size > 1) emitCode(OP_BUILD_SMART_TUPLE, size);
         consume(TK("in"));
@@ -640,7 +664,7 @@ public:
         emitCode(OP_GET_ITER);              // [ptr, list] -> iter
         Loop& loop = enterLoop(true);
         int patch = emitCode(OP_FOR_ITER);
-        compileBlockBody();
+        action(this);
         emitCode(OP_JUMP_ABSOLUTE, loop.start); keepOpcodeLine();
         patchJump(patch);
         exitLoop();
@@ -672,9 +696,9 @@ public:
         } else if (match(TK("if"))) {
             compileIfStatement();
         } else if (match(TK("while"))) {
-            compileWhileStatement();
+            compileWhileLoop();
         } else if (match(TK("for"))) {
-            compileForStatement();
+            __compileForLoop(&Compiler::compileBlockBody);
         } else if(match(TK("assert"))){
             EXPR();
             emitCode(OP_ASSERT);
diff --git a/src/opcodes.h b/src/opcodes.h
index 6c22325f..d6f4f320 100644
--- a/src/opcodes.h
+++ b/src/opcodes.h
@@ -1,5 +1,6 @@
 #ifdef OPCODE
 
+OPCODE(NO_OP)
 OPCODE(LOAD_CONST)
 OPCODE(IMPORT_NAME)
 OPCODE(PRINT_EXPR)
@@ -21,6 +22,8 @@ OPCODE(BUILD_LIST)
 OPCODE(BUILD_MAP)
 OPCODE(BUILD_SLICE)
 
+OPCODE(LIST_APPEND)
+
 OPCODE(GET_ITER)
 OPCODE(FOR_ITER)
 
diff --git a/src/vm.h b/src/vm.h
index a9b9f5e9..00f791a9 100644
--- a/src/vm.h
+++ b/src/vm.h
@@ -162,6 +162,7 @@ public:
 
             switch (byte.op)
             {
+            case OP_NO_OP: break;       // do nothing
             case OP_LOAD_CONST: frame->push(frame->code->co_consts[byte.arg]); break;
             case OP_LOAD_NAME_PTR: {
                 frame->push(PyPointer(frame->code->co_names[byte.arg]));
@@ -220,6 +221,11 @@ public:
             case OP_LOAD_EVAL_FN: {
                 frame->push(builtins->attribs["eval"]);
             } break;
+            case OP_LIST_APPEND: {
+                PyVar obj = frame->popValue(this);
+                PyVar list = frame->topNValue(this, -2);
+                fastCall(list, "append", {list, obj});
+            } break;
             case OP_STORE_FUNCTION:
                 {
                     PyVar obj = frame->popValue(this);