Skip to content

Commit 586cebb

Browse files
committed
FEAT: Port stack to using try_append_array supporting Clone
1 parent 8d43bde commit 586cebb

File tree

1 file changed

+20
-17
lines changed

1 file changed

+20
-17
lines changed

src/stacking.rs

+20-17
Original file line numberDiff line numberDiff line change
@@ -95,9 +95,11 @@ where
9595
res_dim.set_axis(axis, 0);
9696
let mut res;
9797
unsafe {
98+
// create the axis so that 0 is outermost axis
9899
res_dim.slice_mut().swap(axis.index(), 0);
99100
res = Array::from_shape_vec_unchecked(res_dim, Vec::with_capacity(new_len));
100-
res.swap_axes(axis.index(), 0);
101+
// swap axes so that `axis` is outermost axis
102+
res.swap_axes(0, axis.index());
101103
}
102104

103105
for array in arrays {
@@ -137,7 +139,7 @@ pub fn stack_new_axis<A, D>(
137139
arrays: &[ArrayView<A, D>],
138140
) -> Result<Array<A, D::Larger>, ShapeError>
139141
where
140-
A: Copy,
142+
A: Clone,
141143
D: Dimension,
142144
D::Larger: RemoveAxis,
143145
{
@@ -157,24 +159,25 @@ where
157159

158160
res_dim.set_axis(axis, arrays.len());
159161

160-
// we can safely use uninitialized values here because we will
161-
// overwrite every one of them.
162-
let mut res = Array::uninit(res_dim);
163-
164-
res.axis_iter_mut(axis)
165-
.zip(arrays.iter())
166-
.for_each(|(assign_view, array)| {
167-
// assign_view is D::Larger::Smaller which is usually == D
168-
// (but if D is Ix6, we have IxD != Ix6 here; differing types
169-
// but same number of axes).
170-
let assign_view = assign_view.into_dimensionality::<D>()
171-
.expect("same-dimensionality cast");
172-
array.assign_to(assign_view);
173-
});
162+
let new_len = dimension::size_of_shape_checked(&res_dim)?;
174163

164+
// start with empty array with precomputed capacity
165+
res_dim.set_axis(axis, 0);
166+
let mut res;
175167
unsafe {
176-
Ok(res.assume_init())
168+
// create the axis so that 0 is outermost axis
169+
res_dim.slice_mut().swap(axis.index(), 0);
170+
res = Array::from_shape_vec_unchecked(res_dim, Vec::with_capacity(new_len));
171+
// swap axes so that `axis` is outermost axis
172+
res.swap_axes(0, axis.index());
177173
}
174+
175+
for array in arrays {
176+
res.try_append_array(axis, array.clone().insert_axis(axis))?;
177+
}
178+
179+
debug_assert_eq!(res.len_of(axis), arrays.len());
180+
Ok(res)
178181
}
179182

180183
/// Stack arrays along the new axis.

0 commit comments

Comments
 (0)