Skip to content

Commit

Permalink
Take alignment into consideration during malloc
Browse files Browse the repository at this point in the history
  • Loading branch information
daxpedda committed Jun 5, 2023
1 parent 3d78163 commit da0cf24
Show file tree
Hide file tree
Showing 8 changed files with 36 additions and 30 deletions.
14 changes: 10 additions & 4 deletions crates/cli-support/src/js/binding.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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 {
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
14 changes: 7 additions & 7 deletions crates/cli-support/src/js/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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}();
Expand Down Expand Up @@ -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}
Expand Down Expand Up @@ -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]);
Expand All @@ -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]);
Expand Down Expand Up @@ -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;
Expand Down
2 changes: 1 addition & 1 deletion crates/cli/tests/reference/result-string.js
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}

Expand Down
6 changes: 3 additions & 3 deletions crates/cli/tests/reference/string-arg.js
Original file line number Diff line number Diff line change
Expand Up @@ -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();

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

Expand Down
8 changes: 4 additions & 4 deletions crates/cli/tests/reference/string-arg.wat
Original file line number Diff line number Diff line change
@@ -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))
Expand Down
8 changes: 6 additions & 2 deletions crates/threads-xform/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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);
});

Expand Down
4 changes: 2 additions & 2 deletions guide/src/contributing/design/exporting-rust.md
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand All @@ -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);
}
}
```
Expand Down
10 changes: 3 additions & 7 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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) {
Expand All @@ -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);
}
Expand Down

0 comments on commit da0cf24

Please sign in to comment.