fix: type inferance in generate loop

Author: wangl-cc

src/tools.jl

@@ -419,7 +419,7 @@ function gen_loop(inds, inds_ax, fn_args, lhs, rhs, opts...)
     end
     ret_name = lhs.args[1]
     ret_size = Expr[]
-    ret_type = _infer_ret_type(inds, fn_args, rhs)
+    ret_type = _infer_ret_type(inds, inds_ax, fn_args, rhs)
     loop_head = Expr(:block)
     for (i, ind) in enumerate(inds)
         array_name, dims = first(inds_ax[ind])
@@ -480,16 +480,16 @@ function _parse_opt(opt::Expr)
     end
 end
 
-function _infer_ret_type(inds, fn_args, ex)
-    ex_replace = MacroTools.postwalk(ex) do sym
-        if sym isa Symbol && sym in inds
-            return :begin # replace the iteration symbol with first index
-        else
-            return sym
-        end
+function _infer_ret_type(inds, inds_ax, fn_args, ex)
+    ex′ = Expr(:block, ex)
+    for ind in inds
+        array_name, dims = first(inds_ax[ind])
+        dim = first(dims)
+        assignment = :($ind = firstindex($array_name, $dim))
+        pushfirst!(ex′.args, assignment)
     end
     args = collect(fn_args)
-    promote_op_fn = Expr(:(->), Expr(:tuple, args...), ex_replace)
+    promote_op_fn = Expr(:(->), Expr(:tuple, args...), ex′)
     args_type = map(arg -> :(typeof($arg)), args)
     ret_type = :(Base.promote_op($promote_op_fn, $(args_type...)))
     return ret_type # expression for inferring the return type

test/tools.jl

@@ -16,11 +16,16 @@ end
     # those functions is sensitive to the order of arguments, which is undetermined
     f1 = @quickloop R1[i, j] := A[i] * B[j]
     f2 = @quickloop R2[j, i] := A[i] * B[j]
-    f3 = @quickloop A[i, j] # just make a copy
+    f3 = @quickloop ifelse(i == j, zero(eltype(A)), A[i, j]) # make diagonal zero
     f4 = @quickloop R3[i, j, k] := A[i, j] * B[i, k]
     @test (f1(A, B) == A .* B' == transpose(f2(A, B)) ||
            f1(A, B) == B .* A' == transpose(f2(A, B)))
-    @test C == f3(C)
+    F = f3(C)
+    equality_f3 = true
+    for i in 1:4, j in 1:4
+        equality_f3 &= ifelse(i == j, F[i, j] == 0, F[i, j] == C[i, j])
+    end
+    @test equality_f3
     E = f4(C, D)
     equality_f4 = true
     for k in 1:4, j in 1:4, i in 1:4