From 4ecb4099cf3a966bdc6edcb6ffd05e17de019a72 Mon Sep 17 00:00:00 2001
From: Phil Howard <github@gadgetoid.com>
Date: Sat, 17 May 2025 21:42:53 +0100
Subject: [PATCH] ports/rp2: HACK: Bring up modbtree with tracked heap.

Based on pimoroni/micropython commit 0cc56b6.

Signed-off-by: Matthias Blankertz <matthias@blankertz.org>
---
 extmod/extmod.cmake      |  9 +++++++++
 extmod/modbtree_malloc.c | 40 ++++++++++++++++++++++++++++++++++++++++
 ports/rp2/CMakeLists.txt | 12 ++++++++++++
 3 files changed, 61 insertions(+)
 create mode 100644 extmod/modbtree_malloc.c

diff --git a/extmod/extmod.cmake b/extmod/extmod.cmake
index 9ca869582..4afa5d969 100644
--- a/extmod/extmod.cmake
+++ b/extmod/extmod.cmake
@@ -136,6 +136,13 @@ if(MICROPY_PY_BTREE)
         ${MICROPY_LIB_BERKELEY_DIR}/include
     )
 
+    target_link_options(micropy_extmod_btree PRIVATE
+        -Wl,--wrap=malloc
+        -Wl,--wrap=free
+        -Wl,--wrap=realloc
+        -Wl,--wrap=calloc
+    )
+
     if(NOT BERKELEY_DB_CONFIG_FILE)
         set(BERKELEY_DB_CONFIG_FILE "${MICROPY_DIR}/extmod/berkeley-db/berkeley_db_config_port.h")
     endif()
@@ -153,11 +160,13 @@ if(MICROPY_PY_BTREE)
 
     list(APPEND MICROPY_DEF_CORE
         MICROPY_PY_BTREE=1
+        MICROPY_STREAMS_POSIX_API # Required for mp_stream_posix_ functions if !MICROPY_ENABLE_DYNRUNTIME
         BERKELEY_DB_CONFIG_FILE="${BERKELEY_DB_CONFIG_FILE}"
     )
 
     list(APPEND MICROPY_SOURCE_EXTMOD
         ${MICROPY_EXTMOD_DIR}/modbtree.c
+        ${MICROPY_EXTMOD_DIR}/modbtree_malloc.c
     )
 endif()
 
diff --git a/extmod/modbtree_malloc.c b/extmod/modbtree_malloc.c
new file mode 100644
index 000000000..b9a139e84
--- /dev/null
+++ b/extmod/modbtree_malloc.c
@@ -0,0 +1,40 @@
+#include "py/runtime.h"
+#include <string.h>
+
+void *__wrap_malloc(size_t size) {
+    // Allocate an extra sizeof(size_t) bytes
+    size_t *addr = (size_t *)m_tracked_calloc(sizeof(uint8_t), size + sizeof(size_t));
+    // Tag our memory with its allocated size
+    *addr = size;
+    // Skip past the size_t size
+    addr++;
+#ifdef BTREE_DEBUG_MALLOC
+    mp_printf(&mp_plat_print, "malloc %lu %p : %p\n", size, addr, (uint8_t *)addr + size);
+#endif
+    return (void *)addr;
+}
+
+void __wrap_free(void *p) {
+    size_t *pp = (size_t *)p;   // Convert our void pointer to size_t* so we can read the size marker
+    pp--;                       // Skip back to get our real start
+    size_t size __attribute__((unused)) = *pp;
+    m_tracked_free((void *)pp);
+#ifdef BTREE_DEBUG_MALLOC
+    mp_printf(&mp_plat_print, "free %p (size %d)\n", p, size);
+#endif
+}
+
+void *__wrap_realloc(void *p, size_t size) {
+    void *addr = __wrap_malloc(size);
+    size_t old_size = *((size_t *)p - 1);
+    memcpy(addr, p, old_size < size ? old_size : size);
+    __wrap_free(p);
+#ifdef BTREE_DEBUG_MALLOC
+    mp_printf(&mp_plat_print, "realloc %lu -> %lu, %p -> %p : %p\n", old_size, size, p, addr, (uint8_t *)addr + size);
+#endif
+    return addr;
+}
+
+void *__wrap_calloc(size_t nitems, size_t size) {
+    return __wrap_malloc(size * nitems);
+}
diff --git a/ports/rp2/CMakeLists.txt b/ports/rp2/CMakeLists.txt
index e58fc00fd..fb26bfa96 100644
--- a/ports/rp2/CMakeLists.txt
+++ b/ports/rp2/CMakeLists.txt
@@ -95,6 +95,8 @@ include(${PICO_SDK_PATH}/pico_sdk_init.cmake)
 # Define the top-level project
 project(${MICROPY_TARGET})
 
+set(SKIP_PICO_MALLOC 1)
+
 pico_sdk_init()
 
 include(${MICROPY_DIR}/py/usermod.cmake)
@@ -509,6 +511,10 @@ target_link_libraries(${MICROPY_TARGET} micropy_lib_mbedtls)
 
 target_link_libraries(${MICROPY_TARGET} usermod)
 
+if (MICROPY_PY_BTREE)
+    target_link_libraries(${MICROPY_TARGET} micropy_extmod_btree)
+endif()
+
 target_include_directories(${MICROPY_TARGET} PRIVATE
     ${MICROPY_INC_CORE}
     ${MICROPY_INC_USERMOD}
@@ -526,6 +532,11 @@ target_compile_options(${MICROPY_TARGET} PRIVATE
 target_link_options(${MICROPY_TARGET} PRIVATE
     -Wl,--defsym=__micropy_c_heap_size__=${MICROPY_C_HEAP_SIZE}
     -Wl,--wrap=runtime_init_clocks
+
+    -Wl,--wrap=malloc
+    -Wl,--wrap=free
+    -Wl,--wrap=realloc
+    -Wl,--wrap=calloc
 )
 
 if(DEFINED PICO_FLASH_SIZE_BYTES)
@@ -572,6 +583,7 @@ set_source_files_properties(
 )
 
 target_compile_definitions(${MICROPY_TARGET} PRIVATE
+    ${MICROPY_DEF_CORE} # Required for btree module definitions
     ${MICROPY_DEF_BOARD}
     FFCONF_H=\"${MICROPY_OOFATFS_DIR}/ffconf.h\"
     LFS1_NO_MALLOC LFS1_NO_DEBUG LFS1_NO_WARN LFS1_NO_ERROR LFS1_NO_ASSERT