py/persistentcode: Explicitly track native BSS/rodata when needed.

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

3 files changed, 63 insertions, 40 deletions

diff --git a/py/mpconfig.h b/py/mpconfig.h
index b5414312c7..34eafa9e5d 100644
--- a/py/mpconfig.h
+++ b/py/mpconfig.h
@@ -425,18 +425,6 @@
 // Convenience definition for whether any native or inline assembler emitter is enabled
 #define MICROPY_EMIT_MACHINE_CODE (MICROPY_EMIT_NATIVE || MICROPY_EMIT_INLINE_ASM)
 
-// Whether native relocatable code loaded from .mpy files is explicitly tracked
-// so that the GC cannot reclaim it.  Needed on architectures that allocate
-// executable memory on the MicroPython heap and don't explicitly track this
-// data some other way.
-#ifndef MICROPY_PERSISTENT_CODE_TRACK_RELOC_CODE
-#if !MICROPY_EMIT_MACHINE_CODE || defined(MP_PLAT_ALLOC_EXEC) || defined(MP_PLAT_COMMIT_EXEC)
-#define MICROPY_PERSISTENT_CODE_TRACK_RELOC_CODE (0)
-#else
-#define MICROPY_PERSISTENT_CODE_TRACK_RELOC_CODE (1)
-#endif
-#endif
-
 /*****************************************************************************/
 /* Compiler configuration                                                    */
 
@@ -1992,14 +1980,48 @@ typedef double mp_float_t;
 #define MICROPY_MAKE_POINTER_CALLABLE(p) (p)
 #endif
 
-// If these MP_PLAT_*_EXEC macros are overridden then the memory allocated by them
-// must be somehow reachable for marking by the GC, since the native code
-// generators store pointers to GC managed memory in the code.
+// Whether native text/BSS/rodata memory loaded from .mpy files is explicitly tracked
+// so that the GC cannot reclaim it.
+//
+// In general a port should let these options have their defaults, but the defaults here
+// can be overridden if needed by defining both MICROPY_PERSISTENT_CODE_TRACK_FUN_DATA
+// and MICROPY_PERSISTENT_CODE_TRACK_BSS_RODATA.
+#ifndef MICROPY_PERSISTENT_CODE_TRACK_FUN_DATA
+#if MICROPY_EMIT_MACHINE_CODE && MICROPY_PERSISTENT_CODE_LOAD
+// Pointer tracking is required when loading native code is enabled.
+#if defined(MP_PLAT_ALLOC_EXEC) || defined(MP_PLAT_COMMIT_EXEC)
+// If a port defined a custom allocator or commit function for native text, then the
+// text does not need to be tracked (its allocation is managed by the port).  But the
+// BSS/rodata must be tracked (if there is any) because if there are any pointers to it
+// in the function data, they aren't traced by the GC.
+#define MICROPY_PERSISTENT_CODE_TRACK_FUN_DATA (0)
+#define MICROPY_PERSISTENT_CODE_TRACK_BSS_RODATA (1)
+#else
+// If a port uses the default allocator (the GC heap) then all native text is allocated
+// on the GC heap.  But it's not guaranteed that a pointer to the head of the block of
+// native text (which may contain multiple native functions) will be retained for the GC
+// to trace.  This is because native functions can start inside the big block of text
+// and so it's possible that the only GC-reachable pointers are pointers inside.
+// Therefore the big block is explicitly tracked. If there is any BSS/rodata memory,
+// then it does not need to be explicitly tracked because a pointer to it is stored into
+// the function text via `mp_native_relocate()`.
+#define MICROPY_PERSISTENT_CODE_TRACK_FUN_DATA (1)
+#define MICROPY_PERSISTENT_CODE_TRACK_BSS_RODATA (0)
+#endif
+#else // MICROPY_EMIT_MACHINE_CODE && MICROPY_PERSISTENT_CODE_LOAD
+// Pointer tracking not needed when loading native code is disabled.
+#define MICROPY_PERSISTENT_CODE_TRACK_FUN_DATA (0)
+#define MICROPY_PERSISTENT_CODE_TRACK_BSS_RODATA (0)
+#endif
+#endif
+
+// If these macros are defined then the memory allocated by them does not need to be
+// traced by the GC.  But if they are left undefined then the GC heap will be used as
+// the allocator and the memory must be traced by the GC.  See also above logic for
+// enabling MICROPY_PERSISTENT_CODE_TRACK_FUN_DATA and
+// MICROPY_PERSISTENT_CODE_TRACK_BSS_RODATA.
 #ifndef MP_PLAT_ALLOC_EXEC
 #define MP_PLAT_ALLOC_EXEC(min_size, ptr, size) do { *ptr = m_new(byte, min_size); *size = min_size; } while (0)
