Mapping over enumeration of a static array (#1107)

* Mapping over enumeration of a static array

* bump version

* use separate `enumerate_static`

Author: mateuszbaran

Project.toml

@@ -1,6 +1,6 @@
 name = "StaticArrays"
 uuid = "90137ffa-7385-5640-81b9-e52037218182"
-version = "1.5.9"
+version = "1.5.10"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/StaticArrays.jl

@@ -58,6 +58,7 @@ export @MVector, @MMatrix, @MArray
 
 export similar_type
 export push, pop, pushfirst, popfirst, insert, deleteat, setindex
+export enumerate_static
 
 export StaticArraysCore
 

src/mapreduce.jl

@@ -77,6 +77,37 @@ end
     end
 end
 
+struct StaticEnumerate{TA}
+    itr::TA
+end
+
+enumerate_static(a::StaticArray) = StaticEnumerate(a)
+
+@generated function map(f, a::StaticEnumerate{<:StaticArray})
+    S = Size(a.parameters[1])
+    if prod(S) == 0
+        # In the empty case only, use inference to try figuring out a sensible
+        # eltype, as is done in Base.collect and broadcast.
+        # See https://github.com/JuliaArrays/StaticArrays.jl/issues/528
+        return quote
+            @_inline_meta
+            T = Core.Compiler.return_type(f, Tuple{Tuple{Int,$(eltype(a.parameters[1]))}})
+            @inbounds return similar_type(a.itr, T, $S)()
+        end
+    end
+
+    exprs = Vector{Expr}(undef, prod(S))
+    for i ∈ 1:prod(S)
+        exprs[i] = :(f(($i, a.itr[$i])))
+    end
+
+    return quote
+        @_inline_meta
+        @inbounds elements = tuple($(exprs...))
+        @inbounds return similar_type(typeof(a.itr), eltype(elements), $S)(elements)
+    end
+end
+
 @inline function map!(f, dest::StaticArray, a::StaticArray...)
     _map!(f, dest, same_size(dest, a...), a...)
 end

test/mapreduce.jl

@@ -240,4 +240,15 @@ using Statistics: mean
         @test @inferred(reduce(vcat, v2)) === @SVector [1,2,3,4]
         @test @inferred(reduce(hcat, v2)) === @SMatrix [1 3; 2 4]
     end
+    @testset "map over enumerate" begin
+        # issue #1106
+        v = @SVector [1, -2, 3, -4]
+        m = @SMatrix [1 -2; 3 -4]
+        v0 = SVector{0,Float64}()
+        m0 = SMatrix{0,0,Float64}()
+        @test @inferred(map(f -> f[1] * f[2], enumerate_static(v))) === @SVector [1, -4, 9, -16]
+        @test @inferred(map(f -> f[1] * f[2], enumerate_static(m))) === @SMatrix [1 -6; 6 -16]
+        @test @inferred(map(f -> f, enumerate_static(v0))) === SVector{0,Tuple{Int,Float64}}()
+        @test @inferred(map(f -> f, enumerate_static(m0))) === SMatrix{0,0,Tuple{Int,Float64}}()
+    end
 end