Skip to content

Commit 6e56bb6

Browse files
committed
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.
1 parent 371f4ba commit 6e56bb6

File tree

3 files changed

+25
-14
lines changed

3 files changed

+25
-14
lines changed

py/objfun.c

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,7 @@ STATIC void dump_args(const mp_obj_t *a, mp_uint_t sz) {
147147

148148
// With this macro you can tune the maximum number of function state bytes
149149
// that will be allocated on the stack. Any function that needs more
150-
// than this will use the heap.
150+
// than this will try to use the heap, with fallback to stack allocation.
151151
#define VM_MAX_STATE_ON_STACK (11 * sizeof(mp_uint_t))
152152

153153
// 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,
220220

221221
// allocate state for locals and stack
222222
mp_uint_t state_size = n_state * sizeof(mp_obj_t) + n_exc_stack * sizeof(mp_exc_stack_t);
223-
mp_code_state *code_state;
223+
mp_code_state *code_state = NULL;
224224
if (state_size > VM_MAX_STATE_ON_STACK) {
225-
code_state = m_new_obj_var(mp_code_state, byte, state_size);
226-
} else {
225+
code_state = m_new_obj_var_maybe(mp_code_state, byte, state_size);
226+
}
227+
if (code_state == NULL) {
227228
code_state = alloca(sizeof(mp_code_state) + state_size);
229+
state_size = 0; // indicate that we allocated using alloca
228230
}
229231

230232
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,
285287
}
286288

287289
// free the state if it was allocated on the heap
288-
if (state_size > VM_MAX_STATE_ON_STACK) {
290+
if (state_size != 0) {
289291
m_del_var(mp_code_state, byte, state_size, code_state);
290292
}
291293

tests/micropython/heapalloc.py

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,23 +2,31 @@
22

33
import gc
44

5-
def f(a):
5+
def f1(a):
66
print(a)
77

8-
def g(a, b=2):
8+
def f2(a, b=2):
99
print(a, b)
1010

11+
def f3(a, b, c, d):
12+
x1 = x2 = a
13+
x3 = x4 = b
14+
x5 = x6 = c
15+
x7 = x8 = d
16+
print(x1, x3, x5, x7, x2 + x4 + x6 + x8)
17+
1118
global_var = 1
1219

13-
def h():
20+
def test():
1421
global global_var
1522
global_var = 2 # set an existing global variable
1623
for i in range(2): # for loop
17-
f(i) # function call
18-
f(i * 2 + 1) # binary operation with small ints
19-
f(a=i) # keyword arguments
20-
g(i) # default arg (second one)
21-
g(i, i) # 2 args
24+
f1(i) # function call
25+
f1(i * 2 + 1) # binary operation with small ints
26+
f1(a=i) # keyword arguments
27+
f2(i) # default arg (second one)
28+
f2(i, i) # 2 args
29+
f3(1, 2, 3, 4) # function with lots of local state
2230

2331
# call h with heap allocation disabled and all memory used up
2432
gc.disable()
@@ -27,5 +35,5 @@ def h():
2735
'a'.lower # allocates 1 cell for boundmeth
2836
except MemoryError:
2937
pass
30-
h()
38+
test()
3139
gc.enable()

tests/micropython/heapalloc.py.exp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,3 +8,4 @@
88
1
99
1 2
1010
1 1
11+
1 2 3 4 10

0 commit comments

Comments
 (0)