Repository: invenia/NamedDims.jl
Pull request #21 (2019-05-07 19:18:55)

Allow multiple wildcard names
I am not happy with this.
It introduces another (particularly evil) `@generated` function,
and two `@pure` functions that are not really pure. But that are probably acceptable because if the generic function methods that they call get overwritten, then a bunch of things will break everywhere in the standard library.

It does however solve #8 
and it continues to run everything to do with names at compile time.


Patch:
diff --git a/Project.toml b/Project.toml
index 8a312a7..089dbe9 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "NamedDims"
 uuid = "356022a1-0364-5f58-8944-0da4b18d706f"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.1.0"
+version = "0.2.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -9,7 +9,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
 [compat]
-Tracker = "0.2.1"
+Tracker = "0.2.2"
 julia = "1"
 
 [extras]
diff --git a/src/name_core.jl b/src/name_core.jl
index 1d8a0d4..3af6790 100644
--- a/src/name_core.jl
+++ b/src/name_core.jl
@@ -31,33 +31,23 @@ _compile_time_return_hack(::Val{X}) where X = X
 
 
 """
-    dim(dimnames, [name])
+    dim(dimnames, name)
 
 For `dimnames` being a tuple of names (symbols) for the dimensions.
-If called with just the tuple,
-returns a named tuple, with each name mapped to a dimension.
-e.g `dim((:a, :b)) == (a=1, b=2)`.
-
-If the second `name` argument is given, them the dimension corresponding to that `name`,
-is returned.
+and `name` being a one of those names
+This returns the dimension corresponding to that `name`,
 e.g. `dim((:a, :b), :b) == 2`
-If that `name` is not found then `0` is returned.
+If that `name` is not one of the given `dimnames` then an error is thrown.
 """
-function dim(dimnames::Tuple)
-    # 0-Allocations see: `@btime (()->dim((:a, :b)))()`
-    ndims = length(dimnames)
-    return NamedTuple{dimnames, NTuple{ndims, Int}}(1:ndims)
-end
-
-function dim(dimnames::Tuple, name::Symbol)
-    # 0-Allocations see: `@btime  (()->dim((:a, :b), :a))()`
-    this_namemap = NamedTuple{(name,), Tuple{Symbol}}((:notfound,))  # default we will overwrite
-    full_namemap = dim(dimnames)
-    dimnum = first(merge(this_namemap, full_namemap))
-    dimnum isa Int && return dimnum
-    throw(ArgumentError(
-        "Specified name ($(repr(name))) does not match any dimension name ($dimnames)"
-    ))
+function dim(dimnames::Tuple, name::Symbol)::Int
+    # 0-Allocations see: `@btime  (()->dim((:a, :b), :b))()`
+    dimnum = dim_noerror(dimnames, name)
+    if dimnum === 0
+        throw(ArgumentError(
+            "Specified name ($(repr(name))) does not match any dimension name ($dimnames)"
+        ))
+    end
+    return dimnum
 end
 
 function dim(dimnames::Tuple, names)
@@ -65,7 +55,6 @@ function dim(dimnames::Tuple, names)
     return map(name->dim(dimnames, name), names)
 end
 
-
 function dim(dimnames::Tuple, d::Union{Integer, Colon})
     # This is the fallback that allows `NamedDimsArray`'s to be have dimensions
     # referred to by number. This is required to allow functions on `AbstractArray`s
@@ -74,14 +63,16 @@ function dim(dimnames::Tuple, d::Union{Integer, Colon})
     return d
 end
 
+Base.@pure function dim_noerror(dimnames::Tuple{Vararg{Symbol, N}}, name::Symbol) where N
+    # 0-Allocations see: @btime  (()->dim_noerror((:a, :b, :c), :c))()
+    for ii in 1:N
+        getfield(dimnames, ii) === name && return ii
+    end
+    return 0
+end
 
 
 
-function identity_namedtuple(tup::NTuple{N, Symbol}) where N
-    # 0-Allocations
-    return NamedTuple{tup, typeof(tup)}(tup)
-end
-
 """
     permute_dimnames(dimnames, perm)
 
@@ -103,14 +94,21 @@ function permute_dimnames(dimnames::NTuple{N, Symbol}, perm) where N
 end
 
 """
