summaryrefslogtreecommitdiffstats
path: root/third_party/rust/itertools/src/powerset.rs
blob: 4d7685b12a87a36157a927c235927ee7d6c79499 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
use std::fmt;
use std::iter::FusedIterator;
use std::usize;
use alloc::vec::Vec;

use super::combinations::{Combinations, combinations};
use super::size_hint;

/// An iterator to iterate through the powerset of the elements from an iterator.
///
/// See [`.powerset()`](crate::Itertools::powerset) for more
/// information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct Powerset<I: Iterator> {
    combs: Combinations<I>,
    // Iterator `position` (equal to count of yielded elements).
    pos: usize,
}

impl<I> Clone for Powerset<I>
    where I: Clone + Iterator,
          I::Item: Clone,
{
    clone_fields!(combs, pos);
}

impl<I> fmt::Debug for Powerset<I>
    where I: Iterator + fmt::Debug,
          I::Item: fmt::Debug,
{
    debug_fmt_fields!(Powerset, combs, pos);
}

/// Create a new `Powerset` from a clonable iterator.
pub fn powerset<I>(src: I) -> Powerset<I>
    where I: Iterator,
          I::Item: Clone,
{
    Powerset {
        combs: combinations(src, 0),
        pos: 0,
    }
}

impl<I> Iterator for Powerset<I>
    where
        I: Iterator,
        I::Item: Clone,
{
    type Item = Vec<I::Item>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(elt) = self.combs.next() {
            self.pos = self.pos.saturating_add(1);
            Some(elt)
        } else if self.combs.k() < self.combs.n()
            || self.combs.k() == 0
        {
            self.combs.reset(self.combs.k() + 1);
            self.combs.next().map(|elt| {
                self.pos = self.pos.saturating_add(1);
                elt
            })
        } else {
            None
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        // Total bounds for source iterator.
        let src_total = size_hint::add_scalar(self.combs.src().size_hint(), self.combs.n());

        // Total bounds for self ( length(powerset(set) == 2 ^ length(set) )
        let self_total = size_hint::pow_scalar_base(2, src_total);

        if self.pos < usize::MAX {
            // Subtract count of elements already yielded from total.
            size_hint::sub_scalar(self_total, self.pos)
        } else {
            // Fallback: self.pos is saturated and no longer reliable.
            (0, self_total.1)
        }
    }
}

impl<I> FusedIterator for Powerset<I>
    where
        I: Iterator,
        I::Item: Clone,
{}