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

The recently merged 5e9189d now allows
temporary slices to be allocated on the C stack, which is much better than
allocating them on the GC heap.

Unfortunately there are cases where the C-allocated slice can escape and be
retained as an object, which leads to crashes (because that object points
to the C stack which now has other values on it).

The fix here is to add a new `MP_TYPE_FLAG_SUBSCR_ALLOWS_STACK_SLICE`.
Native types should set this flag if their subscr method is guaranteed not
to hold on to a reference of the slice object.

Fixes issue micropython#17733 (see also micropython#17723).

Signed-off-by: Damien George <[email protected]>

Author: dpgeorge

py/obj.h

@@ -558,6 +558,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)
@@ -571,6 +573,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,

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,

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.

tests/basics/slice_optimise.py

@@ -0,0 +1,23 @@
+# Test that the slice-on-stack optimisation does not break various uses of slice
+# (see MP_TYPE_FLAG_SUBSCR_ALLOWS_STACK_SLICE type option).
+#
+# Note: this test has a corresponding .py.exp file because hashing slice objects
+# was not allowed in CPython until 3.12.
+
+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()))

tests/basics/slice_optimise.py.exp

@@ -0,0 +1,2 @@
+KeyError (slice(None, None, None),)
+[slice(None, 'a', None), 'b'] [1, 2]