fix checked_dims (#44)

xref JuliaLang/julia#54244 Fixes #40
JuliaArrays · Apr 29, 2024 · f0dbc65 · f0dbc65
1 parent 6975d6d
commit f0dbc65
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 14 deletions.
diff --git a/src/FixedSizeArrays.jl b/src/FixedSizeArrays.jl
@@ -47,29 +47,32 @@ Base.isassigned(a::FixedSizeArray, i::Int) = isassigned(a.mem, i)
 
 # safe product of a tuple of integers, for calculating dimensions size
 
-checked_dims_impl(a::Int, ::Tuple{}) = a
-function checked_dims_impl(a::Int, t::Tuple{Int,Vararg{Int,N}}) where {N}
+checked_dims_impl(a::Int, ::Tuple{}, have_overflow::Bool) = (a, have_overflow)
+function checked_dims_impl(a::Int, t::Tuple{Int,Vararg{Int,N}}, have_overflow::Bool) where {N}
     b = first(t)
-    tmax = typemax(a)
-    if (a == tmax) || (b == tmax)
-        throw(ArgumentError("array dimension size can't be the maximum representable value"))
-    end
-    if (a < 0) || (b < 0)
-        throw(ArgumentError("array dimension size can't be negative"))
-    end
     (m, o) = Base.Checked.mul_with_overflow(a, b)
-    if o
-        throw(ArgumentError("array dimensions too great, can't represent length"))
-    end
     r = Base.tail(t)::NTuple{N,Int}
-    checked_dims_impl(m, r)::Int
+    checked_dims_impl(m, r, have_overflow | o)::Tuple{Int,Bool}
 end
 
 checked_dims(::Tuple{}) = 1
 function checked_dims(t::Tuple{Int,Vararg{Int,N}}) where {N}
+    any_is_zero     = any(iszero,                 t)::Bool
+    any_is_negative = any((x -> x < false),       t)::Bool
+    any_is_typemax  = any((x -> x == typemax(x)), t)::Bool
     a = first(t)
     r = Base.tail(t)::NTuple{N,Int}
-    checked_dims_impl(a, r)::Int
+    (product, have_overflow) = checked_dims_impl(a, r, false)::Tuple{Int,Bool}
+    if any_is_negative
+        throw(ArgumentError("array dimension size can't be negative"))
+    end
+    if any_is_typemax
+        throw(ArgumentError("array dimension size can't be the maximum representable value"))
+    end
+    if have_overflow & !any_is_zero
+        throw(ArgumentError("array dimensions too great, can't represent length"))
+    end
+    product
 end
 
 # broadcasting

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -60,6 +60,22 @@ end
             @test_throws DimensionMismatch FixedSizeArray{Int}(oa)
             @test_throws DimensionMismatch FixedSizeVector{Int}(oa)
         end
+        @testset "taken from Julia's test/core.jl" begin
+            # inspired by:
+            # https://github.com/JuliaLang/julia/blob/83929ad883c97fa4376618c4559b74f6ada7a3ce/test/core.jl#L7211-L7228
+            b = prevpow(2, typemax(Int))
+            test_inferred(FixedSizeArray{Int}, FixedSizeArray{Int,3}, (undef, 0, b, b))
+            test_inferred(FixedSizeArray{Int}, FixedSizeArray{Int,3}, (undef, b, b, 0))
+            test_inferred(FixedSizeArray{Int}, FixedSizeArray{Int,5}, (undef, b, b, 0, b, b))
+            @test_throws ArgumentError FixedSizeArray{Int}(undef, b, b)
+            @test_throws ArgumentError FixedSizeArray{Int}(undef, 1, b, b)
+            @test_throws ArgumentError FixedSizeArray{Int}(undef, 0, -10)
+            @test_throws ArgumentError FixedSizeArray{Int}(undef, -10, 0)
+            @test_throws ArgumentError FixedSizeArray{Int}(undef, -1, -1)
+            @test_throws ArgumentError FixedSizeArray{Int}(undef, 0, -4, -4)
+            @test_throws ArgumentError FixedSizeArray{Int}(undef, -4, 1, 0)
+            @test_throws ArgumentError FixedSizeArray{Int}(undef, -4, -4, 1)
+        end
         @test_throws ArgumentError FixedSizeArray{Float64,1}(undef, -1)
         @test_throws ArgumentError FixedSizeArray{Float64,1}(undef, (-1,))
         @test_throws ArgumentError FixedSizeArray{Float64,2}(undef, -1, -1)