use std::collections::{BTreeMap, HashSet};
use crate::ast::{BinOp, Expr, Spanned, TopLevel};
use super::untranslate::UntranslatedGoal;
#[derive(Debug, Clone)]
pub enum CalcVerdict {
Lemma(Box<UntranslatedGoal>),
Decline(String),
}
pub struct CalcEnv {
ctors: HashSet<String>,
fn_ret: BTreeMap<String, String>,
}
impl CalcEnv {
pub fn from_items(items: &[TopLevel]) -> Self {
let mut ctors = HashSet::new();
let mut fn_ret = BTreeMap::new();
for it in items {
match it {
TopLevel::TypeDef(crate::ast::TypeDef::Sum { name, variants, .. }) => {
for v in variants {
ctors.insert(v.name.clone());
ctors.insert(format!("{name}.{}", v.name));
}
}
TopLevel::FnDef(f) => {
fn_ret.insert(f.name.clone(), f.return_type.clone());
}
_ => {}
}
}
CalcEnv { ctors, fn_ret }
}
fn is_ctor(&self, name: &str) -> bool {
self.ctors.contains(name)
|| name
.rsplit_once('.')
.is_some_and(|(_, s)| self.ctors.contains(s))
}
fn fn_return(&self, name: &str) -> Option<&str> {
self.fn_ret.get(name).map(String::as_str)
}
}
pub fn calculate(
goal: &UntranslatedGoal,
env: &CalcEnv,
reserved: &HashSet<String>,
) -> CalcVerdict {
if goal.premises.len() > 1 {
return CalcVerdict::Decline(format!(
"{} surviving premises (a `when` is one Bool expression)",
goal.premises.len()
));
}
let mut g = goal.clone();
let substituted = subst_ih_once(&mut g);
if let Some(prem) = g.premises.first().cloned() {
match lift_antiunify(&mut g, &prem, env, reserved) {
Ok(()) => {}
Err(reason) => return CalcVerdict::Decline(reason),
}
}
let lifted_b = lift_opaque_subterms(&mut g, env, reserved);
rebuild_givens(&mut g, goal);
if !substituted && !lifted_b && g.claim == goal.claim && g.premises == goal.premises {
return CalcVerdict::Decline(
"residual is the parent claim itself (no decomposition — an executor gap)".to_string(),
);
}
CalcVerdict::Lemma(Box::new(g))
}
fn subst_ih_once(g: &mut UntranslatedGoal) -> bool {
let Some(prem) = g.premises.first() else {
return false;
};
let Expr::BinOp(BinOp::Eq, l, r) = &prem.node else {
return false;
};
let (l, r) = (l.node.clone(), r.node.clone());
let in_claim = contains_subtree(&g.claim.0.node, &l) || contains_subtree(&g.claim.1.node, &l);
if !in_claim || contains_subtree(&r, &l) {
return false;
}
substitute(&mut g.claim.0.node, &l, &r);
substitute(&mut g.claim.1.node, &l, &r);
g.premises.clear();
true
}
fn lift_antiunify(
g: &mut UntranslatedGoal,
prem: &Spanned<Expr>,
env: &CalcEnv,
reserved: &HashSet<String>,
) -> Result<(), String> {
let Expr::BinOp(BinOp::Eq, pl, pr) = &prem.node else {
return Ok(()); };
let mut existing: HashSet<String> = g.givens.iter().map(|(n, _)| n.clone()).collect();
existing.extend(reserved.iter().cloned());
let mut fresh = FreshVars::new(existing);
let mut subst: Vec<(Expr, String)> = Vec::new();
let new_lhs = antiunify(&g.claim.0.node, &pl.node, env, &mut subst, &mut fresh)?;
let new_rhs = antiunify(&g.claim.1.node, &pr.node, env, &mut subst, &mut fresh)?;
if subst.is_empty() {
return Ok(()); }
g.claim.0.node = new_lhs;
g.claim.1.node = new_rhs;
let mut new_prem = prem.node.clone();
for (concrete, var) in &subst {
substitute(&mut new_prem, concrete, &Expr::Ident(var.clone()));
}
g.premises = vec![Spanned::new(new_prem, 0)];
for (concrete, var) in &subst {
let ty = lifted_type(concrete, env)
.ok_or_else(|| format!("cannot type the lifted variable `{var}`"))?;
g.givens.push((var.clone(), ty));
}
Ok(())
}
fn antiunify(
claim: &Expr,
prem: &Expr,
env: &CalcEnv,
subst: &mut Vec<(Expr, String)>,
fresh: &mut FreshVars,
) -> Result<Expr, String> {
if claim == prem {
if is_liftable_atom_or_call(claim, env) {
return Ok(Expr::Ident(var_for(subst, claim, fresh)));
}
return Ok(claim.clone());
}
if let Some((ctor, inner)) = unary_ctor_app(claim, env) {
let g_inner = antiunify(inner, prem, env, subst, fresh)?;
return Ok(rewrap_ctor(ctor, g_inner));
}
if let (Some((h1, a1)), Some((h2, a2))) = (as_call(claim), as_call(prem))
&& h1 == h2
&& a1.len() == a2.len()
{
let mut args = Vec::with_capacity(a1.len());
for (c, p) in a1.iter().zip(a2.iter()) {
args.push(Spanned::new(
antiunify(&c.node, &p.node, env, subst, fresh)?,
0,
));
}
return Ok(Expr::FnCall(
Box::new(Spanned::new(Expr::Ident(h1.to_string()), 0)),
args,
));
}
if let (Expr::BinOp(o1, l1, r1), Expr::BinOp(o2, l2, r2)) = (claim, prem)
&& o1 == o2
{
let l = antiunify(&l1.node, &l2.node, env, subst, fresh)?;
let r = antiunify(&r1.node, &r2.node, env, subst, fresh)?;
return Ok(Expr::BinOp(
*o1,
Box::new(Spanned::new(l, 0)),
Box::new(Spanned::new(r, 0)),
));
}
Err("anti-unification diff needs a term outside the goal and IH".to_string())
}
fn lift_opaque_subterms(
g: &mut UntranslatedGoal,
env: &CalcEnv,
reserved: &HashSet<String>,
) -> bool {
let mut lhs: Vec<Expr> = Vec::new();
collect_fn_apps(&g.claim.0.node, env, true, &mut lhs);
let mut rhs: Vec<Expr> = Vec::new();
collect_fn_apps(&g.claim.1.node, env, true, &mut rhs);
let mut both: Vec<Expr> = Vec::new();
for e in &lhs {
if rhs.iter().any(|r| r == e) && !both.iter().any(|b| b == e) {
both.push(e.clone());
}
}
let maximal: Vec<Expr> = both
.iter()
.filter(|e| {
!both
.iter()
.any(|other| other != *e && contains_subtree(other, e))
})
.cloned()
.collect();
let mut existing: HashSet<String> = g.givens.iter().map(|(n, _)| n.clone()).collect();
existing.extend(reserved.iter().cloned());
let mut fresh = FreshVars::new(existing);
let mut lifted = false;
for sub in maximal {
let Some(ty) = lifted_type(&sub, env) else {
continue; };
let var = fresh.next();
substitute(&mut g.claim.0.node, &sub, &Expr::Ident(var.clone()));
substitute(&mut g.claim.1.node, &sub, &Expr::Ident(var.clone()));
for p in &mut g.premises {
substitute(&mut p.node, &sub, &Expr::Ident(var.clone()));
}
g.givens.push((var, ty));
lifted = true;
}
lifted
}
fn collect_fn_apps(e: &Expr, env: &CalcEnv, top: bool, out: &mut Vec<Expr>) {
if !top && is_user_fn_call(e, env) {
out.push(e.clone());
}
for c in children(e) {
collect_fn_apps(&c.node, env, false, out);
}
}
struct FreshVars {
used: HashSet<String>,
idx: usize,
}
impl FreshVars {
fn new(used: HashSet<String>) -> Self {
FreshVars { used, idx: 0 }
}
fn next(&mut self) -> String {
loop {
let letter = (b'a' + (self.idx % 26) as u8) as char;
let round = self.idx / 26;
let name = if round == 0 {
letter.to_string()
} else {
format!("{letter}{round}")
};
self.idx += 1;
if !self.used.contains(&name) {
self.used.insert(name.clone());
return name;
}
}
}
}
fn var_for(subst: &mut Vec<(Expr, String)>, subterm: &Expr, fresh: &mut FreshVars) -> String {
if let Some((_, v)) = subst.iter().find(|(e, _)| e == subterm) {
return v.clone();
}
let v = fresh.next();
subst.push((subterm.clone(), v.clone()));
v
}
fn unary_ctor_app<'a>(e: &'a Expr, env: &CalcEnv) -> Option<(&'a str, &'a Expr)> {
if let Expr::FnCall(head, args) = e
&& args.len() == 1
&& let Expr::Ident(name) = &head.node
&& env.is_ctor(name)
{
return Some((name.as_str(), &args[0].node));
}
None
}
fn rewrap_ctor(ctor: &str, inner: Expr) -> Expr {
Expr::FnCall(
Box::new(Spanned::new(Expr::Ident(ctor.to_string()), 0)),
vec![Spanned::new(inner, 0)],
)
}
fn as_call(e: &Expr) -> Option<(&str, &[Spanned<Expr>])> {
if let Expr::FnCall(head, args) = e
&& let Expr::Ident(name) = &head.node
{
return Some((name.as_str(), args.as_slice()));
}
None
}
fn is_user_fn_call(e: &Expr, env: &CalcEnv) -> bool {
matches!(as_call(e), Some((name, _)) if !env.is_ctor(name) && lifted_head(name, env).is_some())
}
fn is_liftable_atom_or_call(e: &Expr, env: &CalcEnv) -> bool {
match e {
Expr::FnCall(..) => lifted_type(e, env).is_some(),
_ => false,
}
}
fn lifted_type(e: &Expr, env: &CalcEnv) -> Option<String> {
let (name, _) = as_call(e)?;
lifted_head(name, env).map(str::to_string)
}
fn lifted_head<'a>(name: &str, env: &'a CalcEnv) -> Option<&'a str> {
env.fn_return(name)
.or_else(|| name.rsplit_once('.').and_then(|(_, s)| env.fn_return(s)))
}
fn children(e: &Expr) -> Vec<&Spanned<Expr>> {
match e {
Expr::FnCall(head, args) => {
let mut v = vec![head.as_ref()];
v.extend(args.iter());
v
}
Expr::BinOp(_, a, b) => vec![a.as_ref(), b.as_ref()],
Expr::Neg(a) => vec![a.as_ref()],
Expr::List(xs) | Expr::Tuple(xs) => xs.iter().collect(),
Expr::Attr(b, _) => vec![b.as_ref()],
_ => vec![],
}
}
fn contains_subtree(hay: &Expr, needle: &Expr) -> bool {
hay == needle
|| children(hay)
.iter()
.any(|c| contains_subtree(&c.node, needle))
}
fn substitute(e: &mut Expr, from: &Expr, to: &Expr) {
if e == from {
*e = to.clone();
return;
}
match e {
Expr::FnCall(head, args) => {
substitute(&mut head.node, from, to);
for a in args {
substitute(&mut a.node, from, to);
}
}
Expr::BinOp(_, a, b) => {
substitute(&mut a.node, from, to);
substitute(&mut b.node, from, to);
}
Expr::Neg(a) => substitute(&mut a.node, from, to),
Expr::List(xs) | Expr::Tuple(xs) => {
for x in xs {
substitute(&mut x.node, from, to);
}
}
Expr::Attr(b, _) => substitute(&mut b.node, from, to),
_ => {}
}
}
fn rebuild_givens(g: &mut UntranslatedGoal, original: &UntranslatedGoal) {
let mut used = HashSet::new();
collect_idents(&g.claim.0.node, &mut used);
collect_idents(&g.claim.1.node, &mut used);
for p in &g.premises {
collect_idents(&p.node, &mut used);
}
let orig_names: HashSet<&String> = original.givens.iter().map(|(n, _)| n).collect();
g.givens.retain(|(n, _)| used.contains(n));
g.givens.sort_by_key(|(n, _)| !orig_names.contains(n));
}
fn collect_idents(e: &Expr, out: &mut HashSet<String>) {
if let Expr::Ident(n) = e {
out.insert(n.clone());
}
for c in children(e) {
collect_idents(&c.node, out);
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ast::unparse;
use crate::codegen::lean::untranslate::{PeanoCtx, UntranslateCtx, untranslate_goal_ctx};
fn render(g: &UntranslatedGoal) -> String {
let mut out = String::new();
for (n, t) in &g.givens {
out.push_str(&format!("given {n}: {t}; "));
}
for p in &g.premises {
let mut buf = String::new();
unparse::write_expr_public(&mut buf, p, 0).unwrap();
out.push_str(&format!("when {buf}; "));
}
let mut l = String::new();
let mut r = String::new();
unparse::write_expr_public(&mut l, &g.claim.0, 0).unwrap();
unparse::write_expr_public(&mut r, &g.claim.1, 0).unwrap();
out.push_str(&format!("{l} => {r}"));
out
}
fn nat_env() -> CalcEnv {
let src = "type Nat\n Z\n S(Nat)\n\n\
fn len(xs: List<Nat>) -> Nat\n Nat.Z\n\n\
fn le(x: Nat, y: Nat) -> Bool\n true\n\n\
fn filterZ(xs: List<Nat>) -> List<Nat>\n xs\n\n\
fn rev(xs: List<Nat>) -> List<Nat>\n xs\n";
CalcEnv::from_items(&crate::source::parse_source(src).expect("parses"))
}
fn peano_nat() -> UntranslateCtx {
UntranslateCtx {
peano: Some(PeanoCtx {
type_name: "Nat".to_string(),
zero_ctor: "Z".to_string(),
succ_ctor: "S".to_string(),
}),
}
}
const DUMP_P66_1: &str = include_str!("testdata/lemma_calc_krok0/p66_1.json");
#[test]
fn antiunify_forces_the_prop_66_successor_lemma() {
let goal = untranslate_goal_ctx(DUMP_P66_1, &peano_nat()).expect("in grammar");
let CalcVerdict::Lemma(l) = calculate(&goal, &nat_env(), &HashSet::new()) else {
panic!("expected a forced lemma");
};
assert_eq!(
render(&l),
"given a: Nat; given b: Nat; when (le(a, b) == true); le(a, Nat.S(b)) => true"
);
}
#[test]
fn fresh_vars_skip_parent_law_given_names() {
let goal = untranslate_goal_ctx(DUMP_P66_1, &peano_nat()).expect("in grammar");
let reserved: HashSet<String> = ["a".to_string()].into_iter().collect();
let CalcVerdict::Lemma(l) = calculate(&goal, &nat_env(), &reserved) else {
panic!("expected a forced lemma");
};
assert_eq!(
render(&l),
"given b: Nat; given c: Nat; when (le(b, c) == true); le(b, Nat.S(c)) => true"
);
}
#[test]
fn subst_consumes_ih_and_lifts_snoc_general_lemma() {
let concat = |a: &str, b: &str| {
format!(r#"{{"app":{{"fn":{{"const":"List.append"}},"args":[{a},{b}]}}}}"#)
};
let rev = |a: &str| format!(r#"{{"app":{{"fn":{{"const":"rev"}},"args":[{a}]}}}}"#);
let filterz = |a: &str| format!(r#"{{"app":{{"fn":{{"const":"filterZ"}},"args":[{a}]}}}}"#);
let tail = r#"{"var":"tail"}"#;
let zlist = r#"{"app":{"fn":{"const":"List.cons"},"args":[{"const":"Nat"},{"app":{"fn":{"const":"OfNat.ofNat"},"args":[{"const":"Nat"},{"nat":"0"},{"opaque":"i"}]}},{"app":{"fn":{"const":"List.nil"},"args":[{"const":"Nat"}]}}]}}"#;
let list_nat = r#"{"app":{"fn":{"const":"List"},"args":[{"const":"Nat"}]}}"#;
let ih = format!(
r#"{{"app":{{"fn":{{"const":"Eq"}},"args":[{list_nat},{},{}]}}}}"#,
rev(&filterz(tail)),
filterz(&rev(tail))
);
let claim = format!(
r#"{{"app":{{"fn":{{"const":"Eq"}},"args":[{list_nat},{},{}]}}}}"#,
concat(&rev(&filterz(tail)), zlist),
filterz(&concat(&rev(tail), zlist))
);
let json = format!(
r#"{{"forall":{{"name":"tail","ty":{list_nat},"body":{{"forall":{{"name":"ih","ty":{ih},"body":{claim}}}}}}}}}"#
);
let goal = untranslate_goal_ctx(&json, &peano_nat()).expect("in grammar");
let CalcVerdict::Lemma(l) = calculate(&goal, &nat_env(), &HashSet::new()) else {
panic!("expected a forced lemma");
};
let r = render(&l);
assert!(!r.contains("when"), "IH should be substituted away: {r}");
assert!(r.contains("given a: List<Nat>"), "{r}");
assert_eq!(
r,
"given a: List<Nat>; List.concat(filterZ(a), List.concat([Nat.Z], [])) \
=> filterZ(List.concat(a, List.concat([Nat.Z], [])))"
);
}
#[test]
fn declines_when_diff_is_not_common_or_ctor_wrap() {
let f = |a: &str| format!(r#"{{"app":{{"fn":{{"const":"le"}},"args":[{a},{a}]}}}}"#);
let ih = format!(
r#"{{"app":{{"fn":{{"const":"Eq"}},"args":[{{"const":"Bool"}},{},{{"const":"Bool.true"}}]}}}}"#,
f(r#"{"var":"x"}"#)
);
let claim = format!(
r#"{{"app":{{"fn":{{"const":"Eq"}},"args":[{{"const":"Bool"}},{},{{"const":"Bool.true"}}]}}}}"#,
f(r#"{"var":"y"}"#)
);
let list_nat = r#"{"app":{"fn":{"const":"List"},"args":[{"const":"Nat"}]}}"#;
let json = format!(
r#"{{"forall":{{"name":"x","ty":{{"const":"Nat"}},"body":{{"forall":{{"name":"ih","ty":{ih},"body":{claim}}}}}}}}}"#
);
let _ = list_nat;
let goal = untranslate_goal_ctx(&json, &peano_nat()).expect("in grammar");
match calculate(&goal, &nat_env(), &HashSet::new()) {
CalcVerdict::Decline(reason) => {
assert!(
reason.contains("outside") || reason.contains("term"),
"{reason}"
);
}
CalcVerdict::Lemma(l) => panic!("should decline, got {}", render(&l)),
}
}
}