Treat a reload requested during the wait state the same as one requested during execution regarding preserving next code info.

Author: cwalther

main.c

@@ -281,6 +281,19 @@ bool cleanup_after_vm(supervisor_allocation* heap) {
     return false;
 }
 
+void maybe_clear_next_code_info(bool new_next_code) {
+    // - 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;
+    }
+}
+
 bool run_code_py(safe_mode_t safe_mode) {
     bool serial_connected_at_start = serial_connected();
     #if CIRCUITPY_AUTORELOAD_DELAY_MS > 0
@@ -304,6 +317,7 @@ bool run_code_py(safe_mode_t safe_mode) {
     result.exception_line = 0;
 
     bool found_main = false;
+    bool new_next_code = false;
 
     if (safe_mode != NO_SAFE_MODE) {
         serial_write_compressed(translate("Running in safe mode! "));
@@ -342,22 +356,14 @@ bool run_code_py(safe_mode_t safe_mode) {
             }
             #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;
-        }
+        new_next_code = cleanup_after_vm(heap);
 
         if (result.return_code == 0 && (next_code_options & SUPERVISOR_NEXT_CODE_OPT_RELOAD_ON_SUCCESS)) {
+            maybe_clear_next_code_info(new_next_code);
             return true;
         }
         if (result.return_code & PYEXEC_FORCED_EXIT) {
+            maybe_clear_next_code_info(new_next_code);
             return reload_requested;
         }
     }
@@ -376,11 +382,13 @@ bool run_code_py(safe_mode_t safe_mode) {
     while (true) {
         RUN_BACKGROUND_TASKS;
         if (reload_requested) {
+            // no need to call maybe_clear_next_code_info(), it would be a no-op
             reload_requested = false;
             return true;
         }
 
         if (serial_connected() && serial_bytes_available()) {
+            maybe_clear_next_code_info(new_next_code);
             // Skip REPL if reload was requested.
             return (serial_read() == CHAR_CTRL_D);
         }