support stack object

include/pocketpy/common.h

@@ -233,12 +233,12 @@ struct PyVar final{
     bool operator!=(std::nullptr_t) const { return (bool)type; }
 
     PyObject* get() const {
-        PK_DEBUG_ASSERT(!is_sso)
+        PK_DEBUG_ASSERT(is_ptr)
         return reinterpret_cast<PyObject*>(_1);
     }
 
     PyObject* operator->() const {
-        PK_DEBUG_ASSERT(!is_sso)
+        PK_DEBUG_ASSERT(is_ptr)
         return reinterpret_cast<PyObject*>(_1);
     }
 

include/pocketpy/iter.h

@@ -51,9 +51,9 @@ struct Generator{
         for(PyVar obj: buffer) s_backup.push_back(obj);
     }
 
-    void _gc_mark(VM* vm) const{
+    void _gc_mark(VM* vm) {
         frame._gc_mark(vm);
-        for(PyVar obj: s_backup) PK_OBJ_MARK(obj);
+        vm->__stack_gc_mark(s_backup.begin(), s_backup.end());
     }
 
     PyVar next(VM* vm);

include/pocketpy/obj.h

@@ -50,6 +50,14 @@ struct Range {
     i64 step = 1;
 };
 
+struct StackMemory{
+    int count;
+    StackMemory(int count) : count(count) {}
+};
+
+template<>
+inline bool constexpr is_sso_v<StackMemory> = true;
+
 struct StarWrapper{
     int level;      // either 1 or 2
     PyVar obj;
@@ -155,8 +163,9 @@ static_assert(sizeof(PyObject) <= PyObject::FIXED_SIZE);
 template<typename T>
 inline constexpr int py_sizeof = PyObject::FIXED_SIZE + sizeof(T);
 
-const int kTpIntIndex = 3;
-const int kTpFloatIndex = 4;
+inline const int kTpIntIndex = 3;
+inline const int kTpFloatIndex = 4;
+inline const int kTpStackMemoryIndex = 27;
 
 inline bool is_tagged(PyVar p) noexcept { return !p.is_ptr; }
 inline bool is_float(PyVar p) noexcept { return p.type.index == kTpFloatIndex; }
@@ -187,14 +196,14 @@ obj_get_t<T> PyVar::obj_get(){
     if constexpr(is_sso_v<T>){
         return as<T>();
     }else{
-        PK_DEBUG_ASSERT(!is_sso)
+        PK_DEBUG_ASSERT(is_ptr)
         void* v = ((PyObject*)_1)->_value_ptr();
         return *reinterpret_cast<T*>(v);
     }
 }
 
 #define PK_OBJ_GET(T, obj) (obj).obj_get<T>()
-#define PK_OBJ_MARK(obj) if((obj).is_ptr) vm->__obj_gc_mark(obj.get());
+#define PK_OBJ_MARK(obj) if((obj).is_ptr) vm->__obj_gc_mark((obj).get());
 
 #define VAR(x) py_var(vm, x)
 #define CAST(T, x) py_cast<T>(vm, x)

include/pocketpy/opcodes.h

@@ -117,6 +117,7 @@ OPCODE(UNARY_STAR)
 OPCODE(UNARY_INVERT)
 /**************************/
 OPCODE(GET_ITER)
+OPCODE(GET_ITER_NEW)
 OPCODE(FOR_ITER)
 OPCODE(FOR_ITER_STORE_FAST)
 OPCODE(FOR_ITER_STORE_GLOBAL)

include/pocketpy/vm.h

@@ -82,7 +82,8 @@ struct PyTypeInfo{
     i64 (*m__hash__)(VM* vm, PyVar) = nullptr;
     i64 (*m__len__)(VM* vm, PyVar) = nullptr;
     PyVar (*m__iter__)(VM* vm, PyVar) = nullptr;
-    unsigned (*m__next__)(VM* vm, PyVar) = nullptr;
+    void (*op__iter__)(VM* vm, PyVar) = nullptr;
+    unsigned (*op__next__)(VM* vm, PyVar) = nullptr;
     PyVar (*m__neg__)(VM* vm, PyVar) = nullptr;
     PyVar (*m__invert__)(VM* vm, PyVar) = nullptr;
 
@@ -207,11 +208,12 @@ public:
     static constexpr Type tp_super=Type(15), tp_exception=Type(16), tp_bytes=Type(17), tp_mappingproxy=Type(18);
     static constexpr Type tp_dict=Type(19), tp_property=Type(20), tp_star_wrapper=Type(21);
     static constexpr Type tp_staticmethod=Type(22), tp_classmethod=Type(23);
-    static constexpr Type tp_none=Type(24), tp_not_implemented=Type(25), tp_ellipsis=Type(26);
+    static constexpr Type tp_none_type=Type(24), tp_not_implemented=Type(25), tp_ellipsis=Type(26);
+    static constexpr Type tp_stack_memory=Type(kTpStackMemoryIndex);
 
     static constexpr PyVar True{const_sso_var(), tp_bool, 1};
     static constexpr PyVar False{const_sso_var(), tp_bool, 0};
-    static constexpr PyVar None{const_sso_var(), tp_none, 0};
+    static constexpr PyVar None{const_sso_var(), tp_none_type, 0};
     static constexpr PyVar NotImplemented{const_sso_var(), tp_not_implemented, 0};
     static constexpr PyVar Ellipsis{const_sso_var(), tp_ellipsis, 0};
 
@@ -440,6 +442,14 @@ public:
         if constexpr(is_sso_v<T>) return PyVar(type, T(std::forward<Args>(args)...));
         else return heap.gcnew<T>(type, std::forward<Args>(args)...);
     }
+
+    template<typename T, typename ...Args>
+    void new_stack_object(Type type, Args&&... args){
+        static_assert(std::is_same_v<T, std::decay_t<T>>);
+        PyObject* p = new(__stack_alloc(py_sizeof<T>)) PyObject(type);
+        p->placement_new<T>(std::forward<Args>(args)...);
+        vm->s_data.emplace(p->type, p);
+    }
 #endif
 
     template<typename T>
@@ -488,6 +498,8 @@ public:
     PyVar __minmax_reduce(bool (VM::*op)(PyVar, PyVar), PyVar args, PyVar key);
     bool __py_bool_non_trivial(PyVar);
     void __obj_gc_mark(PyObject*);
+    void __stack_gc_mark(PyVar* begin, PyVar* end);
+    void* __stack_alloc(int size);
 };
 
 
@@ -514,6 +526,7 @@ template<> constexpr Type _find_type_in_const_cxx_typeid_map<Property>(){ return
 template<> constexpr Type _find_type_in_const_cxx_typeid_map<StarWrapper>(){ return VM::tp_star_wrapper; }
 template<> constexpr Type _find_type_in_const_cxx_typeid_map<StaticMethod>(){ return VM::tp_staticmethod; }
 template<> constexpr Type _find_type_in_const_cxx_typeid_map<ClassMethod>(){ return VM::tp_classmethod; }
+template<> constexpr Type _find_type_in_const_cxx_typeid_map<StackMemory>(){ return VM::tp_stack_memory; }
 
 template<typename __T>
 PyVar py_var(VM* vm, __T&& value){

src/array2d.cpp

@@ -366,8 +366,10 @@ struct Array2dIter{
     PK_ALWAYS_PASS_BY_POINTER(Array2dIter)
 
     PyVar ref;
+    Array2d* a;
     int i;
-    Array2dIter(PyVar ref) : ref(ref), i(0) {}
+    
+    Array2dIter(PyVar ref, Array2d* a): ref(ref), a(a), i(0){}
 
     void _gc_mark(VM* vm) const{ PK_OBJ_MARK(ref); }
 
@@ -375,12 +377,11 @@ struct Array2dIter{
         vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyVar _0) { return _0; });
         vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyVar _0) -> unsigned{
             Array2dIter& self = PK_OBJ_GET(Array2dIter, _0);
-            Array2d& a = PK_OBJ_GET(Array2d, self.ref);
-            if(self.i == a.numel) return 0;
-            std::div_t res = std::div(self.i, a.n_cols);
+            if(self.i == self.a->numel) return 0;
+            std::div_t res = std::div(self.i, self.a->n_cols);
             vm->s_data.emplace(VM::tp_int, res.rem);
             vm->s_data.emplace(VM::tp_int, res.quot);
-            vm->s_data.push(a.data[self.i++]);
+            vm->s_data.push(self.a->data[self.i++]);
             return 3;
         });
     }
@@ -392,9 +393,13 @@ void add_module_array2d(VM* vm){
     vm->register_user_class<Array2d>(mod, "array2d", VM::tp_object, true);
     vm->register_user_class<Array2dIter>(mod, "_array2d_iter");
 
-    vm->bind__iter__(vm->_tp_user<Array2d>(), [](VM* vm, PyVar _0){
-        return vm->new_user_object<Array2dIter>(_0);
+    Type array2d_iter_t = vm->_tp_user<Array2d>();
+    vm->bind__iter__(array2d_iter_t, [](VM* vm, PyVar _0){
+        return vm->new_user_object<Array2dIter>(_0, &_0.obj_get<Array2d>());
     });
+    vm->_all_types[array2d_iter_t].op__iter__ = [](VM* vm, PyVar _0){
+        vm->new_stack_object<Array2dIter>(vm->_tp_user<Array2dIter>(), _0, &_0.obj_get<Array2d>());
+    };
 }
 
 

src/ceval.cpp

@@ -806,6 +806,17 @@ __NEXT_STEP:
     case OP_GET_ITER:
         TOP() = py_iter(TOP());
         DISPATCH()
+    case OP_GET_ITER_NEW: {
+        // This opcode always creates a temporary iterator object
+        const PyTypeInfo* _ti = _tp_info(TOP());
+        if(_ti->op__iter__){
+            PyVar _0 = POPX();
+            _ti->op__iter__(this, _0);
+        }else{
+            TOP() = py_iter(TOP());
+        }
+        DISPATCH()
+    }
     case OP_FOR_ITER:{
         PyVar _0 = py_next(TOP());
         if(_0 == StopIteration){
@@ -849,8 +860,8 @@ __NEXT_STEP:
     case OP_FOR_ITER_UNPACK:{
         PyVar _0 = TOP();
         const PyTypeInfo* _ti = _tp_info(_0);
-        if(_ti->m__next__){
-            unsigned n = _ti->m__next__(this, _0);
+        if(_ti->op__next__){
+            unsigned n = _ti->op__next__(this, _0);
             if(n == 0){
                 // StopIteration
                 int target = frame->prepare_loop_break(&s_data);

src/compiler.cpp

@@ -685,7 +685,7 @@ __EAT_DOTS_END:
         Expr_ vars = EXPR_VARS();
         consume(TK("in"));
         EXPR_TUPLE(); ctx()->emit_expr();
-        ctx()->emit_(OP_GET_ITER, BC_NOARG, BC_KEEPLINE);
+        ctx()->emit_(OP_GET_ITER_NEW, BC_NOARG, BC_KEEPLINE);
         CodeBlock* block = ctx()->enter_block(CodeBlockType::FOR_LOOP);
         int for_codei = ctx()->emit_(OP_FOR_ITER, ctx()->curr_iblock, BC_KEEPLINE);
         bool ok = vars->emit_store(ctx());
@@ -852,7 +852,7 @@ __EAT_DOTS_END:
                 if (contexts.size() <= 1) SyntaxError("'yield from' outside function");
                 EXPR_TUPLE(); ctx()->emit_expr();
 
-                ctx()->emit_(OP_GET_ITER, BC_NOARG, kw_line);
+                ctx()->emit_(OP_GET_ITER_NEW, BC_NOARG, kw_line);
                 ctx()->enter_block(CodeBlockType::FOR_LOOP);
                 ctx()->emit_(OP_FOR_ITER_YIELD_VALUE, BC_NOARG, kw_line);
                 ctx()->emit_(OP_LOOP_CONTINUE, ctx()->get_loop(), kw_line);

src/frame.cpp

@@ -42,7 +42,16 @@ namespace pkpy{
 
     int Frame::_exit_block(ValueStack* _s, int i){
         auto type = co->blocks[i].type;
-        if(type==CodeBlockType::FOR_LOOP || type==CodeBlockType::CONTEXT_MANAGER) _s->pop();
+        if(type==CodeBlockType::FOR_LOOP){
+            _s->pop();  // pop the iterator
+            // pop possible stack memory slots
+            if(_s->top().type == kTpStackMemoryIndex){
+                int count = _s->top().as<StackMemory>().count;
+                _s->_sp -= (count + 2);
+            }
+        }else if(type==CodeBlockType::CONTEXT_MANAGER){
+            _s->pop();
+        }
         return co->blocks[i].parent;
     }
 

src/pocketpy.cpp

@@ -997,6 +997,11 @@ void __init_builtins(VM* _vm) {
         List& self = _CAST(List&, _0);
         return vm->new_user_object<ArrayIter>(_0.get(), self.begin(), self.end());
     });
+    _vm->_all_types[VM::tp_list].op__iter__ = [](VM* vm, PyVar _0){
+        List& self = _CAST(List&, _0);
+        vm->new_stack_object<ArrayIter>(vm->_tp_user<ArrayIter>(), _0.get(), self.begin(), self.end());
+    };
+
     _vm->bind__getitem__(VM::tp_list, PyArrayGetItem<List>);
     _vm->bind__setitem__(VM::tp_list, [](VM* vm, PyVar _0, PyVar _1, PyVar _2){
         List& self = _CAST(List&, _0);

src/vm.cpp

@@ -267,8 +267,8 @@ namespace pkpy{
     }
 
     PyVar VM::_py_next(const PyTypeInfo* ti, PyVar obj){
-        if(ti->m__next__){
-            unsigned n = ti->m__next__(this, obj);
+        if(ti->op__next__){
+            unsigned n = ti->op__next__(this, obj);
             return __pack_next_retval(n);
         }
         return call_method(obj, __next__);
@@ -448,6 +448,29 @@ void VM::__obj_gc_mark(PyObject* obj){
     }
 }
 
+void VM::__stack_gc_mark(PyVar* begin, PyVar* end){
+    for(PyVar* it=begin; it!=end; it++){
+        if(it->is_ptr){
+            __obj_gc_mark(it->get());
+        }else{
+            if(it->type == tp_stack_memory){
+                // [sm:3, _0, _1, _2, sm:-3]
+                int count = it->as<StackMemory>().count;
+                if(count > 0) it += count;
+            }
+        }
+    }
+}
+
+void* VM::__stack_alloc(int size){
+    int count = size / sizeof(PyVar) + 1;
+    s_data.emplace(tp_stack_memory, StackMemory(count));
+    void* out = s_data._sp;
+    s_data._sp += count;
+    s_data.emplace(tp_stack_memory, StackMemory(-count));
+    return out;
+}
+
 List VM::py_list(PyVar it){
     auto _lock = heap.gc_scope_lock();
     it = py_iter(it);
@@ -804,28 +827,37 @@ void VM::__log_s_data(const char* title) {
     for(PyVar* p=s_data.begin(); p!=s_data.end(); p++){
         ss << std::string(sp_bases[p], '|');
         if(sp_bases[p] > 0) ss << " ";
-        PyVar obj = *p;
-        if(obj == nullptr) ss << "(nil)";
-        else if(obj == PY_NULL) ss << "NULL";
-        else if(is_int(obj)) ss << CAST(i64, obj);
-        else if(is_float(obj)) ss << CAST(f64, obj);
-        else if(is_type(obj, tp_str)) ss << CAST(Str, obj).escape();
-        else if(obj == None) ss << "None";
-        else if(obj == True) ss << "True";
-        else if(obj == False) ss << "False";
-        else if(is_type(obj, tp_function)){
-            auto& f = CAST(Function&, obj);
-            ss << f.decl->code->name << "(...)";
-        } else if(is_type(obj, tp_type)){
-            Type t = PK_OBJ_GET(Type, obj);
-            ss << "<class " + _all_types[t].name.escape() + ">";
-        } else if(is_type(obj, tp_list)){
-            auto& t = CAST(List&, obj);
-            ss << "list(size=" << t.size() << ")";
-        } else if(is_type(obj, tp_tuple)){
-            auto& t = CAST(Tuple&, obj);
-            ss << "tuple(size=" << t.size() << ")";
-        } else ss << "(" << _type_name(this, obj.type) << ")";
+        if(*p == PY_NULL) ss << "NULL";
+        else{
+            switch(p->type){
+                case tp_none_type: ss << "None"; break;
+                case tp_int: ss << _CAST(i64, *p); break;
+                case tp_float: ss << _CAST(f64, *p); break;
+                case tp_bool: ss << ((*p == True) ? "True" : "False"); break;
+                case tp_str: ss << _CAST(Str, *p).escape(); break;
+                case tp_function:
+                    ss << p->obj_get<Function>().decl->code->name << "()";
+                    break;
+                case tp_type:
+                    ss << "<class " + _type_name(this, p->obj_get<Type>()).escape() + ">";
+                    break;
+                case tp_list: 
+                    ss << "list(size=" << p->obj_get<List>().size() << ")";
+                    break;
+                case tp_tuple:
+                    ss << "tuple(size=" << p->obj_get<Tuple>().size() << ")";
+                    break;
+                case tp_stack_memory: {
+                    int count = p->obj_get<StackMemory>().count;
+                    ss << "M[" << count << "]";
+                    if(count > 0) p += count;
+                    break;
+                }
+                default:
+                    ss << "(" << _type_name(this, p->type) << ")";
+                    break;
+            }
+        }
         ss << ", ";
     }
     std::string output = ss.str().str();
@@ -874,9 +906,10 @@ void VM::__init_builtin_types(){
     validate(tp_staticmethod, new_type_object<StaticMethod>(nullptr, "staticmethod", tp_object, false));
     validate(tp_classmethod, new_type_object<ClassMethod>(nullptr, "classmethod", tp_object, false));
 
-    validate(tp_none, new_type_object(nullptr, "NoneType", tp_object, false));
+    validate(tp_none_type, new_type_object(nullptr, "NoneType", tp_object, false));
     validate(tp_not_implemented, new_type_object(nullptr, "NotImplementedType", tp_object, false));
     validate(tp_ellipsis, new_type_object(nullptr, "ellipsis", tp_object, false));
+    validate(tp_stack_memory, new_type_object<StackMemory>(nullptr, "stack_memory", tp_object, false));
 
     // SyntaxError and IndentationError must be created here
     PyVar SyntaxError = new_type_object(nullptr, "SyntaxError", tp_exception, true);
@@ -1468,7 +1501,7 @@ PyVar VM::__pack_next_retval(unsigned n){
 }
 
 void VM::bind__next__(Type type, unsigned (*f)(VM*, PyVar)){
-    _all_types[type].m__next__ = f;
+    _all_types[type].op__next__ = f;
     bind_func(type, __next__, 1, [](VM* vm, ArgsView args){
         int n = lambda_get_userdata<unsigned(*)(VM*, PyVar)>(args.begin())(vm, args[0]);
         return vm->__pack_next_retval(n);
@@ -1831,11 +1864,11 @@ void Frame::_gc_mark(VM* vm) const {
 void ManagedHeap::mark() {
     for(PyObject* obj: _no_gc) vm->__obj_gc_mark(obj);
     vm->callstack.apply([this](Frame& frame){ frame._gc_mark(vm); });
-    for(PyVar obj: vm->s_data) PK_OBJ_MARK(obj);
     for(auto [_, co]: vm->__cached_codes) co->_gc_mark(vm);
     if(vm->__last_exception) PK_OBJ_MARK(vm->__last_exception);
     if(vm->__curr_class) PK_OBJ_MARK(vm->__curr_class);
     if(vm->__c.error != nullptr) PK_OBJ_MARK(vm->__c.error);
+    vm->__stack_gc_mark(vm->s_data.begin(), vm->s_data.end());
     if(_gc_marker_ex) _gc_marker_ex(vm);
 }