diff --git a/mypyc/lib-rt/vecs/librt_vecs.h b/mypyc/lib-rt/vecs/librt_vecs.h index 48ceb0a5efa7..04fd06b39855 100644 --- a/mypyc/lib-rt/vecs/librt_vecs.h +++ b/mypyc/lib-rt/vecs/librt_vecs.h @@ -7,6 +7,7 @@ #define PY_SSIZE_T_CLEAN #include #include +#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); diff --git a/mypyc/lib-rt/vecs/vec_nested.c b/mypyc/lib-rt/vecs/vec_nested.c index 18fa2aa33e79..dcb72f527c26 100644 --- a/mypyc/lib-rt/vecs/vec_nested.c +++ b/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)) { diff --git a/mypyc/lib-rt/vecs/vec_t.c b/mypyc/lib-rt/vecs/vec_t.c index 223aa933fb8c..5cc756ab223b 100644 --- a/mypyc/lib-rt/vecs/vec_t.c +++ b/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)) { diff --git a/mypyc/lib-rt/vecs/vec_template.c b/mypyc/lib-rt/vecs/vec_template.c index 8dd88a2c0517..9d34d845ddaf 100644 --- a/mypyc/lib-rt/vecs/vec_template.c +++ b/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);