diff options
Diffstat (limited to 'Python/ast_opt.c')
| -rw-r--r-- | Python/ast_opt.c | 649 |
1 files changed, 649 insertions, 0 deletions
diff --git a/Python/ast_opt.c b/Python/ast_opt.c new file mode 100644 index 0000000..80531cc --- /dev/null +++ b/Python/ast_opt.c @@ -0,0 +1,649 @@ +/* AST Optimizer */ +#include "Python.h" +#include "Python-ast.h" + + +/* TODO: is_const and get_const_value are copied from Python/compile.c. + It should be deduped in the future. Maybe, we can include this file + from compile.c? +*/ +static int +is_const(expr_ty e) +{ + switch (e->kind) { + case Constant_kind: + case Num_kind: + case Str_kind: + case Bytes_kind: + case Ellipsis_kind: + case NameConstant_kind: + return 1; + default: + return 0; + } +} + +static PyObject * +get_const_value(expr_ty e) +{ + switch (e->kind) { + case Constant_kind: + return e->v.Constant.value; + case Num_kind: + return e->v.Num.n; + case Str_kind: + return e->v.Str.s; + case Bytes_kind: + return e->v.Bytes.s; + case Ellipsis_kind: + return Py_Ellipsis; + case NameConstant_kind: + return e->v.NameConstant.value; + default: + Py_UNREACHABLE(); + } +} + +static int +make_const(expr_ty node, PyObject *val, PyArena *arena) +{ + 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.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) +{ + expr_ty arg = node->v.UnaryOp.operand; + + if (!is_const(arg)) { + /* 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. */ + int 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; + default: + op = 0; + } + 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](get_const_value(arg)); + return make_const(node, newval, arena); +} + +static int +fold_binop(expr_ty node, PyArena *arena) +{ + expr_ty lhs, rhs; + lhs = node->v.BinOp.left; + rhs = node->v.BinOp.right; + if (!is_const(lhs) || !is_const(rhs)) { + return 1; + } + + PyObject *lv = get_const_value(lhs); + PyObject *rv = get_const_value(rhs); + PyObject *newval; + + switch (node->v.BinOp.op) { + case Add: + newval = PyNumber_Add(lv, rv); + break; + case Sub: + newval = PyNumber_Subtract(lv, rv); + break; + case Mult: + newval = PyNumber_Multiply(lv, rv); + break; + case Div: + newval = PyNumber_TrueDivide(lv, rv); + break; + case FloorDiv: + newval = PyNumber_FloorDivide(lv, rv); + break; + case Mod: + newval = PyNumber_Remainder(lv, rv); + break; + case Pow: + newval = PyNumber_Power(lv, rv, Py_None); + break; + case LShift: + newval = PyNumber_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; + default: // Unknown operator + return 1; + } + + if (newval == NULL) { + if (PyErr_ExceptionMatches(PyExc_KeyboardInterrupt)) { + return 0; + } + PyErr_Clear(); + return 1; + } + + /* Avoid creating large constants. */ + Py_ssize_t size = PyObject_Size(newval); + if (size == -1) { + if (PyErr_ExceptionMatches(PyExc_KeyboardInterrupt)) { + Py_DECREF(newval); + return 0; + } + PyErr_Clear(); + } + else if (size > 20) { + Py_DECREF(newval); + return 1; + } + return make_const(node, newval, arena); +} + +static PyObject* +make_const_tuple(asdl_seq *elts, int make_set) +{ + for (int i = 0; i < asdl_seq_LEN(elts); i++) { + expr_ty e = (expr_ty)asdl_seq_GET(elts, i); + if (!is_const(e)) { + 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 = get_const_value(e); + Py_INCREF(v); + PyTuple_SET_ITEM(newval, i, v); + } + + /* Need to create frozen_set instead. */ + if (make_set) { + Py_SETREF(newval, PyFrozenSet_New(newval)); + } + return newval; +} + +static int +fold_tuple(expr_ty node, PyArena *arena) +{ + PyObject *newval; + + if (node->v.Tuple.ctx != Load) + return 1; + + newval = make_const_tuple(node->v.Tuple.elts, 0); + return make_const(node, newval, arena); +} + +static int +fold_subscr(expr_ty node, PyArena *arena) +{ + PyObject *newval; + expr_ty arg, idx; + slice_ty slice; + + arg = node->v.Subscript.value; + slice = node->v.Subscript.slice; + if (node->v.Subscript.ctx != Load || + !is_const(arg) || + /* TODO: handle other types of slices */ + slice->kind != Index_kind || + !is_const(slice->v.Index.value)) + { + return 1; + } + + idx = slice->v.Index.value; + newval = PyObject_GetItem(get_const_value(arg), get_const_value(idx)); + return make_const(node, newval, arena); +} + +static int +fold_compare(expr_ty node, PyArena *arena) +{ + asdl_int_seq *ops; + asdl_seq *args; + PyObject *newval; + int i; + + ops = node->v.Compare.ops; + args = node->v.Compare.comparators; + /* TODO: optimize cases with literal arguments. */ + for (i = 0; i < asdl_seq_LEN(ops); i++) { + int op; + expr_ty arg; + asdl_seq *elts; + + op = asdl_seq_GET(ops, i); + arg = (expr_ty)asdl_seq_GET(args, i); + /* Change literal list or set in 'in' or 'not in' into + tuple or frozenset respectively. */ + /* TODO: do the same when list or set is used as iterable + in for loop and comprehensions? */ + if (op != In && op != NotIn) + continue; + if (arg->kind == List_kind) + elts = arg->v.List.elts; + else if (arg->kind == Set_kind) + elts = arg->v.Set.elts; + else continue; + + newval = make_const_tuple(elts, arg->kind == Set_kind); + make_const(arg, newval, arena); + } + return 1; +} + +static int astfold_mod(mod_ty node_, PyArena* ctx_); +static int astfold_stmt(stmt_ty node_, PyArena* ctx_); +static int astfold_expr(expr_ty node_, PyArena* ctx_); +static int astfold_arguments(arguments_ty node_, PyArena* ctx_); +static int astfold_comprehension(comprehension_ty node_, PyArena* ctx_); +static int astfold_keyword(keyword_ty node_, PyArena* ctx_); +static int astfold_slice(slice_ty node_, PyArena* ctx_); +static int astfold_arg(arg_ty node_, PyArena* ctx_); +static int astfold_withitem(withitem_ty node_, PyArena* ctx_); +static int astfold_excepthandler(excepthandler_ty node_, PyArena* ctx_); +#define CALL(FUNC, TYPE, ARG) \ + if (!FUNC((ARG), ctx_)) \ + return 0; + +#define CALL_OPT(FUNC, TYPE, ARG) \ + if ((ARG) != NULL && !FUNC((ARG), ctx_)) \ + return 0; + +#define CALL_SEQ(FUNC, TYPE, ARG) { \ + int i; \ + asdl_seq *seq = (ARG); /* avoid variable capture */ \ + for (i = 0; i < asdl_seq_LEN(seq); i++) { \ + TYPE elt = (TYPE)asdl_seq_GET(seq, i); \ + if (elt != NULL && !FUNC(elt, ctx_)) \ + 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_)) \ + return 0; \ + } \ +} + +static int +astfold_mod(mod_ty node_, PyArena* ctx_) +{ + switch (node_->kind) { + case Module_kind: + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.Module.body); + break; + case Interactive_kind: + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.Interactive.body); + break; + case Expression_kind: + CALL(astfold_expr, expr_ty, node_->v.Expression.body); + break; + case Suite_kind: + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.Suite.body); + break; + default: + break; + } + return 1; +} + +static int +astfold_expr(expr_ty node_, PyArena* ctx_) +{ + switch (node_->kind) { + case BoolOp_kind: + CALL_SEQ(astfold_expr, expr_ty, 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_ty, node_->v.Dict.keys); + CALL_SEQ(astfold_expr, expr_ty, node_->v.Dict.values); + break; + case Set_kind: + CALL_SEQ(astfold_expr, expr_ty, node_->v.Set.elts); + break; + case ListComp_kind: + CALL(astfold_expr, expr_ty, node_->v.ListComp.elt); + CALL_SEQ(astfold_comprehension, comprehension_ty, node_->v.ListComp.generators); + break; + case SetComp_kind: + CALL(astfold_expr, expr_ty, node_->v.SetComp.elt); + CALL_SEQ(astfold_comprehension, comprehension_ty, 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_ty, node_->v.DictComp.generators); + break; + case GeneratorExp_kind: + CALL(astfold_expr, expr_ty, node_->v.GeneratorExp.elt); + CALL_SEQ(astfold_comprehension, comprehension_ty, 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_ty, 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_ty, node_->v.Call.args); + CALL_SEQ(astfold_keyword, keyword_ty, 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_ty, 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_slice, slice_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 List_kind: + CALL_SEQ(astfold_expr, expr_ty, node_->v.List.elts); + break; + case Tuple_kind: + CALL_SEQ(astfold_expr, expr_ty, node_->v.Tuple.elts); + CALL(fold_tuple, expr_ty, node_); + break; + default: + break; + } + return 1; +} + +static int +astfold_slice(slice_ty node_, PyArena* ctx_) +{ + switch (node_->kind) { + 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 ExtSlice_kind: + CALL_SEQ(astfold_slice, slice_ty, node_->v.ExtSlice.dims); + break; + case Index_kind: + CALL(astfold_expr, expr_ty, node_->v.Index.value); + break; + default: + break; + } + return 1; +} + +static int +astfold_keyword(keyword_ty node_, PyArena* ctx_) +{ + CALL(astfold_expr, expr_ty, node_->value); + return 1; +} + +static int +astfold_comprehension(comprehension_ty node_, PyArena* ctx_) +{ + CALL(astfold_expr, expr_ty, node_->target); + CALL(astfold_expr, expr_ty, node_->iter); + CALL_SEQ(astfold_expr, expr_ty, node_->ifs); + return 1; +} + +static int +astfold_arguments(arguments_ty node_, PyArena* ctx_) +{ + CALL_SEQ(astfold_arg, arg_ty, node_->args); + CALL_OPT(astfold_arg, arg_ty, node_->vararg); + CALL_SEQ(astfold_arg, arg_ty, node_->kwonlyargs); + CALL_SEQ(astfold_expr, expr_ty, node_->kw_defaults); + CALL_OPT(astfold_arg, arg_ty, node_->kwarg); + CALL_SEQ(astfold_expr, expr_ty, node_->defaults); + return 1; +} + +static int +astfold_arg(arg_ty node_, PyArena* ctx_) +{ + CALL_OPT(astfold_expr, expr_ty, node_->annotation); + return 1; +} + +static int +astfold_stmt(stmt_ty node_, PyArena* ctx_) +{ + switch (node_->kind) { + case FunctionDef_kind: + CALL(astfold_arguments, arguments_ty, node_->v.FunctionDef.args); + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.FunctionDef.body); + CALL_SEQ(astfold_expr, expr_ty, node_->v.FunctionDef.decorator_list); + CALL_OPT(astfold_expr, expr_ty, node_->v.FunctionDef.returns); + break; + case AsyncFunctionDef_kind: + CALL(astfold_arguments, arguments_ty, node_->v.AsyncFunctionDef.args); + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.AsyncFunctionDef.body); + CALL_SEQ(astfold_expr, expr_ty, node_->v.AsyncFunctionDef.decorator_list); + CALL_OPT(astfold_expr, expr_ty, node_->v.AsyncFunctionDef.returns); + break; + case ClassDef_kind: + CALL_SEQ(astfold_expr, expr_ty, node_->v.ClassDef.bases); + CALL_SEQ(astfold_keyword, keyword_ty, node_->v.ClassDef.keywords); + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.ClassDef.body); + CALL_SEQ(astfold_expr, expr_ty, 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_ty, node_->v.Delete.targets); + break; + case Assign_kind: + CALL_SEQ(astfold_expr, expr_ty, 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(astfold_expr, expr_ty, node_->v.AnnAssign.annotation); + 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_ty, node_->v.For.body); + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.For.orelse); + 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_ty, node_->v.AsyncFor.body); + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.AsyncFor.orelse); + break; + case While_kind: + CALL(astfold_expr, expr_ty, node_->v.While.test); + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.While.body); + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.While.orelse); + break; + case If_kind: + CALL(astfold_expr, expr_ty, node_->v.If.test); + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.If.body); + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.If.orelse); + break; + case With_kind: + CALL_SEQ(astfold_withitem, withitem_ty, node_->v.With.items); + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.With.body); + break; + case AsyncWith_kind: + CALL_SEQ(astfold_withitem, withitem_ty, node_->v.AsyncWith.items); + CALL_SEQ(astfold_stmt, stmt_ty, 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_ty, node_->v.Try.body); + CALL_SEQ(astfold_excepthandler, excepthandler_ty, node_->v.Try.handlers); + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.Try.orelse); + CALL_SEQ(astfold_stmt, stmt_ty, 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; + default: + break; + } + return 1; +} + +static int +astfold_excepthandler(excepthandler_ty node_, PyArena* ctx_) +{ + switch (node_->kind) { + case ExceptHandler_kind: + CALL_OPT(astfold_expr, expr_ty, node_->v.ExceptHandler.type); + CALL_SEQ(astfold_stmt, stmt_ty, node_->v.ExceptHandler.body); + break; + default: + break; + } + return 1; +} + +static int +astfold_withitem(withitem_ty node_, PyArena* ctx_) +{ + CALL(astfold_expr, expr_ty, node_->context_expr); + CALL_OPT(astfold_expr, expr_ty, node_->optional_vars); + return 1; +} + +#undef CALL +#undef CALL_OPT +#undef CALL_SEQ +#undef CALL_INT_SEQ + +int +_PyAST_Optimize(mod_ty mod, PyArena *arena) +{ + int ret = astfold_mod(mod, arena); + assert(ret || PyErr_Occurred()); + return ret; +} |