objarray: Implement array slice assignment.

This is rarely used feature which takes enough code to implement, so is
controlled by MICROPY_PY_ARRAY_SLICE_ASSIGN config setting, default off.
But otherwise it may be useful, as allows to update arbitrary-sized data
buffers in-place.

Slice is yet to implement, and actually, slice assignment implemented in
such a way that RHS of assignment should be array of the exact same item
typecode as LHS. CPython has it more relaxed, where RHS can be any sequence
of compatible types (e.g. it's possible to assign list of int's to a
bytearray slice).

Overall, when all "slice write" features are implemented, it may cost ~1KB
of code.

1 files changed, 37 insertions, 6 deletions

diff --git a/py/objarray.c b/py/objarray.c
index e2fee1f2ed..e030548bbb 100644
--- a/py/objarray.c
+++ b/py/objarray.c
@@ -351,17 +351,48 @@ STATIC mp_obj_t array_subscr(mp_obj_t self_in, mp_obj_t index_in, mp_obj_t value
         if (0) {
 #if MICROPY_PY_BUILTINS_SLICE
         } else if (MP_OBJ_IS_TYPE(index_in, &mp_type_slice)) {
-            if (value != MP_OBJ_SENTINEL) {
-                // Only getting a slice is suported so far, not assignment
-                // TODO: confirmed that both bytearray and array.array support
-                // slice assignment (incl. of different size)
-                return MP_OBJ_NULL; // op not supported
-            }
             mp_bound_slice_t slice;
             if (!mp_seq_get_fast_slice_indexes(o->len, index_in, &slice)) {
                 nlr_raise(mp_obj_new_exception_msg(&mp_type_NotImplementedError,
                     "only slices with step=1 (aka None) are supported"));
             }
+            if (value != MP_OBJ_SENTINEL) {
+                #if MICROPY_PY_ARRAY_SLICE_ASSIGN
+                // Assign
+                if (!MP_OBJ_IS_TYPE(value, &mp_type_array) && !MP_OBJ_IS_TYPE(value, &mp_type_bytearray)) {
+                    mp_not_implemented("array required on right side");
+                }
+                mp_obj_array_t *src_slice = value;
+                int item_sz = mp_binary_get_size('@', o->typecode, NULL);
+                if (item_sz != mp_binary_get_size('@', src_slice->typecode, NULL)) {
+                    mp_not_implemented("arrays should be compatible");
+                }
+
+                // TODO: check src/dst compat
+                mp_int_t len_adj = src_slice->len - (slice.stop - slice.start);
+                if (len_adj > 0) {
+                    if (len_adj > o->free) {
+                        // TODO: alloc policy; at the moment we go conservative
+                        o->items = m_realloc(o->items, (o->len + o->free) * item_sz, (o->len + len_adj) * item_sz);
+                        o->free = 0;
+                    }
+                    mp_seq_replace_slice_grow_inplace(o->items, o->len,
+                        slice.start, slice.stop, src_slice->items, src_slice->len, len_adj, item_sz);
+                } else {
+                    mp_seq_replace_slice_no_grow(o->items, o->len,
+                        slice.start, slice.stop, src_slice->items, src_slice->len, item_sz);
+                    // Clear "freed" elements at the end of list
+                    // TODO: This is actually only needed for typecode=='O'
+                    mp_seq_clear(o->items, o->len + len_adj, o->len, item_sz);
+                    // TODO: alloc policy after shrinking
+                }
+                o->len += len_adj;
+                return mp_const_none;
+                #else
+                return MP_OBJ_NULL; // op not supported
+                #endif
+            }
+
             mp_obj_array_t *res;
             int sz = mp_binary_get_size('@', o->typecode & TYPECODE_MASK, NULL);
             assert(sz > 0);