1use ryo_source::pure::{
19 PureBlock, PureClosureParam, PureExpr, PureMatchArm, PurePattern, PureStmt,
20};
21use ryo_symbol::SymbolId;
22
23use crate::Mutation;
24
25#[derive(Debug, Clone, Default)]
48pub struct ManualMapMutation {
49 pub target_fn: Option<SymbolId>,
51}
52
53impl ManualMapMutation {
54 pub fn new() -> Self {
55 Self::default()
56 }
57
58 pub fn in_function(mut self, id: SymbolId) -> Self {
60 self.target_fn = Some(id);
61 self
62 }
63
64 fn is_some_pattern(pattern: &PurePattern) -> Option<String> {
69 match pattern {
70 PurePattern::Struct { path, fields, .. } => {
72 if (path == "Some" || path.ends_with("::Some")) && fields.len() == 1 {
74 if let Some((_, PurePattern::Ident { name, .. })) = fields.first() {
76 return Some(name.clone());
77 }
78 }
79 None
80 }
81 _ => None,
82 }
83 }
84
85 fn is_none_pattern(pattern: &PurePattern) -> bool {
90 match pattern {
91 PurePattern::Path(p) => p == "None" || p.ends_with("::None"),
92 PurePattern::Ident { name, .. } => name == "None",
93 _ => false,
94 }
95 }
96
97 fn is_ok_pattern(pattern: &PurePattern) -> Option<String> {
99 match pattern {
100 PurePattern::Struct { path, fields, .. } => {
101 if path == "Ok" && fields.len() == 1 {
102 if let (_, PurePattern::Ident { name, .. }) = &fields[0] {
103 return Some(name.clone());
104 }
105 }
106 None
107 }
108 _ => None,
109 }
110 }
111
112 fn is_err_pattern(pattern: &PurePattern) -> Option<String> {
114 match pattern {
115 PurePattern::Struct { path, fields, .. } => {
116 if path == "Err" && fields.len() == 1 {
117 if let (_, PurePattern::Ident { name, .. }) = &fields[0] {
118 return Some(name.clone());
119 }
120 }
121 None
122 }
123 _ => None,
124 }
125 }
126
127 fn is_some_expr(expr: &PureExpr) -> Option<&PureExpr> {
129 match expr {
130 PureExpr::Call { func, args } => {
131 if matches!(func.as_ref(), PureExpr::Path(p) if p == "Some" || p.ends_with("::Some"))
132 && args.len() == 1
133 {
134 return Some(&args[0]);
135 }
136 None
137 }
138 _ => None,
139 }
140 }
141
142 fn is_none_expr(expr: &PureExpr) -> bool {
144 matches!(expr, PureExpr::Path(p) if p == "None" || p.ends_with("::None"))
145 }
146
147 fn is_ok_expr(expr: &PureExpr) -> Option<&PureExpr> {
149 match expr {
150 PureExpr::Call { func, args } => {
151 if matches!(func.as_ref(), PureExpr::Path(p) if p == "Ok") && args.len() == 1 {
152 return Some(&args[0]);
153 }
154 None
155 }
156 _ => None,
157 }
158 }
159
160 fn is_err_passthrough(expr: &PureExpr, err_name: &str) -> bool {
162 match expr {
163 PureExpr::Call { func, args } => {
164 if matches!(func.as_ref(), PureExpr::Path(p) if p == "Err") && args.len() == 1 {
165 return matches!(&args[0], PureExpr::Path(p) if p == err_name);
166 }
167 false
168 }
169 _ => false,
170 }
171 }
172
173 fn try_convert_match(scrutinee: &PureExpr, arms: &[PureMatchArm]) -> Option<PureExpr> {
175 if arms.len() != 2 {
176 return None;
177 }
178
179 if let Some(var_name) = Self::is_some_pattern(&arms[0].pattern) {
181 if Self::is_none_pattern(&arms[1].pattern) && Self::is_none_expr(&arms[1].body) {
182 if let Some(inner) = Self::is_some_expr(&arms[0].body) {
183 return Some(Self::create_map_call(
184 scrutinee.clone(),
185 var_name,
186 inner.clone(),
187 ));
188 }
189 }
190 }
191
192 if Self::is_none_pattern(&arms[0].pattern) && Self::is_none_expr(&arms[0].body) {
194 if let Some(var_name) = Self::is_some_pattern(&arms[1].pattern) {
195 if let Some(inner) = Self::is_some_expr(&arms[1].body) {
196 return Some(Self::create_map_call(
197 scrutinee.clone(),
198 var_name,
199 inner.clone(),
200 ));
201 }
202 }
203 }
204
205 if let Some(ok_var) = Self::is_ok_pattern(&arms[0].pattern) {
207 if let Some(err_var) = Self::is_err_pattern(&arms[1].pattern) {
208 if Self::is_err_passthrough(&arms[1].body, &err_var) {
209 if let Some(inner) = Self::is_ok_expr(&arms[0].body) {
210 return Some(Self::create_map_call(
211 scrutinee.clone(),
212 ok_var,
213 inner.clone(),
214 ));
215 }
216 }
217 }
218 }
219
220 None
221 }
222
223 fn create_map_call(receiver: PureExpr, var_name: String, body: PureExpr) -> PureExpr {
225 PureExpr::MethodCall {
226 receiver: Box::new(receiver),
227 method: "map".to_string(),
228 turbofish: None,
229 args: vec![PureExpr::Closure {
230 is_async: false,
231 is_move: false,
232 params: vec![PureClosureParam::untyped(PurePattern::Ident {
233 name: var_name,
234 is_mut: false,
235 })],
236 ret: None,
237 body: Box::new(body),
238 }],
239 }
240 }
241
242 fn transform_expr(&self, expr: &mut PureExpr) -> usize {
244 let mut changes = 0;
245
246 if let PureExpr::Match {
248 expr: scrutinee,
249 arms,
250 } = expr
251 {
252 if let Some(map_call) = Self::try_convert_match(scrutinee, arms) {
253 *expr = map_call;
254 return 1;
255 }
256 }
257
258 match expr {
264 PureExpr::Binary { left, right, .. } => {
265 changes += self.transform_expr(left);
266 changes += self.transform_expr(right);
267 }
268 PureExpr::Unary { expr: inner, .. } => {
269 changes += self.transform_expr(inner);
270 }
271 PureExpr::Call { func, args } => {
272 changes += self.transform_expr(func);
273 for arg in args {
274 changes += self.transform_expr(arg);
275 }
276 }
277 PureExpr::MethodCall { receiver, args, .. } => {
278 changes += self.transform_expr(receiver);
279 for arg in args {
280 changes += self.transform_expr(arg);
281 }
282 }
283 PureExpr::Block { block, .. } => {
284 changes += self.transform_block(block);
285 }
286 PureExpr::If {
287 cond,
288 then_branch,
289 else_branch,
290 } => {
291 changes += self.transform_expr(cond);
292 changes += self.transform_block(then_branch);
293 if let Some(else_expr) = else_branch {
294 changes += self.transform_expr(else_expr);
295 }
296 }
297 PureExpr::Match { expr: e, arms } => {
298 changes += self.transform_expr(e);
299 for arm in arms {
300 changes += self.transform_expr(&mut arm.body);
301 }
302 }
303 PureExpr::Loop { body: block, .. } | PureExpr::While { body: block, .. } => {
304 changes += self.transform_block(block);
305 }
306 PureExpr::For {
307 expr: iter_expr,
308 body,
309 ..
310 } => {
311 changes += self.transform_expr(iter_expr);
312 changes += self.transform_block(body);
313 }
314 PureExpr::Closure { body, .. } => {
315 changes += self.transform_expr(body);
316 }
317 _ => {}
318 }
319
320 changes
321 }
322
323 pub fn transform_block(&self, block: &mut PureBlock) -> usize {
324 let mut changes = 0;
325 for stmt in &mut block.stmts {
326 changes += self.transform_stmt(stmt);
327 }
328 changes
329 }
330
331 fn transform_stmt(&self, stmt: &mut PureStmt) -> usize {
332 match stmt {
333 PureStmt::Local { init: Some(e), .. } => self.transform_expr(e),
334 PureStmt::Semi(e) | PureStmt::Expr(e) => self.transform_expr(e),
335 _ => 0,
336 }
337 }
338}
339
340impl Mutation for ManualMapMutation {
341 fn describe(&self) -> String {
342 "Convert manual Option/Result map patterns to .map()".to_string()
343 }
344
345 fn mutation_type(&self) -> &'static str {
346 "ManualMap"
347 }
348
349 fn box_clone(&self) -> Box<dyn Mutation> {
350 Box::new(self.clone())
351 }
352}
353
354#[cfg(test)]
355mod tests {
356 use super::*;
357
358 #[test]
359 fn test_is_some_pattern_struct() {
360 let pattern = PurePattern::Struct {
362 path: "Some".to_string(),
363 fields: vec![(
364 "0".to_string(),
365 PurePattern::Ident {
366 name: "x".to_string(),
367 is_mut: false,
368 },
369 )],
370 rest: false,
371 };
372 assert_eq!(
373 ManualMapMutation::is_some_pattern(&pattern),
374 Some("x".to_string())
375 );
376 }
377
378 #[test]
379 fn test_is_none_pattern() {
380 let pattern = PurePattern::Path("None".to_string());
381 assert!(ManualMapMutation::is_none_pattern(&pattern));
382 }
383
384 #[test]
385 fn test_is_some_expr() {
386 let expr = PureExpr::Call {
387 func: Box::new(PureExpr::Path("Some".to_string())),
388 args: vec![PureExpr::Path("value".to_string())],
389 };
390 assert!(ManualMapMutation::is_some_expr(&expr).is_some());
391 }
392
393 #[test]
394 fn test_is_none_expr() {
395 let expr = PureExpr::Path("None".to_string());
396 assert!(ManualMapMutation::is_none_expr(&expr));
397 }
398
399 #[test]
400 fn test_try_convert_match_option() {
401 let scrutinee = PureExpr::Path("opt".to_string());
402 let arms = vec![
403 PureMatchArm {
404 pattern: PurePattern::Struct {
405 path: "Some".to_string(),
406 fields: vec![(
407 "0".to_string(),
408 PurePattern::Ident {
409 name: "x".to_string(),
410 is_mut: false,
411 },
412 )],
413 rest: false,
414 },
415 guard: None,
416 body: PureExpr::Call {
417 func: Box::new(PureExpr::Path("Some".to_string())),
418 args: vec![PureExpr::Binary {
419 op: "+".to_string(),
420 left: Box::new(PureExpr::Path("x".to_string())),
421 right: Box::new(PureExpr::Lit("1".to_string())),
422 }],
423 },
424 },
425 PureMatchArm {
426 pattern: PurePattern::Path("None".to_string()),
427 guard: None,
428 body: PureExpr::Path("None".to_string()),
429 },
430 ];
431
432 let result = ManualMapMutation::try_convert_match(&scrutinee, &arms);
433 assert!(result.is_some());
434
435 if let Some(PureExpr::MethodCall { method, .. }) = result {
436 assert_eq!(method, "map");
437 } else {
438 panic!("Expected MethodCall");
439 }
440 }
441}