Mirror
/
cpython
mirror of https://github.com/python/cpython
4
1
Fork 0
Code Issues Releases Wiki Activity
cpython/Python/ast_opt.c

892 lines
28 KiB
C
Raw Blame History

/* AST Optimizer */
#include "Python.h"
#include "pycore_ast.h" // _PyAST_GetDocString()
#include "pycore_compile.h" // _PyASTOptimizeState
static int
make_const(expr_ty node, PyObject *val, PyArena *arena)
{
// Even if no new value was calculated, make_const may still
// need to clear an error (e.g. for division by zero)
if (val == NULL) {
if (PyErr_ExceptionMatches(PyExc_KeyboardInterrupt)) {
return 0;
}
PyErr_Clear();
return 1;
}
if (_PyArena_AddPyObject(arena, val) < 0) {
Py_DECREF(val);
return 0;
}
node->kind = Constant_kind;
node->v.Constant.kind = NULL;
node->v.Constant.value = val;
return 1;
}
#define COPY_NODE(TO, FROM) (memcpy((TO), (FROM), sizeof(struct _expr)))
static PyObject*
unary_not(PyObject *v)
{
int r = PyObject_IsTrue(v);
if (r < 0)
return NULL;
return PyBool_FromLong(!r);
}
static int
fold_unaryop(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
{
expr_ty arg = node->v.UnaryOp.operand;
if (arg->kind != Constant_kind) {
/* Fold not into comparison */
if (node->v.UnaryOp.op == Not && arg->kind == Compare_kind &&
asdl_seq_LEN(arg->v.Compare.ops) == 1) {
/* Eq and NotEq are often implemented in terms of one another, so
folding not (self == other) into self != other breaks implementation
of !=. Detecting such cases doesn't seem worthwhile.
Python uses </> for 'is subset'/'is superset' operations on sets.
They don't satisfy not folding laws. */
cmpop_ty op = asdl_seq_GET(arg->v.Compare.ops, 0);
switch (op) {
case Is:
op = IsNot;
break;
case IsNot:
op = Is;
break;
case In:
op = NotIn;
break;
case NotIn:
op = In;
break;
// The remaining comparison operators can't be safely inverted
case Eq:
case NotEq:
case Lt:
case LtE:
case Gt:
case GtE:
op = 0; // The AST enums leave "0" free as an "unused" marker
break;
// No default case, so the compiler will emit a warning if new
// comparison operators are added without being handled here
}
if (op) {
asdl_seq_SET(arg->v.Compare.ops, 0, op);
COPY_NODE(node, arg);
return 1;
}
}
return 1;
}
typedef PyObject *(*unary_op)(PyObject*);
static const unary_op ops[] = {
[Invert] = PyNumber_Invert,
[Not] = unary_not,
[UAdd] = PyNumber_Positive,
[USub] = PyNumber_Negative,
};
PyObject *newval = ops[node->v.UnaryOp.op](arg->v.Constant.value);
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;
}
else if (PyFrozenSet_Check(obj)) {
Py_ssize_t i = 0;
PyObject *item;
Py_hash_t hash;
limit -= PySet_GET_SIZE(obj);
while (limit >= 0 && _PySet_NextEntry(obj, &i, &item, &hash)) {
limit = check_complexity(item, 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) && Py_SIZE(v) && Py_SIZE(w)) {
size_t vbits = _PyLong_NumBits(v);
size_t wbits = _PyLong_NumBits(w);
if (vbits == (size_t)-1 || wbits == (size_t)-1) {
return NULL;
}
if (vbits + wbits > MAX_INT_SIZE) {
return NULL;
}
}
else if (PyLong_Check(v) && (PyTuple_Check(w) || PyFrozenSet_Check(w))) {
Py_ssize_t size = PyTuple_Check(w) ? PyTuple_GET_SIZE(w) :
PySet_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) || PyFrozenSet_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) && Py_SIZE(v) && Py_SIZE(w) > 0) {
size_t vbits = _PyLong_NumBits(v);
size_t wbits = PyLong_AsSize_t(w);
if (vbits == (size_t)-1 || wbits == (size_t)-1) {
return NULL;
}
if (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) && Py_SIZE(v) && Py_SIZE(w)) {
size_t vbits = _PyLong_NumBits(v);
size_t wbits = PyLong_AsSize_t(w);
if (vbits == (size_t)-1 || wbits == (size_t)-1) {
return NULL;
}
if (wbits > MAX_INT_SIZE || 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 int
fold_binop(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
{
expr_ty lhs, rhs;
lhs = node->v.BinOp.left;
rhs = node->v.BinOp.right;
if (lhs->kind != Constant_kind || rhs->kind != Constant_kind) {
return 1;
}
PyObject *lv = lhs->v.Constant.value;
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);
}
static PyObject*
make_const_tuple(asdl_expr_seq *elts)
{
for (int i = 0; i < asdl_seq_LEN(elts); i++) {
expr_ty e = (expr_ty)asdl_seq_GET(elts, i);
if (e->kind != Constant_kind) {
return NULL;
}
}
PyObject *newval = PyTuple_New(asdl_seq_LEN(elts));
if (newval == NULL) {
return NULL;
}
for (int i = 0; i < asdl_seq_LEN(elts); i++) {
expr_ty e = (expr_ty)asdl_seq_GET(elts, i);
PyObject *v = e->v.Constant.value;
Py_INCREF(v);
PyTuple_SET_ITEM(newval, i, v);
}
return newval;
}
static int
fold_tuple(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
{
PyObject *newval;
if (node->v.Tuple.ctx != Load)
return 1;
newval = make_const_tuple(node->v.Tuple.elts);
return make_const(node, newval, arena);
}
static int
fold_subscr(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
{
PyObject *newval;
expr_ty arg, idx;
arg = node->v.Subscript.value;
idx = node->v.Subscript.slice;
if (node->v.Subscript.ctx != Load ||
arg->kind != Constant_kind ||
idx->kind != Constant_kind)
{
return 1;
}
newval = PyObject_GetItem(arg->v.Constant.value, idx->v.Constant.value);
return make_const(node, newval, arena);
}
/* Change literal list or set of constants into constant
tuple or frozenset respectively. Change literal list of
non-constants into tuple.
Used for right operand of "in" and "not in" tests and for iterable
in "for" loop and comprehensions.
*/
static int
fold_iter(expr_ty arg, PyArena *arena, _PyASTOptimizeState *state)
{
PyObject *newval;
if (arg->kind == List_kind) {
/* First change a list into tuple. */
asdl_expr_seq *elts = arg->v.List.elts;
Py_ssize_t n = asdl_seq_LEN(elts);
for (Py_ssize_t i = 0; i < n; i++) {
expr_ty e = (expr_ty)asdl_seq_GET(elts, i);
if (e->kind == Starred_kind) {
return 1;
}
}
expr_context_ty ctx = arg->v.List.ctx;
arg->kind = Tuple_kind;
arg->v.Tuple.elts = elts;
arg->v.Tuple.ctx = ctx;
/* Try to create a constant tuple. */
newval = make_const_tuple(elts);
}
else if (arg->kind == Set_kind) {
newval = make_const_tuple(arg->v.Set.elts);
if (newval) {
Py_SETREF(newval, PyFrozenSet_New(newval));
}
}
else {
return 1;
}
return make_const(arg, newval, arena);
}
static int
fold_compare(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
{
asdl_int_seq *ops;
asdl_expr_seq *args;
Py_ssize_t i;
ops = node->v.Compare.ops;
args = node->v.Compare.comparators;
/* TODO: optimize cases with literal arguments. */
/* Change literal list or set in 'in' or 'not in' into
tuple or frozenset respectively. */
i = asdl_seq_LEN(ops) - 1;
int op = asdl_seq_GET(ops, i);
if (op == In || op == NotIn) {
if (!fold_iter((expr_ty)asdl_seq_GET(args, i), arena, state)) {
return 0;
}
}
return 1;
}
static int astfold_mod(mod_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
static int astfold_stmt(stmt_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
static int astfold_expr(expr_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
static int astfold_arguments(arguments_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
static int astfold_comprehension(comprehension_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
static int astfold_keyword(keyword_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
static int astfold_withitem(withitem_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
static int astfold_excepthandler(excepthandler_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
static int astfold_match_case(match_case_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
static int astfold_pattern(expr_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
#define CALL(FUNC, TYPE, ARG) \
if (!FUNC((ARG), ctx_, state)) \
return 0;
#define CALL_OPT(FUNC, TYPE, ARG) \
if ((ARG) != NULL && !FUNC((ARG), ctx_, state)) \
return 0;
#define CALL_SEQ(FUNC, TYPE, ARG) { \
int i; \
asdl_ ## TYPE ## _seq *seq = (ARG); /* avoid variable capture */ \
for (i = 0; i < asdl_seq_LEN(seq); i++) { \
TYPE ## _ty elt = (TYPE ## _ty)asdl_seq_GET(seq, i); \
if (elt != NULL && !FUNC(elt, ctx_, state)) \
return 0; \
} \
}
#define CALL_INT_SEQ(FUNC, TYPE, ARG) { \
int i; \
asdl_int_seq *seq = (ARG); /* avoid variable capture */ \
for (i = 0; i < asdl_seq_LEN(seq); i++) { \
TYPE elt = (TYPE)asdl_seq_GET(seq, i); \
if (!FUNC(elt, ctx_, state)) \
return 0; \
} \
}
static int
astfold_body(asdl_stmt_seq *stmts, PyArena *ctx_, _PyASTOptimizeState *state)
{
int docstring = _PyAST_GetDocString(stmts) != NULL;
CALL_SEQ(astfold_stmt, stmt, stmts);
if (!docstring && _PyAST_GetDocString(stmts) != NULL) {
stmt_ty st = (stmt_ty)asdl_seq_GET(stmts, 0);
asdl_expr_seq *values = _Py_asdl_expr_seq_new(1, ctx_);
if (!values) {
return 0;
}
asdl_seq_SET(values, 0, st->v.Expr.value);
expr_ty expr = _PyAST_JoinedStr(values, st->lineno, st->col_offset,
st->end_lineno, st->end_col_offset,
ctx_);
if (!expr) {
return 0;
}
st->v.Expr.value = expr;
}
return 1;
}
static int
astfold_mod(mod_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
switch (node_->kind) {
case Module_kind:
CALL(astfold_body, asdl_seq, node_->v.Module.body);
break;
case Interactive_kind:
CALL_SEQ(astfold_stmt, stmt, node_->v.Interactive.body);
break;
case Expression_kind:
CALL(astfold_expr, expr_ty, node_->v.Expression.body);
break;
// The following top level nodes don't participate in constant folding
case FunctionType_kind:
break;
// No default case, so the compiler will emit a warning if new top level
// compilation nodes are added without being handled here
}
return 1;
}
static int
astfold_expr(expr_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
switch (node_->kind) {
case BoolOp_kind:
CALL_SEQ(astfold_expr, expr, node_->v.BoolOp.values);
break;
case BinOp_kind:
CALL(astfold_expr, expr_ty, node_->v.BinOp.left);
CALL(astfold_expr, expr_ty, node_->v.BinOp.right);
CALL(fold_binop, expr_ty, node_);
break;
case UnaryOp_kind:
CALL(astfold_expr, expr_ty, node_->v.UnaryOp.operand);
CALL(fold_unaryop, expr_ty, node_);
break;
case Lambda_kind:
CALL(astfold_arguments, arguments_ty, node_->v.Lambda.args);
CALL(astfold_expr, expr_ty, node_->v.Lambda.body);
break;
case IfExp_kind:
CALL(astfold_expr, expr_ty, node_->v.IfExp.test);
CALL(astfold_expr, expr_ty, node_->v.IfExp.body);
CALL(astfold_expr, expr_ty, node_->v.IfExp.orelse);
break;
case Dict_kind:
CALL_SEQ(astfold_expr, expr, node_->v.Dict.keys);
CALL_SEQ(astfold_expr, expr, node_->v.Dict.values);
break;
case Set_kind:
CALL_SEQ(astfold_expr, expr, node_->v.Set.elts);
break;
case ListComp_kind:
CALL(astfold_expr, expr_ty, node_->v.ListComp.elt);
CALL_SEQ(astfold_comprehension, comprehension, node_->v.ListComp.generators);
break;
case SetComp_kind:
CALL(astfold_expr, expr_ty, node_->v.SetComp.elt);
CALL_SEQ(astfold_comprehension, comprehension, node_->v.SetComp.generators);
break;
case DictComp_kind:
CALL(astfold_expr, expr_ty, node_->v.DictComp.key);
CALL(astfold_expr, expr_ty, node_->v.DictComp.value);
CALL_SEQ(astfold_comprehension, comprehension, node_->v.DictComp.generators);
break;
case GeneratorExp_kind:
CALL(astfold_expr, expr_ty, node_->v.GeneratorExp.elt);
CALL_SEQ(astfold_comprehension, comprehension, node_->v.GeneratorExp.generators);
break;
case Await_kind:
CALL(astfold_expr, expr_ty, node_->v.Await.value);
break;
case Yield_kind:
CALL_OPT(astfold_expr, expr_ty, node_->v.Yield.value);
break;
case YieldFrom_kind:
CALL(astfold_expr, expr_ty, node_->v.YieldFrom.value);
break;
case Compare_kind:
CALL(astfold_expr, expr_ty, node_->v.Compare.left);
CALL_SEQ(astfold_expr, expr, node_->v.Compare.comparators);
CALL(fold_compare, expr_ty, node_);
break;
case Call_kind:
CALL(astfold_expr, expr_ty, node_->v.Call.func);
CALL_SEQ(astfold_expr, expr, node_->v.Call.args);
CALL_SEQ(astfold_keyword, keyword, node_->v.Call.keywords);
break;
case FormattedValue_kind:
CALL(astfold_expr, expr_ty, node_->v.FormattedValue.value);
CALL_OPT(astfold_expr, expr_ty, node_->v.FormattedValue.format_spec);
break;
case JoinedStr_kind:
CALL_SEQ(astfold_expr, expr, node_->v.JoinedStr.values);
break;
case Attribute_kind:
CALL(astfold_expr, expr_ty, node_->v.Attribute.value);
break;
case Subscript_kind:
CALL(astfold_expr, expr_ty, node_->v.Subscript.value);
CALL(astfold_expr, expr_ty, node_->v.Subscript.slice);
CALL(fold_subscr, expr_ty, node_);
break;
case Starred_kind:
CALL(astfold_expr, expr_ty, node_->v.Starred.value);
break;
case Slice_kind:
CALL_OPT(astfold_expr, expr_ty, node_->v.Slice.lower);
CALL_OPT(astfold_expr, expr_ty, node_->v.Slice.upper);
CALL_OPT(astfold_expr, expr_ty, node_->v.Slice.step);
break;
case List_kind:
CALL_SEQ(astfold_expr, expr, node_->v.List.elts);
break;
case Tuple_kind:
CALL_SEQ(astfold_expr, expr, node_->v.Tuple.elts);
CALL(fold_tuple, expr_ty, node_);
break;
case Name_kind:
if (node_->v.Name.ctx == Load &&
_PyUnicode_EqualToASCIIString(node_->v.Name.id, "__debug__")) {
return make_const(node_, PyBool_FromLong(!state->optimize), ctx_);
}
break;
case NamedExpr_kind:
CALL(astfold_expr, expr_ty, node_->v.NamedExpr.value);
break;
case Constant_kind:
// Already a constant, nothing further to do
break;
case MatchAs_kind:
case MatchOr_kind:
// These can't occur outside of patterns.
Py_UNREACHABLE();
// No default case, so the compiler will emit a warning if new expression
// kinds are added without being handled here
}
return 1;
}
static int
astfold_keyword(keyword_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
CALL(astfold_expr, expr_ty, node_->value);
return 1;
}
static int
astfold_comprehension(comprehension_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
CALL(astfold_expr, expr_ty, node_->target);
CALL(astfold_expr, expr_ty, node_->iter);
CALL_SEQ(astfold_expr, expr, node_->ifs);
CALL(fold_iter, expr_ty, node_->iter);
return 1;
}
static int
astfold_arguments(arguments_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
CALL_SEQ(astfold_expr, expr, node_->kw_defaults);
CALL_SEQ(astfold_expr, expr, node_->defaults);
return 1;
}
static int
astfold_stmt(stmt_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
switch (node_->kind) {
case FunctionDef_kind:
CALL(astfold_arguments, arguments_ty, node_->v.FunctionDef.args);
CALL(astfold_body, asdl_seq, node_->v.FunctionDef.body);
CALL_SEQ(astfold_expr, expr, node_->v.FunctionDef.decorator_list);
break;
case AsyncFunctionDef_kind:
CALL(astfold_arguments, arguments_ty, node_->v.AsyncFunctionDef.args);
CALL(astfold_body, asdl_seq, node_->v.AsyncFunctionDef.body);
CALL_SEQ(astfold_expr, expr, node_->v.AsyncFunctionDef.decorator_list);
break;
case ClassDef_kind:
CALL_SEQ(astfold_expr, expr, node_->v.ClassDef.bases);
CALL_SEQ(astfold_keyword, keyword, node_->v.ClassDef.keywords);
CALL(astfold_body, asdl_seq, node_->v.ClassDef.body);
CALL_SEQ(astfold_expr, expr, node_->v.ClassDef.decorator_list);
break;
case Return_kind:
CALL_OPT(astfold_expr, expr_ty, node_->v.Return.value);
break;
case Delete_kind:
CALL_SEQ(astfold_expr, expr, node_->v.Delete.targets);
break;
case Assign_kind:
CALL_SEQ(astfold_expr, expr, node_->v.Assign.targets);
CALL(astfold_expr, expr_ty, node_->v.Assign.value);
break;
case AugAssign_kind:
CALL(astfold_expr, expr_ty, node_->v.AugAssign.target);
CALL(astfold_expr, expr_ty, node_->v.AugAssign.value);
break;
case AnnAssign_kind:
CALL(astfold_expr, expr_ty, node_->v.AnnAssign.target);
CALL_OPT(astfold_expr, expr_ty, node_->v.AnnAssign.value);
break;
case For_kind:
CALL(astfold_expr, expr_ty, node_->v.For.target);
CALL(astfold_expr, expr_ty, node_->v.For.iter);
CALL_SEQ(astfold_stmt, stmt, node_->v.For.body);
CALL_SEQ(astfold_stmt, stmt, node_->v.For.orelse);
CALL(fold_iter, expr_ty, node_->v.For.iter);
break;
case AsyncFor_kind:
CALL(astfold_expr, expr_ty, node_->v.AsyncFor.target);
CALL(astfold_expr, expr_ty, node_->v.AsyncFor.iter);
CALL_SEQ(astfold_stmt, stmt, node_->v.AsyncFor.body);
CALL_SEQ(astfold_stmt, stmt, node_->v.AsyncFor.orelse);
break;
case While_kind:
CALL(astfold_expr, expr_ty, node_->v.While.test);
CALL_SEQ(astfold_stmt, stmt, node_->v.While.body);
CALL_SEQ(astfold_stmt, stmt, node_->v.While.orelse);
break;
case If_kind:
CALL(astfold_expr, expr_ty, node_->v.If.test);
CALL_SEQ(astfold_stmt, stmt, node_->v.If.body);
CALL_SEQ(astfold_stmt, stmt, node_->v.If.orelse);
break;
case With_kind:
CALL_SEQ(astfold_withitem, withitem, node_->v.With.items);
CALL_SEQ(astfold_stmt, stmt, node_->v.With.body);
break;
case AsyncWith_kind:
CALL_SEQ(astfold_withitem, withitem, node_->v.AsyncWith.items);
CALL_SEQ(astfold_stmt, stmt, node_->v.AsyncWith.body);
break;
case Raise_kind:
CALL_OPT(astfold_expr, expr_ty, node_->v.Raise.exc);
CALL_OPT(astfold_expr, expr_ty, node_->v.Raise.cause);
break;
case Try_kind:
CALL_SEQ(astfold_stmt, stmt, node_->v.Try.body);
CALL_SEQ(astfold_excepthandler, excepthandler, node_->v.Try.handlers);
CALL_SEQ(astfold_stmt, stmt, node_->v.Try.orelse);
CALL_SEQ(astfold_stmt, stmt, node_->v.Try.finalbody);
break;
case Assert_kind:
CALL(astfold_expr, expr_ty, node_->v.Assert.test);
CALL_OPT(astfold_expr, expr_ty, node_->v.Assert.msg);
break;
case Expr_kind:
CALL(astfold_expr, expr_ty, node_->v.Expr.value);
break;
case Match_kind:
CALL(astfold_expr, expr_ty, node_->v.Match.subject);
CALL_SEQ(astfold_match_case, match_case, node_->v.Match.cases);
break;
// The following statements don't contain any subexpressions to be folded
case Import_kind:
case ImportFrom_kind:
case Global_kind:
case Nonlocal_kind:
case Pass_kind:
case Break_kind:
case Continue_kind:
break;
// No default case, so the compiler will emit a warning if new statement
// kinds are added without being handled here
}
return 1;
}
static int
astfold_excepthandler(excepthandler_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
switch (node_->kind) {
case ExceptHandler_kind:
CALL_OPT(astfold_expr, expr_ty, node_->v.ExceptHandler.type);
CALL_SEQ(astfold_stmt, stmt, node_->v.ExceptHandler.body);
break;
// No default case, so the compiler will emit a warning if new handler
// kinds are added without being handled here
}
return 1;
}
static int
astfold_withitem(withitem_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
CALL(astfold_expr, expr_ty, node_->context_expr);
CALL_OPT(astfold_expr, expr_ty, node_->optional_vars);
return 1;
}
static int
astfold_pattern_negative(expr_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
assert(node_->kind == UnaryOp_kind);
assert(node_->v.UnaryOp.op == USub);
assert(node_->v.UnaryOp.operand->kind == Constant_kind);
PyObject *value = node_->v.UnaryOp.operand->v.Constant.value;
assert(PyComplex_CheckExact(value) ||
PyFloat_CheckExact(value) ||
PyLong_CheckExact(value));
PyObject *negated = PyNumber_Negative(value);
if (negated == NULL) {
return 0;
}
assert(PyComplex_CheckExact(negated) ||
PyFloat_CheckExact(negated) ||
PyLong_CheckExact(negated));
return make_const(node_, negated, ctx_);
}
static int
astfold_pattern_complex(expr_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
expr_ty left = node_->v.BinOp.left;
expr_ty right = node_->v.BinOp.right;
if (left->kind == UnaryOp_kind) {
CALL(astfold_pattern_negative, expr_ty, left);
}
assert(left->kind = Constant_kind);
assert(right->kind = Constant_kind);
// LHS must be real, RHS must be imaginary:
if (!(PyFloat_CheckExact(left->v.Constant.value) ||
PyLong_CheckExact(left->v.Constant.value)) ||
!PyComplex_CheckExact(right->v.Constant.value))
{
// Not actually valid, but it's the compiler's job to complain:
return 1;
}
PyObject *new;
if (node_->v.BinOp.op == Add) {
new = PyNumber_Add(left->v.Constant.value, right->v.Constant.value);
}
else {
assert(node_->v.BinOp.op == Sub);
new = PyNumber_Subtract(left->v.Constant.value, right->v.Constant.value);
}
if (new == NULL) {
return 0;
}
assert(PyComplex_CheckExact(new));
return make_const(node_, new, ctx_);
}
static int
astfold_pattern_keyword(keyword_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
CALL(astfold_pattern, expr_ty, node_->value);
return 1;
}
static int
astfold_pattern(expr_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
// Don't blindly optimize the pattern as an expr; it plays by its own rules!
// Currently, this is only used to form complex/negative numeric constants.
switch (node_->kind) {
case Attribute_kind:
break;
case BinOp_kind:
CALL(astfold_pattern_complex, expr_ty, node_);
break;
case Call_kind:
CALL_SEQ(astfold_pattern, expr, node_->v.Call.args);
CALL_SEQ(astfold_pattern_keyword, keyword, node_->v.Call.keywords);
break;
case Constant_kind:
break;
case Dict_kind:
CALL_SEQ(astfold_pattern, expr, node_->v.Dict.keys);
CALL_SEQ(astfold_pattern, expr, node_->v.Dict.values);
break;
// Not actually valid, but it's the compiler's job to complain:
case JoinedStr_kind:
break;
case List_kind:
CALL_SEQ(astfold_pattern, expr, node_->v.List.elts);
break;
case MatchAs_kind:
CALL(astfold_pattern, expr_ty, node_->v.MatchAs.pattern);
break;
case MatchOr_kind:
CALL_SEQ(astfold_pattern, expr, node_->v.MatchOr.patterns);
break;
case Name_kind:
break;
case Starred_kind:
CALL(astfold_pattern, expr_ty, node_->v.Starred.value);
break;
case Tuple_kind:
CALL_SEQ(astfold_pattern, expr, node_->v.Tuple.elts);
break;
case UnaryOp_kind:
CALL(astfold_pattern_negative, expr_ty, node_);
break;
default:
Py_UNREACHABLE();
}
return 1;
}
static int
astfold_match_case(match_case_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
CALL(astfold_pattern, expr_ty, node_->pattern);
CALL_OPT(astfold_expr, expr_ty, node_->guard);
CALL_SEQ(astfold_stmt, stmt, node_->body);
return 1;
}
#undef CALL
#undef CALL_OPT
#undef CALL_SEQ
#undef CALL_INT_SEQ
int
_PyAST_Optimize(mod_ty mod, PyArena *arena, _PyASTOptimizeState *state)
{
int ret = astfold_mod(mod, arena, state);
assert(ret || PyErr_Occurred());
return ret;
}
Reference in New Issue View Git Blame Copy Permalink