@@ -138,7 +138,7 @@ pub fn stack_new_axis<A, D>(
138
138
arrays : & [ ArrayView < A , D > ] ,
139
139
) -> Result < Array < A , D :: Larger > , ShapeError >
140
140
where
141
- A : Copy ,
141
+ A : Clone ,
142
142
D : Dimension ,
143
143
D :: Larger : RemoveAxis ,
144
144
{
@@ -158,24 +158,22 @@ where
158
158
159
159
res_dim. set_axis ( axis, arrays. len ( ) ) ;
160
160
161
- // we can safely use uninitialized values here because we will
162
- // overwrite every one of them.
163
- let mut res = Array :: uninit ( res_dim) ;
161
+ let new_len = dimension:: size_of_shape_checked ( & res_dim) ?;
164
162
165
- res. axis_iter_mut ( axis)
166
- . zip ( arrays. iter ( ) )
167
- . for_each ( |( assign_view, array) | {
168
- // assign_view is D::Larger::Smaller which is usually == D
169
- // (but if D is Ix6, we have IxD != Ix6 here; differing types
170
- // but same number of axes).
171
- let assign_view = assign_view. into_dimensionality :: < D > ( )
172
- . expect ( "same-dimensionality cast" ) ;
173
- array. assign_to ( assign_view) ;
174
- } ) ;
163
+ // start with empty array with precomputed capacity
164
+ // try_append_array's handling of empty arrays makes sure `axis` is ok for appending
165
+ res_dim. set_axis ( axis, 0 ) ;
166
+ let mut res = unsafe {
167
+ // Safety: dimension is size 0 and vec is empty
168
+ Array :: from_shape_vec_unchecked ( res_dim, Vec :: with_capacity ( new_len) )
169
+ } ;
175
170
176
- unsafe {
177
- Ok ( res. assume_init ( ) )
171
+ for array in arrays {
172
+ res. try_append_array ( axis , array . clone ( ) . insert_axis ( axis ) ) ? ;
178
173
}
174
+
175
+ debug_assert_eq ! ( res. len_of( axis) , arrays. len( ) ) ;
176
+ Ok ( res)
179
177
}
180
178
181
179
/// Stack arrays along the new axis.
0 commit comments