some simplifications and refactors

ronnodas · ronnodas · commit a916d375dceb · 2025-01-09T18:07:38.000+01:00
diff --git a/src/combinations.rs b/src/combinations.rs
@@ -1,9 +1,7 @@
-use core::array;
-use core::borrow::BorrowMut;
 use std::fmt;
 use std::iter::FusedIterator;
 
-use super::lazy_buffer::LazyBuffer;
+use super::lazy_buffer::{ConstUsize, LazyBuffer, MaybeConstUsize as _, PoolIndex};
 use alloc::vec::Vec;
 
 use crate::adaptors::checked_binomial;
@@ -18,15 +16,15 @@ pub fn combinations<I: Iterator>(iter: I, k: usize) -> Combinations<I>
 where
     I::Item: Clone,
 {
-    Combinations::new(iter, (0..k).collect())
+    Combinations::new(iter, PoolIndex::start(k))
 }
 
 /// Create a new `ArrayCombinations` from a clonable iterator.
 pub fn array_combinations<I: Iterator, const K: usize>(iter: I) -> ArrayCombinations<I, K>
 where
     I::Item: Clone,
 {
-    ArrayCombinations::new(iter, array::from_fn(|i| i))
+    ArrayCombinations::new(iter, PoolIndex::start(ConstUsize))
 }
 
 /// An iterator to iterate through all the `k`-length combinations in an iterator.
@@ -39,79 +37,6 @@ pub struct CombinationsGeneric<I: Iterator, Idx> {
     first: bool,
 }
 
-pub trait MaybeConstUsize : Clone + Copy + std::fmt::Debug {
-    /*TODO const*/fn value(self) -> usize;
-}
-
-#[derive(Clone, Copy, Debug)]
-pub struct ConstUsize<const N: usize>;
-impl<const N: usize> MaybeConstUsize for ConstUsize<N> {
-    fn value(self) -> usize {
-        N
-    }
-}
-
-impl MaybeConstUsize for usize {
-    fn value(self) -> usize {
-        self
-    }
-}
-
-/// A type holding indices of elements in a pool or buffer of items from an inner iterator
-/// and used to pick out different combinations in a generic way.
-pub trait PoolIndex: BorrowMut<[usize]> {
-    type Item<T>;
-    type Length: MaybeConstUsize;
-
-    fn extract_item<I: Iterator>(&self, pool: &LazyBuffer<I>) -> Self::Item<I::Item>
-    where
-        I::Item: Clone;
-
-    fn from_fn<T, F: Fn(usize)->T>(k: Self::Length, f: F) -> Self::Item<T>;
-
-    fn len(&self) -> Self::Length;
-}
-
-impl PoolIndex for Vec<usize> {
-    type Item<T> = Vec<T>;
-    type Length = usize;
-
-    fn extract_item<I: Iterator>(&self, pool: &LazyBuffer<I>) -> Self::Item<I::Item>
-    where
-        I::Item: Clone
-    {
-        pool.get_at(self)
-    }
-
-    fn from_fn<T, F: Fn(usize)->T>(k: Self::Length, f: F) -> Self::Item<T> {
-        (0..k).map(f).collect()
-    }
-    
-    fn len(&self) -> Self::Length {
-        self.len()
-    }
-}
-
-impl<const K: usize> PoolIndex for [usize; K] {
-    type Item<T> = [T; K];
-    type Length = ConstUsize<K>;
-
-    fn extract_item<I: Iterator>(&self, pool: &LazyBuffer<I>) -> Self::Item<I::Item>
-    where
-        I::Item: Clone
-    {
-        pool.get_array(*self)
-    }
-
-    fn from_fn<T, F: Fn(usize)->T>(_k: Self::Length, f: F) -> Self::Item<T> {
-        std::array::from_fn(f)
-    }
-
-    fn len(&self) -> Self::Length {
-        ConstUsize::<K>
-    }
-}
-
 impl<I, Idx> Clone for CombinationsGeneric<I, Idx>
 where
     I: Iterator + Clone,
@@ -142,8 +67,8 @@ impl<I: Iterator, Idx: PoolIndex> CombinationsGeneric<I, Idx> {
 
     /// Returns the length of a combination produced by this iterator.
     #[inline]
-    pub fn k(&self) -> Idx::Length {
-        self.indices.len()
+    pub fn k(&self) -> usize {
+        self.indices.len().value()
     }
 
     /// Returns the (current) length of the pool from which combination elements are
diff --git a/src/lazy_buffer.rs b/src/lazy_buffer.rs
@@ -1,4 +1,6 @@
 use alloc::vec::Vec;
+use core::borrow::BorrowMut;
+use core::ops::Deref;
 use std::iter::Fuse;
 use std::ops::Index;
 
@@ -60,7 +62,7 @@ where
         indices.iter().map(|i| self.buffer[*i].clone()).collect()
     }
 
-    pub fn get_array<const K: usize>(&self, indices: [usize; K]) -> [I::Item; K] {
+    pub fn get_array<const K: usize>(&self, indices: &[usize; K]) -> [I::Item; K] {
         indices.map(|i| self.buffer[i].clone())
     }
 }
@@ -77,3 +79,118 @@ where
         self.buffer.index(index)
     }
 }
