Fix generators with multiple `for` nodes

src/runestone.jl

@@ -1127,19 +1127,25 @@ end
 function for_loop_use_in(ctx::Context, node::Node)
     if !(
             (kind(node) === K"for" && !is_leaf(node) && meta_nargs(node) == 4) ||
-                (kind(node) === K"generator" && meta_nargs(node) == 3) # TODO: Unsure about 3.
+                kind(node) === K"generator"
         )
         return nothing
     end
     pos = position(ctx.fmt_io)
     kids = verified_kids(node)
+    kids′ = kids
     for_index = findfirst(c -> kind(c) === K"for" && is_leaf(c), kids)::Int
+    next_index = 1
+    any_for_changed = false
+    # generator can have multiple for nodes
+    while for_index !== nothing
         for_node = kids[for_index]
         @assert kind(for_node) === K"for" && span(for_node) == 3 &&
             is_leaf(for_node) && JuliaSyntax.is_trivia(for_node)
-    for i in 1:for_index
+        for i in next_index:for_index
             accept_node!(ctx, kids[i])
         end
+        while_pos = position(ctx.fmt_io)
         # The for loop specification node can be either K"=" or K"cartesian_iterator"
         for_spec_index = for_index + 1
         for_spec_node = kids[for_spec_index]
@@ -1152,17 +1158,28 @@ function for_loop_use_in(ctx::Context, node::Node)
             @assert kind(for_spec_node) === K"cartesian_iterator"
             for_spec_node′ = replace_with_in_cartesian(ctx, for_spec_node)
         end
-    if for_spec_node′ === nothing
+        if for_spec_node′ !== nothing
-        seek(ctx.fmt_io, pos)
+            @assert position(ctx.fmt_io) == while_pos + span(for_spec_node′)
-        return nothing
+            any_for_changed = true
-    end
-    @assert position(ctx.fmt_io) == pos + mapreduce(span, +, @view(kids[1:for_index])) + span(for_spec_node′)
             # Insert the new for spec node
+            if kids′ === kids
                 kids′ = copy(kids)
+            end
             kids′[for_spec_index] = for_spec_node′
-    # At this point the eq node is done, just accept any remaining nodes
+        end
+        for_index = findnext(c -> kind(c) === K"for" && is_leaf(c), kids, for_spec_index + 1)
+        if for_index !== nothing
+            @assert kind(node) === K"generator"
+        end
+        next_index = for_spec_index + 1
+    end
+    if !any_for_changed
+        seek(ctx.fmt_io, pos)
+        return nothing
+    end
+    # At this point the eq nodes are done, just accept any remaining nodes
     # TODO: Don't need to do this...
-    for i in (for_spec_index + 1):length(kids′)
+    for i in next_index:length(kids′)
         accept_node!(ctx, kids′[i])
     end
     # Construct the full node and return

test/runtests.jl

@@ -426,10 +426,16 @@ end
         # for generators, filter
         for (l, r) in (("[", "]"), ("(", ")"))
             @test format_string("$(l)i for i$(sp)$(op)$(sp)I$(r)") == "$(l)i for i in I$(r)"
+            # cartesian
             @test format_string("$(l)(i, j) for i$(sp)$(op)$(sp)I, j$(sp)$(op)$(sp)J$(r)") ==
                 "$(l)(i, j) for i in I, j in J$(r)"
             @test format_string("$(l)(i, j, k) for i$(sp)$(op)$(sp)I, j$(sp)$(op)$(sp)J, k$(sp)$(op)$(sp)K$(r)") ==
                 "$(l)(i, j, k) for i in I, j in J, k in K$(r)"
+            # multiple for
+            @test format_string("$(l)(i, j) for i$(sp)$(op)$(sp)I for j$(sp)$(op)$(sp)J$(r)") ==
+                "$(l)(i, j) for i in I for j in J$(r)"
+            @test format_string("$(l)(i, j, k) for i$(sp)$(op)$(sp)I for j$(sp)$(op)$(sp)J for k$(sp)$(op)$(sp)K$(r)") ==
+                "$(l)(i, j, k) for i in I for j in J for k in K$(r)"
         end
         # K"filter"
         for (l, r) in (("[", "]"), ("(", ")"))