diff --git a/crates/ail/tests/e2e.rs b/crates/ail/tests/e2e.rs index 477e406..a73aacb 100644 --- a/crates/ail/tests/e2e.rs +++ b/crates/ail/tests/e2e.rs @@ -108,6 +108,26 @@ fn insertion_sort_orders_list() { ); } +/// Iter 12b: polymorphism end-to-end. `id : forall a. (a) -> a` +/// gets called with `42` (Int) and `true` (Bool); each instantiation +/// triggers a separate specialised LLVM fn. Output: 42, then "true". +#[test] +fn polymorphic_id_at_int_and_bool() { + let stdout = build_and_run("poly_id.ail.json"); + let lines: Vec<&str> = stdout.lines().collect(); + assert_eq!(lines, vec!["42", "true"]); +} + +/// Iter 12b: polymorphism with a function-typed parameter. +/// `apply : forall a b. ((a) -> b, a) -> b` invoked as +/// `apply(succ, 41) == 42`. Exercises the closure-pair ABI inside a +/// monomorphised body (a fn-typed param survives substitution). +#[test] +fn polymorphic_apply_with_fn_param() { + let stdout = build_and_run("poly_apply.ail.json"); + assert_eq!(stdout.trim(), "42"); +} + /// Guards `ail diff`: a modified body changes the hash of `sum`, while /// `main` stays unchanged. Expects exit code 1, `changed` contains exactly /// `sum`, `unchanged` contains `main`, `added`/`removed` empty. diff --git a/crates/ailang-codegen/src/lib.rs b/crates/ailang-codegen/src/lib.rs index f7618a8..f5bb55b 100644 --- a/crates/ailang-codegen/src/lib.rs +++ b/crates/ailang-codegen/src/lib.rs @@ -79,22 +79,45 @@ pub fn lower_workspace(ws: &Workspace) -> Result { let mut all_strings: BTreeMap> = 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> = BTreeMap::new(); + let mut module_def_ail_types: BTreeMap> = BTreeMap::new(); + let mut module_polymorphic_fns: BTreeMap> = 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> = 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> = 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 { 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, - /// 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>, + /// 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>, + /// Polymorphic defs per module — full FnDef so we can specialise + /// the body when monomorphising. Mono fns aren't included. + module_polymorphic_fns: &'a BTreeMap>, + /// 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)>, + mono_emitted: BTreeSet<(String, String, String)>, // (module, def, descriptor) /// Import map of the current module (alias/module name → actual module name). import_map: BTreeMap, /// ADT table: type_name -> list of ctors in definition order. @@ -251,6 +297,9 @@ struct CtorRef { type_name: String, tag: u32, fields: Vec, + /// AILang-level field types (parallel to `fields`). Needed for + /// codegen-side type tracking — match field bindings inherit these. + ail_fields: Vec, } #[derive(Debug, Clone)] @@ -264,6 +313,8 @@ impl<'a> Emitter<'a> { module: &'a Module, module_name: &'a str, module_user_fns: &'a BTreeMap>, + module_def_ail_types: &'a BTreeMap>, + module_polymorphic_fns: &'a BTreeMap>, import_map: BTreeMap, ) -> Self { let mut types: BTreeMap> = 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__` and produces the desired symbol. + fn emit_specialised_fn( + &mut self, + owner_module: &str, + def_name: &str, + subst: &BTreeMap, + 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_ 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____`. + 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 = args + .iter() + .map(|a| self.synth_arg_type(a)) + .collect::>()?; + let subst = derive_substitution(&forall_vars, ¶ms, &arg_ail_tys)?; + + // Specialised types and signature. + let mono_params: Vec = params + .iter() + .map(|p| apply_subst_to_type(p, &subst)) + .collect(); + let mono_ret = apply_subst_to_type(&ret, &subst); + let llvm_params: Vec = mono_params.iter().map(llvm_type).collect::>()?; + 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::>() + .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)> = 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)> = 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 { + self.synth_with_extras(t, &[]) + } + + fn synth_with_extras(&self, t: &Term, extras: &[(String, Type)]) -> Result { + 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 = args + .iter() + .map(|a| self.synth_with_extras(a, extras)) + .collect::>()?; + 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, ¶ms, &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 { @@ -1526,6 +1933,254 @@ fn fn_sig_from_type(t: &Type) -> Option { 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 { + 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 { + 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> { + 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 = 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, +) -> 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) -> 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 = 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) -> 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____`. +/// 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 { + let mut parts: Vec = 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 → `FR` (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"), diff --git a/examples/poly_apply.ail.json b/examples/poly_apply.ail.json new file mode 100644 index 0000000..a6e40f9 --- /dev/null +++ b/examples/poly_apply.ail.json @@ -0,0 +1,79 @@ +{ + "schema": "ailang/v0", + "name": "poly_apply", + "imports": [], + "defs": [ + { + "kind": "fn", + "name": "apply", + "type": { + "k": "forall", + "vars": ["a", "b"], + "body": { + "k": "fn", + "params": [ + { + "k": "fn", + "params": [{ "k": "var", "name": "a" }], + "ret": { "k": "var", "name": "b" }, + "effects": [] + }, + { "k": "var", "name": "a" } + ], + "ret": { "k": "var", "name": "b" }, + "effects": [] + } + }, + "params": ["f", "x"], + "body": { + "t": "app", + "fn": { "t": "var", "name": "f" }, + "args": [{ "t": "var", "name": "x" }] + } + }, + { + "kind": "fn", + "name": "succ", + "type": { + "k": "fn", + "params": [{ "k": "con", "name": "Int" }], + "ret": { "k": "con", "name": "Int" }, + "effects": [] + }, + "params": ["n"], + "body": { + "t": "app", + "fn": { "t": "var", "name": "+" }, + "args": [ + { "t": "var", "name": "n" }, + { "t": "lit", "lit": { "kind": "int", "value": 1 } } + ] + } + }, + { + "kind": "fn", + "name": "main", + "type": { + "k": "fn", + "params": [], + "ret": { "k": "con", "name": "Unit" }, + "effects": ["IO"] + }, + "params": [], + "body": { + "t": "do", + "op": "io/print_int", + "args": [ + { + "t": "app", + "fn": { "t": "var", "name": "apply" }, + "args": [ + { "t": "var", "name": "succ" }, + { "t": "lit", "lit": { "kind": "int", "value": 41 } } + ] + } + ] + } + } + ] +} diff --git a/examples/poly_id.ail.json b/examples/poly_id.ail.json new file mode 100644 index 0000000..2cfd580 --- /dev/null +++ b/examples/poly_id.ail.json @@ -0,0 +1,63 @@ +{ + "schema": "ailang/v0", + "name": "poly_id", + "imports": [], + "defs": [ + { + "kind": "fn", + "name": "id", + "type": { + "k": "forall", + "vars": ["a"], + "body": { + "k": "fn", + "params": [{ "k": "var", "name": "a" }], + "ret": { "k": "var", "name": "a" }, + "effects": [] + } + }, + "params": ["x"], + "body": { "t": "var", "name": "x" } + }, + { + "kind": "fn", + "name": "main", + "type": { + "k": "fn", + "params": [], + "ret": { "k": "con", "name": "Unit" }, + "effects": ["IO"] + }, + "params": [], + "body": { + "t": "seq", + "lhs": { + "t": "do", + "op": "io/print_int", + "args": [ + { + "t": "app", + "fn": { "t": "var", "name": "id" }, + "args": [ + { "t": "lit", "lit": { "kind": "int", "value": 42 } } + ] + } + ] + }, + "rhs": { + "t": "do", + "op": "io/print_bool", + "args": [ + { + "t": "app", + "fn": { "t": "var", "name": "id" }, + "args": [ + { "t": "lit", "lit": { "kind": "bool", "value": true } } + ] + } + ] + } + } + } + ] +}