+
+pub trait MaybeConstUsize: Clone + Copy + std::fmt::Debug {
+    /*TODO const*/
+    fn value(self) -> usize;
+}
+
+#[derive(Clone, Copy, Debug)]
+pub struct ConstUsize<const N: usize>;
+impl<const N: usize> MaybeConstUsize for ConstUsize<N> {
+    fn value(self) -> usize {
+        N
+    }
+}
+
+impl MaybeConstUsize for usize {
+    fn value(self) -> usize {
+        self
+    }
+}
+
+/// A type holding indices of elements in a pool or buffer of items from an inner iterator
+/// and used to pick out different combinations in a generic way.
+pub trait PoolIndex: BorrowMut<[usize]> {
+    type Item<T>;
+    type Length: MaybeConstUsize;
+
+    fn extract_item<I: Iterator>(&self, pool: &LazyBuffer<I>) -> Self::Item<I::Item>
+    where
+        I::Item: Clone;
+
+    fn from_fn<F: Fn(usize) -> usize>(k: Self::Length, f: F) -> Self;
+
+    fn item_from_fn<T, F: Fn(usize) -> T>(k: Self::Length, f: F) -> Self::Item<T>;
+
+    fn len(&self) -> Self::Length;
+
+    fn start(len: Self::Length) -> Self
+    where
+        Self: Sized,
+    {
+        Self::from_fn(len, |i| i)
+    }
+}
+
+impl PoolIndex for Vec<usize> {
+    type Item<T> = Vec<T>;
+    type Length = usize;
+
+    fn extract_item<I: Iterator>(&self, pool: &LazyBuffer<I>) -> Self::Item<I::Item>
+    where
+        I::Item: Clone,
+    {
+        pool.get_at(self)
+    }
+
+    fn from_fn<F: Fn(usize) -> usize>(k: Self::Length, f: F) -> Self {
+        (0..k).map(f).collect()
+    }
+
+    fn item_from_fn<T, F: Fn(usize) -> T>(k: Self::Length, f: F) -> Self::Item<T> {
+        (0..k).map(f).collect()
+    }
+
+    fn len(&self) -> Self::Length {
+        self.len()
+    }
+}
+
+impl PoolIndex for Box<[usize]> {
+    type Item<T> = Vec<T>;
+    type Length = usize;
+
+    fn extract_item<I: Iterator>(&self, pool: &LazyBuffer<I>) -> Self::Item<I::Item>
+    where
+        I::Item: Clone,
+    {
+        pool.get_at(self)
+    }
+
+    fn from_fn<F: Fn(usize) -> usize>(k: Self::Length, f: F) -> Self {
+        (0..k).map(f).collect()
+    }
+
+    fn item_from_fn<T, F: Fn(usize) -> T>(k: Self::Length, f: F) -> Self::Item<T> {
+        (0..k).map(f).collect()
+    }
+
+    fn len(&self) -> Self::Length {
+        self.deref().len()
+    }
+}
+
+impl<const K: usize> PoolIndex for [usize; K] {
+    type Item<T> = [T; K];
+    type Length = ConstUsize<K>;
+
+    fn extract_item<I: Iterator>(&self, pool: &LazyBuffer<I>) -> Self::Item<I::Item>
+    where
+        I::Item: Clone,
+    {
+        pool.get_array(self)
+    }
+
+    fn from_fn<F: Fn(usize) -> usize>(_k: Self::Length, f: F) -> Self {
+        core::array::from_fn(f)
+    }
+
+    fn item_from_fn<T, F: Fn(usize) -> T>(_k: Self::Length, f: F) -> Self::Item<T> {
+        core::array::from_fn(f)
+    }
+
+    fn len(&self) -> Self::Length {
+        ConstUsize
+    }
+}
diff --git a/src/permutations.rs b/src/permutations.rs
@@ -1,11 +1,9 @@
 use alloc::boxed::Box;
-use alloc::vec::Vec;
 use std::fmt;
 use std::iter::FusedIterator;
 
-use super::lazy_buffer::LazyBuffer;
+use super::lazy_buffer::{LazyBuffer, MaybeConstUsize, PoolIndex};
 use crate::size_hint::{self, SizeHint};
