py/malloc: Introduce m_tracked_calloc, m_tracked_free functions.

Enabled by MICROPY_TRACKED_ALLOC.

Signed-off-by: Damien George <damien@micropython.org>

ports/unix/coverage.c

@@ -220,6 +220,55 @@ STATIC mp_obj_t extra_coverage(void) {
         mp_printf(&mp_plat_print, "%p\n", gc_nbytes(NULL));
     }
 
+    // tracked allocation
+    {
+        #define NUM_PTRS (8)
+        #define NUM_BYTES (128)
+        #define FLIP_POINTER(p) ((uint8_t *)((uintptr_t)(p) ^ 0x0f))
+
+        mp_printf(&mp_plat_print, "# tracked allocation\n");
+        mp_printf(&mp_plat_print, "m_tracked_head = %p\n", MP_STATE_VM(m_tracked_head));
+
+        uint8_t *ptrs[NUM_PTRS];
+
+        // allocate memory blocks
+        for (size_t i = 0; i < NUM_PTRS; ++i) {
+            ptrs[i] = m_tracked_calloc(1, NUM_BYTES);
+            bool all_zero = true;
+            for (size_t j = 0; j < NUM_BYTES; ++j) {
+                if (ptrs[i][j] != 0) {
+                    all_zero = false;
+                    break;
+                }
+                ptrs[i][j] = j;
+            }
+            mp_printf(&mp_plat_print, "%d %d\n", i, all_zero);
+
+            // hide the pointer from the GC and collect
+            ptrs[i] = FLIP_POINTER(ptrs[i]);
+            gc_collect();
+        }
+
+        // check the memory blocks have the correct content
+        for (size_t i = 0; i < NUM_PTRS; ++i) {
+            bool correct_contents = true;
+            for (size_t j = 0; j < NUM_BYTES; ++j) {
+                if (FLIP_POINTER(ptrs[i])[j] != j) {
+                    correct_contents = false;
+                    break;
+                }
+            }
+            mp_printf(&mp_plat_print, "%d %d\n", i, correct_contents);
+        }
+
+        // free the memory blocks
+        for (size_t i = 0; i < NUM_PTRS; ++i) {
+            m_tracked_free(FLIP_POINTER(ptrs[i]));
+        }
+
+        mp_printf(&mp_plat_print, "m_tracked_head = %p\n", MP_STATE_VM(m_tracked_head));
+    }
+
     // vstr
     {
         mp_printf(&mp_plat_print, "# vstr\n");

ports/unix/variants/coverage/mpconfigvariant.h

@@ -38,6 +38,7 @@
 
 // Enable additional features.
 #define MICROPY_DEBUG_PARSE_RULE_NAME  (1)
+#define MICROPY_TRACKED_ALLOC           (1)
 #define MICROPY_FLOAT_HIGH_QUALITY_HASH (1)
 #define MICROPY_REPL_EMACS_WORDS_MOVE  (1)
 #define MICROPY_REPL_EMACS_EXTRA_WORDS_MOVE (1)

py/malloc.c

@@ -207,6 +207,99 @@ void m_free(void *ptr)
     #endif
 }
 
