Repository: numpy/numpy
File: numpy/_core/src/umath/string_ufuncs.cpp (C++)

Code before:

template <ENCODING enc>
static inline void
string_add(Buffer<enc> buf1, Buffer<enc> buf2, Buffer<enc> out)
{
    size_t len1 = buf1.num_codepoints();
    size_t len2 = buf2.num_codepoints();
    buf1.buffer_memcpy(out, len1);
    buf2.buffer_memcpy(out + len1, len2);
    out.buffer_fill_with_zeros_after_index(len1 + len2);
}


template <ENCODING enc>
static inline void
string_multiply(Buffer<enc> buf1, npy_int64 reps, Buffer<enc> out)
{
    size_t len1 = buf1.num_codepoints();
    if (reps < 1 || len1 == 0) {
        out.buffer_fill_with_zeros_after_index(0);
        return;
    }

    size_t width = out.buffer_width();
    size_t pad = 0;
    if (width < len1*reps) {
        reps = width / len1;
        pad = width % len1;
    }
    if (len1 == 1) {
        out.buffer_memset(*buf1, reps);
        out.buffer_fill_with_zeros_after_index(reps);
    }
    else {
        for (npy_int64 i = 0; i < reps; i++) {
            buf1.buffer_memcpy(out, len1);
            out += len1;
        }
        buf1.buffer_memcpy(out, pad);
        out += pad;
        out.buffer_fill_with_zeros_after_index(0);
    }
}


template <ENCODING enc>
static int
string_add_loop(PyArrayMethod_Context *context,
                char *const data[], npy_intp const dimensions[],
                npy_intp const strides[], NpyAuxData *NPY_UNUSED(auxdata))
{

Review comment (bug):
```suggestion
    if (reps > NPY_MAX_SSIZE_T  // ssize_t is enough here
        || npy_mul_sizes_with_overflow(len1, reps, &width)
        || width > out.buffer_width()) {
            width = out.buffer_width;
```
Admittedly this is tedious and pedantic, but maybe we should do overflow checks?  `npy_mul_sizes_with_overflow` uses `ssize_t` inputs, which is safe here but may require casting to tell the compiler about that.  (Or just using the `size_t` version.)
I suppose `NPY_UNLIKELY()` will be irrelevant here.

Code after:


template <ENCODING enc>
static inline void
string_add(Buffer<enc> buf1, Buffer<enc> buf2, Buffer<enc> out)
{
    size_t len1 = buf1.num_codepoints();
    size_t len2 = buf2.num_codepoints();
    buf1.buffer_memcpy(out, len1);
    buf2.buffer_memcpy(out + len1, len2);
    out.buffer_fill_with_zeros_after_index(len1 + len2);
}


template <ENCODING enc>
static inline int
string_multiply(Buffer<enc> buf1, npy_int64 reps, Buffer<enc> out)
{
    size_t len1 = buf1.num_codepoints();
    if (reps < 1 || len1 == 0) {
        out.buffer_fill_with_zeros_after_index(0);
        return 0;
    }

    if (len1 == 1) {
        out.buffer_memset(*buf1, reps);
        out.buffer_fill_with_zeros_after_index(reps);
        return 0;
    }

    size_t newlen;
    if (NPY_UNLIKELY(npy_mul_with_overflow_size_t(&newlen, reps, len1) != 0) || newlen > PY_SSIZE_T_MAX) {
        return -1;
    }

    size_t pad = 0;
    size_t width = out.buffer_width();
    if (width < newlen) {
        reps = width / len1;
        pad = width % len1;
    }

    for (npy_int64 i = 0; i < reps; i++) {
        buf1.buffer_memcpy(out, len1);
        out += len1;
    }

    buf1.buffer_memcpy(out, pad);
    out += pad;

    out.buffer_fill_with_zeros_after_index(0);