gh-155742: Use PyBytesWriter in Python/assemble.c

[vstinner]

Replace soft deprecated _PyBytes_Resize() with PyBytesWriter.

• Issue: Treat bytes objects as immutable in C code #155742

[vstinner]

This change is a follow-up of PR gh-155747 which already modified a_linetable to use PyBytesWriter.

[methane]

how about PyBytesWriter_Grow(..., a->a_except_table_off + MAX_SIZE_OF_ENTRY)?

[vstinner]

I'm not sure what's the best memory allocation strategy. I wrote a program to stress test the assembler:

import os
code = [
    'x = 1',
    'y = 2',
    'z = 3',
]
lines = [
    'x += y',
    'y += z',
    'z = y - x',
    'try: x / y',
    'except ZeroDivisionError: pass',
]
for _ in range(1024):
    code.extend(lines)
code = '\n'.join(code)
os.uname()
compile(code, '<string>', 'exec')

I also wrote a local hack to display statistics on writer memory allocations:

Details

diff --git a/Python/assemble.c b/Python/assemble.c
index 8b92042345f..2577fa57620 100644
--- a/Python/assemble.c
+++ b/Python/assemble.c
@@ -1,4 +1,5 @@
 #include "Python.h"
+#include "pycore_bytesobject.h"     // PyBytesWriter
 #include "pycore_code.h"            // write_location_entry_start()
 #include "pycore_compile.h"
 #include "pycore_instruction_sequence.h"
@@ -50,14 +51,17 @@ instr_size(instruction *instr)
 struct assembler {
     PyBytesWriter *a_bytecode_writer; /* writer containing bytecode */
     PyObject *a_bytecode;      /* bytes containing bytecode */
+    int a_bytecode_resize;
     int a_offset;              /* offset into bytecode */
     PyBytesWriter *a_except_table_writer; /* writer containing exception table */
     PyObject *a_except_table;  /* bytes containing exception table */
+    int a_except_table_resize;  /* bytes containing exception table */
     int a_except_table_off;    /* offset into exception table */
     /* Location Info */
     int a_lineno;          /* lineno of last emitted instruction */
     PyBytesWriter *a_linetable_writer; /* writer containing location info */
     PyObject *a_linetable; /* bytes object containing location info */
+    int a_linetable_resize; /* bytes object containing location info */
     int a_location_off;    /* offset of last written location info frame */
 };
 
@@ -139,6 +143,7 @@ assemble_emit_exception_table_entry(struct assembler *a, int start, int end,
 {
     Py_ssize_t len = PyBytesWriter_GetSize(a->a_except_table_writer);
     if (a->a_except_table_off + MAX_SIZE_OF_ENTRY >= len) {
+        a->a_except_table_resize++;
         RETURN_IF_ERROR(PyBytesWriter_Resize(a->a_except_table_writer, len * 2));
     }
     int size = end-start;
@@ -292,6 +297,7 @@ write_location_info_entry(struct assembler* a, location loc, int isize)
     Py_ssize_t len = PyBytesWriter_GetSize(a->a_linetable_writer);
     if (a->a_location_off + THEORETICAL_MAX_ENTRY_SIZE >= len) {
         assert(len > THEORETICAL_MAX_ENTRY_SIZE);
+        a->a_linetable_resize++;
         RETURN_IF_ERROR(PyBytesWriter_Resize(a->a_linetable_writer, len * 2));
     }
     if (loc.lineno == NO_LOCATION.lineno) {
@@ -425,6 +431,7 @@ assemble_emit_instr(struct assembler *a, instruction *instr)
             PyErr_NoMemory();
             return ERROR;
         }
+        a->a_bytecode_resize++;
         RETURN_IF_ERROR(PyBytesWriter_Resize(a->a_bytecode_writer, len * 2));
     }
     code = (_Py_CODEUNIT *)PyBytesWriter_GetData(a->a_bytecode_writer) + a->a_offset;
@@ -433,6 +440,17 @@ assemble_emit_instr(struct assembler *a, instruction *instr)
     return SUCCESS;
 }
 
+static Py_ssize_t
+writer_allocated(PyBytesWriter *writer)
+{
+    if (writer->obj) {
+        return PyBytes_GET_SIZE(writer->obj);
+    }
+    else {
+        return sizeof(writer->small_buffer);
+    }
+}
+
 static int
 assemble_emit(struct assembler *a, instr_sequence *instrs,
               int first_lineno, PyObject *const_cache)
