fix function types

include/pocketpy/codeobject.h

@@ -65,7 +65,6 @@ struct CodeObject {
 
     std::shared_ptr<SourceData> src;
     Str name;
-    bool is_generator;
 
     std::vector<Bytecode> codes;
     std::vector<int> iblocks;       // block index for each bytecode
@@ -90,6 +89,14 @@ struct CodeObject {
     void _gc_mark() const;
 };
 
+enum class FuncType{
+    UNSET,
+    NORMAL,
+    SIMPLE,
+    EMPTY,
+    GENERATOR,
+};
+
 struct FuncDecl {
     struct KwArg {
         int index;              // index in co->varnames
@@ -108,8 +115,7 @@ struct FuncDecl {
     Str signature;              // signature of this function
     Str docstring;              // docstring of this function
 
-    bool is_simple;             // whether this function is simple (no *arg, **kwarg, nested)
-    bool is_empty;              // whether this function is empty (no code)
+    FuncType type = FuncType::UNSET;
 
     NameDictInt kw_to_index;
 

src/codeobject.cpp

@@ -3,7 +3,7 @@
 namespace pkpy{
 
     CodeObject::CodeObject(std::shared_ptr<SourceData> src, const Str& name):
-        src(src), name(name), is_generator(false), start_line(-1), end_line(-1) {
+        src(src), name(name), start_line(-1), end_line(-1) {
             blocks.push_back(CodeBlock(CodeBlockType::NO_BLOCK, -1, 0, 0));
         }
 

src/compiler.cpp

@@ -62,27 +62,49 @@ namespace pkpy{
         // pre-compute func->is_simple
         FuncDecl_ func = contexts.top().func;
         if(func){
-            func->is_simple = true;
-            if(func->code->is_generator) func->is_simple = false;
-            if(func->kwargs.size() > 0) func->is_simple = false;
-            if(func->starred_arg >= 0) func->is_simple = false;
-            if(func->starred_kwarg >= 0) func->is_simple = false;
-
-            func->is_empty = false;
-            if(func->code->codes.size() == 1){
-                Bytecode bc = func->code->codes[0];
-                if(bc.op == OP_RETURN_VALUE && bc.arg == 1){
-                    func->is_empty = true;
+            // check generator
+            for(Bytecode bc: func->code->codes){
+                if(bc.op == OP_YIELD_VALUE || bc.op == OP_FOR_ITER_YIELD_VALUE){
+                    func->type = FuncType::GENERATOR;
+                    for(Bytecode bc: func->code->codes){
+                        if(bc.op == OP_RETURN_VALUE && bc.arg == BC_NOARG){
+                            SyntaxError("'return' with argument inside generator function");
+                        }
+                    }
+                    break;
+                }
+            }
+            if(func->type == FuncType::UNSET){
+                bool is_simple = true;
+                if(func->kwargs.size() > 0) is_simple = false;
+                if(func->starred_arg >= 0) is_simple = false;
+                if(func->starred_kwarg >= 0) is_simple = false;
+
+                if(is_simple){
+                    func->type = FuncType::SIMPLE;
+
+                    bool is_empty = false;
+                    if(func->code->codes.size() == 1){
+                        Bytecode bc = func->code->codes[0];
+                        if(bc.op == OP_RETURN_VALUE && bc.arg == 1){
+                            is_empty = true;
+                        }
+                    }
+                    if(is_empty) func->type = FuncType::EMPTY;
                 }
+                else func->type = FuncType::NORMAL;
             }
+
+            PK_ASSERT(func->type != FuncType::UNSET);
         }
         contexts.pop();
     }
 
     void Compiler::init_pratt_rules(){
-        PK_LOCAL_STATIC unsigned int count = 0;
-        if(count > 0) return;
-        count += 1;
+        PK_LOCAL_STATIC bool initialized = false;
+        if(initialized) return;
+        initialized = true;
+
 // http://journal.stuffwithstuff.com/2011/03/19/pratt-parsers-expression-parsing-made-easy/
 #define PK_METHOD(name) &Compiler::name
 #define PK_NO_INFIX nullptr, PREC_LOWEST
@@ -817,16 +839,13 @@ __EAT_DOTS_END:
             case TK("yield"): 
                 if (contexts.size() <= 1) SyntaxError("'yield' outside function");
                 EXPR_TUPLE(); ctx()->emit_expr();
-                // if yield present, mark the function as generator
-                ctx()->co->is_generator = true;
                 ctx()->emit_(OP_YIELD_VALUE, BC_NOARG, kw_line);
                 consume_end_stmt();
                 break;
             case TK("yield from"):
                 if (contexts.size() <= 1) SyntaxError("'yield from' outside function");
                 EXPR_TUPLE(); ctx()->emit_expr();
-                // if yield from present, mark the function as generator
-                ctx()->co->is_generator = true;
+
                 ctx()->emit_(OP_GET_ITER, BC_NOARG, kw_line);
                 ctx()->enter_block(CodeBlockType::FOR_LOOP);
                 ctx()->emit_(OP_FOR_ITER_YIELD_VALUE, BC_NOARG, kw_line);
@@ -840,8 +859,6 @@ __EAT_DOTS_END:
                     ctx()->emit_(OP_RETURN_VALUE, 1, kw_line);
                 }else{
                     EXPR_TUPLE(); ctx()->emit_expr();
-                    // check if it is a generator
-                    if(ctx()->co->is_generator) SyntaxError("'return' with argument inside generator function");
                     consume_end_stmt();
                     ctx()->emit_(OP_RETURN_VALUE, BC_NOARG, kw_line);
                 }

src/vm.cpp

@@ -891,43 +891,38 @@ PyObject* VM::vectorcall(int ARGC, int KWARGC, bool op_call){
         const CodeObject* co = fn.decl->code.get();
         int co_nlocals = co->varnames.size();
 
-        if(fn.decl->is_simple){
-            if(args.size() != fn.decl->args.size()){
-                TypeError(_S(
-                    co->name, "() takes ", fn.decl->args.size(), " positional arguments but ", args.size(), " were given"
-                ));
-            }
-            if(!kwargs.empty()) TypeError(_S(co->name, "() takes no keyword arguments"));
-
-            // fast path for empty function
-            if(fn.decl->is_empty){
+        switch(fn.decl->type){
+            case FuncType::UNSET: PK_FATAL_ERROR(); break;
+            case FuncType::NORMAL:
+                _prepare_py_call(buffer, args, kwargs, fn.decl);
+                // copy buffer back to stack
+                s_data.reset(_base + co_nlocals);
+                for(int j=0; j<co_nlocals; j++) _base[j] = buffer[j];
+                break;
+            case FuncType::SIMPLE:
+                if(args.size() != fn.decl->args.size()) TypeError(_S(co->name, "() takes ", fn.decl->args.size(), " positional arguments but ", args.size(), " were given"));
+                if(!kwargs.empty()) TypeError(_S(co->name, "() takes no keyword arguments"));
+                // [callable, <self>, args..., local_vars...]
+                //      ^p0                    ^p1      ^_sp
+                s_data.reset(_base + co_nlocals);
+                // initialize local variables to PY_NULL
+                for(PyObject** p=p1; p!=s_data._sp; p++) *p = PY_NULL;
+                break;
+            case FuncType::EMPTY:
+                if(args.size() != fn.decl->args.size()) TypeError(_S(co->name, "() takes ", fn.decl->args.size(), " positional arguments but ", args.size(), " were given"));
+                if(!kwargs.empty()) TypeError(_S(co->name, "() takes no keyword arguments"));
                 s_data.reset(p0);
                 return None;
-            }
-
-            // [callable, <self>, args..., local_vars...]
-            //      ^p0                    ^p1      ^_sp
-            s_data.reset(_base + co_nlocals);
-            // initialize local variables to PY_NULL
-            for(PyObject** p=p1; p!=s_data._sp; p++) *p = PY_NULL;
-            goto __FAST_CALL;
-        }
-
-        _prepare_py_call(buffer, args, kwargs, fn.decl);
-        
-        if(co->is_generator){
-            s_data.reset(p0);
-            return _py_generator(
-                Frame(nullptr, co, fn._module, callable, nullptr),
-                ArgsView(buffer, buffer + co_nlocals)
-            );
-        }
-
-        // copy buffer back to stack
-        s_data.reset(_base + co_nlocals);
-        for(int j=0; j<co_nlocals; j++) _base[j] = buffer[j];
+            case FuncType::GENERATOR:
+                _prepare_py_call(buffer, args, kwargs, fn.decl);
+                s_data.reset(p0);
+                return _py_generator(
+                    Frame(nullptr, co, fn._module, callable, nullptr),
+                    ArgsView(buffer, buffer + co_nlocals)
+                );
+        };
 
-__FAST_CALL:
+        // simple or normal
         callstack.emplace(p0, co, fn._module, callable, args.begin());
         if(op_call) return PY_OP_CALL;
         return _run_top_frame();