use wasm_encoder::{
BlockType, CodeSection, DataSection, EntityType, ExportKind, ExportSection, Function,
FunctionSection, ImportSection, Instruction, MemoryType, Module, TypeSection, ValType,
};
use crate::adapter::abi::{WasmEncoderBindgen, WitBridge};
use crate::adapter::func::AdapterFunc;
use crate::adapter::indices::{DispatchIndices, FunctionIndices};
use crate::adapter::names;
use wit_bindgen_core::abi::lift_from_memory;
fn build_task_return_loads(
result_ptr: u32,
result_type_id: cviz::model::ValueTypeId,
bridge: &WitBridge,
indices: &mut FunctionIndices,
) -> Vec<Instruction<'static>> {
let addr_local = indices.alloc_local(ValType::I32);
let result_type = bridge.get(result_type_id);
let mut out: Vec<Instruction<'static>> = vec![
Instruction::I32Const(result_ptr as i32),
Instruction::LocalSet(addr_local),
];
let mut bindgen = WasmEncoderBindgen::new(&bridge.sizes, addr_local, indices);
lift_from_memory(&bridge.resolve, &mut bindgen, (), &result_type);
out.extend(bindgen.into_instructions());
out
}
pub(crate) fn build_mem_module(with_realloc: bool, bump_start: u32) -> Module {
let mut module = Module::new();
if with_realloc {
let mut types = TypeSection::new();
types.ty().function(
[ValType::I32, ValType::I32, ValType::I32, ValType::I32],
[ValType::I32],
);
module.section(&types);
let mut fn_section = FunctionSection::new();
fn_section.function(0);
module.section(&fn_section);
}
{
let mut mem_section = wasm_encoder::MemorySection::new();
mem_section.memory(MemoryType {
minimum: 1,
maximum: None,
memory64: false,
shared: false,
page_size_log2: None,
});
module.section(&mem_section);
}
if with_realloc {
let mut globals = wasm_encoder::GlobalSection::new();
globals.global(
wasm_encoder::GlobalType {
val_type: ValType::I32,
mutable: true,
shared: false,
},
&wasm_encoder::ConstExpr::i32_const(bump_start as i32),
);
module.section(&globals);
}
{
let mut exports = ExportSection::new();
exports.export(names::ENV_MEMORY, ExportKind::Memory, 0);
if with_realloc {
exports.export(names::ENV_REALLOC, ExportKind::Func, 0);
}
module.section(&exports);
}
if with_realloc {
let mut code_section = CodeSection::new();
let mut rf = Function::new(vec![(1u32, ValType::I32)]);
rf.instruction(&Instruction::LocalGet(2)); rf.instruction(&Instruction::I32Const(1));
rf.instruction(&Instruction::I32Sub);
rf.instruction(&Instruction::I32Const(-1));
rf.instruction(&Instruction::I32Xor);
rf.instruction(&Instruction::GlobalGet(0));
rf.instruction(&Instruction::LocalGet(2));
rf.instruction(&Instruction::I32Const(1));
rf.instruction(&Instruction::I32Sub);
rf.instruction(&Instruction::I32Add);
rf.instruction(&Instruction::I32And);
rf.instruction(&Instruction::LocalTee(4));
rf.instruction(&Instruction::LocalGet(3));
rf.instruction(&Instruction::I32Add);
rf.instruction(&Instruction::GlobalSet(0));
rf.instruction(&Instruction::LocalGet(4));
rf.instruction(&Instruction::End);
code_section.function(&rf);
module.section(&code_section);
}
module
}
fn emit_wrapper_types(
types: &mut TypeSection,
indices: &mut DispatchIndices,
funcs: &[AdapterFunc],
) {
for func in funcs {
if func.is_async {
types.ty().function(func.core_params.iter().copied(), []);
} else if func.result_is_complex {
types
.ty()
.function(func.core_params.iter().copied(), [ValType::I32]);
} else {
types.ty().function(
func.core_params.iter().copied(),
func.core_results.iter().copied(),
);
}
indices.alloc_ty();
}
}
fn emit_sync_complex_handler_types(
types: &mut TypeSection,
indices: &mut DispatchIndices,
funcs: &[AdapterFunc],
out: &mut [Option<u32>],
) {
for (i, func) in funcs.iter().enumerate() {
if !func.is_async && func.result_is_complex {
out[i] = Some(indices.alloc_ty());
let mut params: Vec<ValType> = func.core_params.clone();
params.push(ValType::I32); types.ty().function(params, []);
}
}
}
fn emit_async_handler_types(
types: &mut TypeSection,
indices: &mut DispatchIndices,
funcs: &[AdapterFunc],
out: &mut [Option<u32>],
) {
for (i, func) in funcs.iter().enumerate() {
if func.is_async {
out[i] = Some(indices.alloc_ty());
let mut params: Vec<ValType> = func.core_params.clone();
if func.result_type_id.is_some() {
params.push(ValType::I32); }
types.ty().function(params, [ValType::I32]);
}
}
}
fn emit_custom_task_return_types(types: &mut TypeSection, funcs: &[AdapterFunc]) {
for func in funcs {
if func.is_async && func.result_is_complex {
types.ty().function(func.core_results.iter().copied(), []);
}
}
}
#[allow(clippy::too_many_arguments)]
fn emit_handler_imports(
imports: &mut ImportSection,
indices: &mut DispatchIndices,
funcs: &[AdapterFunc],
wrapper_ty_base: u32,
async_ds_tys: &[Option<u32>],
sync_complex_handler_tys: &[Option<u32>],
) -> u32 {
let base = indices.func;
for (i, func) in funcs.iter().enumerate() {
let ty = if func.is_async {
async_ds_tys[i].expect("async_ds_ty must be set for async func")
} else if let Some(handler_ty) = sync_complex_handler_tys[i] {
handler_ty
} else {
wrapper_ty_base + i as u32
};
indices.alloc_func();
imports.import(
names::ENV_INSTANCE,
&names::env_handler_fn(i),
EntityType::Function(ty),
);
}
base
}
#[allow(clippy::too_many_arguments)]
fn emit_task_return_imports(
imports: &mut ImportSection,
indices: &mut DispatchIndices,
funcs: &[AdapterFunc],
void_void_ty: u32,
void_i32_ty: u32,
void_i64_ty: u32,
void_f32_ty: u32,
void_f64_ty: u32,
) -> Vec<Option<u32>> {
let mut trf: Vec<Option<u32>> = vec![None; funcs.len()];
let mut custom_tr_ty_idx = void_void_ty + 1;
for (i, func) in funcs.iter().enumerate() {
if !func.is_async {
continue;
}
let tr_ty = if func.result_type_id.is_none() {
void_void_ty
} else if func.result_is_complex {
let ty = custom_tr_ty_idx;
custom_tr_ty_idx += 1;
ty
} else {
match func.core_results.first() {
Some(ValType::I32) => void_i32_ty,
Some(ValType::I64) => void_i64_ty,
Some(ValType::F32) => void_f32_ty,
Some(ValType::F64) => void_f64_ty,
_ => void_i32_ty,
}
};
trf[i] = Some(indices.alloc_func());
imports.import(
names::ENV_INSTANCE,
&names::env_task_return_fn(i),
EntityType::Function(tr_ty),
);
}
trf
}
struct WaitLoopCtx {
waitable_new_fn: u32,
waitable_join_fn: u32,
waitable_wait_fn: u32,
subtask_drop_fn: u32,
waitable_drop_fn: u32,
event_ptr: u32,
}
fn emit_wait_loop(f: &mut Function, st: u32, ws: u32, ctx: &WaitLoopCtx) {
f.instruction(&Instruction::LocalGet(st));
f.instruction(&Instruction::I32Const(4));
f.instruction(&Instruction::I32ShrU);
f.instruction(&Instruction::LocalSet(st));
f.instruction(&Instruction::LocalGet(st));
f.instruction(&Instruction::If(BlockType::Empty));
f.instruction(&Instruction::Call(ctx.waitable_new_fn));
f.instruction(&Instruction::LocalSet(ws));
f.instruction(&Instruction::LocalGet(st));
f.instruction(&Instruction::LocalGet(ws));
f.instruction(&Instruction::Call(ctx.waitable_join_fn));
f.instruction(&Instruction::LocalGet(ws));
f.instruction(&Instruction::I32Const(ctx.event_ptr as i32));
f.instruction(&Instruction::Call(ctx.waitable_wait_fn));
f.instruction(&Instruction::Drop);
f.instruction(&Instruction::LocalGet(st));
f.instruction(&Instruction::Call(ctx.subtask_drop_fn));
f.instruction(&Instruction::LocalGet(ws));
f.instruction(&Instruction::Call(ctx.waitable_drop_fn));
f.instruction(&Instruction::End);
}
struct DispatchCodeCtx<'a> {
before_import_fn: Option<u32>,
after_import_fn: Option<u32>,
blocking: Option<BlockingConfig>,
handler_import_fn_base: u32,
task_return_fns: &'a [Option<u32>],
wait: WaitLoopCtx,
}
struct BlockingConfig {
import_fn: u32,
result_ptr: u32,
}
fn emit_wrapper_body(
f: &mut Function,
fi: usize,
func: &AdapterFunc,
ctx: &DispatchCodeCtx<'_>,
subtask_local: u32,
ws_local: u32,
task_return_loads: Option<&[Instruction<'static>]>,
) -> anyhow::Result<()> {
let has_result = func.result_type_id.is_some();
if ctx.blocking.is_some() && has_result && !func.is_async {
anyhow::bail!(
"Function '{}' returns a value but the middleware exports \
`should-block-call`. Tier-1 blocking is only supported for \
void-returning functions because the adapter cannot synthesize \
a return value when the call is blocked. Tier-3 (read-write \
interception) will support this in the future.",
func.name
);
}
emit_before_phase(f, func, ctx, subtask_local, ws_local);
emit_blocking_phase(f, fi, func, ctx, subtask_local, ws_local)?;
let result_local_idx = emit_handler_call_phase(f, fi, func, ctx, subtask_local, ws_local);
emit_after_phase(f, func, ctx, subtask_local, ws_local);
emit_return_phase(f, fi, func, ctx, result_local_idx, task_return_loads);
f.instruction(&Instruction::End);
Ok(())
}
fn emit_before_phase(
f: &mut Function,
func: &AdapterFunc,
ctx: &DispatchCodeCtx<'_>,
subtask_local: u32,
ws_local: u32,
) {
if let Some(before_fn) = ctx.before_import_fn {
f.instruction(&Instruction::I32Const(func.name_offset as i32));
f.instruction(&Instruction::I32Const(func.name_len as i32));
f.instruction(&Instruction::Call(before_fn));
f.instruction(&Instruction::LocalSet(subtask_local));
emit_wait_loop(f, subtask_local, ws_local, &ctx.wait);
}
}
fn emit_blocking_phase(
f: &mut Function,
fi: usize,
func: &AdapterFunc,
ctx: &DispatchCodeCtx<'_>,
subtask_local: u32,
ws_local: u32,
) -> anyhow::Result<()> {
let Some(BlockingConfig {
import_fn: block_fn,
result_ptr: blk_ptr,
}) = ctx.blocking
else {
return Ok(());
};
if func.result_type_id.is_some() {
anyhow::bail!(
"Function '{}' is async with a return value and the middleware \
exports `should-block-call`. Tier-1 blocking is not supported \
for async functions with results. Tier-3 (read-write \
interception) will support this in the future.",
func.name
);
}
f.instruction(&Instruction::I32Const(func.name_offset as i32));
f.instruction(&Instruction::I32Const(func.name_len as i32));
f.instruction(&Instruction::I32Const(blk_ptr as i32));
f.instruction(&Instruction::Call(block_fn));
f.instruction(&Instruction::LocalSet(subtask_local));
emit_wait_loop(f, subtask_local, ws_local, &ctx.wait);
f.instruction(&Instruction::I32Const(blk_ptr as i32));
f.instruction(&Instruction::I32Load(wasm_encoder::MemArg {
offset: 0,
align: 0,
memory_index: 0,
}));
f.instruction(&Instruction::If(BlockType::Empty));
if let Some(tr_fn) = ctx.task_return_fns[fi] {
f.instruction(&Instruction::Call(tr_fn));
}
f.instruction(&Instruction::Return);
f.instruction(&Instruction::End);
Ok(())
}
fn emit_handler_call_phase(
f: &mut Function,
fi: usize,
func: &AdapterFunc,
ctx: &DispatchCodeCtx<'_>,
subtask_local: u32,
ws_local: u32,
) -> Option<u32> {
let handler_fn = ctx.handler_import_fn_base + fi as u32;
if func.is_async {
for (pi, _) in func.core_params.iter().enumerate() {
f.instruction(&Instruction::LocalGet(pi as u32));
}
if let Some(result_ptr) = func.async_result_mem_offset {
f.instruction(&Instruction::I32Const(result_ptr as i32));
}
f.instruction(&Instruction::Call(handler_fn));
f.instruction(&Instruction::LocalSet(subtask_local));
emit_wait_loop(f, subtask_local, ws_local, &ctx.wait);
None
} else if func.result_is_complex {
let result_buf = func
.sync_result_mem_offset
.expect("sync_result_mem_offset must be set for sync complex");
for (pi, _) in func.core_params.iter().enumerate() {
f.instruction(&Instruction::LocalGet(pi as u32));
}
f.instruction(&Instruction::I32Const(result_buf as i32));
f.instruction(&Instruction::Call(handler_fn));
None
} else {
let result_local = if ctx.after_import_fn.is_some() && func.result_type_id.is_some() {
Some(func.core_params.len() as u32)
} else {
None
};
for (pi, _) in func.core_params.iter().enumerate() {
f.instruction(&Instruction::LocalGet(pi as u32));
}
f.instruction(&Instruction::Call(handler_fn));
if let Some(local_idx) = result_local {
f.instruction(&Instruction::LocalSet(local_idx));
}
result_local
}
}
fn emit_after_phase(
f: &mut Function,
func: &AdapterFunc,
ctx: &DispatchCodeCtx<'_>,
subtask_local: u32,
ws_local: u32,
) {
if let Some(after_fn) = ctx.after_import_fn {
f.instruction(&Instruction::I32Const(func.name_offset as i32));
f.instruction(&Instruction::I32Const(func.name_len as i32));
f.instruction(&Instruction::Call(after_fn));
f.instruction(&Instruction::LocalSet(subtask_local));
emit_wait_loop(f, subtask_local, ws_local, &ctx.wait);
}
}
fn emit_return_phase(
f: &mut Function,
fi: usize,
func: &AdapterFunc,
ctx: &DispatchCodeCtx<'_>,
result_local_idx: Option<u32>,
task_return_loads: Option<&[Instruction<'static>]>,
) {
if func.is_async {
if func.async_result_mem_offset.is_some() {
if let Some(tr_fn) = ctx.task_return_fns[fi] {
let loads = task_return_loads
.expect("async-with-result wrapper must have pre-built task_return_loads");
for inst in loads {
f.instruction(inst);
}
f.instruction(&Instruction::Call(tr_fn));
}
} else if let Some(tr_fn) = ctx.task_return_fns[fi] {
f.instruction(&Instruction::Call(tr_fn));
}
} else if func.result_is_complex {
let result_buf = func
.sync_result_mem_offset
.expect("sync_result_mem_offset must be set for sync complex");
f.instruction(&Instruction::I32Const(result_buf as i32));
} else if let Some(local_idx) = result_local_idx {
f.instruction(&Instruction::LocalGet(local_idx));
}
}
fn emit_function_name_data(module: &mut Module, funcs: &[AdapterFunc]) {
let mut data_section = DataSection::new();
let all_names: Vec<u8> = funcs
.iter()
.flat_map(|f| f.name.as_bytes().iter().copied())
.collect();
data_section.active(0, &wasm_encoder::ConstExpr::i32_const(0), all_names);
module.section(&data_section);
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn build_dispatch_module(
funcs: &[AdapterFunc],
has_before: bool,
has_after: bool,
has_blocking: bool,
event_ptr: u32,
block_result_ptr: Option<u32>,
bridge: &WitBridge,
) -> anyhow::Result<Vec<u8>> {
let has_async = funcs.iter().any(|f| f.is_async);
let has_async_machinery = has_async || has_before || has_after || has_blocking;
let mut module = Module::new();
let mut indices = DispatchIndices::new();
let hook_ty: u32;
let block_ty: u32;
let wrapper_ty_base: u32;
let mut async_ds_tys: Vec<Option<u32>> = vec![None; funcs.len()];
let mut sync_complex_handler_tys: Vec<Option<u32>> = vec![None; funcs.len()];
let waitable_new_ty: u32;
let waitable_join_ty: u32;
let waitable_wait_ty: u32;
let void_i32_ty: u32; let void_void_ty: u32; let void_i64_ty: u32;
let void_f32_ty: u32;
let void_f64_ty: u32;
{
let mut types = TypeSection::new();
hook_ty = indices.alloc_ty();
types
.ty()
.function([ValType::I32, ValType::I32], [ValType::I32]);
block_ty = indices.alloc_ty();
types
.ty()
.function([ValType::I32, ValType::I32, ValType::I32], [ValType::I32]);
wrapper_ty_base = indices.ty;
emit_wrapper_types(&mut types, &mut indices, funcs);
emit_sync_complex_handler_types(
&mut types,
&mut indices,
funcs,
&mut sync_complex_handler_tys,
);
if has_async_machinery {
emit_async_handler_types(&mut types, &mut indices, funcs, &mut async_ds_tys);
waitable_new_ty = indices.alloc_ty();
types.ty().function([], [ValType::I32]);
waitable_join_ty = indices.alloc_ty();
types.ty().function([ValType::I32, ValType::I32], []);
waitable_wait_ty = indices.alloc_ty();
types
.ty()
.function([ValType::I32, ValType::I32], [ValType::I32]);
void_i32_ty = indices.alloc_ty();
types.ty().function([ValType::I32], []);
void_i64_ty = indices.alloc_ty();
types.ty().function([ValType::I64], []);
void_f32_ty = indices.alloc_ty();
types.ty().function([ValType::F32], []);
void_f64_ty = indices.alloc_ty();
types.ty().function([ValType::F64], []);
void_void_ty = indices.alloc_ty();
types.ty().function([], []);
emit_custom_task_return_types(&mut types, funcs);
} else {
waitable_new_ty = 0;
waitable_join_ty = 0;
waitable_wait_ty = 0;
void_i32_ty = 0;
void_i64_ty = 0;
void_f32_ty = 0;
void_f64_ty = 0;
void_void_ty = 0;
}
module.section(&types);
}
let before_import_fn: Option<u32>;
let after_import_fn: Option<u32>;
let blocking_import_fn: Option<u32>;
let handler_import_fn_base: u32;
let waitable_new_fn: u32;
let waitable_join_fn: u32;
let waitable_wait_fn: u32;
let waitable_drop_fn: u32;
let subtask_drop_fn: u32;
let task_return_fns: Vec<Option<u32>>;
{
use crate::contract::{
TIER1_AFTER_ENV_SLOTS, TIER1_BEFORE_ENV_SLOTS, TIER1_BLOCKING_ENV_SLOTS,
};
let mut imports = ImportSection::new();
imports.import(
names::ENV_INSTANCE,
names::ENV_MEMORY,
EntityType::Memory(MemoryType {
minimum: 1,
maximum: None,
memory64: false,
shared: false,
page_size_log2: None,
}),
);
before_import_fn = if has_before {
let idx = indices.alloc_func();
imports.import(
names::ENV_INSTANCE,
TIER1_BEFORE_ENV_SLOTS[0],
EntityType::Function(hook_ty),
);
Some(idx)
} else {
None
};
after_import_fn = if has_after {
let idx = indices.alloc_func();
imports.import(
names::ENV_INSTANCE,
TIER1_AFTER_ENV_SLOTS[0],
EntityType::Function(hook_ty),
);
Some(idx)
} else {
None
};
blocking_import_fn = if has_blocking {
let idx = indices.alloc_func();
imports.import(
names::ENV_INSTANCE,
TIER1_BLOCKING_ENV_SLOTS[0],
EntityType::Function(block_ty),
);
Some(idx)
} else {
None
};
handler_import_fn_base = emit_handler_imports(
&mut imports,
&mut indices,
funcs,
wrapper_ty_base,
&async_ds_tys,
&sync_complex_handler_tys,
);
if has_async_machinery {
waitable_new_fn = indices.alloc_func();
imports.import(
names::ENV_INSTANCE,
names::ENV_WAITABLE_NEW,
EntityType::Function(waitable_new_ty),
);
waitable_join_fn = indices.alloc_func();
imports.import(
names::ENV_INSTANCE,
names::ENV_WAITABLE_JOIN,
EntityType::Function(waitable_join_ty),
);
waitable_wait_fn = indices.alloc_func();
imports.import(
names::ENV_INSTANCE,
names::ENV_WAITABLE_WAIT,
EntityType::Function(waitable_wait_ty),
);
waitable_drop_fn = indices.alloc_func();
imports.import(
names::ENV_INSTANCE,
names::ENV_WAITABLE_DROP,
EntityType::Function(void_i32_ty),
);
subtask_drop_fn = indices.alloc_func();
imports.import(
names::ENV_INSTANCE,
names::ENV_SUBTASK_DROP,
EntityType::Function(void_i32_ty),
);
let trf = emit_task_return_imports(
&mut imports,
&mut indices,
funcs,
void_void_ty,
void_i32_ty,
void_i64_ty,
void_f32_ty,
void_f64_ty,
);
task_return_fns = trf;
} else {
waitable_new_fn = 0;
waitable_join_fn = 0;
waitable_wait_fn = 0;
waitable_drop_fn = 0;
subtask_drop_fn = 0;
task_return_fns = vec![None; funcs.len()];
}
module.section(&imports);
}
let wrapper_fn_base: u32;
{
let mut fn_section = FunctionSection::new();
wrapper_fn_base = handler_import_fn_base
+ funcs.len() as u32
+ if has_async_machinery {
5 + funcs.iter().filter(|f| f.is_async).count() as u32
} else {
0
};
for (i, _) in funcs.iter().enumerate() {
fn_section.function(wrapper_ty_base + i as u32);
}
module.section(&fn_section);
}
{
let mut exports = ExportSection::new();
for (i, func) in funcs.iter().enumerate() {
exports.export(&func.name, ExportKind::Func, wrapper_fn_base + i as u32);
}
module.section(&exports);
}
let blocking = blocking_import_fn
.zip(block_result_ptr)
.map(|(import_fn, result_ptr)| BlockingConfig {
import_fn,
result_ptr,
});
let code_ctx = DispatchCodeCtx {
before_import_fn,
after_import_fn,
blocking,
handler_import_fn_base,
task_return_fns: &task_return_fns,
wait: WaitLoopCtx {
waitable_new_fn,
waitable_join_fn,
waitable_wait_fn,
subtask_drop_fn,
waitable_drop_fn,
event_ptr,
},
};
{
let mut code_section = CodeSection::new();
for (fi, func) in funcs.iter().enumerate() {
let mut indices = FunctionIndices::new(func.core_params.len() as u32);
let subtask_local = indices.alloc_local(ValType::I32);
let ws_local = indices.alloc_local(ValType::I32);
let task_return_loads: Option<Vec<Instruction<'static>>> = match (
func.is_async,
func.async_result_mem_offset,
func.result_type_id,
) {
(true, Some(result_ptr), Some(rid)) => Some(build_task_return_loads(
result_ptr,
rid,
bridge,
&mut indices,
)),
_ => None,
};
let mut f = Function::new_with_locals_types(indices.into_locals());
emit_wrapper_body(
&mut f,
fi,
func,
&code_ctx,
subtask_local,
ws_local,
task_return_loads.as_deref(),
)?;
code_section.function(&f);
}
module.section(&code_section);
}
emit_function_name_data(&mut module, funcs);
Ok(module.finish().to_vec())
}