staged ind2sub2ind

Author: Jutho

base/abstractarray.jl

@@ -1031,28 +1031,16 @@ function repmat(a::AbstractVector, m::Int)
 end
 
 sub2ind(dims) = 1
-sub2ind(dims, i::Integer) = int(i)
-sub2ind(dims, i::Integer, j::Integer) = sub2ind(dims, int(i), int(j))
-sub2ind(dims, i::Int, j::Int) = (j-1)*dims[1] + i
-sub2ind(dims, i0::Integer, i1::Integer, i2::Integer) = sub2ind(dims, int(i0),int(i1),int(i2))
-sub2ind(dims, i0::Int, i1::Int, i2::Int) =
-    i0 + dims[1]*((i1-1) + dims[2]*(i2-1))
-sub2ind(dims, i0::Integer, i1::Integer, i2::Integer, i3::Integer) =
-    sub2ind(dims, int(i0),int(i1),int(i2),int(i3))
-sub2ind(dims, i0::Int, i1::Int, i2::Int, i3::Int) =
-    i0 + dims[1]*((i1-1) + dims[2]*((i2-1) + dims[3]*(i3-1)))
-
-function sub2ind(dims, I::Integer...)
-    ndims = length(dims)
-    index = int(I[1])
-    stride = 1
-    for k=2:ndims
-        stride = stride * dims[k-1]
-        index += (int(I[k])-1) * stride
+stagedfunction sub2ind{N}(dims::NTuple{N,Integer},I::Integer...)
+    length(I) == N || throw(ArgumentError("number of subscripts should match length of dims"))
+    ex=:(I[$N])
+    for k=N-1:-1:1
+        ex=:(I[$k]+dims[$k]*($ex-1))
     end
-    return index
+    :(Int($ex))
 end
 
+
 function sub2ind{T<:Integer}(dims::Array{T}, sub::Array{T})
     ndims = length(dims)
     ind = sub[1]
@@ -1067,31 +1055,18 @@ end
 sub2ind{T<:Integer}(dims, I::AbstractVector{T}...) =
     [ sub2ind(dims, map(X->X[i], I)...)::Int for i=1:length(I[1]) ]
 
-function ind2sub(dims::(Integer,Integer...), ind::Int)
-    ndims = length(dims)
-    stride = dims[1]
-    for i=2:ndims-1
-        stride *= dims[i]
-    end
 
-    sub = ()
-    for i=(ndims-1):-1:1
-        rest = rem(ind-1, stride) + 1
-        sub = tuple(div(ind - rest, stride) + 1, sub...)
-        ind = rest
-        stride = div(stride, dims[i])
+ind2sub(dims::(), ind::Integer) = ind==1 ? () : throw(BoundsError())
+stagedfunction ind2sub{N}(dims::NTuple{N,Integer}, ind::Int)
+    subsyms = [symbol("i$k") for k=1:N]
+    subex = Expr(:block,[:($(subsyms[k])=rem(ind,dims[$k])+1;ind=div(ind,dims[$k])) for k=1:N-1]...)
+    quote
+        ind-=1
+        $(subex)
+        $(subsyms[N]) = ind+1
+        $(Expr(:tuple,subsyms...))
     end
-    return tuple(ind, sub...)
 end
-
-ind2sub(dims::(Integer...), ind::Integer) = ind2sub(dims, int(ind))
-ind2sub(dims::(), ind::Integer) = ind==1 ? () : throw(BoundsError())
-ind2sub(dims::(Integer,), ind::Int) = (ind,)
-ind2sub(dims::(Integer,Integer), ind::Int) =
-    (rem(ind-1,dims[1])+1, div(ind-1,dims[1])+1)
-ind2sub(dims::(Integer,Integer,Integer), ind::Int) =
-    (rem(ind-1,dims[1])+1, div(rem(ind-1,dims[1]*dims[2]), dims[1])+1,
-     div(rem(ind-1,dims[1]*dims[2]*dims[3]), dims[1]*dims[2])+1)
 ind2sub(a::AbstractArray, ind::Integer) = ind2sub(size(a), Int(ind))
 
 function ind2sub{T<:Integer}(dims::(Integer,Integer...), ind::AbstractVector{T})
@@ -1499,4 +1474,3 @@ function randsubseq!(S::AbstractArray, A::AbstractArray, p::Real)
 end
 
 randsubseq{T}(A::AbstractArray{T}, p::Real) = randsubseq!(T[], A, p)
-