MicroPython: Magic C++ allocator module.

This module redirects all pre-init object memory allocations into a static, fixed-sized, linear, non-freeing memory pool.

It then auto-switches into "MICROPYTHON" mode, whereupon all C++ `new` and `delete` calls are remapped to `m_malloc` and `m_free` respectively.

micropython/modules/badger2040-micropython.cmake

@@ -21,3 +21,6 @@ endif()
 include(plasma/micropython)
 include(ulab/code/micropython)
 include(qrcode/micropython/micropython)
+
+# C++ Magic Memory
+include(cppmem/micropython)

micropython/modules/cppmem/cppmem.c

@@ -0,0 +1,30 @@
+#include "cppmem.h"
+
+STATIC MP_DEFINE_CONST_FUN_OBJ_0(cpp_mem_get_alloc_bytes_obj, cpp_mem_get_alloc_bytes);
+STATIC MP_DEFINE_CONST_FUN_OBJ_0(cpp_mem_get_alloc_count_obj, cpp_mem_get_alloc_count);
+STATIC MP_DEFINE_CONST_FUN_OBJ_0(cpp_mem_get_free_count_obj, cpp_mem_get_free_count);
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(cpp_mem_set_allocator_mode_obj, cpp_mem_set_allocator_mode);
+STATIC MP_DEFINE_CONST_FUN_OBJ_0(cpp_mem_get_allocator_mode_obj, cpp_mem_get_allocator_mode);
+
+STATIC const mp_map_elem_t cppmem_globals_table[] = {
+    { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_cppmem) },
+    { MP_ROM_QSTR(MP_QSTR_alloc_bytes), MP_ROM_PTR(&cpp_mem_get_alloc_bytes_obj) },
+    { MP_ROM_QSTR(MP_QSTR_alloc_count), MP_ROM_PTR(&cpp_mem_get_alloc_count_obj) },
+    { MP_ROM_QSTR(MP_QSTR_free_count), MP_ROM_PTR(&cpp_mem_get_free_count_obj) },
+    { MP_ROM_QSTR(MP_QSTR_set_mode), MP_ROM_PTR(&cpp_mem_set_allocator_mode_obj) },
+    { MP_ROM_QSTR(MP_QSTR_get_mode), MP_ROM_PTR(&cpp_mem_get_allocator_mode_obj) },
+    { MP_ROM_QSTR(MP_QSTR_FIXED_HEAP), MP_ROM_INT(0) },
+    { MP_ROM_QSTR(MP_QSTR_MICROPYTHON), MP_ROM_INT(1) }
+};
+STATIC MP_DEFINE_CONST_DICT(mp_module_cppmem_globals, cppmem_globals_table);
+
+const mp_obj_module_t cppmem_user_cmodule = {
+    .base = { &mp_type_module },
+    .globals = (mp_obj_dict_t*)&mp_module_cppmem_globals,
+};
+
+#if MICROPY_VERSION <= 70144
+MP_REGISTER_MODULE(MP_QSTR_cppmem, cppmem_user_cmodule, MODULE_CPPMEM_ENABLED);
+#else
+MP_REGISTER_MODULE(MP_QSTR_cppmem, cppmem_user_cmodule);
+#endif

micropython/modules/cppmem/cppmem.cpp

@@ -0,0 +1,118 @@
+#include <new>
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+#include <cstdio>
+
+enum allocator_mode {
+    FIXED_HEAP = 0,
+    MICROPYTHON
+};
+
+static uint32_t alloc_bytes = 0;
+static uint32_t alloc_count = 0;
+static uint32_t free_count = 0;
+
+static allocator_mode mode = FIXED_HEAP;
+static constexpr size_t cpp_heap_size = 10 * 1024 / 4;
+static uint32_t cpp_heap[cpp_heap_size];
+static uint32_t ptr = 0;
+static char cpp_err_buf[128] = {0};
+
+extern "C" {
+#include "cppmem.h"
+
+mp_obj_t cpp_mem_set_allocator_mode(mp_obj_t new_mode) {
+    mode = (allocator_mode)mp_obj_get_int(new_mode);
+    return mp_const_none;
+}
+
+mp_obj_t cpp_mem_get_allocator_mode() {
+    return mp_obj_new_int((int)mode);
+}
+
+mp_obj_t cpp_mem_get_alloc_bytes() {
+    return mp_obj_new_int(alloc_bytes);
+}
+
+mp_obj_t cpp_mem_get_alloc_count() {
+    return mp_obj_new_int(alloc_count);
+}
+
+mp_obj_t cpp_mem_get_free_count() {
+    return mp_obj_new_int(free_count);
+}
+
+void mpy_free(void *p) {
+    return m_free(p);
+}
+
+void* mpy_malloc(size_t n) {
+    return m_malloc(n);
+}
+
+}
+
+std::size_t alloc_size(std::size_t n) {
+    std::size_t s = n / 4;
+    if(n & 0b11) s++;
+    return s;
+}
+
+void* stat_new(std::size_t n) {
+    if(mode == MICROPYTHON) {
+        return mpy_malloc(n);
+    }
+    std::size_t s = alloc_size(n);
+    if(ptr + s > cpp_heap_size) {
+        snprintf(cpp_err_buf, sizeof(cpp_err_buf), "Failed to allocate %d bytes.", s * 4);
+        mp_raise_msg(&mp_type_RuntimeError, cpp_err_buf);
+        return nullptr;
+    }
+    alloc_bytes += s * 4;
+    alloc_count ++;
+    void *result = &cpp_heap[ptr];
+    //memset(result, n, 0);
+    ptr += s;
+    return result;
+    //return std::malloc(n);
+}
+
+void stat_free(void *p) {
+    if(mode == MICROPYTHON) {
+        return mpy_free(p);
+    }
+    free_count ++;
+    //return std::free(p);
+    return;
+}
+
+void *operator new(std::size_t n) {
+    return stat_new(n);
+}
+
+void *operator new[](std::size_t n) {
+    return stat_new(n);
+}
+
+void operator delete(void *p) { stat_free(p); }
+
+void operator delete[](void *p) noexcept { stat_free(p); }
+
+#if __cpp_sized_deallocation
+
+void operator delete(void *p, __unused std::size_t n) noexcept { stat_free(p); }
+
+void operator delete[](void *p, __unused std::size_t n) noexcept { stat_free(p); }
+
+#endif
+
+namespace {
+    struct SwitchAllocatorMode {
+        public:
+            SwitchAllocatorMode() {
+                mode = MICROPYTHON;
+            }
+    };
+    SwitchAllocatorMode switchallocatormode __attribute__ ((init_priority (65535)));
+}

