py: Rework bytecode and .mpy file format to be mostly static data.

Background: .mpy files are precompiled .py files, built using mpy-cross,
that contain compiled bytecode functions (and can also contain machine
code). The benefit of using an .mpy file over a .py file is that they are
faster to import and take less memory when importing.  They are also
smaller on disk.

But the real benefit of .mpy files comes when they are frozen into the
firmware.  This is done by loading the .mpy file during compilation of the
firmware and turning it into a set of big C data structures (the job of
mpy-tool.py), which are then compiled and downloaded into the ROM of a
device.  These C data structures can be executed in-place, ie directly from
ROM.  This makes importing even faster because there is very little to do,
and also means such frozen modules take up much less RAM (because their
bytecode stays in ROM).

The downside of frozen code is that it requires recompiling and reflashing
the entire firmware.  This can be a big barrier to entry, slows down
development time, and makes it harder to do OTA updates of frozen code
(because the whole firmware must be updated).

This commit attempts to solve this problem by providing a solution that
sits between loading .mpy files into RAM and freezing them into the
firmware.  The .mpy file format has been reworked so that it consists of
data and bytecode which is mostly static and ready to run in-place.  If
these new .mpy files are located in flash/ROM which is memory addressable,
the .mpy file can be executed (mostly) in-place.

With this approach there is still a small amount of unpacking and linking
of the .mpy file that needs to be done when it's imported, but it's still
much better than loading an .mpy from disk into RAM (although not as good
as freezing .mpy files into the firmware).

The main trick to make static .mpy files is to adjust the bytecode so any
qstrs that it references now go through a lookup table to convert from
local qstr number in the module to global qstr number in the firmware.
That means the bytecode does not need linking/rewriting of qstrs when it's
loaded.  Instead only a small qstr table needs to be built (and put in RAM)
at import time.  This means the bytecode itself is static/constant and can
be used directly if it's in addressable memory.  Also the qstr string data
in the .mpy file, and some constant object data, can be used directly.
Note that the qstr table is global to the module (ie not per function).

In more detail, in the VM what used to be (schematically):

    qst = DECODE_QSTR_VALUE;

is now (schematically):

    idx = DECODE_QSTR_INDEX;
    qst = qstr_table[idx];

That allows the bytecode to be fixed at compile time and not need
relinking/rewriting of the qstr values.  Only qstr_table needs to be linked
when the .mpy is loaded.

Incidentally, this helps to reduce the size of bytecode because what used
to be 2-byte qstr values in the bytecode are now (mostly) 1-byte indices.
If the module uses the same qstr more than two times then the bytecode is
smaller than before.

The following changes are measured for this commit compared to the
previous (the baseline):
- average 7%-9% reduction in size of .mpy files
- frozen code size is reduced by about 5%-7%
- importing .py files uses about 5% less RAM in total
- importing .mpy files uses about 4% less RAM in total
- importing .py and .mpy files takes about the same time as before

The qstr indirection in the bytecode has only a small impact on VM
performance.  For stm32 on PYBv1.0 the performance change of this commit
is:

diff of scores (higher is better)
N=100 M=100             baseline -> this-commit  diff      diff% (error%)
bm_chaos.py               371.07 ->  357.39 :  -13.68 =  -3.687% (+/-0.02%)
bm_fannkuch.py             78.72 ->   77.49 :   -1.23 =  -1.563% (+/-0.01%)
bm_fft.py                2591.73 -> 2539.28 :  -52.45 =  -2.024% (+/-0.00%)
bm_float.py              6034.93 -> 5908.30 : -126.63 =  -2.098% (+/-0.01%)
bm_hexiom.py               48.96 ->   47.93 :   -1.03 =  -2.104% (+/-0.00%)
bm_nqueens.py            4510.63 -> 4459.94 :  -50.69 =  -1.124% (+/-0.00%)
bm_pidigits.py            650.28 ->  644.96 :   -5.32 =  -0.818% (+/-0.23%)
core_import_mpy_multi.py  564.77 ->  581.49 :  +16.72 =  +2.960% (+/-0.01%)
core_import_mpy_single.py  68.67 ->   67.16 :   -1.51 =  -2.199% (+/-0.01%)
core_qstr.py               64.16 ->   64.12 :   -0.04 =  -0.062% (+/-0.00%)
core_yield_from.py        362.58 ->  354.50 :   -8.08 =  -2.228% (+/-0.00%)
misc_aes.py               429.69 ->  405.59 :  -24.10 =  -5.609% (+/-0.01%)
misc_mandel.py           3485.13 -> 3416.51 :  -68.62 =  -1.969% (+/-0.00%)
misc_pystone.py          2496.53 -> 2405.56 :  -90.97 =  -3.644% (+/-0.01%)
misc_raytrace.py          381.47 ->  374.01 :   -7.46 =  -1.956% (+/-0.01%)
viper_call0.py            576.73 ->  572.49 :   -4.24 =  -0.735% (+/-0.04%)
viper_call1a.py           550.37 ->  546.21 :   -4.16 =  -0.756% (+/-0.09%)
viper_call1b.py           438.23 ->  435.68 :   -2.55 =  -0.582% (+/-0.06%)
viper_call1c.py           442.84 ->  440.04 :   -2.80 =  -0.632% (+/-0.08%)
viper_call2a.py           536.31 ->  532.35 :   -3.96 =  -0.738% (+/-0.06%)
viper_call2b.py           382.34 ->  377.07 :   -5.27 =  -1.378% (+/-0.03%)

And for unix on x64:

