From 732407f1bf12364375162e8fb73816532a5d139c Mon Sep 17 00:00:00 2001
From: Damien <damien.p.george@gmail.com>
Date: Sun, 29 Dec 2013 19:33:23 +0000
Subject: [PATCH] Change memory allocation API to require size for free and
 realloc.

---
 py/asmthumb.c  |  4 ++--
 py/asmx64.c    |  2 +-
 py/builtin.c   |  2 +-
 py/compile.c   |  6 +++---
 py/emitpass1.c |  2 +-
 py/lexer.c     |  4 ++--
 py/lexerunix.c |  6 +++---
 py/malloc.c    | 22 +++++++++++-----------
 py/map.c       |  4 ++--
 py/misc.h      | 12 ++++++++----
 py/objclass.c  |  2 +-
 py/objfun.c    |  2 +-
 py/objlist.c   |  2 +-
 py/objstr.c    |  9 +++++----
 py/parse.c     |  2 +-
 py/qstr.c      |  6 +++---
 py/scope.c     |  2 +-
 py/vstr.c      |  9 ++++-----
 stm/lexerstm.c |  2 +-
 stm/main.c     |  2 +-
 stm/pybwlan.c  |  2 +-
 unix/main.c    |  5 +++--
 22 files changed, 57 insertions(+), 52 deletions(-)

diff --git a/py/asmthumb.c b/py/asmthumb.c
index a6fe4be809..ee8041ac9f 100644
--- a/py/asmthumb.c
+++ b/py/asmthumb.c
@@ -44,7 +44,7 @@ asm_thumb_t *asm_thumb_new(uint max_num_labels) {
 
 void asm_thumb_free(asm_thumb_t *as, bool free_code) {
     if (free_code) {
-        m_free(as->code_base);
+        m_del(byte, as->code_base, as->code_size);
     }
     /*
     if (as->label != NULL) {
@@ -58,7 +58,7 @@ void asm_thumb_free(asm_thumb_t *as, bool free_code) {
         g_array_free(as->label, true);
     }
     */
-    m_free(as);
+    m_del_obj(asm_thumb_t, as);
 }
 
 void asm_thumb_start_pass(asm_thumb_t *as, int pass) {
diff --git a/py/asmx64.c b/py/asmx64.c
index c7e3cdc849..2a11de207f 100644
--- a/py/asmx64.c
+++ b/py/asmx64.c
@@ -128,7 +128,7 @@ void asm_x64_free(asm_x64_t* as, bool free_code) {
         g_array_free(as->label, true);
     }
     */
-    m_free(as);
+    m_del_obj(asm_x64_t, as);
 }
 
 void asm_x64_start_pass(asm_x64_t *as, int pass) {
diff --git a/py/builtin.c b/py/builtin.c
index 7d2cc72b77..f6b57109c3 100644
--- a/py/builtin.c
+++ b/py/builtin.c
@@ -158,7 +158,7 @@ mp_obj_t mp_builtin_chr(mp_obj_t o_in) {
         char *str = m_new(char, 2);
         str[0] = ord;
         str[1] = '\0';
-        return mp_obj_new_str(qstr_from_str_take(str));
+        return mp_obj_new_str(qstr_from_str_take(str, 2));
     } else {
         nlr_jump(mp_obj_new_exception_msg(rt_q_ValueError, "chr() arg not in range(0x110000)"));
     }
diff --git a/py/compile.c b/py/compile.c
index 38f14155db..604d5a2437 100644
--- a/py/compile.c
+++ b/py/compile.c
@@ -1184,7 +1184,7 @@ void do_import_name(compiler_t *comp, mp_parse_node_t pn, qstr *q1, qstr *q2) {
                 }
                 strcat(str, qstr_str(MP_PARSE_NODE_LEAF_ARG(pns->nodes[i])));
             }
-            *q2 = qstr_from_str_take(str);
+            *q2 = qstr_from_str_take(str, len + 1);
             EMIT(import_name, *q2);
             if (is_as) {
                 for (int i = 1; i < n; i++) {
@@ -2127,7 +2127,7 @@ void compile_atom_string(compiler_t *comp, mp_parse_node_struct_t *pns) {
         strcat(cat_str, str);
     }
 
-    EMIT(load_const_str, qstr_from_str_take(cat_str), string_kind == MP_PARSE_NODE_BYTES);
+    EMIT(load_const_str, qstr_from_str_take(cat_str, n_bytes + 1), string_kind == MP_PARSE_NODE_BYTES);
 }
 
 // pns needs to have 2 nodes, first is lhs of comprehension, second is PN_comp_for node
@@ -3105,7 +3105,7 @@ bool mp_compile(mp_parse_node_t pn, bool is_repl) {
         }
     }
 
-    m_free(comp);
+    m_del_obj(compiler_t, comp);
 
     return !comp->had_error;
 }
