bpo-32856: Optimize the assignment idiom in comprehensions. (GH-16814)
https://github.com/python/cpython/commit/8c579b1cc86053473eb052b763272794767... commit: 8c579b1cc86053473eb052b76327279476740c9b branch: master author: Serhiy Storchaka <storchaka@gmail.com> committer: GitHub <noreply@github.com> date: 2020-02-12T12:18:59+02:00 summary: bpo-32856: Optimize the assignment idiom in comprehensions. (GH-16814) Now `for y in [expr]` in comprehensions is as fast as a simple assignment `y = expr`. files: A Misc/NEWS.d/next/Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst M Doc/whatsnew/3.9.rst M Lib/test/test_dictcomps.py M Lib/test/test_genexps.py M Lib/test/test_listcomps.py M Lib/test/test_peepholer.py M Lib/test/test_setcomps.py M Python/compile.c diff --git a/Doc/whatsnew/3.9.rst b/Doc/whatsnew/3.9.rst index c0404101d7ce2..ec179845aee7c 100644 --- a/Doc/whatsnew/3.9.rst +++ b/Doc/whatsnew/3.9.rst @@ -315,6 +315,17 @@ case), and one used ``__VENV_NAME__`` instead. Optimizations ============= +* Optimized the idiom for assignment a temporary variable in comprehensions. + Now ``for y in [expr]`` in comprehensions is as fast as a simple assignment + ``y = expr``. For example: + + sums = [s for s in [0] for x in data for s in [s + x]] + + Unlike to the ``:=`` operator this idiom does not leak a variable to the + outer scope. + + (Contributed by Serhiy Storchaka in :issue:`32856`.) + Build and C API Changes ======================= diff --git a/Lib/test/test_dictcomps.py b/Lib/test/test_dictcomps.py index 927e3103e664b..16aa651b93c46 100644 --- a/Lib/test/test_dictcomps.py +++ b/Lib/test/test_dictcomps.py @@ -111,5 +111,22 @@ def add_call(pos, value): self.assertEqual(actual, expected) self.assertEqual(actual_calls, expected_calls) + def test_assignment_idiom_in_comprehensions(self): + expected = {1: 1, 2: 4, 3: 9, 4: 16} + actual = {j: j*j for i in range(4) for j in [i+1]} + self.assertEqual(actual, expected) + expected = {3: 2, 5: 6, 7: 12, 9: 20} + actual = {j+k: j*k for i in range(4) for j in [i+1] for k in [j+1]} + self.assertEqual(actual, expected) + expected = {3: 2, 5: 6, 7: 12, 9: 20} + actual = {j+k: j*k for i in range(4) for j, k in [(i+1, i+2)]} + self.assertEqual(actual, expected) + + def test_star_expression(self): + expected = {0: 0, 1: 1, 2: 4, 3: 9} + self.assertEqual({i: i*i for i in [*range(4)]}, expected) + self.assertEqual({i: i*i for i in (*range(4),)}, expected) + + if __name__ == "__main__": unittest.main() diff --git a/Lib/test/test_genexps.py b/Lib/test/test_genexps.py index fd712bb172d5d..86e4e195f55ec 100644 --- a/Lib/test/test_genexps.py +++ b/Lib/test/test_genexps.py @@ -15,6 +15,22 @@ >>> list((i,j) for i in range(4) for j in range(i) ) [(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2)] +Test the idiom for temporary variable assignment in comprehensions. + + >>> list((j*j for i in range(4) for j in [i+1])) + [1, 4, 9, 16] + >>> list((j*k for i in range(4) for j in [i+1] for k in [j+1])) + [2, 6, 12, 20] + >>> list((j*k for i in range(4) for j, k in [(i+1, i+2)])) + [2, 6, 12, 20] + +Not assignment + + >>> list((i*i for i in [*range(4)])) + [0, 1, 4, 9] + >>> list((i*i for i in (*range(4),))) + [0, 1, 4, 9] + Make sure the induction variable is not exposed >>> i = 20 diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py index ddb169fe58957..62b3319ad936d 100644 --- a/Lib/test/test_listcomps.py +++ b/Lib/test/test_listcomps.py @@ -16,6 +16,22 @@ >>> [(i,j) for i in range(4) for j in range(i)] [(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2)] +Test the idiom for temporary variable assignment in comprehensions. + + >>> [j*j for i in range(4) for j in [i+1]] + [1, 4, 9, 16] + >>> [j*k for i in range(4) for j in [i+1] for k in [j+1]] + [2, 6, 12, 20] + >>> [j*k for i in range(4) for j, k in [(i+1, i+2)]] + [2, 6, 12, 20] + +Not assignment + + >>> [i*i for i in [*range(4)]] + [0, 1, 4, 9] + >>> [i*i for i in (*range(4),)] + [0, 1, 4, 9] + Make sure the induction variable is not exposed >>> i = 20 diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index 567e6a14361d1..7913e91e453ca 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -495,6 +495,20 @@ def f(x): return 6 self.check_lnotab(f) + def test_assignment_idiom_in_comprehensions(self): + def listcomp(): + return [y for x in a for y in [f(x)]] + self.assertEqual(count_instr_recursively(listcomp, 'FOR_ITER'), 1) + def setcomp(): + return {y for x in a for y in [f(x)]} + self.assertEqual(count_instr_recursively(setcomp, 'FOR_ITER'), 1) + def dictcomp(): + return {y: y for x in a for y in [f(x)]} + self.assertEqual(count_instr_recursively(dictcomp, 'FOR_ITER'), 1) + def genexpr(): + return (y for x in a for y in [f(x)]) + self.assertEqual(count_instr_recursively(genexpr, 'FOR_ITER'), 1) + class TestBuglets(unittest.TestCase): diff --git a/Lib/test/test_setcomps.py b/Lib/test/test_setcomps.py index fb7cde03d7823..ecc4fffec0d84 100644 --- a/Lib/test/test_setcomps.py +++ b/Lib/test/test_setcomps.py @@ -21,6 +21,22 @@ >>> list(sorted({(i,j) for i in range(4) for j in range(i)})) [(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2)] +Test the idiom for temporary variable assignment in comprehensions. + + >>> sorted({j*j for i in range(4) for j in [i+1]}) + [1, 4, 9, 16] + >>> sorted({j*k for i in range(4) for j in [i+1] for k in [j+1]}) + [2, 6, 12, 20] + >>> sorted({j*k for i in range(4) for j, k in [(i+1, i+2)]}) + [2, 6, 12, 20] + +Not assignment + + >>> sorted({i*i for i in [*range(4)]}) + [0, 1, 4, 9] + >>> sorted({i*i for i in (*range(4),)}) + [0, 1, 4, 9] + Make sure the induction variable is not exposed >>> i = 20 diff --git a/Misc/NEWS.d/next/Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst b/Misc/NEWS.d/next/Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst new file mode 100644 index 0000000000000..c1cd68f672712 --- /dev/null +++ b/Misc/NEWS.d/next/Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst @@ -0,0 +1,3 @@ +Optimized the idiom for assignment a temporary variable in comprehensions. +Now ``for y in [expr]`` in comprehensions is as fast as a simple assignment +``y = expr``. diff --git a/Python/compile.c b/Python/compile.c index 04b8fe46e194d..bf8c8109d0758 100644 --- a/Python/compile.c +++ b/Python/compile.c @@ -212,11 +212,13 @@ static int compiler_set_qualname(struct compiler *); static int compiler_sync_comprehension_generator( struct compiler *c, asdl_seq *generators, int gen_index, + int depth, expr_ty elt, expr_ty val, int type); static int compiler_async_comprehension_generator( struct compiler *c, asdl_seq *generators, int gen_index, + int depth, expr_ty elt, expr_ty val, int type); static PyCodeObject *assemble(struct compiler *, int addNone); @@ -4343,22 +4345,24 @@ compiler_call_helper(struct compiler *c, static int compiler_comprehension_generator(struct compiler *c, asdl_seq *generators, int gen_index, + int depth, expr_ty elt, expr_ty val, int type) { comprehension_ty gen; gen = (comprehension_ty)asdl_seq_GET(generators, gen_index); if (gen->is_async) { return compiler_async_comprehension_generator( - c, generators, gen_index, elt, val, type); + c, generators, gen_index, depth, elt, val, type); } else { return compiler_sync_comprehension_generator( - c, generators, gen_index, elt, val, type); + c, generators, gen_index, depth, elt, val, type); } } static int compiler_sync_comprehension_generator(struct compiler *c, asdl_seq *generators, int gen_index, + int depth, expr_ty elt, expr_ty val, int type) { /* generate code for the iterator, then each of the ifs, @@ -4386,12 +4390,38 @@ compiler_sync_comprehension_generator(struct compiler *c, } else { /* Sub-iter - calculate on the fly */ - VISIT(c, expr, gen->iter); - ADDOP(c, GET_ITER); + /* Fast path for the temporary variable assignment idiom: + for y in [f(x)] + */ + asdl_seq *elts; + switch (gen->iter->kind) { + case List_kind: + elts = gen->iter->v.List.elts; + break; + case Tuple_kind: + elts = gen->iter->v.Tuple.elts; + break; + default: + elts = NULL; + } + if (asdl_seq_LEN(elts) == 1) { + expr_ty elt = asdl_seq_GET(elts, 0); + if (elt->kind != Starred_kind) { + VISIT(c, expr, elt); + start = NULL; + } + } + if (start) { + VISIT(c, expr, gen->iter); + ADDOP(c, GET_ITER); + } + } + if (start) { + depth++; + compiler_use_next_block(c, start); + ADDOP_JREL(c, FOR_ITER, anchor); + NEXT_BLOCK(c); } - compiler_use_next_block(c, start); - ADDOP_JREL(c, FOR_ITER, anchor); - NEXT_BLOCK(c); VISIT(c, expr, gen->target); /* XXX this needs to be cleaned up...a lot! */ @@ -4405,7 +4435,7 @@ compiler_sync_comprehension_generator(struct compiler *c, if (++gen_index < asdl_seq_LEN(generators)) if (!compiler_comprehension_generator(c, - generators, gen_index, + generators, gen_index, depth, elt, val, type)) return 0; @@ -4420,18 +4450,18 @@ compiler_sync_comprehension_generator(struct compiler *c, break; case COMP_LISTCOMP: VISIT(c, expr, elt); - ADDOP_I(c, LIST_APPEND, gen_index + 1); + ADDOP_I(c, LIST_APPEND, depth + 1); break; case COMP_SETCOMP: VISIT(c, expr, elt); - ADDOP_I(c, SET_ADD, gen_index + 1); + ADDOP_I(c, SET_ADD, depth + 1); break; case COMP_DICTCOMP: /* With '{k: v}', k is evaluated before v, so we do the same. */ VISIT(c, expr, elt); VISIT(c, expr, val); - ADDOP_I(c, MAP_ADD, gen_index + 1); + ADDOP_I(c, MAP_ADD, depth + 1); break; default: return 0; @@ -4440,8 +4470,10 @@ compiler_sync_comprehension_generator(struct compiler *c, compiler_use_next_block(c, skip); } compiler_use_next_block(c, if_cleanup); - ADDOP_JABS(c, JUMP_ABSOLUTE, start); - compiler_use_next_block(c, anchor); + if (start) { + ADDOP_JABS(c, JUMP_ABSOLUTE, start); + compiler_use_next_block(c, anchor); + } return 1; } @@ -4449,6 +4481,7 @@ compiler_sync_comprehension_generator(struct compiler *c, static int compiler_async_comprehension_generator(struct compiler *c, asdl_seq *generators, int gen_index, + int depth, expr_ty elt, expr_ty val, int type) { comprehension_ty gen; @@ -4492,9 +4525,10 @@ compiler_async_comprehension_generator(struct compiler *c, NEXT_BLOCK(c); } + depth++; if (++gen_index < asdl_seq_LEN(generators)) if (!compiler_comprehension_generator(c, - generators, gen_index, + generators, gen_index, depth, elt, val, type)) return 0; @@ -4509,18 +4543,18 @@ compiler_async_comprehension_generator(struct compiler *c, break; case COMP_LISTCOMP: VISIT(c, expr, elt); - ADDOP_I(c, LIST_APPEND, gen_index + 1); + ADDOP_I(c, LIST_APPEND, depth + 1); break; case COMP_SETCOMP: VISIT(c, expr, elt); - ADDOP_I(c, SET_ADD, gen_index + 1); + ADDOP_I(c, SET_ADD, depth + 1); break; case COMP_DICTCOMP: /* With '{k: v}', k is evaluated before v, so we do the same. */ VISIT(c, expr, elt); VISIT(c, expr, val); - ADDOP_I(c, MAP_ADD, gen_index + 1); + ADDOP_I(c, MAP_ADD, depth + 1); break; default: return 0; @@ -4583,7 +4617,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type, ADDOP_I(c, op, 0); } - if (!compiler_comprehension_generator(c, generators, 0, elt, + if (!compiler_comprehension_generator(c, generators, 0, 0, elt, val, type)) goto error_in_scope;
participants (1)
-
Serhiy Storchaka