Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 1 | //! Arithmetic on **Iterator** *.size_hint()* values. |
| 2 | //! |
| 3 | |
| 4 | use std::usize; |
| 5 | use std::cmp; |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 6 | use std::u32; |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 7 | |
| 8 | /// **SizeHint** is the return type of **Iterator::size_hint()**. |
| 9 | pub type SizeHint = (usize, Option<usize>); |
| 10 | |
| 11 | /// Add **SizeHint** correctly. |
| 12 | #[inline] |
| 13 | pub fn add(a: SizeHint, b: SizeHint) -> SizeHint { |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 14 | let min = a.0.saturating_add(b.0); |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 15 | let max = match (a.1, b.1) { |
| 16 | (Some(x), Some(y)) => x.checked_add(y), |
| 17 | _ => None, |
| 18 | }; |
| 19 | |
| 20 | (min, max) |
| 21 | } |
| 22 | |
| 23 | /// Add **x** correctly to a **SizeHint**. |
| 24 | #[inline] |
| 25 | pub fn add_scalar(sh: SizeHint, x: usize) -> SizeHint { |
| 26 | let (mut low, mut hi) = sh; |
| 27 | low = low.saturating_add(x); |
| 28 | hi = hi.and_then(|elt| elt.checked_add(x)); |
| 29 | (low, hi) |
| 30 | } |
| 31 | |
| 32 | /// Sbb **x** correctly to a **SizeHint**. |
| 33 | #[inline] |
| 34 | #[allow(dead_code)] |
| 35 | pub fn sub_scalar(sh: SizeHint, x: usize) -> SizeHint { |
| 36 | let (mut low, mut hi) = sh; |
| 37 | low = low.saturating_sub(x); |
| 38 | hi = hi.map(|elt| elt.saturating_sub(x)); |
| 39 | (low, hi) |
| 40 | } |
| 41 | |
| 42 | |
| 43 | /// Multiply **SizeHint** correctly |
| 44 | /// |
| 45 | /// ```ignore |
| 46 | /// use std::usize; |
| 47 | /// use itertools::size_hint; |
| 48 | /// |
| 49 | /// assert_eq!(size_hint::mul((3, Some(4)), (3, Some(4))), |
| 50 | /// (9, Some(16))); |
| 51 | /// |
| 52 | /// assert_eq!(size_hint::mul((3, Some(4)), (usize::MAX, None)), |
| 53 | /// (usize::MAX, None)); |
| 54 | /// |
| 55 | /// assert_eq!(size_hint::mul((3, None), (0, Some(0))), |
| 56 | /// (0, Some(0))); |
| 57 | /// ``` |
| 58 | #[inline] |
| 59 | pub fn mul(a: SizeHint, b: SizeHint) -> SizeHint { |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 60 | let low = a.0.saturating_mul(b.0); |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 61 | let hi = match (a.1, b.1) { |
| 62 | (Some(x), Some(y)) => x.checked_mul(y), |
| 63 | (Some(0), None) | (None, Some(0)) => Some(0), |
| 64 | _ => None, |
| 65 | }; |
| 66 | (low, hi) |
| 67 | } |
| 68 | |
| 69 | /// Multiply **x** correctly with a **SizeHint**. |
| 70 | #[inline] |
| 71 | pub fn mul_scalar(sh: SizeHint, x: usize) -> SizeHint { |
| 72 | let (mut low, mut hi) = sh; |
| 73 | low = low.saturating_mul(x); |
| 74 | hi = hi.and_then(|elt| elt.checked_mul(x)); |
| 75 | (low, hi) |
| 76 | } |
| 77 | |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 78 | /// Raise `base` correctly by a **`SizeHint`** exponent. |
| 79 | #[inline] |
| 80 | pub fn pow_scalar_base(base: usize, exp: SizeHint) -> SizeHint { |
| 81 | let exp_low = cmp::min(exp.0, u32::MAX as usize) as u32; |
| 82 | let low = base.saturating_pow(exp_low); |
| 83 | |
| 84 | let hi = exp.1.and_then(|exp| { |
| 85 | let exp_hi = cmp::min(exp, u32::MAX as usize) as u32; |
| 86 | base.checked_pow(exp_hi) |
| 87 | }); |
| 88 | |
| 89 | (low, hi) |
| 90 | } |
| 91 | |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 92 | /// Return the maximum |
| 93 | #[inline] |
| 94 | pub fn max(a: SizeHint, b: SizeHint) -> SizeHint { |
| 95 | let (a_lower, a_upper) = a; |
| 96 | let (b_lower, b_upper) = b; |
| 97 | |
| 98 | let lower = cmp::max(a_lower, b_lower); |
| 99 | |
| 100 | let upper = match (a_upper, b_upper) { |
| 101 | (Some(x), Some(y)) => Some(cmp::max(x, y)), |
| 102 | _ => None, |
| 103 | }; |
| 104 | |
| 105 | (lower, upper) |
| 106 | } |
| 107 | |
| 108 | /// Return the minimum |
| 109 | #[inline] |
| 110 | pub fn min(a: SizeHint, b: SizeHint) -> SizeHint { |
| 111 | let (a_lower, a_upper) = a; |
| 112 | let (b_lower, b_upper) = b; |
| 113 | let lower = cmp::min(a_lower, b_lower); |
| 114 | let upper = match (a_upper, b_upper) { |
| 115 | (Some(u1), Some(u2)) => Some(cmp::min(u1, u2)), |
| 116 | _ => a_upper.or(b_upper), |
| 117 | }; |
| 118 | (lower, upper) |
| 119 | } |