Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Rust] use cleanup list as necessary for stream/future writes #1240

Merged
merged 1 commit into from
Mar 24, 2025
Merged

[Rust] use cleanup list as necessary for stream/future writes #1240

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 17 additions & 13 deletions crates/rust/src/bindgen.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,22 @@ impl<'a, 'b> FunctionBindgen<'a, 'b> {
}
}

fn emit_cleanup(&mut self) {
pub(crate) fn flush_cleanup(&mut self) {
if !self.cleanup.is_empty() {
self.needs_cleanup_list = true;
self.push_str("cleanup_list.extend_from_slice(&[");
for (ptr, layout) in mem::take(&mut self.cleanup) {
self.push_str("(");
self.push_str(&ptr);
self.push_str(", ");
self.push_str(&layout);
self.push_str("),");
}
self.push_str("]);\n");
}
}

pub(crate) fn emit_cleanup(&mut self) {
if self.emitted_cleanup {
return;
}
Expand Down Expand Up @@ -244,18 +259,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
}

fn finish_block(&mut self, operands: &mut Vec<String>) {
if !self.cleanup.is_empty() {
self.needs_cleanup_list = true;
self.push_str("cleanup_list.extend_from_slice(&[");
for (ptr, layout) in mem::take(&mut self.cleanup) {
self.push_str("(");
self.push_str(&ptr);
self.push_str(", ");
self.push_str(&layout);
self.push_str("),");
}
self.push_str("]);\n");
}
self.flush_cleanup();
let (prev_src, prev_cleanup) = self.block_storage.pop().unwrap();
let src = mem::replace(&mut self.src, prev_src);
self.cleanup = prev_cleanup;
Expand Down
56 changes: 44 additions & 12 deletions crates/rust/src/interface.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -535,34 +535,40 @@ macro_rules! {macro_name} {{
};
let size = size.size_wasm32();
let align = align.align_wasm32();
let (lower, lift) = if let Some(payload_type) = payload_type {
let lower =
let (lower, cleanup, lift) = if let Some(payload_type) = payload_type {
let (lower, cleanup) =
self.lower_to_memory("address", "&value", &payload_type, &module);
let lift =
self.lift_from_memory("address", "value", &payload_type, &module);
(lower, lift)
(lower, cleanup, lift)
} else {
(String::new(), "let value = ();\n".into())
(String::new(), None, "let value = ();\n".into())
};

let (cleanup_start, cleanup_end) =
cleanup.unwrap_or_else(|| (String::new(), String::new()));

let box_ = self.path_to_box();
let code = format!(
r#"
#[doc(hidden)]
pub mod vtable{ordinal} {{
fn write(future: u32, value: {name}) -> ::core::pin::Pin<{box_}<dyn ::core::future::Future<Output = bool>>> {{
{box_}::pin(async move {{
{cleanup_start}
#[repr(align({align}))]
struct Buffer([::core::mem::MaybeUninit::<u8>; {size}]);
let mut buffer = Buffer([::core::mem::MaybeUninit::uninit(); {size}]);
let address = buffer.0.as_mut_ptr() as *mut u8;
{lower}

match unsafe {{ {async_support}::await_future_result(start_write, future, address).await }} {{
let result = match unsafe {{ {async_support}::await_future_result(start_write, future, address).await }} {{
{async_support}::AsyncWaitResult::Values(_) => true,
{async_support}::AsyncWaitResult::End => false,
{async_support}::AsyncWaitResult::Error(_) => unreachable!("received error while performing write"),
}}
}};
{cleanup_end}
result
}})
}}

Expand Down Expand Up @@ -660,14 +666,15 @@ pub mod vtable{ordinal} {{
let size = size.size_wasm32();
let align = align.align_wasm32();
let alloc = self.path_to_std_alloc_module();
let (lower_address, lower, lift_address, lift) = match payload_type {
let (lower_address, lower, cleanup, lift_address, lift) = match payload_type
{
Some(payload_type) if !stream_direct(payload_type) => {
let address = format!(
"let address = unsafe {{ {alloc}::alloc\
({alloc}::Layout::from_size_align_unchecked\
({size} * values.len(), {align})) }};"
);
let lower = self.lower_to_memory(
let (lower, cleanup) = self.lower_to_memory(
"address",
"value",
&payload_type,
Expand Down Expand Up @@ -696,7 +703,7 @@ for (index, dst) in values.iter_mut().take(count).enumerate() {{
}}
"#
);
(address.clone(), lower, address, lift)
(address.clone(), lower, cleanup, address, lift)
}
_ => {
let lower_address =
Expand All @@ -706,25 +713,29 @@ for (index, dst) in values.iter_mut().take(count).enumerate() {{
(
lower_address,
String::new(),
None,
lift_address,
"let value = ();\n".into(),
)
}
};

let (cleanup_start, cleanup_end) =
cleanup.unwrap_or_else(|| (String::new(), String::new()));

let box_ = self.path_to_box();
let code = format!(
r#"
#[doc(hidden)]
pub mod vtable{ordinal} {{
fn write(stream: u32, values: &[{name}]) -> ::core::pin::Pin<{box_}<dyn ::core::future::Future<Output = usize> + '_>> {{
{box_}::pin(async move {{
{cleanup_start}
{lower_address}
{lower}

let mut total = 0;
while total < values.len() {{

match unsafe {{
{async_support}::await_stream_result(
start_write,
Expand All @@ -738,6 +749,7 @@ pub mod vtable{ordinal} {{
{async_support}::AsyncWaitResult::End => break,
}}
}}
{cleanup_end}
total
}})
}}
Expand Down Expand Up @@ -867,10 +879,30 @@ pub mod vtable{ordinal} {{
}
}

fn lower_to_memory(&mut self, address: &str, value: &str, ty: &Type, module: &str) -> String {
fn lower_to_memory(
&mut self,
address: &str,
value: &str,
ty: &Type,
module: &str,
) -> (String, Option<(String, String)>) {
let mut f = FunctionBindgen::new(self, Vec::new(), true, module, true);
abi::lower_to_memory(f.r#gen.resolve, &mut f, address.into(), value.into(), ty);
format!("unsafe {{ {} }}", String::from(f.src))
f.flush_cleanup();
let lower = format!("unsafe {{ {} }}", String::from(f.src));
let cleanup = if f.needs_cleanup_list {
f.src = Default::default();
f.emit_cleanup();
let body = String::from(f.src);
let vec = self.path_to_vec();
Some((
format!("let mut cleanup_list = {vec}::new();\n"),
format!("unsafe {{ {body} }}"),
))
} else {
None
};
(lower, cleanup)
}

fn lift_from_memory(&mut self, address: &str, value: &str, ty: &Type, module: &str) -> String {
Expand Down
4 changes: 4 additions & 0 deletions tests/codegen/futures.wit
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,10 @@ interface futures {
future-f32-ret: func() -> future<f32>;
future-f64-ret: func() -> future<f64>;

future-result-list-string-ret: func() -> future<result<list<string>>>;
future-result-list-list-u8-ret: func() -> future<result<list<list<u8>>>>;
future-list-list-list-u8-ret: func() -> future<list<list<list<u8>>>>;

tuple-future: func(x: future<tuple<u8, s8>>) -> future<tuple<s64, u32>>;
string-future-arg: func(a: future<string>);
string-future-ret: func() -> future<string>;
Expand Down
4 changes: 4 additions & 0 deletions tests/codegen/streams.wit
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,10 @@ interface streams {
stream-f32-ret: func() -> stream<f32>;
stream-f64-ret: func() -> stream<f64>;

stream-result-list-string-ret: func() -> stream<result<list<string>>>;
stream-result-list-list-u8-ret: func() -> stream<result<list<list<u8>>>>;
stream-list-list-list-u8-ret: func() -> stream<list<list<list<u8>>>>;

tuple-stream: func(x: stream<tuple<u8, s8>>) -> stream<tuple<s64, u32>>;
string-stream-arg: func(a: stream<string>);
string-stream-ret: func() -> stream<string>;
Expand Down