-#endif
-
-#ifndef MP_PLAT_FREE_EXEC
 #define MP_PLAT_FREE_EXEC(ptr, size) m_del(byte, ptr, size)
 #endif
 
diff --git a/py/persistentcode.c b/py/persistentcode.c
index 5088c05f03..be93eaa5b4 100644
--- a/py/persistentcode.c
+++ b/py/persistentcode.c
@@ -72,6 +72,20 @@ typedef struct _bytecode_prelude_t {
 static int read_byte(mp_reader_t *reader);
 static size_t read_uint(mp_reader_t *reader);
 
+#if MICROPY_PERSISTENT_CODE_TRACK_FUN_DATA || MICROPY_PERSISTENT_CODE_TRACK_BSS_RODATA
+
+// An mp_obj_list_t that tracks native text/BSS/rodata to prevent the GC from reclaiming them.
+MP_REGISTER_ROOT_POINTER(mp_obj_t persistent_code_root_pointers);
+
+static void track_root_pointer(void *ptr) {
+    if (MP_STATE_PORT(persistent_code_root_pointers) == MP_OBJ_NULL) {
+        MP_STATE_PORT(persistent_code_root_pointers) = mp_obj_new_list(0, NULL);
+    }
+    mp_obj_list_append(MP_STATE_PORT(persistent_code_root_pointers), MP_OBJ_FROM_PTR(ptr));
+}
+
+#endif
+
 #if MICROPY_EMIT_MACHINE_CODE
 
 typedef struct _reloc_info_t {
@@ -299,11 +313,10 @@ static mp_raw_code_t *load_raw_code(mp_reader_t *reader, mp_module_context_t *co
                 read_bytes(reader, rodata, rodata_size);
             }
 
-            // 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;
+            #if MICROPY_PERSISTENT_CODE_TRACK_BSS_RODATA
+            // Track the BSS/rodata memory so it's not reclaimed by the GC.
+            track_root_pointer(data);
+            #endif
         }
     }
     #endif
@@ -351,16 +364,9 @@ static mp_raw_code_t *load_raw_code(mp_reader_t *reader, mp_module_context_t *co
         fun_data = MP_PLAT_COMMIT_EXEC(fun_data, fun_data_len, opt_ri);
         #else
         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
-            // to trace.  This is because native functions can start inside `buf` and so
-            // it's possible that the only GC-reachable pointers are pointers inside `buf`.
-            // So put this `buf` on a list of reachable root pointers.
-            if (MP_STATE_PORT(track_reloc_code_list) == MP_OBJ_NULL) {
-                MP_STATE_PORT(track_reloc_code_list) = mp_obj_new_list(0, NULL);
-            }
-            mp_obj_list_append(MP_STATE_PORT(track_reloc_code_list), MP_OBJ_FROM_PTR(fun_data));
+            #if MICROPY_PERSISTENT_CODE_TRACK_FUN_DATA
+            // Track the function data memory so it's not reclaimed by the GC.
+            track_root_pointer(fun_data);
             #endif
             // Do the relocations.
             mp_native_relocate(&ri, fun_data, (uintptr_t)fun_data);
@@ -662,8 +668,3 @@ void mp_raw_code_save_file(mp_compiled_module_t *cm, qstr filename) {
 #endif // MICROPY_PERSISTENT_CODE_SAVE_FILE
 
 #endif // MICROPY_PERSISTENT_CODE_SAVE
-
-#if MICROPY_PERSISTENT_CODE_TRACK_RELOC_CODE
-// An mp_obj_list_t that tracks relocated native code to prevent the GC from reclaiming them.
-MP_REGISTER_ROOT_POINTER(mp_obj_t track_reloc_code_list);
-#endif
diff --git a/py/runtime.c b/py/runtime.c
index acb45c94b0..deb55bf283 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -119,8 +119,8 @@ void mp_init(void) {
     MP_STATE_VM(mp_module_builtins_override_dict) = NULL;
     #endif
 
-    #if MICROPY_PERSISTENT_CODE_TRACK_RELOC_CODE
-    MP_STATE_VM(track_reloc_code_list) = MP_OBJ_NULL;
+    #if MICROPY_PERSISTENT_CODE_TRACK_FUN_DATA || MICROPY_PERSISTENT_CODE_TRACK_BSS_RODATA
+    MP_STATE_VM(persistent_code_root_pointers) = MP_OBJ_NULL;
     #endif
 
     #if MICROPY_PY_OS_DUPTERM