From d067bd3f41188e85d18f56c69737cba4c4f3e81b Mon Sep 17 00:00:00 2001
From: Joshua Haberman <jhaberman@gmail.com>
Date: Fri, 10 Dec 2021 07:35:45 -0800
Subject: [PATCH] Addressed PR comments.

---
 python/repeated.c | 38 +++++++++++++++++++++-----------------
 upb/reflection.c  | 14 ++++++++------
 upb/reflection.h  |  6 +++++-
 3 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/python/repeated.c b/python/repeated.c
index fab4bf0729..1286f28c0f 100644
--- a/python/repeated.c
+++ b/python/repeated.c
@@ -243,7 +243,7 @@ PyObject* PyUpb_RepeatedContainer_ToList(PyObject* _self) {
   PyObject* list = PyList_New(n);
   for (size_t i = 0; i < n; i++) {
     PyObject* val = PyUpb_UpbToPy(upb_array_get(arr, i), f, self->arena);
-    PyList_SetItem(list, n, val);
+    PyList_SetItem(list, i, val);
   }
   return list;
 }
@@ -315,7 +315,9 @@ static int PyUpb_RepeatedContainer_SetSubscript(
   // Set range.
   PyObject* seq =
       PySequence_Fast(value, "must assign iterable to extended slice");
+  PyObject* item = NULL;
   if (!seq) return -1;
+  int ret = -1;
   if (PySequence_Size(seq) != count) {
     PyErr_Format(PyExc_ValueError,
                   "attempt to assign sequence of size %zd to extended slice "
@@ -326,35 +328,37 @@ static int PyUpb_RepeatedContainer_SetSubscript(
   }
   for (Py_ssize_t i = 0; i < count; i++, idx += step) {
     upb_msgval msgval;
-    PyObject* item = PySequence_GetItem(seq, i);
+    item = PySequence_GetItem(seq, i);
+    if (!item) goto err;
     // XXX: if this fails we can leave the list partially mutated.
-    bool ok = PyUpb_PyToUpb(item, f, &msgval, arena);
-    Py_DECREF(item);
-    if (!ok) return -1;
+    if (!PyUpb_PyToUpb(item, f, &msgval, arena)) goto err;
     upb_array_set(arr, idx, msgval);
   }
-  return 0;
+  ret = 0;
+
+err:
+  Py_XDECREF(seq);
+  Py_XDECREF(item);
+  return ret;
 }
 
 static int PyUpb_RepeatedContainer_DeleteSubscript(upb_array* arr,
                                                    Py_ssize_t idx,
                                                    Py_ssize_t count,
                                                    Py_ssize_t step) {
-  // This is a delete, not a set.
   // Normalize direction: deletion is order-independent.
-  if (step > 0) {
-    idx += step * (count - 1);
+  Py_ssize_t start = idx;
+  if (step < 0) {
+    Py_ssize_t end = start + step * (count - 1);
+    start = end;
     step = -step;
   }
-  if (step == -1) {
-    // Contiguous range.
-    upb_array_delete(arr, idx - (count - 1), count);
-  } else {
-    // Stepped slice.
-    for (Py_ssize_t i = 0; i < count; i++, idx += step) {
-      upb_array_delete(arr, idx, 1);
-    }
+  size_t dst = start;
+  size_t src = start + 1;
+  for (Py_ssize_t i = 0; i < count; i++, dst += step, src += step + 1) {
+    upb_array_move(arr, dst, src, step);
   }
+  upb_array_resize(arr, upb_array_size(arr) - count, NULL);
   return 0;
 }
 
diff --git a/upb/reflection.c b/upb/reflection.c
index 4bd7865e24..024d1753a9 100644
--- a/upb/reflection.c
+++ b/upb/reflection.c
@@ -373,6 +373,12 @@ bool upb_array_append(upb_array *arr, upb_msgval val, upb_arena *arena) {
   return true;
 }
 
+bool upb_array_move(upb_array* arr, size_t dst_idx, size_t src_idx,
+                    size_t count) {
+  char* data = _upb_array_ptr(arr);
+  int lg2 = arr->data & 7;
+  memmove(&data[dst_idx << lg2], &data[src_idx << lg2], count << lg2);
+}
 
 bool upb_array_insert(upb_array *arr, size_t i, size_t count,
                       upb_arena *arena) {
@@ -382,9 +388,7 @@ bool upb_array_insert(upb_array *arr, size_t i, size_t count,
   if (!upb_array_resize(arr, arr->len + count, arena)) {
     return false;
   }
-  char* data = _upb_array_ptr(arr);
-  int lg2 = arr->data & 7;
-  memmove(&data[(i + count) << lg2], &data[i << lg2], (oldsize - i) << lg2);
+  upb_array_move(arr, i + count, i, oldsize - i);
   return true;
 }
 
@@ -396,9 +400,7 @@ void upb_array_delete(upb_array *arr, size_t i, size_t count) {
   size_t end = i + count;
   UPB_ASSERT(end >= i);
   UPB_ASSERT(end <= arr->len);
-  char* data = _upb_array_ptr(arr);
-  int lg2 = arr->data & 7;
-  memmove(&data[i << lg2], &data[end << lg2], (arr->len - end) << lg2);
+  upb_array_move(arr, i, end, arr->len - end);
   arr->len -= count;
 }
 
diff --git a/upb/reflection.h b/upb/reflection.h
index 146c0ed63d..12ada97f59 100644
--- a/upb/reflection.h
+++ b/upb/reflection.h
@@ -134,11 +134,15 @@ void upb_array_set(upb_array *arr, size_t i, upb_msgval val);
 /* Appends an element to the array.  Returns false on allocation failure. */
 bool upb_array_append(upb_array *array, upb_msgval val, upb_arena *arena);
 
+/* Moves elements within the array using memmove(). Like memmove(), the source
+ * and destination elements may be overlapping. */
+bool upb_array_move(upb_array* array, size_t dst_idx, size_t src_idx,
+                    size_t count);
 
 /* Inserts one or more empty elements into the array.  Existing elements are
  * shifted right.  The new elements have undefined state and must be set with
  * `upb_array_set()`.
- * REQUIRES: i <= `upb_array_size(arr)` */
+ * REQUIRES: `i <= upb_array_size(arr)` */
 bool upb_array_insert(upb_array *array, size_t i, size_t count,
                       upb_arena *arena);