1pub mod detect;
14
15use std::collections::HashSet;
16
17use crate::ast::{Expr, FnBody, MatchArm, Spanned, Stmt, StrPart, TailCallData};
18use crate::codegen::common::expr_to_dotted_name;
19
20pub use detect::analyze_plans;
21
22#[derive(Clone, Debug, PartialEq)]
31pub enum RecursionPlan {
32 IntCountdown { param_index: usize },
35 IntAscending {
40 param_index: usize,
41 bound: Spanned<Expr>,
42 },
43 LinearRecurrence2,
47 ListStructural { param_index: usize },
50 SizeOfStructural,
53 StringPosAdvance,
57 MutualIntCountdown,
60 MutualStringPosAdvance { rank: usize },
64 MutualSizeOfRanked { rank: usize },
67}
68
69#[derive(Clone, Debug, Eq, PartialEq)]
73pub struct ProofModeIssue {
74 pub line: usize,
75 pub message: String,
76}
77
78pub fn fuel_helper_name(name: &str) -> String {
82 format!("{}__fuel", name)
83}
84
85pub fn rewrite_recursive_calls_expr(
90 expr: &Spanned<Expr>,
91 targets: &HashSet<String>,
92 fuel_var: &str,
93) -> Spanned<Expr> {
94 let line = expr.line;
95 let new_node = match &expr.node {
96 Expr::Literal(_) | Expr::Ident(_) | Expr::Resolved { .. } => return expr.clone(),
97 Expr::Attr(obj, field) => Expr::Attr(
98 Box::new(rewrite_recursive_calls_expr(obj, targets, fuel_var)),
99 field.clone(),
100 ),
101 Expr::FnCall(callee, args) => {
102 let rewritten_args: Vec<Spanned<Expr>> = args
103 .iter()
104 .map(|arg| rewrite_recursive_calls_expr(arg, targets, fuel_var))
105 .collect();
106 if let Some(name) = expr_to_dotted_name(&callee.node)
107 && targets.contains(&name)
108 {
109 let mut call_args = Vec::with_capacity(rewritten_args.len() + 1);
110 call_args.push(Spanned::new(Expr::Ident(fuel_var.to_string()), line));
111 call_args.extend(rewritten_args);
112 Expr::FnCall(
113 Box::new(Spanned::new(Expr::Ident(fuel_helper_name(&name)), line)),
114 call_args,
115 )
116 } else {
117 Expr::FnCall(
118 Box::new(rewrite_recursive_calls_expr(callee, targets, fuel_var)),
119 rewritten_args,
120 )
121 }
122 }
123 Expr::BinOp(op, left, right) => Expr::BinOp(
124 *op,
125 Box::new(rewrite_recursive_calls_expr(left, targets, fuel_var)),
126 Box::new(rewrite_recursive_calls_expr(right, targets, fuel_var)),
127 ),
128 Expr::Neg(inner) => Expr::Neg(Box::new(rewrite_recursive_calls_expr(
129 inner, targets, fuel_var,
130 ))),
131 Expr::Match { subject, arms } => Expr::Match {
132 subject: Box::new(rewrite_recursive_calls_expr(subject, targets, fuel_var)),
133 arms: arms
134 .iter()
135 .map(|arm| MatchArm {
136 pattern: arm.pattern.clone(),
137 body: Box::new(rewrite_recursive_calls_expr(&arm.body, targets, fuel_var)),
138 binding_slots: std::sync::OnceLock::new(),
139 })
140 .collect(),
141 },
142 Expr::Constructor(name, arg) => Expr::Constructor(
143 name.clone(),
144 arg.as_ref()
145 .map(|inner| Box::new(rewrite_recursive_calls_expr(inner, targets, fuel_var))),
146 ),
147 Expr::ErrorProp(inner) => Expr::ErrorProp(Box::new(rewrite_recursive_calls_expr(
148 inner, targets, fuel_var,
149 ))),
150 Expr::InterpolatedStr(parts) => Expr::InterpolatedStr(
151 parts
152 .iter()
153 .map(|part| match part {
154 StrPart::Literal(_) => part.clone(),
155 StrPart::Parsed(inner) => StrPart::Parsed(Box::new(
156 rewrite_recursive_calls_expr(inner, targets, fuel_var),
157 )),
158 })
159 .collect(),
160 ),
161 Expr::List(items) => Expr::List(
162 items
163 .iter()
164 .map(|item| rewrite_recursive_calls_expr(item, targets, fuel_var))
165 .collect(),
166 ),
167 Expr::Tuple(items) => Expr::Tuple(
168 items
169 .iter()
170 .map(|item| rewrite_recursive_calls_expr(item, targets, fuel_var))
171 .collect(),
172 ),
173 Expr::IndependentProduct(items, flag) => Expr::IndependentProduct(
174 items
175 .iter()
176 .map(|item| rewrite_recursive_calls_expr(item, targets, fuel_var))
177 .collect(),
178 *flag,
179 ),
180 Expr::MapLiteral(entries) => Expr::MapLiteral(
181 entries
182 .iter()
183 .map(|(k, v)| {
184 (
185 rewrite_recursive_calls_expr(k, targets, fuel_var),
186 rewrite_recursive_calls_expr(v, targets, fuel_var),
187 )
188 })
189 .collect(),
190 ),
191 Expr::RecordCreate { type_name, fields } => Expr::RecordCreate {
192 type_name: type_name.clone(),
193 fields: fields
194 .iter()
195 .map(|(name, value)| {
196 (
197 name.clone(),
198 rewrite_recursive_calls_expr(value, targets, fuel_var),
199 )
200 })
201 .collect(),
202 },
203 Expr::RecordUpdate {
204 type_name,
205 base,
206 updates,
207 } => Expr::RecordUpdate {
208 type_name: type_name.clone(),
209 base: Box::new(rewrite_recursive_calls_expr(base, targets, fuel_var)),
210 updates: updates
211 .iter()
212 .map(|(name, value)| {
213 (
214 name.clone(),
215 rewrite_recursive_calls_expr(value, targets, fuel_var),
216 )
217 })
218 .collect(),
219 },
220 Expr::TailCall(boxed) => {
221 let TailCallData { target, args, .. } = boxed.as_ref();
222 let rewritten_args: Vec<Spanned<Expr>> = args
223 .iter()
224 .map(|arg| rewrite_recursive_calls_expr(arg, targets, fuel_var))
225 .collect();
226 if targets.contains(target) {
227 let mut call_args = Vec::with_capacity(rewritten_args.len() + 1);
228 call_args.push(Spanned::new(Expr::Ident(fuel_var.to_string()), line));
229 call_args.extend(rewritten_args);
230 Expr::FnCall(
231 Box::new(Spanned::new(Expr::Ident(fuel_helper_name(target)), line)),
232 call_args,
233 )
234 } else {
235 Expr::TailCall(Box::new(TailCallData::new(target.clone(), rewritten_args)))
236 }
237 }
238 };
239 Spanned::new(new_node, line)
240}
241
242pub fn rewrite_recursive_calls_body(
244 body: &FnBody,
245 targets: &HashSet<String>,
246 fuel_var: &str,
247) -> FnBody {
248 FnBody::Block(
249 body.stmts()
250 .iter()
251 .map(|stmt| match stmt {
252 Stmt::Binding(name, ty, expr) => Stmt::Binding(
253 name.clone(),
254 ty.clone(),
255 rewrite_recursive_calls_expr(expr, targets, fuel_var),
256 ),
257 Stmt::Expr(expr) => {
258 Stmt::Expr(rewrite_recursive_calls_expr(expr, targets, fuel_var))
259 }
260 })
261 .collect(),
262 )
263}