Brummel/AILang
1
0
Fork 0
Code Issues 29 Pull Requests Actions Packages Projects Releases Wiki Activity

Iter 12b: codegen monomorphisation for polymorphic defs

Browse Source 
Polymorphic defs are now actually emitted: each unique instantiation
gets its own specialised LLVM fn, mangled @ail_<m>_<def>__<descriptor>.
The typechecker's Forall instantiations from Iter 12a now have a real
backend.

Strategy:
- Pass 1 of lower_workspace splits fn-typed defs into mono and poly:
  module_user_fns keeps LLVM-typed FnSig only for monomorphic fns;
  module_polymorphic_fns holds the full FnDef for poly defs;
  module_def_ail_types is a unified AILang-type lookup for both.
- Direct calls to poly defs (current-module or qualified cross-module)
  flow through lower_polymorphic_call: it derives the type
  substitution from the actual arg AILang types, mangles the symbol,
  and queues (def, subst) for specialisation.
- emit_module drains the queue after the regular defs. Each entry
  produces a specialised FnDef with rigid vars substituted in both
  the type and the body (incl. Lam param/ret annotations), then
  emits via the existing emit_fn pipeline. The synthetic name embeds
  the descriptor; standard mangling produces the right symbol.

Codegen-side type tracking:
- locals tuple grew from (name, ssa, llvm_type) to 4-tuple with
  ail_type. Updated all 6 push sites: fn entry, Let, match-ctor field,
  match-open-arm, lambda capture (cap_meta also extended), lambda
  param.
- CtorRef gained ail_fields parallel to fields, used by match-arm
  bindings to inherit AILang types.
- New synth_arg_type / synth_with_extras: a small recursive walker
  that derives the AILang type of an expression in the current scope.
  Used at poly call sites for arg-type inference. The `extras`
  parameter shadows locals during Let recursion without &mut self.

Helpers added:
- derive_substitution(vars, params, arg_tys): unifies the param
  shapes against the actual args, binding rigid vars; reports
  unbound vars as an internal error.
- apply_subst_to_type / apply_subst_to_term: substitute rigid vars
  throughout a Type or Term (the Term variant only matters for Lam,
  the only Term arm carrying types).
- descriptor_for_subst / type_descriptor: stable identifier-safe
  name suffix. `Int → I`, `Bool → B`, `Unit → U`, `Str → S`, ADT
  `Foo → FFoo`, `Fn → Fn_<...>__r_<ret>`. So `id__I` for id at Int,
  `apply__I_I` for apply at (Int, Int).
- builtin_ail_type / builtin_effect_op_ret: AILang types for the
  builtin operators and effect ops, used by the codegen-side type
  tracker.

Examples + tests:
- examples/poly_id.ail.json: id used at Int (42) and Bool (true).
  Output: 42 / true. Two specialised fns + adapters + static
  closures get emitted (visible in the IR).
- examples/poly_apply.ail.json: apply(succ, 41) == 42. Exercises
  the harder case where one of the polymorphic params is itself
  a function value — the closure-pair ABI survives substitution.
- crates/ail/tests/e2e.rs: polymorphic_id_at_int_and_bool and
  polymorphic_apply_with_fn_param.

Tests: 58/58 (was 56/56). Hash invariant holds: sum.ail.json keeps
db33f57cb329935e / d9a916a0ed10a3d3.

Limitations (known, deferred):
- Polymorphic fns can only be DIRECTLY called. Passing a poly fn as
  a value (let f = id in f(42)) fails — resolve_top_level_fn looks
  in module_user_fns which doesn't include poly defs. Adding this
  needs per-instantiation closure-pair globals.
- No higher-rank polymorphism: a polymorphic arg passed to another
  polymorphic call (apply(id, 42)) trips the simple unify_for_subst
  which doesn't recurse into Forall. Acceptable for the MVP.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Brummel
