From aeb3473e2eba4df4537d33f852337191c44b8acc Mon Sep 17 00:00:00 2001 From: Jakob Nybo Nissen Date: Wed, 11 Dec 2024 15:36:28 +0100 Subject: [PATCH 1/7] Fix eltype of flatten of tuple with non-2 length In 4c076c80af, eltype of flatten of tuple was improved by computing a refined eltype at compile time. However, this implementation only worked for length-2 tuples, and errored for all others. Generalize this to all tuples. --- base/iterators.jl | 16 +++++++++++++++- test/iterators.jl | 3 +++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/base/iterators.jl b/base/iterators.jl index 6b8d9fe75e302..1bf3ad8b611e8 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -1202,7 +1202,21 @@ julia> [(x,y) for x in 0:1 for y in 'a':'c'] # collects generators involving It flatten(itr) = Flatten(itr) eltype(::Type{Flatten{I}}) where {I} = eltype(eltype(I)) -eltype(::Type{Flatten{I}}) where {I<:Union{Tuple,NamedTuple}} = promote_typejoin(map(eltype, fieldtypes(I))...) + +# For tuples, we statically know the element type of each index, so we can compute +# this at compile time. +Base.@assume_effects :foldable function eltype( + ::Type{Flatten{I}} +) where {I<:Union{Tuple,NamedTuple}} + T = Union{} + for i in fieldtypes(I) + T = Base.promote_typejoin(T, eltype(i)) + T === Any && return Any + end + T +end + +promote_typejoin(map(eltype, fieldtypes(I))...) eltype(::Type{Flatten{Tuple{}}}) = eltype(Tuple{}) IteratorEltype(::Type{Flatten{I}}) where {I} = _flatteneltype(I, IteratorEltype(I)) IteratorEltype(::Type{Flatten{Tuple{}}}) = IteratorEltype(Tuple{}) diff --git a/test/iterators.jl b/test/iterators.jl index d1e7525c43465..39a860128e7a2 100644 --- a/test/iterators.jl +++ b/test/iterators.jl @@ -515,6 +515,9 @@ end @test eltype(flatten(UnitRange{Int8}[1:2, 3:4])) == Int8 @test eltype(flatten(([1, 2], [3.0, 4.0]))) == Real @test eltype(flatten((a = [1, 2], b = Int8[3, 4]))) == Signed +@test eltype(flatten((Int[], Nothing[], Int[]))) == Union{Int, Nothing} +@test eltype(flatten((String[],))) == String +@test eltype(flatten((Int[], UInt[], Int8[],))) == Integer @test length(flatten(zip(1:3, 4:6))) == 6 @test length(flatten(1:6)) == 6 @test collect(flatten(Any[])) == Any[] From d084db47a9b0b2361b326454dfe2d4473107eab4 Mon Sep 17 00:00:00 2001 From: Jakob Nybo Nissen Date: Wed, 11 Dec 2024 15:57:34 +0100 Subject: [PATCH 2/7] Use foldable meta? --- base/iterators.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/base/iterators.jl b/base/iterators.jl index 1bf3ad8b611e8..5cd53cb6ebc2b 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -11,7 +11,8 @@ const Base = parentmodule(@__MODULE__) using .Base: @inline, Pair, Pairs, AbstractDict, IndexLinear, IndexStyle, AbstractVector, Vector, SizeUnknown, HasLength, HasShape, IsInfinite, EltypeUnknown, HasEltype, OneTo, - @propagate_inbounds, @isdefined, @boundscheck, @inbounds, Generator, IdDict, + @propagate_inbounds, @isdefined, @boundscheck, @inbounds, @_foldable_meta, + Generator, IdDict, AbstractRange, AbstractUnitRange, UnitRange, LinearIndices, TupleOrBottom, (:), |, +, -, *, !==, !, ==, !=, <=, <, >, >=, =>, missing, any, _counttuple, eachindex, ntuple, zero, prod, reduce, in, firstindex, lastindex, @@ -1205,9 +1206,8 @@ eltype(::Type{Flatten{I}}) where {I} = eltype(eltype(I)) # For tuples, we statically know the element type of each index, so we can compute # this at compile time. -Base.@assume_effects :foldable function eltype( - ::Type{Flatten{I}} -) where {I<:Union{Tuple,NamedTuple}} +function eltype(::Type{Flatten{I}}) where {I<:Union{Tuple,NamedTuple}} + @_foldable_meta T = Union{} for i in fieldtypes(I) T = Base.promote_typejoin(T, eltype(i)) From 184153e00def48bb516785fae34c0c75cfa38ed6 Mon Sep 17 00:00:00 2001 From: Jakob Nybo Nissen Date: Wed, 11 Dec 2024 16:03:41 +0100 Subject: [PATCH 3/7] Delete dead code --- base/iterators.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/base/iterators.jl b/base/iterators.jl index 5cd53cb6ebc2b..4eddbb471526a 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -1216,7 +1216,6 @@ function eltype(::Type{Flatten{I}}) where {I<:Union{Tuple,NamedTuple}} T end -promote_typejoin(map(eltype, fieldtypes(I))...) eltype(::Type{Flatten{Tuple{}}}) = eltype(Tuple{}) IteratorEltype(::Type{Flatten{I}}) where {I} = _flatteneltype(I, IteratorEltype(I)) IteratorEltype(::Type{Flatten{Tuple{}}}) = IteratorEltype(Tuple{}) From 6eae199ab75161a14ea98de8f3c71830d714a24a Mon Sep 17 00:00:00 2001 From: Jakob Nybo Nissen Date: Thu, 12 Dec 2024 13:24:44 +0100 Subject: [PATCH 4/7] Use recursion to encourage constant folding --- base/iterators.jl | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/base/iterators.jl b/base/iterators.jl index 4eddbb471526a..178c3fd9d1a59 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -11,12 +11,12 @@ const Base = parentmodule(@__MODULE__) using .Base: @inline, Pair, Pairs, AbstractDict, IndexLinear, IndexStyle, AbstractVector, Vector, SizeUnknown, HasLength, HasShape, IsInfinite, EltypeUnknown, HasEltype, OneTo, - @propagate_inbounds, @isdefined, @boundscheck, @inbounds, @_foldable_meta, - Generator, IdDict, + @propagate_inbounds, @isdefined, @boundscheck, @inbounds, Generator, IdDict, AbstractRange, AbstractUnitRange, UnitRange, LinearIndices, TupleOrBottom, (:), |, +, -, *, !==, !, ==, !=, <=, <, >, >=, =>, missing, any, _counttuple, eachindex, ntuple, zero, prod, reduce, in, firstindex, lastindex, - tail, fieldtypes, min, max, minimum, zero, oneunit, promote, promote_shape, LazyString + tail, fieldtypes, min, max, minimum, zero, oneunit, promote, promote_shape, LazyString, + tuple_type_head, tuple_type_tail using Core: @doc using .Base: @@ -1206,16 +1206,20 @@ eltype(::Type{Flatten{I}}) where {I} = eltype(eltype(I)) # For tuples, we statically know the element type of each index, so we can compute # this at compile time. -function eltype(::Type{Flatten{I}}) where {I<:Union{Tuple,NamedTuple}} - @_foldable_meta - T = Union{} - for i in fieldtypes(I) - T = Base.promote_typejoin(T, eltype(i)) - T === Any && return Any - end - T +eltype(::Type{Flatten{I}}) where {I<:Tuple} = _flatten_eltype(Union{}, I) + +function eltype(::Type{Flatten{I}}) where {I<:NamedTuple{<:Any, T}} where T + _flatten_eltype(Union{}, T) end +function _flatten_eltype(T::Type, I::Type{<:Tuple}) + T2 = promote_typejoin(T, eltype(tuple_type_head(I))) + T2 === Any && return Any + _flatten_eltype(T2, tuple_type_tail(I)) +end + +_flatten_eltype(T::Type, I::Type{Tuple{}}) = T + eltype(::Type{Flatten{Tuple{}}}) = eltype(Tuple{}) IteratorEltype(::Type{Flatten{I}}) where {I} = _flatteneltype(I, IteratorEltype(I)) IteratorEltype(::Type{Flatten{Tuple{}}}) = IteratorEltype(Tuple{}) From a989a70248b07c310f57edc7b0166a795a72ab67 Mon Sep 17 00:00:00 2001 From: Jakob Nybo Nissen Date: Thu, 12 Dec 2024 19:52:48 +0100 Subject: [PATCH 5/7] Back to recursive implementation --- base/iterators.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/base/iterators.jl b/base/iterators.jl index 178c3fd9d1a59..e9e2ec68671d5 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -16,7 +16,7 @@ using .Base: (:), |, +, -, *, !==, !, ==, !=, <=, <, >, >=, =>, missing, any, _counttuple, eachindex, ntuple, zero, prod, reduce, in, firstindex, lastindex, tail, fieldtypes, min, max, minimum, zero, oneunit, promote, promote_shape, LazyString, - tuple_type_head, tuple_type_tail + tuple_type_tail using Core: @doc using .Base: @@ -1212,8 +1212,8 @@ function eltype(::Type{Flatten{I}}) where {I<:NamedTuple{<:Any, T}} where T _flatten_eltype(Union{}, T) end -function _flatten_eltype(T::Type, I::Type{<:Tuple}) - T2 = promote_typejoin(T, eltype(tuple_type_head(I))) +function _flatten_eltype(T::Type, I::Type{<:Tuple{E, Vararg{Any}}}) where E + T2 = promote_typejoin(T, eltype(E)) T2 === Any && return Any _flatten_eltype(T2, tuple_type_tail(I)) end From 56666ebdc9923215f78c0bff35e372d547e47b24 Mon Sep 17 00:00:00 2001 From: Jakob Nybo Nissen Date: Sat, 14 Dec 2024 16:24:36 +0100 Subject: [PATCH 6/7] Use afoldl --- base/iterators.jl | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/base/iterators.jl b/base/iterators.jl index e9e2ec68671d5..c6278e6284d70 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -16,7 +16,7 @@ using .Base: (:), |, +, -, *, !==, !, ==, !=, <=, <, >, >=, =>, missing, any, _counttuple, eachindex, ntuple, zero, prod, reduce, in, firstindex, lastindex, tail, fieldtypes, min, max, minimum, zero, oneunit, promote, promote_shape, LazyString, - tuple_type_tail + afoldl using Core: @doc using .Base: @@ -1206,21 +1206,10 @@ eltype(::Type{Flatten{I}}) where {I} = eltype(eltype(I)) # For tuples, we statically know the element type of each index, so we can compute # this at compile time. -eltype(::Type{Flatten{I}}) where {I<:Tuple} = _flatten_eltype(Union{}, I) - -function eltype(::Type{Flatten{I}}) where {I<:NamedTuple{<:Any, T}} where T - _flatten_eltype(Union{}, T) -end - -function _flatten_eltype(T::Type, I::Type{<:Tuple{E, Vararg{Any}}}) where E - T2 = promote_typejoin(T, eltype(E)) - T2 === Any && return Any - _flatten_eltype(T2, tuple_type_tail(I)) +function eltype(::Type{Flatten{I}}) where {I<:Union{Tuple,NamedTuple}} + afoldl((T, i) -> promote_typejoin(T, eltype(i)), Union{}, fieldtypes(I)...) end -_flatten_eltype(T::Type, I::Type{Tuple{}}) = T - -eltype(::Type{Flatten{Tuple{}}}) = eltype(Tuple{}) IteratorEltype(::Type{Flatten{I}}) where {I} = _flatteneltype(I, IteratorEltype(I)) IteratorEltype(::Type{Flatten{Tuple{}}}) = IteratorEltype(Tuple{}) _flatteneltype(I, ::HasEltype) = IteratorEltype(eltype(I)) From 5c7c0e44ec6d05f7bb37b2301758647e3d3f3880 Mon Sep 17 00:00:00 2001 From: Jakob Nybo Nissen Date: Sat, 14 Dec 2024 17:22:18 +0100 Subject: [PATCH 7/7] Update test/iterators.jl Co-authored-by: Neven Sajko --- test/iterators.jl | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/test/iterators.jl b/test/iterators.jl index 39a860128e7a2..1feccf5fb1d3e 100644 --- a/test/iterators.jl +++ b/test/iterators.jl @@ -518,6 +518,11 @@ end @test eltype(flatten((Int[], Nothing[], Int[]))) == Union{Int, Nothing} @test eltype(flatten((String[],))) == String @test eltype(flatten((Int[], UInt[], Int8[],))) == Integer +@test eltype(flatten((; a = Int[], b = Nothing[], c = Int[]))) == Union{Int, Nothing} +@test eltype(flatten((; a = String[],))) == String +@test eltype(flatten((; a = Int[], b = UInt[], c = Int8[],))) == Integer +@test eltype(flatten(())) == Union{} +@test eltype(flatten((;))) == Union{} @test length(flatten(zip(1:3, 4:6))) == 6 @test length(flatten(1:6)) == 6 @test collect(flatten(Any[])) == Any[]