type stable strides + inferred test

JuliaLang · May 25, 2016 · ebe4c24 · ebe4c24
1 parent 9afb230
commit ebe4c24
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 15 deletions.
diff --git a/base/abstractarray.jl b/base/abstractarray.jl
@@ -38,18 +38,12 @@ function first(itr)
 end
 last(a) = a[end]
 
-function stride(a::AbstractArray, i::Integer)
-    if i > ndims(a)
-        return length(a)
-    end
-    s = 1
-    for n=1:(i-1)
-        s *= size(a, n)
-    end
-    return s
-end
+strides{T}(a::AbstractArray{T,0}) = ()
+strides(a::AbstractArray) = _strides(a,(1,))
+_strides{_,N}(a::AbstractArray{_,N}, s::NTuple{N,Int}) = s
+_strides{_,N,M}(a::AbstractArray{_,N}, s::NTuple{M,Int}) = _strides(a, (s..., s[M]*size(a, M)))
 
-strides(a::AbstractArray) = ntuple(i->stride(a,i), ndims(a))::Dims
+stride(a::AbstractArray, d) = d <= 1 ? 1 : size(a, d-1)*stride(a, d-1)
 
 function isassigned(a::AbstractArray, i::Int...)
     # TODO

diff --git a/base/array.jl b/base/array.jl
@@ -24,10 +24,6 @@ length(a::Array) = arraylen(a)
 elsize{T}(a::Array{T}) = isbits(T) ? sizeof(T) : sizeof(Ptr)
 sizeof(a::Array) = elsize(a) * length(a)
 
-strides{T}(a::Array{T,1}) = (1,)
-strides{T}(a::Array{T,2}) = (1, size(a,1))
-strides{T}(a::Array{T,3}) = (1, size(a,1), size(a,1)*size(a,2))
-
 function isassigned{T}(a::Array{T}, i::Int...)
     ii = sub2ind(size(a), i...)
     1 <= ii <= length(a) || return false

diff --git a/test/abstractarray.jl b/test/abstractarray.jl
@@ -229,6 +229,7 @@ function test_primitives{T}(::Type{T}, shape, ::Type{TestAbstractArray})
     @test last(B) == B[length(B)]
 
     # strides(a::AbstractArray)
+    @inferred strides(B)
     strides_B = strides(B)
     for (i, _stride) in enumerate(collect(strides_B))
         @test _stride == stride(B, i)