Fix Witty memory allocation alignment (#3104)

## Motivation

<!--
Briefly describe the goal(s) of this PR.
-->
When allocating heap memory in Witty, the alignment was incorrectly
hard-coded to `1`. This causes issues when the allocated memory does not
start at an aligned address. The code to store types will automatically
add padding, which offsets the start of the write and can then overflow
the allocated region.

## Proposal

<!--
Summarize the proposed changes and how they address the goal(s) stated
above.
-->
Specify the allocation alignment whenever needed.

## Test Plan

<!--
Explain how you made sure that the changes are correct and that they
perform as intended.

Please describe testing protocols (CI, manual tests, benchmarks, etc) in
a way that others
can reproduce the results.
-->
Unit tests were added, including one that reproduces the issue
`wit_store::test_list_fields` and now passes with the fix.

## Release Plan

<!--
If this PR targets the `main` branch, **keep the applicable lines** to
indicate if you
recommend the changes to be picked in release branches, SDKs, and
hotfixes.

This generally concerns only bug fixes.

Note that altering the public protocol (e.g. transaction format, WASM
syscalls) or storage
formats requires a new deployment.
-->
- Backporting is not possible but we may want to deploy a new `devnet`
and release a new
      SDK soon.

## Links

<!--
Optional section for related PRs, related issues, and other references.

If needed, please create issues to track future improvements and link
them here.
-->
- [reviewer
checklist](https://github.com/linera-io/linera-protocol/blob/main/CONTRIBUTING.md#reviewer-checklist)

Author: jvff

linera-witty/src/imported_function_interface/parameters.rs

@@ -114,7 +114,8 @@ where
         Instance: InstanceWithMemory,
         <Instance::Runtime as Runtime>::Memory: RuntimeMemory<Instance>,
     {
-        let location = memory.allocate(Parameters::SIZE)?;
+        let location =
+            memory.allocate(Parameters::SIZE, <Parameters::Layout as Layout>::ALIGNMENT)?;
 
         parameters.store(memory, location)?;
         location.lower(memory)

linera-witty/src/runtime/error.rs

@@ -14,6 +14,10 @@ pub enum RuntimeError {
     #[error("Requested allocation size is too large")]
     AllocationTooLarge,
 
+    /// Attempt to allocate a buffer that's aligned to an invalid boundary.
+    #[error("Requested allocation alignment is invalid")]
+    InvalidAlignment,
+
     /// Call to `cabi_realloc` returned a negative value instead of a valid address.
     #[error("Memory allocation failed")]
     AllocationFailed,

linera-witty/src/runtime/memory.rs

@@ -120,11 +120,12 @@ where
         self.memory.write(&mut *self.instance, location, bytes)
     }
 
-    /// Returns a newly allocated buffer of `size` bytes in the guest module's memory.
+    /// Returns a newly allocated buffer of `size` bytes in the guest module's memory
+    /// aligned to the requested `alignment`.
     ///
     /// Calls the guest module to allocate the memory, so the resulting allocation is managed by
     /// the guest.
-    pub fn allocate(&mut self, size: u32) -> Result<GuestPointer, RuntimeError> {
+    pub fn allocate(&mut self, size: u32, alignment: u32) -> Result<GuestPointer, RuntimeError> {
         if self.cabi_realloc.is_none() {
             self.cabi_realloc = Some(<Instance as InstanceWithFunction<
                 HList![i32, i32, i32, i32],
@@ -133,14 +134,16 @@ where
         }
 
         let size = i32::try_from(size).map_err(|_| RuntimeError::AllocationTooLarge)?;
+        let alignment = i32::try_from(alignment).map_err(|_| RuntimeError::InvalidAlignment)?;
 
         let cabi_realloc = self
             .cabi_realloc
             .as_ref()
             .expect("`cabi_realloc` function was not loaded before it was called");
 
-        let hlist_pat![allocation_address] =
-            self.instance.call(cabi_realloc, hlist![0, 0, 1, size])?;
+        let hlist_pat![allocation_address] = self
+            .instance
+            .call(cabi_realloc, hlist![0, 0, alignment, size])?;
 
         Ok(GuestPointer(
             allocation_address

linera-witty/src/test.rs

@@ -22,7 +22,7 @@ where
     let mut first_instance = MockInstance::<()>::default();
     let mut first_memory = first_instance.memory()?;
 
-    let first_address = first_memory.allocate(T::SIZE)?;
+    let first_address = first_memory.allocate(T::SIZE, <T::Layout as Layout>::ALIGNMENT)?;
 
     input.store(&mut first_memory, first_address)?;
 
@@ -34,11 +34,11 @@ where
     let mut second_instance = MockInstance::<()>::default();
     let mut second_memory = second_instance.memory()?;
 
-    let second_address = second_memory.allocate(T::SIZE)?;
+    let second_address = second_memory.allocate(T::SIZE, <T::Layout as Layout>::ALIGNMENT)?;
 
     loaded_instance.store(&mut second_memory, second_address)?;
 
-    let total_allocated_memory = first_memory.allocate(0)?.0;
+    let total_allocated_memory = first_memory.allocate(0, 1)?.0;
 
     assert_eq!(
         first_memory.read(first_address, total_allocated_memory)?,
@@ -57,7 +57,7 @@ where
 {
     let mut first_instance = MockInstance::<()>::default();
     let mut first_memory = first_instance.memory()?;
-    let first_start_address = first_memory.allocate(0)?;
+    let first_start_address = first_memory.allocate(0, 1)?;
 
     let first_lowered_layout = input.lower(&mut first_memory)?;
     let lifted_instance = T::lift_from(first_lowered_layout, &first_memory)?;
@@ -67,13 +67,13 @@ where
     // Create a clean separate memory instance
     let mut second_instance = MockInstance::<()>::default();
     let mut second_memory = second_instance.memory()?;
-    let second_start_address = second_memory.allocate(0)?;
+    let second_start_address = second_memory.allocate(0, 1)?;
 
     let second_lowered_layout = lifted_instance.lower(&mut second_memory)?;
 
     assert_eq!(first_lowered_layout, second_lowered_layout);
 
-    let total_allocated_memory = first_memory.allocate(0)?.0;
+    let total_allocated_memory = first_memory.allocate(0, 1)?.0;
 
     assert_eq!(
         first_memory.read(first_start_address, total_allocated_memory)?,

linera-witty/src/type_traits/implementations/std/string.rs

@@ -73,7 +73,7 @@ impl WitStore for String {
         <Instance::Runtime as Runtime>::Memory: RuntimeMemory<Instance>,
     {
         let length = u32::try_from(self.len())?;
-        let destination = memory.allocate(length)?;
+        let destination = memory.allocate(length, 1)?;
 
         destination.store(memory, location)?;
         length.store(memory, location.after::<GuestPointer>())?;
@@ -90,7 +90,7 @@ impl WitStore for String {
         <Instance::Runtime as Runtime>::Memory: RuntimeMemory<Instance>,
     {
         let length = u32::try_from(self.len())?;
-        let destination = memory.allocate(length)?;
+        let destination = memory.allocate(length, 1)?;
 
         memory.write(destination, self.as_bytes())?;
 

linera-witty/src/type_traits/implementations/std/vec.rs

@@ -84,7 +84,7 @@ where
         let length = u32::try_from(self.len())?;
         let size = length * T::SIZE;
 
-        let destination = memory.allocate(size)?;
+        let destination = memory.allocate(size, <T::Layout as Layout>::ALIGNMENT)?;
 
         destination.store(memory, location)?;
         length.store(memory, location.after::<GuestPointer>())?;
@@ -105,7 +105,7 @@ where
         let length = u32::try_from(self.len())?;
         let size = length * T::SIZE;
 
-        let destination = memory.allocate(size)?;
+        let destination = memory.allocate(size, <T::Layout as Layout>::ALIGNMENT)?;
 
         self.iter().zip(0..).try_for_each(|(element, index)| {
             element.store(memory, destination.index::<T>(index))

linera-witty/src/type_traits/implementations/tests.rs

@@ -228,7 +228,9 @@ where
 
     assert_eq!(layout_length, T::SIZE);
 
-    let layout_address = memory.allocate(layout_length).unwrap();
+    let layout_address = memory
+        .allocate(layout_length, <T::Layout as Layout>::ALIGNMENT)
+        .unwrap();
     let heap_address = layout_address.after::<T>();
 
     input.store(&mut memory, layout_address).unwrap();

linera-witty/tests/common/types.rs

@@ -78,3 +78,10 @@ pub struct StructWithHeapFields {
     pub rced: Rc<Leaf>,
     pub arced: Arc<Enum>,
 }
+
+/// A struct that contains fields that should all be represented as lists.
+#[derive(Clone, Debug, Eq, PartialEq, WitType, WitLoad, WitStore)]
+pub struct StructWithLists {
+    pub vec: Vec<SimpleWrapper>,
+    pub second_vec: Vec<TupleWithPadding>,
+}

linera-witty/tests/wit_load.rs

@@ -13,7 +13,8 @@ use linera_witty::{hlist, InstanceWithMemory, Layout, MockInstance, RuntimeError
 
 use self::types::{
     Branch, Enum, Leaf, RecordWithDoublePadding, SimpleWrapper, SpecializedGenericEnum,
-    SpecializedGenericStruct, StructWithHeapFields, TupleWithPadding, TupleWithoutPadding,
+    SpecializedGenericStruct, StructWithHeapFields, StructWithLists, TupleWithPadding,
+    TupleWithoutPadding,
 };
 
 /// Check that a wrapper type is properly loaded from memory and lifted from its flat layout.
@@ -272,7 +273,7 @@ fn test_invalid_discriminant() {
         17,
     ];
 
-    let address = memory.allocate(memory_bytes.len() as u32).unwrap();
+    let address = memory.allocate(memory_bytes.len() as u32, 1).unwrap();
 
     memory.write(address, &memory_bytes).unwrap();
 
@@ -353,6 +354,52 @@ fn test_heap_allocated_fields() {
     );
 }
 
+/// Check that a [`Vec`] type is properly loaded from memory and lifted from its flat
+/// layout.
+#[test]
+fn test_vec() {
+    let expected = vec![SimpleWrapper(false), SimpleWrapper(true)];
+
+    test_load_from_memory(&[8, 0, 0, 0, 2, 0, 0, 0, 0, 1], expected.clone());
+    test_lift_from_flat_layout(hlist![0_i32, 2_i32], expected, &[0, 1]);
+}
+
+/// Check that a type with list fields is properly loaded from memory and lifted from its
+/// flat layout.
+#[test]
+fn test_list_fields() {
+    let expected = StructWithLists {
+        vec: vec![
+            SimpleWrapper(true),
+            SimpleWrapper(true),
+            SimpleWrapper(false),
+        ],
+        second_vec: vec![
+            TupleWithPadding(0x1a19, 0x201f_1e1d, 0x2827_2625_2423_2221),
+            TupleWithPadding(0x2a29, 0x302f_2e2d, 0x3837_3635_3433_3231),
+        ],
+    };
+
+    test_load_from_memory(
+        &[
+            16, 0, 0, 0, 3, 0, 0, 0, 24, 0, 0, 0, 2, 0, 0, 0, 1, 1, 0, 20, 21, 22, 23, 24, 0x19,
+            0x1a, 27, 28, 0x1d, 0x1e, 0x1f, 0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28,
+            0x29, 0x2a, 43, 44, 0x2d, 0x2e, 0x2f, 0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37,
+            0x38,
+        ],
+        expected.clone(),
+    );
+    test_lift_from_flat_layout(
+        hlist![0_i32, 3_i32, 8_i32, 2_i32],
+        expected,
+        &[
+            1, 1, 0, 0, 0, 0, 0, 0, 0x19, 0x1a, 0, 0, 0x1d, 0x1e, 0x1f, 0x20, 0x21, 0x22, 0x23,
+            0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2a, 0, 0, 0x2d, 0x2e, 0x2f, 0x30, 0x31, 0x32,
+            0x33, 0x34, 0x35, 0x36, 0x37, 0x38,
+        ],
+    );
+}
+
 /// Tests that the type `T` and wrapped versions of it can be loaded from an `input` sequence of
 /// bytes in memory and that it matches the `expected` value.
 fn test_load_from_memory<T>(input: &[u8], expected: T)
@@ -374,7 +421,7 @@ where
     let mut instance = MockInstance::<()>::default();
     let mut memory = instance.memory().unwrap();
 
-    let address = memory.allocate(input.len() as u32).unwrap();
+    let address = memory.allocate(input.len() as u32, 1).unwrap();
 
     memory.write(address, input).unwrap();
 
@@ -409,7 +456,7 @@ fn test_single_lift_from_flat_layout<T>(
     let mut instance = MockInstance::<()>::default();
     let mut memory = instance.memory().unwrap();
 
-    let start_address = memory.allocate(initial_memory.len() as u32).unwrap();
+    let start_address = memory.allocate(initial_memory.len() as u32, 1).unwrap();
 
     memory.write(start_address, initial_memory).unwrap();
 

linera-witty/tests/wit_store.rs

@@ -12,7 +12,8 @@ use linera_witty::{hlist, InstanceWithMemory, Layout, MockInstance, WitStore};
 
 use self::types::{
     Branch, Enum, Leaf, RecordWithDoublePadding, SimpleWrapper, SpecializedGenericEnum,
-    SpecializedGenericStruct, StructWithHeapFields, TupleWithPadding, TupleWithoutPadding,
+    SpecializedGenericStruct, StructWithHeapFields, StructWithLists, TupleWithPadding,
+    TupleWithoutPadding,
 };
 
 /// Check that a wrapper type is properly stored in memory and lowered into its flat layout.
@@ -400,6 +401,50 @@ fn test_heap_allocated_fields() {
     );
 }
 
+/// Check that a [`Vec`] type is properly stored in memory and lowered into its flat layout.
+#[test]
+fn test_vec() {
+    let data = vec![
+        SimpleWrapper(true),
+        SimpleWrapper(false),
+        SimpleWrapper(true),
+    ];
+
+    test_store_in_memory(data.clone(), &[8, 0, 0, 0, 3, 0, 0, 0], &[1, 0, 1]);
+    test_lower_to_flat_layout(data, hlist![0_i32, 3_i32,], &[1, 0, 1]);
+}
+
+/// Check that a type with list fields is properly stored in memory and lowered into its
+/// flat layout.
+#[test]
+fn test_list_fields() {
+    let data = StructWithLists {
+        vec: vec![
+            SimpleWrapper(false),
+            SimpleWrapper(true),
+            SimpleWrapper(false),
+            SimpleWrapper(true),
+        ],
+        second_vec: vec![TupleWithPadding(1, 0, -1), TupleWithPadding(10, 11, 12)],
+    };
+
+    let expected_heap = [0, 1, 0, 1]
+        .into_iter()
+        .chain([0; 4])
+        .chain([
+            1, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+        ])
+        .chain([10, 0, 0, 0, 11, 0, 0, 0, 12, 0, 0, 0, 0, 0, 0, 0])
+        .collect::<Vec<_>>();
+
+    test_store_in_memory(
+        data.clone(),
+        &[16, 0, 0, 0, 4, 0, 0, 0, 24, 0, 0, 0, 2, 0, 0, 0],
+        &expected_heap,
+    );
+    test_lower_to_flat_layout(data, hlist![0_i32, 4_i32, 8_i32, 2_i32], &expected_heap);
+}
+
 /// Tests that the `data` of type `T` and wrapped versions of it can be stored as a sequence of
 /// bytes in memory and that they match the `expected` bytes.
 fn test_store_in_memory<T>(
@@ -439,7 +484,7 @@ fn test_single_store_in_memory<T>(
     let mut memory = instance.memory().unwrap();
 
     let length = expected_without_allocation.len() as u32;
-    let address = memory.allocate(length).unwrap();
+    let address = memory.allocate(length, 1).unwrap();
 
     data.store(&mut memory, address).unwrap();