curlywas/src/emit.rs

use std::collections::HashMap;

use wasm_encoder::{
    BlockType, CodeSection, EntityType, Export, ExportSection, Function, FunctionSection,
    GlobalType, ImportSection, Instruction, MemArg, MemoryType, Module, TypeSection, ValType,
};

use crate::ast;

pub fn emit(script: &ast::Script) -> Vec<u8> {
    let mut module = Module::new();

    let function_types = collect_function_types(script);
    {
        let mut types = TypeSection::new();
        let mut type_vec: Vec<_> = function_types.iter().map(|(k, v)| (*v, k)).collect();
        type_vec.sort();
        for (_, (params, result)) in type_vec {
            let params: Vec<_> = params.iter().cloned().map(map_type).collect();
            let result: Vec<_> = result.iter().cloned().map(map_type).collect();
            types.function(params, result);
        }
        module.section(&types);
    }

    let mut globals: HashMap<&str, u32> = HashMap::new();

    {
        let mut imports = ImportSection::new();

        for import in &script.imports {
            let (module, name) = if let Some(dot_index) = import.import.find('.') {
                (
                    &import.import[..dot_index],
                    Some(&import.import[(dot_index + 1)..]),
                )
            } else {
                (import.import, None)
            };
            let type_: EntityType = match import.type_ {
                ast::ImportType::Memory(min_size) => MemoryType {
                    minimum: min_size as u64,
                    maximum: None,
                    memory64: false,
                }
                .into(),
                ast::ImportType::Variable {
                    type_,
                    name,
                    mutable,
                } => {
                    globals.insert(name, globals.len() as u32);
                    GlobalType {
                        val_type: map_type(type_),
                        mutable,
                    }
                    .into()
                }
            };
            imports.import(module, name, type_);
        }

        module.section(&imports);
    }

    {
        let mut functions = FunctionSection::new();
        let mut exports = ExportSection::new();
        let mut code = CodeSection::new();

        for (index, func) in script.functions.iter().enumerate() {
            let type_ = *function_types.get(&function_type_key(func)).unwrap();
            functions.function(type_ as u32);
            if func.export {
                exports.export(func.name, Export::Function(index as u32));
            }

            code.function(&emit_function(func, &globals));
        }

        module.section(&functions);
        module.section(&exports);
        module.section(&code);
    }

    module.finish()
}

type FunctionTypeKey = (Vec<ast::Type>, Option<ast::Type>);

fn collect_function_types(script: &ast::Script) -> HashMap<FunctionTypeKey, usize> {
    let mut types: HashMap<FunctionTypeKey, usize> = HashMap::new();

    for func in &script.functions {
        let index = types.len();
        types
            .entry(function_type_key(func))
            .or_insert_with(|| index);
    }

    types
}

fn function_type_key(func: &ast::Function) -> FunctionTypeKey {
    let param_types: Vec<_> = func.params.iter().map(|(_, type_)| *type_).collect();
    (param_types, func.type_)
}

struct FunctionContext<'a> {
    function: &'a mut Function,
    globals: &'a HashMap<&'a str, u32>,
    locals: &'a HashMap<&'a str, u32>,
    labels: Vec<String>,
    deferred_inits: HashMap<&'a str, &'a ast::Expression<'a>>,
}

fn emit_function(func: &ast::Function, globals: &HashMap<&str, u32>) -> Function {
    let mut locals = Vec::new();
    collect_locals(&func.body, &mut locals);
    locals.sort_by_key(|(_, t)| *t);

    let mut function = Function::new_with_locals_types(locals.iter().map(|(_, t)| map_type(*t)));

    let locals: HashMap<&str, u32> = locals
        .into_iter()
        .enumerate()
        .map(|(index, (name, _))| (name, index as u32))
        .collect();

    let mut context = FunctionContext {
        function: &mut function,
        globals,
        locals: &locals,
        labels: vec![],
        deferred_inits: HashMap::new(),
    };

    emit_block(&mut context, &func.body);
    if func.type_.is_none() && func.body.type_().is_some() {
        function.instruction(&Instruction::Drop);
    }
    function.instruction(&Instruction::End);

    function
}

