move binop folding to cfg

python · iritkatriel · Feb 21, 2025 · Feb 14, 2025 · Feb 14, 2025 · Feb 14, 2025
commit 04b4a847f2c0b98f9fe477491e5695ff07ed3bb9
diff --git a/Python/ast_opt.c b/Python/ast_opt.c
@@ -136,129 +136,6 @@ fold_unaryop(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
     return make_const(node, newval, arena);
 }
 
-/* Check whether a collection doesn't containing too much items (including
-   subcollections).  This protects from creating a constant that needs
-   too much time for calculating a hash.
-   "limit" is the maximal number of items.
-   Returns the negative number if the total number of items exceeds the
-   limit.  Otherwise returns the limit minus the total number of items.
-*/
-
-static Py_ssize_t
-check_complexity(PyObject *obj, Py_ssize_t limit)
-{
-    if (PyTuple_Check(obj)) {
-        Py_ssize_t i;
-        limit -= PyTuple_GET_SIZE(obj);
-        for (i = 0; limit >= 0 && i < PyTuple_GET_SIZE(obj); i++) {
-            limit = check_complexity(PyTuple_GET_ITEM(obj, i), limit);
-        }
-        return limit;
-    }
-    return limit;
-}
-
-#define MAX_INT_SIZE           128  /* bits */
-#define MAX_COLLECTION_SIZE    256  /* items */
-#define MAX_STR_SIZE          4096  /* characters */
-#define MAX_TOTAL_ITEMS       1024  /* including nested collections */
-
-static PyObject *
-safe_multiply(PyObject *v, PyObject *w)
-{
-    if (PyLong_Check(v) && PyLong_Check(w) &&
-        !_PyLong_IsZero((PyLongObject *)v) && !_PyLong_IsZero((PyLongObject *)w)
-    ) {
-        int64_t vbits = _PyLong_NumBits(v);
-        int64_t wbits = _PyLong_NumBits(w);
-        assert(vbits >= 0);
-        assert(wbits >= 0);
-        if (vbits + wbits > MAX_INT_SIZE) {
-            return NULL;
-        }
-    }
-    else if (PyLong_Check(v) && PyTuple_Check(w)) {
-        Py_ssize_t size = PyTuple_GET_SIZE(w);
-        if (size) {
-            long n = PyLong_AsLong(v);
-            if (n < 0 || n > MAX_COLLECTION_SIZE / size) {
-                return NULL;
-            }
-            if (n && check_complexity(w, MAX_TOTAL_ITEMS / n) < 0) {
-                return NULL;
-            }
-        }
-    }
-    else if (PyLong_Check(v) && (PyUnicode_Check(w) || PyBytes_Check(w))) {
-        Py_ssize_t size = PyUnicode_Check(w) ? PyUnicode_GET_LENGTH(w) :
-                                               PyBytes_GET_SIZE(w);
-        if (size) {
-            long n = PyLong_AsLong(v);
-            if (n < 0 || n > MAX_STR_SIZE / size) {
-                return NULL;
-            }
-        }
-    }
-    else if (PyLong_Check(w) &&
-             (PyTuple_Check(v) || PyUnicode_Check(v) || PyBytes_Check(v)))
-    {
-        return safe_multiply(w, v);
-    }
-
-    return PyNumber_Multiply(v, w);
-}
-
-static PyObject *
-safe_power(PyObject *v, PyObject *w)
-{
-    if (PyLong_Check(v) && PyLong_Check(w) &&
-        !_PyLong_IsZero((PyLongObject *)v) && _PyLong_IsPositive((PyLongObject *)w)
-    ) {
-        int64_t vbits = _PyLong_NumBits(v);
-        size_t wbits = PyLong_AsSize_t(w);
-        assert(vbits >= 0);
-        if (wbits == (size_t)-1) {
-            return NULL;
-        }
-        if ((uint64_t)vbits > MAX_INT_SIZE / wbits) {
-            return NULL;
-        }
-    }
-
-    return PyNumber_Power(v, w, Py_None);
-}
-
-static PyObject *
-safe_lshift(PyObject *v, PyObject *w)
-{
-    if (PyLong_Check(v) && PyLong_Check(w) &&
-        !_PyLong_IsZero((PyLongObject *)v) && !_PyLong_IsZero((PyLongObject *)w)
-    ) {
-        int64_t vbits = _PyLong_NumBits(v);
-        size_t wbits = PyLong_AsSize_t(w);
-        assert(vbits >= 0);
-        if (wbits == (size_t)-1) {
-            return NULL;
-        }
-        if (wbits > MAX_INT_SIZE || (uint64_t)vbits > MAX_INT_SIZE - wbits) {
-            return NULL;
-        }
-    }
-
-    return PyNumber_Lshift(v, w);
-}
-
-static PyObject *
-safe_mod(PyObject *v, PyObject *w)
-{
-    if (PyUnicode_Check(v) || PyBytes_Check(v)) {
-        return NULL;
-    }
-
-    return PyNumber_Remainder(v, w);
-}
-
-
 static expr_ty
 parse_literal(PyObject *fmt, Py_ssize_t *ppos, PyArena *arena)
 {
@@ -478,58 +355,7 @@ fold_binop(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
         return optimize_format(node, lv, rhs->v.Tuple.elts, arena);
     }
 
-    if (rhs->kind != Constant_kind) {
-        return 1;
-    }
-
-    PyObject *rv = rhs->v.Constant.value;
-    PyObject *newval = NULL;
-
-    switch (node->v.BinOp.op) {
-    case Add:
-        newval = PyNumber_Add(lv, rv);
-        break;
-    case Sub:
-        newval = PyNumber_Subtract(lv, rv);
-        break;
-    case Mult:
-        newval = safe_multiply(lv, rv);
-        break;
-    case Div:
-        newval = PyNumber_TrueDivide(lv, rv);
-        break;
-    case FloorDiv:
-        newval = PyNumber_FloorDivide(lv, rv);
-        break;
-    case Mod:
-        newval = safe_mod(lv, rv);
-        break;
-    case Pow:
-        newval = safe_power(lv, rv);
-        break;
-    case LShift:
-        newval = safe_lshift(lv, rv);
-        break;
-    case RShift:
-        newval = PyNumber_Rshift(lv, rv);
-        break;
-    case BitOr:
-        newval = PyNumber_Or(lv, rv);
-        break;
-    case BitXor:
-        newval = PyNumber_Xor(lv, rv);
-        break;
-    case BitAnd:
-        newval = PyNumber_And(lv, rv);
-        break;
-    // No builtin constants implement the following operators
-    case MatMult:
-        return 1;
-    // No default case, so the compiler will emit a warning if new binary
-    // operators are added without being handled here
-    }
-
-    return make_const(node, newval, arena);
+    return 1;
 }
 
 static PyObject*
