library/core/src/ops/index_range.rs

   1 use crate::intrinsics::{assert_unsafe_precondition, unchecked_add, unchecked_sub};
   2 use crate::iter::{FusedIterator, TrustedLen};
   3
   4 /// Like a `Range<usize>`, but with a safety invariant that `start <= end`.
   5 ///
   6 /// This means that `end - start` cannot overflow, allowing some μoptimizations.
   7 ///
   8 /// (Normal `Range` code needs to handle degenerate ranges like `10..0`,
   9 ///  which takes extra checks compared to only handling the canonical form.)
  10 #[derive(Clone, Debug, PartialEq, Eq)]
  11 pub(crate) struct IndexRange {
  12     start: usize,
  13     end: usize,
  14 }
  15
  16 impl IndexRange {
  17     /// # Safety
  18     /// - `start <= end`
  19     #[inline]
  20     pub const unsafe fn new_unchecked(start: usize, end: usize) -> Self {
  21         // SAFETY: comparisons on usize are pure
  22         unsafe { assert_unsafe_precondition!((start: usize, end: usize) => start <= end) };
  23         IndexRange { start, end }
  24     }
  25
  26     #[inline]
  27     pub const fn zero_to(end: usize) -> Self {
  28         IndexRange { start: 0, end }
  29     }
  30
  31     #[inline]
  32     pub const fn start(&self) -> usize {
  33         self.start
  34     }
  35
  36     #[inline]
  37     pub const fn end(&self) -> usize {
  38         self.end
  39     }
  40
  41     #[inline]
  42     pub const fn len(&self) -> usize {
  43         // SAFETY: By invariant, this cannot wrap
  44         unsafe { unchecked_sub(self.end, self.start) }
  45     }
  46
  47     /// # Safety
  48     /// - Can only be called when `start < end`, aka when `len > 0`.
  49     #[inline]
  50     unsafe fn next_unchecked(&mut self) -> usize {
  51         debug_assert!(self.start < self.end);
  52
  53         let value = self.start;
  54         // SAFETY: The range isn't empty, so this cannot overflow
  55         self.start = unsafe { unchecked_add(value, 1) };
  56         value
  57     }
  58
  59     /// # Safety
  60     /// - Can only be called when `start < end`, aka when `len > 0`.
  61     #[inline]
  62     unsafe fn next_back_unchecked(&mut self) -> usize {
  63         debug_assert!(self.start < self.end);
  64
  65         // SAFETY: The range isn't empty, so this cannot overflow
  66         let value = unsafe { unchecked_sub(self.end, 1) };
  67         self.end = value;
  68         value
  69     }
  70
  71     /// Removes the first `n` items from this range, returning them as an `IndexRange`.
  72     /// If there are fewer than `n`, then the whole range is returned and
  73     /// `self` is left empty.
  74     ///
  75     /// This is designed to help implement `Iterator::advance_by`.
  76     #[inline]
  77     pub fn take_prefix(&mut self, n: usize) -> Self {
  78         let mid = if n <= self.len() {
  79             // SAFETY: We just checked that this will be between start and end,
  80             // and thus the addition cannot overflow.
  81             unsafe { unchecked_add(self.start, n) }
  82         } else {
  83             self.end
  84         };
  85         let prefix = Self { start: self.start, end: mid };
  86         self.start = mid;
  87         prefix
  88     }
  89
  90     /// Removes the last `n` items from this range, returning them as an `IndexRange`.
  91     /// If there are fewer than `n`, then the whole range is returned and
  92     /// `self` is left empty.
  93     ///
  94     /// This is designed to help implement `Iterator::advance_back_by`.
  95     #[inline]
  96     pub fn take_suffix(&mut self, n: usize) -> Self {
  97         let mid = if n <= self.len() {
  98             // SAFETY: We just checked that this will be between start and end,
  99             // and thus the addition cannot overflow.
 100             unsafe { unchecked_sub(self.end, n) }
 101         } else {
 102             self.start
 103         };
 104         let suffix = Self { start: mid, end: self.end };
 105         self.end = mid;
 106         suffix
 107     }
 108 }
 109
 110 impl Iterator for IndexRange {
 111     type Item = usize;
 112
 113     #[inline]
 114     fn next(&mut self) -> Option<usize> {
 115         if self.len() > 0 {
 116             // SAFETY: We just checked that the range is non-empty
 117             unsafe { Some(self.next_unchecked()) }
 118         } else {
 119             None
 120         }
 121     }
 122
 123     #[inline]
 124     fn size_hint(&self) -> (usize, Option<usize>) {
 125         let len = self.len();
 126         (len, Some(len))
 127     }
 128
 129     #[inline]
 130     fn advance_by(&mut self, n: usize) -> Result<(), usize> {
 131         let original_len = self.len();
 132         self.take_prefix(n);
 133         if n > original_len { Err(original_len) } else { Ok(()) }
 134     }
 135 }
 136
 137 impl DoubleEndedIterator for IndexRange {
 138     #[inline]
 139     fn next_back(&mut self) -> Option<usize> {
 140         if self.len() > 0 {
 141             // SAFETY: We just checked that the range is non-empty
 142             unsafe { Some(self.next_back_unchecked()) }
 143         } else {
 144             None
 145         }
 146     }
 147
 148     #[inline]
 149     fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
 150         let original_len = self.len();
 151         self.take_suffix(n);
 152         if n > original_len { Err(original_len) } else { Ok(()) }
 153     }
 154 }
 155
 156 impl ExactSizeIterator for IndexRange {
 157     #[inline]
 158     fn len(&self) -> usize {
 159         self.len()
 160     }
 161 }
 162
 163 // SAFETY: Because we only deal in `usize`, our `len` is always perfect.
 164 unsafe impl TrustedLen for IndexRange {}
 165
 166 impl FusedIterator for IndexRange {}