use ryo_source::pure::{
PureBlock, PureClosureParam, PureExpr, PureMatchArm, PurePattern, PureStmt,
};
use ryo_symbol::SymbolId;
use crate::Mutation;
#[derive(Debug, Clone, Default)]
pub struct ManualMapMutation {
pub target_fn: Option<SymbolId>,
}
impl ManualMapMutation {
pub fn new() -> Self {
Self::default()
}
pub fn in_function(mut self, id: SymbolId) -> Self {
self.target_fn = Some(id);
self
}
fn is_some_pattern(pattern: &PurePattern) -> Option<String> {
match pattern {
PurePattern::Struct { path, fields, .. } => {
if (path == "Some" || path.ends_with("::Some")) && fields.len() == 1 {
if let Some((_, PurePattern::Ident { name, .. })) = fields.first() {
return Some(name.clone());
}
}
None
}
_ => None,
}
}
fn is_none_pattern(pattern: &PurePattern) -> bool {
match pattern {
PurePattern::Path(p) => p == "None" || p.ends_with("::None"),
PurePattern::Ident { name, .. } => name == "None",
_ => false,
}
}
fn is_ok_pattern(pattern: &PurePattern) -> Option<String> {
match pattern {
PurePattern::Struct { path, fields, .. } => {
if path == "Ok" && fields.len() == 1 {
if let (_, PurePattern::Ident { name, .. }) = &fields[0] {
return Some(name.clone());
}
}
None
}
_ => None,
}
}
fn is_err_pattern(pattern: &PurePattern) -> Option<String> {
match pattern {
PurePattern::Struct { path, fields, .. } => {
if path == "Err" && fields.len() == 1 {
if let (_, PurePattern::Ident { name, .. }) = &fields[0] {
return Some(name.clone());
}
}
None
}
_ => None,
}
}
fn is_some_expr(expr: &PureExpr) -> Option<&PureExpr> {
match expr {
PureExpr::Call { func, args } => {
if matches!(func.as_ref(), PureExpr::Path(p) if p == "Some" || p.ends_with("::Some"))
&& args.len() == 1
{
return Some(&args[0]);
}
None
}
_ => None,
}
}
fn is_none_expr(expr: &PureExpr) -> bool {
matches!(expr, PureExpr::Path(p) if p == "None" || p.ends_with("::None"))
}
fn is_ok_expr(expr: &PureExpr) -> Option<&PureExpr> {
match expr {
PureExpr::Call { func, args } => {
if matches!(func.as_ref(), PureExpr::Path(p) if p == "Ok") && args.len() == 1 {
return Some(&args[0]);
}
None
}
_ => None,
}
}
fn is_err_passthrough(expr: &PureExpr, err_name: &str) -> bool {
match expr {
PureExpr::Call { func, args } => {
if matches!(func.as_ref(), PureExpr::Path(p) if p == "Err") && args.len() == 1 {
return matches!(&args[0], PureExpr::Path(p) if p == err_name);
}
false
}
_ => false,
}
}
fn try_convert_match(scrutinee: &PureExpr, arms: &[PureMatchArm]) -> Option<PureExpr> {
if arms.len() != 2 {
return None;
}
if let Some(var_name) = Self::is_some_pattern(&arms[0].pattern) {
if Self::is_none_pattern(&arms[1].pattern) && Self::is_none_expr(&arms[1].body) {
if let Some(inner) = Self::is_some_expr(&arms[0].body) {
return Some(Self::create_map_call(
scrutinee.clone(),
var_name,
inner.clone(),
));
}
}
}
if Self::is_none_pattern(&arms[0].pattern) && Self::is_none_expr(&arms[0].body) {
if let Some(var_name) = Self::is_some_pattern(&arms[1].pattern) {
if let Some(inner) = Self::is_some_expr(&arms[1].body) {
return Some(Self::create_map_call(
scrutinee.clone(),
var_name,
inner.clone(),
));
}
}
}
if let Some(ok_var) = Self::is_ok_pattern(&arms[0].pattern) {
if let Some(err_var) = Self::is_err_pattern(&arms[1].pattern) {
if Self::is_err_passthrough(&arms[1].body, &err_var) {
if let Some(inner) = Self::is_ok_expr(&arms[0].body) {
return Some(Self::create_map_call(
scrutinee.clone(),
ok_var,
inner.clone(),
));
}
}
}
}
None
}
fn create_map_call(receiver: PureExpr, var_name: String, body: PureExpr) -> PureExpr {
PureExpr::MethodCall {
receiver: Box::new(receiver),
method: "map".to_string(),
turbofish: None,
args: vec![PureExpr::Closure {
is_async: false,
is_move: false,
params: vec![PureClosureParam::untyped(PurePattern::Ident {
name: var_name,
is_mut: false,
})],
ret: None,
body: Box::new(body),
}],
}
}
fn transform_expr(&self, expr: &mut PureExpr) -> usize {
let mut changes = 0;
if let PureExpr::Match {
expr: scrutinee,
arms,
} = expr
{
if let Some(map_call) = Self::try_convert_match(scrutinee, arms) {
*expr = map_call;
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::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);
}
_ => {}
}
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 ManualMapMutation {
fn describe(&self) -> String {
"Convert manual Option/Result map patterns to .map()".to_string()
}
fn mutation_type(&self) -> &'static str {
"ManualMap"
}
fn box_clone(&self) -> Box<dyn Mutation> {
Box::new(self.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_is_some_pattern_struct() {
let pattern = PurePattern::Struct {
path: "Some".to_string(),
fields: vec![(
"0".to_string(),
PurePattern::Ident {
name: "x".to_string(),
is_mut: false,
},
)],
rest: false,
};
assert_eq!(
ManualMapMutation::is_some_pattern(&pattern),
Some("x".to_string())
);
}
#[test]
fn test_is_none_pattern() {
let pattern = PurePattern::Path("None".to_string());
assert!(ManualMapMutation::is_none_pattern(&pattern));
}
#[test]
fn test_is_some_expr() {
let expr = PureExpr::Call {
func: Box::new(PureExpr::Path("Some".to_string())),
args: vec![PureExpr::Path("value".to_string())],
};
assert!(ManualMapMutation::is_some_expr(&expr).is_some());
}
#[test]
fn test_is_none_expr() {
let expr = PureExpr::Path("None".to_string());
assert!(ManualMapMutation::is_none_expr(&expr));
}
#[test]
fn test_try_convert_match_option() {
let scrutinee = PureExpr::Path("opt".to_string());
let arms = vec![
PureMatchArm {
pattern: PurePattern::Struct {
path: "Some".to_string(),
fields: vec![(
"0".to_string(),
PurePattern::Ident {
name: "x".to_string(),
is_mut: false,
},
)],
rest: false,
},
guard: None,
body: PureExpr::Call {
func: Box::new(PureExpr::Path("Some".to_string())),
args: vec![PureExpr::Binary {
op: "+".to_string(),
left: Box::new(PureExpr::Path("x".to_string())),
right: Box::new(PureExpr::Lit("1".to_string())),
}],
},
},
PureMatchArm {
pattern: PurePattern::Path("None".to_string()),
guard: None,
body: PureExpr::Path("None".to_string()),
},
];
let result = ManualMapMutation::try_convert_match(&scrutinee, &arms);
assert!(result.is_some());
if let Some(PureExpr::MethodCall { method, .. }) = result {
assert_eq!(method, "map");
} else {
panic!("Expected MethodCall");
}
}
}