Add supervisor.set_next_code() function (prototype).

Fixes #1084.

Author: cwalther

main.c

@@ -81,6 +81,13 @@
 #include "supervisor/shared/bluetooth.h"
 #endif
 
+typedef struct {
+    uint8_t options;
+    char filename[];
+} next_code_info_t;
+
+static supervisor_allocation* next_code_allocation;
+
 void do_str(const char *src, mp_parse_input_kind_t input_kind) {
     mp_lexer_t *lex = mp_lexer_new_from_str_len(MP_QSTR__lt_stdin_gt_, src, strlen(src), 0);
     if (lex == NULL) {
@@ -211,7 +218,21 @@ bool maybe_run_list(const char * const * filenames, pyexec_result_t* exec_result
     return true;
 }
 
-void cleanup_after_vm(supervisor_allocation* heap) {
+bool cleanup_after_vm(supervisor_allocation* heap) {
+    // Get any info on what code to run next off the heap.
+    next_code_info_t* next_code = NULL;
+    size_t next_code_len = 0;
+    if (MP_STATE_VM(supervisor_next_code)) {
+        size_t len;
+        mp_obj_t* items;
+        mp_obj_tuple_get(MP_STATE_VM(supervisor_next_code), &len, &items);
+        const char* filename = mp_obj_str_get_data(items[0], &len);
+        next_code_len = sizeof(uint8_t) + len; // termination not needed here, saves a byte of stack
+        next_code = alloca(next_code_len);
+        next_code->options = (uint8_t)MP_OBJ_SMALL_INT_VALUE(items[1]);
+        memcpy(&next_code->filename, filename, len);
+        MP_STATE_VM(supervisor_next_code) = NULL;
+    }
     // Reset port-independent devices, like CIRCUITPY_BLEIO_HCI.
     reset_devices();
     // Turn off the display and flush the fileystem before the heap disappears.
@@ -232,6 +253,21 @@ void cleanup_after_vm(supervisor_allocation* heap) {
     #endif
     reset_board();
     reset_status_led();
+
+    // Save next code info.
+    if (next_code) {
+        if (next_code_allocation) {
+            free_memory(next_code_allocation);
+        }
+        // TODO do I need to increase CIRCUITPY_SUPERVISOR_ALLOC_COUNT for this?
+        next_code_allocation = allocate_memory((next_code_len + 1 + 3) & ~3, false);
+        if (next_code_allocation) {
+            memcpy(next_code_allocation->ptr, next_code, next_code_len);
+            ((char*)next_code_allocation->ptr)[next_code_len] = '\0';
+            return true;
+        }
+    }
+    return false;
 }
 
 bool run_code_py(safe_mode_t safe_mode) {
@@ -274,17 +310,42 @@ bool run_code_py(safe_mode_t safe_mode) {
         filesystem_flush();
         supervisor_allocation* heap = allocate_remaining_memory();
         start_mp(heap);
-        found_main = maybe_run_list(supported_filenames, &result);
-        #if CIRCUITPY_FULL_BUILD
-        if (!found_main){
-            found_main = maybe_run_list(double_extension_filenames, &result);
-            if (found_main) {
-                serial_write_compressed(translate("WARNING: Your code filename has two extensions\n"));
+        uint8_t next_code_options = 0;
+        if (next_code_allocation) {
+            next_code_options = ((next_code_info_t*)next_code_allocation->ptr)->options;
+            const char* next_list[] = {((next_code_info_t*)next_code_allocation->ptr)->filename, ""};
+            found_main = maybe_run_list(next_list, &result);
+            if (!found_main) {
+                serial_write(((next_code_info_t*)next_code_allocation->ptr)->filename);
+                serial_write_compressed(translate(" not found.\n"));
             }
         }
-        #endif
-        cleanup_after_vm(heap);
+        if (!found_main) {
+            found_main = maybe_run_list(supported_filenames, &result);
+            #if CIRCUITPY_FULL_BUILD
+            if (!found_main){
+                found_main = maybe_run_list(double_extension_filenames, &result);
+                if (found_main) {
+                    serial_write_compressed(translate("WARNING: Your code filename has two extensions\n"));
+                }
+            }
+            #endif
+        }
+        bool new_next_code = cleanup_after_vm(heap);
+        // - If a new next code file was set, honor that in any case.
+        // - Otherwise: If a reload was requested by a USB file write, we want to run the same file
+        //   again, preserve any next-code info. Currently that also covers a reload requested by
+        //   supervisor.reload(), or do we want to treat that differently by turning
+        //   reload_requested into an enum?
+        // - Otherwise: Clear any leftover next-code info and restart at code.py/main.py.
+        if (!new_next_code && !reload_requested) {
+            free_memory(next_code_allocation);
+            next_code_allocation = NULL;
+        }
 
+        if (result.return_code == 0 && (next_code_options & SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_SUCCESS)) {
+            return true;
+        }
         if (result.return_code & PYEXEC_FORCED_EXIT) {
             return reload_requested;
         }

py/circuitpy_mpconfig.h

@@ -849,6 +849,7 @@ extern const struct _mp_obj_module_t wifi_module;
     FLASH_ROOT_POINTERS \
     MEMORYMONITOR_ROOT_POINTERS \
     NETWORK_ROOT_POINTERS \
+    mp_obj_t supervisor_next_code; \
 
 void supervisor_run_background_tasks_if_tick(void);
 #define RUN_BACKGROUND_TASKS (supervisor_run_background_tasks_if_tick())

shared-bindings/supervisor/__init__.c

@@ -26,6 +26,7 @@
 #include "py/obj.h"
 #include "py/runtime.h"
 #include "py/reload.h"
+#include "py/objstr.h"
 
 #include "lib/utils/interrupt_char.h"
 #include "supervisor/shared/autoreload.h"
@@ -109,6 +110,28 @@ STATIC mp_obj_t supervisor_set_next_stack_limit(mp_obj_t size_obj) {
 }
 MP_DEFINE_CONST_FUN_OBJ_1(supervisor_set_next_stack_limit_obj, supervisor_set_next_stack_limit);
 
+STATIC mp_obj_t supervisor_set_next_code_file(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
+    static const mp_arg_t allowed_args[] = {
+        { MP_QSTR_filename, MP_ARG_REQUIRED | MP_ARG_OBJ, {.u_rom_obj = mp_const_none} },
+        { MP_QSTR_reload_on_success, MP_ARG_KW_ONLY | MP_ARG_BOOL, {.u_bool = false} },
+    };
+    struct {
+        mp_arg_val_t filename;
+        mp_arg_val_t reload_on_success;
+    } args;
+    mp_arg_parse_all(n_args, pos_args, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, (mp_arg_val_t*)&args);
+    if (!MP_OBJ_IS_STR_OR_BYTES(args.filename.u_obj)) {
+        mp_raise_TypeError(translate("argument has wrong type"));
+    }
+    mp_obj_t items[] = {
+        args.filename.u_obj,
+        args.reload_on_success.u_bool ? MP_OBJ_NEW_SMALL_INT(SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_SUCCESS) : MP_OBJ_NEW_SMALL_INT(0)
+    };
+    MP_STATE_VM(supervisor_next_code) = mp_obj_new_tuple(MP_ARRAY_SIZE(items), items);
+    return mp_const_none;
+}
+MP_DEFINE_CONST_FUN_OBJ_KW(supervisor_set_next_code_file_obj, 1, supervisor_set_next_code_file);
+
 STATIC const mp_rom_map_elem_t supervisor_module_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_supervisor) },
     { MP_OBJ_NEW_QSTR(MP_QSTR_enable_autoreload),  MP_ROM_PTR(&supervisor_enable_autoreload_obj) },
@@ -117,7 +140,7 @@ STATIC const mp_rom_map_elem_t supervisor_module_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR_runtime),  MP_ROM_PTR(&common_hal_supervisor_runtime_obj) },
     { MP_ROM_QSTR(MP_QSTR_reload),  MP_ROM_PTR(&supervisor_reload_obj) },
     { MP_ROM_QSTR(MP_QSTR_set_next_stack_limit),  MP_ROM_PTR(&supervisor_set_next_stack_limit_obj) },
-
+    { MP_ROM_QSTR(MP_QSTR_set_next_code_file),  MP_ROM_PTR(&supervisor_set_next_code_file_obj) },
 };
 
 STATIC MP_DEFINE_CONST_DICT(supervisor_module_globals, supervisor_module_globals_table);

supervisor/shared/autoreload.h

@@ -29,6 +29,10 @@
 
 #include <stdbool.h>
 
+enum {
+	SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_SUCCESS = 0x1
+};
+
 extern volatile bool reload_requested;
 
 void autoreload_tick(void);