Specialize zeros/ones/fill for SOneTo (#1234)

* Specialize zeros/ones/fill for SOneTo

* Add comment

* Restrict methods to SOneTo

Author: jishnub

src/abstractarray.jl

@@ -291,3 +291,19 @@ function Base.view(S::SArray, I::Union{Colon, Integer, SOneTo, StaticArray{<:Tup
     V = getindex(S, I...)
     _maybewrapscalar(S, V)
 end
+
+# zeros, ones and fill may return SArrays if all the axes are statically sized
+for (arrf, elf) in ((:zeros, :zero), (:ones, :one))
+    _arrf = Symbol(:_, arrf)
+    @eval begin
+        function $arrf(::Type{T}, ax::Tuple{SOneTo,Vararg{SOneTo}}) where {T}
+            sz = homogenize_shape(ax)
+            similar_type(SArray, T, sz)(ntuple(_->$elf(T), prod(sz)))
+        end
+    end
+end
+
+function fill(v, ax::Tuple{SOneTo,Vararg{SOneTo}})
+    sz = homogenize_shape(ax)
+    similar_type(SArray, typeof(v), sz)(ntuple(_->v, prod(sz)))
+end

test/abstractarray.jl

@@ -418,3 +418,14 @@ end
         @test b isa Vec{0}
     end
 end
+
+@testset "zeros/ones/fill" begin
+    for ax in ((SOneTo(2),), (SOneTo(2),SOneTo(3)))
+        @test @inferred(fill(:abc, ax...)) === @inferred(fill(:abc, ax))
+        @test fill(:abc, ax) == fill(:abc, length.(ax)) == fill(:abc, Base.OneTo.(length.(ax)))
+        for fz in (zeros, ones)
+            @test @inferred(fz(Float32, ax...)) === @inferred(fz(Float32, ax))
+            @test fz(Float32, ax) == fz(Float32, length.(ax)) == fz(Float32, Base.OneTo.(length.(ax)))
+        end
+    end
+end