...

src/ceval.h

@@ -85,7 +85,7 @@ __NEXT_STEP:;
         FuncDecl_ decl = co->func_decls[byte.arg];
         PyObject* obj;
         if(decl->nested){
-            obj = VAR(Function({decl, frame->_module, frame->_locals}));
+            obj = VAR(Function({decl, frame->_module, frame->_locals.to_namedict()}));
         }else{
             obj = VAR(Function({decl, frame->_module}));
         }
@@ -151,7 +151,7 @@ __NEXT_STEP:;
     TARGET(STORE_NAME)
         _name = StrName(byte.arg);
         _0 = POPX();
-        if(frame->_locals.is_valid()){
+        if(frame->_callable != nullptr){
             bool ok = frame->_locals.try_set(_name, _0);
             if(!ok) vm->NameError(_name);
         }else{
@@ -180,7 +180,7 @@ __NEXT_STEP:;
         DISPATCH();
     TARGET(DELETE_NAME)
         _name = StrName(byte.arg);
-        if(frame->_locals.is_valid()){
+        if(frame->_callable != nullptr){
             if(!frame->_locals.contains(_name)) vm->NameError(_name);
             frame->_locals.erase(_name);
         }else{
@@ -362,7 +362,7 @@ __NEXT_STEP:;
         DISPATCH();
     TARGET(GOTO) {
         StrName name(byte.arg);
-        int index = co->labels->try_get(name);
+        int index = co->labels.try_get(name);
         if(index < 0) _error("KeyError", fmt("label ", name.escape(), " not found"));
         frame->jump_abs_break(index);
     } DISPATCH();

src/codeobject.h

@@ -53,19 +53,16 @@ struct CodeObject {
     bool is_generator = false;
 
     CodeObject(shared_ptr<SourceData> src, const Str& name):
-        src(src), name(name) {
-            varnames_inv = make_sp<NameDictInt>();
-            labels = make_sp<NameDictInt>();
-        }
+        src(src), name(name) {}
 
     std::vector<Bytecode> codes;
     std::vector<int> lines; // line number for each bytecode
     List consts;
     std::vector<StrName> varnames;      // local variables
-    NameDictInt_ varnames_inv;
+    NameDictInt varnames_inv;
     std::set<Str> global_names;
     std::vector<CodeBlock> blocks = { CodeBlock(NO_BLOCK, -1, 0, 0) };
-    NameDictInt_ labels;
+    NameDictInt labels;
     std::vector<FuncDecl_> func_decls;
 
     void optimize(VM* vm);

src/common.h

@@ -51,6 +51,10 @@
 #define PK_ENABLE_FILEIO 			0	// TODO: refactor this
 #endif
 
+// This is the maximum number of arguments in a function declaration
+// including positional arguments, keyword-only arguments, and varargs
+#define PK_MAX_CO_VARNAMES			255
+
 #if _MSC_VER
 #define PK_ENABLE_COMPUTED_GOTO		0
 #define UNREACHABLE()				__assume(0)

src/compiler.h

@@ -68,6 +68,9 @@ class Compiler {
             ctx()->emit(OP_RETURN_VALUE, BC_NOARG, BC_KEEPLINE);
         }
         ctx()->co->optimize(vm);
+        if(ctx()->co->varnames.size() > PK_MAX_CO_VARNAMES){
+            SyntaxError("maximum number of local variables exceeded");
+        }
         contexts.pop();
     }
 

src/expr.h

@@ -94,17 +94,17 @@ struct CodeEmitContext{
     }
 
     bool add_label(StrName name){
-        if(co->labels->contains(name)) return false;
-        co->labels->set(name, co->codes.size());
+        if(co->labels.contains(name)) return false;
+        co->labels.set(name, co->codes.size());
         return true;
     }
 
     int add_varname(StrName name){
-        int index = co->varnames_inv->try_get(name);
+        int index = co->varnames_inv.try_get(name);
         if(index >= 0) return index;
         co->varnames.push_back(name);
         index = co->varnames.size() - 1;
-        co->varnames_inv->set(name, index);
+        co->varnames_inv.set(name, index);
         return index;
     }
 
@@ -131,7 +131,7 @@ struct NameExpr: Expr{
     std::string str() const override { return fmt("Name(", name.escape(), ")"); }
 
     void emit(CodeEmitContext* ctx) override {
-        int index = ctx->co->varnames_inv->try_get(name);
+        int index = ctx->co->varnames_inv.try_get(name);
         if(scope == NAME_LOCAL && index >= 0){
             ctx->emit(OP_LOAD_FAST, index, line);
         }else{

src/frame.h

@@ -8,8 +8,10 @@
 
 namespace pkpy{
 
+// weak reference fast locals
 struct FastLocals{
-    NameDictInt_ varnames_inv;
+    // this is a weak reference
+    const NameDictInt* varnames_inv;
     PyObject** a;
 
     int size() const{
@@ -19,19 +21,10 @@ struct FastLocals{
     PyObject*& operator[](int i){ return a[i]; }
     PyObject* operator[](int i) const { return a[i]; }
 
-    FastLocals(): varnames_inv(nullptr), a(nullptr) {}
-    FastLocals(std::nullptr_t): varnames_inv(nullptr), a(nullptr) {}
-
-    FastLocals(const CodeObject* co): varnames_inv(co->varnames_inv){
-        size_t size = this->size() * sizeof(void*);
-        int* counter = (int*)pool128.alloc(sizeof(int) + size);
-        *counter = 1;
-        a = (PyObject**)(counter + 1);
-        memset(a, 0, this->size() * sizeof(void*));
-    }
+    FastLocals(const CodeObject* co, PyObject** a): varnames_inv(&co->varnames_inv), a(a) {}
+    FastLocals(const FastLocals& other): varnames_inv(other.varnames_inv), a(other.a) {}
 
     PyObject* try_get(StrName name){
-        if(!is_valid()) return nullptr;
         int index = varnames_inv->try_get(name);
         if(index == -1) return nullptr;
         return a[index];
@@ -42,95 +35,31 @@ struct FastLocals{
     }
 
     void erase(StrName name){
-        if(!is_valid()) return;
         int index = varnames_inv->try_get(name);
         if(index == -1) FATAL_ERROR();
         a[index] = nullptr;
     }
 
-    bool _try_set(StrName name, PyObject* value){
+    bool try_set(StrName name, PyObject* value){
         int index = varnames_inv->try_get(name);
         if(index == -1) return false;
         a[index] = value;
         return true;
     }
 
-    bool try_set(StrName name, PyObject* value){
-        if(!is_valid()) return false;
-        return _try_set(name, value);
-    }
-
-    FastLocals(const FastLocals& other){
-        varnames_inv = other.varnames_inv;
-        a = other.a;
-        _inc_counter();
-    }
-
-    FastLocals(FastLocals&& other) noexcept{
-        varnames_inv = std::move(other.varnames_inv);
-        a = other.a;
-        other.a = nullptr;
-    }
-
-    FastLocals& operator=(const FastLocals& other){
-        _dec_counter();
-        varnames_inv = other.varnames_inv;
-        a = other.a;
-        _inc_counter();
-        return *this;
-    }
-
-    FastLocals& operator=(FastLocals&& other) noexcept{
-        _dec_counter();
-        varnames_inv = std::move(other.varnames_inv);
-        a = other.a;
-        other.a = nullptr;
-        return *this;
-    }
-
-    bool is_valid() const{ return a != nullptr; }
-
-    void _inc_counter(){
-        if(a == nullptr) return;
-        int* counter = (int*)a - 1;
-        (*counter)++;
-    }
-
-    void _dec_counter(){
-        if(a == nullptr) return;
-        int* counter = (int*)a - 1;
-        (*counter)--;
-        if(*counter == 0){
-            pool128.dealloc(counter);
-        }
-    }
-
-    ~FastLocals(){
-        _dec_counter();
-    }
-
-    void _gc_mark() const{
-        if(a == nullptr) return;
-        for(int i=0; i<size(); i++){
-            if(a[i] != nullptr) OBJ_MARK(a[i]);
+    NameDict_ to_namedict(){
+        NameDict_ dict = make_sp<NameDict>();
+        // TODO: optimize this
+        // NameDict.items() is expensive
+        for(auto& kv: varnames_inv->items()){
+            dict->set(kv.first, a[kv.second]);
         }
+        return dict;
     }
 };
 
-struct Function{
-    FuncDecl_ decl;
-    PyObject* _module;
-    FastLocals _closure;
-};
-
-template<> inline void gc_mark<Function>(Function& t){
-    t.decl->_gc_mark();
-    if(t._module != nullptr) OBJ_MARK(t._module);
-    t._closure._gc_mark();
-}
-
 struct ValueStack {
-    static const size_t MAX_SIZE = 8192;
+    static const size_t MAX_SIZE = 32768;
     // We allocate 512 more bytes to keep `_sp` valid when `is_overflow() == true`.
     PyObject* _begin[MAX_SIZE + 512];
     PyObject** _sp;
@@ -174,33 +103,29 @@ struct Frame {
     int _next_ip = 0;
     ValueStack* _s;
     PyObject** _sp_base;
-    const CodeObject* co;
 
+    const CodeObject* co;
     PyObject* _module;
-    FastLocals _locals;
     PyObject* _callable;
+    FastLocals _locals;
 
     NameDict& f_globals() noexcept { return _module->attr(); }
     
     PyObject* f_closure_try_get(StrName name){
         if(_callable == nullptr) return nullptr;
         Function& fn = OBJ_GET(Function, _callable);
-        return fn._closure.try_get(name);
+        if(fn._closure == nullptr) return nullptr;
+        return fn._closure->try_get(name);
     }
 
-    Frame(ValueStack* _s, PyObject** _sp_base, const CodeObject* co, PyObject* _module, FastLocals&& _locals, PyObject* _callable)
-            : _s(_s), _sp_base(_sp_base), co(co), _module(_module), _locals(std::move(_locals)), _callable(_callable) { }
-
-    Frame(ValueStack* _s, PyObject** _sp_base, const CodeObject* co, PyObject* _module, const FastLocals& _locals, PyObject* _callable)
-            : _s(_s), _sp_base(_sp_base), co(co), _module(_module), _locals(_locals), _callable(_callable) { }
+    Frame(ValueStack* _s, PyObject** p0, const CodeObject* co, PyObject* _module, PyObject* _callable)
+            : _s(_s), _sp_base(p0), co(co), _module(_module), _callable(_callable), _locals(co, p0) { }
 
-    Frame(ValueStack* _s, PyObject** _sp_base, const CodeObject_& co, PyObject* _module)
-            : _s(_s), _sp_base(_sp_base), co(co.get()), _module(_module), _locals(), _callable(nullptr) { }
+    Frame(ValueStack* _s, PyObject** p0, const CodeObject* co, PyObject* _module, PyObject* _callable, FastLocals _locals)
+            : _s(_s), _sp_base(p0), co(co), _module(_module), _callable(_callable), _locals(_locals) { }
 
-    Frame(const Frame& other) = delete;
-    Frame& operator=(const Frame& other) = delete;
-    Frame(Frame&& other) noexcept = default;
-    Frame& operator=(Frame&& other) noexcept = default;
+    Frame(ValueStack* _s, PyObject** p0, const CodeObject_& co, PyObject* _module)
+            : _s(_s), _sp_base(p0), co(co.get()), _module(_module), _callable(nullptr), _locals(co.get(), p0) {}
 
     Bytecode next_bytecode() {
         _ip = _next_ip++;
@@ -255,10 +180,7 @@ struct Frame {
     }
 
     void _gc_mark() const {
-        // do return if this frame has been moved
-        // TODO: fix here
         OBJ_MARK(_module);
-        _locals._gc_mark();
         if(_callable != nullptr) OBJ_MARK(_callable);
         co->_gc_mark();
     }

src/gc.h

@@ -141,6 +141,12 @@ template<> inline void gc_mark<BoundMethod>(BoundMethod& t){
     OBJ_MARK(t.method);
 }
 
+template<> inline void gc_mark<Function>(Function& t){
+    t.decl->_gc_mark();
+    if(t._module != nullptr) OBJ_MARK(t._module);
+    if(t._closure != nullptr) gc_mark<NameDict>(*t._closure);
+}
+
 template<> inline void gc_mark<Super>(Super& t){
     OBJ_MARK(t.first);
 }

src/namedict.h

@@ -208,6 +208,5 @@ while(!_items[i].first.empty()) {       \
 using NameDict = NameDictImpl<PyObject*>;
 using NameDict_ = shared_ptr<NameDict>;
 using NameDictInt = NameDictImpl<int>;
-using NameDictInt_ = shared_ptr<NameDictInt>;
 
 } // namespace pkpy

src/obj.h

@@ -38,6 +38,12 @@ struct FuncDecl {
 
 using FuncDecl_ = shared_ptr<FuncDecl>;
 
+struct Function{
+    FuncDecl_ decl;
+    PyObject* _module;
+    NameDict_ _closure;
+};
+
 struct BoundMethod {
     PyObject* obj;
     PyObject* method;

src/pocketpy.h

@@ -99,13 +99,13 @@ inline void init_builtins(VM* _vm) {
     _vm->bind_builtin_func<1>("eval", [](VM* vm, ArgsView args) {
         CodeObject_ code = vm->compile(CAST(Str&, args[0]), "<eval>", EVAL_MODE, true);
         FrameId frame = vm->top_frame();
-        return vm->_exec(code.get(), frame->_module, frame->_locals, nullptr);
+        return vm->_exec(code.get(), frame->_module, nullptr, frame->_locals);
     });
 
     _vm->bind_builtin_func<1>("exec", [](VM* vm, ArgsView args) {
         CodeObject_ code = vm->compile(CAST(Str&, args[0]), "<exec>", EXEC_MODE, true);
         FrameId frame = vm->top_frame();
-        vm->_exec(code.get(), frame->_module, frame->_locals, nullptr);
+        vm->_exec(code.get(), frame->_module, nullptr, frame->_locals);
         return vm->None;
     });
 

src/vm.h

@@ -864,13 +864,15 @@ inline PyObject* VM::vectorcall(int ARGC, int KWARGC, bool op_call){
     return nullptr;
 }
 
-inline PyObject* VM::_py_call(PyObject** sp_base, PyObject* callable, ArgsView args, ArgsView kwargs){
+inline PyObject* VM::_py_call(PyObject** p0, PyObject* callable, ArgsView args, ArgsView kwargs){
     // callable must be a `function` object
     if(callstack.size() >= recursionlimit) RecursionError();
 
     const Function& fn = CAST(Function&, callable);
     const CodeObject* co = fn.decl->code.get();
-    FastLocals locals(co);
+
+    static THREAD_LOCAL PyObject* buffer[PK_MAX_CO_VARNAMES];
+    for(int i=0; i<co->varnames.size(); i++) buffer[i] = nullptr;
 
     int i = 0;
     if(args.size() < fn.decl->args.size()){
@@ -884,20 +886,20 @@ inline PyObject* VM::_py_call(PyObject** sp_base, PyObject* callable, ArgsView a
     }
 
     // prepare args
-    for(int index: fn.decl->args) locals[index] = args[i++];
+    for(int index: fn.decl->args) buffer[index] = args[i++];
     // prepare kwdefaults
-    for(auto& kv: fn.decl->kwargs) locals[kv.key] = kv.value;
+    for(auto& kv: fn.decl->kwargs) buffer[kv.key] = kv.value;
     
     // handle *args
     if(fn.decl->starred_arg != -1){
         List vargs;        // handle *args
         while(i < args.size()) vargs.push_back(args[i++]);
-        locals[fn.decl->starred_arg] = VAR(Tuple(std::move(vargs)));
+        buffer[fn.decl->starred_arg] = VAR(Tuple(std::move(vargs)));
     }else{
         // kwdefaults override
         for(auto& kv: fn.decl->kwargs){
             if(i < args.size()){
-                locals[kv.key] = args[i++];
+                buffer[kv.key] = args[i++];
             }else{
                 break;
             }
@@ -907,17 +909,21 @@ inline PyObject* VM::_py_call(PyObject** sp_base, PyObject* callable, ArgsView a
     
     for(int i=0; i<kwargs.size(); i+=2){
         StrName key = CAST(int, kwargs[i]);
-        bool ok = locals._try_set(key, kwargs[i+1]);
-        if(!ok) TypeError(fmt(key.escape(), " is an invalid keyword argument for ", co->name, "()"));
+        int index = co->varnames_inv.try_get(key);
+        if(index<0) TypeError(fmt(key.escape(), " is an invalid keyword argument for ", co->name, "()"));
+        buffer[index] = kwargs[i+1];
     }
     PyObject* _module = fn._module != nullptr ? fn._module : top_frame()->_module;
     
-    s_data.reset(sp_base);
+    s_data.reset(p0);
+    // copy buffer to stack
+    for(int i=0; i<co->varnames.size(); i++) PUSH(buffer[i]);
+
     if(co->is_generator){
-        PyObject* ret = PyIter(Generator(this, Frame(&s_data, s_data._sp, co, _module, std::move(locals), callable)));
+        PyObject* ret = PyIter(Generator(this, Frame(&s_data, p0, co, _module, callable)));
         return ret;
     }
-    callstack.emplace(&s_data, s_data._sp, co, _module, std::move(locals), callable);
+    callstack.emplace(&s_data, p0, co, _module, callable);
     return nullptr;
 }