Dynamic dimnames (#238)

* Enable optionally static/dynamic dimnames

Author: Tokazama

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "4.0.2"
+version = "4.0.3"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

docs/src/api.md

@@ -23,6 +23,7 @@ ArrayInterface.ismutable
 ArrayInterface.issingular
 ArrayInterface.isstructured
 ArrayInterface.is_splat_index
+ArrayInterface.known_dimnames
 ArrayInterface.known_first
 ArrayInterface.known_last
 ArrayInterface.known_length

docs/src/index.md

@@ -100,11 +100,22 @@ New types defining dimension names can do something similar to:
 using Static
 using ArrayInterface
 
-struct NewType{dnames} end  # where dnames::Tuple{Vararg{Symbol}}
+struct StaticDimnames{dnames} end  # where dnames::Tuple{Vararg{Symbol}}
+
+ArrayInterface.known_dimnames(::Type{StaticDimnames{dnames}}) where {dnames} = dnames
+ArrayInterface.dimnames(::StaticDimnames{dnames}) where {dnames} = static(dnames)
+
+struct DynamicDimnames{N}
+    dimnames::NTuple{N,Symbol}
+end
+ArrayInterface.known_dimnames(::Type{DynamicDimnames{N}}) where {N} = ntuple(_-> missing, Val(N))
+ArrayInterface.dimnames(x::DynamicDimnames) = getfield(x, :dimnames)
 
-ArrayInterface.dimnames(::Type{NewType{dnames}}) = static(dnames)
 ```
 
+Notice that `DynamicDimnames` returns `missing` instead of a symbol for each dimension.
+This indicates dimension names are present for `DynamicDimnames` but that information is missing at compile time.
+
 Dimension names should be appropriately propagated between nested arrays using `ArrayInterface.to_parent_dims`. 
 This allows types such as `SubArray` and `PermutedDimsArray` to work with named dimensions.
 Similarly, other methods that return information corresponding to dimensions (e.g., `ArrayInterfce.size`, `ArrayInterface.axes`) use `to_parent_dims` to appropriately propagate parent information.

src/axes.jl

@@ -5,20 +5,19 @@
 
 Returns the type of each axis for the `T`, or the type of of the axis along dimension `dim`.
 """
-axes_types(x, dim) = axes_types(typeof(x), dim)
-@inline axes_types(::Type{T}, dim) where {T} = axes_types(T, to_dims(T, dim))
-@inline function axes_types(::Type{T}, dim::StaticInt{D}) where {T,D}
-    if D > ndims(T)
+@inline axes_types(x, dim) = axes_types(x, to_dims(x, dim))
+@inline function axes_types(x, dim::StaticInt{D}) where {D}
+    if D > ndims(x)
         return SOneTo{1}
     else
-        return field_type(axes_types(T), dim)
+        return field_type(axes_types(x), dim)
     end
 end
-@inline function axes_types(::Type{T}, dim::Int) where {T}
-    if dim > ndims(T)
+@inline function axes_types(x, dim::Int)
+    if dim > ndims(x)
         return SOneTo{1}
     else
-        return axes_types(T).parameters[dim]
+        return axes_types(x).parameters[dim]
     end
 end
 axes_types(x) = axes_types(typeof(x))
@@ -288,3 +287,4 @@ lazy_axes(x::CartesianIndices) = axes(x)
 @inline lazy_axes(x::MatAdjTrans) = reverse(lazy_axes(parent(x)))
 @inline lazy_axes(x::VecAdjTrans) = (SOneTo{1}(), first(lazy_axes(parent(x))))
 @inline lazy_axes(x::PermutedDimsArray) = permute(lazy_axes(parent(x)), to_parent_dims(x))
+

src/dimensions.jl

@@ -147,68 +147,76 @@ function to_parent_dims(::Type{T}, ::StaticInt{dim}) where {T,dim}
     end
 end
 
+_nunderscore(::Val{N}) where {N} = ntuple(Compat.Returns(:_), Val(N))
+
 """
-    has_dimnames(::Type{T}) -> Bool
+    has_dimnames(::Type{T}) -> StaticBool
 
 Returns `static(true)` if `x` has on or more named dimensions. If all dimensions correspond
 to `static(:_)`, then `static(false)` is returned.
 """
-has_dimnames(x) = has_dimnames(typeof(x))
-@inline has_dimnames(::Type{T}) where {T} = _has_dimnames(dimnames(T))
-_has_dimnames(::Tuple{Vararg{StaticSymbol{:_}}}) = static(false)
-_has_dimnames(::Tuple) = static(true)
-
-# this takes the place of dimension names that aren't defined
-const SUnderscore = StaticSymbol(:_)
+Compat.@constprop :aggressive has_dimnames(x) = static(_is_named(known_dimnames(x)))
+_is_named(x::NTuple{N,Symbol}) where {N} = x !== _nunderscore(Val(N))
+_is_named(::Any) = true
 
 """
-    dimnames(::Type{T}) -> Tuple{Vararg{StaticSymbol}}
-    dimnames(::Type{T}, dim) -> StaticSymbol
+    known_dimnames(::Type{T}) -> Tuple{Vararg{Union{Symbol,Missing}}}
+    known_dimnames(::Type{T}, dim::Union{Int,StaticInt}) -> Union{Symbol,Missing}
 
-Return the names of the dimensions for `x`. `static(:_)` is used to indicate a dimension
-does not have a name.
+Return the names of the dimensions for `x`. `:_` is used to indicate a dimension does not
+have a name.
 """
-@inline dimnames(x) = dimnames(typeof(x))
-@inline dimnames(::Type{T}) where {T} = _dimnames(has_parent(T), T)
-_dimnames(::False, ::Type{T}) where {T} = ntuple(_->static(:_), Val(ndims(T)))
-@inline function _dimnames(::True, ::Type{T}) where {T}
-    eachop(_perm_dimnames, to_parent_dims(T), dimnames(parent_type(T)))
+@inline known_dimnames(x, dim::Integer) = _known_dimname(known_dimnames(x), canonicalize(dim))
+known_dimnames(x) = known_dimnames(typeof(x))
+known_dimnames(::Type{T}) where {T} = _known_dimnames(T, parent_type(T))
+_known_dimnames(::Type{T}, ::Type{T}) where {T} = _unknown_dimnames(Base.IteratorSize(T))
+_unknown_dimnames(::Base.HasShape{N}) where {N} = _nunderscore(Val(N))
+_unknown_dimnames(::Any) = (:_,)
+function _known_dimnames(::Type{C}, ::Type{P}) where {C,P}
+    eachop(_inbounds_known_dimname, to_parent_dims(C), known_dimnames(P))
+end
+@inline function _known_dimname(x::Tuple{Vararg{Any,N}}, dim::CanonicalInt) where {N}
+    @boundscheck (dim > N || dim < 1) && return :_
+    return @inbounds(x[dim])
 end
+@inline _inbounds_known_dimname(x, dim) = @inbounds(_known_dimname(x, dim))
 
-@inline dimnames(x, dim) = dimnames(typeof(x), dim)
-@inline dimnames(::Type{T}, dim::Integer) where {T} = _perm_dimnames(dimnames(T), dim)
-function _perm_dimnames(dnames::Tuple{Vararg{StaticSymbol,N}}, dim) where {N}
-    if dim > N
-        return static(:_)
-    else
-        return @inbounds(dnames[dim])
-    end
+"""
+    dimnames(x) -> Tuple{Vararg{Union{Symbol,StaticSymbol}}}
+    dimnames(x, dim::Union{Int,StaticInt}) -> Union{Symbol,StaticSymbol}
+
+Return the names of the dimensions for `x`. `:_` is used to indicate a dimension does not
+have a name.
+"""
+@inline dimnames(x, dim::Integer) = _dimname(dimnames(x), canonicalize(dim))
+@inline dimnames(x) = _dimnames(has_parent(x), x)
+@inline function _dimnames(::True, x)
+    eachop(_inbounds_dimname, to_parent_dims(x), dimnames(parent(x)))
+end
+_dimnames(::False, x) = ntuple(_->static(:_), Val(ndims(x)))
+@inline function _dimname(x::Tuple{Vararg{Any,N}}, dim::CanonicalInt) where {N}
+    @boundscheck (dim > N || dim < 1) && return static(:_)
+    return @inbounds(x[dim])
 end
+@inline _inbounds_dimname(x, dim) = @inbounds(_dimname(x, dim))
 
 """
-    to_dims(::Type{T}, dim) -> Union{Int,StaticInt}
+    to_dims(x, dim) -> Union{Int,StaticInt}
 
-This returns the dimension(s) of `x` corresponding to `d`.
+This returns the dimension(s) of `x` corresponding to `dim`.
 """
-to_dims(x, dim) = to_dims(typeof(x), dim)
-to_dims(::Type{T}, dim::StaticInt) where {T} = dim
-to_dims(::Type{T}, dim::Integer) where {T} = Int(dim)
-to_dims(::Type{T}, dim::Colon) where {T} = dim
-function to_dims(::Type{T}, dim::StaticSymbol) where {T}
-    i = find_first_eq(dim, dimnames(T))
-    if i === nothing
-        throw_dim_error(T, dim)
-    end
-    return i
+to_dims(x, dim::Colon) = dim
+to_dims(x, dim::Integer) = canonicalize(dim)
+to_dims(x, dim::Union{StaticSymbol,Symbol}) = _to_dim(dimnames(x), dim)
+function to_dims(x, dims::Tuple{Vararg{Any,N}}) where {N}
+    eachop(_to_dims, nstatic(Val(N)), dimnames(x), dims)
 end
-Compat.@constprop :aggressive function to_dims(::Type{T}, dim::Symbol) where {T}
-    i = find_first_eq(dim, map(Symbol, dimnames(T)))
-    if i === nothing
-        throw_dim_error(T, dim)
-    end
+@inline _to_dims(x::Tuple, d::Tuple, n::StaticInt{N}) where {N} = _to_dim(x, getfield(d, N))
+@inline function _to_dim(x::Tuple, d::Union{Symbol,StaticSymbol})
+    i = find_first_eq(d, x)
+    i === nothing && throw(DimensionMismatch("dimension name $(d) not found"))
     return i
 end
-to_dims(::Type{T}, dims::Tuple) where {T} = map(i -> to_dims(T, i), dims)
 
 #=
     order_named_inds(names, namedtuple)
@@ -224,37 +232,31 @@ An error is thrown if any keywords are used which do not occur in `nda`'s names.
 3. if missing is found use Colon()
 4. if (ndims - ncolon) === nkwargs then all were found, else error
 =#
-order_named_inds(x::Tuple, ::NamedTuple{(),Tuple{}}) = ()
-function order_named_inds(x::Tuple, nd::NamedTuple{L}) where {L}
-    return order_named_inds(x, static(Val(L)), Tuple(nd))
-end
-Compat.@constprop :aggressive function order_named_inds(
-    x::Tuple{Vararg{Any,N}},
-    nd::Tuple,
-    inds::Tuple
-) where {N}
-
-    out = eachop(order_named_inds, nstatic(Val(N)), x, nd, inds)
-    _order_named_inds_check(out, length(nd))
-    return out
-end
-function order_named_inds(x::Tuple, nd::Tuple, inds::Tuple, ::StaticInt{dim}) where {dim}
-    index = find_first_eq(getfield(x, dim), nd)
-    if index === nothing
-        return Colon()
+@generated function find_all_dimnames(x::Tuple{Vararg{Any,ND}}, nd::Tuple{Vararg{Any,NI}}, inds::Tuple, default) where {ND,NI}
+    if NI === 0
+        return :(())
     else
-        return @inbounds(inds[index])
-    end
-end
-
-ncolon(x::Tuple{Colon,Vararg}, n::Int) = ncolon(tail(x), n + 1)
-ncolon(x::Tuple{Any,Vararg}, n::Int) = ncolon(tail(x), n)
-ncolon(x::Tuple{Colon}, n::Int) = n + 1
-ncolon(x::Tuple{Any}, n::Int) = n
-function _order_named_inds_check(inds::Tuple{Vararg{Any,N}}, nkwargs::Int) where {N}
-    if (N - ncolon(inds, 0)) !== nkwargs
-        error("Not all keywords matched dimension names.")
+        out = Expr(:block, Expr(:(=), :names_found, 0))
+        t = Expr(:tuple)
+        for i in 1:ND
+            index_i = Symbol(:index_, i)
+            val_i = Symbol(:val_, i)
+            push!(t.args, val_i)
+            push!(out.args, quote
+                $index_i = find_first_eq(getfield(x, $i), nd)
+                if $index_i === nothing
+                    $val_i = default
+                else
+                    $val_i = @inbounds(inds[$index_i])
+                    names_found += 1
+                end
+            end)
+        end
+        return quote
+            $out
+            @boundscheck names_found === $NI || error("Not all keywords matched dimension names.")
+            return $t
+        end
     end
-    return missing
 end
 

src/indexing.jl

@@ -307,8 +307,8 @@ function getindex(A, args...)
     @boundscheck checkbounds(A, inds...)
     unsafe_getindex(A, inds...)
 end
-function getindex(A; kwargs...)
-    inds = to_indices(A, order_named_inds(dimnames(A), values(kwargs)))
+@propagate_inbounds function getindex(A; kwargs...)
+    inds = to_indices(A, find_all_dimnames(dimnames(A), static(keys(kwargs)), Tuple(values(kwargs)), :))
     @boundscheck checkbounds(A, inds...)
     unsafe_getindex(A, inds...)
 end
@@ -407,7 +407,7 @@ Store the given values at the given key or index within a collection.
 end
 @propagate_inbounds function setindex!(A, val; kwargs...)
     can_setindex(A) || error("Instance of type $(typeof(A)) are not mutable and cannot change elements after construction.")
-    inds = to_indices(A, order_named_inds(dimnames(A), values(kwargs)))
+    inds = to_indices(A, find_all_dimnames(dimnames(A), static(keys(kwargs)), Tuple(values(kwargs)), :))
     @boundscheck checkbounds(A, inds...)
     unsafe_setindex!(A, val, inds...)
 end

test/dimensions.jl

@@ -5,12 +5,18 @@
 ### define wrapper with dimnames
 ###
 
-struct NamedDimsWrapper{L,T,N,P<:AbstractArray{T,N}} <: ArrayInterface.AbstractArray2{T,N}
+struct NamedDimsWrapper{D,T,N,P<:AbstractArray{T,N}} <: ArrayInterface.AbstractArray2{T,N}
+    dimnames::D
     parent::P
-    NamedDimsWrapper{L}(p) where {L} = new{L,eltype(p),ndims(p),typeof(p)}(p)
+    NamedDimsWrapper(d::D, p::P) where {D,P} = new{D,eltype(P),ndims(p),P}(d, p)
 end
+Base.parent(x::NamedDimsWrapper) = getfield(x, :parent)
 ArrayInterface.parent_type(::Type{T}) where {P,T<:NamedDimsWrapper{<:Any,<:Any,<:Any,P}} = P
-ArrayInterface.dimnames(::Type{T}) where {L,T<:NamedDimsWrapper{L}} = static(L)
+ArrayInterface.dimnames(x::NamedDimsWrapper) = getfield(x, :dimnames)
+function ArrayInterface.known_dimnames(::Type{T}) where {L,T<:NamedDimsWrapper{L}}
+    ArrayInterface.Static.known(L)
+end
+
 Base.parent(x::NamedDimsWrapper) = x.parent
 
 @testset "dimension permutations" begin
@@ -70,29 +76,31 @@ end
     n1 = (static(:x),)
     n2 = (n1..., static(:y))
     n3 = (n2..., static(:z))
-    @test @inferred(ArrayInterface.order_named_inds(n1, NamedTuple{(),Tuple{}}(()) )) == ()
-    @test @inferred(ArrayInterface.order_named_inds(n1, (x=2,))) == (2,)
-    @test @inferred(ArrayInterface.order_named_inds(n2, (x=2,))) == (2, :)
-    @test @inferred(ArrayInterface.order_named_inds(n2, (y=2,))) == (:, 2)
-    @test @inferred(ArrayInterface.order_named_inds(n2, (y=20, x=30))) == (30, 20)
-    @test @inferred(ArrayInterface.order_named_inds(n2, (x=30, y=20))) == (30, 20)
-    @test @inferred(ArrayInterface.order_named_inds(n3, (x=30, y=20))) == (30, 20, :)
-
-    @test_throws ErrorException ArrayInterface.order_named_inds(n2, (x=30, y=20, z=40))
+    @test @inferred(ArrayInterface.find_all_dimnames(n1, (), (), :)) == ()
+    @test @inferred(ArrayInterface.find_all_dimnames(n1, (static(:x),), (2,), :)) == (2,)
+    @test @inferred(ArrayInterface.find_all_dimnames(n2, (static(:x),), (2,), :)) == (2,:)
+    @test @inferred(ArrayInterface.find_all_dimnames(n2, (static(:y),), (2,), :)) == (:, 2)
+    @test @inferred(ArrayInterface.find_all_dimnames(n2, (static(:y), static(:x)), (20, 30), :)) == (30, 20)
+    @test @inferred(ArrayInterface.find_all_dimnames(n2, (static(:x), static(:y)), (30, 20), :)) == (30, 20)
+    @test @inferred(ArrayInterface.find_all_dimnames(n3, (static(:x), static(:y)), (30, 20), :)) == (30, 20, :)
+
+    @test_throws ErrorException ArrayInterface.find_all_dimnames(n2, (static(:x), static(:y), static(:z)), (30, 20, 40), :)
 end
 
-
 @testset "dimnames" begin
     d = (static(:x), static(:y))
-    x = NamedDimsWrapper{d}(ones(2,2));
-    y = NamedDimsWrapper{(:x,)}(ones(2));
+    x = NamedDimsWrapper(d, ones(2,2));
+    y = NamedDimsWrapper((static(:x),), ones(2));
+    z = NamedDimsWrapper((:x, static(:y)), ones(2));
     dnums = ntuple(+, length(d))
     @test @inferred(ArrayInterface.has_dimnames(x)) == true
+    @test @inferred(ArrayInterface.has_dimnames(z)) == true
     @test @inferred(ArrayInterface.has_dimnames(ones(2,2))) == false
     @test @inferred(ArrayInterface.has_dimnames(Array{Int,2})) == false
     @test @inferred(ArrayInterface.has_dimnames(typeof(x))) == true
     @test @inferred(ArrayInterface.has_dimnames(typeof(view(x, :, 1, :)))) == true
     @test @inferred(dimnames(x)) === d
+    @test @inferred(ArrayInterface.dimnames(z)) === (:x, static(:y))
     @test @inferred(dimnames(parent(x))) === (static(:_), static(:_))
     @test @inferred(dimnames(x')) === reverse(d)
     @test @inferred(dimnames(y')) === (static(:_), static(:x))
@@ -103,11 +111,14 @@ end
     @test @inferred(dimnames(view(x, :, 1, :))) === (static(:x), static(:_))
     @test @inferred(dimnames(x, ArrayInterface.One())) === static(:x)
     @test @inferred(dimnames(parent(x), ArrayInterface.One())) === static(:_)
+    @test @inferred(ArrayInterface.known_dimnames(Iterators.flatten(1:10))) === (:_,)
+    @test @inferred(ArrayInterface.known_dimnames(Iterators.flatten(1:10), static(1))) === :_
+    @test @inferred(ArrayInterface.known_dimnames(z)) === (missing, :y)
 end
 
 @testset "to_dims" begin
-    x = NamedDimsWrapper{(:x, :y)}(ones(2,2));
-    y = NamedDimsWrapper{(:x, :y, :a, :b, :c, :d)}(ones(6));
+    x = NamedDimsWrapper(static((:x, :y)), ones(2,2));
+    y = NamedDimsWrapper(static((:x, :y, :a, :b, :c, :d)), ones(6));
 
     @test @inferred(ArrayInterface.to_dims(x, :)) == Colon()
     @test @inferred(ArrayInterface.to_dims(x, 1)) == 1
@@ -130,8 +141,8 @@ end
 
 @testset "methods accepting dimnames" begin
     d = (static(:x), static(:y))
-    x = NamedDimsWrapper{d}(ones(2,2));
-    y = NamedDimsWrapper{(:x,)}(ones(2));
+    x = NamedDimsWrapper(d, ones(2,2));
+    y = NamedDimsWrapper((static(:x),), ones(2));
     @test @inferred(size(x, first(d))) == size(parent(x), 1)
     @test @inferred(ArrayInterface.size(y')) == (1, size(parent(x), 1))
     @test @inferred(axes(x, first(d))) == axes(parent(x), 1)
@@ -144,6 +155,9 @@ end
 
     x[x = 1] = [2, 3]
     @test @inferred(getindex(x, x = 1)) == [2, 3]
+    y = NamedDimsWrapper((:x, static(:y)), ones(2, 2));
+    # FIXME this doesn't correctly infer the output because it can't infer 
+    @test getindex(y, x = 1) == [1, 1]
 end
 
 end