blr-lang 0.1.0

//! Nucleus is the final representation of the program as a WASM binary following the component
//! model specification.

pub mod component;

use std::{
    cmp::max,
    collections::{BTreeMap, BTreeSet, HashMap},
    iter::once,
    ops::Index,
};

use tracing::{Level, debug, instrument};
use wasm_encoder::{
    BlockType, CodeSection, CompositeInnerType, CompositeType, ConstExpr, DataSection, DataSegment,
    DataSegmentMode, EntityType, ExportKind, ExportSection, FuncType, Function, FunctionSection,
    ImportSection, Instruction, MemArg, MemoryType, RefType, SubType, TableType, TypeSection,
    ValType,
};

use super::{
    core::{self, Expr, ExternalItem, Item, ItemId, LocalItem, RemoteItem, Type, Var, VarId},
    mantle::Symbol,
};
use crate::{
    external_type::{ExternalType, FunctionType},
    trace_alt,
};

#[derive(Eq, Hash, PartialEq)]
enum PartialTy {
    Func(FuncType),
}

trait AsValTy {
    fn as_val_ty(&self) -> ValType;
}
impl AsValTy for u32 {
    fn as_val_ty(&self) -> ValType {
        ValType::I32
    }
}

struct ClosureTypeIndex {
    func_index: u32,
}
impl AsValTy for ClosureTypeIndex {
    fn as_val_ty(&self) -> ValType {
        ValType::I32
    }
}

#[derive(Default)]
struct EmitType {
    types: Vec<PartialTy>,
}

impl EmitType {
    fn into_type_section(self) -> TypeSection {
        let mut sect = TypeSection::new();
        for ty in self.types.into_iter() {
            let (inner, is_final) = match ty {
                PartialTy::Func(func_type) => (CompositeInnerType::Func(func_type), true),
            };
            sect.ty().subtype(&SubType {
                is_final,
                supertype_idx: None,
                composite_type: CompositeType {
                    shared: false,
                    inner,
                    describes: None,
                    descriptor: None,
                },
            })
        }
        sect
    }

    fn emit_ref_ty(&mut self, key: PartialTy) -> u32 {
        self.types
            .iter()
            .position(|x| x == &key)
            .unwrap_or_else(|| {
                let indx = self.types.len();
                self.types.push(key);
                indx
            })
            .try_into()
            .unwrap()
    }

    fn emit_closure_index(&mut self, arg: &Type, ret: &Type) -> ClosureTypeIndex {
        let arg_valty = self.emit_val_typ(arg);
        let ret_valty = self.emit_val_typ(ret);

        let func_index = self.emit_ref_ty(PartialTy::Func(FuncType::new(
            [ValType::I32, arg_valty],
            [ret_valty],
        )));

        ClosureTypeIndex { func_index }
    }

    fn emit_val_typ(&mut self, ty: &Type) -> ValType {
        match ty {
            Type::Unit => ValType::I64,
            Type::Int => ValType::I64,
            Type::Float => ValType::F64,
            Type::String => ValType::I32,
            Type::DataFrame => ValType::I32,
            Type::Closure(arg, ret) => self.emit_closure_index(arg, ret).as_val_ty(),
            Type::Tuple(fields, _labels) => self.emit_tuple_index(fields),
            Type::Abstraction(_paramters, _body) => todo!(),
            Type::Variant(cases, _labels) => self.emit_variant_index(cases),
        }
    }

    fn emit_tuple_index<'a>(&mut self, fields: impl IntoIterator<Item = &'a Type>) -> ValType {
        fields.into_iter().for_each(|field| {
            self.emit_val_typ(field);
        });

        ValType::I32
    }

    fn emit_variant_index<'a>(&mut self, cases: impl IntoIterator<Item = &'a Type>) -> ValType {
        cases.into_iter().for_each(|field| {
            self.emit_val_typ(field);
        });

        ValType::I32
    }

    fn emit_remote_item_ty(&mut self, item: &RemoteItem) -> u32 {
        let param_ty = self.emit_val_typ(&item.param_typ);
        let ret_ty = self.emit_val_typ(&item.ret_typ);
        let func_ty = FuncType::new([ValType::I32, param_ty], [ret_ty]);
        self.emit_ref_ty(PartialTy::Func(func_ty))
    }
    fn emit_external_item_ty(&mut self, item: &ExternalItem) -> u32 {
        let flat = flatten_function_type(&item.external_typ, Context::Lower);
        let func_ty = FuncType::new(flat.params, flat.ret);
        self.emit_ref_ty(PartialTy::Func(func_ty))
    }
    fn emit_local_item_ty(&mut self, item: &LocalItem) -> u32 {
        let ret_ty = self.emit_val_typ(&item.ret_typ);
        let func_ty = FuncType::new(
            item.params.iter().map(|var| self.emit_val_typ(&var.typ)),
            [ret_ty],
        );

        self.emit_ref_ty(PartialTy::Func(func_ty))
    }
}

