Fix combined cartesian / flattened generators (#263)

c42f · web-flow · commit 8098cb04b4b6 · 2023-05-03T05:34:32.000+10:00
The combined cartesian/flattened generator syntax is a very unusual use
case found in only one or two packages in the general registry. It looks
like this

    ((a,b,c,d) for a=as, b=bs for c=cs, d=ds)

Practically all such generators end up as a flat list (though ... should
they?) but the ordering of tuples (a,b,c,d) in the list depends on which
parts are cartesian (separated by commas) and which parts are flattened
(separated by for's).

Here we fix the `SyntaxNode` representation of these, preferring a single
`generator` head rather than combined `flatten` and `generator`, and
instead grouping the sets of cartesian iterations into
`cartesian_iterator` groups. This keeps the tree structure flat and
closer to the source text while also faithfully representing these
unusual forms.  For example, the above parses as

    (generator
      (tuple a b c d)
      (cartesian_iterator
        (= a as)
        (= b bs))
      (cartesian_iterator
        (= c cs)
        (= d ds)))

In addition, we also make use of the `cartesian_iterator` head in `for`
loops, replacing the `block` which was used previously (but is not
semantically or syntactically a block!). Thus `for i=is, j=js body end`
now parses as

    (for
      (cartesian_iterator
        (= i is)
        (= j js))
      body)

This also improves Expr conversion, removing some special cases which
were necessary when processing `block's` in that situation.
diff --git a/README.md b/README.md
@@ -375,7 +375,7 @@ The children of our trees are strictly in source order. This has many
 consequences in places where `Expr` reorders child expressions.
 
 * Infix and postfix operator calls have the operator name in the *second* child position. `a + b` is parsed as `(call-i a + b)` - where the infix `-i` flag indicates infix child position - rather than `Expr(:call, :+, :a, :b)`.
-* Flattened generators are represented in source order
+* Generators are represented in source order as a single node rather than multiple nested flatten and generator expressions.
 
 ### No `LineNumberNode`s
 
@@ -435,15 +435,16 @@ class of tokenization errors and lets the parser deal with them.
 * Using `try catch else finally end` is parsed with `K"catch"` `K"else"` and `K"finally"` children to avoid the awkwardness of the optional child nodes in the `Expr` representation (#234)
 * The dotted import path syntax as in `import A.b.c` is parsed with a `K"importpath"` kind rather than `K"."`, because a bare `A.b.c` has a very different nested/quoted expression representation (#244)
 * We use flags rather than child nodes to represent the difference between `struct` and `mutable struct`, `module` and `baremodule` (#220)
+* Multiple iterations within the header of a `for`, as in `for a=as, b=bs body end` are represented with a `cartesian_iterator` head rather than a `block`, as these lists of iterators are neither semantically nor syntactically a sequence of statements. Unlike other uses of `block` (see also generators).
 
 
 ## More detail on tree differences
 
-### Flattened generators
+### Generators
 
 Flattened generators are uniquely problematic because the Julia AST doesn't
 respect a key rule we normally expect: that the children of an AST node are a
-*contiguous* range in the source text. This is because the `for`s in
+*contiguous* range in the source text. For example, the `for`s in
 `[xy for x in xs for y in ys]` are parsed in the normal order of a for loop to
 mean
 
@@ -470,16 +471,30 @@ however, note that if this tree were flattened, the order would be
 source order.
 
 However, our green tree is strictly source-ordered, so we must deviate from the
-Julia AST. The natural representation seems to be to remove the generators and
-use a flattened structure:
+Julia AST. We deal with this by grouping cartesian products of iterators
+(separated by commas) within `cartesian_iterator` blocks as in `for` loops, and
+use the presence of multiple iterator blocks rather than the `flatten` head to
+distinguish flattened iterators. The nested flattens and generators of `Expr`
+forms are reconstructed later. In this form the tree structure resembles the
+source much more closely. For example, `(xy for x in xs for y in ys)` is parsed as
 
 ```
-(flatten
+(generator
   xy
   (= x xs)
   (= y ys))
 ```
 
+And the cartesian iteration `(xy for x in xs, y in ys)` is parsed as
+
+```
+(generator
+  xy
+  (cartesian_iterator
+    (= x xs)
+    (= y ys)))
+```
+
 ### Whitespace trivia inside strings
 
 For triple quoted strings, the indentation isn't part of the string data so
diff --git a/src/expr.jl b/src/expr.jl
@@ -126,41 +126,34 @@ function _to_expr(node::SyntaxNode; iteration_spec=false, need_linenodes=true,
     # Convert children
     insert_linenums = (headsym === :block || headsym === :toplevel) && need_linenodes
     args = Vector{Any}(undef, length(node_args)*(insert_linenums ? 2 : 1))
-    if headsym === :for && length(node_args) == 2
-        # No line numbers in for loop iteration spec
-        args[1] = _to_expr(node_args[1], iteration_spec=true, need_linenodes=false)
-        args[2] = _to_expr(node_args[2])
-    elseif headsym === :let && length(node_args) == 2
-        # No line numbers in let statement binding list
-        args[1] = _to_expr(node_args[1], need_linenodes=false)
-        args[2] = _to_expr(node_args[2])
+    eq_to_kw_in_call =
+        ((headsym === :call || headsym === :dotcall) && is_prefix_call(node)) ||
+        headsym === :ref
+    eq_to_kw_all = (headsym === :parameters && !map_kw_in_params) ||
+                   (headsym === :parens && eq_to_kw)
+    in_vcbr = headsym === :vect || headsym === :curly || headsym === :braces || headsym === :ref
+    if insert_linenums && isempty(node_args)
+        push!(args, source_location(LineNumberNode, node.source, node.position))
     else
-        eq_to_kw_in_call =
-            ((headsym === :call || headsym === :dotcall) && is_prefix_call(node)) ||
-            headsym === :ref
-        eq_to_kw_all = (headsym === :parameters && !map_kw_in_params) ||
-                       (headsym === :parens && eq_to_kw)
-        in_vcbr = headsym === :vect || headsym === :curly || headsym === :braces || headsym === :ref
-        if insert_linenums && isempty(node_args)
-            push!(args, source_location(LineNumberNode, node.source, node.position))
-        else
-            for i in 1:length(node_args)
-                n = node_args[i]
-                if insert_linenums
-                    args[2*i-1] = source_location(LineNumberNode, n.source, n.position)
-                end
-                eq_to_kw = eq_to_kw_in_call && i > 1 || eq_to_kw_all
-                coalesce_dot_with_ops = i==1 &&
-                    (nodekind in KSet"call dotcall curly" ||
-                     nodekind == K"quote" && flags(node) == COLON_QUOTE)
-                args[insert_linenums ? 2*i : i] =
-                    _to_expr(n, eq_to_kw=eq_to_kw,
-                             map_kw_in_params=in_vcbr,
-                             coalesce_dot=coalesce_dot_with_ops)
-            end
-            if nodekind == K"block" && has_flags(node, PARENS_FLAG)
-                popfirst!(args)
+        for i in 1:length(node_args)
+            n = node_args[i]
+            if insert_linenums
+                args[2*i-1] = source_location(LineNumberNode, n.source, n.position)
             end
+            eq_to_kw = eq_to_kw_in_call && i > 1 || eq_to_kw_all
+            coalesce_dot_with_ops = i==1 &&
+                (nodekind in KSet"call dotcall curly" ||
+                 nodekind == K"quote" && flags(node) == COLON_QUOTE)
+            args[insert_linenums ? 2*i : i] =
+                _to_expr(n, eq_to_kw=eq_to_kw,
+                         map_kw_in_params=in_vcbr,
+                         coalesce_dot=coalesce_dot_with_ops,
+                         iteration_spec=(headsym == :for && i == 1),
+                         need_linenodes=!(headsym == :let && i == 1)
+                        )
+        end
+        if nodekind == K"block" && has_flags(node, PARENS_FLAG)
+            popfirst!(args)
         end
     end
 
@@ -202,6 +195,10 @@ function _to_expr(node::SyntaxNode; iteration_spec=false, need_linenodes=true,
     elseif headsym in (:ref, :curly)
         # Move parameters blocks to args[2]
         reorder_parameters!(args, 2)
+    elseif headsym === :for
+        if Meta.isexpr(args[1], :cartesian_iterator)
+            args[1] = Expr(:block, args[1].args...)
+        end
     elseif headsym in (:tuple, :vect, :braces)
         # Move parameters blocks to args[1]
         reorder_parameters!(args, 1)
@@ -262,18 +259,32 @@ function _to_expr(node::SyntaxNode; iteration_spec=false, need_linenodes=true,
                 push!(args, else_)
             end
         end
-    elseif headsym === :filter
-        pushfirst!(args, last(args))
-        pop!(args)
-    elseif headsym === :flatten
-        # The order of nodes inside the generators in Julia's flatten AST
-        # is noncontiguous in the source text, so need to reconstruct
-        # Julia's AST here from our alternative `flatten` expression.
-        gen = Expr(:generator, args[1], args[end])
-        for i in length(args)-1:-1:2
-            gen = Expr(:flatten, Expr(:generator, gen, args[i]))
+    elseif headsym === :generator
+        # Reconstruct the nested Expr form for generator from our flatter
+        # source-ordered `generator` format.
+        gen = args[1]
+        for j = length(args):-1:2
+            if Meta.isexpr(args[j], :cartesian_iterator)
+                gen = Expr(:generator, gen, args[j].args...)
+            else
+                gen = Expr(:generator, gen, args[j])
+            end
+            if j < length(args)
+                # Additional `for`s flatten the inner generator
+                gen = Expr(:flatten, gen)
+            end
         end
         return gen
+    elseif headsym == :filter
+        @assert length(args) == 2
+        iterspec = args[1]
+        outargs = Any[args[2]]
+        if Meta.isexpr(iterspec, :cartesian_iterator)
+            append!(outargs, iterspec.args)
+        else
+            push!(outargs, iterspec)
+        end
+        args = outargs
     elseif headsym in (:nrow, :ncat)
         # For lack of a better place, the dimension argument to nrow/ncat
         # is stored in the flags
diff --git a/src/kinds.jl b/src/kinds.jl
@@ -905,7 +905,7 @@ const _kind_names =
         # Comprehensions
         "generator"
         "filter"
-        "flatten"
+        "cartesian_iterator"
         "comprehension"
         "typed_comprehension"
     "END_SYNTAX_KINDS"
diff --git a/src/parser.jl b/src/parser.jl
@@ -1783,13 +1783,9 @@ function parse_resword(ps::ParseState)
         emit(ps, mark, K"while")
     elseif word == K"for"
         # for x in xs end  ==>  (for (= x xs) (block))
-        # for x in xs, y in ys \n a \n end ==> (for (block (= x xs) (= y ys)) (block a))
+        # for x in xs, y in ys \n a \n end ==> (for (cartesian_iterator (= x xs) (= y ys)) (block a))
         bump(ps, TRIVIA_FLAG)
-        m = position(ps)
-        n_subexprs = parse_comma_separated(ps, parse_iteration_spec)
-        if n_subexprs > 1
-            emit(ps, m, K"block")
-        end
+        parse_iteration_specs(ps)
         parse_block(ps)
         bump_closing_token(ps, K"end")
         emit(ps, mark, K"for")
@@ -2625,6 +2621,16 @@ function parse_iteration_spec(ps::ParseState)
     emit(ps, mark, K"=")
 end
 
+# Parse an iteration spec, or a comma separate list of such for for loops and
+# generators
+function parse_iteration_specs(ps::ParseState)
+    mark = position(ps)
+    n_iters = parse_comma_separated(ps, parse_iteration_spec)
+    if n_iters > 1
+        emit(ps, mark, K"cartesian_iterator")
+    end
+end
+
 # flisp: parse-space-separated-exprs
 function parse_space_separated_exprs(ps::ParseState)
     ps = with_space_sensitive(ps)
@@ -2669,61 +2675,41 @@ function parse_vect(ps::ParseState, closer)
     return (K"vect", EMPTY_FLAGS)
 end
 
-# Flattened generators are hard because the Julia AST doesn't respect a key
-# rule we normally expect: that the children of an AST node are a contiguous
-# range in the source text. This is because the `for`s in
-# `[xy for x in xs for y in ys]` are parsed in the normal order of a for as
+# Parse generators
 #
-# (flatten
-#  (generator
-#   (generator
-#    xy
-#    y in ys)
-#   x in xs))
+# We represent generators quite differently from `Expr`:
+# * Cartesian products of iterators are grouped within cartesian_iterator
+#   nodes, as in the short form of `for` loops.
+# * The `generator` kind is used for both cartesian and flattened generators
 #
-# We deal with this by only emitting the flatten:
-#
-# (flatten xy (= x xs) (= y ys))
-#
-# then reconstructing the nested flattens and generators when converting to Expr.
-#
-# [x for a = as for b = bs if cond1 for c = cs if cond2] ==> (comprehension (flatten x (= a as) (filter (= b bs) cond1) (filter (= c cs) cond2)))
-# [x for a = as if begin cond2 end]  =>  (comprehension (generator x (filter (= a as) (block cond2))))
+# (x for a in as for b in bs) ==> (parens (generator x (= a as) (= b bs)))
+# (x for a in as, b in bs) ==> (parens (generator x (cartesian_iterator (= a as) (= b bs))))
+# (x for a in as, b in bs if z)  ==> (parens (generator x (filter (cartesian_iterator (= a as) (= b bs)) z)))
 #
 # flisp: parse-generator
-function parse_generator(ps::ParseState, mark, flatten=false)
-    t = peek_token(ps)
-    if !preceding_whitespace(t)
-        # [(x)for x in xs]  ==>  (comprehension (generator (parens x) (error) (= x xs)))
-        bump_invisible(ps, K"error", TRIVIA_FLAG,
-                       error="Expected space before `for` in generator")
-    end
-    @check kind(t) == K"for"
-    bump(ps, TRIVIA_FLAG)
-    filter_mark = position(ps)
-    parse_comma_separated(ps, parse_iteration_spec)
-    if peek(ps) == K"if"
-        # (a for x in xs if cond) ==> (parens (generator a (filter (= x xs) cond)))
+function parse_generator(ps::ParseState, mark)
+    while (t = peek_token(ps); kind(t) == K"for")
+        if !preceding_whitespace(t)
+            # ((x)for x in xs)  ==>  (parens (generator (parens x) (error) (= x xs)))
+            bump_invisible(ps, K"error", TRIVIA_FLAG,
+                           error="Expected space before `for` in generator")
+        end
         bump(ps, TRIVIA_FLAG)
-        parse_cond(ps)
-        emit(ps, filter_mark, K"filter")
-    end
-    t = peek_token(ps)
-    if kind(t) == K"for"
-        # (xy for x in xs for y in ys)  ==> (parens (flatten xy (= x xs) (= y ys)))
-        # (xy for x in xs for y in ys for z in zs)  ==> (parens (flatten xy (= x xs) (= y ys) (= z zs)))
-        parse_generator(ps, mark, true)
-        if !flatten
-            emit(ps, mark, K"flatten")
+        iter_mark = position(ps)
+        parse_iteration_specs(ps)
+        if peek(ps) == K"if"
+            # (x for a in as if z) ==> (parens (generator x (filter (= a as) z)))
+            bump(ps, TRIVIA_FLAG)
+            parse_cond(ps)
+            emit(ps, iter_mark, K"filter")
         end
-    elseif !flatten
-        # (x for a in as)  ==>  (parens (generator x (= a as)))
-        emit(ps, mark, K"generator")
     end
+    emit(ps, mark, K"generator")
 end
 
 # flisp: parse-comprehension
 function parse_comprehension(ps::ParseState, mark, closer)
+    # [x for a in as] ==> (comprehension (generator x a in as))
     ps = ParseState(ps, whitespace_newline=true,
                     space_sensitive=false,
                     end_symbol=false)
diff --git a/test/expr.jl b/test/expr.jl
diff --git a/test/parser.jl b/test/parser.jl
