lean_agentic/
typechecker.rs

1//! Trusted type checking kernel
2//!
3//! This is the minimal trusted core that verifies all terms.
4//! No term is accepted into the environment unless it passes
5//! these checks, ensuring logical soundness.
6
7use crate::arena::Arena;
8use crate::context::Context;
9use crate::conversion::Converter;
10use crate::environment::Environment;
11use crate::level::{LevelArena, LevelId};
12use crate::term::{TermId, TermKind};
13
14/// Type checker (trusted kernel)
15pub struct TypeChecker {
16    /// Conversion checker
17    converter: Converter,
18}
19
20impl TypeChecker {
21    /// Create a new type checker
22    pub fn new() -> Self {
23        Self {
24            converter: Converter::new(),
25        }
26    }
27
28    /// Infer the type of a term
29    ///
30    /// This is the heart of the type checker: Γ ⊢ t : ?
31    pub fn infer(
32        &mut self,
33        arena: &mut Arena,
34        levels: &mut LevelArena,
35        env: &Environment,
36        ctx: &Context,
37        term: TermId,
38    ) -> crate::Result<TermId> {
39        let kind = arena.kind(term).ok_or_else(|| {
40            crate::Error::Internal(format!("Invalid term ID: {:?}", term))
41        })?.clone();
42
43        match kind {
44            // Γ ⊢ Type u : Type (u+1)
45            TermKind::Sort(level_id) => {
46                let _level = levels.get(level_id).ok_or_else(|| {
47                    crate::Error::Internal("Invalid level ID".to_string())
48                })?;
49
50                let succ_level = levels.succ(level_id);
51                Ok(arena.mk_sort(succ_level))
52            }
53
54            // Γ ⊢ c : T if c : T in environment
55            TermKind::Const(name, _level_args) => {
56                let decl = env.get_decl(name).ok_or_else(|| {
57                    crate::Error::NotFound(format!("Constant not found: {:?}", name))
58                })?;
59
60                // TODO: Instantiate universe parameters with level_args
61                // For now, just return the declared type
62                Ok(decl.ty)
63            }
64
65            // Γ ⊢ #i : Γ(i)
66            TermKind::Var(idx) => {
67                ctx.type_of(idx).ok_or_else(|| {
68                    crate::Error::TypeError(format!("Variable #{} not in context", idx))
69                })
70            }
71
72            // Γ ⊢ f a : B[x := a] if Γ ⊢ f : Πx:A.B and Γ ⊢ a : A
73            TermKind::App(func, arg) => {
74                let func_ty = self.infer(arena, levels, env, ctx, func)?;
75
76                // Reduce function type to WHNF to expose Pi
77                let func_ty_whnf = self.converter.whnf(arena, env, ctx, func_ty)?;
78
79                if let Some(TermKind::Pi(binder, body)) = arena.kind(func_ty_whnf).cloned() {
80                    // Check argument has correct type
81                    self.check(arena, levels, env, ctx, arg, binder.ty)?;
82
83                    // Return body with variable substituted
84                    // B[x := a]
85                    self.converter.substitute(arena, body, 0, arg)
86                } else {
87                    Err(crate::Error::TypeError(format!(
88                        "Expected function type, got: {:?}",
89                        func_ty_whnf
90                    )))
91                }
92            }
93
94            // Γ ⊢ λx:A.b : Πx:A.B if Γ,x:A ⊢ b : B
95            TermKind::Lam(binder, body) => {
96                // Check binder type is well-formed
97                let binder_ty_sort = self.infer(arena, levels, env, ctx, binder.ty)?;
98                self.ensure_sort(arena, levels, env, ctx, binder_ty_sort)?;
99
100                // Check body under extended context
101                let mut new_ctx = ctx.clone();
102                new_ctx.push_var(binder.name, binder.ty);
103
104                let body_ty = self.infer(arena, levels, env, &new_ctx, body)?;
105
106                // Result type is Πx:A.B
107                Ok(arena.mk_pi(binder, body_ty))
108            }
109
110            // Γ ⊢ Πx:A.B : Type (imax u v)
111            // if Γ ⊢ A : Type u and Γ,x:A ⊢ B : Type v
112            TermKind::Pi(binder, body) => {
113                // Check domain is well-typed
114                let domain_ty = self.infer(arena, levels, env, ctx, binder.ty)?;
115                let domain_level = self.extract_level(arena, levels, env, ctx, domain_ty)?;
116
117                // Check codomain under extended context
118                let mut new_ctx = ctx.clone();
119                new_ctx.push_var(binder.name, binder.ty);
120
121                let codomain_ty = self.infer(arena, levels, env, &new_ctx, body)?;
122                let codomain_level = self.extract_level(arena, levels, env, &new_ctx, codomain_ty)?;
123
124                // Result universe is imax of domain and codomain
125                let result_level = levels.imax(domain_level, codomain_level);
126                Ok(arena.mk_sort(result_level))
127            }
128
129            // Γ ⊢ (let x:A := v in b) : B[x := v]
130            // if Γ ⊢ v : A and Γ,x:A ⊢ b : B
131            TermKind::Let(binder, value, body) => {
132                // Check value has declared type
133                self.check(arena, levels, env, ctx, value, binder.ty)?;
134
135                // Check body under extended context with let binding
136                let mut new_ctx = ctx.clone();
137                new_ctx.push(crate::context::ContextEntry::with_value(
138                    binder.name,
139                    binder.ty,
140                    value,
141                ));
142
143                let body_ty = self.infer(arena, levels, env, &new_ctx, body)?;
144
145                // Substitute value in body type
146                self.converter.substitute(arena, body_ty, 0, value)
147            }
148
149            // Metavariables: use their assigned type (set during elaboration)
150            TermKind::MVar(_) => {
151                // In the kernel, metavariables should already be resolved
152                Err(crate::Error::TypeError(
153                    "Metavariables not allowed in kernel".to_string(),
154                ))
155            }
156
157            // Literals
158            TermKind::Lit(lit) => {
159                match lit {
160                    crate::term::Literal::Nat(_) => {
161                        // TODO: Return Nat type when we have it
162                        // For now, return Type 0
163                        let zero = levels.zero();
164                        Ok(arena.mk_sort(zero))
165                    }
166                    crate::term::Literal::String(_) => {
167                        // TODO: Return String type
168                        let zero = levels.zero();
169                        Ok(arena.mk_sort(zero))
170                    }
171                }
172            }
173        }
174    }
175
176    /// Check that a term has an expected type
177    ///
178    /// Γ ⊢ t : T (checking mode)
179    pub fn check(
180        &mut self,
181        arena: &mut Arena,
182        levels: &mut LevelArena,
183        env: &Environment,
184        ctx: &Context,
185        term: TermId,
186        expected_ty: TermId,
187    ) -> crate::Result<()> {
188        let inferred_ty = self.infer(arena, levels, env, ctx, term)?;
189
190        if self.converter.is_def_eq(arena, env, ctx, inferred_ty, expected_ty)? {
191            Ok(())
192        } else {
193            Err(crate::Error::ConversionError {
194                expected: format!("{:?}", expected_ty),
195                actual: format!("{:?}", inferred_ty),
196            })
197        }
198    }
199
200    /// Ensure a term is a sort (Type u)
201    fn ensure_sort(
202        &mut self,
203        arena: &mut Arena,
204        levels: &mut LevelArena,
205        env: &Environment,
206        ctx: &Context,
207        term: TermId,
208    ) -> crate::Result<LevelId> {
209        let whnf = self.converter.whnf(arena, env, ctx, term)?;
210
211        if let Some(TermKind::Sort(level)) = arena.kind(whnf) {
212            Ok(*level)
213        } else {
214            Err(crate::Error::TypeError(format!(
215                "Expected sort, got: {:?}",
216                whnf
217            )))
218        }
219    }
220
221    /// Extract universe level from a sort
222    fn extract_level(
223        &mut self,
224        arena: &mut Arena,
225        levels: &mut LevelArena,
226        env: &Environment,
227        ctx: &Context,
228        ty: TermId,
229    ) -> crate::Result<LevelId> {
230        self.ensure_sort(arena, levels, env, ctx, ty)
231    }
232
233    /// Verify a declaration is well-typed before adding to environment
234    pub fn check_declaration(
235        &mut self,
236        arena: &mut Arena,
237        levels: &mut LevelArena,
238        env: &Environment,
239        decl: &crate::environment::Declaration,
240    ) -> crate::Result<()> {
241        let ctx = Context::new();
242
243        // Check type is well-formed
244        let ty_sort = self.infer(arena, levels, env, &ctx, decl.ty)?;
245        self.ensure_sort(arena, levels, env, &ctx, ty_sort)?;
246
247        // If there's a value, check it has the declared type
248        if let Some(value) = decl.value {
249            self.check(arena, levels, env, &ctx, value, decl.ty)?;
250        }
251
252        Ok(())
253    }
254
255    /// Get the converter (for advanced use)
256    pub fn converter(&mut self) -> &mut Converter {
257        &mut self.converter
258    }
259}
260
261impl Default for TypeChecker {
262    fn default() -> Self {
263        Self::new()
264    }
265}
266
267#[cfg(test)]
268mod tests {
269    use super::*;
270    use crate::symbol::SymbolId;
271    use crate::term::Binder;
272
273    #[test]
274    fn test_var_typing() {
275        let mut arena = Arena::new();
276        let mut levels = LevelArena::new();
277        let env = Environment::new();
278        let mut ctx = Context::new();
279        let mut tc = TypeChecker::new();
280
281        // Add x : Type 0
282        let zero = levels.zero();
283        let type0 = arena.mk_sort(zero);
284        ctx.push_var(SymbolId::new(0), type0);
285
286        // Check #0 has type Type 0
287        let var0 = arena.mk_var(0);
288        let ty = tc.infer(&mut arena, &mut levels, &env, &ctx, var0).unwrap();
289
290        assert_eq!(ty, type0);
291    }
292
293    #[test]
294    fn test_sort_typing() {
295        let mut arena = Arena::new();
296        let mut levels = LevelArena::new();
297        let env = Environment::new();
298        let ctx = Context::new();
299        let mut tc = TypeChecker::new();
300
301        // Type 0 : Type 1
302        let zero = levels.zero();
303        let type0 = arena.mk_sort(zero);
304
305        let ty = tc.infer(&mut arena, &mut levels, &env, &ctx, type0).unwrap();
306
307        // Check it's a sort (which is the expected result)
308        match arena.kind(ty) {
309            Some(TermKind::Sort(_)) => (), // Success
310            _ => panic!("Expected sort type"),
311        }
312    }
313
314    #[test]
315    fn test_lambda_typing() {
316        let mut arena = Arena::new();
317        let mut levels = LevelArena::new();
318        let env = Environment::new();
319        let ctx = Context::new();
320        let mut tc = TypeChecker::new();
321
322        // λx:Type 0. x  has type  Πx:Type 0. Type 0
323        let zero = levels.zero();
324        let type0 = arena.mk_sort(zero);
325
326        let x = arena.mk_var(0);
327        let binder = Binder::new(SymbolId::new(0), type0);
328        let lam = arena.mk_lam(binder, x);
329
330        let ty = tc.infer(&mut arena, &mut levels, &env, &ctx, lam).unwrap();
331
332        // Should be a Pi type
333        if let Some(TermKind::Pi(_, _)) = arena.kind(ty) {
334            // Correct
335        } else {
336            panic!("Expected Pi type");
337        }
338    }
339}