From 43edfe27f464568322168c951a120f98adb279ae Mon Sep 17 00:00:00 2001 From: David Frank Date: Mon, 2 Dec 2024 15:08:44 +0000 Subject: [PATCH] Support read_unsafe in VirtualMemory and RestrictedMemory --- canbench_results.yml | 14 +++---- src/lib.rs | 5 +++ src/memory_manager.rs | 86 ++++++++++++++++++++++++++++++------------- 3 files changed, 73 insertions(+), 32 deletions(-) diff --git a/canbench_results.yml b/canbench_results.yml index 7697b660..577e4c46 100644 --- a/canbench_results.yml +++ b/canbench_results.yml @@ -73,7 +73,7 @@ benches: scopes: {} btreemap_get_blob_512_1024_v2_mem_manager: total: - instructions: 3093928908 + instructions: 2755579551 heap_increase: 0 stable_memory_increase: 0 scopes: { } @@ -139,7 +139,7 @@ benches: scopes: { } btreemap_get_u64_u64_v2_mem_manager: total: - instructions: 619036864 + instructions: 522662446 heap_increase: 0 stable_memory_increase: 0 scopes: {} @@ -223,7 +223,7 @@ benches: scopes: { } btreemap_insert_blob_1024_512_v2_mem_manager: total: - instructions: 6188501297 + instructions: 5594888576 heap_increase: 0 stable_memory_increase: 256 scopes: {} @@ -379,7 +379,7 @@ benches: scopes: {} btreemap_insert_u64_u64_mem_manager: total: - instructions: 937850849 + instructions: 830923183 heap_increase: 0 stable_memory_increase: 0 scopes: { } @@ -631,7 +631,7 @@ benches: scopes: {} memory_manager_overhead: total: - instructions: 1182056741 + instructions: 1182056386 heap_increase: 0 stable_memory_increase: 8320 scopes: {} @@ -661,7 +661,7 @@ benches: scopes: {} vec_get_blob_4_mem_manager: total: - instructions: 12856122 + instructions: 12194073 heap_increase: 0 stable_memory_increase: 0 scopes: { } @@ -673,7 +673,7 @@ benches: scopes: { } vec_get_blob_64_mem_manager: total: - instructions: 22526389 + instructions: 20656466 heap_increase: 0 stable_memory_increase: 0 scopes: {} diff --git a/src/lib.rs b/src/lib.rs index 05be22a3..e1d0e1d2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -270,6 +270,11 @@ impl Memory for RestrictedMemory { .read(self.page_range.start * WASM_PAGE_SIZE + offset, dst) } + unsafe fn read_unsafe(&self, offset: u64, dst: *mut u8, count: usize) { + self.memory + .read_unsafe(self.page_range.start * WASM_PAGE_SIZE + offset, dst, count) + } + fn write(&self, offset: u64, src: &[u8]) { self.memory .write(self.page_range.start * WASM_PAGE_SIZE + offset, src) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index 1aafc32a..8e0396c7 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -171,16 +171,16 @@ struct Header { version: u8, - // The number of buckets allocated by the memory manager. + /// The number of buckets allocated by the memory manager. num_allocated_buckets: u16, - // The size of a bucket in Wasm pages. + /// The size of a bucket in Wasm pages. bucket_size_in_pages: u16, - // Reserved bytes for future extensions + /// Reserved bytes for future extensions _reserved: [u8; HEADER_RESERVED_BYTES], - // The size of each individual memory that can be created by the memory manager. + /// The size of each individual memory that can be created by the memory manager. memory_sizes_in_pages: [u64; MAX_NUM_MEMORIES as usize], } @@ -209,6 +209,12 @@ impl Memory for VirtualMemory { self.memory_manager.borrow().read(self.id, offset, dst) } + unsafe fn read_unsafe(&self, offset: u64, dst: *mut u8, count: usize) { + self.memory_manager + .borrow() + .read_unsafe(self.id, offset, dst, count) + } + fn write(&self, offset: u64, src: &[u8]) { self.memory_manager.borrow().write(self.id, offset, src) } @@ -218,15 +224,15 @@ impl Memory for VirtualMemory { struct MemoryManagerInner { memory: M, - // The number of buckets that have been allocated. + /// The number of buckets that have been allocated. allocated_buckets: u16, bucket_size_in_pages: u16, - // An array storing the size (in pages) of each of the managed memories. + /// An array storing the size (in pages) of each of the managed memories. memory_sizes_in_pages: [u64; MAX_NUM_MEMORIES as usize], - // A map mapping each managed memory to the bucket ids that are allocated to it. + /// A map mapping each managed memory to the bucket ids that are allocated to it. memory_buckets: Vec>, } @@ -313,12 +319,12 @@ impl MemoryManagerInner { write_struct(&header, Address::from(0), &self.memory); } - // Returns the size of a memory (in pages). + /// Returns the size of a memory (in pages). fn memory_size(&self, id: MemoryId) -> u64 { self.memory_sizes_in_pages[id.0 as usize] } - // Grows the memory with the given id by the given number of pages. + /// Grows the memory with the given id by the given number of pages. fn grow(&mut self, id: MemoryId, pages: u64) -> i64 { // Compute how many additional buckets are needed. let old_size = self.memory_size(id); @@ -384,23 +390,35 @@ impl MemoryManagerInner { } } + #[inline] fn read(&self, id: MemoryId, offset: u64, dst: &mut [u8]) { - if (offset + dst.len() as u64) > self.memory_size(id) * WASM_PAGE_SIZE { + // SAFETY: this is trivially safe because dst has dst.len() space. + unsafe { self.read_unsafe(id, offset, dst.as_mut_ptr(), dst.len()) } + } + + /// # 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`. + unsafe fn read_unsafe(&self, id: MemoryId, offset: u64, dst: *mut u8, count: usize) { + if (offset + count as u64) > self.memory_size(id) * WASM_PAGE_SIZE { panic!("{id:?}: read out of bounds"); } - let mut bytes_read = 0; - for Segment { address, length } in self.bucket_iter(id, offset, dst.len()) { - self.memory.read( - address.get(), - &mut dst[bytes_read as usize..(bytes_read + length.get()) as usize], - ); + let mut bytes_read: usize = 0; + for Segment { address, length } in self.bucket_iter(id, offset, count) { + let length = length.get().try_into().expect("Length overflows usize"); + self.memory + .read_unsafe(address.get(), dst.add(bytes_read), length); - bytes_read += length.get(); + bytes_read = bytes_read + .checked_add(length) + .expect("Bytes read overflowed usize"); } } - // Initializes a [`BucketIterator`]. + /// Initializes a [`BucketIterator`]. fn bucket_iter(&self, MemoryId(id): MemoryId, offset: u64, length: usize) -> BucketIterator { // Get the buckets allocated to the given memory id. let buckets = self.memory_buckets[id as usize].as_slice(); @@ -419,13 +437,13 @@ impl MemoryManagerInner { Bytes::from(self.bucket_size_in_pages as u64 * WASM_PAGE_SIZE) } - // Returns the number of buckets needed to accommodate the given number of pages. + /// Returns the number of buckets needed to accommodate the given number of pages. fn num_buckets_needed(&self, num_pages: u64) -> u64 { // Ceiling division. (num_pages + self.bucket_size_in_pages as u64 - 1) / self.bucket_size_in_pages as u64 } - // Returns the underlying memory. + /// Returns the underlying memory. pub fn into_memory(self) -> M { self.memory } @@ -509,7 +527,7 @@ impl Iterator for BucketIterator<'_> { } impl<'a> BucketIterator<'a> { - // Returns the address of a given bucket. + /// Returns the address of a given bucket. fn bucket_address(&self, id: BucketId) -> Address { Address::from(BUCKETS_OFFSET_IN_BYTES) + self.bucket_size_in_bytes * Bytes::from(id.0) } @@ -815,10 +833,28 @@ mod test { // Write a random blob into the memory, growing the memory as it needs to. write(memory, offset, &data); - // Verify the blob can be read back. - let mut bytes = vec![0; data.len()]; - memory.read(offset, &mut bytes); - assert_eq!(bytes, data); + { + // Verify the blob can be read back using read. + let mut bytes = vec![0; data.len()]; + memory.read(offset, &mut bytes); + assert_eq!(bytes, data); + } + + { + // Verify the blob can be read back using read_to_vec. + let mut bytes = vec![]; + read_to_vec(memory, offset.into(), &mut bytes, data.len()); + assert_eq!(bytes, data); + } + + { + // Verify the blob can be read back using read_unsafe. + let mut bytes = vec![0; data.len()]; + unsafe { + memory.read_unsafe(offset, bytes.as_mut_ptr(), data.len()); + } + assert_eq!(bytes, data); + } } }); }