ndarray/
slice.rs

1// Copyright 2014-2016 bluss and ndarray developers.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8use crate::dimension::slices_intersect;
9use crate::error::{ErrorKind, ShapeError};
10use crate::{ArrayViewMut, DimAdd, Dimension, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
11use alloc::vec::Vec;
12use std::convert::TryFrom;
13use std::fmt;
14use std::marker::PhantomData;
15use std::ops::{Deref, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive};
16
17/// A slice (range with step size).
18///
19/// `end` is an exclusive index. Negative `start` or `end` indexes are counted
20/// from the back of the axis. If `end` is `None`, the slice extends to the end
21/// of the axis.
22///
23/// See also the [`s![]`](s!) macro.
24///
25/// ## Examples
26///
27/// `Slice::new(0, None, 1)` is the full range of an axis. It can also be
28/// created with `Slice::from(..)`. The Python equivalent is `[:]`.
29///
30/// `Slice::new(a, b, 2)` is every second element from `a` until `b`. It can
31/// also be created with `Slice::from(a..b).step_by(2)`. The Python equivalent
32/// is `[a:b:2]`.
33///
34/// `Slice::new(a, None, -1)` is every element, from `a` until the end, in
35/// reverse order. It can also be created with `Slice::from(a..).step_by(-1)`.
36/// The Python equivalent is `[a::-1]`.
37#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
38pub struct Slice {
39    /// start index; negative are counted from the back of the axis
40    pub start: isize,
41    /// end index; negative are counted from the back of the axis; when not present
42    /// the default is the full length of the axis.
43    pub end: Option<isize>,
44    /// step size in elements; the default is 1, for every element.
45    pub step: isize,
46}
47
48impl Slice {
49    /// Create a new `Slice` with the given extents.
50    ///
51    /// See also the `From` impls, converting from ranges; for example
52    /// `Slice::from(i..)` or `Slice::from(j..k)`.
53    ///
54    /// `step` must be nonzero.
55    /// (This method checks with a debug assertion that `step` is not zero.)
56    pub fn new(start: isize, end: Option<isize>, step: isize) -> Slice {
57        debug_assert_ne!(step, 0, "Slice::new: step must be nonzero");
58        Slice { start, end, step }
59    }
60
61    /// Create a new `Slice` with the given step size (multiplied with the
62    /// previous step size).
63    ///
64    /// `step` must be nonzero.
65    /// (This method checks with a debug assertion that `step` is not zero.)
66    #[inline]
67    pub fn step_by(self, step: isize) -> Self {
68        debug_assert_ne!(step, 0, "Slice::step_by: step must be nonzero");
69        Slice {
70            step: self.step * step,
71            ..self
72        }
73    }
74}
75
76/// Token to represent a new axis in a slice description.
77///
78/// See also the [`s![]`](s!) macro.
79#[derive(Clone, Copy, Debug)]
80pub struct NewAxis;
81
82/// A slice (range with step), an index, or a new axis token.
83///
84/// See also the [`s![]`](s!) macro for a convenient way to create a
85/// `SliceInfo<[SliceInfoElem; n], Din, Dout>`.
86///
87/// ## Examples
88///
89/// `SliceInfoElem::Index(a)` is the index `a`. It can also be created with
90/// `SliceInfoElem::from(a)`. The Python equivalent is `[a]`. The macro
91/// equivalent is `s![a]`.
92///
93/// `SliceInfoElem::Slice { start: 0, end: None, step: 1 }` is the full range
94/// of an axis. It can also be created with `SliceInfoElem::from(..)`. The
95/// Python equivalent is `[:]`. The macro equivalent is `s![..]`.
96///
97/// `SliceInfoElem::Slice { start: a, end: Some(b), step: 2 }` is every second
98/// element from `a` until `b`. It can also be created with
99/// `SliceInfoElem::from(Slice::from(a..b).step_by(2))`. The Python equivalent
100/// is `[a:b:2]`. The macro equivalent is `s![a..b;2]`.
101///
102/// `SliceInfoElem::Slice { start: a, end: None, step: -1 }` is every element,
103/// from `a` until the end, in reverse order. It can also be created with
104/// `SliceInfoElem::from(Slice::from(a..).step_by(-1))`. The Python equivalent
105/// is `[a::-1]`. The macro equivalent is `s![a..;-1]`.
106///
107/// `SliceInfoElem::NewAxis` is a new axis of length 1. It can also be created
108/// with `SliceInfoElem::from(NewAxis)`. The Python equivalent is
109/// `[np.newaxis]`. The macro equivalent is `s![NewAxis]`.
110#[derive(Debug, PartialEq, Eq, Hash)]
111pub enum SliceInfoElem {
112    /// A range with step size. `end` is an exclusive index. Negative `start`
113    /// or `end` indexes are counted from the back of the axis. If `end` is
114    /// `None`, the slice extends to the end of the axis.
115    Slice {
116        /// start index; negative are counted from the back of the axis
117        start: isize,
118        /// end index; negative are counted from the back of the axis; when not present
119        /// the default is the full length of the axis.
120        end: Option<isize>,
121        /// step size in elements; the default is 1, for every element.
122        step: isize,
123    },
124    /// A single index.
125    Index(isize),
126    /// A new axis of length 1.
127    NewAxis,
128}
129
130copy_and_clone! {SliceInfoElem}
131
132impl SliceInfoElem {
133    /// Returns `true` if `self` is a `Slice` value.
134    pub fn is_slice(&self) -> bool {
135        matches!(self, SliceInfoElem::Slice { .. })
136    }
137
138    /// Returns `true` if `self` is an `Index` value.
139    pub fn is_index(&self) -> bool {
140        matches!(self, SliceInfoElem::Index(_))
141    }
142
143    /// Returns `true` if `self` is a `NewAxis` value.
144    pub fn is_new_axis(&self) -> bool {
145        matches!(self, SliceInfoElem::NewAxis)
146    }
147}
148
149impl fmt::Display for SliceInfoElem {
150    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
151        match *self {
152            SliceInfoElem::Index(index) => write!(f, "{}", index)?,
153            SliceInfoElem::Slice { start, end, step } => {
154                if start != 0 {
155                    write!(f, "{}", start)?;
156                }
157                write!(f, "..")?;
158                if let Some(i) = end {
159                    write!(f, "{}", i)?;
160                }
161                if step != 1 {
162                    write!(f, ";{}", step)?;
163                }
164            }
165            SliceInfoElem::NewAxis => write!(f, stringify!(NewAxis))?,
166        }
167        Ok(())
168    }
169}
170
171macro_rules! impl_slice_variant_from_range {
172    ($self:ty, $constructor:path, $index:ty) => {
173        impl From<Range<$index>> for $self {
174            #[inline]
175            fn from(r: Range<$index>) -> $self {
176                $constructor {
177                    start: r.start as isize,
178                    end: Some(r.end as isize),
179                    step: 1,
180                }
181            }
182        }
183
184        impl From<RangeInclusive<$index>> for $self {
185            #[inline]
186            fn from(r: RangeInclusive<$index>) -> $self {
187                let end = *r.end() as isize;
188                $constructor {
189                    start: *r.start() as isize,
190                    end: if end == -1 { None } else { Some(end + 1) },
191                    step: 1,
192                }
193            }
194        }
195
196        impl From<RangeFrom<$index>> for $self {
197            #[inline]
198            fn from(r: RangeFrom<$index>) -> $self {
199                $constructor {
200                    start: r.start as isize,
201                    end: None,
202                    step: 1,
203                }
204            }
205        }
206
207        impl From<RangeTo<$index>> for $self {
208            #[inline]
209            fn from(r: RangeTo<$index>) -> $self {
210                $constructor {
211                    start: 0,
212                    end: Some(r.end as isize),
213                    step: 1,
214                }
215            }
216        }
217
218        impl From<RangeToInclusive<$index>> for $self {
219            #[inline]
220            fn from(r: RangeToInclusive<$index>) -> $self {
221                let end = r.end as isize;
222                $constructor {
223                    start: 0,
224                    end: if end == -1 { None } else { Some(end + 1) },
225                    step: 1,
226                }
227            }
228        }
229    };
230}
231impl_slice_variant_from_range!(Slice, Slice, isize);
232impl_slice_variant_from_range!(Slice, Slice, usize);
233impl_slice_variant_from_range!(Slice, Slice, i32);
234impl_slice_variant_from_range!(SliceInfoElem, SliceInfoElem::Slice, isize);
235impl_slice_variant_from_range!(SliceInfoElem, SliceInfoElem::Slice, usize);
236impl_slice_variant_from_range!(SliceInfoElem, SliceInfoElem::Slice, i32);
237
238impl From<RangeFull> for Slice {
239    #[inline]
240    fn from(_: RangeFull) -> Slice {
241        Slice {
242            start: 0,
243            end: None,
244            step: 1,
245        }
246    }
247}
248
249impl From<RangeFull> for SliceInfoElem {
250    #[inline]
251    fn from(_: RangeFull) -> SliceInfoElem {
252        SliceInfoElem::Slice {
253            start: 0,
254            end: None,
255            step: 1,
256        }
257    }
258}
259
260impl From<Slice> for SliceInfoElem {
261    #[inline]
262    fn from(s: Slice) -> SliceInfoElem {
263        SliceInfoElem::Slice {
264            start: s.start,
265            end: s.end,
266            step: s.step,
267        }
268    }
269}
270
271macro_rules! impl_sliceinfoelem_from_index {
272    ($index:ty) => {
273        impl From<$index> for SliceInfoElem {
274            #[inline]
275            fn from(r: $index) -> SliceInfoElem {
276                SliceInfoElem::Index(r as isize)
277            }
278        }
279    };
280}
281impl_sliceinfoelem_from_index!(isize);
282impl_sliceinfoelem_from_index!(usize);
283impl_sliceinfoelem_from_index!(i32);
284
285impl From<NewAxis> for SliceInfoElem {
286    #[inline]
287    fn from(_: NewAxis) -> SliceInfoElem {
288        SliceInfoElem::NewAxis
289    }
290}
291
292/// A type that can slice an array of dimension `D`.
293///
294/// This trait is unsafe to implement because the implementation must ensure
295/// that `D`, `Self::OutDim`, `self.in_dim()`, and `self.out_ndim()` are
296/// consistent with the `&[SliceInfoElem]` returned by `self.as_ref()` and that
297/// `self.as_ref()` always returns the same value when called multiple times.
298#[allow(clippy::missing_safety_doc)] // not implementable downstream
299pub unsafe trait SliceArg<D: Dimension>: AsRef<[SliceInfoElem]> {
300    /// Dimensionality of the output array.
301    type OutDim: Dimension;
302
303    /// Returns the number of axes in the input array.
304    fn in_ndim(&self) -> usize;
305
306    /// Returns the number of axes in the output array.
307    fn out_ndim(&self) -> usize;
308
309    private_decl! {}
310}
311
312unsafe impl<T, D> SliceArg<D> for &T
313where
314    T: SliceArg<D> + ?Sized,
315    D: Dimension,
316{
317    type OutDim = T::OutDim;
318
319    fn in_ndim(&self) -> usize {
320        T::in_ndim(self)
321    }
322
323    fn out_ndim(&self) -> usize {
324        T::out_ndim(self)
325    }
326
327    private_impl! {}
328}
329
330macro_rules! impl_slicearg_samedim {
331    ($in_dim:ty) => {
332        unsafe impl<T, Dout> SliceArg<$in_dim> for SliceInfo<T, $in_dim, Dout>
333        where
334            T: AsRef<[SliceInfoElem]>,
335            Dout: Dimension,
336        {
337            type OutDim = Dout;
338
339            fn in_ndim(&self) -> usize {
340                self.in_ndim()
341            }
342
343            fn out_ndim(&self) -> usize {
344                self.out_ndim()
345            }
346
347            private_impl! {}
348        }
349    };
350}
351impl_slicearg_samedim!(Ix0);
352impl_slicearg_samedim!(Ix1);
353impl_slicearg_samedim!(Ix2);
354impl_slicearg_samedim!(Ix3);
355impl_slicearg_samedim!(Ix4);
356impl_slicearg_samedim!(Ix5);
357impl_slicearg_samedim!(Ix6);
358
359unsafe impl<T, Din, Dout> SliceArg<IxDyn> for SliceInfo<T, Din, Dout>
360where
361    T: AsRef<[SliceInfoElem]>,
362    Din: Dimension,
363    Dout: Dimension,
364{
365    type OutDim = Dout;
366
367    fn in_ndim(&self) -> usize {
368        self.in_ndim()
369    }
370
371    fn out_ndim(&self) -> usize {
372        self.out_ndim()
373    }
374
375    private_impl! {}
376}
377
378unsafe impl SliceArg<IxDyn> for [SliceInfoElem] {
379    type OutDim = IxDyn;
380
381    fn in_ndim(&self) -> usize {
382        self.iter().filter(|s| !s.is_new_axis()).count()
383    }
384
385    fn out_ndim(&self) -> usize {
386        self.iter().filter(|s| !s.is_index()).count()
387    }
388
389    private_impl! {}
390}
391
392/// Represents all of the necessary information to perform a slice.
393///
394/// The type `T` is typically `[SliceInfoElem; n]`, `&[SliceInfoElem]`, or
395/// `Vec<SliceInfoElem>`. The type `Din` is the dimension of the array to be
396/// sliced, and `Dout` is the output dimension after calling [`.slice()`]. Note
397/// that if `Din` is a fixed dimension type (`Ix0`, `Ix1`, `Ix2`, etc.), the
398/// `SliceInfo` instance can still be used to slice an array with dimension
399/// `IxDyn` as long as the number of axes matches.
400///
401/// [`.slice()`]: crate::ArrayBase::slice
402#[derive(Debug)]
403pub struct SliceInfo<T, Din: Dimension, Dout: Dimension> {
404    in_dim: PhantomData<Din>,
405    out_dim: PhantomData<Dout>,
406    indices: T,
407}
408
409impl<T, Din, Dout> Deref for SliceInfo<T, Din, Dout>
410where
411    Din: Dimension,
412    Dout: Dimension,
413{
414    type Target = T;
415    fn deref(&self) -> &Self::Target {
416        &self.indices
417    }
418}
419
420fn check_dims_for_sliceinfo<Din, Dout>(indices: &[SliceInfoElem]) -> Result<(), ShapeError>
421where
422    Din: Dimension,
423    Dout: Dimension,
424{
425    if let Some(in_ndim) = Din::NDIM {
426        if in_ndim != indices.in_ndim() {
427            return Err(ShapeError::from_kind(ErrorKind::IncompatibleShape));
428        }
429    }
430    if let Some(out_ndim) = Dout::NDIM {
431        if out_ndim != indices.out_ndim() {
432            return Err(ShapeError::from_kind(ErrorKind::IncompatibleShape));
433        }
434    }
435    Ok(())
436}
437
438impl<T, Din, Dout> SliceInfo<T, Din, Dout>
439where
440    T: AsRef<[SliceInfoElem]>,
441    Din: Dimension,
442    Dout: Dimension,
443{
444    /// Returns a new `SliceInfo` instance.
445    ///
446    /// **Note:** only unchecked for non-debug builds of `ndarray`.
447    ///
448    /// # Safety
449    ///
450    /// The caller must ensure that `in_dim` and `out_dim` are consistent with
451    /// `indices` and that `indices.as_ref()` always returns the same value
452    /// when called multiple times.
453    #[doc(hidden)]
454    pub unsafe fn new_unchecked(
455        indices: T,
456        in_dim: PhantomData<Din>,
457        out_dim: PhantomData<Dout>,
458    ) -> SliceInfo<T, Din, Dout> {
459        if cfg!(debug_assertions) {
460            check_dims_for_sliceinfo::<Din, Dout>(indices.as_ref())
461                .expect("`Din` and `Dout` must be consistent with `indices`.");
462        }
463        SliceInfo {
464            in_dim,
465            out_dim,
466            indices,
467        }
468    }
469
470    /// Returns a new `SliceInfo` instance.
471    ///
472    /// Errors if `Din` or `Dout` is not consistent with `indices`.
473    ///
474    /// For common types, a safe alternative is to use `TryFrom` instead.
475    ///
476    /// # Safety
477    ///
478    /// The caller must ensure `indices.as_ref()` always returns the same value
479    /// when called multiple times.
480    pub unsafe fn new(indices: T) -> Result<SliceInfo<T, Din, Dout>, ShapeError> {
481        check_dims_for_sliceinfo::<Din, Dout>(indices.as_ref())?;
482        Ok(SliceInfo {
483            in_dim: PhantomData,
484            out_dim: PhantomData,
485            indices,
486        })
487    }
488
489    /// Returns the number of dimensions of the input array for
490    /// [`.slice()`](crate::ArrayBase::slice).
491    ///
492    /// If `Din` is a fixed-size dimension type, then this is equivalent to
493    /// `Din::NDIM.unwrap()`. Otherwise, the value is calculated by iterating
494    /// over the `SliceInfoElem` elements.
495    pub fn in_ndim(&self) -> usize {
496        if let Some(ndim) = Din::NDIM {
497            ndim
498        } else {
499            self.indices.as_ref().in_ndim()
500        }
501    }
502
503    /// Returns the number of dimensions after calling
504    /// [`.slice()`](crate::ArrayBase::slice) (including taking
505    /// subviews).
506    ///
507    /// If `Dout` is a fixed-size dimension type, then this is equivalent to
508    /// `Dout::NDIM.unwrap()`. Otherwise, the value is calculated by iterating
509    /// over the `SliceInfoElem` elements.
510    pub fn out_ndim(&self) -> usize {
511        if let Some(ndim) = Dout::NDIM {
512            ndim
513        } else {
514            self.indices.as_ref().out_ndim()
515        }
516    }
517}
518
519impl<'a, Din, Dout> TryFrom<&'a [SliceInfoElem]> for SliceInfo<&'a [SliceInfoElem], Din, Dout>
520where
521    Din: Dimension,
522    Dout: Dimension,
523{
524    type Error = ShapeError;
525
526    fn try_from(
527        indices: &'a [SliceInfoElem],
528    ) -> Result<SliceInfo<&'a [SliceInfoElem], Din, Dout>, ShapeError> {
529        unsafe {
530            // This is okay because `&[SliceInfoElem]` always returns the same
531            // value for `.as_ref()`.
532            Self::new(indices)
533        }
534    }
535}
536
537impl<Din, Dout> TryFrom<Vec<SliceInfoElem>> for SliceInfo<Vec<SliceInfoElem>, Din, Dout>
538where
539    Din: Dimension,
540    Dout: Dimension,
541{
542    type Error = ShapeError;
543
544    fn try_from(
545        indices: Vec<SliceInfoElem>,
546    ) -> Result<SliceInfo<Vec<SliceInfoElem>, Din, Dout>, ShapeError> {
547        unsafe {
548            // This is okay because `Vec` always returns the same value for
549            // `.as_ref()`.
550            Self::new(indices)
551        }
552    }
553}
554
555macro_rules! impl_tryfrom_array_for_sliceinfo {
556    ($len:expr) => {
557        impl<Din, Dout> TryFrom<[SliceInfoElem; $len]>
558            for SliceInfo<[SliceInfoElem; $len], Din, Dout>
559        where
560            Din: Dimension,
561            Dout: Dimension,
562        {
563            type Error = ShapeError;
564
565            fn try_from(
566                indices: [SliceInfoElem; $len],
567            ) -> Result<SliceInfo<[SliceInfoElem; $len], Din, Dout>, ShapeError> {
568                unsafe {
569                    // This is okay because `[SliceInfoElem; N]` always returns
570                    // the same value for `.as_ref()`.
571                    Self::new(indices)
572                }
573            }
574        }
575    };
576}
577impl_tryfrom_array_for_sliceinfo!(0);
578impl_tryfrom_array_for_sliceinfo!(1);
579impl_tryfrom_array_for_sliceinfo!(2);
580impl_tryfrom_array_for_sliceinfo!(3);
581impl_tryfrom_array_for_sliceinfo!(4);
582impl_tryfrom_array_for_sliceinfo!(5);
583impl_tryfrom_array_for_sliceinfo!(6);
584impl_tryfrom_array_for_sliceinfo!(7);
585impl_tryfrom_array_for_sliceinfo!(8);
586
587impl<T, Din, Dout> AsRef<[SliceInfoElem]> for SliceInfo<T, Din, Dout>
588where
589    T: AsRef<[SliceInfoElem]>,
590    Din: Dimension,
591    Dout: Dimension,
592{
593    fn as_ref(&self) -> &[SliceInfoElem] {
594        self.indices.as_ref()
595    }
596}
597
598impl<'a, T, Din, Dout> From<&'a SliceInfo<T, Din, Dout>>
599    for SliceInfo<&'a [SliceInfoElem], Din, Dout>
600where
601    T: AsRef<[SliceInfoElem]>,
602    Din: Dimension,
603    Dout: Dimension,
604{
605    fn from(info: &'a SliceInfo<T, Din, Dout>) -> SliceInfo<&'a [SliceInfoElem], Din, Dout> {
606        SliceInfo {
607            in_dim: info.in_dim,
608            out_dim: info.out_dim,
609            indices: info.indices.as_ref(),
610        }
611    }
612}
613
614impl<T, Din, Dout> Copy for SliceInfo<T, Din, Dout>
615where
616    T: Copy,
617    Din: Dimension,
618    Dout: Dimension,
619{
620}
621
622impl<T, Din, Dout> Clone for SliceInfo<T, Din, Dout>
623where
624    T: Clone,
625    Din: Dimension,
626    Dout: Dimension,
627{
628    fn clone(&self) -> Self {
629        SliceInfo {
630            in_dim: PhantomData,
631            out_dim: PhantomData,
632            indices: self.indices.clone(),
633        }
634    }
635}
636
637/// Trait for determining dimensionality of input and output for [`s!`] macro.
638#[doc(hidden)]
639pub trait SliceNextDim {
640    /// Number of dimensions that this slicing argument consumes in the input array.
641    type InDim: Dimension;
642    /// Number of dimensions that this slicing argument produces in the output array.
643    type OutDim: Dimension;
644
645    fn next_in_dim<D>(&self, _: PhantomData<D>) -> PhantomData<<D as DimAdd<Self::InDim>>::Output>
646    where
647        D: Dimension + DimAdd<Self::InDim>,
648    {
649        PhantomData
650    }
651
652    fn next_out_dim<D>(&self, _: PhantomData<D>) -> PhantomData<<D as DimAdd<Self::OutDim>>::Output>
653    where
654        D: Dimension + DimAdd<Self::OutDim>,
655    {
656        PhantomData
657    }
658}
659
660macro_rules! impl_slicenextdim {
661    (($($generics:tt)*), $self:ty, $in:ty, $out:ty) => {
662        impl<$($generics)*> SliceNextDim for $self {
663            type InDim = $in;
664            type OutDim = $out;
665        }
666    };
667}
668
669impl_slicenextdim!((), isize, Ix1, Ix0);
670impl_slicenextdim!((), usize, Ix1, Ix0);
671impl_slicenextdim!((), i32, Ix1, Ix0);
672
673impl_slicenextdim!((T), Range<T>, Ix1, Ix1);
674impl_slicenextdim!((T), RangeInclusive<T>, Ix1, Ix1);
675impl_slicenextdim!((T), RangeFrom<T>, Ix1, Ix1);
676impl_slicenextdim!((T), RangeTo<T>, Ix1, Ix1);
677impl_slicenextdim!((T), RangeToInclusive<T>, Ix1, Ix1);
678impl_slicenextdim!((), RangeFull, Ix1, Ix1);
679impl_slicenextdim!((), Slice, Ix1, Ix1);
680
681impl_slicenextdim!((), NewAxis, Ix0, Ix1);
682
683/// Slice argument constructor.
684///
685/// `s![]` takes a list of ranges/slices/indices/new-axes, separated by comma,
686/// with optional step sizes that are separated from the range by a semicolon.
687/// It is converted into a [`SliceInfo`] instance.
688///
689/// Each range/slice/index uses signed indices, where a negative value is
690/// counted from the end of the axis. Step sizes are also signed and may be
691/// negative, but must not be zero.
692///
693/// The syntax is `s![` *[ elem [, elem [ , ... ] ] ]* `]`, where *elem* is any
694/// of the following:
695///
696/// * *index*: an index to use for taking a subview with respect to that axis.
697///   (The index is selected. The axis is removed except with
698///   [`.slice_collapse()`].)
699/// * *range*: a range with step size 1 to use for slicing that axis.
700/// * *range* `;` *step*: a range with step size *step* to use for slicing that axis.
701/// * *slice*: a [`Slice`] instance to use for slicing that axis.
702/// * *slice* `;` *step*: a range constructed from a [`Slice`] instance,
703///   multiplying the step size by *step*, to use for slicing that axis.
704/// * *new-axis*: a [`NewAxis`] instance that represents the creation of a new axis.
705///   (Except for [`.slice_collapse()`], which panics on [`NewAxis`] elements.)
706///
707/// The number of *elem*, not including *new-axis*, must match the
708/// number of axes in the array. *index*, *range*, *slice*, *step*, and
709/// *new-axis* can be expressions. *index* must be of type `isize`, `usize`, or
710/// `i32`. *range* must be of type `Range<I>`, `RangeTo<I>`, `RangeFrom<I>`, or
711/// `RangeFull` where `I` is `isize`, `usize`, or `i32`. *step* must be a type
712/// that can be converted to `isize` with the `as` keyword.
713///
714/// For example, `s![0..4;2, 6, 1..5, NewAxis]` is a slice of the first axis
715/// for 0..4 with step size 2, a subview of the second axis at index 6, a slice
716/// of the third axis for 1..5 with default step size 1, and a new axis of
717/// length 1 at the end of the shape. The input array must have 3 dimensions.
718/// The resulting slice would have shape `[2, 4, 1]` for [`.slice()`],
719/// [`.slice_mut()`], and [`.slice_move()`], while [`.slice_collapse()`] would
720/// panic. Without the `NewAxis`, i.e. `s![0..4;2, 6, 1..5]`,
721/// [`.slice_collapse()`] would result in an array of shape `[2, 1, 4]`.
722///
723/// [`.slice()`]: crate::ArrayBase::slice
724/// [`.slice_mut()`]: crate::ArrayBase::slice_mut
725/// [`.slice_move()`]: crate::ArrayBase::slice_move
726/// [`.slice_collapse()`]: crate::ArrayBase::slice_collapse
727///
728/// See also [*Slicing*](crate::ArrayBase#slicing).
729///
730/// # Example
731///
732/// ```
733/// use ndarray::{s, Array2, ArrayView2};
734///
735/// fn laplacian(v: &ArrayView2<f32>) -> Array2<f32> {
736///     -4. * &v.slice(s![1..-1, 1..-1])
737///     + v.slice(s![ ..-2, 1..-1])
738///     + v.slice(s![1..-1,  ..-2])
739///     + v.slice(s![1..-1, 2..  ])
740///     + v.slice(s![2..  , 1..-1])
741/// }
742/// # fn main() { let _ = laplacian; }
743/// ```
744///
745/// # Negative *step*
746///
747/// The behavior of negative *step* arguments is most easily understood with
748/// slicing as a two-step process:
749///
750/// 1. First, perform a slice with *range*.
751///
752/// 2. If *step* is positive, start with the front of the slice; if *step* is
753///    negative, start with the back of the slice. Then, add *step* until
754///    reaching the other end of the slice (inclusive).
755///
756/// An equivalent way to think about step 2 is, "If *step* is negative, reverse
757/// the slice. Start at the front of the (possibly reversed) slice, and add
758/// *step.abs()* until reaching the back of the slice (inclusive)."
759///
760/// For example,
761///
762/// ```
763/// # use ndarray::prelude::*;
764/// #
765/// # fn main() {
766/// let arr = array![0, 1, 2, 3];
767/// assert_eq!(arr.slice(s![1..3;-1]), array![2, 1]);
768/// assert_eq!(arr.slice(s![1..;-2]), array![3, 1]);
769/// assert_eq!(arr.slice(s![0..4;-2]), array![3, 1]);
770/// assert_eq!(arr.slice(s![0..;-2]), array![3, 1]);
771/// assert_eq!(arr.slice(s![..;-2]), array![3, 1]);
772/// # }
773/// ```
774#[macro_export]
775macro_rules! s(
776    // convert a..b;c into @convert(a..b, c), final item
777    (@parse $in_dim:expr, $out_dim:expr, [$($stack:tt)*] $r:expr;$s:expr) => {
778        match $r {
779            r => {
780                let in_dim = $crate::SliceNextDim::next_in_dim(&r, $in_dim);
781                let out_dim = $crate::SliceNextDim::next_out_dim(&r, $out_dim);
782                #[allow(unsafe_code)]
783                unsafe {
784                    $crate::SliceInfo::new_unchecked(
785                        [$($stack)* $crate::s!(@convert r, $s)],
786                        in_dim,
787                        out_dim,
788                    )
789                }
790            }
791        }
792    };
793    // convert a..b into @convert(a..b), final item
794    (@parse $in_dim:expr, $out_dim:expr, [$($stack:tt)*] $r:expr) => {
795        match $r {
796            r => {
797                let in_dim = $crate::SliceNextDim::next_in_dim(&r, $in_dim);
798                let out_dim = $crate::SliceNextDim::next_out_dim(&r, $out_dim);
799                #[allow(unsafe_code)]
800                unsafe {
801                    $crate::SliceInfo::new_unchecked(
802                        [$($stack)* $crate::s!(@convert r)],
803                        in_dim,
804                        out_dim,
805                    )
806                }
807            }
808        }
809    };
810    // convert a..b;c into @convert(a..b, c), final item, trailing comma
811    (@parse $in_dim:expr, $out_dim:expr, [$($stack:tt)*] $r:expr;$s:expr ,) => {
812        $crate::s![@parse $in_dim, $out_dim, [$($stack)*] $r;$s]
813    };
814    // convert a..b into @convert(a..b), final item, trailing comma
815    (@parse $in_dim:expr, $out_dim:expr, [$($stack:tt)*] $r:expr ,) => {
816        $crate::s![@parse $in_dim, $out_dim, [$($stack)*] $r]
817    };
818    // convert a..b;c into @convert(a..b, c)
819    (@parse $in_dim:expr, $out_dim:expr, [$($stack:tt)*] $r:expr;$s:expr, $($t:tt)*) => {
820        match $r {
821            r => {
822                $crate::s![@parse
823                   $crate::SliceNextDim::next_in_dim(&r, $in_dim),
824                   $crate::SliceNextDim::next_out_dim(&r, $out_dim),
825                   [$($stack)* $crate::s!(@convert r, $s),]
826                   $($t)*
827                ]
828            }
829        }
830    };
831    // convert a..b into @convert(a..b)
832    (@parse $in_dim:expr, $out_dim:expr, [$($stack:tt)*] $r:expr, $($t:tt)*) => {
833        match $r {
834            r => {
835                $crate::s![@parse
836                   $crate::SliceNextDim::next_in_dim(&r, $in_dim),
837                   $crate::SliceNextDim::next_out_dim(&r, $out_dim),
838                   [$($stack)* $crate::s!(@convert r),]
839                   $($t)*
840                ]
841            }
842        }
843    };
844    // empty call, i.e. `s![]`
845    (@parse ::core::marker::PhantomData::<$crate::Ix0>, ::core::marker::PhantomData::<$crate::Ix0>, []) => {
846        {
847            #[allow(unsafe_code)]
848            unsafe {
849                $crate::SliceInfo::new_unchecked(
850                    [],
851                    ::core::marker::PhantomData::<$crate::Ix0>,
852                    ::core::marker::PhantomData::<$crate::Ix0>,
853                )
854            }
855        }
856    };
857    // Catch-all clause for syntax errors
858    (@parse $($t:tt)*) => { compile_error!("Invalid syntax in s![] call.") };
859    // convert range/index/new-axis into SliceInfoElem
860    (@convert $r:expr) => {
861        <$crate::SliceInfoElem as ::core::convert::From<_>>::from($r)
862    };
863    // convert range/index/new-axis and step into SliceInfoElem
864    (@convert $r:expr, $s:expr) => {
865        <$crate::SliceInfoElem as ::core::convert::From<_>>::from(
866            <$crate::Slice as ::core::convert::From<_>>::from($r).step_by($s as isize)
867        )
868    };
869    ($($t:tt)*) => {
870        $crate::s![@parse
871              ::core::marker::PhantomData::<$crate::Ix0>,
872              ::core::marker::PhantomData::<$crate::Ix0>,
873              []
874              $($t)*
875        ]
876    };
877);
878
879/// Slicing information describing multiple mutable, disjoint slices.
880///
881/// It's unfortunate that we need `'a` and `A` to be parameters of the trait,
882/// but they're necessary until Rust supports generic associated types.
883pub trait MultiSliceArg<'a, A, D>
884where
885    A: 'a,
886    D: Dimension,
887{
888    /// The type of the slices created by `.multi_slice_move()`.
889    type Output;
890
891    /// Split the view into multiple disjoint slices.
892    ///
893    /// **Panics** if performing any individual slice panics or if the slices
894    /// are not disjoint (i.e. if they intersect).
895    fn multi_slice_move(&self, view: ArrayViewMut<'a, A, D>) -> Self::Output;
896
897    private_decl! {}
898}
899
900impl<'a, A, D> MultiSliceArg<'a, A, D> for ()
901where
902    A: 'a,
903    D: Dimension,
904{
905    type Output = ();
906
907    fn multi_slice_move(&self, _view: ArrayViewMut<'a, A, D>) -> Self::Output {}
908
909    private_impl! {}
910}
911
912impl<'a, A, D, I0> MultiSliceArg<'a, A, D> for (I0,)
913where
914    A: 'a,
915    D: Dimension,
916    I0: SliceArg<D>,
917{
918    type Output = (ArrayViewMut<'a, A, I0::OutDim>,);
919
920    fn multi_slice_move(&self, view: ArrayViewMut<'a, A, D>) -> Self::Output {
921        (view.slice_move(&self.0),)
922    }
923
924    private_impl! {}
925}
926
927macro_rules! impl_multislice_tuple {
928    ([$($but_last:ident)*] $last:ident) => {
929        impl_multislice_tuple!(@def_impl ($($but_last,)* $last,), [$($but_last)*] $last);
930    };
931    (@def_impl ($($all:ident,)*), [$($but_last:ident)*] $last:ident) => {
932        impl<'a, A, D, $($all,)*> MultiSliceArg<'a, A, D> for ($($all,)*)
933        where
934            A: 'a,
935            D: Dimension,
936            $($all: SliceArg<D>,)*
937        {
938            type Output = ($(ArrayViewMut<'a, A, $all::OutDim>,)*);
939
940            fn multi_slice_move(&self, view: ArrayViewMut<'a, A, D>) -> Self::Output {
941                #[allow(non_snake_case)]
942                let ($($all,)*) = self;
943
944                let shape = view.raw_dim();
945                assert!(!impl_multislice_tuple!(@intersects_self &shape, ($($all,)*)));
946
947                let raw_view = view.into_raw_view_mut();
948                unsafe {
949                    (
950                        $(raw_view.clone().slice_move($but_last).deref_into_view_mut(),)*
951                        raw_view.slice_move($last).deref_into_view_mut(),
952                    )
953                }
954            }
955
956            private_impl! {}
957        }
958    };
959    (@intersects_self $shape:expr, ($head:expr,)) => {
960        false
961    };
962    (@intersects_self $shape:expr, ($head:expr, $($tail:expr,)*)) => {
963        $(slices_intersect($shape, $head, $tail)) ||*
964            || impl_multislice_tuple!(@intersects_self $shape, ($($tail,)*))
965    };
966}
967
968impl_multislice_tuple!([I0] I1);
969impl_multislice_tuple!([I0 I1] I2);
970impl_multislice_tuple!([I0 I1 I2] I3);
971impl_multislice_tuple!([I0 I1 I2 I3] I4);
972impl_multislice_tuple!([I0 I1 I2 I3 I4] I5);
973
974impl<'a, A, D, T> MultiSliceArg<'a, A, D> for &T
975where
976    A: 'a,
977    D: Dimension,
978    T: MultiSliceArg<'a, A, D>,
979{
980    type Output = T::Output;
981
982    fn multi_slice_move(&self, view: ArrayViewMut<'a, A, D>) -> Self::Output {
983        T::multi_slice_move(self, view)
984    }
985
986    private_impl! {}
987}
ndarray/slice.rs

ndarray/
slice.rs
