improve inferrabilities by telling the compiler relational invariants

Our compiler doesn't understand these relations automatically yet.
JuliaLang · Apr 25, 2021 · ca13b0d · ca13b0d
1 parent 248c02f
commit ca13b0d
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 14 deletions.
diff --git a/base/multidimensional.jl b/base/multidimensional.jl
@@ -125,9 +125,11 @@ module IteratorsMD
 
     # comparison
     @inline isless(I1::CartesianIndex{N}, I2::CartesianIndex{N}) where {N} = _isless(0, I1.I, I2.I)
-    @inline function _isless(ret, I1::NTuple{N,Int}, I2::NTuple{N,Int}) where N
-        newret = ifelse(ret==0, icmp(I1[N], I2[N]), ret)
-        _isless(newret, Base.front(I1), Base.front(I2))
+    @inline function _isless(ret, I1::Tuple{Int,Vararg{Int,N}}, I2::Tuple{Int,Vararg{Int,N}}) where {N}
+        newret = ifelse(ret==0, icmp(last(I1), last(I2)), ret)
+        t1, t2 = Base.front(I1), Base.front(I2)
+        # avoid dynamic dispatch by telling the compiler relational invariants
+        return isa(t1, Tuple{}) ? _isless(newret, (), ()) : _isless(newret, t1, t2::Tuple{Int,Vararg{Int}})
     end
     _isless(ret, ::Tuple{}, ::Tuple{}) = ifelse(ret==1, true, false)
     icmp(a, b) = ifelse(isless(a,b), 1, ifelse(a==b, 0, -1))
@@ -394,19 +396,21 @@ module IteratorsMD
     # `iterate` returns `Union{Nothing, Tuple}`, we explicitly pass a `valid` flag to eliminate
     # the type instability inside the core `__inc` logic, and this gives better runtime performance.
     __inc(::Tuple{}, ::Tuple{}) = false, ()
-    @inline function __inc(state::Tuple{Int}, indices::Tuple{<:OrdinalRange})
+    @inline function __inc(state::Tuple{Int}, indices::Tuple{T}) where {T<:OrdinalRange{Int,Int}}
         rng = indices[1]
         I = state[1] + step(rng)
         valid = __is_valid_range(I, rng) && state[1] != last(rng)
         return valid, (I, )
     end
-    @inline function __inc(state, indices)
+    @inline function __inc(state::Tuple{Int,Int,Vararg{Int,N}}, indices::Tuple{Any,Any,Vararg{Any,N}}) where {N}
         rng = indices[1]
         I = state[1] + step(rng)
         if __is_valid_range(I, rng) && state[1] != last(rng)
             return true, (I, tail(state)...)
         end
-        valid, I = __inc(tail(state), tail(indices))
+        t1, t2 = tail(state), tail(indices)
+        # avoid dynamic dispatch by telling the compiler relational invariants
+        valid, I = isa(t1, Tuple{Int}) ? __inc(t1, t2::Tuple{Any}) : __inc(t1, t2::Tuple{Any,Any,Vararg{Any}})
         return valid, (first(rng), I...)
     end
 
@@ -505,20 +509,21 @@ module IteratorsMD
 
     # decrement post check to avoid integer overflow
     @inline __dec(::Tuple{}, ::Tuple{}) = false, ()
-    @inline function __dec(state::Tuple{Int}, indices::Tuple{<:OrdinalRange})
+    @inline function __dec(state::Tuple{Int}, indices::Tuple{T}) where {T<:OrdinalRange{Int,Int}}
         rng = indices[1]
         I = state[1] - step(rng)
         valid = __is_valid_range(I, rng) && state[1] != first(rng)
         return valid, (I,)
     end
-
-    @inline function __dec(state, indices)
+    @inline function __dec(state::Tuple{Int,Int,Vararg{Int,N}}, indices::Tuple{Any,Any,Vararg{Any,N}}) where{N}
         rng = indices[1]
         I = state[1] - step(rng)
         if __is_valid_range(I, rng) && state[1] != first(rng)
             return true, (I, tail(state)...)
         end
-        valid, I = __dec(tail(state), tail(indices))
+        t1, t2 = tail(state), tail(indices)
+        # avoid dynamic dispatch by telling the compiler relational invariants
+        valid, I = isa(t1, Tuple{Int}) ? __dec(t1, t2::Tuple{Any}) : __dec(t1, t2::Tuple{Any,Any,Vararg{Any}})
         return valid, (last(rng), I...)
     end
 

diff --git a/base/tuple.jl b/base/tuple.jl
@@ -357,10 +357,14 @@ filter(f, t::Any32) = Tuple(filter(f, collect(t)))
 
 ## comparison ##
 
-isequal(t1::Tuple, t2::Tuple) = (length(t1) == length(t2)) && _isequal(t1, t2)
-_isequal(t1::Tuple{}, t2::Tuple{}) = true
-_isequal(t1::Tuple{Any}, t2::Tuple{Any}) = isequal(t1[1], t2[1])
-_isequal(t1::Tuple, t2::Tuple) = isequal(t1[1], t2[1]) && _isequal(tail(t1), tail(t2))
+isequal(t1::Tuple, t2::Tuple) = length(t1) == length(t2) && _isequal(t1, t2)
+_isequal(::Tuple{}, ::Tuple{}) = true
+function _isequal(t1::Tuple{Any,Vararg{Any,N}}, t2::Tuple{Any,Vararg{Any,N}}) where {N}
+    isequal(t1[1], t2[1]) || return false
+    t1, t2 = tail(t1), tail(t2)
+    # avoid dynamic dispatch by telling the compiler relational invariants
+    return isa(t1, Tuple{}) ? true : _isequal(t1, t2::Tuple{Any,Vararg{Any}})
+end
 function _isequal(t1::Any32, t2::Any32)
     for i = 1:length(t1)
         if !isequal(t1[i], t2[i])