From 2680f811cfeb1e33ca9dc6cb5d0c2ebd36a9be63 Mon Sep 17 00:00:00 2001
From: Joshua Simmons <josh@nega.tv>
Date: Thu, 3 Nov 2022 22:18:31 +0100
Subject: [PATCH] Expand arena functionality

Add `alloc_slice_fill_with`
Add `alloc_slice_fill_copy`
Add `alloc_slice_fill_clone`
Add `alloc_slice_fill_iter`
---
 narcissus-core/src/arena.rs | 96 +++++++++++++++++++++++++++++++++++++
 1 file changed, 96 insertions(+)

diff --git a/narcissus-core/src/arena.rs b/narcissus-core/src/arena.rs
index 7b678a0..d25826e 100644
--- a/narcissus-core/src/arena.rs
+++ b/narcissus-core/src/arena.rs
@@ -389,6 +389,54 @@ impl Arena {
         // Safety: We've just copied this string from a valid `&str`, so it must be valid too.
         unsafe { std::str::from_utf8_unchecked_mut(str) }
     }
+
+    #[inline(always)]
+    #[allow(clippy::mut_from_ref)]
+    pub fn alloc_slice_fill_with<T, F>(&self, len: usize, mut f: F) -> &mut [T]
+    where
+        F: FnMut(usize) -> T,
+    {
+        let layout = Layout::array::<T>(len).unwrap_or_else(|_| oom());
+        let dst = self.alloc_layout(layout).cast::<T>();
+
+        // Safety: We allocated an array of len elements of T above.
+        unsafe {
+            for i in 0..len {
+                std::ptr::write(dst.as_ptr().add(i), f(i))
+            }
+
+            std::slice::from_raw_parts_mut(dst.as_ptr(), len)
+        }
+    }
+
+    #[inline(always)]
+    #[allow(clippy::mut_from_ref)]
+    pub fn alloc_slice_fill_copy<T>(&self, len: usize, value: T) -> &mut [T]
+    where
+        T: Copy,
+    {
+        self.alloc_slice_fill_with(len, |_| value)
+    }
+
+    #[inline(always)]
+    #[allow(clippy::mut_from_ref)]
+    pub fn alloc_slice_fill_clone<T>(&self, len: usize, value: T) -> &mut [T]
+    where
+        T: Clone,
+    {
+        self.alloc_slice_fill_with(len, |_| value.clone())
+    }
+
+    #[inline(always)]
+    #[allow(clippy::mut_from_ref)]
+    pub fn alloc_slice_fill_iter<T, I>(&self, iter: I) -> &mut [T]
+    where
+        I: IntoIterator<Item = T>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        let mut iter = iter.into_iter();
+        self.alloc_slice_fill_with(iter.len(), |_| iter.next().unwrap())
+    }
 }
 
 impl Default for Arena {
@@ -621,6 +669,54 @@ impl<const STACK_CAP: usize> HybridArena<STACK_CAP> {
         // Safety: We've just copied this string from a valid `&str`, so it must be valid too.
         unsafe { std::str::from_utf8_unchecked_mut(str) }
     }
+
+    #[inline(always)]
+    #[allow(clippy::mut_from_ref)]
+    pub fn alloc_slice_fill_with<T, F>(&self, len: usize, mut f: F) -> &mut [T]
+    where
+        F: FnMut(usize) -> T,
+    {
+        let layout = Layout::array::<T>(len).unwrap_or_else(|_| oom());
+        let dst = self.alloc_layout(layout).cast::<T>();
+
+        // Safety: We allocated an array of len elements of T above.
+        unsafe {
+            for i in 0..len {
+                std::ptr::write(dst.as_ptr().add(i), f(i))
+            }
+
+            std::slice::from_raw_parts_mut(dst.as_ptr(), len)
+        }
+    }
+
+    #[inline(always)]
+    #[allow(clippy::mut_from_ref)]
+    pub fn alloc_slice_fill_copy<T>(&self, len: usize, value: T) -> &mut [T]
+    where
+        T: Copy,
+    {
+        self.alloc_slice_fill_with(len, |_| value)
+    }
+
+    #[inline(always)]
+    #[allow(clippy::mut_from_ref)]
+    pub fn alloc_slice_fill_clone<T>(&self, len: usize, value: T) -> &mut [T]
+    where
+        T: Clone,
+    {
+        self.alloc_slice_fill_with(len, |_| value.clone())
+    }
+
+    #[inline(always)]
+    #[allow(clippy::mut_from_ref)]
+    pub fn alloc_slice_fill_iter<T, I>(&self, iter: I) -> &mut [T]
+    where
+        I: IntoIterator<Item = T>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        let mut iter = iter.into_iter();
+        self.alloc_slice_fill_with(iter.len(), |_| iter.next().unwrap())
+    }
 }
 
 impl<const STACK_CAP: usize> Default for HybridArena<STACK_CAP> {
-- 
2.49.0