Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 1 | use std::fmt; |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 2 | use std::iter::FusedIterator; |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 3 | use std::usize; |
| 4 | use alloc::vec::Vec; |
| 5 | |
| 6 | use super::combinations::{Combinations, combinations}; |
| 7 | use super::size_hint; |
| 8 | |
| 9 | /// An iterator to iterate through the powerset of the elements from an iterator. |
| 10 | /// |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 11 | /// See [`.powerset()`](crate::Itertools::powerset) for more |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 12 | /// information. |
| 13 | #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] |
| 14 | pub struct Powerset<I: Iterator> { |
| 15 | combs: Combinations<I>, |
| 16 | // Iterator `position` (equal to count of yielded elements). |
| 17 | pos: usize, |
| 18 | } |
| 19 | |
| 20 | impl<I> Clone for Powerset<I> |
| 21 | where I: Clone + Iterator, |
| 22 | I::Item: Clone, |
| 23 | { |
| 24 | clone_fields!(combs, pos); |
| 25 | } |
| 26 | |
| 27 | impl<I> fmt::Debug for Powerset<I> |
| 28 | where I: Iterator + fmt::Debug, |
| 29 | I::Item: fmt::Debug, |
| 30 | { |
| 31 | debug_fmt_fields!(Powerset, combs, pos); |
| 32 | } |
| 33 | |
| 34 | /// Create a new `Powerset` from a clonable iterator. |
| 35 | pub fn powerset<I>(src: I) -> Powerset<I> |
| 36 | where I: Iterator, |
| 37 | I::Item: Clone, |
| 38 | { |
| 39 | Powerset { |
| 40 | combs: combinations(src, 0), |
| 41 | pos: 0, |
| 42 | } |
| 43 | } |
| 44 | |
| 45 | impl<I> Iterator for Powerset<I> |
| 46 | where |
| 47 | I: Iterator, |
| 48 | I::Item: Clone, |
| 49 | { |
| 50 | type Item = Vec<I::Item>; |
| 51 | |
| 52 | fn next(&mut self) -> Option<Self::Item> { |
| 53 | if let Some(elt) = self.combs.next() { |
| 54 | self.pos = self.pos.saturating_add(1); |
| 55 | Some(elt) |
| 56 | } else if self.combs.k() < self.combs.n() |
| 57 | || self.combs.k() == 0 |
| 58 | { |
| 59 | self.combs.reset(self.combs.k() + 1); |
| 60 | self.combs.next().map(|elt| { |
| 61 | self.pos = self.pos.saturating_add(1); |
| 62 | elt |
| 63 | }) |
| 64 | } else { |
| 65 | None |
| 66 | } |
| 67 | } |
| 68 | |
| 69 | fn size_hint(&self) -> (usize, Option<usize>) { |
| 70 | // Total bounds for source iterator. |
| 71 | let src_total = size_hint::add_scalar(self.combs.src().size_hint(), self.combs.n()); |
| 72 | |
| 73 | // Total bounds for self ( length(powerset(set) == 2 ^ length(set) ) |
| 74 | let self_total = size_hint::pow_scalar_base(2, src_total); |
| 75 | |
| 76 | if self.pos < usize::MAX { |
| 77 | // Subtract count of elements already yielded from total. |
| 78 | size_hint::sub_scalar(self_total, self.pos) |
| 79 | } else { |
| 80 | // Fallback: self.pos is saturated and no longer reliable. |
| 81 | (0, self_total.1) |
| 82 | } |
| 83 | } |
| 84 | } |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 85 | |
| 86 | impl<I> FusedIterator for Powerset<I> |
| 87 | where |
| 88 | I: Iterator, |
| 89 | I::Item: Clone, |
| 90 | {} |