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

   1 use crate::assist_context::{AssistContext, Assists};
   2 use hir::{HasVisibility, HirDisplay, Module};
   3 use ide_db::{
   4     assists::{AssistId, AssistKind},
   5     base_db::{FileId, Upcast},
   6     defs::{Definition, NameRefClass},
   7 };
   8 use syntax::{
   9     ast::{self, edit::IndentLevel, NameRef},
  10     AstNode, Direction, SyntaxKind, TextSize,
  11 };
  12
  13 // Assist: generate_constant
  14 //
  15 // Generate a named constant.
  16 //
  17 // ```
  18 // struct S { i: usize }
  19 // impl S { pub fn new(n: usize) {} }
  20 // fn main() {
  21 //     let v = S::new(CAPA$0CITY);
  22 // }
  23 // ```
  24 // ->
  25 // ```
  26 // struct S { i: usize }
  27 // impl S { pub fn new(n: usize) {} }
  28 // fn main() {
  29 //     const CAPACITY: usize = $0;
  30 //     let v = S::new(CAPACITY);
  31 // }
  32 // ```
  33
  34 pub(crate) fn generate_constant(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  35     let constant_token = ctx.find_node_at_offset::<ast::NameRef>()?;
  36     if constant_token.to_string().chars().any(|it| !(it.is_uppercase() || it == '_')) {
  37         cov_mark::hit!(not_constant_name);
  38         return None;
  39     }
  40     if NameRefClass::classify(&ctx.sema, &constant_token).is_some() {
  41         cov_mark::hit!(already_defined);
  42         return None;
  43     }
  44     let expr = constant_token.syntax().ancestors().find_map(ast::Expr::cast)?;
  45     let statement = expr.syntax().ancestors().find_map(ast::Stmt::cast)?;
  46     let ty = ctx.sema.type_of_expr(&expr)?;
  47     let scope = ctx.sema.scope(statement.syntax())?;
  48     let constant_module = scope.module();
  49     let type_name = ty.original().display_source_code(ctx.db(), constant_module.into()).ok()?;
  50     let target = statement.syntax().parent()?.text_range();
  51     let path = constant_token.syntax().ancestors().find_map(ast::Path::cast)?;
  52
  53     let name_refs = path.segments().map(|s| s.name_ref());
  54     let mut outer_exists = false;
  55     let mut not_exist_name_ref = Vec::new();
  56     let mut current_module = constant_module;
  57     for name_ref in name_refs {
  58         let name_ref_value = name_ref?;
  59         let name_ref_class = NameRefClass::classify(&ctx.sema, &name_ref_value);
  60         match name_ref_class {
  61             Some(NameRefClass::Definition(Definition::Module(m))) => {
  62                 if !m.visibility(ctx.sema.db).is_visible_from(ctx.sema.db, constant_module.into()) {
  63                     return None;
  64                 }
  65                 outer_exists = true;
  66                 current_module = m;
  67             }
  68             Some(_) => {
  69                 return None;
  70             }
  71             None => {
  72                 not_exist_name_ref.push(name_ref_value);
  73             }
  74         }
  75     }
  76     let (offset, indent, file_id, post_string) =
  77         target_data_for_generate_constant(ctx, current_module, constant_module).unwrap_or_else(
  78             || {
  79                 let indent = IndentLevel::from_node(statement.syntax());
  80                 (statement.syntax().text_range().start(), indent, None, format!("\n{}", indent))
  81             },
  82         );
  83
  84     let text = get_text_for_generate_constant(not_exist_name_ref, indent, outer_exists, type_name)?;
  85     acc.add(
  86         AssistId("generate_constant", AssistKind::QuickFix),
  87         "Generate constant",
  88         target,
  89         |builder| {
  90             if let Some(file_id) = file_id {
  91                 builder.edit_file(file_id);
  92             }
  93             builder.insert(offset, format!("{}{}", text, post_string));
  94         },
  95     )
  96 }
  97
  98 fn get_text_for_generate_constant(
  99     mut not_exist_name_ref: Vec<NameRef>,
 100     indent: IndentLevel,
 101     outer_exists: bool,
 102     type_name: String,
 103 ) -> Option<String> {
 104     let constant_token = not_exist_name_ref.pop()?;
 105     let vis = if not_exist_name_ref.len() == 0 && !outer_exists { "" } else { "\npub " };
 106     let mut text = format!("{}const {}: {} = $0;", vis, constant_token, type_name);
 107     while let Some(name_ref) = not_exist_name_ref.pop() {
 108         let vis = if not_exist_name_ref.len() == 0 && !outer_exists { "" } else { "\npub " };
 109         text = text.replace("\n", "\n    ");
 110         text = format!("{}mod {} {{{}\n}}", vis, name_ref.to_string(), text);
 111     }
 112     Some(text.replace("\n", &format!("\n{}", indent)))
 113 }
 114
 115 fn target_data_for_generate_constant(
 116     ctx: &AssistContext<'_>,
 117     current_module: Module,
 118     constant_module: Module,
 119 ) -> Option<(TextSize, IndentLevel, Option<FileId>, String)> {
 120     if current_module == constant_module {
 121         // insert in current file
 122         return None;
 123     }
 124     let in_file_source = current_module.definition_source(ctx.sema.db);
 125     let file_id = in_file_source.file_id.original_file(ctx.sema.db.upcast());
 126     match in_file_source.value {
 127         hir::ModuleSource::Module(module_node) => {
 128             let indent = IndentLevel::from_node(module_node.syntax());
 129             let l_curly_token = module_node.item_list()?.l_curly_token()?;
 130             let offset = l_curly_token.text_range().end();
 131
 132             let siblings_has_newline = l_curly_token
 133                 .siblings_with_tokens(Direction::Next)
 134                 .find(|it| it.kind() == SyntaxKind::WHITESPACE && it.to_string().contains("\n"))
 135                 .is_some();
 136             let post_string =
 137                 if siblings_has_newline { format!("{}", indent) } else { format!("\n{}", indent) };
 138             Some((offset, indent + 1, Some(file_id), post_string))
 139         }
 140         _ => Some((TextSize::from(0), 0.into(), Some(file_id), "\n".into())),
 141     }
 142 }
 143
 144 #[cfg(test)]
 145 mod tests {
 146     use super::*;
 147     use crate::tests::{check_assist, check_assist_not_applicable};
 148
 149     #[test]
 150     fn test_trivial() {
 151         check_assist(
 152             generate_constant,
 153             r#"struct S { i: usize }
 154 impl S {
 155     pub fn new(n: usize) {}
 156 }
 157 fn main() {
 158     let v = S::new(CAPA$0CITY);
 159 }"#,
 160             r#"struct S { i: usize }
 161 impl S {
 162     pub fn new(n: usize) {}
 163 }
 164 fn main() {
 165     const CAPACITY: usize = $0;
 166     let v = S::new(CAPACITY);
 167 }"#,
 168         );
 169     }
 170     #[test]
 171     fn test_wont_apply_when_defined() {
 172         cov_mark::check!(already_defined);
 173         check_assist_not_applicable(
 174             generate_constant,
 175             r#"struct S { i: usize }
 176 impl S {
 177     pub fn new(n: usize) {}
 178 }
 179 fn main() {
 180     const CAPACITY: usize = 10;
 181     let v = S::new(CAPAC$0ITY);
 182 }"#,
 183         );
 184     }
 185     #[test]
 186     fn test_wont_apply_when_maybe_not_constant() {
 187         cov_mark::check!(not_constant_name);
 188         check_assist_not_applicable(
 189             generate_constant,
 190             r#"struct S { i: usize }
 191 impl S {
 192     pub fn new(n: usize) {}
 193 }
 194 fn main() {
 195     let v = S::new(capa$0city);
 196 }"#,
 197         );
 198     }
 199
 200     #[test]
 201     fn test_constant_with_path() {
 202         check_assist(
 203             generate_constant,
 204             r#"mod foo {}
 205 fn bar() -> i32 {
 206     foo::A_CON$0STANT
 207 }"#,
 208             r#"mod foo {
 209     pub const A_CONSTANT: i32 = $0;
 210 }
 211 fn bar() -> i32 {
 212     foo::A_CONSTANT
 213 }"#,
 214         );
 215     }
 216
 217     #[test]
 218     fn test_constant_with_longer_path() {
 219         check_assist(
 220             generate_constant,
 221             r#"mod foo {
 222     pub mod goo {}
 223 }
 224 fn bar() -> i32 {
 225     foo::goo::A_CON$0STANT
 226 }"#,
 227             r#"mod foo {
 228     pub mod goo {
 229         pub const A_CONSTANT: i32 = $0;
 230     }
 231 }
 232 fn bar() -> i32 {
 233     foo::goo::A_CONSTANT
 234 }"#,
 235         );
 236     }
 237
 238     #[test]
 239     fn test_constant_with_not_exist_longer_path() {
 240         check_assist(
 241             generate_constant,
 242             r#"fn bar() -> i32 {
 243     foo::goo::A_CON$0STANT
 244 }"#,
 245             r#"mod foo {
 246     pub mod goo {
 247         pub const A_CONSTANT: i32 = $0;
 248     }
 249 }
 250 fn bar() -> i32 {
 251     foo::goo::A_CONSTANT
 252 }"#,
 253         );
 254     }
 255 }