-use crate::combinations::{MaybeConstUsize, PoolIndex};
 
 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
 pub struct PermutationsGeneric<I: Iterator, Idx: PoolIndex> {
@@ -18,7 +16,7 @@ pub struct PermutationsGeneric<I: Iterator, Idx: PoolIndex> {
 ///
 /// See [`.permutations()`](crate::Itertools::permutations) for
 /// more information.
-pub type Permutations<I> = PermutationsGeneric<I, Vec<usize>>;
+pub type Permutations<I> = PermutationsGeneric<I, Box<[usize]>>;
 
 impl<I, Idx> Clone for PermutationsGeneric<I, Idx>
 where
@@ -34,13 +32,9 @@ enum PermutationState<Idx: PoolIndex> {
     /// No permutation generated yet.
     Start { k: Idx::Length },
     /// Values from the iterator are not fully loaded yet so `n` is still unknown.
-    Buffered { k: Idx::Length, min_n: usize },
+    Buffered { indices: Idx, min_n: usize },
     /// All values from the iterator are known so `n` is known.
-    Loaded {
-        indices: Box<[usize]>,
-        cycles: Box<[usize]>, // TODO Should be Idx::Item<usize>
-        k: Idx::Length, // TODO Should be inferred from cycles
-    },
+    Loaded { indices: Box<[usize]>, cycles: Idx },
     /// No permutation left to generate.
     End,
 }
@@ -61,10 +55,11 @@ pub fn permutations<I: Iterator>(iter: I, k: usize) -> Permutations<I> {
     }
 }
 
-impl<I, Idx: PoolIndex> Iterator for PermutationsGeneric<I, Idx>
+impl<I, Idx> Iterator for PermutationsGeneric<I, Idx>
 where
     I: Iterator,
     I::Item: Clone,
+    Idx: PoolIndex,
 {
     type Item = Idx::Item<I::Item>;
 
@@ -80,45 +75,44 @@ where
                         *state = PermutationState::End;
                         return None;
                     }
-                    *state = PermutationState::Buffered { k, min_n: k.value() };
+                    *state = PermutationState::Buffered {
+                        indices: Idx::start(k),
+                        min_n: k.value(),
+                    };
                 }
-                Some(Idx::from_fn(k, |i| vals[i].clone()))
+                Some(Idx::start(k).extract_item(vals))
             }
-            PermutationState::Buffered { k, min_n } => {
+            PermutationState::Buffered { indices, min_n } => {
+                let k = indices.len();
                 if vals.get_next() {
-                    // TODO This is ugly. Maybe working on indices is better?
-                    let item = Idx::from_fn(*k, |i| {
-                        vals[if i==k.value()-1 {
-                            *min_n
-                        } else {
-                            i
-                        }].clone()
-                    });
+                    indices.borrow_mut()[k.value() - 1] += 1;
                     *min_n += 1;
-                    Some(item)
+                    Some(indices.extract_item(vals))
                 } else {
                     let n = *min_n;
                     let prev_iteration_count = n - k.value() + 1;
                     let mut indices: Box<[_]> = (0..n).collect();
-                    let mut cycles: Box<[_]> = (n - k.value()..n).rev().collect();
+                    let mut cycles = Idx::from_fn(k, |i| n - 1 - i);
                     // Advance the state to the correct point.
                     for _ in 0..prev_iteration_count {
-                        if advance(&mut indices, &mut cycles) {
+                        if advance(&mut indices, cycles.borrow_mut()) {
                             *state = PermutationState::End;
                             return None;
                         }
                     }
-                    let item = Idx::from_fn(*k, |i| vals[indices[i]].clone());
-                    *state = PermutationState::Loaded { indices, cycles, k:*k };
+                    let item = Idx::item_from_fn(k, |i| vals[indices[i]].clone());
+                    *state = PermutationState::Loaded { indices, cycles };
                     Some(item)
                 }
             }
-            PermutationState::Loaded { indices, cycles, k} => {
-                if advance(indices, cycles) {
+            PermutationState::Loaded { indices, cycles } => {
+                if advance(indices, cycles.borrow_mut()) {
                     *state = PermutationState::End;
                     return None;
                 }
-                Some(Idx::from_fn(*k, |i| vals[indices[i]].clone()))
+                Some(Idx::item_from_fn(cycles.len(), |i| {
+                    vals[indices[i]].clone()
+                }))
             }
             PermutationState::End => None,
         }
@@ -174,19 +168,23 @@ impl<Idx: PoolIndex> PermutationState<Idx> {
         match *self {
             Self::Start { k } if n < k.value() => (0, Some(0)),
             Self::Start { k } => at_start(n, k),
-            Self::Buffered { k, min_n } => {
+            Self::Buffered { ref indices, min_n } => {
+                let k = indices.len();
                 // Same as `Start` minus the previously generated items.
                 size_hint::sub_scalar(at_start(n, k), min_n - k.value() + 1)
             }
             Self::Loaded {
                 ref indices,
                 ref cycles,
-                k: _,
             } => {
-                let count = cycles.iter().enumerate().try_fold(0usize, |acc, (i, &c)| {
-                    acc.checked_mul(indices.len() - i)
-                        .and_then(|count| count.checked_add(c))
-                });
+                let count = cycles
+                    .borrow()
+                    .iter()
+                    .enumerate()
+                    .try_fold(0usize, |acc, (i, &c)| {
+                        acc.checked_mul(indices.len() - i)
+                            .and_then(|count| count.checked_add(c))
+                    });
                 (count.unwrap_or(usize::MAX), count)
             }
             Self::End => (0, Some(0)),
