Skip to content

Commit f94d18c

Browse files
authored
[mypyc] Faster vec init and extend with bytes arguments (#21355)
Special case `vec[u8](<bytes>)` and `extend(<vec_u8>, <bytes>)`. These are expected to be common operations and the specialization makes them significantly faster with short `bytes` objects. Existing test coverage is sufficient. I used microbenchmarks to check that the performance impact is non-trivial.
1 parent 38c3fcb commit f94d18c

1 file changed

Lines changed: 19 additions & 0 deletions

File tree

mypyc/lib-rt/vecs/vec_template.c

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,18 @@ VEC FUNC(FromIterable)(PyObject *iterable, int64_t cap) {
141141
return vec_error();
142142
}
143143

144+
if (ITEM_TYPE_MAGIC == VEC_ITEM_TYPE_U8 && PyBytes_CheckExact(iterable)) {
145+
Py_ssize_t n = PyBytes_GET_SIZE(iterable);
146+
Py_ssize_t alloc_size = n > cap ? n : cap;
147+
VEC v = vec_alloc(alloc_size);
148+
if (VEC_IS_ERROR(v))
149+
return vec_error();
150+
if (n > 0)
151+
memcpy(v.buf->items, PyBytes_AS_STRING(iterable), n);
152+
v.len = n;
153+
return v;
154+
}
155+
144156
#ifdef BUFFER_FORMAT_CHAR_OK
145157
Py_buffer view;
146158
int buf_ok = vec_get_buffer(iterable, &view);
@@ -443,6 +455,13 @@ VEC FUNC(Extend)(VEC vec, PyObject *iterable) {
443455
return FUNC(ExtendVec)(vec, ((VEC_OBJECT *)iterable)->vec);
444456
}
445457

458+
if (ITEM_TYPE_MAGIC == VEC_ITEM_TYPE_U8 && PyBytes_CheckExact(iterable)) {
459+
Py_ssize_t n = PyBytes_GET_SIZE(iterable);
460+
if (n > 0)
461+
return vec_extend_items(vec, (const ITEM_C_TYPE *)PyBytes_AS_STRING(iterable), n, 0);
462+
return vec;
463+
}
464+
446465
#ifdef BUFFER_FORMAT_CHAR_OK
447466
Py_buffer view;
448467
int buf_ok = vec_get_buffer(iterable, &view);

0 commit comments

Comments
 (0)