WIP: adds AxisArray setindex! methods

Requires JuliaLang/julia#10331. But I think it is also running into a MAX_TUPLE_LEN issue, which is messing up my workaround for JuliaLang/julia#10191.

Needs tests, too.

Author: mbauman

src/indexing.jl

@@ -22,10 +22,12 @@ Base.getindex{T}(A::AxisArray{T,1}, idx::Colon) = A
 
 # Linear indexing with an array
 Base.getindex{S<:Int}(A::AxisArray, idx::AbstractArray{S}) = A.data[idx]
+Base.setindex!{S<:Int}(A::AxisArray, v, idx::AbstractArray{S}) = (A.data[idx] = v)
 
 # Cartesian iteration
 Base.eachindex(A::AxisArray) = eachindex(A.data)
 Base.getindex(A::AxisArray, idx::Base.IteratorsMD.CartesianIndex) = A.data[idx]
+Base.setindex!(A::AxisArray, v, idx::Base.IteratorsMD.CartesianIndex) = (A.data[idx] = v)
 
 # More complicated cases where we must create a subindexed AxisArray
 # TODO: do we want to be dogmatic about using views? For the data? For the axes?
@@ -51,15 +53,38 @@ stagedfunction Base.getindex{T,N,D,Ax}(A::AxisArray{T,N,D,Ax}, idxs::Idx...)
         AxisArray(data, $axes) # TODO: avoid checking the axes here
     end
 end
+# Setindex is so much simpler. Just assign it to the data:
+Base.setindex!(A::AxisArray, v, idxs::Idx...) = (A.data[idxs...] = v)
 
 ### Fancier indexing capabilities provided only by AxisArrays ###
 
-# First is the ability to index by named axis.
+# Defining the fallbacks on get/setindex are tricky due to ambiguities with 
+# AbstractArray definitions... but they simply punt to to_index to convert the
+# special indexing forms to integers and integer ranges.
+# Even though all these splats look scary, they get inlined and don't allocate.
+Base.getindex(A::AxisArray, idx::AbstractArray) = A[to_index(A,idx)...]
+Base.setindex!(A::AxisArray, v, idx::AbstractArray) = (A[to_index(A,idx)...] = v)
+let rargs = Expr[], aargs = Expr[], idxs = Symbol[]
+    for i = 1:4
+        isym = symbol("i$i")
+        push!(rargs, :($isym::Real))
+        push!(aargs, :($isym::Any))
+        push!(idxs, isym)
+        @eval Base.getindex(A::AxisArray, $(rargs...)) = A[to_index(A,$(idxs...))...]
+        @eval Base.setindex!(A::AxisArray, v, $(rargs...)) = (A[to_index(A,$(idxs...))...] = v)
+        @eval Base.getindex(A::AxisArray, $(aargs...)) = A[to_index(A,$(idxs...))...]
+        @eval Base.setindex!(A::AxisArray, v, $(aargs...)) = (A[to_index(A,$(idxs...))...] = v)
+    end
+end
+Base.getindex(A::AxisArray, idxs...) = A[to_index(A,idxs...)...]
+Base.setindex!(A::AxisArray, v, idxs...) = (A[to_index(A,idxs...)...] = v)
+
+# First is indexing by named axis. We simply sort the axes and re-dispatch.
 # When indexing by named axis the shapes of omitted dimensions are preserved
 # TODO: should we handle multidimensional Axis indexes? It could be interpreted
 #       as adding dimensions in the middle of an AxisArray.
 # TODO: should we allow repeated axes? As a union of indices of the duplicates?
-stagedfunction Base.getindex{T,N,D,Ax}(A::AxisArray{T,N,D,Ax}, I::Axis...)
+stagedfunction to_index{T,N,D,Ax}(A::AxisArray{T,N,D,Ax}, I::Axis...)
     dims = Int[axisdim(A, ax) for ax in I]
     idxs = Expr[:(Colon()) for d = 1:N]
     names = axisnames(A)
@@ -68,7 +93,8 @@ stagedfunction Base.getindex{T,N,D,Ax}(A::AxisArray{T,N,D,Ax}, I::Axis...)
         idxs[dims[i]] = :(I[$i].val)
     end
 
