py/emitglue: Introduce mp_proto_fun_t as a more general mp_raw_code_t.

Allows bytecode itself to be used instead of an mp_raw_code_t in the simple and common cases of a bytecode function without any children. This can be used to further reduce frozen code size, and has the potential to optimise other areas like importing. Signed-off-by: Damien George <[email protected]>
picoscratch · Feb 16, 2024 · e2ff00e · e2ff00e
1 parent 5e3006f
commit e2ff00e
Show file tree

Hide file tree

Showing 15 changed files with 63 additions and 39 deletions.
diff --git a/docs/develop/compiler.rst b/docs/develop/compiler.rst
@@ -157,7 +157,7 @@ The most relevant method you should know about is this:
        mp_compile_to_raw_code(parse_tree, source_file, is_repl, &cm);
 
        // Create and return a function object that executes the outer module.
-       return mp_make_function_from_raw_code(cm.rc, cm.context, NULL);
+       return mp_make_function_from_proto_fun(cm.rc, cm.context, NULL);
    }
 
 The compiler compiles the code in four passes: scope, stack size, code size and emit.

diff --git a/ports/embed/port/embed_util.c b/ports/embed/port/embed_util.c
@@ -69,7 +69,7 @@ void mp_embed_exec_mpy(const uint8_t *mpy, size_t len) {
         mp_compiled_module_t cm;
         cm.context = ctx;
         mp_raw_code_load_mem(mpy, len, &cm);
-        mp_obj_t f = mp_make_function_from_raw_code(cm.rc, ctx, MP_OBJ_NULL);
+        mp_obj_t f = mp_make_function_from_proto_fun(cm.rc, ctx, MP_OBJ_NULL);
         mp_call_function_0(f);
         nlr_pop();
     } else {

diff --git a/py/bc.h b/py/bc.h
@@ -44,11 +44,11 @@
 //  prelude size    : var uint
 //      contains two values interleaved bit-wise as: xIIIIIIC repeated
 //          x = extension           another byte follows
-//          I = n_info              number of bytes in source info section
+//          I = n_info              number of bytes in source info section (always > 0)
 //          C = n_cells             number of bytes/cells in closure section
 //
 //  source info section:
-//      simple_name : var qstr
+//      simple_name : var qstr      always exists
 //      argname0    : var qstr
 //      ...         : var qstr
 //      argnameN    : var qstr      N = num_pos_args + num_kwonly_args - 1
@@ -226,7 +226,7 @@ typedef struct _mp_compiled_module_t {
 // Outer level struct defining a frozen module.
 typedef struct _mp_frozen_module_t {
     const mp_module_constants_t constants;
-    const struct _mp_raw_code_t *rc;
+    const void *proto_fun;
 } mp_frozen_module_t;
 
 // State for an executing function.

diff --git a/py/builtinimport.c b/py/builtinimport.c
@@ -165,7 +165,7 @@ STATIC void do_load_from_lexer(mp_module_context_t *context, mp_lexer_t *lex) {
 #endif
 
 #if (MICROPY_HAS_FILE_READER && MICROPY_PERSISTENT_CODE_LOAD) || MICROPY_MODULE_FROZEN_MPY
-STATIC void do_execute_raw_code(const mp_module_context_t *context, const mp_raw_code_t *rc, qstr source_name) {
+STATIC void do_execute_proto_fun(const mp_module_context_t *context, mp_proto_fun_t proto_fun, qstr source_name) {
     #if MICROPY_PY___FILE__
     mp_store_attr(MP_OBJ_FROM_PTR(&context->module), MP_QSTR___file__, MP_OBJ_NEW_QSTR(source_name));
     #else
@@ -188,7 +188,7 @@ STATIC void do_execute_raw_code(const mp_module_context_t *context, const mp_raw
     nlr_push_jump_callback(&ctx.callback, mp_globals_locals_set_from_nlr_jump_callback);
 
     // make and execute the function
-    mp_obj_t module_fun = mp_make_function_from_raw_code(rc, context, NULL);
+    mp_obj_t module_fun = mp_make_function_from_proto_fun(proto_fun, context, NULL);
     mp_call_function_0(module_fun);
 
     // deregister exception handler and restore context
@@ -230,7 +230,7 @@ STATIC void do_load(mp_module_context_t *module_obj, vstr_t *file) {
             #else
             qstr frozen_file_qstr = MP_QSTRnull;
             #endif
-            do_execute_raw_code(module_obj, frozen->rc, frozen_file_qstr);
+            do_execute_proto_fun(module_obj, frozen->proto_fun, frozen_file_qstr);
             return;
         }
         #endif
@@ -247,7 +247,7 @@ STATIC void do_load(mp_module_context_t *module_obj, vstr_t *file) {
         mp_compiled_module_t cm;
         cm.context = module_obj;
         mp_raw_code_load_file(file_qstr, &cm);
-        do_execute_raw_code(cm.context, cm.rc, file_qstr);
+        do_execute_proto_fun(cm.context, cm.rc, file_qstr);
         return;
     }
     #endif

diff --git a/py/compile.c b/py/compile.c
@@ -3667,7 +3667,7 @@ mp_obj_t mp_compile(mp_parse_tree_t *parse_tree, qstr source_file, bool is_repl)
     cm.context->module.globals = mp_globals_get();
     mp_compile_to_raw_code(parse_tree, source_file, is_repl, &cm);
     // return function that executes the outer module
-    return mp_make_function_from_raw_code(cm.rc, cm.context, NULL);
+    return mp_make_function_from_proto_fun(cm.rc, cm.context, NULL);
 }
 
 #endif // MICROPY_ENABLE_COMPILER
diff --git a/py/dynruntime.h b/py/dynruntime.h
@@ -196,8 +196,8 @@ static inline void *mp_obj_malloc_helper_dyn(size_t num_bytes, const mp_obj_type
 #define mp_unary_op(op, obj)        (mp_fun_table.unary_op((op), (obj)))
 #define mp_binary_op(op, lhs, rhs)  (mp_fun_table.binary_op((op), (lhs), (rhs)))
 
-#define mp_make_function_from_raw_code(rc, context, def_args) \
-    (mp_fun_table.make_function_from_raw_code((rc), (context), (def_args)))
+#define mp_make_function_from_proto_fun(rc, context, def_args) \
+    (mp_fun_table.make_function_from_proto_fun((rc), (context), (def_args)))
 
 #define mp_call_function_n_kw(fun, n_args, n_kw, args) \
     (mp_fun_table.call_function_n_kw((fun), (n_args) | ((n_kw) << 8), args))
@@ -208,6 +208,8 @@ static inline void *mp_obj_malloc_helper_dyn(size_t num_bytes, const mp_obj_type
 #define MP_DYNRUNTIME_INIT_ENTRY \
     mp_obj_t old_globals = mp_fun_table.swap_globals(self->context->module.globals); \
     mp_raw_code_truncated_t rc; \
+    rc.proto_fun_indicator[0] = MP_PROTO_FUN_INDICATOR_RAW_CODE_0; \
+    rc.proto_fun_indicator[1] = MP_PROTO_FUN_INDICATOR_RAW_CODE_1; \
     rc.kind = MP_CODE_NATIVE_VIPER; \
     rc.is_generator = 0; \
     (void)rc;
@@ -217,7 +219,7 @@ static inline void *mp_obj_malloc_helper_dyn(size_t num_bytes, const mp_obj_type
     return mp_const_none;
 
 #define MP_DYNRUNTIME_MAKE_FUNCTION(f) \
-    (mp_make_function_from_raw_code((rc.fun_data = (f), (const mp_raw_code_t *)&rc), self->context, NULL))
+    (mp_make_function_from_proto_fun((rc.fun_data = (f), (const mp_raw_code_t *)&rc), self->context, NULL))
 
 #define mp_import_name(name, fromlist, level) \
     (mp_fun_table.import_name((name), (fromlist), (level)))

diff --git a/py/emitglue.c b/py/emitglue.c
@@ -173,16 +173,19 @@ void mp_emit_glue_assign_native(mp_raw_code_t *rc, mp_raw_code_kind_t kind, void
 }
 #endif
 
-mp_obj_t mp_make_function_from_raw_code(const mp_raw_code_t *rc, const mp_module_context_t *context, const mp_obj_t *def_args) {
-    DEBUG_OP_printf("make_function_from_raw_code %p\n", rc);
-    assert(rc != NULL);
+mp_obj_t mp_make_function_from_proto_fun(mp_proto_fun_t proto_fun, const mp_module_context_t *context, const mp_obj_t *def_args) {
+    DEBUG_OP_printf("make_function_from_proto_fun %p\n", proto_fun);
+    assert(proto_fun != NULL);
 
     // def_args must be MP_OBJ_NULL or a tuple
     assert(def_args == NULL || def_args[0] == MP_OBJ_NULL || mp_obj_is_type(def_args[0], &mp_type_tuple));
 
     // def_kw_args must be MP_OBJ_NULL or a dict
     assert(def_args == NULL || def_args[1] == MP_OBJ_NULL || mp_obj_is_type(def_args[1], &mp_type_dict));
 
+    // the proto-function is a mp_raw_code_t
+    const mp_raw_code_t *rc = proto_fun;
+
     // make the function, depending on the raw code kind
     mp_obj_t fun;
     switch (rc->kind) {
@@ -221,16 +224,16 @@ mp_obj_t mp_make_function_from_raw_code(const mp_raw_code_t *rc, const mp_module
     return fun;
 }
 
-mp_obj_t mp_make_closure_from_raw_code(const mp_raw_code_t *rc, const mp_module_context_t *context, mp_uint_t n_closed_over, const mp_obj_t *args) {
-    DEBUG_OP_printf("make_closure_from_raw_code %p " UINT_FMT " %p\n", rc, n_closed_over, args);
+mp_obj_t mp_make_closure_from_proto_fun(mp_proto_fun_t proto_fun, const mp_module_context_t *context, mp_uint_t n_closed_over, const mp_obj_t *args) {
+    DEBUG_OP_printf("make_closure_from_proto_fun %p " UINT_FMT " %p\n", proto_fun, n_closed_over, args);
     // make function object
     mp_obj_t ffun;
     if (n_closed_over & 0x100) {
         // default positional and keyword args given
-        ffun = mp_make_function_from_raw_code(rc, context, args);
+        ffun = mp_make_function_from_proto_fun(proto_fun, context, args);
     } else {
         // default positional and keyword args not given
-        ffun = mp_make_function_from_raw_code(rc, context, NULL);
+        ffun = mp_make_function_from_proto_fun(proto_fun, context, NULL);
     }
     // wrap function in closure object
     return mp_obj_new_closure(ffun, n_closed_over & 0xff, args + ((n_closed_over >> 7) & 2));

diff --git a/py/emitglue.h b/py/emitglue.h
@@ -31,6 +31,11 @@
 
 // These variables and functions glue the code emitters to the runtime.
 
+// Used with mp_raw_code_t::proto_fun_indicator to detect if a mp_proto_fun_t is a
+// mp_raw_code_t struct or a direct pointer to bytecode.
+#define MP_PROTO_FUN_INDICATOR_RAW_CODE_0 (0)
+#define MP_PROTO_FUN_INDICATOR_RAW_CODE_1 (0)
+
 // These must fit in 8 bits; see scope.h
 enum {
     MP_EMIT_OPT_NONE,
@@ -49,14 +54,25 @@ typedef enum {
     MP_CODE_NATIVE_ASM,
 } mp_raw_code_kind_t;
 
-// This mp_raw_code_t struct holds static information about a non-instantiated function.
+// An mp_proto_fun_t points to static information about a non-instantiated function.
 // A function object is created from this information, and that object can then be executed.
-//
-// This struct appears in the following places:
+// It points either to bytecode, or an mp_raw_code_t struct.
+typedef const void *mp_proto_fun_t;
+
+// Bytecode is distinguished from an mp_raw_code_t struct by the first two bytes: bytecode
+// is guaranteed to have either its first or second byte non-zero.  So if both bytes are
+// zero then the mp_proto_fun_t pointer must be an mp_raw_code_t.
+static inline bool mp_proto_fun_is_bytecode(mp_proto_fun_t proto_fun) {
+    const uint8_t *header = proto_fun;
+    return (header[0] | (header[1] << 8)) != (MP_PROTO_FUN_INDICATOR_RAW_CODE_0 | (MP_PROTO_FUN_INDICATOR_RAW_CODE_1 << 8));
+}
+
+// The mp_raw_code_t struct appears in the following places:
 // compiled bytecode: instance in RAM, referenced by outer scope, usually freed after first (and only) use
 // mpy file: instance in RAM, created when .mpy file is loaded (same comments as above)
 // frozen: instance in ROM
 typedef struct _mp_raw_code_t {
+    uint8_t proto_fun_indicator[2];
     uint8_t kind; // of type mp_raw_code_kind_t; only 3 bits used
     bool is_generator;
     const void *fun_data;
@@ -88,6 +104,7 @@ typedef struct _mp_raw_code_t {
 // only needed when the kind is MP_CODE_NATIVE_ASM.  So this struct can be used when the
 // kind is MP_CODE_BYTECODE, MP_CODE_NATIVE_PY or MP_CODE_NATIVE_VIPER, to reduce its size.
 typedef struct _mp_raw_code_truncated_t {
+    uint8_t proto_fun_indicator[2];
     uint8_t kind;
     bool is_generator;
     const void *fun_data;
@@ -127,7 +144,7 @@ void mp_emit_glue_assign_native(mp_raw_code_t *rc, mp_raw_code_kind_t kind, void
     #endif
     uint16_t scope_flags, uint32_t asm_n_pos_args, uint32_t asm_type_sig);
 
-mp_obj_t mp_make_function_from_raw_code(const mp_raw_code_t *rc, const mp_module_context_t *context, const mp_obj_t *def_args);
-mp_obj_t mp_make_closure_from_raw_code(const mp_raw_code_t *rc, const mp_module_context_t *context, mp_uint_t n_closed_over, const mp_obj_t *args);
+mp_obj_t mp_make_function_from_proto_fun(mp_proto_fun_t proto_fun, const mp_module_context_t *context, const mp_obj_t *def_args);
+mp_obj_t mp_make_closure_from_proto_fun(mp_proto_fun_t proto_fun, const mp_module_context_t *context, mp_uint_t n_closed_over, const mp_obj_t *args);
 
 #endif // MICROPY_INCLUDED_PY_EMITGLUE_H
diff --git a/py/emitnative.c b/py/emitnative.c
@@ -2657,7 +2657,7 @@ STATIC void emit_native_make_function(emit_t *emit, scope_t *scope, mp_uint_t n_
         need_reg_all(emit);
     }
     emit_load_reg_with_child(emit, REG_ARG_1, scope->raw_code);
-    ASM_CALL_IND(emit->as, MP_F_MAKE_FUNCTION_FROM_RAW_CODE);
+    ASM_CALL_IND(emit->as, MP_F_MAKE_FUNCTION_FROM_PROTO_FUN);
     emit_post_push_reg(emit, VTYPE_PYOBJ, REG_RET);
 }
 
@@ -2675,7 +2675,7 @@ STATIC void emit_native_make_closure(emit_t *emit, scope_t *scope, mp_uint_t n_c
         need_reg_all(emit);
     }
     emit_load_reg_with_child(emit, REG_ARG_1, scope->raw_code);
-    ASM_CALL_IND(emit->as, MP_F_MAKE_FUNCTION_FROM_RAW_CODE);
+    ASM_CALL_IND(emit->as, MP_F_MAKE_FUNCTION_FROM_PROTO_FUN);
 
     // make closure
     #if REG_ARG_1 != REG_RET

diff --git a/py/emitnx86.c b/py/emitnx86.c
@@ -37,7 +37,7 @@ STATIC byte mp_f_n_args[MP_F_NUMBER_OF] = {
     [MP_F_STORE_SET] = 2,
     [MP_F_LIST_APPEND] = 2,
     [MP_F_STORE_MAP] = 3,
-    [MP_F_MAKE_FUNCTION_FROM_RAW_CODE] = 3,
+    [MP_F_MAKE_FUNCTION_FROM_PROTO_FUN] = 3,
     [MP_F_NATIVE_CALL_FUNCTION_N_KW] = 3,
     [MP_F_CALL_METHOD_N_KW] = 3,
     [MP_F_CALL_METHOD_N_KW_VAR] = 3,

diff --git a/py/nativeglue.c b/py/nativeglue.c
@@ -281,7 +281,7 @@ const mp_fun_table_t mp_fun_table = {
     mp_obj_set_store,
     mp_obj_list_append,
     mp_obj_dict_store,
-    mp_make_function_from_raw_code,
+    mp_make_function_from_proto_fun,
     mp_native_call_function_n_kw,
     mp_call_method_n_kw,
     mp_call_method_n_kw_var,

diff --git a/py/nativeglue.h b/py/nativeglue.h
@@ -58,7 +58,7 @@ typedef enum {
     MP_F_STORE_SET,
     MP_F_LIST_APPEND,
     MP_F_STORE_MAP,
-    MP_F_MAKE_FUNCTION_FROM_RAW_CODE,
+    MP_F_MAKE_FUNCTION_FROM_PROTO_FUN,
     MP_F_NATIVE_CALL_FUNCTION_N_KW,
     MP_F_CALL_METHOD_N_KW,
     MP_F_CALL_METHOD_N_KW_VAR,
@@ -112,7 +112,7 @@ typedef struct _mp_fun_table_t {
     void (*set_store)(mp_obj_t self_in, mp_obj_t item);
     mp_obj_t (*list_append)(mp_obj_t self_in, mp_obj_t arg);
     mp_obj_t (*dict_store)(mp_obj_t self_in, mp_obj_t key, mp_obj_t value);
-    mp_obj_t (*make_function_from_raw_code)(const mp_raw_code_t *rc, const mp_module_context_t *cm, const mp_obj_t *def_args);
+    mp_obj_t (*make_function_from_proto_fun)(mp_proto_fun_t proto_fun, const mp_module_context_t *cm, const mp_obj_t *def_args);
     mp_obj_t (*call_function_n_kw)(mp_obj_t fun_in, size_t n_args_kw, const mp_obj_t *args);
     mp_obj_t (*call_method_n_kw)(size_t n_args, size_t n_kw, const mp_obj_t *args);
     mp_obj_t (*call_method_n_kw_var)(bool have_self, size_t n_args_n_kw, const mp_obj_t *args);

diff --git a/py/vm.c b/py/vm.c
@@ -893,15 +893,15 @@ unwind_jump:;
 
                 ENTRY(MP_BC_MAKE_FUNCTION): {
                     DECODE_PTR;
-                    PUSH(mp_make_function_from_raw_code(ptr, code_state->fun_bc->context, NULL));
+                    PUSH(mp_make_function_from_proto_fun(ptr, code_state->fun_bc->context, NULL));
                     DISPATCH();
                 }
 
                 ENTRY(MP_BC_MAKE_FUNCTION_DEFARGS): {
                     DECODE_PTR;
                     // Stack layout: def_tuple def_dict <- TOS
                     sp -= 1;
-                    SET_TOP(mp_make_function_from_raw_code(ptr, code_state->fun_bc->context, sp));
+                    SET_TOP(mp_make_function_from_proto_fun(ptr, code_state->fun_bc->context, sp));
                     DISPATCH();
                 }
 
@@ -910,7 +910,7 @@ unwind_jump:;
                     size_t n_closed_over = *ip++;
                     // Stack layout: closed_overs <- TOS
                     sp -= n_closed_over - 1;
-                    SET_TOP(mp_make_closure_from_raw_code(ptr, code_state->fun_bc->context, n_closed_over, sp));
+                    SET_TOP(mp_make_closure_from_proto_fun(ptr, code_state->fun_bc->context, n_closed_over, sp));
                     DISPATCH();
                 }
 
@@ -919,7 +919,7 @@ unwind_jump:;
                     size_t n_closed_over = *ip++;
                     // Stack layout: def_tuple def_dict closed_overs <- TOS
                     sp -= 2 + n_closed_over - 1;
-                    SET_TOP(mp_make_closure_from_raw_code(ptr, code_state->fun_bc->context, 0x100 | n_closed_over, sp));
+                    SET_TOP(mp_make_closure_from_proto_fun(ptr, code_state->fun_bc->context, 0x100 | n_closed_over, sp));
                     DISPATCH();
                 }
 

diff --git a/shared/runtime/pyexec.c b/shared/runtime/pyexec.c
@@ -88,7 +88,7 @@ STATIC int parse_compile_execute(const void *source, mp_parse_input_kind_t input
             mp_module_context_t *ctx = m_new_obj(mp_module_context_t);
             ctx->module.globals = mp_globals_get();
             ctx->constants = frozen->constants;
-            module_fun = mp_make_function_from_raw_code(frozen->rc, ctx, NULL);
+            module_fun = mp_make_function_from_proto_fun(frozen->proto_fun, ctx, NULL);
         } else
         #endif
         {

diff --git a/tools/mpy-tool.py b/tools/mpy-tool.py
@@ -684,7 +684,7 @@ def freeze(self, compiled_module_index):
         else:
             print("        .obj_table = NULL,")
         print("    },")
-        print("    .rc = (const mp_raw_code_t *)&raw_code_%s," % self.raw_code.escaped_name)
+        print("    .proto_fun = &proto_fun_%s," % self.raw_code.escaped_name)
         print("};")
 
     def freeze_constant_obj(self, obj_name, obj):
@@ -899,7 +899,7 @@ def freeze_children(self, prelude_ptr=None):
                 print()
             print("static const mp_raw_code_t *const children_%s[] = {" % self.escaped_name)
             for rc in self.children:
-                print("    (const mp_raw_code_t *)&raw_code_%s," % rc.escaped_name)
+                print("    (const mp_raw_code_t *)&proto_fun_%s," % rc.escaped_name)
             if prelude_ptr:
                 print("    (void *)%s," % prelude_ptr)
             print("};")
@@ -911,7 +911,9 @@ def freeze_raw_code(self, prelude_ptr=None, type_sig=0):
             raw_code_type = "mp_raw_code_t"
         else:
             raw_code_type = "mp_raw_code_truncated_t"
-        print("static const %s raw_code_%s = {" % (raw_code_type, self.escaped_name))
+        print("static const %s proto_fun_%s = {" % (raw_code_type, self.escaped_name))
+        print("    .proto_fun_indicator[0] = MP_PROTO_FUN_INDICATOR_RAW_CODE_0,")
+        print("    .proto_fun_indicator[1] = MP_PROTO_FUN_INDICATOR_RAW_CODE_1,")
         print("    .kind = %s," % RawCode.code_kind_str[self.code_kind])
         print("    .is_generator = %d," % bool(self.scope_flags & MP_SCOPE_FLAG_GENERATOR))
         print("    .fun_data = fun_data_%s," % self.escaped_name)