py/vm.c: Avoid heap-allocating slices for built-ins.

jimmo · jimmo · commit ec15a0b61f28 · 2023-10-05T14:34:02.000+11:00
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 &lt;jim.mussared@gmail.com&gt;
diff --git a/py/vm.c b/py/vm.c
@@ -850,8 +850,26 @@ unwind_jump:;
                         step = POP();
                     }
                     mp_obj_t stop = POP();
-                    mp_obj_t start = TOP();
-                    SET_TOP(mp_obj_new_slice(start, stop, step));
+                    mp_obj_t start = POP();
+                    if ((*ip == MP_BC_LOAD_SUBSCR || *ip == MP_BC_STORE_SUBSCR) && mp_obj_is_native_type(mp_obj_get_type(TOP()))) {
+                        // 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++;
+                        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;
+                        }
+                    } else {
+                        PUSH(mp_obj_new_slice(start, stop, step));
+                    }
                     DISPATCH();
                 }
                 #endif
diff --git a/tests/basics/builtin_slice.py b/tests/basics/builtin_slice.py
@@ -1,11 +1,44 @@
 # test builtin slice
 
+# ensures that slices passed to user types are heap-allocated and can be
+# safely stored as well as not overriden by subsequent slices.
+
 # print slice
 class A:
     def __getitem__(self, idx):
-        print(idx)
+        print("get", idx)
+        print("abc"[1:])
+        print("get", idx)
+        return idx
+
+    def __setitem__(self, idx, value):
+        print("set", idx)
+        print("abc"[1:])
+        print("set", idx)
+        self.saved_idx = idx
+        return idx
+
+    def __delitem__(self, idx):
+        print("del", idx)
+        print("abc"[1:])
+        print("del", idx)
         return idx
-s = A()[1:2:3]
+
+a = A()
+s = a[1:2:3]
+a[4:5:6] = s
+del a[7:8:9]
+
+print(a.saved_idx)
+
+# nested slicing
+print(A()[1:A()[A()[2:3:4]:5]])
+
+# tuple slicing
+a[1:2,4:5,7:8]
+a[1,4:5,7:8,2]
+a[1:2, a[3:4], 5:6]
 
 # check type
 print(type(s) is slice)
+
diff --git a/tests/micropython/heapalloc_slice.py b/tests/micropython/heapalloc_slice.py
@@ -0,0 +1,18 @@
+# slice operations that don't require allocation
+try:
+    from micropython import heap_lock, heap_unlock
+except (ImportError, AttributeError):
+    heap_lock = heap_unlock = lambda: 0
+
+b = bytearray(range(10))
+
+m = memoryview(b)
+
+heap_lock()
+
+b[3:5] = b"aa"
+m[5:7] = b"bb"
+
+heap_unlock()
+
+print(b)
diff --git a/tests/run-tests.py b/tests/run-tests.py
@@ -639,6 +639,9 @@ def run_tests(pyb, tests, args, result_dir, num_threads=1):
         skip_tests.add(
             "micropython/emg_exc.py"
         )  # because native doesn't have proper traceback info
+        skip_tests.add(
+            "micropython/heapalloc_slice.py"
+        )  # because native doesn't do the stack-allocated slice optimisation
         skip_tests.add(
             "micropython/heapalloc_traceback.py"
         )  # because native doesn't have proper traceback info
