gh-126835: Avoid creating unnecessary tuple when looking for constant sequence during constant folding (#131054)

1 files changed, 125 insertions, 78 deletions

diff --git a/Python/flowgraph.c b/Python/flowgraph.c
index 286f8999902..fdafafd7661 100644
--- a/Python/flowgraph.c
+++ b/Python/flowgraph.c
@@ -1344,64 +1344,45 @@ add_const(PyObject *newconst, PyObject *consts, PyObject *const_cache)
 }
 
 /*
-   Walk basic block backwards starting from "start" trying to collect "size" number of
-   subsequent constants from instructions loading constants into new tuple ignoring NOP's in between.
+  Traverse the instructions of the basic block backwards from index "start", skipping over NOPs.
+  Try to collect "size" number of consecutive instructions that load constants into the array "instrs".
+  Caller must make sure that length of "instrs" is sufficient to fit in at least "size" instructions.
 
-   Returns ERROR on error and sets "seq" to NULL.
-   Returns SUCCESS on success and sets "seq" to NULL if failed to collect requested number of constants.
-   Returns SUCCESS on success and sets "seq" to resulting tuple if succeeded to collect requested number of constants.
+  Return boolean indicating whether "size" such instructions were found.
 */
-static int
-get_constant_sequence(basicblock *bb, int start, int size,
-                      PyObject *consts, PyObject **seq)
+static bool
+get_const_loading_instrs(basicblock *bb, int start, cfg_instr **instrs, int size)
 {
     assert(start < bb->b_iused);
-    *seq = NULL;
-    PyObject *res = PyTuple_New((Py_ssize_t)size);
-    if (res == NULL) {
-        return ERROR;
-    }
+    assert(size >= 0);
+    assert(size <= _PY_STACK_USE_GUIDELINE);
+
     for (; start >= 0 && size > 0; start--) {
         cfg_instr *instr = &bb->b_instr[start];
         if (instr->i_opcode == NOP) {
             continue;
         }
         if (!loads_const(instr->i_opcode)) {
-            break;
-        }
-        PyObject *constant = get_const_value(instr->i_opcode, instr->i_oparg, consts);
-        if (constant == NULL) {
-            Py_DECREF(res);
-            return ERROR;
+            return false;
         }
-        PyTuple_SET_ITEM(res, --size, constant);
+        instrs[--size] = instr;
     }
-    if (size > 0) {
-        Py_DECREF(res);
-    }
-    else {
-        *seq = res;
-    }
-    return SUCCESS;
+
+    return size == 0;
 }
 
 /*
-  Walk basic block backwards starting from "start" and change "count" number of
-  non-NOP instructions to NOP's and set their location to NO_LOCATION.
+  Change every instruction in "instrs" NOP and set its location to NO_LOCATION.
+  Caller must make sure "instrs" has at least "size" elements.
 */
 static void