fn collect_locals<'a>(block: &ast::Block<'a>, locals: &mut Vec<(&'a str, ast::Type)>) {
    for stmt in &block.statements {
        match stmt {
            ast::Statement::LocalVariable(v) => {
                locals.push((v.name, v.type_.unwrap()));
                if let Some(ref value) = v.value {
                    collect_locals_expr(value, locals);
                }
            }
            ast::Statement::Expression(e) => collect_locals_expr(e, locals),
            ast::Statement::Poke {
                mem_location,
                value,
                ..
            } => {
                collect_locals_expr(&mem_location.left, locals);
                collect_locals_expr(value, locals);
            }
        }
    }
    if let Some(ref expr) = block.final_expression {
        collect_locals_expr(expr, locals);
    }
}

fn collect_locals_expr<'a>(expr: &ast::Expression<'a>, locals: &mut Vec<(&'a str, ast::Type)>) {
    match &expr.expr {
        ast::Expr::Variable { .. } | ast::Expr::I32Const(_) | ast::Expr::F32Const(_) => (),
        ast::Expr::BinOp { left, right, .. } => {
            collect_locals_expr(left, locals);
            collect_locals_expr(right, locals);
        }
        ast::Expr::BranchIf { condition, .. } => collect_locals_expr(condition, locals),
        ast::Expr::LocalTee { value, .. } => collect_locals_expr(value, locals),
        ast::Expr::Loop { block, .. } => collect_locals(block, locals),
        ast::Expr::Cast { value, .. } => collect_locals_expr(value, locals),
    }
}

fn emit_block<'a>(ctx: &mut FunctionContext<'a>, block: &'a ast::Block) {
    for stmt in &block.statements {
        match stmt {
            ast::Statement::Expression(e) => {
                emit_expression(ctx, e);
                if e.type_.is_some() {
                    ctx.function.instruction(&Instruction::Drop);
                }
            }
            ast::Statement::LocalVariable(v) => {
                if let Some(ref val) = v.value {
                    if v.defer {
                        ctx.deferred_inits.insert(v.name, val);
                    } else {
                        emit_expression(ctx, val);
                        ctx.function
                            .instruction(&Instruction::LocalSet(*ctx.locals.get(v.name).unwrap()));
                    }
                }
            }
            ast::Statement::Poke {
                mem_location,
                value,
                ..
            } => {
                emit_expression(ctx, &mem_location.left);
                emit_expression(ctx, value);
                let offset = if let ast::Expr::I32Const(v) = mem_location.right.expr {
                    v as u32 as u64
                } else {
                    unreachable!()
                };
                ctx.function.instruction(&match mem_location.size {
                    ast::MemSize::Byte => Instruction::I32Store8(MemArg {
                        align: 0,
                        memory_index: 0,
                        offset,
                    }),
                    ast::MemSize::Word => Instruction::I32Store(MemArg {
                        align: 2,
                        memory_index: 0,
                        offset,
                    }),
                });
            }
        }
    }
    if let Some(ref expr) = block.final_expression {
        emit_expression(ctx, expr);
    }
}

