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`.
This commit is contained in:
Serhiy Storchaka 2020-02-12 12:18:59 +02:00 committed by GitHub
parent 0cc6b5e559
commit 8c579b1cc8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 145 additions and 18 deletions

View file

@ -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 @@ ex_call:
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;