Skip to main content

logicaffeine_kernel/dag/sem/
whnf.rs

1//! Weak-head reduction of values: ι (recursors, with K and struct-η), quotients, projections,
2//! validated `Nat` literal operations, and δ through constant heads.
3//!
4//! β is already done by evaluation, so a value's weak head can only be blocked by a recursor or
5//! quotient eliminator, a stuck projection, a literal operation, or a definition. Every step is
6//! remembered on the value that took it: the unfolding of a definition application, the
7//! `whnf_core` and `whnf` of any value.
8
9use super::value::{Head, Kind, Sp, Val, V};
10use super::{reject, Sem};
11use crate::dag::env::{Kind as CKind, QuotKind, RecInfo};
12use crate::dag::error::KResult;
13use crate::dag::expr::StrId;
14use crate::dag::level::{LevelId, LevelsId};
15use crate::dag::name::NameId;
16use crate::dag::nat::Nat;
17
18impl Sem<'_> {
19    /// The constant heading a neutral, with its arguments.
20    #[inline]
21    pub(crate) fn const_app(&self, v: V) -> Option<(NameId, LevelsId, Kind, Sp)> {
22        match self.st.arena.get(v) {
23            Val::Neu { kind, head, spine } => match self.st.arena.head(head) {
24                Head::Const { name, levels } => Some((name, levels, kind, spine)),
25                _ => None,
26            },
27            _ => None,
28        }
29    }
30
31    #[inline]
32    pub(crate) fn is_const_app_of(&self, v: V, n: NameId, arity: usize) -> bool {
33        match self.const_app(v) {
34            Some((m, _, _, sp)) => m == n && self.st.arena.spine_len(sp) == arity,
35            None => false,
36        }
37    }
38
39    /// Weak-head normal form without δ. With `cheap_proj`, a projection's structure is only
40    /// brought to `whnf_core` (not unfolded), as Lean's `whnfCore` does inside conversion.
41    pub fn whnf_core(&mut self, v: V, cheap_proj: bool) -> KResult<V> {
42        let Val::Neu { kind, .. } = self.st.arena.get(v) else {
43            return Ok(v);
44        };
45        if !matches!(kind, Kind::Rec | Kind::Quot | Kind::Proj) {
46            return Ok(v);
47        }
48        if let Some(r) = self.st.arena.whnf_core_memo(v) {
49            return Ok(r);
50        }
51        self.enter()?;
52        let r = self.whnf_core_step(v, kind, cheap_proj);
53        self.leave();
54        let r = r?;
55        if !cheap_proj {
56            self.st.arena.set_whnf_core_memo(v, r);
57        }
58        Ok(r)
59    }
60
61    fn whnf_core_step(&mut self, v: V, kind: Kind, cheap_proj: bool) -> KResult<V> {
62        let reduced = match kind {
63            Kind::Rec => self.inductive_reduce_rec(v)?,
64            Kind::Quot => {
65                if self.k.env.quot_init {
66                    self.quot_reduce_rec(v)?
67                } else {
68                    None
69                }
70            }
71            Kind::Proj => self.reduce_proj_head(v, cheap_proj)?,
72            _ => None,
73        };
74        match reduced {
75            Some(r) => self.whnf_core(r, cheap_proj),
76            None => Ok(v),
77        }
78    }
79
80    /// A stuck projection at the head of `v`, reduced if its structure now reduces to a
81    /// constructor; the arguments applied to the projection are applied to the field.
82    fn reduce_proj_head(&mut self, v: V, cheap_proj: bool) -> KResult<Option<V>> {
83        let Val::Neu { head, spine, .. } = self.st.arena.get(v) else {
84            return Ok(None);
85        };
86        let Head::Proj { idx, x, .. } = self.st.arena.head(head) else {
87            return Ok(None);
88        };
89        let Some(field) = self.reduce_proj(idx, x, cheap_proj)? else {
90            return Ok(None);
91        };
92        let args = self.st.arena.args(spine);
93        self.apply_all(field, &args).map(Some)
94    }
95
96    /// Field `idx` of `x` if `x` reduces to a constructor application.
97    pub(crate) fn reduce_proj(&mut self, idx: u32, x: V, cheap_proj: bool) -> KResult<Option<V>> {
98        let c = if cheap_proj {
99            self.whnf_core(x, true)?
100        } else {
101            self.whnf(x)?
102        };
103        self.proj_field_of(idx, c)
104    }
105
106    /// Field `idx` of `c` if `c` is a constructor application (a string literal counting as the
107    /// constructor application it abbreviates).
108    pub(crate) fn proj_field_of(&mut self, idx: u32, mut c: V) -> KResult<Option<V>> {
109        if let Val::Str(s) = self.st.arena.get(c) {
110            let e = self.str_to_ctor(s)?;
111            c = self.whnf(e)?;
112        }
113        Ok(self.ctor_field(c, idx))
114    }
115
116    /// Field `idx` of a constructor application.
117    fn ctor_field(&self, c: V, idx: u32) -> Option<V> {
118        let (mk, _, kind, spine) = self.const_app(c)?;
119        if kind != Kind::Ctor {
120            return None;
121        }
122        let ci = self.k.env.ctor(mk)?;
123        self.st.arena.arg(spine, (ci.num_params + idx) as usize)
124    }
125
126    /// `x.idx` as evaluation builds it: the field at once if `x` is a constructor application,
127    /// otherwise a stuck projection (the way a `proj` term waits for `whnf_core`).
128    pub(crate) fn proj(&mut self, s: NameId, idx: u32, x: V) -> KResult<V> {
129        if let Some(f) = self.ctor_field(x, idx) {
130            return Ok(f);
131        }
132        Ok(self
133            .st
134            .arena
135            .head_value(Kind::Proj, Head::Proj { s, idx, x }))
136    }
137
138    // ---- δ ----
139
140    /// A definition's body at universe levels, as a value (evaluated once per session).
141    pub(crate) fn const_body(&mut self, n: NameId, ls: LevelsId) -> KResult<Option<V>> {
142        if let Some(v) = self.st.const_body.get(&(n.0, ls.0)) {
143            return Ok(Some(V(v)));
144        }
145        let Some(c) = self.k.env.get(n) else {
146            return Ok(None);
147        };
148        let Some(body) = c.delta_value() else {
149            return Ok(None);
150        };
151        if self.k.s.levels.get_list(ls).len() != c.lparams.len() {
152            return Ok(None);
153        }
154        let base = self.instance_base(n, ls)?;
155        let v = self.eval(base, body)?;
156        self.st.const_body.insert((n.0, ls.0), v.0);
157        Ok(Some(v))
158    }
159
160    /// One δ-step: the definition at the head of `v` unfolded and applied to the arguments.
161    pub fn unfold(&mut self, v: V) -> KResult<Option<V>> {
162        let Some((n, ls, Kind::Delta, spine)) = self.const_app(v) else {
163            return Ok(None);
164        };
165        if let Some(r) = self.st.arena.delta_memo(v) {
166            return Ok(Some(r));
167        }
168        let Some(body) = self.const_body(n, ls)? else {
169            return Ok(None);
170        };
171        let args = self.st.arena.args(spine);
172        let r = self.apply_all(body, &args)?;
173        self.st.arena.set_delta_memo(v, r);
174        Ok(Some(r))
175    }
176
177    /// Full weak-head normal form.
178    pub fn whnf(&mut self, v: V) -> KResult<V> {
179        if !matches!(self.st.arena.get(v), Val::Neu { .. }) {
180            return Ok(v);
181        }
182        if let Some(r) = self.st.arena.whnf_memo(v) {
183            return Ok(r);
184        }
185        self.enter()?;
186        let r = self.whnf_loop(v);
187        self.leave();
188        let r = r?;
189        self.st.arena.set_whnf_memo(v, r);
190        Ok(r)
191    }
192
193    fn whnf_loop(&mut self, v: V) -> KResult<V> {
194        let mut t = v;
195        loop {
196            t = self.whnf_core(t, false)?;
197            if let Some(r) = self.reduce_nat(t)? {
198                return Ok(r);
199            }
200            match self.unfold(t)? {
201                Some(u) => t = u,
202                None => return Ok(t),
203            }
204        }
205    }
206
207    // ---- literals ----
208
209    pub(crate) fn nat_value(&mut self, n: Nat) -> V {
210        let id = self.k.s.nat_id(n);
211        self.st.arena.nat(id)
212    }
213
214    /// `0` as `Nat.zero`, `n + 1` as `Nat.succ n`.
215    pub(crate) fn nat_to_ctor(&mut self, n: &Nat) -> KResult<V> {
216        if n.is_zero() {
217            return self.const0(self.k.wk.nat_zero);
218        }
219        let succ = self.const0(self.k.wk.nat_succ)?;
220        let pred = self.nat_value(n.sub(&Nat::from_u64(1)));
221        Ok(self.st.arena.neutral_app(succ, pred))
222    }
223
224    /// A string literal as `String.ofList [Char.ofNat c₀, …]`.
225    pub(crate) fn str_to_ctor(&mut self, s: StrId) -> KResult<V> {
226        let text = self.k.s.string(s);
227        let zero = self.k.s.levels.list(&[LevelId::ZERO]);
228        let ch = self.const0(self.k.wk.char_)?;
229        let nil_c = self.const_value(self.k.wk.list_nil, zero)?;
230        let nil = self.st.arena.neutral_app(nil_c, ch);
231        let cons_c = self.const_value(self.k.wk.list_cons, zero)?;
232        let cons = self.st.arena.neutral_app(cons_c, ch);
233        let of_nat = self.const0(self.k.wk.char_of_nat)?;
234        let mut list = nil;
235        for c in text.chars().rev() {
236            let lit = self.nat_value(Nat::from_u64(c as u64));
237            let chr = self.st.arena.neutral_app(of_nat, lit);
238            let cell = self.st.arena.neutral_app(cons, chr);
239            list = self.st.arena.neutral_app(cell, list);
240        }
241        let of_list = self.const0(self.k.wk.string_of_list)?;
242        Ok(self.st.arena.neutral_app(of_list, list))
243    }
244
245    /// `Nat.zero` or a literal, as a number.
246    pub(crate) fn raw_nat(&self, v: V) -> Option<Nat> {
247        match self.st.arena.get(v) {
248            Val::Nat(n) => Some(self.k.s.nat(n).clone()),
249            _ => match self.const_app(v) {
250                Some((n, ls, _, sp))
251                    if n == self.k.wk.nat_zero && ls == LevelsId::EMPTY && sp == Sp::NIL =>
252                {
253                    Some(Nat::zero())
254                }
255                _ => None,
256            },
257        }
258    }
259
260    fn bool_value(&mut self, b: bool) -> KResult<V> {
261        let n = if b {
262            self.k.wk.bool_true
263        } else {
264            self.k.wk.bool_false
265        };
266        self.const0(n)
267    }
268
269    /// Lean's `reduceNat`: `Nat.succ` and the validated binary operations on literals.
270    pub(crate) fn reduce_nat(&mut self, v: V) -> KResult<Option<V>> {
271        if !self.k.nat_literal_rules() {
272            return Ok(None);
273        }
274        let Some((n, ls, _, spine)) = self.const_app(v) else {
275            return Ok(None);
276        };
277        let nargs = self.st.arena.spine_len(spine);
278        if nargs == 1 {
279            if n != self.k.wk.nat_succ || ls != LevelsId::EMPTY {
280                return Ok(None);
281            }
282            let a = self.st.arena.spine_node(spine).arg;
283            let w = self.whnf(a)?;
284            return Ok(self.raw_nat(w).map(|x| self.nat_value(x.succ())));
285        }
286        if nargs != 2 {
287            return Ok(None);
288        }
289        let w = &self.k.wk;
290        let op = if n == w.nat_add {
291            0
292        } else if n == w.nat_sub {
293            1
294        } else if n == w.nat_mul {
295            2
296        } else if n == w.nat_pow {
297            3
298        } else if n == w.nat_gcd {
299            4
300        } else if n == w.nat_mod {
301            5
302        } else if n == w.nat_div {
303            6
304        } else if n == w.nat_beq {
305            7
306        } else if n == w.nat_ble {
307            8
308        } else if n == w.nat_land {
309            9
310        } else if n == w.nat_lor {
311            10
312        } else if n == w.nat_xor {
313            11
314        } else if n == w.nat_shift_left {
315            12
316        } else if n == w.nat_shift_right {
317            13
318        } else {
319            return Ok(None);
320        };
321        if !self.k.validated_primitives.contains(&n) {
322            return Ok(None);
323        }
324        let last = self.st.arena.spine_node(spine);
325        let (a0, a1) = (self.st.arena.spine_node(last.prev).arg, last.arg);
326        let wa = self.whnf(a0)?;
327        let Some(x) = self.raw_nat(wa) else {
328            return Ok(None);
329        };
330        let wb = self.whnf(a1)?;
331        let Some(y) = self.raw_nat(wb) else {
332            return Ok(None);
333        };
334        let r = match op {
335            0 => x.add(&y),
336            1 => x.sub(&y),
337            2 => x.mul(&y),
338            3 => match y.to_u64() {
339                Some(k) if k <= 1 << 24 => x.pow(k as u32),
340                _ => return Ok(None),
341            },
342            4 => x.gcd(&y),
343            5 => x.rem(&y),
344            6 => x.div(&y),
345            7 => return self.bool_value(x == y).map(Some),
346            8 => return self.bool_value(x <= y).map(Some),
347            9 => x.land(&y),
348            10 => x.lor(&y),
349            11 => x.xor(&y),
350            12 => match y.to_u64() {
351                Some(k) if k <= 1 << 24 => x.shl(k),
352                _ => return Ok(None),
353            },
354            13 => match y.to_u64() {
355                Some(k) => x.shr(k),
356                None => Nat::zero(),
357            },
358            _ => unreachable!(),
359        };
360        Ok(Some(self.nat_value(r)))
361    }
362
363    // ---- quotients ----
364
365    fn quot_reduce_rec(&mut self, v: V) -> KResult<Option<V>> {
366        let Some((n, _, _, spine)) = self.const_app(v) else {
367            return Ok(None);
368        };
369        let (mk_pos, arg_pos) = match self.k.env.get(n).map(|c| &c.kind) {
370            Some(CKind::Quot(QuotKind::Lift)) => (5, 3),
371            Some(CKind::Quot(QuotKind::Ind)) => (4, 3),
372            _ => return Ok(None),
373        };
374        let args = self.st.arena.args(spine);
375        if mk_pos >= args.len() {
376            return Ok(None);
377        }
378        let Some(q) = self
379            .k
380            .env
381            .quotients
382            .iter()
383            .find(|q| q.lift == n || q.ind == n)
384            .copied()
385        else {
386            return Ok(None);
387        };
388        let mk = self.whnf(args[mk_pos])?;
389        if !self.is_const_app_of(mk, q.mk, 3) {
390            return Ok(None);
391        }
392        let Val::Neu { spine: mks, .. } = self.st.arena.get(mk) else {
393            unreachable!()
394        };
395        let mk_arg = self.st.arena.spine_node(mks).arg;
396        let r = self.apply(args[arg_pos], mk_arg)?;
397        self.apply_all(r, &args[mk_pos + 1..]).map(Some)
398    }
399
400    // ---- recursors ----
401
402    fn inductive_reduce_rec(&mut self, v: V) -> KResult<Option<V>> {
403        let Some((rec, ls, _, spine)) = self.const_app(v) else {
404            return Ok(None);
405        };
406        let Some(c) = self.k.env.get(rec).cloned() else {
407            return Ok(None);
408        };
409        let CKind::Rec(info) = &c.kind else {
410            return Ok(None);
411        };
412        let args = self.st.arena.args(spine);
413        let major_idx = info.major_idx();
414        let Some(&major0) = args.get(major_idx) else {
415            return Ok(None);
416        };
417        let mut major = major0;
418        if info.k {
419            major = self.to_ctor_when_k(info, major)?;
420        }
421        major = match self.eta_major_first(info, major)? {
422            Some(expanded) => expanded,
423            None => {
424                let wm = self.whnf(major)?;
425                match self.st.arena.get(wm) {
426                    Val::Nat(n) if self.k.reads_nat_literal(self.k.s.nat(n)) => {
427                        let x = self.k.s.nat(n).clone();
428                        self.nat_to_ctor(&x)?
429                    }
430                    Val::Str(s) => {
431                        let e = self.str_to_ctor(s)?;
432                        self.whnf(e)?
433                    }
434                    _ => self.to_ctor_when_struct(info.major_induct, wm)?,
435                }
436            }
437        };
438        let Some((ctor, _, Kind::Ctor, mspine)) = self.const_app(major) else {
439            return Ok(None);
440        };
441        let Some(rule_idx) = info.rules.iter().position(|r| r.ctor == ctor) else {
442            return Ok(None);
443        };
444        let rule = &info.rules[rule_idx];
445        let nfields = rule.nfields as usize;
446        let margs = self.st.arena.args(mspine);
447        if nfields > margs.len() {
448            return Ok(None);
449        }
450        if self.k.s.levels.get_list(ls).len() != c.lparams.len() {
451            return Ok(None);
452        }
453        let rhs = match self.st.rule_rhs.get(&(rec.0, rule_idx as u32, ls.0)) {
454            Some(r) => V(r),
455            None => {
456                let base = self.instance_base(rec, ls)?;
457                let r = self.eval(base, rule.rhs)?;
458                self.st.rule_rhs.insert((rec.0, rule_idx as u32, ls.0), r.0);
459                r
460            }
461        };
462        let mut r = self.apply_all(rhs, &args[..info.first_index_idx()])?;
463        r = self.apply_all(r, &margs[margs.len() - nfields..])?;
464        r = self.apply_all(r, &args[major_idx + 1..])?;
465        Ok(Some(r))
466    }
467
468    /// The inductive's only constructor applied to the parameters of `ty`.
469    fn mk_nullary_ctor(&mut self, ty: V, nparams: usize) -> KResult<Option<V>> {
470        let Some((d, ls, _, sp)) = self.const_app(ty) else {
471            return Ok(None);
472        };
473        let Some(ctor) = self
474            .k
475            .env
476            .inductive(d)
477            .and_then(|i| i.ctors.first().copied())
478        else {
479            return Ok(None);
480        };
481        let c = self.const_value(ctor, ls)?;
482        let args = self.st.arena.args(sp);
483        self.apply_all(c, &args[..nparams.min(args.len())])
484            .map(Some)
485    }
486
487    fn to_ctor_when_k(&mut self, info: &RecInfo, e: V) -> KResult<V> {
488        let et = self.type_of(e)?;
489        let app_type = self.whnf(et)?;
490        let Some((ti, _, _, _)) = self.const_app(app_type) else {
491            return Ok(e);
492        };
493        if ti != info.major_induct {
494            return Ok(e);
495        }
496        let Some(new_ctor) = self.mk_nullary_ctor(app_type, info.num_params as usize)? else {
497            return Ok(e);
498        };
499        let nt = self.type_of(new_ctor)?;
500        if !self.is_def_eq(app_type, nt)? {
501            return Ok(e);
502        }
503        Ok(new_ctor)
504    }
505
506    fn is_constructor_app(&self, v: V) -> bool {
507        matches!(self.const_app(v), Some((_, _, Kind::Ctor, _)))
508    }
509
510    /// The major premise of a structure's recursor as its constructor applied to its projections, taken before the
511    /// major is reduced; `None` where that does not apply. By η a structure's value is that application, so the rule
512    /// fires without the major's weak head normal form, which a proof need not have
513    /// (`corner-cases/eta-struct-major-no-whnf`). It applies to a structure in a universe that is never `Prop`, and to
514    /// one in `Prop` whose own recursor eliminates into every universe (`And`), whose fields are all proofs.
515    fn eta_major_first(&mut self, info: &RecInfo, major: V) -> KResult<Option<V>> {
516        if !self.k.env.is_non_rec_structure(info.major_induct)
517            || !matches!(self.st.arena.get(major), Val::Neu { .. })
518            || self.is_constructor_app(major)
519        {
520            return Ok(None);
521        }
522        let et = self.type_of(major)?;
523        let ety = self.whnf(et)?;
524        let Some((i, ls, _, sp)) = self.const_app(ety) else {
525            return Ok(None);
526        };
527        if i != info.major_induct {
528            return Ok(None);
529        }
530        let tt = self.type_of(ety)?;
531        let sort = self.whnf(tt)?;
532        let Val::Sort(u) = self.st.arena.get(sort) else {
533            return Ok(None);
534        };
535        let levels = &self.k.s.levels;
536        if !(levels.is_never_zero(u) || (levels.is_always_zero(u) && self.k.env.eliminates_into_every_sort(info.major_induct))) {
537            return Ok(None);
538        }
539        let args = self.st.arena.args(sp);
540        self.expand_eta_struct(i, ls, &args, major).map(Some)
541    }
542
543    fn to_ctor_when_struct(&mut self, induct: NameId, e: V) -> KResult<V> {
544        if !self.k.env.is_non_rec_structure(induct) || self.is_constructor_app(e) {
545            return Ok(e);
546        }
547        let et = self.type_of(e)?;
548        let ety = self.whnf(et)?;
549        let Some((i, ls, _, sp)) = self.const_app(ety) else {
550            return Ok(e);
551        };
552        if i != induct {
553            return Ok(e);
554        }
555        let tt = self.type_of(ety)?;
556        let sort = self.whnf(tt)?;
557        let Val::Sort(u) = self.st.arena.get(sort) else {
558            return Ok(e);
559        };
560        if !self.k.s.levels.is_never_zero(u) {
561            return Ok(e);
562        }
563        let args = self.st.arena.args(sp);
564        self.expand_eta_struct(i, ls, &args, e)
565    }
566
567    /// `mk params (x.1) … (x.n)` for a structure value `x`.
568    pub(crate) fn expand_eta_struct(
569        &mut self,
570        i: NameId,
571        ls: LevelsId,
572        args: &[V],
573        x: V,
574    ) -> KResult<V> {
575        let ctor = self
576            .k
577            .env
578            .inductive(i)
579            .ok_or_else(|| reject("not an inductive type"))?
580            .ctors[0];
581        let ci = self
582            .k
583            .env
584            .ctor(ctor)
585            .ok_or_else(|| reject("not a constructor"))?
586            .clone();
587        let c = self.const_value(ctor, ls)?;
588        let mut r = self.apply_all(c, &args[..(ci.num_params as usize).min(args.len())])?;
589        for f in 0..ci.num_fields {
590            let p = self.proj(i, f, x)?;
591            r = self.apply(r, p)?;
592        }
593        Ok(r)
594    }
595}