-    return :(A[$(idxs...)])
+    meta = Expr(:meta, :inline)
+    return :($meta; to_index(A, $(idxs...)))
 end
 
 ### Indexing along values of the axes ###
@@ -93,58 +119,24 @@ function axisindexes{T}(::Type{Categorical}, ax::AbstractVector{T}, idx::Abstrac
     res
 end
 
-# Defining the fallbacks on getindex are tricky due to ambiguities with
-# AbstractArray definitions -
-let args = Expr[], idxs = Symbol[]
-    for i = 1:4
-        isym = symbol("i$i")
-        push!(args, :($isym::Real))
-        push!(idxs, isym)
-        @eval Base.getindex{T,N,D,Ax}(A::AxisArray{T,N,D,Ax}, $(args...)) = fallback_getindex(A, $(idxs...))
-    end
-end
-Base.getindex{T,N,D,Ax}(A::AxisArray{T,N,D,Ax}, idx::AbstractArray) = fallback_getindex(A, idx)
-Base.getindex{T,N,D,Ax}(A::AxisArray{T,N,D,Ax}, idxs...) = fallback_getindex(A, idxs...)
-
 # These catch-all methods attempt to convert any axis-specific non-standard
-# indexing types to their integer or integer range equivalents using the
-# They are separate from the `Base.getindex` function to help alleviate
+# indexing types to their integer or integer range equivalents using axisindexes
+# They are separate from the `Base.getindex` function to help alleviate 
 # ambiguity warnings from, e.g., `getindex(::AbstractArray, ::Real...)`.
-# TODO: These could be generated with meta-meta-programming
-stagedfunction fallback_getindex{T,N,D,Ax}(A::AxisArray{T,N,D,Ax}, I1)
-    ex = :(getindex(A))
-    push!(ex.args, I1 <: Idx || length(Ax) < 1 ? :(I1) : :(axisindexes(A.axes[1], I1)))
-    for _=2:N
-        push!(ex.args, :(Colon()))
-    end
-    ex
-end
-stagedfunction fallback_getindex{T,N,D,Ax}(A::AxisArray{T,N,D,Ax}, I1, I2)
-    ex = :(getindex(A))
-    push!(ex.args, I1 <: Idx || length(Ax) < 1 ? :(I1) : :(axisindexes(A.axes[1], I1)))
-    push!(ex.args, I2 <: Idx || length(Ax) < 2 ? :(I2) : :(axisindexes(A.axes[2], I2)))
-    for _=3:N
-        push!(ex.args, :(Colon()))
-    end
-    ex
-end
-stagedfunction fallback_getindex{T,N,D,Ax}(A::AxisArray{T,N,D,Ax}, I1, I2, I3)
-    ex = :(getindex(A))
-    push!(ex.args, I1 <: Idx || length(Ax) < 1 ? :(I1) : :(axisindexes(A.axes[1], I1)))
-    push!(ex.args, I2 <: Idx || length(Ax) < 2 ? :(I2) : :(axisindexes(A.axes[2], I2)))
-    push!(ex.args, I3 <: Idx || length(Ax) < 3 ? :(I3) : :(axisindexes(A.axes[3], I3)))
-    for _=4:N
-        push!(ex.args, :(Colon()))
-    end
-    ex
-end
-stagedfunction fallback_getindex{T,N,D,Ax}(A::AxisArray{T,N,D,Ax}, I...)
-    ex = :(getindex(A))
+stagedfunction to_index{T,N,D,Ax}(A::AxisArray{T,N,D,Ax}, I...)
+    ex = Expr(:tuple)
     for i=1:length(I)
-        push!(ex.args, I[i] <: Idx || length(Ax) < i ? :(I[$i]) : :(axisindexes(A.axes[$i], I[$i])))
+        if I[i] <: Idx
+            push!(ex.args, :(I[$i]))
+        elseif i <= length(Ax)
+            push!(ex.args, :(axisindexes(A.axes[$i], I[$i])))
+        else
+            push!(ex.args, :(error("dimension ", $i, " does not have an axis to index")))
+        end
     end
     for _=length(I)+1:N
         push!(ex.args, :(Colon()))
     end
-    ex
+    meta = Expr(:meta, :inline)
+    return :($meta; $ex)
 end