|
| 1 | +use either::Either; |
| 2 | +use itertools::Itertools; |
| 3 | +use std::iter::successors; |
| 4 | +use syntax::{ |
| 5 | + AstNode, NodeOrToken, SyntaxKind, SyntaxNode, SyntaxToken, T, |
| 6 | + ast::{self, edit::AstNodeEdit, make}, |
| 7 | +}; |
| 8 | + |
| 9 | +use crate::{AssistContext, AssistId, Assists}; |
| 10 | + |
| 11 | +// Assist: replace_if_matches_to_if_let |
| 12 | +// |
| 13 | +// Replace matches!() to let-chain in if condition. |
| 14 | +// |
| 15 | +// ``` |
| 16 | +// fn foo(x: Option<i32>) { |
| 17 | +// if let Some(n) = x && $0matches!(n.checked_div(2), Some(5..8)) { |
| 18 | +// } |
| 19 | +// } |
| 20 | +// ``` |
| 21 | +// -> |
| 22 | +// ``` |
| 23 | +// fn foo(x: Option<i32>) { |
| 24 | +// if let Some(n) = x && let Some(5..8) = n.checked_div(2) { |
| 25 | +// } |
| 26 | +// } |
| 27 | +// ``` |
| 28 | +pub(crate) fn replace_if_matches_to_if_let( |
| 29 | + acc: &mut Assists, |
| 30 | + ctx: &AssistContext<'_>, |
| 31 | +) -> Option<()> { |
| 32 | + if !ctx.has_empty_selection() { |
| 33 | + return None; |
| 34 | + } |
| 35 | + let macro_call_expr = ctx.find_node_at_offset::<ast::MacroExpr>()?; |
| 36 | + let macro_call = macro_call_expr.macro_call()?; |
| 37 | + |
| 38 | + if macro_call.path()?.segment()?.name_ref()?.text() != "matches" { |
| 39 | + return None; |
| 40 | + } |
| 41 | + let has_cond_expr = has_cond_expr_of_let_chain(macro_call_expr.syntax())?; |
| 42 | + let condition = either::for_both!(&has_cond_expr, it => it.condition())?; |
| 43 | + |
| 44 | + let token_tree = macro_call.token_tree()?.clone_for_update(); |
| 45 | + let tts = tt_content(token_tree); |
| 46 | + let (expr_tts, pat_tts, guard_tts) = split_matches_args(&tts)?; |
| 47 | + |
| 48 | + let target = macro_call.syntax().text_range(); |
| 49 | + acc.add( |
| 50 | + AssistId::refactor_rewrite("replace_if_matches_to_if_let"), |
| 51 | + "Replace matches to let-chain", |
| 52 | + target, |
| 53 | + |builder| { |
| 54 | + let mut edit = builder.make_editor(macro_call.syntax()); |
| 55 | + |
| 56 | + let mut new_tts = |
| 57 | + vec![make::token(T![let]).into(), make::tokens::whitespace(" ").into()]; |
| 58 | + new_tts.extend(pat_tts.iter().map(to_syntax_element)); |
| 59 | + new_tts.extend([ |
| 60 | + make::tokens::whitespace(" ").into(), |
| 61 | + make::token(T![=]).into(), |
| 62 | + make::tokens::whitespace(" ").into(), |
| 63 | + ]); |
| 64 | + new_tts.extend(expr_tts.iter().map(to_syntax_element)); |
| 65 | + |
| 66 | + if let Some(guard_tts) = guard_tts { |
| 67 | + let whitespace = if condition.syntax().text().contains_char('\n') { |
| 68 | + let indent = has_cond_expr.indent_level() + 1; |
| 69 | + format!("\n{indent}") |
| 70 | + } else { |
| 71 | + " ".to_string() |
| 72 | + }; |
| 73 | + new_tts.extend( |
| 74 | + [ |
| 75 | + make::tokens::whitespace(&whitespace).into(), |
| 76 | + make::token(T![&&]).into(), |
| 77 | + make::tokens::whitespace(" ").into(), |
| 78 | + ] |
| 79 | + .into_iter() |
| 80 | + .chain(guard_tts.iter().map(to_syntax_element)), |
| 81 | + ); |
| 82 | + } |
| 83 | + |
| 84 | + edit.replace_with_many(macro_call.syntax(), new_tts); |
| 85 | + |
| 86 | + builder.add_file_edits(ctx.vfs_file_id(), edit); |
| 87 | + }, |
| 88 | + ) |
| 89 | +} |
| 90 | + |
| 91 | +type TT = NodeOrToken<ast::TokenTree, SyntaxToken>; |
| 92 | + |
| 93 | +fn split_matches_args(tts: &[TT]) -> Option<(&[TT], &[TT], Option<&[TT]>)> { |
| 94 | + let (expr_tts, rest_tts) = |
| 95 | + tts.split(|tt| tt.as_token().is_some_and(|it| it.kind() == T![,])).next_tuple()?; |
| 96 | + let (pat_tts, guard_tts) = rest_tts |
| 97 | + .split(|tt| tt.as_token().is_some_and(|it| it.kind() == T![if])) |
| 98 | + .next_tuple() |
| 99 | + .unzip(); |
| 100 | + let pat_tts = pat_tts.unwrap_or(rest_tts); |
| 101 | + Some((expr_tts, trim_tts(pat_tts), guard_tts.map(trim_tts))) |
| 102 | +} |
| 103 | + |
| 104 | +fn trim_tts(mut tts: &[TT]) -> &[TT] { |
| 105 | + let is_whitespace: fn(&(&TT, &[TT])) -> bool = |
| 106 | + |&(tt, _)| tt.as_token().is_some_and(|it| it.kind() == SyntaxKind::WHITESPACE); |
| 107 | + while let Some((_, rest)) = |
| 108 | + tts.split_first().filter(is_whitespace).or_else(|| tts.split_last().filter(is_whitespace)) |
| 109 | + { |
| 110 | + tts = rest |
| 111 | + } |
| 112 | + tts |
| 113 | +} |
| 114 | + |
| 115 | +fn tt_content(token_tree: ast::TokenTree) -> Vec<TT> { |
| 116 | + token_tree |
| 117 | + .token_trees_and_tokens() |
| 118 | + .skip(1) |
| 119 | + .take_while(|it| { |
| 120 | + it.as_token().is_none_or(|t| { |
| 121 | + !matches!(t.kind(), SyntaxKind::R_PAREN | SyntaxKind::R_CURLY | SyntaxKind::R_BRACK) |
| 122 | + }) |
| 123 | + }) |
| 124 | + .collect() |
| 125 | +} |
| 126 | + |
| 127 | +fn to_syntax_element( |
| 128 | + tt: &NodeOrToken<ast::TokenTree, SyntaxToken>, |
| 129 | +) -> NodeOrToken<SyntaxNode, SyntaxToken> { |
| 130 | + match tt { |
| 131 | + NodeOrToken::Node(node) => NodeOrToken::Node(node.syntax().clone()), |
| 132 | + NodeOrToken::Token(tok) => NodeOrToken::Token(tok.clone()), |
| 133 | + } |
| 134 | +} |
| 135 | + |
| 136 | +fn has_cond_expr_of_let_chain(node: &SyntaxNode) -> Option<Either<ast::IfExpr, ast::WhileExpr>> { |
| 137 | + let condition = successors(Some(node.clone()), |node| { |
| 138 | + let parent = node.parent()?; |
| 139 | + let bin_expr = ast::BinExpr::cast(parent)?; |
| 140 | + let ast::BinaryOp::LogicOp(ast::LogicOp::And) = bin_expr.op_kind()? else { return None }; |
| 141 | + Some(bin_expr.syntax().clone()) |
| 142 | + }) |
| 143 | + .last()?; |
| 144 | + AstNode::cast(condition.parent()?) |
| 145 | +} |
| 146 | + |
| 147 | +#[cfg(test)] |
| 148 | +mod tests { |
| 149 | + use crate::tests::check_assist; |
| 150 | + |
| 151 | + use super::*; |
| 152 | + |
| 153 | + #[test] |
| 154 | + fn test_replace_if_matches_to_if_let() { |
| 155 | + check_assist( |
| 156 | + replace_if_matches_to_if_let, |
| 157 | + " |
| 158 | +fn foo(x: Option<i32>) { |
| 159 | + if $0matches!(n.checked_div(2), Some(5..8)) { |
| 160 | + } |
| 161 | +} |
| 162 | + ", |
| 163 | + " |
| 164 | +fn foo(x: Option<i32>) { |
| 165 | + if let Some(5..8) = n.checked_div(2) { |
| 166 | + } |
| 167 | +} |
| 168 | + ", |
| 169 | + ); |
| 170 | + } |
| 171 | + |
| 172 | + #[test] |
| 173 | + fn test_replace_if_matches_to_if_let_has_guard() { |
| 174 | + check_assist( |
| 175 | + replace_if_matches_to_if_let, |
| 176 | + " |
| 177 | +fn foo(x: Option<i32>) { |
| 178 | + if $0matches!(n.checked_div(2), Some(m) if m > 8) { |
| 179 | + } |
| 180 | +} |
| 181 | + ", |
| 182 | + " |
| 183 | +fn foo(x: Option<i32>) { |
| 184 | + if let Some(m) = n.checked_div(2) && m > 8 { |
| 185 | + } |
| 186 | +} |
| 187 | + ", |
| 188 | + ); |
| 189 | + |
| 190 | + check_assist( |
| 191 | + replace_if_matches_to_if_let, |
| 192 | + " |
| 193 | +fn foo(x: Option<i32>) { |
| 194 | + if true && $0matches!(n.checked_div(2), Some(m) if m > 8) { |
| 195 | + } |
| 196 | +} |
| 197 | + ", |
| 198 | + " |
| 199 | +fn foo(x: Option<i32>) { |
| 200 | + if true && let Some(m) = n.checked_div(2) && m > 8 { |
| 201 | + } |
| 202 | +} |
| 203 | + ", |
| 204 | + ); |
| 205 | + |
| 206 | + check_assist( |
| 207 | + replace_if_matches_to_if_let, |
| 208 | + " |
| 209 | +fn foo(x: Option<i32>) { |
| 210 | + if true |
| 211 | + && $0matches!(n.checked_div(2), Some(m) if m > 8) |
| 212 | + { |
| 213 | + } |
| 214 | +} |
| 215 | + ", |
| 216 | + " |
| 217 | +fn foo(x: Option<i32>) { |
| 218 | + if true |
| 219 | + && let Some(m) = n.checked_div(2) |
| 220 | + && m > 8 |
| 221 | + { |
| 222 | + } |
| 223 | +} |
| 224 | + ", |
| 225 | + ); |
| 226 | + } |
| 227 | + |
| 228 | + #[test] |
| 229 | + fn test_replace_if_matches_to_if_let_in_let_chain() { |
| 230 | + check_assist( |
| 231 | + replace_if_matches_to_if_let, |
| 232 | + " |
| 233 | +fn foo(x: Option<i32>) { |
| 234 | + if let Some(n) = x && $0matches!(n.checked_div(2), Some(5..8)) { |
| 235 | + } |
| 236 | +} |
| 237 | + ", |
| 238 | + " |
| 239 | +fn foo(x: Option<i32>) { |
| 240 | + if let Some(n) = x && let Some(5..8) = n.checked_div(2) { |
| 241 | + } |
| 242 | +} |
| 243 | + ", |
| 244 | + ); |
| 245 | + |
| 246 | + check_assist( |
| 247 | + replace_if_matches_to_if_let, |
| 248 | + " |
| 249 | +fn foo(x: Option<i32>) { |
| 250 | + if let Some(n) = x && true && $0matches!(n.checked_div(2), Some(5..8)) { |
| 251 | + } |
| 252 | +} |
| 253 | + ", |
| 254 | + " |
| 255 | +fn foo(x: Option<i32>) { |
| 256 | + if let Some(n) = x && true && let Some(5..8) = n.checked_div(2) { |
| 257 | + } |
| 258 | +} |
| 259 | + ", |
| 260 | + ); |
| 261 | + } |
| 262 | +} |
0 commit comments