1 files changed, 17 insertions, 16 deletions

diff --git a/ports/esp32/main.c b/ports/esp32/main.c
index c9b033d62..d21dc9f75 100644
--- a/ports/esp32/main.c
+++ b/ports/esp32/main.c
@@ -97,11 +97,12 @@ void mp_task(void *pvParameter) {
     #endif
     machine_init();
 
-    // TODO: CONFIG_SPIRAM_SUPPORT is for 3.3 compatibility, remove after move to 4.0.
-    #if CONFIG_ESP32_SPIRAM_SUPPORT || CONFIG_SPIRAM_SUPPORT
-    // Try to use the entire external SPIRAM directly for the heap
     size_t mp_task_heap_size;
-    void *mp_task_heap = (void *)SOC_EXTRAM_DATA_LOW;
+    void *mp_task_heap = NULL;
+
+    #if CONFIG_ESP32_SPIRAM_SUPPORT
+    // Try to use the entire external SPIRAM directly for the heap
+    mp_task_heap = (void *)SOC_EXTRAM_DATA_LOW;
     switch (esp_spiram_get_chip_size()) {
         case ESP_SPIRAM_SIZE_16MBITS:
             mp_task_heap_size = 2 * 1024 * 1024;
@@ -112,28 +113,28 @@ void mp_task(void *pvParameter) {
             break;
         default:
             // No SPIRAM, fallback to normal allocation
-            mp_task_heap_size = heap_caps_get_largest_free_block(MALLOC_CAP_8BIT);
-            mp_task_heap = malloc(mp_task_heap_size);
+            mp_task_heap = NULL;
             break;
     }
     #elif CONFIG_ESP32S2_SPIRAM_SUPPORT || CONFIG_ESP32S3_SPIRAM_SUPPORT
     // Try to use the entire external SPIRAM directly for the heap
-    size_t mp_task_heap_size;
     size_t esp_spiram_size = esp_spiram_get_size();
-    void *mp_task_heap = (void *)SOC_EXTRAM_DATA_HIGH - esp_spiram_size;
     if (esp_spiram_size > 0) {
+        mp_task_heap = (void *)SOC_EXTRAM_DATA_HIGH - esp_spiram_size;
         mp_task_heap_size = esp_spiram_size;
-    } else {
-        // No SPIRAM, fallback to normal allocation
-        mp_task_heap_size = heap_caps_get_largest_free_block(MALLOC_CAP_8BIT);
-        mp_task_heap = malloc(mp_task_heap_size);
     }
-    #else
-    // Allocate the uPy heap using malloc and get the largest available region
-    size_t mp_task_heap_size = heap_caps_get_largest_free_block(MALLOC_CAP_8BIT);
-    void *mp_task_heap = malloc(mp_task_heap_size);
     #endif
 
+    if (mp_task_heap == NULL) {
+        // Allocate the uPy heap using malloc and get the largest available region,
+        // limiting to 1/2 total available memory to leave memory for the OS.
+        mp_task_heap_size = MIN(
+            heap_caps_get_largest_free_block(MALLOC_CAP_8BIT),
+            heap_caps_get_total_size(MALLOC_CAP_8BIT) / 2
+            );
+        mp_task_heap = malloc(mp_task_heap_size);
+    }
+
 soft_reset:
     // initialise the stack pointer for the main thread
     mp_stack_set_top((void *)sp);