Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

gh-108113: Make it possible to optimize an AST #108282

Merged
merged 2 commits into from
Aug 23, 2023
Merged

gh-108113: Make it possible to optimize an AST #108282

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b214797
gh-108113: Make it possible to optimize and AST
iritkatriel Aug 22, 2023
37a52a6
do not allocate an arena for a NOP
iritkatriel Aug 22, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions Include/internal/pycore_compile.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,14 @@ PyAPI_FUNC(PyCodeObject*) _PyAST_Compile(
int optimize,
struct _arena *arena);

/* AST optimizations */
PyAPI_FUNC(int) _PyCompile_AstOptimize(
struct _mod *mod,
PyObject *filename,
PyCompilerFlags *flags,
int optimize,
struct _arena *arena);

static const _PyCompilerSrcLocation NO_LOCATION = {-1, -1, -1, -1};

extern int _PyAST_Optimize(
Expand Down
26 changes: 15 additions & 11 deletions Lib/test/test_ast.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -361,14 +361,16 @@ def test_optimization_levels__debug__(self):
cases = [(-1, '__debug__'), (0, '__debug__'), (1, False), (2, False)]
for (optval, expected) in cases:
with self.subTest(optval=optval, expected=expected):
res = ast.parse("__debug__", optimize=optval)
self.assertIsInstance(res.body[0], ast.Expr)
if isinstance(expected, bool):
self.assertIsInstance(res.body[0].value, ast.Constant)
self.assertEqual(res.body[0].value.value, expected)
else:
self.assertIsInstance(res.body[0].value, ast.Name)
self.assertEqual(res.body[0].value.id, expected)
res1 = ast.parse("__debug__", optimize=optval)
res2 = ast.parse(ast.parse("__debug__"), optimize=optval)
for res in [res1, res2]:
self.assertIsInstance(res.body[0], ast.Expr)
if isinstance(expected, bool):
self.assertIsInstance(res.body[0].value, ast.Constant)
self.assertEqual(res.body[0].value.value, expected)
else:
self.assertIsInstance(res.body[0].value, ast.Name)
self.assertEqual(res.body[0].value.id, expected)

def test_optimization_levels_const_folding(self):
folded = ('Expr', (1, 0, 1, 5), ('Constant', (1, 0, 1, 5), 3, None))
Expand All @@ -381,9 +383,11 @@ def test_optimization_levels_const_folding(self):
cases = [(-1, not_folded), (0, not_folded), (1, folded), (2, folded)]
for (optval, expected) in cases:
with self.subTest(optval=optval):
tree = ast.parse("1 + 2", optimize=optval)
res = to_tuple(tree.body[0])
self.assertEqual(res, expected)
tree1 = ast.parse("1 + 2", optimize=optval)
tree2 = ast.parse(ast.parse("1 + 2"), optimize=optval)
for tree in [tree1, tree2]:
res = to_tuple(tree.body[0])
self.assertEqual(res, expected)

def test_invalid_position_information(self):
invalid_linenos = [
Expand Down
12 changes: 7 additions & 5 deletions Lib/test/test_builtin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -521,9 +521,10 @@ def test_compile_async_generator(self):
def test_compile_ast(self):
args = ("a*(1+2)", "f.py", "exec")
raw = compile(*args, flags = ast.PyCF_ONLY_AST).body[0]
opt = compile(*args, flags = ast.PyCF_OPTIMIZED_AST).body[0]
opt1 = compile(*args, flags = ast.PyCF_OPTIMIZED_AST).body[0]
opt2 = compile(ast.parse(args[0]), *args[1:], flags = ast.PyCF_OPTIMIZED_AST).body[0]

for tree in (raw, opt):
for tree in (raw, opt1, opt2):
self.assertIsInstance(tree.value, ast.BinOp)
self.assertIsInstance(tree.value.op, ast.Mult)
self.assertIsInstance(tree.value.left, ast.Name)
Expand All @@ -536,9 +537,10 @@ def test_compile_ast(self):
self.assertIsInstance(raw_right.right, ast.Constant)
self.assertEqual(raw_right.right.value, 2)

opt_right = opt.value.right # expect Constant(3)
self.assertIsInstance(opt_right, ast.Constant)
self.assertEqual(opt_right.value, 3)
for opt in [opt1, opt2]:
opt_right = opt.value.right # expect Constant(3)
self.assertIsInstance(opt_right, ast.Constant)
self.assertEqual(opt_right.value, 3)

def test_delattr(self):
sys.spam = 1
Expand Down
32 changes: 23 additions & 9 deletions Python/bltinmodule.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -804,25 +804,39 @@ builtin_compile_impl(PyObject *module, PyObject *source, PyObject *filename,
if (is_ast == -1)
goto error;
if (is_ast) {
PyArena *arena = _PyArena_New();
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if anyone has perf-sensitive code that does ast.parse on an AST as a no-op, and what the impact on that code will be of allocating and freeing an unused arena. Probably not an issue? I think it would be possible, but slightly more complex, to still avoid arena allocation in the no-op case.

if (arena == NULL) {
goto error;
}

if (flags & PyCF_ONLY_AST) {
result = Py_NewRef(source);
if ((flags & PyCF_OPTIMIZED_AST) == PyCF_OPTIMIZED_AST) {
mod_ty mod = PyAST_obj2mod(source, arena, compile_mode);
if (mod == NULL || !_PyAST_Validate(mod)) {
_PyArena_Free(arena);
goto error;
}
if (_PyCompile_AstOptimize(mod, filename, &cf, optimize,
arena) < 0) {
_PyArena_Free(arena);
goto error;
}
result = PyAST_mod2obj(mod);
}
else {
result = Py_NewRef(source);
}
}
else {
PyArena *arena;
mod_ty mod;

arena = _PyArena_New();
if (arena == NULL)
goto error;
mod = PyAST_obj2mod(source, arena, compile_mode);
mod_ty mod = PyAST_obj2mod(source, arena, compile_mode);
if (mod == NULL || !_PyAST_Validate(mod)) {
_PyArena_Free(arena);
goto error;
}
result = (PyObject*)_PyAST_Compile(mod, filename,
&cf, optimize, arena);
_PyArena_Free(arena);
}
_PyArena_Free(arena);
goto finally;
}

Expand Down
18 changes: 18 additions & 0 deletions Python/compile.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -557,6 +557,24 @@ _PyAST_Compile(mod_ty mod, PyObject *filename, PyCompilerFlags *pflags,
return co;
}

int
_PyCompile_AstOptimize(mod_ty mod, PyObject *filename, PyCompilerFlags *cf,
int optimize, PyArena *arena)
{
PyFutureFeatures future;
if (!_PyFuture_FromAST(mod, filename, &future)) {
return -1;
}
int flags = future.ff_features | cf->cf_flags;
if (optimize == -1) {
optimize = _Py_GetConfig()->optimization_level;
}
if (!_PyAST_Optimize(mod, arena, optimize, flags)) {
return -1;
}
return 0;
}

static void
compiler_free(struct compiler *c)
{
Expand Down
22 changes: 1 addition & 21 deletions Python/pythonrun.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@
#include "pycore_pyerrors.h" // _PyErr_GetRaisedException, _Py_Offer_Suggestions
#include "pycore_pylifecycle.h" // _Py_UnhandledKeyboardInterrupt
#include "pycore_pystate.h" // _PyInterpreterState_GET()
#include "pycore_symtable.h" // _PyFuture_FromAST()
#include "pycore_sysmodule.h" // _PySys_Audit()
#include "pycore_traceback.h" // _PyTraceBack_Print_Indented()

Expand Down Expand Up @@ -1791,24 +1790,6 @@ run_pyc_file(FILE *fp, PyObject *globals, PyObject *locals,
return NULL;
}

static int
ast_optimize(mod_ty mod, PyObject *filename, PyCompilerFlags *cf,
int optimize, PyArena *arena)
{
PyFutureFeatures future;
if (!_PyFuture_FromAST(mod, filename, &future)) {
return -1;
}
int flags = future.ff_features | cf->cf_flags;
if (optimize == -1) {
optimize = _Py_GetConfig()->optimization_level;
}
if (!_PyAST_Optimize(mod, arena, optimize, flags)) {
return -1;
}
return 0;
}

PyObject *
Py_CompileStringObject(const char *str, PyObject *filename, int start,
PyCompilerFlags *flags, int optimize)
Expand All @@ -1826,8 +1807,7 @@ Py_CompileStringObject(const char *str, PyObject *filename, int start,
}
if (flags && (flags->cf_flags & PyCF_ONLY_AST)) {
if ((flags->cf_flags & PyCF_OPTIMIZED_AST) == PyCF_OPTIMIZED_AST) {
if (ast_optimize(mod, filename, flags, optimize, arena) < 0) {
_PyArena_Free(arena);
if (_PyCompile_AstOptimize(mod, filename, flags, optimize, arena) < 0) {
return NULL;
}
}
Expand Down