improve inferrability of tuple slicing (#39074)

* RFC: improve inferrability of tuple slicing

This makes slicing tuples type stable in a lot more cases. This now only calls `ntuple` if the resulting tuple has no more than 10 elements, as just relying on the fallback in `ntuple` did have some compile-time overhead in my artificial benchmarks. I also widened the signature to `AbstractUnitRange`, since I don't see why we shouldn't have this for ranges like `OneTo` as well.

* add check for offset ranges and more tests

* add guard if OffsetArrays is already imported

Co-authored-by: Matt Bauman <mbauman@juliacomputing.com>

* import OffsetArrays correctly

Co-authored-by: Matt Bauman <mbauman@juliacomputing.com>

Author: simeonschaub

base/range.jl

@@ -299,16 +299,18 @@ unitrange(x) = UnitRange(x)
 
 if isdefined(Main, :Base)
     # Constant-fold-able indexing into tuples to functionally expose Base.tail and Base.front
-    function getindex(@nospecialize(t::Tuple), r::UnitRange)
+    function getindex(@nospecialize(t::Tuple), r::AbstractUnitRange)
         @_inline_meta
-        r.start > r.stop && return ()
-        if r.start == 1
-            r.stop == length(t)   && return t
-            r.stop == length(t)-1 && return front(t)
-            r.stop == length(t)-2 && return front(front(t))
-        elseif r.stop == length(t)
-            r.start == 2 && return tail(t)
-            r.start == 3 && return tail(tail(t))
+        require_one_based_indexing(r)
+        if length(r) <= 10
+            return ntuple(i -> t[i + first(r) - 1], length(r))
+        elseif first(r) == 1
+            last(r) == length(t)   && return t
+            last(r) == length(t)-1 && return front(t)
+            last(r) == length(t)-2 && return front(front(t))
+        elseif last(r) == length(t)
+            first(r) == 2 && return tail(t)
+            first(r) == 3 && return tail(tail(t))
         end
         return (eltype(t)[t[ri] for ri in r]...,)
     end

test/tuple.jl

@@ -1,5 +1,8 @@
 # This file is a part of Julia. License is MIT: https://julialang.org/license
 
+isdefined(Main, :OffsetArrays) || @eval Main include("testhelpers/OffsetArrays.jl")
+using .Main.OffsetArrays
+
 struct BitPerm_19352
     p::NTuple{8,UInt8}
     function BitPerm(p::NTuple{8,UInt8})
@@ -598,3 +601,24 @@ end
 # issue #38837
 f38837(xs) = map((F,x)->F(x), (Float32, Float64), xs)
 @test @inferred(f38837((1,2))) === (1.0f0, 2.0)
+
+@testset "indexing with UnitRanges" begin
+    f(t) = t[3:end-2]
+    @test @inferred(f(Tuple(1:10))) === Tuple(3:8)
+    @test @inferred(f((true, 2., 3, 4f0, 0x05, 6, 7.))) === (3, 4f0, 0x05)
+
+    f(t) = t[Base.OneTo(5)]
+    @test @inferred(f(Tuple(1:10))) === Tuple(1:5)
+    @test @inferred(f((true, 2., 3, 4f0, 0x05, 6, 7.))) === (true, 2., 3, 4f0, 0x05)
+
+    @test @inferred((t -> t[1:end])(Tuple(1:15))) === Tuple(1:15)
+    @test @inferred((t -> t[2:end])(Tuple(1:15))) === Tuple(2:15)
+    @test @inferred((t -> t[3:end])(Tuple(1:15))) === Tuple(3:15)
+    @test @inferred((t -> t[1:end-1])(Tuple(1:15))) === Tuple(1:14)
+    @test @inferred((t -> t[1:end-2])(Tuple(1:15))) === Tuple(1:13)
+    @test @inferred((t -> t[3:2])(Tuple(1:15))) === ()
+
+    @test_throws BoundsError (1, 2)[1:4]
+    @test_throws BoundsError (1, 2)[0:2]
+    @test_throws ArgumentError (1, 2)[OffsetArrays.IdOffsetRange(1:2, -1)]
+end