#[derive(Default)]
struct EmitData {
    data: Vec<(u32, String)>,
}
impl EmitData {
    fn emit_data_segment(&mut self, offset: u32, data: String) {
        self.data.push((offset, data))
    }
    fn into_data_section(self) -> DataSection {
        let mut section = DataSection::new();
        for (offset, data) in self.data {
            section.segment(DataSegment {
                mode: DataSegmentMode::Active {
                    memory_index: 0,
                    offset: &ConstExpr::i32_const(offset as i32),
                },
                data: data.bytes(),
            });
        }
        section
    }
}

struct EmitLocals {
    next_local: u32,
    local_tys: Vec<ValType>,
    locals: HashMap<VarId, u32>,
}

impl EmitLocals {
    fn param_for(&mut self, id: VarId) -> u32 {
        let local = self.next_local;
        self.next_local += 1;
        self.locals.insert(id, local);
        local
    }

    fn local_for(&mut self, id: VarId, ty: ValType) -> u32 {
        let local = self.next_local;
        self.next_local += 1;
        self.local_tys.push(ty);
        self.locals.insert(id, local);
        local
    }

    fn anon_local(&mut self, ty: ValType) -> u32 {
        let local = self.next_local;
        self.next_local += 1;
        self.local_tys.push(ty);
        local
    }
}

#[derive(Default)]
struct EmitImports {
    imports: ImportSection,
    map: HashMap<Symbol, u32>,
    func_count: u32,
}

impl EmitImports {
    fn func_count(&self) -> u32 {
        self.func_count
    }
    fn import_memory(&mut self) {
        self.import(
            &Symbol {
                module: "blr_runtime".to_string(),
                field: "mem".to_string(),
            },
            EntityType::Memory(MemoryType {
                minimum: 1,
                maximum: None,
                memory64: false,
                shared: false,
                page_size_log2: None,
            }),
        );
    }

    pub fn import_alloc_func(&mut self, alloc_typ_idx: u32) -> u32 {
        self.import(
            &Symbol {
                module: "blr_runtime".to_string(),
                field: "alloc".to_string(),
            },
            EntityType::Function(alloc_typ_idx),
        )
    }
    pub fn import_func_table(&mut self) -> u32 {
        self.import(
            &Symbol {
                module: "blr_runtime".to_string(),
                field: "func_tbl".to_string(),
            },
            EntityType::Table(TableType {
                element_type: RefType::FUNCREF,
                table64: false,
                minimum: 1_000,
                maximum: None,
                shared: false,
            }),
        );
        // We only ever have a single table
        0
    }
    pub fn import(&mut self, symbol: &Symbol, ty: impl Into<EntityType>) -> u32 {
        let ty: EntityType = ty.into();
        *self.map.entry(symbol.clone()).or_insert_with(|| {
            let func_idx = self.func_count;
            if let EntityType::Function(_) = ty {
                self.func_count += 1;
            }
            self.imports.import(&symbol.module, &symbol.field, ty);
            func_idx
        })
    }
}

impl Index<&VarId> for EmitLocals {
    type Output = u32;

    fn index(&self, index: &VarId) -> &Self::Output {
        &self.locals[index]
    }
}

struct EmitModule<'a> {
    types: EmitType,
    data: EmitData,
    functions: HashMap<ItemId, u32>,
    shared_runtime: &'a mut SharedRuntime,
    alloc_idx: u32,
}

impl<'a> EmitModule<'a> {
    fn emit_local_item(&mut self, item: LocalItem) -> Function {
        let (inss, local_tys) = self.emit_body(&item.params, item.body);

        let mut function = Function::new_with_locals_types(local_tys);
        for ins in inss {
            function.instruction(&ins);
        }
        function.instruction(&Instruction::Return);
        function.instruction(&Instruction::End);
        function
    }

    fn emit_body(
        &mut self,
        params: &[Var],
        body: Expr,
    ) -> (Vec<Instruction<'static>>, Vec<ValType>) {
        let mut locals = EmitLocals {
            next_local: 0,
            local_tys: vec![],
            locals: HashMap::default(),
        };
        for param in params {
            locals.param_for(param.id);
        }
        let mut inss: Vec<Instruction> = vec![];

        self.emit_expr(body, &mut locals, &mut inss);

        (inss, locals.local_tys)
    }

