Skip to main content

lisette_semantics/checker/
type_env.rs

1//! Union-find-style binding table for `Type::Var` handles.
2//!
3//! `Type::Var(TypeVarId)` is a handle; the binding (Unbound vs Bound-to-a-Type)
4//! lives here in `entries`, indexed by id. Cloning a `Type` clones just the
5//! handle, so `Type` is a pure value (Clone / Eq / Hash / Serialize friendly)
6//! with no shared mutable state.
7//!
8//! Speculative unification works through the `undo_log`: a fresh log is
9//! pushed when entering speculation, bindings are recorded as they happen,
10//! and on Err the originals are restored in reverse order. On Ok the log is
11//! either discarded (no enclosing speculation) or appended to the parent log
12//! (nested speculation — the bindings are committed to the parent, but still
13//! reversible if the parent fails).
14
15use ecow::EcoString;
16use syntax::types::{Bound, Type, TypeVarId};
17
18#[derive(Debug, Clone)]
19pub enum VarState {
20    Unbound { hint: Option<EcoString> },
21    Bound(Type),
22}
23
24pub struct TypeEnv {
25    entries: Vec<VarState>,
26    /// When Some, bindings performed during the current speculative region
27    /// are logged here as `(id, prior_state)` so they can be reverted.
28    undo_log: Option<Vec<(TypeVarId, VarState)>>,
29}
30
31impl Default for TypeEnv {
32    fn default() -> Self {
33        Self::new()
34    }
35}
36
37impl TypeEnv {
38    pub fn new() -> Self {
39        Self {
40            entries: Vec::new(),
41            undo_log: None,
42        }
43    }
44
45    /// Allocate a fresh unbound variable and return its handle.
46    pub fn fresh(&mut self, hint: Option<EcoString>) -> TypeVarId {
47        let id = TypeVarId(self.entries.len() as u32);
48        self.entries.push(VarState::Unbound { hint });
49        id
50    }
51
52    fn slot(id: TypeVarId) -> usize {
53        debug_assert!(
54            !id.is_reserved(),
55            "TypeEnv should not be queried for reserved ids"
56        );
57        id.0 as usize
58    }
59
60    pub fn state(&self, id: TypeVarId) -> &VarState {
61        &self.entries[Self::slot(id)]
62    }
63
64    pub fn is_unbound(&self, id: TypeVarId) -> bool {
65        if id.is_reserved() {
66            return true;
67        }
68        matches!(self.entries[Self::slot(id)], VarState::Unbound { .. })
69    }
70
71    /// Bind `id` to `ty`. Reserved sentinel ids (ignored/uninferred) are
72    /// silently accepted: they unify with anything without storing anything.
73    pub fn bind(&mut self, id: TypeVarId, ty: Type) {
74        if id.is_reserved() {
75            return;
76        }
77        let slot = Self::slot(id);
78        let old = std::mem::replace(&mut self.entries[slot], VarState::Bound(ty));
79        if let Some(log) = &mut self.undo_log {
80            log.push((id, old));
81        }
82    }
83
84    /// Follow a `Type::Var` chain one step at a time until we reach either
85    /// an unbound variable or a non-Var type.
86    pub fn shallow_resolve(&self, ty: &Type) -> Type {
87        let mut current = ty.clone();
88        loop {
89            match &current {
90                Type::Var { id, .. } if !id.is_reserved() => match &self.entries[Self::slot(*id)] {
91                    VarState::Unbound { .. } => return current,
92                    VarState::Bound(bound) => current = bound.clone(),
93                },
94                _ => return current,
95            }
96        }
97    }
98
99    /// Deep resolve: chase `Type::Var` chains, substitute every bound var with
100    /// its chased value, and recurse into composites. Unbound vars (including
101    /// reserved sentinel ids like `IGNORED`/`UNINFERRED`) are preserved as-is.
102    /// Used both during inference and as the post-inference freeze pass.
103    pub fn resolve(&self, ty: &Type) -> Type {
104        self.resolve_changed(ty).unwrap_or_else(|| ty.clone())
105    }
106
107    /// Resolve `ty` in place; skips the clone-and-rebuild when nothing is bound.
108    /// Returns whether anything changed.
109    pub fn resolve_in_place(&self, ty: &mut Type) -> bool {
110        if let Some(resolved) = self.resolve_changed(ty) {
111            *ty = resolved;
112            true
113        } else {
114            false
115        }
116    }
117
118    /// Returns `Some` only when resolving `ty` would change it (some bound var
119    /// is reachable), allocating just the changed spine. `None` means unchanged.
120    fn resolve_changed(&self, ty: &Type) -> Option<Type> {
121        match ty {
122            Type::Var { id, .. } if !id.is_reserved() => match &self.entries[Self::slot(*id)] {
123                VarState::Unbound { .. } => None,
124                VarState::Bound(bound) => Some(self.resolve(bound)),
125            },
126            Type::Nominal {
127                id,
128                params,
129                underlying_ty,
130            } => {
131                let new_params = self.resolve_slice(params);
132                let new_underlying = underlying_ty
133                    .as_ref()
134                    .and_then(|u| self.resolve_changed(u).map(Box::new));
135                if new_params.is_none() && new_underlying.is_none() {
136                    return None;
137                }
138                Some(Type::Nominal {
139                    id: id.clone(),
140                    params: new_params.unwrap_or_else(|| params.clone()),
141                    underlying_ty: new_underlying
142                        .map(Some)
143                        .unwrap_or_else(|| underlying_ty.clone()),
144                })
145            }
146            Type::Compound { kind, args } => self
147                .resolve_slice(args)
148                .map(|args| Type::Compound { kind: *kind, args }),
149            Type::Function(f) => {
150                let new_params = self.resolve_slice(&f.params);
151                let new_return = self.resolve_changed(&f.return_type).map(Box::new);
152                let new_bounds = self.resolve_bounds(&f.bounds);
153                if new_params.is_none() && new_return.is_none() && new_bounds.is_none() {
154                    return None;
155                }
156                Some(Type::function(
157                    new_params.unwrap_or_else(|| f.params.clone()),
158                    f.param_mutability.clone(),
159                    new_bounds.unwrap_or_else(|| f.bounds.clone()),
160                    new_return.unwrap_or_else(|| f.return_type.clone()),
161                ))
162            }
163            Type::Forall { vars, body } => self.resolve_changed(body).map(|body| Type::Forall {
164                vars: vars.clone(),
165                body: Box::new(body),
166            }),
167            Type::Tuple(elements) => self.resolve_slice(elements).map(Type::Tuple),
168            _ => None,
169        }
170    }
171
172    /// [`resolve_changed`] over a slice; `Some` only if an element changed.
173    fn resolve_slice(&self, items: &[Type]) -> Option<Vec<Type>> {
174        let mut out: Option<Vec<Type>> = None;
175        for (i, item) in items.iter().enumerate() {
176            match self.resolve_changed(item) {
177                Some(resolved) => {
178                    out.get_or_insert_with(|| items[..i].to_vec())
179                        .push(resolved);
180                }
181                None => {
182                    if let Some(v) = out.as_mut() {
183                        v.push(item.clone());
184                    }
185                }
186            }
187        }
188        out
189    }
190
191    /// Bound-list variant of [`resolve_slice`].
192    fn resolve_bounds(&self, bounds: &[Bound]) -> Option<Vec<Bound>> {
193        let mut out: Option<Vec<Bound>> = None;
194        for (i, b) in bounds.iter().enumerate() {
195            let new_generic = self.resolve_changed(&b.generic);
196            let new_ty = self.resolve_changed(&b.ty);
197            if new_generic.is_none() && new_ty.is_none() {
198                if let Some(v) = out.as_mut() {
199                    v.push(b.clone());
200                }
201                continue;
202            }
203            out.get_or_insert_with(|| bounds[..i].to_vec()).push(Bound {
204                param_name: b.param_name.clone(),
205                generic: new_generic.unwrap_or_else(|| b.generic.clone()),
206                ty: new_ty.unwrap_or_else(|| b.ty.clone()),
207            });
208        }
209        out
210    }
211
212    /// Occurs check: does `id` appear anywhere inside `ty` (following Var
213    /// chains but stopping at unbound Vars)?
214    pub fn occurs(&self, id: TypeVarId, ty: &Type) -> bool {
215        match ty {
216            Type::Var { id: other, .. } => {
217                if *other == id {
218                    return true;
219                }
220                if other.is_reserved() {
221                    return false;
222                }
223                match &self.entries[Self::slot(*other)] {
224                    VarState::Unbound { .. } => false,
225                    VarState::Bound(bound) => self.occurs(id, bound),
226                }
227            }
228            Type::Nominal { params, .. } => params.iter().any(|p| self.occurs(id, p)),
229            Type::Compound { args, .. } => args.iter().any(|a| self.occurs(id, a)),
230            Type::Function(f) => {
231                f.params.iter().any(|p| self.occurs(id, p)) || self.occurs(id, &f.return_type)
232            }
233            Type::Forall { body, .. } => self.occurs(id, body),
234            Type::Tuple(elements) => elements.iter().any(|e| self.occurs(id, e)),
235            _ => false,
236        }
237    }
238
239    /// Begin a speculative region. Caller holds the returned handle and
240    /// passes it back to `end_speculation` with the region's outcome.
241    pub fn begin_speculation(&mut self) -> Speculation {
242        let prev = self.undo_log.take();
243        self.undo_log = Some(Vec::new());
244        Speculation { prev }
245    }
246
247    /// End a speculative region. If `is_err`, revert all bindings made
248    /// during the region. Otherwise, either commit them (no enclosing
249    /// region) or append them to the enclosing region's log (so it can
250    /// still revert them if it fails).
251    pub fn end_speculation(&mut self, spec: Speculation, is_err: bool) {
252        let log = self.undo_log.take().expect("speculation log must exist");
253        self.undo_log = spec.prev;
254        if is_err {
255            for (id, original) in log.into_iter().rev() {
256                self.entries[Self::slot(id)] = original;
257            }
258        } else if let Some(parent_log) = &mut self.undo_log {
259            parent_log.extend(log);
260        }
261    }
262}
263
264/// Handle returned by `begin_speculation`, consumed by `end_speculation`.
265/// Not clonable — ensures each region is ended exactly once.
266#[must_use]
267pub struct Speculation {
268    prev: Option<Vec<(TypeVarId, VarState)>>,
269}
270
271/// Extension trait for `Type` giving env-aware resolve convenience methods.
272/// Call-site sugar for `env.resolve(&ty)` written as `ty.resolve_in(&env)`.
273pub trait EnvResolve {
274    fn resolve_in(&self, env: &TypeEnv) -> Type;
275    fn shallow_resolve_in(&self, env: &TypeEnv) -> Type;
276}
277
278impl EnvResolve for Type {
279    fn resolve_in(&self, env: &TypeEnv) -> Type {
280        env.resolve(self)
281    }
282    fn shallow_resolve_in(&self, env: &TypeEnv) -> Type {
283        env.shallow_resolve(self)
284    }
285}