py: Allow to stat and import frozen mpy files using new frozen "VFS".

Freezing mpy files using mpy-tool.py now works again.

2 files changed, 73 insertions, 38 deletions

diff --git a/py/builtinimport.c b/py/builtinimport.c
index dc9bf75f07..d3670858e5 100644
--- a/py/builtinimport.c
+++ b/py/builtinimport.c
@@ -63,7 +63,7 @@ bool mp_obj_is_package(mp_obj_t module) {
 // Stat either frozen or normal module by a given path
 // (whatever is available, if at all).
 STATIC mp_import_stat_t mp_import_stat_any(const char *path) {
-    #if MICROPY_MODULE_FROZEN_STR
+    #if MICROPY_MODULE_FROZEN
     mp_import_stat_t st = mp_frozen_stat(path);
     if (st != MP_IMPORT_STAT_NO_EXIST) {
         return st;
@@ -194,10 +194,37 @@ STATIC void do_execute_raw_code(mp_obj_t module_obj, mp_raw_code_t *raw_code) {
 #endif
 
 STATIC void do_load(mp_obj_t module_obj, vstr_t *file) {
-    #if MICROPY_PERSISTENT_CODE_LOAD || MICROPY_ENABLE_COMPILER
+    #if MICROPY_MODULE_FROZEN || MICROPY_PERSISTENT_CODE_LOAD || MICROPY_ENABLE_COMPILER
     char *file_str = vstr_null_terminated_str(file);
     #endif
 
+    // If we support frozen modules (either as str or mpy) then try to find the
+    // requested filename in the list of frozen module filenames.
+    #if MICROPY_MODULE_FROZEN
+    void *modref;
+    int frozen_type = mp_find_frozen_module(file_str, file->len, &modref);
+    #endif
+
+    // If we support frozen str modules and the compiler is enabled, and we
+    // found the filename in the list of frozen files, then load and execute it.
+    #if MICROPY_MODULE_FROZEN_STR
+    if (frozen_type == MP_FROZEN_STR) {
+        do_load_from_lexer(module_obj, modref, file_str);
+        return;
+    }
+    #endif
+
+    // If we support frozen mpy modules and we found a corresponding file (and
+    // its data) in the list of frozen files, execute it.
+    #if MICROPY_MODULE_FROZEN_MPY
+    if (frozen_type == MP_FROZEN_MPY) {
+        do_execute_raw_code(module_obj, modref);
+        return;
+    }
+    #endif
+
+    // If we support loading .mpy files then check if the file extension is of
+    // the correct format and, if so, load and execute the file.
     #if MICROPY_PERSISTENT_CODE_LOAD
     if (file_str[file->len - 3] == 'm') {
         mp_raw_code_t *raw_code = mp_raw_code_load_file(file_str);
@@ -206,29 +233,18 @@ STATIC void do_load(mp_obj_t module_obj, vstr_t *file) {
     }
     #endif
 
+    // If we can compile scripts then load the file and compile and execute it.
     #if MICROPY_ENABLE_COMPILER
     {
-        void *modref;
-        #if MICROPY_MODULE_FROZEN
-        int frozen_type = mp_find_frozen_module(file_str, file->len, &modref);
-        #else
-        int frozen_type = MP_FROZEN_NONE;
-        #endif
-        #if MICROPY_PERSISTENT_CODE_LOAD || MICROPY_MODULE_FROZEN_MPY
-        if (frozen_type == MP_FROZEN_MPY) {
-            do_execute_raw_code(module_obj, modref);
-            return;
-        }
-        #endif
-        if (frozen_type == MP_FROZEN_NONE) {
-            modref = mp_lexer_new_from_file(file_str);
-        }
-        do_load_from_lexer(module_obj, modref, file_str);
+        mp_lexer_t *lex = mp_lexer_new_from_file(file_str);
+        do_load_from_lexer(module_obj, lex, file_str);
+        return;
     }
-    #else
+    #endif
+
+    // If we get here then the file was not frozen and we can't compile scripts.
     nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_ImportError,
         "script compilation not supported"));
-    #endif
 }
 
 STATIC void chop_component(const char *start, const char **end) {
diff --git a/py/frozenmod.c b/py/frozenmod.c
index 0fabb06a98..660167eed4 100644
--- a/py/frozenmod.c
+++ b/py/frozenmod.c
@@ -43,24 +43,6 @@ extern const char mp_frozen_str_names[];
 extern const uint32_t mp_frozen_str_sizes[];
 extern const char mp_frozen_str_content[];
 
-mp_import_stat_t mp_frozen_stat(const char *str) {
-    size_t len = strlen(str);
-    const char *name = mp_frozen_str_names;
-
-    for (int i = 0; *name != 0; i++) {
-        size_t l = strlen(name);
-        if (l >= len && !memcmp(str, name, len)) {
-            if (name[len] == 0) {
-                return MP_IMPORT_STAT_FILE;
-            } else if (name[len] == '/') {
-                return MP_IMPORT_STAT_DIR;
-            }
-        }
-        name += l + 1;
-    }
-    return MP_IMPORT_STAT_NO_EXIST;
-}
-
 STATIC mp_lexer_t *mp_find_frozen_str(const char *str, size_t len) {
     const char *name = mp_frozen_str_names;
 
@@ -103,6 +85,43 @@ STATIC const mp_raw_code_t *mp_find_frozen_mpy(const char *str, size_t len) {
 
 #if MICROPY_MODULE_FROZEN
 
+STATIC mp_import_stat_t mp_frozen_stat_helper(const char *name, const char *str) {
+    size_t len = strlen(str);
+
+    for (int i = 0; *name != 0; i++) {
+        size_t l = strlen(name);
+        if (l >= len && !memcmp(str, name, len)) {
+            if (name[len] == 0) {
+                return MP_IMPORT_STAT_FILE;
+            } else if (name[len] == '/') {
+                return MP_IMPORT_STAT_DIR;
+            }
+        }
+        name += l + 1;
+    }
+    return MP_IMPORT_STAT_NO_EXIST;
+}
+
+mp_import_stat_t mp_frozen_stat(const char *str) {
+    mp_import_stat_t stat;
+
+    #if MICROPY_MODULE_FROZEN_STR
+    stat = mp_frozen_stat_helper(mp_frozen_str_names, str);
+    if (stat != MP_IMPORT_STAT_NO_EXIST) {
+        return stat;
+    }
+    #endif
+
+    #if MICROPY_MODULE_FROZEN_MPY
+    stat = mp_frozen_stat_helper(mp_frozen_mpy_names, str);
+    if (stat != MP_IMPORT_STAT_NO_EXIST) {
+        return stat;
+    }
+    #endif
+
+    return MP_IMPORT_STAT_NO_EXIST;
+}
+
 int mp_find_frozen_module(const char *str, size_t len, void **data) {
     #if MICROPY_MODULE_FROZEN_STR
     mp_lexer_t *lex = mp_find_frozen_str(str, len);