[move][move-vm][rewrite][cleanup 2/n] Use arena allocation across the entire runtime AST

external-crates/move/crates/move-vm-runtime/src/cache/arena.rs

@@ -14,6 +14,16 @@ use bumpalo::Bump;
 
 pub struct Arena(Bump);
 
+/// An arena-allocated vector. Notably, `Drop` does not drop the elements it holds, as that is the
+/// perview of the arena it was allocated in.
+#[derive(Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
+pub struct ArenaVec<T>(std::mem::ManuallyDrop<Vec<T>>);
+
+/// An arena-allocated vector. Notably, `Drop` does not drop the value it holds, as that is the
+/// perview of the arena it was allocated in.
+#[derive(Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
+pub struct ArenaBox<T>(std::mem::ManuallyDrop<Box<T>>);
+
 /// Size of a package arena.
 /// This is 10 megabytes, which should be more than enough room for any pacakge on chain.
 /// FIXME: Test this limit and validate. See how large packages are in backtesting and double that
@@ -37,28 +47,102 @@ impl Arena {
         Arena(Bump::new())
     }
 
-    /// SAFETY: it is the caller's responsibility to ensure that `self` is not shared across
-    /// threads during this call. This should be fine as the translation step that uses an arena
-    /// should happen in a thread that holds that arena, with no other contention for allocation
-    /// into it, and nothing should allocate into a LoadedModule after it is loaded.
-    pub fn alloc_slice<T>(
+    /// SAFETY:
+    /// 1. It is the caller's responsibility to ensure that `self` is not shared across threads
+    ///    during this call.
+    /// 2. This vector is allocated in the arena, and thus even it is dropped its memory will not
+    ///    be reclaimed until the arena is discarded.
+    pub fn alloc_vec<T>(
         &self,
         items: impl ExactSizeIterator<Item = T>,
-    ) -> PartialVMResult<*mut [T]> {
+    ) -> PartialVMResult<ArenaVec<T>> {
+        let size = items.len();
         if let Ok(slice) = self.0.try_alloc_slice_fill_iter(items) {
-            Ok(slice)
+            Ok(ArenaVec(unsafe {
+                std::mem::ManuallyDrop::new(Vec::from_raw_parts(slice.as_mut_ptr(), size, size))
+            }))
+        } else {
+            Err(PartialVMError::new(StatusCode::PACKAGE_ARENA_LIMIT_REACHED))
+        }
+    }
+
+    /// SAFETY:
+    /// 1. It is the caller's responsibility to ensure that `self` is not shared across threads
+    ///    during this call.
+    /// 2. This box is allocated in the arena, and thus even it is dropped its memory will not be
+    ///    reclaimed until the arena is discarded.
+    pub fn alloc_box<T>(&self, item: T) -> PartialVMResult<ArenaBox<T>> {
+        if let Ok(slice) = self.0.try_alloc(item) {
+            Ok(ArenaBox(unsafe {
+                std::mem::ManuallyDrop::new(Box::from_raw(slice as *mut T))
+            }))
+        } else {
+            Err(PartialVMError::new(StatusCode::PACKAGE_ARENA_LIMIT_REACHED))
+        }
+    }
+
+    pub fn alloc_item<T>(&self, item: T) -> PartialVMResult<*const T> {
+        if let Ok(value) = self.0.try_alloc(item) {
+            Ok(value)
         } else {
             Err(PartialVMError::new(StatusCode::PACKAGE_ARENA_LIMIT_REACHED))
         }
     }
 }
 
-// -----------------------------------------------
+impl<T> ArenaVec<T> {
+    pub fn iter(&self) -> std::slice::Iter<T> {
+        self.0.iter()
+    }
+
+    /// Returns an iterator over mutable references.
+    /// Crate-only because nobody else should be modifying arena values.
+    pub(crate) fn iter_mut(&mut self) -> std::slice::IterMut<T> {
+        self.0.iter_mut()
+    }
+
+    /// Make an empty ArenaVec
+    pub fn empty() -> Self {
+        ArenaVec(std::mem::ManuallyDrop::new(vec![]))
+    }
+}
+
+impl<T> ArenaBox<T> {
+    pub fn as_ref(&self) -> &T {
+        &self.0
+    }
+}
+
+// -------------------------------------------------------------------------------------------------
 // Trait Implementations
-// -----------------------------------------------
+// -------------------------------------------------------------------------------------------------
 
 // SAFETY: these are okay, if callers follow the documented safety requirements for `Arena`'s
 // unsafe methods.
 
 unsafe impl Send for Arena {}
 unsafe impl Sync for Arena {}
+
+// -----------------------------------------------
+// ArenaVec Trait Implementations
+// -----------------------------------------------
+
+impl<T> std::ops::Deref for ArenaVec<T> {
+    type Target = [T];
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+// -----------------------------------------------
+// ArenaVec Trait Implementations
+// -----------------------------------------------
+
+impl<T> std::ops::Deref for ArenaBox<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}

external-crates/move/crates/move-vm-runtime/src/cache/move_cache.rs

@@ -103,7 +103,7 @@ impl Package {
     /// Used for testing that the correct number of types are loaded
     #[allow(dead_code)]
     pub(crate) fn loaded_types_len(&self) -> usize {
-        self.runtime.vtable.types.cached_types.len()
+        self.runtime.vtable.types.len()
     }
 }