powerset.rs

use std::fmt;
use std::usize;
use std::iter::{self, Empty};
use super::combinations::Combinations;
use super::size_hint;

/// An iterator over the powerset of all elements from its source iterator.
///
/// See [`.powerset()`](../trait.Itertools.html#method.powerset) for more information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct Powerset<I: Iterator> {
    src: I,
    buf: Vec<I::Item>,
    // Only used when once a source containing two or more items is exhausted.
    combs: Option<SetCombinations<I>>,
    // Used for detecting special cases and size_hint() calculation
    pos: usize,
}

type SetCombinations<I> = Combinations<iter::Empty<<I as Iterator>::Item>>;

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

/// Create a new `Powerset` from a cloneable iterator.
pub fn powerset<I>(src: I) -> Powerset<I>
    where I: Iterator
{
    let (src_low, _) = src.size_hint();

    Powerset {
        src,
        buf: Vec::with_capacity(src_low),
        combs: None,
        pos: 0,
    }
}

impl<I> Powerset<I>
    where
        I: Iterator,
        I::Item: Clone
{
    #[inline]
    /// Creates and stores the next Combination, returning the first element or None if the
    /// iterator is done.
    fn next_from_new_combs(&mut self) -> Option<Vec<I::Item>> {
        debug_assert!(self.buf.len() > 1);

        let new_len = match &self.combs {
            Some(combs) => combs.n() + 1,
            None => 2
        };

        let mut combs = Combinations::<Empty<I::Item>>::from_slice(&self.buf[..], new_len);
        let result = combs.next();
        self.combs = Some(combs);

        result
    }
}

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

    fn next(&mut self) -> Option<Vec<I::Item>> {
        let mut new_combs = false;

        let mut result = match &mut self.combs {
            None if self.pos == 0 => {
                // First item, return the empty set
                Some(Vec::new())
            }
            None => {
                // Still draining from source iterator
                if let Some(elt) = self.src.next() {
                    self.buf.push(elt.clone());
                    Some(vec!(elt))
                } else {
                    new_combs = self.buf.len() >= 2;
                    None
                }
            }
            Some(combs) if combs.n() < self.buf.len() => {
                // Generating elements from Combinations
                match combs.next() {
                    Some(elt) => Some(elt),
                    None => {
                        new_combs = self.buf.len() >= 2;
                        None
                    }
                }
            }
            Some(_) => {
                // Iteration is done
                None
            }
        };

        // Merge these lines into the relevant match arms above when NLLs are available
        if new_combs { result = self.next_from_new_combs(); }
        if result.is_some() { self.pos = self.pos.saturating_add(1); }

        result
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        // Aggregate bounds for 'unwalked' source iterator
        let src_total = size_hint::add_scalar(self.src.size_hint(), self.buf.len());

        // Aggregate bounds for self ( length(powerset(set) == 2 ^ length(set) )
        let self_total = size_hint::two_raised_pow(src_total);

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