Skip to content

Commit

Permalink
Fix Witty memory allocation alignment (linera-io#3104)
Browse files Browse the repository at this point in the history
## 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)
  • Loading branch information
jvff authored Jan 9, 2025
1 parent d6825b7 commit 52d2a75
Show file tree
Hide file tree
Showing 11 changed files with 173 additions and 23 deletions.
3 changes: 2 additions & 1 deletion linera-witty/src/imported_function_interface/parameters.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
4 changes: 4 additions & 0 deletions linera-witty/src/runtime/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
11 changes: 7 additions & 4 deletions linera-witty/src/runtime/memory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand All @@ -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
Expand Down
12 changes: 6 additions & 6 deletions linera-witty/src/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)?;

Expand All @@ -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)?,
Expand All @@ -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)?;
Expand All @@ -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)?,
Expand Down
4 changes: 2 additions & 2 deletions linera-witty/src/type_traits/implementations/std/string.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>())?;
Expand All @@ -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())?;

Expand Down
4 changes: 2 additions & 2 deletions linera-witty/src/type_traits/implementations/std/vec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>())?;
Expand All @@ -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))
Expand Down
4 changes: 3 additions & 1 deletion linera-witty/src/type_traits/implementations/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
7 changes: 7 additions & 0 deletions linera-witty/tests/common/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
}
55 changes: 51 additions & 4 deletions linera-witty/tests/wit_load.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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();

Expand Down Expand Up @@ -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)
Expand All @@ -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();

Expand Down Expand Up @@ -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();

Expand Down
49 changes: 47 additions & 2 deletions linera-witty/tests/wit_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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>(
Expand Down Expand Up @@ -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();

Expand Down
43 changes: 42 additions & 1 deletion linera-witty/tests/wit_type.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@ use linera_witty::{HList, Layout, RegisterWitTypes, WitType};

use self::types::{
Branch, Enum, Leaf, RecordWithDoublePadding, SimpleWrapper, SpecializedGenericEnum,
SpecializedGenericStruct, StructWithHeapFields, TupleWithPadding, TupleWithoutPadding,
SpecializedGenericStruct, StructWithHeapFields, StructWithLists, TupleWithPadding,
TupleWithoutPadding,
};

/// Check the memory size, layout and WIT type declaration derived for a wrapper type.
Expand Down Expand Up @@ -209,6 +210,46 @@ fn test_heap_allocated_fields() {
});
}

/// Check the memory size, layout and WIT declaration derived for a [`Vec`] type.
#[test]
fn test_vec() {
test_wit_type_implementation::<Vec<SimpleWrapper>>(ExpectedMetadata {
size: 8,
alignment: 4,
flat_layout_length: 2,
declaration: concat!(
" record simple-wrapper {\n",
" inner0: bool,\n",
" }\n"
),
});
}

/// Check the memory size, layout and WIT declaration derived for a type that has list
/// fields.
#[test]
fn test_list_fields() {
test_wit_type_implementation::<StructWithLists>(ExpectedMetadata {
size: 16,
alignment: 4,
flat_layout_length: 4,
declaration: concat!(
" record simple-wrapper {\n",
" inner0: bool,\n",
" }\n\n",
" record struct-with-lists {\n",
" vec: list<simple-wrapper>,\n",
" second-vec: list<tuple-with-padding>,\n",
" }\n\n",
" record tuple-with-padding {\n",
" inner0: u16,\n",
" inner1: u32,\n",
" inner2: s64,\n",
" }\n"
),
});
}

/// Helper type to make visible what each metadata value is.
#[derive(Clone, Copy, Debug)]
struct ExpectedMetadata {
Expand Down

0 comments on commit 52d2a75

Please sign in to comment.