use std::collections::HashSet;
use super::expr::aver_name_to_dafny;
use crate::ast::{FnDef, TypeDef};
use crate::codegen::CodegenContext;
use crate::codegen::common::parse_type_annotation;
use crate::codegen::recursion::{fuel_helper_name, rewrite_recursive_calls_body};
use crate::types::Type;
pub fn emit_mutual_fuel_group(fns: &[&FnDef], ctx: &CodegenContext) -> Option<String> {
for fd in fns {
dafny_default_value(&fd.return_type, ctx)?;
}
let scc_size = fns.len();
let targets: HashSet<String> = fns.iter().map(|fd| fd.name.clone()).collect();
let mut helper_lines: Vec<String> = Vec::new();
let mut wrapper_lines: Vec<String> = Vec::new();
for fd in fns {
let helper_name = fuel_helper_name(&fd.name);
let fn_name = aver_name_to_dafny(&fd.name);
let params_str = emit_dafny_params(&fd.params);
let ret_type_str = super::toplevel::emit_type(&fd.return_type);
let default_val = dafny_default_value(&fd.return_type, ctx)
.expect("default value presence is checked above");
let arg_names = emit_dafny_arg_names(&fd.params);
let metric = emit_fuel_metric(fd, ctx, scc_size);
let lowered_body = crate::types::checker::effect_lifting::lower_pure_question_bang_fn(fd)
.ok()
.flatten()
.map(|lowered| lowered.body.as_ref().clone())
.unwrap_or_else(|| fd.body.as_ref().clone());
let rewritten_body = rewrite_recursive_calls_body(&lowered_body, &targets, "fuel'");
let body_str = super::toplevel::emit_fn_body(&rewritten_body, ctx);
if let Some(desc) = &fd.desc {
helper_lines.push(format!("// {}", desc));
}
helper_lines.push(format!(
"function {}(fuel: nat, {}): {}",
helper_name, params_str, ret_type_str
));
helper_lines.push(" decreases fuel".to_string());
helper_lines.push("{".to_string());
helper_lines.push(format!(" if fuel == 0 then {}", default_val));
helper_lines.push(format!(" else var fuel' := fuel - 1; {}", body_str));
helper_lines.push("}\n".to_string());
wrapper_lines.push(format!(
"function {}({}): {}",
fn_name, params_str, ret_type_str
));
wrapper_lines.push("{".to_string());
wrapper_lines.push(format!(" {}({}, {})", helper_name, metric, arg_names));
wrapper_lines.push("}\n".to_string());
}
Some(
[helper_lines, wrapper_lines]
.into_iter()
.flatten()
.collect::<Vec<_>>()
.join("\n"),
)
}
pub fn emit_mutual_native_decreases_group(fns: &[&FnDef], ctx: &CodegenContext) -> Option<String> {
let mut ranks: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
for fd in fns {
let contract = crate::codegen::common::find_fn_contract_for_fn(ctx, fd)?;
match contract.recursion.as_ref()? {
crate::ir::RecursionContract::Fuel {
fuel_metric: crate::ir::FuelMetric::Lex { params, rank },
} if params.is_empty() => {
ranks.insert(fd.name.clone(), *rank);
}
_ => return None,
}
}
let mut measures: std::collections::HashMap<String, String> = std::collections::HashMap::new();
for fd in fns {
let measure = emit_sizeof_measure_expr_dafny(fd)?;
measures.insert(fd.name.clone(), measure);
}
if crate::codegen::recursion::detect::scc_has_growing_accumulator(fns) {
return None;
}
let mut lines: Vec<String> = Vec::new();
for fd in fns {
let measure = measures.get(&fd.name).unwrap();
let rank = ranks.get(&fd.name).unwrap();
let fn_name = aver_name_to_dafny(&fd.name);
let params_str = emit_dafny_params(&fd.params);
let ret_type_str = super::toplevel::emit_type(&fd.return_type);
let lowered_body = crate::types::checker::effect_lifting::lower_pure_question_bang_fn(fd)
.ok()
.flatten()
.map(|lowered| lowered.body.as_ref().clone())
.unwrap_or_else(|| fd.body.as_ref().clone());
let body_str = super::toplevel::emit_fn_body(&lowered_body, ctx);
if let Some(desc) = &fd.desc {
lines.push(format!("// {}", desc));
}
lines.push(format!(
"function {}({}): {}",
fn_name, params_str, ret_type_str
));
lines.push(format!(" decreases {}, {}", measure, rank));
lines.push("{".to_string());
lines.push(format!(" {}", body_str));
lines.push("}\n".to_string());
}
Some(lines.join("\n"))
}
fn emit_sizeof_measure_expr_dafny(fd: &FnDef) -> Option<String> {
let mut terms: Vec<String> = Vec::new();
for (name, ty) in &fd.params {
let parsed = crate::codegen::common::parse_type_annotation(ty);
match parsed {
crate::types::Type::List(_)
| crate::types::Type::Vector(_)
| crate::types::Type::Str => {
terms.push(format!("|{}|", aver_name_to_dafny(name)));
}
_ => {}
}
}
(!terms.is_empty()).then(|| terms.join(" + "))
}
fn emit_dafny_params(params: &[(String, String)]) -> String {
params
.iter()
.map(|(pname, ptype)| {
format!(
"{}: {}",
aver_name_to_dafny(pname),
super::toplevel::emit_type(ptype)
)
})
.collect::<Vec<_>>()
.join(", ")
}
fn emit_dafny_arg_names(params: &[(String, String)]) -> String {
params
.iter()
.map(|(pname, _)| aver_name_to_dafny(pname))
.collect::<Vec<_>>()
.join(", ")
}
fn emit_fuel_metric(fd: &FnDef, ctx: &CodegenContext, scc_size: usize) -> String {
let Some(contract) = crate::codegen::common::find_fn_contract_for_fn(ctx, fd) else {
return "1".to_string();
};
let Some(crate::ir::RecursionContract::Fuel {
fuel_metric: crate::ir::FuelMetric::Lex { params, rank },
}) = contract.recursion.as_ref()
else {
return "1".to_string();
};
match params.as_slice() {
[p] if *rank == 0 => {
let name = aver_name_to_dafny(p);
format!("(if {n} >= 0 then {n} else 0) + 1", n = name)
}
[_s, _pos] => {
let Some(name) = first_seq_or_string_param(fd) else {
return format!("{}", rank.max(&1));
};
format!(
"(|{n}| + 1) * {budget}",
n = aver_name_to_dafny(name),
budget = rank * scc_size + 1
)
}
[] => {
let Some(name) = first_seq_or_string_param(fd) else {
return format!("{}", rank.max(&1));
};
format!(
"(|{n}| + 1) * {budget}",
n = aver_name_to_dafny(name),
budget = rank * scc_size + 1
)
}
_ => "1".to_string(),
}
}
fn first_seq_or_string_param(fd: &FnDef) -> Option<&String> {
fd.params
.iter()
.find(|(_, t)| {
let ty = parse_type_annotation(t);
matches!(ty, Type::List(_) | Type::Vector(_) | Type::Str)
})
.map(|(n, _)| n)
}
pub fn dafny_default_value(type_str: &str, ctx: &CodegenContext) -> Option<String> {
let mut visiting = HashSet::new();
type_default(&parse_type_annotation(type_str), ctx, &mut visiting)
}
fn type_default(ty: &Type, ctx: &CodegenContext, visiting: &mut HashSet<String>) -> Option<String> {
Some(match ty {
Type::Int => "0".to_string(),
Type::Float => "0.0".to_string(),
Type::Str => "\"\"".to_string(),
Type::Bool => "false".to_string(),
Type::Unit => "()".to_string(),
Type::List(_) | Type::Vector(_) => "[]".to_string(),
Type::Map(_, _) => "map[]".to_string(),
Type::Option(_) => "Option.None".to_string(),
Type::Result(_, err) => {
format!("Result.Err({})", type_default(err, ctx, visiting)?)
}
Type::Tuple(items) => {
let parts: Vec<String> = items
.iter()
.map(|t| type_default(t, ctx, visiting))
.collect::<Option<_>>()?;
format!("({})", parts.join(", "))
}
Type::Named { .. } => {
let key = crate::codegen::common::backend_named_type_key(ctx, ty)?;
named_type_default(&key, ctx, visiting)?
}
Type::Fn(_, _, _) | Type::Var(_) | Type::Invalid => return None,
})
}
fn named_type_default(
name: &str,
ctx: &CodegenContext,
visiting: &mut HashSet<String>,
) -> Option<String> {
if visiting.contains(name) {
return None;
}
visiting.insert(name.to_string());
let td = find_type_def(ctx, name)?;
let result = type_def_default(td, ctx, visiting);
visiting.remove(name);
result
}
fn find_type_def<'a>(ctx: &'a CodegenContext, target: &str) -> Option<&'a TypeDef> {
if let Some((prefix, bare)) = target.rsplit_once('.') {
for m in &ctx.modules {
if m.prefix == prefix {
return m
.type_defs
.iter()
.find(|td| crate::codegen::common::type_def_name(td) == bare);
}
}
}
ctx.type_defs
.iter()
.chain(ctx.modules.iter().flat_map(|m| m.type_defs.iter()))
.find(|td| crate::codegen::common::type_def_name(td) == target)
}
fn type_def_default(
td: &TypeDef,
ctx: &CodegenContext,
visiting: &mut HashSet<String>,
) -> Option<String> {
match td {
TypeDef::Sum { name, variants, .. } => {
let variant = variants.first()?;
if variant.fields.is_empty() {
Some(format!("{}.{}", name, variant.name))
} else {
let args: Vec<String> = variant
.fields
.iter()
.map(|ft| type_default(&parse_type_annotation(ft), ctx, visiting))
.collect::<Option<_>>()?;
Some(format!("{}.{}({})", name, variant.name, args.join(", ")))
}
}
TypeDef::Product { name, fields, .. } => {
let args: Vec<String> = fields
.iter()
.map(|(fname, ftype)| {
type_default(&parse_type_annotation(ftype), ctx, visiting)
.map(|v| format!("{} := {}", aver_name_to_dafny(fname), v))
})
.collect::<Option<_>>()?;
Some(format!("{}({})", name, args.join(", ")))
}
}
}