re-write indexing of key vectors

Author: Michael Abbott

src/struct.jl

@@ -58,53 +58,78 @@ Base.eachindex(A::KeyedArray) = eachindex(parent(A))
 
 Base.keys(A::KeyedArray) = error("Base.keys(::KeyedArray) not defined, please open an issue if this happens unexpectedly.")
 
-for (bget, rget, cpy) in [(:getindex, :keys_getindex, :copy), (:view, :keys_view, :identity)]
+for (get_or_view, key_get, maybe_copy) in [
+        (:getindex, :keys_getindex, :copy),
+        (:view, :keys_view, :identity)
+    ]
     @eval begin
 
-        @inline function Base.$bget(A::KeyedArray, I::Integer...)
-            @boundscheck checkbounds(parent(A), I...)
-            @inbounds Base.$bget(parent(A), I...)
-        end
+        @inline function Base.$get_or_view(A::KeyedArray, raw_inds...)
+            inds = to_indices(A, raw_inds)
+            @boundscheck checkbounds(parent(A), inds...)
+            data = @inbounds $get_or_view(parent(A), inds...)
+            data isa AbstractArray || return data # scalar output
+
+            raw_keys = $key_get(axiskeys(A), inds)
+            raw_keys === () && return data # things like A[A .> 0]
 
-        @inline function Base.$bget(A::KeyedArray, I::Union{Colon, CartesianIndex, BitArray})
-            @boundscheck checkbounds(parent(A), I)
-            @inbounds Base.$bget(parent(A), I)
+            new_keys = ntuple(ndims(data)) do d
+                isnothing(raw_keys) && return axes(data, d)
+                raw_keys[d]
+            end
+            KeyedArray(data, new_keys)
         end
 
-        @inline @propagate_inbounds function Base.$bget(A::KeyedArray, Iraw...)
-            I = Base.to_indices(A, Iraw) # allows InvertedIndices.jl
+        # drop all, for A[:] and A[A .> 0] with ndims>=2
+        @inline $key_get(keys::Tuple{Any, Any, Vararg{Any}}, inds::Tuple{Base.LogicalIndex}) = ()
+        @inline $key_get(keys::Tuple{Any, Any, Vararg{Any}}, inds::Tuple{Base.Slice}) = ()
 
-            @boundscheck checkbounds(parent(A), I...)
-            data = @inbounds Base.$bget(parent(A), I...)
+        # drop one, for integer index
+        @inline $key_get(keys::Tuple, inds::Tuple{Integer, Vararg{Any}}) =
+            $key_get(tail(keys), tail(inds))
 
-            @boundscheck map(checkbounds, axiskeys(A), I)
-            new_keys = $rget(axiskeys(A), I)
+        # from a Colon, k[:] would copy too, but this avoids view([1,2,3], :)
+        @inline $key_get(keys::Tuple, inds::Tuple{Base.Slice, Vararg{Any}}) =
+            ($maybe_copy(first(keys)), $key_get(tail(keys), tail(inds))...)
 
-            new_keys isa Tuple{} ? data : KeyedArray(data, new_keys)
+        # this avoids making views of 1:10 etc, they are immutable anyway
+        @inline function $key_get(keys::Tuple, inds::Tuple{AbstractVector, Vararg{Any}})
+            got = if first(keys) isa AbstractRange
+                @inbounds getindex(first(keys), first(inds))
+            else
+                @inbounds $get_or_view(first(keys), first(inds))
+            end
+            (got, $key_get(tail(keys), tail(inds))...)
         end
 
-        @inline function $rget(keys, inds)
-            got = map(keys, inds) do r,i
-                i isa Integer       && return nothing
-                i isa Colon         && return $cpy(r)        # avoids view([1,2,3], :)
-                r isa AbstractRange && return getindex(r,i)  # don't make views of 1:10
-                return @inbounds $bget(r,i)
-            end
-            filter(r -> r isa AbstractArray, got)
+        # newindex=[CartesianIndex{0}()], uses up one ind, sets N keys to default i.e. axes
+        @inline function $key_get(keys::Tuple, inds::Tuple{AbstractVector{CartesianIndex{0}}, Vararg{Any}})
+            (OneTo(1), $key_get(keys, tail(inds))...)
         end
+        @inline function $key_get(keys::Tuple, inds::Tuple{AbstractVector{CartesianIndex{N}}, Vararg{Any}}) where {N}
+            _, keys_left = Base.IteratorsMD.split(keys, Val(N))
+            (ntuple(_->nothing, N)..., $key_get(keys_left, tail(inds))...)
+        end
+
+        # terminating case, trailing 1s (already checked) could be left over
+        @inline $key_get(keys::Tuple{}, inds::Tuple{}) = ()
+        @inline $key_get(keys::Tuple{}, inds::Tuple{Integer, Vararg{Any}}) = ()
 
     end
 end
 
-@inline function Base.setindex!(A::KeyedArray, val, I...)
+@inline function Base.setindex!(A::KeyedArray, val, raw_inds...)
+    I = Base.to_indices(A, raw_inds)
     @boundscheck checkbounds(A, I...)
     @inbounds setindex!(parent(A), val, I...)
     val
 end
 
-@inline function Base.Broadcast.dotview(A::KeyedArray, I...)
+@inline function Base.dotview(A::KeyedArray, raw_inds...)
+    I = Base.to_indices(A, raw_inds)
     @boundscheck checkbounds(A, I...)
-    @inbounds Base.dotview(parent(A), I...)
+    @inbounds setindex!(parent(A), val, I...)
+    val
 end
 
 """

test/_basic.jl

@@ -20,9 +20,11 @@ using Test, AxisKeys
     @test R[1,1] == 321
 
     @test R[:] == vec(R.data)
-    @test_broken R[1:2, 1, 1] == R.data[1:2, 1, 1] # trailing 1s are broken
+    @test R[1:2, 1, 1] == R.data[1:2, 1, 1]
     @test axiskeys(R[:, [0.9,0.1,0.9,0.1] .> 0.5],2) == [10,30]
     @test ndims(R[R.data .> 0.5]) == 1
+    newaxis = [CartesianIndex{0}()]
+    @test axiskeys(R[[1,3],newaxis,:]) == (['a', 'c'], Base.OneTo(1), 10:10:40)
 
     @test_throws Exception R(:nope) # ideally ArgumentError
     @test_throws Exception R('z')   # ideally BoundsError?