use inlined generator expression for sorted(genexpr)

CHANGES.rst

@@ -45,6 +45,9 @@ Features added
   for-loops when called on generator expressions to avoid the generator
   iteration overhead.
 
+* The builtin function ``sorted()`` is optimised when called on a
+  generator expression.
+
 * Keyword argument dicts are no longer copied on function entry when they
   are not being used or only passed through to other function calls (e.g.
   in wrapper functions).

Cython/Compiler/ExprNodes.py

@@ -7316,32 +7316,42 @@ class DictComprehensionAppendNode(ComprehensionAppendNode):
 
 
 class InlinedGeneratorExpressionNode(ExprNode):
-    # An inlined generator expression for which the result is
-    # calculated inside of the loop.  This will only be created by
-    # transforms when replacing builtin calls on generator
-    # expressions.
+    # An inlined generator expression for which the result is calculated
+    # inside of the loop and returned as a single, first and only Generator
+    # return value.
+    # This will only be created by transforms when replacing safe builtin
+    # calls on generator expressions.
     #
     # gen            GeneratorExpressionNode      the generator, not containing any YieldExprNodes
     # orig_func      String                       the name of the builtin function this node replaces
+    # target         ExprNode or None             a 'target' for a ComprehensionAppend node
 
     subexprs = ["gen"]
     orig_func = None
+    target = None
+    is_temp = True
     type = py_object_type
 
-    def __init__(self, pos, gen, **kwargs):
-        gen.def_node.gbody.is_inlined = True
-        kwargs['gen'] = gen
-        super(InlinedGeneratorExpressionNode, self).__init__(pos, **kwargs)
+    def __init__(self, pos, gen, comprehension_type=None, **kwargs):
+        gbody = gen.def_node.gbody
+        gbody.is_inlined = True
+        if comprehension_type is not None:
+            assert comprehension_type in (list_type, set_type, dict_type), comprehension_type
+            gbody.inlined_comprehension_type = comprehension_type
+            kwargs.update(
+                target=RawCNameExprNode(pos, comprehension_type, Naming.retval_cname),
+                type=comprehension_type,
+            )
+        super(InlinedGeneratorExpressionNode, self).__init__(pos, gen=gen, **kwargs)
 
     def may_be_none(self):
-        return self.orig_func not in ('any', 'all')
+        return self.orig_func not in ('any', 'all', 'sorted')
 
     def infer_type(self, env):
-        return py_object_type
+        return self.type
 
     def analyse_types(self, env):
         self.gen = self.gen.analyse_expressions(env)
-        self.is_temp = True
         return self
 
     def generate_result_code(self, code):
@@ -7351,62 +7361,6 @@ class InlinedGeneratorExpressionNode(ExprNode):
         code.put_gotref(self.result())
 
 
