66// option. This file may not be copied, modified, or distributed
77// except according to those terms.
88
9+ use alloc:: vec:: Vec ;
10+
11+ use crate :: dimension;
912use crate :: error:: { from_kind, ErrorKind , ShapeError } ;
1013use crate :: imp_prelude:: * ;
1114
6871/// ```
6972pub fn concatenate < A , D > ( axis : Axis , arrays : & [ ArrayView < A , D > ] ) -> Result < Array < A , D > , ShapeError >
7073where
71- A : Copy ,
74+ A : Clone ,
7275 D : RemoveAxis ,
7376{
7477 if arrays. is_empty ( ) {
@@ -88,24 +91,21 @@ where
8891
8992 let stacked_dim = arrays. iter ( ) . fold ( 0 , |acc, a| acc + a. len_of ( axis) ) ;
9093 res_dim. set_axis ( axis, stacked_dim) ;
94+ let new_len = dimension:: size_of_shape_checked ( & res_dim) ?;
9195
92- // we can safely use uninitialized values here because we will
93- // overwrite every one of them.
94- let mut res = Array :: uninit ( res_dim) ;
96+ // start with empty array with precomputed capacity
97+ // try_append_array's handling of empty arrays makes sure `axis` is ok for appending
98+ res_dim. set_axis ( axis, 0 ) ;
99+ let mut res = unsafe {
100+ // Safety: dimension is size 0 and vec is empty
101+ Array :: from_shape_vec_unchecked ( res_dim, Vec :: with_capacity ( new_len) )
102+ } ;
95103
96- {
97- let mut assign_view = res. view_mut ( ) ;
98- for array in arrays {
99- let len = array. len_of ( axis) ;
100- let ( front, rest) = assign_view. split_at ( axis, len) ;
101- array. assign_to ( front) ;
102- assign_view = rest;
103- }
104- debug_assert_eq ! ( assign_view. len( ) , 0 ) ;
105- }
106- unsafe {
107- Ok ( res. assume_init ( ) )
104+ for array in arrays {
105+ res. try_append_array ( axis, array. clone ( ) ) ?;
108106 }
107+ debug_assert_eq ! ( res. len_of( axis) , stacked_dim) ;
108+ Ok ( res)
109109}
110110
111111#[ deprecated( note="Use under the name stack instead." , since="0.15.0" ) ]
0 commit comments