-    default_inds(dimnames::Tuple)
-This is the default value for all indexing expressions using the given dimnames.
-Which is to say: take a full slice on everything
+    tuple_issubset
+A version of `is_subset` sepecifically for `Tuple`s of `Symbol`s, that is `@pure`.
+This helps it get optimised out of existance. It is less of an abuse of `@pure` than
+most of the stuff for making `NamedTuples` work.
 """
-function default_inds(dimnames::NTuple{N}) where N
-    # 0-Allocations
-    values = ntuple(_->Colon(), N)
-    return NamedTuple{dimnames, NTuple{N, Colon}}(values)
+Base.@pure function tuple_issubset(lhs::Tuple{Vararg{Symbol,N}}, rhs::Tuple{Vararg{Symbol,M}}) where {N,M}
+    N <= M || return false
+    for a in lhs
+        found = false
+        for b in rhs
+            found |= a === b
+        end
+        found || return false
+    end
+    return true
 end
 
 """
@@ -120,18 +118,21 @@ Returns the values of the `named_inds`, sorted as per the order they appear in `
 with any missing dimnames, having there value set to `:`.
 An error is thrown if any dimnames are given in `named_inds` that do not occur in `dimnames`.
 """
-function order_named_inds(dimnames::Tuple; named_inds...)
-    # 0-Allocations
+function order_named_inds(dimnames::Tuple{Vararg{Symbol,N}}; named_inds...) where {N}
+    # 0-Allocations: see `@code_typed (()->order_named_inds((:a, :b, :c), (b=1, c=2)))()`
+    if !tuple_issubset(keys(named_inds), dimnames)
+        throw(DimensionMismatch("Expected subset of $(dimnames), got $(keys(named_inds))"))
+    end
 
-    slice_everything = default_inds(dimnames)
-    full_named_inds = merge(slice_everything, named_inds)
-    if length(full_named_inds) != length(dimnames)
-        throw(DimensionMismatch("Expected $(dimnames), got $(keys(named_inds))"))
+    full_inds = ntuple(N) do ii
+        name_ii = dimnames[ii]
+        get(named_inds, name_ii, :)
     end
-    inds = Tuple(full_named_inds)
-    return inds
+
+    return full_inds
 end
 
+
 """
     incompatible_dimension_error(names_a, names_b)
 Throws a `DimensionMismatch` explaining that these dimension names are not compatible.
@@ -202,10 +203,9 @@ function unify_names_longest(names_a, names_b)
 
     length(names_a) == length(names_b) && return unify_names(names_a, names_b)
     long, short = length(names_a) > length(names_b) ? (names_a, names_b) : (names_b, names_a)
-    short_names = identity_namedtuple(short)
     ret = ntuple(length(long)) do ii
         a = getfield(long, ii)
-        b = get(short_names, ii, :_)
+        b = ii <= length(short) ? short[ii] : :_
         a === :_ && return b
         b === :_ && return a
         a === b && return a
@@ -244,22 +244,20 @@ Given a tuple of dimension names, and either a collection of dimensions,
 or a single dimension, expressed as a number,
 Returns the dimension names with those dimensions dropped.
 """
-function remaining_dimnames_after_dropping(dimnames::Tuple, dropped_dim::Integer)
+function remaining_dimnames_after_dropping(dimnames::Tuple, dropped_dim::Int)
     # 0 allocations. See `@btime remaining_dimnames_after_dropping((:a,:b,:c,:d,:e), 4)`
-    return remaining_dimnames_after_dropping(dimnames, (dropped_dim,))
+    return _remaining_dimnames_after_dropping(dimnames, Tuple{dropped_dim})
 end
 
-function remaining_dimnames_after_dropping(dimnames::Tuple, dropped_dims)
-    # 0-Allocations see: `@btime remaining_dimnames_after_dropping((:a,:b,:c,:d,:e), (1,2,))
-
-    anti_names = identity_namedtuple(map(x->dimnames[x], dropped_dims))
-    full_names = identity_namedtuple(dimnames)
+function remaining_dimnames_after_dropping(dimnames::NTuple{N, Symbol}, dropped_dims::Tuple{Vararg{Int}}) where N
+    # 0-Allocations see:
+    # `@code_typed (()->remaining_dimnames_after_dropping((:a,:b,:c,:d,:e), (1,3)))()`
+    return _remaining_dimnames_after_dropping(dimnames, Tuple{dropped_dims...})
+end
 
