From b1bb0a6ed3d78d787dead891d44af91316372fa6 Mon Sep 17 00:00:00 2001 From: N5N3 <2642243996@qq.com> Date: Thu, 12 May 2022 09:24:35 +0800 Subject: [PATCH] Fix for `length(::StepRange{T})` where `T` isa `Union` And define `firstindex` accordingly. Co-Authored-By: Jameson Nash --- base/range.jl | 46 ++++++++++++++++++++++------------------------ test/ranges.jl | 10 +++++++++- 2 files changed, 31 insertions(+), 25 deletions(-) diff --git a/base/range.jl b/base/range.jl index 7ba7807f80ff8..4b2bf18e9f634 100644 --- a/base/range.jl +++ b/base/range.jl @@ -691,10 +691,6 @@ axes(r::AbstractRange) = (oneto(length(r)),) # Needed to ensure `has_offset_axes` can constant-fold. has_offset_axes(::StepRange) = false -let baseints = Union{Int8,UInt8,Int16,UInt16,Int32,UInt32,Int64,UInt64,Int128,UInt128} - global firstindex - firstindex(::StepRange{T,<:baseints}) where {T<:baseints} = sizeof(T) < sizeof(Int) ? 1 : one(T) -end # n.b. checked_length for these is defined iff checked_add and checked_sub are # defined between the relevant types @@ -757,57 +753,58 @@ length(r::OneTo) = Integer(r.stop - zero(r.stop)) length(r::StepRangeLen) = r.len length(r::LinRange) = r.len -let bigints = Union{Int, UInt, Int64, UInt64, Int128, UInt128} - global length, checked_length +let bigints = Union{Int, UInt, Int64, UInt64, Int128, UInt128}, + smallints = (Int === Int64 ? + Union{Int8, UInt8, Int16, UInt16, Int32, UInt32} : + Union{Int8, UInt8, Int16, UInt16}), + bitints = Union{bigints, smallints} + global length, checked_length, firstindex # compile optimization for which promote_type(T, Int) == T length(r::OneTo{T}) where {T<:bigints} = r.stop # slightly more accurate length and checked_length in extreme cases # (near typemax) for types with known `unsigned` functions function length(r::OrdinalRange{T}) where T<:bigints s = step(r) - isempty(r) && return zero(T) diff = last(r) - first(r) + isempty(r) && return zero(diff) # if |s| > 1, diff might have overflowed, but unsigned(diff)÷s should # therefore still be valid (if the result is representable at all) # n.b. !(s isa T) if s isa Unsigned || -1 <= s <= 1 || s == -s - a = div(diff, s) % T + a = div(diff, s) % typeof(diff) elseif s < 0 - a = div(unsigned(-diff), -s) % T + a = div(unsigned(-diff), -s) % typeof(diff) else - a = div(unsigned(diff), s) % T + a = div(unsigned(diff), s) % typeof(diff) end - return a + oneunit(T) + return a + oneunit(a) end function checked_length(r::OrdinalRange{T}) where T<:bigints s = step(r) - isempty(r) && return zero(T) stop, start = last(r), first(r) + ET = promote_type(typeof(stop), typeof(start)) + isempty(r) && return zero(ET) # n.b. !(s isa T) if s > 1 diff = stop - start - a = convert(T, div(unsigned(diff), s)) + a = convert(ET, div(unsigned(diff), s)) elseif s < -1 diff = start - stop - a = convert(T, div(unsigned(diff), -s)) + a = convert(ET, div(unsigned(diff), -s)) elseif s > 0 - a = div(checked_sub(stop, start), s) + a = convert(ET, div(checked_sub(stop, start), s)) else - a = div(checked_sub(start, stop), -s) + a = convert(ET, div(checked_sub(start, stop), -s)) end - return checked_add(convert(T, a), oneunit(T)) + return checked_add(a, oneunit(a)) end -end + firstindex(r::StepRange{<:bigints,<:bitints}) = one(last(r)-first(r)) -# some special cases to favor default Int type -let smallints = (Int === Int64 ? - Union{Int8, UInt8, Int16, UInt16, Int32, UInt32} : - Union{Int8, UInt8, Int16, UInt16}) - global length, checked_length - # n.b. !(step isa T) + # some special cases to favor default Int type function length(r::OrdinalRange{<:smallints}) s = step(r) isempty(r) && return 0 + # n.b. !(step isa T) return Int(div(Int(last(r)) - Int(first(r)), s)) + 1 end length(r::AbstractUnitRange{<:smallints}) = Int(last(r)) - Int(first(r)) + 1 @@ -815,6 +812,7 @@ let smallints = (Int === Int64 ? checked_length(r::OrdinalRange{<:smallints}) = length(r) checked_length(r::AbstractUnitRange{<:smallints}) = length(r) checked_length(r::OneTo{<:smallints}) = length(r) + firstindex(::StepRange{<:smallints,<:bitints}) = 1 end first(r::OrdinalRange{T}) where {T} = convert(T, r.start) diff --git a/test/ranges.jl b/test/ranges.jl index 975991befc354..404a908adf6b3 100644 --- a/test/ranges.jl +++ b/test/ranges.jl @@ -2036,6 +2036,12 @@ end @test typeof(length(r1)) == typeof(checked_length(r1)) == typeof(length(r2)) == typeof(checked_length(r2)) end + SR = StepRange{Union{Int64,Int128},Int} + test_length(r, l) = length(r) === checked_length(r) === l + @test test_length(SR(Int64(1), 1, Int128(1)), Int128(1)) + @test test_length(SR(Int64(1), 1, Int128(0)), Int128(0)) + @test test_length(SR(Int64(1), 1, Int64(1)), Int64(1)) + @test test_length(SR(Int64(1), 1, Int64(0)), Int64(0)) end @testset "LinRange eltype for element types that wrap integers" begin @@ -2350,10 +2356,12 @@ end @test length(range(1, length=typemax(Int128))) === typemax(Int128) -@testset "firstindex(::StepRange{T,T})" begin +@testset "firstindex(::StepRange{<:Base.BitInteger})" begin test_firstindex(x) = firstindex(x) === first(Base.axes1(x)) for T in Base.BitInteger_types, S in Base.BitInteger_types @test test_firstindex(StepRange{T,S}(1, 1, 1)) @test test_firstindex(StepRange{T,S}(1, 1, 0)) end + @test test_firstindex(StepRange{Union{Int64,Int128},Int}(Int64(1), 1, Int128(1))) + @test test_firstindex(StepRange{Union{Int64,Int128},Int}(Int64(1), 1, Int128(0))) end