From 2129c23fac32b7da3cc55a0c65a6a066dd4bdd2a Mon Sep 17 00:00:00 2001 From: David Frank Date: Fri, 22 Nov 2024 13:06:18 +0100 Subject: [PATCH] perf: Add new efficient APIs read_unsafe and read_to_vec (#248) I found that a source of significant performance loss is the read method of `Memory`. The `read` method takes a mutable buffer which it fills with values read from the stable memory. According to Rust rules, the buffer passed to read must be initialized before it's passed to read (buffers containing uninitialized values are unsound and can cause UB). The usual pattern is to create a properly sized Vec, eg. by using `vec![0; size]` or `vec.resize(size, 0)` and pass that to `read`. However, initializing the bytes with values that get overwritten by `read` is only necessary in order to be sound and requires significant number of instructions. This PR introduces a new method `read_unsafe` which allows passing in a raw pointer and a `count` parameter. Implementations can be more efficient by reading directly and skipping initialization. This can lead to instruction reductions of up to 40%. The PR also introduces a helper method `read_to_vec` which is a safe wrapper around `read_unsafe` for the most common use-case: reading into a `Vec`. Clients can for example pass an empty `Vec` and profit from the extra efficiency without having to call unsafe methods. --------- Co-authored-by: Andriy Berestovskyy --- canbench_results.yml | 212 ++++++++++++++++++++-------------------- src/base_vec.rs | 9 +- src/btreemap/node.rs | 13 ++- src/btreemap/node/io.rs | 32 ++++-- src/btreemap/node/v1.rs | 7 +- src/btreemap/node/v2.rs | 4 +- src/cell.rs | 6 +- src/ic0_memory.rs | 6 ++ src/lib.rs | 79 +++++++++++---- src/log.rs | 5 +- src/memory_manager.rs | 15 ++- src/tests.rs | 58 +++++++++++ src/vec_mem.rs | 19 ++++ 13 files changed, 305 insertions(+), 160 deletions(-) diff --git a/canbench_results.yml b/canbench_results.yml index d80974e5..d4a0e687 100644 --- a/canbench_results.yml +++ b/canbench_results.yml @@ -1,595 +1,595 @@ benches: btreemap_get_blob_128_1024: total: - instructions: 868028993 + instructions: 816917041 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_128_1024_v2: total: - instructions: 970754065 + instructions: 891561192 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_16_1024: total: - instructions: 304555831 + instructions: 288610553 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_16_1024_v2: total: - instructions: 409386098 + instructions: 357151536 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_256_1024: total: - instructions: 1399133930 + instructions: 1318848488 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_256_1024_v2: total: - instructions: 1499851078 + instructions: 1396573492 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_32_1024: total: - instructions: 347918329 + instructions: 318958937 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_32_1024_v2: total: - instructions: 452317746 + instructions: 390388108 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_4_1024: total: - instructions: 194406501 + instructions: 178025469 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_4_1024_v2: total: - instructions: 289345080 + instructions: 256537092 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_512_1024: total: - instructions: 2454262332 + instructions: 2328426328 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_512_1024_v2: total: - instructions: 2556213321 + instructions: 2400844220 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_64_1024: total: - instructions: 599847298 + instructions: 568565131 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_64_1024_v2: total: - instructions: 712044951 + instructions: 642026306 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_8_1024: total: - instructions: 231572100 + instructions: 209549111 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_8_1024_v2: total: - instructions: 323197606 + instructions: 285966738 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_8_u64: total: - instructions: 204171772 + instructions: 193326958 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_blob_8_u64_v2: total: - instructions: 311528236 + instructions: 284356404 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_u64_blob_8: total: - instructions: 186152093 + instructions: 175234112 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_u64_blob_8_v2: total: - instructions: 273254410 + instructions: 244592514 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_u64_u64: total: - instructions: 187408768 + instructions: 176021844 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_get_u64_u64_v2: total: - instructions: 279916538 + instructions: 251380137 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_insert_10mib_values: total: - instructions: 141252389 + instructions: 82015558 heap_increase: 0 stable_memory_increase: 32 scopes: {} btreemap_insert_blob_1024_128: total: - instructions: 5083095809 + instructions: 4875355783 heap_increase: 0 stable_memory_increase: 262 scopes: {} btreemap_insert_blob_1024_128_v2: total: - instructions: 5196791890 + instructions: 4956719432 heap_increase: 0 stable_memory_increase: 196 scopes: {} btreemap_insert_blob_1024_16: total: - instructions: 5039894252 + instructions: 4863260567 heap_increase: 0 stable_memory_increase: 241 scopes: {} btreemap_insert_blob_1024_16_v2: total: - instructions: 5154652816 + instructions: 4940158553 heap_increase: 0 stable_memory_increase: 181 scopes: {} btreemap_insert_blob_1024_256: total: - instructions: 5109831287 + instructions: 4901893503 heap_increase: 0 stable_memory_increase: 292 scopes: {} btreemap_insert_blob_1024_256_v2: total: - instructions: 5222715261 + instructions: 4985868123 heap_increase: 0 stable_memory_increase: 219 scopes: {} btreemap_insert_blob_1024_32: total: - instructions: 5048498828 + instructions: 4869006447 heap_increase: 0 stable_memory_increase: 239 scopes: {} btreemap_insert_blob_1024_32_v2: total: - instructions: 5165315603 + instructions: 4957385696 heap_increase: 0 stable_memory_increase: 180 scopes: {} btreemap_insert_blob_1024_4: total: - instructions: 4938758192 + instructions: 4767021647 heap_increase: 0 stable_memory_increase: 235 scopes: {} btreemap_insert_blob_1024_4_v2: total: - instructions: 5053526974 + instructions: 4836048896 heap_increase: 0 stable_memory_increase: 176 scopes: {} btreemap_insert_blob_1024_512: total: - instructions: 5232116884 + instructions: 4978502221 heap_increase: 0 stable_memory_increase: 348 scopes: {} btreemap_insert_blob_1024_512_v2: total: - instructions: 5346066549 + instructions: 5068330764 heap_increase: 0 stable_memory_increase: 261 scopes: {} btreemap_insert_blob_1024_64: total: - instructions: 5080908184 + instructions: 4887214834 heap_increase: 0 stable_memory_increase: 250 scopes: {} btreemap_insert_blob_1024_64_v2: total: - instructions: 5196471697 + instructions: 4976282466 heap_increase: 0 stable_memory_increase: 188 scopes: {} btreemap_insert_blob_1024_8: total: - instructions: 5022590452 + instructions: 4832986464 heap_increase: 0 stable_memory_increase: 237 scopes: {} btreemap_insert_blob_1024_8_v2: total: - instructions: 5136760933 + instructions: 4925472061 heap_increase: 0 stable_memory_increase: 178 scopes: {} btreemap_insert_blob_128_1024: total: - instructions: 1437432650 + instructions: 1251468811 heap_increase: 0 stable_memory_increase: 260 scopes: {} btreemap_insert_blob_128_1024_v2: total: - instructions: 1543573613 + instructions: 1342067516 heap_increase: 0 stable_memory_increase: 195 scopes: {} btreemap_insert_blob_16_1024: total: - instructions: 819015110 + instructions: 677272206 heap_increase: 0 stable_memory_increase: 215 scopes: {} btreemap_insert_blob_16_1024_v2: total: - instructions: 929386468 + instructions: 761510359 heap_increase: 0 stable_memory_increase: 161 scopes: {} btreemap_insert_blob_256_1024: total: - instructions: 2007916344 + instructions: 1805606715 heap_increase: 0 stable_memory_increase: 292 scopes: {} btreemap_insert_blob_256_1024_v2: total: - instructions: 2117490329 + instructions: 1893976279 heap_increase: 0 stable_memory_increase: 219 scopes: {} btreemap_insert_blob_32_1024: total: - instructions: 870306996 + instructions: 712489642 heap_increase: 0 stable_memory_increase: 230 scopes: {} btreemap_insert_blob_32_1024_v2: total: - instructions: 980735900 + instructions: 800214505 heap_increase: 0 stable_memory_increase: 173 scopes: {} btreemap_insert_blob_4_1024: total: - instructions: 615105126 + instructions: 483550320 heap_increase: 0 stable_memory_increase: 123 scopes: {} btreemap_insert_blob_4_1024_v2: total: - instructions: 714364544 + instructions: 571897395 heap_increase: 0 stable_memory_increase: 92 scopes: {} btreemap_insert_blob_512_1024: total: - instructions: 3140701710 + instructions: 2894814130 heap_increase: 0 stable_memory_increase: 351 scopes: {} btreemap_insert_blob_512_1024_v2: total: - instructions: 3248440359 + instructions: 2979049543 heap_increase: 0 stable_memory_increase: 263 scopes: {} btreemap_insert_blob_64_1024: total: - instructions: 1136783783 + instructions: 980721257 heap_increase: 0 stable_memory_increase: 245 scopes: {} btreemap_insert_blob_64_1024_v2: total: - instructions: 1254935822 + instructions: 1067083136 heap_increase: 0 stable_memory_increase: 183 scopes: {} btreemap_insert_blob_8_1024: total: - instructions: 737978721 + instructions: 595149775 heap_increase: 0 stable_memory_increase: 183 scopes: {} btreemap_insert_blob_8_1024_v2: total: - instructions: 840828948 + instructions: 688041135 heap_increase: 0 stable_memory_increase: 138 scopes: {} btreemap_insert_blob_8_u64: total: - instructions: 333681568 + instructions: 320065029 heap_increase: 0 stable_memory_increase: 6 scopes: {} btreemap_insert_blob_8_u64_v2: total: - instructions: 448187595 + instructions: 422061780 heap_increase: 0 stable_memory_increase: 4 scopes: {} btreemap_insert_u64_blob_8: total: - instructions: 343283004 + instructions: 331086826 heap_increase: 0 stable_memory_increase: 7 scopes: {} btreemap_insert_u64_blob_8_v2: total: - instructions: 429045170 + instructions: 403842403 heap_increase: 0 stable_memory_increase: 5 scopes: {} btreemap_insert_u64_u64: total: - instructions: 353426417 + instructions: 336638169 heap_increase: 0 stable_memory_increase: 7 scopes: {} btreemap_insert_u64_u64_v2: total: - instructions: 441738753 + instructions: 412448876 heap_increase: 0 stable_memory_increase: 6 scopes: {} btreemap_iter_10mib_values: total: - instructions: 17054815 + instructions: 11384376 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_iter_count_10mib_values: total: - instructions: 492818 + instructions: 472523 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_iter_count_small_values: total: - instructions: 9854951 + instructions: 9316829 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_iter_rev_10mib_values: total: - instructions: 17052720 + instructions: 11381349 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_iter_rev_small_values: total: - instructions: 13980402 + instructions: 13437581 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_iter_small_values: total: - instructions: 13983601 + instructions: 13444863 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_keys_10mib_values: total: - instructions: 482392 + instructions: 461429 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_keys_rev_10mib_values: total: - instructions: 484745 + instructions: 462826 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_keys_rev_small_values: total: - instructions: 10226743 + instructions: 9683922 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_keys_small_values: total: - instructions: 9990906 + instructions: 9452168 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_read_every_third_value_from_range: total: - instructions: 112238714 + instructions: 82482266 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_read_keys_from_range: total: - instructions: 112238714 + instructions: 82482266 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_128_1024: total: - instructions: 1783414773 + instructions: 1506697673 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_128_1024_v2: total: - instructions: 1942678397 + instructions: 1637209205 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_16_1024: total: - instructions: 1017479433 + instructions: 785362788 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_16_1024_v2: total: - instructions: 1170176089 + instructions: 910340295 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_256_1024: total: - instructions: 2454002394 + instructions: 2148746627 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_256_1024_v2: total: - instructions: 2606407435 + instructions: 2275040941 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_32_1024: total: - instructions: 1095183445 + instructions: 848733142 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_32_1024_v2: total: - instructions: 1252648744 + instructions: 975603769 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_4_1024: total: - instructions: 617902738 + instructions: 469755304 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_4_1024_v2: total: - instructions: 735874430 + instructions: 574320850 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_512_1024: total: - instructions: 3849259424 + instructions: 3494883108 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_512_1024_v2: total: - instructions: 4003887427 + instructions: 3611250969 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_64_1024: total: - instructions: 1427599410 + instructions: 1176535123 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_64_1024_v2: total: - instructions: 1592814332 + instructions: 1303890927 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_8_1024: total: - instructions: 816476142 + instructions: 619861550 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_8_1024_v2: total: - instructions: 952832423 + instructions: 741193163 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_8_u64: total: - instructions: 438199023 + instructions: 419913395 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_blob_8_u64_v2: total: - instructions: 591309106 + instructions: 558059334 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_u64_blob_8: total: - instructions: 485820315 + instructions: 470317117 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_u64_blob_8_v2: total: - instructions: 606724491 + instructions: 575650396 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_u64_u64: total: - instructions: 506491112 + instructions: 483160741 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_remove_u64_u64_v2: total: - instructions: 637006929 + instructions: 597771476 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_values_10mib_values: total: - instructions: 17078596 + instructions: 11408157 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_values_rev_10mib_values: total: - instructions: 17077387 + instructions: 11406016 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_values_rev_small_values: total: - instructions: 15212366 + instructions: 14669545 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_values_small_values: total: - instructions: 15171561 + instructions: 14632823 heap_increase: 0 stable_memory_increase: 0 scopes: {} @@ -613,43 +613,43 @@ benches: scopes: {} vec_get_blob_128: total: - instructions: 21384965 + instructions: 17779013 heap_increase: 0 stable_memory_increase: 0 scopes: {} vec_get_blob_16: total: - instructions: 9821962 + instructions: 7783358 heap_increase: 0 stable_memory_increase: 0 scopes: {} vec_get_blob_32: total: - instructions: 10694461 + instructions: 8410711 heap_increase: 0 stable_memory_increase: 0 scopes: {} vec_get_blob_4: total: - instructions: 5814638 + instructions: 4717509 heap_increase: 0 stable_memory_increase: 0 scopes: {} vec_get_blob_64: total: - instructions: 15512039 + instructions: 12783244 heap_increase: 0 stable_memory_increase: 0 scopes: {} vec_get_blob_8: total: - instructions: 6947467 + instructions: 5434935 heap_increase: 0 stable_memory_increase: 0 scopes: {} vec_get_u64: total: - instructions: 6220307 + instructions: 5240307 heap_increase: 0 stable_memory_increase: 0 scopes: {} diff --git a/src/base_vec.rs b/src/base_vec.rs index 8a3a3b45..cc5b102f 100644 --- a/src/base_vec.rs +++ b/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 BaseVec { } /// Reads the item at the specified index without any bound checks. - fn read_entry_to(&self, index: u64, buf: &mut std::vec::Vec) { + fn read_entry_to(&self, index: u64, buf: &mut Vec) { let offset = DATA_OFFSET + slot_size::() 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. diff --git a/src/btreemap/node.rs b/src/btreemap/node.rs index ff70ccef..98288f64 100644 --- a/src/btreemap/node.rs +++ b/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 Node { 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(&self, offset: Bytes, memory: &M) -> Vec { let reader = NodeReader { address: self.address, @@ -200,8 +200,13 @@ impl Node { }; 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 } diff --git a/src/btreemap/node/io.rs b/src/btreemap/node/io.rs index 45ebec27..f73cb9bc 100644 --- a/src/btreemap/node/io.rs +++ b/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") } diff --git a/src/btreemap/node/v1.rs b/src/btreemap/node/v1.rs index 9bb5b411..2dc02857 100644 --- a/src/btreemap/node/v1.rs +++ b/src/btreemap/node/v1.rs @@ -65,15 +65,14 @@ impl Node { // 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 Node { 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(); diff --git a/src/btreemap/node/v2.rs b/src/btreemap/node/v2.rs index 5f8699ab..f559d777 100644 --- a/src/btreemap/node/v2.rs +++ b/src/btreemap/node/v2.rs @@ -141,6 +141,7 @@ impl Node { 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 Node { }; // 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)))); diff --git a/src/cell.rs b/src/cell.rs index 47b09a37..33f672dc 100644 --- a/src/cell.rs +++ b/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 Cell { /// /// 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)) } diff --git a/src/ic0_memory.rs b/src/ic0_memory.rs index 00078dd0..1e3eeb81 100644 --- a/src/ic0_memory.rs +++ b/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. diff --git a/src/lib.rs b/src/lib.rs index efe8d1b1..05be22a3 100644 --- a/src/lib.rs +++ b/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: &M, addr: Address, dst: &mut std::vec::Vec, 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: &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: &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: &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: &M, addr: Address, val: u64) { write(m, addr.get(), &val.to_le_bytes()); } @@ -148,17 +188,20 @@ fn write(memory: &M, offset: u64, bytes: &[u8]) { } } -// Reads a struct from memory. +/// Reads a struct from memory. fn read_struct(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::()) - }; - memory.read(addr.get(), t_slice); - t + let mut value = MaybeUninit::::uninit(); + unsafe { + memory.read_unsafe( + addr.get(), + value.as_mut_ptr() as *mut u8, + core::mem::size_of::(), + ); + value.assume_init() + } } -// Writes a struct to memory. +/// Writes a struct to memory. fn write_struct(t: &T, addr: Address, memory: &M) { let slice = unsafe { core::slice::from_raw_parts(t as *const _ as *const u8, core::mem::size_of::()) diff --git a/src/log.rs b/src/log.rs index 9a273f0a..f6d1d284 100644 --- a/src/log.rs +++ b/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 Log { /// ignores the result. pub fn read_entry(&self, idx: u64, buf: &mut Vec) -> 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(()) } diff --git a/src/memory_manager.rs b/src/memory_manager.rs index a84075be..fd30182a 100644 --- a/src/memory_manager.rs +++ b/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 MemoryManagerInner { } // 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 MemoryManagerInner { 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() { diff --git a/src/tests.rs b/src/tests.rs index 3d138d68..faa01d33 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -66,3 +66,61 @@ fn should_fail_to_recover_memory_from_memory_manager_if_memory_is_in_use() { let recovered_memory = memory_manager.into_memory(); assert!(recovered_memory.is_none()); } + +#[test] +fn test_read_to_vec_roundtrip() { + let memory = DefaultMemoryImpl::default(); + memory.grow(1); + memory.write(0, &[5, 6, 7, 8, 9]); + + let mut out = vec![]; + read_to_vec(&memory, Address::from(0), &mut out, 5); + assert_eq!(out, vec![5, 6, 7, 8, 9]); +} + +#[test] +fn test_read_write_struct_roundtrip() { + #[derive(Eq, PartialEq, Debug)] + struct Foo { + a: i32, + b: [char; 5], + } + + let foo = Foo { + a: 42, + b: ['a', 'b', 'c', 'd', 'e'], + }; + + let memory = DefaultMemoryImpl::default(); + memory.grow(1); + write_struct(&foo, Address::from(3), &memory); + + assert_eq!( + read_struct::(Address::from(3), &memory), + foo + ) +} + +#[test] +fn test_vector_memory_read() { + let memory = VectorMemory::default(); + memory.grow(1); + memory.write(1, &[4, 6, 8]); + + { + let mut buffer = [0; 3]; + memory.read(1, &mut buffer[..]); + assert_eq!(buffer, [4, 6, 8]); + } + + { + let mut buffer = std::vec::Vec::with_capacity(3); + unsafe { + let ptr = buffer.as_mut_ptr(); + memory.read_unsafe(1, buffer.as_mut_ptr(), 3); + assert_eq!(std::ptr::read(ptr), 4); + assert_eq!(std::ptr::read(ptr.add(1)), 6); + assert_eq!(std::ptr::read(ptr.add(2)), 8); + } + } +} diff --git a/src/vec_mem.rs b/src/vec_mem.rs index 4f1c5919..90b3bd98 100644 --- a/src/vec_mem.rs +++ b/src/vec_mem.rs @@ -39,6 +39,22 @@ impl Memory for RefCell> { dst.copy_from_slice(&self.borrow()[offset as usize..n as usize]); } + unsafe fn read_unsafe(&self, offset: u64, dst: *mut u8, count: usize) { + let n = offset + .checked_add(count as u64) + .expect("read: out of bounds"); + + if n as usize > self.borrow().len() { + panic!("read: out of bounds"); + } + + // SAFETY: + // - we just checked that self is long enough + // - the caller guarantees that there are at least count byte space after dst + // - we are copying bytes so the pointers are automatically aligned + std::ptr::copy(self.borrow().as_ptr().add(offset as usize), dst, count); + } + fn write(&self, offset: u64, src: &[u8]) { let n = offset .checked_add(src.len() as u64) @@ -61,6 +77,9 @@ impl Memory for Rc { fn read(&self, offset: u64, dst: &mut [u8]) { self.deref().read(offset, dst) } + unsafe fn read_unsafe(&self, offset: u64, dst: *mut u8, count: usize) { + self.deref().read_unsafe(offset, dst, count) + } fn write(&self, offset: u64, src: &[u8]) { self.deref().write(offset, src) }