    fn emit_expr(&mut self, body: Expr, locals: &mut EmitLocals, inss: &mut Vec<Instruction>) {
        match body {
            Expr::Var(var) => {
                inss.push(Instruction::LocalGet(locals[&var.id]));
            }
            Expr::Unit => inss.push(Instruction::I32Const(-1)),
            Expr::Int(i) => inss.push(Instruction::I64Const(i)),
            Expr::Float(f) => inss.push(Instruction::F64Const(f.into())),
            Expr::String(s) => {
                let str_tuple = locals.anon_local(ValType::I32);
                let ptr = self.shared_runtime.alloc_data(&s);
                let len = s.len();
                self.data.emit_data_segment(ptr, s);
                inss.extend([
                    // Manually construct tuple of ptr length
                    Instruction::I32Const(4),
                    Instruction::I32Const(8),
                    Instruction::Call(self.alloc_idx),
                    Instruction::LocalTee(str_tuple),
                    Instruction::I32Const(ptr as i32),
                    Instruction::I32Store(MemArg {
                        offset: 0,
                        align: 2,
                        memory_index: 0,
                    }),
                    Instruction::LocalGet(str_tuple),
                    Instruction::I32Const(len as i32),
                    Instruction::I32Store(MemArg {
                        offset: 4,
                        align: 2,
                        memory_index: 0,
                    }),
                    Instruction::LocalGet(str_tuple),
                ]);
            }
            Expr::Closure(ty, item_id, vars) => {
                trace_alt!(item_id, ty, vars, "emit_expr closure");
                let func_index = self.functions[&item_id];
                let func_ptr = self.shared_runtime.func_ptr_for(item_id);
                let ValType::I32 = self.types.emit_val_typ(&ty) else {
                    panic!("ICE: Closure assigned to variable with non closure type");
                };
                let Type::Tuple(field_typs, _) =
                    Type::closure_env(ty.clone(), vars.iter().map(|var| var.typ.clone()))
                else {
                    unreachable!("should always be a tuple type")
                };
                let (align, offsets, size) = fields_align_offsets_size(&field_typs);

                let ptr = locals.anon_local(ValType::I32);
                inss.extend([
                    Instruction::I32Const(func_ptr as i32),
                    Instruction::RefFunc(func_index),
                    Instruction::TableSet(self.shared_runtime.func_table_id()),
                    Instruction::I32Const(align as i32),
                    Instruction::I32Const(size as i32),
                    Instruction::Call(self.alloc_idx),
                    Instruction::LocalTee(ptr),
                    // Store function pointer as first field of tuple
                    Instruction::I32Const(func_ptr as i32),
                    Instruction::I32Store(MemArg {
                        offset: 0,
                        align: 2,
                        memory_index: 0,
                    }),
                ]);

                vars.into_iter()
                    .zip(offsets.into_iter().skip(1))
                    .for_each(|(var, offset)| {
                        inss.push(Instruction::LocalGet(ptr));
                        self.emit_store_expr(
                            Expr::Var(var.clone()),
                            &var.typ,
                            offset,
                            locals,
                            inss,
                        );
                    });
                inss.push(Instruction::LocalGet(ptr));
            }
            Expr::AppClosure(fun, param) => {
                let local_ty = fun.type_of();
                let Type::Closure(arg_ty, ret_ty) = local_ty else {
                    panic!("ICE: Expected closure type for function of apply");
                };
                let closure_indices = self.types.emit_closure_index(&arg_ty, &ret_ty);
                self.emit_expr(*fun, locals, inss);
                let fun_local = locals.anon_local(closure_indices.as_val_ty());
                inss.push(Instruction::LocalTee(fun_local));
                self.emit_expr(*param, locals, inss);
                inss.extend([
                    Instruction::LocalGet(fun_local),
                    // We know the code field is always the first field and an i32
                    Instruction::I32Load(MemArg {
                        offset: 0,
                        align: 2,
                        memory_index: 0,
                    }),
                    Instruction::CallIndirect {
                        type_index: closure_indices.func_index,
                        table_index: self.shared_runtime.func_table_id(),
                    },
                ]);
            }
            Expr::AppExternal(typ, ft, item_id, parameters) => {
                let ft = flatten_function_type(&ft, Context::Lower);
                // Add logic re flattening.
                // I.e. alloc params and results as needed
                parameters.into_iter().for_each(|param| {
                    if ft.params_on_heap {
                        todo!()
                    } else {
                        let typ = param.type_of();
                        self.emit_expr(param, locals, inss);
                        match typ {
                            Type::Unit | Type::Int | Type::Float | Type::DataFrame => {}
                            Type::String => {
                                let ptr = locals.anon_local(ValType::I32);
                                // Flatten string, get both i32 from the heap onto the stack
                                inss.extend([
                                    Instruction::LocalTee(ptr),
                                    Instruction::I32Load(MemArg {
                                        offset: 0,
                                        align: 2,
                                        memory_index: 0,
                                    }),
                                    Instruction::LocalGet(ptr),
                                    Instruction::I32Load(MemArg {
                                        offset: 4,
                                        align: 2,
                                        memory_index: 0,
                                    }),
                                ]);
                            }
                            Type::Closure(_, _) => todo!(),
                            Type::Abstraction(_, _) => todo!(),
                            Type::Tuple(fields, _) => {
                                let fields = flatten_fields(fields);
                                let (_align, offsets, _size) = fields_align_offsets_size(&fields);
                                let ptr = locals.anon_local(ValType::I32);
                                inss.push(Instruction::LocalSet(ptr));
                                fields.into_iter().enumerate().for_each(|(i, field)| {
                                    inss.push(Instruction::LocalGet(ptr));
                                    emit_load(&field, offsets[i], inss);
                                });
                            }
                            Type::Variant(_items, _items1) => todo!(),
                        };
                    }
                });
                if ft.ret_on_heap {
                    match typ {
                        Type::Unit => todo!(),
                        Type::Int => todo!(),
                        Type::Float => todo!(),
                        Type::String => {
                            let ptr = locals.anon_local(ValType::I32);
                            inss.extend([
                                Instruction::I32Const(4),
                                Instruction::I32Const(8),
                                Instruction::Call(self.alloc_idx),
                                Instruction::LocalTee(ptr),
                                // Call external function and leave allocated pointer on the stack
                                Instruction::Call(self.functions[&item_id]),
                                Instruction::LocalGet(ptr),
                            ]);
                        }
                        Type::Closure(_, _) => todo!(),
                        Type::Abstraction(_, _) => todo!(),
                        Type::Tuple(fields, _) => {
                            let fields = flatten_fields(fields);
                            let (align, _offsets, size) = fields_align_offsets_size(&fields);
                            let ptr = locals.anon_local(ValType::I32);
                            inss.extend([
                                Instruction::I32Const(align as i32),
                                Instruction::I32Const(size as i32),
                                Instruction::Call(self.alloc_idx),
                                Instruction::LocalTee(ptr),
                                // Call external function and leave allocated pointer on the stack
                                Instruction::Call(self.functions[&item_id]),
                                Instruction::LocalGet(ptr),
                            ]);
                        }
                        Type::DataFrame => todo!(),
                        Type::Variant(_items, _items1) => todo!(),
                    }
                } else {
                    inss.push(Instruction::Call(self.functions[&item_id]));
                }
            }
            Expr::Local(var, defn, body) => {
                self.emit_expr(*defn, locals, inss);
                let val_typ = self.types.emit_val_typ(&var.typ);
                let local = locals.local_for(var.id, val_typ);
                inss.push(Instruction::LocalSet(local));
                self.emit_expr(*body, locals, inss);
            }
            Expr::Access(tuple, field) => {
                let typ = tuple.type_of();
                match &typ {
                    Type::Tuple(fields, _labels) => {
                        self.emit_expr(*tuple, locals, inss);
                        let (_align, offsets, _size) = fields_align_offsets_size(&fields[..]);
                        emit_load(&fields[field], offsets[field], inss);
                    }
                    _ => panic!("ICE: Expected tuple type for field access"),
                }
            }
            Expr::Tuple(fields, _labels) => {
                debug!(?fields, ?_labels, "tuple raw");
                let field_typs: Vec<_> = fields.iter().map(|field| field.type_of()).collect();
                let (align, offsets, size) = fields_align_offsets_size(&field_typs);
                debug!(?fields, align, ?offsets, size, "tuple");

                let ptr = locals.anon_local(ValType::I32);
                inss.extend([
                    Instruction::I32Const(align as i32),
                    Instruction::I32Const(size as i32),
                    Instruction::Call(self.alloc_idx),
                    Instruction::LocalSet(ptr),
                ]);
                fields.into_iter().zip(field_typs).zip(offsets).for_each(
                    |((field, field_typ), offset)| {
                        inss.push(Instruction::LocalGet(ptr));
                        self.emit_store_expr(field, &field_typ, offset, locals, inss);
                    },
                );
                inss.push(Instruction::LocalGet(ptr));
            }
            Expr::Tag(typ, idx, expr) => {
                let Type::Variant(cases, _) = typ else {
                    unreachable!("tag type does not have a variant type {typ:?}");
                };
                let (dt, size, align, case_align) = cases_dt_size_align(&cases);
                debug!(?dt, size, align, "tag");
                let (dt_size, _) = dt.size_align();

                let ptr = locals.anon_local(ValType::I32);
                inss.extend([
                    Instruction::I32Const(align as i32),
                    Instruction::I32Const(size as i32),
                    Instruction::Call(self.alloc_idx),
                    Instruction::LocalTee(ptr),
                ]);
                dt.emit_store(idx as u32, inss);
                inss.extend([
                    Instruction::LocalGet(ptr),
                    Instruction::I32Const(dt_size as i32),
                    Instruction::I32Add,
                ]);
                self.emit_align_to(case_align, locals, inss);
                self.emit_store_expr(*expr, &cases[idx], 0, locals, inss);
                inss.push(Instruction::LocalGet(ptr));
            }
            // TODO: Case needs both the Variant type and the return type
            Expr::Case(_typ, scrutinee, branches) => {
                let Type::Variant(cases, _) = scrutinee.type_of() else {
                    unreachable!("scrutinee should always be a variant type")
                };
                debug_assert_eq!(cases.len(), branches.len());
                let (dt, size, _align, case_align) = cases_dt_size_align(&cases);
                self.emit_expr(*scrutinee, locals, inss);
                let ptr = locals.anon_local(ValType::I32);
                //  Load discriminant from ptr to variant value
                match dt {
                    DiscriminantType::U8 => inss.push(Instruction::I32Load8U(MemArg {
                        offset: 0,
                        align: 0,
                        memory_index: 0,
                    })),
                    DiscriminantType::U16 => inss.push(Instruction::I32Load16U(MemArg {
                        offset: 0,
                        align: 1,
                        memory_index: 0,
                    })),
                    DiscriminantType::U32 => inss.push(Instruction::I32Load(MemArg {
                        offset: 0,
                        align: 2,
                        memory_index: 0,
                    })),
                };
                let discriminant = locals.anon_local(ValType::I32);
                inss.push(Instruction::LocalSet(discriminant));

                // Adjust ptr on stack to be the value ptr reusing the same local
                inss.extend([
                    Instruction::LocalGet(ptr),
                    Instruction::I32Const(size as i32),
                    Instruction::I32Add,
                ]);
                self.emit_align_to(case_align, locals, inss);
                inss.push(Instruction::LocalSet(ptr));

                // Emit nested conditional block for each branch
                branches.into_iter().enumerate().for_each(|(idx, branch)| {
                    let Type::Closure(param_typ, ret) = branch.type_of() else {
                        unreachable!("branch should always be a closure")
                    };
                    inss.extend([
                        Instruction::LocalGet(discriminant),
                        Instruction::I32Const(idx as i32),
                        Instruction::I32Eq,
                        Instruction::If(BlockType::Result(self.types.emit_val_typ(&ret))),
                        Instruction::LocalGet(ptr),
                    ]);
                    emit_load(&param_typ, 0, inss);
                    // HACK creating fake VarId
                    let param = locals.local_for(VarId(999), self.types.emit_val_typ(&param_typ));
                    inss.push(Instruction::LocalSet(param));
                    self.emit_expr(
                        Expr::apply_closure(
                            branch,
                            Expr::Var(Var {
                                id: VarId(999),
                                typ: *param_typ,
                            }),
                        ),
                        locals,
                        inss,
                    );
                    inss.push(Instruction::Else);
                });
                // The final else block is unreachable
                inss.push(Instruction::Unreachable);
                inss.extend(once(Instruction::End).cycle().take(cases.len()));
            }
        }
    }
    fn emit_store_expr(
        &mut self,
        expr: Expr,
        typ: &Type,
        offset: u32,
        locals: &mut EmitLocals,
        inss: &mut Vec<Instruction>,
    ) {
        let offset = offset as u64;
        match typ {
            Type::Unit => {}
            Type::Int => {
                self.emit_expr(expr, locals, inss);
                inss.push(Instruction::I64Store(MemArg {
                    offset,
                    align: 3,
                    memory_index: 0,
                }))
            }
            Type::Float => {
                self.emit_expr(expr, locals, inss);
                inss.push(Instruction::F64Store(MemArg {
                    offset,
                    align: 3,
                    memory_index: 0,
                }))
            }
            Type::String => todo!(),
            Type::Closure(_, _) => {
                self.emit_expr(expr, locals, inss);
                inss.push(Instruction::I32Store(MemArg {
                    offset,
                    align: 2,
                    memory_index: 0,
                }))
            }
            Type::Abstraction(_, _) => todo!(),
            Type::Tuple(fields, _) => {
                // Instead of recusing we can copy the entire tuple as opaque bytes
                let (_align, _offsets, size) = fields_align_offsets_size(fields);
                // Destination is nested tuple pointer plus offset
                inss.extend([Instruction::I32Const(offset as i32), Instruction::I32Add]);
                // Source
                self.emit_expr(expr, locals, inss);
                inss.extend([
                    // Size
                    Instruction::I32Const(size as i32),
                    Instruction::MemoryCopy {
                        src_mem: 0,
                        dst_mem: 0,
                    },
                ]);
            }
            Type::DataFrame => todo!(),
            Type::Variant(_items, _items1) => todo!(),
        }
    }

    // Emit instructions inline that read an i32 off the stack and align it to alignment
    fn emit_align_to(&self, alignment: u32, locals: &mut EmitLocals, inss: &mut Vec<Instruction>) {
        let mask = locals.anon_local(ValType::I32);
        inss.extend([
            Instruction::I32Const(alignment as i32),
            Instruction::I32Const(-1),
            Instruction::I32Add,
            Instruction::LocalTee(mask),
            Instruction::I32Add,
            Instruction::LocalGet(mask),
            Instruction::I32Const(-1),
            Instruction::I32Xor,
            Instruction::I32And,
        ]);
    }
}

fn align_to(ptr: u32, alignment: u32) -> u32 {
    let align_mask = alignment - 1;
    (ptr + align_mask) & !(align_mask)
}

fn flatten_fields(fields: Vec<Type>) -> Vec<Type> {
    fields
        .into_iter()
        .flat_map(|field| match field {
            Type::Tuple(fields, _) => flatten_fields(fields),
            f => vec![f],
        })
        .collect()
}
// Determine the alignment offsets and size of a set of fields.
#[instrument(ret( level = Level::DEBUG ))]
fn fields_align_offsets_size(fields: &[Type]) -> (u32, Vec<u32>, u32) {
    let (align, offsets, end) = fields.iter().fold(
        (0u32, Vec::with_capacity(fields.len()), 0u32),
        |(align, mut offsets, end), field_typ| {
            let (f_size, f_align) = size_align(field_typ);
            let offset = align_to(end, f_align);
            offsets.push(offset);
            (max(align, f_align), offsets, offset + f_size)
        },
    );
    let size = align_to(end, align);
    (align, offsets, size)
}

