use super::WasmLowerError;
type R<T> = Result<T, WasmLowerError>;
fn leb_u(mut v: u64, out: &mut Vec<u8>) {
loop {
let b = (v & 0x7f) as u8;
v >>= 7;
if v != 0 {
out.push(b | 0x80);
} else {
out.push(b);
break;
}
}
}
fn leb_u32_padded(v: u32, out: &mut Vec<u8>) {
let mut v = v;
for i in 0..5 {
let mut b = (v & 0x7f) as u8;
v >>= 7;
if i < 4 {
b |= 0x80;
}
out.push(b);
}
}
fn read_uleb(bytes: &[u8], i: &mut usize) -> R<u64> {
let mut r = 0u64;
let mut s = 0;
loop {
let b = *bytes.get(*i).ok_or(WasmLowerError::Unsupported("truncated LEB"))?;
*i += 1;
r |= u64::from(b & 0x7f) << s;
if b & 0x80 == 0 {
break;
}
s += 7;
if s >= 64 {
return Err(WasmLowerError::Unsupported("overlong LEB"));
}
}
Ok(r)
}
fn skip_sleb(bytes: &[u8], i: &mut usize) -> R<()> {
loop {
let b = *bytes.get(*i).ok_or(WasmLowerError::Unsupported("truncated SLEB"))?;
*i += 1;
if b & 0x80 == 0 {
break;
}
}
Ok(())
}
fn skip_blocktype(bytes: &[u8], i: &mut usize) -> R<()> {
let b = *bytes.get(*i).ok_or(WasmLowerError::Unsupported("truncated blocktype"))?;
if matches!(b, 0x40 | 0x7f | 0x7e | 0x7d | 0x7c | 0x7b | 0x70 | 0x6f) {
*i += 1;
Ok(())
} else {
skip_sleb(bytes, i) }
}
fn rewrite_body(body: &[u8]) -> R<(Vec<u8>, Vec<(u32, u32)>, Vec<(u32, u32)>)> {
let mut out = Vec::with_capacity(body.len() + 8);
let mut calls = Vec::new();
let mut globals = Vec::new();
let mut i = 0usize;
let start = i;
let n_groups = read_uleb(body, &mut i)?;
for _ in 0..n_groups {
read_uleb(body, &mut i)?; i += 1; if i > body.len() {
return Err(WasmLowerError::Unsupported("truncated locals"));
}
}
out.extend_from_slice(&body[start..i]);
let mut depth = 0i32;
loop {
let op_start = i;
let op = *body.get(i).ok_or(WasmLowerError::Unsupported("truncated body"))?;
i += 1;
match op {
0x02 | 0x03 | 0x04 => {
skip_blocktype(body, &mut i)?;
depth += 1;
out.extend_from_slice(&body[op_start..i]);
}
0x05 => out.extend_from_slice(&body[op_start..i]), 0x0b => {
out.extend_from_slice(&body[op_start..i]); if depth == 0 {
break;
}
depth -= 1;
}
0x0c | 0x0d => {
read_uleb(body, &mut i)?; out.extend_from_slice(&body[op_start..i]);
}
0x0e => {
let n = read_uleb(body, &mut i)?; for _ in 0..=n {
read_uleb(body, &mut i)?;
}
out.extend_from_slice(&body[op_start..i]);
}
0x00 | 0x01 | 0x0f | 0x1a | 0x1b => out.extend_from_slice(&body[op_start..i]), 0x10 => {
let func = read_uleb(body, &mut i)? as u32;
out.push(0x10);
calls.push((out.len() as u32, func));
leb_u32_padded(func, &mut out);
}
0x11 => return Err(WasmLowerError::Unsupported("call_indirect needs a type relocation")),
0x20 | 0x21 | 0x22 => {
read_uleb(body, &mut i)?;
out.extend_from_slice(&body[op_start..i]);
}
0x23 | 0x24 => {
let g = read_uleb(body, &mut i)? as u32;
out.push(op);
globals.push((out.len() as u32, g));
leb_u32_padded(g, &mut out);
}
0x28..=0x3e => {
read_uleb(body, &mut i)?;
read_uleb(body, &mut i)?;
out.extend_from_slice(&body[op_start..i]);
}
0x3f | 0x40 => {
i += 1; out.extend_from_slice(&body[op_start..i]);
}
0x41 | 0x42 => {
skip_sleb(body, &mut i)?; out.extend_from_slice(&body[op_start..i]);
}
0x43 => {
i += 4; out.extend_from_slice(&body[op_start..i]);
}
0x44 => {
i += 8; out.extend_from_slice(&body[op_start..i]);
}
0x45..=0xc4 => out.extend_from_slice(&body[op_start..i]),
0xfc => {
let sub = read_uleb(body, &mut i)?;
if sub > 7 {
return Err(WasmLowerError::Unsupported("unsupported 0xFC subopcode"));
}
out.extend_from_slice(&body[op_start..i]);
}
_ => return Err(WasmLowerError::Unsupported("unknown opcode in relocatable rewrite")),
}
if i > body.len() {
return Err(WasmLowerError::Unsupported("body overran"));
}
}
Ok((out, calls, globals))
}
struct Sec {
id: u8,
start: usize,
end: usize,
}
fn sections(module: &[u8]) -> R<Vec<Sec>> {
if module.len() < 8 || &module[0..4] != b"\0asm" {
return Err(WasmLowerError::Unsupported("not a wasm module"));
}
let mut secs = Vec::new();
let mut i = 8;
while i < module.len() {
let id = module[i];
i += 1;
let size = read_uleb(module, &mut i)? as usize;
let start = i;
let end = start + size;
if end > module.len() {
return Err(WasmLowerError::Unsupported("section overruns module"));
}
secs.push(Sec { id, start, end });
i = end;
}
Ok(secs)
}
fn parse_imports(body: &[u8]) -> R<Vec<String>> {
let mut names = Vec::new();
let mut i = 0;
let count = read_uleb(body, &mut i)?;
for _ in 0..count {
let ml = read_uleb(body, &mut i)? as usize;
i += ml;
let fl = read_uleb(body, &mut i)? as usize;
let field = String::from_utf8_lossy(&body[i..i + fl]).to_string();
i += fl;
let kind = body[i];
i += 1;
match kind {
0x00 => {
read_uleb(body, &mut i)?; names.push(field);
}
0x01 => {
i += 1;
let f = body[i];
i += 1;
read_uleb(body, &mut i)?;
if f & 1 != 0 {
read_uleb(body, &mut i)?;
}
}
0x02 => {
let f = body[i];
i += 1;
read_uleb(body, &mut i)?;
if f & 1 != 0 {
read_uleb(body, &mut i)?;
}
}
0x03 => i += 2,
_ => return Err(WasmLowerError::Unsupported("unknown import kind")),
}
}
Ok(names)
}
fn put_section(out: &mut Vec<u8>, id: u8, body: &[u8]) {
out.push(id);
leb_u(body.len() as u64, out);
out.extend_from_slice(body);
}
fn put_custom(out: &mut Vec<u8>, name: &str, body: &[u8]) {
let mut payload = Vec::new();
leb_u(name.len() as u64, &mut payload);
payload.extend_from_slice(name.as_bytes());
payload.extend_from_slice(body);
put_section(out, 0, &payload);
}
pub(crate) fn module_to_relocatable(module: &[u8]) -> R<Vec<u8>> {
let secs = sections(module)?;
for id in [5u8, 9, 11] {
if secs.iter().any(|s| s.id == id) {
return Err(WasmLowerError::Unsupported("relocatable transform: program needs memory/data/element relocations (not yet)"));
}
}
let sec = |id: u8| secs.iter().find(|s| s.id == id).map(|s| &module[s.start..s.end]);
let imports = sec(2).map(parse_imports).transpose()?.unwrap_or_default();
let k = imports.len() as u32;
let func_body = sec(3).ok_or(WasmLowerError::Unsupported("no function section"))?;
let mut fi = 0;
let num_defined = read_uleb(func_body, &mut fi)? as u32;
let num_globals = match sec(6) {
Some(g) => {
let mut gi = 0;
read_uleb(g, &mut gi)? as u32
}
None => 0,
};
let num_func_symbols = k + num_defined;
let code_body = sec(10).ok_or(WasmLowerError::Unsupported("no code section"))?;
let mut ci = 0;
let n_bodies = read_uleb(code_body, &mut ci)?;
if n_bodies != u64::from(num_defined) {
return Err(WasmLowerError::Unsupported("code/function count mismatch"));
}
let mut new_code = Vec::new();
leb_u(n_bodies, &mut new_code);
let mut func_relocs: Vec<(u32, u32)> = Vec::new();
let mut global_relocs: Vec<(u32, u32)> = Vec::new();
for _ in 0..n_bodies {
let sz = read_uleb(code_body, &mut ci)? as usize;
let body = &code_body[ci..ci + sz];
ci += sz;
let (new_body, calls, globals) = rewrite_body(body)?;
leb_u(new_body.len() as u64, &mut new_code);
let body_start = new_code.len() as u32;
new_code.extend_from_slice(&new_body);
for (off, func) in calls {
func_relocs.push((body_start + off, func));
}
for (off, g) in globals {
global_relocs.push((body_start + off, num_func_symbols + g));
}
}
const WASM_SYM_UNDEFINED: u32 = 0x10;
const WASM_SYM_EXPORTED: u32 = 0x04;
const WASM_SYM_BINDING_LOCAL: u32 = 0x02;
let mut symtab = Vec::new();
leb_u(u64::from(num_func_symbols + num_globals), &mut symtab);
for (idx, name) in imports.iter().enumerate() {
symtab.push(0x00); leb_u(u64::from(WASM_SYM_UNDEFINED), &mut symtab);
leb_u(idx as u64, &mut symtab); let _ = name; }
for d in 0..num_defined {
let fidx = k + d;
let (name, flags) = if d == 0 { ("main".to_string(), WASM_SYM_EXPORTED) } else { (format!("fn{d}"), 0) };
symtab.push(0x00);
leb_u(u64::from(flags), &mut symtab);
leb_u(u64::from(fidx), &mut symtab);
leb_u(name.len() as u64, &mut symtab);
symtab.extend_from_slice(name.as_bytes());
}
for g in 0..num_globals {
symtab.push(0x02); leb_u(u64::from(WASM_SYM_BINDING_LOCAL), &mut symtab); leb_u(u64::from(g), &mut symtab); let name = format!("g{g}");
leb_u(name.len() as u64, &mut symtab);
symtab.extend_from_slice(name.as_bytes());
}
let mut linking = Vec::new();
leb_u(2, &mut linking); linking.push(0x08); leb_u(symtab.len() as u64, &mut linking);
linking.extend_from_slice(&symtab);
let reemit = |id: u8| id != 0 && id != 7;
let emitted: Vec<u8> = secs.iter().filter(|s| reemit(s.id)).map(|s| s.id).collect();
let code_index =
emitted.iter().position(|&id| id == 10).ok_or(WasmLowerError::Unsupported("no code section"))? as u32;
const R_WASM_FUNCTION_INDEX_LEB: u8 = 0x00;
const R_WASM_GLOBAL_INDEX_LEB: u8 = 0x07;
let mut all: Vec<(u32, u8, u32)> = Vec::new();
all.extend(func_relocs.iter().map(|&(o, s)| (o, R_WASM_FUNCTION_INDEX_LEB, s)));
all.extend(global_relocs.iter().map(|&(o, s)| (o, R_WASM_GLOBAL_INDEX_LEB, s)));
all.sort_by_key(|&(o, _, _)| o);
let mut reloc = Vec::new();
leb_u(u64::from(code_index), &mut reloc);
leb_u(all.len() as u64, &mut reloc);
for (off, ty, sym) in &all {
reloc.push(*ty);
leb_u(u64::from(*off), &mut reloc);
leb_u(u64::from(*sym), &mut reloc);
}
let mut out = Vec::new();
out.extend_from_slice(b"\0asm");
out.extend_from_slice(&[0x01, 0x00, 0x00, 0x00]);
for s in &secs {
if !reemit(s.id) {
continue; }
if s.id == 10 {
put_section(&mut out, 10, &new_code);
} else {
put_section(&mut out, s.id, &module[s.start..s.end]);
}
}
put_custom(&mut out, "linking", &linking);
put_custom(&mut out, "reloc.CODE", &reloc);
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::vm::wasm::link;
use std::cell::RefCell;
use std::rc::Rc;
fn run(bytes: &[u8]) -> Vec<i64> {
let engine = wasmi::Engine::default();
let module = wasmi::Module::new(&engine, bytes).expect("valid wasm module");
let out: Rc<RefCell<Vec<i64>>> = Rc::new(RefCell::new(Vec::new()));
let mut store = wasmi::Store::new(&engine, out.clone());
let mut linker = wasmi::Linker::new(&engine);
linker
.func_wrap("env", "print_i64", |c: wasmi::Caller<'_, Rc<RefCell<Vec<i64>>>>, v: i64| {
c.data().borrow_mut().push(v);
})
.unwrap();
let instance = linker
.instantiate(&mut store, &module)
.unwrap()
.start(&mut store)
.unwrap();
instance.get_typed_func::<(), ()>(&store, "main").unwrap().call(&mut store, ()).unwrap();
let r = out.borrow().clone();
r
}
#[test]
fn relocatable_links_and_matches_standalone_over_a_corpus() {
if !link::toolchain_available() {
eprintln!("SKIP relocatable_links_and_matches_standalone_over_a_corpus: toolchain unavailable");
return;
}
let corpus: &[(&str, &str)] = &[
(
"calls",
"## To dbl (n: Int) -> Int:\n Return n + n.\n## Main\n Show 5.\n Show 2 + 3.\n Show dbl(21).\n",
),
(
"recursion",
"## To fib (n: Int) -> Int:\n If n is less than 2:\n Return n.\n Return fib(n - 1) + fib(n - 2).\n## Main\n Show fib(10).\n",
),
(
"recursion_control_flow",
"## To collatz (n: Int) -> Int:\n If n is equal to 1:\n Return 0.\n If n % 2 is equal to 0:\n Return 1 + collatz(n / 2).\n Return 1 + collatz(3 * n + 1).\n## Main\n Show collatz(27).\n",
),
(
"mutual_arith",
"## To f (n: Int) -> Int:\n Return n * n - n.\n## To g (n: Int) -> Int:\n Return f(n) + f(n - 1).\n## Main\n Show g(10).\n Show f(7) + g(3).\n",
),
(
"mutable_global",
"## Main\n Let mutable x be 0.\n Set x to 10.\n Set x to x + 32.\n Show x.\n Let mutable y be 100.\n Set y to y - x.\n Show y.\n",
),
];
let dir = std::env::temp_dir().join("logos_reloc_corpus");
for (name, src) in corpus {
let standalone = crate::compile::compile_to_wasm(src).unwrap_or_else(|e| panic!("compile {name}: {e:?}"));
let obj = module_to_relocatable(&standalone).unwrap_or_else(|e| panic!("relocatable {name}: {e:?}"));
let linked = link::link_objects(&[&obj], &dir).unwrap_or_else(|e| panic!("link {name}: {e:?}"));
let standalone_out = run(&standalone);
let linked_out = run(&linked);
assert_eq!(linked_out, standalone_out, "relocatable-linked output must equal standalone for {name}");
assert!(!standalone_out.is_empty(), "{name} should Show something");
}
}
fn fuzz_expr(state: &mut u64, depth: u32) -> (String, i64) {
let next = |s: &mut u64| {
*s ^= *s << 13;
*s ^= *s >> 7;
*s ^= *s << 17;
*s
};
if depth == 0 || next(state) % 3 == 0 {
let v = (next(state) % 20) as i64; return (v.to_string(), v);
}
let (ls, lv) = fuzz_expr(state, depth - 1);
let (rs, rv) = fuzz_expr(state, depth - 1);
match next(state) % 3 {
0 => (format!("({ls} + {rs})"), lv + rv),
1 => (format!("({ls} * {rs})"), lv * rv),
_ if lv >= rv => (format!("({ls} - {rs})"), lv - rv),
_ => (format!("({rs} - {ls})"), rv - lv),
}
}
#[test]
fn fuzz_relocatable_matches_standalone() {
if !link::toolchain_available() {
eprintln!("SKIP fuzz_relocatable_matches_standalone: toolchain unavailable");
return;
}
let dir = std::env::temp_dir().join("logos_reloc_fuzz");
for seed in 0..60u64 {
let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ 0x2545_F491_4F6C_DD1D;
let (expr, expected) = fuzz_expr(&mut state, 4);
let src = format!("## Main\n Show {expr}.\n");
let standalone = crate::compile::compile_to_wasm(&src).unwrap_or_else(|e| panic!("compile seed {seed}: {e:?}\n{src}"));
let obj = match module_to_relocatable(&standalone) {
Ok(o) => o,
Err(_) => continue, };
let linked = link::link_objects(&[&obj], &dir).unwrap_or_else(|e| panic!("link seed {seed}: {e:?}"));
let standalone_out = run(&standalone);
let linked_out = run(&linked);
assert_eq!(linked_out, standalone_out, "seed {seed}: linked != standalone for {src}");
assert_eq!(standalone_out, vec![expected], "seed {seed}: wrong value for {src}");
}
}
}