src/tools/rust-analyzer/crates/ide-assists/src/handlers/move_format_string_arg.rs

   1 use crate::{AssistContext, Assists};
   2 use ide_db::{
   3     assists::{AssistId, AssistKind},
   4     syntax_helpers::{
   5         format_string::is_format_string,
   6         format_string_exprs::{parse_format_exprs, Arg},
   7     },
   8 };
   9 use itertools::Itertools;
  10 use stdx::format_to;
  11 use syntax::{ast, AstNode, AstToken, NodeOrToken, SyntaxKind::COMMA, TextRange};
  12
  13 // Assist: move_format_string_arg
  14 //
  15 // Move an expression out of a format string.
  16 //
  17 // ```
  18 // macro_rules! format_args {
  19 //     ($lit:literal $(tt:tt)*) => { 0 },
  20 // }
  21 // macro_rules! print {
  22 //     ($($arg:tt)*) => (std::io::_print(format_args!($($arg)*)));
  23 // }
  24 //
  25 // fn main() {
  26 //     print!("{x + 1}$0");
  27 // }
  28 // ```
  29 // ->
  30 // ```
  31 // macro_rules! format_args {
  32 //     ($lit:literal $(tt:tt)*) => { 0 },
  33 // }
  34 // macro_rules! print {
  35 //     ($($arg:tt)*) => (std::io::_print(format_args!($($arg)*)));
  36 // }
  37 //
  38 // fn main() {
  39 //     print!("{}"$0, x + 1);
  40 // }
  41 // ```
  42
  43 pub(crate) fn move_format_string_arg(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  44     let fmt_string = ctx.find_token_at_offset::<ast::String>()?;
  45     let tt = fmt_string.syntax().parent().and_then(ast::TokenTree::cast)?;
  46
  47     let expanded_t = ast::String::cast(
  48         ctx.sema.descend_into_macros_with_kind_preference(fmt_string.syntax().clone()),
  49     )?;
  50     if !is_format_string(&expanded_t) {
  51         return None;
  52     }
  53
  54     let (new_fmt, extracted_args) = parse_format_exprs(fmt_string.text()).ok()?;
  55     if extracted_args.is_empty() {
  56         return None;
  57     }
  58
  59     acc.add(
  60         AssistId(
  61             "move_format_string_arg",
  62             // if there aren't any expressions, then make the assist a RefactorExtract
  63             if extracted_args.iter().filter(|f| matches!(f, Arg::Expr(_))).count() == 0 {
  64                 AssistKind::RefactorExtract
  65             } else {
  66                 AssistKind::QuickFix
  67             },
  68         ),
  69         "Extract format args",
  70         tt.syntax().text_range(),
  71         |edit| {
  72             let fmt_range = fmt_string.syntax().text_range();
  73
  74             // Replace old format string with new format string whose arguments have been extracted
  75             edit.replace(fmt_range, new_fmt);
  76
  77             // Insert cursor at end of format string
  78             edit.insert(fmt_range.end(), "$0");
  79
  80             // Extract existing arguments in macro
  81             let tokens =
  82                 tt.token_trees_and_tokens().collect_vec();
  83
  84             let mut existing_args: Vec<String> = vec![];
  85
  86             let mut current_arg = String::new();
  87             if let [_opening_bracket, NodeOrToken::Token(format_string), _args_start_comma, tokens @ .., NodeOrToken::Token(end_bracket)] =
  88                 tokens.as_slice()
  89             {
  90                 for t in tokens {
  91                     match t {
  92                         NodeOrToken::Node(n) => {
  93                             format_to!(current_arg, "{n}");
  94                         },
  95                         NodeOrToken::Token(t) if t.kind() == COMMA=> {
  96                             existing_args.push(current_arg.trim().into());
  97                             current_arg.clear();
  98                         },
  99                         NodeOrToken::Token(t) => {
 100                             current_arg.push_str(t.text());
 101                         },
 102                     }
 103                 }
 104                 existing_args.push(current_arg.trim().into());
 105
 106                 // delete everything after the format string till end bracket
 107                 // we're going to insert the new arguments later
 108                 edit.delete(TextRange::new(
 109                     format_string.text_range().end(),
 110                     end_bracket.text_range().start(),
 111                 ));
 112             }
 113
 114             // Start building the new args
 115             let mut existing_args = existing_args.into_iter();
 116             let mut args = String::new();
 117
 118             let mut placeholder_idx = 1;
 119
 120             for extracted_args in extracted_args {
 121                 // remove expr from format string
 122                 args.push_str(", ");
 123
 124                 match extracted_args {
 125                     Arg::Ident(s) | Arg::Expr(s) => {
 126                         // insert arg
 127                         args.push_str(&s);
 128                     }
 129                     Arg::Placeholder => {
 130                         // try matching with existing argument
 131                         match existing_args.next() {
 132                             Some(ea) => {
 133                                 args.push_str(&ea);
 134                             }
 135                             None => {
 136                                 // insert placeholder
 137                                 args.push_str(&format!("${placeholder_idx}"));
 138                                 placeholder_idx += 1;
 139                             }
 140                         }
 141                     }
 142                 }
 143             }
 144
 145             // Insert new args
 146             edit.insert(fmt_range.end(), args);
 147         },
 148     );
 149
 150     Some(())
 151 }
 152
 153 #[cfg(test)]
 154 mod tests {
 155     use super::*;
 156     use crate::tests::check_assist;
 157
 158     const MACRO_DECL: &'static str = r#"
 159 macro_rules! format_args {
 160     ($lit:literal $(tt:tt)*) => { 0 },
 161 }
 162 macro_rules! print {
 163     ($($arg:tt)*) => (std::io::_print(format_args!($($arg)*)));
 164 }
 165 "#;
 166
 167     fn add_macro_decl(s: &'static str) -> String {
 168         MACRO_DECL.to_string() + s
 169     }
 170
 171     #[test]
 172     fn multiple_middle_arg() {
 173         check_assist(
 174             move_format_string_arg,
 175             &add_macro_decl(
 176                 r#"
 177 fn main() {
 178     print!("{} {x + 1:b} {}$0", y + 2, 2);
 179 }
 180 "#,
 181             ),
 182             &add_macro_decl(
 183                 r#"
 184 fn main() {
 185     print!("{} {:b} {}"$0, y + 2, x + 1, 2);
 186 }
 187 "#,
 188             ),
 189         );
 190     }
 191
 192     #[test]
 193     fn single_arg() {
 194         check_assist(
 195             move_format_string_arg,
 196             &add_macro_decl(
 197                 r#"
 198 fn main() {
 199     print!("{obj.value:b}$0",);
 200 }
 201 "#,
 202             ),
 203             &add_macro_decl(
 204                 r#"
 205 fn main() {
 206     print!("{:b}"$0, obj.value);
 207 }
 208 "#,
 209             ),
 210         );
 211     }
 212
 213     #[test]
 214     fn multiple_middle_placeholders_arg() {
 215         check_assist(
 216             move_format_string_arg,
 217             &add_macro_decl(
 218                 r#"
 219 fn main() {
 220     print!("{} {x + 1:b} {} {}$0", y + 2, 2);
 221 }
 222 "#,
 223             ),
 224             &add_macro_decl(
 225                 r#"
 226 fn main() {
 227     print!("{} {:b} {} {}"$0, y + 2, x + 1, 2, $1);
 228 }
 229 "#,
 230             ),
 231         );
 232     }
 233
 234     #[test]
 235     fn multiple_trailing_args() {
 236         check_assist(
 237             move_format_string_arg,
 238             &add_macro_decl(
 239                 r#"
 240 fn main() {
 241     print!("{} {x + 1:b} {Struct(1, 2)}$0", 1);
 242 }
 243 "#,
 244             ),
 245             &add_macro_decl(
 246                 r#"
 247 fn main() {
 248     print!("{} {:b} {}"$0, 1, x + 1, Struct(1, 2));
 249 }
 250 "#,
 251             ),
 252         );
 253     }
 254
 255     #[test]
 256     fn improper_commas() {
 257         check_assist(
 258             move_format_string_arg,
 259             &add_macro_decl(
 260                 r#"
 261 fn main() {
 262     print!("{} {x + 1:b} {Struct(1, 2)}$0", 1,);
 263 }
 264 "#,
 265             ),
 266             &add_macro_decl(
 267                 r#"
 268 fn main() {
 269     print!("{} {:b} {}"$0, 1, x + 1, Struct(1, 2));
 270 }
 271 "#,
 272             ),
 273         );
 274     }
 275
 276     #[test]
 277     fn nested_tt() {
 278         check_assist(
 279             move_format_string_arg,
 280             &add_macro_decl(
 281                 r#"
 282 fn main() {
 283     print!("My name is {} {x$0 + x}", stringify!(Paperino))
 284 }
 285 "#,
 286             ),
 287             &add_macro_decl(
 288                 r#"
 289 fn main() {
 290     print!("My name is {} {}"$0, stringify!(Paperino), x + x)
 291 }
 292 "#,
 293             ),
 294         );
 295     }
 296 }