librt: Support specifying vec capacity during construction (#21304)

Use e.g. `vec[i64](capacity=n)` to reserve capacity in the buffer during
construction. This can be used to reduce the number of buffer reallocs
when appending items. Capacity is the length of buffer, while length is
the number of items currently stored in the buffer.

I used coding agent assist but did this in multiple smaller increments.

Author: JukkaL

mypy/typeshed/stubs/librt/librt/vecs.pyi

@@ -5,9 +5,9 @@ T = TypeVar("T")
 
 class vec(Generic[T]):
     @overload
-    def __init__(self) -> None: ...
+    def __init__(self, *, capacity: i64 = ...) -> None: ...
     @overload
-    def __init__(self, items: Iterable[T], /) -> None: ...
+    def __init__(self, items: Iterable[T], /, *, capacity: i64 = ...) -> None: ...
     def __len__(self) -> i64: ...
     @overload
     def __getitem__(self, i: i64, /) -> T: ...

mypyc/irbuild/expression.py

@@ -364,10 +364,17 @@ def transform_call_expr(builder: IRBuilder, expr: CallExpr) -> Value:
         ):
             item_type = builder.type_to_rtype(analyzed.types[0])
             vec_type = RVec(item_type)
-            if len(expr.args) == 0:
-                return vec_create(builder.builder, vec_type, 0, expr.line)
-            elif len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]:
-                return translate_vec_create_from_iterable(builder, vec_type, expr.args[0])
+            capacity = _get_vec_capacity(builder, expr)
+            if len(expr.args) == 0 or (len(expr.args) == 1 and expr.arg_kinds == [ARG_NAMED]):
+                # vec[T]() or vec[T](capacity=N)
+                return vec_create(builder.builder, vec_type, 0, expr.line, capacity=capacity)
+            elif (len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]) or (
+                len(expr.args) == 2 and expr.arg_kinds == [ARG_POS, ARG_NAMED]
+            ):
+                # vec[T](items) or vec[T](items, capacity=N)
+                return translate_vec_create_from_iterable(
+                    builder, vec_type, expr.args[0], capacity=capacity
+                )
         callee = analyzed.expr  # Unwrap type application
 
     if isinstance(callee, MemberExpr):
