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

malachite_base/iterators/
iterator_cache.rs
