mz_repr/row/
iter.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Defines a "lending iterator" for [`Row`]
11
12use std::fmt::Debug;
13use std::sync::Arc;
14
15use crate::row::{Row, RowRef};
16
17/// An iterator that can borrow from `self` and yield [`RowRef`]s.
18///
19/// This trait is a "lending iterator" for [`Row`]s, in other words, an iterator that borrows from
20/// self (e.g. an underlying memory buffer) to return a [`RowRef`]. The [`std::iter::Iterator`]
21/// trait does not currently support this pattern because there is no way to name the lifetime of
22/// the borrow on its associated `Item` type. Generic Associated Types (GATs) would allow this but
23/// so far no new trait has been introduced with this API.
24///
25/// There are a few open source crates that provide a trait:
26///
27/// * [`streaming_iterator`](https://docs.rs/streaming-iterator/latest/streaming_iterator/)
28/// * [`lending-iterator`](https://docs.rs/lending-iterator/latest/lending_iterator/)
29///
30/// Neither have an `IntoLendingIterator` trait that is useful for our interface, nor do they work
31/// well with trait objects.
32pub trait RowIterator: Debug {
33    /// Returns the next [`RowRef`] advancing the iterator.
34    fn next(&mut self) -> Option<&RowRef>;
35
36    /// Returns the next [`RowRef`] without advancing the iterator.
37    fn peek(&mut self) -> Option<&RowRef>;
38
39    /// The total number of [`Row`]s this iterator could ever yield.
40    ///
41    /// Note: it _does not_ return the number of rows _remaining_, in otherwords calling `.next()`
42    /// will not change the value returned from this method.
43    fn count(&self) -> usize;
44
45    /// Returns a clone of `self` as a `Box<dyn RowIterator>`.
46    fn box_clone(&self) -> Box<dyn RowIterator>;
47
48    /// Maps the returned [`RowRef`]s from this [`RowIterator`].
49    fn map<T, F>(self, f: F) -> MappedRowIterator<Self, F>
50    where
51        Self: Sized,
52        F: FnMut(&RowRef) -> T,
53    {
54        MappedRowIterator {
55            inner: self,
56            func: f,
57        }
58    }
59}
60
61impl<I: RowIterator + ?Sized> RowIterator for Box<I> {
62    fn next(&mut self) -> Option<&RowRef> {
63        (**self).next()
64    }
65
66    fn peek(&mut self) -> Option<&RowRef> {
67        (**self).peek()
68    }
69
70    fn count(&self) -> usize {
71        (**self).count()
72    }
73
74    fn box_clone(&self) -> Box<dyn RowIterator> {
75        (**self).box_clone()
76    }
77}
78
79impl<I: RowIterator + ?Sized> RowIterator for &mut I {
80    fn next(&mut self) -> Option<&RowRef> {
81        (**self).next()
82    }
83
84    fn peek(&mut self) -> Option<&RowRef> {
85        (**self).peek()
86    }
87
88    fn count(&self) -> usize {
89        (**self).count()
90    }
91
92    fn box_clone(&self) -> Box<dyn RowIterator> {
93        (**self).box_clone()
94    }
95}
96
97#[derive(Debug)]
98pub struct MappedRowIterator<I: RowIterator, F> {
99    inner: I,
100    func: F,
101}
102
103impl<T, F, I: RowIterator> Iterator for MappedRowIterator<I, F>
104where
105    F: FnMut(&RowRef) -> T,
106{
107    type Item = T;
108
109    fn next(&mut self) -> Option<Self::Item> {
110        let row_ref = self.inner.next()?;
111        Some((self.func)(row_ref))
112    }
113}
114
115/// Convert a type into a [`RowIterator`].
116pub trait IntoRowIterator {
117    type Iter: RowIterator;
118    fn into_row_iter(self) -> Self::Iter;
119}
120
121impl<T: RowIterator> IntoRowIterator for T {
122    type Iter = Self;
123    fn into_row_iter(self) -> Self::Iter {
124        self
125    }
126}
127
128/// A [`RowIterator`] for a single [`Row`].
129#[derive(Debug, Clone)]
130pub struct SingleRowIter {
131    row: Arc<Row>,
132    finished: bool,
133}
134
135impl RowIterator for SingleRowIter {
136    fn next(&mut self) -> Option<&RowRef> {
137        if self.finished {
138            None
139        } else {
140            self.finished = true;
141            Some(self.row.as_ref())
142        }
143    }
144
145    fn peek(&mut self) -> Option<&RowRef> {
146        if self.finished {
147            None
148        } else {
149            Some(self.row.as_ref())
150        }
151    }
152
153    fn count(&self) -> usize {
154        1
155    }
156
157    fn box_clone(&self) -> Box<dyn RowIterator> {
158        Box::new(self.clone())
159    }
160}
161
162impl IntoRowIterator for Row {
163    type Iter = SingleRowIter;
164
165    fn into_row_iter(self) -> Self::Iter {
166        SingleRowIter {
167            row: Arc::new(self),
168            finished: false,
169        }
170    }
171}
172
173/// A [`RowIterator`] for a [`Vec`] of [`Row`]s.
174#[derive(Debug, Clone)]
175pub struct VecRowIter {
176    rows: Arc<[Row]>,
177    index: usize,
178}
179
180impl RowIterator for VecRowIter {
181    fn next(&mut self) -> Option<&RowRef> {
182        let row = self.rows.get(self.index).map(|r| r.as_ref())?;
183        self.index = self.index.saturating_add(1);
184
185        Some(row)
186    }
187
188    fn peek(&mut self) -> Option<&RowRef> {
189        self.rows.get(self.index).map(|r| r.as_ref())
190    }
191
192    fn count(&self) -> usize {
193        self.rows.len()
194    }
195
196    fn box_clone(&self) -> Box<dyn RowIterator> {
197        Box::new(self.clone())
198    }
199}
200
201impl IntoRowIterator for Vec<Row> {
202    type Iter = VecRowIter;
203
204    fn into_row_iter(self) -> Self::Iter {
205        VecRowIter {
206            rows: self.into(),
207            index: 0,
208        }
209    }
210}