diff --git a/py/vm.c b/py/vm.c
index 7b3a0b3241..8abe65e43b 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -911,21 +911,22 @@ unwind_jump:;
                         code_state->sp = sp;
                         code_state->exc_sp = MP_TAGPTR_MAKE(exc_sp, currently_in_except_block);
                         mp_code_state_t *new_state = mp_obj_fun_bc_prepare_codestate(*sp, unum & 0xff, (unum >> 8) & 0xff, sp + 1);
-                        if (new_state) {
+                        #if !MICROPY_ENABLE_PYSTACK
+                        if (new_state == NULL) {
+                            // Couldn't allocate codestate on heap: in the strict case raise
+                            // an exception, otherwise just fall through to stack allocation.
+                            #if MICROPY_STACKLESS_STRICT
+                        deep_recursion_error:
+                            mp_raise_recursion_depth();
+                            #endif
+                        } else
+                        #endif
+                        {
                             new_state->prev = code_state;
                             code_state = new_state;
                             nlr_pop();
                             goto run_code_state;
                         }
-                        #if MICROPY_STACKLESS_STRICT
-                        else {
-                        deep_recursion_error:
-                            mp_raise_recursion_depth();
-                        }
-                        #else
-                        // If we couldn't allocate codestate on heap, in
-                        // non non-strict case fall thru to stack allocation.
-                        #endif
                     }
                     #endif
                     SET_TOP(mp_call_function_n_kw(*sp, unum & 0xff, (unum >> 8) & 0xff, sp + 1));
@@ -956,20 +957,21 @@ unwind_jump:;
                         // pystack is not enabled.  For pystack, they are freed when code_state is.
                         mp_nonlocal_free(out_args.args, out_args.n_alloc * sizeof(mp_obj_t));
                         #endif
-                        if (new_state) {
+                        #if !MICROPY_ENABLE_PYSTACK
+                        if (new_state == NULL) {
+                            // Couldn't allocate codestate on heap: in the strict case raise
+                            // an exception, otherwise just fall through to stack allocation.
+                            #if MICROPY_STACKLESS_STRICT
+                            goto deep_recursion_error;
+                            #endif
+                        } else
+                        #endif
+                        {
                             new_state->prev = code_state;
                             code_state = new_state;
                             nlr_pop();
                             goto run_code_state;
                         }
-                        #if MICROPY_STACKLESS_STRICT
-                        else {
-                            goto deep_recursion_error;
-                        }
-                        #else
-                        // If we couldn't allocate codestate on heap, in
-                        // non non-strict case fall thru to stack allocation.
-                        #endif
                     }
                     #endif
                     SET_TOP(mp_call_method_n_kw_var(false, unum, sp));
@@ -993,20 +995,21 @@ unwind_jump:;
                         int adjust = (sp[1] == MP_OBJ_NULL) ? 0 : 1;
 
                         mp_code_state_t *new_state = mp_obj_fun_bc_prepare_codestate(*sp, n_args + adjust, n_kw, sp + 2 - adjust);
-                        if (new_state) {
+                        #if !MICROPY_ENABLE_PYSTACK
+                        if (new_state == NULL) {
+                            // Couldn't allocate codestate on heap: in the strict case raise
+                            // an exception, otherwise just fall through to stack allocation.
+                            #if MICROPY_STACKLESS_STRICT
+                            goto deep_recursion_error;
+                            #endif
+                        } else
+                        #endif
+                        {
                             new_state->prev = code_state;
                             code_state = new_state;
                             nlr_pop();
                             goto run_code_state;
                         }
-                        #if MICROPY_STACKLESS_STRICT
-                        else {
-                            goto deep_recursion_error;
-                        }
-                        #else
-                        // If we couldn't allocate codestate on heap, in
-                        // non non-strict case fall thru to stack allocation.
-                        #endif
                     }
                     #endif
                     SET_TOP(mp_call_method_n_kw(unum & 0xff, (unum >> 8) & 0xff, sp));
@@ -1037,20 +1040,21 @@ unwind_jump:;
                         // pystack is not enabled.  For pystack, they are freed when code_state is.
                         mp_nonlocal_free(out_args.args, out_args.n_alloc * sizeof(mp_obj_t));
                         #endif
-                        if (new_state) {
+                        #if !MICROPY_ENABLE_PYSTACK
+                        if (new_state == NULL) {
+                            // Couldn't allocate codestate on heap: in the strict case raise
+                            // an exception, otherwise just fall through to stack allocation.
+                            #if MICROPY_STACKLESS_STRICT
+                            goto deep_recursion_error;
+                            #endif
+                        } else
+                        #endif
+                        {
                             new_state->prev = code_state;
                             code_state = new_state;
                             nlr_pop();
                             goto run_code_state;
                         }
-                        #if MICROPY_STACKLESS_STRICT
-                        else {
-                            goto deep_recursion_error;
-                        }
-                        #else
-                        // If we couldn't allocate codestate on heap, in
-                        // non non-strict case fall thru to stack allocation.
-                        #endif
                     }
                     #endif
                     SET_TOP(mp_call_method_n_kw_var(true, unum, sp));