diff --git a/src/collections/inner_group.rs b/src/collections/inner_group.rs index 136e008..554b66b 100644 --- a/src/collections/inner_group.rs +++ b/src/collections/inner_group.rs @@ -11,7 +11,9 @@ use smallvec::{smallvec, SmallVec}; use crate::utils::{PollVec, WakerVec}; -const GROUP_GROWTH_FACTOR: usize = 2; +const fn grow_group_capacity(cap: usize) -> usize { + cap * 2 + 1 +} #[pin_project::pin_project] pub struct InnerGroup { @@ -56,7 +58,7 @@ impl InnerGroup { pub fn insert(&mut self, item: A) -> Key { if !self.has_capacity() { - self.resize((self.cap + 1) * GROUP_GROWTH_FACTOR); + self.reserve(grow_group_capacity(self.cap)); } let index = self.items.insert(item); @@ -75,7 +77,7 @@ impl InnerGroup { if !self.has_capacity() { let r = unsafe { &mut self.as_mut().get_unchecked_mut() }; - r.resize((r.cap + 1) * GROUP_GROWTH_FACTOR); + r.reserve(grow_group_capacity(r.cap)); } let mut this = self.project(); @@ -102,15 +104,17 @@ impl InnerGroup { Some(item) } - // todo: rename to reserve - pub fn resize(&mut self, cap: usize) { - if self.len + cap < self.cap { + /// Reserve `additional` capacity for new items + /// Does nothing if the capacity is already sufficient + pub fn reserve(&mut self, additional: usize) { + if self.len + additional < self.cap { return; } - self.wakers.resize(cap); - self.states.resize(cap); - self.items.reserve_exact(cap); - self.cap = cap; + let new_cap = self.cap + additional; + self.wakers.resize(new_cap); + self.states.resize(new_cap); + self.items.reserve_exact(new_cap); + self.cap = new_cap; } pub fn set_top_waker(&mut self, waker: &Waker) { diff --git a/src/future/future_group.rs b/src/future/future_group.rs index 415faa0..5f3c46c 100644 --- a/src/future/future_group.rs +++ b/src/future/future_group.rs @@ -186,6 +186,25 @@ impl FutureGroup { pub fn contains_key(&mut self, key: Key) -> bool { self.inner.contains_key(key) } + + /// Reserves capacity for `additional` more futures to be inserted. + /// Does nothing if the capacity is already sufficient. + /// + /// # Example + /// + /// ```rust + /// use futures_concurrency::future::FutureGroup; + /// use std::future; + /// # futures_lite::future::block_on(async { + /// let mut group = FutureGroup::with_capacity(0); + /// assert_eq!(group.capacity(), 0); + /// group.reserve(10); + /// assert_eq!(group.capacity(), 10); + /// # }) + /// ``` + pub fn reserve(&mut self, additional: usize) { + self.inner.reserve(additional); + } } impl Default for FutureGroup { diff --git a/src/stream/stream_group.rs b/src/stream/stream_group.rs index 628f768..58136e1 100644 --- a/src/stream/stream_group.rs +++ b/src/stream/stream_group.rs @@ -183,6 +183,25 @@ impl StreamGroup { pub fn contains_key(&mut self, key: Key) -> bool { self.inner.contains_key(key) } + + /// Reserves capacity for `additional` more streams to be inserted. + /// Does nothing if the capacity is already sufficient. + /// + /// # Example + /// + /// ```rust + /// use futures_concurrency::future::StreamGroup; + /// use std::future; + /// # futures_lite::future::block_on(async { + /// let mut group = StreamGroup::with_capacity(0); + /// assert_eq!(group.capacity(), 0); + /// group.reserve(10); + /// assert_eq!(group.capacity(), 10); + /// # }) + /// ``` + pub fn reserve(&mut self, additional: usize) { + self.inner.reserve(additional); + } } impl StreamGroup {