Skip to content

Commit 7c7d055

Browse files
committed
FEAT: Port stack to try_append_array, supporting Clone
1 parent 7b0e5f7 commit 7c7d055

File tree

1 file changed

+14
-16
lines changed

1 file changed

+14
-16
lines changed

src/stacking.rs

Lines changed: 14 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ pub fn stack_new_axis<A, D>(
138138
arrays: &[ArrayView<A, D>],
139139
) -> Result<Array<A, D::Larger>, ShapeError>
140140
where
141-
A: Copy,
141+
A: Clone,
142142
D: Dimension,
143143
D::Larger: RemoveAxis,
144144
{
@@ -158,24 +158,22 @@ where
158158

159159
res_dim.set_axis(axis, arrays.len());
160160

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)?;
164162

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+
};
175170

176-
unsafe {
177-
Ok(res.assume_init())
171+
for array in arrays {
172+
res.try_append_array(axis, array.clone().insert_axis(axis))?;
178173
}
174+
175+
debug_assert_eq!(res.len_of(axis), arrays.len());
176+
Ok(res)
179177
}
180178

181179
/// Stack arrays along the new axis.

0 commit comments

Comments
 (0)