@@ -448,6 +466,11 @@ assemble_emit(struct assembler *a, instr_sequence *instrs,
 
     RETURN_IF_ERROR(assemble_exception_table(a, instrs));
 
+printf("assembler stats:\n");
+printf("- a_except_table: %i/%zd bytes; %i resize\n", a->a_except_table_off, writer_allocated(a->a_except_table_writer), a->a_except_table_resize);
+printf("- a_linetable: %i/%zd bytes; %i resize\n", a->a_location_off, writer_allocated(a->a_linetable_writer), a->a_linetable_resize);
+printf("- a_bytecode: %zd/%zd bytes; %i resize\n", a->a_offset * sizeof(_Py_CODEUNIT), writer_allocated(a->a_bytecode_writer), a->a_bytecode_resize);
+
     a->a_except_table = PyBytesWriter_FinishWithSize(a->a_except_table_writer,
                                                      a->a_except_table_off);
     a->a_except_table_writer = NULL;

Result with the current PR:

Assembler stats (size/allocated):

• a_except_table: 32,654/40,960 bytes; 11 resize
• a_linetable: 49,177/81,920 bytes; 11 resize
• a_bytecode: 102,420/163,840 bytes; 10 resize

---

I modified the 3 PyBytesWriter_Resize() calls to allocate just one chunk instead of len * 2, patch:

Details

vstinner@vstinner-thinkpadp1gen3$ git diff Python/
diff --git a/Python/assemble.c b/Python/assemble.c
index 2577fa57620..55ef6a22382 100644
--- a/Python/assemble.c
+++ b/Python/assemble.c
@@ -144,7 +144,7 @@ assemble_emit_exception_table_entry(struct assembler *a, int start, int end,
     Py_ssize_t len = PyBytesWriter_GetSize(a->a_except_table_writer);
     if (a->a_except_table_off + MAX_SIZE_OF_ENTRY >= len) {
         a->a_except_table_resize++;
-        RETURN_IF_ERROR(PyBytesWriter_Resize(a->a_except_table_writer, len * 2));
+        RETURN_IF_ERROR(PyBytesWriter_Resize(a->a_except_table_writer, a->a_except_table_off + MAX_SIZE_OF_ENTRY));
     }
     int size = end-start;
     assert(end > start);
@@ -298,7 +298,7 @@ write_location_info_entry(struct assembler* a, location loc, int isize)
     if (a->a_location_off + THEORETICAL_MAX_ENTRY_SIZE >= len) {
         assert(len > THEORETICAL_MAX_ENTRY_SIZE);
         a->a_linetable_resize++;
-        RETURN_IF_ERROR(PyBytesWriter_Resize(a->a_linetable_writer, len * 2));
+        RETURN_IF_ERROR(PyBytesWriter_Resize(a->a_linetable_writer, a->a_location_off + THEORETICAL_MAX_ENTRY_SIZE));
     }
     if (loc.lineno == NO_LOCATION.lineno) {
         write_location_info_none(a, isize);
@@ -432,7 +432,7 @@ assemble_emit_instr(struct assembler *a, instruction *instr)
             return ERROR;
         }
         a->a_bytecode_resize++;
-        RETURN_IF_ERROR(PyBytesWriter_Resize(a->a_bytecode_writer, len * 2));
+        RETURN_IF_ERROR(PyBytesWriter_Resize(a->a_bytecode_writer, (a->a_offset + size) * sizeof(_Py_CODEUNIT)));
     }
     code = (_Py_CODEUNIT *)PyBytesWriter_GetData(a->a_bytecode_writer) + a->a_offset;
     a->a_offset += size;

Assembler stats (size/allocated):

• a_except_table: 32,654/37,182 bytes; 4,096 resize
• a_linetable: 49,177/55,201 bytes; 2,0487 resize
• a_bytecode: 102,420/118,447 bytes; 28,650 resize

There between 372x and 2865x more calls to PyBytesWriter_Resize().

Well, internally PyBytesWriter_Resize() overallocates the buffer, so a call can be cheap. But least, that sounds less efficient to call PyBytesWriter_Resize() so often.

IMO it's ok use len * 2 allocatation strategy to call PyBytesWriter_Resize() less often. At the end, PyBytesWriter_FinishWithSize() truncates the buffer to create a bytes object with the exact size anyway.

---

Obviously, on a trivial code, sizes and smaller and there are only a few PyBytesWriter_Resize() calls.

Example with compile("x+y", "<string>", "exec"):

• a_except_table: 0/256 bytes; 0 resize
• a_linetable: 14/256 bytes; 3 resize
• a_bytecode: 26/256 bytes; 0 resize