From 4810cd3681f921cd0057ace89c230d024093a3b2 Mon Sep 17 00:00:00 2001
From: Jorge Aparicio <jorge@japaric.io>
Date: Mon, 6 May 2019 18:56:12 +0200
Subject: [PATCH] add Pool.grow_exact API

---
 src/pool/mod.rs       | 46 ++++++++++++++++++++++++++++++++++++++-----
 src/pool/singleton.rs | 15 ++++++++++++--
 2 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/src/pool/mod.rs b/src/pool/mod.rs
index c9f8a94b..6972c57c 100644
--- a/src/pool/mod.rs
+++ b/src/pool/mod.rs
@@ -154,17 +154,18 @@
 //!
 //! [1]: https://static.docs.arm.com/ddi0403/eb/DDI0403E_B_armv7m_arm.pdf
 
+#[cfg(not(armv6m))]
+use core::{any::TypeId, mem, sync::atomic::Ordering};
 use core::{
     cell::UnsafeCell,
     cmp, fmt,
     hash::{Hash, Hasher},
     marker::PhantomData,
+    mem::MaybeUninit,
     ops::{Deref, DerefMut},
     ptr::{self, NonNull},
     sync::atomic::AtomicPtr,
 };
-#[cfg(not(armv6m))]
-use core::{any::TypeId, mem, sync::atomic::Ordering};
 
 use as_slice::{AsMutSlice, AsSlice};
 
@@ -268,6 +269,22 @@ impl<T> Pool<T> {
         n
     }
 
+    /// Increases the capacity of the pool
+    ///
+    /// Unlike [`Pool.grow`](struct.Pool.html#method.grow) this method fully utilizes the given
+    /// memory block
+    pub fn grow_exact<A>(&self, memory: &'static mut MaybeUninit<A>) -> usize
+    where
+        A: AsMutSlice<Element = Node<T>>,
+    {
+        let nodes = unsafe { (*memory.as_mut_ptr()).as_mut_slice() };
+        let cap = nodes.len();
+        for p in nodes {
+            self.push(NonNull::from(p))
+        }
+        cap
+    }
+
     fn pop(&self) -> Option<NonNull<Node<T>>> {
         // NOTE `Ordering`s come from crossbeam's (v0.6.0) `TreiberStack`
 
@@ -314,7 +331,9 @@ impl<T> Pool<T> {
     }
 }
 
-struct Node<T> {
+/// Unfortunate implementation detail required to use the
+/// [`Pool.grow_exact`](struct.Pool.html#method.grow_exact) method
+pub struct Node<T> {
     data: UnsafeCell<T>,
     next: *mut Node<T>,
 }
@@ -445,11 +464,11 @@ where
 #[cfg(test)]
 mod tests {
     use core::{
-        mem,
+        mem::{self, MaybeUninit},
         sync::atomic::{AtomicUsize, Ordering},
     };
 
-    use super::Pool;
+    use super::{Node, Pool};
 
     #[test]
     fn grow() {
@@ -466,6 +485,23 @@ mod tests {
         }
     }
 
+    #[test]
+    fn grow_exact() {
+        const SZ: usize = 8;
+        static mut MEMORY: MaybeUninit<[Node<[u8; 128]>; SZ]> = MaybeUninit::uninit();
+
+        static POOL: Pool<[u8; 128]> = Pool::new();
+
+        unsafe {
+            POOL.grow_exact(&mut MEMORY);
+        }
+
+        for _ in 0..SZ {
+            assert!(POOL.alloc().is_some());
+        }
+        assert!(POOL.alloc().is_none());
+    }
+
     #[test]
     fn sanity() {
         static mut MEMORY: [u8; 31] = [0; 31];
diff --git a/src/pool/singleton.rs b/src/pool/singleton.rs
index 43ef6827..671f7ade 100644
--- a/src/pool/singleton.rs
+++ b/src/pool/singleton.rs
@@ -5,14 +5,14 @@ use core::{
     cmp, fmt,
     hash::{Hash, Hasher},
     marker::PhantomData,
-    mem,
+    mem::{self, MaybeUninit},
     ops::{Deref, DerefMut},
     ptr,
 };
 
 use as_slice::{AsMutSlice, AsSlice};
 
-use super::{Init, Uninit};
+use super::{Init, Node, Uninit};
 
 /// Instantiates a pool as a global singleton
 #[cfg(any(armv7m, armv7r, test))]
@@ -65,6 +65,17 @@ pub trait Pool {
     fn grow(memory: &'static mut [u8]) -> usize {
         Self::ptr().grow(memory)
     }
+
+    /// Increases the capacity of the pool
+    ///
+    /// Unlike [`Pool.grow`](trait.Pool.html#method.grow_exact) this method fully utilizes the given
+    /// memory block
+    fn grow_exact<A>(memory: &'static mut MaybeUninit<A>) -> usize
+    where
+        A: AsMutSlice<Element = Node<Self::Data>>,
+    {
+        Self::ptr().grow_exact(memory)
+    }
 }
 
 /// A memory block that belongs to the global memory pool, `POOL`