@@ -971,6 +797,115 @@ astfold_withitem(withitem_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
     return 1;
 }
 
+#define IS_CONST_EXPR(N) \
+    ((N)->kind == Constant_kind)
+
+#define CONST_EXPR_VALUE(N) \
+    ((N)->v.Constant.value)
+
+#define IS_COMPLEX_CONST_EXPR(N) \
+    (IS_CONST_EXPR(N) && PyComplex_CheckExact(CONST_EXPR_VALUE(N)))
+
+#define IS_NUMERIC_CONST_EXPR(N) \
+    (IS_CONST_EXPR(N) && (PyLong_CheckExact(CONST_EXPR_VALUE(N)) || PyFloat_CheckExact(CONST_EXPR_VALUE(N))))
+
+#define IS_UNARY_EXPR(N) \
+    ((N)->kind == UnaryOp_kind)
+
+#define UNARY_EXPR_OP(N) \
+    ((N)->v.UnaryOp.op)
+
+#define UNARY_EXPR_OPERAND(N) \
+    ((N)->v.UnaryOp.operand)
+
+#define UNARY_EXPR_OPERAND_CONST_VALUE(N) \
+    (CONST_EXPR_VALUE(UNARY_EXPR_OPERAND(N)))
+
+#define IS_UNARY_SUB_EXPR(N) \
+    (IS_UNARY_EXPR(N) && UNARY_EXPR_OP(N) == USub)
+
+#define IS_NUMERIC_UNARY_CONST_EXPR(N) \
+    (IS_UNARY_SUB_EXPR(N) && IS_NUMERIC_CONST_EXPR(UNARY_EXPR_OPERAND(N)))
+
+#define IS_COMPLEX_UNARY_CONST_EXPR(N) \
+    (IS_UNARY_SUB_EXPR(N) && IS_COMPLEX_CONST_EXPR(UNARY_EXPR_OPERAND(N)))
+
+#define BINARY_EXPR(N) \
+    ((N)->v.BinOp)
+
+#define BINARY_EXPR_OP(N) \
+    (BINARY_EXPR(N).op)
+
+#define BINARY_EXPR_LEFT(N) \
+    (BINARY_EXPR(N).left)
+
+#define BINARY_EXPR_RIGHT(N) \
+    (BINARY_EXPR(N).right)
+
+#define IS_BINARY_EXPR(N) \
+    ((N)->kind == BinOp_kind)
+
+#define IS_BINARY_ADD_EXPR(N) \
+    (IS_BINARY_EXPR(N) && BINARY_EXPR_OP(N) == Add)
+
+#define IS_BINARY_SUB_EXPR(N) \
+    (IS_BINARY_EXPR(N) && BINARY_EXPR_OP(N) == Sub)
+
+#define IS_MATCH_NUMERIC_OR_COMPLEX_UNARY_CONST_EXPR(N) \
+    (IS_NUMERIC_UNARY_CONST_EXPR(N) || IS_COMPLEX_UNARY_CONST_EXPR(N))
+
+#define IS_MATCH_COMPLEX_BINARY_CONST_EXPR(N) \
+    ( \
+        (IS_BINARY_ADD_EXPR(N) || IS_BINARY_SUB_EXPR(N)) \
+        && (IS_NUMERIC_UNARY_CONST_EXPR(BINARY_EXPR_LEFT(N)) || IS_CONST_EXPR(BINARY_EXPR_LEFT(N))) \
+        && IS_COMPLEX_CONST_EXPR(BINARY_EXPR_RIGHT(N)) \
+    )
+
+
+static int
+fold_const_unary_or_complex_expr(expr_ty e, PyArena *arena)
+{
+    assert(IS_MATCH_NUMERIC_OR_COMPLEX_UNARY_CONST_EXPR(e));
+    PyObject *constant = UNARY_EXPR_OPERAND_CONST_VALUE(e);
+    assert(UNARY_EXPR_OP(e) == USub);
+    PyObject* folded = PyNumber_Negative(constant);
+    return make_const(e, folded, arena);
+}
+
+static int
+fold_const_binary_complex_expr(expr_ty e, PyArena *arena)
+{
+    assert(IS_MATCH_COMPLEX_BINARY_CONST_EXPR(e));
+    expr_ty left_expr = BINARY_EXPR_LEFT(e);
+    if (IS_NUMERIC_UNARY_CONST_EXPR(left_expr)) {
+        if (!fold_const_unary_or_complex_expr(left_expr, arena)) {
+            return 0;
+        }
+    }
+    assert(IS_CONST_EXPR(BINARY_EXPR_LEFT(e)));
+    operator_ty op = BINARY_EXPR_OP(e);
+    PyObject *left = CONST_EXPR_VALUE(BINARY_EXPR_LEFT(e));
+    PyObject *right = CONST_EXPR_VALUE(BINARY_EXPR_RIGHT(e));
+    assert(op == Add || op == Sub);
+    PyObject *folded = op == Add ? PyNumber_Add(left, right) : PyNumber_Subtract(left, right);
+    return make_const(e, folded, arena);
+}
+
+static int
+fold_pattern_match_value(expr_ty node, PyArena *arena, _PyASTOptimizeState *Py_UNUSED(state))
+{
+    switch (node->kind)
+    {
+        case UnaryOp_kind:
+            return fold_const_unary_or_complex_expr(node, arena);
+        case BinOp_kind:
+            return fold_const_binary_complex_expr(node, arena);
+        default:
+            break;
+    }
+    return 1;
+}
+
 static int
 astfold_pattern(pattern_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
 {
@@ -980,15 +915,15 @@ astfold_pattern(pattern_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
     ENTER_RECURSIVE(state);
     switch (node_->kind) {
         case MatchValue_kind:
-            CALL(astfold_expr, expr_ty, node_->v.MatchValue.value);
+            CALL(fold_pattern_match_value, expr_ty, node_->v.MatchValue.value);
             break;
         case MatchSingleton_kind:
             break;
         case MatchSequence_kind:
             CALL_SEQ(astfold_pattern, pattern, node_->v.MatchSequence.patterns);
             break;
         case MatchMapping_kind:
-            CALL_SEQ(astfold_expr, expr, node_->v.MatchMapping.keys);
+            CALL_SEQ(fold_pattern_match_value, expr, node_->v.MatchMapping.keys);
             CALL_SEQ(astfold_pattern, pattern, node_->v.MatchMapping.patterns);
             break;
         case MatchClass_kind: