malachite_base/iterators/
iterator_cache.rs

1// Copyright © 2025 Mikhail Hogrefe
2//
3// This file is part of Malachite.
4//
5// Malachite is free software: you can redistribute it and/or modify it under the terms of the GNU
6// Lesser General Public License (LGPL) as published by the Free Software Foundation; either version
7// 3 of the License, or (at your option) any later version. See <https://www.gnu.org/licenses/>.
8
9use alloc::vec::Vec;
10/// Remembers values produced by an iterator.
11///
12/// After wrapping an iterator with an `IteratorCache`, you can retrieve a reference to the $n$th
13/// element of an iterator, and then retrieve a reference to the $m$th element in constant time for
14/// any $m \leq n$ (not counting the time it took to first get the $n$th element).
15#[derive(Clone, Debug)]
16pub struct IteratorCache<I: Iterator> {
17    xs: I,
18    cache: Vec<I::Item>,
19    done: bool,
20}
21
22impl<I: Iterator> IteratorCache<I> {
23    /// Creates a new `IteratorCache`.
24    ///
25    /// This function does not allocate any memory or advance the iterator.
26    ///
27    /// # Worst-case complexity
28    /// Constant time and additional memory.
29    ///
30    /// # Examples
31    /// ```
32    /// use malachite_base::iterators::iterator_cache::IteratorCache;
33    ///
34    /// IteratorCache::new([1, 2, 3].iter());
35    /// ```
36    pub const fn new(xs: I) -> IteratorCache<I> {
37        IteratorCache {
38            xs,
39            cache: Vec::new(),
40            done: false,
41        }
42    }
43
44    /// Retrieves the $n$th element of an iterator. Indexing starts at 0.
45    ///
46    /// If the index is higher than any other previously-requested index, the iterator is advanced
47    /// to that index, or until it runs out. If the iterator has previously been advanced past the
48    /// index, the requested element is returned from the cache in constant time. If the iterator is
49    /// too short to have an element at the index, `None` is returned.
50    ///
51    /// If you know that the element is present, and want to only take an immutable reference to
52    /// `self`, consider using [`assert_get`](Self::assert_get) instead.
53    ///
54    /// # Worst-case complexity
55    /// $T(n) = O(n)$
56    ///
57    /// $M(n) = O(n)$
58    ///
59    /// where $T$ is time, $M$ is additional memory, and $n$ is 1 if `get` has previously been
60    /// called with an index at least this large, or `index` otherwise.
61    ///
62    /// # Examples
63    /// ```
64    /// use malachite_base::iterators::iterator_cache::IteratorCache;
65    ///
66    /// let mut xs = IteratorCache::new([1, 2, 3].iter().cloned());
67    /// assert_eq!(xs.get(1), Some(&2));
68    /// assert_eq!(xs.get(0), Some(&1));
69    /// assert_eq!(xs.get(3), None);
70    /// assert_eq!(xs.get(2), Some(&3));
71    /// ```
72    pub fn get(&mut self, index: usize) -> Option<&I::Item> {
73        if !self.done && index >= self.cache.len() {
74            self.cache
75                .extend((&mut self.xs).take(index - self.cache.len() + 1));
76            if index >= self.cache.len() {
77                self.done = true;
78                return None;
79            }
80        }
81        self.cache.get(index)
82    }
83
84    /// Retrieves the $n$th element of an iterator, while asserting that the iterator has already
85    /// reached that element. Indexing starts at 0.
86    ///
87    /// If the iterator has not advanced that far, or if it has fewer than $n + 1$ elements, this
88    /// function panics.
89    ///
90    /// The purpose of this function is to allow the caller to get an element immutably, assuming
91    /// that the caller knows that the element is present. The [`get`](Self::get) function has to
92    /// take a mutable reference.
93    ///
94    /// # Worst-case complexity
95    /// Constant time and additional memory.
96    ///
97    /// # Examples
98    /// ```
99    /// use malachite_base::iterators::iterator_cache::IteratorCache;
100    ///
101    /// let mut xs = IteratorCache::new([1, 2, 3].iter().cloned());
102    /// // Force the iterator to iterate to completion
103    /// xs.get(3);
104    /// assert_eq!(xs.assert_get(1), &2);
105    /// assert_eq!(xs.assert_get(0), &1);
106    /// assert_eq!(xs.assert_get(2), &3);
107    /// ```
108    pub fn assert_get(&self, index: usize) -> &I::Item {
109        self.cache.get(index).unwrap()
110    }
111
112    /// Returns the total number of elements in the iterator, if the iterator has been completely
113    /// consumed.
114    ///
115    /// If the iterator has not been completely consumed yet, returns `None`.
116    ///
117    /// # Worst-case complexity
118    /// Constant time and additional memory.
119    ///
120    /// # Examples
121    /// ```
122    /// use malachite_base::iterators::iterator_cache::IteratorCache;
123    ///
124    /// let mut xs = IteratorCache::new([1, 2, 3].iter().cloned());
125    /// assert_eq!(xs.known_len(), None);
126    /// assert_eq!(xs.get(1), Some(&2));
127    /// assert_eq!(xs.known_len(), None);
128    /// assert_eq!(xs.get(0), Some(&1));
129    /// assert_eq!(xs.get(3), None);
130    /// assert_eq!(xs.known_len(), Some(3));
131    /// assert_eq!(xs.get(2), Some(&3));
132    /// ```
133    pub fn known_len(&self) -> Option<usize> {
134        if self.done {
135            Some(self.cache.len())
136        } else {
137            None
138        }
139    }
140}
malachite_base/iterators/iterator_cache.rs

malachite_base/iterators/
iterator_cache.rs