2026-05-07 14:10:51 +02:00
parent 1c8c1a0441
commit 1fd4763fad
4 changed files with 845 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/ailang-codegen/src/lib.rs
+683 -28
View File
@@ -79,22 +79,45 @@ pub fn lower_workspace(ws: &Workspace) -> Result<String> {
let mut all_strings: BTreeMap<String, Vec<(String, String)>> = BTreeMap::new();
// ^ per module: list of (global-name, content). Order = insertion order.
// Pass 1: per-module top-level symbol table (for local var lookups).
// Pass 1: per-module top-level symbol tables.
// - `module_user_fns`: LLVM-typed FnSig for monomorphic fns. Used by
// the call resolver. Polymorphic fns are deliberately excluded —
// they don't have a single LLVM signature; specialised entries
// appear here on demand during monomorphisation (Iter 12b).
// - `module_def_ail_types`: AILang `Type` for every fn-typed def
// (poly or mono). Used by the codegen-side type tracker to derive
// substitutions at polymorphic call sites and to look up the
// original `Forall` body when specialising.
let mut module_user_fns: BTreeMap<String, BTreeMap<String, FnSig>> = BTreeMap::new();
let mut module_def_ail_types: BTreeMap<String, BTreeMap<String, Type>> = BTreeMap::new();
let mut module_polymorphic_fns: BTreeMap<String, BTreeMap<String, FnDef>> = BTreeMap::new();
for (mname, m) in &ws.modules {
let mut user_fns = BTreeMap::new();
let mut ail_types = BTreeMap::new();
let mut poly_fns = BTreeMap::new();
for def in &m.defs {
if let Def::Fn(f) = def {
if let Type::Fn { params, ret, .. } = &f.ty {
let psig: Result<Vec<String>> = params.iter().map(llvm_type).collect();
let rsig = llvm_type(ret);
if let (Ok(params), Ok(ret)) = (psig, rsig) {
user_fns.insert(f.name.clone(), FnSig { params, ret });
ail_types.insert(f.name.clone(), f.ty.clone());
match &f.ty {
Type::Fn { params, ret, .. } => {
let psig: Result<Vec<String>> = params.iter().map(llvm_type).collect();
let rsig = llvm_type(ret);
if let (Ok(params), Ok(ret)) = (psig, rsig) {
user_fns.insert(f.name.clone(), FnSig { params, ret });
}
}
Type::Forall { .. } => {
// Polymorphic def — no LLVM sig now; specialised
// versions get queued as call sites are lowered.
poly_fns.insert(f.name.clone(), f.clone());
}
_ => {}
}
}
}
module_user_fns.insert(mname.clone(), user_fns);
module_def_ail_types.insert(mname.clone(), ail_types);
module_polymorphic_fns.insert(mname.clone(), poly_fns);
}
// Pass 2: lower per module. Globals/strings are accumulated per module,
@@ -109,7 +132,14 @@ pub fn lower_workspace(ws: &Workspace) -> Result<String> {
import_map.insert(key, imp.module.clone());
}
let mut emitter = Emitter::new(m, mname, &module_user_fns, import_map);
let mut emitter = Emitter::new(
m,
mname,
&module_user_fns,
&module_def_ail_types,
&module_polymorphic_fns,
import_map,
);
emitter
.emit_module()
.map_err(|e| CodegenError::InModule(mname.clone(), Box::new(e)))?;
@@ -201,14 +231,30 @@ struct Emitter<'a> {
body: String,
/// String constants: content -> (global name (without `@`), llvm type length incl. \0)
strings: BTreeMap<String, (String, usize)>,
/// Local symbol table per function: name -> (ssa name incl. `%`, llvm type).
locals: Vec<(String, String, String)>,
/// Local symbol table per function: (name, ssa, llvm_type, ail_type).
/// The AILang type is recorded so that the codegen-side type tracker
/// can derive substitutions at polymorphic call sites without
/// re-running the typechecker (Iter 12b).
locals: Vec<(String, String, String, Type)>,
/// Monotonic counter for SSA values and labels.
counter: u64,
/// Monotonic counter for global string names (per module).
str_counter: u64,
/// Top-level functions per module of the workspace, for call resolution.
module_user_fns: &'a BTreeMap<String, BTreeMap<String, FnSig>>,
/// AILang types of every fn-typed top-level def, per module. Carries
/// `Forall` for polymorphic defs (used to derive substitutions at
/// monomorphic call sites). Populated in pass 1 of `lower_workspace`.
module_def_ail_types: &'a BTreeMap<String, BTreeMap<String, Type>>,
/// Polymorphic defs per module — full FnDef so we can specialise
/// the body when monomorphising. Mono fns aren't included.
module_polymorphic_fns: &'a BTreeMap<String, BTreeMap<String, FnDef>>,
/// Iter 12b: queue of (module, def, substitution) tuples that need
/// to be emitted as specialised versions of polymorphic defs.
/// `mono_emitted` tracks the same keys to deduplicate. The descriptor
/// string is the deterministic name suffix (`Int`, `Int_Bool`, ...).
mono_queue: Vec<(String, String, BTreeMap<String, Type>, String)>,
mono_emitted: BTreeSet<(String, String, String)>, // (module, def, descriptor)
/// Import map of the current module (alias/module name → actual module name).
import_map: BTreeMap<String, String>,
/// ADT table: type_name -> list of ctors in definition order.
@@ -251,6 +297,9 @@ struct CtorRef {
type_name: String,
tag: u32,
fields: Vec<String>,
/// AILang-level field types (parallel to `fields`). Needed for
/// codegen-side type tracking — match field bindings inherit these.
ail_fields: Vec<Type>,
}
#[derive(Debug, Clone)]
@@ -264,6 +313,8 @@ impl<'a> Emitter<'a> {
module: &'a Module,
module_name: &'a str,
module_user_fns: &'a BTreeMap<String, BTreeMap<String, FnSig>>,
module_def_ail_types: &'a BTreeMap<String, BTreeMap<String, Type>>,
module_polymorphic_fns: &'a BTreeMap<String, BTreeMap<String, FnDef>>,
import_map: BTreeMap<String, String>,
) -> Self {
let mut types: BTreeMap<String, Vec<CtorInfo>> = BTreeMap::new();
@@ -287,6 +338,7 @@ impl<'a> Emitter<'a> {
type_name: td.name.clone(),
tag: i as u32,
fields,
ail_fields: c.fields.clone(),
},
);
}
@@ -304,6 +356,10 @@ impl<'a> Emitter<'a> {
counter: 0,
str_counter: 0,
module_user_fns,
module_def_ail_types,
module_polymorphic_fns,
mono_queue: Vec::new(),
mono_emitted: BTreeSet::new(),
import_map,
types,
ctor_index,
@@ -326,6 +382,13 @@ impl<'a> Emitter<'a> {
for def in defs {
match def {
Def::Fn(f) => {
// Polymorphic defs aren't emitted in their original
// form — they are specialised on demand at call sites
// (Iter 12b). Skip them here; the drain pass below
// emits the specialised versions.
if matches!(&f.ty, Type::Forall { .. }) {
continue;
}
self.emit_fn(f)
.map_err(|e| CodegenError::Def(f.name.clone(), Box::new(e)))?;
}
@@ -340,9 +403,84 @@ impl<'a> Emitter<'a> {
}
}
}
// Drain the monomorphisation queue. Specialised fns may
// themselves invoke polymorphic defs and queue further entries,
// so iterate until empty.
while let Some((owner_module, def_name, subst, descriptor)) = self.mono_queue.pop() {
self.emit_specialised_fn(&owner_module, &def_name, &subst, &descriptor)
.map_err(|e| {
CodegenError::Def(
format!("{def_name}__{descriptor}"),
Box::new(e),
)
})?;
}
Ok(())
}
/// Iter 12b: emit one specialised version of a polymorphic def.
/// Substitutes rigid vars in both the type and the body, then
/// calls `emit_fn` against a synthetic FnDef whose `name` already
/// contains the descriptor — the existing mangling concatenates
/// `ail_<module>_<name>` and produces the desired symbol.
fn emit_specialised_fn(
&mut self,
owner_module: &str,
def_name: &str,
subst: &BTreeMap<String, Type>,
descriptor: &str,
) -> Result<()> {
let fdef = self
.module_polymorphic_fns
.get(owner_module)
.and_then(|m| m.get(def_name))
.cloned()
.ok_or_else(|| {
CodegenError::Internal(format!(
"emit_specialised_fn: `{owner_module}.{def_name}` not registered"
))
})?;
let inner_ty = match &fdef.ty {
Type::Forall { body, .. } => (**body).clone(),
other => other.clone(),
};
let mono_ty = apply_subst_to_type(&inner_ty, subst);
let mono_body = apply_subst_to_term(&fdef.body, subst);
let synthetic = FnDef {
name: format!("{def_name}__{descriptor}"),
ty: mono_ty,
params: fdef.params.clone(),
body: mono_body,
doc: fdef.doc.clone(),
};
// Specialised def belongs to the polymorphic def's owner
// module, not necessarily self.module_name. Swap module_name
// briefly so mangling stays correct.
let saved_module = self.module_name;
// SAFETY: we rebind module_name through a raw pointer cast
// because the field is `&'a str`. Equivalent: hand
// emit_fn the mangling target via a parameter. Simpler to
// restore.
// Instead of unsafe, we just call emit_fn directly — the
// mangling uses self.module_name which is &'a str but the
// owner_module string lives inside self.module_polymorphic_fns,
// also borrowed for 'a, so we can re-borrow it.
let owner_ref: &'a str = self
.module_polymorphic_fns
.keys()
.find(|k| k.as_str() == owner_module)
.map(|s| s.as_str())
.ok_or_else(|| {
CodegenError::Internal(format!(
"owner module `{owner_module}` not in module_polymorphic_fns"
))
})?;
self.module_name = owner_ref;
let r = self.emit_fn(&synthetic);
self.module_name = saved_module;
r
}
fn emit_const(&mut self, c: &ConstDef) -> Result<()> {
let lty = llvm_type(&c.ty)?;
let lit = match &c.value {
@@ -424,7 +562,12 @@ impl<'a> Emitter<'a> {
// SSA argument name: %arg_<name>
let pssa = format!("%arg_{}", pname);
sig.push_str(&format!("{pty} {pssa}"));
self.locals.push((pname.clone(), pssa.clone(), pty.clone()));
self.locals.push((
pname.clone(),
pssa.clone(),
pty.clone(),
pty_ail.clone(),
));
// Iter 7: if this param is a function value, record its sig
// so that `f(args)` inside the body can emit an indirect call.
if let Some(fs) = fn_sig_from_type(pty_ail) {
@@ -521,7 +664,9 @@ impl<'a> Emitter<'a> {
Literal::Unit => ("0".into(), "i8".into()),
}),
Term::Var { name } => {
if let Some((_, ssa, ty)) = self.locals.iter().rev().find(|(n, _, _)| n == name) {
if let Some((_, ssa, ty, _)) =
self.locals.iter().rev().find(|(n, _, _, _)| n == name)
{
return Ok((ssa.clone(), ty.clone()));
}
// Iter 7: bare reference to a top-level fn yields a fn-pointer
@@ -534,8 +679,9 @@ impl<'a> Emitter<'a> {
Err(CodegenError::UnknownVar(name.clone()))
}
Term::Let { name, value, body } => {
let val_ail = self.synth_arg_type(value)?;
let (val_ssa, val_ty) = self.lower_term(value)?;
self.locals.push((name.clone(), val_ssa, val_ty));
self.locals.push((name.clone(), val_ssa, val_ty, val_ail));
let r = self.lower_term(body);
self.locals.pop();
r
@@ -606,7 +752,7 @@ impl<'a> Emitter<'a> {
// which lowers the callee to a fn-pointer and looks up
// its sig in the sidetable.
if let Term::Var { name } = callee.as_ref() {
let shadowed = self.locals.iter().any(|(n, _, _)| n == name);
let shadowed = self.locals.iter().any(|(n, _, _, _)| n == name);
if !shadowed && self.is_static_callee(name) {
return self.lower_app(name, args);
}
@@ -712,6 +858,7 @@ impl<'a> Emitter<'a> {
scrutinee: &Term,
arms: &[Arm],
) -> Result<(String, String)> {
let s_ail = self.synth_arg_type(scrutinee)?;
let (s_val, s_ty) = self.lower_term(scrutinee)?;
if s_ty != "ptr" {
return Err(CodegenError::Internal(format!(
@@ -803,8 +950,9 @@ impl<'a> Emitter<'a> {
self.body.push_str(&format!(
" {v} = load {fty}, ptr {addr}, align 8\n"
));
let fail_ail = cref.ail_fields.get(idx).cloned().unwrap_or(Type::unit());
self.locals
.push((bname.clone(), v, fty.clone()));
.push((bname.clone(), v, fty.clone(), fail_ail));
pushed += 1;
}
}
@@ -832,7 +980,8 @@ impl<'a> Emitter<'a> {
if let Some(arm) = open_arm {
// set up var binding if needed
let pushed = if let Some(name) = open_var.take() {
self.locals.push((name, s_val.clone(), "ptr".into()));
self.locals
.push((name, s_val.clone(), "ptr".into(), s_ail.clone()));
1
} else {
0
@@ -909,6 +1058,14 @@ impl<'a> Emitter<'a> {
"cross-module call `{name}`: prefix `{prefix}` not in import map"
))
})?;
// Polymorphic def in target module? Monomorphise on demand.
if self
.module_polymorphic_fns
.get(&target_module)
.is_some_and(|m| m.contains_key(suffix))
{
return self.lower_polymorphic_call(&target_module, suffix, args);
}
let target_fns = self
.module_user_fns
.get(&target_module)
@@ -928,6 +1085,16 @@ impl<'a> Emitter<'a> {
return self.emit_call(&target_module, suffix, &sig, args);
}
// Polymorphic def in the current module?
if self
.module_polymorphic_fns
.get(self.module_name)
.is_some_and(|m| m.contains_key(name))
{
let owner = self.module_name.to_string();
return self.lower_polymorphic_call(&owner, name, args);
}
// User function in the current module?
if let Some(sig) = self
.module_user_fns
@@ -943,6 +1110,98 @@ impl<'a> Emitter<'a> {
)))
}
/// Iter 12b: lower a direct call to a polymorphic def. Derives the
/// substitution from the actual arg types, queues the (def,
/// substitution) pair for specialisation if not yet emitted, and
/// emits a direct call to the mangled name `@ail_<m>_<def>__<descr>`.
fn lower_polymorphic_call(
&mut self,
owner_module: &str,
def_name: &str,
args: &[Term],
) -> Result<(String, String)> {
let fdef = self
.module_polymorphic_fns
.get(owner_module)
.and_then(|m| m.get(def_name))
.cloned()
.ok_or_else(|| {
CodegenError::Internal(format!(
"lower_polymorphic_call: `{owner_module}.{def_name}` not registered"
))
})?;
let (forall_vars, params, ret) = match &fdef.ty {
Type::Forall { vars, body } => match body.as_ref() {
Type::Fn { params, ret, .. } => {
(vars.clone(), params.clone(), (**ret).clone())
}
_ => {
return Err(CodegenError::Internal(format!(
"lower_polymorphic_call: `{def_name}` Forall body is not Fn"
)));
}
},
_ => {
return Err(CodegenError::Internal(format!(
"lower_polymorphic_call: `{def_name}` is not polymorphic"
)));
}
};
// Derive the substitution by comparing the declared param
// types against the actual arg types.
let arg_ail_tys: Vec<Type> = args
.iter()
.map(|a| self.synth_arg_type(a))
.collect::<Result<_>>()?;
let subst = derive_substitution(&forall_vars, &params, &arg_ail_tys)?;
// Specialised types and signature.
let mono_params: Vec<Type> = params
.iter()
.map(|p| apply_subst_to_type(p, &subst))
.collect();
let mono_ret = apply_subst_to_type(&ret, &subst);
let llvm_params: Vec<String> = mono_params.iter().map(llvm_type).collect::<Result<_>>()?;
let llvm_ret = llvm_type(&mono_ret)?;
let descriptor = descriptor_for_subst(&forall_vars, &subst);
let mangled = format!("ail_{owner_module}_{def_name}__{descriptor}");
// Queue specialisation (deduplicated).
let key = (owner_module.to_string(), def_name.to_string(), descriptor.clone());
if self.mono_emitted.insert(key) {
self.mono_queue.push((
owner_module.to_string(),
def_name.to_string(),
subst.clone(),
descriptor,
));
}
// Lower args and emit a direct call to the mangled name.
let mut compiled_args = Vec::new();
for (a, exp_ty) in args.iter().zip(llvm_params.iter()) {
let (v, vty) = self.lower_term(a)?;
if &vty != exp_ty {
return Err(CodegenError::Internal(format!(
"poly call `{owner_module}.{def_name}` arg type mismatch: expected {exp_ty}, got {vty}"
)));
}
compiled_args.push((v, vty));
}
let arglist = compiled_args
.iter()
.map(|(v, t)| format!("{t} {v}"))
.collect::<Vec<_>>()
.join(", ");
let dst = self.fresh_ssa();
self.body.push_str(&format!(
" {dst} = call {ret} @{mangled}({arglist})\n",
ret = llvm_ret,
));
Ok((dst, llvm_ret))
}
fn emit_call(
&mut self,
target_module: &str,
@@ -1092,14 +1351,14 @@ impl<'a> Emitter<'a> {
// pull SSA + LLVM type from `self.locals`. Unknown captures
// would mean the typechecker let through an unbound var, so a
// hard internal error is right.
let mut cap_meta: Vec<(String, String, String, Option<FnSig>)> = Vec::new();
// (name, outer_ssa, llvm_type, optional_sig_for_fn_typed_capture)
// Tuple: (name, outer_ssa, llvm_type, ail_type, optional_fn_sig).
let mut cap_meta: Vec<(String, String, String, Type, Option<FnSig>)> = Vec::new();
for c in &captures {
let (_, outer_ssa, lty) = self
let (_, outer_ssa, lty, ail_ty) = self
.locals
.iter()
.rev()
.find(|(n, _, _)| n == c)
.find(|(n, _, _, _)| n == c)
.ok_or_else(|| {
CodegenError::Internal(format!(
"lambda capture `{c}` not in scope (typechecker bug?)"
@@ -1107,7 +1366,7 @@ impl<'a> Emitter<'a> {
})?
.clone();
let sig = self.ssa_fn_sigs.get(&outer_ssa).cloned();
cap_meta.push((c.clone(), outer_ssa, lty, sig));
cap_meta.push((c.clone(), outer_ssa, lty, ail_ty, sig));
}
// 2. Pick a thunk name and switch the emitter into "thunk
@@ -1141,7 +1400,7 @@ impl<'a> Emitter<'a> {
// body. Layout: 8 bytes per slot, regardless of LLVM type
// (typed load reads only the needed bytes; padding is wasted
// but uniform).
for (i, (cname, _outer, cty, sig_opt)) in cap_meta.iter().enumerate() {
for (i, (cname, _outer, cty, c_ail, sig_opt)) in cap_meta.iter().enumerate() {
let offset = (i * 8) as i64;
let slot = self.fresh_ssa();
let val = self.fresh_ssa();
@@ -1150,8 +1409,12 @@ impl<'a> Emitter<'a> {
));
self.body
.push_str(&format!(" {val} = load {cty}, ptr {slot}\n"));
self.locals
.push((cname.clone(), val.clone(), cty.clone()));
self.locals.push((
cname.clone(),
val.clone(),
cty.clone(),
c_ail.clone(),
));
if let Some(sig) = sig_opt {
self.ssa_fn_sigs.insert(val, sig.clone());
}
@@ -1165,7 +1428,12 @@ impl<'a> Emitter<'a> {
.zip(lam_param_tys.iter())
{
let pssa = format!("%arg_{pname}");
self.locals.push((pname.clone(), pssa.clone(), pty.clone()));
self.locals.push((
pname.clone(),
pssa.clone(),
pty.clone(),
pty_ail.clone(),
));
if let Some(fs) = fn_sig_from_type(pty_ail) {
self.ssa_fn_sigs.insert(pssa, fs);
}
@@ -1198,7 +1466,7 @@ impl<'a> Emitter<'a> {
let env = self.fresh_ssa();
self.body
.push_str(&format!(" {env} = call ptr @malloc(i64 {env_size})\n"));
for (i, (_cname, outer_ssa, cty, _sig)) in cap_meta.iter().enumerate() {
for (i, (_cname, outer_ssa, cty, _c_ail, _sig)) in cap_meta.iter().enumerate() {
let offset = (i * 8) as i64;
let slot = self.fresh_ssa();
self.body.push_str(&format!(
@@ -1351,8 +1619,9 @@ impl<'a> Emitter<'a> {
/// Iter 7: is `name` a callee that can be resolved at compile time
/// (no fn-pointer needed)? True for builtin operators, the
/// current-module top-level fns, and qualified `prefix.def`.
/// Locals shadow this — the caller checks for that first.
/// current-module top-level fns (mono or poly), and qualified
/// `prefix.def`. Locals shadow this — the caller checks for
/// that first.
fn is_static_callee(&self, name: &str) -> bool {
if builtin_binop(name).is_some() || name == "not" {
return true;
@@ -1360,7 +1629,14 @@ impl<'a> Emitter<'a> {
if name.matches('.').count() == 1 {
return true;
}
self.module_user_fns
if self
.module_user_fns
.get(self.module_name)
.is_some_and(|m| m.contains_key(name))
{
return true;
}
self.module_polymorphic_fns
.get(self.module_name)
.is_some_and(|m| m.contains_key(name))
}
@@ -1487,6 +1763,137 @@ impl<'a> Emitter<'a> {
.insert(content.to_string(), (name.clone(), len));
name
}
/// Iter 12b: lightweight AILang-type computation for an expression
/// in the current scope. Mirrors what the typechecker already
/// derived; we replay it here only because the typechecker doesn't
/// hand its annotations down.
///
/// Used at polymorphic call sites to derive the type substitution
/// from the actual argument types, at let-bindings / match-arm
/// scrutinees to populate the AILang-type slot of locals. Trusts
/// the typechecker for well-formedness — failures here are internal
/// errors (e.g. unbound var that the checker should have rejected).
///
/// Limitations: nested polymorphic instantiations (an arg that is
/// itself a polymorphic call) work via `synth_with_extras`'s
/// recursion; the substitution is derived on-the-fly and applied
/// to the return type. The body of a let is walked with the
/// let-bound name added to a small `extras` shadow stack so we
/// don't need `&mut self`.
fn synth_arg_type(&self, t: &Term) -> Result<Type> {
self.synth_with_extras(t, &[])
}
fn synth_with_extras(&self, t: &Term, extras: &[(String, Type)]) -> Result<Type> {
match t {
Term::Lit { lit } => Ok(match lit {
Literal::Int { .. } => Type::int(),
Literal::Bool { .. } => Type::bool_(),
Literal::Str { .. } => Type::str_(),
Literal::Unit => Type::unit(),
}),
Term::Var { name } => {
// Lookup precedence: extras (let-bindings introduced
// during this synth walk) → emitter locals → globals
// → builtins. Mirrors typechecker shadowing.
for (n, ty) in extras.iter().rev() {
if n == name {
return Ok(ty.clone());
}
}
if let Some((_, _, _, ail)) =
self.locals.iter().rev().find(|(n, _, _, _)| n == name)
{
return Ok(ail.clone());
}
if name.matches('.').count() == 1 {
let (prefix, suffix) = name.split_once('.').expect("checked");
if let Some(target) = self.import_map.get(prefix) {
if let Some(ty) = self
.module_def_ail_types
.get(target)
.and_then(|m| m.get(suffix))
{
return Ok(ty.clone());
}
}
}
if let Some(ty) = self
.module_def_ail_types
.get(self.module_name)
.and_then(|m| m.get(name))
{
return Ok(ty.clone());
}
if let Some(t) = builtin_ail_type(name) {
return Ok(t);
}
Err(CodegenError::UnknownVar(name.clone()))
}
Term::Lam {
param_tys,
ret_ty,
effects,
..
} => Ok(Type::Fn {
params: param_tys.clone(),
ret: ret_ty.clone(),
effects: effects.clone(),
}),
Term::App { callee, args } => {
let cty = self.synth_with_extras(callee, extras)?;
match cty {
Type::Fn { ret, .. } => Ok(*ret),
Type::Forall { vars, body } => {
let arg_tys: Vec<Type> = args
.iter()
.map(|a| self.synth_with_extras(a, extras))
.collect::<Result<_>>()?;
let (params, ret) = match body.as_ref() {
Type::Fn { params, ret, .. } => (params.clone(), (**ret).clone()),
_ => {
return Err(CodegenError::Internal(
"synth_arg_type: forall body is not Fn".into(),
));
}
};
let subst = derive_substitution(&vars, &params, &arg_tys)?;
Ok(apply_subst_to_type(&ret, &subst))
}
other => Err(CodegenError::Internal(format!(
"synth_arg_type: callee not a fn type: {}",
ailang_core::pretty::type_to_string(&other)
))),
}
}
Term::Let { name, value, body } => {
let v_ail = self.synth_with_extras(value, extras)?;
let mut new_extras: Vec<(String, Type)> = extras.to_vec();
new_extras.push((name.clone(), v_ail));
self.synth_with_extras(body, &new_extras)
}
Term::If { then, .. } => self.synth_with_extras(then, extras),
Term::Do { op, .. } => builtin_effect_op_ret(op).ok_or_else(|| {
CodegenError::Internal(format!(
"synth_arg_type: unknown effect op `{op}`"
))
}),
Term::Ctor { type_name, .. } => Ok(Type::Con {
name: type_name.clone(),
}),
Term::Match { arms, .. } => {
if let Some(first) = arms.first() {
self.synth_with_extras(&first.body, extras)
} else {
Err(CodegenError::Internal(
"synth_arg_type: empty match".into(),
))
}
}
Term::Seq { rhs, .. } => self.synth_with_extras(rhs, extras),
}
}
}
fn llvm_type(t: &Type) -> Result<String> {
@@ -1526,6 +1933,254 @@ fn fn_sig_from_type(t: &Type) -> Option<FnSig> {
None
}
/// Iter 12b: AILang type of a builtin operator. Used by
/// `synth_arg_type` for arg-type inference at polymorphic call sites.
/// Mirrors what the typechecker installs in its env via `builtins`.
fn builtin_ail_type(name: &str) -> Option<Type> {
let int_int_int = || Type::Fn {
params: vec![Type::int(), Type::int()],
ret: Box::new(Type::int()),
effects: vec![],
};
let int_int_bool = || Type::Fn {
params: vec![Type::int(), Type::int()],
ret: Box::new(Type::bool_()),
effects: vec![],
};
Some(match name {
"+" | "-" | "*" | "/" | "%" => int_int_int(),
"==" | "!=" | "<" | "<=" | ">" | ">=" => int_int_bool(),
"not" => Type::Fn {
params: vec![Type::bool_()],
ret: Box::new(Type::bool_()),
effects: vec![],
},
_ => return None,
})
}
/// Iter 12b: AILang return type of a built-in effect op. The op's
/// param signature is irrelevant here since we only consume the ret.
fn builtin_effect_op_ret(op: &str) -> Option<Type> {
Some(match op {
"io/print_int" | "io/print_bool" | "io/print_str" => Type::unit(),
_ => return None,
})
}
/// Iter 12b: derive a name → concrete-type substitution from the
/// declared params of a `Forall` body and the actual arg types at a
/// call site. Walks both sides in parallel; whenever a `Type::Var`
/// (rigid name) appears on the params side, binds it to the
/// corresponding concrete type. Conflicts (same var bound to two
/// different types) surface as an internal error — the typechecker
/// would already have rejected such a call.
fn derive_substitution(
vars: &[String],
params: &[Type],
arg_tys: &[Type],
) -> Result<BTreeMap<String, Type>> {
if params.len() != arg_tys.len() {
return Err(CodegenError::Internal(format!(
"derive_substitution: arity mismatch ({} params vs {} args)",
params.len(),
arg_tys.len(),
)));
}
let var_set: BTreeSet<&str> = vars.iter().map(|s| s.as_str()).collect();
let mut subst: BTreeMap<String, Type> = BTreeMap::new();
for (p, a) in params.iter().zip(arg_tys.iter()) {
unify_for_subst(p, a, &var_set, &mut subst)?;
}
// Any forall var not pinned by the args is left unbound. For the
// MVP this is an error — we can't specialise without a concrete
// type. The typechecker's body should have constrained it already
// through return-type unification, but at the call site we only
// see args; if needed, callers can extend this with expected-ret
// info.
for v in vars {
if !subst.contains_key(v) {
return Err(CodegenError::Internal(format!(
"monomorphisation: type var `{v}` not pinned by call args"
)));
}
}
Ok(subst)
}
/// Walks `param` and `arg` in parallel, treating any `Type::Var { name }`
/// on the param side whose name is in `vars` as an unknown to be bound
/// in `subst`. Identical concrete shapes pass through; structural
/// mismatches yield an internal error.
fn unify_for_subst(
param: &Type,
arg: &Type,
vars: &BTreeSet<&str>,
subst: &mut BTreeMap<String, Type>,
) -> Result<()> {
match (param, arg) {
(Type::Var { name }, _) if vars.contains(name.as_str()) => {
if let Some(prev) = subst.get(name) {
if prev != arg {
return Err(CodegenError::Internal(format!(
"monomorphisation: var `{name}` bound to two distinct types"
)));
}
return Ok(());
}
subst.insert(name.clone(), arg.clone());
Ok(())
}
(Type::Con { name: pn }, Type::Con { name: an }) if pn == an => Ok(()),
(
Type::Fn { params: pp, ret: pr, .. },
Type::Fn { params: ap, ret: ar, .. },
) => {
if pp.len() != ap.len() {
return Err(CodegenError::Internal(
"monomorphisation: fn arity mismatch in arg".into(),
));
}
for (p, a) in pp.iter().zip(ap.iter()) {
unify_for_subst(p, a, vars, subst)?;
}
unify_for_subst(pr, ar, vars, subst)
}
(Type::Var { name: pn }, Type::Var { name: an }) if pn == an => Ok(()),
_ => Err(CodegenError::Internal(format!(
"monomorphisation: cannot match param `{}` to arg `{}`",
ailang_core::pretty::type_to_string(param),
ailang_core::pretty::type_to_string(arg),
))),
}
}
/// Iter 12b: substitute rigid type vars in `t` according to `subst`.
/// Used to specialise the type of a polymorphic def for a given
/// instantiation.
fn apply_subst_to_type(t: &Type, subst: &BTreeMap<String, Type>) -> Type {
match t {
Type::Var { name } => subst.get(name).cloned().unwrap_or_else(|| t.clone()),
Type::Con { .. } => t.clone(),
Type::Fn { params, ret, effects } => Type::Fn {
params: params.iter().map(|p| apply_subst_to_type(p, subst)).collect(),
ret: Box::new(apply_subst_to_type(ret, subst)),
effects: effects.clone(),
},
Type::Forall { vars, body } => {
// Inner forall shadows: don't substitute re-bound names.
let inner: BTreeMap<String, Type> = subst
.iter()
.filter(|(k, _)| !vars.contains(k))
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
Type::Forall {
vars: vars.clone(),
body: Box::new(apply_subst_to_type(body, &inner)),
}
}
}
}
/// Iter 12b: substitute rigid type vars throughout a Term. Only
/// `Term::Lam` carries types in the AST (params/ret), so most arms
/// just recurse. `Term::Var` contains a name string only and is
/// left untouched.
fn apply_subst_to_term(t: &Term, subst: &BTreeMap<String, Type>) -> Term {
match t {
Term::Lit { .. } | Term::Var { .. } => t.clone(),
Term::App { callee, args } => Term::App {
callee: Box::new(apply_subst_to_term(callee, subst)),
args: args.iter().map(|a| apply_subst_to_term(a, subst)).collect(),
},
Term::Let { name, value, body } => Term::Let {
name: name.clone(),
value: Box::new(apply_subst_to_term(value, subst)),
body: Box::new(apply_subst_to_term(body, subst)),
},
Term::If { cond, then, else_ } => Term::If {
cond: Box::new(apply_subst_to_term(cond, subst)),
then: Box::new(apply_subst_to_term(then, subst)),
else_: Box::new(apply_subst_to_term(else_, subst)),
},
Term::Do { op, args } => Term::Do {
op: op.clone(),
args: args.iter().map(|a| apply_subst_to_term(a, subst)).collect(),
},
Term::Ctor { type_name, ctor, args } => Term::Ctor {
type_name: type_name.clone(),
ctor: ctor.clone(),
args: args.iter().map(|a| apply_subst_to_term(a, subst)).collect(),
},
Term::Match { scrutinee, arms } => Term::Match {
scrutinee: Box::new(apply_subst_to_term(scrutinee, subst)),
arms: arms
.iter()
.map(|a| Arm {
pat: a.pat.clone(),
body: apply_subst_to_term(&a.body, subst),
})
.collect(),
},
Term::Lam { params, param_tys, ret_ty, effects, body } => Term::Lam {
params: params.clone(),
param_tys: param_tys
.iter()
.map(|t| apply_subst_to_type(t, subst))
.collect(),
ret_ty: Box::new(apply_subst_to_type(ret_ty, subst)),
effects: effects.clone(),
body: Box::new(apply_subst_to_term(body, subst)),
},
Term::Seq { lhs, rhs } => Term::Seq {
lhs: Box::new(apply_subst_to_term(lhs, subst)),
rhs: Box::new(apply_subst_to_term(rhs, subst)),
},
}
}
/// Iter 12b: deterministic descriptor string for a substitution. Used
/// as the suffix in the mangled name `@ail_<m>_<def>__<descriptor>`.
/// Vars are emitted in the order given by the FnDef's forall vars
/// (so two call sites with the same instantiation map to the same
/// descriptor regardless of internal BTreeMap ordering).
fn descriptor_for_subst(vars: &[String], subst: &BTreeMap<String, Type>) -> String {
let mut parts: Vec<String> = Vec::with_capacity(vars.len());
for v in vars {
let ty = subst.get(v).cloned().unwrap_or_else(|| Type::unit());
parts.push(type_descriptor(&ty));
}
parts.join("_")
}
/// Iter 12b: a stable, identifier-safe descriptor for a `Type`.
/// Maps `Int → I`, `Bool → B`, `Unit → U`, `Str → S`, ADT name `Foo →
/// FFoo`, fn → `F<params...>R<ret>` (no recursion guard since types in
/// the MVP are non-recursive at the type level).
fn type_descriptor(t: &Type) -> String {
match t {
Type::Con { name } => match name.as_str() {
"Int" => "I".into(),
"Bool" => "B".into(),
"Unit" => "U".into(),
"Str" => "S".into(),
other => format!("F{other}"),
},
Type::Fn { params, ret, .. } => {
let mut s = String::from("Fn");
for p in params {
s.push('_');
s.push_str(&type_descriptor(p));
}
s.push_str("__r_");
s.push_str(&type_descriptor(ret));
s
}
Type::Var { name } => format!("V{name}"),
Type::Forall { .. } => "FORALL".into(),
}
}
fn builtin_binop(name: &str) -> Option<(&'static str, &'static str)> {
Some(match name {
"+" => ("add", "i64"),