use std::cell::Cell;
use std::num::NonZeroU32;
use std::rc::Rc;
use rustc_hash::FxHashMap;
use super::projection::ProjectionTable;
use super::ty::{Bound, GenericId, Type, TypeVar};
use crate::core::{BindingId, PrimitiveType};
use crate::source::Span;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LiteralKind {
Integer,
Float,
}
impl LiteralKind {
pub fn default_type(self) -> PrimitiveType {
match self {
LiteralKind::Float => PrimitiveType::F64,
LiteralKind::Integer => PrimitiveType::I64,
}
}
pub fn accepts(self, primitive: PrimitiveType) -> bool {
match self {
LiteralKind::Integer => primitive.is_integer() || primitive.is_float(),
LiteralKind::Float => primitive.is_float(),
}
}
fn merge(left: Option<LiteralKind>, right: Option<LiteralKind>) -> Option<LiteralKind> {
match (left, right) {
(Some(LiteralKind::Float), _) | (_, Some(LiteralKind::Float)) => {
Some(LiteralKind::Float)
}
(Some(LiteralKind::Integer), _) | (_, Some(LiteralKind::Integer)) => {
Some(LiteralKind::Integer)
}
(None, None) => None,
}
}
}
pub enum UnifyError {
Mismatch,
Infinite,
NotNumeric,
}
#[derive(Clone, Copy)]
pub struct ObligationSite {
pub span: Span,
pub bound_site: Option<Span>,
}
pub struct Obligation {
pub bound: Bound,
pub implementor: Type,
pub site: ObligationSite,
pub inside_operand: bool,
}
pub struct DeferredEquality {
pub projected: Type,
pub other: Type,
pub span: Span,
}
#[derive(Default)]
struct VarState {
binding: Option<Type>,
shortcut: Cell<Option<TypeVar>>,
literal: Option<LiteralKind>,
bounds: Vec<Bound>,
bound_site: Option<BoundSiteId>,
}
#[derive(Clone, Copy)]
struct BoundSiteId(NonZeroU32);
#[derive(Default)]
pub struct Unifier {
vars: Vec<VarState>,
bound_sites: Vec<Span>,
obligations: Vec<Obligation>,
generic_bounds: FxHashMap<GenericId, Vec<Bound>>,
projections: Rc<ProjectionTable>,
deferred: Vec<DeferredEquality>,
inside_operand: bool,
}
impl Unifier {
pub fn adopt_projections(&mut self, projections: Rc<ProjectionTable>) {
self.projections = projections;
}
pub fn take_deferred(&mut self) -> Vec<DeferredEquality> {
std::mem::take(&mut self.deferred)
}
pub fn fresh(&mut self) -> Type {
self.fresh_with(None, Vec::new())
}
pub fn fresh_literal(&mut self, kind: LiteralKind) -> Type {
self.fresh_with(Some(kind), Vec::new())
}
pub fn register_generic(&mut self, generic: GenericId, bounds: Vec<Bound>) {
self.generic_bounds.insert(generic, bounds);
}
pub fn generic_bounds(&self, generic: GenericId) -> &[Bound] {
self.generic_bounds
.get(&generic)
.map(Vec::as_slice)
.unwrap_or_default()
}
pub fn literal_kind(&self, var: TypeVar) -> Option<LiteralKind> {
self.vars[var.index()].literal
}
pub fn holds_open_literal(&self, value: &Type) -> bool {
match value {
Type::Var(start) => self.var_holds_open_literal(*start),
Type::Named(_, arguments) => self.any_holds_open_literal(arguments),
Type::Function(_, parameters, _) => self.any_holds_open_literal(parameters),
Type::Ref(_, inner) | Type::Mut(_, inner) => self.holds_open_literal(inner),
_ => false,
}
}
fn var_holds_open_literal(&self, start: TypeVar) -> bool {
let last = &self.vars[self.last_var(start).index()];
match &last.binding {
Some(bound) => self.holds_open_literal(bound),
None => last.literal.is_some(),
}
}
fn any_holds_open_literal(&self, values: &[Type]) -> bool {
values.iter().any(|value| self.holds_open_literal(value))
}
pub fn bounds(&self, var: TypeVar) -> &[Bound] {
&self.vars[var.index()].bounds
}
pub fn set_bounds(&mut self, var: TypeVar, bounds: Vec<Bound>, site: Span) {
let bound_site = (!bounds.is_empty()).then(|| self.intern_bound_site(site));
let state = &mut self.vars[var.index()];
state.bound_site = bound_site;
state.bounds = bounds;
}
fn intern_bound_site(&mut self, site: Span) -> BoundSiteId {
if self.bound_sites.last() != Some(&site) {
self.bound_sites.push(site);
}
let count = u32::try_from(self.bound_sites.len()).expect("fewer sites than variables");
BoundSiteId(NonZeroU32::new(count).expect("the last site exists"))
}
fn bound_site(&self, var: TypeVar) -> Option<Span> {
let BoundSiteId(id) = self.vars[var.index()].bound_site?;
self.bound_sites.get(id.get() as usize - 1).copied()
}
fn add_bound_from(&mut self, var: TypeVar, bound: Bound, site: Option<BoundSiteId>) {
self.add_bound(var, bound);
let bound_site = &mut self.vars[var.index()].bound_site;
*bound_site = bound_site.or(site);
}
pub fn add_bound(&mut self, var: TypeVar, bound: Bound) {
let bounds = &mut self.vars[var.index()].bounds;
match bounds.iter_mut().find(|known| **known == bound) {
Some(known) => known.merge_origins(bound),
None => bounds.push(bound),
}
}
pub fn bind(&mut self, var: TypeVar, value: Type) {
self.vars[var.index()].binding = Some(value);
}
pub fn inside_operand(&self) -> bool {
self.inside_operand
}
pub fn set_inside_operand(&mut self, inside_operand: bool) {
self.inside_operand = inside_operand;
}
pub fn defer(&mut self, obligation: Obligation) {
self.obligations.push(obligation);
}
pub fn take_obligations(&mut self) -> Vec<Obligation> {
std::mem::take(&mut self.obligations)
}
pub fn resolve(&self, value: &Type) -> Type {
match self.shallow(value) {
Type::Named(binding, arguments) => Type::Named(binding, self.resolve_all(&arguments)),
Type::Function(kind, parameters, result) => Type::Function(
kind,
self.resolve_all(¶meters),
Box::new(self.resolve(&result)),
),
Type::Ref(lifetime, inner) => Type::Ref(lifetime, Box::new(self.resolve(&inner))),
Type::Mut(lifetime, inner) => Type::Mut(lifetime, Box::new(self.resolve(&inner))),
other => other,
}
}
fn resolve_all(&self, values: &[Type]) -> Vec<Type> {
values.iter().map(|value| self.resolve(value)).collect()
}
pub fn shallow(&self, value: &Type) -> Type {
let Type::Var(start) = value else {
return value.clone();
};
let last = self.last_var(*start);
self.shorten_chain(*start, last);
match &self.vars[last.index()].binding {
Some(bound) => bound.clone(),
None => Type::Var(last),
}
}
fn last_var(&self, start: TypeVar) -> TypeVar {
let mut current = start;
while let Some(next) = self.next_var(current) {
current = next;
}
current
}
fn shorten_chain(&self, start: TypeVar, last: TypeVar) {
let mut current = start;
while let Some(next) = self.next_var(current).filter(|_| current != last) {
self.vars[current.index()].shortcut.set(Some(last));
current = next;
}
}
fn next_var(&self, var: TypeVar) -> Option<TypeVar> {
let state = &self.vars[var.index()];
state.shortcut.get().or(match &state.binding {
Some(Type::Var(next)) => Some(*next),
_ => None,
})
}
pub fn unify(&mut self, expected: &Type, found: &Type, span: Span) -> Result<(), UnifyError> {
match (self.shallow(expected), self.shallow(found)) {
(Type::Var(left), Type::Var(right)) => {
self.unify_vars(left, right);
Ok(())
}
(Type::Var(var), other) | (other, Type::Var(var)) => {
self.bind_checked(var, other, span)
}
(
Type::Function(_, left_parameters, left_result),
Type::Function(_, right_parameters, right_result),
) => self.unify_functions(
&left_parameters,
&left_result,
&right_parameters,
&right_result,
span,
),
(Type::Ref(_, left), Type::Ref(_, right))
| (Type::Mut(_, left), Type::Mut(_, right)) => self.unify(&left, &right, span),
(left, right) => self.unify_rigid(&left, &right, span),
}
}
fn unify_rigid(&mut self, left: &Type, right: &Type, span: Span) -> Result<(), UnifyError> {
if let Some((left, right)) = self.normalized_pair(left, right) {
return self.unify(&left, &right, span);
}
match (left, right) {
(Type::Named(projection, _), Type::Named(other, _))
if projection == other && self.awaits_shared_owner(*projection, left, right) =>
{
self.defer_once(left, right, span);
Ok(())
}
(Type::Named(left, left_arguments), Type::Named(right, right_arguments))
if left == right =>
{
self.unify_pairwise(left_arguments, right_arguments, span)
}
_ if left == right => Ok(()),
_ => self.defer_projection(left, right, span),
}
}
fn normalized_pair(&self, left: &Type, right: &Type) -> Option<(Type, Type)> {
if self.is_open_pair(left, right) && self.shares_rigid_head(left, right) {
return None;
}
match (self.normalized(left), self.normalized(right)) {
(None, None) => None,
(normalized_left, normalized_right) => Some((
normalized_left.unwrap_or_else(|| left.clone()),
normalized_right.unwrap_or_else(|| right.clone()),
)),
}
}
fn awaits_shared_owner(&self, projection: BindingId, left: &Type, right: &Type) -> bool {
self.projections.is_projection(projection) && self.is_open_pair(left, right)
}
fn is_open_pair(&self, left: &Type, right: &Type) -> bool {
self.resolve(left).has_vars() || self.resolve(right).has_vars()
}
fn shares_rigid_head(&self, left: &Type, right: &Type) -> bool {
let projected =
matches!(left, Type::Named(binding, _) if self.projections.is_projection(*binding));
left.shares_head_with(right) && !projected
}
fn normalized(&self, value: &Type) -> Option<Type> {
self.projections.normalize(&self.resolve(value))
}
fn defer_projection(
&mut self,
left: &Type,
right: &Type,
span: Span,
) -> Result<(), UnifyError> {
let (projected, other) = if self.awaits_owner(left, right) {
(left, right)
} else if self.awaits_owner(right, left) {
(right, left)
} else {
return Err(UnifyError::Mismatch);
};
self.defer_once(projected, other, span);
Ok(())
}
fn defer_once(&mut self, projected: &Type, other: &Type, span: Span) {
let known = self.deferred.iter().any(|deferred| {
deferred.projected == *projected && deferred.other == *other && deferred.span == span
});
if !known {
self.deferred.push(DeferredEquality {
projected: projected.clone(),
other: other.clone(),
span,
});
}
}
fn awaits_owner(&self, projected: &Type, other: &Type) -> bool {
self.projections
.awaits_owner(&self.resolve(projected), &self.resolve(other))
}
fn unify_pairwise(
&mut self,
left: &[Type],
right: &[Type],
span: Span,
) -> Result<(), UnifyError> {
if left.len() != right.len() {
return Err(UnifyError::Mismatch);
}
for (left_part, right_part) in left.iter().zip(right) {
self.unify(left_part, right_part, span)?;
}
Ok(())
}
fn unify_vars(&mut self, left: TypeVar, right: TypeVar) {
if left != right {
self.merge_vars(left, right);
}
}
fn unify_functions(
&mut self,
left: &[Type],
left_result: &Type,
right: &[Type],
right_result: &Type,
span: Span,
) -> Result<(), UnifyError> {
self.unify_pairwise(left, right, span)?;
self.unify(left_result, right_result, span)
}
fn fresh_with(&mut self, literal: Option<LiteralKind>, bounds: Vec<Bound>) -> Type {
let var = TypeVar::new(self.vars.len());
self.vars.push(VarState {
binding: None,
shortcut: Cell::new(None),
literal,
bounds,
bound_site: None,
});
Type::Var(var)
}
fn merge_vars(&mut self, from: TypeVar, into: TypeVar) {
let moved_bounds = std::mem::take(&mut self.vars[from.index()].bounds);
let literal = LiteralKind::merge(
self.vars[from.index()].literal,
self.vars[into.index()].literal,
);
self.vars[into.index()].literal = literal;
let bound_site = self.vars[from.index()].bound_site;
moved_bounds
.into_iter()
.for_each(|bound| self.add_bound_from(into, bound, bound_site));
self.bind(from, Type::Var(into));
}
fn bind_checked(&mut self, var: TypeVar, value: Type, span: Span) -> Result<(), UnifyError> {
if value.contains_var(var) {
return Err(UnifyError::Infinite);
}
if let Some(literal) = self.vars[var.index()].literal {
self.check_literal(literal, &value)?;
}
for bound in self.vars[var.index()].bounds.clone() {
self.obligations.push(Obligation {
bound,
implementor: value.clone(),
site: ObligationSite {
span,
bound_site: self.bound_site(var),
},
inside_operand: self.inside_operand,
});
}
self.bind(var, value);
Ok(())
}
fn check_literal(&self, literal: LiteralKind, value: &Type) -> Result<(), UnifyError> {
match value {
Type::Primitive(primitive) if literal.accepts(*primitive) => Ok(()),
Type::Generic(_) => Ok(()),
_ => Err(UnifyError::NotNumeric),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::{BindingArena, BindingInfo, BindingKind};
use crate::source::SourceId;
struct SiteFixture {
unifier: Unifier,
bound: Bound,
}
impl SiteFixture {
fn new() -> Self {
let mut bindings = BindingArena::default();
let binding = bindings.add(BindingInfo::new(
"trait".into(),
BindingKind::TypeVariable,
None,
));
Self {
unifier: Unifier::default(),
bound: Bound::plain(binding),
}
}
fn span(&self, start: usize) -> Span {
Span::new(SourceId::new(0), start..start + 1)
}
fn bounded_var(&mut self, site: Span) -> TypeVar {
let Type::Var(var) = self.unifier.fresh() else {
unreachable!("a fresh type is a variable");
};
self.unifier.set_bounds(var, vec![self.bound.clone()], site);
var
}
}
#[test]
fn variables_of_one_call_share_one_bound_site() {
let mut fixture = SiteFixture::new();
let call = fixture.span(0);
fixture.bounded_var(call);
let second = fixture.bounded_var(call);
assert_eq!(fixture.unifier.bound_sites, vec![call]);
assert_eq!(fixture.unifier.bound_site(second), Some(call));
}
#[test]
fn another_call_keeps_its_own_bound_site() {
let mut fixture = SiteFixture::new();
let (first_call, later_call) = (fixture.span(0), fixture.span(4));
let first = fixture.bounded_var(first_call);
let later = fixture.bounded_var(later_call);
assert_eq!(fixture.unifier.bound_site(first), Some(first_call));
assert_eq!(fixture.unifier.bound_site(later), Some(later_call));
}
}