py/vm: Avoid heap-allocating slices when subscripting built-ins.

This commit adds a fast-path optimisation for when a BUILD_SLICE is
immediately followed by a LOAD/STORE_SUBSCR for a native type, to avoid
needing to allocate the slice on the heap.

In some cases (e.g. `a[1:3] = x`) this can result in no allocations at all.

We can't do this for instance types because the get/set/delattr
implementation may keep a reference to the slice.

Adds more tests to the basic slice tests to ensure that a stack-allocated
slice never makes it to Python, and also a heapalloc test that verifies
(when using bytecode) that assigning to a slice is no-alloc.

This work was funded through GitHub Sponsors.

Signed-off-by: Jim Mussared <jim.mussared@gmail.com>
Signed-off-by: Damien George <damien@micropython.org>

1 files changed, 29 insertions, 3 deletions

diff --git a/py/vm.c b/py/vm.c
index f87e52c92..6f1179721 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -195,6 +195,22 @@
 #define TRACE_TICK(current_ip, current_sp, is_exception)
 #endif // MICROPY_PY_SYS_SETTRACE
 
+#if MICROPY_PY_BUILTINS_SLICE
+// This function is marked "no inline" so it doesn't increase the C stack usage of the main VM function.
+MP_NOINLINE static mp_obj_t *build_slice_stack_allocated(byte op, mp_obj_t *sp, mp_obj_t step) {
+    mp_obj_t stop = sp[2];
+    mp_obj_t start = sp[1];
+    mp_obj_slice_t slice = { .base = { .type = &mp_type_slice }, .start = start, .stop = stop, .step = step };
+    if (op == MP_BC_LOAD_SUBSCR) {
+        SET_TOP(mp_obj_subscr(TOP(), MP_OBJ_FROM_PTR(&slice), MP_OBJ_SENTINEL));
+    } else { // MP_BC_STORE_SUBSCR
+        mp_obj_subscr(TOP(), MP_OBJ_FROM_PTR(&slice), sp[-1]);
+        sp -= 2;
+    }
+    return sp;
+}
+#endif
+
 // fastn has items in reverse order (fastn[0] is local[0], fastn[-1] is local[1], etc)
 // sp points to bottom of stack which grows up
 // returns:
@@ -849,9 +865,19 @@ unwind_jump:;
                         // 3-argument slice includes step
                         step = POP();
                     }
-                    mp_obj_t stop = POP();
-                    mp_obj_t start = TOP();
-                    SET_TOP(mp_obj_new_slice(start, stop, step));
+                    if ((*ip == MP_BC_LOAD_SUBSCR || *ip == MP_BC_STORE_SUBSCR) && mp_obj_is_native_type(mp_obj_get_type(sp[-2]))) {
+                        // Fast path optimisation for when the BUILD_SLICE is immediately followed
+                        // by a LOAD/STORE_SUBSCR for a native type to avoid needing to allocate
+                        // the slice on the heap.  In some cases (e.g. a[1:3] = x) this can result
+                        // in no allocations at all.  We can't do this for instance types because
+                        // the get/set/delattr implementation may keep a reference to the slice.
+                        byte op = *ip++;
+                        sp = build_slice_stack_allocated(op, sp - 2, step);
+                    } else {
+                        mp_obj_t stop = POP();
+                        mp_obj_t start = TOP();
+                        SET_TOP(mp_obj_new_slice(start, stop, step));
+                    }
                     DISPATCH();
                 }
                 #endif