diff --git a/py/emitpass1.c b/py/emitpass1.c
index 551f30eb4c..4ed0549727 100644
--- a/py/emitpass1.c
+++ b/py/emitpass1.c
@@ -26,7 +26,7 @@ emit_t *emit_pass1_new(qstr qstr___class__) {
 }
 
 void emit_pass1_free(emit_t *emit) {
-    m_free(emit);
+    m_del_obj(emit_t, emit);
 }
 
 static void emit_pass1_dummy(emit_t *emit) {
diff --git a/py/lexer.c b/py/lexer.c
index cf9eae5531..6e43c7469a 100644
--- a/py/lexer.c
+++ b/py/lexer.c
@@ -183,8 +183,8 @@ static void next_char(mp_lexer_t *lex) {
 
 void indent_push(mp_lexer_t *lex, uint indent) {
     if (lex->num_indent_level >= lex->alloc_indent_level) {
+        lex->indent_level = m_renew(uint16_t, lex->indent_level, lex->alloc_indent_level, lex->alloc_indent_level * 2);
         lex->alloc_indent_level *= 2;
-        lex->indent_level = m_renew(uint16_t, lex->indent_level, lex->alloc_indent_level);
     }
     lex->indent_level[lex->num_indent_level++] = indent;
 }
@@ -639,7 +639,7 @@ void mp_lexer_free(mp_lexer_t *lex) {
             lex->stream_close(lex->stream_data);
         }
         vstr_clear(&lex->vstr);
-        m_free(lex);
+        m_del_obj(mp_lexer_t, lex);
     }
 }
 
