Lines
99.39 %
Functions
57.62 %
Branches
100 %
use std::collections::HashMap;
use anodized::spec;
use tracing::debug;
use super::layout::OutputSerializer;
use crate::ast::WasmType;
use wasm_encoder::{
ArrayType, CodeSection, CompositeInnerType, CompositeType, DataCountSection, DataSection,
EntityType as WasmEntityType, ExportKind, ExportSection, FieldType, Function, FunctionSection,
HeapType, ImportSection, Instruction, MemorySection, MemoryType, Module, RefType, StorageType,
StructType, SubType, TypeSection, ValType,
};
/// Base index for user-allocatable locals (after pre-allocated slots)
pub const LOCAL_POOL_BASE: u32 = 5;
pub struct CompileContext {
types: TypeSection,
imports: ImportSection,
functions: FunctionSection,
memories: MemorySection,
exports: ExportSection,
data: DataSection,
codes: CodeSection,
data_count: u32,
type_count: u32,
import_func_count: u32,
local_func_count: u32,
type_cache: HashMap<Vec<ValType>, HashMap<Vec<ValType>, u32>>,
type_names: HashMap<String, u32>,
func_names: HashMap<String, u32>,
pending_helpers: Vec<Function>,
next_local: u32,
local_types: Vec<(WasmType, u32)>,
serializer: OutputSerializer,
}
impl CompileContext {
pub fn new() -> Self {
debug!("initializing compile context");
let mut ctx = Self {
types: TypeSection::new(),
imports: ImportSection::new(),
functions: FunctionSection::new(),
memories: MemorySection::new(),
exports: ExportSection::new(),
data: DataSection::new(),
codes: CodeSection::new(),
data_count: 0,
type_count: 0,
import_func_count: 0,
local_func_count: 0,
type_cache: HashMap::new(),
type_names: HashMap::new(),
func_names: HashMap::new(),
pending_helpers: Vec::new(),
next_local: LOCAL_POOL_BASE,
local_types: Vec::new(),
serializer: OutputSerializer::new(super::expr::LOCAL_OUTPUT_BASE),
ctx.register_type("i8_array");
ctx.register_struct_type("ratio", &[ValType::I64, ValType::I64]);
ctx.register_cons_type();
ctx.register_import("env", "get_output_offset", &[], &[ValType::I32]);
ctx.register_import("env", "symbol_resolve", &[ValType::I32, ValType::I32], &[]);
ctx.register_import(
"env",
"log",
&[ValType::I32, ValType::I32, ValType::I32],
&[],
);
ctx.register_import("env", "get_input_offset", &[], &[ValType::I32]);
ctx.register_import("env", "get_strings_offset", &[], &[ValType::I32]);
ctx.register_import("env", "get_input_entities_count", &[], &[ValType::I32]);
ctx.register_import("env", "get_timestamp", &[], &[ValType::I64]);
ctx.register_import("env", "generate_uuid", &[ValType::I32], &[]);
"write_string",
&[ValType::I32, ValType::I32],
&[ValType::I32],
"write_bytes",
ctx.register_function("should_apply", &[], &[ValType::I32]);
ctx.register_ratio_helpers();
ctx.register_string_eq_helper();
ctx.register_function("process", &[], &[]);
ctx.export_func("should_apply");
ctx.export_func("process");
ctx.memories.memory(MemoryType {
minimum: 1,
maximum: None,
memory64: false,
shared: false,
page_size_log2: None,
});
ctx.exports.export("memory", ExportKind::Memory, 0);
debug!("compile context initialized");
ctx
pub fn ratio_ref(&self) -> ValType {
ValType::Ref(RefType {
nullable: true,
heap_type: HeapType::Concrete(self.type_idx("ratio")),
})
pub fn cons_ref(&self) -> ValType {
heap_type: HeapType::Concrete(self.type_idx("cons")),
pub fn alloc_local(&mut self, ty: WasmType) -> u32 {
let idx = self.next_local;
self.local_types.push((ty, idx));
self.next_local += 1;
idx
pub fn reset_locals(&mut self) {
self.next_local = LOCAL_POOL_BASE;
self.local_types.clear();
self.serializer = OutputSerializer::new(super::expr::LOCAL_OUTPUT_BASE);
pub fn build_locals_declaration(&self) -> Vec<(u32, ValType)> {
let preallocated: [(u32, ValType); 5] = [
(1, self.string_ref()),
(1, ValType::I32),
(1, self.ratio_ref()),
];
let mut locals: Vec<(u32, ValType)> = preallocated.to_vec();
for &(ty, _) in &self.local_types {
locals.push((1, self.wasm_val_type(ty)));
locals
pub fn string_ref(&self) -> ValType {
heap_type: HeapType::Concrete(self.type_idx("i8_array")),
pub fn wasm_val_type(&self, ty: WasmType) -> ValType {
match ty {
WasmType::I32 => ValType::I32,
WasmType::Ratio => self.ratio_ref(),
WasmType::ConsRef => self.cons_ref(),
WasmType::StringRef => self.string_ref(),
#[spec(
captures: [self.type_count as prev_count],
ensures: [self.type_count == prev_count + 1, *output == prev_count],
)]
pub fn register_type(&mut self, name: &str) -> u32 {
let idx = self.type_count;
self.types.ty().subtype(&SubType {
is_final: true,
supertype_idx: None,
composite_type: CompositeType {
inner: CompositeInnerType::Array(ArrayType(FieldType {
element_type: StorageType::I8,
mutable: true,
})),
describes: None,
descriptor: None,
},
self.type_count += 1;
self.type_names.insert(name.to_string(), idx);
pub fn register_struct_type(&mut self, name: &str, fields: &[ValType]) -> u32 {
let struct_fields: Vec<FieldType> = fields
.iter()
.map(|vt| FieldType {
element_type: StorageType::Val(*vt),
mutable: false,
.collect();
inner: CompositeInnerType::Struct(StructType {
fields: struct_fields.into_boxed_slice(),
}),
captures: [self.import_func_count as prev_count],
ensures: [self.import_func_count == prev_count + 1, *output == prev_count],
pub fn register_import(
&mut self,
module: &str,
name: &str,
params: &[ValType],
results: &[ValType],
) -> u32 {
let type_idx = self.get_or_create_func_type(params, results);
self.imports
.import(module, name, WasmEntityType::Function(type_idx));
let func_idx = self.import_func_count;
self.import_func_count += 1;
self.func_names.insert(name.to_string(), func_idx);
func_idx
captures: [self.local_func_count as prev_count],
ensures: [self.local_func_count == prev_count + 1],
pub fn register_function(&mut self, name: &str, params: &[ValType], results: &[ValType]) {
self.functions.function(type_idx);
let func_idx = self.import_func_count + self.local_func_count;
self.local_func_count += 1;
pub fn export_func(&mut self, name: &str) {
let idx = self.func_names[name];
self.exports.export(name, ExportKind::Func, idx);
pub fn func(&self, name: &str) -> u32 {
self.func_names[name]
pub fn type_idx(&self, name: &str) -> u32 {
self.type_names[name]
pub fn serializer(&mut self) -> &mut OutputSerializer {
&mut self.serializer
#[spec(ensures: [*output < self.type_count])]
fn get_or_create_func_type(&mut self, params: &[ValType], results: &[ValType]) -> u32 {
if let Some(inner) = self.type_cache.get(params)
&& let Some(&idx) = inner.get(results)
{
return idx;
self.types
.ty()
.function(params.iter().copied(), results.iter().copied());
self.type_cache
.entry(params.to_vec())
.or_default()
.insert(results.to_vec(), idx);
captures: [self.data_count as prev_count],
ensures: [self.data_count == prev_count + 1, *output == prev_count],
pub fn add_data(&mut self, bytes: &[u8]) -> u32 {
let idx = self.data_count;
debug!(idx, len = bytes.len(), "adding data segment");
self.data.passive(bytes.iter().copied());
self.data_count += 1;
fn register_cons_type(&mut self) {
let cons_idx = self.type_count;
let cons_nullable_ref = ValType::Ref(RefType {
heap_type: HeapType::Concrete(cons_idx),
let struct_fields = vec![
FieldType {
element_type: StorageType::Val(ValType::I32),
element_type: StorageType::Val(cons_nullable_ref),
self.type_names.insert("cons".to_string(), cons_idx);
fn register_string_eq_helper(&mut self) {
let string_ref = self.string_ref();
self.register_function("string_eq", &[string_ref, string_ref], &[ValType::I32]);
self.build_string_eq_body();
fn build_string_eq_body(&mut self) {
let i8_array_idx = self.type_idx("i8_array");
// params: $a=0 (ref null $i8_array), $b=1 (ref null $i8_array)
// locals: $len=2 (i32), $i=3 (i32)
let mut f = Function::new([(1, ValType::I32), (1, ValType::I32)]);
// len = array.len(a)
f.instruction(&Instruction::LocalGet(0));
f.instruction(&Instruction::ArrayLen);
f.instruction(&Instruction::LocalSet(2));
// if array.len(b) != len → return 0
f.instruction(&Instruction::LocalGet(1));
f.instruction(&Instruction::LocalGet(2));
f.instruction(&Instruction::I32Ne);
f.instruction(&Instruction::If(wasm_encoder::BlockType::Empty));
f.instruction(&Instruction::I32Const(0));
f.instruction(&Instruction::Return);
f.instruction(&Instruction::End);
// i = 0
f.instruction(&Instruction::LocalSet(3));
// block $exit
f.instruction(&Instruction::Block(wasm_encoder::BlockType::Empty));
// loop $cmp
f.instruction(&Instruction::Loop(wasm_encoder::BlockType::Empty));
// if i >= len → break (equal)
f.instruction(&Instruction::LocalGet(3));
f.instruction(&Instruction::I32GeU);
f.instruction(&Instruction::BrIf(1));
// if a[i] != b[i] → return 0
f.instruction(&Instruction::ArrayGetU(i8_array_idx));
// i++
f.instruction(&Instruction::I32Const(1));
f.instruction(&Instruction::I32Add);
f.instruction(&Instruction::Br(0));
f.instruction(&Instruction::End); // end loop
f.instruction(&Instruction::End); // end block
// return 1 (equal)
self.pending_helpers.push(f);
fn register_ratio_helpers(&mut self) {
let ratio_ref = self.ratio_ref();
self.register_function("gcd", &[ValType::I64, ValType::I64], &[ValType::I64]);
self.register_function("ratio_new", &[ValType::I64, ValType::I64], &[ratio_ref]);
self.register_function("ratio_add", &[ratio_ref, ratio_ref], &[ratio_ref]);
self.register_function("ratio_sub", &[ratio_ref, ratio_ref], &[ratio_ref]);
self.register_function("ratio_mul", &[ratio_ref, ratio_ref], &[ratio_ref]);
self.register_function("ratio_div", &[ratio_ref, ratio_ref], &[ratio_ref]);
self.register_function("ratio_eq", &[ratio_ref, ratio_ref], &[ValType::I32]);
self.register_function("ratio_lt", &[ratio_ref, ratio_ref], &[ValType::I32]);
self.register_function("ratio_from_i64", &[ValType::I64], &[ratio_ref]);
self.build_gcd_body();
self.build_ratio_new_body();
self.build_ratio_binop_bodies();
self.build_ratio_cmp_bodies();
self.build_ratio_from_i64_body();
fn build_gcd_body(&mut self) {
// Euclidean GCD: gcd(a, b) while b != 0 { t = b; b = a % b; a = t }; abs(a)
let mut f = Function::new([(1, ValType::I64)]); // local $t
// params: $a=0, $b=1, local: $t=2
// if b == 0, break
f.instruction(&Instruction::I64Eqz);
// t = b
// b = a % b
f.instruction(&Instruction::I64RemS);
f.instruction(&Instruction::LocalSet(1));
// a = t
f.instruction(&Instruction::LocalSet(0));
f.instruction(&Instruction::End); // loop
f.instruction(&Instruction::End); // block
// return abs(a)
f.instruction(&Instruction::I64Const(0));
f.instruction(&Instruction::I64Sub);
f.instruction(&Instruction::I64LtS);
f.instruction(&Instruction::Select);
fn build_ratio_new_body(&mut self) {
let ratio_idx = self.type_idx("ratio");
let gcd_func = self.func("gcd");
// ratio_new(num, denom) -> reduce by GCD, normalize sign
// params: $num=0, $denom=1, locals: $g=2
let mut f = Function::new([(1, ValType::I64)]);
// g = gcd(num, denom)
f.instruction(&Instruction::Call(gcd_func));
// num = num / g
f.instruction(&Instruction::I64DivS);
// denom = denom / g
// if denom < 0 { num = -num; denom = -denom }
// struct.new $ratio (num, denom)
f.instruction(&Instruction::StructNew(ratio_idx));
fn build_ratio_binop_bodies(&mut self) {
let ratio_new = self.func("ratio_new");
// ratio_add(a, b): (a.num * b.denom + b.num * a.denom, a.denom * b.denom)
let mut f = Function::new([]);
// a.num * b.denom
f.instruction(&Instruction::StructGet {
struct_type_index: ratio_idx,
field_index: 0,
field_index: 1,
f.instruction(&Instruction::I64Mul);
// b.num * a.denom
f.instruction(&Instruction::I64Add);
// a.denom * b.denom
f.instruction(&Instruction::Call(ratio_new));
// ratio_sub(a, b): (a.num * b.denom - b.num * a.denom, a.denom * b.denom)
// ratio_mul(a, b): (a.num * b.num, a.denom * b.denom)
// ratio_div(a, b): (a.num * b.denom, a.denom * b.num)
fn build_ratio_cmp_bodies(&mut self) {
// ratio_eq(a, b): a.num * b.denom == b.num * a.denom
f.instruction(&Instruction::I64Eq);
// ratio_lt(a, b): a.num * b.denom < b.num * a.denom
fn build_ratio_from_i64_body(&mut self) {
// ratio_from_i64(n) -> struct.new $ratio (n, 1)
f.instruction(&Instruction::I64Const(1));
pub fn add_should_apply(&mut self, f: Function) {
debug!("emitting should_apply function");
self.codes.function(&f);
for helper in self.pending_helpers.drain(..) {
self.codes.function(&helper);
pub fn default_should_apply() -> Function {
f
pub fn add_process(&mut self, f: Function) {
debug!("emitting process function");
pub fn finish(self) -> Vec<u8> {
debug!(data_segments = self.data_count, "assembling WASM module");
let mut module = Module::new();
module.section(&self.types);
module.section(&self.imports);
module.section(&self.functions);
module.section(&self.memories);
module.section(&self.exports);
module.section(&DataCountSection {
count: self.data_count,
module.section(&self.codes);
module.section(&self.data);
module.finish()
impl Default for CompileContext {
fn default() -> Self {
Self::new()
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_ratio_type_registered() {
let ctx = CompileContext::new();
assert!(ctx.type_names.contains_key("ratio"));
assert!(ctx.type_names.contains_key("i8_array"));
fn test_ratio_helpers_registered() {
for name in [
"gcd",
"ratio_new",
"ratio_add",
"ratio_sub",
"ratio_mul",
"ratio_div",
"ratio_eq",
"ratio_lt",
"ratio_from_i64",
] {
assert!(
ctx.func_names.contains_key(name),
"missing helper function: {name}"
fn test_host_imports_registered() {
"get_output_offset",
"get_input_offset",
"get_strings_offset",
"get_input_entities_count",
"get_timestamp",
"generate_uuid",
"symbol_resolve",
"missing host import: {name}"
fn test_ratio_ref_type() {
let ratio_ref = ctx.ratio_ref();
match ratio_ref {
ValType::Ref(r) => {
assert!(r.nullable);
assert_eq!(r.heap_type, HeapType::Concrete(ctx.type_idx("ratio")));
_ => panic!("expected Ref type"),
fn test_valid_wasm_with_ratio_helpers() {
let mut ctx = CompileContext::new();
ctx.add_should_apply(CompileContext::default_should_apply());
ctx.add_process(f);
let wasm = ctx.finish();
assert_eq!(&wasm[0..4], b"\0asm");