/// Emits intructions to load a value for a type.
/// The stack must have the ptr value on top.
fn emit_load(typ: &Type, offset: u32, inss: &mut Vec<Instruction>) {
    let offset = offset as u64;
    match typ {
        Type::Unit => inss.push(Instruction::I64Load(MemArg {
            offset,
            align: 3,
            memory_index: 0,
        })),
        Type::Int => inss.push(Instruction::I64Load(MemArg {
            offset,
            align: 3,
            memory_index: 0,
        })),
        Type::Float => inss.push(Instruction::F64Load(MemArg {
            offset,
            align: 3,
            memory_index: 0,
        })),
        Type::String => todo!(),
        Type::Closure(_, _) => inss.push(Instruction::I32Load(MemArg {
            offset,
            align: 2,
            memory_index: 0,
        })),
        Type::Abstraction(_, _) => todo!(),
        Type::Tuple(_, _) => {
            // For a nested tuple adjust the pointer on the stack by the offset leaving a new
            // pointer to this tuple.
            inss.extend([Instruction::I32Const(offset as i32), Instruction::I32Add])
        }
        Type::DataFrame => todo!(),
        Type::Variant(_items, _items1) => todo!(),
    }
}

// Determine the alignment offsets and size of a set of fields.
fn cases_dt_size_align(cases: &[Type]) -> (DiscriminantType, u32, u32, u32) {
    let (case_size, case_align) = cases.iter().fold((0, 0), |(size, align), case| {
        let (s, a) = size_align(case);
        (max(size, s), max(align, a))
    });
    let dt = discriminant_bits(cases);
    let (dt_size, dt_align) = dt.size_align();
    let align = max(dt_align, case_align);
    let size = align_to(align_to(dt_size, case_align) + case_size, align);
    (dt, size, align, case_align)
}

