micropython · jimmo · Dec 5, 2022
diff --git a/py/mpstate.h b/py/mpstate.h
@@ -268,6 +268,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/obj.c b/py/obj.c
@@ -534,6 +534,13 @@ mp_obj_t mp_obj_subscr(mp_obj_t base, mp_obj_t index, mp_obj_t value) {
     const mp_obj_type_t *type = mp_obj_get_type(base);
     if (MP_OBJ_TYPE_HAS_SLOT(type, subscr)) {
         mp_obj_t ret = MP_OBJ_TYPE_GET_SLOT(type, subscr)(base, index, value);
+
+        // If the slice was a thread-local slice (see mp_obj_new_slice), then
+        // mark it as no longer "pending" (i.e. it's now safe to re-use the
+        // thread-local slot).
+        mp_obj_slice_t *o = &MP_STATE_THREAD(slice);
+        o->base.type = NULL;
+
         if (ret != MP_OBJ_NULL) {
             return ret;
         }

diff --git a/py/objslice.c b/py/objslice.c
@@ -116,7 +116,24 @@ 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);
+    if (o->base.type) {
+        // There's already a thread-local slice "pending" (waiting to be
+        // passed to LOAD_SUBSCR or STORE_SUBSCR), so make this one
+        // heap-allocated. This happens for example with `a[1:2,3:4]`.
+        o = m_new_obj(mp_obj_slice_t);
+    }
+    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
@@ -826,6 +826,41 @@ 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
+        // It's possible that:
+        // a) The index is a slice (in which case it will be the thread-local
+        //    slice object).
+        // b) The index is a tuple (in which case the first slice-type
+        //    element, if any, will be the thread-local slice).
+        // c) The index is something else (no-thread-local slice to deal with).
+
+        // For built-in subscr methods, it's safe to just use the thread-local
+        // slice directly. However in this case we're calling a user-defined
+        // {get,set,del}item handler and cannot make any assumptions (e.g.
+        // they may store the slice), so copy it to the heap before passing it
+        // to Python code. See mp_obj_new_slice for details about the
+        // thread-local slice.
+        mp_obj_t *slice_check = &member[2];
+        if (mp_obj_is_type(index, &mp_type_tuple)) {
+            // If the index is a tuple, then at most one of the elements may
+            // be the thread-local slice.
+            mp_obj_tuple_t *t = MP_OBJ_TO_PTR(index);
+            for (size_t i = 0; i < t->len; ++i) {
+                if (t->items[i] == MP_OBJ_FROM_PTR(&MP_STATE_THREAD(slice))) {
+                    slice_check = &t->items[i];
+                }
+            }
+        }
+        if (*slice_check == MP_OBJ_FROM_PTR(&MP_STATE_THREAD(slice))) {
+            // By default the slice object will be thread-local under the
+            // assumption that built-in subscr methods will be safe.
+            mp_obj_slice_t *o = m_new(mp_obj_slice_t, 1);
+            memcpy(o, MP_OBJ_TO_PTR(*slice_check), sizeof(mp_obj_slice_t));
+            *slice_check = 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,43 @@
 # 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]
 
 # 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)