vendor/itertools/src/permutations.rs

   1 use std::fmt;
   2 use std::iter::once;
   3
   4 use super::lazy_buffer::LazyBuffer;
   5
   6 /// An iterator adaptor that iterates through all the `k`-permutations of the
   7 /// elements from an iterator.
   8 ///
   9 /// See [`.permutations()`](../trait.Itertools.html#method.permutations) for
  10 /// more information.
  11 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
  12 pub struct Permutations<I: Iterator> {
  13     vals: LazyBuffer<I>,
  14     state: PermutationState,
  15 }
  16
  17 impl<I> Clone for Permutations<I>
  18     where I: Clone + Iterator,
  19           I::Item: Clone,
  20 {
  21     clone_fields!(vals, state);
  22 }
  23
  24 #[derive(Clone, Debug)]
  25 enum PermutationState {
  26     StartUnknownLen {
  27         k: usize,
  28     },
  29     OngoingUnknownLen {
  30         k: usize,
  31         min_n: usize,
  32     },
  33     Complete(CompleteState),
  34     Empty,
  35 }
  36
  37 #[derive(Clone, Debug)]
  38 enum CompleteState {
  39     Start {
  40         n: usize,
  41         k: usize,
  42     },
  43     Ongoing {
  44         indices: Vec<usize>,
  45         cycles: Vec<usize>,
  46     }
  47 }
  48
  49 enum CompleteStateRemaining {
  50     Known(usize),
  51     Overflow,
  52 }
  53
  54 impl<I> fmt::Debug for Permutations<I>
  55     where I: Iterator + fmt::Debug,
  56           I::Item: fmt::Debug,
  57 {
  58     debug_fmt_fields!(Permutations, vals, state);
  59 }
  60
  61 pub fn permutations<I: Iterator>(iter: I, k: usize) -> Permutations<I> {
  62     let mut vals = LazyBuffer::new(iter);
  63
  64     if k == 0 {
  65         // Special case, yields single empty vec; `n` is irrelevant
  66         let state = PermutationState::Complete(CompleteState::Start { n: 0, k: 0 });
  67
  68         return Permutations {
  69             vals,
  70             state
  71         };
  72     }
  73
  74     let mut enough_vals = true;
  75
  76     while vals.len() < k {
  77         if !vals.get_next() {
  78             enough_vals = false;
  79             break;
  80         }
  81     }
  82
  83     let state = if enough_vals {
  84         PermutationState::StartUnknownLen { k }
  85     } else {
  86         PermutationState::Empty
  87     };
  88
  89     Permutations {
  90         vals,
  91         state
  92     }
  93 }
  94
  95 impl<I> Iterator for Permutations<I>
  96 where
  97     I: Iterator,
  98     I::Item: Clone
  99 {
 100     type Item = Vec<I::Item>;
 101
 102     fn next(&mut self) -> Option<Self::Item> {
 103         self.advance();
 104
 105         let &mut Permutations { ref vals, ref state } = self;
 106
 107         match state {
 108             &PermutationState::StartUnknownLen { .. } => panic!("unexpected iterator state"),
 109             &PermutationState::OngoingUnknownLen { k, min_n } => {
 110                 let latest_idx = min_n - 1;
 111                 let indices = (0..(k - 1)).chain(once(latest_idx));
 112
 113                 Some(indices.map(|i| vals[i].clone()).collect())
 114             }
 115             &PermutationState::Complete(CompleteState::Start { .. }) => None,
 116             &PermutationState::Complete(CompleteState::Ongoing { ref indices, ref cycles }) => {
 117                 let k = cycles.len();
 118
 119                 Some(indices[0..k].iter().map(|&i| vals[i].clone()).collect())
 120             },
 121             &PermutationState::Empty => None
 122         }
 123     }
 124
 125     fn count(self) -> usize {
 126         let Permutations { vals, state } = self;
 127
 128         fn from_complete(complete_state: CompleteState) -> usize {
 129             match complete_state.remaining() {
 130                 CompleteStateRemaining::Known(count) => count,
 131                 CompleteStateRemaining::Overflow => {
 132                     panic!("Iterator count greater than usize::MAX");
 133                 }
 134             }
 135         }
 136
 137         match state {
 138             PermutationState::StartUnknownLen { k } => {
 139                 let n = vals.len() + vals.it.count();
 140                 let complete_state = CompleteState::Start { n, k };
 141
 142                 from_complete(complete_state)
 143             }
 144             PermutationState::OngoingUnknownLen { k, min_n } => {
 145                 let prev_iteration_count = min_n - k + 1;
 146                 let n = vals.len() + vals.it.count();
 147                 let complete_state = CompleteState::Start { n, k };
 148
 149                 from_complete(complete_state) - prev_iteration_count
 150             },
 151             PermutationState::Complete(state) => from_complete(state),
 152             PermutationState::Empty => 0
 153         }
 154     }
 155
 156     fn size_hint(&self) -> (usize, Option<usize>) {
 157         match self.state {
 158             PermutationState::StartUnknownLen { .. } |
 159             PermutationState::OngoingUnknownLen { .. } => (0, None), // TODO can we improve this lower bound?
 160             PermutationState::Complete(ref state) => match state.remaining() {
 161                 CompleteStateRemaining::Known(count) => (count, Some(count)),
 162                 CompleteStateRemaining::Overflow => (::std::usize::MAX, None)
 163             }
 164             PermutationState::Empty => (0, Some(0))
 165         }
 166     }
 167 }
 168
 169 impl<I> Permutations<I>
 170 where
 171     I: Iterator,
 172     I::Item: Clone
 173 {
 174     fn advance(&mut self) {
 175         let &mut Permutations { ref mut vals, ref mut state } = self;
 176
 177         *state = match state {
 178             &mut PermutationState::StartUnknownLen { k } => {
 179                 PermutationState::OngoingUnknownLen { k, min_n: k }
 180             }
 181             &mut PermutationState::OngoingUnknownLen { k, min_n } => {
 182                 if vals.get_next() {
 183                     PermutationState::OngoingUnknownLen { k, min_n: min_n + 1 }
 184                 } else {
 185                     let n = min_n;
 186                     let prev_iteration_count = n - k + 1;
 187                     let mut complete_state = CompleteState::Start { n, k };
 188
 189                     // Advance the complete-state iterator to the correct point
 190                     for _ in 0..(prev_iteration_count + 1) {
 191                         complete_state.advance();
 192                     }
 193
 194                     PermutationState::Complete(complete_state)
 195                 }
 196             }
 197             &mut PermutationState::Complete(ref mut state) => {
 198                 state.advance();
 199
 200                 return;
 201             }
 202             &mut PermutationState::Empty => { return; }
 203         };
 204     }
 205 }
 206
 207 impl CompleteState {
 208     fn advance(&mut self) {
 209         *self = match self {
 210             &mut CompleteState::Start { n, k } => {
 211                 let indices = (0..n).collect();
 212                 let cycles = ((n - k)..n).rev().collect();
 213
 214                 CompleteState::Ongoing {
 215                     cycles,
 216                     indices
 217                 }
 218             },
 219             &mut CompleteState::Ongoing { ref mut indices, ref mut cycles } => {
 220                 let n = indices.len();
 221                 let k = cycles.len();
 222
 223                 for i in (0..k).rev() {
 224                     if cycles[i] == 0 {
 225                         cycles[i] = n - i - 1;
 226
 227                         let to_push = indices.remove(i);
 228                         indices.push(to_push);
 229                     } else {
 230                         let swap_index = n - cycles[i];
 231                         indices.swap(i, swap_index);
 232
 233                         cycles[i] -= 1;
 234                         return;
 235                     }
 236                 }
 237
 238                 CompleteState::Start { n, k }
 239             }
 240         }
 241     }
 242
 243     fn remaining(&self) -> CompleteStateRemaining {
 244         use self::CompleteStateRemaining::{Known, Overflow};
 245
 246         match self {
 247             &CompleteState::Start { n, k } => {
 248                 if n < k {
 249                     return Known(0);
 250                 }
 251
 252                 let count: Option<usize> = (n - k + 1..n + 1).fold(Some(1), |acc, i| {
 253                     acc.and_then(|acc| acc.checked_mul(i))
 254                 });
 255
 256                 match count {
 257                     Some(count) => Known(count),
 258                     None => Overflow
 259                 }
 260             }
 261             &CompleteState::Ongoing { ref indices, ref cycles } => {
 262                 let mut count: usize = 0;
 263
 264                 for (i, &c) in cycles.iter().enumerate() {
 265                     let radix = indices.len() - i;
 266                     let next_count = count.checked_mul(radix)
 267                         .and_then(|count| count.checked_add(c));
 268
 269                     count = match next_count {
 270                         Some(count) => count,
 271                         None => { return Overflow; }
 272                     };
 273                 }
 274
 275                 Known(count)
 276             }
 277         }
 278     }
 279 }