py/obj: Add new type flag to indicate subscr accepts slice-on-stack.

dpgeorge · dpgeorge · commit 8b65d69b6f9d · 2025-07-22T13:06:28.000+10:00
Fixes issue #17733. Signed-off-by: Damien George <damien@micropython.org>
diff --git a/py/obj.h b/py/obj.h
@@ -554,6 +554,8 @@ typedef mp_obj_t (*mp_fun_kw_t)(size_t n, const mp_obj_t *, mp_map_t *);
 // If MP_TYPE_FLAG_ITER_IS_STREAM is set then the type implicitly gets a "return self"
 //   getiter, and mp_stream_unbuffered_iter for iternext.
 // If MP_TYPE_FLAG_INSTANCE_TYPE is set then this is an instance type (i.e. defined in Python).
+// If MP_TYPE_FLAG_SUBSCR_ALLOWS_STACK_SLICE is set then the "subscr" slot allows a stack
+//   allocated slice to be passed in (no references to it will be retained after the call).
 #define MP_TYPE_FLAG_NONE (0x0000)
 #define MP_TYPE_FLAG_IS_SUBCLASSED (0x0001)
 #define MP_TYPE_FLAG_HAS_SPECIAL_ACCESSORS (0x0002)
@@ -567,6 +569,7 @@ typedef mp_obj_t (*mp_fun_kw_t)(size_t n, const mp_obj_t *, mp_map_t *);
 #define MP_TYPE_FLAG_ITER_IS_CUSTOM (0x0100)
 #define MP_TYPE_FLAG_ITER_IS_STREAM (MP_TYPE_FLAG_ITER_IS_ITERNEXT | MP_TYPE_FLAG_ITER_IS_CUSTOM)
 #define MP_TYPE_FLAG_INSTANCE_TYPE (0x0200)
+#define MP_TYPE_FLAG_SUBSCR_ALLOWS_STACK_SLICE (0x0400)
 
 typedef enum {
     PRINT_STR = 0,
diff --git a/py/objarray.c b/py/objarray.c
@@ -626,7 +626,7 @@ MP_DEFINE_CONST_OBJ_TYPE(
 MP_DEFINE_CONST_OBJ_TYPE(
     mp_type_bytearray,
     MP_QSTR_bytearray,
-    MP_TYPE_FLAG_EQ_CHECKS_OTHER_TYPE | MP_TYPE_FLAG_ITER_IS_GETITER,
+    MP_TYPE_FLAG_EQ_CHECKS_OTHER_TYPE | MP_TYPE_FLAG_ITER_IS_GETITER | MP_TYPE_FLAG_SUBSCR_ALLOWS_STACK_SLICE,
     make_new, bytearray_make_new,
     print, array_print,
     iter, array_iterator_new,
@@ -654,7 +654,7 @@ MP_DEFINE_CONST_OBJ_TYPE(
 MP_DEFINE_CONST_OBJ_TYPE(
     mp_type_memoryview,
     MP_QSTR_memoryview,
-    MP_TYPE_FLAG_EQ_CHECKS_OTHER_TYPE | MP_TYPE_FLAG_ITER_IS_GETITER,
+    MP_TYPE_FLAG_EQ_CHECKS_OTHER_TYPE | MP_TYPE_FLAG_ITER_IS_GETITER | MP_TYPE_FLAG_SUBSCR_ALLOWS_STACK_SLICE,
     make_new, memoryview_make_new,
     iter, array_iterator_new,
     unary_op, array_unary_op,
diff --git a/py/vm.c b/py/vm.c
@@ -865,9 +865,10 @@ unwind_jump:;
                         // 3-argument slice includes step
                         step = POP();
                     }
-                    if ((*ip == MP_BC_LOAD_SUBSCR || *ip == MP_BC_STORE_SUBSCR) && mp_obj_is_native_type(mp_obj_get_type(sp[-2]))) {
+                    if ((*ip == MP_BC_LOAD_SUBSCR || *ip == MP_BC_STORE_SUBSCR)
+                        && (mp_obj_get_type(sp[-2])->flags & MP_TYPE_FLAG_SUBSCR_ALLOWS_STACK_SLICE)) {
                         // 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
+                        // by a LOAD/STORE_SUBSCR for an accepting 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.
diff --git a/tests/basics/slice_optimse.py b/tests/basics/slice_optimse.py
@@ -0,0 +1,20 @@
+# test that the slice-on-stack optimisation does not break various uses of slice
+# (see MP_TYPE_FLAG_SUBSCR_ALLOWS_STACK_SLICE type option)
+
+try:
+    from collections import OrderedDict
+except ImportError:
+    print("SKIP")
+    raise SystemExit
+
+# attempt to index with a slice, error should contain the slice (failed key)
+try:
+    dict()[:]
+except KeyError as e:
+    print("KeyError", e.args)
+
+# put a slice and another object into an OrderedDict, and retrieve them
+x = OrderedDict()
+x[:"a"] = 1
+x["b"] = 2
+print(list(x.keys()), list(x.values()))