diff --git a/py/lexerunix.c b/py/lexerunix.c
index 80daf6009a..398cb792a7 100644
--- a/py/lexerunix.c
+++ b/py/lexerunix.c
@@ -24,9 +24,9 @@ unichar str_buf_next_char(str_buf_t *sb) {
 void str_buf_free(str_buf_t *sb) {
     if (sb) {
         if (sb->free) {
-            m_free((char*)sb->src_beg);
+            m_del(char, (char*)sb->src_beg, 0 /* unknown size of src_beg */);
         }
-        m_free(sb);
+        m_del_obj(str_buf_t, sb);
     }
 }
 
@@ -52,7 +52,7 @@ mp_lexer_t *mp_lexer_new_from_file(const char *filename) {
     close(fd);
     if (read_size != size) {
         printf("error reading file %s\n", filename);
-        m_free(data);
+        m_del(char, data, size);
         return NULL;
     }
 
diff --git a/py/malloc.c b/py/malloc.c
index 2f8b5f78b5..c65d38a968 100644
--- a/py/malloc.c
+++ b/py/malloc.c
@@ -5,12 +5,6 @@
 
 static int total_bytes_allocated = 0;
 
-void m_free(void *ptr) {
-    if (ptr != NULL) {
-        free(ptr);
-    }
-}
-
 void *m_malloc(int num_bytes) {
     if (num_bytes == 0) {
         return NULL;
@@ -37,20 +31,26 @@ void *m_malloc0(int num_bytes) {
     return ptr;
 }
 
-void *m_realloc(void *ptr, int num_bytes) {
-    if (num_bytes == 0) {
+void *m_realloc(void *ptr, int old_num_bytes, int new_num_bytes) {
+    if (new_num_bytes == 0) {
         free(ptr);
         return NULL;
     }
-    ptr = realloc(ptr, num_bytes);
+    ptr = realloc(ptr, new_num_bytes);
     if (ptr == NULL) {
-        printf("could not allocate memory, reallocating %d bytes\n", num_bytes);
+        printf("could not allocate memory, reallocating %d bytes\n", new_num_bytes);
         return NULL;
     }
-    total_bytes_allocated += num_bytes;
+    total_bytes_allocated += new_num_bytes;
     return ptr;
 }
 
+void m_free(void *ptr, int num_bytes) {
+    if (ptr != NULL) {
+        free(ptr);
+    }
+}
+
 int m_get_total_bytes_allocated(void) {
     return total_bytes_allocated;
 }
diff --git a/py/map.c b/py/map.c
index 3c2d1fbb55..01209c9b74 100644
--- a/py/map.c
+++ b/py/map.c
@@ -64,7 +64,7 @@ mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_n
                             mp_map_lookup_helper(map, old_table[i].key, true)->value = old_table[i].value;
                         }
                     }
-                    m_free(old_table);
+                    m_del(mp_map_elem_t, old_table, old_alloc);
                     // restart the search for the new element
                     pos = hash % map->alloc;
                 } else {
@@ -124,7 +124,7 @@ mp_obj_t mp_set_lookup(mp_set_t *set, mp_obj_t index, bool add_if_not_found) {
                             mp_set_lookup(set, old_table[i], true);
                         }
                     }
-                    m_free(old_table);
+                    m_del(mp_obj_t, old_table, old_alloc);
                     // restart the search for the new element
                     pos = hash % set->alloc;
                 } else {
diff --git a/py/misc.h b/py/misc.h
index 6c0e18ca0f..1a33f0505d 100644
--- a/py/misc.h
+++ b/py/misc.h
@@ -16,16 +16,20 @@ typedef unsigned int uint;
 
 /** memomry allocation ******************************************/
 
+// TODO make a lazy m_renew that can increase by a smaller amount than requested (but by at least 1 more element)
+
 #define m_new(type, num) ((type*)(m_malloc(sizeof(type) * (num))))
 #define m_new0(type, num) ((type*)(m_malloc0(sizeof(type) * (num))))
-#define m_renew(type, ptr, num) ((type*)(m_realloc((ptr), sizeof(type) * (num))))
 #define m_new_obj(type) (m_new(type, 1))
 #define m_new_obj_var(obj_type, var_type, var_num) ((obj_type*)m_malloc(sizeof(obj_type) + sizeof(var_type) * (var_num)))
+#define m_renew(type, ptr, old_num, new_num) ((type*)(m_realloc((ptr), sizeof(type) * (old_num), sizeof(type) * (new_num))))
+#define m_del(type, ptr, num) m_free(ptr, sizeof(type) * (num))
+#define m_del_obj(type, ptr) (m_del(type, ptr, 1))
 
-void m_free(void *ptr);
 void *m_malloc(int num_bytes);
 void *m_malloc0(int num_bytes);
-void *m_realloc(void *ptr, int num_bytes);
+void *m_realloc(void *ptr, int old_num_bytes, int new_num_bytes);
+void m_free(void *ptr, int num_bytes);
 
 int m_get_total_bytes_allocated(void);
 
@@ -96,7 +100,7 @@ typedef unsigned int qstr;
 
 void qstr_init(void);
 qstr qstr_from_str_static(const char *str);
-qstr qstr_from_str_take(char *str);
+qstr qstr_from_str_take(char *str, int alloc_len);
 qstr qstr_from_strn_copy(const char *str, int len);
 const char* qstr_str(qstr qstr);
 
diff --git a/py/objclass.c b/py/objclass.c
index 52e43724f7..a4bde3f001 100644
--- a/py/objclass.c
+++ b/py/objclass.c
@@ -37,7 +37,7 @@ mp_obj_t class_call_n(mp_obj_t self_in, int n_args, const mp_obj_t *args) {
             memcpy(args2, args, n_args * sizeof(mp_obj_t));
             args2[n_args] = o;
             init_ret = rt_call_function_n(init_fn->value, n_args + 1, args2);
-            m_free(args2);
+            m_del(mp_obj_t, args2, n_args + 1);
         }
         if (init_ret != mp_const_none) {
             nlr_jump(mp_obj_new_exception_msg_1_arg(rt_q_TypeError, "__init__() should return None, not '%s'", mp_obj_get_type_str(init_ret)));
diff --git a/py/objfun.c b/py/objfun.c
index bf11b5e6ab..cefc9a95fe 100644
--- a/py/objfun.c
+++ b/py/objfun.c
@@ -58,7 +58,7 @@ mp_obj_t fun_native_call_n(mp_obj_t self_in, int n_args, const mp_obj_t *args) {
         }
 
         mp_obj_t res = ((mp_fun_var_t)self->fun)(n_args, args_ordered);
-        m_free(args_ordered);
+        m_del(mp_obj_t, args_ordered, n_args);
 
         return res;
     }
diff --git a/py/objlist.c b/py/objlist.c
index 4d70c37eb3..ce16ab6fe9 100644
--- a/py/objlist.c
+++ b/py/objlist.c
@@ -57,8 +57,8 @@ mp_obj_t mp_obj_list_append(mp_obj_t self_in, mp_obj_t arg) {
     assert(MP_OBJ_IS_TYPE(self_in, &list_type));
     mp_obj_list_t *self = self_in;
     if (self->len >= self->alloc) {
+        self->items = m_renew(mp_obj_t, self->items, self->alloc, self->alloc * 2);
         self->alloc *= 2;
-        self->items = m_renew(mp_obj_t, self->items, self->alloc);
     }
     self->items[self->len++] = arg;
     return mp_const_none; // return None, as per CPython
diff --git a/py/objstr.c b/py/objstr.c
index d82ea3d3ce..104a54bb32 100644
--- a/py/objstr.c
+++ b/py/objstr.c
@@ -36,9 +36,10 @@ mp_obj_t str_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
             if (MP_OBJ_IS_TYPE(rhs_in, &str_type)) {
                 // add 2 strings
                 const char *rhs_str = qstr_str(((mp_obj_str_t*)rhs_in)->qstr);
-                char *val = m_new(char, strlen(lhs_str) + strlen(rhs_str) + 1);
+                int alloc_len = strlen(lhs_str) + strlen(rhs_str) + 1;
+                char *val = m_new(char, alloc_len);
                 stpcpy(stpcpy(val, lhs_str), rhs_str);
-                return mp_obj_new_str(qstr_from_str_take(val));
+                return mp_obj_new_str(qstr_from_str_take(val, alloc_len));
             }
             break;
     }
@@ -78,7 +79,7 @@ mp_obj_t str_join(mp_obj_t self_in, mp_obj_t arg) {
         }
         strcat(joined_str, s2);
     }
-    return mp_obj_new_str(qstr_from_str_take(joined_str));
+    return mp_obj_new_str(qstr_from_str_take(joined_str, required_len + 1));
 
 bad_arg:
     nlr_jump(mp_obj_new_exception_msg(rt_q_TypeError, "?str.join expecting a list of str's"));
@@ -115,7 +116,7 @@ mp_obj_t str_format(int n_args, const mp_obj_t *args) {
         }
     }
 
-    return mp_obj_new_str(qstr_from_str_take(vstr->buf));
+    return mp_obj_new_str(qstr_from_str_take(vstr->buf, vstr->alloc));
 }
 
 static MP_DEFINE_CONST_FUN_OBJ_2(str_join_obj, str_join);
