Take alignment into consideration during malloc

Author: daxpedda

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

@@ -550,11 +550,13 @@ fn instruction(js: &mut JsBuilder, instr: &Instruction, log_error: &mut bool) ->
         | Instruction::CallTableElement(_)
         | Instruction::DeferCallCore(_) => {
             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() {
+                // substract alignment
+                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 +566,8 @@ fn instruction(js: &mut JsBuilder, instr: &Instruction, log_error: &mut bool) ->
                     writeln!(js.prelude, "{name} = {arg};").unwrap();
                     args.push(name);
                 }
+                // add alignment
+                args.push(String::from("4"));
             } else {
                 // Otherwise, pop off the number of parameters for the function we're calling.
                 for _ in 0..params {
@@ -813,12 +817,14 @@ fn instruction(js: &mut JsBuilder, instr: &Instruction, log_error: &mut bool) ->
             let func = js.cx.pass_to_wasm_function(kind.clone(), *mem)?;
             let malloc = js.cx.export_name_of(*malloc);
             let i = js.tmp();
+            let align = std::cmp::max(kind.size(), 4);
             js.prelude(&format!(
-                "const ptr{i} = {f}({0}, wasm.{malloc});",
+                "const ptr{i} = {f}({0}, wasm.{malloc}, {align});",
                 val,
                 i = i,
                 f = func,
                 malloc = malloc,
+                align = align,
             ));
             js.prelude(&format!("const len{} = WASM_VECTOR_LEN;", i));
             js.push(format!("ptr{}", i));
@@ -922,7 +928,7 @@ fn instruction(js: &mut JsBuilder, instr: &Instruction, log_error: &mut bool) ->
             let malloc = js.cx.export_name_of(*malloc);
             let val = js.pop();
             js.prelude(&format!(
-                "var ptr{i} = isLikeNone({0}) ? 0 : {f}({0}, wasm.{malloc});",
+                "var ptr{i} = isLikeNone({0}) ? 0 : {f}({0}, wasm.{malloc}, 4);",
                 val,
                 i = i,
                 f = func,
@@ -940,7 +946,7 @@ fn instruction(js: &mut JsBuilder, instr: &Instruction, log_error: &mut bool) ->
             let malloc = js.cx.export_name_of(*malloc);
             let i = js.tmp();
             js.prelude(&format!(
-                "var ptr{i} = {f}({val}, wasm.{malloc});",
+                "var ptr{i} = {f}({val}, wasm.{malloc}, 4);",
                 val = val,
                 i = i,
                 f = func,

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, 4) >>> 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, 4) >>> 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, 4) >>> 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]);
@@ -1415,8 +1415,8 @@ impl<'a> Context<'a> {
         self.expose_wasm_vector_len();
         self.global(&format!(
             "
-            function {}(arg, malloc) {{
-                const ptr = malloc(arg.length * {size}) >>> 0;
+            function {}(arg, malloc, align) {{
+                const ptr = malloc(arg.length * {size}, align) >>> 0;
                 {}().set(arg, ptr / {size});
                 WASM_VECTOR_LEN = arg.length;
                 return ptr;

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, 4);
     }
 }
 

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, 4) >>> 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, 4) >>> 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, 4) >>> 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 (;0;) (func (param i32 i32)))
+  (type (;1;) (func (param i32 i32) (result 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))
+  (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(4)
                         .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(4)
                 .call(free);
         },
         |body| {
             with_temp_stack(body, memory, stack, |body| {
                 body.global_get(stack.alloc)
                     .i32_const(stack.size as i32)
+                    .i32_const(4)
                     .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, 4);
   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, 4);
   }
 }
 ```

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);
         }