@@ -2729,44 +2729,41 @@ _typed_stack(dims::Integer, ::Type{T}, ::Type{S}, ::IteratorSize, A) where {T,S}
2729
2729
_vec_axis (A, ax= _axes (A)) = length (ax) == 1 ? only (ax) : OneTo (prod (length, ax; init= 1 ))
2730
2730
2731
2731
function _dim_stack (dims:: Integer , :: Type{T} , :: Type{S} , A) where {T,S}
2732
- xit = iterate (A)
2732
+ xit = Iterators . peel (A)
2733
2733
nothing === xit && return _empty_stack (dims, T, S, A)
2734
- x1, _ = xit
2734
+ x1, xrest = xit
2735
2735
ax1 = _axes (x1)
2736
2736
N1 = length (ax1)+ 1
2737
2737
dims in 1 : N1 || throw (ArgumentError (" cannot stack slices ndims(x) = $(N1- 1 ) along dims = $dims " ))
2738
2738
2739
2739
newaxis = _vec_axis (A)
2740
- outax = ntuple (d -> d== dims ? newaxis : _axes (x1) [d - (d> dims)], N1)
2740
+ outax = ntuple (d -> d== dims ? newaxis : ax1 [d - (d> dims)], N1)
2741
2741
B = similar (_first_array (x1, A), T, outax... )
2742
2742
2743
- iit = iterate (newaxis)
2744
- while xit != = nothing
2745
- x, state = xit
2746
- i, istate = iit
2747
- _stack_size_check (x, ax1)
2748
- @inbounds if dims== 1
2749
- inds1 = ntuple (d -> d== 1 ? i : Colon (), N1)
2750
- if x isa AbstractArray
2751
- B[inds1... ] = x
2752
- else
2753
- copyto! (view (B, inds1... ), x)
2754
- end
2755
- else
2756
- inds = ntuple (d -> d== dims ? i : Colon (), N1)
2757
- if x isa AbstractArray
2758
- B[inds... ] = x
2759
- else
2760
- # This is where the type-instability of inds hurts, but it is pretty exotic:
2761
- copyto! (view (B, inds... ), x)
2762
- end
2763
- end
2764
- xit = iterate (A, state)
2765
- iit = iterate (newaxis, istate)
2743
+ if dims == 1
2744
+ _dim_stack! (Val (1 ), B, x1, xrest)
2745
+ elseif dims == 2
2746
+ _dim_stack! (Val (2 ), B, x1, xrest)
2747
+ else
2748
+ _dim_stack! (Val (dims), B, x1, xrest)
2766
2749
end
2767
2750
B
2768
2751
end
2769
2752
2753
+ function _dim_stack! (:: Val{dims} , B:: AbstractArray , x1, xrest) where {dims}
2754
+ before = ntuple (d -> Colon (), dims - 1 )
2755
+ after = ntuple (d -> Colon (), ndims (B) - dims)
2756
+
2757
+ i = firstindex (B, dims)
2758
+ copyto! (view (B, before... , i, after... ), x1)
2759
+
2760
+ for x in xrest
2761
+ _stack_size_check (x, _axes (x1))
2762
+ i += 1
2763
+ @inbounds copyto! (view (B, before... , i, after... ), x)
2764
+ end
2765
+ end
2766
+
2770
2767
@inline function _stack_size_check (x, ax1:: Tuple )
2771
2768
if _axes (x) != ax1
2772
2769
uax1 = UnitRange .(ax1)
0 commit comments