pub mod detect;
use std::collections::HashSet;
use crate::ast::{Expr, FnBody, MatchArm, Spanned, Stmt, StrPart, TailCallData};
use crate::codegen::common::expr_to_dotted_name;
pub use detect::analyze_plans_in_scope;
#[derive(Clone, Debug, PartialEq)]
pub enum RecursionPlan {
IntCountdown { param_index: usize },
IntCountdownGuarded {
param_index: usize,
base_arm_literal: i64,
base_arm_body: Spanned<Expr>,
wildcard_arm_body: Spanned<Expr>,
precondition: Vec<Spanned<Expr>>,
},
IntAscending {
param_index: usize,
bound: Spanned<Expr>,
},
LinearRecurrence2,
ListStructural { param_index: usize },
SizeOfStructural,
StringPosAdvance,
MutualIntCountdown,
MutualStringPosAdvance { rank: usize },
MutualSizeOfRanked { rank: usize },
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ProofModeIssue {
pub line: usize,
pub message: String,
}
pub fn fuel_helper_name(name: &str) -> String {
format!("{}__fuel", name)
}
pub fn native_aux_name(name: &str) -> String {
format!("{}__aux", name)
}
pub const OMEGA_PROOF_SENTINEL: &str = "__aver_omega_proof__";
pub fn flip_comparison_binop(expr: &Spanned<Expr>) -> Option<Spanned<Expr>> {
let Expr::BinOp(op, left, right) = &expr.node else {
return None;
};
use crate::ast::BinOp::*;
let flipped = match op {
Lt => Gte,
Gt => Lte,
Lte => Gt,
Gte => Lt,
Eq => Neq,
Neq => Eq,
_ => return None,
};
Some(Spanned::new(
Expr::BinOp(flipped, left.clone(), right.clone()),
expr.line,
))
}
pub use crate::codegen::common::substitute_ident_in_expr;
pub fn expr_references_ident(expr: &Spanned<Expr>, name: &str) -> bool {
match &expr.node {
Expr::Ident(n) | Expr::Resolved { name: n, .. } => n == name,
Expr::Literal(_) => false,
Expr::Attr(obj, _) => expr_references_ident(obj, name),
Expr::FnCall(callee, args) => {
expr_references_ident(callee, name)
|| args.iter().any(|a| expr_references_ident(a, name))
}
Expr::BinOp(_, l, r) => expr_references_ident(l, name) || expr_references_ident(r, name),
Expr::Neg(inner) => expr_references_ident(inner, name),
Expr::Match { subject, arms } => {
expr_references_ident(subject, name)
|| arms.iter().any(|a| expr_references_ident(&a.body, name))
}
Expr::Constructor(_, arg) => arg
.as_deref()
.is_some_and(|a| expr_references_ident(a, name)),
Expr::ErrorProp(inner) => expr_references_ident(inner, name),
Expr::InterpolatedStr(parts) => parts.iter().any(|p| match p {
StrPart::Parsed(inner) => expr_references_ident(inner, name),
_ => false,
}),
Expr::List(items) | Expr::Tuple(items) | Expr::IndependentProduct(items, _) => {
items.iter().any(|i| expr_references_ident(i, name))
}
Expr::MapLiteral(entries) => entries
.iter()
.any(|(k, v)| expr_references_ident(k, name) || expr_references_ident(v, name)),
Expr::RecordCreate { fields, .. } => {
fields.iter().any(|(_, v)| expr_references_ident(v, name))
}
Expr::RecordUpdate { base, updates, .. } => {
expr_references_ident(base, name)
|| updates.iter().any(|(_, v)| expr_references_ident(v, name))
}
Expr::TailCall(boxed) => boxed.args.iter().any(|a| expr_references_ident(a, name)),
}
}
pub fn rewrite_recursive_calls_expr(
expr: &Spanned<Expr>,
targets: &HashSet<String>,
fuel_var: &str,
) -> Spanned<Expr> {
let line = expr.line;
let new_node = match &expr.node {
Expr::Literal(_) | Expr::Ident(_) | Expr::Resolved { .. } => return expr.clone(),
Expr::Attr(obj, field) => Expr::Attr(
Box::new(rewrite_recursive_calls_expr(obj, targets, fuel_var)),
field.clone(),
),
Expr::FnCall(callee, args) => {
let rewritten_args: Vec<Spanned<Expr>> = args
.iter()
.map(|arg| rewrite_recursive_calls_expr(arg, targets, fuel_var))
.collect();
if let Some(name) = expr_to_dotted_name(&callee.node)
&& targets.contains(&name)
{
let mut call_args = Vec::with_capacity(rewritten_args.len() + 1);
call_args.push(Spanned::new(Expr::Ident(fuel_var.to_string()), line));
call_args.extend(rewritten_args);
Expr::FnCall(
Box::new(Spanned::new(Expr::Ident(fuel_helper_name(&name)), line)),
call_args,
)
} else {
Expr::FnCall(
Box::new(rewrite_recursive_calls_expr(callee, targets, fuel_var)),
rewritten_args,
)
}
}
Expr::BinOp(op, left, right) => Expr::BinOp(
*op,
Box::new(rewrite_recursive_calls_expr(left, targets, fuel_var)),
Box::new(rewrite_recursive_calls_expr(right, targets, fuel_var)),
),
Expr::Neg(inner) => Expr::Neg(Box::new(rewrite_recursive_calls_expr(
inner, targets, fuel_var,
))),
Expr::Match { subject, arms } => Expr::Match {
subject: Box::new(rewrite_recursive_calls_expr(subject, targets, fuel_var)),
arms: arms
.iter()
.map(|arm| MatchArm {
pattern: arm.pattern.clone(),
body: Box::new(rewrite_recursive_calls_expr(&arm.body, targets, fuel_var)),
binding_slots: std::sync::OnceLock::new(),
})
.collect(),
},
Expr::Constructor(name, arg) => Expr::Constructor(
name.clone(),
arg.as_ref()
.map(|inner| Box::new(rewrite_recursive_calls_expr(inner, targets, fuel_var))),
),
Expr::ErrorProp(inner) => Expr::ErrorProp(Box::new(rewrite_recursive_calls_expr(
inner, targets, fuel_var,
))),
Expr::InterpolatedStr(parts) => Expr::InterpolatedStr(
parts
.iter()
.map(|part| match part {
StrPart::Literal(_) => part.clone(),
StrPart::Parsed(inner) => StrPart::Parsed(Box::new(
rewrite_recursive_calls_expr(inner, targets, fuel_var),
)),
})
.collect(),
),
Expr::List(items) => Expr::List(
items
.iter()
.map(|item| rewrite_recursive_calls_expr(item, targets, fuel_var))
.collect(),
),
Expr::Tuple(items) => Expr::Tuple(
items
.iter()
.map(|item| rewrite_recursive_calls_expr(item, targets, fuel_var))
.collect(),
),
Expr::IndependentProduct(items, flag) => Expr::IndependentProduct(
items
.iter()
.map(|item| rewrite_recursive_calls_expr(item, targets, fuel_var))
.collect(),
*flag,
),
Expr::MapLiteral(entries) => Expr::MapLiteral(
entries
.iter()
.map(|(k, v)| {
(
rewrite_recursive_calls_expr(k, targets, fuel_var),
rewrite_recursive_calls_expr(v, targets, fuel_var),
)
})
.collect(),
),
Expr::RecordCreate { type_name, fields } => Expr::RecordCreate {
type_name: type_name.clone(),
fields: fields
.iter()
.map(|(name, value)| {
(
name.clone(),
rewrite_recursive_calls_expr(value, targets, fuel_var),
)
})
.collect(),
},
Expr::RecordUpdate {
type_name,
base,
updates,
} => Expr::RecordUpdate {
type_name: type_name.clone(),
base: Box::new(rewrite_recursive_calls_expr(base, targets, fuel_var)),
updates: updates
.iter()
.map(|(name, value)| {
(
name.clone(),
rewrite_recursive_calls_expr(value, targets, fuel_var),
)
})
.collect(),
},
Expr::TailCall(boxed) => {
let TailCallData { target, args, .. } = boxed.as_ref();
let rewritten_args: Vec<Spanned<Expr>> = args
.iter()
.map(|arg| rewrite_recursive_calls_expr(arg, targets, fuel_var))
.collect();
if targets.contains(target) {
let mut call_args = Vec::with_capacity(rewritten_args.len() + 1);
call_args.push(Spanned::new(Expr::Ident(fuel_var.to_string()), line));
call_args.extend(rewritten_args);
Expr::FnCall(
Box::new(Spanned::new(Expr::Ident(fuel_helper_name(target)), line)),
call_args,
)
} else {
Expr::TailCall(Box::new(TailCallData::new(target.clone(), rewritten_args)))
}
}
};
Spanned::new(new_node, line)
}
pub fn rewrite_native_guarded_calls_expr(
expr: &Spanned<Expr>,
fn_name: &str,
aux_name: &str,
) -> Spanned<Expr> {
let line = expr.line;
let new_node = match &expr.node {
Expr::Literal(_) | Expr::Ident(_) | Expr::Resolved { .. } => return expr.clone(),
Expr::Attr(obj, field) => Expr::Attr(
Box::new(rewrite_native_guarded_calls_expr(obj, fn_name, aux_name)),
field.clone(),
),
Expr::FnCall(callee, args) => {
let rewritten_args: Vec<Spanned<Expr>> = args
.iter()
.map(|arg| rewrite_native_guarded_calls_expr(arg, fn_name, aux_name))
.collect();
if let Some(name) = expr_to_dotted_name(&callee.node)
&& name == fn_name
{
let mut call_args = rewritten_args;
call_args.push(Spanned::new(
Expr::Ident(OMEGA_PROOF_SENTINEL.to_string()),
line,
));
Expr::FnCall(
Box::new(Spanned::new(Expr::Ident(aux_name.to_string()), line)),
call_args,
)
} else {
Expr::FnCall(
Box::new(rewrite_native_guarded_calls_expr(callee, fn_name, aux_name)),
rewritten_args,
)
}
}
Expr::BinOp(op, left, right) => Expr::BinOp(
*op,
Box::new(rewrite_native_guarded_calls_expr(left, fn_name, aux_name)),
Box::new(rewrite_native_guarded_calls_expr(right, fn_name, aux_name)),
),
Expr::Neg(inner) => Expr::Neg(Box::new(rewrite_native_guarded_calls_expr(
inner, fn_name, aux_name,
))),
Expr::Match { subject, arms } => Expr::Match {
subject: Box::new(rewrite_native_guarded_calls_expr(
subject, fn_name, aux_name,
)),
arms: arms
.iter()
.map(|arm| MatchArm {
pattern: arm.pattern.clone(),
body: Box::new(rewrite_native_guarded_calls_expr(
&arm.body, fn_name, aux_name,
)),
binding_slots: std::sync::OnceLock::new(),
})
.collect(),
},
Expr::Constructor(name, arg) => Expr::Constructor(
name.clone(),
arg.as_ref()
.map(|inner| Box::new(rewrite_native_guarded_calls_expr(inner, fn_name, aux_name))),
),
Expr::ErrorProp(inner) => Expr::ErrorProp(Box::new(rewrite_native_guarded_calls_expr(
inner, fn_name, aux_name,
))),
Expr::InterpolatedStr(parts) => Expr::InterpolatedStr(
parts
.iter()
.map(|part| match part {
StrPart::Literal(_) => part.clone(),
StrPart::Parsed(inner) => StrPart::Parsed(Box::new(
rewrite_native_guarded_calls_expr(inner, fn_name, aux_name),
)),
})
.collect(),
),
Expr::List(items) => Expr::List(
items
.iter()
.map(|item| rewrite_native_guarded_calls_expr(item, fn_name, aux_name))
.collect(),
),
Expr::Tuple(items) => Expr::Tuple(
items
.iter()
.map(|item| rewrite_native_guarded_calls_expr(item, fn_name, aux_name))
.collect(),
),
Expr::IndependentProduct(items, flag) => Expr::IndependentProduct(
items
.iter()
.map(|item| rewrite_native_guarded_calls_expr(item, fn_name, aux_name))
.collect(),
*flag,
),
Expr::MapLiteral(entries) => Expr::MapLiteral(
entries
.iter()
.map(|(k, v)| {
(
rewrite_native_guarded_calls_expr(k, fn_name, aux_name),
rewrite_native_guarded_calls_expr(v, fn_name, aux_name),
)
})
.collect(),
),
Expr::RecordCreate { type_name, fields } => Expr::RecordCreate {
type_name: type_name.clone(),
fields: fields
.iter()
.map(|(name, value)| {
(
name.clone(),
rewrite_native_guarded_calls_expr(value, fn_name, aux_name),
)
})
.collect(),
},
Expr::RecordUpdate {
type_name,
base,
updates,
} => Expr::RecordUpdate {
type_name: type_name.clone(),
base: Box::new(rewrite_native_guarded_calls_expr(base, fn_name, aux_name)),
updates: updates
.iter()
.map(|(name, value)| {
(
name.clone(),
rewrite_native_guarded_calls_expr(value, fn_name, aux_name),
)
})
.collect(),
},
Expr::TailCall(boxed) => {
let TailCallData { target, args, .. } = boxed.as_ref();
let rewritten_args: Vec<Spanned<Expr>> = args
.iter()
.map(|arg| rewrite_native_guarded_calls_expr(arg, fn_name, aux_name))
.collect();
if target == fn_name {
let mut call_args = rewritten_args;
call_args.push(Spanned::new(
Expr::Ident(OMEGA_PROOF_SENTINEL.to_string()),
line,
));
Expr::FnCall(
Box::new(Spanned::new(Expr::Ident(aux_name.to_string()), line)),
call_args,
)
} else {
Expr::TailCall(Box::new(TailCallData::new(target.clone(), rewritten_args)))
}
}
};
Spanned::new(new_node, line)
}
pub fn rewrite_native_guarded_calls_resolved_expr(
expr: &Spanned<crate::ir::hir::ResolvedExpr>,
target_fn_id: crate::ir::FnId,
aux_name: &str,
) -> Spanned<crate::ir::hir::ResolvedExpr> {
use crate::ir::hir::{ResolvedCallee, ResolvedExpr, ResolvedMatchArm, ResolvedStrPart};
let line = expr.line;
let rec = |e: &Spanned<ResolvedExpr>| {
rewrite_native_guarded_calls_resolved_expr(e, target_fn_id, aux_name)
};
let synth_aux_callee = |line: usize| -> ResolvedCallee {
ResolvedCallee::Unresolved {
callee: Box::new(Spanned::new(
ResolvedExpr::Ident(aux_name.to_string()),
line,
)),
}
};
let new_node = match &expr.node {
ResolvedExpr::Literal(_) | ResolvedExpr::Ident(_) | ResolvedExpr::Resolved { .. } => {
return expr.clone();
}
ResolvedExpr::Attr(obj, field) => ResolvedExpr::Attr(Box::new(rec(obj)), field.clone()),
ResolvedExpr::Call(callee, args) => {
let rewritten_args: Vec<Spanned<ResolvedExpr>> = args.iter().map(&rec).collect();
let target_matches =
matches!(callee, ResolvedCallee::Fn(callee_id) if *callee_id == target_fn_id);
if target_matches {
let mut call_args = rewritten_args;
call_args.push(Spanned::new(
ResolvedExpr::Ident(OMEGA_PROOF_SENTINEL.to_string()),
line,
));
ResolvedExpr::Call(synth_aux_callee(line), call_args)
} else {
ResolvedExpr::Call(callee.clone(), rewritten_args)
}
}
ResolvedExpr::BinOp(op, left, right) => {
ResolvedExpr::BinOp(*op, Box::new(rec(left)), Box::new(rec(right)))
}
ResolvedExpr::Neg(inner) => ResolvedExpr::Neg(Box::new(rec(inner))),
ResolvedExpr::Match { subject, arms } => ResolvedExpr::Match {
subject: Box::new(rec(subject)),
arms: arms
.iter()
.map(|arm| ResolvedMatchArm {
pattern: arm.pattern.clone(),
body: Box::new(rec(&arm.body)),
binding_slots: std::sync::OnceLock::new(),
})
.collect(),
},
ResolvedExpr::Ctor(ctor, args) => {
ResolvedExpr::Ctor(ctor.clone(), args.iter().map(&rec).collect())
}
ResolvedExpr::ErrorProp(inner) => ResolvedExpr::ErrorProp(Box::new(rec(inner))),
ResolvedExpr::InterpolatedStr(parts) => ResolvedExpr::InterpolatedStr(
parts
.iter()
.map(|p| match p {
ResolvedStrPart::Literal(_) => p.clone(),
ResolvedStrPart::Parsed(inner) => ResolvedStrPart::Parsed(Box::new(rec(inner))),
})
.collect(),
),
ResolvedExpr::List(items) => ResolvedExpr::List(items.iter().map(&rec).collect()),
ResolvedExpr::Tuple(items) => ResolvedExpr::Tuple(items.iter().map(&rec).collect()),
ResolvedExpr::IndependentProduct(items, flag) => {
ResolvedExpr::IndependentProduct(items.iter().map(&rec).collect(), *flag)
}
ResolvedExpr::MapLiteral(entries) => {
ResolvedExpr::MapLiteral(entries.iter().map(|(k, v)| (rec(k), rec(v))).collect())
}
ResolvedExpr::RecordCreate {
type_id,
type_name,
fields,
} => ResolvedExpr::RecordCreate {
type_id: *type_id,
type_name: type_name.clone(),
fields: fields.iter().map(|(n, v)| (n.clone(), rec(v))).collect(),
},
ResolvedExpr::RecordUpdate {
type_id,
type_name,
base,
updates,
} => ResolvedExpr::RecordUpdate {
type_id: *type_id,
type_name: type_name.clone(),
base: Box::new(rec(base)),
updates: updates.iter().map(|(n, v)| (n.clone(), rec(v))).collect(),
},
ResolvedExpr::TailCall { target, args } => {
let rewritten_args: Vec<Spanned<ResolvedExpr>> = args.iter().map(&rec).collect();
if *target == target_fn_id {
let mut call_args = rewritten_args;
call_args.push(Spanned::new(
ResolvedExpr::Ident(OMEGA_PROOF_SENTINEL.to_string()),
line,
));
ResolvedExpr::Call(synth_aux_callee(line), call_args)
} else {
ResolvedExpr::TailCall {
target: *target,
args: rewritten_args,
}
}
}
};
Spanned::new(new_node, line)
}
pub fn rewrite_native_guarded_calls_body(body: &FnBody, fn_name: &str, aux_name: &str) -> FnBody {
FnBody::Block(
body.stmts()
.iter()
.map(|stmt| match stmt {
Stmt::Binding(name, ty, expr) => Stmt::Binding(
name.clone(),
ty.clone(),
rewrite_native_guarded_calls_expr(expr, fn_name, aux_name),
),
Stmt::Expr(expr) => {
Stmt::Expr(rewrite_native_guarded_calls_expr(expr, fn_name, aux_name))
}
})
.collect(),
)
}
pub fn rewrite_recursive_calls_body(
body: &FnBody,
targets: &HashSet<String>,
fuel_var: &str,
) -> FnBody {
FnBody::Block(
body.stmts()
.iter()
.map(|stmt| match stmt {
Stmt::Binding(name, ty, expr) => Stmt::Binding(
name.clone(),
ty.clone(),
rewrite_recursive_calls_expr(expr, targets, fuel_var),
),
Stmt::Expr(expr) => {
Stmt::Expr(rewrite_recursive_calls_expr(expr, targets, fuel_var))
}
})
.collect(),
)
}