-nop_out(basicblock *bb, int start, int count)
+nop_out(cfg_instr **instrs, int size)
 {
-    assert(start < bb->b_iused);
-    for (; count > 0; start--) {
-        assert(start >= 0);
-        cfg_instr *instr = &bb->b_instr[start];
-        if (instr->i_opcode == NOP) {
-            continue;
-        }
+    for (int i = 0; i < size; i++) {
+        cfg_instr *instr = instrs[i];
+        assert(instr->i_opcode != NOP);
         INSTR_SET_OP0(instr, NOP);
         INSTR_SET_LOC(instr, NO_LOCATION);
-        count--;
     }
 }
 
@@ -1437,19 +1418,39 @@ fold_tuple_of_constants(basicblock *bb, int i, PyObject *consts, PyObject *const
     /* Pre-conditions */
     assert(PyDict_CheckExact(const_cache));
     assert(PyList_CheckExact(consts));
+
     cfg_instr *instr = &bb->b_instr[i];
     assert(instr->i_opcode == BUILD_TUPLE);
+
     int seq_size = instr->i_oparg;
-    PyObject *newconst;
-    RETURN_IF_ERROR(get_constant_sequence(bb, i-1, seq_size, consts, &newconst));
-    if (newconst == NULL) {
+    if (seq_size > _PY_STACK_USE_GUIDELINE) {
+        return SUCCESS;
+    }
+
+    cfg_instr *const_instrs[_PY_STACK_USE_GUIDELINE];
+    if (!get_const_loading_instrs(bb, i-1, const_instrs, seq_size)) {
         /* not a const sequence */
         return SUCCESS;
     }
-    assert(PyTuple_Size(newconst) == seq_size);
-    RETURN_IF_ERROR(instr_make_load_const(instr, newconst, consts, const_cache));
-    nop_out(bb, i-1, seq_size);
-    return SUCCESS;
+
+    PyObject *const_tuple = PyTuple_New((Py_ssize_t)seq_size);
+    if (const_tuple == NULL) {
+        return ERROR;
+    }
+
+    for (int i = 0; i < seq_size; i++) {
+        cfg_instr *inst = const_instrs[i];
+        assert(loads_const(inst->i_opcode));
+        PyObject *element = get_const_value(inst->i_opcode, inst->i_oparg, consts);
+        if (element == NULL) {
+            Py_DECREF(const_tuple);
+            return ERROR;
+        }
+        PyTuple_SET_ITEM(const_tuple, i, element);
+    }
+
+    nop_out(const_instrs, seq_size);
+    return instr_make_load_const(instr, const_tuple, consts, const_cache);
 }
 
 #define MIN_CONST_SEQUENCE_SIZE 3
@@ -1470,34 +1471,56 @@ optimize_lists_and_sets(basicblock *bb, int i, int nextop,
 {
     assert(PyDict_CheckExact(const_cache));
     assert(PyList_CheckExact(consts));
+
     cfg_instr *instr = &bb->b_instr[i];
     assert(instr->i_opcode == BUILD_LIST || instr->i_opcode == BUILD_SET);
+
     bool contains_or_iter = nextop == GET_ITER || nextop == CONTAINS_OP;
     int seq_size = instr->i_oparg;
-    if (seq_size < MIN_CONST_SEQUENCE_SIZE && !contains_or_iter) {
+    if (seq_size > _PY_STACK_USE_GUIDELINE ||
+        (seq_size < MIN_CONST_SEQUENCE_SIZE && !contains_or_iter))
+    {
         return SUCCESS;
     }
-    PyObject *newconst;
-    RETURN_IF_ERROR(get_constant_sequence(bb, i-1, seq_size, consts, &newconst));
-    if (newconst == NULL) {    /* not a const sequence */
+
+    cfg_instr *const_instrs[_PY_STACK_USE_GUIDELINE];
+    if (!get_const_loading_instrs(bb, i-1, const_instrs, seq_size)) {  /* not a const sequence */
         if (contains_or_iter && instr->i_opcode == BUILD_LIST) {
             /* iterate over a tuple instead of list */
             INSTR_SET_OP1(instr, BUILD_TUPLE, instr->i_oparg);
         }
         return SUCCESS;
     }
-    assert(PyTuple_Size(newconst) == seq_size);
+
+    PyObject *const_result = PyTuple_New((Py_ssize_t)seq_size);
+    if (const_result == NULL) {
+        return ERROR;
+    }
+
+    for (int i = 0; i < seq_size; i++) {
+        cfg_instr *inst = const_instrs[i];
+        assert(loads_const(inst->i_opcode));
+        PyObject *element = get_const_value(inst->i_opcode, inst->i_oparg, consts);
+        if (element == NULL) {
+            Py_DECREF(const_result);
+            return ERROR;
+        }
+        PyTuple_SET_ITEM(const_result, i, element);
+    }
+
     if (instr->i_opcode == BUILD_SET) {
-        PyObject *frozenset = PyFrozenSet_New(newconst);
+        PyObject *frozenset = PyFrozenSet_New(const_result);
         if (frozenset == NULL) {
-            Py_DECREF(newconst);
+            Py_DECREF(const_result);
             return ERROR;
         }
-        Py_SETREF(newconst, frozenset);
+        Py_SETREF(const_result, frozenset);
     }
-    int index = add_const(newconst, consts, const_cache);
+
+    int index = add_const(const_result, consts, const_cache);
     RETURN_IF_ERROR(index);
-    nop_out(bb, i-1, seq_size);
+    nop_out(const_instrs, seq_size);
+
     if (contains_or_iter) {
         INSTR_SET_OP1(instr, LOAD_CONST, index);
     }
@@ -1696,19 +1719,34 @@ fold_const_binop(basicblock *bb, int i, PyObject *consts, PyObject *const_cache)
     #define BINOP_OPERAND_COUNT 2
     assert(PyDict_CheckExact(const_cache));
     assert(PyList_CheckExact(consts));
+
     cfg_instr *binop = &bb->b_instr[i];
     assert(binop->i_opcode == BINARY_OP);
-    PyObject *pair;
-    RETURN_IF_ERROR(get_constant_sequence(bb, i-1, BINOP_OPERAND_COUNT, consts, &pair));
-    if (pair == NULL) {
+
+    cfg_instr *operands_instrs[BINOP_OPERAND_COUNT];
+    if (!get_const_loading_instrs(bb, i-1, operands_instrs, BINOP_OPERAND_COUNT)) {
         /* not a const sequence */
         return SUCCESS;
     }
-    assert(PyTuple_Size(pair) == BINOP_OPERAND_COUNT);
-    PyObject *left = PyTuple_GET_ITEM(pair, 0);
-    PyObject *right = PyTuple_GET_ITEM(pair, 1);
-    PyObject *newconst = eval_const_binop(left, binop->i_oparg, right);
-    Py_DECREF(pair);
+
+    cfg_instr *lhs_instr = operands_instrs[0];
+    assert(loads_const(lhs_instr->i_opcode));
+    PyObject *lhs = get_const_value(lhs_instr->i_opcode, lhs_instr->i_oparg, consts);
+    if (lhs == NULL) {
+        return ERROR;
+    }
+
+    cfg_instr *rhs_instr = operands_instrs[1];
+    assert(loads_const(rhs_instr->i_opcode));
+    PyObject *rhs = get_const_value(rhs_instr->i_opcode, rhs_instr->i_oparg, consts);
+    if (rhs == NULL) {
+        Py_DECREF(lhs);
+        return ERROR;
+    }
+
+    PyObject *newconst = eval_const_binop(lhs, binop->i_oparg, rhs);
+    Py_DECREF(lhs);
+    Py_DECREF(rhs);
     if (newconst == NULL) {
         if (PyErr_ExceptionMatches(PyExc_KeyboardInterrupt)) {
             return ERROR;
@@ -1716,9 +1754,9 @@ fold_const_binop(basicblock *bb, int i, PyObject *consts, PyObject *const_cache)
         PyErr_Clear();
         return SUCCESS;
     }
-    RETURN_IF_ERROR(instr_make_load_const(binop, newconst, consts, const_cache));
-    nop_out(bb, i-1, BINOP_OPERAND_COUNT);
-    return SUCCESS;
+
+    nop_out(operands_instrs, BINOP_OPERAND_COUNT);
+    return instr_make_load_const(binop, newconst, consts, const_cache);
 }
 
 static PyObject *
@@ -1765,17 +1803,26 @@ fold_const_unaryop(basicblock *bb, int i, PyObject *consts, PyObject *const_cach
     #define UNARYOP_OPERAND_COUNT 1
     assert(PyDict_CheckExact(const_cache));
     assert(PyList_CheckExact(consts));
-    cfg_instr *instr = &bb->b_instr[i];
-    PyObject *seq;
-    RETURN_IF_ERROR(get_constant_sequence(bb, i-1, UNARYOP_OPERAND_COUNT, consts, &seq));
-    if (seq == NULL) {
+    cfg_instr *unaryop = &bb->b_instr[i];
+
+    cfg_instr *operand_instr;
+    if (!get_const_loading_instrs(bb, i-1, &operand_instr, UNARYOP_OPERAND_COUNT)) {
         /* not a const */
         return SUCCESS;
     }
-    assert(PyTuple_Size(seq) == UNARYOP_OPERAND_COUNT);
-    PyObject *operand = PyTuple_GET_ITEM(seq, 0);
-    PyObject *newconst = eval_const_unaryop(operand, instr->i_opcode, instr->i_oparg);
-    Py_DECREF(seq);
+
+    assert(loads_const(operand_instr->i_opcode));
+    PyObject *operand = get_const_value(
+        operand_instr->i_opcode,
+        operand_instr->i_oparg,
+        consts
+    );
+    if (operand == NULL) {
+        return ERROR;
+    }
+
+    PyObject *newconst = eval_const_unaryop(operand, unaryop->i_opcode, unaryop->i_oparg);
+    Py_DECREF(operand);
     if (newconst == NULL) {
         if (PyErr_ExceptionMatches(PyExc_KeyboardInterrupt)) {
             return ERROR;
@@ -1783,12 +1830,12 @@ fold_const_unaryop(basicblock *bb, int i, PyObject *consts, PyObject *const_cach
         PyErr_Clear();
         return SUCCESS;
     }
-    if (instr->i_opcode == UNARY_NOT) {
+
+    if (unaryop->i_opcode == UNARY_NOT) {
         assert(PyBool_Check(newconst));
     }
-    RETURN_IF_ERROR(instr_make_load_const(instr, newconst, consts, const_cache));
-    nop_out(bb, i-1, UNARYOP_OPERAND_COUNT);
-    return SUCCESS;
+    nop_out(&operand_instr, UNARYOP_OPERAND_COUNT);
+    return instr_make_load_const(unaryop, newconst, consts, const_cache);
 }
 
 #define VISITED (-1)