From 0ba0064931e6fa4d511086e1b551e71a4a0720a6 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Fri, 8 Nov 2024 19:24:02 +0000 Subject: [PATCH 01/20] Free virtual memory WIP --- src/memory_manager.rs | 166 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 153 insertions(+), 13 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index 34f6efc4..ca86b080 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -48,7 +48,9 @@ use crate::{ write, write_struct, Memory, WASM_PAGE_SIZE, }; use std::cmp::min; +use std::collections::btree_map::Entry::Vacant; use std::collections::BTreeMap; +use std::mem::size_of; use std::rc::Rc; use std::{cell::RefCell, collections::BTreeSet}; @@ -186,6 +188,54 @@ const BUCKET_ID_LEN_IN_BITS: usize = 15; /// -------------------------------------------------- <- Page ((MAX_NUM_BUCKETS - 1) * N + 1) /// Bucket MAX_NUM_BUCKETS ↕ N pages /// ``` +/// # V2.1 layout +/// +/// ```text +/// -------------------------------------------------- <- Address 0 +/// Magic "MGR" ↕ 3 bytes +/// -------------------------------------------------- +/// Layout version ↕ 1 byte +/// -------------------------------------------------- +/// Number of allocated buckets ↕ 2 bytes +/// -------------------------------------------------- +/// Bucket size (in pages) = N ↕ 2 bytes +/// -------------------------------------------------- +/// Reserved space ↕ 32 bytes +/// -------------------------------------------------- +/// Size of memory 0 (in pages) ↕ 8 bytes +/// -------------------------------------------------- +/// Size of memory 1 (in pages) ↕ 8 bytes +/// -------------------------------------------------- +/// ... +/// -------------------------------------------------- +/// Size of memory 254 (in pages) ↕ 8 bytes +/// -------------------------------------------------- <- IDs of buckets +/// Bucket 1 belonging to memory 0 ↕ 2 bytes +/// -------------------------------------------------- +/// Bucket 1 belonging to memory 1 ↕ 2 bytes +/// -------------------------------------------------- +/// ... +/// -------------------------------------------------- +/// Bucket 1 belonging to memory 254 ↕ 2 bytes +/// -------------------------------------------------- +/// Next bucket in linked list after bucket 1 ↕ 15 bits +/// -------------------------------------------------- +/// Next bucket in linked list after bucket 2 ↕ 15 bits +/// -------------------------------------------------- +/// ... +/// --------------------------------------------------- +/// Next bucket in linked list after bucket MAX_NUM_BUCKETS ↕ 15 bits +/// -------------------------------------------------- +/// Unallocated space ↕ 1'506 bytes +/// -------------------------------------------------- <- Buckets (Page 1) +/// Bucket 1 ↕ N pages +/// -------------------------------------------------- <- Page N + 1 +/// Bucket 2 ↕ N pages +/// -------------------------------------------------- +/// ... +/// -------------------------------------------------- <- Page ((MAX_NUM_BUCKETS - 1) * N + 1) +/// Bucket MAX_NUM_BUCKETS ↕ N pages +/// ``` pub struct MemoryManager { inner: Rc>>, } @@ -258,6 +308,9 @@ struct Header { // The size of each individual memory that can be created by the memory manager. memory_sizes_in_pages: [u64; MAX_NUM_MEMORIES as usize], + + // The first bucket assigned to each memory. + first_bucket_per_memory: [u16; MAX_NUM_MEMORIES as usize], } impl Header { @@ -308,6 +361,8 @@ struct MemoryManagerInner { // Tracks the buckets that were freed to be reused in future calls to `grow`. // NOTE: A BTreeSet is used so that bucket IDs are maintained in sorted order. freed_buckets: BTreeSet, + + bucket_bits: BucketBits, } impl MemoryManagerInner { @@ -337,6 +392,7 @@ impl MemoryManagerInner { memory_buckets: BTreeMap::new(), bucket_size_in_pages, freed_buckets: BTreeSet::new(), + bucket_bits: BucketBits::default(), }; mem_mgr.save_header(); @@ -372,6 +428,7 @@ impl MemoryManagerInner { memory_buckets: BTreeMap::new(), bucket_size_in_pages, freed_buckets: BTreeSet::new(), + bucket_bits: BucketBits::default(), }; mem_mgr.save_header_v1(); @@ -400,6 +457,13 @@ impl MemoryManagerInner { } } + let mut bucket_bits = BucketBits::default(); + for (_, buckets) in memory_buckets.iter() { + for (bucket, next_bucket) in buckets.iter().zip(buckets.iter().skip(1)) { + bucket_bits.set_next(*bucket, *next_bucket); + } + } + Self { memory, allocated_buckets: header.num_allocated_buckets, @@ -407,6 +471,7 @@ impl MemoryManagerInner { memory_sizes_in_pages: header.memory_sizes_in_pages, memory_buckets, freed_buckets: BTreeSet::new(), + bucket_bits, } } @@ -414,11 +479,27 @@ impl MemoryManagerInner { let mut memory_size_in_buckets = vec![]; let mut number_of_used_buckets = 0; + // Map of all memories with their assigned buckets. + let mut memory_buckets = BTreeMap::new(); + + let bucket_bits = BucketBits::load(&memory); + // Translate memory sizes expressed in pages to sizes expressed in buckets. - for memory_size_in_pages in header.memory_sizes_in_pages.into_iter() { + for (index, memory_size_in_pages) in header.memory_sizes_in_pages.into_iter().enumerate() { + let memory_id = MemoryId(index as u8); let size_in_buckets = memory_size_in_pages.div_ceil(header.bucket_size_in_pages as u64); memory_size_in_buckets.push(size_in_buckets); number_of_used_buckets += size_in_buckets; + + if size_in_buckets > 0 { + let mut bucket = BucketId(header.first_bucket_per_memory[index]); + let mut buckets = vec![bucket]; + for _ in 1..size_in_buckets { + bucket = bucket_bits.get_next(bucket); + buckets.push(bucket); + } + memory_buckets.insert(memory_id, buckets); + } } // Load the buckets. @@ -434,9 +515,6 @@ impl MemoryManagerInner { bytes_to_bucket_indexes(&buckets) }; - // Map of all memories with their assigned buckets. - let mut memory_buckets = BTreeMap::new(); - // The last bucket that's accessed. let mut max_bucket_id: u16 = 0; @@ -470,6 +548,7 @@ impl MemoryManagerInner { memory_sizes_in_pages: header.memory_sizes_in_pages, memory_buckets, freed_buckets, + bucket_bits, } } @@ -485,6 +564,13 @@ impl MemoryManagerInner { } fn save_header(&self) { + let mut first_bucket_per_memory = [0; MAX_NUM_MEMORIES as usize]; + for (memory_id, buckets) in self.memory_buckets.iter() { + if let Some(first_bucket) = buckets.first() { + first_bucket_per_memory[memory_id.0 as usize] = first_bucket.0; + } + } + let header = Header { magic: *MAGIC, version: LAYOUT_VERSION_V2, @@ -492,6 +578,7 @@ impl MemoryManagerInner { bucket_size_in_pages: self.bucket_size_in_pages, _reserved: [0; HEADER_RESERVED_BYTES], memory_sizes_in_pages: self.memory_sizes_in_pages, + first_bucket_per_memory, }; write_struct(&header, Address::from(0), &self.memory); @@ -529,10 +616,11 @@ impl MemoryManagerInner { } }; - self.memory_buckets - .entry(id) - .or_default() - .push(new_bucket_id); + let buckets = self.memory_buckets.entry(id).or_default(); + if let Some(last_bucket) = buckets.last() { + self.bucket_bits.set_next(*last_bucket, new_bucket_id); + } + buckets.push(new_bucket_id); } // Grow the underlying memory if necessary. @@ -553,11 +641,7 @@ impl MemoryManagerInner { self.save_header(); // Write in stable store that this bucket belongs to the memory with the provided `id`. - write( - &self.memory, - bucket_indexes_offset().get(), - self.get_bucket_ids_in_bytes().as_ref(), - ); + self.bucket_bits.save(&self.memory); // Return the old size. old_size as i64 @@ -664,6 +748,13 @@ impl MemoryManagerInner { #[cfg(test)] fn save_header_v1(&self) { + let mut first_bucket_per_memory = [0; MAX_NUM_MEMORIES as usize]; + for (memory_id, buckets) in self.memory_buckets.iter() { + if let Some(first_bucket) = buckets.first() { + first_bucket_per_memory[memory_id.0 as usize] = first_bucket.0; + } + } + let header = Header { magic: *MAGIC, version: LAYOUT_VERSION_V1, @@ -671,6 +762,7 @@ impl MemoryManagerInner { bucket_size_in_pages: self.bucket_size_in_pages, _reserved: [0; HEADER_RESERVED_BYTES], memory_sizes_in_pages: self.memory_sizes_in_pages, + first_bucket_per_memory, }; write_struct(&header, Address::from(0), &self.memory); @@ -803,6 +895,54 @@ fn bucket_indexes_offset() -> Address { Address::from(0) + Header::size() } +#[derive(Clone)] +struct BucketBits { + bits: [u8; (15 * MAX_NUM_BUCKETS / 8) as usize], + dirty_bytes: Vec, +} + +impl BucketBits { + fn load(memory: &M) -> Self { + let mut bucket_bits = BucketBits::default(); + let offset = size_of::
() as u64; + memory.read(offset, &mut bucket_bits.bits); + bucket_bits + } + + fn get_next(&self, bucket: BucketId) -> BucketId { + let start_bit_index = bucket.0 * 15; + let end_bit_index = start_bit_index + 14; + let start_byte_index = start_bit_index / 8; + let end_byte_index = end_bit_index / 8; + + let mut result = [0u8; 2]; + todo!() + } + + fn set_next(&mut self, bucket: BucketId, next: BucketId) { + todo!() + } + + fn save(&self, memory: &M) { + if !self.dirty_bytes.is_empty() { + let min = *self.dirty_bytes.iter().min().unwrap() as usize; + let max = *self.dirty_bytes.iter().max().unwrap() as usize; + let offset = (size_of::
() + min) as u64; + let segment = &self.bits[min..=max]; + memory.write(offset, segment); + } + } +} + +impl Default for BucketBits { + fn default() -> Self { + BucketBits { + bits: [0; (15 * MAX_NUM_BUCKETS / 8) as usize], + dirty_bytes: Vec::new(), + } + } +} + #[cfg(test)] mod test { use super::*; From fcec6d3d33f89d96bd6f877d42e7ec841df42940 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Fri, 8 Nov 2024 23:11:19 +0000 Subject: [PATCH 02/20] Getting there --- Cargo.lock | 41 ++++++++++- Cargo.toml | 2 +- src/memory_manager.rs | 163 +++++++++++------------------------------- 3 files changed, 84 insertions(+), 122 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a9bb11e8..fb655646 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -84,6 +84,18 @@ version = "2.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed570934406eb16438a4e976b1b4500774099c13b8cb96eec99f620f05090ddf" +[[package]] +name = "bitvec" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc2832c24239b0141d5674bb9174f9d68a8b5b3f2753311927c172ca46f7e9c" +dependencies = [ + "funty", + "radium", + "tap", + "wyz", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -244,6 +256,12 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "funty" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" + [[package]] name = "generic-array" version = "0.14.7" @@ -302,7 +320,7 @@ dependencies = [ name = "ic-stable-structures" version = "0.6.5" dependencies = [ - "bit-vec", + "bitvec", "canbench-rs", "candid", "hex", @@ -478,6 +496,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "radium" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" + [[package]] name = "rand" version = "0.8.5" @@ -672,6 +696,12 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "tap" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" + [[package]] name = "tempfile" version = "3.9.0" @@ -861,3 +891,12 @@ name = "windows_x86_64_msvc" version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dff9641d1cd4be8d1a070daf9e3773c5f67e78b4d9d42263020c057706765c04" + +[[package]] +name = "wyz" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05f360fc0b24296329c78fda852a1e9ae82de9cf7b27dae4b7f62f118f77b9ed" +dependencies = [ + "tap", +] diff --git a/Cargo.toml b/Cargo.toml index 441ea38f..0dfc7d5e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ include = ["src", "Cargo.toml", "LICENSE", "README.md"] repository = "https://github.com/dfinity/stable-structures" [dependencies] -bit-vec = "0.6" +bitvec = "1.0.1" ic_principal = { version = "0.1.1", default-features = false } # An optional dependency to benchmark parts of the code. canbench-rs = { version = "0.1.7", optional = true } diff --git a/src/memory_manager.rs b/src/memory_manager.rs index ca86b080..0faed97b 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -40,15 +40,11 @@ //! memory_1.read(0, &mut bytes); //! assert_eq!(bytes, vec![4, 5, 6]); //! ``` -use bit_vec::BitVec; -use crate::{ - read_struct, - types::{Address, Bytes}, - write, write_struct, Memory, WASM_PAGE_SIZE, -}; +use crate::{read_struct, types::{Address, Bytes}, write, write_struct, Memory, WASM_PAGE_SIZE}; +use bitvec::array::BitArray; +use bitvec::macros::internal::funty::Fundamental; use std::cmp::min; -use std::collections::btree_map::Entry::Vacant; use std::collections::BTreeMap; use std::mem::size_of; use std::rc::Rc; @@ -77,9 +73,6 @@ const BUCKETS_OFFSET_IN_BYTES: u64 = BUCKETS_OFFSET_IN_PAGES * WASM_PAGE_SIZE; // Reserved bytes in the header for future extensions. const HEADER_RESERVED_BYTES: usize = 32; -// Size of the bucket ID in the header. -const BUCKET_ID_LEN_IN_BITS: usize = 15; - /// A memory manager simulates multiple memories within a single memory. /// /// The memory manager can return up to 255 unique instances of [`VirtualMemory`], and each can be @@ -477,11 +470,14 @@ impl MemoryManagerInner { fn load_layout_v2(memory: M, header: Header) -> Self { let mut memory_size_in_buckets = vec![]; - let mut number_of_used_buckets = 0; // Map of all memories with their assigned buckets. let mut memory_buckets = BTreeMap::new(); + // Set of all buckets with ID smaller than 'max_bucket_id' which were allocated and freed. + let mut freed_buckets: BTreeSet = + (0..header.num_allocated_buckets).map(BucketId).collect(); + let bucket_bits = BucketBits::load(&memory); // Translate memory sizes expressed in pages to sizes expressed in buckets. @@ -489,58 +485,20 @@ impl MemoryManagerInner { let memory_id = MemoryId(index as u8); let size_in_buckets = memory_size_in_pages.div_ceil(header.bucket_size_in_pages as u64); memory_size_in_buckets.push(size_in_buckets); - number_of_used_buckets += size_in_buckets; if size_in_buckets > 0 { let mut bucket = BucketId(header.first_bucket_per_memory[index]); let mut buckets = vec![bucket]; + freed_buckets.remove(&bucket); for _ in 1..size_in_buckets { bucket = bucket_bits.get_next(bucket); + freed_buckets.remove(&bucket); buckets.push(bucket); } memory_buckets.insert(memory_id, buckets); } } - // Load the buckets. - let buckets = { - const BYTE_SIZE_IN_BITS: usize = 8; - let buckets_index_size_in_bytes: usize = (number_of_used_buckets as usize - * BUCKET_ID_LEN_IN_BITS) - .div_ceil(BYTE_SIZE_IN_BITS); - - let mut buckets = vec![0; buckets_index_size_in_bytes]; - memory.read(bucket_indexes_offset().get(), &mut buckets); - - bytes_to_bucket_indexes(&buckets) - }; - - // The last bucket that's accessed. - let mut max_bucket_id: u16 = 0; - - let mut bucket_idx: usize = 0; - - // Assign buckets to the memories they are part of. - for (memory, size_in_buckets) in memory_size_in_buckets.into_iter().enumerate() { - let mut vec_buckets = vec![]; - for _ in 0..size_in_buckets { - let bucket = buckets[bucket_idx]; - max_bucket_id = std::cmp::max(bucket.0, max_bucket_id); - vec_buckets.push(bucket); - bucket_idx += 1; - } - memory_buckets - .entry(MemoryId(memory as u8)) - .or_insert(vec_buckets); - } - - // Set of all buckets with ID smaller than 'max_bucket_id' which were allocated and freed. - let mut freed_buckets: BTreeSet = (0..max_bucket_id).map(BucketId).collect(); - - for id in buckets.iter() { - freed_buckets.remove(id); - } - Self { memory, allocated_buckets: header.num_allocated_buckets, @@ -647,21 +605,6 @@ impl MemoryManagerInner { old_size as i64 } - fn get_bucket_ids_in_bytes(&self) -> Vec { - let mut bit_vec = BitVec::new(); - for memory in self.memory_buckets.iter() { - for bucket in memory.1 { - let bucket_ind = bucket.0; - // Splits bit_vec returning the slice [1, .., bit_vec.size() - 1]. - // This is precisely what we need since the BucketId can be represented - // using only 15 bits, instead of 16. - let mut bit_vec_temp = BitVec::from_bytes(&bucket_ind.to_be_bytes()).split_off(1); - bit_vec.append(&mut bit_vec_temp); - } - } - bit_vec.to_bytes() - } - fn write(&self, id: MemoryId, offset: u64, src: &[u8]) { if (offset + src.len() as u64) > self.memory_size(id) * WASM_PAGE_SIZE { panic!("{id:?}: write out of bounds"); @@ -732,13 +675,6 @@ impl MemoryManagerInner { } // Update the header. self.save_header(); - - // Write in stable store that no bucket belongs to the memory with the provided `id`. - write( - &self.memory, - bucket_indexes_offset().get(), - self.get_bucket_ids_in_bytes().as_ref(), - ); } // Returns the underlying memory. @@ -769,23 +705,6 @@ impl MemoryManagerInner { } } -fn bytes_to_bucket_indexes(input: &[u8]) -> Vec { - let mut bucket_ids = vec![]; - let bit_vec = BitVec::from_bytes(input); - for bucket_order_number in 0..bit_vec.len() / BUCKET_ID_LEN_IN_BITS { - let mut bucket_id: u16 = 0; - for bucket_id_bit in 0..BUCKET_ID_LEN_IN_BITS { - let next_bit = BUCKET_ID_LEN_IN_BITS * bucket_order_number + bucket_id_bit; - bucket_id <<= 1; - if bit_vec.get(next_bit) == Some(true) { - bucket_id |= 1; - } - } - bucket_ids.push(BucketId(bucket_id)); - } - bucket_ids -} - struct Segment { address: Address, length: Bytes, @@ -891,36 +810,53 @@ fn bucket_allocations_address(id: BucketId) -> Address { Address::from(0) + Header::size() + Bytes::from(id.0) } -fn bucket_indexes_offset() -> Address { - Address::from(0) + Header::size() -} +const BUCKET_BITS_LEN: usize = 15 * MAX_NUM_BUCKETS as usize; -#[derive(Clone)] +#[derive(Clone, Default)] struct BucketBits { - bits: [u8; (15 * MAX_NUM_BUCKETS / 8) as usize], + bits: BitArray<[u8; BUCKET_BITS_LEN / 8]>, dirty_bytes: Vec, } impl BucketBits { fn load(memory: &M) -> Self { - let mut bucket_bits = BucketBits::default(); + let mut bytes = [0; BUCKET_BITS_LEN / 8]; let offset = size_of::
() as u64; - memory.read(offset, &mut bucket_bits.bits); - bucket_bits + memory.read(offset, &mut bytes); + + BucketBits { + bits: BitArray::from(bytes), + ..Default::default() + } } fn get_next(&self, bucket: BucketId) -> BucketId { - let start_bit_index = bucket.0 * 15; - let end_bit_index = start_bit_index + 14; - let start_byte_index = start_bit_index / 8; - let end_byte_index = end_bit_index / 8; + let start_bit_index = (bucket.0 * 15) as usize; + let mut next_bits: BitArray<[u8; 2]> = BitArray::new([0u8; 2]); + + for i in 0..15 { + next_bits.set(i + 1, self.bits.get(start_bit_index + i).unwrap().as_bool()); + } - let mut result = [0u8; 2]; - todo!() + BucketId(u16::from_be_bytes(next_bits.data)) } fn set_next(&mut self, bucket: BucketId, next: BucketId) { - todo!() + let start_bit_index = (bucket.0 * 15) as usize; + let next_bits: BitArray<[u8; 2]> = BitArray::from(next.0.to_be_bytes()); + + for (index, bit) in next_bits.iter().skip(1).enumerate() { + self.bits.set(start_bit_index + index, bit.as_bool()); + } + + let start_byte_index = start_bit_index / 8; + let end_byte_index = (start_bit_index + 14) / 8; + + for index in (start_byte_index..=end_byte_index).map(|i| i as u16) { + if !self.dirty_bytes.contains(&index) { + self.dirty_bytes.push(index); + } + } } fn save(&self, memory: &M) { @@ -928,21 +864,12 @@ impl BucketBits { let min = *self.dirty_bytes.iter().min().unwrap() as usize; let max = *self.dirty_bytes.iter().max().unwrap() as usize; let offset = (size_of::
() + min) as u64; - let segment = &self.bits[min..=max]; + let segment = &self.bits.data[min..=max]; memory.write(offset, segment); } } } -impl Default for BucketBits { - fn default() -> Self { - BucketBits { - bits: [0; (15 * MAX_NUM_BUCKETS / 8) as usize], - dirty_bytes: Vec::new(), - } - } -} - #[cfg(test)] mod test { use super::*; @@ -1488,13 +1415,9 @@ mod test { mem_mgr.free(MemoryId(4)); let mem_mgr = MemoryManager::init(mem); - // Only Memory 0 and 2 buckets should be counted as freed_buckets. - // The bucket belonging to Memory 4 should not be counted as - // freed because it has the biggest Bucket ID of all allocated - // buckets hence it should become part of unallocated buckets. assert_eq!( mem_mgr.inner.borrow().freed_buckets, - maplit::btreeset! { BucketId(0), BucketId(2) } + maplit::btreeset! { BucketId(0), BucketId(2), BucketId(4) } ); } From 772c24adaaeafd918de358193f6e0b673658269a Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Sat, 9 Nov 2024 00:15:31 +0000 Subject: [PATCH 03/20] Move new field out of header --- src/memory_manager.rs | 102 +++++++++++++++++------------------------- 1 file changed, 42 insertions(+), 60 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index 0faed97b..851d895b 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, types::{Address, Bytes}, write, write_struct, Memory, WASM_PAGE_SIZE}; +use crate::{read_struct, types::{Address, Bytes}, write_struct, Memory, WASM_PAGE_SIZE}; use bitvec::array::BitArray; use bitvec::macros::internal::funty::Fundamental; use std::cmp::min; @@ -70,6 +70,8 @@ const UNALLOCATED_BUCKET_MARKER: u8 = MAX_NUM_MEMORIES; const BUCKETS_OFFSET_IN_PAGES: u64 = 1; const BUCKETS_OFFSET_IN_BYTES: u64 = BUCKETS_OFFSET_IN_PAGES * WASM_PAGE_SIZE; +const BUCKET_BITS_OFFSET: u64 = size_of::
() as u64; + // Reserved bytes in the header for future extensions. const HEADER_RESERVED_BYTES: usize = 32; @@ -301,9 +303,6 @@ struct Header { // The size of each individual memory that can be created by the memory manager. memory_sizes_in_pages: [u64; MAX_NUM_MEMORIES as usize], - - // The first bucket assigned to each memory. - first_bucket_per_memory: [u16; MAX_NUM_MEMORIES as usize], } impl Header { @@ -397,7 +396,7 @@ impl MemoryManagerInner { fn init_v1(memory: M, bucket_size_in_pages: u16) -> Self { if memory.size() == 0 { // Memory is empty. Create a new map. - return Self::new(memory, bucket_size_in_pages); + return Self::new_v1(memory, bucket_size_in_pages); } // Check if the magic in the memory corresponds to this object. @@ -427,7 +426,7 @@ impl MemoryManagerInner { mem_mgr.save_header_v1(); // Mark all the buckets as unallocated. - write( + crate::write( &mem_mgr.memory, bucket_allocations_address(BucketId(0)).get(), &[UNALLOCATED_BUCKET_MARKER; MAX_NUM_BUCKETS as usize], @@ -478,7 +477,7 @@ impl MemoryManagerInner { let mut freed_buckets: BTreeSet = (0..header.num_allocated_buckets).map(BucketId).collect(); - let bucket_bits = BucketBits::load(&memory); + let bucket_bits: BucketBits = read_struct(Address::from(BUCKET_BITS_OFFSET), &memory); // Translate memory sizes expressed in pages to sizes expressed in buckets. for (index, memory_size_in_pages) in header.memory_sizes_in_pages.into_iter().enumerate() { @@ -487,7 +486,7 @@ impl MemoryManagerInner { memory_size_in_buckets.push(size_in_buckets); if size_in_buckets > 0 { - let mut bucket = BucketId(header.first_bucket_per_memory[index]); + let mut bucket = bucket_bits.get_first(memory_id); let mut buckets = vec![bucket]; freed_buckets.remove(&bucket); for _ in 1..size_in_buckets { @@ -522,13 +521,6 @@ impl MemoryManagerInner { } fn save_header(&self) { - let mut first_bucket_per_memory = [0; MAX_NUM_MEMORIES as usize]; - for (memory_id, buckets) in self.memory_buckets.iter() { - if let Some(first_bucket) = buckets.first() { - first_bucket_per_memory[memory_id.0 as usize] = first_bucket.0; - } - } - let header = Header { magic: *MAGIC, version: LAYOUT_VERSION_V2, @@ -536,7 +528,6 @@ impl MemoryManagerInner { bucket_size_in_pages: self.bucket_size_in_pages, _reserved: [0; HEADER_RESERVED_BYTES], memory_sizes_in_pages: self.memory_sizes_in_pages, - first_bucket_per_memory, }; write_struct(&header, Address::from(0), &self.memory); @@ -577,6 +568,8 @@ impl MemoryManagerInner { let buckets = self.memory_buckets.entry(id).or_default(); if let Some(last_bucket) = buckets.last() { self.bucket_bits.set_next(*last_bucket, new_bucket_id); + } else { + self.bucket_bits.set_first(id, new_bucket_id); } buckets.push(new_bucket_id); } @@ -599,7 +592,7 @@ impl MemoryManagerInner { self.save_header(); // Write in stable store that this bucket belongs to the memory with the provided `id`. - self.bucket_bits.save(&self.memory); + write_struct(&self.bucket_bits, Address::from(BUCKET_BITS_OFFSET), &self.memory); // Return the old size. old_size as i64 @@ -684,13 +677,6 @@ impl MemoryManagerInner { #[cfg(test)] fn save_header_v1(&self) { - let mut first_bucket_per_memory = [0; MAX_NUM_MEMORIES as usize]; - for (memory_id, buckets) in self.memory_buckets.iter() { - if let Some(first_bucket) = buckets.first() { - first_bucket_per_memory[memory_id.0 as usize] = first_bucket.0; - } - } - let header = Header { magic: *MAGIC, version: LAYOUT_VERSION_V1, @@ -698,7 +684,6 @@ impl MemoryManagerInner { bucket_size_in_pages: self.bucket_size_in_pages, _reserved: [0; HEADER_RESERVED_BYTES], memory_sizes_in_pages: self.memory_sizes_in_pages, - first_bucket_per_memory, }; write_struct(&header, Address::from(0), &self.memory); @@ -812,22 +797,20 @@ fn bucket_allocations_address(id: BucketId) -> Address { const BUCKET_BITS_LEN: usize = 15 * MAX_NUM_BUCKETS as usize; -#[derive(Clone, Default)] +#[repr(C, packed)] +#[derive(Clone)] struct BucketBits { - bits: BitArray<[u8; BUCKET_BITS_LEN / 8]>, - dirty_bytes: Vec, + first_bucket_per_memory: [u16; MAX_NUM_MEMORIES as usize], + bucket_links: BitArray<[u8; BUCKET_BITS_LEN / 8]>, } impl BucketBits { - fn load(memory: &M) -> Self { - let mut bytes = [0; BUCKET_BITS_LEN / 8]; - let offset = size_of::
() as u64; - memory.read(offset, &mut bytes); + fn get_first(&self, memory_id: MemoryId) -> BucketId { + BucketId(self.first_bucket_per_memory[memory_id.0 as usize]) + } - BucketBits { - bits: BitArray::from(bytes), - ..Default::default() - } + fn set_first(&mut self, memory_id: MemoryId, value: BucketId) { + self.first_bucket_per_memory[memory_id.0 as usize] = value.0; } fn get_next(&self, bucket: BucketId) -> BucketId { @@ -835,7 +818,7 @@ impl BucketBits { let mut next_bits: BitArray<[u8; 2]> = BitArray::new([0u8; 2]); for i in 0..15 { - next_bits.set(i + 1, self.bits.get(start_bit_index + i).unwrap().as_bool()); + next_bits.set(i + 1, self.bucket_links.get(start_bit_index + i).unwrap().as_bool()); } BucketId(u16::from_be_bytes(next_bits.data)) @@ -846,26 +829,25 @@ impl BucketBits { let next_bits: BitArray<[u8; 2]> = BitArray::from(next.0.to_be_bytes()); for (index, bit) in next_bits.iter().skip(1).enumerate() { - self.bits.set(start_bit_index + index, bit.as_bool()); + self.bucket_links.set(start_bit_index + index, bit.as_bool()); } - let start_byte_index = start_bit_index / 8; - let end_byte_index = (start_bit_index + 14) / 8; - - for index in (start_byte_index..=end_byte_index).map(|i| i as u16) { - if !self.dirty_bytes.contains(&index) { - self.dirty_bytes.push(index); - } - } + // let start_byte_index = start_bit_index / 8; + // let end_byte_index = (start_bit_index + 14) / 8; + // + // for index in (start_byte_index..=end_byte_index).map(|i| i as u16) { + // if !self.dirty_bytes.contains(&index) { + // self.dirty_bytes.push(index); + // } + // } } +} - fn save(&self, memory: &M) { - if !self.dirty_bytes.is_empty() { - let min = *self.dirty_bytes.iter().min().unwrap() as usize; - let max = *self.dirty_bytes.iter().max().unwrap() as usize; - let offset = (size_of::
() + min) as u64; - let segment = &self.bits.data[min..=max]; - memory.write(offset, segment); +impl Default for BucketBits { + fn default() -> Self { + BucketBits { + first_bucket_per_memory: [0; MAX_NUM_MEMORIES as usize], + bucket_links: BitArray::default(), } } } @@ -1148,7 +1130,7 @@ mod test { )| { for memory in memories.iter().take(num_memories) { // Write a random blob into the memory, growing the memory as it needs to. - write(memory, offset, &data); + crate::write(memory, offset, &data); // Verify the blob can be read back. let mut bytes = vec![0; data.len()]; @@ -1403,11 +1385,11 @@ mod test { #[test] fn freed_memories_are_tracked() { let mem = make_memory(); - let mut mem_mgr = MemoryManager::init(mem.clone()); + let mut mem_mgr = MemoryManager::init_with_bucket_size(mem.clone(), 1); mem_mgr.get(MemoryId(0)).grow(1); - mem_mgr.get(MemoryId(1)).grow(1); - mem_mgr.get(MemoryId(2)).grow(1); - mem_mgr.get(MemoryId(3)).grow(1); + mem_mgr.get(MemoryId(1)).grow(2); + mem_mgr.get(MemoryId(2)).grow(3); + mem_mgr.get(MemoryId(3)).grow(2); mem_mgr.get(MemoryId(4)).grow(1); mem_mgr.free(MemoryId(0)); @@ -1417,7 +1399,7 @@ mod test { let mem_mgr = MemoryManager::init(mem); assert_eq!( mem_mgr.inner.borrow().freed_buckets, - maplit::btreeset! { BucketId(0), BucketId(2), BucketId(4) } + maplit::btreeset! { BucketId(0), BucketId(3), BucketId(4), BucketId(5), BucketId(8) } ); } @@ -1438,7 +1420,7 @@ mod test { )| { for memory_v1 in memories_v1.iter().take(num_memories) { // Write a random blob into the memory, growing the memory as it needs to. - write(memory_v1, offset, &data); + crate::write(memory_v1, offset, &data); } // Load layout v1 and convert it to layout v2. From 6516ed066aae271449a222da077d260a3c7cc06d Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 11 Nov 2024 09:14:23 +0000 Subject: [PATCH 04/20] Clean it up a bit --- src/memory_manager.rs | 35 +++++++++++++++++++++++------------ 1 file changed, 23 insertions(+), 12 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index 851d895b..b93ddb97 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -473,10 +473,6 @@ impl MemoryManagerInner { // Map of all memories with their assigned buckets. let mut memory_buckets = BTreeMap::new(); - // Set of all buckets with ID smaller than 'max_bucket_id' which were allocated and freed. - let mut freed_buckets: BTreeSet = - (0..header.num_allocated_buckets).map(BucketId).collect(); - let bucket_bits: BucketBits = read_struct(Address::from(BUCKET_BITS_OFFSET), &memory); // Translate memory sizes expressed in pages to sizes expressed in buckets. @@ -486,18 +482,19 @@ impl MemoryManagerInner { memory_size_in_buckets.push(size_in_buckets); if size_in_buckets > 0 { - let mut bucket = bucket_bits.get_first(memory_id); - let mut buckets = vec![bucket]; - freed_buckets.remove(&bucket); - for _ in 1..size_in_buckets { - bucket = bucket_bits.get_next(bucket); - freed_buckets.remove(&bucket); - buckets.push(bucket); - } + let buckets = bucket_bits.buckets_for_memory(memory_id, size_in_buckets as u16); memory_buckets.insert(memory_id, buckets); } } + // Set of all buckets with ID smaller than 'max_bucket_id' which were allocated and freed. + let mut freed_buckets: BTreeSet = + (0..header.num_allocated_buckets).map(BucketId).collect(); + + for bucket_id in memory_buckets.values().flat_map(|buckets| buckets.iter()) { + freed_buckets.remove(bucket_id); + } + Self { memory, allocated_buckets: header.num_allocated_buckets, @@ -841,6 +838,20 @@ impl BucketBits { // } // } } + + fn buckets_for_memory(&self, memory_id: MemoryId, count: u16) -> Vec { + if count == 0 { + return Vec::new(); + } + + let mut bucket = self.get_first(memory_id); + let mut buckets = vec![bucket]; + for _ in 1..count { + bucket = self.get_next(bucket); + buckets.push(bucket); + } + buckets + } } impl Default for BucketBits { From fd8e45646106b2138bf64156f09b2ef0498ca8d0 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 11 Nov 2024 09:36:11 +0000 Subject: [PATCH 05/20] Track dirty bytes --- src/memory_manager.rs | 86 ++++++++++++++++++++++++++++++++++--------- 1 file changed, 69 insertions(+), 17 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index b93ddb97..e2d16664 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -41,7 +41,11 @@ //! assert_eq!(bytes, vec![4, 5, 6]); //! ``` -use crate::{read_struct, types::{Address, Bytes}, write_struct, Memory, WASM_PAGE_SIZE}; +use crate::{ + read_struct, + types::{Address, Bytes}, + write_struct, Memory, WASM_PAGE_SIZE, +}; use bitvec::array::BitArray; use bitvec::macros::internal::funty::Fundamental; use std::cmp::min; @@ -473,7 +477,8 @@ impl MemoryManagerInner { // Map of all memories with their assigned buckets. let mut memory_buckets = BTreeMap::new(); - let bucket_bits: BucketBits = read_struct(Address::from(BUCKET_BITS_OFFSET), &memory); + let bucket_bits: BucketBits = + read_struct::(Address::from(BUCKET_BITS_OFFSET), &memory).into(); // Translate memory sizes expressed in pages to sizes expressed in buckets. for (index, memory_size_in_pages) in header.memory_sizes_in_pages.into_iter().enumerate() { @@ -589,7 +594,8 @@ impl MemoryManagerInner { self.save_header(); // Write in stable store that this bucket belongs to the memory with the provided `id`. - write_struct(&self.bucket_bits, Address::from(BUCKET_BITS_OFFSET), &self.memory); + self.bucket_bits + .flush_dirty_bytes(&self.memory, size_of::
() as u64); // Return the old size. old_size as i64 @@ -794,20 +800,28 @@ fn bucket_allocations_address(id: BucketId) -> Address { const BUCKET_BITS_LEN: usize = 15 * MAX_NUM_BUCKETS as usize; -#[repr(C, packed)] #[derive(Clone)] struct BucketBits { + inner: BucketBitsPacked, + first_buckets_dirty: bool, + dirty_bucket_link_bytes: BTreeSet, +} + +#[derive(Clone)] +#[repr(C, packed)] +struct BucketBitsPacked { first_bucket_per_memory: [u16; MAX_NUM_MEMORIES as usize], bucket_links: BitArray<[u8; BUCKET_BITS_LEN / 8]>, } impl BucketBits { fn get_first(&self, memory_id: MemoryId) -> BucketId { - BucketId(self.first_bucket_per_memory[memory_id.0 as usize]) + BucketId(self.inner.first_bucket_per_memory[memory_id.0 as usize]) } fn set_first(&mut self, memory_id: MemoryId, value: BucketId) { - self.first_bucket_per_memory[memory_id.0 as usize] = value.0; + self.inner.first_bucket_per_memory[memory_id.0 as usize] = value.0; + self.first_buckets_dirty = true; } fn get_next(&self, bucket: BucketId) -> BucketId { @@ -815,7 +829,14 @@ impl BucketBits { let mut next_bits: BitArray<[u8; 2]> = BitArray::new([0u8; 2]); for i in 0..15 { - next_bits.set(i + 1, self.bucket_links.get(start_bit_index + i).unwrap().as_bool()); + next_bits.set( + i + 1, + self.inner + .bucket_links + .get(start_bit_index + i) + .unwrap() + .as_bool(), + ); } BucketId(u16::from_be_bytes(next_bits.data)) @@ -826,17 +847,16 @@ impl BucketBits { let next_bits: BitArray<[u8; 2]> = BitArray::from(next.0.to_be_bytes()); for (index, bit) in next_bits.iter().skip(1).enumerate() { - self.bucket_links.set(start_bit_index + index, bit.as_bool()); + self.inner + .bucket_links + .set(start_bit_index + index, bit.as_bool()); } - // let start_byte_index = start_bit_index / 8; - // let end_byte_index = (start_bit_index + 14) / 8; - // - // for index in (start_byte_index..=end_byte_index).map(|i| i as u16) { - // if !self.dirty_bytes.contains(&index) { - // self.dirty_bytes.push(index); - // } - // } + let start_byte_index = start_bit_index / 8; + let end_byte_index = (start_bit_index + 14) / 8; + + self.dirty_bucket_link_bytes + .extend((start_byte_index..=end_byte_index).map(|i| i as u16)) } fn buckets_for_memory(&self, memory_id: MemoryId, count: u16) -> Vec { @@ -852,13 +872,45 @@ impl BucketBits { } buckets } + + fn flush_dirty_bytes(&mut self, memory: &M, start_offset: u64) { + if !self.first_buckets_dirty && self.dirty_bucket_link_bytes.is_empty() { + return; + } + + // Improve this to only write dirty bytes + write_struct(&self.inner, Address::from(start_offset), memory); + + self.first_buckets_dirty = false; + self.dirty_bucket_link_bytes.clear(); + } } impl Default for BucketBits { fn default() -> Self { BucketBits { + inner: BucketBitsPacked::default(), + first_buckets_dirty: false, + dirty_bucket_link_bytes: BTreeSet::new(), + } + } +} + +impl Default for BucketBitsPacked { + fn default() -> Self { + BucketBitsPacked { first_bucket_per_memory: [0; MAX_NUM_MEMORIES as usize], - bucket_links: BitArray::default(), + bucket_links: BitArray::new([0u8; BUCKET_BITS_LEN / 8]), + } + } +} + +impl From for BucketBits { + fn from(value: BucketBitsPacked) -> Self { + BucketBits { + inner: value, + first_buckets_dirty: false, + dirty_bucket_link_bytes: BTreeSet::new(), } } } From 2ac69baea539d1154c7df76df2f5bb158a23c941 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 11 Nov 2024 10:36:47 +0000 Subject: [PATCH 06/20] Only flush dirty bytes --- src/memory_manager.rs | 39 +++++++++++++++++++++++---------------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index e2d16664..641eee78 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -41,16 +41,12 @@ //! assert_eq!(bytes, vec![4, 5, 6]); //! ``` -use crate::{ - read_struct, - types::{Address, Bytes}, - write_struct, Memory, WASM_PAGE_SIZE, -}; +use crate::{read_struct, types::{Address, Bytes}, write, write_struct, Memory, WASM_PAGE_SIZE}; use bitvec::array::BitArray; use bitvec::macros::internal::funty::Fundamental; use std::cmp::min; use std::collections::BTreeMap; -use std::mem::size_of; +use std::mem::{size_of, transmute}; use std::rc::Rc; use std::{cell::RefCell, collections::BTreeSet}; @@ -803,7 +799,7 @@ const BUCKET_BITS_LEN: usize = 15 * MAX_NUM_BUCKETS as usize; #[derive(Clone)] struct BucketBits { inner: BucketBitsPacked, - first_buckets_dirty: bool, + dirty_first_buckets: BTreeSet, dirty_bucket_link_bytes: BTreeSet, } @@ -821,7 +817,7 @@ impl BucketBits { fn set_first(&mut self, memory_id: MemoryId, value: BucketId) { self.inner.first_bucket_per_memory[memory_id.0 as usize] = value.0; - self.first_buckets_dirty = true; + self.dirty_first_buckets.insert(memory_id); } fn get_next(&self, bucket: BucketId) -> BucketId { @@ -874,15 +870,26 @@ impl BucketBits { } fn flush_dirty_bytes(&mut self, memory: &M, start_offset: u64) { - if !self.first_buckets_dirty && self.dirty_bucket_link_bytes.is_empty() { - return; + if !self.dirty_first_buckets.is_empty() { + let bytes: [u8; MAX_NUM_MEMORIES as usize * 2] = unsafe { transmute(self.inner.first_bucket_per_memory) }; + + // Multiply by 2 since we've converted from [u16] to [u8]. + let min = 2 * self.dirty_first_buckets.first().unwrap().0 as usize; + let max = 2 * self.dirty_first_buckets.last().unwrap().0 as usize + 1; + + let slice = &bytes[min..=max]; + write(memory, start_offset + min as u64, slice); + self.dirty_first_buckets.clear(); } - // Improve this to only write dirty bytes - write_struct(&self.inner, Address::from(start_offset), memory); + if !self.dirty_bucket_link_bytes.is_empty() { + let min = *self.dirty_bucket_link_bytes.first().unwrap() as usize; + let max = *self.dirty_bucket_link_bytes.last().unwrap() as usize; - self.first_buckets_dirty = false; - self.dirty_bucket_link_bytes.clear(); + let slice = &self.inner.bucket_links.data[min..=max]; + write(memory, start_offset + min as u64, slice); + self.dirty_bucket_link_bytes.clear(); + } } } @@ -890,7 +897,7 @@ impl Default for BucketBits { fn default() -> Self { BucketBits { inner: BucketBitsPacked::default(), - first_buckets_dirty: false, + dirty_first_buckets: BTreeSet::new(), dirty_bucket_link_bytes: BTreeSet::new(), } } @@ -909,7 +916,7 @@ impl From for BucketBits { fn from(value: BucketBitsPacked) -> Self { BucketBits { inner: value, - first_buckets_dirty: false, + dirty_first_buckets: BTreeSet::new(), dirty_bucket_link_bytes: BTreeSet::new(), } } From 6817fc9b017ed6455f515fdca26295c7af67083d Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 11 Nov 2024 10:45:48 +0000 Subject: [PATCH 07/20] Fix --- src/memory_manager.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index 641eee78..b45027f5 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -821,7 +821,7 @@ impl BucketBits { } fn get_next(&self, bucket: BucketId) -> BucketId { - let start_bit_index = (bucket.0 * 15) as usize; + let start_bit_index = bucket.0 as usize * 15; let mut next_bits: BitArray<[u8; 2]> = BitArray::new([0u8; 2]); for i in 0..15 { @@ -870,8 +870,10 @@ impl BucketBits { } fn flush_dirty_bytes(&mut self, memory: &M, start_offset: u64) { + const FIRST_BUCKET_PER_MEMORY_LEN: usize = 2 * MAX_NUM_MEMORIES as usize; + if !self.dirty_first_buckets.is_empty() { - let bytes: [u8; MAX_NUM_MEMORIES as usize * 2] = unsafe { transmute(self.inner.first_bucket_per_memory) }; + let bytes: [u8; FIRST_BUCKET_PER_MEMORY_LEN] = unsafe { transmute(self.inner.first_bucket_per_memory) }; // Multiply by 2 since we've converted from [u16] to [u8]. let min = 2 * self.dirty_first_buckets.first().unwrap().0 as usize; @@ -887,7 +889,7 @@ impl BucketBits { let max = *self.dirty_bucket_link_bytes.last().unwrap() as usize; let slice = &self.inner.bucket_links.data[min..=max]; - write(memory, start_offset + min as u64, slice); + write(memory, start_offset + (FIRST_BUCKET_PER_MEMORY_LEN + min) as u64, slice); self.dirty_bucket_link_bytes.clear(); } } From d78af60686a76c16952947af987e78b9922ebcf3 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 11 Nov 2024 11:46:13 +0000 Subject: [PATCH 08/20] Update layout diagram --- src/memory_manager.rs | 58 +++---------------------------------------- 1 file changed, 3 insertions(+), 55 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index b45027f5..7c212ddd 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -153,65 +153,13 @@ const HEADER_RESERVED_BYTES: usize = 32; /// -------------------------------------------------- /// Size of memory 254 (in pages) ↕ 8 bytes /// -------------------------------------------------- <- IDs of buckets -/// Bucket 1 belonging to memory 0 ↕ 15 bits +/// First bucket belonging to memory 0 ↕ 2 bytes /// -------------------------------------------------- -/// Bucket 2 belonging to memory 0 ↕ 15 bits +/// First bucket belonging to memory 1 ↕ 2 bytes /// -------------------------------------------------- /// ... /// -------------------------------------------------- -/// Bucket 1 belonging to memory 1 ↕ 15 bits -/// -------------------------------------------------- -/// Bucket 2 belonging to memory 1 ↕ 15 bits -/// -------------------------------------------------- -/// ... -/// -------------------------------------------------- -/// ... -/// --------------------------------------------------- -/// Bucket 1 belonging to memory 254 ↕ 15 bits -/// -------------------------------------------------- -/// Bucket 2 belonging to memory 254 ↕ 15 bits -/// -------------------------------------------------- -/// ... -/// -------------------------------------------------- -/// Unallocated space ↕ 2'016 bytes -/// -------------------------------------------------- <- Buckets (Page 1) -/// Bucket 1 ↕ N pages -/// -------------------------------------------------- <- Page N + 1 -/// Bucket 2 ↕ N pages -/// -------------------------------------------------- -/// ... -/// -------------------------------------------------- <- Page ((MAX_NUM_BUCKETS - 1) * N + 1) -/// Bucket MAX_NUM_BUCKETS ↕ N pages -/// ``` -/// # V2.1 layout -/// -/// ```text -/// -------------------------------------------------- <- Address 0 -/// Magic "MGR" ↕ 3 bytes -/// -------------------------------------------------- -/// Layout version ↕ 1 byte -/// -------------------------------------------------- -/// Number of allocated buckets ↕ 2 bytes -/// -------------------------------------------------- -/// Bucket size (in pages) = N ↕ 2 bytes -/// -------------------------------------------------- -/// Reserved space ↕ 32 bytes -/// -------------------------------------------------- -/// Size of memory 0 (in pages) ↕ 8 bytes -/// -------------------------------------------------- -/// Size of memory 1 (in pages) ↕ 8 bytes -/// -------------------------------------------------- -/// ... -/// -------------------------------------------------- -/// Size of memory 254 (in pages) ↕ 8 bytes -/// -------------------------------------------------- <- IDs of buckets -/// Bucket 1 belonging to memory 0 ↕ 2 bytes -/// -------------------------------------------------- -/// Bucket 1 belonging to memory 1 ↕ 2 bytes -/// -------------------------------------------------- -/// ... -/// -------------------------------------------------- -/// Bucket 1 belonging to memory 254 ↕ 2 bytes +/// First bucket belonging to memory 254 ↕ 2 bytes /// -------------------------------------------------- /// Next bucket in linked list after bucket 1 ↕ 15 bits /// -------------------------------------------------- From 4a035e684c2ef165b7857e0ee0808fb55106888a Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 11 Nov 2024 12:40:28 +0000 Subject: [PATCH 09/20] Save header and bucket bits after loading v1 layout --- src/memory_manager.rs | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index 7c212ddd..ed46f0c2 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -398,13 +398,21 @@ impl MemoryManagerInner { } let mut bucket_bits = BucketBits::default(); - for (_, buckets) in memory_buckets.iter() { - for (bucket, next_bucket) in buckets.iter().zip(buckets.iter().skip(1)) { - bucket_bits.set_next(*bucket, *next_bucket); + for (memory_id, buckets) in memory_buckets.iter() { + let mut previous = BucketId(0); + for (index, bucket) in buckets.iter().enumerate() { + if index == 0 { + bucket_bits.set_first(*memory_id, *bucket) + } else { + bucket_bits.set_next(previous, *bucket); + } + previous = *bucket; } } - Self { + write_struct(&bucket_bits.inner, Address::from(BUCKET_BITS_OFFSET), &memory); + + let mem_mgr = Self { memory, allocated_buckets: header.num_allocated_buckets, bucket_size_in_pages: header.bucket_size_in_pages, @@ -412,7 +420,9 @@ impl MemoryManagerInner { memory_buckets, freed_buckets: BTreeSet::new(), bucket_bits, - } + }; + mem_mgr.save_header(); + mem_mgr } fn load_layout_v2(memory: M, header: Header) -> Self { @@ -539,7 +549,7 @@ impl MemoryManagerInner { // Write in stable store that this bucket belongs to the memory with the provided `id`. self.bucket_bits - .flush_dirty_bytes(&self.memory, size_of::
() as u64); + .flush_dirty_bytes(&self.memory, BUCKET_BITS_OFFSET); // Return the old size. old_size as i64 From 3be6acb8466dd6de0245d560e3c4f89a24d8217f Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 11 Nov 2024 12:52:26 +0000 Subject: [PATCH 10/20] Clean up --- src/memory_manager.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index ed46f0c2..7436bd9b 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -410,7 +410,7 @@ impl MemoryManagerInner { } } - write_struct(&bucket_bits.inner, Address::from(BUCKET_BITS_OFFSET), &memory); + bucket_bits.flush_all(&memory, BUCKET_BITS_OFFSET); let mem_mgr = Self { memory, @@ -851,6 +851,12 @@ impl BucketBits { self.dirty_bucket_link_bytes.clear(); } } + + fn flush_all(&mut self, memory: &M, start_offset: u64) { + write_struct(&self.inner, Address::from(start_offset), memory); + self.dirty_first_buckets.clear(); + self.dirty_bucket_link_bytes.clear(); + } } impl Default for BucketBits { From 99cbf5076d5a03c8f7ebae18f59cc4bd244c3357 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 11 Nov 2024 15:06:08 +0000 Subject: [PATCH 11/20] Fix upgrade test --- src/memory_manager.rs | 76 ++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 72 insertions(+), 4 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index 7436bd9b..485c81f0 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -287,6 +287,8 @@ impl Memory for VirtualMemory { struct MemoryManagerInner { memory: M, + version: u8, + // The number of buckets that have been allocated. allocated_buckets: u16, @@ -327,6 +329,7 @@ impl MemoryManagerInner { fn new(memory: M, bucket_size_in_pages: u16) -> Self { let mem_mgr = Self { memory, + version: LAYOUT_VERSION_V2, allocated_buckets: 0, memory_sizes_in_pages: [0; MAX_NUM_MEMORIES as usize], memory_buckets: BTreeMap::new(), @@ -363,6 +366,7 @@ impl MemoryManagerInner { fn new_v1(memory: M, bucket_size_in_pages: u16) -> Self { let mem_mgr = Self { memory, + version: LAYOUT_VERSION_V1, allocated_buckets: 0, memory_sizes_in_pages: [0; MAX_NUM_MEMORIES as usize], memory_buckets: BTreeMap::new(), @@ -414,6 +418,7 @@ impl MemoryManagerInner { let mem_mgr = Self { memory, + version: LAYOUT_VERSION_V2, allocated_buckets: header.num_allocated_buckets, bucket_size_in_pages: header.bucket_size_in_pages, memory_sizes_in_pages: header.memory_sizes_in_pages, @@ -456,6 +461,7 @@ impl MemoryManagerInner { Self { memory, + version: LAYOUT_VERSION_V2, allocated_buckets: header.num_allocated_buckets, bucket_size_in_pages: header.bucket_size_in_pages, memory_sizes_in_pages: header.memory_sizes_in_pages, @@ -479,7 +485,7 @@ impl MemoryManagerInner { fn save_header(&self) { let header = Header { magic: *MAGIC, - version: LAYOUT_VERSION_V2, + version: self.version, num_allocated_buckets: self.allocated_buckets, bucket_size_in_pages: self.bucket_size_in_pages, _reserved: [0; HEADER_RESERVED_BYTES], @@ -496,6 +502,13 @@ impl MemoryManagerInner { // Grows the memory with the given id by the given number of pages. fn grow(&mut self, id: MemoryId, pages: u64) -> i64 { + #[cfg(test)] + if self.version == LAYOUT_VERSION_V1 { + return self.grow_v1(id, pages); + } + + debug_assert_eq!(self.version, LAYOUT_VERSION_V2); + // Compute how many additional buckets are needed. let old_size = self.memory_size(id); let new_size = old_size + pages; @@ -555,6 +568,58 @@ impl MemoryManagerInner { old_size as i64 } + #[cfg(test)] + fn grow_v1(&mut self, id: MemoryId, pages: u64) -> i64 { + // Compute how many additional buckets are needed. + let old_size = self.memory_size(id); + let new_size = old_size + pages; + let current_buckets = self.num_buckets_needed(old_size); + let required_buckets = self.num_buckets_needed(new_size); + let new_buckets_needed = required_buckets - current_buckets; + + if new_buckets_needed + self.allocated_buckets as u64 > MAX_NUM_BUCKETS { + // Exceeded the memory that can be managed. + return -1; + } + + // Allocate new buckets as needed. + for _ in 0..new_buckets_needed { + let new_bucket_id = BucketId(self.allocated_buckets); + + self.memory_buckets + .entry(id) + .or_default() + .push(new_bucket_id); + + // Write in stable store that this bucket belongs to the memory with the provided `id`. + write( + &self.memory, + bucket_allocations_address(new_bucket_id).get(), + &[id.0], + ); + + self.allocated_buckets += 1; + } + + // Grow the underlying memory if necessary. + let pages_needed = BUCKETS_OFFSET_IN_PAGES + + self.bucket_size_in_pages as u64 * self.allocated_buckets as u64; + if pages_needed > self.memory.size() { + let additional_pages_needed = pages_needed - self.memory.size(); + let prev_pages = self.memory.grow(additional_pages_needed); + if prev_pages == -1 { + panic!("{id:?}: grow failed"); + } + } + + // Update the memory with the new size. + self.memory_sizes_in_pages[id.0 as usize] = new_size; + + // Update the header and return the old size. + self.save_header(); + old_size as i64 + } + fn write(&self, id: MemoryId, offset: u64, src: &[u8]) { if (offset + src.len() as u64) > self.memory_size(id) * WASM_PAGE_SIZE { panic!("{id:?}: write out of bounds"); @@ -797,7 +862,7 @@ impl BucketBits { } fn set_next(&mut self, bucket: BucketId, next: BucketId) { - let start_bit_index = (bucket.0 * 15) as usize; + let start_bit_index = bucket.0 as usize * 15; let next_bits: BitArray<[u8; 2]> = BitArray::from(next.0.to_be_bytes()); for (index, bit) in next_bits.iter().skip(1).enumerate() { @@ -1450,7 +1515,7 @@ mod test { .collect(); proptest!(|( - num_memories in 0..255usize, + num_memories in 0..10usize, data in proptest::collection::vec(0..u8::MAX, 0..2*WASM_PAGE_SIZE as usize), offset in 0..10*WASM_PAGE_SIZE )| { @@ -1459,8 +1524,11 @@ mod test { crate::write(memory_v1, offset, &data); } + // Copy the underlying memory because loading a v1 memory manager will convert it to v2 + let mem_clone = Rc::new(RefCell::new(mem.borrow().clone())); + // Load layout v1 and convert it to layout v2. - let mem_mgr_v2 = MemoryManager::init_with_bucket_size(mem.clone(), 1); + let mem_mgr_v2 = MemoryManager::init_with_bucket_size(mem_clone, 1); let memories_v2: Vec<_> = (0..MAX_NUM_MEMORIES) .map(|id| mem_mgr_v2.get(MemoryId(id))) .collect(); From 074e07383a13ad4a3bb58b764bbb9b1e69e25198 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 11 Nov 2024 16:47:45 +0000 Subject: [PATCH 12/20] fmt --- src/memory_manager.rs | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index 485c81f0..aa22514b 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -41,7 +41,11 @@ //! assert_eq!(bytes, vec![4, 5, 6]); //! ``` -use crate::{read_struct, types::{Address, Bytes}, write, write_struct, Memory, WASM_PAGE_SIZE}; +use crate::{ + read_struct, + types::{Address, Bytes}, + write, write_struct, Memory, WASM_PAGE_SIZE, +}; use bitvec::array::BitArray; use bitvec::macros::internal::funty::Fundamental; use std::cmp::min; @@ -896,7 +900,8 @@ impl BucketBits { const FIRST_BUCKET_PER_MEMORY_LEN: usize = 2 * MAX_NUM_MEMORIES as usize; if !self.dirty_first_buckets.is_empty() { - let bytes: [u8; FIRST_BUCKET_PER_MEMORY_LEN] = unsafe { transmute(self.inner.first_bucket_per_memory) }; + let bytes: [u8; FIRST_BUCKET_PER_MEMORY_LEN] = + unsafe { transmute(self.inner.first_bucket_per_memory) }; // Multiply by 2 since we've converted from [u16] to [u8]. let min = 2 * self.dirty_first_buckets.first().unwrap().0 as usize; @@ -912,7 +917,11 @@ impl BucketBits { let max = *self.dirty_bucket_link_bytes.last().unwrap() as usize; let slice = &self.inner.bucket_links.data[min..=max]; - write(memory, start_offset + (FIRST_BUCKET_PER_MEMORY_LEN + min) as u64, slice); + write( + memory, + start_offset + (FIRST_BUCKET_PER_MEMORY_LEN + min) as u64, + slice, + ); self.dirty_bucket_link_bytes.clear(); } } From 619e4d1c4fb68a40e516844ec398a33e7391636a Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 11 Nov 2024 16:49:53 +0000 Subject: [PATCH 13/20] Add comment about safety --- src/memory_manager.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index aa22514b..e028e7b8 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -900,6 +900,7 @@ impl BucketBits { const FIRST_BUCKET_PER_MEMORY_LEN: usize = 2 * MAX_NUM_MEMORIES as usize; if !self.dirty_first_buckets.is_empty() { + // SAFETY: This is safe because we simply cast from [u16] to [u8] and double the length. let bytes: [u8; FIRST_BUCKET_PER_MEMORY_LEN] = unsafe { transmute(self.inner.first_bucket_per_memory) }; From bc7376d9ecf35dd27c86cab8f22aa78b65cfcf6b Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 11 Nov 2024 21:52:33 +0000 Subject: [PATCH 14/20] clippy --- src/memory_manager.rs | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index e028e7b8..02545ffb 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -823,7 +823,7 @@ fn bucket_allocations_address(id: BucketId) -> Address { const BUCKET_BITS_LEN: usize = 15 * MAX_NUM_BUCKETS as usize; -#[derive(Clone)] +#[derive(Clone, Default)] struct BucketBits { inner: BucketBitsPacked, dirty_first_buckets: BTreeSet, @@ -934,16 +934,6 @@ impl BucketBits { } } -impl Default for BucketBits { - fn default() -> Self { - BucketBits { - inner: BucketBitsPacked::default(), - dirty_first_buckets: BTreeSet::new(), - dirty_bucket_link_bytes: BTreeSet::new(), - } - } -} - impl Default for BucketBitsPacked { fn default() -> Self { BucketBitsPacked { From 7fad329bc87392cbe3cd1d185ee1a488e297e780 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 11 Nov 2024 21:53:45 +0000 Subject: [PATCH 15/20] Reduce size of changeset --- src/memory_manager.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index 02545ffb..fab67625 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -221,13 +221,6 @@ impl MemoryManager { } } - /// Frees the specified memory. - /// Note that the underlying physical memory doesn't shrink, but the space previously - /// occupied by the given memory will be reused. - pub fn free(&mut self, id: MemoryId) { - self.inner.borrow_mut().free(id); - } - /// Returns the underlying memory. /// /// # Returns @@ -236,6 +229,13 @@ impl MemoryManager { pub fn into_memory(self) -> Option { Rc::into_inner(self.inner).map(|inner| inner.into_inner().into_memory()) } + + /// Frees the specified memory. + /// Note that the underlying physical memory doesn't shrink, but the space previously + /// occupied by the given memory will be reused. + pub fn free(&mut self, id: MemoryId) { + self.inner.borrow_mut().free(id); + } } #[repr(C, packed)] From 3edded2615ac7a462df317208fb1bebc566461b0 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 11 Nov 2024 21:58:18 +0000 Subject: [PATCH 16/20] Reduce a bit more --- src/memory_manager.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index fab67625..dcca4234 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -684,6 +684,11 @@ impl MemoryManagerInner { (num_pages + self.bucket_size_in_pages as u64 - 1) / self.bucket_size_in_pages as u64 } + // Returns the underlying memory. + pub fn into_memory(self) -> M { + self.memory + } + fn free(&mut self, id: MemoryId) { self.memory_sizes_in_pages[id.0 as usize] = 0; let buckets = self.memory_buckets.remove(&id); @@ -696,11 +701,6 @@ impl MemoryManagerInner { self.save_header(); } - // Returns the underlying memory. - pub fn into_memory(self) -> M { - self.memory - } - #[cfg(test)] fn save_header_v1(&self) { let header = Header { From 72b291343877b49880f2e5966c9569ce97b02fec Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Tue, 12 Nov 2024 00:05:27 +0000 Subject: [PATCH 17/20] Add comments --- src/memory_manager.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index dcca4234..ae0d06d5 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -838,15 +838,19 @@ struct BucketBitsPacked { } impl BucketBits { + // Gets the first bucket assigned to a memory. + // Only call this if you know that there are buckets assigned to the memory. fn get_first(&self, memory_id: MemoryId) -> BucketId { BucketId(self.inner.first_bucket_per_memory[memory_id.0 as usize]) } + // Sets the first bucket assigned to a memory fn set_first(&mut self, memory_id: MemoryId, value: BucketId) { self.inner.first_bucket_per_memory[memory_id.0 as usize] = value.0; self.dirty_first_buckets.insert(memory_id); } + // Gets the next bucket in the linked list of buckets fn get_next(&self, bucket: BucketId) -> BucketId { let start_bit_index = bucket.0 as usize * 15; let mut next_bits: BitArray<[u8; 2]> = BitArray::new([0u8; 2]); @@ -865,6 +869,7 @@ impl BucketBits { BucketId(u16::from_be_bytes(next_bits.data)) } + // Sets the next bucket in the linked list of buckets fn set_next(&mut self, bucket: BucketId, next: BucketId) { let start_bit_index = bucket.0 as usize * 15; let next_bits: BitArray<[u8; 2]> = BitArray::from(next.0.to_be_bytes()); @@ -882,6 +887,7 @@ impl BucketBits { .extend((start_byte_index..=end_byte_index).map(|i| i as u16)) } + // Calculates the buckets for a given memory by iterating over its linked list fn buckets_for_memory(&self, memory_id: MemoryId, count: u16) -> Vec { if count == 0 { return Vec::new(); @@ -896,6 +902,7 @@ impl BucketBits { buckets } + // Flushes only the dirty bytes to memory fn flush_dirty_bytes(&mut self, memory: &M, start_offset: u64) { const FIRST_BUCKET_PER_MEMORY_LEN: usize = 2 * MAX_NUM_MEMORIES as usize; @@ -927,6 +934,7 @@ impl BucketBits { } } + // Flushes all bytes to memory fn flush_all(&mut self, memory: &M, start_offset: u64) { write_struct(&self.inner, Address::from(start_offset), memory); self.dirty_first_buckets.clear(); From 77fb4a2ba5d5ddd57c9eae93c0af99261370a0e9 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Wed, 13 Nov 2024 15:39:44 +0000 Subject: [PATCH 18/20] Wrap `version` in `cfg(test)` --- src/memory_manager.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index ae0d06d5..92578c87 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -291,6 +291,7 @@ impl Memory for VirtualMemory { struct MemoryManagerInner { memory: M, + #[cfg(test)] version: u8, // The number of buckets that have been allocated. @@ -333,6 +334,7 @@ impl MemoryManagerInner { fn new(memory: M, bucket_size_in_pages: u16) -> Self { let mem_mgr = Self { memory, + #[cfg(test)] version: LAYOUT_VERSION_V2, allocated_buckets: 0, memory_sizes_in_pages: [0; MAX_NUM_MEMORIES as usize], @@ -422,6 +424,7 @@ impl MemoryManagerInner { let mem_mgr = Self { memory, + #[cfg(test)] version: LAYOUT_VERSION_V2, allocated_buckets: header.num_allocated_buckets, bucket_size_in_pages: header.bucket_size_in_pages, @@ -465,6 +468,7 @@ impl MemoryManagerInner { Self { memory, + #[cfg(test)] version: LAYOUT_VERSION_V2, allocated_buckets: header.num_allocated_buckets, bucket_size_in_pages: header.bucket_size_in_pages, @@ -489,7 +493,10 @@ impl MemoryManagerInner { fn save_header(&self) { let header = Header { magic: *MAGIC, + #[cfg(test)] version: self.version, + #[cfg(not(test))] + version: LAYOUT_VERSION_V2, num_allocated_buckets: self.allocated_buckets, bucket_size_in_pages: self.bucket_size_in_pages, _reserved: [0; HEADER_RESERVED_BYTES], @@ -511,8 +518,6 @@ impl MemoryManagerInner { return self.grow_v1(id, pages); } - debug_assert_eq!(self.version, LAYOUT_VERSION_V2); - // Compute how many additional buckets are needed. let old_size = self.memory_size(id); let new_size = old_size + pages; From 4e34b014276aaa5fb80447d354773175a5cdd6dd Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Wed, 13 Nov 2024 17:46:35 +0000 Subject: [PATCH 19/20] Convert bucket linked list to 16-bit after reading --- src/memory_manager.rs | 184 +++++++++++++++++++++++++++--------------- 1 file changed, 120 insertions(+), 64 deletions(-) diff --git a/src/memory_manager.rs b/src/memory_manager.rs index 92578c87..2cc82b5e 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -48,6 +48,8 @@ use crate::{ }; use bitvec::array::BitArray; use bitvec::macros::internal::funty::Fundamental; +use bitvec::order::Msb0; +use bitvec::vec::BitVec; use std::cmp::min; use std::collections::BTreeMap; use std::mem::{size_of, transmute}; @@ -309,7 +311,7 @@ struct MemoryManagerInner { // NOTE: A BTreeSet is used so that bucket IDs are maintained in sorted order. freed_buckets: BTreeSet, - bucket_bits: BucketBits, + bucket_bits: BucketLinks, } impl MemoryManagerInner { @@ -341,7 +343,7 @@ impl MemoryManagerInner { memory_buckets: BTreeMap::new(), bucket_size_in_pages, freed_buckets: BTreeSet::new(), - bucket_bits: BucketBits::default(), + bucket_bits: BucketLinks::default(), }; mem_mgr.save_header(); @@ -378,7 +380,7 @@ impl MemoryManagerInner { memory_buckets: BTreeMap::new(), bucket_size_in_pages, freed_buckets: BTreeSet::new(), - bucket_bits: BucketBits::default(), + bucket_bits: BucketLinks::default(), }; mem_mgr.save_header_v1(); @@ -407,7 +409,7 @@ impl MemoryManagerInner { } } - let mut bucket_bits = BucketBits::default(); + let mut bucket_bits = BucketLinks::default(); for (memory_id, buckets) in memory_buckets.iter() { let mut previous = BucketId(0); for (index, bucket) in buckets.iter().enumerate() { @@ -443,8 +445,8 @@ impl MemoryManagerInner { // Map of all memories with their assigned buckets. let mut memory_buckets = BTreeMap::new(); - let bucket_bits: BucketBits = - read_struct::(Address::from(BUCKET_BITS_OFFSET), &memory).into(); + let bucket_bits: BucketLinks = + read_struct::(Address::from(BUCKET_BITS_OFFSET), &memory).into(); // Translate memory sizes expressed in pages to sizes expressed in buckets. for (index, memory_size_in_pages) in header.memory_sizes_in_pages.into_iter().enumerate() { @@ -828,68 +830,45 @@ fn bucket_allocations_address(id: BucketId) -> Address { const BUCKET_BITS_LEN: usize = 15 * MAX_NUM_BUCKETS as usize; -#[derive(Clone, Default)] -struct BucketBits { - inner: BucketBitsPacked, +#[derive(Clone)] +struct BucketLinks { + first_bucket_per_memory: [u16; MAX_NUM_MEMORIES as usize], + bucket_links: [u16; MAX_NUM_BUCKETS as usize], dirty_first_buckets: BTreeSet, - dirty_bucket_link_bytes: BTreeSet, + dirty_bucket_links: BTreeSet, } #[derive(Clone)] #[repr(C, packed)] -struct BucketBitsPacked { +struct BucketLinksPacked { first_bucket_per_memory: [u16; MAX_NUM_MEMORIES as usize], - bucket_links: BitArray<[u8; BUCKET_BITS_LEN / 8]>, + bucket_links: BitArray<[u8; BUCKET_BITS_LEN / 8], Msb0>, } -impl BucketBits { +const FIRST_BUCKET_PER_MEMORY_BYTES_LEN: usize = 2 * MAX_NUM_MEMORIES as usize; + +impl BucketLinks { // Gets the first bucket assigned to a memory. // Only call this if you know that there are buckets assigned to the memory. fn get_first(&self, memory_id: MemoryId) -> BucketId { - BucketId(self.inner.first_bucket_per_memory[memory_id.0 as usize]) + BucketId(self.first_bucket_per_memory[memory_id.0 as usize]) } // Sets the first bucket assigned to a memory fn set_first(&mut self, memory_id: MemoryId, value: BucketId) { - self.inner.first_bucket_per_memory[memory_id.0 as usize] = value.0; + self.first_bucket_per_memory[memory_id.0 as usize] = value.0; self.dirty_first_buckets.insert(memory_id); } // Gets the next bucket in the linked list of buckets fn get_next(&self, bucket: BucketId) -> BucketId { - let start_bit_index = bucket.0 as usize * 15; - let mut next_bits: BitArray<[u8; 2]> = BitArray::new([0u8; 2]); - - for i in 0..15 { - next_bits.set( - i + 1, - self.inner - .bucket_links - .get(start_bit_index + i) - .unwrap() - .as_bool(), - ); - } - - BucketId(u16::from_be_bytes(next_bits.data)) + BucketId(self.bucket_links[bucket.0 as usize]) } // Sets the next bucket in the linked list of buckets fn set_next(&mut self, bucket: BucketId, next: BucketId) { - let start_bit_index = bucket.0 as usize * 15; - let next_bits: BitArray<[u8; 2]> = BitArray::from(next.0.to_be_bytes()); - - for (index, bit) in next_bits.iter().skip(1).enumerate() { - self.inner - .bucket_links - .set(start_bit_index + index, bit.as_bool()); - } - - let start_byte_index = start_bit_index / 8; - let end_byte_index = (start_bit_index + 14) / 8; - - self.dirty_bucket_link_bytes - .extend((start_byte_index..=end_byte_index).map(|i| i as u16)) + self.bucket_links[bucket.0 as usize] = next.0; + self.dirty_bucket_links.insert(bucket); } // Calculates the buckets for a given memory by iterating over its linked list @@ -909,12 +888,10 @@ impl BucketBits { // Flushes only the dirty bytes to memory fn flush_dirty_bytes(&mut self, memory: &M, start_offset: u64) { - const FIRST_BUCKET_PER_MEMORY_LEN: usize = 2 * MAX_NUM_MEMORIES as usize; - if !self.dirty_first_buckets.is_empty() { // SAFETY: This is safe because we simply cast from [u16] to [u8] and double the length. - let bytes: [u8; FIRST_BUCKET_PER_MEMORY_LEN] = - unsafe { transmute(self.inner.first_bucket_per_memory) }; + let bytes: [u8; FIRST_BUCKET_PER_MEMORY_BYTES_LEN] = + unsafe { transmute(self.first_bucket_per_memory) }; // Multiply by 2 since we've converted from [u16] to [u8]. let min = 2 * self.dirty_first_buckets.first().unwrap().0 as usize; @@ -925,43 +902,122 @@ impl BucketBits { self.dirty_first_buckets.clear(); } - if !self.dirty_bucket_link_bytes.is_empty() { - let min = *self.dirty_bucket_link_bytes.first().unwrap() as usize; - let max = *self.dirty_bucket_link_bytes.last().unwrap() as usize; + if !self.dirty_bucket_links.is_empty() { + let min = self.dirty_bucket_links.first().unwrap().0 as usize; + let max = self.dirty_bucket_links.last().unwrap().0 as usize; + + let start_bit_index = 15 * min; + let start_byte_index = start_bit_index / 8; + let end_bit_index = 15 * max + 14; + let end_byte_index = end_bit_index / 8; + + let mut bits: BitVec = + BitVec::with_capacity(8 * (end_byte_index - start_byte_index + 1)); + + let prefix_bits = start_bit_index % 8; + if prefix_bits > 0 { + let previous: BitArray<[u16; 1], Msb0> = + BitArray::new([self.bucket_links[min - 1]]); + + for i in 0..prefix_bits { + bits.push(previous.get(15 - prefix_bits + i).unwrap().as_bool()); + } + } + + for bucket_id in &self.bucket_links[min..=max] { + let bucket_bits: BitArray<[u16; 1], Msb0> = BitArray::new([*bucket_id]); + for bit in bucket_bits.iter().by_vals().skip(1) { + bits.push(bit); + } + } + + let suffix_bits = end_bit_index % 8; + if prefix_bits > 0 { + let next: BitArray<[u16; 1], Msb0> = BitArray::new([self.bucket_links[max + 1]]); + + for i in 0..suffix_bits { + bits.push(next.get(i).unwrap().as_bool()); + } + } - let slice = &self.inner.bucket_links.data[min..=max]; write( memory, - start_offset + (FIRST_BUCKET_PER_MEMORY_LEN + min) as u64, - slice, + start_offset + FIRST_BUCKET_PER_MEMORY_BYTES_LEN as u64 + start_byte_index as u64, + bits.as_raw_slice(), ); - self.dirty_bucket_link_bytes.clear(); + self.dirty_bucket_links.clear(); } } // Flushes all bytes to memory fn flush_all(&mut self, memory: &M, start_offset: u64) { - write_struct(&self.inner, Address::from(start_offset), memory); + // SAFETY: This is safe because we simply cast from [u16] to [u8] and double the length. + let first_bucket_per_memory_bytes: [u8; FIRST_BUCKET_PER_MEMORY_BYTES_LEN] = + unsafe { transmute(self.first_bucket_per_memory) }; + + write(memory, start_offset, &first_bucket_per_memory_bytes); + + let mut bits: BitArray<[u8; 15 * MAX_NUM_BUCKETS as usize / 8], Msb0> = BitArray::default(); + let mut bit_index = 0; + for next in self.bucket_links.iter() { + let next_bits: BitArray<[u16; 1], Msb0> = BitArray::new([*next]); + for bit in next_bits.iter().by_vals().skip(1) { + bits.set(bit_index, bit); + bit_index += 1; + } + } + write( + memory, + start_offset + first_bucket_per_memory_bytes.len() as u64, + &bits.data, + ); + self.dirty_first_buckets.clear(); - self.dirty_bucket_link_bytes.clear(); + self.dirty_bucket_links.clear(); + } +} + +impl Default for BucketLinks { + fn default() -> Self { + BucketLinks { + first_bucket_per_memory: [0; MAX_NUM_MEMORIES as usize], + bucket_links: [0; MAX_NUM_BUCKETS as usize], + dirty_first_buckets: BTreeSet::new(), + dirty_bucket_links: BTreeSet::new(), + } } } -impl Default for BucketBitsPacked { +impl Default for BucketLinksPacked { fn default() -> Self { - BucketBitsPacked { + BucketLinksPacked { first_bucket_per_memory: [0; MAX_NUM_MEMORIES as usize], bucket_links: BitArray::new([0u8; BUCKET_BITS_LEN / 8]), } } } -impl From for BucketBits { - fn from(value: BucketBitsPacked) -> Self { - BucketBits { - inner: value, +impl From for BucketLinks { + fn from(value: BucketLinksPacked) -> Self { + let mut bucket_links = [0u16; MAX_NUM_BUCKETS as usize]; + + let mut bucket_id = 0; + let mut next_bucket: BitArray<[u16; 1], Msb0> = BitArray::new([0u16]); + for (index, bit) in value.bucket_links.iter().by_vals().enumerate() { + let bit_index = (index % 15) + 1; + next_bucket.set(bit_index, bit); + + if bit_index == 15 { + bucket_links[bucket_id] = next_bucket.data[0]; + bucket_id += 1; + } + } + + BucketLinks { + first_bucket_per_memory: value.first_bucket_per_memory, + bucket_links, dirty_first_buckets: BTreeSet::new(), - dirty_bucket_link_bytes: BTreeSet::new(), + dirty_bucket_links: BTreeSet::new(), } } } From f174c63d482bcbccaaef751983a2b0cde8343ef0 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Thu, 14 Nov 2024 13:56:04 +0000 Subject: [PATCH 20/20] Manually implement 15bit -> 16bit using rotations --- Cargo.lock | 41 +--------- Cargo.toml | 2 +- src/memory_manager.rs | 170 ++++++++++++++++++++++++++---------------- 3 files changed, 109 insertions(+), 104 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4f4352d6..952ad9c9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -78,18 +78,6 @@ version = "2.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" -[[package]] -name = "bitvec" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bc2832c24239b0141d5674bb9174f9d68a8b5b3f2753311927c172ca46f7e9c" -dependencies = [ - "funty", - "radium", - "tap", - "wyz", -] - [[package]] name = "block-buffer" version = "0.10.4" @@ -250,12 +238,6 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" -[[package]] -name = "funty" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" - [[package]] name = "generic-array" version = "0.14.7" @@ -314,7 +296,6 @@ dependencies = [ name = "ic-stable-structures" version = "0.6.5" dependencies = [ - "bitvec", "canbench-rs", "candid", "hex", @@ -323,6 +304,7 @@ dependencies = [ "ic_principal", "maplit", "proptest", + "rand", "tempfile", "test-strategy", ] @@ -497,12 +479,6 @@ dependencies = [ "proc-macro2", ] -[[package]] -name = "radium" -version = "0.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" - [[package]] name = "rand" version = "0.8.5" @@ -694,12 +670,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "tap" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" - [[package]] name = "tempfile" version = "3.14.0" @@ -884,15 +854,6 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" -[[package]] -name = "wyz" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05f360fc0b24296329c78fda852a1e9ae82de9cf7b27dae4b7f62f118f77b9ed" -dependencies = [ - "tap", -] - [[package]] name = "zerocopy" version = "0.7.35" diff --git a/Cargo.toml b/Cargo.toml index 0dfc7d5e..ae2bf33f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,6 @@ include = ["src", "Cargo.toml", "LICENSE", "README.md"] repository = "https://github.com/dfinity/stable-structures" [dependencies] -bitvec = "1.0.1" ic_principal = { version = "0.1.1", default-features = false } # An optional dependency to benchmark parts of the code. canbench-rs = { version = "0.1.7", optional = true } @@ -24,6 +23,7 @@ ic-cdk.workspace = true ic-cdk-macros.workspace = true maplit = "1.0.2" proptest = "1" +rand = "0.8.5" tempfile = "3.3.0" test-strategy = "0.3.1" diff --git a/src/memory_manager.rs b/src/memory_manager.rs index 2cc82b5e..2fd34163 100644 --- a/src/memory_manager.rs +++ b/src/memory_manager.rs @@ -46,10 +46,6 @@ use crate::{ types::{Address, Bytes}, write, write_struct, Memory, WASM_PAGE_SIZE, }; -use bitvec::array::BitArray; -use bitvec::macros::internal::funty::Fundamental; -use bitvec::order::Msb0; -use bitvec::vec::BitVec; use std::cmp::min; use std::collections::BTreeMap; use std::mem::{size_of, transmute}; @@ -828,8 +824,6 @@ fn bucket_allocations_address(id: BucketId) -> Address { Address::from(0) + Header::size() + Bytes::from(id.0) } -const BUCKET_BITS_LEN: usize = 15 * MAX_NUM_BUCKETS as usize; - #[derive(Clone)] struct BucketLinks { first_bucket_per_memory: [u16; MAX_NUM_MEMORIES as usize], @@ -842,7 +836,7 @@ struct BucketLinks { #[repr(C, packed)] struct BucketLinksPacked { first_bucket_per_memory: [u16; MAX_NUM_MEMORIES as usize], - bucket_links: BitArray<[u8; BUCKET_BITS_LEN / 8], Msb0>, + bucket_links: [u8; 15 * MAX_NUM_BUCKETS as usize / 8], } const FIRST_BUCKET_PER_MEMORY_BYTES_LEN: usize = 2 * MAX_NUM_MEMORIES as usize; @@ -890,15 +884,15 @@ impl BucketLinks { fn flush_dirty_bytes(&mut self, memory: &M, start_offset: u64) { if !self.dirty_first_buckets.is_empty() { // SAFETY: This is safe because we simply cast from [u16] to [u8] and double the length. - let bytes: [u8; FIRST_BUCKET_PER_MEMORY_BYTES_LEN] = + let first_bucket_per_memory_bytes: [u8; FIRST_BUCKET_PER_MEMORY_BYTES_LEN] = unsafe { transmute(self.first_bucket_per_memory) }; // Multiply by 2 since we've converted from [u16] to [u8]. let min = 2 * self.dirty_first_buckets.first().unwrap().0 as usize; let max = 2 * self.dirty_first_buckets.last().unwrap().0 as usize + 1; - let slice = &bytes[min..=max]; - write(memory, start_offset + min as u64, slice); + let dirty_bytes = &first_bucket_per_memory_bytes[min..=max]; + write(memory, start_offset + min as u64, dirty_bytes); self.dirty_first_buckets.clear(); } @@ -906,44 +900,21 @@ impl BucketLinks { let min = self.dirty_bucket_links.first().unwrap().0 as usize; let max = self.dirty_bucket_links.last().unwrap().0 as usize; - let start_bit_index = 15 * min; - let start_byte_index = start_bit_index / 8; - let end_bit_index = 15 * max + 14; - let end_byte_index = end_bit_index / 8; - - let mut bits: BitVec = - BitVec::with_capacity(8 * (end_byte_index - start_byte_index + 1)); - - let prefix_bits = start_bit_index % 8; - if prefix_bits > 0 { - let previous: BitArray<[u16; 1], Msb0> = - BitArray::new([self.bucket_links[min - 1]]); - - for i in 0..prefix_bits { - bits.push(previous.get(15 - prefix_bits + i).unwrap().as_bool()); - } - } - - for bucket_id in &self.bucket_links[min..=max] { - let bucket_bits: BitArray<[u16; 1], Msb0> = BitArray::new([*bucket_id]); - for bit in bucket_bits.iter().by_vals().skip(1) { - bits.push(bit); - } - } - - let suffix_bits = end_bit_index % 8; - if prefix_bits > 0 { - let next: BitArray<[u16; 1], Msb0> = BitArray::new([self.bucket_links[max + 1]]); + let start_segment = min / 8; + let end_segment = (max / 8) + 1; - for i in 0..suffix_bits { - bits.push(next.get(i).unwrap().as_bool()); - } + let mut dirty_bucket_link_bytes_15bit = + Vec::with_capacity(15 * (end_segment - start_segment + 1)); + for segment in self.bucket_links[8 * start_segment..8 * end_segment].chunks(8) { + dirty_bucket_link_bytes_15bit.extend_from_slice(&convert_16bit_to_15bit(segment)); } write( memory, - start_offset + FIRST_BUCKET_PER_MEMORY_BYTES_LEN as u64 + start_byte_index as u64, - bits.as_raw_slice(), + start_offset + + FIRST_BUCKET_PER_MEMORY_BYTES_LEN as u64 + + (8 * start_segment) as u64, + &dirty_bucket_link_bytes_15bit, ); self.dirty_bucket_links.clear(); } @@ -957,19 +928,17 @@ impl BucketLinks { write(memory, start_offset, &first_bucket_per_memory_bytes); - let mut bits: BitArray<[u8; 15 * MAX_NUM_BUCKETS as usize / 8], Msb0> = BitArray::default(); - let mut bit_index = 0; - for next in self.bucket_links.iter() { - let next_bits: BitArray<[u16; 1], Msb0> = BitArray::new([*next]); - for bit in next_bits.iter().by_vals().skip(1) { - bits.set(bit_index, bit); - bit_index += 1; - } + let mut bucket_link_bytes_15bit = [0; 15 * MAX_NUM_BUCKETS as usize / 8]; + + for (index, segment) in self.bucket_links.chunks(8).enumerate() { + let start = 15 * index; + let end = start + 15; + bucket_link_bytes_15bit[start..end].copy_from_slice(&convert_16bit_to_15bit(segment)); } write( memory, start_offset + first_bucket_per_memory_bytes.len() as u64, - &bits.data, + &bucket_link_bytes_15bit, ); self.dirty_first_buckets.clear(); @@ -992,25 +961,87 @@ impl Default for BucketLinksPacked { fn default() -> Self { BucketLinksPacked { first_bucket_per_memory: [0; MAX_NUM_MEMORIES as usize], - bucket_links: BitArray::new([0u8; BUCKET_BITS_LEN / 8]), + bucket_links: [0; 15 * MAX_NUM_BUCKETS as usize / 8], } } } +fn convert_15bit_to_16bit(i: &[u8]) -> [u16; 8] { + assert_eq!(i.len(), 15); + + let mut o = [0u16; 8]; + + o[0] = u16::from_be_bytes([ + (i[0] & 0b11111110).rotate_right(1), + ((i[0] & 0b00000001) + (i[1] & 0b11111110)).rotate_right(1), + ]); + o[1] = u16::from_be_bytes([ + ((i[1] & 0b00000001) + (i[2] & 0b11111100)).rotate_right(2), + ((i[2] & 0b00000011) + (i[3] & 0b11111100)).rotate_right(2), + ]); + o[2] = u16::from_be_bytes([ + ((i[3] & 0b00000011) + (i[4] & 0b11111000)).rotate_right(3), + ((i[4] & 0b00000111) + (i[5] & 0b11111000)).rotate_right(3), + ]); + o[3] = u16::from_be_bytes([ + ((i[5] & 0b00000111) + (i[6] & 0b11110000)).rotate_right(4), + ((i[6] & 0b00001111) + (i[7] & 0b11110000)).rotate_right(4), + ]); + o[4] = u16::from_be_bytes([ + ((i[7] & 0b00001111) + (i[8] & 0b11100000)).rotate_right(5), + ((i[8] & 0b00011111) + (i[9] & 0b11100000)).rotate_right(5), + ]); + o[5] = u16::from_be_bytes([ + ((i[9] & 0b00011111) + (i[10] & 0b11000000)).rotate_right(6), + ((i[10] & 0b00111111) + (i[11] & 0b11000000)).rotate_right(6), + ]); + o[6] = u16::from_be_bytes([ + ((i[11] & 0b00111111) + (i[12] & 0b10000000)).rotate_right(7), + ((i[12] & 0b01111111) + (i[13] & 0b10000000)).rotate_right(7), + ]); + o[7] = u16::from_be_bytes([i[13] & 0b01111111, i[14] & 0b11111111]); + o +} + +fn convert_16bit_to_15bit(i_u16: &[u16]) -> [u8; 15] { + assert_eq!(i_u16.len(), 8); + + let mut o = [0u8; 15]; + + let mut i = [0u8; 16]; + for (index, value) in i_u16.iter().enumerate() { + let start = 2 * index; + let end = start + 2; + i[start..end].copy_from_slice(&value.to_be_bytes()); + } + + o[0] = ((i[0] & 0b01111111) + (i[1] & 0b10000000)).rotate_left(1); + o[1] = (i[1] & 0b01111111).rotate_left(1) + (i[2] & 0b01000000).rotate_left(2); + o[2] = ((i[2] & 0b00111111) + (i[3] & 0b11000000)).rotate_left(2); + o[3] = (i[3] & 0b00111111).rotate_left(2) + (i[4] & 0b01100000).rotate_left(3); + o[4] = ((i[4] & 0b00011111) + (i[5] & 0b11100000)).rotate_left(3); + o[5] = (i[5] & 0b00011111).rotate_left(3) + (i[6] & 0b01110000).rotate_left(4); + o[6] = ((i[6] & 0b00001111) + (i[7] & 0b11110000)).rotate_left(4); + o[7] = (i[7] & 0b00001111).rotate_left(4) + (i[8] & 0b01111000).rotate_left(5); + o[8] = ((i[8] & 0b00000111) + (i[9] & 0b11111000)).rotate_left(5); + o[9] = (i[9] & 0b00000111).rotate_left(5) + (i[10] & 0b01111100).rotate_left(6); + o[10] = ((i[10] & 0b00000011) + (i[11] & 0b11111100)).rotate_left(6); + o[11] = (i[11] & 0b00000011).rotate_left(6) + (i[12] & 0b01111110).rotate_left(7); + o[12] = ((i[12] & 0b00000001) + (i[13] & 0b11111110)).rotate_left(7); + o[13] = (i[13] & 0b00000001).rotate_left(7) + (i[14] & 0b01111111); + o[14] = i[15]; + o +} + impl From for BucketLinks { fn from(value: BucketLinksPacked) -> Self { + let bucket_link_bytes_15bit = value.bucket_links; let mut bucket_links = [0u16; MAX_NUM_BUCKETS as usize]; - let mut bucket_id = 0; - let mut next_bucket: BitArray<[u16; 1], Msb0> = BitArray::new([0u16]); - for (index, bit) in value.bucket_links.iter().by_vals().enumerate() { - let bit_index = (index % 15) + 1; - next_bucket.set(bit_index, bit); - - if bit_index == 15 { - bucket_links[bucket_id] = next_bucket.data[0]; - bucket_id += 1; - } + for (index, segment) in bucket_link_bytes_15bit.chunks(15).enumerate() { + let start = 8 * index; + let end = start + 8; + bucket_links[start..end].copy_from_slice(&convert_15bit_to_16bit(segment)); } BucketLinks { @@ -1027,6 +1058,7 @@ mod test { use super::*; use maplit::btreemap; use proptest::prelude::*; + use rand::random; const MAX_MEMORY_IN_PAGES: u64 = MAX_NUM_BUCKETS * BUCKET_SIZE_IN_PAGES; @@ -1610,4 +1642,16 @@ mod test { } }); } + + #[test] + fn roundtrip_16bit_to_15bit() { + for _ in 0..10 { + let mut input = [0u16; 8]; + input.copy_from_slice(&(0..8).map(|_| random::() % 32768).collect::>()); + + let as_15bit = convert_16bit_to_15bit(&input); + let output = convert_15bit_to_16bit(&as_15bit); + assert_eq!(input, output); + } + } }