1 files changed, 35 insertions, 16 deletions

diff --git a/py/builtinevex.c b/py/builtinevex.c
index e25cbd4d08..74a4640492 100644
--- a/py/builtinevex.c
+++ b/py/builtinevex.c
@@ -26,6 +26,7 @@
 
 #include <stdint.h>
 
+#include "py/objcode.h"
 #include "py/objfun.h"
 #include "py/compile.h"
 #include "py/runtime.h"
@@ -33,17 +34,6 @@
 
 #if MICROPY_PY_BUILTINS_COMPILE
 
-typedef struct _mp_obj_code_t {
-    mp_obj_base_t base;
-    mp_obj_t module_fun;
-} mp_obj_code_t;
-
-static MP_DEFINE_CONST_OBJ_TYPE(
-    mp_type_code,
-    MP_QSTR_code,
-    MP_TYPE_FLAG_NONE
-    );
-
 static mp_obj_t code_execute(mp_obj_code_t *self, mp_obj_dict_t *globals, mp_obj_dict_t *locals) {
     // save context
     nlr_jump_callback_node_globals_locals_t ctx;
@@ -57,19 +47,28 @@ static mp_obj_t code_execute(mp_obj_code_t *self, mp_obj_dict_t *globals, mp_obj
     // set exception handler to restore context if an exception is raised
     nlr_push_jump_callback(&ctx.callback, mp_globals_locals_set_from_nlr_jump_callback);
 
+    #if MICROPY_PY_BUILTINS_CODE >= MICROPY_PY_BUILTINS_CODE_BASIC
+    mp_module_context_t *module_context = m_new_obj(mp_module_context_t);
+    module_context->module.base.type = &mp_type_module;
+    module_context->module.globals = globals;
+    module_context->constants = *mp_code_get_constants(self);
+    mp_obj_t module_fun = mp_make_function_from_proto_fun(mp_code_get_proto_fun(self), module_context, NULL);
+    #else
     // The call to mp_parse_compile_execute() in mp_builtin_compile() below passes
     // NULL for the globals, so repopulate that entry now with the correct globals.
+    mp_obj_t module_fun = self->module_fun;
     if (mp_obj_is_type(self->module_fun, &mp_type_fun_bc)
         #if MICROPY_EMIT_NATIVE
         || mp_obj_is_type(self->module_fun, &mp_type_fun_native)
         #endif
         ) {
-        mp_obj_fun_bc_t *fun_bc = MP_OBJ_TO_PTR(self->module_fun);
+        mp_obj_fun_bc_t *fun_bc = MP_OBJ_TO_PTR(module_fun);
         ((mp_module_context_t *)fun_bc->context)->module.globals = globals;
     }
+    #endif
 
     // execute code
-    mp_obj_t ret = mp_call_function_0(self->module_fun);
+    mp_obj_t ret = mp_call_function_0(module_fun);
 
     // deregister exception handler and restore context
     nlr_pop_jump_callback(true);
@@ -108,9 +107,29 @@ static mp_obj_t mp_builtin_compile(size_t n_args, const mp_obj_t *args) {
             mp_raise_ValueError(MP_ERROR_TEXT("bad compile mode"));
     }
 
-    mp_obj_code_t *code = mp_obj_malloc(mp_obj_code_t, &mp_type_code);
-    code->module_fun = mp_parse_compile_execute(lex, parse_input_kind, NULL, NULL);
-    return MP_OBJ_FROM_PTR(code);
+    #if MICROPY_PY_BUILTINS_CODE >= MICROPY_PY_BUILTINS_CODE_BASIC
+
+    mp_parse_tree_t parse_tree = mp_parse(lex, parse_input_kind);
+    mp_module_context_t ctx;
+    ctx.module.globals = NULL;
+    mp_compiled_module_t cm;
+    cm.context = &ctx;
+    mp_compile_to_raw_code(&parse_tree, lex->source_name, parse_input_kind == MP_PARSE_SINGLE_INPUT, &cm);
+
+    #if MICROPY_PY_BUILTINS_CODE >= MICROPY_PY_BUILTINS_CODE_FULL
+    mp_module_context_t *ctx_ptr = m_new_obj(mp_module_context_t);
+    *ctx_ptr = ctx;
+    return mp_obj_new_code(ctx_ptr, cm.rc, true);
+    #else
+    return mp_obj_new_code(ctx.constants, cm.rc);
+    #endif
+
+    #else
+
+    mp_obj_t module_fun = mp_parse_compile_execute(lex, parse_input_kind, NULL, NULL);
+    return mp_obj_new_code(module_fun);
+
+    #endif
 }
 MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_builtin_compile_obj, 3, 6, mp_builtin_compile);