// Return size and alignment in bytes
fn size_align(typ: &Type) -> (u32, u32) {
    match typ {
        Type::Unit => (8, 8),
        Type::Int => (8, 8),
        Type::Float => (8, 8),
        Type::String => (4, 4),
        Type::Closure(_, _) => (4, 4),
        Type::Abstraction(_, _) => (4, 4),
        Type::Tuple(nested_fields, _) => {
            let (nested_align, _nested_offsets, nested_size) =
                fields_align_offsets_size(nested_fields);
            (nested_size, nested_align)
        }
        Type::DataFrame => todo!(),
        Type::Variant(cases, _) => {
            let (_dt, size, align, _) = cases_dt_size_align(cases);
            (size, align)
        }
    }
}

// TODO: Remove this once blr has native u8,u16 and u32 types
#[derive(Debug)]
enum DiscriminantType {
    U8,
    U16,
    U32,
}
impl DiscriminantType {
    fn size_align(&self) -> (u32, u32) {
        match self {
            DiscriminantType::U8 => (1, 1),
            DiscriminantType::U16 => (2, 2),
            DiscriminantType::U32 => (4, 4),
        }
    }

    fn emit_store(&self, idx: u32, inss: &mut Vec<Instruction>) {
        match self {
            DiscriminantType::U8 => {
                inss.extend([
                    Instruction::I32Const(idx as i32),
                    Instruction::I32Store8(MemArg {
                        offset: 0,
                        align: 0,
                        memory_index: 0,
                    }),
                ]);
            }
            DiscriminantType::U16 => {
                inss.extend([
                    Instruction::I32Const(idx as i32),
                    Instruction::I32Store16(MemArg {
                        offset: 0,
                        align: 1,
                        memory_index: 0,
                    }),
                ]);
            }
            DiscriminantType::U32 => {
                inss.extend([
                    Instruction::I32Const(idx as i32),
                    Instruction::I32Store(MemArg {
                        offset: 0,
                        align: 2,
                        memory_index: 0,
                    }),
                ]);
            }
        }
    }
}
fn discriminant_bits(cases: &[Type]) -> DiscriminantType {
    let l = cases.len();
    if l <= u8::MAX.into() {
        DiscriminantType::U8
    } else if l <= u16::MAX.into() {
        DiscriminantType::U16
    } else {
        DiscriminantType::U32
    }
}
#[allow(dead_code)]
fn max_case_alignment(cases: &[Type]) -> u32 {
    cases
        .iter()
        .fold(0, |align, typ| max(align, size_align(typ).1))
}

