Skip to content

Commit

Permalink
FEAT: Add new methods .assign_into() and .map_assign_into
Browse files Browse the repository at this point in the history
These do essentially the same job as Zip::from(a).map_assign_into(f, b)

Comparison with existing assign:

- Reverse argument order: source, destination
- These new methods don't broadcast.
- These new methods are generic over AssignElem (assign into &Cell, &mut
ManuallyUninit etc).
  • Loading branch information
bluss committed Mar 17, 2021
1 parent 09884fc commit b22e2b1
Show file tree
Hide file tree
Showing 5 changed files with 41 additions and 49 deletions.
20 changes: 2 additions & 18 deletions src/impl_constructors.rs
Expand Up @@ -531,30 +531,14 @@ where
/// // two first columns in b are two last in a
/// // rest of columns in b are the initial columns in a
///
/// assign_to(a.slice(s![.., -2..]), b.slice_mut(s![.., ..2]));
/// assign_to(a.slice(s![.., 2..]), b.slice_mut(s![.., ..-2]));
/// a.slice(s![.., -2..]).assign_into(b.slice_mut(s![.., ..2]));
/// a.slice(s![.., 2..]).assign_into(b.slice_mut(s![.., ..-2]));
///
/// // Now we can promise that `b` is safe to use with all operations
/// unsafe {
/// b.assume_init()
/// }
/// }
///
/// use ndarray::{IntoNdProducer, AssignElem};
///
/// // This function clones elements from the first input to the second;
/// // the two producers must have the same shape
/// fn assign_to<'a, P1, P2, A>(from: P1, to: P2)
/// where P1: IntoNdProducer<Item = &'a A>,
/// P2: IntoNdProducer<Dim = P1::Dim>,
/// P2::Item: AssignElem<A>,
/// A: Clone + 'a
/// {
/// Zip::from(from)
/// .map_assign_into(to, A::clone);
/// }
///
/// # shift_by_two(&Array2::zeros((8, 8)));
/// ```
pub fn uninit<Sh>(shape: Sh) -> ArrayBase<S::MaybeUninit, D>
where
Expand Down
38 changes: 37 additions & 1 deletion src/impl_methods.rs
Expand Up @@ -15,6 +15,7 @@ use rawpointer::PointerExt;
use crate::imp_prelude::*;

use crate::{arraytraits, DimMax};
use crate::argument_traits::AssignElem;
use crate::dimension;
use crate::dimension::IntoDimension;
use crate::dimension::{
Expand All @@ -25,7 +26,7 @@ use crate::dimension::broadcast::co_broadcast;
use crate::error::{self, ErrorKind, ShapeError, from_kind};
use crate::math_cell::MathCell;
use crate::itertools::zip;
use crate::zip::Zip;
use crate::zip::{IntoNdProducer, Zip};
use crate::AxisDescription;

use crate::iter::{
Expand Down Expand Up @@ -2049,6 +2050,41 @@ where
self.zip_mut_with(rhs, |x, y| *x = y.clone());
}

/// Perform an elementwise assigment of values from `self` into array or producer `to`.
///
/// `to` can be another array or a producer of assignable elements.
/// [`AssignElem`] determines how elements are assigned.
///
/// **Panics** if shapes disagree.
pub fn assign_into<P>(&self, to: P)
where
S: Data,
P: IntoNdProducer<Dim = D>,
P::Item: AssignElem<A>,
A: Clone,
{
self.map_assign_into(A::clone, to)
}

/// Perform an elementwise assigment of values mapped with `f` from `self` into array or
/// producer `to`.
///
/// `to` can be another array or a producer of assignable elements.
/// [`AssignElem`] determines how elements are assigned.
///
/// **Panics** if shapes disagree.
pub fn map_assign_into<'a, F, B, P>(&'a self, f: F, to: P)
where
S: Data,
F: FnMut(&'a A) -> B,
P: IntoNdProducer<Dim = D>,
P::Item: AssignElem<B>,
A: 'a,
{
Zip::from(self)
.map_assign_into(to, f);
}

/// Perform an elementwise assigment to `self` from element `x`.
pub fn fill(&mut self, x: A)
where
Expand Down
1 change: 0 additions & 1 deletion src/lib.rs
Expand Up @@ -206,7 +206,6 @@ mod shape_builder;
mod slice;
mod split_at;
mod stacking;
mod traversal_utils;
#[macro_use]
mod zip;

Expand Down
5 changes: 2 additions & 3 deletions src/stacking.rs
Expand Up @@ -8,7 +8,6 @@

use crate::error::{from_kind, ErrorKind, ShapeError};
use crate::imp_prelude::*;
use crate::traversal_utils::assign_to;

/// Stack arrays along the new axis.
///
Expand Down Expand Up @@ -99,7 +98,7 @@ where
for array in arrays {
let len = array.len_of(axis);
let (front, rest) = assign_view.split_at(axis, len);
assign_to(array, front);
array.assign_into(front);
assign_view = rest;
}
debug_assert_eq!(assign_view.len(), 0);
Expand Down Expand Up @@ -171,7 +170,7 @@ where
// but same number of axes).
let assign_view = assign_view.into_dimensionality::<D>()
.expect("same-dimensionality cast");
assign_to(array, assign_view);
array.assign_into(assign_view);
});

unsafe {
Expand Down
26 changes: 0 additions & 26 deletions src/traversal_utils.rs

This file was deleted.

0 comments on commit b22e2b1

Please sign in to comment.