Tag aware unification type inference

src/inference/mod.rs

@@ -16,4 +16,5 @@ pub enum InferError {
     ExpectedTypeWithTag,
     DoesNotFitTag(TypeTag, TaggedType),
     ConfilictingBind,
+    ConfilictingTags,
 }

src/inference/recursive.rs

@@ -34,7 +34,7 @@ fn infer_type_debrujin_int(
         }
         DeBrujinAst::Application(lhs, rhs) => {
             let left_type = infer_type_debrujin_int(gamma_free, gamma_bound.clone(), *lhs)?;
-            let Some((in_type, out_type)) = left_type.arrow() else {
+            let Some((in_type, out_type)) = left_type.split_arrow() else {
                 return Err(InferError::NotAFunction);
             };
             let Some(right_type) =

src/inference/test.rs

@@ -79,6 +79,29 @@ fn infer_add_nat_partial_rec() {
     );
 }
 
+#[test]
+fn subtype_constraints() {
+    let typ = TaggedType::Tagged(
+        TypeTag::Num,
+        "a".to_string(),
+        Box::new(Type::arrow("a", "a").into()),
+    );
+
+    let mut st_constraints = Rc::new(RefCell::new(MultiMap::new()));
+    let ctx = TypeVarCtx::new();
+
+    let res = typ.make_constraints(st_constraints.clone(), &ctx);
+
+    assert_eq!(
+        st_constraints
+            .borrow()
+            .get(&Type::TypeVariable("?typ_0".to_string())),
+        vec![&TypeTag::Num]
+    );
+
+    assert_eq!(res, Type::arrow("?typ_0", "?typ_0"));
+}
+
 #[test]
 fn infer_add_nat_uni() {
     let ast: DeBrujinAst = Ast::Application(
@@ -93,13 +116,17 @@ fn infer_add_nat_uni() {
     let mut gamma = HashMap::new();
     gamma.insert(
         "add".to_string(),
-        TaggedType::Concrete(Type::arrow(
+        TaggedType::Tagged(
-            PrimitiveType::Nat,
+            TypeTag::Num,
-            Type::arrow(PrimitiveType::Nat, PrimitiveType::Nat),
+            "a".to_string(),
-        )),
+            Box::new(TaggedType::Concrete(Type::arrow(
+                "a",
+                Type::arrow("a", "a"),
+            ))),
+        ),
     );
 
-    let infered = infer_type_rec(&gamma, ast).unwrap();
+    let infered = infer_type_uni(&gamma, ast).unwrap();
     assert_eq!(
         infered,
         TaggedType::Concrete(Type::Primitive(PrimitiveType::Nat))

src/inference/unification.rs

@@ -1,4 +1,10 @@
-use std::{cell::RefCell, clone, collections::HashMap, rc::Rc};
+use std::{
+    cell::RefCell,
+    clone,
+    collections::HashMap,
+    fmt::{Debug, DebugMap},
+    rc::Rc,
+};
 
 use crate::{
     Ast, DeBrujinAst,
@@ -10,7 +16,7 @@ use crate::{
 
 use super::InferError;
 
-type TypeVar = TaggedType;
+type TypeVar = Type;
 
 pub struct TypeVarCtx {
     counter: RefCell<usize>,
@@ -23,45 +29,59 @@ impl TypeVarCtx {
         }
     }
 
-    pub fn get_var(&self) -> TaggedType {
+    pub fn get_var(&self) -> Type {
         let mut num = self.counter.borrow_mut();
         let res = format!("?typ_{num}");
         *num += 1;
-        Type::TypeVariable(res).into()
+        Type::TypeVariable(res)
     }
 }
 
 #[derive(Debug, Clone, PartialEq)]
 pub enum TypeVarAst {
-    Abstraction(TypeVar, Ident, TaggedType, Box<TypeVarAst>), // \0:1.2
+    Abstraction(TypeVar, Ident, Type, Box<TypeVarAst>), // \0:1.2
     Application(TypeVar, Box<TypeVarAst>, Box<TypeVarAst>), // 0 1
     FreeVariable(TypeVar, Ident),                       // x
     BoundVariable(TypeVar, usize),                      // 1
     Constant(TypeVar, Constant),                        // true | false | n
 }
 
-pub type Constraints = MultiMap<TypeVar, TaggedType>;
+pub type Constraints<T> = MultiMap<TypeVar, T>;
 
 pub(super) fn step_1(
     ast: DeBrujinAst,
     gamma_free: &HashMap<Ident, TaggedType>,
-    mut gamma_bound: Rc<VecMap<usize, TaggedType>>,
+    mut gamma_bound: Rc<VecMap<usize, Type>>,
-    constraints: Rc<RefCell<Constraints>>,
+    eq_constraints: Rc<RefCell<Constraints<Type>>>,
+    st_constraints: Rc<RefCell<Constraints<TypeTag>>>,
     ctx: &TypeVarCtx,
 ) -> Result<(TypeVarAst, TypeVar), InferError> {
     match ast {
         DeBrujinAst::Abstraction(i, Some(typ), ast) => {
-            let var = ctx.get_var();
+            // let var = ctx.get_var();
+
+            let typ = typ.make_constraints(st_constraints.clone(), ctx);
 
             let gamma_ref = Rc::make_mut(&mut gamma_bound);
             gamma_ref.map_keys(|i| *i += 1);
             gamma_ref.insert(1, typ.clone());
 
-            let (ast, rhs_var) = step_1(*ast, gamma_free, gamma_bound, constraints.clone(), ctx)?;
+            let (ast, rhs_var) = step_1(
-            // RefCell::borrow_mut(&constraints).insert(var.clone(), typ.clone().make_arrow(rhs_var));
+                *ast,
+                gamma_free,
+                gamma_bound,
+                eq_constraints.clone(),
+                st_constraints.clone(),
+                ctx,
+            )?;
             Ok((
-                TypeVarAst::Abstraction(var.clone(), i, typ, Box::new(ast)),
+                TypeVarAst::Abstraction(
-                var,
+                    Type::arrow(typ.clone(), rhs_var.clone()),
+                    i,
+                    typ.clone(),
+                    Box::new(ast),
+                ),
+                Type::arrow(typ, rhs_var.clone()),
             ))
         }
         DeBrujinAst::Abstraction(i, None, ast) => {
@@ -72,11 +92,22 @@ pub(super) fn step_1(
             gamma_ref.map_keys(|i| *i += 1);
             gamma_ref.insert(1, typ.clone());
 
-            let (ast, rhs_var) = step_1(*ast, gamma_free, gamma_bound, constraints.clone(), ctx)?;
+            let (ast, rhs_var) = step_1(
-            // RefCell::borrow_mut(&constraints).insert(var.clone(), typ.clone().make_arrow(rhs_var));
+                *ast,
+                gamma_free,
+                gamma_bound,
+                eq_constraints.clone(),
+                st_constraints.clone(),
+                ctx,
+            )?;
             Ok((
-                TypeVarAst::Abstraction(var.clone(), i, typ, Box::new(ast)),
+                TypeVarAst::Abstraction(
-                var,
+                    Type::arrow(typ.clone(), rhs_var.clone()),
+                    i,
+                    typ.clone(),
+                    Box::new(ast),
+                ),
+                Type::arrow(typ, rhs_var.clone()),
             ))
         }
         DeBrujinAst::Application(lhs, rhs) => {
@@ -85,17 +116,19 @@ pub(super) fn step_1(
                 *lhs,
                 gamma_free,
                 gamma_bound.clone(),
-                constraints.clone(),
+                eq_constraints.clone(),
+                st_constraints.clone(),
                 ctx,
             )?;
             let (rhs, rhs_var) = step_1(
                 *rhs,
                 gamma_free,
                 gamma_bound.clone(),
-                constraints.clone(),
+                eq_constraints.clone(),
+                st_constraints.clone(),
                 ctx,
             )?;
-            RefCell::borrow_mut(&constraints).insert(lhs_var, rhs_var.make_arrow(var.clone()));
+            RefCell::borrow_mut(&eq_constraints).insert(lhs_var, Type::arrow(rhs_var, var.clone()));
             Ok((
                 TypeVarAst::Application(var.clone(), Box::new(lhs), Box::new(rhs)),
                 var,
@@ -108,7 +141,9 @@ pub(super) fn step_1(
                 .cloned()
                 .ok_or(InferError::NotInContext)?;
 
-            RefCell::borrow_mut(&constraints).insert(var.clone(), typ);
+            let typ = typ.make_constraints(st_constraints.clone(), ctx);
+
+            RefCell::borrow_mut(&eq_constraints).insert(var.clone(), typ);
 
             Ok((TypeVarAst::FreeVariable(var.clone(), v), var))
         }
@@ -119,7 +154,7 @@ pub(super) fn step_1(
                 .cloned()
                 .ok_or(InferError::NotInContext)?;
 
-            RefCell::borrow_mut(&constraints).insert(var.clone(), typ);
+            RefCell::borrow_mut(&eq_constraints).insert(var.clone(), typ);
 
             Ok((TypeVarAst::BoundVariable(var.clone(), i), var))
         }
@@ -132,48 +167,40 @@ pub(super) fn step_1(
             }
             .into();
 
-            RefCell::borrow_mut(&constraints).insert(var.clone(), typ);
+            RefCell::borrow_mut(&eq_constraints).insert(var.clone(), typ);
 
             Ok((TypeVarAst::Constant(var.clone(), constant), var))
         }
     }
 }
 
-pub(super) fn step_2(
+pub(super) fn step_2(mut eq_constraints: Constraints<Type>) -> Result<Option<SubstFn>, InferError> {
-    mut constraints: Constraints,
+    if let Some((s, t)) = eq_constraints.pop() {
-) -> Result<Option<Box<dyn FnOnce(TypeVarAst) -> Result<TypeVarAst, InferError>>>, InferError> {
-    if let Some((s, t)) = constraints.pop() {
         if s == t {
-            step_2(constraints)
+            step_2(eq_constraints)
         } else if s.type_var().is_some_and(|x| !t.name_used(&x)) {
             let Some(x) = s.type_var() else {
                 unreachable!()
             };
-            constraints.try_map(|(k, v)| {
+            eq_constraints
-                Ok((
+                .try_map(|(k, v)| Ok((k.subst_typevar(&x, &t)?, v.subst_typevar(&x, &t)?)))?;
-                    k.subst_typevar(&x, 0, t.clone())?,
+            let subst = step_2(eq_constraints)?;
-                    v.subst_typevar(&x, 0, t.clone())?,
+            Ok(Some(SubstFn::new(x, t, subst)))
-                ))
-            })?;
-            let subst = step_2(constraints)?;
-            Ok(Some(subst_comp(x, t, subst)))
         } else if t.type_var().is_some_and(|x| !s.name_used(&x)) {
             let Some(x) = t.type_var() else {
                 unreachable!()
             };
-            constraints.try_map(|(k, v)| {
+            eq_constraints
-                Ok((
+                .try_map(|(k, v)| Ok((k.subst_typevar(&x, &s)?, v.subst_typevar(&x, &s)?)))?;
-                    k.subst_typevar(&x, 0, s.clone())?,
+
-                    v.subst_typevar(&x, 0, s.clone())?,
+            let subst = step_2(eq_constraints)?;
-                ))
+            Ok(Some(SubstFn::new(x, s, subst)))
-            })?;
+        } else if let (Some((s_lhs, s_rhs)), Some((t_lhs, t_rhs))) =
-
+            (s.split_arrow(), t.split_arrow())
-            let subst = step_2(constraints)?;
+        {
-            Ok(Some(subst_comp(x, t, subst)))
+            eq_constraints.insert(s_lhs, t_lhs);
-        } else if let (Some((s_lhs, s_rhs)), Some((t_lhs, t_rhs))) = (s.arrow(), t.arrow()) {
+            eq_constraints.insert(s_rhs, t_rhs);
-            constraints.insert(s_lhs, t_lhs);
+            step_2(eq_constraints)
-            constraints.insert(s_rhs, t_rhs);
-            step_2(constraints)
         } else {
             panic!()
         }
@@ -182,50 +209,89 @@ pub(super) fn step_2(
     }
 }
 
-impl TypeVarAst {
+pub trait Subst {
-    pub fn subst(self, var: &str, subst: TaggedType) -> Result<TypeVarAst, InferError> {
+    fn subst(self, var: &str, subst: &Type) -> Result<Self, InferError>
+    where
+        Self: Sized;
+}
+
+impl Subst for TypeVarAst {
+    fn subst(self, var: &str, subst: &Type) -> Result<TypeVarAst, InferError> {
         match self {
-            TypeVarAst::Abstraction(tagged_type1, ident, tagged_type2, ast) => {
+            TypeVarAst::Abstraction(type1, ident, type2, ast) => Ok(TypeVarAst::Abstraction(
-                Ok(TypeVarAst::Abstraction(
+                type1.subst_typevar(var, subst)?,
-                    tagged_type1.subst_typevar(var, 0, subst.clone())?,
                 ident,
-                    tagged_type2.subst_typevar(var, 0, subst.clone())?,
+                type2.subst_typevar(var, subst)?,
                 Box::new(ast.subst(var, subst)?),
-                ))
-            }
-            TypeVarAst::Application(tagged_type, lhs, rhs) => Ok(TypeVarAst::Application(
-                tagged_type.subst_typevar(var, 0, subst.clone())?,
-                Box::new(lhs.subst(var, subst.clone())?),
-                Box::new(rhs.subst(var, subst)?),
             )),
-            TypeVarAst::FreeVariable(tagged_type, x) => Ok(TypeVarAst::FreeVariable(
+            TypeVarAst::Application(typ, lhs, rhs) => Ok(TypeVarAst::Application(
-                tagged_type.subst_typevar(var, 0, subst)?,
+                typ.subst_typevar(var, subst)?,
-                x,
+                Box::new(lhs.subst(var, subst)?),
+                Box::new(rhs.subst(var, subst)?),
             )),
-            TypeVarAst::BoundVariable(tagged_type, i) => Ok(TypeVarAst::BoundVariable(
+            TypeVarAst::FreeVariable(typ, x) => {
-                tagged_type.subst_typevar(var, 0, subst)?,
+                Ok(TypeVarAst::FreeVariable(typ.subst_typevar(var, &subst)?, x))
+            }
+            TypeVarAst::BoundVariable(typ, i) => Ok(TypeVarAst::BoundVariable(
+                typ.subst_typevar(var, &subst)?,
                 i,
             )),
-            TypeVarAst::Constant(tagged_type, constant) => Ok(TypeVarAst::Constant(
+            TypeVarAst::Constant(typ, constant) => Ok(TypeVarAst::Constant(
-                tagged_type.subst_typevar(var, 0, subst)?,
+                typ.subst_typevar(var, &subst)?,
                 constant,
             )),
         }
     }
 }
 
-pub fn subst_comp<'a, F>(
+impl Subst for Constraints<TypeTag> {
-    var: String,
+    fn subst(mut self, var: &str, subst: &Type) -> Result<Constraints<TypeTag>, InferError> {
-    subst: TaggedType,
+        self.into_iter()
-    then: Option<F>,
+            .map::<Result<(Type, TypeTag), InferError>, _>(|(k, v)| {
-) -> Box<dyn FnOnce(TypeVarAst) -> Result<TypeVarAst, InferError> + 'a>
+                Ok((k.subst_typevar(var, &subst)?, v))
-where
-    F: FnOnce(TypeVarAst) -> Result<TypeVarAst, InferError> + 'a,
-{
-    Box::new(move |ast: TypeVarAst| -> Result<TypeVarAst, InferError> {
-        let ast = ast.subst(&var, subst)?;
-        if let Some(f) = then { f(ast) } else { Ok(ast) }
             })
+            .collect()
+    }
+}
+
+pub struct SubstFn {
+    var: String,
+    subst: Type,
+    then: Option<Box<Self>>,
+}
+
+impl Debug for SubstFn {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut f = f.debug_map();
+        self.debug_map(&mut f);
+        f.finish()
+    }
+}
+
+impl SubstFn {
+    pub fn new(var: String, subst: Type, then: Option<Self>) -> Self {
+        Self {
+            var,
+            subst,
+            then: then.map(Box::new),
+        }
+    }
+
+    pub fn apply<T: Subst>(&self, target: T) -> Result<T, InferError> {
+        let target = target.subst(&self.var, &self.subst);
+        if let Some(s) = &self.then {
+            s.apply(target?)
+        } else {
+            target
+        }
+    }
+
+    fn debug_map(&self, f: &mut DebugMap) {
+        f.entry(&self.var, &self.subst);
+        if let Some(s) = &self.then {
+            s.debug_map(f);
+        }
+    }
 }
 
 pub fn infer_type(
@@ -233,37 +299,91 @@ pub fn infer_type(
     ast: DeBrujinAst,
 ) -> Result<TaggedType, InferError> {
     let gamma_bound = Rc::new(VecMap::new());
-    let constraints = Rc::new(RefCell::new(MultiMap::new()));
+    let eq_constraints = Rc::new(RefCell::new(MultiMap::new()));
+    let st_constraints = Rc::new(RefCell::new(MultiMap::new()));
     let ctx = TypeVarCtx::new();
 
-    let (ast, _) = step_1(ast, gamma, gamma_bound, constraints.clone(), &ctx)?;
+    let (ast, _) = step_1(
-    constraints.clone();
+        ast,
-    let res = step_2(constraints.take())?.unwrap();
+        gamma,
+        gamma_bound,
+        eq_constraints.clone(),
+        st_constraints.clone(),
+        &ctx,
+    )?;
+    let res = step_2(eq_constraints.take())?.unwrap();
 
-    let ast = res(ast)?;
+    let ast = res.apply(ast)?;
 
-    fn get_type(ast: TypeVarAst) -> TaggedType {
+    fn get_type(ast: TypeVarAst) -> Type {
         match ast {
-            TypeVarAst::Abstraction(_, _, typ, ast) => typ.make_arrow(get_type(*ast)),
+            TypeVarAst::Abstraction(_, _, typ, ast) => Type::arrow(typ, get_type(*ast)),
-            TypeVarAst::Application(tagged_type, _, _) => tagged_type,
+            TypeVarAst::Application(typ, _, _) => typ,
-            TypeVarAst::FreeVariable(tagged_type, _) => tagged_type,
+            TypeVarAst::FreeVariable(typ, _) => typ,
-            TypeVarAst::BoundVariable(tagged_type, _) => tagged_type,
+            TypeVarAst::BoundVariable(typ, _) => typ,
-            TypeVarAst::Constant(tagged_type, constant) => tagged_type,
+            TypeVarAst::Constant(typ, constant) => typ,
         }
     }
 
     let typ = get_type(ast);
 
-    let mut typ = typ;
+    let st_constraints = res.apply(st_constraints.take())?;
+    let st_constraints = st_constraints
+        .into_iter()
+        .map(|(typ, tag)| {
+            if typ.has_tag(&tag) {
+                Ok((typ, tag))
+            } else {
+                Err(InferError::DoesNotFitTag(tag, typ.into()))
+            }
+        })
+        .collect::<Result<MultiMap<_, _>, _>>()?;
+
+    let mut typ: TaggedType = typ.into();
 
     for free_var in typ.free_vars() {
+        let tags = st_constraints.get(&Type::TypeVariable(free_var.clone()));
+        if tags.is_empty() {
             typ = TaggedType::Tagged(TypeTag::Any, free_var, Box::new(typ));
+        } else {
+            let tag = tags
+                .into_iter()
+                .fold(Some(TypeTag::Any), |acc, t| {
+                    acc.map(|acc| acc.tightens(t.clone())).flatten()
+                })
+                .ok_or(InferError::ConfilictingTags)?;
+            typ = TaggedType::Tagged(tag, free_var, Box::new(typ));
+        }
     }
 
-    Ok(typ.normalise())
+    Ok(typ)
 }
 
 impl TaggedType {
+    pub fn make_constraints(
+        self,
+        st_constraints: Rc<RefCell<Constraints<TypeTag>>>,
+        ctx: &TypeVarCtx,
+    ) -> Type {
+        match self {
+            TaggedType::Tagged(type_tag, ident, mut tagged_type) => {
+                let var = ctx.get_var();
+                let name = var.type_var().unwrap();
+
+                tagged_type.map_name(|n| {
+                    if *n == ident {
+                        *n = name.clone()
+                    }
+                });
+
+                st_constraints.borrow_mut().insert(var, type_tag);
+
+                tagged_type.make_constraints(st_constraints, ctx)
+            }
+            TaggedType::Concrete(c) => c,
+        }
+    }
+
     pub fn subst_typevar(
         self,
         var: &str,

src/multi_map.rs

@@ -1,4 +1,4 @@
-use std::{fmt::Debug, mem};
+use std::{fmt::Debug, mem, vec};
 
 pub struct MultiMap<K: Eq, V> {
     map: Vec<(K, V)>,
@@ -61,6 +61,24 @@ impl<K: Eq, V> MultiMap<K, V> {
     }
 }
 
+impl<K: Eq, V> IntoIterator for MultiMap<K, V> {
+    type Item = (K, V);
+
+    type IntoIter = vec::IntoIter<(K, V)>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.map.into_iter()
+    }
+}
+
+impl<K: Eq, V> FromIterator<(K, V)> for MultiMap<K, V> {
+    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
+        Self {
+            map: iter.into_iter().collect(),
+        }
+    }
+}
+
 impl<K: Eq, V> Default for MultiMap<K, V> {
     fn default() -> Self {
         Self {

src/parse/test.rs

@@ -251,4 +251,4 @@ fn parse_application() {
     );
 }
 
-// (\x:Any a => a -> a.\y:Num a => a. + 1 (x y)) (\x:Any a => a.x) 2
+// (\x:Any a => a -> a.\y:Num a => a. add 1 (x y)) (\x:Any a => a.x) 2

src/types/mod.rs

@@ -42,7 +42,7 @@ impl From<Type> for TaggedType {
 
 impl TypeTag {
     // If one tag tightens the other return that tag, otherwise none,
-    fn tightens(self, other: Self) -> Option<Self> {
+    pub fn tightens(self, other: Self) -> Option<Self> {
         match (self, other) {
             (TypeTag::Num, TypeTag::Num) => Some(Self::Num),
             (TypeTag::Num, TypeTag::Any) => Some(Self::Num),
@@ -87,10 +87,10 @@ impl TaggedType {
         }
     }
 
-    pub fn arrow(self) -> Option<(TaggedType, TaggedType)> {
+    pub fn split_arrow(self) -> Option<(TaggedType, TaggedType)> {
         match self {
             TaggedType::Tagged(type_tag, ident, tagged_type) => {
-                let (lhs, rhs) = tagged_type.arrow()?;
+                let (lhs, rhs) = tagged_type.split_arrow()?;
                 Some((
                     TaggedType::Tagged(type_tag.clone(), ident.clone(), Box::new(lhs))
                         .clear_unused_names(),
@@ -237,6 +237,13 @@ impl Type {
         Self::Arrow(Box::new(t1.into()), Box::new(t2.into()))
     }
 
+    pub fn split_arrow(self) -> Option<(Type, Type)> {
+        match self {
+            Type::Arrow(lhs, rhs) => Some((*lhs, *rhs)),
+            _ => None,
+        }
+    }
+
     pub fn is_concrete(&self) -> bool {
         match self {
             Type::TypeVariable(_) => false,
@@ -247,7 +254,7 @@ impl Type {
 
     pub fn has_tag(&self, tag: &TypeTag) -> bool {
         match self {
-            Type::TypeVariable(_) => false,
+            Type::TypeVariable(_) => true,
             Type::Primitive(primitive_type) => match (primitive_type, tag) {
                 (_, TypeTag::Any) => true,
                 (PrimitiveType::Nat, TypeTag::Num) => true,