add test & optimize performance

1. add more test cases for widen definition
2. make ind2sub faster via proper fallback

Author: N5N3

base/multidimensional.jl

@@ -610,25 +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()),)
     @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)
-        @inline function _splitlinear(i::Int)
-            ci = Base._to_subscript_indices(piter, i)
+        function _splitlinear(i::Int)
+            ci = _ind2sub_rs(axes(piter), iter.parent.mi, i)
             ci[1], Base._to_linear_index(oiter, tail(ci)...)
         end
         fl, vl = _splitlinear(first(vindʷ))
         fr, vr = _splitlinear(last(vindʷ))
-        @inline function _view(oiter::CartesianIndices, ind::UnitRange)
-            # we dont have #40344 for 1.6 and 1.7, force this return a CartesianIndices{1}
-            ndims(oiter) == 1 && return @inbounds CartesianIndices((oiter.indices[1][ind],))
-            # there's no need to make outer range fast-indexable
-            oiter′ = ReshapedArray(oiter, (length(oiter),), ())
-            @inbounds view(oiter′, ind)
-        end
+        # 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)

test/iterators.jl

@@ -544,7 +544,7 @@ function index_elements(iter)
     return vals
 end
 
-@testset "CartesianPartition optimizations" for dims in ((1,), (64,), (101,),
+@testset "CartesianPartition optimizations" for dims in ((), (1,), (64,), (101,),
                                                          (1,1), (8,8), (11, 13),
                                                          (1,1,1), (8, 4, 2), (11, 13, 17)),
                                                 part in (1, 7, 8, 11, 63, 64, 65, 142, 143, 144)
@@ -556,6 +556,13 @@ 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))
+        iter′ = Base.ReshapedArray(iter, (length(iter),), ())
+        P′ = partition(iter′, part)
+        for I in P′
+            @test length(I) == iterate_length(I) == simd_iterate_length(I) == simd_trip_count(I)
+            @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′))
     end
 end
 @testset "empty/invalid partitions" begin