diff --git a/py/vm.c b/py/vm.c index 385d13ee40..69e698e3ce 100644 --- 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 index df84d5c57b..48dd3304d7 100644 --- 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 -s = A()[1:2:3] + + 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 + +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 new file mode 100644 index 0000000000..62d96595c7 --- /dev/null +++ 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 index 2d511c20ca..f34bcd5cb8 100755 --- 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