1use 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
14pub struct TypeChecker {
16 converter: Converter,
18}
19
20impl TypeChecker {
21 pub fn new() -> Self {
23 Self {
24 converter: Converter::new(),
25 }
26 }
27
28 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 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 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 Ok(decl.ty)
63 }
64
65 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 TermKind::App(func, arg) => {
74 let func_ty = self.infer(arena, levels, env, ctx, func)?;
75
76 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 self.check(arena, levels, env, ctx, arg, binder.ty)?;
82
83 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 TermKind::Lam(binder, body) => {
96 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 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 Ok(arena.mk_pi(binder, body_ty))
108 }
109
110 TermKind::Pi(binder, body) => {
113 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 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 let result_level = levels.imax(domain_level, codomain_level);
126 Ok(arena.mk_sort(result_level))
127 }
128
129 TermKind::Let(binder, value, body) => {
132 self.check(arena, levels, env, ctx, value, binder.ty)?;
134
135 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 self.converter.substitute(arena, body_ty, 0, value)
147 }
148
149 TermKind::MVar(_) => {
151 Err(crate::Error::TypeError(
153 "Metavariables not allowed in kernel".to_string(),
154 ))
155 }
156
157 TermKind::Lit(lit) => {
159 match lit {
160 crate::term::Literal::Nat(_) => {
161 let zero = levels.zero();
164 Ok(arena.mk_sort(zero))
165 }
166 crate::term::Literal::String(_) => {
167 let zero = levels.zero();
169 Ok(arena.mk_sort(zero))
170 }
171 }
172 }
173 }
174 }
175
176 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 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 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 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 let ty_sort = self.infer(arena, levels, env, &ctx, decl.ty)?;
245 self.ensure_sort(arena, levels, env, &ctx, ty_sort)?;
246
247 if let Some(value) = decl.value {
249 self.check(arena, levels, env, &ctx, value, decl.ty)?;
250 }
251
252 Ok(())
253 }
254
255 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 let zero = levels.zero();
283 let type0 = arena.mk_sort(zero);
284 ctx.push_var(SymbolId::new(0), type0);
285
286 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 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 match arena.kind(ty) {
309 Some(TermKind::Sort(_)) => (), _ => 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 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 if let Some(TermKind::Pi(_, _)) = arena.kind(ty) {
334 } else {
336 panic!("Expected Pi type");
337 }
338 }
339}