use ryo_source::pure::{PureBlock, PureExpr, PureStmt};
use ryo_symbol::SymbolId;
use crate::Mutation;
#[derive(Debug, Clone, Default)]
pub struct MapUnwrapOrMutation {
pub target_fn: Option<SymbolId>,
}
impl MapUnwrapOrMutation {
pub fn new() -> Self {
Self::default()
}
pub fn in_function(mut self, id: SymbolId) -> Self {
self.target_fn = Some(id);
self
}
fn is_map_unwrap_or_pattern(expr: &PureExpr) -> Option<(PureExpr, PureExpr, PureExpr, bool)> {
if let PureExpr::MethodCall {
receiver,
method,
args,
..
} = expr
{
let (is_lazy, default_arg) = if method == "unwrap_or" && args.len() == 1 {
(false, &args[0])
} else if method == "unwrap_or_else" && args.len() == 1 {
(true, &args[0])
} else {
return None;
};
if let PureExpr::MethodCall {
receiver: map_receiver,
method: map_method,
args: map_args,
..
} = receiver.as_ref()
{
if map_method == "map" && map_args.len() == 1 {
return Some((
map_receiver.as_ref().clone(),
map_args[0].clone(),
default_arg.clone(),
is_lazy,
));
}
}
}
None
}
fn is_map_or_none_pattern(expr: &PureExpr) -> Option<(PureExpr, PureExpr)> {
if let PureExpr::MethodCall {
receiver,
method,
args,
..
} = expr
{
if method == "unwrap_or"
&& args.len() == 1
&& matches!(&args[0], PureExpr::Path(p) if p == "None")
{
if let PureExpr::MethodCall {
receiver: map_receiver,
method: map_method,
args: map_args,
..
} = receiver.as_ref()
{
if map_method == "map" && map_args.len() == 1 {
return Some((map_receiver.as_ref().clone(), map_args[0].clone()));
}
}
}
}
None
}
fn transform_expr(&self, expr: &mut PureExpr) -> usize {
let mut changes = 0;
if let Some((receiver, map_fn)) = Self::is_map_or_none_pattern(expr) {
*expr = PureExpr::MethodCall {
receiver: Box::new(receiver),
method: "and_then".to_string(),
turbofish: None,
args: vec![map_fn],
};
return 1;
}
if let Some((receiver, map_fn, default_value, is_lazy)) =
Self::is_map_unwrap_or_pattern(expr)
{
let method = if is_lazy { "map_or_else" } else { "map_or" };
*expr = PureExpr::MethodCall {
receiver: Box::new(receiver),
method: method.to_string(),
turbofish: None,
args: vec![default_value, map_fn],
};
return 1;
}
match expr {
PureExpr::Binary { left, right, .. } => {
changes += self.transform_expr(left);
changes += self.transform_expr(right);
}
PureExpr::Unary { expr: inner, .. } => {
changes += self.transform_expr(inner);
}
PureExpr::Call { func, args } => {
changes += self.transform_expr(func);
for arg in args {
changes += self.transform_expr(arg);
}
}
PureExpr::MethodCall { receiver, args, .. } => {
changes += self.transform_expr(receiver);
for arg in args {
changes += self.transform_expr(arg);
}
}
PureExpr::Field { expr: inner, .. } => {
changes += self.transform_expr(inner);
}
PureExpr::Index { expr: inner, index } => {
changes += self.transform_expr(inner);
changes += self.transform_expr(index);
}
PureExpr::Block { block, .. } => {
changes += self.transform_block(block);
}
PureExpr::If {
cond,
then_branch,
else_branch,
} => {
changes += self.transform_expr(cond);
changes += self.transform_block(then_branch);
if let Some(else_expr) = else_branch {
changes += self.transform_expr(else_expr);
}
}
PureExpr::Match { expr: e, arms } => {
changes += self.transform_expr(e);
for arm in arms {
changes += self.transform_expr(&mut arm.body);
}
}
PureExpr::Loop { body: block, .. } | PureExpr::While { body: block, .. } => {
changes += self.transform_block(block);
}
PureExpr::For {
expr: iter_expr,
body,
..
} => {
changes += self.transform_expr(iter_expr);
changes += self.transform_block(body);
}
PureExpr::Closure { body, .. } => {
changes += self.transform_expr(body);
}
PureExpr::Tuple(exprs) | PureExpr::Array(exprs) => {
for e in exprs {
changes += self.transform_expr(e);
}
}
PureExpr::Struct { fields, .. } => {
for (_, e) in fields {
changes += self.transform_expr(e);
}
}
PureExpr::Ref { expr: inner, .. } => {
changes += self.transform_expr(inner);
}
PureExpr::Return(Some(inner)) => {
changes += self.transform_expr(inner);
}
PureExpr::Try(inner) | PureExpr::Await(inner) => {
changes += self.transform_expr(inner);
}
_ => {}
}
changes
}
pub fn transform_block(&self, block: &mut PureBlock) -> usize {
let mut changes = 0;
for stmt in &mut block.stmts {
changes += self.transform_stmt(stmt);
}
changes
}
fn transform_stmt(&self, stmt: &mut PureStmt) -> usize {
match stmt {
PureStmt::Local { init: Some(e), .. } => self.transform_expr(e),
PureStmt::Semi(e) | PureStmt::Expr(e) => self.transform_expr(e),
_ => 0,
}
}
}
impl Mutation for MapUnwrapOrMutation {
fn describe(&self) -> String {
"Convert .map().unwrap_or() to .map_or()".to_string()
}
fn mutation_type(&self) -> &'static str {
"MapUnwrapOr"
}
fn box_clone(&self) -> Box<dyn Mutation> {
Box::new(self.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_is_map_unwrap_or_pattern() {
let expr = PureExpr::MethodCall {
receiver: Box::new(PureExpr::MethodCall {
receiver: Box::new(PureExpr::Path("opt".to_string())),
method: "map".to_string(),
turbofish: None,
args: vec![PureExpr::Path("f".to_string())],
}),
method: "unwrap_or".to_string(),
turbofish: None,
args: vec![PureExpr::Lit("0".to_string())],
};
let result = MapUnwrapOrMutation::is_map_unwrap_or_pattern(&expr);
assert!(result.is_some());
let (_, _, _, is_lazy) = result.unwrap();
assert!(!is_lazy);
}
#[test]
fn test_is_map_unwrap_or_else_pattern() {
let expr = PureExpr::MethodCall {
receiver: Box::new(PureExpr::MethodCall {
receiver: Box::new(PureExpr::Path("opt".to_string())),
method: "map".to_string(),
turbofish: None,
args: vec![PureExpr::Path("f".to_string())],
}),
method: "unwrap_or_else".to_string(),
turbofish: None,
args: vec![PureExpr::Path("g".to_string())],
};
let result = MapUnwrapOrMutation::is_map_unwrap_or_pattern(&expr);
assert!(result.is_some());
let (_, _, _, is_lazy) = result.unwrap();
assert!(is_lazy);
}
#[test]
fn test_is_map_or_none_pattern() {
let expr = PureExpr::MethodCall {
receiver: Box::new(PureExpr::MethodCall {
receiver: Box::new(PureExpr::Path("opt".to_string())),
method: "map".to_string(),
turbofish: None,
args: vec![PureExpr::Path("f".to_string())],
}),
method: "unwrap_or".to_string(),
turbofish: None,
args: vec![PureExpr::Path("None".to_string())],
};
assert!(MapUnwrapOrMutation::is_map_or_none_pattern(&expr).is_some());
}
}