From 7deb36a6de279e5468a2ceaa54256f5075b8c4ee Mon Sep 17 00:00:00 2001
From: aps <62445385+apsz3@users.noreply.github.com>
Date: Fri, 17 Feb 2023 16:45:38 -0500
Subject: [PATCH] F

---
 src/compiler.h | 52 ++++++++++++++++++++++++++++++++++++++++++++++----
 src/vm.h       | 24 ++++++++++++++---------
 2 files changed, 63 insertions(+), 13 deletions(-)

diff --git a/src/compiler.h b/src/compiler.h
index 1d6e7ecf..6aef19b8 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -176,7 +176,7 @@ private:
             }
         }catch(std::exception& _){
             SyntaxError("invalid number literal");
-        } 
+        }
     }
 
     void lex_token(){
@@ -283,7 +283,7 @@ private:
                         eat_number();
                         return;
                     }
-                    
+
                     switch (parser->eat_name())
                     {
                         case 0: break;
@@ -532,7 +532,7 @@ __LISTCOMP:
         emit(OP_BUILD_LIST, 0);
         EXPR_FOR_VARS();consume(TK("in"));EXPR_TUPLE();
         match_newlines(mode()==REPL_MODE);
-        
+
         int _skipPatch = emit(OP_JUMP_ABSOLUTE);
         int _cond_start = co()->codes.size();
         int _cond_end_return = -1;
@@ -567,6 +567,8 @@ __LISTCOMP:
     }
 
     void exprMap() {
+        int _patch = emit(OP_NO_OP);\
+        int _body_start = co()->codes.size();
         bool parsing_dict = false;
         int size = 0;
         do {
@@ -577,7 +579,7 @@ __LISTCOMP:
             if(parsing_dict){
                 consume(TK(":"));
                 EXPR();
-            }
+            } else {if (match(TK("for"))) goto __SETCOMP;}
             size++;
             match_newlines(mode()==REPL_MODE);
         } while (match(TK(",")));
@@ -585,6 +587,48 @@ __LISTCOMP:
 
         if(size == 0 || parsing_dict) emit(OP_BUILD_MAP, size);
         else emit(OP_BUILD_SET, size);
+        return;
+__SETCOMP:
+        int _body_end_return = emit(OP_JUMP_ABSOLUTE, -1);
+        int _body_end = co()->codes.size();
+        co()->codes[_patch].op = OP_JUMP_ABSOLUTE;
+        co()->codes[_patch].arg = _body_end;
+        emit(OP_BUILD_LIST, 0);
+        EXPR_FOR_VARS();consume(TK("in"));EXPR_TUPLE();
+        match_newlines(mode()==REPL_MODE);
+
+        int _skipPatch = emit(OP_JUMP_ABSOLUTE);
+        int _cond_start = co()->codes.size();
+        int _cond_end_return = -1;
+        if(match(TK("if"))) {
+            EXPR_TUPLE();
+            _cond_end_return = emit(OP_JUMP_ABSOLUTE, -1);
+        }
+        patch_jump(_skipPatch);
+
+        emit(OP_GET_ITER);
+        co()->_enter_block(FOR_LOOP);
+        emit(OP_FOR_ITER);
+
+        if(_cond_end_return != -1) {      // there is an if condition
+            emit(OP_JUMP_ABSOLUTE, _cond_start);
+            patch_jump(_cond_end_return);
+            int ifpatch = emit(OP_POP_JUMP_IF_FALSE);
+            emit(OP_JUMP_ABSOLUTE, _body_start);
+            patch_jump(_body_end_return);
+            emit(OP_LIST_APPEND);
+            patch_jump(ifpatch);
+        }else{
+            emit(OP_JUMP_ABSOLUTE, _body_start);
+            patch_jump(_body_end_return);
+            emit(OP_LIST_APPEND);
+        }
+
+        emit(OP_LOOP_CONTINUE, -1, true);
+        co()->_exit_block();
+        match_newlines(mode()==REPL_MODE);
+        consume(TK("}"));
+        emit(OP_BUILD_SET, -1);
     }
 
     void exprCall() {
diff --git a/src/vm.h b/src/vm.h
index 7428f12c..dc8f8747 100644
--- a/src/vm.h
+++ b/src/vm.h
@@ -217,9 +217,15 @@ public:
                 } break;
             case OP_BUILD_SET:
                 {
-                    PyVar list = PyList(
+                    PyVar list;
+                    if (byte.arg == -1) {
+                        list = frame->pop_value(this);
+                    }
+                    else {
+                        list = PyList(
                         frame->pop_n_values_reversed(this, byte.arg).to_list()
                     );
+                    }
                     PyVar obj = call(builtins->attr("set"), pkpy::one_arg(list));
                     frame->push(obj);
                 } break;
@@ -354,7 +360,7 @@ public:
     bool use_stdio;
     std::ostream* _stdout;
     std::ostream* _stderr;
-    
+
     PyVar builtins;         // builtins module
     PyVar _main;            // __main__ module
 
@@ -454,7 +460,7 @@ public:
             callable = &bm.method;      // get unbound method
             args.extend_self(bm.obj);
         }
-        
+
         if((*callable)->is_type(tp_native_function)){
             const auto& f = OBJ_GET(pkpy::NativeFunc, *callable);
             if(kwargs.size() != 0) TypeError("native_function does not accept keyword arguments");
@@ -491,7 +497,7 @@ public:
                 }
                 if(i < args.size()) TypeError("too many arguments");
             }
-            
+
             for(int i=0; i<kwargs.size(); i+=2){
                 const Str& key = PyStr_AS_C(kwargs[i]);
                 if(!fn->kwArgs.contains(key)){
@@ -666,7 +672,7 @@ public:
             for(int i=0; i<depth; i++) cls = cls->attr(__base__).get();
 
             it = (*root)->attr().find(name);
-            if(it != (*root)->attr().end()) return it->second;        
+            if(it != (*root)->attr().end()) return it->second;
         }else{
             if(obj->is_attr_valid()){
                 it = obj->attr().find(name);
@@ -712,7 +718,7 @@ public:
 
     template<int ARGC>
     void bind_func(Str typeName, Str funcName, NativeFuncRaw fn) {
-        bind_func<ARGC>(_types[typeName], funcName, fn);     
+        bind_func<ARGC>(_types[typeName], funcName, fn);
     }
 
     template<int ARGC>
@@ -869,7 +875,7 @@ public:
     DEF_NATIVE(Range, pkpy::Range, tp_range)
     DEF_NATIVE(Slice, pkpy::Slice, tp_slice)
     DEF_NATIVE(Exception, pkpy::Exception, tp_exception)
-    
+
     // there is only one True/False, so no need to copy them!
     inline bool PyBool_AS_C(const PyVar& obj){return obj == True;}
     inline const PyVar& PyBool(bool value){return value ? True : False;}
@@ -894,7 +900,7 @@ public:
         tp_range = _new_type_object("range");
         tp_module = _new_type_object("module");
         tp_ref = _new_type_object("_ref");
-        
+
         tp_function = _new_type_object("function");
         tp_native_function = _new_type_object("native_function");
         tp_native_iterator = _new_type_object("native_iterator");
@@ -913,7 +919,7 @@ public:
 
         setattr(_t(tp_type), __base__, _t(tp_object));
         setattr(_t(tp_object), __base__, None);
-        
+
         for (auto& [name, type] : _types) {
             setattr(type, __name__, PyStr(name));
         }