Skip to content

Commit efe4535

Browse files
authored
split_rest for Tuple: slighly improve inference (#59588)
* Introduce function `_split_tuple` as the core of the implementation of `split_rest` for `Tuple`. * Improves inference for `split_rest` in some cases because `_split_tuple` only accepts `i::Int`. * Avoids a bit of code duplication in `split_rest`. * Apply `_split_tuple` within the generic method making `Fix` callable, where the same pattern is implemented inline.
1 parent c6ccf20 commit efe4535

File tree

3 files changed

+28
-2
lines changed

3 files changed

+28
-2
lines changed

base/operators.jl

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1194,7 +1194,8 @@ end
11941194

11951195
function (f::Fix{N})(args::Vararg{Any,M}; kws...) where {N,M}
11961196
M < N-1 && throw(ArgumentError(LazyString("expected at least ", N-1, " arguments to `Fix{", N, "}`, but got ", M)))
1197-
return f.f(args[begin:begin+(N-2)]..., f.x, args[begin+(N-1):end]...; kws...)
1197+
(left, right) = _split_tuple(args, N-1)
1198+
return f.f(left..., f.x, right...; kws...)
11981199
end
11991200

12001201
# Special cases for improved constant propagation

base/tuple.jl

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -261,7 +261,9 @@ function _split_rest(a::Union{AbstractArray, Core.SimpleVector}, n::Int)
261261
return a[begin:end-n], a[end-n+1:end]
262262
end
263263

264-
@eval split_rest(t::Tuple, n::Int, i=1) = ($(Expr(:meta, :aggressive_constprop)); (t[i:end-n], t[end-n+1:end]))
264+
@eval _split_tuple(t::Tuple, n::Int, i::Int=1) = ($(Expr(:meta, :aggressive_constprop)); (t[i:n], t[n+1:end]))
265+
266+
@eval split_rest(t::Tuple, n::Int, i=1) = ($(Expr(:meta, :aggressive_constprop)); _split_tuple(t, length(t)-n, Int(i)))
265267

266268
# Use dispatch to avoid a branch in first
267269
first(::Tuple{}) = throw(ArgumentError("tuple must be non-empty"))

test/tuple.jl

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -815,6 +815,29 @@ namedtup = (;a=1, b=2, c=3)
815815
NamedTuple{(:c,), Tuple{Int}},
816816
}
817817

818+
@testset "`Base.split_rest(::Tuple, ::Vararg)` return type inference" begin
819+
let f(t) = Base.split_rest(t, 3)
820+
tuple_types_of_length(n::Int) = (NTuple{n, Any}, NTuple{n}, NTuple{n, Float32})
821+
@testset "inferred return type must subtype `NTuple{2, Tuple}`" begin
822+
for T in (
823+
Tuple, Tuple{Vararg{Float32}}, # any length
824+
Tuple{Any, Vararg{Any}}, (Tuple{T, Vararg{T}} where {T}), Tuple{Float32, Vararg{Float32}}, # length greater than one
825+
tuple_types_of_length(5)..., # length five
826+
)
827+
@test Base.infer_return_type(f, Tuple{T}) <: NTuple{2, Tuple}
828+
for S in (Tuple{T, Any, Any}, Tuple{T, Any}, Tuple{T, Int, Any}, Tuple{T, Int}, Tuple{T, Int, Int})
829+
@test Base.infer_return_type(Base.split_rest, S) <: NTuple{2, Tuple}
830+
end
831+
end
832+
end
833+
@testset "with exactly-known length: `5 == 2 + 3`" begin
834+
for T in tuple_types_of_length(5)
835+
@test Base.infer_return_type(f, Tuple{T}) <: Tuple{Tuple{Any, Any}, Tuple{Any, Any, Any}}
836+
end
837+
end
838+
end
839+
end
840+
818841
# Make sure that tuple iteration is foldable
819842
@test Core.Compiler.is_foldable(Base.infer_effects(iterate, Tuple{NTuple{4, Float64}, Int}))
820843
@test Core.Compiler.is_foldable(Base.infer_effects(eltype, Tuple{Tuple}))

0 commit comments

Comments
 (0)