Skip to content

Commit 2179f14

Browse files
authored
better type inference for several functions taking NTuple args (#55124)
Improves the abstract return type inference for these functions, for homogeneous tuple arguments: * `tail` * `front` * `reverse` * `circshift` Example: ```julia f(g, t::Type{<:Tuple}) = println(Core.Compiler.return_type(g, t)) f(Base.tail, Tuple{NTuple}) f(Base.front, Tuple{NTuple}) f(reverse, Tuple{NTuple}) f(circshift, Tuple{NTuple,Int}) ``` Results before: ```julia Tuple Tuple Tuple Tuple ``` Results after: ```julia NTuple{N, T} where {N, T} NTuple{N, T} where {N, T} NTuple{N, T} where {N, T} NTuple{N, T} where {N, T} ``` Updates #54495
1 parent 5092379 commit 2179f14

File tree

4 files changed

+28
-5
lines changed

4 files changed

+28
-5
lines changed

base/essentials.jl

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -501,7 +501,13 @@ julia> Base.tail(())
501501
ERROR: ArgumentError: Cannot call tail on an empty tuple.
502502
```
503503
"""
504-
tail(x::Tuple) = argtail(x...)
504+
function tail(x::Tuple{Any,Vararg})
505+
y = argtail(x...)::Tuple
506+
if x isa NTuple # help the type inference
507+
y = y::NTuple
508+
end
509+
y
510+
end
505511
tail(::Tuple{}) = throw(ArgumentError("Cannot call tail on an empty tuple."))
506512

507513
function unwrap_unionall(@nospecialize(a))

base/ntuple.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,5 +95,5 @@ end
9595
function reverse(t::NTuple{N}) where N
9696
ntuple(Val{N}()) do i
9797
t[end+1-i]
98-
end
98+
end::NTuple
9999
end

base/tuple.jl

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -340,9 +340,15 @@ ERROR: ArgumentError: Cannot call front on an empty tuple.
340340
"""
341341
function front(t::Tuple)
342342
@inline
343-
_front(t...)
343+
if t === ()
344+
throw(ArgumentError("Cannot call front on an empty tuple."))
345+
end
346+
r = _front(t...)::Tuple
347+
if t isa NTuple # help the type inference
348+
r = r::NTuple
349+
end
350+
r
344351
end
345-
_front() = throw(ArgumentError("Cannot call front on an empty tuple."))
346352
_front(v) = ()
347353
function _front(v, t...)
348354
@inline
@@ -699,5 +705,9 @@ function circshift(x::Tuple{Any,Any,Any,Vararg{Any,N}}, shift::Integer) where {N
699705
@inline
700706
len = N + 3
701707
j = mod1(shift, len)
702-
ntuple(k -> getindex(x, k-j+ifelse(k>j,0,len)), Val(len))::Tuple
708+
y = ntuple(k -> getindex(x, k-j+ifelse(k>j,0,len)), Val(len))::Tuple
709+
if x isa NTuple # help the type inference
710+
y = y::NTuple
711+
end
712+
y
703713
end

test/tuple.jl

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -845,3 +845,10 @@ end
845845
end
846846
end
847847
end
848+
849+
@testset "abstract return type inference for homogeneous tuples" begin
850+
@test NTuple == Core.Compiler.return_type(Base.tail, Tuple{NTuple})
851+
@test NTuple == Core.Compiler.return_type(Base.front, Tuple{NTuple})
852+
@test NTuple == Core.Compiler.return_type(reverse, Tuple{NTuple})
853+
@test NTuple == Core.Compiler.return_type(circshift, Tuple{NTuple,Int})
854+
end

0 commit comments

Comments
 (0)