-class __InlinedGeneratorExpressionNode(ScopedExprNode):
-    # An inlined generator expression for which the result is
-    # calculated inside of the loop.  This will only be created by
-    # transforms when replacing builtin calls on generator
-    # expressions.
-    #
-    # loop           ForStatNode      the for-loop, not containing any YieldExprNodes
-    # result_node    ResultRefNode    the reference to the result value temp
-    # orig_func      String           the name of the builtin function this node replaces
-
-    child_attrs = ["loop"]
-    loop_analysed = False
-    type = py_object_type
-
-    def analyse_scoped_declarations(self, env):
-        self.loop.analyse_declarations(env)
-
-    def may_be_none(self):
-        return False
-
-    def annotate(self, code):
-        self.loop.annotate(code)
-
-    def infer_type(self, env):
-        return self.result_node.infer_type(env)
-
-    def analyse_types(self, env):
-        if not self.has_local_scope:
-            self.loop_analysed = True
-            self.loop = self.loop.analyse_expressions(env)
-        self.type = self.result_node.type
-        self.is_temp = True
-        return self
-
-    def analyse_scoped_expressions(self, env):
-        self.loop_analysed = True
-        if self.has_local_scope:
-            self.loop = self.loop.analyse_expressions(env)
-        return self
-
-    def coerce_to(self, dst_type, env):
-        if self.orig_func == 'sum' and dst_type.is_numeric and not self.loop_analysed:
-            # We can optimise by dropping the aggregation variable and
-            # the add operations into C.  This can only be done safely
-            # before analysing the loop body, after that, the result
-            # reference type will have infected expressions and
-            # assignments.
-            self.result_node.type = self.type = dst_type
-            return self
-        return super(InlinedGeneratorExpressionNode, self).coerce_to(dst_type, env)
-
-    def generate_result_code(self, code):
-        self.result_node.result_code = self.result()
-        self.loop.generate_execution_code(code)
-
-
 class MergedSequenceNode(ExprNode):
     """
     Merge a sequence of iterables into a set/list/tuple.

Cython/Compiler/Nodes.py

@@ -3996,6 +3996,7 @@ class GeneratorBodyDefNode(DefNode):
 
     is_generator_body = True
     is_inlined = False
+    inlined_comprehension_type = None  # container type for inlined comprehensions
 
     def __init__(self, pos=None, name=None, body=None):
         super(GeneratorBodyDefNode, self).__init__(
@@ -4058,6 +4059,23 @@ class GeneratorBodyDefNode(DefNode):
         code.putln('%s' %
                    (code.error_goto_if_null(Naming.sent_value_cname, self.pos)))
 
+        # ----- prepare target container for inlined comprehension
+        if self.is_inlined and self.inlined_comprehension_type is not None:
+            target_type = self.inlined_comprehension_type
+            if target_type is Builtin.list_type:
+                comp_init = 'PyList_New(0)'
+            elif target_type is Builtin.set_type:
+                comp_init = 'PySet_New(NULL)'
+            elif target_type is Builtin.dict_type:
+                comp_init = 'PyDict_New()'
+            else:
+                raise InternalError(
+                    "invalid type of inlined comprehension: %s" % target_type)
+            code.putln("%s = %s; %s" % (
+                Naming.retval_cname, comp_init,
+                code.error_goto_if_null(Naming.retval_cname, self.pos)))
+            code.put_gotref(Naming.retval_cname)
+
         # ----- Function body
         self.generate_function_body(env, code)
         # ----- Closure initialization
@@ -4073,13 +4091,15 @@ class GeneratorBodyDefNode(DefNode):
 
         # on normal generator termination, we do not take the exception propagation
         # path: no traceback info is required and not creating it is much faster
-        if not self.body.is_terminator:
+        if not self.is_inlined and not self.body.is_terminator:
             code.putln('PyErr_SetNone(PyExc_StopIteration);')
         # ----- Error cleanup
         if code.error_label in code.labels_used:
             if not self.body.is_terminator:
                 code.put_goto(code.return_label)
             code.put_label(code.error_label)
+            if self.is_inlined and self.inlined_comprehension_type is not None:
+                code.put_xdecref_clear(Naming.retval_cname, py_object_type)
             if Future.generator_stop in env.global_scope().context.future_directives:
                 # PEP 479: turn accidental StopIteration exceptions into a RuntimeError
                 code.globalstate.use_utility_code(UtilityCode.load_cached("pep479", "Coroutine.c"))

Cython/Compiler/Optimize.py

@@ -1579,60 +1579,60 @@ class EarlyReplaceBuiltinCalls(Visitor.EnvTransform):
         """
         if len(pos_args) != 1:
             return node
-        if isinstance(pos_args[0], ExprNodes.ComprehensionNode) \
-               and pos_args[0].type is Builtin.list_type:
-            listcomp_node = pos_args[0]
-            loop_node = listcomp_node.loop
-        elif isinstance(pos_args[0], ExprNodes.GeneratorExpressionNode):
-            gen_expr_node = pos_args[0]
+
+        arg = pos_args[0]
+        if isinstance(arg, ExprNodes.ComprehensionNode) and arg.type is Builtin.list_type:
+            list_node = pos_args[0]
+            loop_node = list_node.loop
+
+        elif isinstance(arg, ExprNodes.GeneratorExpressionNode):
+            gen_expr_node = arg
             loop_node = gen_expr_node.loop
             yield_expression, yield_stat_node = self._find_single_yield_expression(loop_node)
-            # FIXME: currently nonfunctional
-            yield_expression = None
             if yield_expression is None:
                 return node
 
