upgraded TupleVector implementation

Author: chakravala

Project.toml

@@ -1,7 +1,7 @@
 name = "AbstractTensors"
 uuid = "a8e43f4a-99b7-5565-8bf1-0165161caaea"
 authors = ["Michael Reed"]
-version = "0.6.6"
+version = "0.6.7"
 
 [deps]
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

src/AbstractTensors.jl

@@ -262,7 +262,7 @@ Base.cosc(t::T) where T<:TensorAlgebra = iszero(t) ? zero(Manifold(t)) : (x=(1π
 # absolute value norm
 
 @inline Base.abs(t::T) where T<:TensorAlgebra = Base.sqrt(Base.abs2(t))
-@inline Base.abs2(t::T) where T<:TensorAlgebra = (a=contraction(t,t); isscalar(a) ? scalar(a) : a)
+@inline Base.abs2(t::T) where T<:TensorAlgebra = (a=(~t)*t; isscalar(a) ? scalar(a) : a)
 @inline Base.abs2(t::T) where T<:TensorGraded = contraction(t,t)
 @inline norm(z) = LinearAlgebra.norm(z)
 @inline LinearAlgebra.norm(t::T) where T<:TensorAlgebra = norm(value(t))

src/FixedVector.jl

@@ -2,22 +2,31 @@
 # This file is adapted from JuliaArrays/StaticArrays.jl License is MIT:
 # https://github.com/JuliaArrays/StaticArrays.jl/blob/master/LICENSE.md
 
-struct FixedVector{N,T} <: TupleVector{N,T}
-    v::Vector{T}
-    function FixedVector{N,T}(a::Vector) where {N,T}
+require_one_based_indexing(A...) = !Base.has_offset_axes(A...) ||
+    throw(ArgumentError("offset arrays are not supported but got an array with index other than 1"))
+
+struct FixedVector{N,T,TData<:AbstractVector{T}} <: TupleVector{N,T}
+    v::TData
+    function FixedVector{N,T,TData}(a::TData) where {N,T,TData<:AbstractVector{T}}
+        require_one_based_indexing(a)
         if length(a) != N
             throw(DimensionMismatch("Dimensions $(size(a)) don't match static size $S"))
         end
-        new{N,T}(a)
+        new{N,T,TData}(a)
     end
-    function FixedVector{N,T}(::UndefInitializer) where {N,T}
-        new{N,T}(Vector{T}(undef,N))
+    function FixedVector{N,T,TData}(::UndefInitializer) where {N,T,TData<:AbstractVector{T}}
+        new{N,T,TData}(TData(undef,N))
     end
 end
 
-@inline FixedVector{N}(a::Vector{T}) where {N,T} = FixedVector{N,T}(a)
-
-@generated function FixedVector{N,T}(x::NTuple{N,Any}) where {N,T}
+@inline function FixedVector{N,T}(a::TData) where {N,T,TData<:AbstractVector{T}}
+    return FixedVector{N,T,TData}(a)
+end
+@inline function FixedVector{N}(a::TData) where {N,T,TData<:AbstractVector{T}}
+    return FixedVector{N,T,TData}(a)
+end
+@inline FixedVector{N,T}(::UndefInitializer) where {N,T} = FixedVector{N,T,Vector{T}}(undef)
+@generated function (::Type{FixedVector{N,T,TData}})(x::NTuple{N,Any}) where {N,T,TData<:AbstractVector{T}}
     exprs = [:(a[$i] = x[$i]) for i = 1:N]
     return quote
         $(Expr(:meta, :inline))
@@ -27,27 +36,79 @@ end
     end
 end
 
-@inline FixedVector{N,T}(x::Tuple) where {N,T} = FixedVector{N,T}(x)
+@inline FixedVector{N,T}(x::Tuple) where {N,T} = FixedVector{N,T,Vector{T}}(x)
 @inline FixedVector{N}(x::NTuple{N,T}) where {N,T} = FixedVector{N,T}(x)
 
 # Overide some problematic default behaviour
-@inline Base.convert(::Type{SA}, sa::FixedVector) where {SA<:FixedVector} = SA(sa.v)
-@inline Base.convert(::Type{SA}, sa::SA) where {SA<:FixedVector} = sa
+@inline Base.convert(::Type{TV}, tv::FixedVector) where {TV<:FixedVector} = TV(tv.v)
+@inline Base.convert(::Type{TV}, tv::TV) where {TV<:FixedVector} = tv
 
-# Back to Array (unfortunately need both convert and construct to overide other methods)
-@inline Base.Array(sa::FixedVector) = Vector(sa.v)
-@inline Base.Array{T}(sa::FixedVector{N,T}) where {N,T} = Vector{T}(sa.v)
-@inline Base.Array{T,1}(sa::FixedVector{N,T}) where {N,T} = Vector{T}(sa.v)
+# Back to Vector (unfortunately need both convert and construct to overide other methods)
+@inline Base.Vector(tv::FixedVector) = Vector(tv.v)
+@inline Base.Vector{T}(tv::FixedVector{N,T}) where {N,T} = Vector{T}(tv.v)
 
-@inline Base.convert(::Type{Array}, sa::FixedVector) = sa.v
-@inline Base.convert(::Type{Array{T}}, sa::FixedVector{N,T}) where {N,T} = sa.v
-@inline Base.convert(::Type{Array{T,1}}, sa::FixedVector{N,T}) where {N,T} = sa.v
+@inline function Base.convert(::Type{Vector}, tv::FixedVector{N}) where N
+    return Vector(tv.v)
+end
+@inline function Base.convert(::Type{Vector}, tv::FixedVector{N,T,Vector{T}}) where {N,T}
+    return tv.v
+end
+@inline function Base.convert(::Type{Vector{T}}, tv::FixedVector{N,T}) where {N,T}
+    return Vector(tv.v)
+end
+@inline function Base.convert(::Type{Vector{T}}, tv::FixedVector{N,T,Vector{T}}) where {N,T}
+    return tv.v
+end
 
 @propagate_inbounds Base.getindex(a::FixedVector, i::Int) = getindex(a.v, i)
 @propagate_inbounds Base.setindex!(a::FixedVector, v, i::Int) = setindex!(a.v, v, i)
 
+Base.parent(tv::FixedVector) = tv.v
+
+Base.pointer(tv::FixedVector) = pointer(tv.v)
+
+Base.unsafe_convert(::Type{Ptr{T}}, tv::FixedVector) where T = Base.unsafe_convert(Ptr{T}, tv.v)
+Base.elsize(::Type{FixedVector{S,T}}) where {S,T} = Base.elsize(A)
+
+FixedVector(a::TupleVector{N,T}) where {N,T} = FixedVector{N,T}(a)
+
 Base.dataids(sa::FixedVector) = Base.dataids(sa.v)
 
-function Base.promote_rule(::Type{<:FixedVector{N,T}}, ::Type{<:FixedVector{N,U}}) where {N,T,U}
-    FixedVector{N,promote_type(T,U)}
+function Base.promote_rule(
+    ::Type{FixedVector{N,T,TDataA}},
+    ::Type{FixedVector{N,U,TDataB}},
+) where {N,T,U,TDataA,TDataB}
+    return FixedVector{N, promote_type(T, U), promote_type(TDataA, TDataB)}
+end
+
+function promote_rule(::Type{FixedVector{N,T}},::Type{FixedVector{N,U}}) where {N,T,U}
+    return FixedVector{N, promote_type(T, U)}
+end
+
+### Code that makes views of statically sized arrays also statically sized (where possible)
+
+# Note, _get_tuple_vector_length is used in a generated function so it's strictly internal and can't be extended
+_get_tuple_vector_length(::Type{<:TupleVector{N}}) where N = N
+
+@generated function new_out_size(::Type{Val{N}}, ind) where N
+    if ind <: Integer
+        Tuple{1} # dimension is fixed
+    elseif ind <: TupleVector
+        Tuple{_get_tuple_vector_length(ind)}
+    elseif ind == Colon || ind <: Base.Slice
+        Tuple{N}
+    elseif ind <: SOneTo
+        Tuple{ind.parameters[1]}
+    else
+        error("Unknown index type: $ind")
+    end
+end
+
+@generated new_out_size(::Type{Val{N}}, ::Colon) where N = Tuple{N}
+
+function Base.view(
+    a::FixedVector{N},
+    index::Union{Integer, Colon, TupleVector, Base.Slice, SOneTo},
+) where N
+    return FixedVector{new_out_size(Val(N), index)}(view(a.v, index))
 end

src/Values.jl

@@ -17,21 +17,77 @@ end
     end
 end
 
+@noinline function generator_too_short_error(inds::CartesianIndices, i::CartesianIndex)
+    error("Generator produced too few elements: Expected exactly $(shape_string(inds)) elements, but generator stopped at $(shape_string(i))")
+end
+@noinline function generator_too_long_error(inds::CartesianIndices)
+    error("Generator produced too many elements: Expected exactly $(shape_string(inds)) elements, but generator yields more")
+end
+
+shape_string(inds::CartesianIndices) = join(length.(inds.indices), '×')
+shape_string(inds::CartesianIndex) = join(Tuple(inds), '×')
+
+@inline throw_if_nothing(x, inds, i) =
+    (x === nothing && generator_too_short_error(inds, i); x)
+
+@generated function tvcollect(::Type{TV}, gen) where {TV <: TupleVector{N}} where N
+    stmts = [:(Base.@_inline_meta)]
+    args = []
+    iter = :(iterate(gen))
+    inds = CartesianIndices(Tuple(N))
+    for i in inds
+        el = Symbol(:el, i)
+        push!(stmts, :(($el,st) = throw_if_nothing($iter, $inds, $i)))
+        push!(args, el)
+        iter = :(iterate(gen,st))
+    end
+    push!(stmts, :($iter === nothing || generator_too_long_error($inds)))
+    push!(stmts, :(TV($(args...))))
+    Expr(:block, stmts...)
+end
+"""
+   tvcollect(TV, gen)
+
+Construct a statically-sized vector of type `TV`.from a generator
+`gen`. `TV` needs to have a size parameter since the length of `vec`
+is unknown to the compiler. `TV` can optionally specify the element
+type as well.
+
+Example:
+
+    tvcollect(Values{3, Int}, 2i+1 for i in 1:3)
+
+This creates the same statically-sized vector as if the generator were
+collected in an array, but is more efficient since no array is
+allocated.
+
+Equivalent:
+
+    Values{3, Int}([2i+1 for i in 1:3])
+"""
+tvcollect
+
+@inline (::Type{TV})(gen::Base.Generator) where {TV <: TupleVector} =
+    tvcollect(TV, gen)
+
 @inline Values(a::TupleVector{N}) where N = Values{N}(Tuple(a))
 @propagate_inbounds Base.getindex(v::Values, i::Int) = v.v[i]
 @inline Tuple(v::Values) = v.v
 Base.dataids(::Values) = ()
 
 # See #53
 Base.cconvert(::Type{Ptr{T}}, a::Values) where {T} = Base.RefValue(a)
-Base.unsafe_convert(::Type{Ptr{T}}, a::Base.RefValue{SA}) where {N,T,SA<:Values{N,T}} = Ptr{T}(Base.unsafe_convert(Ptr{Values{N,T}}, a))
+Base.unsafe_convert(::Type{Ptr{T}}, a::Base.RefValue{TV}) where {N,T,TV<:Values{N,T}} = Ptr{T}(Base.unsafe_convert(Ptr{Values{N,T}}, a))
 
 # SVector.jl
 
 @inline Values(x::NTuple{N,Any}) where N = Values{N}(x)
 @inline Values{N}(x::NTuple{N,T}) where {N,T} = Values{N,T}(x)
 @inline Values{N}(x::T) where {N,T<:Tuple} = Values{N,promote_tuple_eltype(T)}(x)
 
+@inline Values{N, T}(gen::Base.Generator) where {N, T} = tvcollect(Values{N, T}, gen)
+@inline Values{N}(gen::Base.Generator) where {N} = tvcollect(Values{N}, gen)
+
 # Some more advanced constructor-like functions
 @pure @inline Base.zeros(::Type{Values{N}}) where N = zeros(Values{N,Float64})
 @pure @inline Base.ones(::Type{Values{N}}) where N = ones(Values{N,Float64})

src/Variables.jl

@@ -11,7 +11,7 @@ mutable struct Variables{N,T} <: TupleVector{N,T}
     Variables{N,T}(::UndefInitializer) where {N,T} = new{N,T}()
 end
 
-@inline Variables(a::TupleVector{N}) where N = Variables{N}(Tuple(a))
+@inline Variables(a::TupleVector{N,T}) where {N,T} = Variables{N,T}(Tuple(a))
 @generated function (::Type{Variables{N,T}})(x::Tuple) where {N,T}
     return quote
         $(Expr(:meta, :inline))
@@ -43,7 +43,7 @@ end
         # unsafe_store!(Base.unsafe_convert(Ptr{Ptr{Nothing}}, pointer_from_objref(v.data)), pointer_from_objref(val), i)
         error("setindex!() with non-isbitstype eltype is not supported by TupleVectors. Consider using FixedVector.")
     end
-    return val
+    return v
 end
 
 @inline Base.Tuple(v::Variables) = v.v
@@ -57,8 +57,19 @@ function Base.promote_rule(::Type{<:Variables{N,T}}, ::Type{<:Variables{N,U}}) w
     Variables{N,promote_type(T,U)}
 end
 
+function Base.view(
+    a::Variables{N},
+    index::Union{Integer, Colon, TupleVector, Base.Slice, SOneTo},
+) where N
+    view_from_invoke = invoke(view, Tuple{AbstractVector, typeof(index)}, a, index)
+    return FixedVector{new_out_size(Val(N), index)}(view_from_invoke)
+end
+
+Base.elsize(::Type{<:Variables{S,T}}) where {S,T} = sizeof(T)
+
 # MVector.jl
 
+@inline Variables(x::Values{N,T}) where {N,T} = Variables{N,T}(x)
 @inline Variables(x::NTuple{N,Any}) where N = Variables{N}(x)
 @inline Variables{N}(x::NTuple{N,T}) where {N,T} = Variables{N,T}(x)
 @inline Variables{N}(x::NTuple{N,Any}) where N = Variables{N, promote_tuple_eltype(typeof(x))}(x)

src/abstractvector.jl

@@ -21,26 +21,26 @@ similar_type(::Type{SA}) where {SA<:TupleVector} = similar_type(SA,eltype(SA))
 similar_type(::SA,::Type{T}) where {SA<:TupleVector{N},T} where N = similar_type(SA,T,Val(N))
 similar_type(::Type{SA},::Type{T}) where {SA<:TupleVector{N},T} where N = similar_type(SA,T,Val(N))
 
-similar_type(::A,n::Val) where {A<:AbstractArray} = similar_type(A,eltype(A),n)
-similar_type(::Type{A},n::Val) where {A<:AbstractArray} = similar_type(A,eltype(A),n)
+similar_type(::A,n::Val) where {A<:AbstractVector} = similar_type(A,eltype(A),n)
+similar_type(::Type{A},n::Val) where {A<:AbstractVector} = similar_type(A,eltype(A),n)
 
-similar_type(::A,::Type{T},n::Val) where {A<:AbstractArray,T} = similar_type(A,T,n)
+similar_type(::A,::Type{T},n::Val) where {A<:AbstractVector,T} = similar_type(A,T,n)
 
 # We should be able to deal with SOneTo axes
 @pure similar_type(s::SOneTo) = similar_type(typeof(s))
 @pure similar_type(::Type{SOneTo{n}}) where n = similar_type(SOneTo{n}, Int, Val(n))
 
 # Default types
-# Generally, use SArray
-similar_type(::Type{A},::Type{T},n::Val) where {A<:AbstractArray,T} = default_similar_type(T,n)
+# Generally, use TupleVector
+similar_type(::Type{A},::Type{T},n::Val) where {A<:AbstractVector,T} = default_similar_type(T,n)
 default_similar_type(::Type{T},::Val{N}) where {N,T} = Values{N,T}
 
 similar_type(::Type{SA},::Type{T},n::Val) where {SA<:Variables,T} = mutable_similar_type(T,n)
 
 mutable_similar_type(::Type{T},::Val{N}) where {N,T} = Variables{N,T}
 
 similar_type(::Type{<:FixedVector},::Type{T},n::Val) where T = sizedarray_similar_type(T,n)
-# Should FixedVector also be used for normal Array?
+# Should FixedVector also be used for normal Vector?
 #similar_type(::Type{<:Array},::Type{T},n::Val) where T = sizedarray_similar_type(T,n)
 
 sizedarray_similar_type(::Type{T},::Val{N}) where {N,T} = FixedVector{N,T}
@@ -52,17 +52,17 @@ Base.similar(::SA,::Type{T}) where {SA<:TupleVector{N},T} where N = similar(SA,T
 Base.similar(::Type{SA},::Type{T}) where {SA<:TupleVector{N},T} where N = similar(SA,T,Val(N))
 
 # Cases where a Val is given as the dimensions
-Base.similar(::A,n::Val) where A<:AbstractArray = similar(A,eltype(A),n)
-Base.similar(::Type{A},n::Val) where A<:AbstractArray = similar(A,eltype(A),n)
+Base.similar(::A,n::Val) where A<:AbstractVector = similar(A,eltype(A),n)
+Base.similar(::Type{A},n::Val) where A<:AbstractVector = similar(A,eltype(A),n)
 
-Base.similar(::A,::Type{T},n::Val) where {A<:AbstractArray,T} = similar(A,T,n)
+Base.similar(::A,::Type{T},n::Val) where {A<:AbstractVector,T} = similar(A,T,n)
 
 # defaults to built-in mutable types
-Base.similar(::Type{A},::Type{T},n::Val) where {A<:AbstractArray,T} = mutable_similar_type(T,n)(undef)
+Base.similar(::Type{A},::Type{T},n::Val) where {A<:AbstractVector,T} = mutable_similar_type(T,n)(undef)
 
 # both FixedVector and Array return FixedVector
 Base.similar(::Type{SA},::Type{T},n::Val) where {SA<:FixedVector,T} = sizedarray_similar_type(T,n)(undef)
-Base.similar(::Type{A},::Type{T},n::Val) where {A<:Array,T} = sizedarray_similar_type(T,n)(undef)
+Base.similar(::Type{A},::Type{T},n::Val) where {A<:Vector,T} = sizedarray_similar_type(T,n)(undef)
 
 # Support tuples of mixtures of `SOneTo`s alongside the normal `Integer` and `OneTo` options
 # by simply converting them to either a tuple of Ints or a Val, re-dispatching to either one
@@ -105,3 +105,17 @@ end
         @inbounds return similar_type(a, promote_type(eltype(a), eltype(b)), Val($Snew))(tuple($(exprs...)))
     end
 end
+
+if VERSION >= v"1.6.0-DEV.1334"
+    # FIXME: This always assumes one-based linear indexing and that subtypes of TupleVector
+    # don't overload iterate
+    @inline function Base.rest(a::TupleVector{N}, (_, i) = (nothing, 0)) where N
+        return similar_type(typeof(a), Val(N - i))(Base.rest(Tuple(a), i + 1))
+    end
+end
+
+# SArrays may avoid the SubArray wrapper and consequently an additional level of indirection
+# The output may use the broadcasting machinery defined for StaticArrays (see issue #892)
+function Base.view(S::Values, I::Union{Colon, Integer, SOneTo, TupleVector{N, Int} where N, CartesianIndex})
+    getindex(S, I)
+end

src/convert.jl

@@ -7,27 +7,25 @@
 
 @inline (::Type{SA})(x...) where {SA <: TupleVector} = SA(x)
 @inline (::Type{SA})(a::TupleVector) where {SA<:TupleVector} = SA(Tuple(a))
-@propagate_inbounds (::Type{SA})(a::AbstractArray) where {SA <: TupleVector} = convert(SA, a)
+@propagate_inbounds (::Type{SA})(a::AbstractVector) where {SA <: TupleVector} = convert(SA, a)
 
 # this covers most conversions and "statically-sized reshapes"
 #@inline Base.convert(::Type{SA}, sa::TupleVector) where {SA<:TupleVector} = SA(Tuple(sa))
 #@inline Base.convert(::Type{SA}, sa::SA) where {SA<:TupleVector} = sa
 @inline Base.convert(::Type{SA}, x::Tuple) where {SA<:TupleVector} = SA(x) # convert -> constructor. Hopefully no loops...
 
-# support conversion to AbstractArray
-Base.AbstractArray{T}(sa::TupleVector{N,T}) where {N,T} = sa
-Base.AbstractArray{T,1}(sa::TupleVector{N,T}) where {N,T} = sa
-Base.AbstractArray{T}(sa::TupleVector{N,U}) where {N,T,U} = similar_type(typeof(sa),T,Val(N))(sa)
-Base.AbstractArray{T,1}(sa::TupleVector{N,U}) where {N,T,U} = similar_type(typeof(sa),T,Val(N))(sa)
+# support conversion to AbstractVector
+Base.AbstractVector{T}(sa::TupleVector{N,T}) where {N,T} = sa
+Base.AbstractVector{T}(sa::TupleVector{N,U}) where {N,T,U} = similar_type(typeof(sa),T,Val(N))(sa)
 
 # Constructing a Tuple from a TupleVector
 @inline Base.Tuple(a::TupleVector{N}) where N = unroll_tuple(a, Val(N))
 
-@noinline function dimension_mismatch_fail(SA::Type, a::AbstractArray)
-    throw(DimensionMismatch("expected input array of length $(length(SA)), got length $(length(a))"))
+@noinline function dimension_mismatch_fail(SA::Type, a::AbstractVector)
+    throw(DimensionMismatch("expected input vector of length $(length(SA)), got length $(length(a))"))
 end
 
-@propagate_inbounds function Base.convert(::Type{SA}, a::AbstractArray) where {SA<:TupleVector{N}} where N
+@propagate_inbounds function Base.convert(::Type{SA}, a::AbstractVector) where {SA<:TupleVector{N}} where N
     @boundscheck if length(a) != length(SA)
         dimension_mismatch_fail(SA, a)
     end
@@ -45,10 +43,14 @@ end
 @pure length_val(a::T) where {T <: TupleMatrixLike{S}} where S = Val(S)
 @pure length_val(a::Type{T}) where {T<:TupleMatrixLike{S}} where S = Val(S)
 
-@generated function unroll_tuple(a::AbstractArray, ::Val{N}) where N
+@generated function unroll_tuple(a::AbstractVector, ::Val{N}) where N
     exprs = [:(a[$j]) for j = 1:N]
     quote
         @_inline_meta
         @inbounds return $(Expr(:tuple, exprs...))
     end
 end
+
+# `float` and `real` of TupleVector types, analogously to application to scalars (issue 935)
+Base.float(::Type{TV}) where TV<:TupleVector{T,_N} where {T,_N} = similar_type(TV, float(T))
+Base.real(::Type{TV}) where TV<:TupleVector{T,_N} where {T,_N} = similar_type(TV, real(T))