[mypyc] Faster set creation with generator expression (#10261)

Closes mypyc/mypyc#771 and #9707.

Changes:

* Move code from tranform_set_comprehension to translate_set_comprehension
* Fix the unnecessary set creation (#9707 (comment))

Author: 97littleleaf11

mypyc/irbuild/expression.py

@@ -30,12 +30,15 @@
 from mypyc.primitives.list_ops import list_append_op, list_extend_op, list_slice_op
 from mypyc.primitives.tuple_ops import list_tuple_op, tuple_slice_op
 from mypyc.primitives.dict_ops import dict_new_op, dict_set_item_op, dict_get_item_op
-from mypyc.primitives.set_ops import new_set_op, set_add_op, set_update_op
+from mypyc.primitives.set_ops import set_add_op, set_update_op
 from mypyc.primitives.str_ops import str_slice_op
 from mypyc.primitives.int_ops import int_comparison_op_mapping
 from mypyc.irbuild.specialize import specializers
 from mypyc.irbuild.builder import IRBuilder
-from mypyc.irbuild.for_helpers import translate_list_comprehension, comprehension_helper
+from mypyc.irbuild.for_helpers import (
+    translate_list_comprehension, translate_set_comprehension,
+    comprehension_helper
+)
 
 
 # Name and attribute references
@@ -656,16 +659,7 @@ def transform_list_comprehension(builder: IRBuilder, o: ListComprehension) -> Va
 
 
 def transform_set_comprehension(builder: IRBuilder, o: SetComprehension) -> Value:
-    gen = o.generator
-    set_ops = builder.call_c(new_set_op, [], o.line)
-    loop_params = list(zip(gen.indices, gen.sequences, gen.condlists))
-
-    def gen_inner_stmts() -> None:
-        e = builder.accept(gen.left_expr)
-        builder.call_c(set_add_op, [set_ops, e], o.line)
-
-    comprehension_helper(builder, loop_params, gen_inner_stmts, o.line)
-    return set_ops
+    return translate_set_comprehension(builder, o.generator)
 
 
 def transform_dictionary_comprehension(builder: IRBuilder, o: DictionaryComprehension) -> Value:

mypyc/irbuild/for_helpers.py

@@ -24,6 +24,7 @@
     dict_key_iter_op, dict_value_iter_op, dict_item_iter_op
 )
 from mypyc.primitives.list_ops import list_append_op, list_get_item_unsafe_op
+from mypyc.primitives.set_ops import set_add_op
 from mypyc.primitives.generic_ops import iter_op, next_op
 from mypyc.primitives.exc_ops import no_err_occurred_op
 from mypyc.irbuild.builder import IRBuilder
@@ -142,6 +143,18 @@ def gen_inner_stmts() -> None:
     return list_ops
 
 
+def translate_set_comprehension(builder: IRBuilder, gen: GeneratorExpr) -> Value:
+    set_ops = builder.new_set_op([], gen.line)
+    loop_params = list(zip(gen.indices, gen.sequences, gen.condlists))
+
+    def gen_inner_stmts() -> None:
+        e = builder.accept(gen.left_expr)
+        builder.call_c(set_add_op, [set_ops, e], gen.line)
+
+    comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line)
+    return set_ops
+
+
 def comprehension_helper(builder: IRBuilder,
                          loop_params: List[Tuple[Lvalue, Expression, List[Expression]]],
                          gen_inner_stmts: Callable[[], None],

mypyc/irbuild/specialize.py

@@ -28,8 +28,8 @@
 from mypyc.primitives.tuple_ops import new_tuple_set_item_op
 from mypyc.irbuild.builder import IRBuilder
 from mypyc.irbuild.for_helpers import (
-    translate_list_comprehension, comprehension_helper,
-    for_loop_helper_with_index
+    translate_list_comprehension, translate_set_comprehension,
+    comprehension_helper, for_loop_helper_with_index,
 )
 
 
@@ -112,8 +112,19 @@ def dict_methods_fast_path(
         return builder.call_c(dict_items_op, [obj], expr.line)
 
 
-@specialize_function('builtins.tuple')
 @specialize_function('builtins.set')
+def translate_set_from_generator_call(
+        builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Optional[Value]:
+    # Special case for set creation from a generator：
+    #     set(f(...) for ... in iterator/nested_generators...)
+    if (len(expr.args) == 1
+            and expr.arg_kinds[0] == ARG_POS
+            and isinstance(expr.args[0], GeneratorExpr)):
+        return translate_set_comprehension(builder, expr.args[0])
+    return None
+
+
+@specialize_function('builtins.tuple')
 @specialize_function('builtins.frozenset')
 @specialize_function('builtins.dict')
 @specialize_function('builtins.sum')
@@ -143,6 +154,7 @@ def translate_safe_generator_call(
                 val = tuple_from_generator_helper(builder, expr.args[0])
                 if val is not None:
                     return val
+
             return builder.call_refexpr_with_args(
                 expr, callee,
                 ([translate_list_comprehension(builder, expr.args[0])]