#[derive(Debug)]
pub struct SharedRuntime {
    func_table_id: u32,
    next_func_idx: u32,
    func_table: HashMap<ItemId, u32>,
    data_offset: u32,
}
impl Default for SharedRuntime {
    fn default() -> Self {
        Self {
            func_table_id: Default::default(),
            // HACK: Leave room for hand written wat builtins
            next_func_idx: 100,
            func_table: Default::default(),
            data_offset: Default::default(),
        }
    }
}

impl SharedRuntime {
    pub fn data_size(&self) -> u32 {
        self.data_offset + 1
    }
    fn alloc_data(&mut self, s: &str) -> u32 {
        let ptr = self.data_offset;
        self.data_offset += s.len() as u32;
        ptr
    }

    // Enter a new scope for the table with the provided table index for this scope.
    // This causes all previous function mappings to be forgotten
    pub fn enter(&mut self, id: u32) {
        self.func_table_id = id;
        self.func_table.clear();
    }
    // Associate an item with a new entry in the function table.
    fn func_ptr_for(&mut self, item_id: ItemId) -> u32 {
        *self.func_table.entry(item_id).or_insert_with(|| {
            let ptr = self.next_func_idx;
            self.next_func_idx += 1;
            ptr
        })
    }

    fn func_table_id(&self) -> u32 {
        self.func_table_id
    }
}

pub struct Module {
    pub module: Vec<u8>,
    pub imports: BTreeSet<String>,
    pub exports: BTreeMap<String, Type>,
    pub externals: BTreeMap<Symbol, FunctionType>,
}

pub fn emit_wasm_core_module(module: core::Module, shared_runtime: &mut SharedRuntime) -> Module {
    let mut types = EmitType::default();
    let data = EmitData::default();
    let mut import = EmitImports::default();
    let mut func = FunctionSection::new();
    let mut export = ExportSection::new();

    shared_runtime.enter(import.import_func_table());

    let func_ty = FuncType::new([ValType::I32, ValType::I32], [ValType::I32]);
    let alloc_typ_idx = types.emit_ref_ty(PartialTy::Func(func_ty));
    let alloc_idx = import.import_alloc_func(alloc_typ_idx);
    import.import_memory();

    let mut imports = BTreeSet::new();
    let mut exports = BTreeMap::new();

    // Build import section first as it shares an index space with functions
    let import_functions: HashMap<ItemId, u32> = module
        .closure_items
        .iter()
        .filter_map(|(item_id, item)| match item {
            Item::Local(_) => None,
            Item::Remote(item) => {
                let type_index = types.emit_remote_item_ty(item);
                let func_index = import.import(&item.symbol, EntityType::Function(type_index));
                imports.insert(item.symbol.module.clone());
                export.export(
                    &format!("remote{}", item_id.0),
                    ExportKind::Func,
                    func_index,
                );
                Some((*item_id, func_index))
            }
            Item::External(item) => {
                let type_index = types.emit_external_item_ty(item);
                let func_index = import.import(&item.symbol, EntityType::Function(type_index));
                export.export(
                    &format!("external{}", item_id.0),
                    ExportKind::Func,
                    func_index,
                );
                Some((*item_id, func_index))
            }
        })
        .collect();

    let mut functions: HashMap<ItemId, u32> = module
        .closure_items
        .iter()
        .filter_map(|(item_id, item)| {
            if let Item::Local(item) = item {
                let func_index = import.func_count() + func.len();
                let type_index = types.emit_local_item_ty(item);
                func.function(type_index);
                export.export(&item.name, ExportKind::Func, func_index);
                exports.insert(
                    item.name.clone(),
                    Type::abs(
                        item.params.iter().map(|param| param.typ.clone()),
                        item.ret_typ.clone(),
                    ),
                );
                Some((*item_id, func_index))
            } else {
                None
            }
        })
        .collect();
    functions.extend(import_functions);

    // Collect externals metadata
    let externals = module
        .closure_items
        .values()
        .filter_map(|item| {
            if let Item::External(item) = item {
                Some((item.symbol.clone(), item.external_typ.clone()))
            } else {
                None
            }
        })
        .collect();

    let mut emitter = EmitModule {
        types,
        data,
        functions,
        shared_runtime,
        alloc_idx,
    };
    let mut code = CodeSection::default();
    for (_, item) in module.closure_items {
        match item {
            Item::Local(local_item) => {
                code.function(&emitter.emit_local_item(local_item));
            }
            Item::Remote(_) => {}
            Item::External(_) => {}
        };
    }

    debug_assert_eq!(code.len(), func.len());
    let mut module = wasm_encoder::Module::default();
    module
        .section(&emitter.types.into_type_section())
        .section(&import.imports)
        .section(&func)
        .section(&export)
        .section(&code)
        .section(&emitter.data.into_data_section());

    imports.insert("blr_runtime".to_string());
    Module {
        module: module.finish(),
        imports,
        exports,
        externals,
    }
}

