perf: Add new efficient APIs read_unsafe and read_to_vec

src/base_vec.rs

@@ -33,8 +33,8 @@
 //! bytes required to represent integers up to that max size.
 use crate::storable::{bounds, bytes_to_store_size_bounded};
 use crate::{
-    read_u32, read_u64, safe_write, write, write_u32, write_u64, Address, GrowFailed, Memory,
-    Storable,
+    read_to_vec, read_u32, read_u64, safe_write, write, write_u32, write_u64, Address, GrowFailed,
+    Memory, Storable,
 };
 use std::borrow::{Borrow, Cow};
 use std::cmp::min;
@@ -245,11 +245,10 @@ impl<T: Storable, M: Memory> BaseVec<T, M> {
     }
 
     /// Reads the item at the specified index without any bound checks.
-    fn read_entry_to(&self, index: u64, buf: &mut std::vec::Vec<u8>) {
+    fn read_entry_to(&self, index: u64, buf: &mut Vec<u8>) {
         let offset = DATA_OFFSET + slot_size::<T>() as u64 * index;
         let (data_offset, data_size) = self.read_entry_size(offset);
-        buf.resize(data_size, 0);
-        self.memory.read(data_offset, &mut buf[..]);
+        read_to_vec(&self.memory, data_offset.into(), buf, data_size);
     }
 
     /// Sets the vector's length.

src/btreemap/node.rs

@@ -1,6 +1,6 @@
 use crate::{
     btreemap::Allocator,
-    read_struct, read_u32, read_u64,
+    read_struct, read_to_vec, read_u32, read_u64,
     storable::Storable,
     types::{Address, Bytes},
     write, write_struct, write_u32, Memory,
@@ -190,7 +190,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
         value.take_or_load(|offset| self.load_value_from_memory(offset, memory))
     }
 
-    /// Loads a value from stable memory at the given offset.
+    /// Loads a value from stable memory at the given offset of this node.
     fn load_value_from_memory<M: Memory>(&self, offset: Bytes, memory: &M) -> Vec<u8> {
         let reader = NodeReader {
             address: self.address,
@@ -200,8 +200,13 @@ impl<K: Storable + Ord + Clone> Node<K> {
         };
 
         let value_len = read_u32(&reader, Address::from(offset.get())) as usize;
-        let mut bytes = vec![0; value_len];
-        reader.read((offset + U32_SIZE).get(), &mut bytes);
+        let mut bytes = vec![];
+        read_to_vec(
+            &reader,
+            Address::from((offset + U32_SIZE).get()),
+            &mut bytes,
+            value_len,
+        );
 
         bytes
     }

src/btreemap/node/io.rs

@@ -18,20 +18,21 @@ pub struct NodeReader<'a, M: Memory> {
 // Note: The `Memory` interface is implemented so that helper methods such `read_u32`,
 // `read_struct`, etc. can be used with a `NodeReader` directly.
 impl<'a, M: Memory> Memory for NodeReader<'a, M> {
-    fn read(&self, offset: u64, dst: &mut [u8]) {
+    unsafe fn read_unsafe(&self, offset: u64, dst: *mut u8, count: usize) {
         // If the read is only in the initial page, then read it directly in one go.
         // This is a performance enhancement to avoid the cost of creating a `NodeIterator`.
-        if (offset + dst.len() as u64) < self.page_size.get() as u64 {
-            self.memory.read(self.address.get() + offset, dst);
+        if (offset + count as u64) < self.page_size.get() as u64 {
+            self.memory
+                .read_unsafe(self.address.get() + offset, dst, count);
             return;
         }
 
-        // The read is split across several pages. Create a `NodeIterator` to to read from
+        // The read is split across several pages. Create a `NodeIterator` to read from
         // each of the individual pages.
         let iter = NodeIterator::new(
             VirtualSegment {
                 address: Address::from(offset),
-                length: Bytes::from(dst.len() as u64),
+                length: Bytes::from(count as u64),
             },
             Bytes::from(self.page_size.get()),
         );
@@ -43,22 +44,33 @@ impl<'a, M: Memory> Memory for NodeReader<'a, M> {
             length,
         } in iter
         {
+            // SAFETY: read_unsafe() is safe to call iff bytes_read + length <= count since the
+            // caller guarantees that we can write `count` number of bytes to `dst`.
+            assert!(bytes_read + length.get() as usize <= count);
             if page_idx == 0 {
-                self.memory.read(
+                self.memory.read_unsafe(
                     (self.address + offset).get(),
-                    &mut dst[bytes_read as usize..(bytes_read + length.get()) as usize],
+                    dst.add(bytes_read),
+                    length.get() as usize,
                 );
             } else {
-                self.memory.read(
+                self.memory.read_unsafe(
                     (self.overflows[page_idx - 1] + offset).get(),
-                    &mut dst[bytes_read as usize..(bytes_read + length.get()) as usize],
+                    dst.add(bytes_read),
+                    length.get() as usize,
                 );
             }
 
-            bytes_read += length.get();
+            bytes_read += length.get() as usize;
         }
     }
 
+    #[inline]
+    fn read(&self, offset: u64, dst: &mut [u8]) {
+        // SAFETY: since dst is dst.len() long, it fulfills the safety requirements of read_unsafe.
+        unsafe { self.read_unsafe(offset, dst.as_mut_ptr(), dst.len()) }
+    }
+
     fn write(&self, _: u64, _: &[u8]) {
         unreachable!("NodeReader does not support write")
     }

src/btreemap/node/v1.rs

@@ -65,15 +65,14 @@ impl<K: Storable + Ord + Clone> Node<K> {
         // Load the entries.
         let mut keys_encoded_values = Vec::with_capacity(header.num_entries as usize);
         let mut offset = NodeHeader::size();
-        let mut buf = Vec::with_capacity(max_key_size.max(max_value_size) as usize);
+        let mut buf = vec![];
         for _ in 0..header.num_entries {
             // Read the key's size.
             let key_size = read_u32(memory, address + offset);
             offset += U32_SIZE;
 
             // Read the key.
-            buf.resize(key_size as usize, 0);
-            memory.read((address + offset).get(), &mut buf);
+            read_to_vec(memory, address + offset, &mut buf, key_size as usize);
             offset += Bytes::from(max_key_size);
             let key = K::from_bytes(Cow::Borrowed(&buf));
             // Values are loaded lazily. Store a reference and skip loading it.
@@ -86,7 +85,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
         let mut children = vec![];
         if header.node_type == INTERNAL_NODE_TYPE {
             // The number of children is equal to the number of entries + 1.
-            children.reserve(header.num_entries as usize + 1);
+            children.reserve_exact(header.num_entries as usize + 1);
             for _ in 0..header.num_entries + 1 {
                 let child = Address::from(read_u64(memory, address + offset));
                 offset += Address::size();

src/btreemap/node/v2.rs

@@ -141,6 +141,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
         let mut children = vec![];
         if node_type == NodeType::Internal {
             // The number of children is equal to the number of entries + 1.
+            children.reserve_exact(num_entries + 1);
             for _ in 0..num_entries + 1 {
                 let child = Address::from(read_u64(&reader, offset));
                 offset += Address::size();
@@ -164,8 +165,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
             };
 
             // Load the key.
-            buf.resize(key_size as usize, 0);
-            reader.read(offset.get(), &mut buf);
+            read_to_vec(&reader, offset, &mut buf, key_size as usize);
             let key = K::from_bytes(Cow::Borrowed(&buf));
             offset += Bytes::from(key_size);
             keys_encoded_values.push((key, Value::by_ref(Bytes::from(0usize))));

src/cell.rs

@@ -1,6 +1,6 @@
 //! A serializable value stored in the stable memory.
 use crate::storable::Storable;
-use crate::{Memory, WASM_PAGE_SIZE};
+use crate::{read_to_vec, Memory, WASM_PAGE_SIZE};
 use std::borrow::{Borrow, Cow};
 use std::fmt;
 
@@ -132,8 +132,8 @@ impl<T: Storable, M: Memory> Cell<T, M> {
     ///
     /// PRECONDITION: memory is large enough to contain the value.
     fn read_value(memory: &M, len: u32) -> T {
-        let mut buf = vec![0; len as usize];
-        memory.read(HEADER_V1_SIZE, &mut buf);
+        let mut buf = vec![];
+        read_to_vec(memory, HEADER_V1_SIZE.into(), &mut buf, len as usize);
         T::from_bytes(Cow::Owned(buf))
     }
 

src/ic0_memory.rs

@@ -30,6 +30,12 @@ impl Memory for Ic0StableMemory {
         unsafe { stable64_read(dst.as_ptr() as u64, offset, dst.len() as u64) }
     }
 
+    #[inline]
+    unsafe fn read_unsafe(&self, offset: u64, dst: *mut u8, count: usize) {
+        // SAFETY: This is safe because of the ic0 api guarantees.
+        stable64_read(dst as u64, offset, count as u64);
+    }
+
     #[inline]
     fn write(&self, offset: u64, src: &[u8]) {
         // SAFETY: This is safe because of the ic0 api guarantees.

src/lib.rs

@@ -26,6 +26,7 @@ pub use file_mem::FileMemory;
 pub use ic0_memory::Ic0StableMemory;
 use std::error;
 use std::fmt::{Display, Formatter};
+use std::mem::MaybeUninit;
 pub use storable::Storable;
 use types::Address;
 pub use vec_mem::VectorMemory;
@@ -46,42 +47,81 @@ pub trait Memory {
     /// pages. (One WebAssembly page is 64Ki bytes.)
     fn size(&self) -> u64;
 
-    /// Tries to grow the memory by new_pages many pages containing
+    /// Tries to grow the memory by `pages` many pages containing
     /// zeroes.  If successful, returns the previous size of the
     /// memory (in pages).  Otherwise, returns -1.
     fn grow(&self, pages: u64) -> i64;
 
-    /// Copies the data referred to by offset out of the stable memory
-    /// and replaces the corresponding bytes in dst.
+    /// Copies the data referred to by `offset` out of the stable memory
+    /// and replaces the corresponding bytes in `dst`.
     fn read(&self, offset: u64, dst: &mut [u8]);
 
-    /// Copies the data referred to by src and replaces the
+    /// Copies `count` number of bytes of the data starting from `offset` out of the stable memory
+    /// into the buffer starting at `dst`.
+    ///
+    /// This method is an alternative to `read` which does not require initializing a buffer and may
+    /// therefore be faster.
+    ///
+    /// # Safety
+    ///
+    /// Callers must guarantee that
+    ///   * it is valid to write `count` number of bytes starting from `dst`,
+    ///   * `dst..dst + count` does not overlap with `self`.
+    ///
+    /// Implementations must guarantee that before the method returns, `count` number of bytes
+    /// starting from `dst` will be initialized.
+    #[inline]
+    unsafe fn read_unsafe(&self, offset: u64, dst: *mut u8, count: usize) {
+        // Initialize the buffer to make the slice valid.
+        std::ptr::write_bytes(dst, 0, count);
+        let slice = std::slice::from_raw_parts_mut(dst, count);
+        self.read(offset, slice)
+    }
+
+    /// Copies the data referred to by `src` and replaces the
     /// corresponding segment starting at offset in the stable memory.
     fn write(&self, offset: u64, src: &[u8]);
 }
 
-// A helper function that reads a single 32bit integer encoded as
-// little-endian from the specified memory at the specified offset.
+/// Copies `count` bytes of data starting from `addr` out of the stable memory into `dst`.
+///
+/// Callers are allowed to pass vectors in any state (e.g. empty vectors).
+/// After the method returns, `dst.len() == count`.
+/// This method is an alternative to `read` which does not require initializing a buffer and may
+/// therefore be faster.
+#[inline]
+fn read_to_vec<M: Memory>(m: &M, addr: Address, dst: &mut std::vec::Vec<u8>, count: usize) {
+    dst.clear();
+    dst.reserve_exact(count);
+    unsafe {
+        m.read_unsafe(addr.get(), dst.as_mut_ptr(), count);
+        // SAFETY: read_unsafe guarantees to initialize the first `count` bytes
+        dst.set_len(count);
+    }
+}
+
+/// A helper function that reads a single 32bit integer encoded as
+/// little-endian from the specified memory at the specified offset.
 fn read_u32<M: Memory>(m: &M, addr: Address) -> u32 {
     let mut buf: [u8; 4] = [0; 4];
     m.read(addr.get(), &mut buf);
     u32::from_le_bytes(buf)
 }
 
-// A helper function that reads a single 64bit integer encoded as
-// little-endian from the specified memory at the specified offset.
+/// A helper function that reads a single 64bit integer encoded as
+/// little-endian from the specified memory at the specified offset.
 fn read_u64<M: Memory>(m: &M, addr: Address) -> u64 {
     let mut buf: [u8; 8] = [0; 8];
     m.read(addr.get(), &mut buf);
     u64::from_le_bytes(buf)
 }
 
-// Writes a single 32-bit integer encoded as little-endian.
+/// Writes a single 32-bit integer encoded as little-endian.
 fn write_u32<M: Memory>(m: &M, addr: Address, val: u32) {
     write(m, addr.get(), &val.to_le_bytes());
 }
 
-// Writes a single 64-bit integer encoded as little-endian.
+/// Writes a single 64-bit integer encoded as little-endian.
 fn write_u64<M: Memory>(m: &M, addr: Address, val: u64) {
     write(m, addr.get(), &val.to_le_bytes());
 }
@@ -148,17 +188,20 @@ fn write<M: Memory>(memory: &M, offset: u64, bytes: &[u8]) {
     }
 }
 
-// Reads a struct from memory.
+/// Reads a struct from memory.
 fn read_struct<T, M: Memory>(addr: Address, memory: &M) -> T {
-    let mut t: T = unsafe { core::mem::zeroed() };
-    let t_slice = unsafe {
-        core::slice::from_raw_parts_mut(&mut t as *mut _ as *mut u8, core::mem::size_of::<T>())
-    };
-    memory.read(addr.get(), t_slice);
-    t
+    let mut value = MaybeUninit::<T>::uninit();
+    unsafe {
+        memory.read_unsafe(
+            addr.get(),
+            value.as_mut_ptr() as *mut u8,
+            core::mem::size_of::<T>(),
+        );
+        value.assume_init()
+    }
 }
 
-// Writes a struct to memory.
+/// Writes a struct to memory.
 fn write_struct<T, M: Memory>(t: &T, addr: Address, memory: &M) {
     let slice = unsafe {
         core::slice::from_raw_parts(t as *const _ as *const u8, core::mem::size_of::<T>())

src/log.rs

@@ -54,7 +54,7 @@
 //! ----------------------------------------
 //! Unallocated space
 //! ```
-use crate::{read_u64, safe_write, write_u64, Address, GrowFailed, Memory, Storable};
+use crate::{read_to_vec, read_u64, safe_write, write_u64, Address, GrowFailed, Memory, Storable};
 use std::borrow::Cow;
 use std::cell::RefCell;
 use std::fmt;
@@ -331,8 +331,7 @@ impl<T: Storable, INDEX: Memory, DATA: Memory> Log<T, INDEX, DATA> {
     /// ignores the result.
     pub fn read_entry(&self, idx: u64, buf: &mut Vec<u8>) -> Result<(), NoSuchEntry> {
         let (offset, len) = self.entry_meta(idx).ok_or(NoSuchEntry)?;
-        buf.resize(len, 0);
-        self.data_memory.read(HEADER_OFFSET + offset, buf);
+        read_to_vec(&self.data_memory, (HEADER_OFFSET + offset).into(), buf, len);
         Ok(())
     }
 

src/memory_manager.rs

@@ -41,7 +41,7 @@
 //! assert_eq!(bytes, vec![4, 5, 6]);
 //! ```
 use crate::{
-    read_struct,
+    read_struct, read_to_vec,
     types::{Address, Bytes},
     write, write_struct, Memory, WASM_PAGE_SIZE,
 };
@@ -239,9 +239,9 @@ impl<M: Memory> MemoryManagerInner<M> {
         }
 
         // Check if the magic in the memory corresponds to this object.
-        let mut dst = vec![0; 3];
+        let mut dst = [0; 3];
         memory.read(0, &mut dst);
-        if dst != MAGIC {
+        if &dst != MAGIC {
             // No memory manager found. Create a new instance.
             MemoryManagerInner::new(memory, bucket_size_in_pages)
         } else {
@@ -277,8 +277,13 @@ impl<M: Memory> MemoryManagerInner<M> {
         assert_eq!(&header.magic, MAGIC, "Bad magic.");
         assert_eq!(header.version, LAYOUT_VERSION, "Unsupported version.");
 
-        let mut buckets = vec![0; MAX_NUM_BUCKETS as usize];
-        memory.read(bucket_allocations_address(BucketId(0)).get(), &mut buckets);
+        let mut buckets = vec![];
+        read_to_vec(
+            &memory,
+            bucket_allocations_address(BucketId(0)),
+            &mut buckets,
+            MAX_NUM_BUCKETS as usize,
+        );
 
         let mut memory_buckets = BTreeMap::new();
         for (bucket_idx, memory) in buckets.into_iter().enumerate() {