zipslices.rs 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188
  1. use std::cmp;
  2. // Note: There are different ways to implement ZipSlices.
  3. // This version performed the best in benchmarks.
  4. //
  5. // I also implemented a version with three pointes (tptr, tend, uptr),
  6. // that mimiced slice::Iter and only checked bounds by using tptr == tend,
  7. // but that was inferior to this solution.
  8. /// An iterator which iterates two slices simultaneously.
  9. ///
  10. /// `ZipSlices` acts like a double-ended `.zip()` iterator.
  11. ///
  12. /// It was intended to be more efficient than `.zip()`, and it was, then
  13. /// rustc changed how it optimizes so it can not promise improved performance
  14. /// at this time.
  15. ///
  16. /// Note that elements past the end of the shortest of the two slices are ignored.
  17. ///
  18. /// Iterator element type for `ZipSlices<T, U>` is `(T::Item, U::Item)`. For example,
  19. /// for a `ZipSlices<&'a [A], &'b mut [B]>`, the element type is `(&'a A, &'b mut B)`.
  20. #[derive(Clone)]
  21. pub struct ZipSlices<T, U> {
  22. t: T,
  23. u: U,
  24. len: usize,
  25. index: usize,
  26. }
  27. impl<'a, 'b, A, B> ZipSlices<&'a [A], &'b [B]> {
  28. /// Create a new `ZipSlices` from slices `a` and `b`.
  29. ///
  30. /// Act like a double-ended `.zip()` iterator, but more efficiently.
  31. ///
  32. /// Note that elements past the end of the shortest of the two slices are ignored.
  33. #[inline(always)]
  34. pub fn new(a: &'a [A], b: &'b [B]) -> Self {
  35. let minl = cmp::min(a.len(), b.len());
  36. ZipSlices {
  37. t: a,
  38. u: b,
  39. len: minl,
  40. index: 0,
  41. }
  42. }
  43. }
  44. impl<T, U> ZipSlices<T, U>
  45. where T: Slice,
  46. U: Slice
  47. {
  48. /// Create a new `ZipSlices` from slices `a` and `b`.
  49. ///
  50. /// Act like a double-ended `.zip()` iterator, but more efficiently.
  51. ///
  52. /// Note that elements past the end of the shortest of the two slices are ignored.
  53. #[inline(always)]
  54. pub fn from_slices(a: T, b: U) -> Self {
  55. let minl = cmp::min(a.len(), b.len());
  56. ZipSlices {
  57. t: a,
  58. u: b,
  59. len: minl,
  60. index: 0,
  61. }
  62. }
  63. }
  64. impl<T, U> Iterator for ZipSlices<T, U>
  65. where T: Slice,
  66. U: Slice
  67. {
  68. type Item = (T::Item, U::Item);
  69. #[inline(always)]
  70. fn next(&mut self) -> Option<Self::Item> {
  71. unsafe {
  72. if self.index >= self.len {
  73. None
  74. } else {
  75. let i = self.index;
  76. self.index += 1;
  77. Some((
  78. self.t.get_unchecked(i),
  79. self.u.get_unchecked(i)))
  80. }
  81. }
  82. }
  83. #[inline]
  84. fn size_hint(&self) -> (usize, Option<usize>) {
  85. let len = self.len - self.index;
  86. (len, Some(len))
  87. }
  88. }
  89. impl<T, U> DoubleEndedIterator for ZipSlices<T, U>
  90. where T: Slice,
  91. U: Slice
  92. {
  93. #[inline(always)]
  94. fn next_back(&mut self) -> Option<Self::Item> {
  95. unsafe {
  96. if self.index >= self.len {
  97. None
  98. } else {
  99. self.len -= 1;
  100. let i = self.len;
  101. Some((
  102. self.t.get_unchecked(i),
  103. self.u.get_unchecked(i)))
  104. }
  105. }
  106. }
  107. }
  108. impl<T, U> ExactSizeIterator for ZipSlices<T, U>
  109. where T: Slice,
  110. U: Slice
  111. {}
  112. unsafe impl<T, U> Slice for ZipSlices<T, U>
  113. where T: Slice,
  114. U: Slice
  115. {
  116. type Item = (T::Item, U::Item);
  117. fn len(&self) -> usize {
  118. self.len - self.index
  119. }
  120. unsafe fn get_unchecked(&mut self, i: usize) -> Self::Item {
  121. (self.t.get_unchecked(i),
  122. self.u.get_unchecked(i))
  123. }
  124. }
  125. /// A helper trait to let `ZipSlices` accept both `&[T]` and `&mut [T]`.
  126. ///
  127. /// Unsafe trait because:
  128. ///
  129. /// - Implementors must guarantee that `get_unchecked` is valid for all indices `0..len()`.
  130. pub unsafe trait Slice {
  131. /// The type of a reference to the slice's elements
  132. type Item;
  133. #[doc(hidden)]
  134. fn len(&self) -> usize;
  135. #[doc(hidden)]
  136. unsafe fn get_unchecked(&mut self, i: usize) -> Self::Item;
  137. }
  138. unsafe impl<'a, T> Slice for &'a [T] {
  139. type Item = &'a T;
  140. #[inline(always)]
  141. fn len(&self) -> usize { (**self).len() }
  142. #[inline(always)]
  143. unsafe fn get_unchecked(&mut self, i: usize) -> &'a T {
  144. debug_assert!(i < self.len());
  145. (**self).get_unchecked(i)
  146. }
  147. }
  148. unsafe impl<'a, T> Slice for &'a mut [T] {
  149. type Item = &'a mut T;
  150. #[inline(always)]
  151. fn len(&self) -> usize { (**self).len() }
  152. #[inline(always)]
  153. unsafe fn get_unchecked(&mut self, i: usize) -> &'a mut T {
  154. debug_assert!(i < self.len());
  155. // override the lifetime constraints of &mut &'a mut [T]
  156. (*(*self as *mut [T])).get_unchecked_mut(i)
  157. }
  158. }
  159. #[test]
  160. fn zipslices() {
  161. let xs = [1, 2, 3, 4, 5, 6];
  162. let ys = [1, 2, 3, 7];
  163. ::itertools::assert_equal(ZipSlices::new(&xs, &ys), xs.iter().zip(&ys));
  164. let xs = [1, 2, 3, 4, 5, 6];
  165. let mut ys = [0; 6];
  166. for (x, y) in ZipSlices::from_slices(&xs[..], &mut ys[..]) {
  167. *y = *x;
  168. }
  169. ::itertools::assert_equal(&xs, &ys);
  170. }