diff --git a/py/parse.c b/py/parse.c
index b939ab8db7..ad9a47ffbc 100644
--- a/py/parse.c
+++ b/py/parse.c
@@ -94,8 +94,8 @@ typedef struct _parser_t {
 
 static void push_rule(parser_t *parser, const rule_t *rule, int arg_i) {
     if (parser->rule_stack_top >= parser->rule_stack_alloc) {
+        parser->rule_stack = m_renew(rule_stack_t, parser->rule_stack, parser->rule_stack_alloc, parser->rule_stack_alloc * 2);
         parser->rule_stack_alloc *= 2;
-        parser->rule_stack = m_renew(rule_stack_t, parser->rule_stack, parser->rule_stack_alloc);
     }
     parser->rule_stack[parser->rule_stack_top].rule_id = rule->rule_id;
     parser->rule_stack[parser->rule_stack_top].arg_i = arg_i;
diff --git a/py/qstr.c b/py/qstr.c
index 6fa69ddb95..0dd8a04b70 100644
--- a/py/qstr.c
+++ b/py/qstr.c
@@ -16,8 +16,8 @@ void qstr_init(void) {
 
 static qstr qstr_add(const char *str) {
     if (qstrs_len >= qstrs_alloc) {
+        qstrs = m_renew(const char*, qstrs, qstrs_alloc, qstrs_alloc * 2);
         qstrs_alloc *= 2;
-        qstrs = m_renew(const char*, qstrs, qstrs_alloc);
     }
     qstrs[qstrs_len++] = str;
     return qstrs_len - 1;
@@ -32,10 +32,10 @@ qstr qstr_from_str_static(const char *str) {
     return qstr_add(str);
 }
 
-qstr qstr_from_str_take(char *str) {
+qstr qstr_from_str_take(char *str, int alloc_len) {
     for (int i = 0; i < qstrs_len; i++) {
         if (strcmp(qstrs[i], str) == 0) {
-            m_free(str);
+            m_del(char, str, alloc_len);
             return i;
         }
     }
diff --git a/py/scope.c b/py/scope.c
index 7d55d3923f..38ea5a9e2f 100644
--- a/py/scope.c
+++ b/py/scope.c
@@ -69,8 +69,8 @@ id_info_t *scope_find_or_add_id(scope_t *scope, qstr qstr, bool *added) {
 
     // make sure we have enough memory
     if (scope->id_info_len >= scope->id_info_alloc) {
+        scope->id_info = m_renew(id_info_t, scope->id_info, scope->id_info_alloc, scope->id_info_alloc * 2);
         scope->id_info_alloc *= 2;
-        scope->id_info = m_renew(id_info_t, scope->id_info, scope->id_info_alloc);
     }
 
     id_info_t *id_info;
diff --git a/py/vstr.c b/py/vstr.c
index f6652e6f10..98cf027250 100644
--- a/py/vstr.c
+++ b/py/vstr.c
@@ -11,7 +11,6 @@ void vstr_init(vstr_t *vstr) {
     vstr->len = 0;
     vstr->buf = m_new(char, vstr->alloc);
     if (vstr->buf == NULL) {
-        m_free(vstr);
         vstr->had_error = true;
         return;
     }
@@ -20,7 +19,7 @@ void vstr_init(vstr_t *vstr) {
 }
 
 void vstr_clear(vstr_t *vstr) {
-    m_free(vstr->buf);
+    m_del(char, vstr->buf, vstr->alloc);
     vstr->buf = NULL;
 }
 
@@ -35,8 +34,8 @@ vstr_t *vstr_new(void) {
 
 void vstr_free(vstr_t *vstr) {
     if (vstr != NULL) {
-        m_free(vstr->buf);
-        m_free(vstr);
+        m_del(char, vstr->buf, vstr->alloc);
+        m_del_obj(vstr_t, vstr);
     }
 }
 
@@ -67,7 +66,7 @@ int vstr_len(vstr_t *vstr) {
 bool vstr_ensure_extra(vstr_t *vstr, int size) {
     if (vstr->len + size + 1 > vstr->alloc) {
         int new_alloc = ROUND_ALLOC((vstr->len + size + 1) * 2);
-        char *new_buf = m_renew(char, vstr->buf, new_alloc);
+        char *new_buf = m_renew(char, vstr->buf, vstr->alloc, new_alloc);
         if (new_buf == NULL) {
             vstr->had_error = true;
             return false;
diff --git a/stm/lexerstm.c b/stm/lexerstm.c
index 4e99052242..dfb84cca13 100644
--- a/stm/lexerstm.c
+++ b/stm/lexerstm.c
@@ -17,7 +17,7 @@ unichar str_buf_next_char(mp_lexer_str_buf_t *sb) {
 
 void str_buf_free(mp_lexer_str_buf_t *sb) {
     if (sb->free) {
-        m_free((char*)sb->src_beg);
+        m_del(char, (char*)sb->src_beg, 0 /* don't know allocated size of src */);
     }
 }
 
diff --git a/stm/main.c b/stm/main.c
index 22ce46f256..8c2c96af2a 100644
--- a/stm/main.c
+++ b/stm/main.c
@@ -706,7 +706,7 @@ mp_obj_t file_obj_read(mp_obj_t self_in, mp_obj_t arg) {
     UINT n_out;
     f_read(&self->fp, buf, n, &n_out);
     buf[n_out] = 0;
-    return mp_obj_new_str(qstr_from_str_take(buf));
+    return mp_obj_new_str(qstr_from_str_take(buf, n + 1));
 }
 
 mp_obj_t file_obj_write(mp_obj_t self_in, mp_obj_t arg) {
diff --git a/stm/pybwlan.c b/stm/pybwlan.c
index c078ecfe14..73aa8273ff 100644
--- a/stm/pybwlan.c
+++ b/stm/pybwlan.c
@@ -191,7 +191,7 @@ mp_obj_t pyb_wlan_http_get(mp_obj_t host_name, mp_obj_t host_path) {
             vstr_add_strn(vstr, buf, ret);
         }
 
-        mp_ret = mp_obj_new_str(qstr_from_str_take(vstr_str(vstr)));
+        mp_ret = mp_obj_new_str(qstr_from_str_take(vstr->buf, vstr->alloc));
     }
 
     closesocket(sd);
diff --git a/unix/main.c b/unix/main.c
index 16012e7181..73da1ecfc2 100644
--- a/unix/main.c
+++ b/unix/main.c
@@ -1,6 +1,7 @@
 #include <stdint.h>
 #include <stdio.h>
 #include <string.h>
+#include <malloc.h>
 
 #include "nlr.h"
 #include "misc.h"
@@ -44,8 +45,8 @@ static void do_repl(void) {
                     break;
                 }
                 char *line3 = str_join(line, '\n', line2);
-                m_free(line);
-                m_free(line2);
+                free(line);
+                free(line2);
                 line = line3;
             }
         }