Skip to content

Commit

Permalink
Use checked_mul in length(::ProductIterator) (#43429)
Browse files Browse the repository at this point in the history
Co-authored-by: Steven G. Johnson <stevenj@mit.edu>
  • Loading branch information
fonsp and stevengj authored Nov 30, 2022
1 parent c46834b commit 04a4edb
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 2 deletions.
8 changes: 6 additions & 2 deletions base/iterators.jl
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,17 @@ using .Base:
@propagate_inbounds, @isdefined, @boundscheck, @inbounds, Generator,
AbstractRange, AbstractUnitRange, UnitRange, LinearIndices,
(:), |, +, -, *, !==, !, ==, !=, <=, <, >, >=, missing,
any, _counttuple, eachindex, ntuple, zero, prod, in, firstindex, lastindex,
any, _counttuple, eachindex, ntuple, zero, prod, reduce, in, firstindex, lastindex,
tail, fieldtypes, min, max, minimum, zero, oneunit, promote, promote_shape
using Core: @doc

if Base !== Core.Compiler
using .Base:
cld, fld, SubArray, view, resize!, IndexCartesian
using .Base.Checked: checked_mul
else
# Checked.checked_mul is not available during bootstrapping:
const checked_mul = *
end

import .Base:
Expand Down Expand Up @@ -1055,7 +1059,7 @@ _prod_axes1(a, A) =
throw(ArgumentError("Cannot compute indices for object of type $(typeof(a))"))

ndims(p::ProductIterator) = length(axes(p))
length(P::ProductIterator) = prod(size(P))
length(P::ProductIterator) = reduce(checked_mul, size(P); init=1)

IteratorEltype(::Type{ProductIterator{Tuple{}}}) = HasEltype()
IteratorEltype(::Type{ProductIterator{Tuple{I}}}) where {I} = IteratorEltype(I)
Expand Down
6 changes: 6 additions & 0 deletions test/iterators.jl
Original file line number Diff line number Diff line change
Expand Up @@ -336,21 +336,25 @@ let a = 1:2,
c = Int32(1):Int32(0)

# length
@test length(product()) == 1
@test length(product(a)) == 2
@test length(product(a, b)) == 20
@test length(product(a, b, c)) == 0

# size
@test size(product()) == tuple()
@test size(product(a)) == (2,)
@test size(product(a, b)) == (2, 10)
@test size(product(a, b, c)) == (2, 10, 0)

# eltype
@test eltype(product()) == Tuple{}
@test eltype(product(a)) == Tuple{Int}
@test eltype(product(a, b)) == Tuple{Int, Float64}
@test eltype(product(a, b, c)) == Tuple{Int, Float64, Int32}

# ndims
@test ndims(product()) == 0
@test ndims(product(a)) == 1
@test ndims(product(a, b)) == 2
@test ndims(product(a, b, c)) == 3
Expand Down Expand Up @@ -425,6 +429,8 @@ let a = 1:2,
@test_throws ArgumentError size(product(itr))
@test_throws ArgumentError ndims(product(itr))
end

@test_throws OverflowError length(product(1:typemax(Int), 1:typemax(Int)))
end

# IteratorSize trait business
Expand Down

0 comments on commit 04a4edb

Please sign in to comment.