@@ -2662,44 +2662,41 @@ _typed_stack(dims::Integer, ::Type{T}, ::Type{S}, ::IteratorSize, A) where {T,S}
2662
2662
_vec_axis (A, ax= _axes (A)) = length (ax) == 1 ? only (ax) : OneTo (prod (length, ax; init= 1 ))
2663
2663
2664
2664
function _dim_stack (dims:: Integer , :: Type{T} , :: Type{S} , A) where {T,S}
2665
- xit = iterate (A)
2665
+ xit = Iterators . peel (A)
2666
2666
nothing === xit && return _empty_stack (dims, T, S, A)
2667
- x1, _ = xit
2667
+ x1, xrest = xit
2668
2668
ax1 = _axes (x1)
2669
2669
N1 = length (ax1)+ 1
2670
2670
dims in 1 : N1 || throw (ArgumentError (" cannot stack slices ndims(x) = $(N1- 1 ) along dims = $dims " ))
2671
2671
2672
2672
newaxis = _vec_axis (A)
2673
- outax = ntuple (d -> d== dims ? newaxis : _axes (x1) [d - (d> dims)], N1)
2673
+ outax = ntuple (d -> d== dims ? newaxis : ax1 [d - (d> dims)], N1)
2674
2674
B = similar (_first_array (x1, A), T, outax... )
2675
2675
2676
- iit = iterate (newaxis)
2677
- while xit != = nothing
2678
- x, state = xit
2679
- i, istate = iit
2680
- _stack_size_check (x, ax1)
2681
- @inbounds if dims== 1
2682
- inds1 = ntuple (d -> d== 1 ? i : Colon (), N1)
2683
- if x isa AbstractArray
2684
- B[inds1... ] = x
2685
- else
2686
- copyto! (view (B, inds1... ), x)
2687
- end
2688
- else
2689
- inds = ntuple (d -> d== dims ? i : Colon (), N1)
2690
- if x isa AbstractArray
2691
- B[inds... ] = x
2692
- else
2693
- # This is where the type-instability of inds hurts, but it is pretty exotic:
2694
- copyto! (view (B, inds... ), x)
2695
- end
2696
- end
2697
- xit = iterate (A, state)
2698
- iit = iterate (newaxis, istate)
2676
+ if dims == 1
2677
+ _dim_stack! (Val (1 ), B, x1, xrest)
2678
+ elseif dims == 2
2679
+ _dim_stack! (Val (2 ), B, x1, xrest)
2680
+ else
2681
+ _dim_stack! (Val (dims), B, x1, xrest)
2699
2682
end
2700
2683
B
2701
2684
end
2702
2685
2686
+ function _dim_stack! (:: Val{dims} , B:: AbstractArray , x1, xrest) where {dims}
2687
+ before = ntuple (d -> Colon (), dims - 1 )
2688
+ after = ntuple (d -> Colon (), ndims (B) - dims)
2689
+
2690
+ i = firstindex (B, dims)
2691
+ copyto! (view (B, before... , i, after... ), x1)
2692
+
2693
+ for x in xrest
2694
+ _stack_size_check (x, _axes (x1))
2695
+ i += 1
2696
+ @inbounds copyto! (view (B, before... , i, after... ), x)
2697
+ end
2698
+ end
2699
+
2703
2700
@inline function _stack_size_check (x, ax1:: Tuple )
2704
2701
if _axes (x) != ax1
2705
2702
uax1 = UnitRange .(ax1)
0 commit comments