+#if MICROPY_TRACKED_ALLOC
+
+#define MICROPY_TRACKED_ALLOC_STORE_SIZE (!MICROPY_ENABLE_GC)
+
+typedef struct _m_tracked_node_t {
+    struct _m_tracked_node_t *prev;
+    struct _m_tracked_node_t *next;
+    #if MICROPY_TRACKED_ALLOC_STORE_SIZE
+    uintptr_t size;
+    #endif
+    uint8_t data[];
+} m_tracked_node_t;
+
+#if MICROPY_DEBUG_VERBOSE
+STATIC size_t m_tracked_count_links(size_t *nb) {
+    m_tracked_node_t *node = MP_STATE_VM(m_tracked_head);
+    size_t n = 0;
+    *nb = 0;
+    while (node != NULL) {
+        ++n;
+        #if MICROPY_TRACKED_ALLOC_STORE_SIZE
+        *nb += node->size;
+        #else
+        *nb += gc_nbytes(node);
+        #endif
+        node = node->next;
+    }
+    return n;
+}
+#endif
+
+void *m_tracked_calloc(size_t nmemb, size_t size) {
+    m_tracked_node_t *node = m_malloc_maybe(sizeof(m_tracked_node_t) + nmemb * size);
+    if (node == NULL) {
+        return NULL;
+    }
+    #if MICROPY_DEBUG_VERBOSE
+    size_t nb;
+    size_t n = m_tracked_count_links(&nb);
+    DEBUG_printf("m_tracked_calloc(%u, %u) -> (%u;%u) %p\n", (int)nmemb, (int)size, (int)n, (int)nb, node);
+    #endif
+    if (MP_STATE_VM(m_tracked_head) != NULL) {
+        MP_STATE_VM(m_tracked_head)->prev = node;
+    }
+    node->prev = NULL;
+    node->next = MP_STATE_VM(m_tracked_head);
+    MP_STATE_VM(m_tracked_head) = node;
+    #if MICROPY_TRACKED_ALLOC_STORE_SIZE
+    node->size = nmemb * size;
+    #endif
+    #if !MICROPY_GC_CONSERVATIVE_CLEAR
+    memset(&node->data[0], 0, nmemb * size);
+    #endif
+    return &node->data[0];
+}
+
+void m_tracked_free(void *ptr_in) {
+    if (ptr_in == NULL) {
+        return;
+    }
+    m_tracked_node_t *node = (m_tracked_node_t *)((uint8_t *)ptr_in - sizeof(m_tracked_node_t));
+    #if MICROPY_DEBUG_VERBOSE
+    size_t data_bytes;
+    #if MICROPY_TRACKED_ALLOC_STORE_SIZE
+    data_bytes = node->size;
+    #else
+    data_bytes = gc_nbytes(node);
+    #endif
+    size_t nb;
+    size_t n = m_tracked_count_links(&nb);
+    DEBUG_printf("m_tracked_free(%p, [%p, %p], nbytes=%u, links=%u;%u)\n", node, node->prev, node->next, (int)data_bytes, (int)n, (int)nb);
+    #endif
+    if (node->next != NULL) {
+        node->next->prev = node->prev;
+    }
+    if (node->prev != NULL) {
+        node->prev->next = node->next;
+    } else {
+        MP_STATE_VM(m_tracked_head) = node->next;
+    }
+    m_free(node
+        #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+        #if MICROPY_TRACKED_ALLOC_STORE_SIZE
+        , node->size
+        #else
+        , gc_nbytes(node)
+        #endif
+        #endif
+        );
+}
+
+#endif // MICROPY_TRACKED_ALLOC
+
 #if MICROPY_MEM_STATS
 size_t m_get_total_bytes_allocated(void) {
     return MP_STATE_MEM(total_bytes_allocated);

py/misc.h

@@ -103,6 +103,13 @@ void m_free(void *ptr);
 #endif
 NORETURN void m_malloc_fail(size_t num_bytes);
 
+#if MICROPY_TRACKED_ALLOC
+// These alloc/free functions track the pointers in a linked list so the GC does not reclaim
+// them.  They can be used by code that requires traditional C malloc/free semantics.
+void *m_tracked_calloc(size_t nmemb, size_t size);
+void m_tracked_free(void *ptr_in);
+#endif
+
 #if MICROPY_MEM_STATS
 size_t m_get_total_bytes_allocated(void);
 size_t m_get_current_bytes_allocated(void);

py/mpconfig.h

@@ -621,6 +621,11 @@
 #define MICROPY_GC_HOOK_LOOP
 #endif
 
+// Whether to provide m_tracked_calloc, m_tracked_free functions
+#ifndef MICROPY_TRACKED_ALLOC
+#define MICROPY_TRACKED_ALLOC (0)
+#endif
+
 // Whether to enable finalisers in the garbage collector (ie call __del__)
 #ifndef MICROPY_ENABLE_FINALISER
 #define MICROPY_ENABLE_FINALISER (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES)

py/mpstate.h

@@ -125,6 +125,10 @@ typedef struct _mp_state_vm_t {
 
     qstr_pool_t *last_pool;
 
+    #if MICROPY_TRACKED_ALLOC
+    struct _m_tracked_node_t *m_tracked_head;
+    #endif
+
     // non-heap memory for creating an exception if we can't allocate RAM
     mp_obj_exception_t mp_emergency_exception_obj;
 

tests/unix/extra_coverage.py.exp

@@ -17,6 +17,25 @@ abc
 # GC
 0
 0
+# tracked allocation
+m_tracked_head = 0
+0 1
+1 1
+2 1
+3 1
+4 1
+5 1
+6 1
+7 1
+0 1
+1 1
+2 1
+3 1
+4 1
+5 1
+6 1
+7 1
+m_tracked_head = 0
 # vstr
 tests
 sts