From b488a4a8480533a6a3c9468c2f8bd359c94d4d02 Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Fri, 29 Jun 2018 16:32:58 +1000
Subject: [PATCH] py/objgenerator: Eliminate need for mp_obj_gen_wrap wrapper
 instances.

For generating functions there is no need to wrap the bytecode function in
a generator wrapper instance.  Instead the type of the bytecode function
can be changed to mp_type_gen_wrap.  This reduces code size and saves a
block of GC heap RAM for each generator.
---
 py/emitglue.c     |  9 ++++-----
 py/obj.h          |  1 +
 py/objgenerator.c | 17 ++---------------
 3 files changed, 7 insertions(+), 20 deletions(-)

diff --git a/py/emitglue.c b/py/emitglue.c
index 74bf8ddca2..a7fff0e0eb 100644
--- a/py/emitglue.c
+++ b/py/emitglue.c
@@ -146,14 +146,13 @@ mp_obj_t mp_make_function_from_raw_code(const mp_raw_code_t *rc, mp_obj_t def_ar
             // rc->kind should always be set and BYTECODE is the only remaining case
             assert(rc->kind == MP_CODE_BYTECODE);
             fun = mp_obj_new_fun_bc(def_args, def_kw_args, rc->data.u_byte.bytecode, rc->data.u_byte.const_table);
+            // check for generator functions and if so change the type of the object
+            if ((rc->scope_flags & MP_SCOPE_FLAG_GENERATOR) != 0) {
+                ((mp_obj_base_t*)MP_OBJ_TO_PTR(fun))->type = &mp_type_gen_wrap;
+            }
             break;
     }
 
-    // check for generator functions and if so wrap in generator object
-    if ((rc->scope_flags & MP_SCOPE_FLAG_GENERATOR) != 0) {
-        fun = mp_obj_new_gen_wrap(fun);
-    }
-
     return fun;
 }
 
diff --git a/py/obj.h b/py/obj.h
index a64cd0f694..54bc980055 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -547,6 +547,7 @@ extern const mp_obj_type_t mp_type_slice;
 extern const mp_obj_type_t mp_type_zip;
 extern const mp_obj_type_t mp_type_array;
 extern const mp_obj_type_t mp_type_super;
+extern const mp_obj_type_t mp_type_gen_wrap;
 extern const mp_obj_type_t mp_type_gen_instance;
 extern const mp_obj_type_t mp_type_fun_builtin_0;
 extern const mp_obj_type_t mp_type_fun_builtin_1;
diff --git a/py/objgenerator.c b/py/objgenerator.c
index a2ad490d63..8b898c9374 100644
--- a/py/objgenerator.c
+++ b/py/objgenerator.c
@@ -37,11 +37,6 @@
 /******************************************************************************/
 /* generator wrapper                                                          */
 
-typedef struct _mp_obj_gen_wrap_t {
-    mp_obj_base_t base;
-    mp_obj_t *fun;
-} mp_obj_gen_wrap_t;
-
 typedef struct _mp_obj_gen_instance_t {
     mp_obj_base_t base;
     mp_obj_dict_t *globals;
@@ -49,9 +44,8 @@ typedef struct _mp_obj_gen_instance_t {
 } mp_obj_gen_instance_t;
 
 STATIC mp_obj_t gen_wrap_call(mp_obj_t self_in, size_t n_args, size_t n_kw, const mp_obj_t *args) {
-    mp_obj_gen_wrap_t *self = MP_OBJ_TO_PTR(self_in);
-    mp_obj_fun_bc_t *self_fun = (mp_obj_fun_bc_t*)self->fun;
-    assert(self_fun->base.type == &mp_type_fun_bc);
+    // A generating function is just a bytecode function with type mp_type_gen_wrap
+    mp_obj_fun_bc_t *self_fun = MP_OBJ_TO_PTR(self_in);
 
     // bytecode prelude: get state size and exception stack size
     size_t n_state = mp_decode_uint_value(self_fun->bytecode);
@@ -76,13 +70,6 @@ const mp_obj_type_t mp_type_gen_wrap = {
     .unary_op = mp_generic_unary_op,
 };
 
-mp_obj_t mp_obj_new_gen_wrap(mp_obj_t fun) {
-    mp_obj_gen_wrap_t *o = m_new_obj(mp_obj_gen_wrap_t);
-    o->base.type = &mp_type_gen_wrap;
-    o->fun = MP_OBJ_TO_PTR(fun);
-    return MP_OBJ_FROM_PTR(o);
-}
-
 /******************************************************************************/
 /* generator instance                                                         */