Skip to content

Commit

Permalink
better type inference for several functions taking NTuple args (Jul…
Browse files Browse the repository at this point in the history
…iaLang#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 JuliaLang#54495
  • Loading branch information
nsajko authored and lazarusA committed Aug 17, 2024
1 parent b3d6a51 commit 177437b
Show file tree
Hide file tree
Showing 4 changed files with 28 additions and 5 deletions.
8 changes: 7 additions & 1 deletion base/essentials.jl
Original file line number Diff line number Diff line change
Expand Up @@ -501,7 +501,13 @@ julia> Base.tail(())
ERROR: ArgumentError: Cannot call tail on an empty tuple.
```
"""
tail(x::Tuple) = argtail(x...)
function tail(x::Tuple{Any,Vararg})
y = argtail(x...)::Tuple
if x isa NTuple # help the type inference
y = y::NTuple
end
y
end
tail(::Tuple{}) = throw(ArgumentError("Cannot call tail on an empty tuple."))

function unwrap_unionall(@nospecialize(a))
Expand Down
2 changes: 1 addition & 1 deletion base/ntuple.jl
Original file line number Diff line number Diff line change
Expand Up @@ -95,5 +95,5 @@ end
function reverse(t::NTuple{N}) where N
ntuple(Val{N}()) do i
t[end+1-i]
end
end::NTuple
end
16 changes: 13 additions & 3 deletions base/tuple.jl
Original file line number Diff line number Diff line change
Expand Up @@ -340,9 +340,15 @@ ERROR: ArgumentError: Cannot call front on an empty tuple.
"""
function front(t::Tuple)
@inline
_front(t...)
if t === ()
throw(ArgumentError("Cannot call front on an empty tuple."))
end
r = _front(t...)::Tuple
if t isa NTuple # help the type inference
r = r::NTuple
end
r
end
_front() = throw(ArgumentError("Cannot call front on an empty tuple."))
_front(v) = ()
function _front(v, t...)
@inline
Expand Down Expand Up @@ -699,5 +705,9 @@ function circshift(x::Tuple{Any,Any,Any,Vararg{Any,N}}, shift::Integer) where {N
@inline
len = N + 3
j = mod1(shift, len)
ntuple(k -> getindex(x, k-j+ifelse(k>j,0,len)), Val(len))::Tuple
y = ntuple(k -> getindex(x, k-j+ifelse(k>j,0,len)), Val(len))::Tuple
if x isa NTuple # help the type inference
y = y::NTuple
end
y
end
7 changes: 7 additions & 0 deletions test/tuple.jl
Original file line number Diff line number Diff line change
Expand Up @@ -845,3 +845,10 @@ end
end
end
end

@testset "abstract return type inference for homogeneous tuples" begin
@test NTuple == Core.Compiler.return_type(Base.tail, Tuple{NTuple})
@test NTuple == Core.Compiler.return_type(Base.front, Tuple{NTuple})
@test NTuple == Core.Compiler.return_type(reverse, Tuple{NTuple})
@test NTuple == Core.Compiler.return_type(circshift, Tuple{NTuple,Int})
end

0 comments on commit 177437b

Please sign in to comment.