use std::collections::{BTreeMap, BTreeSet};
use crate::ast::{TopLevel, VerifyKind};
use crate::codegen::proof_lower::{LawProofCone, ProofLowerInputs};
use crate::ir::proof_ir::ProofIR;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LemmaProvenance {
Conjectured,
Enumerated,
Recognized,
Calculated,
}
impl LemmaProvenance {
pub fn as_tag(&self) -> &'static str {
match self {
LemmaProvenance::Conjectured => "conjectured",
LemmaProvenance::Enumerated => "enumerated",
LemmaProvenance::Recognized => "recognized",
LemmaProvenance::Calculated => "calculated",
}
}
}
impl std::fmt::Display for LemmaProvenance {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_tag())
}
}
#[derive(Debug, Clone)]
pub struct CommittedLemma {
pub name: String,
pub text: String,
pub embed: bool,
pub provenance: LemmaProvenance,
}
impl CommittedLemma {
pub fn reference(name: String, text: String) -> Self {
Self {
name,
text,
embed: false,
provenance: LemmaProvenance::Enumerated,
}
}
}
pub fn parse_committed_lemmas(content: &str) -> Vec<CommittedLemma> {
let mut lemmas: Vec<CommittedLemma> = Vec::new();
let mut current: Option<CommittedLemma> = None;
for line in content.lines() {
if let Some(rest) = line.strip_prefix("theorem ") {
if let Some(mut done) = current.take() {
done.text.truncate(done.text.trim_end().len());
lemmas.push(done);
}
let name = rest
.split_whitespace()
.next()
.unwrap_or("")
.trim_end_matches(':')
.to_string();
current = Some(CommittedLemma {
name,
text: line.to_string(),
embed: true,
provenance: LemmaProvenance::Enumerated,
});
} else if let Some(block) = current.as_mut() {
block.text.push('\n');
block.text.push_str(line);
}
}
if let Some(mut done) = current.take() {
done.text.truncate(done.text.trim_end().len());
lemmas.push(done);
}
lemmas.retain(|l| !l.name.is_empty());
lemmas
}
pub fn forbidden_token_in_lemma(text: &str) -> Option<&'static str> {
const DENY: [&str; 30] = [
"axiom",
"opaque",
"unsafe",
"macro",
"macro_rules",
"notation",
"syntax",
"elab",
"attribute",
"set_option",
"instance",
"structure",
"inductive",
"class",
"def",
"abbrev",
"example",
"import",
"open",
"namespace",
"section",
"end",
"mutual",
"initialize",
"run_cmd",
"partial",
"noncomputable",
"deriving",
"theorem",
"sorry",
];
for (line_idx, line) in text.lines().enumerate() {
let code = line.split("--").next().unwrap_or("");
for (tok_idx, tok) in code
.split(|c: char| !(c.is_alphanumeric() || c == '_' || c == '.' || c == '\''))
.filter(|t| !t.is_empty())
.enumerate()
{
if line_idx == 0 && tok_idx == 0 && tok == "theorem" {
continue;
}
if let Some(hit) = DENY.iter().find(|d| **d == tok) {
return Some(hit);
}
}
}
None
}
pub fn mentioned_fns(text: &str, lean_index: &BTreeMap<String, String>) -> BTreeSet<String> {
let mut out = BTreeSet::new();
for token in text.split(|c: char| !(c.is_alphanumeric() || c == '_' || c == '.' || c == '\'')) {
if let Some(v) = lean_index.get(token) {
out.insert(v.clone());
}
}
out
}
pub fn lemma_lhs_fns(text: &str, lean_index: &BTreeMap<String, String>) -> BTreeSet<String> {
let lhs = statement_body(text)
.and_then(|stmt| {
split_after_top_eq(stmt).map(|rhs| {
let end = stmt.len() - rhs.len() - 1; stmt[..end].trim()
})
})
.unwrap_or(text);
mentioned_fns(lhs, lean_index)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SimpDirection {
Forward,
Reversed,
}
pub fn simp_orientation(text: &str, program_fns: &BTreeSet<String>) -> Option<SimpDirection> {
let stmt = statement_body(text)?;
let rhs = split_after_top_eq(stmt);
let lhs = rhs.map(|r| {
let end = stmt.len() - r.len() - 1; stmt[..end].trim()
});
let forward_grows = matches!((lhs, rhs), (Some(l), Some(r)) if !l.is_empty() && r.contains(l));
if program_fns.contains(&head_token(stmt)) && !forward_grows {
return Some(SimpDirection::Forward);
}
let rhs = rhs?;
let reversed_grows = matches!(lhs, Some(l) if !rhs.trim().is_empty() && l.contains(rhs.trim()));
if program_fns.contains(&head_token(rhs)) && !reversed_grows {
return Some(SimpDirection::Reversed);
}
None
}
pub fn simp_entries(lemmas: &[&CommittedLemma], program_fns: &BTreeSet<String>) -> Vec<String> {
let classified: Vec<(&CommittedLemma, SimpDirection)> = lemmas
.iter()
.filter_map(|l| simp_orientation(&l.text, program_fns).map(|d| (*l, d)))
.collect();
let reversed_heads: BTreeSet<String> = classified
.iter()
.filter(|(_, d)| *d == SimpDirection::Reversed)
.filter_map(|(l, _)| {
let rhs = split_after_top_eq(statement_body(&l.text)?)?;
Some(head_token(rhs))
})
.collect();
classified
.into_iter()
.filter_map(|(l, d)| match d {
SimpDirection::Forward => {
let rhs = split_after_top_eq(statement_body(&l.text)?)?;
let mentions_unfolded = rhs
.split(|c: char| !(c.is_alphanumeric() || c == '_' || c == '.' || c == '\''))
.any(|tok| reversed_heads.contains(tok));
if mentions_unfolded {
None
} else {
Some(l.name.clone())
}
}
SimpDirection::Reversed => Some(format!("← {}", l.name)),
})
.collect()
}
fn statement_body(text: &str) -> Option<&str> {
let stmt = statement_of(text)?.trim_start();
let body = if let Some(rest) = stmt.strip_prefix('∀') {
split_after_depth0(rest, ',')?
} else {
stmt
};
Some(strip_implication_premises(body))
}
fn strip_implication_premises(text: &str) -> &str {
let mut depth = 0i32;
let mut after_last_arrow = 0usize;
let bytes = text.as_bytes();
for (i, c) in text.char_indices() {
match c {
'(' | '[' | '{' => depth += 1,
')' | ']' | '}' => depth -= 1,
'-' if depth == 0 && bytes.get(i + 1) == Some(&b'>') => {
after_last_arrow = i + 2;
}
_ => {}
}
}
text[after_last_arrow..].trim_start()
}
fn head_token(text: &str) -> String {
text.chars()
.skip_while(|c| c.is_whitespace() || *c == '(')
.take_while(|c| c.is_alphanumeric() || *c == '_' || *c == '.' || *c == '\'')
.collect()
}
fn split_after_top_eq(text: &str) -> Option<&str> {
let mut depth = 0i32;
let mut prev = ' ';
let bytes = text.as_bytes();
for (i, c) in text.char_indices() {
match c {
'(' | '[' | '{' => depth += 1,
')' | ']' | '}' => depth -= 1,
'=' if depth == 0 => {
let next_eq = bytes.get(i + 1) == Some(&b'=');
if !matches!(prev, '<' | '>' | '!' | '=') && !next_eq {
return Some(&text[i + 1..]);
}
}
_ => {}
}
prev = c;
}
None
}
fn statement_of(text: &str) -> Option<&str> {
let mut depth = 0i32;
let mut start = None;
let mut prev_colon = false;
for (i, c) in text.char_indices() {
match c {
'(' | '[' | '{' => depth += 1,
')' | ']' | '}' => depth -= 1,
':' if depth == 0 && start.is_none() => {
start = Some(i + 1);
}
'=' if depth == 0 && prev_colon => {
let s = start?;
if i > s {
return Some(&text[s..i - 1]);
}
return None;
}
_ => {}
}
prev_colon = c == ':' && depth == 0;
}
None
}
fn split_after_depth0(text: &str, sep: char) -> Option<&str> {
let mut depth = 0i32;
for (i, c) in text.char_indices() {
match c {
'(' | '[' | '{' => depth += 1,
')' | ']' | '}' => depth -= 1,
c2 if c2 == sep && depth == 0 => return Some(&text[i + c.len_utf8()..]),
_ => {}
}
}
None
}
pub type SimpOverLemmaPin = (crate::ir::FnId, String, Vec<String>);
pub fn plan_simp_over_lemma_pins(
inputs: &ProofLowerInputs,
ir: &ProofIR,
lemmas: &[CommittedLemma],
) -> Vec<SimpOverLemmaPin> {
use crate::codegen::lean::aver_name_to_lean;
if lemmas.is_empty() {
return Vec::new();
}
let all_fns: BTreeMap<String, String> = inputs
.pure_fns()
.iter()
.map(|fd| {
let lean = aver_name_to_lean(&fd.name);
(lean.clone(), lean)
})
.collect();
let all_fn_names: BTreeSet<String> = all_fns.keys().cloned().collect();
let mentions: Vec<BTreeSet<String>> = lemmas
.iter()
.map(|l| mentioned_fns(&l.text, &all_fns))
.collect();
let oriented: Vec<bool> = lemmas
.iter()
.map(|l| simp_orientation(&l.text, &all_fn_names).is_some())
.collect();
let mut plan = Vec::new();
for item in inputs.entry_items {
let TopLevel::Verify(vb) = item else { continue };
let VerifyKind::Law(law) = &vb.kind else {
continue;
};
let Some(fn_id) = inputs
.symbol_table
.fn_id_of(&crate::ir::FnKey::entry(&vb.fn_name))
else {
continue;
};
let Some(thm) = ir
.law_theorems
.iter()
.find(|t| t.fn_id == fn_id && t.law_name == law.name)
else {
continue;
};
if !matches!(thm.strategy, crate::ir::ProofStrategy::Induction { .. }) {
continue;
}
let cone = LawProofCone::compute(law, &vb.fn_name, inputs);
let mut scope: BTreeSet<String> = cone
.pure_fns()
.iter()
.map(|fd| aver_name_to_lean(&fd.name))
.collect();
scope.insert(aver_name_to_lean(&vb.fn_name));
let mut any_oriented = false;
let mut selected: BTreeSet<usize> = BTreeSet::new();
for (i, (m, o)) in mentions.iter().zip(&oriented).enumerate() {
if !m.is_empty() && m.is_subset(&scope) {
selected.insert(i);
any_oriented |= *o;
}
}
if !any_oriented {
continue;
}
loop {
let added: Vec<usize> = lemmas
.iter()
.enumerate()
.filter(|(j, lj)| {
!selected.contains(j)
&& selected.iter().any(|&i| lemmas[i].text.contains(&lj.name))
})
.map(|(j, _)| j)
.collect();
if added.is_empty() {
break;
}
selected.extend(added);
}
let names: Vec<String> = selected.iter().map(|&i| lemmas[i].name.clone()).collect();
plan.push((fn_id, law.name.clone(), names));
}
plan
}
pub fn apply_simp_over_lemma_pins(ir: &mut ProofIR, plan: &[SimpOverLemmaPin]) {
for (fn_id, law_name, names) in plan {
if let Some(t) = ir
.law_theorems
.iter_mut()
.find(|t| t.fn_id == *fn_id && t.law_name == *law_name)
{
t.strategy = crate::ir::ProofStrategy::SimpOverLemmas(names.clone());
}
}
}
pub fn discovery_surface_hash(inputs: &ProofLowerInputs) -> String {
let mut sigs: Vec<String> = inputs
.pure_fns()
.iter()
.map(|fd| {
let params: Vec<String> = fd.params.iter().map(|(n, t)| format!("{n}:{t}")).collect();
format!("{}({})->{}", fd.name, params.join(","), fd.return_type)
})
.collect();
sigs.sort();
let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
for byte in sigs.join(";").bytes() {
hash ^= u64::from(byte);
hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
}
format!("{hash:016x}")
}
#[cfg(test)]
mod tests {
use super::*;
const SRC: &str = r#"
type Nat
Z
S(Nat)
fn eqNat(x: Nat, y: Nat) -> Bool
match x
Nat.Z -> match y
Nat.Z -> true
Nat.S(z) -> false
Nat.S(x2) -> match y
Nat.Z -> false
Nat.S(y2) -> eqNat(x2, y2)
fn count(x: Nat, y: List<Nat>) -> Nat
match y
[] -> Nat.Z
[z, ..ys] -> match eqNat(x, z)
true -> Nat.S(count(x, ys))
false -> count(x, ys)
fn plus(x: Nat, y: Nat) -> Nat
match x
Nat.Z -> y
Nat.S(z) -> Nat.S(plus(z, y))
fn appendNat(xs: List<Nat>, ys: List<Nat>) -> List<Nat>
List.concat(xs, ys)
fn orphan(x: Nat) -> Nat
x
verify count law countPlusConcat
given n: Nat = [Nat.Z, Nat.S(Nat.Z)]
given xs: List<Nat> = [[], [Nat.Z]]
given ys: List<Nat> = [[], [Nat.S(Nat.Z)]]
plus(count(n, xs), count(n, ys)) => count(n, appendNat(xs, ys))
"#;
const COMMITTED: &str = "-- Discovered lemmas for prop_02.av — `aver proof --discover`\n\
-- cone-hash: 00deadbeef00\n\
-- Each theorem below was discovered and kernel-proved.\n\
\n\
theorem aver_helper_succ (n : Int) : Int.natAbs (n + 1) = Int.natAbs n + 1 := by\n\
\x20 omega\n\
\n\
theorem aver_discovered_lemma_0 (x0 : List Nat) (x1 : List Nat) (x2 : Nat) : count x2 (x0 ++ x1) = plus (count x2 x0) (count x2 x1) := by\n\
\x20 induction x0 with\n\
\x20 | nil => first | (simp [count]; done) | (simp [count, aver_helper_succ]; omega)\n\
\x20 | cons head tail ih => first | (simp_all [count]; done) | (simp_all [count]; omega)\n\
\n\
theorem aver_discovered_lemma_1 (x0 : Nat) : orphan (plus x0 x0) = plus x0 x0 := by\n\
\x20 simp [orphan]\n";
fn with_inputs<R>(src: &str, f: impl FnOnce(&ProofLowerInputs) -> R) -> R {
let mut lexer = crate::lexer::Lexer::new(src);
let tokens = lexer.tokenize().expect("lex");
let mut items = crate::parser::Parser::new(tokens).parse().expect("parse");
crate::ir::pipeline::tco(&mut items);
crate::ir::pipeline::resolve(&mut items);
let symbols = crate::ir::SymbolTable::build(&items, &[]);
let prefixes: std::collections::HashSet<String> = std::collections::HashSet::new();
let recursive: std::collections::HashSet<crate::ir::FnId> =
std::collections::HashSet::new();
let no_modules: &[crate::codegen::ModuleInfo] = &[];
let inputs = ProofLowerInputs {
entry_items: &items,
dep_modules: no_modules,
module_prefixes: &prefixes,
recursive_fns: &recursive,
symbol_table: &symbols,
program_shape: None,
};
f(&inputs)
}
#[test]
fn parses_committed_theorem_blocks() {
let lemmas = parse_committed_lemmas(COMMITTED);
assert_eq!(lemmas.len(), 3);
assert_eq!(lemmas[0].name, "aver_helper_succ");
assert_eq!(lemmas[1].name, "aver_discovered_lemma_0");
assert_eq!(lemmas[2].name, "aver_discovered_lemma_1");
assert!(
lemmas[1]
.text
.starts_with("theorem aver_discovered_lemma_0 ")
);
assert!(lemmas[1].text.contains("induction x0 with"));
assert!(!lemmas[1].text.contains("aver_discovered_lemma_1"));
assert!(lemmas[2].text.ends_with("simp [orphan]"));
assert!(lemmas.iter().all(|l| !l.text.contains("cone-hash")));
}
#[test]
fn plan_pins_in_scope_lemma_and_rejects_out_of_cone() {
with_inputs(SRC, |inputs| {
let mut ir = ProofIR::default();
crate::codegen::proof_lower::populate_law_theorems(inputs, &mut ir);
assert_eq!(ir.law_theorems.len(), 1);
assert!(matches!(
ir.law_theorems[0].strategy,
crate::ir::ProofStrategy::Induction { .. }
));
let lemmas = parse_committed_lemmas(COMMITTED);
let plan = plan_simp_over_lemma_pins(inputs, &ir, &lemmas);
assert_eq!(plan.len(), 1);
assert_eq!(plan[0].1, "countPlusConcat");
assert_eq!(
plan[0].2,
vec![
"aver_helper_succ".to_string(),
"aver_discovered_lemma_0".to_string()
]
);
apply_simp_over_lemma_pins(&mut ir, &plan);
match &ir.law_theorems[0].strategy {
crate::ir::ProofStrategy::SimpOverLemmas(names) => {
assert_eq!(names.len(), 2);
}
other => panic!("expected SimpOverLemmas pin, got {other:?}"),
}
});
}
#[test]
fn simp_orientation_classifies_rewrite_direction() {
let fns: BTreeSet<String> = ["count", "plus", "appendNat", "decode", "encode"]
.iter()
.map(|s| s.to_string())
.collect();
assert_eq!(
simp_orientation(
"theorem t0 (x0 : List Nat) (x2 : Nat) : (count x2 (x0 ++ x1)) = (plus (count x2 x0) (count x2 x1)) := by\n simp",
&fns
),
Some(SimpDirection::Forward)
);
assert_eq!(
simp_orientation(
"theorem t1 (xs : List String) : decode (encode xs) = xs := by\n simp",
&fns
),
Some(SimpDirection::Forward)
);
assert_eq!(
simp_orientation(
"theorem t2 (x0 : List Nat) : (x0 ++ x0) = (appendNat x0 x0) := by\n simp",
&fns
),
Some(SimpDirection::Reversed)
);
assert_eq!(
simp_orientation(
"theorem t3 : ∀ (list : List Int) (acc : Int), plus list acc = acc := by\n simp",
&fns
),
Some(SimpDirection::Forward)
);
assert_eq!(
simp_orientation(
"theorem t4 (acc : Acc) (x : Int) : 0 <= (count acc x) := by\n simp",
&fns
),
None
);
assert_eq!(
simp_orientation(
"theorem t5 (x0 : List Nat) : ((x0 ++ x0) ++ x0) = (x0 ++ (x0 ++ x0)) := by\n simp",
&fns
),
None
);
let dfns: BTreeSet<String> = ["dbl", "idNat"].iter().map(|s| s.to_string()).collect();
assert_eq!(
simp_orientation(
"theorem t6 (x : Nat) : dbl x = idNat (dbl x) := by\n simp",
&dfns
),
Some(SimpDirection::Reversed)
);
let efns: BTreeSet<String> = ["loopy"].iter().map(|s| s.to_string()).collect();
assert_eq!(
simp_orientation(
"theorem t7 (x : Nat) : loopy x = loopy x := by\n rfl",
&efns
),
None
);
assert_eq!(
simp_orientation(
"theorem t8 (a b : Nat) (xs : List Nat) : natEq a b = true -> count a (xs ++ [b]) = count a xs := by\n simp",
&fns
),
Some(SimpDirection::Forward)
);
}
#[test]
fn forbidden_tokens_reject_smuggled_declarations() {
let lemmas = parse_committed_lemmas(COMMITTED);
assert!(
lemmas
.iter()
.all(|l| forbidden_token_in_lemma(&l.text).is_none()),
"discovery-emitted blocks must validate clean"
);
assert_eq!(
forbidden_token_in_lemma("theorem t : True := by\n trivial\naxiom cheat : False"),
Some("axiom")
);
assert_eq!(
forbidden_token_in_lemma("theorem t : True := by\n trivial\n set_option foo true"),
Some("set_option")
);
assert_eq!(
forbidden_token_in_lemma("theorem t : P := by\n first | simp | sorry"),
Some("sorry")
);
assert_eq!(
forbidden_token_in_lemma("theorem t : True := by\n trivial -- no axiom here"),
None
);
assert_eq!(
forbidden_token_in_lemma("theorem t : True := by\n trivial\n theorem u : True"),
Some("theorem")
);
}
#[test]
fn plan_ignores_lemmas_with_no_program_connection() {
with_inputs(SRC, |inputs| {
let mut ir = ProofIR::default();
crate::codegen::proof_lower::populate_law_theorems(inputs, &mut ir);
let lemmas = vec![CommittedLemma {
name: "free_floating".to_string(),
text: "theorem free_floating (a : Nat) : a + 0 = a := by simp".to_string(),
embed: true,
provenance: LemmaProvenance::Enumerated,
}];
assert!(plan_simp_over_lemma_pins(inputs, &ir, &lemmas).is_empty());
});
}
}