diff of scores (higher is better)
N=2000 M=2000        baseline -> this-commit     diff      diff% (error%)
bm_chaos.py          13594.20 ->  13073.84 :  -520.36 =  -3.828% (+/-5.44%)
bm_fannkuch.py          60.63 ->     59.58 :    -1.05 =  -1.732% (+/-3.01%)
bm_fft.py           112009.15 -> 111603.32 :  -405.83 =  -0.362% (+/-4.03%)
bm_float.py         246202.55 -> 247923.81 : +1721.26 =  +0.699% (+/-2.79%)
bm_hexiom.py           615.65 ->    617.21 :    +1.56 =  +0.253% (+/-1.64%)
bm_nqueens.py       215807.95 -> 215600.96 :  -206.99 =  -0.096% (+/-3.52%)
bm_pidigits.py        8246.74 ->   8422.82 :  +176.08 =  +2.135% (+/-3.64%)
misc_aes.py          16133.00 ->  16452.74 :  +319.74 =  +1.982% (+/-1.50%)
misc_mandel.py      128146.69 -> 130796.43 : +2649.74 =  +2.068% (+/-3.18%)
misc_pystone.py      83811.49 ->  83124.85 :  -686.64 =  -0.819% (+/-1.03%)
misc_raytrace.py     21688.02 ->  21385.10 :  -302.92 =  -1.397% (+/-3.20%)

The code size change is (firmware with a lot of frozen code benefits the
most):

       bare-arm:  +396 +0.697%
    minimal x86: +1595 +0.979% [incl +32(data)]
       unix x64: +2408 +0.470% [incl +800(data)]
    unix nanbox: +1396 +0.309% [incl -96(data)]
          stm32: -1256 -0.318% PYBV10
         cc3200:  +288 +0.157%
        esp8266:  -260 -0.037% GENERIC
          esp32:  -216 -0.014% GENERIC[incl -1072(data)]
            nrf:  +116 +0.067% pca10040
            rp2:  -664 -0.135% PICO
           samd:  +844 +0.607% ADAFRUIT_ITSYBITSY_M4_EXPRESS

As part of this change the .mpy file format version is bumped to version 6.
And mpy-tool.py has been improved to provide a good visualisation of the
contents of .mpy files.

In summary: this commit changes the bytecode to use qstr indirection, and
reworks the .mpy file format to be simpler and allow .mpy files to be
executed in-place.  Performance is not impacted too much.  Eventually it
will be possible to store such .mpy files in a linear, read-only, memory-
mappable filesystem so they can be executed from flash/ROM.  This will
essentially be able to replace frozen code for most applications.

Signed-off-by: Damien George <damien@micropython.org>

1 files changed, 182 insertions, 372 deletions

diff --git a/py/persistentcode.c b/py/persistentcode.c
index ac523990c..b473f1830 100644
--- a/py/persistentcode.c
+++ b/py/persistentcode.c
@@ -63,57 +63,6 @@ STATIC int mp_small_int_bits(void) {
 }
 #endif
 
