gruel_air/inference/
types.rs

1//! Type variable infrastructure for Hindley-Milner type inference.
2//!
3//! This module provides the core type representations used during inference:
4//! - [`TypeVarId`] - Unique identifier for type variables
5//! - [`InferType`] - Type representation during inference (supports variables)
6//! - [`TypeVarAllocator`] - Allocates fresh type variables
7
8use crate::Type;
9
10/// Unique identifier for a type variable.
11///
12/// Type variables represent unknown types during constraint generation.
13/// They are resolved to concrete types during unification.
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
15pub struct TypeVarId(u32);
16
17impl TypeVarId {
18    /// Create a new type variable ID with the given index.
19    #[inline]
20    pub fn new(index: u32) -> Self {
21        TypeVarId(index)
22    }
23
24    /// Get the underlying index.
25    #[inline]
26    pub fn index(self) -> u32 {
27        self.0
28    }
29}
30
31impl std::fmt::Display for TypeVarId {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        write!(f, "?{}", self.0)
34    }
35}
36
37/// Internal type representation during inference.
38///
39/// This is separate from the final [`Type`] enum to support type variables
40/// and the special `IntLiteral` type. After unification, all `InferType`s
41/// are resolved to concrete `Type`s.
42#[derive(Debug, Clone, PartialEq, Eq)]
43pub enum InferType {
44    /// A concrete type (maps directly to [`Type`] enum).
45    Concrete(Type),
46
47    /// A type variable (unknown, to be solved).
48    Var(TypeVarId),
49
50    /// An integer literal type.
51    ///
52    /// Integer literals can unify with any integer type (i8-i64, u8-u64).
53    /// When unified with a concrete integer type, the literal takes that type.
54    /// If unconstrained at the end of inference, defaults to `i32`.
55    IntLiteral,
56
57    /// A floating-point literal type.
58    ///
59    /// Float literals can unify with any float type (f16, f32, f64).
60    /// If unconstrained at the end of inference, defaults to `f64`.
61    FloatLiteral,
62
63    /// An array type during inference.
64    ///
65    /// Unlike `Concrete(Type::new_array(id))`, this stores the element type as an
66    /// `InferType` so we can handle cases where the element type is still a
67    /// type variable. After unification, these are converted to `Type::Array`.
68    Array {
69        element: Box<InferType>,
70        length: u64,
71    },
72}
73
74impl InferType {
75    /// Create an `InferType` from a concrete `Type`.
76    #[inline]
77    pub fn concrete(ty: Type) -> Self {
78        InferType::Concrete(ty)
79    }
80
81    /// Create a type variable.
82    #[inline]
83    pub fn var(id: TypeVarId) -> Self {
84        InferType::Var(id)
85    }
86
87    /// Create an integer literal type.
88    #[inline]
89    pub fn int_literal() -> Self {
90        InferType::IntLiteral
91    }
92
93    /// Check if this is a concrete type.
94    pub fn is_concrete(&self) -> bool {
95        matches!(self, InferType::Concrete(_))
96    }
97
98    /// Check if this is a type variable.
99    pub fn is_var(&self) -> bool {
100        matches!(self, InferType::Var(_))
101    }
102
103    /// Check if this is an integer literal type.
104    pub fn is_int_literal(&self) -> bool {
105        matches!(self, InferType::IntLiteral)
106    }
107
108    /// Get the concrete type if this is `Concrete`.
109    pub fn as_concrete(&self) -> Option<Type> {
110        match self {
111            InferType::Concrete(ty) => Some(*ty),
112            _ => None,
113        }
114    }
115
116    /// Get the type variable ID if this is `Var`.
117    pub fn as_var(&self) -> Option<TypeVarId> {
118        match self {
119            InferType::Var(id) => Some(*id),
120            _ => None,
121        }
122    }
123}
124
125impl std::fmt::Display for InferType {
126    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
127        match self {
128            InferType::Concrete(ty) => write!(f, "{ty}"),
129            InferType::Var(id) => write!(f, "{id}"),
130            InferType::IntLiteral => write!(f, "{{integer}}"),
131            InferType::FloatLiteral => write!(f, "{{float}}"),
132            InferType::Array { element, length } => write!(f, "[{element}; {length}]"),
133        }
134    }
135}
136
137impl From<Type> for InferType {
138    fn from(ty: Type) -> Self {
139        InferType::Concrete(ty)
140    }
141}
142
143/// Allocator for fresh type variables.
144///
145/// Each function gets its own allocator to generate unique type variable IDs.
146#[derive(Debug, Default)]
147pub struct TypeVarAllocator {
148    next_id: u32,
149}
150
151impl TypeVarAllocator {
152    /// Create a new allocator starting from ID 0.
153    pub fn new() -> Self {
154        TypeVarAllocator { next_id: 0 }
155    }
156
157    /// Allocate a fresh type variable.
158    pub fn fresh(&mut self) -> TypeVarId {
159        let id = TypeVarId::new(self.next_id);
160        self.next_id += 1;
161        id
162    }
163
164    /// Get the number of type variables allocated.
165    pub fn count(&self) -> u32 {
166        self.next_id
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn test_type_var_allocator() {
176        let mut alloc = TypeVarAllocator::new();
177        let v0 = alloc.fresh();
178        let v1 = alloc.fresh();
179        let v2 = alloc.fresh();
180
181        assert_eq!(v0.index(), 0);
182        assert_eq!(v1.index(), 1);
183        assert_eq!(v2.index(), 2);
184        assert_eq!(alloc.count(), 3);
185    }
186
187    #[test]
188    fn test_infer_type_display() {
189        assert_eq!(format!("{}", InferType::Concrete(Type::I32)), "i32");
190        assert_eq!(format!("{}", InferType::Var(TypeVarId::new(5))), "?5");
191        assert_eq!(format!("{}", InferType::IntLiteral), "{integer}");
192    }
193}