seqc/
unification.rs

1//! Type unification for Seq
2//!
3//! Implements Hindley-Milner style unification with support for:
4//! - Type variables (T, U, V)
5//! - Row variables (..a, ..rest)
6//! - Concrete types (Int, Bool, String)
7
8use crate::types::{StackType, Type};
9use std::collections::HashMap;
10
11/// Substitutions for type variables
12pub type TypeSubst = HashMap<String, Type>;
13
14/// Substitutions for row variables (stack type variables)
15pub type RowSubst = HashMap<String, StackType>;
16
17/// Combined substitution environment
18#[derive(Debug, Clone, PartialEq)]
19pub struct Subst {
20    pub types: TypeSubst,
21    pub rows: RowSubst,
22}
23
24impl Subst {
25    /// Create an empty substitution
26    pub fn empty() -> Self {
27        Subst {
28            types: HashMap::new(),
29            rows: HashMap::new(),
30        }
31    }
32
33    /// Apply substitutions to a Type
34    pub fn apply_type(&self, ty: &Type) -> Type {
35        match ty {
36            Type::Var(name) => self.types.get(name).cloned().unwrap_or(ty.clone()),
37            _ => ty.clone(),
38        }
39    }
40
41    /// Apply substitutions to a StackType
42    pub fn apply_stack(&self, stack: &StackType) -> StackType {
43        match stack {
44            StackType::Empty => StackType::Empty,
45            StackType::Cons { rest, top } => {
46                let new_rest = self.apply_stack(rest);
47                let new_top = self.apply_type(top);
48                StackType::Cons {
49                    rest: Box::new(new_rest),
50                    top: new_top,
51                }
52            }
53            StackType::RowVar(name) => self.rows.get(name).cloned().unwrap_or(stack.clone()),
54        }
55    }
56
57    /// Compose two substitutions (apply other after self)
58    /// Result: (other ∘ self) where self is applied first, then other
59    pub fn compose(&self, other: &Subst) -> Subst {
60        let mut types = HashMap::new();
61        let mut rows = HashMap::new();
62
63        // Apply other to all of self's type substitutions
64        for (k, v) in &self.types {
65            types.insert(k.clone(), other.apply_type(v));
66        }
67
68        // Add other's type substitutions (applying self to other's values)
69        for (k, v) in &other.types {
70            let v_subst = self.apply_type(v);
71            types.insert(k.clone(), v_subst);
72        }
73
74        // Apply other to all of self's row substitutions
75        for (k, v) in &self.rows {
76            rows.insert(k.clone(), other.apply_stack(v));
77        }
78
79        // Add other's row substitutions (applying self to other's values)
80        for (k, v) in &other.rows {
81            let v_subst = self.apply_stack(v);
82            rows.insert(k.clone(), v_subst);
83        }
84
85        Subst { types, rows }
86    }
87}
88
89/// Check if a type variable occurs in a type (for occurs check)
90///
91/// Prevents infinite types like: T = List<T>
92///
93/// NOTE: Currently we only have simple types (Int, String, Bool).
94/// When parametric types are added (e.g., List<T>, Option<T>), this function
95/// must be extended to recursively check type arguments:
96///
97/// ```ignore
98/// Type::Named { name: _, args } => {
99///     args.iter().any(|arg| occurs_in_type(var, arg))
100/// }
101/// ```
102fn occurs_in_type(var: &str, ty: &Type) -> bool {
103    match ty {
104        Type::Var(name) => name == var,
105        Type::Int | Type::Float | Type::Bool | Type::String => false,
106        Type::Quotation(effect) => {
107            // Check if var occurs in quotation's input or output stack types
108            occurs_in_stack(var, &effect.inputs) || occurs_in_stack(var, &effect.outputs)
109        }
110        Type::Closure { effect, captures } => {
111            // Check if var occurs in closure's effect or any captured types
112            occurs_in_stack(var, &effect.inputs)
113                || occurs_in_stack(var, &effect.outputs)
114                || captures.iter().any(|t| occurs_in_type(var, t))
115        }
116    }
117}
118
119/// Check if a row variable occurs in a stack type (for occurs check)
120fn occurs_in_stack(var: &str, stack: &StackType) -> bool {
121    match stack {
122        StackType::Empty => false,
123        StackType::RowVar(name) => name == var,
124        StackType::Cons { rest, top: _ } => {
125            // Row variables only occur in stack positions, not in type positions
126            // So we only need to check the rest of the stack
127            occurs_in_stack(var, rest)
128        }
129    }
130}
131
132/// Unify two types, returning a substitution or an error
133pub fn unify_types(t1: &Type, t2: &Type) -> Result<Subst, String> {
134    match (t1, t2) {
135        // Same concrete types unify
136        (Type::Int, Type::Int)
137        | (Type::Float, Type::Float)
138        | (Type::Bool, Type::Bool)
139        | (Type::String, Type::String) => Ok(Subst::empty()),
140
141        // Type variable unifies with anything (with occurs check)
142        (Type::Var(name), ty) | (ty, Type::Var(name)) => {
143            // If unifying a variable with itself, no substitution needed
144            if matches!(ty, Type::Var(ty_name) if ty_name == name) {
145                return Ok(Subst::empty());
146            }
147
148            // Occurs check: prevent infinite types
149            if occurs_in_type(name, ty) {
150                return Err(format!(
151                    "Occurs check failed: cannot unify {:?} with {:?} (would create infinite type)",
152                    Type::Var(name.clone()),
153                    ty
154                ));
155            }
156
157            let mut subst = Subst::empty();
158            subst.types.insert(name.clone(), ty.clone());
159            Ok(subst)
160        }
161
162        // Quotation types unify if their effects unify
163        (Type::Quotation(effect1), Type::Quotation(effect2)) => {
164            // Unify inputs
165            let s_in = unify_stacks(&effect1.inputs, &effect2.inputs)?;
166
167            // Apply substitution to outputs and unify
168            let out1 = s_in.apply_stack(&effect1.outputs);
169            let out2 = s_in.apply_stack(&effect2.outputs);
170            let s_out = unify_stacks(&out1, &out2)?;
171
172            // Compose substitutions
173            Ok(s_in.compose(&s_out))
174        }
175
176        // Closure types unify if their effects unify (ignoring captures)
177        // Captures are an implementation detail determined by the type checker,
178        // not part of the user-visible type
179        (
180            Type::Closure {
181                effect: effect1, ..
182            },
183            Type::Closure {
184                effect: effect2, ..
185            },
186        ) => {
187            // Unify inputs
188            let s_in = unify_stacks(&effect1.inputs, &effect2.inputs)?;
189
190            // Apply substitution to outputs and unify
191            let out1 = s_in.apply_stack(&effect1.outputs);
192            let out2 = s_in.apply_stack(&effect2.outputs);
193            let s_out = unify_stacks(&out1, &out2)?;
194
195            // Compose substitutions
196            Ok(s_in.compose(&s_out))
197        }
198
199        // Closure <: Quotation (subtyping)
200        // A Closure can be used where a Quotation is expected
201        // The runtime will dispatch appropriately
202        (Type::Quotation(quot_effect), Type::Closure { effect, .. })
203        | (Type::Closure { effect, .. }, Type::Quotation(quot_effect)) => {
204            // Unify the effects (ignoring captures - they're an implementation detail)
205            let s_in = unify_stacks(&quot_effect.inputs, &effect.inputs)?;
206
207            // Apply substitution to outputs and unify
208            let out1 = s_in.apply_stack(&quot_effect.outputs);
209            let out2 = s_in.apply_stack(&effect.outputs);
210            let s_out = unify_stacks(&out1, &out2)?;
211
212            // Compose substitutions
213            Ok(s_in.compose(&s_out))
214        }
215
216        // Different concrete types don't unify
217        _ => Err(format!(
218            "Type mismatch: cannot unify {:?} with {:?}",
219            t1, t2
220        )),
221    }
222}
223
224/// Unify two stack types, returning a substitution or an error
225pub fn unify_stacks(s1: &StackType, s2: &StackType) -> Result<Subst, String> {
226    match (s1, s2) {
227        // Empty stacks unify
228        (StackType::Empty, StackType::Empty) => Ok(Subst::empty()),
229
230        // Row variable unifies with any stack (with occurs check)
231        (StackType::RowVar(name), stack) | (stack, StackType::RowVar(name)) => {
232            // If unifying a row var with itself, no substitution needed
233            if matches!(stack, StackType::RowVar(stack_name) if stack_name == name) {
234                return Ok(Subst::empty());
235            }
236
237            // Occurs check: prevent infinite stack types
238            if occurs_in_stack(name, stack) {
239                return Err(format!(
240                    "Occurs check failed: cannot unify {:?} with {:?} (would create infinite stack type)",
241                    StackType::RowVar(name.clone()),
242                    stack
243                ));
244            }
245
246            let mut subst = Subst::empty();
247            subst.rows.insert(name.clone(), stack.clone());
248            Ok(subst)
249        }
250
251        // Cons cells unify if tops and rests unify
252        (
253            StackType::Cons {
254                rest: rest1,
255                top: top1,
256            },
257            StackType::Cons {
258                rest: rest2,
259                top: top2,
260            },
261        ) => {
262            // Unify the tops
263            let s_top = unify_types(top1, top2)?;
264
265            // Apply substitution to rests and unify
266            let rest1_subst = s_top.apply_stack(rest1);
267            let rest2_subst = s_top.apply_stack(rest2);
268            let s_rest = unify_stacks(&rest1_subst, &rest2_subst)?;
269
270            // Compose substitutions
271            Ok(s_top.compose(&s_rest))
272        }
273
274        // Empty doesn't unify with Cons
275        _ => Err(format!(
276            "Stack shape mismatch: cannot unify {:?} with {:?}",
277            s1, s2
278        )),
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285
286    #[test]
287    fn test_unify_concrete_types() {
288        assert!(unify_types(&Type::Int, &Type::Int).is_ok());
289        assert!(unify_types(&Type::Bool, &Type::Bool).is_ok());
290        assert!(unify_types(&Type::String, &Type::String).is_ok());
291
292        assert!(unify_types(&Type::Int, &Type::Bool).is_err());
293    }
294
295    #[test]
296    fn test_unify_type_variable() {
297        let subst = unify_types(&Type::Var("T".to_string()), &Type::Int).unwrap();
298        assert_eq!(subst.types.get("T"), Some(&Type::Int));
299
300        let subst = unify_types(&Type::Bool, &Type::Var("U".to_string())).unwrap();
301        assert_eq!(subst.types.get("U"), Some(&Type::Bool));
302    }
303
304    #[test]
305    fn test_unify_empty_stacks() {
306        assert!(unify_stacks(&StackType::Empty, &StackType::Empty).is_ok());
307    }
308
309    #[test]
310    fn test_unify_row_variable() {
311        let subst = unify_stacks(
312            &StackType::RowVar("a".to_string()),
313            &StackType::singleton(Type::Int),
314        )
315        .unwrap();
316
317        assert_eq!(subst.rows.get("a"), Some(&StackType::singleton(Type::Int)));
318    }
319
320    #[test]
321    fn test_unify_cons_stacks() {
322        // ( Int ) unifies with ( Int )
323        let s1 = StackType::singleton(Type::Int);
324        let s2 = StackType::singleton(Type::Int);
325
326        assert!(unify_stacks(&s1, &s2).is_ok());
327    }
328
329    #[test]
330    fn test_unify_cons_with_type_var() {
331        // ( T ) unifies with ( Int ), producing T := Int
332        let s1 = StackType::singleton(Type::Var("T".to_string()));
333        let s2 = StackType::singleton(Type::Int);
334
335        let subst = unify_stacks(&s1, &s2).unwrap();
336        assert_eq!(subst.types.get("T"), Some(&Type::Int));
337    }
338
339    #[test]
340    fn test_unify_row_poly_stack() {
341        // ( ..a Int ) unifies with ( Bool Int ), producing ..a := ( Bool )
342        let s1 = StackType::RowVar("a".to_string()).push(Type::Int);
343        let s2 = StackType::Empty.push(Type::Bool).push(Type::Int);
344
345        let subst = unify_stacks(&s1, &s2).unwrap();
346
347        assert_eq!(subst.rows.get("a"), Some(&StackType::singleton(Type::Bool)));
348    }
349
350    #[test]
351    fn test_unify_polymorphic_dup() {
352        // dup: ( ..a T -- ..a T T )
353        // Applied to: ( Int ) should work with ..a := Empty, T := Int
354
355        let input_actual = StackType::singleton(Type::Int);
356        let input_declared = StackType::RowVar("a".to_string()).push(Type::Var("T".to_string()));
357
358        let subst = unify_stacks(&input_declared, &input_actual).unwrap();
359
360        assert_eq!(subst.rows.get("a"), Some(&StackType::Empty));
361        assert_eq!(subst.types.get("T"), Some(&Type::Int));
362
363        // Apply substitution to output: ( ..a T T )
364        let output_declared = StackType::RowVar("a".to_string())
365            .push(Type::Var("T".to_string()))
366            .push(Type::Var("T".to_string()));
367
368        let output_actual = subst.apply_stack(&output_declared);
369
370        // Should be ( Int Int )
371        assert_eq!(
372            output_actual,
373            StackType::Empty.push(Type::Int).push(Type::Int)
374        );
375    }
376
377    #[test]
378    fn test_subst_compose() {
379        // s1: T := Int
380        let mut s1 = Subst::empty();
381        s1.types.insert("T".to_string(), Type::Int);
382
383        // s2: U := T
384        let mut s2 = Subst::empty();
385        s2.types.insert("U".to_string(), Type::Var("T".to_string()));
386
387        // Compose: should give U := Int, T := Int
388        let composed = s1.compose(&s2);
389
390        assert_eq!(composed.types.get("T"), Some(&Type::Int));
391        assert_eq!(composed.types.get("U"), Some(&Type::Int));
392    }
393
394    #[test]
395    fn test_occurs_check_type_var_with_itself() {
396        // Unifying T with T should succeed (no substitution needed)
397        let result = unify_types(&Type::Var("T".to_string()), &Type::Var("T".to_string()));
398        assert!(result.is_ok());
399        let subst = result.unwrap();
400        // Should be empty - no substitution needed when unifying var with itself
401        assert!(subst.types.is_empty());
402    }
403
404    #[test]
405    fn test_occurs_check_row_var_with_itself() {
406        // Unifying ..a with ..a should succeed (no substitution needed)
407        let result = unify_stacks(
408            &StackType::RowVar("a".to_string()),
409            &StackType::RowVar("a".to_string()),
410        );
411        assert!(result.is_ok());
412        let subst = result.unwrap();
413        // Should be empty - no substitution needed when unifying var with itself
414        assert!(subst.rows.is_empty());
415    }
416
417    #[test]
418    fn test_occurs_check_prevents_infinite_stack() {
419        // Attempting to unify ..a with (..a Int) should fail
420        // This would create an infinite type: ..a = (..a Int) = ((..a Int) Int) = ...
421        let row_var = StackType::RowVar("a".to_string());
422        let infinite_stack = StackType::RowVar("a".to_string()).push(Type::Int);
423
424        let result = unify_stacks(&row_var, &infinite_stack);
425        assert!(result.is_err());
426        let err = result.unwrap_err();
427        assert!(err.contains("Occurs check failed"));
428        assert!(err.contains("infinite"));
429    }
430
431    #[test]
432    fn test_occurs_check_allows_different_row_vars() {
433        // Unifying ..a with ..b should succeed (different variables)
434        let result = unify_stacks(
435            &StackType::RowVar("a".to_string()),
436            &StackType::RowVar("b".to_string()),
437        );
438        assert!(result.is_ok());
439        let subst = result.unwrap();
440        assert_eq!(
441            subst.rows.get("a"),
442            Some(&StackType::RowVar("b".to_string()))
443        );
444    }
445
446    #[test]
447    fn test_occurs_check_allows_concrete_stack() {
448        // Unifying ..a with (Int String) should succeed (no occurs)
449        let row_var = StackType::RowVar("a".to_string());
450        let concrete = StackType::Empty.push(Type::Int).push(Type::String);
451
452        let result = unify_stacks(&row_var, &concrete);
453        assert!(result.is_ok());
454        let subst = result.unwrap();
455        assert_eq!(subst.rows.get("a"), Some(&concrete));
456    }
457
458    #[test]
459    fn test_occurs_in_type() {
460        // T occurs in T
461        assert!(occurs_in_type("T", &Type::Var("T".to_string())));
462
463        // T does not occur in U
464        assert!(!occurs_in_type("T", &Type::Var("U".to_string())));
465
466        // T does not occur in Int
467        assert!(!occurs_in_type("T", &Type::Int));
468        assert!(!occurs_in_type("T", &Type::String));
469        assert!(!occurs_in_type("T", &Type::Bool));
470    }
471
472    #[test]
473    fn test_occurs_in_stack() {
474        // ..a occurs in ..a
475        assert!(occurs_in_stack("a", &StackType::RowVar("a".to_string())));
476
477        // ..a does not occur in ..b
478        assert!(!occurs_in_stack("a", &StackType::RowVar("b".to_string())));
479
480        // ..a does not occur in Empty
481        assert!(!occurs_in_stack("a", &StackType::Empty));
482
483        // ..a occurs in (..a Int)
484        let stack = StackType::RowVar("a".to_string()).push(Type::Int);
485        assert!(occurs_in_stack("a", &stack));
486
487        // ..a does not occur in (..b Int)
488        let stack = StackType::RowVar("b".to_string()).push(Type::Int);
489        assert!(!occurs_in_stack("a", &stack));
490
491        // ..a does not occur in (Int String)
492        let stack = StackType::Empty.push(Type::Int).push(Type::String);
493        assert!(!occurs_in_stack("a", &stack));
494    }
495}