diff --git a/mypy/typeshed/stubs/librt/librt/vecs.pyi b/mypy/typeshed/stubs/librt/librt/vecs.pyi index dbf199da9465..b1320eccadaa 100644 --- a/mypy/typeshed/stubs/librt/librt/vecs.pyi +++ b/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: ... diff --git a/mypyc/irbuild/expression.py b/mypyc/irbuild/expression.py index dab794aeb82a..652279936eff 100644 --- a/mypyc/irbuild/expression.py +++ b/mypyc/irbuild/expression.py @@ -364,10 +364,20 @@ 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 +571,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 +599,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]([ for 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) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 89d201f5b823..33a716624b17 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/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,7 +374,7 @@ 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, ) @@ -380,7 +382,9 @@ def set_item(x: Value, y: Value, z: Value, line: int) -> 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: diff --git a/mypyc/irbuild/vec.py b/mypyc/irbuild/vec.py index d17504150395..6b78c197e5d1 100644 --- a/mypyc/irbuild/vec.py +++ b/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) diff --git a/mypyc/lib-rt/vecs/librt_vecs.c b/mypyc/lib-rt/vecs/librt_vecs.c index 13d540a1dc45..9bb93b9562b0 100644 --- a/mypyc/lib-rt/vecs/librt_vecs.c +++ b/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); } } } diff --git a/mypyc/lib-rt/vecs/librt_vecs.h b/mypyc/lib-rt/vecs/librt_vecs.h index ff289ce2633b..97d8a11e1e58 100644 --- a/mypyc/lib-rt/vecs/librt_vecs.h +++ b/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); diff --git a/mypyc/lib-rt/vecs/vec_nested.c b/mypyc/lib-rt/vecs/vec_nested.c index 6743935b96da..ba09b1bba697 100644 --- a/mypyc/lib-rt/vecs/vec_nested.c +++ b/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 > 0 ? cap : 0, 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); diff --git a/mypyc/lib-rt/vecs/vec_t.c b/mypyc/lib-rt/vecs/vec_t.c index a68712d69b7b..a73f2ff00adf 100644 --- a/mypyc/lib-rt/vecs/vec_t.c +++ b/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 > 0 ? cap : 0, 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); diff --git a/mypyc/lib-rt/vecs/vec_template.c b/mypyc/lib-rt/vecs/vec_template.c index 5978f4af5fb2..b5d61065b2e8 100644 --- a/mypyc/lib-rt/vecs/vec_template.c +++ b/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 > 0 ? cap : 0); 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); } } diff --git a/mypyc/test-data/irbuild-vec-i64.test b/mypyc/test-data/irbuild-vec-i64.test index 22687d759c44..0aa948663f53 100644 --- a/mypyc/test-data/irbuild-vec-i64.test +++ b/mypyc/test-data/irbuild-vec-i64.test @@ -820,3 +820,82 @@ L1: r2 = r1.len L2: return 1 + +[case testVecI64CreateWithCap_64bit] +from librt.vecs import vec +from mypy_extensions import i64 + +def cap_only() -> vec[i64]: + return vec[i64](capacity=5) + +def cap_variable(n: i64) -> vec[i64]: + return vec[i64](capacity=n) +[out] +def cap_only(): + r0 :: vec[i64] +L0: + r0 = VecI64Api.alloc(0, 5) + return r0 +def cap_variable(n): + n :: i64 + r0 :: vec[i64] +L0: + r0 = VecI64Api.alloc(0, n) + return r0 + +[case testVecI64CreateFromListWithCap_64bit] +from librt.vecs import vec +from mypy_extensions import i64 + +def list_with_cap() -> vec[i64]: + return vec[i64]([1, 2], capacity=5) +[out] +def list_with_cap(): + r0 :: vec[i64] + r1 :: object + r2, r3, r4 :: ptr +L0: + r0 = VecI64Api.alloc(2, 5) + r1 = r0.buf + r2 = get_element_ptr r1 items :: VecI64BufObject + set_mem r2, 1 :: i64* + r3 = r2 + 8 + set_mem r3, 2 :: i64* + r4 = r3 + 8 + keep_alive r0 + return r0 + +[case testVecI64CreateFromListMultiplyWithCap_64bit] +from librt.vecs import vec +from mypy_extensions import i64 + +def repeated_with_cap(n: i64) -> vec[i64]: + return vec[i64]([3] * n, capacity=10) +[out] +def repeated_with_cap(n): + n :: i64 + r0 :: vec[i64] + r1 :: object + r2 :: ptr + r3 :: i64 + r4, r5 :: ptr + r6 :: bit + r7 :: ptr +L0: + r0 = VecI64Api.alloc(n, 10) + r1 = r0.buf + r2 = get_element_ptr r1 items :: VecI64BufObject + r3 = n * 8 + r4 = r2 + r3 + r5 = r2 +L1: + r6 = r5 < r4 :: unsigned + if r6 goto L2 else goto L3 :: bool +L2: + set_mem r5, 3 :: i64* + r7 = r5 + 8 + r5 = r7 + goto L1 +L3: + keep_alive r0 + return r0 diff --git a/mypyc/test-data/irbuild-vec-nested.test b/mypyc/test-data/irbuild-vec-nested.test index 6b3389394d93..269017d55e4c 100644 --- a/mypyc/test-data/irbuild-vec-nested.test +++ b/mypyc/test-data/irbuild-vec-nested.test @@ -309,3 +309,41 @@ L5: r10 = load_mem r9 :: vec[vec[str]]* keep_alive v return r10 + +[case testVecNestedCreateWithCap_64bit] +from librt.vecs import vec +from mypy_extensions import i64 + +def cap_str() -> vec[vec[str]]: + return vec[vec[str]](capacity=5) + +def cap_i64() -> vec[vec[i64]]: + return vec[vec[i64]](capacity=5) + +def cap_variable(n: i64) -> vec[vec[str]]: + return vec[vec[str]](capacity=n) +[out] +def cap_str(): + r0 :: object + r1 :: ptr + r2 :: vec[vec[str]] +L0: + r0 = load_address PyUnicode_Type + r1 = r0 + r2 = VecNestedApi.alloc(0, 5, r1, 1) + return r2 +def cap_i64(): + r0 :: vec[vec[i64]] +L0: + r0 = VecNestedApi.alloc(0, 5, 2, 1) + return r0 +def cap_variable(n): + n :: i64 + r0 :: object + r1 :: ptr + r2 :: vec[vec[str]] +L0: + r0 = load_address PyUnicode_Type + r1 = r0 + r2 = VecNestedApi.alloc(0, n, r1, 1) + return r2 diff --git a/mypyc/test-data/irbuild-vec-t.test b/mypyc/test-data/irbuild-vec-t.test index 7e5c31fd0ded..60d1f035c7fb 100644 --- a/mypyc/test-data/irbuild-vec-t.test +++ b/mypyc/test-data/irbuild-vec-t.test @@ -359,3 +359,62 @@ def f() -> None: main:5: error: Invalid item type for "vec" main:6: error: Invalid item type for "vec" main:7: error: Invalid item type for "vec" + +[case testVecTCreateWithCap_64bit] +from librt.vecs import vec +from mypy_extensions import i64 + +def cap_only() -> vec[str]: + return vec[str](capacity=5) + +def cap_variable(n: i64) -> vec[str]: + return vec[str](capacity=n) +[out] +def cap_only(): + r0 :: object + r1 :: ptr + r2 :: vec[str] +L0: + r0 = load_address PyUnicode_Type + r1 = r0 + r2 = VecTApi.alloc(0, 5, r1) + return r2 +def cap_variable(n): + n :: i64 + r0 :: object + r1 :: ptr + r2 :: vec[str] +L0: + r0 = load_address PyUnicode_Type + r1 = r0 + r2 = VecTApi.alloc(0, n, r1) + return r2 + +[case testVecTCreateFromListWithCap_64bit] +from librt.vecs import vec +from mypy_extensions import i64 + +def list_with_cap() -> vec[str]: + return vec[str](['a', 'b'], capacity=5) +[out] +def list_with_cap(): + r0, r1 :: str + r2 :: object + r3 :: ptr + r4 :: vec[str] + r5 :: object + r6, r7, r8 :: ptr +L0: + r0 = 'a' + r1 = 'b' + r2 = load_address PyUnicode_Type + r3 = r2 + r4 = VecTApi.alloc(2, 5, r3) + r5 = r4.buf + r6 = get_element_ptr r5 items :: VecTBufObject + set_mem r6, r0 :: builtins.str* + r7 = r6 + 8 + set_mem r7, r1 :: builtins.str* + r8 = r7 + 8 + keep_alive r4 + return r4 diff --git a/mypyc/test-data/run-vecs-i64-interp.test b/mypyc/test-data/run-vecs-i64-interp.test index 5fd2c00d984f..965e1cd04b58 100644 --- a/mypyc/test-data/run-vecs-i64-interp.test +++ b/mypyc/test-data/run-vecs-i64-interp.test @@ -417,3 +417,53 @@ def test_pop_index() -> None: v, n = pop(v, 0) assert n == 15 assert v == vec[i64]() + +def test_cap_basic() -> None: + v = vec[i64](capacity=3) + assert len(v) == 0 + assert v == vec[i64]() + v = vec[i64](capacity=0) + assert len(v) == 0 + v = vec[i64]([10, 20], capacity=3) + assert v == vec[i64]([10, 20]) + v = vec[i64]([5], capacity=0) + assert v == vec[i64]([5]) + v = vec[i64]([10, 20, 30], capacity=1) + assert len(v) == 3 + assert v[0] == 10 + +def test_cap_negative() -> None: + with assertRaises(ValueError): + vec[i64](capacity=-1) + with assertRaises(ValueError): + vec[i64]([1], capacity=-1) + +def test_cap_buffer_reuse() -> None: + v = vec[i64](capacity=3) + old = v + v = append(v, 1) + v = append(v, 2) + v = append(v, 3) + assert len(v) == 3 + old = append(old, 99) + assert v[0] == 99 + + v = vec[i64]([10, 20], capacity=3) + old = v + v = append(v, 30) + v[0] = 99 + assert old[0] == 99 + +def test_cap_reallocation() -> None: + v = vec[i64]([10, 20], capacity=2) + old = v + v = append(v, 30) + old = append(old, 99) + assert v[0] == 10 + + v = vec[i64]([10, 20, 30], capacity=3) + old = v + v = append(v, 40) + v[0] = 99 + assert old[0] == 10 + assert old[2] == 30 diff --git a/mypyc/test-data/run-vecs-i64.test b/mypyc/test-data/run-vecs-i64.test index 0c1fb2f83e54..afc0ca28fcf9 100644 --- a/mypyc/test-data/run-vecs-i64.test +++ b/mypyc/test-data/run-vecs-i64.test @@ -488,3 +488,50 @@ def test_tuple() -> None: assert b == t[1] with assertRaises(RuntimeError): ftuple(True) + +def test_cap_empty() -> None: + v = vec[i64](capacity=3) + assert len(v) == 0 + assert v == vec[i64]() + +def test_cap_with_initializer() -> None: + v = vec[i64]([1, 2], capacity=3) + assert len(v) == 2 + assert v[0] == 1 + assert v[1] == 2 + +def test_cap_from_range() -> None: + v = vec[i64](range(7), capacity=10) + old = v + v = append(v, 7) + v = append(v, 8) + v = append(v, 9) + v[0] = 99 + assert old[0] == 99 + assert v == vec[i64]([99, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + +def test_cap_variable(n: i64 = 3) -> None: + v = vec[i64](capacity=n) + assert len(v) == 0 + v = append(v, 5) + assert v[0] == 5 + +def test_cap_buffer_reuse() -> None: + v = vec[i64]([1, 2], capacity=3) + old = v + v = append(v, 3) + v[0] = 99 + assert old[0] == 99 # shared buffer + +def test_cap_below_initializer_length() -> None: + v = vec[i64]([10, 20, 30], capacity=1) + assert len(v) == 3 + assert v[0] == 10 + assert v[1] == 20 + assert v[2] == 30 + +def test_cap_negative() -> None: + with assertRaises(ValueError): + vec[i64](capacity=-1) + with assertRaises(ValueError): + vec[i64]([1], capacity=-1) diff --git a/mypyc/test-data/run-vecs-misc-interp.test b/mypyc/test-data/run-vecs-misc-interp.test index 07ba78c3f4a2..318248d1b153 100644 --- a/mypyc/test-data/run-vecs-misc-interp.test +++ b/mypyc/test-data/run-vecs-misc-interp.test @@ -453,6 +453,58 @@ def test_iterator_protocol() -> None: with assertRaises(StopIteration): next(it) +def test_cap_empty() -> None: + for t in ITEM_TYPES: + v = vec[t](capacity=3) + assert len(v) == 0 + assert v == vec[t]() + +def test_cap_with_initializer() -> None: + v = vec[float]([1.5, 2.5], capacity=3) + assert len(v) == 2 + assert v[0] == 1.5 + assert v[1] == 2.5 + + v = vec[u8]([10, 20], capacity=3) + assert len(v) == 2 + assert v[0] == 10 + + v = vec[i16]([-100, 100], capacity=3) + assert len(v) == 2 + + v = vec[i32]([1000, 2000], capacity=3) + assert len(v) == 2 + + v = vec[bool]([True, False], capacity=3) + assert len(v) == 2 + assert v[0] is True + assert v[1] is False + +def test_cap_negative() -> None: + for t in ITEM_TYPES: + with assertRaises(ValueError): + vec[t](capacity=-1) + +def test_cap_buffer_reuse() -> None: + # Verify buffer reuse within cap for each type + v = vec[float]([1.5], capacity=3) + old = v + v = append(v, 2.5) + v[0] = 9.5 + assert old[0] == 9.5 # shared buffer + + v = vec[u8]([10], capacity=3) + old = v + v = append(v, 20) + v[0] = 99 + assert old[0] == 99 + + v = vec[bool]([True], capacity=3) + old = v + v = append(v, False) + v[0] = False + assert old[0] is False + [case testLibrtVecsFeaturesNotAvailableInNonExperimentalBuild_librt] # This also ensures librt.vecs can be built without experimental features import librt.vecs diff --git a/mypyc/test-data/run-vecs-misc.test b/mypyc/test-data/run-vecs-misc.test index 1caedbf7da87..8f96007c92dc 100644 --- a/mypyc/test-data/run-vecs-misc.test +++ b/mypyc/test-data/run-vecs-misc.test @@ -236,3 +236,30 @@ def test_nested_misc() -> None: v[0][0] = 3.5 assert v[0] == vec[float]([3.5]) assert v[:5] == vec[vec[float]]([vec[float]([3.5])]) + +def test_cap_misc() -> None: + v_f = vec[float](capacity=3) + assert len(v_f) == 0 + v_f = append(v_f, 1.5) + assert v_f[0] == 1.5 + + v_u8 = vec[u8]([1, 2], capacity=5) + assert len(v_u8) == 2 + assert v_u8[0] == 1 + + v_i16 = vec[i16](capacity=2) + v_i16 = append(v_i16, i16(10)) + assert v_i16[0] == 10 + + v_i32 = vec[i32]([i32(7)], capacity=3) + assert v_i32[0] == 7 + + v_bool = vec[bool](capacity=4) + assert len(v_bool) == 0 + v_bool = append(v_bool, True) + assert v_bool[0] == True + + v_nested = vec[vec[float]](capacity=2) + assert len(v_nested) == 0 + v_nested = append(v_nested, vec[float]([1.0])) + assert v_nested[0] == vec[float]([1.0]) diff --git a/mypyc/test-data/run-vecs-nested-interp.test b/mypyc/test-data/run-vecs-nested-interp.test index d4b0af15aa18..89d7a21e1cba 100644 --- a/mypyc/test-data/run-vecs-nested-interp.test +++ b/mypyc/test-data/run-vecs-nested-interp.test @@ -467,3 +467,52 @@ def test_append_with_alias() -> None: assert len(alias) == 1 assert v[0] == vec[i64]([1, 2, 3]) assert len(v) == 4 + +def test_cap_basic() -> None: + v = vec[vec[str]](capacity=3) + assert len(v) == 0 + assert v == vec[vec[str]]() + v = vec[vec[str]](capacity=0) + assert len(v) == 0 + v = vec[vec[str]]([vec[str](['a']), vec[str](['b'])], capacity=3) + assert v[0] == vec[str](['a']) + v = vec[vec[str]]([vec[str](['x'])], capacity=0) + assert v == vec[vec[str]]([vec[str](['x'])]) + v = vec[vec[str]]([vec[str](['a']), vec[str](['b']), vec[str](['c'])], capacity=1) + assert len(v) == 3 + +def test_cap_negative() -> None: + with assertRaises(ValueError): + vec[vec[str]](capacity=-1) + with assertRaises(ValueError): + vec[vec[str]]([vec[str]()], capacity=-1) + +def test_cap_buffer_reuse() -> None: + v = vec[vec[str]](capacity=3) + old = v + v = append(v, vec[str](['a'])) + v = append(v, vec[str](['b'])) + v = append(v, vec[str](['c'])) + assert len(v) == 3 + old = append(old, vec[str](['Q'])) + assert v[0] == vec[str](['Q']) + + v = vec[vec[str]]([vec[str](['a']), vec[str](['b'])], capacity=3) + old = v + v = append(v, vec[str](['c'])) + v[0] = vec[str](['z']) + assert old[0] == vec[str](['z']) + +def test_cap_reallocation() -> None: + v = vec[vec[str]]([vec[str](['a']), vec[str](['b'])], capacity=2) + old = v + v = append(v, vec[str](['c'])) + old = append(old, vec[str](['Q'])) + assert v[0] == vec[str](['a']) + + v = vec[vec[str]]([vec[str](['a']), vec[str](['b']), vec[str](['c'])], capacity=3) + old = v + v = append(v, vec[str](['d'])) + v[0] = vec[str](['z']) + assert old[0] == vec[str](['a']) + assert old[2] == vec[str](['c']) diff --git a/mypyc/test-data/run-vecs-nested.test b/mypyc/test-data/run-vecs-nested.test index 49a6b892f20b..d1f8ee350827 100644 --- a/mypyc/test-data/run-vecs-nested.test +++ b/mypyc/test-data/run-vecs-nested.test @@ -518,3 +518,40 @@ def test_append_with_alias() -> None: assert len(alias) == 1 assert v[0] == vec[i64]([1, 2, 3]) assert len(v) == 4 + +def test_cap_empty() -> None: + v = vec[vec[i64]](capacity=3) + assert len(v) == 0 + assert v == vec[vec[i64]]() + +def test_cap_with_initializer() -> None: + v = vec[vec[i64]]([vec[i64]([1]), vec[i64]([2])], capacity=3) + assert len(v) == 2 + assert v[0] == vec[i64]([1]) + assert v[1] == vec[i64]([2]) + +def test_cap_variable(n: i64 = 3) -> None: + v = vec[vec[str]](capacity=n) + assert len(v) == 0 + v = append(v, vec[str](["x"])) + assert v[0] == vec[str](["x"]) + +def test_cap_buffer_reuse() -> None: + v = vec[vec[i64]]([vec[i64]([1]), vec[i64]([2])], capacity=3) + old = v + v = append(v, vec[i64]([3])) + v[0] = vec[i64]([99]) + assert old[0] == vec[i64]([99]) # shared buffer + +def test_cap_below_initializer_length() -> None: + v = vec[vec[i64]]([vec[i64]([10]), vec[i64]([20]), vec[i64]([30])], capacity=1) + assert len(v) == 3 + assert v[0] == vec[i64]([10]) + assert v[1] == vec[i64]([20]) + assert v[2] == vec[i64]([30]) + +def test_cap_negative() -> None: + with assertRaises(ValueError): + vec[vec[i64]](capacity=-1) + with assertRaises(ValueError): + vec[vec[i64]]([vec[i64]([1])], capacity=-1) diff --git a/mypyc/test-data/run-vecs-t-interp.test b/mypyc/test-data/run-vecs-t-interp.test index aab522acc9ce..5158aa3e78f8 100644 --- a/mypyc/test-data/run-vecs-t-interp.test +++ b/mypyc/test-data/run-vecs-t-interp.test @@ -545,6 +545,68 @@ def test_append_with_alias() -> None: assert v[0] == 'a' assert len(v) == 4 +def test_cap_basic() -> None: + v = vec[str](capacity=3) + assert len(v) == 0 + assert v == vec[str]() + v = vec[str](capacity=0) + assert len(v) == 0 + v = vec[str](['a', 'b'], capacity=3) + assert v == vec[str](['a', 'b']) + v = vec[str](['x'], capacity=0) + assert v == vec[str](['x']) + v = vec[str](['a', 'b', 'c'], capacity=1) + assert len(v) == 3 + assert v[0] == 'a' + +def test_cap_negative() -> None: + with assertRaises(ValueError): + vec[str](capacity=-1) + with assertRaises(ValueError): + vec[str](['x'], capacity=-1) + +def test_cap_buffer_reuse() -> None: + v = vec[str](capacity=3) + old = v + v = append(v, 'a') + v = append(v, 'b') + v = append(v, 'c') + assert len(v) == 3 + old = append(old, 'Q') + assert v[0] == 'Q' + + v = vec[str](['a', 'b'], capacity=3) + old = v + v = append(v, 'c') + v[0] = 'z' + assert old[0] == 'z' + +def test_cap_reallocation() -> None: + # At capacity -> next append reallocates, aliases become independent + v = vec[str](['a', 'b'], capacity=2) + old = v + v = append(v, 'c') + old = append(old, 'Q') + assert v[0] == 'a' + + # Old alias retains original items after reallocation + v = vec[str](['a', 'b', 'c'], capacity=3) + old = v + v = append(v, 'd') + v[0] = 'z' + assert old[0] == 'a' + assert old[2] == 'c' + +def test_cap_optional() -> None: + v = vec[Optional[str]](capacity=3) + v = append(v, 'x') + v = append(v, None) + assert v[0] == 'x' + assert v[1] is None + + v = vec[Optional[str]](['x', None], capacity=3) + assert v[0] == 'x' + assert v[1] is None [file helpers.py] # Test helper classes diff --git a/mypyc/test-data/run-vecs-t.test b/mypyc/test-data/run-vecs-t.test index 6855aba68d57..bd45fc94a364 100644 --- a/mypyc/test-data/run-vecs-t.test +++ b/mypyc/test-data/run-vecs-t.test @@ -448,3 +448,55 @@ def test_append_with_alias() -> None: assert len(alias) == 1 assert v[0] == 'a' assert len(v) == 4 + +def test_cap_empty() -> None: + v = vec[str](capacity=3) + assert len(v) == 0 + assert v == vec[str]() + +def test_cap_with_initializer() -> None: + v = vec[str](['a', 'b'], capacity=3) + assert len(v) == 2 + assert v[0] == 'a' + assert v[1] == 'b' + +def test_cap_from_comprehension_and_iterable() -> None: + v = vec[str]([str(i) for i in range(2)], capacity=3) + old = v + v = append(v, '2') + v[0] = 'x' + assert old[0] == 'x' + assert v == vec[str](['x', '1', '2']) + + v = vec[str]((str(i) for i in range(2)), capacity=3) + old = v + v = append(v, '2') + v[0] = 'y' + assert old[0] == 'y' + assert v == vec[str](['y', '1', '2']) + +def test_cap_variable(n: i64 = 3) -> None: + v = vec[str](capacity=n) + assert len(v) == 0 + v = append(v, 'x') + assert v[0] == 'x' + +def test_cap_buffer_reuse() -> None: + v = vec[str](['a', 'b'], capacity=3) + old = v + v = append(v, 'c') + v[0] = 'z' + assert old[0] == 'z' # shared buffer + +def test_cap_below_initializer_length() -> None: + v = vec[str](['a', 'b', 'c'], capacity=1) + assert len(v) == 3 + assert v[0] == 'a' + assert v[1] == 'b' + assert v[2] == 'c' + +def test_cap_negative() -> None: + with assertRaises(ValueError): + vec[str](capacity=-1) + with assertRaises(ValueError): + vec[str](['x'], capacity=-1)