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.

importpyperfrunner=pyperf.Runner()
count1=1_000count2=100count3=10_000CHUNK_A=b'a'*count1CHUNK_B=b'b'*count2CHUNK_C=b'c'*count3defbuild_bytes():
# Bytes not uniquely referenced.ba=bytearray()
ba+=CHUNK_Aba+=CHUNK_Bba+=CHUNK_Cdefbuild_bytes_unique():
ba=bytearray()
# Repeat inline results in uniquely referenced bytes.ba+=b'a'*count1ba+=b'b'*count2ba+=b'c'*count3defbuild_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' * 1000discuss 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

[encukou]

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

// make some bytesPyObject*bytes=PyBytes_FromString("aaB");
assert(bytes);
// make an empty bytearrayPyObject*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 bytearrayPy_DECREF(PyObject_CallMethod(ba, "pop", ""));
// check that our bytes was not modifiedassert(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)) {
	returnPyStackRef_IsHeapSafe(*stackpointer);
	}
	}
	return0;
	}

[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.

[github-actions]

This PR is stale because it has been open for 30 days with no activity.