Get AllocationAlarm working

10 files changed, 40 insertions, 40 deletions

diff --git a/main.c b/main.c
index f928d0d62..5ad146c9c 100755
--- a/main.c
+++ b/main.c
@@ -64,6 +64,10 @@
 #include "shared-module/displayio/__init__.h"
 #endif
 
+#if CIRCUITPY_MEMORYMONITOR
+#include "shared-module/memorymonitor/__init__.h"
+#endif
+
 #if CIRCUITPY_NETWORK
 #include "shared-module/network/__init__.h"
 #endif
@@ -198,6 +202,9 @@ void cleanup_after_vm(supervisor_allocation* heap) {
     #if CIRCUITPY_DISPLAYIO
     reset_displays();
     #endif
+    #if CIRCUITPY_MEMORYMONITOR
+    memorymonitor_reset();
+    #endif
     filesystem_flush();
     stop_mp();
     free_memory(heap);
diff --git a/py/gc.c b/py/gc.c
index 0f08ffb2e..69327060f 100755
--- a/py/gc.c
+++ b/py/gc.c
@@ -34,7 +34,7 @@
 #include "supervisor/shared/safe_mode.h"
 
 #if CIRCUITPY_MEMORYMONITOR
-#include "shared-module/memorymonitor/AllocationSize.h"
+#include "shared-module/memorymonitor/__init__.h"
 #endif
 
 #if MICROPY_ENABLE_GC
@@ -658,7 +658,7 @@ void *gc_alloc(size_t n_bytes, bool has_finaliser, bool long_lived) {
     #endif
 
     #if CIRCUITPY_MEMORYMONITOR
-    memorymonitor_allocationsizes_track_allocation(end_block - start_block + 1);
+    memorymonitor_track_allocation(end_block - start_block + 1);
     #endif
 
     return ret_ptr;
@@ -915,7 +915,7 @@ void *gc_realloc(void *ptr_in, size_t n_bytes, bool allow_move) {
         #endif
 
         #if CIRCUITPY_MEMORYMONITOR
-        memorymonitor_allocationsizes_track_allocation(new_blocks);
+        memorymonitor_track_allocation(new_blocks);
         #endif
 
         return ptr_in;
@@ -948,7 +948,7 @@ void *gc_realloc(void *ptr_in, size_t n_bytes, bool allow_move) {
         #endif
 
         #if CIRCUITPY_MEMORYMONITOR
-        memorymonitor_allocationsizes_track_allocation(new_blocks);
+        memorymonitor_track_allocation(new_blocks);
         #endif
 
         return ptr_in;
diff --git a/shared-bindings/memorymonitor/AllocationAlarm.c b/shared-bindings/memorymonitor/AllocationAlarm.c
index 836bf7833..71a156f32 100644
--- a/shared-bindings/memorymonitor/AllocationAlarm.c
+++ b/shared-bindings/memorymonitor/AllocationAlarm.c
@@ -76,6 +76,8 @@ STATIC mp_obj_t memorymonitor_allocationalarm_make_new(const mp_obj_type_t *type
     return MP_OBJ_FROM_PTR(self);
 }
 
+// TODO: Add .countdown(count) to skip allocations and alarm on something after the first.
+
 //|     def __enter__(self) -> memorymonitor.AllocationAlarm:
 //|         """Enables the alarm."""
 //|         ...
diff --git a/shared-bindings/memorymonitor/AllocationSize.c b/shared-bindings/memorymonitor/AllocationSize.c
index 411e27e15..25ecae97b 100644
--- a/shared-bindings/memorymonitor/AllocationSize.c
+++ b/shared-bindings/memorymonitor/AllocationSize.c
@@ -82,6 +82,7 @@ STATIC mp_obj_t memorymonitor_allocationsize_make_new(const mp_obj_type_t *type,
 //|         ...
 //|
 STATIC mp_obj_t memorymonitor_allocationsize_obj___enter__(mp_obj_t self_in) {
+    common_hal_memorymonitor_allocationsize_clear(self_in);
     common_hal_memorymonitor_allocationsize_resume(self_in);
     return self_in;
 }
@@ -99,48 +100,13 @@ STATIC mp_obj_t memorymonitor_allocationsize_obj___exit__(size_t n_args, const m
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(memorymonitor_allocationsize___exit___obj, 4, 4, memorymonitor_allocationsize_obj___exit__);
 
-//|     def pause(self) -> None:
-//|         """Pause allocation tracking"""
-//|         ...
-//|
-STATIC mp_obj_t memorymonitor_allocationsize_obj_pause(mp_obj_t self_in) {
-    memorymonitor_allocationsize_obj_t *self = MP_OBJ_TO_PTR(self_in);
-
-    common_hal_memorymonitor_allocationsize_pause(self);
-    return mp_const_none;
-}
-MP_DEFINE_CONST_FUN_OBJ_1(memorymonitor_allocationsize_pause_obj, memorymonitor_allocationsize_obj_pause);
-
-//|     def resume(self) -> None:
-//|         """Resumes allocation tracking."""
-//|         ...
-//|
-STATIC mp_obj_t memorymonitor_allocationsize_obj_resume(mp_obj_t self_in) {
-    common_hal_memorymonitor_allocationsize_resume(self_in);
-    return mp_const_none;
-}
-MP_DEFINE_CONST_FUN_OBJ_1(memorymonitor_allocationsize_resume_obj, memorymonitor_allocationsize_obj_resume);
-
-//|     def clear(self) -> Any:
-//|         """Clears all captured pulses"""
-//|         ...
-//|
-STATIC mp_obj_t memorymonitor_allocationsize_obj_clear(mp_obj_t self_in) {
-    memorymonitor_allocationsize_obj_t *self = MP_OBJ_TO_PTR(self_in);
-
-    common_hal_memorymonitor_allocationsize_clear(self);
-    return mp_const_none;
-}
-MP_DEFINE_CONST_FUN_OBJ_1(memorymonitor_allocationsize_clear_obj, memorymonitor_allocationsize_obj_clear);
-
-
 //|     bytes_per_block: int = ...
 //|     """Number of bytes per block"""
 //|
 STATIC mp_obj_t memorymonitor_allocationsize_obj_get_bytes_per_block(mp_obj_t self_in) {
     memorymonitor_allocationsize_obj_t *self = MP_OBJ_TO_PTR(self_in);
 
-    return mp_obj_new_bool(common_hal_memorymonitor_allocationsize_get_bytes_per_block(self));
+    return MP_OBJ_NEW_SMALL_INT(common_hal_memorymonitor_allocationsize_get_bytes_per_block(self));
 }
 MP_DEFINE_CONST_FUN_OBJ_1(memorymonitor_allocationsize_get_bytes_per_block_obj, memorymonitor_allocationsize_obj_get_bytes_per_block);
 
diff --git a/shared-module/memorymonitor/AllocationAlarm.c b/shared-module/memorymonitor/AllocationAlarm.c
index bbd33f2b0..ed7ab75ab 100644
--- a/shared-module/memorymonitor/AllocationAlarm.c
+++ b/shared-module/memorymonitor/AllocationAlarm.c
@@ -38,6 +38,11 @@ void common_hal_memorymonitor_allocationalarm_construct(memorymonitor_allocation
 }
 
 void common_hal_memorymonitor_allocationalarm_pause(memorymonitor_allocationalarm_obj_t* self) {
+    // Check to make sure we aren't already paused. We can be if we're exiting from an exception we
+    // caused.
+    if (self->previous == NULL) {
+        return;
+    }
     *self->previous = self->next;
     self->next = NULL;
     self->previous = NULL;
@@ -62,6 +67,10 @@ void memorymonitor_allocationalarms_allocation(size_t block_count) {
         // Hold onto next in case we remove the alarm from the list.
         memorymonitor_allocationalarm_obj_t* next = alarm->next;
         if (block_count >= alarm->minimum_block_count) {
+            // Uncomment the breakpoint below if you want to use a C debugger to figure out the C
+            // call stack for an allocation.
+            // asm("bkpt");
+            // Pause now because we may alert when throwing the exception too.
             common_hal_memorymonitor_allocationalarm_pause(alarm);
             alert_count++;
         }
@@ -71,3 +80,7 @@ void memorymonitor_allocationalarms_allocation(size_t block_count) {
         mp_raise_memorymonitor_AllocationError(translate("Attempt to allocate %d blocks"), block_count);
     }
 }
+
+void memorymonitor_allocationalarms_reset(void) {
+    MP_STATE_VM(active_allocationalarms) = NULL;
+}
diff --git a/shared-module/memorymonitor/AllocationAlarm.h b/shared-module/memorymonitor/AllocationAlarm.h
index ebdc0d801..95381c660 100644
--- a/shared-module/memorymonitor/AllocationAlarm.h
+++ b/shared-module/memorymonitor/AllocationAlarm.h
@@ -45,5 +45,6 @@ typedef struct _memorymonitor_allocationalarm_obj_t {
 } memorymonitor_allocationalarm_obj_t;
 
 void memorymonitor_allocationalarms_allocation(size_t block_count);
+void memorymonitor_allocationalarms_reset(void);
 
 #endif // MICROPY_INCLUDED_SHARED_MODULE_MEMORYMONITOR_ALLOCATIONALARM_H
diff --git a/shared-module/memorymonitor/AllocationSize.c b/shared-module/memorymonitor/AllocationSize.c
index fb11471fd..c28e65592 100644
--- a/shared-module/memorymonitor/AllocationSize.c
+++ b/shared-module/memorymonitor/AllocationSize.c
@@ -85,3 +85,7 @@ void memorymonitor_allocationsizes_track_allocation(size_t block_count) {
         as = as->next;
     }
 }
+
+void memorymonitor_allocationsizes_reset(void) {
+    MP_STATE_VM(active_allocationsizes) = NULL;
+}
diff --git a/shared-module/memorymonitor/AllocationSize.h b/shared-module/memorymonitor/AllocationSize.h
index 3fa4e0465..3baab2213 100644
--- a/shared-module/memorymonitor/AllocationSize.h
+++ b/shared-module/memorymonitor/AllocationSize.h
@@ -46,5 +46,6 @@ typedef struct _memorymonitor_allocationsize_obj_t {
 } memorymonitor_allocationsize_obj_t;
 
 void memorymonitor_allocationsizes_track_allocation(size_t block_count);
+void memorymonitor_allocationsizes_reset(void);
 
 #endif // MICROPY_INCLUDED_SHARED_MODULE_MEMORYMONITOR_ALLOCATIONSIZE_H
diff --git a/shared-module/memorymonitor/__init__.c b/shared-module/memorymonitor/__init__.c
index 4c24c2587..6cb424153 100644
--- a/shared-module/memorymonitor/__init__.c
+++ b/shared-module/memorymonitor/__init__.c
@@ -32,3 +32,8 @@ void memorymonitor_track_allocation(size_t block_count) {
     memorymonitor_allocationalarms_allocation(block_count);
     memorymonitor_allocationsizes_track_allocation(block_count);
 }
+
+void memorymonitor_reset(void) {
+    memorymonitor_allocationalarms_reset();
+    memorymonitor_allocationsizes_reset();
+}
diff --git a/shared-module/memorymonitor/__init__.h b/shared-module/memorymonitor/__init__.h
index cf76f88f8..f47f6434b 100644
--- a/shared-module/memorymonitor/__init__.h
+++ b/shared-module/memorymonitor/__init__.h
@@ -30,5 +30,6 @@
 #include <stddef.h>
 
 void memorymonitor_track_allocation(size_t block_count);
+void memorymonitor_reset(void);
 
 #endif  // MICROPY_INCLUDED_MEMORYMONITOR___INIT___H