py: Add m_malloc_fail function to handle memory allocation error.

dpgeorge · dpgeorge · commit 6902eeda259a · 2014-04-04T10:52:59.000Z
A malloc/realloc fail now throws MemoryError.
diff --git a/py/malloc.c b/py/malloc.c
@@ -41,8 +41,7 @@ void *m_malloc(int num_bytes) {
     }
     void *ptr = malloc(num_bytes);
     if (ptr == NULL) {
-        printf("could not allocate memory, allocating %d bytes\n", num_bytes);
-        return NULL;
+        return m_malloc_fail(num_bytes);
     }
 #if MICROPY_MEM_STATS
     total_bytes_allocated += num_bytes;
@@ -68,8 +67,7 @@ void *m_realloc(void *ptr, int old_num_bytes, int new_num_bytes) {
     }
     void *new_ptr = realloc(ptr, new_num_bytes);
     if (new_ptr == NULL) {
-        printf("could not allocate memory, reallocating %d bytes\n", new_num_bytes);
-        return NULL;
+        return m_malloc_fail(new_num_bytes);
     }
 #if MICROPY_MEM_STATS
     // At first thought, "Total bytes allocated" should only grow,
diff --git a/py/misc.h b/py/misc.h
@@ -36,6 +36,7 @@ void *m_malloc(int num_bytes);
 void *m_malloc0(int num_bytes);
 void *m_realloc(void *ptr, int old_num_bytes, int new_num_bytes);
 void m_free(void *ptr, int num_bytes);
+void *m_malloc_fail(int num_bytes);
 
 int m_get_total_bytes_allocated(void);
 int m_get_current_bytes_allocated(void);
diff --git a/py/obj.h b/py/obj.h
@@ -306,6 +306,7 @@ extern const struct _mp_obj_bool_t mp_const_false_obj;
 extern const struct _mp_obj_bool_t mp_const_true_obj;
 extern const struct _mp_obj_tuple_t mp_const_empty_tuple_obj;
 extern const struct _mp_obj_ellipsis_t mp_const_ellipsis_obj;
+extern const struct _mp_obj_exception_t mp_const_MemoryError_obj;
 extern const struct _mp_obj_exception_t mp_const_GeneratorExit_obj;
 
 // General API for objects
diff --git a/py/objexcept.c b/py/objexcept.c
@@ -17,6 +17,9 @@ typedef struct _mp_obj_exception_t {
     mp_obj_tuple_t args;
 } mp_obj_exception_t;
 
+// Instance of MemoryError exception - needed by mp_malloc_fail
+const mp_obj_exception_t mp_const_MemoryError_obj = {{&mp_type_MemoryError}, MP_OBJ_NULL, {{&mp_type_tuple}, 0}};
+
 // Instance of GeneratorExit exception - needed by generator.close()
 // This would belong to objgenerator.c, but to keep mp_obj_exception_t
 // definition module-private so far, have it here.
diff --git a/py/runtime.c b/py/runtime.c
@@ -1034,6 +1034,11 @@ void mp_globals_set(mp_map_t *m) {
     map_globals = m;
 }
 
+void *m_malloc_fail(int num_bytes) {
+    DEBUG_printf("memory allocation failed, allocating %d bytes\n", num_bytes);
+    nlr_jump((mp_obj_t)&mp_const_MemoryError_obj);
+}
+
 // these must correspond to the respective enum
 void *const mp_fun_table[MP_F_NUMBER_OF] = {
     mp_load_const_dec,
diff --git a/py/vm.c b/py/vm.c
@@ -849,8 +849,8 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
             // set file and line number that the exception occurred at
             // TODO: don't set traceback for exceptions re-raised by END_FINALLY.
             // But consider how to handle nested exceptions.
-            // TODO need a better way of not adding traceback to constant objects (right now, just GeneratorExit_obj)
-            if (mp_obj_is_exception_instance(nlr.ret_val) && nlr.ret_val != &mp_const_GeneratorExit_obj) {
+            // TODO need a better way of not adding traceback to constant objects (right now, just GeneratorExit_obj and MemoryError_obj)
+            if (mp_obj_is_exception_instance(nlr.ret_val) && nlr.ret_val != &mp_const_GeneratorExit_obj && nlr.ret_val != &mp_const_MemoryError_obj) {
                 machine_uint_t code_info_size = code_info[0] | (code_info[1] << 8) | (code_info[2] << 16) | (code_info[3] << 24);
                 qstr source_file = code_info[4] | (code_info[5] << 8) | (code_info[6] << 16) | (code_info[7] << 24);
                 qstr block_name = code_info[8] | (code_info[9] << 8) | (code_info[10] << 16) | (code_info[11] << 24);
diff --git a/tests/basics/memoryerror.py b/tests/basics/memoryerror.py
@@ -0,0 +1,6 @@
+l = list(range(10000))
+try:
+    100000000 * l
+except MemoryError:
+    print('MemoryError')
+print(len(l), l[0], l[-1])

Original file line number	Diff line number	Diff line change
`@@ -41,8 +41,7 @@ void *m_malloc(int num_bytes) {`
`41`	`41`	`}`
`42`	`42`	`void *ptr = malloc(num_bytes);`
`43`	`43`	`if (ptr == NULL) {`
`44`		`- printf("could not allocate memory, allocating %d bytes\n", num_bytes);`
`45`		`- return NULL;`
	`44`	`+ return m_malloc_fail(num_bytes);`
`46`	`45`	`}`
`47`	`46`	`#if MICROPY_MEM_STATS`
`48`	`47`	`total_bytes_allocated += num_bytes;`
`@@ -68,8 +67,7 @@ void m_realloc(void ptr, int old_num_bytes, int new_num_bytes) {`
`68`	`67`	`}`
`69`	`68`	`void *new_ptr = realloc(ptr, new_num_bytes);`
`70`	`69`	`if (new_ptr == NULL) {`
`71`		`- printf("could not allocate memory, reallocating %d bytes\n", new_num_bytes);`
`72`		`- return NULL;`
	`70`	`+ return m_malloc_fail(new_num_bytes);`
`73`	`71`	`}`
`74`	`72`	`#if MICROPY_MEM_STATS`
`75`	`73`	`// At first thought, "Total bytes allocated" should only grow,`