Something went wrong. Try again.
An experiment in static, refcounted, procedural language with ownership
Something went wrong. Try again.
23 kB · 760 lines
Rust
at master
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761use std::{collections::VecDeque, mem};
use compiler_data_structures::index_vec::Key;use rustc_hash::FxHashSet;use ustr::{ustr, Ustr};
use compiler_core::{ db::{AdtField, AdtKind, Db, DefId, StructDef, UnionDef}, mangle, middle::{BinOp, CallConv, CmpOp, NamePat, Pat}, span::{Span, Spanned as _}, sym, ty::{fold::TyFolder, FnTy, FnTyFlags, FnTyParam, Instantiation, Subst, SubstTy, Ty, TyKind}, word::Word,};
use crate::{ BlockId, Body, Const, Fn, FnParam, FnSig, FnSigId, FxHashMap, GlobalId, GlobalKind, IdMap, Inst, Mir, ValueId, ValueKind,};
pub fn monomorphize(db: &Db, mir: &mut Mir) { Monomorphize::new(db).run(mir); ExpandDrops::new(db, mir).run(mir);}
struct Monomorphize<'db> { db: &'db Db, work: Work, used_fns: FxHashSet<FnSigId>, used_globals: FxHashSet<GlobalId>,
// Functions that have already been monomorphized and should be re-used monomorphized_fns: FxHashMap<MonoFn, FnSigId>,
// Functions that are a result of monomorphization, and should never be monomorphized again dont_monomorphize: FxHashSet<FnSigId>,}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]pub struct MonoFn { pub id: FnSigId, pub ty: Ty,}
impl<'db> Monomorphize<'db> { fn new(db: &'db Db) -> Self { Self { db, work: Work::new(), used_fns: FxHashSet::default(), used_globals: FxHashSet::default(), monomorphized_fns: FxHashMap::default(), dont_monomorphize: FxHashSet::default(), } }
fn run(mut self, mir: &mut Mir) { self.monomorphize_bodies(mir); self.retain_used_mir(mir); }
fn monomorphize_bodies(&mut self, mir: &mut Mir) { let main_fn = mir.main_fn.expect("to have a main fn"); self.work.push(Job { target: JobTarget::Fn(main_fn) });
while let Some(job) = self.work.pop() { MonomorphizeBody::new(self, mir).run(mir, &job); } }
fn retain_used_mir(&mut self, mir: &mut Mir) { mir.fn_sigs.inner_mut().retain(|id, _| self.used_fns.contains(id)); mir.fns.retain(|id, _| self.used_fns.contains(id)); mir.globals.inner_mut().retain(|id, _| self.used_globals.contains(id)); }
fn mark_used(&mut self, target: JobTarget) { match target { JobTarget::Fn(id) => { self.used_fns.insert(id); } JobTarget::Global(id) => { self.used_globals.insert(id); } } }}
#[derive(Debug)]struct Work { queue: VecDeque<Job>, done: FxHashSet<JobTarget>,}
impl Work { fn new() -> Self { Self { queue: VecDeque::new(), done: FxHashSet::default() } }
pub fn push(&mut self, job: Job) { if !self.done.contains(&job.target) { self.queue.push_back(job); } }
pub fn pop(&mut self) -> Option<Job> { self.queue.pop_front() }
pub fn mark_done(&mut self, target: JobTarget) { self.done.insert(target); }}
#[derive(Debug)]struct Job { target: JobTarget,}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]enum JobTarget { Fn(FnSigId), Global(GlobalId),}
struct MonomorphizeBody<'db, 'cx> { cx: &'cx mut Monomorphize<'db>, mono_mir: MonoMir,}
impl<'db, 'cx> MonomorphizeBody<'db, 'cx> { fn new(cx: &'cx mut Monomorphize<'db>, mir: &Mir) -> Self { Self { cx, mono_mir: MonoMir::new(mir) } }
fn run(mut self, mir: &mut Mir, job: &Job) { self.cx.mark_used(job.target); self.monomorphize(mir, job.target); self.mono_mir.merge_into(mir); self.cx.work.mark_done(job.target); }
fn monomorphize(&mut self, mir: &mut Mir, target: JobTarget) { match target { JobTarget::Fn(id) => { let Some(fun) = mir.fns.get(&id) else { debug_assert!( mir.fn_sigs.contains_key(&id), "function must have an existing signature" ); return; };
let body_subst = self.monomorphize_body(mir, &fun.body); body_subst.subst(&mut mir.fns.get_mut(&id).unwrap().body); } JobTarget::Global(id) => { let global = mir.globals.get(&id).expect("global to exist");
if let GlobalKind::Static(body) = &global.kind { let body_subst = self.monomorphize_body(mir, body); let GlobalKind::Static(body_mut) = &mut mir.globals[id].kind else { unreachable!() }; body_subst.subst(body_mut); } } } }
fn monomorphize_body(&mut self, mir: &Mir, body: &Body) -> BodySubst { let mut body_subst = BodySubst::new();
for value in body.values() { if let Some(instantiation) = body.instantiation(value.id) { // This is a polymorphic value which requires specialization if let Some(new_kind) = self.monomorphize_value(mir, &value.kind, instantiation) { body_subst.insert_value(value.id, new_kind); } } else { // This is a monomorphic value, we can enqueue it safely match value.kind { ValueKind::Fn(id) => { self.cx.work.push(Job { target: JobTarget::Fn(id) }); } ValueKind::Global(id) => { self.cx.work.push(Job { target: JobTarget::Global(id) }); } _ => (), } } }
body_subst }
fn monomorphize_value( &mut self, mir: &Mir, value_kind: &ValueKind, instantiation: &Instantiation, ) -> Option<ValueKind> { match value_kind { &ValueKind::Fn(id) => { // We ignore builtin functions, as they're expanded during mir::lower if self.cx.db.builtins.contains_key(&mir.fn_sigs[id].def_id) { return None; }
if self.cx.dont_monomorphize.contains(&id) { return None; }
let ty = instantiation.fold(mir.fn_sigs[id].ty); let monomorphized_fn = MonoFn { id, ty };
let monomorphized_sig_id = self.monomorphize_fn(mir, monomorphized_fn, instantiation);
Some(ValueKind::Fn(monomorphized_sig_id)) } _ => { // This is a polymorphic type. Doesn't require specialization... None } } }
#[must_use] fn monomorphize_fn( &mut self, mir: &Mir, monomorphized_fn: MonoFn, instantiation: &Instantiation, ) -> FnSigId { if let Some(sig_id) = self.cx.monomorphized_fns.get(&monomorphized_fn).copied() { return sig_id; }
let old_sig_id = monomorphized_fn.id; let new_sig_id = self.monomorphize_fn_sig(mir, &monomorphized_fn, instantiation);
self.cx.monomorphized_fns.insert(monomorphized_fn, new_sig_id); self.cx.dont_monomorphize.insert(new_sig_id);
let mut fun = mir.fns.get(&old_sig_id).expect("fn to exist").clone();
fun.sig = new_sig_id; fun.subst(&mut MonoParamFolder { instantiation });
self.mono_mir.fns.insert(new_sig_id, fun); self.cx.work.push(Job { target: JobTarget::Fn(new_sig_id) });
new_sig_id }
#[must_use] fn monomorphize_fn_sig( &mut self, mir: &Mir, monomorphized_fn: &MonoFn, instantiation: &Instantiation, ) -> FnSigId { let mut sig = mir.fn_sigs[monomorphized_fn.id].clone(); sig.subst(&mut MonoParamFolder { instantiation });
let instantation_str = instantiation .tys() .map(|ty| mangle::ty_name(self.cx.db, ty)) .collect::<Vec<_>>() .join("_");
sig.mangled_name = ustr(&format!("{}$${}", sig.mangled_name, instantation_str));
self.mono_mir.fn_sigs.insert_with_key(|id| { sig.id = id; sig }) }}
#[derive(Debug)]struct BodySubst { value_subst: FxHashMap<ValueId, ValueKind>,}
impl BodySubst { fn new() -> Self { Self { value_subst: FxHashMap::default() } }
fn insert_value(&mut self, id: ValueId, kind: ValueKind) { self.value_subst.insert(id, kind); }
fn subst(self, body: &mut Body) { for (id, kind) in self.value_subst { body.value_mut(id).kind = kind; } }}
#[derive(Debug)]struct MonoMir { fn_sigs: IdMap<FnSigId, FnSig>, fns: FxHashMap<FnSigId, Fn>,}
impl MonoMir { fn new(mir: &Mir) -> Self { Self { fn_sigs: Self::fn_sigs_from_mir(mir), fns: FxHashMap::default() } }
fn merge_into(&mut self, mir: &mut Mir) { mir.fn_sigs.extend(mem::take(&mut self.fn_sigs)); mir.fns.extend(mem::take(&mut self.fns)); }
fn fn_sigs_from_mir(mir: &Mir) -> IdMap<FnSigId, FnSig> { IdMap::new_with_counter(*mir.fn_sigs.counter()) }}
#[derive(Debug)]struct ExpandDrops<'db> { db: &'db Db, mono_mir: MonoMir,
// Generated free functions for adt types free_fns: FxHashMap<Ty, FnSigId>,}
impl<'db> ExpandDrops<'db> { fn new(db: &'db Db, mir: &Mir) -> Self { Self { db, mono_mir: MonoMir::new(mir), free_fns: FxHashMap::default() } }
fn run(mut self, mir: &mut Mir) { for fun in mir.fns.values_mut() { self.body(&mut fun.body); }
for global in mir.globals.values_mut() { if let GlobalKind::Static(body) = &mut global.kind { self.body(body); } }
self.mono_mir.merge_into(mir); }
fn body(&mut self, body: &mut Body) { self.remove_unused_drops(body); self.expand_drop_glue(body); }
fn remove_unused_drops(&self, body: &mut Body) { let value_tys: FxHashMap<ValueId, Ty> = body.values().iter().map(|v| (v.id, v.ty)).collect();
for block in body.blocks_mut() { block.insts.retain_mut(|inst| match inst { Inst::Drop { value, drop_glue, span } => { let ty = value_tys[&*value];
if !self.should_drop_ty(ty, *drop_glue) { return false; }
if ty.is_ref() { *inst = Inst::DecRef { value: *value, span: *span }; }
true } Inst::IncRef { value, .. } | Inst::DecRef { value, .. } => { self.should_refcount_ty(value_tys[&*value]) } _ => true, }); } }
fn should_drop_ty(&self, ty: Ty, drop_glue: bool) -> bool { if self.should_refcount_ty(ty) { return true; }
if !ty.is_rc(self.db) && !drop_glue { return false; }
ty.needs_free(self.db) }
fn should_refcount_ty(&self, ty: Ty) -> bool { matches!(ty.kind(), TyKind::Ref(ty, _) if ty.is_rc(self.db)) }
fn expand_drop_glue(&mut self, body: &mut Body) { let mut expanded: Vec<(BlockId, usize, Option<FnSigId>)> = vec![];
for block in body.blocks() { for (idx, inst) in block.insts.iter().enumerate() { if let Inst::Drop { value, drop_glue, .. } = inst { let free_fn = if *drop_glue { let ty = body.value(*value).ty; Some(self.get_or_create_free_fn(ty)) } else { None }; expanded.push((block.id, idx, free_fn)); } } }
for (block, inst_idx, free_fn) in expanded { let Inst::Drop { value, span, .. } = body.block(block).insts[inst_idx] else { unreachable!() };
if let Some(free_fn) = free_fn { let free_fn_ty = self.mono_mir.fn_sigs[free_fn].ty; let result = body.create_register(self.db.types.unit); let callee = body.create_value(free_fn_ty, ValueKind::Fn(free_fn));
body.block_mut(block).insts[inst_idx] = Inst::Call { value: result, callee, args: vec![value], span }; } else { body.block_mut(block).insts[inst_idx] = Inst::Free { value, traced: true, span }; } } }
fn get_free_call_values(&mut self, body: &mut Body, ty: Ty) -> (ValueId, ValueId) { let free_fn = self.get_or_create_free_fn(ty);
let free_fn_ty = self.mono_mir.fn_sigs[free_fn].ty; let result = body.create_register(self.db.types.unit); let callee = body.create_value(free_fn_ty, ValueKind::Fn(free_fn));
(result, callee) }
fn get_or_create_free_fn(&mut self, ty: Ty) -> FnSigId { if let Some(sig_id) = self.free_fns.get(&ty) { return *sig_id; }
match ty.kind() { TyKind::Adt(..) => CreateAdtFree::new(self).create(ty), TyKind::Slice(..) | TyKind::Str => CreateSliceFree::new(self).create(ty), ty => unreachable!("unexpected ty {ty:?}"), } }}
struct MonoParamFolder<'a> { pub instantiation: &'a Instantiation,}
impl SubstTy for MonoParamFolder<'_> { fn subst_ty(&mut self, ty: Ty, _: Span) -> Ty { self.fold(ty) }}
impl TyFolder for MonoParamFolder<'_> { fn fold(&mut self, ty: Ty) -> Ty { match ty.kind() { TyKind::Param(p) => match self.instantiation.get(p.var) { Some(ty) => ty, None => ty, }, _ => self.super_fold(ty), } }}
#[derive(Debug)]struct CreateAdtFree<'cx, 'db> { cx: &'cx mut ExpandDrops<'db>, body: Body,}
impl<'cx, 'db> CreateAdtFree<'cx, 'db> { fn new(cx: &'cx mut ExpandDrops<'db>) -> Self { Self { cx, body: Body::new() } }
fn create(mut self, ty: Ty) -> FnSigId { let TyKind::Adt(adt_id, targs) = ty.kind() else { unreachable!() };
let adt = &self.cx.db[*adt_id]; let instantiation = adt.instantiation(targs); let adt_ty = instantiation.fold(adt.ty()); let adt_span = adt.name.span();
let adt_name = mangle::adt_name(self.cx.db, adt, targs); let mangled_name = ustr(&format!("{adt_name}_drop")); let display_name = ustr(&self.cx.db[adt.def_id].qpath.clone().child(ustr("$drop")).join()); let sig = create_drop_sig(self.cx, adt_ty, mangled_name, display_name, adt_span);
let self_value = self.body.create_value(adt_ty, ValueKind::Param(DefId::null(), 0));
let join_block = match &adt.kind { AdtKind::Union(union_def) => { self.lower_union_free(union_def, self_value, &instantiation, adt_span) } AdtKind::Struct(struct_def) => { self.lower_struct_free(struct_def, self_value, &instantiation, adt_span) } };
if adt.is_ref() { self.body.ins(join_block).free(self_value, false, adt_span); }
let unit_value = self.body.create_value(self.cx.db.types.unit, ValueKind::Const(Const::Unit)); self.body.ins(join_block).ret(unit_value, adt_span);
self.cx.mono_mir.fns.insert(sig, Fn { sig, body: self.body });
sig }
fn lower_struct_free( &mut self, struct_def: &StructDef, self_value: ValueId, instantiation: &Instantiation, span: Span, ) -> BlockId { let start_block = self.body.create_block("start"); self.free_adt_fields(start_block, self_value, &struct_def.fields, instantiation, span); start_block }
fn lower_union_free( &mut self, union_def: &UnionDef, self_value: ValueId, instantiation: &Instantiation, span: Span, ) -> BlockId { let start_block = self.body.create_block("start"); let join_block = self.body.create_block("join"); let self_ty = self.body.value(self_value).ty;
let mut blocks = vec![];
for &variant_id in &union_def.variants { let variant = &self.cx.db[variant_id];
let variant_value = self.body.create_value(self_ty, ValueKind::Variant(self_value, variant.id));
let block = self.body.create_block(format!("case_{}", variant.name));
self.free_adt_fields(block, variant_value, &variant.fields, instantiation, span);
self.body.ins(block).br(join_block); blocks.push(block); }
let uint = self.cx.db.types.uint; let tag_field = self.body.create_value(uint, ValueKind::Field(self_value, ustr("tag"))); self.body.ins(start_block).switch(tag_field, blocks, span);
join_block }
fn free_adt_fields( &mut self, block: BlockId, self_value: ValueId, fields: &[AdtField], instantiation: &Instantiation, span: Span, ) { for field in fields { let ty = instantiation.fold(field.ty); let field_value = self.body.create_value(field.ty, ValueKind::Field(self_value, field.name.name()));
match ty.kind() { TyKind::Adt(..) => { let (result, callee) = self.cx.get_free_call_values(&mut self.body, ty);
self.body.ins(block).call(result, callee, vec![field_value], span); } TyKind::Ref(..) if self.cx.should_refcount_ty(ty) => { self.body.ins(block).decref(field_value, field.span()); } _ => (), } } }}
#[derive(Debug)]struct CreateSliceFree<'cx, 'db> { cx: &'cx mut ExpandDrops<'db>, body: Body,}
impl<'cx, 'db> CreateSliceFree<'cx, 'db> { fn new(cx: &'cx mut ExpandDrops<'db>) -> Self { Self { cx, body: Body::new() } }
fn create(mut self, ty: Ty) -> FnSigId { let elem_ty = ty.slice_elem().unwrap();
let main_span = Span::uniform(self.cx.db.main_source_id(), 0);
let tyname = mangle::ty_name(self.cx.db, ty); let mangled_name = ustr(&format!("{tyname}_drop")); let display_name = ustr(&format!("{}$drop", ty.display(self.cx.db))); let sig = create_drop_sig(self.cx, ty, mangled_name, display_name, main_span);
let slice = self.body.create_value(ty, ValueKind::Param(DefId::null(), 0));
let start = self.body.create_block("start");
let end_block = if self.cx.should_drop_ty(elem_ty, true) { self.free_slice_elems(start, slice, elem_ty, main_span) } else { start };
let unit_value = self.body.create_value(self.cx.db.types.unit, ValueKind::Const(Const::Unit)); self.body.ins(end_block).free(slice, false, main_span).ret(unit_value, main_span);
self.cx.mono_mir.fns.insert(sig, Fn { sig, body: self.body });
sig }
// Generates a loop which goes over all slice elements, droping them. // The generated code in psuedo-code: // // ```rust // let i = 0 // loop { // if i == slice.len { break } // free(slice[i]) // } // ``` fn free_slice_elems( &mut self, start: BlockId, slice: ValueId, elem_ty: Ty, span: Span, ) -> BlockId { let uint = self.cx.db.types.uint;
let loop_start = self.body.create_block("loop_start"); let loop_end = self.body.create_block("loop_end");
// uint i = 0; let index = self.body.create_register(uint); let zero = self.body.create_value(uint, ValueKind::Const(Const::Int(0))); self.body.ins(start).stackalloc(index, zero, span).br(loop_start);
// slice.len let slice_len = self.body.create_value(uint, ValueKind::Field(slice, ustr(sym::field::LEN)));
// i == slice.len let cond = self.body.create_register(self.cx.db.types.bool); self.body.ins(loop_start).binary_unchecked( cond, index, slice_len, BinOp::Cmp(CmpOp::Eq), span, );
self.body.ins(loop_start).brif(cond, loop_end, None, span); self.free_elem(loop_start, slice, index, elem_ty, span);
// i += 1 let next_index = self.body.create_register(uint); let one = self.body.create_value(uint, ValueKind::Const(Const::Int(1))); self.body .ins(loop_start) .binary_unchecked(next_index, index, one, BinOp::Add, span) .store(next_index, index, span);
self.body.ins(loop_start).br(loop_start);
loop_end }
fn free_elem( &mut self, block: BlockId, slice: ValueId, index: ValueId, elem_ty: Ty, span: Span, ) { let elem = self.body.create_register(elem_ty); self.body.ins(block).slice_index_unchecked(elem, slice, index, span);
match elem_ty.kind() { TyKind::Adt(..) | TyKind::Slice(..) | TyKind::Str => { let (result, callee) = self.cx.get_free_call_values(&mut self.body, elem_ty); self.body.ins(block).call(result, callee, vec![elem], span); } TyKind::Ref(..) if self.cx.should_refcount_ty(elem_ty) => { self.body.ins(block).decref(elem, span); } ty => unreachable!("{ty:?}"), } }}
fn create_drop_sig( cx: &mut ExpandDrops, ty: Ty, mangled_name: Ustr, display_name: Ustr, span: Span,) -> FnSigId { let params = vec![FnParam { pat: Pat::Name(NamePat::new_anon(Word::new(ustr("self"), Span::unknown()), ty)), ty, }];
let fn_ty_params = params.iter().map(|p| FnTyParam { name: p.pat.param_name(), ty: p.ty }).collect();
let fn_ty = Ty::new(TyKind::Fn(FnTy { params: fn_ty_params, ret: cx.db.types.unit, callconv: CallConv::default(), flags: FnTyFlags::empty(), }));
let sig = cx.mono_mir.fn_sigs.insert_with_key(|id| FnSig { id, def_id: DefId::null(), mangled_name, display_name, params, ty: fn_ty, is_inline: true, span, });
cx.free_fns.insert(ty, sig);
sig}