Merge pull request #10094 from SamantazFox/fix-malloc-uses-allocated-…

…size

Fix m_realloc/m_free when MICROPY_MALLOC_USES_ALLOCATED_SIZE is set

ports/atmel-samd/audio_dma.c

@@ -200,15 +200,27 @@ audio_dma_result audio_dma_setup_playback(audio_dma_t *dma,
     }
 
 
-    dma->buffer[0] = (uint8_t *)m_realloc(dma->buffer[0], max_buffer_length);
+    dma->buffer[0] = (uint8_t *)m_realloc(dma->buffer[0],
+        #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+        dma->buffer_length[0], // Old size
+        #endif
+        max_buffer_length);
+
     dma->buffer_length[0] = max_buffer_length;
+
     if (dma->buffer[0] == NULL) {
         return AUDIO_DMA_MEMORY_ERROR;
     }
 
     if (!dma->single_buffer) {
-        dma->buffer[1] = (uint8_t *)m_realloc(dma->buffer[1], max_buffer_length);
+        dma->buffer[1] = (uint8_t *)m_realloc(dma->buffer[1],
+            #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+            dma->buffer_length[1], // Old size
+            #endif
+            max_buffer_length);
+
         dma->buffer_length[1] = max_buffer_length;
+
         if (dma->buffer[1] == NULL) {
             return AUDIO_DMA_MEMORY_ERROR;
         }

ports/nordic/common-hal/audiopwmio/PWMAudioOut.c

@@ -188,6 +188,26 @@ bool common_hal_audiopwmio_pwmaudioout_deinited(audiopwmio_pwmaudioout_obj_t *se
     return !self->pwm;
 }
 
+static void free_buffers(audiopwmio_pwmaudioout_obj_t *self) {
+    #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+    m_free(self->buffers[0], self->buffer_size[0]);
+    self->buffer_size[0] = 0;
+    #else
+    m_free(self->buffers[0]);
+    #endif
+
+    self->buffers[0] = NULL;
+
+    #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+    m_free(self->buffers[1], self->buffer_size[1]);
+    self->buffer_size[1] = 0;
+    #else
+    m_free(self->buffers[1]);
+    #endif
+
+    self->buffers[1] = NULL;
+}
+
 void common_hal_audiopwmio_pwmaudioout_deinit(audiopwmio_pwmaudioout_obj_t *self) {
     if (common_hal_audiopwmio_pwmaudioout_deinited(self)) {
         return;
@@ -209,11 +229,7 @@ void common_hal_audiopwmio_pwmaudioout_deinit(audiopwmio_pwmaudioout_obj_t *self
 
     self->pwm = NULL;
 
-    m_free(self->buffers[0]);
-    self->buffers[0] = NULL;
-
-    m_free(self->buffers[1]);
-    self->buffers[1] = NULL;
+    free_buffers(self);
 }
 
 void common_hal_audiopwmio_pwmaudioout_play(audiopwmio_pwmaudioout_obj_t *self, mp_obj_t sample, bool loop) {
@@ -235,10 +251,18 @@ void common_hal_audiopwmio_pwmaudioout_play(audiopwmio_pwmaudioout_obj_t *self,
 
     mp_arg_validate_length_max(max_buffer_length, UINT16_MAX, MP_QSTR_buffer);
 
-    uint16_t buffer_length = (uint16_t)max_buffer_length;
-    self->buffers[0] = m_malloc(buffer_length * 2 * sizeof(uint16_t));
+    size_t buffer_size = (size_t)max_buffer_length * 2 * sizeof(uint16_t);
+
+    self->buffers[0] = m_malloc(buffer_size);
+    #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+    self->buffer_size[0] = buffer_size;
+    #endif
+
     if (!self->single_buffer) {
-        self->buffers[1] = m_malloc(buffer_length * 2 * sizeof(uint16_t));
+        self->buffers[1] = m_malloc(buffer_size);
+        #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+        self->buffer_size[1] = buffer_size;
+        #endif
     }
 
 
@@ -274,11 +298,7 @@ void common_hal_audiopwmio_pwmaudioout_stop(audiopwmio_pwmaudioout_obj_t *self)
     self->stopping = false;
     self->paused = false;
 
-    m_free(self->buffers[0]);
-    self->buffers[0] = NULL;
-
-    m_free(self->buffers[1]);
-    self->buffers[1] = NULL;
+    free_buffers(self);
 }
 
 bool common_hal_audiopwmio_pwmaudioout_get_playing(audiopwmio_pwmaudioout_obj_t *self) {

ports/nordic/common-hal/audiopwmio/PWMAudioOut.h

@@ -12,7 +12,11 @@ typedef struct {
     mp_obj_base_t base;
     mp_obj_t *sample;
     NRF_PWM_Type *pwm;
+
     uint16_t *buffers[2];
+    #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+    size_t buffer_size[2]; // Keeps track of allocated size
+    #endif
 
     uint16_t quiescent_value;
     uint16_t scale;

ports/raspberrypi/audio_dma.c

@@ -227,15 +227,27 @@ audio_dma_result audio_dma_setup_playback(
         max_buffer_length /= dma->sample_spacing;
     }
 
-    dma->buffer[0] = (uint8_t *)m_realloc(dma->buffer[0], max_buffer_length);
+    dma->buffer[0] = (uint8_t *)m_realloc(dma->buffer[0],
+        #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+        dma->buffer_length[0], // Old size
+        #endif
+        max_buffer_length);
+
     dma->buffer_length[0] = max_buffer_length;
+
     if (dma->buffer[0] == NULL) {
         return AUDIO_DMA_MEMORY_ERROR;
     }
 
     if (!single_buffer) {
-        dma->buffer[1] = (uint8_t *)m_realloc(dma->buffer[1], max_buffer_length);
+        dma->buffer[1] = (uint8_t *)m_realloc(dma->buffer[1],
+            #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+            dma->buffer_length[1], // Old size
+            #endif
+            max_buffer_length);
+
         dma->buffer_length[1] = max_buffer_length;
+
         if (dma->buffer[1] == NULL) {
             return AUDIO_DMA_MEMORY_ERROR;
         }
@@ -419,16 +431,31 @@ void audio_dma_init(audio_dma_t *dma) {
     dma->buffer[0] = NULL;
     dma->buffer[1] = NULL;
 
+    dma->buffer_length[0] = 0;
+    dma->buffer_length[1] = 0;
+
     dma->channel[0] = NUM_DMA_CHANNELS;
     dma->channel[1] = NUM_DMA_CHANNELS;
 }
 
 void audio_dma_deinit(audio_dma_t *dma) {
+    #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+    m_free(dma->buffer[0], dma->buffer_length[0]);
+    #else
     m_free(dma->buffer[0]);
+    #endif
+
     dma->buffer[0] = NULL;
+    dma->buffer_length[0] = 0;
 
+    #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+    m_free(dma->buffer[1], dma->buffer_length[1]);
+    #else
     m_free(dma->buffer[1]);
+    #endif
+
     dma->buffer[1] = NULL;
+    dma->buffer_length[1] = 0;
 }
 
 bool audio_dma_get_playing(audio_dma_t *dma) {

ports/stm/common-hal/audiopwmio/PWMAudioOut.c

@@ -206,6 +206,11 @@ void common_hal_audiopwmio_pwmaudioout_construct(audiopwmio_pwmaudioout_obj_t *s
     self->buffer[0] = NULL;
     self->buffer[1] = NULL;
 
+    #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+    self->buffer_size[0] = 0;
+    self->buffer_size[1] = 0;
+    #endif
+
     self->quiescent_value = quiescent_value;
 }
 
@@ -214,9 +219,22 @@ bool common_hal_audiopwmio_pwmaudioout_deinited(audiopwmio_pwmaudioout_obj_t *se
 }
 
 static void free_buffers(audiopwmio_pwmaudioout_obj_t *self) {
+    #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+    m_free(self->buffer[0], self->buffer_size[0]);
+    self->buffer_size[0] = 0;
+    #else
     m_free(self->buffer[0]);
+    #endif
+
     self->buffer[0] = NULL;
+
+    #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+    m_free(self->buffer[1], self->buffer_size[1]);
+    self->buffer_size[1] = 0;
+    #else
     m_free(self->buffer[1]);
+    #endif
+
     self->buffer[1] = NULL;
 }
 
@@ -257,11 +275,21 @@ void common_hal_audiopwmio_pwmaudioout_play(audiopwmio_pwmaudioout_obj_t *self,
     if (max_buffer_length > UINT16_MAX) {
         mp_raise_ValueError_varg(MP_ERROR_TEXT("Buffer length %d too big. It must be less than %d"), max_buffer_length, UINT16_MAX);
     }
+
     uint16_t buffer_length = (uint16_t)max_buffer_length / self->bytes_per_sample;
-    self->buffer[0] = m_malloc(buffer_length * sizeof(uint16_t));
+    size_t buffer_size = buffer_length * sizeof(uint16_t);
+
+    self->buffer[0] = m_malloc(buffer_size);
+    #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+    self->buffer_size[0] = buffer_size;
+    #endif
     self->buffer_ptr[0] = self->buffer_length[0] = 0;
+
     if (self->pin[1]) {
-        self->buffer[1] = m_malloc(buffer_length * sizeof(uint16_t));
+        self->buffer[1] = m_malloc(buffer_size);
+        #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+        self->buffer_size[1] = buffer_size;
+        #endif
         self->buffer_ptr[1] = self->buffer_length[1] = 0;
     }
 

ports/stm/common-hal/audiopwmio/PWMAudioOut.h

@@ -14,6 +14,9 @@ typedef struct {
     uint16_t quiescent_value;
 
     uint16_t *buffer[2];
+    #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+    uint16_t buffer_size[2]; // Keeps track of allocated size
+    #endif
     uint16_t buffer_length[2];
     uint16_t buffer_ptr[2];
 

ports/zephyr-cp/common-hal/wifi/ScannedNetworks.c

@@ -116,7 +116,12 @@ void wifi_scannednetworks_deinit(wifi_scannednetworks_obj_t *self) {
     // Free any results we don't need.
     while (!k_fifo_is_empty(&self->fifo)) {
         wifi_network_obj_t *entry = k_fifo_get(&self->fifo, K_NO_WAIT);
+
+        #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+        m_free(entry, sizeof(wifi_network_obj_t));
+        #else
         m_free(entry);
+        #endif
     }
     wifi_scannednetworks_done(self);
 }

shared-module/atexit/__init__.c

@@ -12,8 +12,13 @@ static size_t callback_len = 0;
 static atexit_callback_t *callback = NULL;
 
 void atexit_reset(void) {
-    callback_len = 0;
+    #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+    m_free(callback, callback_len * sizeof(atexit_callback_t));
+    #else
     m_free(callback);
+    #endif
+
+    callback_len = 0;
     callback = NULL;
 }
 
@@ -39,7 +44,13 @@ void shared_module_atexit_register(mp_obj_t *func, size_t n_args, const mp_obj_t
         cb.args[i] = kw_args->table[cb.n_kw].key;
         cb.args[i += 1] = kw_args->table[cb.n_kw].value;
     }
-    callback = (atexit_callback_t *)m_realloc(callback, (callback_len + 1) * sizeof(cb));
+
+    callback = (atexit_callback_t *)m_realloc(callback,
+        #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+        callback_len *sizeof(cb), // Old size
+        #endif
+        (callback_len + 1) * sizeof(cb));
+
     callback[callback_len++] = cb;
 }
 

shared-module/displayio/OnDiskBitmap.c

@@ -129,7 +129,12 @@ void common_hal_displayio_ondiskbitmap_construct(displayio_ondiskbitmap_t *self,
             for (uint16_t i = 0; i < number_of_colors; i++) {
                 common_hal_displayio_palette_set_color(palette, i, palette_data[i]);
             }
+
+            #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+            m_free(palette_data, palette_size);
+            #else
             m_free(palette_data);
+            #endif
         } else {
             common_hal_displayio_palette_set_color(palette, 0, 0x0);
             common_hal_displayio_palette_set_color(palette, 1, 0xffffff);

shared-module/keypad/EventQueue.c

@@ -39,7 +39,13 @@ mp_obj_t common_hal_keypad_eventqueue_get(keypad_eventqueue_obj_t *self) {
     if (result) {
         return event;
     }
+
+    #if MICROPY_MALLOC_USES_ALLOCATED_SIZE
+    m_free(event, sizeof(keypad_event_obj_t));
+    #else
     m_free(event);
+    #endif
+
     return MP_ROM_NONE;
 }