use ryo_source::pure::{PureBlock, PureExpr, PureStmt};
use ryo_symbol::SymbolId;
use crate::Mutation;
#[derive(Debug, Clone, Default)]
pub struct UnwrapToQuestionMutation {
pub include_expect: bool,
pub target_fn: Option<SymbolId>,
}
impl UnwrapToQuestionMutation {
pub fn new() -> Self {
Self {
include_expect: true,
target_fn: None,
}
}
pub fn unwrap_only(mut self) -> Self {
self.include_expect = false;
self
}
pub fn in_function(mut self, id: SymbolId) -> Self {
self.target_fn = Some(id);
self
}
fn is_unwrap_call(&self, method: &str, args: &[PureExpr]) -> bool {
match method {
"unwrap" if args.is_empty() => true,
"expect" if args.len() == 1 && self.include_expect => true,
"unwrap_or_else" if args.len() == 1 => {
if let PureExpr::Closure { body, .. } = &args[0] {
Self::is_panic_expr(body)
} else {
false
}
}
_ => false,
}
}
fn is_panic_expr(expr: &PureExpr) -> bool {
match expr {
PureExpr::Macro { name, .. } => {
name == "panic"
|| name == "unreachable"
|| name == "todo"
|| name == "unimplemented"
}
PureExpr::Block { block, .. } if block.stmts.len() == 1 => match &block.stmts[0] {
PureStmt::Expr(e) | PureStmt::Semi(e) => Self::is_panic_expr(e),
_ => false,
},
_ => false,
}
}
fn transform_expr(&self, expr: &mut PureExpr) -> usize {
let mut changes = 0;
if let PureExpr::MethodCall {
receiver,
method,
args,
..
} = expr
{
if self.is_unwrap_call(method, args) {
changes += self.transform_expr(receiver);
let inner = std::mem::replace(
receiver.as_mut(),
PureExpr::Path("__placeholder".to_string()),
);
*expr = PureExpr::Try(Box::new(inner));
return changes + 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) => {
changes += self.transform_expr(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 UnwrapToQuestionMutation {
fn describe(&self) -> String {
"Convert .unwrap()/.expect() to ? operator".to_string()
}
fn mutation_type(&self) -> &'static str {
"UnwrapToQuestion"
}
fn box_clone(&self) -> Box<dyn Mutation> {
Box::new(self.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_unwrap_expr() -> PureExpr {
PureExpr::MethodCall {
receiver: Box::new(PureExpr::Path("x".to_string())),
method: "unwrap".to_string(),
turbofish: None,
args: vec![],
}
}
fn make_expect_expr() -> PureExpr {
PureExpr::MethodCall {
receiver: Box::new(PureExpr::Path("x".to_string())),
method: "expect".to_string(),
turbofish: None,
args: vec![PureExpr::Lit("\"error\"".to_string())],
}
}
#[test]
fn test_is_unwrap_call() {
let mutation = UnwrapToQuestionMutation::new();
assert!(mutation.is_unwrap_call("unwrap", &[]));
assert!(mutation.is_unwrap_call("expect", &[PureExpr::Lit("\"msg\"".to_string())]));
assert!(!mutation.is_unwrap_call("map", &[]));
}
#[test]
fn test_transform_unwrap() {
let mutation = UnwrapToQuestionMutation::new();
let mut expr = make_unwrap_expr();
let changes = mutation.transform_expr(&mut expr);
assert_eq!(changes, 1);
assert!(matches!(expr, PureExpr::Try(_)));
}
#[test]
fn test_transform_expect() {
let mutation = UnwrapToQuestionMutation::new();
let mut expr = make_expect_expr();
let changes = mutation.transform_expr(&mut expr);
assert_eq!(changes, 1);
assert!(matches!(expr, PureExpr::Try(_)));
}
#[test]
fn test_skip_expect_when_unwrap_only() {
let mutation = UnwrapToQuestionMutation::new().unwrap_only();
let mut expr = make_expect_expr();
let changes = mutation.transform_expr(&mut expr);
assert_eq!(changes, 0);
assert!(matches!(expr, PureExpr::MethodCall { .. }));
}
}