Skip to content

Commit

Permalink
Merge pull request #78 from michaelwoerister/simple2
Browse files Browse the repository at this point in the history
Avoid allocation for small length in PollState tracking
  • Loading branch information
yoshuawuyts authored Nov 15, 2022
2 parents ff31355 + ed1fb5c commit 7f398ba
Show file tree
Hide file tree
Showing 3 changed files with 97 additions and 6 deletions.
11 changes: 6 additions & 5 deletions src/future/join/vec.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use super::Join as JoinTrait;
use crate::utils::{iter_pin_mut_vec, PollState};
use crate::utils::{iter_pin_mut_vec, PollState, PollStates};

use core::fmt;
use core::future::{Future, IntoFuture};
Expand All @@ -26,7 +26,7 @@ where
consumed: bool,
pending: usize,
items: Vec<MaybeUninit<<Fut as Future>::Output>>,
state: Vec<PollState>,
state: PollStates,
#[pin]
futures: Vec<Fut>,
}
Expand All @@ -42,7 +42,7 @@ where
items: std::iter::repeat_with(MaybeUninit::uninit)
.take(futures.len())
.collect(),
state: vec![PollState::default(); futures.len()],
state: PollStates::new(futures.len()),
futures,
}
}
Expand Down Expand Up @@ -86,11 +86,12 @@ where

// Poll all futures
let futures = this.futures.as_mut();
let states = &mut this.state[..];
for (i, fut) in iter_pin_mut_vec(futures).enumerate() {
if this.state[i].is_pending() {
if states[i].is_pending() {
if let Poll::Ready(value) = fut.poll(cx) {
this.items[i] = MaybeUninit::new(value);
this.state[i] = PollState::Done;
states[i] = PollState::Done;
*this.pending -= 1;
}
}
Expand Down
2 changes: 1 addition & 1 deletion src/utils/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ mod wakers;
pub(crate) use fuse::Fuse;
pub(crate) use maybe_done::MaybeDone;
pub(crate) use pin::{get_pin_mut, get_pin_mut_from_vec, iter_pin_mut, iter_pin_mut_vec};
pub(crate) use poll_state::PollState;
pub(crate) use poll_state::{PollState, PollStates};
pub(crate) use rng::RandomGenerator;
pub(crate) use tuple::{gen_conditions, permutations, tuple_len};
pub(crate) use wakers::{InlineWaker, Readiness, WakerList};
Expand Down
90 changes: 90 additions & 0 deletions src/utils/poll_state.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
use std::ops::{Deref, DerefMut};

/// Enumerate the current poll state.
#[derive(Debug, Clone, Copy, Default)]
#[repr(u8)]
pub(crate) enum PollState {
/// Polling the underlying future.
#[default]
Expand Down Expand Up @@ -37,3 +40,90 @@ impl PollState {
matches!(self, Self::Consumed)
}
}

/// The maximum number of entries that `PollStates` can store without
/// dynamic memory allocation.
///
/// The `Boxed` variant is the minimum size the data structure can have.
/// It consists of a boxed slice (=2 usizes) and space for the enum
/// tag (another usize because of padding), so 3 usizes.
/// The inline variant then consists of `3 * size_of(usize) - 2` entries.
/// Each entry is a byte and we subtract one byte for a length field,
/// and another byte for the enum tag.
///
/// ```txt
/// Boxed
/// vvvvv
/// tag
/// | <-------padding----> <--- Box<[T]>::len ---> <--- Box<[T]>::ptr --->
/// 00 01 02 03 04 05 06 07 08 09 10 11 12 13 14 15 16 17 18 19 20 21 22 23 <bytes
/// | | <------------------- entries ----------------------------------->
/// tag |
/// len ^^^^^
/// Inline
/// ```
const MAX_INLINE_ENTRIES: usize = std::mem::size_of::<usize>() * 3 - 2;

pub(crate) enum PollStates {
Inline(u8, [PollState; MAX_INLINE_ENTRIES]),
Boxed(Box<[PollState]>),
}

impl PollStates {
pub(crate) fn new(len: usize) -> Self {
assert!(MAX_INLINE_ENTRIES <= u8::MAX as usize);

if len <= MAX_INLINE_ENTRIES {
Self::Inline(len as u8, Default::default())
} else {
// Make sure that we don't reallocate the vec's memory
// during `Vec::into_boxed_slice()`.
let mut states = Vec::new();
debug_assert_eq!(states.capacity(), 0);
states.reserve_exact(len);
debug_assert_eq!(states.capacity(), len);
states.resize(len, PollState::default());
debug_assert_eq!(states.capacity(), len);
Self::Boxed(states.into_boxed_slice())
}
}
}

impl Deref for PollStates {
type Target = [PollState];

fn deref(&self) -> &Self::Target {
match self {
PollStates::Inline(len, states) => &states[..*len as usize],
Self::Boxed(states) => &states[..],
}
}
}

impl DerefMut for PollStates {
fn deref_mut(&mut self) -> &mut Self::Target {
match self {
PollStates::Inline(len, states) => &mut states[..*len as usize],
Self::Boxed(states) => &mut states[..],
}
}
}

#[cfg(test)]
mod tests {
use super::{PollStates, MAX_INLINE_ENTRIES};

#[test]
fn type_size() {
assert_eq!(
std::mem::size_of::<PollStates>(),
std::mem::size_of::<usize>() * 3
);
}

#[test]
fn boxed_does_not_allocate_twice() {
// Make sure the debug_assertions in PollStates::new() don't fail.
let _ = PollStates::new(MAX_INLINE_ENTRIES + 10);
}
}

0 comments on commit 7f398ba

Please sign in to comment.