fix: name iter and disable single itersym

Author: wangl-cc

src/constructors.jl

@@ -78,8 +78,9 @@ end
 
 
 function LoopSet(q::Expr, mod::Symbol = :Main)
-  contract_pass!(q)
   ls = LoopSet(mod)
+  check_inputs!(q, ls.prepreamble)
+  contract_pass!(q)
   copyto!(ls, q)
   resize!(ls.loop_order, num_loops(ls))
   ls
@@ -158,32 +159,32 @@ function process_args(
 end
 # check if the body of loop is a block, if not convert it to a block issue#395
 # and check if the range of loop is an enumerate, if it is replace it, issue#393
-function check_inputs!(q)
+function check_inputs!(q, prepreamble)
   if Meta.isexpr(q, :for)
     if !Meta.isexpr(q.args[2], :block)
       q.args[2] = Expr(:block, q.args[2])
-      replace_enumerate!(q) # must after warp block
+      replace_enumerate!(q, prepreamble) # must after warp block
     else # maybe inner loops in block
-      replace_enumerate!(q)
+      replace_enumerate!(q, prepreamble)
       for arg in q.args[2].args
-        check_inputs!(arg) # check recursively for inner loop
+        check_inputs!(arg, prepreamble) # check recursively for inner loop
       end
     end
   end
   return q
 end
-function replace_enumerate!(q)
+function replace_enumerate!(q, prepreamble)
   looprange = q.args[1]
   if Meta.isexpr(looprange, :block)
     for i in 1:length(looprange.args)
-      convert_single_enumerate!(q, i)
+      replace_single_enumerate!(q, prepreamble, i)
     end
   else
-    convert_single_enumerate!(q)
+    replace_single_enumerate!(q, prepreamble)
   end
   return q
 end
-function convert_single_enumerate!(q, i=nothing)
+function replace_single_enumerate!(q, prepreamble, i=nothing)
   if isnothing(i) # not nest loop
     looprange, body = q.args[1], q.args[2]
   else # nest loop
@@ -192,7 +193,13 @@ function convert_single_enumerate!(q, i=nothing)
   @assert Meta.isexpr(looprange, :(=), 2)
   itersyms, r = looprange.args
   if Meta.isexpr(r, :call, 2) && r.args[1] == :enumerate
-    iter = r.args[2]
+    _iter = r.args[2]
+    if _iter isa Symbol
+      iter = _iter
+    else # name complex expr
+      iter = gensym(:iter)
+      push!(prepreamble.args, :($iter = $_iter))
+    end
     if Meta.isexpr(itersyms, :tuple, 2)
       indsym, varsym = itersyms.args[1]::Symbol, itersyms.args[2]::Symbol
       _replace_looprange!(q, i, indsym, iter)
@@ -201,14 +208,8 @@ function convert_single_enumerate!(q, i=nothing)
       indsym = itersyms.args[1]::Symbol
       _replace_looprange!(q, i, indsym, iter)
     elseif itersyms isa Symbol # if itersyms are not unbox in loop range
-      # generate new symbols to avoid name conflict
-      indsym = gensym(Symbol(itersyms, :_ind))
-      varsym = gensym(Symbol(itersyms, :_var))
-      _replace_looprange!(q, i, indsym, iter)
-      pushfirst!(body.args,
-        :($varsym = $iter[$indsym + firstindex($iter) - 1]),
-        :($itersyms = ($indsym, $varsym)), # regroud the indsym and varsym for user
-      )
+      throw(ArgumentError("`for $itersyms in enumerate($r)` is not supported,
+        please use `for ($(itersyms)_i, $(itersyms)_v) in enumerate($r)` instead."))
     else
       throw(ArgumentError("Don't know how to handle expression `$itersyms`."))
     end
@@ -221,7 +222,6 @@ _replace_looprange!(q, i::Int, indsym, iter) = q.args[1].args[i] = :($indsym = B
 function turbo_macro(mod, src, q, args...)
   q = macroexpand(mod, q)
   if q.head === :for
-    check_inputs!(q)
     ls = LoopSet(q, mod)
     inline, check_empty, u₁, u₂, v, threads, warncheckarg = process_args(args)
     esc(setup_call(ls, q, src, inline, check_empty, u₁, u₂, v, threads, warncheckarg))

test/parsing_inputs.jl

@@ -1,4 +1,5 @@
 using LoopVectorization, Test, ArrayInterface
+using LoopVectorization: check_inputs!
 
 # macros for generate loops whose body is not a block
 macro gen_loop_issue395(ex)
@@ -52,18 +53,27 @@ end
 @testset "enumerate, #393" begin
   A = zeros(4)
   B = zeros(4)
-  C = zeros(4)
-  @turbo for (i, x) in enumerate(A)
-    A[i] = i + x
+  C = zeros(4, 4)
+  D = zeros(4, 4)
+  @turbo for (i, x) in enumerate(1:4)
+    A[i] = x
   end
   @turbo for (i,) in enumerate(B)
     B[i] += 1
   end
-  @turbo for ix in enumerate(C)
-    C[ix[1]] = ix[1] + ix[2]
+  @turbo for (j, Aj) in enumerate(A), (i, Bi) in enumerate(B)
+    C[i, j] = Aj * Bi
+  end
+  @turbo for (j, Bj) in enumerate(B)
+    for (i, Ai) in enumerate(A)
+      D[i, j] = Ai * Bj
+    end
   end
-  @test_throws ArgumentError @turbo for () in enumerate(A) end
   @test A == 1:4
-  @test B == 1:4
-  @test C == 1:4
+  @test B == ones(4)
+  @test A .* B' == C' == D
+  @test_throws ArgumentError check_inputs!(:(for ix in enumerate(A)
+    A[ix[1]] = ix[1] + ix[2]
+  end), Any[])
+  @test_throws ArgumentError check_inputs!(:(for () in enumerate(A); end), Any[])
 end