@@ -561,8 +568,16 @@ def translate_super_method_call(builder: IRBuilder, expr: CallExpr, callee: Supe
     return builder.builder.call(decl, arg_values, arg_kinds, arg_names, expr.line)
 
 
+def _get_vec_capacity(builder: IRBuilder, expr: CallExpr) -> Value | None:
+    """Extract the 'capacity' keyword argument value from a vec() call, or None."""
+    for i, (kind, name) in enumerate(zip(expr.arg_kinds, expr.arg_names)):
+        if kind == ARG_NAMED and name == "capacity":
+            return builder.accept(expr.args[i])
+    return None
+
+
 def translate_vec_create_from_iterable(
-    builder: IRBuilder, vec_type: RVec, arg: Expression
+    builder: IRBuilder, vec_type: RVec, arg: Expression, *, capacity: Value | None = None
 ) -> Value:
     line = arg.line
     item_type = vec_type.item_type
@@ -581,28 +596,35 @@ def translate_vec_create_from_iterable(
         if is_int64_rprimitive(other_type) or is_int_rprimitive(other_type):
             length = builder.accept(other)
             init = builder.accept(lst.items[0])
-            return vec_create_initialized(builder.builder, vec_type, length, init, line)
+            return vec_create_initialized(
+                builder.builder, vec_type, length, init, line, capacity=capacity
+            )
         assert False, other_type
     if isinstance(arg, ListExpr):
         items = []
         for item in arg.items:
             value = builder.accept(item)
             items.append(builder.coerce(value, item_type, line))
-        return vec_create_from_values(builder.builder, vec_type, items, line)
+        return vec_create_from_values(builder.builder, vec_type, items, line, capacity=capacity)
     if isinstance(arg, ListComprehension):
-        return translate_vec_comprehension(builder, vec_type, arg.generator)
-    return vec_from_iterable(builder, vec_type, arg, line)
+        return translate_vec_comprehension(builder, vec_type, arg.generator, capacity=capacity)
+    return vec_from_iterable(builder, vec_type, arg, line, capacity=capacity)
 
 
 def vec_from_iterable(
-    builder: IRBuilder, vec_type: RVec, iterable: Expression, line: int
+    builder: IRBuilder,
+    vec_type: RVec,
+    iterable: Expression,
+    line: int,
+    *,
+    capacity: Value | None = None,
 ) -> Value:
     """Construct a vec from an arbitrary iterable."""
     # Translate it as a vec comprehension vec[t]([<name> for <name> in
     # iterable]). This way we can use various special casing supported
     # by for loops and comprehensions.
     vec = Register(vec_type)
-    builder.assign(vec, vec_create(builder.builder, vec_type, 0, line), line)
+    builder.assign(vec, vec_create(builder.builder, vec_type, 0, line, capacity=capacity), line)
     name = f"___tmp_{line}"
     var = Var(name)
     reg = builder.add_local(var, vec_type.item_type)

mypyc/irbuild/for_helpers.py

@@ -363,7 +363,9 @@ def gen_inner_stmts() -> None:
     return builder.read(set_ops, gen.line)
 
 
-def translate_vec_comprehension(builder: IRBuilder, vec_type: RVec, gen: GeneratorExpr) -> Value:
+def translate_vec_comprehension(
+    builder: IRBuilder, vec_type: RVec, gen: GeneratorExpr, *, capacity: Value | None = None
+) -> Value:
     def set_item(x: Value, y: Value, z: Value, line: int) -> None:
         vec_init_item_unsafe(builder.builder, x, y, z, line)
 
@@ -372,15 +374,17 @@ def set_item(x: Value, y: Value, z: Value, line: int) -> None:
         builder,
         gen,
         empty_op_llbuilder=lambda length, line: vec_create(
-            builder.builder, vec_type, length, line
+            builder.builder, vec_type, length, line, capacity=capacity
         ),
         set_item_op=set_item,
     )
     if val is not None:
         return val
 
     vec = Register(vec_type)
-    builder.assign(vec, vec_create(builder.builder, vec_type, 0, gen.line), gen.line)
+    builder.assign(
+        vec, vec_create(builder.builder, vec_type, 0, gen.line, capacity=capacity), gen.line
+    )
     loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
 
     def gen_inner_stmts() -> None:

mypyc/irbuild/vec.py

@@ -74,17 +74,28 @@ def as_platform_int(builder: LowLevelIRBuilder, v: Value, line: int) -> Value:
     return builder.coerce(v, c_pyssize_t_rprimitive, line)
 
 
-def vec_create(builder: LowLevelIRBuilder, vtype: RVec, length: int | Value, line: int) -> Value:
+def vec_create(
+    builder: LowLevelIRBuilder,
+    vtype: RVec,
+    length: int | Value,
+    line: int,
+    *,
+    capacity: Value | None = None,
+) -> Value:
     if isinstance(length, int):
         length = Integer(length, c_pyssize_t_rprimitive)
     length = as_platform_int(builder, length, line)
+    if capacity is not None:
+        capacity = as_platform_int(builder, capacity, line)
+    else:
+        capacity = length
 
     item_type = vtype.item_type
     api_name = vec_api_by_item_type.get(item_type)
     if api_name is not None:
         call = CallC(
             f"{api_name}.alloc",
-            [length, length],
+            [length, capacity],
             vtype,
             False,
             False,
@@ -110,7 +121,7 @@ def vec_create(builder: LowLevelIRBuilder, vtype: RVec, length: int | Value, lin
         if depth == 0:
             call = CallC(
                 "VecTApi.alloc",
-                [length, length, typeval],
+                [length, capacity, typeval],
                 vtype,
                 False,
                 False,
@@ -121,7 +132,7 @@ def vec_create(builder: LowLevelIRBuilder, vtype: RVec, length: int | Value, lin
         else:
             call = CallC(
                 "VecNestedApi.alloc",
-                [length, length, typeval, Integer(depth, int32_rprimitive)],
+                [length, capacity, typeval, Integer(depth, int32_rprimitive)],
                 vtype,
                 False,
                 False,
@@ -134,7 +145,13 @@ def vec_create(builder: LowLevelIRBuilder, vtype: RVec, length: int | Value, lin
 
 
 def vec_create_initialized(
-    builder: LowLevelIRBuilder, vtype: RVec, length: int | Value, init: Value, line: int
+    builder: LowLevelIRBuilder,
+    vtype: RVec,
+    length: int | Value,
+    init: Value,
+    line: int,
+    *,
+    capacity: Value | None = None,
 ) -> Value:
     """Create vec with items initialized to the given value."""
     if isinstance(length, int):
@@ -143,7 +160,7 @@ def vec_create_initialized(
 
     item_type = vtype.item_type
     init = builder.coerce(init, item_type, line)
-    vec = vec_create(builder, vtype, length, line)
+    vec = vec_create(builder, vtype, length, line, capacity=capacity)
 
     items_start = vec_items(builder, vec)
     step = step_size(item_type)
@@ -160,9 +177,14 @@ def vec_create_initialized(
 
 
 def vec_create_from_values(
-    builder: LowLevelIRBuilder, vtype: RVec, values: list[Value], line: int
+    builder: LowLevelIRBuilder,
+    vtype: RVec,
+    values: list[Value],
+    line: int,
+    *,
+    capacity: Value | None = None,
 ) -> Value:
-    vec = vec_create(builder, vtype, len(values), line)
+    vec = vec_create(builder, vtype, len(values), line, capacity=capacity)
     ptr = vec_items(builder, vec)
     item_type = vtype.item_type
     step = step_size(item_type)

mypyc/lib-rt/vecs/librt_vecs.c

@@ -96,29 +96,34 @@ typedef struct {
 
 static PyObject *vec_generic_alias_call(PyObject *self, PyObject *args, PyObject *kw)
 {
-    static char *kwlist[] = {"", NULL};
+    static char *kwlist[] = {"", "capacity", NULL};
     PyObject *init = NULL;
-    if (!PyArg_ParseTupleAndKeywords(args, kw, "|O:vec", kwlist, &init)) {
+    int64_t cap = 0;
+    if (!PyArg_ParseTupleAndKeywords(args, kw, "|OL:vec", kwlist, &init, &cap)) {
+        return NULL;
+    }
+    if (cap < 0) {
+        PyErr_SetString(PyExc_ValueError, "capacity must not be negative");
         return NULL;
     }
     VecGenericAlias *p = (VecGenericAlias *)self;
     if (p->depth == 0) {
         if (init == NULL) {
-            VecT vec = VecT_New(0, 0, p->item_type);
+            VecT vec = VecT_New(0, cap, p->item_type);
             if (VEC_IS_ERROR(vec))
                 return NULL;
             return VecT_Box(vec, p->item_type);
         } else {
-            return VecT_FromIterable(p->item_type, init);
+            return VecT_FromIterable(p->item_type, init, cap);
         }
     } else {
         if (init == NULL) {
-            VecNested vec = VecNested_New(0, 0, p->item_type, p->depth);
+            VecNested vec = VecNested_New(0, cap, p->item_type, p->depth);
             if (VEC_IS_ERROR(vec))
                 return NULL;
             return VecNested_Box(vec);
         } else {
-            return VecNested_FromIterable(p->item_type, p->depth, init);
+            return VecNested_FromIterable(p->item_type, p->depth, init, cap);
         }
     }
 }

mypyc/lib-rt/vecs/librt_vecs.h

@@ -713,7 +713,7 @@ static inline int VecT_ItemCheck(VecT v, PyObject *item, size_t item_type) {
 }
 
 VecT VecT_New(Py_ssize_t size, Py_ssize_t cap, size_t item_type);
-PyObject *VecT_FromIterable(size_t item_type, PyObject *iterable);
+PyObject *VecT_FromIterable(size_t item_type, PyObject *iterable, int64_t cap);
 PyObject *VecT_Box(VecT vec, size_t item_type);
 VecT VecT_Append(VecT vec, PyObject *x, size_t item_type);
 VecT VecT_Remove(VecT vec, PyObject *x);
@@ -726,7 +726,7 @@ static inline int VecNested_Check(PyObject *o) {
 }
 
 VecNested VecNested_New(Py_ssize_t size, Py_ssize_t cap, size_t item_type, size_t depth);
-PyObject *VecNested_FromIterable(size_t item_type, size_t depth, PyObject *iterable);
+PyObject *VecNested_FromIterable(size_t item_type, size_t depth, PyObject *iterable, int64_t cap);
 PyObject *VecNested_Box(VecNested);
 VecNested VecNested_Append(VecNested vec, VecNestedBufItem x);
 VecNested VecNested_Remove(VecNested vec, VecNestedBufItem x);

mypyc/lib-rt/vecs/vec_nested.c

@@ -65,6 +65,10 @@ VecNested VecNested_ConvertFromNested(VecNestedBufItem item) {
 }
 
 VecNested VecNested_New(Py_ssize_t size, Py_ssize_t cap, size_t item_type, size_t depth) {
+    if (cap < 0) {
+        PyErr_SetString(PyExc_ValueError, "capacity must not be negative");
+        return vec_error();
+    }
     if (cap < size)
         cap = size;
     VecNested vec = vec_alloc(cap, item_type, depth);
@@ -564,10 +568,16 @@ PyTypeObject VecNestedType = {
     // TODO: free
 };
 
-PyObject *VecNested_FromIterable(size_t item_type, size_t depth, PyObject *iterable) {
-    VecNested v = vec_alloc(0, item_type, depth);
+PyObject *VecNested_FromIterable(size_t item_type, size_t depth, PyObject *iterable, int64_t cap) {
+    VecNested v = vec_alloc(cap, item_type, depth);
     if (VEC_IS_ERROR(v))
         return NULL;
+    if (cap > 0) {
+        for (int64_t i = 0; i < cap; i++) {
+            v.buf->items[i].len = -1;
+            v.buf->items[i].buf = NULL;
+        }
+    }
     v.len = 0;
 
     PyObject *iter = PyObject_GetIter(iterable);

mypyc/lib-rt/vecs/vec_t.c

@@ -81,6 +81,10 @@ VecT VecT_ConvertFromNested(VecNestedBufItem item) {
 }
 
 VecT VecT_New(Py_ssize_t size, Py_ssize_t cap, size_t item_type) {
+    if (cap < 0) {
+        PyErr_SetString(PyExc_ValueError, "capacity must not be negative");
+        return vec_error();
+    }
     if (cap < size)
         cap = size;
     VecT vec = vec_alloc(cap, item_type);
@@ -557,10 +561,14 @@ PyTypeObject VecTType = {
     // TODO: free
 };
 
-PyObject *VecT_FromIterable(size_t item_type, PyObject *iterable) {
-    VecT v = vec_alloc(0, item_type);
+PyObject *VecT_FromIterable(size_t item_type, PyObject *iterable, int64_t cap) {
+    VecT v = vec_alloc(cap, item_type);
     if (VEC_IS_ERROR(v))
         return NULL;
+    if (cap > 0) {
+        for (int64_t i = 0; i < cap; i++)
+            v.buf->items[i] = NULL;
+    }
     v.len = 0;
 
     PyObject *iter = PyObject_GetIter(iterable);

mypyc/lib-rt/vecs/vec_template.c

@@ -84,8 +84,12 @@ VEC FUNC(ConvertFromNested)(VecNestedBufItem item) {
 }
 
 VEC FUNC(New)(Py_ssize_t size, Py_ssize_t cap) {
+    if (cap < 0) {
+        PyErr_SetString(PyExc_ValueError, "capacity must not be negative");
+        return vec_error();
+    }
     if (cap < size)
-        size = cap;
+        cap = size;
     VEC vec = vec_alloc(cap);
     if (VEC_IS_ERROR(vec))
         return vec;
@@ -96,10 +100,13 @@ VEC FUNC(New)(Py_ssize_t size, Py_ssize_t cap) {
     return vec;
 }
 
-PyObject *FUNC(FromIterable)(PyObject *iterable) {
-    VEC v = vec_alloc(0);
+PyObject *FUNC(FromIterable)(PyObject *iterable, int64_t cap) {
+    VEC v = vec_alloc(cap);
     if (VEC_IS_ERROR(v))
         return NULL;
+    if (cap > 0) {
+        memset(v.buf->items, 0, sizeof(ITEM_C_TYPE) * cap);
+    }
     v.len = 0;
 
     PyObject *iter = PyObject_GetIter(iterable);
@@ -132,15 +139,20 @@ PyObject *FUNC(FromIterable)(PyObject *iterable) {
 }
 
 static PyObject *vec_new(PyTypeObject *self, PyObject *args, PyObject *kw) {
-    static char *kwlist[] = {"", NULL};
+    static char *kwlist[] = {"", "capacity", NULL};
     PyObject *init = NULL;
-    if (!PyArg_ParseTupleAndKeywords(args, kw, "|O:vec", kwlist, &init)) {
+    int64_t cap = 0;
+    if (!PyArg_ParseTupleAndKeywords(args, kw, "|OL:vec", kwlist, &init, &cap)) {
+        return NULL;
+    }
+    if (cap < 0) {
+        PyErr_SetString(PyExc_ValueError, "capacity must not be negative");
         return NULL;
     }
     if (init == NULL) {
-        return FUNC(Box)(FUNC(New)(0, 0));
+        return FUNC(Box)(FUNC(New)(0, cap));
     } else {
-        return (PyObject *)FUNC(FromIterable)(init);
+        return (PyObject *)FUNC(FromIterable)(init, cap);
     }
 }