1 files changed, 8 insertions, 6 deletions

diff --git a/main.c b/main.c
index e137f7d47..b17c11137 100755
--- a/main.c
+++ b/main.c
@@ -44,6 +44,7 @@
 #include "lib/utils/pyexec.h"
 
 #include "mpconfigboard.h"
+#include "shared-module/displayio/__init__.h"
 #include "supervisor/cpu.h"
 #include "supervisor/memory.h"
 #include "supervisor/port.h"
@@ -61,10 +62,6 @@
 #include "shared-module/network/__init__.h"
 #endif
 
-#ifdef CIRCUITPY_DISPLAYIO
-#include "shared-module/displayio/__init__.h"
-#endif
-
 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) {
@@ -203,12 +200,11 @@ bool run_code_py(safe_mode_t safe_mode) {
                 serial_write_compressed(translate("WARNING: Your code filename has two extensions\n"));
             }
         }
-        #ifdef CIRCUITPY_DISPLAYIO
         // Turn off the display before the heap disappears.
         reset_displays();
-        #endif
         stop_mp();
         free_memory(heap);
+        supervisor_move_memory();
 
         reset_port();
         reset_board_busses();
@@ -221,6 +217,10 @@ bool run_code_py(safe_mode_t safe_mode) {
     }
 
     // Wait for connection or character.
+    if (!serial_connected_at_start) {
+        serial_write_compressed(translate("\nCode done running. Waiting for reload.\n"));
+    }
+
     bool serial_connected_before_animation = false;
     rgb_status_animation_t animation;
     prep_rgb_status_animation(&result, found_main, safe_mode, &animation);
@@ -342,6 +342,7 @@ void __attribute__ ((noinline)) run_boot_py(safe_mode_t safe_mode) {
         reset_board();
         stop_mp();
         free_memory(heap);
+        supervisor_move_memory();
     }
 }
 
@@ -362,6 +363,7 @@ int run_repl(void) {
     reset_board();
     stop_mp();
     free_memory(heap);
+    supervisor_move_memory();
     autoreload_resume();
     return exit_code;
 }