-    # Now we construct a new named tuple, with all the names we want to remove at the start
-    combined_names = merge(anti_names, full_names)
-    n_skip = length(anti_names)
-    ret = ntuple(length(full_names) - n_skip) do ii
-        combined_names[ii + n_skip]  # Skip over the ones we left as the start
-    end
-    return compile_time_return_hack(ret)
+# dropped_dims must be a Tuple type where the values are Int literal for the dimensions being dropped
+@generated function _remaining_dimnames_after_dropping(dimnames::NTuple{N,Symbol}, dropped_dims::Type) where N
+    dropped_dims_vals = dropped_dims.parameters[1].parameters
+    keep_names = [:(getfield(dimnames, $ii)) for ii in 1:N if ii ∉ dropped_dims_vals]
+    return Expr(:call, :compile_time_return_hack, Expr(:tuple, keep_names...))
 end
diff --git a/src/wrapper_array.jl b/src/wrapper_array.jl
index f7f4124..dd3b249 100644
--- a/src/wrapper_array.jl
+++ b/src/wrapper_array.jl
@@ -68,7 +68,7 @@ Base.parent(x::NamedDimsArray) = x.data
 Returns a tuple of containing the names of all the dimensions of the array `A`.
 """
 names(::Type{<:NamedDimsArray{L}}) where L = L
-names(::Type{<:AbstractArray{T, N}}) where {T,N} = ntuple(_->:_, N)
+names(::Type{<:AbstractArray{T, N}}) where {T, N} = ntuple(_->:_, N)
 names(x::T) where T<:AbstractArray = names(T)
 
 dim(a::NamedDimsArray{L}, name) where L = dim(L, name)


Test patch:
diff --git a/test/broadcasting.jl b/test/broadcasting.jl
index bb2118f..56eb3e2 100644
--- a/test/broadcasting.jl
+++ b/test/broadcasting.jl
@@ -81,6 +81,14 @@ using Tracker
             @test names(total) == (:foo, :bar)
         end
     end
+
+    @testset "Regression test again #8b" begin
+        # https://github.com/invenia/NamedDims.jl/issues/8#issuecomment-490124369
+        nda = NamedDimsArray{(:x,:y,:z)}(ones(10,20,30))
+        @test nda .+ ones(1,20) == 2ones(10,20,30)
+        @test names(nda .+ ones(1,20)) == (:x, :y, :z)
+    end
+
 end
 
 @testset "Competing Wrappers" begin
diff --git a/test/name_core.jl b/test/name_core.jl
index d54c641..eb3cde0 100644
--- a/test/name_core.jl
+++ b/test/name_core.jl
@@ -11,14 +11,6 @@ using Test
 
 
 @testset "dim" begin
-    @testset "get map only" begin
-        @test dim((:x, :y)) == (x=1, y=2)
-
-        manynames = Tuple(Symbol.('A':'z'))
-        namemap = dim(manynames)
-        @test keys(namemap) == manynames
-        @test values(namemap) == Tuple(1:length(manynames))
-    end
     @testset "small case" begin
         @test dim((:x, :y), :x) == 1
         @test dim((:x, :y), :y) == 2
diff --git a/test/wrapper_array.jl b/test/wrapper_array.jl
index ade5403..498772b 100644
--- a/test/wrapper_array.jl
+++ b/test/wrapper_array.jl
@@ -42,6 +42,13 @@ end
         @test names(nda[y=1]) == (:x,)
         @test names(nda[y=1:1]) == (:x, :y)
     end
+
+    # https://github.com/invenia/NamedDims.jl/issues/8
+    @testset "with multiple-wildcards" begin
+        nda_mw = NamedDimsArray{(:_,:_,:c)}(ones(10,20,30));
+        @test nda_mw[c=2] == ones(10,20)
+        @test names(nda_mw[c=2]) == (:_, :_)
+    end
 end