General indexstyle clean up (#82)

* refactor bestindex and indexstyle

* switch to index_type

* more clean up

Author: Pietro Vertechi

src/StructArrays.jl

@@ -1,5 +1,7 @@
 module StructArrays
 
+using Base: tuple_type_cons, tuple_type_head, tuple_type_tail, tail
+
 import Requires
 using PooledArrays: PooledArray
 

src/lazy.jl

@@ -44,8 +44,10 @@ Base.size(v::LazyRows) = size(parent(v))
 Base.getindex(v::LazyRows{<:Any, <:Any, <:Any, Int}, i::Int) = LazyRow(parent(v), i)
 Base.getindex(v::LazyRows{<:Any, <:Any, <:Any, CartesianIndex{N}}, i::Vararg{Int, N}) where {N} = LazyRow(parent(v), CartesianIndex(i))
 
-_best_index(::Type{LazyRows{T, N, C, I}}) where {T, N, C, I} = I
-Base.IndexStyle(::Type{L}) where {L<:LazyRows} = _indexstyle(_best_index(L))
+index_type(::Type{LazyRows{T, N, C, I}}) where {T, N, C, I} = I
+function Base.IndexStyle(::Type{L}) where {L<:LazyRows}
+    index_type(L) === Int ? IndexLinear() : IndexCartesian()
+end
 
 function Base.showarg(io::IO, s::LazyRows{T}, toplevel) where T
     print(io, "LazyRows")

src/structarray.jl

@@ -14,24 +14,23 @@ struct StructArray{T, N, C<:Tup, I} <: AbstractArray{T, N}
                 axes(c[i]) == ax || error("all field arrays must have same shape")
             end
         end
-        new{T, N, C, _best_index(c...)}(c)
+        new{T, N, C, index_type(C)}(c)
     end
 end
 
-_best_index() = Int
-_best_index(col::AbstractArray, cols::AbstractArray...) = _best_index(IndexStyle(col, cols...), col)
-_best_index(::IndexLinear, ::AbstractArray) = Int
-_best_index(::IndexCartesian, ::AbstractArray{T, N}) where {T, N} = CartesianIndex{N}
-_best_index(::Type{StructArray{T, N, C, I}}) where {T, N, C, I} = I
-_indexstyle(::Type{Int}) = IndexLinear()
-_indexstyle(::Type{CartesianIndex{N}}) where {N} = IndexCartesian()
+index_type(::Type{NamedTuple{names, types}}) where {names, types} = index_type(types)
+index_type(::Type{Tuple{}}) = Int
+function index_type(::Type{T}) where {T<:Tuple}
+    S, U = tuple_type_head(T), tuple_type_tail(T)
+    IndexStyle(S) isa IndexCartesian ? CartesianIndex{ndims(S)} : index_type(U) 
+end
 
-_dims(c::Tup) = length(axes(c[1]))
-_dims(c::EmptyTup) = 1
+index_type(::Type{StructArray{T, N, C, I}}) where {T, N, C, I} = I
 
 function StructArray{T}(c::C) where {T, C<:Tup}
     cols = strip_params(staticschema(T))(c)
-    StructArray{T, _dims(cols), typeof(cols)}(cols)
+    N = isempty(cols) ? 1 : ndims(cols[1]) 
+    StructArray{T, N, typeof(cols)}(cols)
 end
 
 StructArray(c::C) where {C<:NamedTuple} = StructArray{eltypes(C)}(c)
@@ -49,7 +48,9 @@ const StructVector{T, C<:Tup, I} = StructArray{T, 1, C, I}
 StructVector{T}(args...; kwargs...) where {T} = StructArray{T}(args...; kwargs...)
 StructVector(args...; kwargs...) = StructArray(args...; kwargs...)
 
-Base.IndexStyle(::Type{S}) where {S<:StructArray} = _indexstyle(_best_index(S))
+function Base.IndexStyle(::Type{S}) where {S<:StructArray}
+    index_type(S) === Int ? IndexLinear() : IndexCartesian()
+end
 
 function _undef_array(::Type{T}, sz; unwrap = t -> false) where {T}
     if unwrap(T)

src/utils.jl

@@ -1,5 +1,3 @@
-import Base: tuple_type_cons, tuple_type_head, tuple_type_tail, tail
-
 eltypes(::Type{T}) where {T} = map_params(eltype, T)
 
 map_params(f, ::Type{Tuple{}}) = Tuple{}

test/runtests.jl

@@ -46,7 +46,7 @@ end
     s = StructArray(a=rand(10,10), b=view(rand(100,100), 1:10, 1:10))
     T = typeof(s)
     @test IndexStyle(T) === IndexCartesian()
-    @test StructArrays._best_index(T) == CartesianIndex{2}
+    @test StructArrays.index_type(T) == CartesianIndex{2}
     @test s[100] == s[10, 10] == (a=s.a[10,10], b=s.b[10,10])
     s[100] = (a=1, b=1)
     @test s[100] == s[10, 10] == (a=1, b=1)
@@ -56,7 +56,7 @@ end
     s = StructArray(a=rand(10,10), b=rand(10,10))
     T = typeof(s)
     @test IndexStyle(T) === IndexLinear()
-    @test StructArrays._best_index(T) == Int
+    @test StructArrays.index_type(T) == Int
     @test s[100] == s[10, 10] == (a=s.a[10,10], b=s.b[10,10])
     s[100] = (a=1, b=1)
     @test s[100] == s[10, 10] == (a=1, b=1)