micropython · dpgeorge · Dec 19, 2023 · Dec 18, 2023 · Dec 18, 2023
diff --git a/ports/esp32/main.c b/ports/esp32/main.c
@@ -79,10 +79,6 @@
 #define MP_TASK_STACK_LIMIT_MARGIN (1024)
 #endif
 
-// Initial Python heap size. This starts small but adds new heap areas on
-// demand due to settings MICROPY_GC_SPLIT_HEAP & MICROPY_GC_SPLIT_HEAP_AUTO
-#define MP_TASK_HEAP_SIZE (64 * 1024)
-
 int vprintf_null(const char *format, va_list ap) {
     // do nothing: this is used as a log target during raw repl mode
     return 0;
@@ -120,13 +116,13 @@ void mp_task(void *pvParameter) {
         ESP_LOGE("esp_init", "can't create event loop: 0x%x\n", err);
     }
 
-    void *mp_task_heap = MP_PLAT_ALLOC_HEAP(MP_TASK_HEAP_SIZE);
+    void *mp_task_heap = MP_PLAT_ALLOC_HEAP(MICROPY_GC_INITIAL_HEAP_SIZE);
 
 soft_reset:
     // initialise the stack pointer for the main thread
     mp_stack_set_top((void *)sp);
     mp_stack_set_limit(MICROPY_TASK_STACK_SIZE - MP_TASK_STACK_LIMIT_MARGIN);
-    gc_init(mp_task_heap, mp_task_heap + MP_TASK_HEAP_SIZE);
+    gc_init(mp_task_heap, mp_task_heap + MICROPY_GC_INITIAL_HEAP_SIZE);
     mp_init();
     mp_obj_list_append(mp_sys_path, MP_OBJ_NEW_QSTR(MP_QSTR__slash_lib));
     readline_init0();

diff --git a/ports/esp32/mpconfigport.h b/ports/esp32/mpconfigport.h
@@ -26,6 +26,20 @@
 // memory allocation policies
 #define MICROPY_ALLOC_PATH_MAX              (128)
 
+// Initial Python heap size.  This starts small but adds new heap areas on demand due to
+// the settings MICROPY_GC_SPLIT_HEAP and MICROPY_GC_SPLIT_HEAP_AUTO.  The value is
+// different for different MCUs and is chosen so they can grow the heap once (double it)
+// and still have enough internal RAM to start WiFi and make a HTTPS request.
+#ifndef MICROPY_GC_INITIAL_HEAP_SIZE
+#if CONFIG_IDF_TARGET_ESP32
+#define MICROPY_GC_INITIAL_HEAP_SIZE        (56 * 1024)
+#elif CONFIG_IDF_TARGET_ESP32S2 && !CONFIG_SPIRAM
+#define MICROPY_GC_INITIAL_HEAP_SIZE        (36 * 1024)
+#else
+#define MICROPY_GC_INITIAL_HEAP_SIZE        (64 * 1024)
+#endif
+#endif
+
 // emitters
 #define MICROPY_PERSISTENT_CODE_LOAD        (1)
 #if !CONFIG_IDF_TARGET_ESP32C3

diff --git a/py/gc.c b/py/gc.c
@@ -282,14 +282,28 @@ STATIC bool gc_try_add_heap(size_t failed_alloc) {
     // - If the new heap won't fit in the available free space, add the largest
     //   new heap that will fit (this may lead to failed system heap allocations
     //   elsewhere, but some allocation will likely fail in this circumstance!)
-    size_t total_heap = 0;
+
+    // Compute total number of blocks in the current heap.
+    size_t total_blocks = 0;
     for (mp_state_mem_area_t *area = &MP_STATE_MEM(area);
          area != NULL;
          area = NEXT_AREA(area)) {
-        total_heap += area->gc_pool_end - area->gc_alloc_table_start;
-        total_heap += ALLOC_TABLE_GAP_BYTE + sizeof(mp_state_mem_area_t);
+        total_blocks += area->gc_alloc_table_byte_len * BLOCKS_PER_ATB;
     }
 
+    // Compute bytes needed to build a heap with total_blocks blocks.
+    size_t total_heap =
+        total_blocks / BLOCKS_PER_ATB
+        #if MICROPY_ENABLE_FINALISER
+        + total_blocks / BLOCKS_PER_FTB
+        #endif
+        + total_blocks * BYTES_PER_BLOCK
+        + ALLOC_TABLE_GAP_BYTE
+        + sizeof(mp_state_mem_area_t);
+
+    // Round up size to the nearest multiple of BYTES_PER_BLOCK.
+    total_heap = (total_heap + BYTES_PER_BLOCK - 1) & (~(BYTES_PER_BLOCK - 1));
+
     DEBUG_printf("total_heap " UINT_FMT " bytes\n", total_heap);
 
     size_t to_alloc = MIN(avail, MAX(total_heap, needed));