Skip to content

Commit ddd900c

Browse files
committed
make event iterator exact sized
1 parent c3ba29c commit ddd900c

File tree

1 file changed

+88
-31
lines changed

1 file changed

+88
-31
lines changed

crates/bevy_ecs/src/event.rs

Lines changed: 88 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -212,21 +212,25 @@ impl<T> Default for ManualEventReader<T> {
212212
#[allow(clippy::len_without_is_empty)] // Check fails since the is_empty implementation has a signature other than `(&self) -> bool`
213213
impl<T: Event> ManualEventReader<T> {
214214
/// See [`EventReader::iter`]
215-
pub fn iter<'a>(&'a mut self, events: &'a Events<T>) -> impl DoubleEndedIterator<Item = &'a T> {
215+
pub fn iter<'a>(
216+
&'a mut self,
217+
events: &'a Events<T>,
218+
) -> impl DoubleEndedIterator<Item = &'a T> + ExactSizeIterator<Item = &'a T> {
216219
internal_event_reader(&mut self.last_event_count, events).map(|(e, _)| e)
217220
}
218221

219222
/// See [`EventReader::iter_with_id`]
220223
pub fn iter_with_id<'a>(
221224
&'a mut self,
222225
events: &'a Events<T>,
223-
) -> impl DoubleEndedIterator<Item = (&'a T, EventId<T>)> {
226+
) -> impl DoubleEndedIterator<Item = (&'a T, EventId<T>)>
227+
+ ExactSizeIterator<Item = (&'a T, EventId<T>)> {
224228
internal_event_reader(&mut self.last_event_count, events)
225229
}
226230

227231
/// See [`EventReader::len`]
228232
pub fn len(&self, events: &Events<T>) -> usize {
229-
events.event_reader_len(self.last_event_count)
233+
internal_event_reader(&mut self.last_event_count.clone(), events).len()
230234
}
231235

232236
/// See [`EventReader::is_empty`]
@@ -240,7 +244,8 @@ impl<T: Event> ManualEventReader<T> {
240244
fn internal_event_reader<'a, T>(
241245
last_event_count: &'a mut usize,
242246
events: &'a Events<T>,
243-
) -> impl DoubleEndedIterator<Item = (&'a T, EventId<T>)> {
247+
) -> impl DoubleEndedIterator<Item = (&'a T, EventId<T>)> + ExactSizeIterator<Item = (&'a T, EventId<T>)>
248+
{
244249
// if the reader has seen some of the events in a buffer, find the proper index offset.
245250
// otherwise read all events in the buffer
246251
let a_index = if *last_event_count > events.a_start_event_count {
@@ -263,19 +268,77 @@ fn internal_event_reader<'a, T>(
263268
};
264269
iterator
265270
.map(map_instance_event_with_id)
271+
.with_exact_size(unread_count)
266272
.inspect(move |(_, id)| *last_event_count = (id.id + 1).max(*last_event_count))
267273
}
268274

275+
trait IteratorExt {
276+
fn with_exact_size(self, len: usize) -> ExactSize<Self>
277+
where
278+
Self: Sized,
279+
{
280+
ExactSize::new(self, len)
281+
}
282+
}
283+
impl<I> IteratorExt for I where I: Iterator {}
284+
285+
#[must_use = "iterators are lazy and do nothing unless consumed"]
286+
#[derive(Clone)]
287+
struct ExactSize<I> {
288+
iter: I,
289+
len: usize,
290+
}
291+
impl<I> ExactSize<I> {
292+
fn new(iter: I, len: usize) -> Self {
293+
ExactSize { iter, len }
294+
}
295+
}
296+
297+
impl<I: Iterator> Iterator for ExactSize<I> {
298+
type Item = I::Item;
299+
300+
#[inline]
301+
fn next(&mut self) -> Option<I::Item> {
302+
self.iter.next().map(|e| {
303+
self.len -= 1;
304+
e
305+
})
306+
}
307+
308+
#[inline]
309+
fn size_hint(&self) -> (usize, Option<usize>) {
310+
(self.len, Some(self.len))
311+
}
312+
}
313+
314+
impl<I: DoubleEndedIterator> DoubleEndedIterator for ExactSize<I> {
315+
#[inline]
316+
fn next_back(&mut self) -> Option<I::Item> {
317+
self.iter.next_back().map(|e| {
318+
self.len -= 1;
319+
e
320+
})
321+
}
322+
}
323+
impl<I: Iterator> ExactSizeIterator for ExactSize<I> {
324+
fn len(&self) -> usize {
325+
self.len
326+
}
327+
}
328+
269329
impl<'w, 's, T: Event> EventReader<'w, 's, T> {
270330
/// Iterates over the events this [`EventReader`] has not seen yet. This updates the
271331
/// [`EventReader`]'s event counter, which means subsequent event reads will not include events
272332
/// that happened before now.
273-
pub fn iter(&mut self) -> impl DoubleEndedIterator<Item = &T> {
333+
pub fn iter(&mut self) -> impl DoubleEndedIterator<Item = &T> + ExactSizeIterator<Item = &T> {
274334
self.iter_with_id().map(|(event, _id)| event)
275335
}
276336

277337
/// Like [`iter`](Self::iter), except also returning the [`EventId`] of the events.
278-
pub fn iter_with_id(&mut self) -> impl DoubleEndedIterator<Item = (&T, EventId<T>)> {
338+
pub fn iter_with_id(
339+
&mut self,
340+
) -> impl DoubleEndedIterator<Item = (&T, EventId<T>)> + ExactSizeIterator<Item = (&T, EventId<T>)>
341+
{
279342
internal_event_reader(&mut self.last_event_count.0, &self.events).map(|(event, id)| {
280343
trace!("EventReader::iter() -> {}", id);
281344
(event, id)
@@ -284,7 +347,7 @@ impl<'w, 's, T: Event> EventReader<'w, 's, T> {
284347

285348
/// Determines the number of events available to be read from this [`EventReader`] without consuming any.
286349
pub fn len(&self) -> usize {
287-
self.events.event_reader_len(self.last_event_count.0)
350+
internal_event_reader(&mut self.last_event_count.0.clone(), &self.events).len()
288351
}
289352

290353
/// Determines if are any events available to be read without consuming any.
@@ -405,35 +468,14 @@ impl<T: Event> Events<T> {
405468
/// between the last `update()` call and your call to `iter_current_update_events`.
406469
/// If events happen outside that window, they will not be handled. For example, any events that
407470
/// happen after this call and before the next `update()` call will be dropped.
408-
pub fn iter_current_update_events(&self) -> impl DoubleEndedIterator<Item = &T> {
471+
pub fn iter_current_update_events(
472+
&self,
473+
) -> impl DoubleEndedIterator<Item = &T> + ExactSizeIterator<Item = &T> {
409474
match self.state {
410475
State::A => self.events_a.iter().map(map_instance_event),
411476
State::B => self.events_b.iter().map(map_instance_event),
412477
}
413478
}
414-
415-
/// Determines how many events are in the reader after the given `last_event_count` parameter
416-
fn event_reader_len(&self, last_event_count: usize) -> usize {
417-
let a_count = if last_event_count <= self.a_start_event_count {
418-
self.events_a.len()
419-
} else {
420-
self.events_a
421-
.len()
422-
.checked_sub(last_event_count - self.a_start_event_count)
423-
.unwrap_or_default()
424-
};
425-
426-
let b_count = if last_event_count <= self.b_start_event_count {
427-
self.events_b.len()
428-
} else {
429-
self.events_b
430-
.len()
431-
.checked_sub(last_event_count - self.b_start_event_count)
432-
.unwrap_or_default()
433-
};
434-
435-
a_count + b_count
436-
}
437479
}
438480

439481
impl<T> std::iter::Extend<T> for Events<T> {
@@ -651,6 +693,21 @@ mod tests {
651693
assert!(!events.get_reader().is_empty(&events));
652694
}
653695

696+
#[test]
697+
fn test_event_iter_len_updated() {
698+
let mut events = Events::<TestEvent>::default();
699+
events.send(TestEvent { i: 0 });
700+
events.send(TestEvent { i: 1 });
701+
events.send(TestEvent { i: 2 });
702+
let mut reader = events.get_reader();
703+
let mut iter = reader.iter(&events);
704+
assert_eq!(iter.len(), 3);
705+
iter.next();
706+
assert_eq!(iter.len(), 2);
707+
iter.next_back();
708+
assert_eq!(iter.len(), 1);
709+
}
710+
654711
#[test]
655712
fn test_event_reader_len_current() {
656713
let mut events = Events::<TestEvent>::default();

0 commit comments

Comments
 (0)