Take alignment into consideration during malloc (#3463)

* Take alignment into consideration during `malloc`

* Use smallest possible alignment

Co-Authored-By: Liam Murphy <[email protected]>

* Rework `DeferCallCore` to `DeferFree`

Co-Authored-By: Liam Murphy <[email protected]>

* Address review

Co-Authored-By: Liam Murphy <[email protected]>

---------

Co-authored-by: Liam Murphy <[email protected]>

Author: daxpedda

crates/cli-support/src/js/binding.rs

@@ -548,13 +548,17 @@ fn instruction(js: &mut JsBuilder, instr: &Instruction, log_error: &mut bool) ->
         | Instruction::CallExport(_)
         | Instruction::CallAdapter(_)
         | Instruction::CallTableElement(_)
-        | Instruction::DeferCallCore(_) => {
+        | Instruction::DeferFree { .. } => {
             let invoc = Invocation::from(instr, js.cx.module)?;
-            let (params, results) = invoc.params_results(js.cx);
+            let (mut params, results) = invoc.params_results(js.cx);
 
             let mut args = Vec::new();
             let tmp = js.tmp();
             if invoc.defer() {
+                if let Instruction::DeferFree { .. } = instr {
+                    // Ignore `free`'s final `align` argument, since that's manually inserted later.
+                    params -= 1;
+                }
                 // If the call is deferred, the arguments to the function still need to be
                 // accessible in the `finally` block, so we declare variables to hold the args
                 // outside of the try-finally block and then set those to the args.
@@ -564,6 +568,10 @@ fn instruction(js: &mut JsBuilder, instr: &Instruction, log_error: &mut bool) ->
                     writeln!(js.prelude, "{name} = {arg};").unwrap();
                     args.push(name);
                 }
+                if let Instruction::DeferFree { align, .. } = instr {
+                    // add alignment
+                    args.push(align.to_string());
+                }
             } else {
                 // Otherwise, pop off the number of parameters for the function we're calling.
                 for _ in 0..params {
@@ -1190,8 +1198,8 @@ impl Invocation {
                 defer: false,
             },
 
-            DeferCallCore(f) => Invocation::Core {
-                id: *f,
+            DeferFree { free, .. } => Invocation::Core {
+                id: *free,
                 defer: true,
             },
 

crates/cli-support/src/js/mod.rs

@@ -1259,14 +1259,14 @@ impl<'a> Context<'a> {
             "\
                 if (realloc === undefined) {{
                     const buf = cachedTextEncoder.encode(arg);
-                    const ptr = malloc(buf.length) >>> 0;
+                    const ptr = malloc(buf.length, 1) >>> 0;
                     {mem}().subarray(ptr, ptr + buf.length).set(buf);
                     WASM_VECTOR_LEN = buf.length;
                     return ptr;
                 }}
 
                 let len = arg.length;
-                let ptr = malloc(len) >>> 0;
+                let ptr = malloc(len, 1) >>> 0;
 
                 const mem = {mem}();
 
@@ -1294,7 +1294,7 @@ impl<'a> Context<'a> {
                     if (offset !== 0) {{
                         arg = arg.slice(offset);
                     }}
-                    ptr = realloc(ptr, len, len = offset + arg.length * 3) >>> 0;
+                    ptr = realloc(ptr, len, len = offset + arg.length * 3, 1) >>> 0;
                     const view = {mem}().subarray(ptr + offset, ptr + len);
                     const ret = encodeString(arg, view);
                     {debug_end}
@@ -1366,7 +1366,7 @@ impl<'a> Context<'a> {
                 self.global(&format!(
                     "
                         function {}(array, malloc) {{
-                            const ptr = malloc(array.length * 4) >>> 0;
+                            const ptr = malloc(array.length * 4, 4) >>> 0;
                             const mem = {}();
                             for (let i = 0; i < array.length; i++) {{
                                 mem[ptr / 4 + i] = {}(array[i]);
@@ -1383,7 +1383,7 @@ impl<'a> Context<'a> {
                 self.global(&format!(
                     "
                         function {}(array, malloc) {{
-                            const ptr = malloc(array.length * 4) >>> 0;
+                            const ptr = malloc(array.length * 4, 4) >>> 0;
                             const mem = {}();
                             for (let i = 0; i < array.length; i++) {{
                                 mem[ptr / 4 + i] = addHeapObject(array[i]);
@@ -1416,7 +1416,7 @@ impl<'a> Context<'a> {
         self.global(&format!(
             "
             function {}(arg, malloc) {{
-                const ptr = malloc(arg.length * {size}) >>> 0;
+                const ptr = malloc(arg.length * {size}, {size}) >>> 0;
                 {}().set(arg, ptr / {size});
                 WASM_VECTOR_LEN = arg.length;
                 return ptr;

crates/cli-support/src/wit/outgoing.rs

@@ -105,7 +105,7 @@ impl InstructionBuilder<'_, '_> {
                 // ... then defer a call to `free` to happen later
                 let free = self.cx.free()?;
                 self.instructions.push(InstructionData {
-                    instr: Instruction::DeferCallCore(free),
+                    instr: Instruction::DeferFree { free, align: 1 },
                     stack_change: StackChange::Modified {
                         popped: 2,
                         pushed: 2,
@@ -389,7 +389,7 @@ impl InstructionBuilder<'_, '_> {
                 // special case it.
                 assert!(!self.instructions[len..]
                     .iter()
-                    .any(|idata| matches!(idata.instr, Instruction::DeferCallCore(_))));
+                    .any(|idata| matches!(idata.instr, Instruction::DeferFree { .. })));
 
                 // Finally, we add the two inputs to UnwrapResult, and everything checks out
                 //
@@ -429,7 +429,7 @@ impl InstructionBuilder<'_, '_> {
                 // implementation.
                 let free = self.cx.free()?;
                 self.instructions.push(InstructionData {
-                    instr: Instruction::DeferCallCore(free),
+                    instr: Instruction::DeferFree { free, align: 1 },
                     stack_change: StackChange::Modified {
                         popped: 2,
                         pushed: 2,

crates/cli-support/src/wit/standard.rs

@@ -93,9 +93,11 @@ pub enum AdapterType {
 pub enum Instruction {
     /// Calls a function by its id.
     CallCore(walrus::FunctionId),
-    /// Schedules a function to be called after the whole lift/lower cycle is
-    /// finished, e.g. to deallocate a string or something.
-    DeferCallCore(walrus::FunctionId),
+    /// Call the deallocation function.
+    DeferFree {
+        free: walrus::FunctionId,
+        align: usize,
+    },
     /// A call to one of our own defined adapters, similar to the standard
     /// call-adapter instruction
     CallAdapter(AdapterId),
@@ -423,7 +425,7 @@ impl walrus::CustomSection for NonstandardWitSection {
             };
             for instr in instrs {
                 match instr.instr {
-                    DeferCallCore(f) | CallCore(f) => {
+                    DeferFree { free: f, .. } | CallCore(f) => {
                         roots.push_func(f);
                     }
                     StoreRetptr { mem, .. }

crates/cli/tests/reference/result-string.js

@@ -85,7 +85,7 @@ export function exported() {
         return getStringFromWasm0(ptr1, len1);
     } finally {
         wasm.__wbindgen_add_to_stack_pointer(16);
-        wasm.__wbindgen_free(deferred2_0, deferred2_1);
+        wasm.__wbindgen_free(deferred2_0, deferred2_1, 1);
     }
 }
 

crates/cli/tests/reference/result-string.wat

@@ -1,9 +1,9 @@
 (module
   (type (;0;) (func (param i32)))
   (type (;1;) (func (param i32) (result i32)))
-  (type (;2;) (func (param i32 i32)))
+  (type (;2;) (func (param i32 i32 i32)))
   (func $exported (;0;) (type 0) (param i32))
-  (func $__wbindgen_free (;1;) (type 2) (param i32 i32))
+  (func $__wbindgen_free (;1;) (type 2) (param i32 i32 i32))
   (func $__wbindgen_add_to_stack_pointer (;2;) (type 1) (param i32) (result i32))
   (memory (;0;) 17)
   (export "memory" (memory 0))

crates/cli/tests/reference/string-arg.js

@@ -47,14 +47,14 @@ function passStringToWasm0(arg, malloc, realloc) {
 
     if (realloc === undefined) {
         const buf = cachedTextEncoder.encode(arg);
-        const ptr = malloc(buf.length) >>> 0;
+        const ptr = malloc(buf.length, 1) >>> 0;
         getUint8Memory0().subarray(ptr, ptr + buf.length).set(buf);
         WASM_VECTOR_LEN = buf.length;
         return ptr;
     }
 
     let len = arg.length;
-    let ptr = malloc(len) >>> 0;
+    let ptr = malloc(len, 1) >>> 0;
 
     const mem = getUint8Memory0();
 
@@ -70,7 +70,7 @@ function passStringToWasm0(arg, malloc, realloc) {
         if (offset !== 0) {
             arg = arg.slice(offset);
         }
-        ptr = realloc(ptr, len, len = offset + arg.length * 3) >>> 0;
+        ptr = realloc(ptr, len, len = offset + arg.length * 3, 1) >>> 0;
         const view = getUint8Memory0().subarray(ptr + offset, ptr + len);
         const ret = encodeString(arg, view);
 

crates/cli/tests/reference/string-arg.wat

@@ -1,10 +1,10 @@
 (module
-  (type (;0;) (func (param i32) (result i32)))
-  (type (;1;) (func (param i32 i32)))
-  (type (;2;) (func (param i32 i32 i32) (result i32)))
-  (func $__wbindgen_realloc (;0;) (type 2) (param i32 i32 i32) (result i32))
-  (func $__wbindgen_malloc (;1;) (type 0) (param i32) (result i32))
-  (func $foo (;2;) (type 1) (param i32 i32))
+  (type (;0;) (func (param i32 i32)))
+  (type (;1;) (func (param i32 i32) (result i32)))
+  (type (;2;) (func (param i32 i32 i32 i32) (result i32)))
+  (func $__wbindgen_realloc (;0;) (type 2) (param i32 i32 i32 i32) (result i32))
+  (func $__wbindgen_malloc (;1;) (type 1) (param i32 i32) (result i32))
+  (func $foo (;2;) (type 0) (param i32 i32))
   (memory (;0;) 17)
   (export "memory" (memory 0))
   (export "foo" (func $foo))

crates/threads-xform/src/lib.rs

@@ -344,9 +344,10 @@ fn inject_start(
             // we give ourselves a stack and we update our stack
             // pointer as the default stack pointer is surely wrong for us.
             |body| {
-                // local = malloc(stack.size) [aka base]
+                // local = malloc(stack.size, align) [aka base]
                 with_temp_stack(body, memory, stack, |body| {
                     body.i32_const(stack.size as i32)
+                        .i32_const(16)
                         .call(malloc)
                         .local_tee(local);
                 });
@@ -368,7 +369,6 @@ fn inject_start(
     // Afterwards we need to initialize our thread-local state.
     body.i32_const(tls.size as i32)
         .i32_const(tls.align as i32)
-        .drop() // TODO: need to actually respect alignment
         .call(malloc)
         .global_set(tls.base)
         .global_get(tls.base)
@@ -406,11 +406,13 @@ fn inject_destroy(
         |body| {
             body.local_get(tls_base)
                 .i32_const(tls.size as i32)
+                .i32_const(tls.align as i32)
                 .call(free);
         },
         |body| {
             body.global_get(tls.base)
                 .i32_const(tls.size as i32)
+                .i32_const(tls.align as i32)
                 .call(free);
 
             // set tls.base = i32::MIN to trigger invalid memory
@@ -425,12 +427,14 @@ fn inject_destroy(
             // we're destroying somebody else's stack, so we can use our own
             body.local_get(stack_alloc)
                 .i32_const(stack.size as i32)
+                .i32_const(16)
                 .call(free);
         },
         |body| {
             with_temp_stack(body, memory, stack, |body| {
                 body.global_get(stack.alloc)
                     .i32_const(stack.size as i32)
+                    .i32_const(16)
                     .call(free);
             });
 

guide/src/contributing/design/exporting-rust.md

@@ -33,7 +33,7 @@ import * as wasm from './foo_bg';
 function passStringToWasm(arg) {
   const buf = new TextEncoder('utf-8').encode(arg);
   const len = buf.length;
-  const ptr = wasm.__wbindgen_malloc(len);
+  const ptr = wasm.__wbindgen_malloc(len, 1);
   let array = new Uint8Array(wasm.memory.buffer);
   array.set(buf, ptr);
   return [ptr, len];
@@ -56,7 +56,7 @@ export function greet(arg0) {
     wasm.__wbindgen_boxed_str_free(ret);
     return realRet;
   } finally {
-    wasm.__wbindgen_free(ptr0, len0);
+    wasm.__wbindgen_free(ptr0, len0, 1);
   }
 }
 ```

src/lib.rs

@@ -1566,11 +1566,9 @@ pub mod __rt {
 
     if_std! {
         use std::alloc::{alloc, dealloc, realloc, Layout};
-        use std::mem;
 
         #[no_mangle]
-        pub extern "C" fn __wbindgen_malloc(size: usize) -> *mut u8 {
-            let align = mem::align_of::<usize>();
+        pub extern "C" fn __wbindgen_malloc(size: usize, align: usize) -> *mut u8 {
             if let Ok(layout) = Layout::from_size_align(size, align) {
                 unsafe {
                     if layout.size() > 0 {
@@ -1588,8 +1586,7 @@ pub mod __rt {
         }
 
         #[no_mangle]
-        pub unsafe extern "C" fn __wbindgen_realloc(ptr: *mut u8, old_size: usize, new_size: usize) -> *mut u8 {
-            let align = mem::align_of::<usize>();
+        pub unsafe extern "C" fn __wbindgen_realloc(ptr: *mut u8, old_size: usize, new_size: usize, align: usize) -> *mut u8 {
             debug_assert!(old_size > 0);
             debug_assert!(new_size > 0);
             if let Ok(layout) = Layout::from_size_align(old_size, align) {
@@ -1611,13 +1608,12 @@ pub mod __rt {
         }
 
         #[no_mangle]
-        pub unsafe extern "C" fn __wbindgen_free(ptr: *mut u8, size: usize) {
+        pub unsafe extern "C" fn __wbindgen_free(ptr: *mut u8, size: usize, align: usize) {
             // This happens for zero-length slices, and in that case `ptr` is
             // likely bogus so don't actually send this to the system allocator
             if size == 0 {
                 return
             }
-            let align = mem::align_of::<usize>();
             let layout = Layout::from_size_align_unchecked(size, align);
             dealloc(ptr, layout);
         }