test: Add unit tests for light_utils::offset (#832)

* test: Add unit tests for `light_utils::offset`

* Reorganize the code to make clear distinction between zero-copy and
  copy utilities.
* Cover the whole module with tests.

* fix: Use `ptr::write` instead of `get_mut` on uninitialized fields

The former works fine with uninitialized regions. The latter does not.

merkle-tree/bounded-vec/src/lib.rs

@@ -51,6 +51,10 @@ impl BoundedVecMetadata {
         }
     }
 
+    pub fn new_with_length(capacity: usize, length: usize) -> Self {
+        Self { capacity, length }
+    }
+
     pub fn from_ne_bytes(bytes: [u8; mem::size_of::<Self>()]) -> Self {
         Self {
             capacity: usize::from_ne_bytes(bytes[span_of!(Self, capacity)].try_into().unwrap()),
@@ -92,7 +96,7 @@ where
     T: Clone,
 {
     #[inline]
-    pub fn with_capacity(capacity: usize) -> Self {
+    fn metadata_with_capacity(capacity: usize) -> *mut BoundedVecMetadata {
         let layout = Layout::new::<BoundedVecMetadata>();
         let metadata = unsafe { alloc::alloc(layout) as *mut BoundedVecMetadata };
         if metadata.is_null() {
@@ -105,12 +109,52 @@ where
             };
         }
 
+        metadata
+    }
+
+    #[inline]
+    fn metadata_from(src_metadata: &BoundedVecMetadata) -> *mut BoundedVecMetadata {
+        let layout = Layout::new::<BoundedVecMetadata>();
+        let metadata = unsafe { alloc::alloc(layout) as *mut BoundedVecMetadata };
+        if metadata.is_null() {
+            handle_alloc_error(layout);
+        }
+        unsafe { (*metadata).clone_from(src_metadata) };
+
+        metadata
+    }
+
+    #[inline]
+    fn data_with_capacity(capacity: usize) -> NonNull<T> {
         let layout = Layout::array::<T>(capacity).unwrap();
         let data_ptr = unsafe { alloc::alloc(layout) as *mut T };
         if data_ptr.is_null() {
             handle_alloc_error(layout);
         }
-        let data = NonNull::new(data_ptr).unwrap();
+        // PANICS: We ensured that the pointer is not NULL.
+        NonNull::new(data_ptr).unwrap()
+    }
+
+    #[inline]
+    pub fn with_capacity(capacity: usize) -> Self {
+        let metadata = Self::metadata_with_capacity(capacity);
+        let data = Self::data_with_capacity(capacity);
+
+        Self { metadata, data }
+    }
+
+    /// Creates a `BoundedVec<T>` with the given `metadata`.
+    ///
+    /// # Safety
+    ///
+    /// This method is unsafe, as it does not guarantee the correctness of
+    /// provided parameters (other than `capacity`). The full responisibility
+    /// is on the caller.
+    #[inline]
+    pub unsafe fn with_metadata(metadata: &BoundedVecMetadata) -> Self {
+        let capacity = metadata.capacity();
+        let metadata = Self::metadata_from(metadata);
+        let data = Self::data_with_capacity(capacity);
 
         Self { metadata, data }
     }
@@ -244,6 +288,12 @@ where
         Some(cell)
     }
 
+    /// Returns a mutable pointer to `BoundedVec`'s buffer.
+    #[inline(always)]
+    pub fn as_mut_ptr(&mut self) -> *mut T {
+        self.data.as_ptr()
+    }
+
     #[inline]
     pub fn iter(&self) -> Iter<'_, T> {
         self.as_slice().iter()
@@ -460,6 +510,20 @@ impl CyclicBoundedVecMetadata {
         }
     }
 
+    pub fn new_with_indices(
+        capacity: usize,
+        length: usize,
+        first_index: usize,
+        last_index: usize,
+    ) -> Self {
+        Self {
+            capacity,
+            length,
+            first_index,
+            last_index,
+        }
+    }
+
     pub fn from_ne_bytes(bytes: [u8; mem::size_of::<CyclicBoundedVecMetadata>()]) -> Self {
         Self {
             capacity: usize::from_ne_bytes(bytes[span_of!(Self, capacity)].try_into().unwrap()),
@@ -509,7 +573,7 @@ where
     T: Clone,
 {
     #[inline]
-    pub fn with_capacity(capacity: usize) -> Self {
+    fn metadata_with_capacity(capacity: usize) -> *mut CyclicBoundedVecMetadata {
         let layout = Layout::new::<CyclicBoundedVecMetadata>();
         let metadata = unsafe { alloc::alloc(layout) as *mut CyclicBoundedVecMetadata };
         if metadata.is_null() {
@@ -524,12 +588,52 @@ where
             };
         }
 
+        metadata
+    }
+
+    #[inline]
+    fn metadata_from(src_metadata: &CyclicBoundedVecMetadata) -> *mut CyclicBoundedVecMetadata {
+        let layout = Layout::new::<CyclicBoundedVecMetadata>();
+        let metadata = unsafe { alloc::alloc(layout) as *mut CyclicBoundedVecMetadata };
+        if metadata.is_null() {
+            handle_alloc_error(layout);
+        }
+        unsafe { (*metadata).clone_from(src_metadata) };
+
+        metadata
+    }
+
+    #[inline]
+    fn data_with_capacity(capacity: usize) -> NonNull<T> {
         let layout = Layout::array::<T>(capacity).unwrap();
         let data_ptr = unsafe { alloc::alloc(layout) as *mut T };
         if data_ptr.is_null() {
             handle_alloc_error(layout);
         }
-        let data = NonNull::new(data_ptr).unwrap();
+        // PANICS: We ensured that the pointer is not NULL.
+        NonNull::new(data_ptr).unwrap()
+    }
+
+    #[inline]
+    pub fn with_capacity(capacity: usize) -> Self {
+        let metadata = Self::metadata_with_capacity(capacity);
+        let data = Self::data_with_capacity(capacity);
+
+        Self { metadata, data }
+    }
+
+    /// Creates a `CyclicBoundedVec<T>` with the given `metadata`.
+    ///
+    /// # Safety
+    ///
+    /// This method is unsafe, as it does not guarantee the correctness of
+    /// provided parameters (other than `capacity`). The full responisibility
+    /// is on the caller.
+    #[inline]
+    pub unsafe fn with_metadata(metadata: &CyclicBoundedVecMetadata) -> Self {
+        let capacity = metadata.capacity();
+        let metadata = Self::metadata_from(metadata);
+        let data = Self::data_with_capacity(capacity);
 
         Self { metadata, data }
     }
@@ -634,6 +738,12 @@ where
         Some(cell)
     }
 
+    /// Returns a mutable pointer to `BoundedVec`'s buffer.
+    #[inline(always)]
+    pub fn as_mut_ptr(&mut self) -> *mut T {
+        self.data.as_ptr()
+    }
+
     #[inline]
     pub fn iter(&self) -> CyclicBoundedVecIterator<'_, T> {
         CyclicBoundedVecIterator {

merkle-tree/concurrent/src/copy.rs

@@ -2,7 +2,7 @@ use std::ops::Deref;
 
 use light_bounded_vec::{BoundedVecMetadata, CyclicBoundedVecMetadata};
 use light_hasher::Hasher;
-use light_utils::offset::{read_bounded_vec_at, read_cyclic_bounded_vec_at, read_value_at};
+use light_utils::offset::copy::{read_bounded_vec_at, read_cyclic_bounded_vec_at, read_value_at};
 use memoffset::{offset_of, span_of};
 
 use crate::{errors::ConcurrentMerkleTreeError, ConcurrentMerkleTree};

merkle-tree/concurrent/src/zero_copy.rs

@@ -8,7 +8,7 @@ use light_bounded_vec::{
     BoundedVec, BoundedVecMetadata, CyclicBoundedVec, CyclicBoundedVecMetadata,
 };
 use light_hasher::Hasher;
-use light_utils::offset::{read_array_like_ptr_at, read_ptr_at, write_at};
+use light_utils::offset::zero_copy::{read_array_like_ptr_at, read_ptr_at, write_at};
 use memoffset::{offset_of, span_of};
 
 use crate::{errors::ConcurrentMerkleTreeError, ConcurrentMerkleTree};

merkle-tree/indexed/src/copy.rs

@@ -5,7 +5,7 @@ use light_concurrent_merkle_tree::{
     copy::ConcurrentMerkleTreeCopy, errors::ConcurrentMerkleTreeError,
 };
 use light_hasher::Hasher;
-use light_utils::offset::{read_cyclic_bounded_vec_at, read_value_at};
+use light_utils::offset::copy::{read_cyclic_bounded_vec_at, read_value_at};
 use num_traits::{CheckedAdd, CheckedSub, ToBytes, Unsigned};
 
 use crate::{errors::IndexedMerkleTreeError, IndexedMerkleTree};

merkle-tree/indexed/src/zero_copy.rs

@@ -12,7 +12,7 @@ use light_concurrent_merkle_tree::{
     ConcurrentMerkleTree,
 };
 use light_hasher::Hasher;
-use light_utils::offset::{read_array_like_ptr_at, read_ptr_at, write_at};
+use light_utils::offset::zero_copy::{read_array_like_ptr_at, read_ptr_at, write_at};
 use num_traits::{CheckedAdd, CheckedSub, ToBytes, Unsigned};
 
 use crate::{errors::IndexedMerkleTreeError, IndexedMerkleTree};

utils/Cargo.toml

@@ -15,3 +15,7 @@ thiserror = "1.0"
 solana-program = { workspace = true }
 ark-bn254 = "0.4.0"
 rand = "0.8"
+
+[dev-dependencies]
+bytemuck = "1.16"
+memoffset = "0.9"