micropython/modules/cppmem/cppmem.h

@@ -0,0 +1,8 @@
+#include "py/runtime.h"
+#include "py/objstr.h"
+
+extern mp_obj_t cpp_mem_get_alloc_bytes();
+extern mp_obj_t cpp_mem_get_alloc_count();
+extern mp_obj_t cpp_mem_get_free_count();
+extern mp_obj_t cpp_mem_set_allocator_mode(mp_obj_t mode);
+extern mp_obj_t cpp_mem_get_allocator_mode();

micropython/modules/cppmem/micropython.cmake

@@ -0,0 +1,24 @@
+add_library(usermod_cppmem INTERFACE)
+
+target_sources(usermod_cppmem INTERFACE
+    ${CMAKE_CURRENT_LIST_DIR}/cppmem.c
+    ${CMAKE_CURRENT_LIST_DIR}/cppmem.cpp
+)
+
+target_include_directories(usermod_cppmem INTERFACE
+    ${CMAKE_CURRENT_LIST_DIR}
+)
+
+target_compile_definitions(usermod_cppmem INTERFACE
+    MODULE_CPPMEM_ENABLED=1
+)
+
+target_link_libraries(usermod INTERFACE usermod_cppmem)
+
+target_compile_options(usermod INTERFACE "-DPICO_CXX_DISABLE_ALLOCATION_OVERRIDES")
+
+set_source_files_properties(
+    ${CMAKE_CURRENT_LIST_DIR}/cppmem.c
+    PROPERTIES COMPILE_FLAGS
+    "-Wno-discarded-qualifiers"
+)

micropython/modules/micropython-badger2040w.cmake

@@ -39,4 +39,7 @@ include(motor/micropython)
 
 # include(micropython-common)
 
-include(modules_py/modules_py)
+include(modules_py/modules_py)
+
+# C++ Magic Memory
+include(cppmem/micropython)

micropython/modules/micropython-pico.cmake

@@ -8,3 +8,6 @@ set(CMAKE_C_STANDARD 11)
 set(CMAKE_CXX_STANDARD 17)
 
 include(micropython-common)
+
+# C++ Magic Memory
+include(cppmem/micropython)

micropython/modules/micropython-picolipo_16mb.cmake

@@ -10,3 +10,6 @@ set(CMAKE_CXX_STANDARD 17)
 include(micropython-common)
 
 enable_ulab()
+
+# C++ Magic Memory
+include(cppmem/micropython)

micropython/modules/micropython-picolipo_4mb.cmake

@@ -10,3 +10,6 @@ set(CMAKE_CXX_STANDARD 17)
 include(micropython-common)
 
 enable_ulab()
+
+# C++ Magic Memory
+include(cppmem/micropython)

micropython/modules/micropython-picow.cmake

@@ -8,3 +8,6 @@ set(CMAKE_C_STANDARD 11)
 set(CMAKE_CXX_STANDARD 17)
 
 include(micropython-common)
+
+# C++ Magic Memory
+include(cppmem/micropython)

micropython/modules/micropython-picow_cosmic_unicorn.cmake

@@ -35,3 +35,6 @@ include(micropython-common-ulab)
 enable_ulab()
 
 include(modules_py/modules_py)
+
+# C++ Magic Memory
+include(cppmem/micropython)

micropython/modules/micropython-picow_enviro.cmake

@@ -44,4 +44,7 @@ include(motor/micropython)
 
 # include(micropython-common)
 
-include(modules_py/modules_py)
+include(modules_py/modules_py)
+
+# C++ Magic Memory
+include(cppmem/micropython)

micropython/modules/micropython-picow_galactic_unicorn.cmake

@@ -35,3 +35,6 @@ include(micropython-common-ulab)
 enable_ulab()
 
 include(modules_py/modules_py)
+
+# C++ Magic Memory
+include(cppmem/micropython)

micropython/modules/micropython-picow_inky_frame.cmake

@@ -46,3 +46,6 @@ include(motor/micropython)
 include(modules_py/modules_py)
 
 copy_module(inky_frame.py)
+
+# C++ Magic Memory
+include(cppmem/micropython)

micropython/modules/micropython-tiny2040.cmake

@@ -8,3 +8,6 @@ set(CMAKE_C_STANDARD 11)
 set(CMAKE_CXX_STANDARD 17)
 
 include(micropython-common)
+
+# C++ Magic Memory
+include(cppmem/micropython)

micropython/modules/micropython-tufty2040.cmake

@@ -39,4 +39,7 @@ include(motor/micropython)
 include(micropython-common-ulab)
 enable_ulab()
 
-include(modules_py/modules_py)
+include(modules_py/modules_py)
+
+# C++ Magic Memory
+include(cppmem/micropython)