fn emit_expression<'a>(ctx: &mut FunctionContext<'a>, expr: &'a ast::Expression) {
    match &expr.expr {
        ast::Expr::BinOp {
            left, op, right, ..
        } => {
            emit_expression(ctx, left);
            emit_expression(ctx, right);
            use ast::BinOp::*;
            use ast::Type::*;
            ctx.function.instruction(&match (left.type_.unwrap(), op) {
                (I32, Add) => Instruction::I32Add,
                (I32, Sub) => Instruction::I32Sub,
                (I32, Mul) => Instruction::I32Mul,
                (I32, Div) => Instruction::I32DivS,
                (I32, Rem) => Instruction::I32RemS,
                (I32, And) => Instruction::I32And,
                (I32, Or) => Instruction::I32Or,
                (I32, Xor) => Instruction::I32Xor,
                (I32, Eq) => Instruction::I32Eq,
                (I32, Ne) => Instruction::I32Neq,
                (I32, Lt) => Instruction::I32LtS,
                (I32, Le) => Instruction::I32LeS,
                (I32, Gt) => Instruction::I32GtS,
                (I32, Ge) => Instruction::I32GeS,

                (F32, Add) => Instruction::F32Add,
                (F32, Sub) => Instruction::F32Sub,
                (F32, Mul) => Instruction::F32Mul,
                (F32, Div) => Instruction::F32Div,
                (F32, Rem | And | Or | Xor) => unreachable!(),
                (F32, Eq) => Instruction::F32Eq,
                (F32, Ne) => Instruction::F32Neq,
                (F32, Lt) => Instruction::F32Lt,
                (F32, Le) => Instruction::F32Le,
                (F32, Gt) => Instruction::F32Gt,
                (F32, Ge) => Instruction::F32Ge,

                (I64, _) => todo!(),
                (F64, _) => todo!(),
            });
        }
        ast::Expr::BranchIf {
            condition, label, ..
        } => {
            emit_expression(ctx, condition);
            let depth = ctx
                .labels
                .iter()
                .rev()
                .enumerate()
                .find(|(_, l)| l == label)
                .unwrap()
                .0;
            ctx.function.instruction(&Instruction::BrIf(depth as u32));
        }
        ast::Expr::I32Const(v) => {
            ctx.function.instruction(&Instruction::I32Const(*v));
        }
        ast::Expr::F32Const(v) => {
            ctx.function.instruction(&Instruction::F32Const(*v));
        }
        ast::Expr::LocalTee { name, value, .. } => {
            emit_expression(ctx, value);
            let index = ctx.locals.get(*name).unwrap();
            ctx.function.instruction(&Instruction::LocalTee(*index));
        }
        ast::Expr::Loop { label, block, .. } => {
            ctx.labels.push(label.to_string());
            ctx.function
                .instruction(&Instruction::Loop(map_block_type(block.type_())));
            emit_block(ctx, block);
            ctx.labels.pop();
            ctx.function.instruction(&Instruction::End);
        }
        ast::Expr::Variable { name, .. } => {
            if let Some(index) = ctx.locals.get(*name) {
                if let Some(expr) = ctx.deferred_inits.remove(*name) {
                    emit_expression(ctx, expr);
                    ctx.function.instruction(&Instruction::LocalTee(*index));
                } else {
                    ctx.function.instruction(&Instruction::LocalGet(*index));
                }
            } else if let Some(index) = ctx.globals.get(*name) {
                ctx.function.instruction(&Instruction::GlobalGet(*index));
            } else {
                unreachable!()
            }
        }
        ast::Expr::Cast { value, type_, .. } => {
            emit_expression(ctx, value);
            use ast::Type::*;
            let inst = match (value.type_.unwrap(), *type_) {
                (t1, t2) if t1 == t2 => None,
                (I32, F32) => Some(Instruction::F32ConvertI32S),
                (F32, I32) => Some(Instruction::I32TruncF32S),
                _ => todo!(),
            };
            if let Some(inst) = inst {
                ctx.function.instruction(&inst);
            }
        }
    }
}

fn map_type(t: ast::Type) -> ValType {
    match t {
        ast::Type::I32 => ValType::I32,
        ast::Type::I64 => ValType::I64,
        ast::Type::F32 => ValType::F32,
        ast::Type::F64 => ValType::F64,
    }
}

fn map_block_type(t: Option<ast::Type>) -> BlockType {
    if let Some(t) = t {
        BlockType::Result(map_type(t))
    } else {
        BlockType::Empty
    }
}