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

jimmo · jimmo · commit 6c1776fbe56f · 2022-12-05T17:01:35.000+11:00
Assume that a slice object (created by the `BUILD_SLICE` op-code or the
native emitter via `mp_obj_new_slice`) will be used to slice a built-in
type. This assumes that all built-in type-&gt;subscr functions will:
1. Not cause another slice to be allocated in the meantime.
2. Not save a reference to the slice.

This prevents a heap allocation, which for some cases (e.g. memoryview
assignment) prevents allocation altogether.

Only when this assumption is incorrect (i.e. slicing a user-defined type
via the `__{get,set,del}item__` methods), then heap-allocate the slice
object.

This work was funded through GitHub Sponsors.

Signed-off-by: Jim Mussared &lt;jim.mussared@gmail.com&gt;
diff --git a/py/mpstate.h b/py/mpstate.h
@@ -256,6 +256,8 @@ typedef struct _mp_state_thread_t {
     // Locking of the GC is done per thread.
     uint16_t gc_lock_depth;
 
+    mp_obj_slice_t slice;
+
     ////////////////////////////////////////////////////////////
     // START ROOT POINTER SECTION
     // Everything that needs GC scanning must start here, and
diff --git a/py/objslice.c b/py/objslice.c
@@ -109,7 +109,18 @@ MP_DEFINE_CONST_OBJ_TYPE(
     );
 
 mp_obj_t mp_obj_new_slice(mp_obj_t ostart, mp_obj_t ostop, mp_obj_t ostep) {
-    mp_obj_slice_t *o = mp_obj_malloc(mp_obj_slice_t, &mp_type_slice);
+    // Use the thread-local slice object, assuming that the likely case is
+    // that this will be passed to a built-in type->subscr which will:
+    // 1. Not cause another slice to be allocated in the meantime.
+    // 2. Not save a reference to the slice.
+    // The only exception to this is instance_subscr which will copy the slice
+    // object into the heap before calling the user-defined __
+    // {get,set,del}item__.
+    // This prevents allocation of slice objects for all built-in types, as
+    // well as making some specific cases (e.g. memoryview or bytearray slice
+    // assignment) completely allocation-free.
+    mp_obj_slice_t *o = &MP_STATE_THREAD(slice);
+    o->base.type = &mp_type_slice;
     o->start = ostart;
     o->stop = ostop;
     o->step = ostep;
diff --git a/py/objtype.c b/py/objtype.c
@@ -837,6 +837,20 @@ STATIC mp_obj_t instance_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value
         return mp_obj_subscr(self->subobj[0], index, value);
     } else if (member[0] != MP_OBJ_NULL) {
         size_t n_args = value == MP_OBJ_NULL || value == MP_OBJ_SENTINEL ? 1 : 2;
+
+        #if MICROPY_PY_BUILTINS_SLICE
+        if (mp_obj_is_type(index, &mp_type_slice)) {
+            // By default the slice object will be thread-local under the
+            // assumption that built-in subscr methods will be safe. However
+            // in this case we're calling a user-defined {get,set,del}item
+            // handler and cannot make any assumptions, so copy it to the heap.
+            // See mp_obj_new_slice for details.
+            mp_obj_slice_t *o = m_new(mp_obj_slice_t, 1);
+            memcpy(o, MP_OBJ_TO_PTR(index), sizeof(mp_obj_slice_t));
+            member[2] = MP_OBJ_FROM_PTR(o);
+        }
+        #endif
+
         mp_obj_t ret = mp_call_method_n_kw(n_args, 0, member);
         if (value == MP_OBJ_SENTINEL) {
             return ret;
diff --git a/tests/basics/builtin_slice.py b/tests/basics/builtin_slice.py
@@ -1,11 +1,36 @@
 # 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)
+
 
 # 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)
