@nLoops: support pre/post expressions

src/Cartesian.jl

@@ -40,21 +40,29 @@ function cartesian_linear(A::AbstractArray, I::Vector{Int})
 end
 
 # Generate nested loops
-macro nloops(N, itersym, rangeexpr, ex)
-    _nloops(N, itersym, rangeexpr, ex)
+macro nloops(N, itersym, rangeexpr, args...)
+    _nloops(N, itersym, rangeexpr, args...)
 end
 
 # Range of each loop is determined by an array's size,
 #   for i2 = 1:size(A,2)
 #     for i1 = 1:size(A,1)
 #       ...
-function _nloops(N::Int, itersym::Symbol, arrayforsize::Symbol, body::Expr)
+function _nloops(N::Int, itersym::Symbol, arrayforsize::Symbol, args::Expr...)
+    if !(1 <= length(args) <= 3)
+        error("Too many arguments")
+    end
+    body = args[end]
     ex = Expr(:escape, body)
     for dim = 1:N
         itervar = namedvar(itersym, dim)
+        preexpr = length(args) > 1 ? inlineanonymous(args[1], dim) : (:(nothing))
+        postexpr = length(args) > 2 ? inlineanonymous(args[2], dim) : (:(nothing))
         ex = quote
             for $(esc(itervar)) = 1:size($(esc(arrayforsize)),$dim)
+                $(esc(preexpr))
                 $ex
+                $(esc(postexpr))
             end
         end
     end
@@ -67,17 +75,25 @@ end
 # anonymous function expression.
 # It's possible to make the range depend on a set of indexed variables, using the
 # notation i_d which gets translated into i3 for d=3.
-function _nloops(N::Int, itersym::Symbol, rangeexpr::Expr, body::Expr)
+function _nloops(N::Int, itersym::Symbol, rangeexpr::Expr, args::Expr...)
     if rangeexpr.head != :->
         error("Second argument must be an anonymous function expression to compute the range")
     end
+    if !(1 <= length(args) <= 3)
+        error("Too many arguments")
+    end
+    body = args[end]
     ex = Expr(:escape, body)
     for dim = 1:N
         itervar = namedvar(itersym, dim)
         rng = inlineanonymous(rangeexpr, dim)
+        preexpr = length(args) > 1 ? inlineanonymous(args[1], dim) : (:(nothing))
+        postexpr = length(args) > 2 ? inlineanonymous(args[2], dim) : (:(nothing))
         ex = quote
             for $(esc(itervar)) = $(esc(rng))
+                $(esc(preexpr))
                 $ex
+                $(esc(postexpr))
             end
         end
     end

test/testloops.jl

@@ -67,6 +67,29 @@ t = @elapsed loopsum(S)
 t = @elapsed loopsum2(S)
 @assert t < 3tbase
 
+# @nloops with pre/post expression
+for N = 1:4
+    @eval begin
+        function maxoverdims{T}(A::AbstractArray{T,$N}, region)
+            szout = [size(A,d) for d = 1:$N]
+            szout[[region...]] = 1
+            B = fill(typemin(T), szout...)::Array{T,$N}
+            Cartesian.@nextract $N szout szout
+            Cartesian.@nloops $N i A d->(j_d = szout_d==1 ? 1 : i_d) begin
+                (Cartesian.@nref $N B j) = max((Cartesian.@nref $N B j), (Cartesian.@nref $N A i))
+            end
+            B
+        end
+    end
+end
+
+A = reshape(1:10,5,2)
+A1 = maxoverdims(A, 1)
+@assert A1 == [5,10]'
+A2 = maxoverdims(A, 2)
+@assert A2 == reshape(6:10,5,1)
+
+
 # @nref, @nrefshift, @nextract, and @nlookup
 A = reshape(1:15, 3, 5)
 i1 = 2
@@ -111,11 +134,10 @@ end
 @assert pairs == {(1,2),(2,2),(1,3),(2,3),(1,4),(2,4)}
 
 # @nexprs
-println("About to check nexprs")
 A = reshape(1:20*7, 20, 7)
 indexes = (2:5:20,3:7)
 strds = strides(A)
 i1 = 2
 i2 = 3
 ind = 1
-@assert (Cartesian.@nexprs 2 d->(ind += (indexes[d][i_d]-1)*strds[d])) == A[indexes[1][i1],indexes[2][i2]])
+@assert (Cartesian.@nexprs 2 d->(ind += (indexes[d][i_d]-1)*strds[d])) == A[indexes[1][i1],indexes[2][i2]]