use std::collections::{BTreeMap, BTreeSet, HashMap};
use wasm_encoder::CanonicalOption;
use super::core;
mod emit {
use std::collections::HashMap;
use wasm_encoder::{
Alias, CanonicalFunctionSection, CanonicalOption, Component, ComponentAliasSection,
ComponentExportKind, ComponentExportSection, ComponentImportSection, ComponentSection,
ComponentSectionId, ComponentTypeRef, ComponentTypeSection, ComponentValType, Encode,
ExportKind, InstanceSection, InstanceType, ModuleArg, PrimitiveValType, TypeBounds,
};
use crate::external_type;
#[derive(Default, Debug)]
pub struct EmitComponent {
component: Component,
core_module_idx: CoreModuleIdx,
core_function_idx: CoreFunctionIdx,
core_instance_idx: CoreInstanceIdx,
typ_idx: TypIdx,
instance_idx: InstanceIdx,
function_idx: FunctionIdx,
realloc_idx: Option<CoreFunctionIdx>,
exported_record_types: HashMap<Vec<super::core::Type>, TypIdx>,
}
#[derive(Default, Debug, Clone, Copy)]
pub struct CoreModuleIdx(u32);
impl std::ops::AddAssign<u32> for CoreModuleIdx {
fn add_assign(&mut self, rhs: u32) {
self.0 += rhs;
}
}
#[derive(Default, Debug, Clone, Copy)]
pub struct CoreFunctionIdx(u32);
impl std::ops::AddAssign<u32> for CoreFunctionIdx {
fn add_assign(&mut self, rhs: u32) {
self.0 += rhs;
}
}
#[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
pub struct CoreInstanceIdx(u32);
impl std::ops::AddAssign<u32> for CoreInstanceIdx {
fn add_assign(&mut self, rhs: u32) {
self.0 += rhs;
}
}
#[derive(Default, Debug, Clone, Copy)]
pub struct TypIdx(u32);
impl std::ops::AddAssign<u32> for TypIdx {
fn add_assign(&mut self, rhs: u32) {
self.0 += rhs;
}
}
#[derive(Default, Debug, Clone, Copy)]
pub struct InstanceIdx(u32);
impl std::ops::AddAssign<u32> for InstanceIdx {
fn add_assign(&mut self, rhs: u32) {
self.0 += rhs;
}
}
#[derive(Default, Debug, Clone, Copy)]
pub struct FunctionIdx(u32);
impl std::ops::AddAssign<u32> for FunctionIdx {
fn add_assign(&mut self, rhs: u32) {
self.0 += rhs;
}
}
impl EmitComponent {
pub fn new() -> Self {
Self::default()
}
pub fn convert_ty_to_ctyp(&mut self, ty: &super::core::Type) -> ComponentValType {
match ty {
super::core::Type::Unit => ComponentValType::Primitive(PrimitiveValType::S64),
super::core::Type::Int => ComponentValType::Primitive(PrimitiveValType::S64),
super::core::Type::Float => ComponentValType::Primitive(PrimitiveValType::F64),
super::core::Type::String => ComponentValType::Primitive(PrimitiveValType::String),
super::core::Type::Closure(_, _) => todo!(),
super::core::Type::Abstraction(_params, _ret) => todo!(),
super::core::Type::Variant(_, _) => todo!(),
super::core::Type::DataFrame => todo!(),
super::core::Type::Tuple(field_types, labels) => {
let typ_idx = if let Some(typ_idx) = self.exported_record_types.get(field_types)
{
*typ_idx
} else {
if let Some(labels) = labels {
let fields = labels
.iter()
.zip(field_types)
.map(|(name, typ)| (name.as_str(), self.convert_ty_to_ctyp(typ)))
.collect::<Vec<_>>();
let typ_idx = self.emit_record_type(fields);
let typ_idx =
self.emit_type_export(&format!("record{}", typ_idx.0), typ_idx);
self.exported_record_types
.insert(field_types.clone(), typ_idx);
typ_idx
} else {
panic!("must have labels for tuple component type")
}
};
ComponentValType::Type(typ_idx.0)
}
}
}
pub fn emit_core_module(&mut self, data: &[u8]) -> CoreModuleIdx {
let idx = self.core_module_idx;
self.core_module_idx += 1;
self.component.section(&RawModuleSection(data));
idx
}
pub fn emit_core_function_alias(
&mut self,
core_instance: CoreInstanceIdx,
name: &str,
) -> CoreFunctionIdx {
let idx = self.core_function_idx;
self.core_function_idx += 1;
self.component.section(
ComponentAliasSection::new().alias(Alias::CoreInstanceExport {
instance: core_instance.0,
kind: ExportKind::Func,
name,
}),
);
idx
}
pub fn emit_instance_func_alias(
&mut self,
instance: InstanceIdx,
name: &str,
) -> FunctionIdx {
let idx = self.function_idx;
self.function_idx += 1;
self.component
.section(ComponentAliasSection::new().alias(Alias::InstanceExport {
instance: instance.0,
kind: ComponentExportKind::Func,
name,
}));
idx
}
pub fn emit_memory_alias(&mut self, id: CoreInstanceIdx) {
self.component.section(
ComponentAliasSection::new().alias(Alias::CoreInstanceExport {
instance: id.0,
kind: wasm_encoder::ExportKind::Memory,
name: "mem",
}),
);
}
pub fn emit_realloc_alias(&mut self, id: CoreInstanceIdx) {
self.realloc_idx = Some(self.emit_core_function_alias(id, "realloc"))
}
pub fn emit_instance_type(&mut self, typ: &InstanceType) -> TypIdx {
let idx = self.typ_idx;
self.typ_idx += 1;
self.component
.section(ComponentTypeSection::new().instance(typ));
idx
}
pub fn emit_function_type<'b, P, T>(
&mut self,
params: P,
result: ComponentValType,
) -> TypIdx
where
P: IntoIterator<Item = (&'b str, T)>,
P::IntoIter: ExactSizeIterator,
T: Into<ComponentValType>,
{
let idx = self.typ_idx;
self.typ_idx += 1;
let mut types = ComponentTypeSection::new();
types.function().params(params).result(Some(result));
self.component.section(&types);
idx
}
pub fn emit_record_type<'a, F, T>(&mut self, fields: F) -> TypIdx
where
F: IntoIterator<Item = (&'a str, T)>,
F::IntoIter: ExactSizeIterator,
T: Into<ComponentValType>,
{
let idx = self.typ_idx;
self.typ_idx += 1;
let mut types = ComponentTypeSection::new();
types.defined_type().record(fields);
self.component.section(&types);
idx
}
pub fn emit_type_export(&mut self, name: &str, typ_idx: TypIdx) -> TypIdx {
let idx = self.typ_idx;
self.typ_idx += 1;
self.component.section(ComponentExportSection::new().export(
name,
ComponentExportKind::Type,
typ_idx.0,
None,
));
idx
}
pub fn emit_import_instance(&mut self, name: &str, typ: TypIdx) -> InstanceIdx {
let idx = self.instance_idx;
self.instance_idx += 1;
self.component.section(
ComponentImportSection::new().import(name, ComponentTypeRef::Instance(typ.0)),
);
idx
}
pub fn emit_core_instance<A, S>(
&mut self,
core_module: CoreModuleIdx,
with_instances: A,
) -> CoreInstanceIdx
where
A: IntoIterator<Item = (S, CoreInstanceIdx)>,
A::IntoIter: ExactSizeIterator,
S: AsRef<str>,
{
let idx = self.core_instance_idx;
self.core_instance_idx += 1;
self.component.section(
InstanceSection::new().instantiate(
core_module.0,
with_instances
.into_iter()
.map(|(name, idx)| (name, ModuleArg::Instance(idx.0))),
),
);
idx
}
pub fn emit_core_instance_from_exports<E, S>(&mut self, exports: E) -> CoreInstanceIdx
where
E: IntoIterator<Item = (S, CoreFunctionIdx)>,
E::IntoIter: ExactSizeIterator,
S: AsRef<str>,
{
let idx = self.core_instance_idx;
self.core_instance_idx += 1;
self.component.section(
InstanceSection::new().export_items(
exports
.into_iter()
.map(|(name, core_function)| (name, ExportKind::Func, core_function.0)),
),
);
idx
}
pub fn emit_function_lift<O>(
&mut self,
core_function_idx: CoreFunctionIdx,
typ_idx: TypIdx,
options: O,
) -> FunctionIdx
where
O: IntoIterator<Item = CanonicalOption>,
O::IntoIter: ExactSizeIterator,
{
let idx = self.function_idx;
self.function_idx += 1;
self.component.section(CanonicalFunctionSection::new().lift(
core_function_idx.0,
typ_idx.0,
options,
));
idx
}
pub fn emit_core_function_lower<O>(
&mut self,
function: FunctionIdx,
options: O,
) -> CoreFunctionIdx
where
O: IntoIterator<Item = CanonicalOption>,
O::IntoIter: ExactSizeIterator,
{
let idx = self.core_function_idx;
self.core_function_idx += 1;
self.component
.section(CanonicalFunctionSection::new().lower(function.0, options));
idx
}
pub fn emit_function_export(&mut self, name: &str, idx: FunctionIdx, typ_idx: TypIdx) {
self.component.section(ComponentExportSection::new().export(
name,
ComponentExportKind::Func,
idx.0,
Some(ComponentTypeRef::Func(typ_idx.0)),
));
}
pub fn finish(self) -> Vec<u8> {
self.component.finish()
}
pub fn emit_result_type_export(&mut self, result: ComponentValType) -> ComponentValType {
if let ComponentValType::Type(idx) = result {
ComponentValType::Type(self.emit_type_export("ret", TypIdx(idx)).0)
} else {
result
}
}
pub fn memory_option(&self) -> CanonicalOption {
CanonicalOption::Memory(0)
}
pub fn realloc_option(&self) -> CanonicalOption {
CanonicalOption::Realloc(
self.realloc_idx
.expect("emit_realloc_alias should have been called")
.0,
)
}
}
#[derive(Clone, Debug)]
pub struct RawModuleSection<'a>(pub &'a [u8]);
impl Encode for RawModuleSection<'_> {
fn encode(&self, sink: &mut Vec<u8>) {
self.0.encode(sink);
}
}
impl ComponentSection for RawModuleSection<'_> {
fn id(&self) -> u8 {
ComponentSectionId::CoreModule.into()
}
}
#[derive(Clone, Debug, Default)]
pub struct EmitInstanceType {
typ_idx: TypIdx,
func_idx: FunctionIdx,
instance_typ: InstanceType,
type_names: HashMap<String, TypIdx>,
}
impl From<EmitInstanceType> for InstanceType {
fn from(value: EmitInstanceType) -> Self {
value.instance_typ
}
}
impl EmitInstanceType {
pub fn emit_function_type<'b, P, T>(
&mut self,
params: P,
result: ComponentValType,
) -> TypIdx
where
P: IntoIterator<Item = (&'b str, T)>,
P::IntoIter: ExactSizeIterator,
T: Into<ComponentValType>,
{
let idx = self.typ_idx;
self.typ_idx += 1;
self.instance_typ
.ty()
.function()
.params(params)
.result(Some(result));
idx
}
pub fn emit_function_export(&mut self, name: &str, typ_idx: TypIdx) -> FunctionIdx {
let idx = self.func_idx;
self.func_idx += 1;
self.instance_typ
.export(name, ComponentTypeRef::Func(typ_idx.0));
idx
}
pub fn convert_ext_ty_to_ctyp(
&mut self,
ty: &external_type::ExternalType,
) -> ComponentValType {
match ty {
external_type::ExternalType::Unit => {
ComponentValType::Primitive(PrimitiveValType::S64)
}
external_type::ExternalType::Int => {
ComponentValType::Primitive(PrimitiveValType::S64)
}
external_type::ExternalType::Float => {
ComponentValType::Primitive(PrimitiveValType::F64)
}
external_type::ExternalType::String => {
ComponentValType::Primitive(PrimitiveValType::String)
}
external_type::ExternalType::Resource(name) => {
if let Some(typ_idx) = self.type_names.get(name) {
ComponentValType::Type(typ_idx.0)
} else {
let typ_idx = self.emit_resource_export(name);
let typ_idx = self.emit_owned_resource_type(typ_idx);
self.type_names.insert(name.clone(), typ_idx);
ComponentValType::Type(typ_idx.0)
}
}
external_type::ExternalType::Fun(_function_type) => todo!(),
external_type::ExternalType::Record(name, fields) => {
if let Some(typ_idx) = self.type_names.get(name) {
ComponentValType::Type(typ_idx.0)
} else {
let fields = fields
.iter()
.map(|(name, typ)| (name.as_str(), self.convert_ext_ty_to_ctyp(typ)))
.collect::<Vec<_>>();
let typ_idx = self.emit_record_type(fields);
let typ_idx = self.emit_type_export(name, typ_idx);
self.type_names.insert(name.clone(), typ_idx);
ComponentValType::Type(typ_idx.0)
}
}
}
}
pub fn emit_resource_export(&mut self, name: &str) -> TypIdx {
let idx = self.typ_idx;
self.typ_idx += 1;
self.instance_typ
.export(name, ComponentTypeRef::Type(TypeBounds::SubResource));
idx
}
fn emit_owned_resource_type(&mut self, typ_idx: TypIdx) -> TypIdx {
let idx = self.typ_idx;
self.typ_idx += 1;
self.instance_typ.ty().defined_type().own(typ_idx.0);
idx
}
pub fn emit_record_type<'a, F, T>(&mut self, fields: F) -> TypIdx
where
F: IntoIterator<Item = (&'a str, T)>,
F::IntoIter: ExactSizeIterator,
T: Into<ComponentValType>,
{
let idx = self.typ_idx;
self.typ_idx += 1;
self.instance_typ.ty().defined_type().record(fields);
idx
}
pub fn emit_type_export(&mut self, name: &str, typ_idx: TypIdx) -> TypIdx {
let idx = self.typ_idx;
self.typ_idx += 1;
self.instance_typ
.export(name, ComponentTypeRef::Type(TypeBounds::Eq(typ_idx.0)));
idx
}
}
}
pub fn emit_root_component(
modules: Vec<(String, String, super::Module)>,
) -> (Vec<u8>, BTreeSet<String>) {
let core::Type::Abstraction(params, ret) = &modules[modules.len() - 1].2.exports["main"] else {
panic!("main function should have function type")
};
let mut component = emit::EmitComponent::new();
let params: Vec<_> = params
.iter()
.enumerate()
.map(|(i, ty)| (PARAM_NAMES[i], component.convert_ty_to_ctyp(ty)))
.collect();
let result = component.convert_ty_to_ctyp(ret);
let result = component.emit_result_type_export(result);
let mut path_to_core_instance: HashMap<String, emit::CoreInstanceIdx> = Default::default();
let mut external_modules = BTreeSet::default();
for (path, cpath, module) in modules.iter() {
let mut instantiate_args: BTreeMap<_, _> = Default::default();
if !module.externals.is_empty() {
let mut instance_typ = emit::EmitInstanceType::default();
module.externals.iter().for_each(|(symbol, typ)| {
let params = typ
.parameter_names
.iter()
.zip(typ.parameter_typs.iter())
.map(|(name, typ)| (name.as_str(), instance_typ.convert_ext_ty_to_ctyp(typ)))
.collect::<Vec<_>>();
let result = instance_typ.convert_ext_ty_to_ctyp(&typ.ret);
let idx = instance_typ.emit_function_type(params, result);
instance_typ.emit_function_export(&symbol.field, idx);
});
let typ_idx = component.emit_instance_type(&instance_typ.into());
let instance_idx = component.emit_import_instance(cpath, typ_idx);
let exports: Vec<_> = module
.externals
.keys()
.map(|symbol| {
let function_idx =
component.emit_instance_func_alias(instance_idx, &symbol.field);
let core_function_idx = component.emit_core_function_lower(
function_idx,
[
CanonicalOption::UTF8,
component.memory_option(),
component.realloc_option(),
],
);
(&symbol.field, core_function_idx)
})
.collect();
let core_instance_idx = component.emit_core_instance_from_exports(exports);
module.externals.iter().for_each(|(symbol, _typ)| {
debug_assert_eq!(
core_instance_idx,
*instantiate_args
.entry(symbol.module.to_string())
.or_insert(core_instance_idx),
"should be only a single instance of each core module"
);
external_modules.insert(symbol.module.to_string());
});
}
let core_module_idx = component.emit_core_module(&module.module);
instantiate_args.extend(module.imports.iter().map(|name| {
let idx = path_to_core_instance[name];
(name.to_string(), idx)
}));
let core_instance_idx = component.emit_core_instance(core_module_idx, instantiate_args);
path_to_core_instance.insert(path.clone(), core_instance_idx);
if path == "blr_runtime" {
component.emit_memory_alias(core_instance_idx);
component.emit_realloc_alias(core_instance_idx);
}
}
let main_core_function =
component.emit_core_function_alias(path_to_core_instance["main"], "main");
let main_typ = component.emit_function_type(params, result);
let main_function =
component.emit_function_lift(main_core_function, main_typ, [CanonicalOption::Memory(0)]);
component.emit_function_export("main", main_function, main_typ);
(component.finish(), external_modules)
}
const PARAM_NAMES: [&str; 5] = ["a", "b", "c", "d", "e"];