From d2eb16575996ac6fb66d755f6388aba7ac3879ce Mon Sep 17 00:00:00 2001
From: John <j@soft.fish>
Date: Sat, 4 May 2024 22:17:34 -0500
Subject: [PATCH] cl-arena: Add iterator allocation to TypedArena

This is copied almost verbatim from rustc-arena https://github.com/rust-lang/rust/blob/e82c861d7e5ecd766cb0dab0bf622445dec999dc/compiler/rustc_arena/src/lib.rs#L203

TODO: Add unit tests, run unit tests in Miri
---
 compiler/cl-arena/src/lib.rs | 51 ++++++++++++++++++++++++++++++++++--
 1 file changed, 49 insertions(+), 2 deletions(-)

diff --git a/compiler/cl-arena/src/lib.rs b/compiler/cl-arena/src/lib.rs
index 5bca8e0..1d19c80 100644
--- a/compiler/cl-arena/src/lib.rs
+++ b/compiler/cl-arena/src/lib.rs
@@ -77,12 +77,14 @@ mod chunk {
 
 pub mod typed_arena {
     //! A [TypedArena] can hold many instances of a single type, and will properly [Drop] them.
+    #![allow(clippy::mut_from_ref)]
+
     use crate::{chunk::ArenaChunk, constants::*};
     use alloc::vec::Vec;
     use core::{
         cell::{Cell, RefCell},
         marker::PhantomData,
-        mem, ptr,
+        mem, ptr, slice,
     };
 
     /// A [TypedArena] can hold many instances of a single type, and will properly [Drop] them when
@@ -112,7 +114,6 @@ pub mod typed_arena {
             }
         }
 
-        #[allow(clippy::mut_from_ref)]
         pub fn alloc(&'arena self, value: T) -> &'arena mut T {
             if self.head == self.tail {
                 self.grow(1);
@@ -134,6 +135,52 @@ pub mod typed_arena {
             }
         }
 
+        fn can_allocate(&self, len: usize) -> bool {
+            len <= unsafe { self.tail.get().offset_from(self.head.get()) as usize }
+        }
+
+        /// # Panics
+        /// Panics if size_of::<T> == 0 || len == 0
+        #[inline]
+        fn alloc_raw_slice(&self, len: usize) -> *mut T {
+            assert!(mem::size_of::<T>() != 0);
+            assert!(len != 0);
+
+            if !self.can_allocate(len) {
+                self.grow(len)
+            }
+
+            let out = self.head.get();
+
+            unsafe { self.head.set(out.add(len)) };
+            out
+        }
+
+        pub fn alloc_from_iter<I>(&'arena self, iter: I) -> &'arena mut [T]
+        where I: IntoIterator<Item = T> {
+            // Collect them all into a buffer so they're allocated contiguously
+            let mut buf = iter.into_iter().collect::<Vec<_>>();
+            if buf.is_empty() {
+                return &mut [];
+            }
+
+            let len = buf.len();
+            // If T is a ZST, calling alloc_raw_slice will panic
+            let slice = if mem::size_of::<T>() == 0 {
+                self.head
+                    .set(ptr::without_provenance_mut(self.head.get().addr() + len));
+                ptr::NonNull::dangling().as_ptr()
+            } else {
+                self.alloc_raw_slice(len)
+            };
+
+            unsafe {
+                buf.as_ptr().copy_to_nonoverlapping(slice, len);
+                buf.set_len(0);
+                slice::from_raw_parts_mut(slice, len)
+            }
+        }
+
         #[cold]
         #[inline(never)]
         fn grow(&self, len: usize) {