py/persistentcode: Add support for loading .mpy files from a ROM reader.

This adds an optimisation for loading .mpy files from a reader that points
to ROM.  In such a case qstr, str and bytes data, along with bytecode, are
all referenced in-place in ROM.

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

1 files changed, 63 insertions, 7 deletions

diff --git a/py/persistentcode.c b/py/persistentcode.c
index be93eaa5b..0843d1a2c 100644
--- a/py/persistentcode.c
+++ b/py/persistentcode.c
@@ -181,6 +181,15 @@ static qstr load_qstr(mp_reader_t *reader) {
         return len >> 1;
     }
     len >>= 1;
+
+    #if MICROPY_VFS_ROM
+    // If possible, create the qstr from the memory-mapped string data.
+    const uint8_t *memmap = mp_reader_try_read_rom(reader, len + 1);
+    if (memmap != NULL) {
+        return qstr_from_strn_static((const char *)memmap, len);
+    }
+    #endif
+
     char *str = m_new(char, len);
     read_bytes(reader, (byte *)str, len);
     read_byte(reader); // read and discard null terminator
@@ -189,6 +198,24 @@ static qstr load_qstr(mp_reader_t *reader) {
     return qst;
 }
 
+#if MICROPY_VFS_ROM
+// Create a str/bytes object that can forever reference the given data.
+static mp_obj_t mp_obj_new_str_static(const mp_obj_type_t *type, const byte *data, size_t len) {
+    if (type == &mp_type_str) {
+        qstr q = qstr_find_strn((const char *)data, len);
+        if (q != MP_QSTRnull) {
+            return MP_OBJ_NEW_QSTR(q);
+        }
+    }
+    assert(data[len] == '\0');
+    mp_obj_str_t *o = mp_obj_malloc(mp_obj_str_t, type);
+    o->len = len;
+    o->hash = qstr_compute_hash(data, len);
+    o->data = data;
+    return MP_OBJ_FROM_PTR(o);
+}
+#endif
+
 static mp_obj_t load_obj(mp_reader_t *reader) {
     byte obj_type = read_byte(reader);
     #if MICROPY_EMIT_MACHINE_CODE
@@ -206,6 +233,8 @@ static mp_obj_t load_obj(mp_reader_t *reader) {
         return MP_OBJ_FROM_PTR(&mp_const_ellipsis_obj);
     } else {
         size_t len = read_uint(reader);
+
+        // Handle empty bytes object, and tuple objects.
         if (len == 0 && obj_type == MP_PERSISTENT_OBJ_BYTES) {
             read_byte(reader); // skip null terminator
             return mp_const_empty_bytes;
@@ -216,11 +245,31 @@ static mp_obj_t load_obj(mp_reader_t *reader) {
             }
             return MP_OBJ_FROM_PTR(tuple);
         }
+
+        // Read in the object's data, either from ROM or into RAM.
+        const uint8_t *memmap = NULL;
         vstr_t vstr;
-        vstr_init_len(&vstr, len);
-        read_bytes(reader, (byte *)vstr.buf, len);
+        #if MICROPY_VFS_ROM
+        memmap = mp_reader_try_read_rom(reader, len);
+        vstr.buf = (void *)memmap;
+        vstr.len = len;
+        #endif
+        if (memmap == NULL) {
+            // Data could not be memory-mapped, so allocate it in RAM and read it in.
+            vstr_init_len(&vstr, len);
+            read_bytes(reader, (byte *)vstr.buf, len);
+        }
+
+        // Create and return the object.
         if (obj_type == MP_PERSISTENT_OBJ_STR || obj_type == MP_PERSISTENT_OBJ_BYTES) {
-            read_byte(reader); // skip null terminator
+            read_byte(reader); // skip null terminator (it needs to be there for ROM str objects)
+            #if MICROPY_VFS_ROM
+            if (memmap != NULL) {
+                // Create a str/bytes that references the memory-mapped data.
+                const mp_obj_type_t *t = obj_type == MP_PERSISTENT_OBJ_STR ? &mp_type_str : &mp_type_bytes;
+                return mp_obj_new_str_static(t, memmap, len);
+            }
+            #endif
             if (obj_type == MP_PERSISTENT_OBJ_STR) {
                 return mp_obj_new_str_from_utf8_vstr(&vstr);
             } else {
@@ -257,10 +306,17 @@ static mp_raw_code_t *load_raw_code(mp_reader_t *reader, mp_module_context_t *co
     #endif
 
     if (kind == MP_CODE_BYTECODE) {
-        // Allocate memory for the bytecode
-        fun_data = m_new(uint8_t, fun_data_len);
-        // Load bytecode
-        read_bytes(reader, fun_data, fun_data_len);
+        #if MICROPY_VFS_ROM
+        // Try to reference memory-mapped data for the bytecode.
+        fun_data = (uint8_t *)mp_reader_try_read_rom(reader, fun_data_len);
+        #endif
+
+        if (fun_data == NULL) {
+            // Allocate memory for the bytecode.
+            fun_data = m_new(uint8_t, fun_data_len);
+            // Load bytecode.
+            read_bytes(reader, fun_data, fun_data_len);
+        }
 
     #if MICROPY_EMIT_MACHINE_CODE
     } else {