introduce function collect_as for construction from an iterator (#48)

* introduce function `collect_as` for construction from an iterator

Makes constructing `FixedSizeArray`s more convenient!

Inspired by
JuliaLang/julia#36288

This currently ignores `Base.IteratorElType`, xref
https://discourse.julialang.org/t/i-dont-get-base-iteratoreltype/113604

The allocations in some code paths are probably excessive/could be
optimized. But I guess this is good for a start.

Fixes #20

* also test an iterator with `BigInt`-valued `size` and `length`

* remove the premature optimization for `AbstractArray`

* improve tests

* add to the Readme

* whitespace/formatting fix

* delete two useless `nothing` lines

One of these wasn't being recorded by code coverage (another Julia
coverage bug, I guess).

* simplify a bit

Author: nsajko

README.md

@@ -13,6 +13,46 @@ Main differences between `FixedSizeArray` and `MArray` are:
 * `FixedSizeArray` is based on the `Memory` type introduced in Julia v1.11, `MArray` is backed by tuples;
 * the size of the array is part of the type parameters of `MArray`, this isn't the case for `FixedSizeArray`, where the size is only a constant field of the data structure.
 
+FixedSizeArrays supports the usual array interfaces, so things like broadcasting, matrix
+multiplication, other linear algebra operations, `similar`, `copyto!` or `map` should just work.
+
+Use the constructors to convert from other array types. Use `collect_as` to convert from
+arbitrary iterators.
+
+```julia-repl
+julia> arr = [10 20; 30 14]
+2×2 Matrix{Int64}:
+ 10  20
+ 30  14
+
+julia> iter = (i for i ∈ 7:9 if i≠8);
+
+julia> using FixedSizeArrays
+
+julia> FixedSizeArray(arr)  # construct from an `AbstractArray` value
+2×2 FixedSizeMatrix{Int64}:
+ 10  20
+ 30  14
+
+julia> FixedSizeArray{Float64}(arr)  # construct from an `AbstractArray` value while converting element type
+2×2 FixedSizeMatrix{Float64}:
+ 10.0  20.0
+ 30.0  14.0
+
+julia> const ca = FixedSizeArrays.collect_as
+collect_as (generic function with 1 method)
+
+julia> ca(FixedSizeArray, iter)  # construct from an arbitrary iterator
+2-element FixedSizeVector{Int64}:
+ 7
+ 9
+
+julia> ca(FixedSizeArray{Float64}, iter)  # construct from an arbitrary iterator while converting element type
+2-element FixedSizeVector{Float64}:
+ 7.0
+ 9.0
+```
+
 Note: `FixedSizeArray`s are not guaranteed to be stack-allocated, in fact they will more likely *not* be stack-allocated.
 However, in some *extremely* simple cases the compiler may be able to completely elide their allocations:
 ```julia

src/FixedSizeArrays.jl

@@ -2,6 +2,8 @@ module FixedSizeArrays
 
 export FixedSizeArray, FixedSizeVector, FixedSizeMatrix
 
+public collect_as
+
 """
     Internal()
 
@@ -94,6 +96,41 @@ end
 
 # helper functions
 
+dimension_count_of(::Base.SizeUnknown) = 1
+dimension_count_of(::Base.HasLength) = 1
+dimension_count_of(::Base.HasShape{N}) where {N} = convert(Int, N)::Int
+
+struct LengthIsUnknown end
+struct LengthIsKnown end
+length_status(::Base.SizeUnknown) = LengthIsUnknown()
+length_status(::Base.HasLength) = LengthIsKnown()
+length_status(::Base.HasShape) = LengthIsKnown()
+
+function check_count_value(n::Int)
+    if n < 0
+        throw(ArgumentError("count can't be negative"))
+    end
+end
+function check_count_value(n)
+    throw(ArgumentError("count must be an `Int`"))
+end
+
+struct SpecFSA{T,N} end
+function fsa_spec_from_type(::Type{FixedSizeArray})
+    SpecFSA{nothing,nothing}()
+end
+function fsa_spec_from_type(::Type{FixedSizeArray{<:Any,M}}) where {M}
+    check_count_value(M)
+    SpecFSA{nothing,M}()
+end
+function fsa_spec_from_type(::Type{FixedSizeArray{E}}) where {E}
+    SpecFSA{E::Type,nothing}()
+end
+function fsa_spec_from_type(::Type{FixedSizeArray{E,M}}) where {E,M}
+    check_count_value(M)
+    SpecFSA{E::Type,M}()
+end
+
 parent_type(::Type{<:FixedSizeArray{T}}) where {T} = Memory{T}
 
 underlying_storage(m) = m
@@ -168,4 +205,115 @@ function Base.reshape(a::FixedSizeArray{T}, size::NTuple{N,Int}) where {T,N}
     FixedSizeArray{T,N}(Internal(), a.mem, size)
 end
 
+# `collect_as`
+
+function collect_as_fsa0(iterator, ::Val{nothing})
+    x = only(iterator)
+    ret = FixedSizeArray{typeof(x),0}(undef)
+    ret[] = x
+    ret
+end
+
+function collect_as_fsa0(iterator, ::Val{E}) where {E}
+    E::Type
+    x = only(iterator)
+    ret = FixedSizeArray{E,0}(undef)
+    ret[] = x
+    ret
+end
+
+function fill_fsa_from_iterator!(a, iterator)
+    actual_count = 0
+    for e ∈ iterator
+        actual_count += 1
+        a[actual_count] = e
+    end
+    if actual_count != length(a)
+        throw(ArgumentError("`size`-`length` inconsistency"))
+    end
+end
+
+function collect_as_fsam_with_shape(
+    iterator, ::SpecFSA{nothing,M}, shape::Tuple{Vararg{Int}},
+) where {M}
+    E = eltype(iterator)::Type
+    ret = FixedSizeArray{E,M}(undef, shape)
+    fill_fsa_from_iterator!(ret, iterator)
+    map(identity, ret)::FixedSizeArray{<:Any,M}
+end
+
+function collect_as_fsam_with_shape(
+    iterator, ::SpecFSA{E,M}, shape::Tuple{Vararg{Int}},
+) where {E,M}
+    E::Type
+    ret = FixedSizeArray{E,M}(undef, shape)
+    fill_fsa_from_iterator!(ret, iterator)
+    ret::FixedSizeArray{E,M}
+end
+
+function collect_as_fsam(iterator, spec::SpecFSA{<:Any,M}) where {M}
+    check_count_value(M)
+    shape = if isone(M)
+        (length(iterator),)
+    else
+        size(iterator)
+    end::NTuple{M,Any}
+    shap = map(Int, shape)::NTuple{M,Int}
+    collect_as_fsam_with_shape(iterator, spec, shap)::FixedSizeArray{<:Any,M}
+end
+
+function collect_as_fsa1_from_unknown_length(iterator, ::Val{nothing})
+    v = collect(iterator)::AbstractVector
+    T = FixedSizeVector
+    map(identity, T(v))::T
+end
+
+function collect_as_fsa1_from_unknown_length(iterator, ::Val{E}) where {E}
+    E::Type
+    v = collect(E, iterator)::AbstractVector{E}
+    T = FixedSizeVector{E}
+    T(v)::T
+end
+
+function collect_as_fsa_impl(iterator, ::SpecFSA{E,0}, ::LengthIsKnown) where {E}
+    collect_as_fsa0(iterator, Val(E))::FixedSizeArray{<:Any,0}
+end
+
+function collect_as_fsa_impl(iterator, spec::SpecFSA, ::LengthIsKnown)
+    collect_as_fsam(iterator, spec)::FixedSizeArray
+end
+
+function collect_as_fsa_impl(iterator, ::SpecFSA{E,1}, ::LengthIsUnknown) where {E}
+    collect_as_fsa1_from_unknown_length(iterator, Val(E))::FixedSizeVector
+end
+
+function collect_as_fsa_checked(iterator, ::SpecFSA{E,nothing}, ::Val{M}, length_status) where {E,M}
+    check_count_value(M)
+    collect_as_fsa_impl(iterator, SpecFSA{E,M}(), length_status)::FixedSizeArray{<:Any,M}
+end
+
+function collect_as_fsa_checked(iterator, ::SpecFSA{E,M}, ::Val{M}, length_status) where {E,M}
+    check_count_value(M)
+    collect_as_fsa_impl(iterator, SpecFSA{E,M}(), length_status)::FixedSizeArray{<:Any,M}
+end
+
+"""
+    collect_as(t::Type{<:FixedSizeArray}, iterator)
+
+Tries to construct a value of type `t` from the iterator `iterator`. The type `t`
+must either be concrete, or a `UnionAll` without constraints.
+"""
+function collect_as(::Type{T}, iterator) where {T<:FixedSizeArray}
+    spec = fsa_spec_from_type(T)::SpecFSA
+    size_class = Base.IteratorSize(iterator)
+    if size_class == Base.IsInfinite()
+        throw(ArgumentError("iterator is infinite, can't fit infinitely many elements into a `FixedSizeArray`"))
+    end
+    dim_count_int = dimension_count_of(size_class)
+    check_count_value(dim_count_int)
+    dim_count = Val(dim_count_int)::Val
+    len_stat = length_status(size_class)
+    collect_as_fsa_checked(iterator, spec, dim_count, len_stat)::T
+end
+
 end # module FixedSizeArrays

test/runtests.jl

@@ -4,6 +4,7 @@ using OffsetArrays: OffsetArray
 import Aqua
 
 const checked_dims = FixedSizeArrays.checked_dims
+const collect_as = FixedSizeArrays.collect_as
 
 # helpers for testing for allocation or suboptimal inference
 
@@ -353,4 +354,80 @@ end
             end
         end
     end
+
+    @testset "`collect_as`" begin
+        for T ∈ (FixedSizeArray, FixedSizeVector, FixedSizeArray{Int}, FixedSizeVector{Int})
+            for iterator ∈ (Iterators.repeated(7), Iterators.cycle(7))
+                @test_throws ArgumentError collect_as(T, iterator)
+            end
+        end
+        for T ∈ (FixedSizeArray{<:Any,-1}, FixedSizeArray{Int,-1}, FixedSizeArray{Int,3.1})
+            iterator = (7:8, (7, 8))
+            @test_throws ArgumentError collect_as(T, iterator)
+        end
+        for T ∈ (FixedSizeArray{3}, FixedSizeVector{3})
+            iterator = (7:8, (7, 8))
+            @test_throws TypeError collect_as(T, iterator)
+        end
+        struct Iter{E,N,I<:Integer}
+            size::NTuple{N,I}
+            length::I
+            val::E
+        end
+        function Base.iterate(i::Iter)
+            l = i.length
+            iterate(i, max(zero(l), l))
+        end
+        function Base.iterate(i::Iter, state::Int)
+            if iszero(state)
+                nothing
+            else
+                (i.val, state - 1)
+            end
+        end
+        Base.IteratorSize(::Type{<:Iter{<:Any,N}}) where {N} = Base.HasShape{N}()
+        Base.length(i::Iter) = i.length
+        Base.size(i::Iter) = i.size
+        Base.eltype(::Type{<:Iter{E}}) where {E} = E
+        @testset "buggy iterator with mismatched `size` and `length" begin
+            for iterator ∈ (Iter((), 0, 7), Iter((3, 2), 5, 7))
+                E = eltype(iterator)
+                dim_count = length(size(iterator))
+                for T ∈ (FixedSizeArray, FixedSizeArray{E}, FixedSizeArray{<:Any,dim_count}, FixedSizeArray{E,dim_count})
+                    @test_throws ArgumentError collect_as(T, iterator)
+                end
+            end
+        end
+        iterators = (
+            (), (7,), (7, 8), 7, (7 => 8), Ref(7), fill(7),
+            (i for i ∈ 1:3), ((i + 100*j) for i ∈ 1:3, j ∈ 1:2), Iterators.repeated(7, 2),
+            (i for i ∈ 7:9 if i==8), 7:8, 8:7, map(BigInt, 7:8), Int[], [7], [7 8],
+            Iter((), 1, 7), Iter((3,), 3, 7), Iter((3, 2), 6, 7),
+        )
+        abstract_array_params(::AbstractArray{T,N}) where {T,N} = (T, N)
+        @testset "iterator: $iterator" for iterator ∈ iterators
+            a = collect(iterator)
+            (E, dim_count) = abstract_array_params(a)
+            af = collect(Float64, iterator)
+            @test abstract_array_params(af) == (Float64, dim_count)  # meta
+            @test_throws MethodError collect_as(FixedSizeArray{E,dim_count+1}, iterator)
+            for T ∈ (FixedSizeArray, FixedSizeArray{<:Any,dim_count})
+                fsa = collect_as(T, iterator)
+                @test a == fsa
+                @test first(abstract_array_params(fsa)) <: E
+            end
+            for T ∈ (FixedSizeArray{E}, FixedSizeArray{E,dim_count})
+                test_inferred(collect_as, FixedSizeArray{E,dim_count}, (T, iterator))
+                fsa = collect_as(T, iterator)
+                @test a == fsa
+                @test first(abstract_array_params(fsa)) <: E
+            end
+            for T ∈ (FixedSizeArray{Float64}, FixedSizeArray{Float64,dim_count})
+                test_inferred(collect_as, FixedSizeArray{Float64,dim_count}, (T, iterator))
+                fsa = collect_as(T, iterator)
+                @test af == fsa
+                @test first(abstract_array_params(fsa)) <: Float64
+            end
+        end
+    end
 end