-#define QSTR_WINDOW_SIZE (32)
-
-typedef struct _qstr_window_t {
-    uint16_t idx; // indexes the head of the window
-    uint16_t window[QSTR_WINDOW_SIZE];
-} qstr_window_t;
-
-// Push a qstr to the head of the window, and the tail qstr is overwritten
-STATIC void qstr_window_push(qstr_window_t *qw, qstr qst) {
-    qw->idx = (qw->idx + 1) % QSTR_WINDOW_SIZE;
-    qw->window[qw->idx] = qst;
-}
-
-// Pull an existing qstr from within the window to the head of the window
-STATIC qstr qstr_window_pull(qstr_window_t *qw, size_t idx) {
-    qstr qst = qw->window[idx];
-    if (idx > qw->idx) {
-        memmove(&qw->window[idx], &qw->window[idx + 1], (QSTR_WINDOW_SIZE - idx - 1) * sizeof(uint16_t));
-        qw->window[QSTR_WINDOW_SIZE - 1] = qw->window[0];
-        idx = 0;
-    }
-    memmove(&qw->window[idx], &qw->window[idx + 1], (qw->idx - idx) * sizeof(uint16_t));
-    qw->window[qw->idx] = qst;
-    return qst;
-}
-
-#if MICROPY_PERSISTENT_CODE_LOAD
-
-// Access a qstr at the given index, relative to the head of the window (0=head)
-STATIC qstr qstr_window_access(qstr_window_t *qw, size_t idx) {
-    return qstr_window_pull(qw, (qw->idx + QSTR_WINDOW_SIZE - idx) % QSTR_WINDOW_SIZE);
-}
-
-#endif
-
-#if MICROPY_PERSISTENT_CODE_SAVE
-
-// Insert a qstr at the head of the window, either by pulling an existing one or pushing a new one
-STATIC size_t qstr_window_insert(qstr_window_t *qw, qstr qst) {
-    for (size_t idx = 0; idx < QSTR_WINDOW_SIZE; ++idx) {
-        if (qw->window[idx] == qst) {
-            qstr_window_pull(qw, idx);
-            return (qw->idx + QSTR_WINDOW_SIZE - idx) % QSTR_WINDOW_SIZE;
-        }
-    }
-    qstr_window_push(qw, qst);
-    return QSTR_WINDOW_SIZE;
-}
-
-#endif
-
 typedef struct _bytecode_prelude_t {
     uint n_state;
     uint n_exc_stack;
@@ -124,23 +73,6 @@ typedef struct _bytecode_prelude_t {
     uint code_info_size;
 } bytecode_prelude_t;
 
-// ip will point to start of opcodes
-// return value will point to simple_name, source_file qstrs
-STATIC byte *extract_prelude(const byte **ip, bytecode_prelude_t *prelude) {
-    MP_BC_PRELUDE_SIG_DECODE(*ip);
-    prelude->n_state = n_state;
-    prelude->n_exc_stack = n_exc_stack;
-    prelude->scope_flags = scope_flags;
-    prelude->n_pos_args = n_pos_args;
-    prelude->n_kwonly_args = n_kwonly_args;
-    prelude->n_def_pos_args = n_def_pos_args;
-    MP_BC_PRELUDE_SIZE_DECODE(*ip);
-    byte *ip_info = (byte *)*ip;
-    *ip += n_info;
-    *ip += n_cell;
-    return ip_info;
-}
-
 #endif // MICROPY_PERSISTENT_CODE_LOAD || MICROPY_PERSISTENT_CODE_SAVE
 
 #if MICROPY_PERSISTENT_CODE_LOAD
@@ -148,13 +80,14 @@ STATIC byte *extract_prelude(const byte **ip, bytecode_prelude_t *prelude) {
 #include "py/parsenum.h"
 
 STATIC int read_byte(mp_reader_t *reader);
-STATIC size_t read_uint(mp_reader_t *reader, byte **out);
+STATIC size_t read_uint(mp_reader_t *reader);
 
 #if MICROPY_EMIT_MACHINE_CODE
 
 typedef struct _reloc_info_t {
     mp_reader_t *reader;
-    mp_uint_t *const_table;
+    uint8_t *rodata;
+    uint8_t *bss;
 } reloc_info_t;
 
 #if MICROPY_EMIT_THUMB
@@ -197,13 +130,13 @@ void mp_native_relocate(void *ri_in, uint8_t *text, uintptr_t reloc_text) {
     while ((op = read_byte(ri->reader)) != 0xff) {
         if (op & 1) {
             // Point to new location to make adjustments
-            size_t addr = read_uint(ri->reader, NULL);
+            size_t addr = read_uint(ri->reader);
             if ((addr & 1) == 0) {
                 // Point to somewhere in text
                 addr_to_adjust = &((uintptr_t *)text)[addr >> 1];
             } else {
                 // Point to somewhere in rodata
-                addr_to_adjust = &((uintptr_t *)ri->const_table[1])[addr >> 1];
+                addr_to_adjust = &((uintptr_t *)ri->rodata)[addr >> 1];
             }
         }
         op >>= 1;
@@ -212,15 +145,18 @@ void mp_native_relocate(void *ri_in, uint8_t *text, uintptr_t reloc_text) {
         if (op <= 5) {
             if (op & 1) {
                 // Read in number of adjustments to make
-                n = read_uint(ri->reader, NULL);
+                n = read_uint(ri->reader);
             }
             op >>= 1;
             if (op == 0) {
                 // Destination is text
                 dest = reloc_text;
+            } else if (op == 1) {
+                // Destination is rodata
+                dest = (uintptr_t)ri->rodata;
             } else {
-                // Destination is rodata (op=1) or bss (op=1 if no rodata, else op=2)
-                dest = ri->const_table[op];
+                // Destination is bss
+                dest = (uintptr_t)ri->bss;
             }
         } else if (op == 6) {
             // Destination is mp_fun_table itself
@@ -247,14 +183,10 @@ STATIC void read_bytes(mp_reader_t *reader, byte *buf, size_t len) {
     }
 }
 
-STATIC size_t read_uint(mp_reader_t *reader, byte **out) {
+STATIC size_t read_uint(mp_reader_t *reader) {
     size_t unum = 0;
     for (;;) {
         byte b = reader->readbyte(reader->data);
-        if (out != NULL) {
-            **out = b;
-            ++*out;
-        }
         unum = (unum << 7) | (b & 0x7f);
         if ((b & 0x80) == 0) {
             break;
@@ -263,35 +195,41 @@ STATIC size_t read_uint(mp_reader_t *reader, byte **out) {
     return unum;
 }
 
-STATIC qstr load_qstr(mp_reader_t *reader, qstr_window_t *qw) {
-    size_t len = read_uint(reader, NULL);
-    if (len == 0) {
-        // static qstr
-        return read_byte(reader);
-    }
+STATIC qstr load_qstr(mp_reader_t *reader) {
+    size_t len = read_uint(reader);
     if (len & 1) {
-        // qstr in window
-        return qstr_window_access(qw, len >> 1);
+        // static qstr
+        return len >> 1;
     }
     len >>= 1;
     char *str = m_new(char, len);
     read_bytes(reader, (byte *)str, len);
+    read_byte(reader); // read and discard null terminator
     qstr qst = qstr_from_strn(str, len);
     m_del(char, str, len);
-    qstr_window_push(qw, qst);
     return qst;
 }
 
 STATIC mp_obj_t load_obj(mp_reader_t *reader) {
     byte obj_type = read_byte(reader);
+    #if MICROPY_EMIT_MACHINE_CODE
+    if (obj_type == 't') {
+        return MP_OBJ_FROM_PTR(&mp_fun_table);
+    } else
+    #endif
     if (obj_type == 'e') {
         return MP_OBJ_FROM_PTR(&mp_const_ellipsis_obj);
     } else {
-        size_t len = read_uint(reader, NULL);
+        size_t len = read_uint(reader);
+        if (len == 0 && obj_type == 'b') {
+            read_byte(reader); // skip null terminator
+            return mp_const_empty_bytes;
+        }
         vstr_t vstr;
         vstr_init_len(&vstr, len);
         read_bytes(reader, (byte *)vstr.buf, len);
         if (obj_type == 's' || obj_type == 'b') {
+            read_byte(reader); // skip null terminator
             return mp_obj_new_str_from_vstr(obj_type == 's' ? &mp_type_str : &mp_type_bytes, &vstr);
         } else if (obj_type == 'i') {
             return mp_parse_num_integer(vstr.buf, vstr.len, 10, NULL);
@@ -302,58 +240,12 @@ STATIC mp_obj_t load_obj(mp_reader_t *reader) {
     }
 }
 
-STATIC void load_prelude_qstrs(mp_reader_t *reader, qstr_window_t *qw, byte *ip) {
-    qstr simple_name = load_qstr(reader, qw);
-    ip[0] = simple_name;
-    ip[1] = simple_name >> 8;
-    qstr source_file = load_qstr(reader, qw);
-    ip[2] = source_file;
-    ip[3] = source_file >> 8;
-}
-
-STATIC void load_prelude(mp_reader_t *reader, qstr_window_t *qw, byte **ip, bytecode_prelude_t *prelude) {
-    // Read in the prelude header
-    byte *ip_read = *ip;
-    read_uint(reader, &ip_read);                    // read in n_state/etc (is effectively a var-uint)
-    read_uint(reader, &ip_read);                    // read in n_info/n_cell (is effectively a var-uint)
-
-    // Prelude header has been read into *ip, now decode and extract values from it
-    extract_prelude((const byte **)ip, prelude);
-
-    // Load qstrs in prelude
-    load_prelude_qstrs(reader, qw, ip_read);
-    ip_read += 4;
-
-    // Read remaining code info
-    read_bytes(reader, ip_read, *ip - ip_read);
-}
-
-STATIC void load_bytecode(mp_reader_t *reader, qstr_window_t *qw, byte *ip, byte *ip_top) {
-    while (ip < ip_top) {
-        *ip = read_byte(reader);
-        size_t sz;
-        uint f = mp_opcode_format(ip, &sz, false);
-        ++ip;
-        --sz;
-        if (f == MP_BC_FORMAT_QSTR) {
-            qstr qst = load_qstr(reader, qw);
-            *ip++ = qst;
-            *ip++ = qst >> 8;
-            sz -= 2;
-        } else if (f == MP_BC_FORMAT_VAR_UINT) {
-            while ((*ip++ = read_byte(reader)) & 0x80) {
-            }
-        }
-        read_bytes(reader, ip, sz);
-        ip += sz;
-    }
-}
-
-STATIC mp_raw_code_t *load_raw_code(mp_reader_t *reader, qstr_window_t *qw) {
+STATIC mp_raw_code_t *load_raw_code(mp_reader_t *reader) {
     // Load function kind and data length
-    size_t kind_len = read_uint(reader, NULL);
+    size_t kind_len = read_uint(reader);
     int kind = (kind_len & 3) + MP_CODE_BYTECODE;
-    size_t fun_data_len = kind_len >> 2;
+    bool has_children = !!(kind_len & 4);
+    size_t fun_data_len = kind_len >> 3;
 
     #if !MICROPY_EMIT_MACHINE_CODE
     if (kind != MP_CODE_BYTECODE) {
@@ -362,23 +254,18 @@ STATIC mp_raw_code_t *load_raw_code(mp_reader_t *reader, qstr_window_t *qw) {
     #endif
 
     uint8_t *fun_data = NULL;
-    bytecode_prelude_t prelude = {0};
     #if MICROPY_EMIT_MACHINE_CODE
     size_t prelude_offset = 0;
-    mp_uint_t type_sig = 0;
-    size_t n_qstr_link = 0;
+    mp_uint_t native_scope_flags = 0;
+    mp_uint_t native_n_pos_args = 0;
+    mp_uint_t native_type_sig = 0;
     #endif
 
     if (kind == MP_CODE_BYTECODE) {
         // Allocate memory for the bytecode
         fun_data = m_new(uint8_t, fun_data_len);
-
-        // Load prelude
-        byte *ip = fun_data;
-        load_prelude(reader, qw, &ip, &prelude);
-
         // Load bytecode
-        load_bytecode(reader, qw, ip, fun_data + fun_data_len);
+        read_bytes(reader, fun_data, fun_data_len);
 
     #if MICROPY_EMIT_MACHINE_CODE
     } else {
@@ -389,10 +276,10 @@ STATIC mp_raw_code_t *load_raw_code(mp_reader_t *reader, qstr_window_t *qw) {
 
         if (kind == MP_CODE_NATIVE_PY || kind == MP_CODE_NATIVE_VIPER) {
             // Parse qstr link table and link native code
-            n_qstr_link = read_uint(reader, NULL);
+            size_t n_qstr_link = read_uint(reader);
             for (size_t i = 0; i < n_qstr_link; ++i) {
-                size_t off = read_uint(reader, NULL);
-                qstr qst = load_qstr(reader, qw);
+                size_t off = read_uint(reader);
+                qstr qst = load_qstr(reader);
                 uint8_t *dest = fun_data + (off >> 2);
                 if ((off & 3) == 0) {
                     // Generic 16-bit link
@@ -409,113 +296,92 @@ STATIC mp_raw_code_t *load_raw_code(mp_reader_t *reader, qstr_window_t *qw) {
         }
 
         if (kind == MP_CODE_NATIVE_PY) {
-            // Extract prelude for later use
-            prelude_offset = read_uint(reader, NULL);
+            // Read prelude offset within fun_data, and extract scope flags.
+            prelude_offset = read_uint(reader);
             const byte *ip = fun_data + prelude_offset;
-            byte *ip_info = extract_prelude(&ip, &prelude);
-            // Load qstrs in prelude
-            load_prelude_qstrs(reader, qw, ip_info);
+            MP_BC_PRELUDE_SIG_DECODE(ip);
+            native_scope_flags = scope_flags;
         } else {
-            // Load basic scope info for viper and asm
-            prelude.scope_flags = read_uint(reader, NULL);
-            prelude.n_pos_args = 0;
-            prelude.n_kwonly_args = 0;
+            // Load basic scope info for viper and asm.
+            native_scope_flags = read_uint(reader);
             if (kind == MP_CODE_NATIVE_ASM) {
-                prelude.n_pos_args = read_uint(reader, NULL);
-                type_sig = read_uint(reader, NULL);
+                native_n_pos_args = read_uint(reader);
+                native_type_sig = read_uint(reader);
             }
         }
     #endif
     }
 
-    size_t n_obj = 0;
-    size_t n_raw_code = 0;
-    mp_uint_t *const_table = NULL;
+    size_t n_children = 0;
+    mp_raw_code_t **children = NULL;
 
-    if (kind != MP_CODE_NATIVE_ASM) {
-        // Load constant table for bytecode, native and viper
-
-        // Number of entries in constant table
-        n_obj = read_uint(reader, NULL);
-        n_raw_code = read_uint(reader, NULL);
-
-        // Allocate constant table
-        size_t n_alloc = prelude.n_pos_args + prelude.n_kwonly_args + n_obj + n_raw_code;
-        #if MICROPY_EMIT_MACHINE_CODE
-        if (kind != MP_CODE_BYTECODE) {
-            ++n_alloc; // additional entry for mp_fun_table
-            if (prelude.scope_flags & MP_SCOPE_FLAG_VIPERRODATA) {
-                ++n_alloc; // additional entry for rodata
-            }
-            if (prelude.scope_flags & MP_SCOPE_FLAG_VIPERBSS) {
-                ++n_alloc; // additional entry for BSS
-            }
+    #if MICROPY_EMIT_MACHINE_CODE
+    // Load optional BSS/rodata for viper.
+    uint8_t *rodata = NULL;
+    uint8_t *bss = NULL;
+    if (kind == MP_CODE_NATIVE_VIPER) {
+        size_t rodata_size = 0;
+        if (native_scope_flags & MP_SCOPE_FLAG_VIPERRODATA) {
+            rodata_size = read_uint(reader);
         }
-        #endif
 
-        const_table = m_new(mp_uint_t, n_alloc);
-        mp_uint_t *ct = const_table;
-
-        // Load function argument names (initial entries in const_table)
-        // (viper has n_pos_args=n_kwonly_args=0 so doesn't load any qstrs here)
-        for (size_t i = 0; i < prelude.n_pos_args + prelude.n_kwonly_args; ++i) {
-            *ct++ = (mp_uint_t)MP_OBJ_NEW_QSTR(load_qstr(reader, qw));
+        size_t bss_size = 0;
+        if (native_scope_flags & MP_SCOPE_FLAG_VIPERBSS) {
+            bss_size = read_uint(reader);
         }
 
-        #if MICROPY_EMIT_MACHINE_CODE
-        if (kind != MP_CODE_BYTECODE) {
-            // Populate mp_fun_table entry
-            *ct++ = (mp_uint_t)(uintptr_t)&mp_fun_table;
-
-            // Allocate and load rodata if needed
-            if (prelude.scope_flags & MP_SCOPE_FLAG_VIPERRODATA) {
-                size_t size = read_uint(reader, NULL);
-                uint8_t *rodata = m_new(uint8_t, size);
-                read_bytes(reader, rodata, size);
-                *ct++ = (uintptr_t)rodata;
+        if (rodata_size + bss_size != 0) {
+            bss_size = (uintptr_t)MP_ALIGN(bss_size, sizeof(uintptr_t));
+            uint8_t *data = m_new0(uint8_t, bss_size + rodata_size);
+            bss = data;
+            rodata = bss + bss_size;
+            if (native_scope_flags & MP_SCOPE_FLAG_VIPERRODATA) {
+                read_bytes(reader, rodata, rodata_size);
             }
 
-            // Allocate BSS if needed
-            if (prelude.scope_flags & MP_SCOPE_FLAG_VIPERBSS) {
-                size_t size = read_uint(reader, NULL);
-                uint8_t *bss = m_new0(uint8_t, size);
-                *ct++ = (uintptr_t)bss;
-            }
+            // Viper code with BSS/rodata should not have any children.
+            // Reuse the children pointer to reference the BSS/rodata
+            // memory so that it is not reclaimed by the GC.
+            assert(!has_children);
+            children = (void *)data;
         }
-        #endif
+    }
+    #endif
 
-        // Load constant objects and raw code children
-        for (size_t i = 0; i < n_obj; ++i) {
-            *ct++ = (mp_uint_t)load_obj(reader);
-        }
-        for (size_t i = 0; i < n_raw_code; ++i) {
-            *ct++ = (mp_uint_t)(uintptr_t)load_raw_code(reader, qw);
+    // Load children if any.
+    if (has_children) {
+        n_children = read_uint(reader);
+        children = m_new(mp_raw_code_t *, n_children);
+        for (size_t i = 0; i < n_children; ++i) {
+            children[i] = load_raw_code(reader);
         }
     }
 
     // Create raw_code and return it
     mp_raw_code_t *rc = mp_emit_glue_new_raw_code();
     if (kind == MP_CODE_BYTECODE) {
+        const byte *ip = fun_data;
+        MP_BC_PRELUDE_SIG_DECODE(ip);
         // Assign bytecode to raw code object
         mp_emit_glue_assign_bytecode(rc, fun_data,
             #if MICROPY_PERSISTENT_CODE_SAVE || MICROPY_DEBUG_PRINTERS
             fun_data_len,
             #endif
-            const_table,
+            children,
             #if MICROPY_PERSISTENT_CODE_SAVE
-            n_obj, n_raw_code,
+            n_children,
             #endif
-            prelude.scope_flags);
+            scope_flags);
 
     #if MICROPY_EMIT_MACHINE_CODE
     } else {
         // Relocate and commit code to executable address space
-        reloc_info_t ri = {reader, const_table};
+        reloc_info_t ri = {reader, rodata, bss};
         #if defined(MP_PLAT_COMMIT_EXEC)
-        void *opt_ri = (prelude.scope_flags & MP_SCOPE_FLAG_VIPERRELOC) ? &ri : NULL;
+        void *opt_ri = (native_scope_flags & MP_SCOPE_FLAG_VIPERRELOC) ? &ri : NULL;
         fun_data = MP_PLAT_COMMIT_EXEC(fun_data, fun_data_len, opt_ri);
         #else
-        if (prelude.scope_flags & MP_SCOPE_FLAG_VIPERRELOC) {
+        if (native_scope_flags & MP_SCOPE_FLAG_VIPERRELOC) {
             #if MICROPY_PERSISTENT_CODE_TRACK_RELOC_CODE
             // If native code needs relocations then it's not guaranteed that a pointer to
             // the head of `buf` (containing the machine code) will be retained for the GC
@@ -534,26 +400,27 @@ STATIC mp_raw_code_t *load_raw_code(mp_reader_t *reader, qstr_window_t *qw) {
 
         // Assign native code to raw code object
         mp_emit_glue_assign_native(rc, kind,
-            fun_data, fun_data_len, const_table,
+            fun_data, fun_data_len,
+            children,
             #if MICROPY_PERSISTENT_CODE_SAVE
+            n_children,
             prelude_offset,
-            n_obj, n_raw_code,
-            n_qstr_link, NULL,
+            0, NULL,
             #endif
-            prelude.n_pos_args, prelude.scope_flags, type_sig);
+            native_scope_flags, native_n_pos_args, native_type_sig
+            );
     #endif
     }
     return rc;
 }
 
-mp_raw_code_t *mp_raw_code_load(mp_reader_t *reader) {
+mp_compiled_module_t mp_raw_code_load(mp_reader_t *reader, mp_module_context_t *context) {
     byte header[4];
     read_bytes(reader, header, sizeof(header));
     if (header[0] != 'M'
         || header[1] != MPY_VERSION
         || MPY_FEATURE_DECODE_FLAGS(header[2]) != MPY_FEATURE_FLAGS
-        || header[3] > mp_small_int_bits()
-        || read_uint(reader, NULL) > QSTR_WINDOW_SIZE) {
+        || header[3] > mp_small_int_bits()) {
         mp_raise_ValueError(MP_ERROR_TEXT("incompatible .mpy file"));
     }
     if (MPY_FEATURE_DECODE_ARCH(header[2]) != MP_NATIVE_ARCH_NONE) {
@@ -562,25 +429,49 @@ mp_raw_code_t *mp_raw_code_load(mp_reader_t *reader) {
             mp_raise_ValueError(MP_ERROR_TEXT("incompatible .mpy arch"));
         }
     }
-    qstr_window_t qw;
-    qw.idx = 0;
-    mp_raw_code_t *rc = load_raw_code(reader, &qw);
+
+    size_t n_qstr = read_uint(reader);
+    size_t n_obj = read_uint(reader);
+    mp_module_context_alloc_tables(context, n_qstr, n_obj);
+
+    // Load qstrs.
+    for (size_t i = 0; i < n_qstr; ++i) {
+        context->constants.qstr_table[i] = load_qstr(reader);
+    }
+
+    // Load constant objects.
+    for (size_t i = 0; i < n_obj; ++i) {
+        context->constants.obj_table[i] = load_obj(reader);
+    }
+
+    // Load top-level module.
+    mp_compiled_module_t cm2;
+    cm2.rc = load_raw_code(reader);
+    cm2.context = context;
+
+    #if MICROPY_PERSISTENT_CODE_SAVE
+    cm2.has_native = MPY_FEATURE_DECODE_ARCH(header[2]) != MP_NATIVE_ARCH_NONE;
+    cm2.n_qstr = n_qstr;
+    cm2.n_obj = n_obj;
+    #endif
+
     reader->close(reader->data);
-    return rc;
+
+    return cm2;
 }
 
-mp_raw_code_t *mp_raw_code_load_mem(const byte *buf, size_t len) {
+mp_compiled_module_t mp_raw_code_load_mem(const byte *buf, size_t len, mp_module_context_t *context) {
     mp_reader_t reader;
     mp_reader_new_mem(&reader, buf, len, 0);
-    return mp_raw_code_load(&reader);
+    return mp_raw_code_load(&reader, context);
 }
 
 #if MICROPY_HAS_FILE_READER
 
-mp_raw_code_t *mp_raw_code_load_file(const char *filename) {
+mp_compiled_module_t mp_raw_code_load_file(const char *filename, mp_module_context_t *context) {
     mp_reader_t reader;
     mp_reader_new_file(&reader, filename);
-    return mp_raw_code_load(&reader);
+    return mp_raw_code_load(&reader, context);
 }
 
 #endif // MICROPY_HAS_FILE_READER
@@ -607,26 +498,25 @@ STATIC void mp_print_uint(mp_print_t *print, size_t n) {
     print->print_strn(print->data, (char *)p, buf + sizeof(buf) - p);
 }
 
-STATIC void save_qstr(mp_print_t *print, qstr_window_t *qw, qstr qst) {
+STATIC void save_qstr(mp_print_t *print, qstr qst) {
     if (qst <= QSTR_LAST_STATIC) {
         // encode static qstr
-        byte buf[2] = {0, qst & 0xff};
-        mp_print_bytes(print, buf, 2);
-        return;
-    }
-    size_t idx = qstr_window_insert(qw, qst);
-    if (idx < QSTR_WINDOW_SIZE) {
-        // qstr found in window, encode index to it
-        mp_print_uint(print, idx << 1 | 1);
+        mp_print_uint(print, qst << 1 | 1);
         return;
     }
     size_t len;
     const byte *str = qstr_data(qst, &len);
     mp_print_uint(print, len << 1);
-    mp_print_bytes(print, str, len);
+    mp_print_bytes(print, str, len + 1); // +1 to store null terminator
 }
 
 STATIC void save_obj(mp_print_t *print, mp_obj_t o) {
+    #if MICROPY_EMIT_MACHINE_CODE
+    if (o == MP_OBJ_FROM_PTR(&mp_fun_table)) {
+        byte obj_type = 't';
+        mp_print_bytes(print, &obj_type, 1);
+    } else
+    #endif
     if (mp_obj_is_str_or_bytes(o)) {
         byte obj_type;
         if (mp_obj_is_str(o)) {
@@ -638,7 +528,7 @@ STATIC void save_obj(mp_print_t *print, mp_obj_t o) {
         const char *str = mp_obj_str_get_data(o, &len);
         mp_print_bytes(print, &obj_type, 1);
         mp_print_uint(print, len);
-        mp_print_bytes(print, (const byte *)str, len);
+        mp_print_bytes(print, (const byte *)str, len + 1); // +1 to store null terminator
     } else if (MP_OBJ_TO_PTR(o) == &mp_const_ellipsis_obj) {
         byte obj_type = 'e';
         mp_print_bytes(print, &obj_type, 1);
@@ -667,136 +557,45 @@ STATIC void save_obj(mp_print_t *print, mp_obj_t o) {
     }
 }
 
-STATIC void save_prelude_qstrs(mp_print_t *print, qstr_window_t *qw, const byte *ip) {
-    save_qstr(print, qw, ip[0] | (ip[1] << 8)); // simple_name
-    save_qstr(print, qw, ip[2] | (ip[3] << 8)); // source_file
-}
-
-STATIC void save_bytecode(mp_print_t *print, qstr_window_t *qw, const byte *ip, const byte *ip_top) {
-    while (ip < ip_top) {
-        size_t sz;
-        uint f = mp_opcode_format(ip, &sz, true);
-        if (f == MP_BC_FORMAT_QSTR) {
-            mp_print_bytes(print, ip, 1);
-            qstr qst = ip[1] | (ip[2] << 8);
-            save_qstr(print, qw, qst);
-            ip += 3;
-            sz -= 3;
-        }
-        mp_print_bytes(print, ip, sz);
-        ip += sz;
-    }
-}
-
-STATIC void save_raw_code(mp_print_t *print, mp_raw_code_t *rc, qstr_window_t *qstr_window) {
+STATIC void save_raw_code(mp_print_t *print, const mp_raw_code_t *rc) {
     // Save function kind and data length
-    mp_print_uint(print, (rc->fun_data_len << 2) | (rc->kind - MP_CODE_BYTECODE));
+    mp_print_uint(print, (rc->fun_data_len << 3) | ((rc->n_children != 0) << 2) | (rc->kind - MP_CODE_BYTECODE));
 
-    bytecode_prelude_t prelude;
+    // Save function code.
+    mp_print_bytes(print, rc->fun_data, rc->fun_data_len);
 
-    if (rc->kind == MP_CODE_BYTECODE) {
-        // Extract prelude
-        const byte *ip = rc->fun_data;
-        const byte *ip_info = extract_prelude(&ip, &prelude);
-
-        // Save prelude
-        mp_print_bytes(print, rc->fun_data, ip_info - (const byte *)rc->fun_data);
-        save_prelude_qstrs(print, qstr_window, ip_info);
-        ip_info += 4;
-        mp_print_bytes(print, ip_info, ip - ip_info);
-
-        // Save bytecode
-        const byte *ip_top = (const byte *)rc->fun_data + rc->fun_data_len;
-        save_bytecode(print, qstr_window, ip, ip_top);
     #if MICROPY_EMIT_MACHINE_CODE
-    } else {
-        // Save native code
-        mp_print_bytes(print, rc->fun_data, rc->fun_data_len);
-
-        if (rc->kind == MP_CODE_NATIVE_PY || rc->kind == MP_CODE_NATIVE_VIPER) {
-            // Save qstr link table for native code
-            mp_print_uint(print, rc->n_qstr);
-            for (size_t i = 0; i < rc->n_qstr; ++i) {
-                mp_print_uint(print, rc->qstr_link[i].off);
-                save_qstr(print, qstr_window, rc->qstr_link[i].qst);
-            }
+    if (rc->kind == MP_CODE_NATIVE_PY || rc->kind == MP_CODE_NATIVE_VIPER) {
+        // Save qstr link table for native code
+        mp_print_uint(print, rc->n_qstr);
+        for (size_t i = 0; i < rc->n_qstr; ++i) {
+            mp_print_uint(print, rc->qstr_link[i].off);
+            save_qstr(print, rc->qstr_link[i].qst);
         }
-
-        if (rc->kind == MP_CODE_NATIVE_PY) {
-            // Save prelude size
-            mp_print_uint(print, rc->prelude_offset);
-
-            // Extract prelude and save qstrs in prelude
-            const byte *ip = (const byte *)rc->fun_data + rc->prelude_offset;
-            const byte *ip_info = extract_prelude(&ip, &prelude);
-            save_prelude_qstrs(print, qstr_window, ip_info);
-        } else {
-            // Save basic scope info for viper and asm
-            mp_print_uint(print, rc->scope_flags & MP_SCOPE_FLAG_ALL_SIG);
-            prelude.n_pos_args = 0;
-            prelude.n_kwonly_args = 0;
-            if (rc->kind == MP_CODE_NATIVE_ASM) {
-                mp_print_uint(print, rc->n_pos_args);
-                mp_print_uint(print, rc->type_sig);
-            }
-        }
-    #endif
     }
 
-    if (rc->kind != MP_CODE_NATIVE_ASM) {
-        // Save constant table for bytecode, native and viper
-
-        // Number of entries in constant table
-        mp_print_uint(print, rc->n_obj);
-        mp_print_uint(print, rc->n_raw_code);
-
-        const mp_uint_t *const_table = rc->const_table;
-
-        // Save function argument names (initial entries in const_table)
-        // (viper has n_pos_args=n_kwonly_args=0 so doesn't save any qstrs here)
-        for (size_t i = 0; i < prelude.n_pos_args + prelude.n_kwonly_args; ++i) {
-            mp_obj_t o = (mp_obj_t)*const_table++;
-            save_qstr(print, qstr_window, MP_OBJ_QSTR_VALUE(o));
-        }
-
-        if (rc->kind != MP_CODE_BYTECODE) {
-            // Skip saving mp_fun_table entry
-            ++const_table;
-        }
-
-        // Save constant objects and raw code children
-        for (size_t i = 0; i < rc->n_obj; ++i) {
-            save_obj(print, (mp_obj_t)*const_table++);
+    if (rc->kind == MP_CODE_NATIVE_PY) {
+        // Save prelude size
+        mp_print_uint(print, rc->prelude_offset);
+    } else if (rc->kind == MP_CODE_NATIVE_VIPER || rc->kind == MP_CODE_NATIVE_ASM) {
+        // Save basic scope info for viper and asm
+        mp_print_uint(print, rc->scope_flags & MP_SCOPE_FLAG_ALL_SIG);
+        if (rc->kind == MP_CODE_NATIVE_ASM) {
+            mp_print_uint(print, rc->n_pos_args);
+            mp_print_uint(print, rc->type_sig);
         }
-        for (size_t i = 0; i < rc->n_raw_code; ++i) {
-            save_raw_code(print, (mp_raw_code_t *)(uintptr_t)*const_table++, qstr_window);
-        }
-    }
-}
-
-STATIC bool mp_raw_code_has_native(mp_raw_code_t *rc) {
-    if (rc->kind != MP_CODE_BYTECODE) {
-        return true;
     }
+    #endif
 
-    const byte *ip = rc->fun_data;
-    bytecode_prelude_t prelude;
-    extract_prelude(&ip, &prelude);
-
-    const mp_uint_t *const_table = rc->const_table
-        + prelude.n_pos_args + prelude.n_kwonly_args
-        + rc->n_obj;
-
-    for (size_t i = 0; i < rc->n_raw_code; ++i) {
-        if (mp_raw_code_has_native((mp_raw_code_t *)(uintptr_t)*const_table++)) {
-            return true;
+    if (rc->n_children) {
+        mp_print_uint(print, rc->n_children);
+        for (size_t i = 0; i < rc->n_children; ++i) {
+            save_raw_code(print, rc->children[i]);
         }
     }
-
-    return false;
 }
 
-void mp_raw_code_save(mp_raw_code_t *rc, mp_print_t *print) {
+void mp_raw_code_save(mp_compiled_module_t *cm, mp_print_t *print) {
     // header contains:
     //  byte  'M'
     //  byte  version
@@ -813,16 +612,27 @@ void mp_raw_code_save(mp_raw_code_t *rc, mp_print_t *print) {
         mp_small_int_bits(),
         #endif
     };
-    if (mp_raw_code_has_native(rc)) {
+    if (cm->has_native) {
         header[2] |= MPY_FEATURE_ENCODE_ARCH(MPY_FEATURE_ARCH_DYNAMIC);
     }
     mp_print_bytes(print, header, sizeof(header));
-    mp_print_uint(print, QSTR_WINDOW_SIZE);
 
-    qstr_window_t qw;
-    qw.idx = 0;
-    memset(qw.window, 0, sizeof(qw.window));
-    save_raw_code(print, rc, &qw);
+    // Number of entries in constant table.
+    mp_print_uint(print, cm->n_qstr);
+    mp_print_uint(print, cm->n_obj);
+
+    // Save qstrs.
+    for (size_t i = 0; i < cm->n_qstr; ++i) {
+        save_qstr(print, cm->context->constants.qstr_table[i]);
+    }
+
+    // Save constant objects.
+    for (size_t i = 0; i < cm->n_obj; ++i) {
+        save_obj(print, (mp_obj_t)cm->context->constants.obj_table[i]);
+    }
+
+    // Save outer raw code, which will save all its child raw codes.
+    save_raw_code(print, cm->rc);
 }
 
 #if MICROPY_PERSISTENT_CODE_SAVE_FILE
@@ -839,12 +649,12 @@ STATIC void fd_print_strn(void *env, const char *str, size_t len) {
     (void)ret;
 }
 
-void mp_raw_code_save_file(mp_raw_code_t *rc, const char *filename) {
+void mp_raw_code_save_file(mp_compiled_module_t *cm, const char *filename) {
     MP_THREAD_GIL_EXIT();
     int fd = open(filename, O_WRONLY | O_CREAT | O_TRUNC, 0644);
     MP_THREAD_GIL_ENTER();
     mp_print_t fd_print = {(void *)(intptr_t)fd, fd_print_strn};
-    mp_raw_code_save(rc, &fd_print);
+    mp_raw_code_save(cm, &fd_print);
     MP_THREAD_GIL_EXIT();
     close(fd);
     MP_THREAD_GIL_ENTER();