code clean

make `_ind2sub_rs` fallback to existing `Base` index transformation

Author: N5N3

base/multidimensional.jl

@@ -610,28 +610,27 @@ module IteratorsMD
         return I, (I, n+1)
     end
 
-    # make _ind2sub faster when possible
-    _ind2sub_rs(axp, mi::Tuple, i) = Base.ind2sub_rs(axp, mi, i)
-    _ind2sub_rs(axp, mi::Tuple{}, i) = Base._ind2sub(axp, i)
-    _ind2sub_rs(axp::Base.Indices{1}, mi::Tuple{}, i) = (first(axp[1]) + i - 1,)
-    @inline simd_outer_range(iter::CartesianPartition{<:CartesianIndex{0}}) = ((0,1,CartesianIndex()),)
+    # use to ReshapedArray's index machinery when possible
+    @inline _ind2sub_rs(ax::Tuple, mi::Tuple, i::Int) = Base.ind2sub_rs(ax, mi, Base.offset_if_vec(i, ax))
+    @inline _ind2sub_rs(ax::Tuple{Any,Any,Vararg{Any}}, mi::Tuple{}, i::Int) = Base._ind2sub(ax, i)
+    @inline function _splitlinear(rs::ReshapedArray, i::Int)
+        axp = axes(rs.parent)
+        ci = _ind2sub_rs(axp, rs.mi, i)
+        length(ci) == 2 ? ci : ci[1], Base._sub2ind(tail(axp), tail(ci)...)
+    end
+    simd_outer_range(iter::CartesianPartition{CartesianIndex{0}}) = ((0,1,iter[]),)
     @inline function simd_outer_range(iter::CartesianPartition)
         # CartesianPartition might start and stop in the middle of the outer
         # dimensions, thus the outer range itself is a CartesianPartition.
-        piter = iter.parent.parent
-        ax1, oiter = split(piter, Val(1))
-        vindʷ = only(iter.indices)
-        function _splitlinear(i::Int)
-            ci = _ind2sub_rs(axes(piter), iter.parent.mi, i)
-            ci[1], Base._to_linear_index(oiter, tail(ci)...)
+        rs = iter.parent
+        ax1, oiter = split(rs.parent, Val(1))
+        fl, vl = _splitlinear(rs, first(iter.indices[1]))
+        fr, vr = _splitlinear(rs, last(iter.indices[1]))
+        outer = @inbounds if ndims(oiter) == 1
+            CartesianIndices((x.indices[1][vl:vr],))
+        else
+            view(ReshapedArray(x, (length(x),), ()), vl:vr)
         end
-        fl, vl = _splitlinear(first(vindʷ))
-        fr, vr = _splitlinear(last(vindʷ))
-        # we dont have #40344 for 1.6 and 1.7, force this return a CartesianIndices{1}
-        _view(x::CartesianIndices{1}, ind::UnitRange) = @inbounds CartesianIndices((x.indices[1][ind],))
-        # there's no need to make outer range fast-indexable
-        _view(x::CartesianIndices, ind::UnitRange) = @inbounds view(ReshapedArray(x, (length(x),), ()), ind)
-        outer = _view(oiter, vl:vr)
         # Use Generator to make inner loop branchless
         @inline function genouter(i::Int, I::CartesianIndex)
             l = i == 1 ? fl : firstindex(ax1)

test/iterators.jl

@@ -556,6 +556,7 @@ end
             @test collect(I) == iterate_elements(I) == simd_iterate_elements(I) == index_elements(I)
         end
         @test all(Base.splat(==), zip(Iterators.flatten(map(collect, P)), iter))
+        # test for slowindex cases
         iter′ = Base.ReshapedArray(iter, (length(iter),), ())
         P′ = partition(iter′, part)
         for I in P′