Skip to main content

submilli_engine/codegen/
function_adapters.rs

1//! Collects top-level functions used in value position and emits closure-shaped adapter bodies for them.
2
3use std::collections::{BTreeMap, BTreeSet};
4
5use wasm_encoder::{
6    CodeSection, ConstExpr, EntityType, GlobalSection, GlobalType, HeapType, ImportSection,
7    Instruction, RefType, ValType,
8};
9
10use crate::codegen::CodegenCtx;
11use crate::codegen::function_emitter::{FunctionEmitter, cast};
12use crate::codegen::symbol_table::SymbolTable;
13use crate::{Ident, MangledName, Span, Type, TypedAst};
14
15#[derive(Clone, Debug)]
16pub struct AdapterMeta {
17    pub name: String,
18    pub mangled: crate::MangledName,
19    pub signature: Type,
20}
21
22/// The closure caches this module exports, from the name a consumer knows the
23/// function by to the function it defines: its non-generic exported functions,
24/// under each name they are exported as, and its classes' public static
25/// methods. A consumer that reads one as a value imports the cache instead of
26/// keeping its own, so the function is one closure across packages, as in
27/// JavaScript.
28pub fn exported_closure_caches(ta: &TypedAst) -> BTreeMap<MangledName, MangledName> {
29    let defined: BTreeSet<&MangledName> = ta
30        .functions
31        .iter()
32        .filter(|function| function.generics.is_empty())
33        .map(|function| &function.mangled_name)
34        .collect();
35    let functions = ta
36        .exports
37        .iter()
38        .filter(|entry| {
39            entry.kind == crate::ExportKind::Function && defined.contains(&entry.target)
40        })
41        .map(|entry| (entry.public_name.clone(), entry.target.clone()));
42    let static_methods = ta.types.iter().flat_map(|decl| {
43        let crate::TypedTypeDecl::Class(class) = decl else {
44            return Vec::new();
45        };
46        class
47            .static_methods
48            .iter()
49            .filter(|(_, visibility)| **visibility == crate::Visibility::Public)
50            .map(|(name, _)| {
51                let method = crate::mangle::static_member(&class.mangled_name, name);
52                (method.clone(), method)
53            })
54            .collect()
55    });
56    functions.chain(static_methods).collect()
57}
58
59/// A closure cache another module may fill: typed `anyref`, because each
60/// module reads it as the closure struct of its own view of the signature.
61fn shared_cache_type() -> GlobalType {
62    GlobalType {
63        val_type: ValType::Ref(RefType {
64            nullable: true,
65            heap_type: HeapType::ANY,
66        }),
67        mutable: true,
68        shared: false,
69    }
70}
71
72/// Imports the closure cache of each dependency function this module reads as
73/// a value, where the function's package exports one.
74pub fn import_shared_closure_caches(
75    metas: &[AdapterMeta],
76    dependencies: &[&crate::PackageDeclaration],
77    imports: &mut ImportSection,
78    symbols: &mut SymbolTable,
79    next_global_idx: &mut u32,
80) -> Result<(), crate::compiler_error::CompilerFailure> {
81    for meta in metas {
82        let Some(package) = dependencies
83            .iter()
84            .find(|package| package.closure_caches.contains(&meta.mangled))
85        else {
86            continue;
87        };
88        imports.import(
89            &package.package_name,
90            crate::mangle::closure_cache(&meta.mangled).as_str(),
91            EntityType::Global(shared_cache_type()),
92        );
93        symbols.record_shared_closure_global_idx(meta.mangled.clone(), *next_global_idx);
94        crate::codegen::next_index(next_global_idx)?;
95    }
96    Ok(())
97}
98
99/// Allocates the globals that cache a function's closure, so every read of a
100/// top-level function yields the same value and `f === f` holds. The closure
101/// has no environment, so one instance serves every read. Each starts null and
102/// is filled on first read. Each function this module exports gets one shared
103/// `anyref` cache; any other function read here gets a typed cache of its own,
104/// unless it imported a shared one. Returns how many globals were added.
105pub fn allocate_closure_globals(
106    metas: &[AdapterMeta],
107    exported: &BTreeMap<MangledName, MangledName>,
108    globals: &mut GlobalSection,
109    symbols: &mut SymbolTable,
110    next_global_idx: &mut u32,
111) -> Result<u32, crate::compiler_error::CompilerFailure> {
112    let shared: BTreeSet<&MangledName> = exported.values().collect();
113    let mut count = 0usize;
114    for meta in metas {
115        let has_imported_cache = symbols.adapter_closure_global_idx(&meta.mangled).is_some();
116        if shared.contains(&meta.mangled) || has_imported_cache {
117            continue;
118        }
119        let closure_struct_idx = symbols
120            .closure_struct_type_idx(super::closures::classify(&meta.signature)?)
121            .ok_or_else(|| {
122                crate::codegen::internal_failure(
123                    "closure struct type registered for every function-as-value",
124                )
125            })?;
126        let heap_type = HeapType::Concrete(closure_struct_idx);
127        globals.global(
128            GlobalType {
129                val_type: ValType::Ref(RefType {
130                    nullable: true,
131                    heap_type,
132                }),
133                mutable: true,
134                shared: false,
135            },
136            &ConstExpr::ref_null(heap_type),
137        );
138        symbols.record_adapter_closure_global_idx(meta.mangled.clone(), *next_global_idx);
139        crate::codegen::next_index(next_global_idx)?;
140        count += 1;
141    }
142    for mangled in shared {
143        globals.global(shared_cache_type(), &ConstExpr::ref_null(HeapType::ANY));
144        symbols.record_shared_closure_global_idx(mangled.clone(), *next_global_idx);
145        crate::codegen::next_index(next_global_idx)?;
146        count += 1;
147    }
148    crate::codegen::wasm_u32(count)
149}
150
151pub fn emit_bodies(
152    metas: &[AdapterMeta],
153    code: &mut CodeSection,
154    ctx: &CodegenCtx<'_>,
155) -> Result<(), crate::compiler_error::CompilerFailure> {
156    let intrinsics = ctx
157        .symbols
158        .intrinsic_type_indices()
159        .ok_or_else(|| crate::codegen::internal_failure("intrinsics declared by codegen entry"))?;
160    // adapter erased-arg slots are nullable to match the
161    // closure ABI.
162    let object_ref = ValType::Ref(RefType {
163        nullable: true,
164        heap_type: HeapType::Concrete(intrinsics.object),
165    });
166    let any_ref = ValType::Ref(RefType {
167        nullable: false,
168        heap_type: HeapType::ANY,
169    });
170    for meta in metas {
171        let target_idx = ctx.symbols.func_idx(&meta.mangled).ok_or_else(|| {
172            crate::codegen::internal_failure(
173                "adapter target function recorded during user-function pre-pass",
174            )
175        })?;
176        let Type::Function { params, .. } = &meta.signature else {
177            return Err(super::internal_failure(
178                "adapter signature requires a function type",
179            ));
180        };
181
182        super::closures::classify(&meta.signature)?;
183
184        // names are unused; adapter bodies access params by Wasm slot index
185        let env_name = Ident {
186            name: "$__env__".to_string(),
187            span: Span::at(ctx.file),
188        };
189        let mut wasm_params: Vec<(Ident, ValType)> = vec![(env_name, any_ref)];
190        for (i, _) in params.iter().enumerate() {
191            wasm_params.push((
192                Ident {
193                    name: format!("$__arg{i}__"),
194                    span: Span::at(ctx.file),
195                },
196                object_ref,
197            ));
198        }
199        let mut emitter = FunctionEmitter::new(ctx, &wasm_params)?;
200
201        let target = ctx
202            .symbols
203            .top_level_fn(&meta.mangled)
204            .ok_or_else(|| crate::codegen::internal_failure("adapter target signature recorded"))?;
205        if target.params.len() != params.len() {
206            return Err(super::internal_failure(
207                "adapter signature disagrees with its target",
208            ));
209        }
210        for (i, p_ty) in target.params.iter().enumerate() {
211            emitter.instruction(Instruction::LocalGet(super::parameter_local(i)?));
212            let default = ctx
213                .symbols
214                .host_call_defaults(target_idx, target.params.len())
215                .get(i)
216                .and_then(Option::as_ref);
217            super::argument_defaults::emit_argument(
218                &mut emitter,
219                ctx,
220                &Type::Unknown,
221                object_ref,
222                default,
223            )?;
224            crate::codegen::cast_check::emit_checked_parameter_cast_on_stack(
225                &mut emitter,
226                ctx,
227                &crate::Type::Unknown,
228                p_ty,
229            )?;
230        }
231        emitter.instruction(Instruction::Call(target_idx));
232        if ctx.symbols.resultless_functions.contains(&target_idx) {
233            super::function_emitter::expr::emit_undefined(&mut emitter, ctx)?;
234        } else {
235            cast::emit_box(&mut emitter, ctx, &target.ret)?;
236        }
237        let built = emitter.build()?;
238        code.function(&built);
239    }
240    Ok(())
241}