Refactor reaxis to be slightly less generated (#79)

* Refactor reaxis to be slightly less generated

Mostly use dispatch now, except in cases where we change axis names.

* Add tests

Test indexing beyond the dimensionality and adding dimensions. Add some inferred tests.

* 0.5 compatibility

Author: mbauman

src/combine.jl

@@ -122,7 +122,7 @@ Combines AxisArrays with matching axis names into a single AxisArray. Unlike `me
 If an array value in the output array is not defined in any of the input arrays (i.e. in the case of a left, right, or outer join), it takes the value of the optional `fillvalue` keyword argument (default zero).
 """
 function Base.join{T,N,D,Ax}(As::AxisArray{T,N,D,Ax}...; fillvalue::T=zero(T),
-                             newaxis::Axis=_nextaxistype(As[1].data, As[1].axes)(1:length(As)),
+                             newaxis::Axis=_nextaxistype(As[1].axes)(1:length(As)),
                              method::Symbol=:outer)
 
     prejoin_resultaxes = map(as -> axismerge(method, as...), map(tuple, axes.(As)...))

src/core.jl

@@ -166,9 +166,9 @@ _defaultdimname(i) = i == 1 ? (:row) : i == 2 ? (:col) : i == 3 ? (:page) : Symb
 default_axes(A::AbstractArray) = _default_axes(A, indices(A), ())
 _default_axes{T,N}(A::AbstractArray{T,N}, inds, axs::NTuple{N,Axis}) = axs
 @inline _default_axes{T,N,M}(A::AbstractArray{T,N}, inds, axs::NTuple{M,Axis}) =
-    _default_axes(A, inds, (axs..., _nextaxistype(A, axs)(inds[M+1])))
+    _default_axes(A, inds, (axs..., _nextaxistype(axs)(inds[M+1])))
 # Why doesn't @pure work here?
-@generated function _nextaxistype{T,M}(A::AbstractArray{T}, axs::NTuple{M,Axis})
+@generated function _nextaxistype{M}(axs::NTuple{M,Axis})
     name = _defaultdimname(M+1)
     :(Axis{$(Expr(:quote, name))})
 end

src/indexing.jl

@@ -1,6 +1,6 @@
 const Idx = Union{Real,Colon,AbstractArray{Int}}
 
-using Base: ViewIndex, @propagate_inbounds
+using Base: ViewIndex, @propagate_inbounds, tail
 
 # Defer IndexStyle to the wrapped array
 @compat Base.IndexStyle{T,N,D,Ax}(::Type{AxisArray{T,N,D,Ax}}) = IndexStyle(D)
@@ -12,46 +12,58 @@ using Base: ViewIndex, @propagate_inbounds
 # Cartesian iteration
 Base.eachindex(A::AxisArray) = eachindex(A.data)
 
-@generated function reaxis(A::AxisArray, I::Idx...)
-    N = length(I)
-    # Determine the new axes:
-    # Drop linear indexing over multiple axes
-    droplastaxis = ndims(A) > N && !(I[end] <: Real) ? 1 : 0
-    # Drop trailing scalar dimensions
-    lastnonscalar = N
-    while lastnonscalar > 0 && I[lastnonscalar] <: Real
-        lastnonscalar -= 1
+"""
+    reaxis(A::AxisArray, I...)
+
+This internal function determines the new set of axes that are constructed upon
+indexing with I.
+"""
+reaxis(A::AxisArray, I::Idx...) = _reaxis(make_axes_match(axes(A), I), I)
+# Ensure the number of axes matches the number of indexing dimensions
+@inline make_axes_match(axs, idxs) = _make_axes_match((), axs, Base.index_ndims(idxs...))
+# Move the axes into newaxes, until we run out of both simultaneously
+@inline _make_axes_match(newaxes, axs::Tuple, nidxs::Tuple) =
+    _make_axes_match((newaxes..., axs[1]), tail(axs), tail(nidxs))
+@inline _make_axes_match(newaxes, axs::Tuple{}, nidxs::Tuple{}) = newaxes
+# Drop trailing axes, replacing it with a default name for the linear span
+@inline _make_axes_match(newaxes, axs::Tuple, nidxs::Tuple{}) =
+    (maybefront(newaxes)..., _nextaxistype(newaxes)(Base.OneTo(length(newaxes[end]) * prod(map(length, axs)))))
+# Insert phony singleton trailing axes
+@inline _make_axes_match(newaxes, axs::Tuple{}, nidxs::Tuple) =
+    _make_axes_match((newaxes..., _nextaxistype(newaxes)(Base.OneTo(1))), (), tail(nidxs))
+
+@inline maybefront(::Tuple{}) = ()
+@inline maybefront(t::Tuple) = Base.front(t)
+
+# Now we can reaxis without worrying about mismatched axes/indices
+@inline _reaxis(axs::Tuple{}, idxs::Tuple{}) = ()
+# Scalars are dropped
+const ScalarIndex = @compat Union{Real, AbstractArray{<:Any, 0}}
+@inline _reaxis(axs::Tuple, idxs::Tuple{ScalarIndex, Vararg{Any}}) = _reaxis(tail(axs), tail(idxs))
+# Colon passes straight through
+@inline _reaxis(axs::Tuple, idxs::Tuple{Colon, Vararg{Any}}) = (axs[1], _reaxis(tail(axs), tail(idxs))...)
+# But arrays can add or change dimensions and accompanying axis names
+@inline _reaxis(axs::Tuple, idxs::Tuple{AbstractArray, Vararg{Any}}) =
+    (_new_axes(axs[1], idxs[1])..., _reaxis(tail(axs), tail(idxs))...)
+
+# Vectors simply create new axes with the same name; just subsetted by their value
+@inline _new_axes{name}(ax::Axis{name}, idx::AbstractVector) = (Axis{name}(ax.val[idx]),)
+# Arrays create multiple axes with _N appended to the axis name containing their indices
+@generated function _new_axes{name, N}(ax::Axis{name}, idx::@compat(AbstractArray{<:Any,N}))
+    newaxes = Expr(:tuple)
+    for i=1:N
+        push!(newaxes.args, :($(Axis{Symbol(name, "_", i)})(indices(idx, $i))))
     end
-    names = axisnames(A)
-    newaxes = Expr[]
-    drange = 1:lastnonscalar-droplastaxis
-    for d=drange
-        if I[d] <: AxisArray
-            # Indexing with an AxisArray joins the axis names
-            idxnames = axisnames(I[d])
-            for i=1:ndims(I[d])
-                push!(newaxes, :($(Axis{Symbol(names[d], "_", idxnames[i])})(I[$d].axes[$i].val)))
-            end
-        elseif I[d] <: Real
-        elseif I[d] <: AbstractVector
-            push!(newaxes, :($(Axis{names[d]})(A.axes[$d].val[Base.to_index(I[$d])])))
-        elseif I[d] <: Colon
-            if d < length(I) || d <= ndims(A)
-                push!(newaxes, :($(Axis{names[d]})(A.axes[$d].val)))
-            else
-                dimname = _defaultdimname(d)
-                push!(newaxes, :($(Axis{dimname})(Base.OneTo(Base.trailingsize(A, $d)))))
-            end
-        elseif I[d] <: AbstractArray
-            for i=1:ndims(I[d])
-                # When we index with non-vector arrays, we *add* dimensions.
-                push!(newaxes, :($(Axis{Symbol(names[d], "_", i)})(indices(I[$d], $i))))
-            end
-        end
-    end
-    quote
-        ($(newaxes...),)
+    newaxes
+end
+# And indexing with an AxisArray joins the name and overrides the values
+@generated function _new_axes{name, N}(ax::Axis{name}, idx::@compat(AxisArray{<:Any, N}))
+    newaxes = Expr(:tuple)
+    idxnames = axisnames(idx)
+    for i=1:N
+        push!(newaxes.args, :($(Axis{Symbol(name, "_", idxnames[i])})(idx.axes[$i].val)))
     end
+    newaxes
 end
 
 @propagate_inbounds function Base.getindex(A::AxisArray, idxs::Idx...)

test/indexing.jl

@@ -1,16 +1,17 @@
 A = AxisArray(reshape(1:24, 2,3,4), .1:.1:.2, .1:.1:.3, .1:.1:.4)
 D = similar(A)
 D[1,1,1,1,1] = 10
-@test D[1,1,1,1,1] == D[1] == D.data[1] == 10
+@test @inferred(D[1,1,1,1,1]) == @inferred(D[1]) == D.data[1] == 10
+@test @inferred(D[1,1,1,:]) == @inferred(D[1,1,1,1:1]) == @inferred(D[1,1,1,[1]]) == AxisArray([10], Axis{:dim_4}(Base.OneTo(1)))
 
 # Test slices
 
 @test A == A.data
 @test A[:,:,:] == A[Axis{:row}(:)] == A[Axis{:col}(:)] == A[Axis{:page}(:)] == A.data[:,:,:]
 # Test UnitRange slices
-@test A[1:2,:,:] == A.data[1:2,:,:] == A[Axis{:row}(1:2)]  == A[Axis{1}(1:2)] == A[Axis{:row}(ClosedInterval(-Inf,Inf))] == A[[true,true],:,:]
-@test @view(A[1:2,:,:]) == A.data[1:2,:,:] == @view(A[Axis{:row}(1:2)]) == @view(A[Axis{1}(1:2)]) == @view(A[Axis{:row}(ClosedInterval(-Inf,Inf))]) == @view(A[[true,true],:,:])
-@test A[:,1:2,:] == A.data[:,1:2,:] == A[Axis{:col}(1:2)]  == A[Axis{2}(1:2)] == A[Axis{:col}(ClosedInterval(0.0, .25))] == A[:,[true,true,false],:]
+@test @inferred(A[1:2,:,:]) == A.data[1:2,:,:] == @inferred(A[Axis{:row}(1:2)])  == @inferred(A[Axis{1}(1:2)]) == @inferred(A[Axis{:row}(ClosedInterval(-Inf,Inf))]) == @inferred(A[[true,true],:,:])
+@test @inferred(view(A,1:2,:,:)) == A.data[1:2,:,:] == @inferred(view(A,Axis{:row}(1:2))) == @inferred(view(A,Axis{1}(1:2))) == @inferred(view(A,Axis{:row}(ClosedInterval(-Inf,Inf)))) == @inferred(view(A,[true,true],:,:))
+@test @inferred(A[:,1:2,:]) == A.data[:,1:2,:] == @inferred(A[Axis{:col}(1:2)])  == @inferred(A[Axis{2}(1:2)]) == @inferred(A[Axis{:col}(ClosedInterval(0.0, .25))]) == @inferred(A[:,[true,true,false],:])
 @test @view(A[:,1:2,:]) == A.data[:,1:2,:] == @view(A[Axis{:col}(1:2)])  == @view(A[Axis{2}(1:2)]) == @view(A[Axis{:col}(ClosedInterval(0.0, .25))]) == @view(A[:,[true,true,false],:])
 @test A[:,:,1:2] == A.data[:,:,1:2] == A[Axis{:page}(1:2)] == A[Axis{3}(1:2)] == A[Axis{:page}(ClosedInterval(-1., .22))] == A[:,:,[true,true,false,false]]
 @test @view(A[:,:,1:2]) == @view(A.data[:,:,1:2]) == @view(A[Axis{:page}(1:2)]) == @view(A[Axis{3}(1:2)]) == @view(A[Axis{:page}(ClosedInterval(-1., .22))]) == @view(A[:,:,[true,true,false,false]])