Rename MmapVec::drain to split_off (#4673)

* Rename `MmapVec::drain` to `split_off`

As suggested on #4609

* Fix tests

* Make MmapVec::split_off work like Vec::split_off

Co-authored-by: Jamey Sharp <jsharp@fastly.com>

crates/runtime/src/mmap_vec.rs

@@ -1,7 +1,7 @@
 use crate::Mmap;
 use anyhow::{Context, Result};
 use std::fs::File;
-use std::ops::{Deref, DerefMut, Range, RangeTo};
+use std::ops::{Deref, DerefMut, Range};
 use std::path::Path;
 use std::sync::Arc;
 
@@ -73,36 +73,25 @@ impl MmapVec {
         self.mmap.is_readonly()
     }
 
-    /// "Drains" leading bytes up to the end specified in `range` from this
+    /// Splits the collection into two at the given index.
-    /// `MmapVec`, returning a separately owned `MmapVec` which retains access
-    /// to the bytes.
     ///
-    /// This method is similar to the `Vec` type's `drain` method, except that
+    /// Returns a separate `MmapVec` which shares the underlying mapping, but
-    /// the return value is not an iterator but rather a new `MmapVec`. The
+    /// only has access to elements in the range `[at, len)`. After the call,
-    /// purpose of this method is the ability to split-off new `MmapVec` values
+    /// the original `MmapVec` will be left with access to the elements in the
-    /// which are sub-slices of the original one.
+    /// range `[0, at)`.
-    ///
-    /// Once data has been drained from an `MmapVec` it is no longer accessible
-    /// from the original `MmapVec`, it's only accessible from the returned
-    /// `MmapVec`. In other words ownership of the drain'd bytes is returned
-    /// through the `MmapVec` return value.
-    ///
-    /// This `MmapVec` will shrink by `range.end` bytes, and it will only refer
-    /// to the bytes that come after the drain range.
     ///
     /// This is an `O(1)` operation which does not involve copies.
-    pub fn drain(&mut self, range: RangeTo<usize>) -> MmapVec {
+    pub fn split_off(&mut self, at: usize) -> MmapVec {
-        let amt = range.end;
+        assert!(at <= self.range.len());
-        assert!(amt <= (self.range.end - self.range.start));
 
         // Create a new `MmapVec` which refers to the same underlying mmap, but
         // has a disjoint range from ours. Our own range is adjusted to be
         // disjoint just after `ret` is created.
         let ret = MmapVec {
             mmap: self.mmap.clone(),
-            range: self.range.start..self.range.start + amt,
+            range: at..self.range.end,
         };
-        self.range.start += amt;
+        self.range.end = self.range.start + at;
         return ret;
     }
 
@@ -173,29 +162,24 @@ mod tests {
     }
 
     #[test]
-    fn drain() {
+    fn split_off() {
-        let mut mmap = MmapVec::from_slice(&[1, 2, 3, 4]).unwrap();
+        let mut vec = Vec::from([1, 2, 3, 4]);
-        assert_eq!(mmap.len(), 4);
+        let mut mmap = MmapVec::from_slice(&vec).unwrap();
-        assert!(mmap.drain(..0).is_empty());
+        assert_eq!(&mmap[..], &vec[..]);
-        assert_eq!(mmap.len(), 4);
+        // remove nothing; vec length remains 4
-        let one = mmap.drain(..1);
+        assert_eq!(&mmap.split_off(4)[..], &vec.split_off(4)[..]);
-        assert_eq!(one.len(), 1);
+        assert_eq!(&mmap[..], &vec[..]);
-        assert_eq!(one[0], 1);
+        // remove 1 element; vec length is now 3
-        assert_eq!(mmap.len(), 3);
+        assert_eq!(&mmap.split_off(3)[..], &vec.split_off(3)[..]);
-        assert_eq!(&mmap[..], &[2, 3, 4]);
+        assert_eq!(&mmap[..], &vec[..]);
-        drop(one);
+        // remove 2 elements; vec length is now 1
-        assert_eq!(mmap.len(), 3);
+        assert_eq!(&mmap.split_off(1)[..], &vec.split_off(1)[..]);
-
+        assert_eq!(&mmap[..], &vec[..]);
-        let two = mmap.drain(..2);
+        // remove last element; vec length is now 0
-        assert_eq!(two.len(), 2);
+        assert_eq!(&mmap.split_off(0)[..], &vec.split_off(0)[..]);
-        assert_eq!(two[0], 2);
+        assert_eq!(&mmap[..], &vec[..]);
-        assert_eq!(two[1], 3);
+        // nothing left to remove, but that's okay
-        assert_eq!(mmap.len(), 1);
+        assert_eq!(&mmap.split_off(0)[..], &vec.split_off(0)[..]);
-        assert_eq!(mmap[0], 4);
+        assert_eq!(&mmap[..], &vec[..]);
-        drop(two);
-        assert!(mmap.drain(..0).is_empty());
-        assert!(mmap.drain(..1).len() == 1);
-        assert!(mmap.is_empty());
-        assert!(mmap.drain(..0).is_empty());
     }
 }

crates/wasmtime/src/module/serialization.rs

@@ -268,11 +268,11 @@ impl<'a> SerializedModule<'a> {
         // First validate that this is at least somewhat an elf file within
         // `mmap` and additionally skip to the end of the elf file to find our
         // metadata.
-        let elf = take_first_elf(&mut mmap)?;
+        let metadata = data_after_elf(&mut mmap)?;
 
         // The metadata has a few guards up front which we process first, and
         // eventually this bottoms out in a `bincode::deserialize` call.
-        let metadata = mmap
+        let metadata = metadata
             .strip_prefix(HEADER)
             .ok_or_else(|| anyhow!("bytes are not a compatible serialized wasmtime module"))?;
         if metadata.is_empty() {
@@ -309,13 +309,13 @@ impl<'a> SerializedModule<'a> {
             .context("deserialize compilation artifacts")?;
 
         return Ok(SerializedModule {
-            artifacts: MyCow::Owned(elf),
+            artifacts: MyCow::Owned(mmap),
             metadata,
         });
 
         /// This function will return the trailing data behind the ELF file
         /// parsed from `data` which is where we find our metadata section.
-        fn take_first_elf(mmap: &mut MmapVec) -> Result<MmapVec> {
+        fn data_after_elf(mmap: &mut MmapVec) -> Result<MmapVec> {
             use object::NativeEndian as NE;
             // There's not actually a great utility for figuring out where
             // the end of an ELF file is in the `object` crate. In lieu of that
@@ -335,7 +335,7 @@ impl<'a> SerializedModule<'a> {
                 .section_headers(NE, data)
                 .context("failed to read section headers")?;
             let range = subslice_range(object::bytes_of_slice(sections), data);
-            Ok(mmap.drain(..range.end))
+            Ok(mmap.split_off(range.end))
         }
     }