pub enum Context {
    Lift,
    Lower,
}

pub struct FlatFunctionType {
    pub params: Vec<ValType>,
    pub ret: Vec<ValType>,
    /// Indicates if the parameters should be stored in linear memory and a pointer passed as a
    /// parameter.
    pub params_on_heap: bool,
    /// Indicates if the return value should be preallocated in linear memory and its destiniation
    /// address passed as the final parameter.
    pub ret_on_heap: bool,
}

/// Flatten a function type according to the CanonicalABI
/// [specificatiobn](https://github.com/WebAssembly/component-model/blob/main/design/mvp/CanonicalABI.md#flattening).
pub fn flatten_function_type(ft: &FunctionType, context: Context) -> FlatFunctionType {
    const MAX_FLAT_PARAMS: usize = 16;
    const MAX_FLAT_RESULTS: usize = 1;
    let mut params_on_heap = false;
    let mut ret_on_heap = false;
    let mut flat_params = flatten_types(&ft.parameter_typs);
    let mut flat_results = flatten_type(&ft.ret);
    if flat_params.len() > MAX_FLAT_PARAMS {
        flat_params = vec![ValType::I32];
        params_on_heap = true
    }
    if flat_results.len() > MAX_FLAT_RESULTS {
        match context {
            Context::Lift => todo!(),
            Context::Lower => {
                flat_params.push(ValType::I32);
                flat_results.clear();
                ret_on_heap = true
            }
        }
    }
    FlatFunctionType {
        params: flat_params,
        ret: flat_results,
        params_on_heap,
        ret_on_heap,
    }
}
pub fn flatten_types(ft: &[ExternalType]) -> Vec<ValType> {
    ft.iter().flat_map(flatten_type).collect()
}
pub fn flatten_type(ft: &ExternalType) -> Vec<ValType> {
    match ft {
        ExternalType::Unit => vec![ValType::I32],
        ExternalType::Int => vec![ValType::I64],
        ExternalType::Float => vec![ValType::F64],
        ExternalType::String => vec![ValType::I32, ValType::I32],
        ExternalType::Resource(_) => vec![ValType::I32],
        ExternalType::Fun(_function_type) => todo!(),
        ExternalType::Record(_, fields) => fields
            .iter()
            .flat_map(|(_label, field_typ)| flatten_type(field_typ))
            .collect(),
    }
}

#[cfg(test)]
mod tests {
    use expect_test::expect;
    use test_log::test;

    use crate::compiler::core::Type;

    use super::{fields_align_offsets_size, size_align};

    #[test]
    fn test_fields_align_offsets_size() {
        expect![[r#"
            (
                8,
                [
                    0,
                    8,
                ],
                24,
            )
        "#]]
        .assert_debug_eq(&fields_align_offsets_size(&vec![
            Type::closure(Type::Int, Type::tuple([Type::Int, Type::Int, Type::Int])),
            Type::tuple([Type::Int, Type::Int]),
        ]))
    }

    #[test]
    fn test_size_align() {
        expect![[r#"
            (
                24,
                8,
            )
        "#]]
        .assert_debug_eq(&size_align(&Type::Tuple(
            vec![
                Type::closure(Type::Int, Type::tuple([Type::Int, Type::Int, Type::Int])),
                Type::tuple([Type::Int, Type::Int]),
            ],
            None,
        )));
        expect![[r#"
            (
                24,
                8,
            )
        "#]]
        .assert_debug_eq(&size_align(&Type::Variant(
            vec![Type::Int, Type::Float, Type::tuple([Type::Int, Type::Int])],
            None,
        )));
    }
}