use std::collections::{HashMap, HashSet};
use crate::ast::{Expr, FnDef, Literal, Spanned, TopLevel, TypeDef};
use crate::codegen::common::expr_to_dotted_name;
use crate::codegen::recursion::RecursionPlan;
use crate::codegen::{CodegenContext, ModuleInfo};
use crate::ir::proof_ir::{
DecreaseProof, FnContract, Measure, NativeIntCountdownBody, Predicate, PreservationProof,
ProofIR, QuantifierType, RecursionContract, RefinedTypeDecl,
};
pub struct ProofLowerInputs<'a> {
pub entry_items: &'a [TopLevel],
pub dep_modules: &'a [ModuleInfo],
pub module_prefixes: &'a HashSet<String>,
pub recursive_fns: &'a HashSet<crate::ir::FnId>,
pub symbol_table: &'a crate::ir::SymbolTable,
pub program_shape: Option<&'a crate::analysis::shape::ProgramShape>,
}
impl<'a> ProofLowerInputs<'a> {
pub fn from_ctx(ctx: &'a CodegenContext) -> Self {
Self {
entry_items: &ctx.items,
dep_modules: &ctx.modules,
module_prefixes: &ctx.module_prefixes,
recursive_fns: &ctx.recursive_fns,
symbol_table: &ctx.symbol_table,
program_shape: ctx.program_shape.as_ref(),
}
}
pub fn pure_fns(&self) -> Vec<&'a FnDef> {
self.dep_modules
.iter()
.flat_map(|m| m.fn_defs.iter())
.chain(self.entry_items.iter().filter_map(|item| match item {
TopLevel::FnDef(fd) => Some(fd),
_ => None,
}))
.filter(|fd| crate::codegen::common::is_pure_fn(fd))
.collect()
}
pub fn recursive_pure_fn_names(&self) -> HashSet<String> {
let symbols = self.symbol_table;
let pure_ids: HashSet<crate::ir::FnId> = self
.pure_fns()
.into_iter()
.filter_map(|fd| {
let scope = self
.dep_modules
.iter()
.find(|m| m.fn_defs.iter().any(|d| std::ptr::eq(d, fd)))
.map(|m| m.prefix.as_str());
let key = match scope {
Some(prefix) => crate::ir::FnKey::in_module(prefix.to_string(), &fd.name),
None => crate::ir::FnKey::entry(&fd.name),
};
symbols.fn_id_of(&key)
})
.collect();
self.recursive_fns
.intersection(&pure_ids)
.map(|id| symbols.fn_entry(*id).key.name.clone())
.collect()
}
pub fn pure_fns_in_scope(&self, scope: Option<&str>) -> Vec<&'a FnDef> {
match scope {
None => self
.entry_items
.iter()
.filter_map(|item| match item {
TopLevel::FnDef(fd) => Some(fd),
_ => None,
})
.filter(|fd| crate::codegen::common::is_pure_fn(fd))
.collect(),
Some(prefix) => self
.dep_modules
.iter()
.filter(|m| m.prefix == prefix)
.flat_map(|m| m.fn_defs.iter())
.filter(|fd| crate::codegen::common::is_pure_fn(fd))
.collect(),
}
}
pub fn recursive_pure_fn_names_in_scope(&self, scope: Option<&str>) -> HashSet<String> {
let symbols = self.symbol_table;
let pure_ids: HashSet<crate::ir::FnId> = self
.pure_fns_in_scope(scope)
.into_iter()
.filter_map(|fd| {
let key = match scope {
Some(prefix) => crate::ir::FnKey::in_module(prefix.to_string(), &fd.name),
None => crate::ir::FnKey::entry(&fd.name),
};
symbols.fn_id_of(&key)
})
.collect();
self.recursive_fns
.intersection(&pure_ids)
.map(|id| symbols.fn_entry(*id).key.name.clone())
.collect()
}
pub fn scopes(&self) -> Vec<Option<String>> {
let mut out = vec![None];
for m in self.dep_modules {
out.push(Some(m.prefix.clone()));
}
out
}
pub fn fn_owning_scope(&self, fd: &FnDef) -> Option<&'a str> {
for m in self.dep_modules {
for f in &m.fn_defs {
if std::ptr::eq(f, fd) {
return Some(m.prefix.as_str());
}
}
}
None
}
pub fn resolve_expr(
&self,
expr: &crate::ast::Spanned<crate::ast::Expr>,
scope: Option<&str>,
) -> crate::ast::Spanned<crate::ir::hir::ResolvedExpr> {
use crate::ir::hir::{ResolveCtx, ResolvedStmt};
let mut rctx = ResolveCtx::new(self.symbol_table);
rctx.current_module = scope.map(String::from);
let stmt = crate::ast::Stmt::Expr(expr.clone());
match crate::ir::hir::resolve::resolve_stmt_external(&rctx, &stmt) {
ResolvedStmt::Expr(s) => s,
ResolvedStmt::Binding { value, .. } => value,
}
}
pub fn recursive_type_names(&self) -> HashSet<String> {
self.entry_items
.iter()
.filter_map(|item| match item {
TopLevel::TypeDef(td) => Some(td),
_ => None,
})
.chain(self.dep_modules.iter().flat_map(|m| m.type_defs.iter()))
.filter(|td| crate::codegen::common::is_recursive_type_def(td))
.map(|td| crate::codegen::common::type_def_name(td).to_string())
.collect()
}
pub fn find_fn_def_by_call_name(&self, call_name: &str) -> Option<&'a FnDef> {
let find_exact = |name: &str| -> Option<&'a FnDef> {
self.dep_modules
.iter()
.flat_map(|m| m.fn_defs.iter())
.chain(self.entry_items.iter().filter_map(|item| match item {
TopLevel::FnDef(fd) => Some(fd),
_ => None,
}))
.find(|fd| fd.name == name)
};
find_exact(call_name).or_else(|| {
let short = call_name.rsplit('.').next()?;
find_exact(short)
})
}
pub fn find_type_def(&self, type_name: &str) -> Option<&'a TypeDef> {
self.entry_items
.iter()
.filter_map(|item| match item {
TopLevel::TypeDef(td) => Some(td),
_ => None,
})
.chain(self.dep_modules.iter().flat_map(|m| m.type_defs.iter()))
.find(|td| crate::codegen::common::type_def_name(td) == type_name)
}
}
pub fn lower(inputs: &ProofLowerInputs) -> ProofIR {
let mut ir = ProofIR::default();
populate_refined_types(inputs, &mut ir);
populate_fn_contracts(inputs, &mut ir);
populate_law_theorems(inputs, &mut ir);
ir
}
pub fn populate_refined_types(inputs: &ProofLowerInputs, ir: &mut ProofIR) {
let symbols = inputs.symbol_table;
let entry_typedefs = inputs.entry_items.iter().filter_map(|item| match item {
TopLevel::TypeDef(td) => Some((None::<&str>, td)),
_ => None,
});
let module_typedefs = inputs.dep_modules.iter().flat_map(|m| {
m.type_defs
.iter()
.map(move |td| (Some(m.prefix.as_str()), td))
});
for (module_prefix, td) in entry_typedefs.chain(module_typedefs) {
let TypeDef::Product { name, fields, .. } = td else {
continue;
};
if fields.len() != 1 {
continue;
}
let type_key = match module_prefix {
Some(prefix) => crate::ir::TypeKey::in_module(prefix.to_string(), name),
None => crate::ir::TypeKey::entry(name),
};
let Some(canonical_key) = symbols.type_id_of(&type_key) else {
continue;
};
if ir.refined_types.contains_key(&canonical_key) {
continue;
}
let Some(info) =
crate::codegen::common::refinement_info_for_in_scope(name, inputs, module_prefix)
else {
continue;
};
let invariant = Predicate {
free_vars: vec![(
info.param_name.to_string(),
crate::ir::proof_ir::QuantifierType::Plain(info.carrier_type.to_string()),
)],
expr: inputs.resolve_expr(info.predicate, module_prefix),
};
let witness = pick_witness(
name,
canonical_key,
inputs,
info.predicate,
info.param_name,
module_prefix,
);
let Some(witness) = witness else {
continue;
};
ir.refined_types.insert(
canonical_key,
RefinedTypeDecl {
name: name.clone(),
carrier_type: info.carrier_type.to_string(),
carrier_field: info.carrier_field.to_string(),
predicate_param: info.param_name.to_string(),
invariant,
witness: Some(witness),
},
);
}
}
pub fn populate_fn_contracts(inputs: &ProofLowerInputs, ir: &mut ProofIR) {
for scope in inputs.scopes() {
let (plans, issues) =
crate::codegen::recursion::analyze_plans_in_scope(inputs, scope.as_deref(), false);
ir.unclassified_fns
.extend(issues.into_iter().map(|issue| crate::ir::UnclassifiedFn {
line: issue.line,
message: issue.message,
}));
populate_fn_contracts_for_scope(inputs, ir, scope.as_deref(), &plans);
}
}
fn populate_fn_contracts_for_scope(
inputs: &ProofLowerInputs,
ir: &mut ProofIR,
scope: Option<&str>,
plans: &HashMap<String, RecursionPlan>,
) {
let scoped_fns: Vec<&FnDef> = inputs.pure_fns_in_scope(scope);
let qualify = |bare: &str| -> crate::ir::FnKey {
match scope {
Some(prefix) => crate::ir::FnKey::in_module(prefix.to_string(), bare),
None => crate::ir::FnKey::entry(bare),
}
};
let symbols = inputs.symbol_table;
for (fn_name, plan) in plans {
let Some(fd) = scoped_fns.iter().find(|fd| fd.name == *fn_name) else {
continue;
};
let fn_key = qualify(fn_name);
let Some(canonical_key) = symbols.fn_id_of(&fn_key) else {
continue;
};
if let RecursionPlan::IntCountdown { param_index } = plan {
if let Some((param_name, _)) = fd.params.get(*param_index) {
ir.fn_contracts.insert(
canonical_key,
FnContract {
source_name: fn_name.clone(),
recursion: Some(RecursionContract::Fuel {
fuel_metric: crate::ir::FuelMetric::NatAbsPlusOne {
param: param_name.clone(),
},
}),
},
);
}
continue;
}
if let RecursionPlan::IntFloorDivCountdown {
param_index,
divisor,
helper_fn,
} = plan
{
if let Some((param_name, _)) = fd.params.get(*param_index) {
ir.fn_contracts.insert(
canonical_key,
FnContract {
source_name: fn_name.clone(),
recursion: Some(RecursionContract::WellFoundedToNat {
param: param_name.clone(),
floor_div: Some(crate::ir::FloorDivShrink {
divisor: *divisor,
helper_fn: helper_fn.clone(),
}),
}),
},
);
}
continue;
}
if let RecursionPlan::IntAscending { param_index, bound } = plan {
if let Some((param_name, _)) = fd.params.get(*param_index) {
ir.fn_contracts.insert(
canonical_key,
FnContract {
source_name: fn_name.clone(),
recursion: Some(RecursionContract::Fuel {
fuel_metric: crate::ir::FuelMetric::BoundMinusParamNatAbsPlusOne {
param: param_name.clone(),
bound: inputs.resolve_expr(bound, scope),
},
}),
},
);
}
continue;
}
if let RecursionPlan::ListStructural { param_index } = plan {
if let Some((param_name, _)) = fd.params.get(*param_index) {
ir.fn_contracts.insert(
canonical_key,
FnContract {
source_name: fn_name.clone(),
recursion: Some(RecursionContract::Fuel {
fuel_metric: crate::ir::FuelMetric::SeqLenPlusOne {
param: param_name.clone(),
},
}),
},
);
}
continue;
}
if matches!(plan, RecursionPlan::SizeOfStructural) {
ir.fn_contracts.insert(
canonical_key,
FnContract {
source_name: fn_name.clone(),
recursion: Some(RecursionContract::Fuel {
fuel_metric: crate::ir::FuelMetric::SizeOfPlusOne,
}),
},
);
continue;
}
if matches!(plan, RecursionPlan::StringPosAdvance) {
if let (Some((string_param, _)), Some((pos_param, _))) =
(fd.params.first(), fd.params.get(1))
{
ir.fn_contracts.insert(
canonical_key,
FnContract {
source_name: fn_name.clone(),
recursion: Some(RecursionContract::Fuel {
fuel_metric: crate::ir::FuelMetric::StringLenMinusPos {
string_param: string_param.clone(),
pos_param: pos_param.clone(),
},
}),
},
);
}
continue;
}
match plan {
RecursionPlan::MutualIntCountdown => {
let params = fd
.params
.first()
.map(|(n, _)| vec![n.clone()])
.unwrap_or_default();
ir.fn_contracts.insert(
canonical_key,
FnContract {
source_name: fn_name.clone(),
recursion: Some(RecursionContract::Fuel {
fuel_metric: crate::ir::FuelMetric::Lex { params, rank: 0 },
}),
},
);
continue;
}
RecursionPlan::MutualStringPosAdvance { rank } => {
let params = fd.params.iter().take(2).map(|(n, _)| n.clone()).collect();
ir.fn_contracts.insert(
canonical_key,
FnContract {
source_name: fn_name.clone(),
recursion: Some(RecursionContract::Fuel {
fuel_metric: crate::ir::FuelMetric::Lex {
params,
rank: *rank,
},
}),
},
);
continue;
}
RecursionPlan::MutualSizeOfRanked { rank } => {
ir.fn_contracts.insert(
canonical_key,
FnContract {
source_name: fn_name.clone(),
recursion: Some(RecursionContract::Fuel {
fuel_metric: crate::ir::FuelMetric::Lex {
params: Vec::new(),
rank: *rank,
},
}),
},
);
continue;
}
RecursionPlan::LinearRecurrence2 => {
ir.fn_contracts.insert(
canonical_key,
FnContract {
source_name: fn_name.clone(),
recursion: Some(RecursionContract::LinearRecurrence2),
},
);
continue;
}
_ => {}
}
let RecursionPlan::IntCountdownGuarded {
param_index,
base_arm_literal,
base_arm_body,
wildcard_arm_body,
precondition,
} = plan
else {
continue;
};
let Some((countdown_param_name, _)) = fd.params.get(*param_index) else {
continue;
};
let precondition_predicates: Vec<Predicate> = precondition
.iter()
.map(|clause| Predicate {
free_vars: vec![(
countdown_param_name.clone(),
QuantifierType::Plain("Int".to_string()),
)],
expr: inputs.resolve_expr(clause, scope),
})
.collect();
ir.fn_contracts.insert(
canonical_key,
FnContract {
source_name: fn_name.clone(),
recursion: Some(RecursionContract::Native {
precondition: precondition_predicates,
measure: Measure::NatAbsInt {
param: countdown_param_name.clone(),
},
preservation: PreservationProof::IntCountdownLiteralZero,
decrease: DecreaseProof::NatAbsCountdown,
body: NativeIntCountdownBody {
base_arm_literal: *base_arm_literal,
base_arm_body: inputs.resolve_expr(base_arm_body, scope),
wildcard_arm_body: inputs.resolve_expr(wildcard_arm_body, scope),
},
}),
},
);
}
}
pub fn populate_law_theorems(inputs: &ProofLowerInputs, ir: &mut ProofIR) {
use crate::ast::{TopLevel, VerifyKind};
use crate::ir::{LawTheorem, Predicate, Quantifier, QuantifierType};
let symbols = inputs.symbol_table;
let entry_verifies = inputs.entry_items.iter().filter_map(|item| match item {
TopLevel::Verify(vb) => Some(vb),
_ => None,
});
for vb in entry_verifies {
let VerifyKind::Law(law) = &vb.kind else {
continue;
};
let quantifiers: Vec<Quantifier> = law
.givens
.iter()
.map(|g| Quantifier {
name: g.name.clone(),
binder_type: QuantifierType::Plain(g.type_name.clone()),
})
.collect();
let law_scope: Option<String> = symbols
.fn_id_of(&crate::ir::FnKey::entry(&vb.fn_name))
.or_else(|| {
inputs.dep_modules.iter().find_map(|m| {
symbols.fn_id_of(&crate::ir::FnKey::in_module(m.prefix.clone(), &vb.fn_name))
})
})
.and_then(|id| symbols.fn_entry(id).key.scope_str().map(|s| s.to_string()));
let law_scope_ref = law_scope.as_deref();
let premises: Vec<Predicate> = match &law.when {
Some(when_expr) => vec![Predicate {
free_vars: quantifiers
.iter()
.map(|q| (q.name.clone(), q.binder_type.clone()))
.collect(),
expr: inputs.resolve_expr(when_expr, law_scope_ref),
}],
None => Vec::new(),
};
let strategy = classify_law_strategy(
law,
&vb.fn_name,
inputs,
&ir.refined_types,
&ir.fn_contracts,
law_scope_ref,
);
let Some(fn_id) = symbols.fn_id_of(&crate::ir::FnKey::entry(&vb.fn_name)) else {
continue;
};
ir.law_theorems.push(LawTheorem {
fn_id,
law_name: law.name.clone(),
quantifiers,
premises,
claim_lhs: inputs.resolve_expr(&law.lhs, law_scope_ref),
claim_rhs: inputs.resolve_expr(&law.rhs, law_scope_ref),
strategy,
});
}
let window_pow_fns: HashSet<String> = ir
.law_theorems
.iter()
.filter_map(|t| match &t.strategy {
crate::ir::ProofStrategy::FloorDivWindow { figure } => Some(match figure {
crate::ir::FloorWindowFigure::PowPositive { pow_fn } => pow_fn.clone(),
crate::ir::FloorWindowFigure::PowSumSplit { pow_fn } => pow_fn.clone(),
crate::ir::FloorWindowFigure::SigWindow { pow_fn, .. } => pow_fn.clone(),
crate::ir::FloorWindowFigure::ProductWindow { pow_fn, .. } => pow_fn.clone(),
}),
_ => None,
})
.collect();
for pow_fn in window_pow_fns {
let Some(fn_id) = symbols.fn_id_of(&crate::ir::FnKey::entry(&pow_fn)) else {
continue;
};
let Some(contract) = ir.fn_contracts.get_mut(&fn_id) else {
continue;
};
if let Some(crate::ir::RecursionContract::Fuel {
fuel_metric: crate::ir::FuelMetric::NatAbsPlusOne { param },
}) = &contract.recursion
{
contract.recursion = Some(crate::ir::RecursionContract::WellFoundedToNat {
param: param.clone(),
floor_div: None,
});
}
}
}
fn classify_law_strategy(
law: &crate::ast::VerifyLaw,
fn_name: &str,
inputs: &ProofLowerInputs,
refined_types: &std::collections::HashMap<crate::ir::TypeId, crate::ir::RefinedTypeDecl>,
fn_contracts: &std::collections::HashMap<crate::ir::FnId, crate::ir::FnContract>,
scope: Option<&str>,
) -> crate::ir::ProofStrategy {
use crate::ir::ProofStrategy;
if law.when.is_none()
&& let Some(s) = detect_match_dispatcher_fold_equivalence(law, fn_name, inputs)
{
return s;
}
if law.when.is_none()
&& let Some(s) = detect_result_pipeline_chain_equivalence(law, fn_name, inputs)
{
return s;
}
if law.when.is_none()
&& let Some(s) = detect_wrapper_over_recursion(law, fn_name, inputs)
{
return s;
}
if law.when.is_none()
&& let Some(param) = detect_induction_target(law, inputs)
{
return ProofStrategy::Induction { param };
}
if law.lhs == law.rhs {
return ProofStrategy::Reflexive;
}
if let Some(op) = wrapper_binop(fn_name, inputs) {
if detect_wrapper_commutative(law, fn_name, op) {
return ProofStrategy::Commutative { op };
}
if detect_wrapper_associative(law, fn_name, op) {
return ProofStrategy::Associative { op };
}
if detect_wrapper_identity(law, fn_name, op) {
return ProofStrategy::IdentityElement { op };
}
if matches!(op, crate::ast::BinOp::Sub) && detect_wrapper_sub_right_identity(law, fn_name) {
return ProofStrategy::IdentityElement { op };
}
if matches!(op, crate::ast::BinOp::Sub)
&& let Some(neg_on_rhs) = detect_wrapper_sub_anti_commutative(law, fn_name)
{
return ProofStrategy::AntiCommutative { op, neg_on_rhs };
}
}
if let Some(inner_fn) = detect_wrapper_unary_equivalence(law, fn_name, inputs) {
return ProofStrategy::UnaryEqualsBinary { inner_fn };
}
if let Some((axiom, args)) = detect_map_set_axiom(law) {
let resolved_args: Vec<_> = args.iter().map(|a| inputs.resolve_expr(a, scope)).collect();
return ProofStrategy::LibraryAxiom {
axiom,
args: resolved_args,
};
}
if let Some(inc) = detect_map_key_tracked_increment(law, fn_name, inputs) {
return ProofStrategy::MapKeyTrackedIncrement {
outer_fn: inc.outer_fn,
map_arg: inputs.resolve_expr(&inc.map_arg, scope),
key_arg: inputs.resolve_expr(&inc.key_arg, scope),
};
}
if let Some(post) = detect_map_update_postcondition(law, fn_name, inputs) {
return ProofStrategy::MapUpdatePostcondition {
outer_fn: post.outer_fn,
kind: post.kind,
map_arg: inputs.resolve_expr(&post.map_arg, scope),
key_arg: inputs.resolve_expr(&post.key_arg, scope),
extra_unfolds: post.extra_unfolds,
};
}
if let Some(extra_unfolds) = detect_spec_equivalence(law, fn_name, inputs) {
return ProofStrategy::SpecEquivalence { extra_unfolds };
}
if let Some(extra_unfolds) = detect_simp_normalized_spec_equivalence(law, fn_name, inputs) {
return ProofStrategy::SpecEquivalenceSimpNormalized { extra_unfolds };
}
if let Some((unfolded_impl, unfolded_spec)) =
detect_linear_int_spec_equivalence(law, fn_name, inputs)
{
return ProofStrategy::LinearIntSpecEquivalence {
unfolded_impl: inputs.resolve_expr(&unfolded_impl, scope),
unfolded_spec: inputs.resolve_expr(&unfolded_spec, scope),
};
}
if let Some(spec_fn) = detect_effectful_spec_equivalence(law, fn_name, inputs) {
return ProofStrategy::EffectfulSpecEquivalence {
impl_fn: fn_name.to_string(),
spec_fn,
};
}
if let Some((spec_fn, helper_fn)) =
detect_linear_recurrence2_spec_equivalence(law, fn_name, inputs)
{
return ProofStrategy::LinearRecurrence2SpecEquivalence {
impl_fn: fn_name.to_string(),
spec_fn,
helper_fn,
};
}
if let Some(plan) = detect_simp_omega_unfold(law, fn_name, inputs, refined_types) {
return ProofStrategy::LinearArithmetic {
unfold_fns: plan.unfold_fns,
wrapper_return: plan.wrapper_return,
smart_guard: plan.smart_guard,
lifted: plan.lifted,
};
}
if law.when.is_none()
&& let Some(unfold_fns) = detect_enum_constant_fold(law, fn_name, inputs)
{
return ProofStrategy::EnumConstantFold { unfold_fns };
}
if law.when.is_none()
&& let Some(givens) = detect_finite_domain_cases(law, inputs)
{
return ProofStrategy::FiniteDomainCases { givens };
}
if law.when.is_none()
&& let Some(unfold_fns) = detect_ring_identity(law, fn_name, inputs)
{
return ProofStrategy::RingIdentity { unfold_fns };
}
if law.when.is_none()
&& let Some(s) = detect_int_decimal_roundtrip(law, fn_name, inputs, fn_contracts)
{
return s;
}
if law.when.is_none()
&& let Some(s) = detect_string_escape_roundtrip(law, inputs, fn_contracts)
{
return s;
}
if let Some(figure) = detect_floor_window(law, fn_name, inputs, fn_contracts) {
return ProofStrategy::FloorDivWindow { figure };
}
if law.when.is_none()
&& let Some(s) = detect_simp_over_prelude_lemmas(law, fn_name, inputs, fn_contracts)
{
return s;
}
ProofStrategy::BackendDispatch
}
mod finite_domain;
mod floor_window;
mod induction;
mod int_decimal_roundtrip;
mod map_laws;
mod refinement;
mod ring;
mod simp;
mod spec_equivalence;
mod string_escape_roundtrip;
mod wrapper_laws;
pub(crate) use induction::LawProofCone;
use finite_domain::*;
use floor_window::*;
use induction::*;
use int_decimal_roundtrip::*;
use map_laws::*;
use refinement::*;
use ring::*;
use simp::*;
use spec_equivalence::*;
use string_escape_roundtrip::*;
use wrapper_laws::*;