py: Fallback to stack alloca for Python-stack if heap alloc fails.

If heap allocation for the Python-stack of a function fails then we may
as well allocate the Python-stack on the C stack.  This will allow to
run more code without using the heap.

3 files changed, 25 insertions, 14 deletions

diff --git a/py/objfun.c b/py/objfun.c
index 487c432a61..ff37da1463 100644
--- a/py/objfun.c
+++ b/py/objfun.c
@@ -147,7 +147,7 @@ STATIC void dump_args(const mp_obj_t *a, mp_uint_t sz) {
 
 // With this macro you can tune the maximum number of function state bytes
 // that will be allocated on the stack.  Any function that needs more
-// than this will use the heap.
+// than this will try to use the heap, with fallback to stack allocation.
 #define VM_MAX_STATE_ON_STACK (11 * sizeof(mp_uint_t))
 
 // Set this to enable a simple stack overflow check.
@@ -220,11 +220,13 @@ STATIC mp_obj_t fun_bc_call(mp_obj_t self_in, mp_uint_t n_args, mp_uint_t n_kw,
 
     // allocate state for locals and stack
     mp_uint_t state_size = n_state * sizeof(mp_obj_t) + n_exc_stack * sizeof(mp_exc_stack_t);
-    mp_code_state *code_state;
+    mp_code_state *code_state = NULL;
     if (state_size > VM_MAX_STATE_ON_STACK) {
-        code_state = m_new_obj_var(mp_code_state, byte, state_size);
-    } else {
+        code_state = m_new_obj_var_maybe(mp_code_state, byte, state_size);
+    }
+    if (code_state == NULL) {
         code_state = alloca(sizeof(mp_code_state) + state_size);
+        state_size = 0; // indicate that we allocated using alloca
     }
 
     code_state->n_state = n_state;
@@ -285,7 +287,7 @@ STATIC mp_obj_t fun_bc_call(mp_obj_t self_in, mp_uint_t n_args, mp_uint_t n_kw,
     }
 
     // free the state if it was allocated on the heap
-    if (state_size > VM_MAX_STATE_ON_STACK) {
+    if (state_size != 0) {
         m_del_var(mp_code_state, byte, state_size, code_state);
     }
 
diff --git a/tests/micropython/heapalloc.py b/tests/micropython/heapalloc.py
index b4b27d19b7..2dc7fa5e7e 100644
--- a/tests/micropython/heapalloc.py
+++ b/tests/micropython/heapalloc.py
@@ -2,23 +2,31 @@
 
 import gc
 
-def f(a):
+def f1(a):
     print(a)
 
-def g(a, b=2):
+def f2(a, b=2):
     print(a, b)
 
+def f3(a, b, c, d):
+    x1 = x2 = a
+    x3 = x4 = b
+    x5 = x6 = c
+    x7 = x8 = d
+    print(x1, x3, x5, x7, x2 + x4 + x6 + x8)
+
 global_var = 1
 
-def h():
+def test():
     global global_var
     global_var = 2      # set an existing global variable
     for i in range(2):  # for loop
-        f(i)            # function call
-        f(i * 2 + 1)    # binary operation with small ints
-        f(a=i)          # keyword arguments
-        g(i)            # default arg (second one)
-        g(i, i)         # 2 args
+        f1(i)           # function call
+        f1(i * 2 + 1)   # binary operation with small ints
+        f1(a=i)         # keyword arguments
+        f2(i)           # default arg (second one)
+        f2(i, i)        # 2 args
+    f3(1, 2, 3, 4)  # function with lots of local state
 
 # call h with heap allocation disabled and all memory used up
 gc.disable()
@@ -27,5 +35,5 @@ try:
         'a'.lower # allocates 1 cell for boundmeth
 except MemoryError:
     pass
-h()
+test()
 gc.enable()
diff --git a/tests/micropython/heapalloc.py.exp b/tests/micropython/heapalloc.py.exp
index f0487b9f62..c8cffe183f 100644
--- a/tests/micropython/heapalloc.py.exp
+++ b/tests/micropython/heapalloc.py.exp
@@ -8,3 +8,4 @@
 1
 1 2
 1 1
+1 2 3 4 10