Make tuple t[2:end] inferrable through constant prop (#31138)

Spurred by #31132, I was curious if our constant propagation was robust enough to handle a peephole optimization here -- and it is! This removes the need to export front/tail by simply allowing folks to get inferrable results when they spell these operations as `t[2:end]` and `t[1:end-1]`.  I additionally allowed for `t[3:end]`, `t[1:end-2]`, and the trivial `t[1:end]` as I have used `tail(tail(t))` in the past.

Author: mbauman

base/range.jl

@@ -289,16 +289,19 @@ unitrange_last(start::T, stop::T) where {T} =
                           convert(T,start-oneunit(stop-start)))
 
 if isdefined(Main, :Base)
-    function getindex(t::Tuple, r::AbstractUnitRange{<:Real})
-        n = length(r)
-        n == 0 && return ()
-        a = Vector{eltype(t)}(undef, n)
-        o = first(r) - 1
-        for i = 1:n
-            el = t[o + i]
-            @inbounds a[i] = el
+    # Constant-fold-able indexing into tuples to functionally expose Base.tail and Base.front
+    function getindex(@nospecialize(t::Tuple), r::UnitRange)
+        @_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))
         end
-        (a...,)
+        return (eltype(t)[t[ri] for ri in r]...,)
     end
 end
 

base/tuple.jl

@@ -23,7 +23,7 @@ size(@nospecialize(t::Tuple), d::Integer) = (d == 1) ? length(t) : throw(Argumen
 axes(@nospecialize t::Tuple) = (OneTo(length(t)),)
 @eval getindex(@nospecialize(t::Tuple), i::Int) = getfield(t, i, $(Expr(:boundscheck)))
 @eval getindex(@nospecialize(t::Tuple), i::Real) = getfield(t, convert(Int, i), $(Expr(:boundscheck)))
-getindex(t::Tuple, r::AbstractArray{<:Any,1}) = ([t[ri] for ri in r]...,)
+getindex(t::Tuple, r::AbstractArray{<:Any,1}) = (eltype(t)[t[ri] for ri in r]...,)
 getindex(t::Tuple, b::AbstractArray{Bool,1}) = length(b) == length(t) ? getindex(t, findall(b)) : throw(BoundsError(t, b))
 getindex(t::Tuple, c::Colon) = t
 

test/tuple.jl

@@ -512,4 +512,60 @@ end
 
     @test Base.setindex((1, 2, 4), 4, true) === (4, 2, 4)
     @test_throws BoundsError Base.setindex((1, 2), 2, false)
-end
+end
+
+@testset "inferrable range indexing with constant values" begin
+    whole(t) = t[1:end]
+    tail(t) = t[2:end]
+    ttail(t) = t[3:end]
+    front(t) = t[1:end-1]
+    ffront(t) = t[1:end-2]
+
+    @test @inferred( whole(())) == ()
+    @test @inferred(  tail(())) == ()
+    @test @inferred( ttail(())) == ()
+    @test @inferred( front(())) == ()
+    @test @inferred(ffront(())) == ()
+
+    @test @inferred( whole((1,))) == (1,)
+    @test @inferred(  tail((1,))) == ()
+    @test @inferred( ttail((1,))) == ()
+    @test @inferred( front((1,))) == ()
+    @test @inferred(ffront((1,))) == ()
+
+    @test @inferred( whole((1,2.0))) == (1,2.0)
+    @test @inferred(  tail((1,2.0))) == (2.0,)
+    @test @inferred( ttail((1,2.0))) == ()
+    @test @inferred( front((1,2.0))) == (1.0,)
+    @test @inferred(ffront((1,2.0))) == ()
+
+    @test @inferred( whole((1,2.0,3//1))) == (1,2.0,3//1)
+    @test @inferred(  tail((1,2.0,3//1))) == (2.0,3//1)
+    @test @inferred( ttail((1,2.0,3//1))) == (3//1,)
+    @test @inferred( front((1,2.0,3//1))) == (1,2.0)
+    @test @inferred(ffront((1,2.0,3//1))) == (1,)
+
+    @test @inferred( whole((1,2.0,3//1,0x04))) == (1,2.0,3//1,0x04)
+    @test @inferred(  tail((1,2.0,3//1,0x04))) == (2.0,3//1,0x04)
+    @test @inferred( ttail((1,2.0,3//1,0x04))) == (3//1,0x04)
+    @test @inferred( front((1,2.0,3//1,0x04))) == (1,2.0,3//1)
+    @test @inferred(ffront((1,2.0,3//1,0x04))) == (1,2.0)
+
+    @test (1,)[0:-1] == ()
+    @test (1,)[1:0] == ()
+    @test (1,)[2:1] == ()
+    @test (1,2.0)[0:-1] == ()
+    @test (1,2.0)[1:0] == ()
+    @test (1,2.0)[2:1] == ()
+    @test (1,2.0)[3:2] == ()
+
+    @test_throws BoundsError (1,)[2:2]
+    @test_throws BoundsError (1,)[1:2]
+    @test_throws BoundsError (1,)[0:1]
+    @test_throws BoundsError (1,)[0:0]
+    @test_throws BoundsError (1,2.0)[3:3]
+    @test_throws BoundsError (1,2.0)[1:3]
+    @test_throws BoundsError (1,2.0)[0:2]
+    @test_throws BoundsError (1,2.0)[0:1]
+    @test_throws BoundsError (1,2.0)[0:0]
+end