gh-139871: Optimize bytearray unique bytes iconcat

[cmaloney]

If the bytearray is empty and a uniquely referenced bytes object is being concatenated (ex. one just received from read), just use its storage as the backing for the bytearray rather than copying it. The bigger the bytes the bigger the saving.

build_bytes_unique: Mean +- std dev: [base] 383 ns +- 11 ns -> [iconcat_opt] 342 ns +- 5 ns: 1.12x faster
build_bytearray: Mean +- std dev: [base] 496 ns +- 8 ns -> [iconcat_opt] 471 ns +- 13 ns: 1.05x faster
encode: Mean +- std dev: [base] 482 us +- 2 us -> [iconcat_opt] 13.8 us +- 0.1 us: 34.78x faster

Benchmark hidden because not significant (1): build_bytes

Geometric mean: 2.53x faster

note: Performance of build_bytes is expected to stay constant.

import pyperf

runner = pyperf.Runner()

count1 = 1_000
count2 = 100
count3 = 10_000

CHUNK_A = b'a' * count1
CHUNK_B = b'b' * count2
CHUNK_C = b'c' * count3

def build_bytes():
    # Bytes not uniquely referenced.
    ba = bytearray()
    ba += CHUNK_A
    ba += CHUNK_B
    ba += CHUNK_C

def build_bytes_unique():
    ba = bytearray()
    # Repeat inline results in uniquely referenced bytes.
    ba += b'a' * count1
    ba += b'b' * count2
    ba += b'c' * count3

def build_bytearray():
    # Each bytearray appended is uniquely referenced.
    ba = bytearray()
    ba += bytearray(CHUNK_A)
    ba += bytearray(CHUNK_B)
    ba += bytearray(CHUNK_C)

runner.bench_func('build_bytes', build_bytes)
runner.bench_func('build_bytes_unique', build_bytes_unique)
runner.bench_func('build_bytearray', build_bytearray)
runner.timeit(
    name="encode",
    setup="a = 'a' * 1_000_000",
    stmt="bytearray(a, encoding='utf8')")

---

From my understanding of reference counting I think this is safe to do for iconcat (and would be safe to do for ba[:] = b'\0' * 1000 discuss topic). The briefly refcount 2 isn't ideal but I think good enough for the performance delta. I'm hoping if I can ship an implementation of gh-87613 can do the same optimization for bytearray(b'\0' * 4096).

If the iconcat refcount 2 part isn't good, can tweak to keep the enecode + bytearray performance improvement without changing iconcat generally.

cc: @vstinner , @encukou

• Issue: Add .take_bytes([n]) a zero-copy path from bytearray to bytes #139871

[vstinner]

Can you replace the comment with an assertion?

[vstinner]

Why not also run this check when (self->ob_alloc == 0 && _PyObject_IsUniquelyReferenced(other)) is false? So move it before the outer if?

[cmaloney]

For the general case it's currently checked by bytearray_resize_lock_held.

Checking here technically changes the order of operations where an error can happen (the main codepath the buffer protocol always is called and can error before _canresize). Not certain how important that is to keep in this case. Keeping order would mean moving _canresize check inside the type checks here.

[encukou]

Here's a test that should pass, but doesn't:

    // make some bytes
    PyObject *bytes = PyBytes_FromString("aaB");
    assert(bytes);
    // make an empty bytearray
    PyObject *ba = PyByteArray_FromStringAndSize("", 0);
    assert(ba);
    // append bytes to bytearray (in place, getting a new reference)
    PyObject *new_ba = PySequence_InPlaceConcat(ba, bytes);
    assert(new_ba == ba);
    Py_DECREF(new_ba);
    // pop from bytearray
    Py_DECREF(PyObject_CallMethod(ba, "pop", ""));

    // check that our bytes was not modified
    assert(memcmp(PyBytes_AsString(bytes), "aaB", 3) == 0);

    Py_DECREF(bytes);
    Py_DECREF(ba);

AFAIK, you need to use PyUnstable_Object_IsUniqueReferencedTemporary.

[colesbury]

This looks unsafe to me. If you call a method, you may invalidate the assumptions you verified earlier

[vstinner]

Maybe call bytearray_take_bytes_impl() directly to reduce the risk of side effects? And you can check again _PyObject_IsUniquelyReferenced() in an assertion.

[cmaloney]

(Iterating on this locally; should have updates next week)

[cmaloney]

@encukou : When I try using PyUnstable_Object_IsUniqueReferencedTemporary then use ./python -m test test_bytes -W I get: python: ../cpython/Include/internal/pycore_stackref.h:695: PyObject *PyStackRef_AsPyObjectBorrow(_PyStackRef): Assertion !PyStackRef_IsTaggedInt(ref)' failed.`. Not sure exactly why.

I was modeling this off of PyBytes_Concat which does a _PyObject_IsUniquelyReferenced but looking more closely that isn't the implementation behind the sequence operations (bytes_concat) nor does it modify the right hand side in any way.

I don't see a way to do the generalized optimization currently; still think there should be a way just not sure the path and suspect it's a number of steps. I'll probably close this PR and open one for just the encoding case (bytearray('test', encoding='utf-8')) and put more general byearray iconcat, extend, and construction optimization in the back of my head for the moment.

[colesbury]

That seems like a bug in PyUnstable_Object_IsUniqueReferencedTemporary. We should skip over tagged ints when checking variables on the stack:

cpython/Objects/object.c

Lines 2758 to 2767 in c0c6514

	_PyStackRef *base = _PyFrame_Stackbase(frame);
	_PyStackRef *stackpointer = frame->stackpointer;
	while (stackpointer > base) {
	stackpointer--;
	if (op == PyStackRef_AsPyObjectBorrow(*stackpointer)) {
	return PyStackRef_IsHeapSafe(*stackpointer);
	}
	}
	return 0;
	}

[cmaloney]

Created GH-142243 doing just the bytearray('test', encoding='utf-8') portion of this. My time is becoming more limited shortly (starting a new job) so won't have quite as much time to explore new to me corners of CPython. Definitely happy to revisit more generally (there's a lot of optimizations both here and in bytes), maybe at PyConUS this year.

[cmaloney]

Tested with if (PyStackRef_IsTaggedInt(*stackpointer)) { continue; } in PyUnstable_Object_IsUniqueReferencedTemporary and that seems to work; that change I think needs a separate issue + news + test. Added to my backlog but not sure when I'll be able to get to.