[mypyc] Refactor: use helper functions for vec growth

mypyc/lib-rt/vecs/librt_vecs.h

@@ -7,6 +7,7 @@
 #define PY_SSIZE_T_CLEAN
 #include <Python.h>
 #include <stdint.h>
+#include "mypyc_util.h"
 
 #ifdef MYPYC_EXPERIMENTAL
 
@@ -832,6 +833,27 @@ static inline PyObject *VecNested_BoxItem(VecNested v, VecNestedBufItem item) {
     }
 }
 
+// Growth helpers
+
+static inline Py_ssize_t Vec_GrowCapacity(Py_ssize_t cap) {
+    if (unlikely(cap > (PY_SSIZE_T_MAX - 1) / 2)) {
+        // Allocation will fail at this size, but avoid overflow
+        return PY_SSIZE_T_MAX;
+    }
+    return 2 * cap + 1;
+}
+
+static inline Py_ssize_t Vec_GrowCapacityTo(Py_ssize_t cap, Py_ssize_t min_cap) {
+    while (cap < min_cap) {
+        if (unlikely(cap > (PY_SSIZE_T_MAX - 1) / 2)) {
+            cap = min_cap;
+            break;
+        }
+        cap = 2 * cap + 1;
+    }
+    return cap;
+}
+
 // Misc helpers
 
 PyObject *Vec_TypeToStr(size_t item_type, size_t depth);

mypyc/lib-rt/vecs/vec_nested.c

@@ -255,7 +255,7 @@ VecNested VecNested_Append(VecNested vec, VecNestedBufItem x) {
         vec.len++;
         return vec;
     } else {
-        Py_ssize_t new_size = 2 * cap + 1;
+        Py_ssize_t new_size = Vec_GrowCapacity(cap);
         // TODO: Avoid initializing to zero here
         VecNested new = vec_alloc(new_size, vec.buf->item_type, vec.buf->depth);
         if (VEC_IS_ERROR(new)) {
@@ -350,14 +350,7 @@ VecNested VecNested_ExtendVec(VecNested dst, VecNested src) {
         return dst;
     }
     // Need to reallocate (or dst and src share a buffer)
-    Py_ssize_t new_cap = cap;
-    while (new_cap < new_len) {
-        if (new_cap > (PY_SSIZE_T_MAX - 1) / 2) {
-            new_cap = new_len;
-            break;
-        }
-        new_cap = 2 * new_cap + 1;
-    }
+    Py_ssize_t new_cap = Vec_GrowCapacityTo(cap, new_len);
     int aliased = dst.buf == src.buf;
     VecNested new = vec_alloc(new_cap, dst.buf->item_type, dst.buf->depth);
     if (VEC_IS_ERROR(new)) {

mypyc/lib-rt/vecs/vec_t.c

@@ -270,7 +270,7 @@ VecT VecT_Append(VecT vec, PyObject *x, size_t item_type) {
         vec.len++;
         return vec;
     } else {
-        Py_ssize_t new_size = 2 * cap + 1;
+        Py_ssize_t new_size = Vec_GrowCapacity(cap);
         // TODO: Avoid initializing to zero here
         VecT new = vec_alloc(new_size, vec.buf->item_type);
         if (VEC_IS_ERROR(new)) {
@@ -375,14 +375,7 @@ VecT VecT_ExtendVec(VecT dst, VecT src, size_t item_type) {
         return dst;
     }
     // Need to reallocate (or dst and src share a buffer)
-    Py_ssize_t new_cap = cap;
-    while (new_cap < new_len) {
-        if (new_cap > (PY_SSIZE_T_MAX - 1) / 2) {
-            new_cap = new_len;
-            break;
-        }
-        new_cap = 2 * new_cap + 1;
-    }
+    Py_ssize_t new_cap = Vec_GrowCapacityTo(cap, new_len);
     int aliased = dst.buf == src.buf;
     VecT new = vec_alloc(new_cap, dst.buf->item_type);
     if (VEC_IS_ERROR(new)) {

mypyc/lib-rt/vecs/vec_template.c

@@ -392,7 +392,7 @@ VEC FUNC(Append)(VEC vec, ITEM_C_TYPE x) {
         return vec;
     } else {
         Py_ssize_t cap = vec.buf ? VEC_CAP(vec) : 0;
-        Py_ssize_t new_size = 2 * cap + 1;
+        Py_ssize_t new_size = Vec_GrowCapacity(cap);
         VEC new = vec_alloc(new_size);
         if (VEC_IS_ERROR(new)) {
             // The input v is being consumed/stolen by this function, so on error
@@ -437,14 +437,7 @@ inline static VEC vec_extend_items(
         dst.len = new_len;
         return dst;
     }
-    Py_ssize_t new_cap = cap;
-    while (new_cap < new_len) {
-        if (unlikely(new_cap > (PY_SSIZE_T_MAX - 1) / 2)) {
-            new_cap = new_len;
-            break;
-        }
-        new_cap = 2 * new_cap + 1;
-    }
+    Py_ssize_t new_cap = Vec_GrowCapacityTo(cap, new_len);
     VEC new = vec_alloc(new_cap);
     if (VEC_IS_ERROR(new)) {
         VEC_DECREF(dst);