Skip to content

Commit 7cd6adf

Browse files
committed
FEAT: Port stack to using try_append_array supporting Clone
1 parent f569af6 commit 7cd6adf

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
@@ -97,9 +97,11 @@ where
9797
res_dim.set_axis(axis, 0);
9898
let mut res;
9999
unsafe {
100+
// create the axis so that 0 is outermost axis
100101
res_dim.slice_mut().swap(axis.index(), 0);
101102
res = Array::from_shape_vec_unchecked(res_dim, Vec::with_capacity(new_len));
102-
res.swap_axes(axis.index(), 0);
103+
// swap axes so that `axis` is outermost axis
104+
res.swap_axes(0, axis.index());
103105
}
104106

105107
for array in arrays {
@@ -139,7 +141,7 @@ pub fn stack_new_axis<A, D>(
139141
arrays: &[ArrayView<A, D>],
140142
) -> Result<Array<A, D::Larger>, ShapeError>
141143
where
142-
A: Copy,
144+
A: Clone,
143145
D: Dimension,
144146
D::Larger: RemoveAxis,
145147
{
@@ -159,24 +161,25 @@ where
159161

160162
res_dim.set_axis(axis, arrays.len());
161163

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

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

182185
/// Stack arrays along the new axis.

0 commit comments

Comments
 (0)