use std::collections::BTreeMap;
use rustc_hash::{FxHashMap, FxHashSet};
use super::display::TypeDisplay;
use super::graph::DefinitionGraph;
use super::ty::{Bound, GenericId, GenericParam, Scheme, Type};
use super::typed::{TypedDefinition, TypedExpr, TypedExprKind, TypedLeaf, TypedProgram};
use crate::core::{BindingId, ItemKind, PrimitiveType};
use crate::expand::Program;
#[derive(Clone, Copy, PartialEq, Eq)]
enum UseKind {
Call,
Value,
}
struct UseTypes {
types: Vec<Option<Type>>,
resolved: bool,
kind: UseKind,
}
struct UseSite {
enclosing: BindingId,
actual: Type,
kind: UseKind,
}
struct FunctionUses {
params: Vec<GenericId>,
literal: Vec<(GenericId, PrimitiveType)>,
body: Type,
sites: Vec<UseSite>,
uses: Vec<UseTypes>,
written: bool,
}
struct Target {
params: Vec<GenericId>,
literal: Vec<(GenericId, PrimitiveType)>,
body: Type,
}
pub struct InstantiationReport {
entries: BTreeMap<BindingId, FunctionUses>,
substitutions: FxHashMap<BindingId, FxHashMap<GenericId, Type>>,
}
impl InstantiationReport {
pub fn of(program: &Program, graph: &DefinitionGraph, typed: &TypedProgram) -> Self {
let written = Self::written_signatures(program);
let groups = Self::recursive_groups(graph);
let targets = Self::targets(typed);
let mut entries = Self::entries(&targets, &written);
Self::collect_uses(typed, &targets, &groups, &mut entries);
let substitutions = Self::group_substitutions(graph, &mut entries);
Self {
entries,
substitutions,
}
}
pub fn one_type(&self, binding: BindingId) -> Option<FxHashMap<GenericId, Type>> {
self.substitutions.get(&binding).cloned()
}
pub fn uses_of(&self, binding: BindingId, generic: GenericId) -> Vec<Option<Type>> {
let Some(entry) = self.entries.get(&binding) else {
return Vec::new();
};
let Some(index) = entry.params.iter().position(|param| *param == generic) else {
return Vec::new();
};
entry
.uses
.iter()
.map(|used| used.types[index].clone())
.collect()
}
pub fn is_used(&self, binding: BindingId) -> bool {
self.entries
.get(&binding)
.is_some_and(|entry| !entry.sites.is_empty())
}
fn written_signatures(program: &Program) -> FxHashSet<BindingId> {
program
.modules
.iter()
.flat_map(|module| &module.items)
.filter_map(|item| match &item.kind {
ItemKind::Signature(binding, _) => Some(*binding),
_ => None,
})
.collect()
}
fn recursive_groups(graph: &DefinitionGraph) -> FxHashMap<BindingId, usize> {
let mut groups = FxHashMap::default();
for (index, component) in graph.components().iter().enumerate() {
for binding in component {
groups.insert(*binding, index);
}
}
groups
}
fn targets(typed: &TypedProgram) -> FxHashMap<BindingId, Target> {
typed
.modules
.iter()
.flat_map(|module| &module.definitions)
.filter(|definition| {
matches!(definition.value.kind, TypedExprKind::Fn(_))
&& !definition.scheme.params.is_empty()
})
.map(|definition| {
let target = Target {
params: definition
.scheme
.params
.iter()
.map(|param| param.id)
.collect(),
literal: definition
.scheme
.params
.iter()
.filter_map(|param| {
typed
.literal_generics
.get(¶m.id)
.map(|literal| (param.id, literal.default_type()))
})
.collect(),
body: definition.scheme.body.clone(),
};
(definition.binding, target)
})
.collect()
}
fn entries(
targets: &FxHashMap<BindingId, Target>,
written: &FxHashSet<BindingId>,
) -> BTreeMap<BindingId, FunctionUses> {
targets
.iter()
.map(|(binding, target)| {
let entry = FunctionUses {
params: target.params.clone(),
literal: target.literal.clone(),
body: target.body.clone(),
sites: Vec::new(),
uses: Vec::new(),
written: written.contains(binding),
};
(*binding, entry)
})
.collect()
}
fn collect_uses(
typed: &TypedProgram,
targets: &FxHashMap<BindingId, Target>,
groups: &FxHashMap<BindingId, usize>,
entries: &mut BTreeMap<BindingId, FunctionUses>,
) {
for module in &typed.modules {
for definition in &module.definitions {
let operands = typed.type_texts.operands_of(definition.binding);
for checked in std::iter::once(&definition.value).chain(operands) {
Self::collect_expr(checked, definition.binding, targets, groups, entries);
}
}
}
}
fn collect_expr(
expr: &TypedExpr,
enclosing: BindingId,
targets: &FxHashMap<BindingId, Target>,
groups: &FxHashMap<BindingId, usize>,
entries: &mut BTreeMap<BindingId, FunctionUses>,
) {
if let TypedExprKind::Call(callee, arguments) = &expr.kind {
Self::collect_callee(callee, arguments, enclosing, targets, groups, entries);
return;
}
if let Some((binding, actual)) = expr.function_value() {
let site = UseSite {
enclosing,
actual: actual.clone(),
kind: UseKind::Value,
};
Self::record(binding, site, targets, groups, entries);
return;
}
for child in expr.children() {
Self::collect_expr(child, enclosing, targets, groups, entries);
}
}
fn collect_callee(
callee: &TypedExpr,
arguments: &[super::typed::TypedArgument],
enclosing: BindingId,
targets: &FxHashMap<BindingId, Target>,
groups: &FxHashMap<BindingId, usize>,
entries: &mut BTreeMap<BindingId, FunctionUses>,
) {
match &callee.kind {
TypedExprKind::Leaf(TypedLeaf::Global(binding)) => {
let site = UseSite {
enclosing,
actual: callee.ty.clone(),
kind: UseKind::Call,
};
Self::record(*binding, site, targets, groups, entries);
}
_ => Self::collect_expr(callee, enclosing, targets, groups, entries),
}
for argument in arguments {
Self::collect_expr(&argument.value, enclosing, targets, groups, entries);
}
}
fn record(
target_binding: BindingId,
site: UseSite,
targets: &FxHashMap<BindingId, Target>,
groups: &FxHashMap<BindingId, usize>,
entries: &mut BTreeMap<BindingId, FunctionUses>,
) {
if groups.get(&site.enclosing) == groups.get(&target_binding) {
return;
}
if !targets.contains_key(&target_binding) {
return;
}
let entry = entries
.get_mut(&target_binding)
.expect("every target has an entry");
entry.sites.push(site);
}
fn use_types(body: &Type, params: &[GenericId], actual: &Type, kind: UseKind) -> UseTypes {
let mut bound = FxHashMap::default();
let matched = Self::bind(body, actual, &mut bound);
let types: Vec<Option<Type>> = params
.iter()
.map(|param| bound.get(param).cloned())
.collect();
let resolved = matched
&& types
.iter()
.all(|ty| ty.as_ref().is_some_and(Type::is_concrete));
UseTypes {
types,
resolved,
kind,
}
}
fn bind(scheme: &Type, actual: &Type, bound: &mut FxHashMap<GenericId, Type>) -> bool {
match scheme {
Type::Generic(generic) => {
bound.insert(*generic, actual.clone());
true
}
Type::Named(..) => Self::bind_named(scheme, actual, bound),
Type::Function(..) => Self::bind_function(scheme, actual, bound),
Type::Ref(..) | Type::Mut(..) => Self::bind_reference(scheme, actual, bound),
Type::Primitive(primitive) => {
matches!(actual, Type::Primitive(other) if primitive == other)
}
Type::Var(_) | Type::TraitObject(_) => false,
}
}
fn bind_named(scheme: &Type, actual: &Type, bound: &mut FxHashMap<GenericId, Type>) -> bool {
let (Type::Named(left, left_arguments), Type::Named(right, right_arguments)) =
(scheme, actual)
else {
return false;
};
left == right
&& left_arguments.len() == right_arguments.len()
&& Self::bind_all(left_arguments, right_arguments, bound)
}
fn bind_function(scheme: &Type, actual: &Type, bound: &mut FxHashMap<GenericId, Type>) -> bool {
let (
Type::Function(_, left_parameters, left_result),
Type::Function(_, right_parameters, right_result),
) = (scheme, actual)
else {
return false;
};
left_parameters.len() == right_parameters.len()
&& Self::bind_all(left_parameters, right_parameters, bound)
&& Self::bind(left_result, right_result, bound)
}
fn bind_reference(
scheme: &Type,
actual: &Type,
bound: &mut FxHashMap<GenericId, Type>,
) -> bool {
match (scheme, actual) {
(Type::Ref(_, left), Type::Ref(_, right))
| (Type::Mut(_, left), Type::Mut(_, right)) => Self::bind(left, right, bound),
_ => false,
}
}
fn bind_all(scheme: &[Type], actual: &[Type], bound: &mut FxHashMap<GenericId, Type>) -> bool {
scheme
.iter()
.zip(actual)
.all(|(scheme, actual)| Self::bind(scheme, actual, bound))
}
fn group_substitutions(
graph: &DefinitionGraph,
entries: &mut BTreeMap<BindingId, FunctionUses>,
) -> FxHashMap<BindingId, FxHashMap<GenericId, Type>> {
let mut substitutions = FxHashMap::default();
for component in graph.components().iter().rev() {
let members: Vec<BindingId> = component
.iter()
.filter(|binding| entries.contains_key(binding))
.copied()
.collect();
Self::resolve_group_uses(&members, entries, &substitutions);
if let Some(substitution) = Self::decide_group(&members, entries)
&& !Self::only_literal_defaults(&members, entries, &substitution)
{
for member in &members {
substitutions.insert(*member, substitution.clone());
}
}
}
substitutions
}
fn resolve_group_uses(
members: &[BindingId],
entries: &mut BTreeMap<BindingId, FunctionUses>,
callers: &FxHashMap<BindingId, FxHashMap<GenericId, Type>>,
) {
for member in members {
let entry = entries.get_mut(member).expect("every member has an entry");
entry.uses = entry
.sites
.iter()
.map(|site| Self::resolved_use(entry, site, callers))
.collect();
}
}
fn resolved_use(
entry: &FunctionUses,
site: &UseSite,
callers: &FxHashMap<BindingId, FxHashMap<GenericId, Type>>,
) -> UseTypes {
match callers.get(&site.enclosing) {
Some(caller) => {
let actual = site.actual.substitute(caller);
Self::use_types(&entry.body, &entry.params, &actual, site.kind)
}
None => Self::use_types(&entry.body, &entry.params, &site.actual, site.kind),
}
}
fn only_literal_defaults(
members: &[BindingId],
entries: &BTreeMap<BindingId, FunctionUses>,
substitution: &FxHashMap<GenericId, Type>,
) -> bool {
substitution.iter().all(|(generic, ty)| {
members.iter().any(|member| {
entries[member]
.literal
.iter()
.any(|(id, primitive)| id == generic && *ty == Type::Primitive(*primitive))
})
})
}
fn decide_group(
members: &[BindingId],
entries: &BTreeMap<BindingId, FunctionUses>,
) -> Option<FxHashMap<GenericId, Type>> {
if members.iter().any(|member| entries[member].written) {
return None;
}
let combined = Self::combined_uses(members, entries)?;
members
.iter()
.all(|member| Self::covers_params(&entries[member], &combined))
.then_some(combined)
}
fn combined_uses(
members: &[BindingId],
entries: &BTreeMap<BindingId, FunctionUses>,
) -> Option<FxHashMap<GenericId, Type>> {
let mut combined: FxHashMap<GenericId, Type> = FxHashMap::default();
let mut uses = 0;
for member in members {
for used in &entries[member].uses {
if !used.resolved {
return None;
}
uses += 1;
Self::merge_use(&mut combined, &entries[member].params, used)?;
}
}
(uses > 0).then_some(combined)
}
fn merge_use(
combined: &mut FxHashMap<GenericId, Type>,
params: &[GenericId],
used: &UseTypes,
) -> Option<()> {
for (param, ty) in params.iter().zip(&used.types) {
let ty = ty.clone().expect("a resolved use names every parameter");
match combined.entry(*param) {
std::collections::hash_map::Entry::Occupied(existing) => {
if *existing.get() != ty {
return None;
}
}
std::collections::hash_map::Entry::Vacant(entry) => {
entry.insert(ty);
}
}
}
Some(())
}
fn covers_params(entry: &FunctionUses, combined: &FxHashMap<GenericId, Type>) -> bool {
entry
.params
.iter()
.all(|param| combined.contains_key(param))
}
}
#[derive(Clone, Default)]
pub struct LiteralDefaults {
entries: FxHashMap<(BindingId, usize), LiteralDefault>,
}
#[derive(Clone)]
pub struct LiteralDefault {
pub parameter: String,
pub default: String,
pub calls: Vec<String>,
pub values: Vec<String>,
pub signature: String,
pub generic_signature: String,
}
impl LiteralDefaults {
pub fn of(report: &InstantiationReport, typed: &TypedProgram) -> Self {
let display = TypeDisplay::new(&typed.bindings, &typed.generics)
.with_projections(typed.traits.projections());
let mut defaults = Self::default();
for definition in typed.modules.iter().flat_map(|module| &module.definitions) {
defaults.collect(report, typed, definition, &display);
}
defaults
}
pub fn get(&self, binding: BindingId, position: usize) -> Option<&LiteralDefault> {
self.entries.get(&(binding, position))
}
fn collect(
&mut self,
report: &InstantiationReport,
typed: &TypedProgram,
definition: &TypedDefinition,
display: &TypeDisplay<'_>,
) {
let TypedExprKind::Fn(lambda) = &definition.value.kind else {
return;
};
let Type::Function(_, parameters, result) = &definition.scheme.body else {
return;
};
let Some(entry) = report.entries.get(&definition.binding) else {
return;
};
if entry.written {
return;
}
for (position, parameter) in parameters.iter().enumerate() {
let Type::Generic(generic) = parameter else {
continue;
};
let Some(literal) = typed.literal_generics.get(generic) else {
continue;
};
let default = literal.default_type();
let Some(generic_index) = Self::generic_index(definition, *generic) else {
continue;
};
if Self::distinct_use_types(entry, generic_index, display).len() < 2 {
continue;
}
let Some(binder) = lambda.params.get(position) else {
continue;
};
let signature = LiteralSignature {
definition,
entry,
typed,
display,
};
let concrete = signature.concrete(parameters, result);
self.entries.insert(
(definition.binding, position),
LiteralDefault {
parameter: typed.bindings.get(binder.binding).name().to_string(),
default: default.name().to_string(),
calls: Self::kind_types(entry, generic_index, UseKind::Call, display),
values: Self::kind_types(entry, generic_index, UseKind::Value, display),
generic_signature: signature.generic(),
signature: concrete,
},
);
}
}
fn generic_index(definition: &TypedDefinition, generic: GenericId) -> Option<usize> {
definition
.scheme
.params
.iter()
.position(|param| param.id == generic)
}
fn distinct_use_types(
entry: &FunctionUses,
generic_index: usize,
display: &TypeDisplay<'_>,
) -> Vec<String> {
let values = entry.uses.iter().filter(|used| used.kind == UseKind::Value);
let calls = entry.uses.iter().filter(|used| used.kind == UseKind::Call);
Self::distinct_types(values.chain(calls), generic_index, display)
}
fn kind_types(
entry: &FunctionUses,
generic_index: usize,
kind: UseKind,
display: &TypeDisplay<'_>,
) -> Vec<String> {
let uses = entry.uses.iter().filter(|used| used.kind == kind);
Self::distinct_types(uses, generic_index, display)
}
fn distinct_types<'b>(
uses: impl Iterator<Item = &'b UseTypes>,
generic_index: usize,
display: &TypeDisplay<'_>,
) -> Vec<String> {
let mut types = Vec::new();
for used in uses {
let Some(Some(ty)) = used.types.get(generic_index) else {
continue;
};
let shown = display.show(ty);
if !types.contains(&shown) {
types.push(shown);
}
}
types
}
}
struct LiteralSignature<'a> {
definition: &'a TypedDefinition,
entry: &'a FunctionUses,
typed: &'a TypedProgram,
display: &'a TypeDisplay<'a>,
}
impl LiteralSignature<'_> {
fn concrete(&self, parameters: &[Type], result: &Type) -> String {
let parameters: Vec<String> = parameters
.iter()
.map(|parameter| self.shown_type(parameter))
.collect();
format!(
"(fn ({}) {})",
parameters.join(" "),
self.shown_type(result)
)
}
fn generic(&self) -> String {
let scheme = Scheme {
params: self
.definition
.scheme
.params
.iter()
.map(|param| self.with_literal_bound(param))
.collect(),
..self.definition.scheme.clone()
};
self.display.scheme(&scheme)
}
fn with_literal_bound(&self, param: &GenericParam) -> GenericParam {
let mut bounded = param.clone();
let literal_bound = self.literal_bound(param.id).filter(|bound| {
!param
.bounds
.iter()
.any(|known| known.binding == bound.binding)
});
bounded.bounds.extend(literal_bound);
bounded
}
fn literal_bound(&self, generic: GenericId) -> Option<Bound> {
let literal = self.typed.literal_generics.get(&generic)?;
let accepted: Vec<Type> = PrimitiveType::all()
.filter(|primitive| literal.accepts(*primitive))
.map(Type::Primitive)
.collect();
Some(Bound {
binding: self.typed.traits.implemented_exactly_by(&accepted)?,
argument: None,
origins: Vec::new(),
})
}
fn shown_type(&self, shown: &Type) -> String {
let Type::Generic(generic) = shown else {
return self.display.show(shown);
};
let Some(literal) = self.typed.literal_generics.get(generic) else {
return self.display.show(shown);
};
let Some(generic_index) = LiteralDefaults::generic_index(self.definition, *generic) else {
return self.display.show(shown);
};
LiteralDefaults::distinct_use_types(self.entry, generic_index, self.display)
.first()
.cloned()
.unwrap_or_else(|| literal.default_type().name().to_string())
}
}
#[derive(Clone, Default)]
pub struct Pinning {
pins: FxHashMap<BindingId, Type>,
literal_pins: FxHashMap<BindingId, Type>,
}
impl Pinning {
pub fn of(report: &InstantiationReport) -> Self {
let mut pinning = Self::default();
for (binding, entry) in &report.entries {
if let Some(substitution) = report.one_type(*binding) {
pinning.pin(*binding, entry, &substitution);
}
}
pinning
}
fn pin(
&mut self,
binding: BindingId,
entry: &FunctionUses,
substitution: &FxHashMap<GenericId, Type>,
) {
let literal_substitution = substitution
.iter()
.filter(|(generic, _)| entry.literal.iter().any(|(literal, _)| literal == *generic))
.map(|(generic, ty)| (*generic, ty.clone()))
.collect();
self.pins
.insert(binding, entry.body.substitute(substitution));
self.literal_pins
.insert(binding, entry.body.substitute(&literal_substitution));
}
pub fn pins(&self) -> &FxHashMap<BindingId, Type> {
&self.pins
}
pub fn literal_pins(&self) -> &FxHashMap<BindingId, Type> {
&self.literal_pins
}
pub fn is_empty(&self) -> bool {
self.pins.is_empty()
}
pub fn pins_type_parameters(&self) -> bool {
self.pins != self.literal_pins
}
pub fn unknown_to(&self, known: &Pinning) -> Pinning {
let unknown = |pins: &FxHashMap<BindingId, Type>| {
pins.iter()
.filter(|(binding, _)| !known.pins.contains_key(binding))
.map(|(binding, ty)| (*binding, ty.clone()))
.collect()
};
Pinning {
pins: unknown(&self.pins),
literal_pins: unknown(&self.literal_pins),
}
}
pub fn extend(&mut self, other: Pinning) {
self.pins.extend(other.pins);
self.literal_pins.extend(other.literal_pins);
}
}
#[cfg(test)]
mod tests {
use rustc_hash::{FxHashMap, FxHashSet};
use super::{InstantiationReport, Pinning};
use crate::core::{
BindingArena, BindingId, BindingInfo, BindingKind, ClosureKind, ModuleId, PrimitiveType,
};
use crate::types::{DefinitionGraph, GenericId, Type, TypeVar, TypedSource};
struct BindFixture {
bound: FxHashMap<GenericId, Type>,
named: BindingId,
}
impl BindFixture {
fn new() -> Self {
let mut bindings = BindingArena::default();
let named = bindings.add(BindingInfo::new(
"item".into(),
BindingKind::Global(ModuleId::new(0)),
None,
));
Self {
bound: FxHashMap::default(),
named,
}
}
fn bind(&mut self, scheme: &Type, actual: &Type) -> bool {
InstantiationReport::bind(scheme, actual, &mut self.bound)
}
fn primitive(&self) -> Type {
Type::Primitive(PrimitiveType::I64)
}
}
#[test]
fn a_type_variable_scheme_does_not_bind() {
let mut fixture = BindFixture::new();
let scheme = Type::Var(TypeVar::new(0));
assert!(!fixture.bind(&scheme, &fixture.primitive()));
}
#[test]
fn a_named_scheme_does_not_bind_a_primitive() {
let mut fixture = BindFixture::new();
let scheme = Type::Named(fixture.named, Vec::new());
assert!(!fixture.bind(&scheme, &fixture.primitive()));
}
#[test]
fn a_function_scheme_does_not_bind_a_primitive() {
let mut fixture = BindFixture::new();
let scheme = Type::Function(ClosureKind::Fn, Vec::new(), Box::new(Type::unit()));
assert!(!fixture.bind(&scheme, &fixture.primitive()));
}
#[test]
fn a_reference_scheme_does_not_bind_a_primitive() {
let mut fixture = BindFixture::new();
let scheme = Type::shared(Type::Primitive(PrimitiveType::I64));
assert!(!fixture.bind(&scheme, &fixture.primitive()));
}
#[test]
fn a_pin_unknown_to_the_known_pins_keeps_its_literal_pin() {
let binding = BindFixture::new().named;
let pinned = Type::Primitive(PrimitiveType::U8);
let found = Pinning {
pins: FxHashMap::from_iter([(binding, pinned.clone())]),
literal_pins: FxHashMap::from_iter([(binding, pinned.clone())]),
};
let unknown = found.unknown_to(&Pinning::default());
assert_eq!(unknown.literal_pins().get(&binding), Some(&pinned));
}
struct ChainFixture {
source: TypedSource,
}
impl ChainFixture {
const WRAPPER_CHAIN: &str = "
(defn inner (x) x)
(defn middle (x) (inner x))
(defn outer (x) (middle x))
(defn main () (println (outer (the u16 1))))";
fn new() -> Self {
Self {
source: TypedSource::expand(Self::WRAPPER_CHAIN),
}
}
fn pinned_after_one_run(&self) -> FxHashSet<String> {
let program = self.source.program();
let graph = DefinitionGraph::of(program);
let typed = self.source.check(&graph);
let report = InstantiationReport::of(program, &graph, &typed);
Pinning::of(&report)
.pins()
.keys()
.map(|binding| program.bindings.get(*binding).name().to_string())
.collect()
}
}
#[test]
fn one_run_pins_every_wrapper_that_a_pinned_caller_reaches() {
let pinned = ChainFixture::new().pinned_after_one_run();
let expected = FxHashSet::from_iter(["inner", "middle", "outer"].map(String::from));
assert_eq!(pinned, expected);
}
}