+            list_node = ExprNodes.InlinedGeneratorExpressionNode(
+                node.pos, gen_expr_node, orig_func='sorted',
+                comprehension_type=Builtin.list_type)
+
             append_node = ExprNodes.ComprehensionAppendNode(
-                yield_expression.pos, expr = yield_expression)
+                yield_expression.pos,
+                expr=yield_expression,
+                target=list_node.target)
 
             Visitor.recursively_replace_node(loop_node, yield_stat_node, append_node)
 
-            listcomp_node = ExprNodes.ComprehensionNode(
-                gen_expr_node.pos, loop = loop_node,
-                append = append_node, type = Builtin.list_type,
-                expr_scope = gen_expr_node.expr_scope,
-                has_local_scope = True)
-            append_node.target = listcomp_node
-        elif isinstance(pos_args[0], (ExprNodes.ListNode, ExprNodes.TupleNode)):
-            # sorted([a, b, c]) or sorted((a, b, c)). The result of the latter
-            # is a list in CPython, so change it into one.
-            expr = pos_args[0].as_list()
-            listcomp_node = loop_node = expr
+        elif arg.is_sequence_constructor:
+            # sorted([a, b, c]) or sorted((a, b, c)).  The result is always a list,
+            # so starting off with a fresh one is more efficient.
+            list_node = loop_node = arg.as_list()
+
         else:
             # Interestingly, PySequence_List works on a lot of non-sequence
             # things as well.
-            listcomp_node = loop_node = ExprNodes.PythonCapiCallNode(
+            list_node = loop_node = ExprNodes.PythonCapiCallNode(
                 node.pos, "PySequence_List", self.PySequence_List_func_type,
                 args=pos_args, is_temp=True)
 
         result_node = UtilNodes.ResultRefNode(
-            pos = loop_node.pos, type = Builtin.list_type, may_hold_none=False)
-        listcomp_assign_node = Nodes.SingleAssignmentNode(
-            node.pos, lhs = result_node, rhs = listcomp_node, first = True)
+            pos=loop_node.pos, type=Builtin.list_type, may_hold_none=False)
+        list_assign_node = Nodes.SingleAssignmentNode(
+            node.pos, lhs=result_node, rhs=list_node, first=True)
 
         sort_method = ExprNodes.AttributeNode(
-            node.pos, obj = result_node, attribute = EncodedString('sort'),
+            node.pos, obj=result_node, attribute=EncodedString('sort'),
             # entry ? type ?
-            needs_none_check = False)
+            needs_none_check=False)
         sort_node = Nodes.ExprStatNode(
-            node.pos, expr = ExprNodes.SimpleCallNode(
-                node.pos, function = sort_method, args = []))
+            node.pos, expr=ExprNodes.SimpleCallNode(
+                node.pos, function=sort_method, args=[]))
 
         sort_node.analyse_declarations(self.current_env())
 
         return UtilNodes.TempResultFromStatNode(
             result_node,
-            Nodes.StatListNode(node.pos, stats = [ listcomp_assign_node, sort_node ]))
+            Nodes.StatListNode(node.pos, stats=[list_assign_node, sort_node]))
 
     def __handle_simple_function_sum(self, node, pos_args):
         """Transform sum(genexpr) into an equivalent inlined aggregation loop.

tests/run/builtin_sorted.pyx

@@ -6,7 +6,7 @@ def generator():
     yield 3
 
 def returns_set():
-    return set(["foo", "bar", "baz"])
+    return {"foo", "bar", "baz"}
 
 def returns_tuple():
     return (1, 2, 3, 0)
@@ -31,9 +31,9 @@ def sorted_arg(x):
     """
     return sorted(x)
 
-#@cython.test_fail_if_path_exists("//GeneratorExpressionNode",
-#                                 "//ComprehensionNode//NoneCheckNode")
-#@cython.test_assert_path_exists("//ComprehensionNode")
+@cython.test_fail_if_path_exists("//YieldExprNode",
+                                 "//NoneCheckNode")
+@cython.test_assert_path_exists("//InlinedGeneratorExpressionNode")
 def sorted_genexp():
     """
     >>> sorted_genexp()