els/
code_action.rs

1use std::collections::HashMap;
2
3use erg_common::consts::{ERG_MODE, PYTHON_MODE};
4use erg_common::deepen_indent;
5use erg_common::traits::{Locational, Stream};
6use erg_compiler::artifact::BuildRunnable;
7use erg_compiler::erg_parser::parse::Parsable;
8use erg_compiler::erg_parser::token::{Token, TokenKind};
9use erg_compiler::hir::Expr;
10use erg_compiler::ty::HasType;
11
12use lsp_types::{
13    CodeAction, CodeActionKind, CodeActionOrCommand, CodeActionParams, CodeActionResponse,
14    Position, Range, TextEdit, Url, WorkspaceEdit,
15};
16
17use crate::server::{ELSResult, RedirectableStdout, Server};
18use crate::util::{self, NormalizedUrl};
19
20impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
21    fn gen_eliminate_unused_vars_action(
22        &self,
23        params: &CodeActionParams,
24    ) -> ELSResult<Option<CodeAction>> {
25        let uri = NormalizedUrl::new(params.text_document.uri.clone());
26        let diags = &params.context.diagnostics;
27        let Some(diag) = diags.first().cloned() else {
28            return Ok(None);
29        };
30        let mut map = HashMap::new();
31        let Some(visitor) = self.get_visitor(&uri) else {
32            self.send_log("visitor not found")?;
33            return Ok(None);
34        };
35        let Some(warns) = self.get_warns(&uri) else {
36            self.send_log("artifact not found")?;
37            return Ok(None);
38        };
39        let warns = warns
40            .iter()
41            .filter(|warn| warn.core.main_message.ends_with("is not used"))
42            .collect::<Vec<_>>();
43        for warn in warns {
44            let uri = NormalizedUrl::new(Url::from_file_path(warn.input.full_path()).unwrap());
45            let Some(pos) = util::loc_to_pos(warn.core.loc) else {
46                continue;
47            };
48            match visitor.get_min_expr(pos) {
49                Some(Expr::Def(def)) => {
50                    let Some(mut range) = util::loc_to_range(def.loc()) else {
51                        self.send_log("range not found")?;
52                        continue;
53                    };
54                    let next = lsp_types::Range {
55                        start: lsp_types::Position {
56                            line: range.end.line,
57                            character: range.end.character,
58                        },
59                        end: lsp_types::Position {
60                            line: range.end.line,
61                            character: range.end.character + 1,
62                        },
63                    };
64                    let code = self.file_cache.get_ranged(&uri, next)?;
65                    match code.as_ref().map(|s| &s[..]) {
66                        None => {
67                            // \n
68                            range.end.line += 1;
69                            range.end.character = 0;
70                        }
71                        Some(";") => range.end.character += 1,
72                        Some(other) => {
73                            crate::_log!(self, "? {other}");
74                        }
75                    }
76                    let edit = TextEdit::new(range, "".to_string());
77                    map.entry(uri.clone().raw()).or_insert(vec![]).push(edit);
78                }
79                Some(_) => {}
80                None => {
81                    let Some(token) = self.file_cache.get_token(&uri, diag.range.start) else {
82                        continue;
83                    };
84                    let Some(vi) = visitor.get_info(&token) else {
85                        continue;
86                    };
87                    if vi.kind.is_parameter() {
88                        let edit = TextEdit::new(diag.range, "_".to_string());
89                        map.entry(uri.clone().raw()).or_insert(vec![]).push(edit);
90                    }
91                }
92            }
93        }
94        let edit = WorkspaceEdit::new(map);
95        let action = CodeAction {
96            title: "Eliminate unused variables".to_string(),
97            kind: Some(CodeActionKind::QUICKFIX),
98            diagnostics: Some(vec![diag]),
99            edit: Some(edit),
100            ..Default::default()
101        };
102        Ok(Some(action))
103    }
104
105    fn gen_change_case_action(
106        &self,
107        token: Token,
108        uri: &NormalizedUrl,
109        params: CodeActionParams,
110    ) -> Option<CodeAction> {
111        let new_text = token.content.to_snake_case().to_string();
112        let mut map = HashMap::new();
113        let visitor = self.get_visitor(uri)?;
114        let def_loc = visitor.get_info(&token)?.def_loc;
115        let edit = TextEdit::new(util::loc_to_range(def_loc.loc)?, new_text.clone());
116        map.insert(uri.clone().raw(), vec![edit]);
117        if let Some(value) = self.shared.index.get_refs(&def_loc) {
118            for refer in value.referrers.iter() {
119                let url = Url::from_file_path(refer.module.as_ref()?).ok()?;
120                let range = util::loc_to_range(refer.loc)?;
121                let edit = TextEdit::new(range, new_text.clone());
122                map.entry(url).or_insert(vec![]).push(edit);
123            }
124        }
125        let edit = WorkspaceEdit::new(map);
126        let action = CodeAction {
127            title: format!("Convert to snake case ({} -> {})", token.content, new_text),
128            kind: Some(CodeActionKind::REFACTOR),
129            diagnostics: Some(params.context.diagnostics),
130            edit: Some(edit),
131            ..Default::default()
132        };
133        Some(action)
134    }
135
136    fn gen_extract_action(&self, params: &CodeActionParams) -> Vec<CodeAction> {
137        let mut actions = vec![];
138        if params.range.start.line == params.range.end.line {
139            if params.range.start.character == params.range.end.character {
140                return vec![];
141            }
142            actions.push(CodeAction {
143                title: "Extract into variable".to_string(),
144                kind: Some(CodeActionKind::REFACTOR_EXTRACT),
145                data: Some(serde_json::to_value(params.clone()).unwrap()),
146                ..Default::default()
147            });
148        }
149        actions.push(CodeAction {
150            title: "Extract into function".to_string(),
151            kind: Some(CodeActionKind::REFACTOR_EXTRACT),
152            data: Some(serde_json::to_value(params.clone()).unwrap()),
153            ..Default::default()
154        });
155        actions
156    }
157
158    fn gen_inline_action(&self, params: &CodeActionParams) -> Option<CodeAction> {
159        let uri = NormalizedUrl::new(params.text_document.uri.clone());
160        let visitor = self.get_visitor(&uri)?;
161        match visitor.get_min_expr(params.range.start)? {
162            Expr::Def(def) => {
163                let title = if def.sig.is_subr() {
164                    "Inline function"
165                } else {
166                    "Inline variable"
167                };
168                let action = CodeAction {
169                    title: title.to_string(),
170                    kind: Some(CodeActionKind::REFACTOR_INLINE),
171                    data: Some(serde_json::to_value(params.clone()).unwrap()),
172                    ..Default::default()
173                };
174                Some(action)
175            }
176            Expr::Accessor(acc) => {
177                let title = if acc.ref_t().is_subr() {
178                    "Inline function"
179                } else {
180                    "Inline variable"
181                };
182                let action = CodeAction {
183                    title: title.to_string(),
184                    kind: Some(CodeActionKind::REFACTOR_INLINE),
185                    data: Some(serde_json::to_value(params.clone()).unwrap()),
186                    ..Default::default()
187                };
188                Some(action)
189            }
190            _ => None,
191        }
192    }
193
194    fn send_normal_action(&self, params: &CodeActionParams) -> ELSResult<Vec<CodeAction>> {
195        let mut actions = vec![];
196        let uri = NormalizedUrl::new(params.text_document.uri.clone());
197        if let Some(token) = self.file_cache.get_token(&uri, params.range.start) {
198            if token.is(TokenKind::Symbol) && !token.is_const() && !token.content.is_snake_case() {
199                let action = self.gen_change_case_action(token, &uri, params.clone());
200                actions.extend(action);
201            }
202        }
203        actions.extend(self.send_quick_fix(params)?);
204        actions.extend(self.gen_extract_action(params));
205        actions.extend(self.gen_inline_action(params));
206        Ok(actions)
207    }
208
209    fn send_quick_fix(&self, params: &CodeActionParams) -> ELSResult<Vec<CodeAction>> {
210        let mut result: Vec<CodeAction> = vec![];
211        let diags = &params.context.diagnostics;
212        if diags.is_empty() {
213            return Ok(result);
214        }
215        if diags
216            .first()
217            .is_some_and(|diag| diag.message.ends_with("is not used"))
218        {
219            let actions = self.gen_eliminate_unused_vars_action(params)?;
220            result.extend(actions);
221        }
222        Ok(result)
223    }
224
225    pub(crate) fn handle_code_action(
226        &mut self,
227        params: CodeActionParams,
228    ) -> ELSResult<Option<CodeActionResponse>> {
229        self.send_log(format!("code action requested: {params:?}"))?;
230        let result = match params
231            .context
232            .only
233            .as_ref()
234            .and_then(|kinds| kinds.first().map(|s| s.as_str()))
235        {
236            Some("quickfix") => self.send_quick_fix(&params)?,
237            None => self.send_normal_action(&params)?,
238            Some(other) => {
239                self.send_log(format!("Unknown code action requested: {other}"))?;
240                vec![]
241            }
242        };
243        Ok(Some(
244            result
245                .into_iter()
246                .map(CodeActionOrCommand::CodeAction)
247                .collect(),
248        ))
249    }
250
251    pub(crate) fn handle_code_action_resolve(
252        &mut self,
253        action: CodeAction,
254    ) -> ELSResult<CodeAction> {
255        self.send_log(format!("code action resolve requested: {action:?}"))?;
256        match &action.title[..] {
257            "Extract into function" | "Extract into variable" => {
258                self.resolve_extract_action(action)
259            }
260            "Inline variable" => self.resolve_inline_variable_action(action),
261            _ => Ok(action),
262        }
263    }
264
265    fn resolve_extract_action(&self, mut action: CodeAction) -> ELSResult<CodeAction> {
266        let params = action
267            .data
268            .take()
269            .and_then(|v| serde_json::from_value::<CodeActionParams>(v).ok())
270            .ok_or("invalid params")?;
271        let extract_function = action.title == "Extract into function";
272        let uri = NormalizedUrl::new(params.text_document.uri);
273        let start = Position::new(params.range.start.line, 0);
274        let mut range = Range::new(start, params.range.end);
275        let indented_code = self.file_cache.get_ranged(&uri, range)?.unwrap_or_default();
276        let indent_len = indented_code.chars().take_while(|c| *c == ' ').count();
277        range.start.character = params.range.start.character;
278        let code = self.file_cache.get_ranged(&uri, range)?.unwrap_or_default();
279        // `    |foo|` (|...| is the selected range) -> `|    foo|`
280        let diff = indented_code.trim_end_matches(&code);
281        let diff_indent_len = diff.chars().take_while(|c| *c == ' ').count();
282        let diff_is_indent = diff.trim().is_empty();
283        let code = if diff_is_indent {
284            range.start.character = 0;
285            indented_code
286        } else {
287            code
288        };
289        let body = if extract_function {
290            let mut code = deepen_indent(code);
291            let should_insert_indent = code.lines().count() == 1 && !diff_is_indent;
292            if should_insert_indent {
293                code = format!("{}{code}", " ".repeat(diff_indent_len));
294            }
295            // add `return` to the last line
296            if PYTHON_MODE {
297                let mut lines = code.lines().collect::<Vec<_>>();
298                let mut last_line = lines.last().unwrap().to_string();
299                let code_start = last_line.chars().position(|c| c != ' ').unwrap();
300                last_line.insert_str(code_start, "return ");
301                lines.pop();
302                lines.push(&last_line);
303                lines.join("\n")
304            } else {
305                code
306            }
307        } else {
308            code.trim_start().to_string()
309        };
310        let sig = match (ERG_MODE, extract_function) {
311            (true, true) => "new_func() =\n",
312            (true, false) => "new_var = ",
313            (false, true) => "def new_func():\n",
314            (false, false) => "new_var = ",
315        };
316        let expanded = if extract_function {
317            "new_func()"
318        } else {
319            "new_var"
320        };
321        let expanded = if range.start.character == 0 {
322            format!("{}{expanded}", " ".repeat(indent_len))
323        } else {
324            expanded.to_string()
325        };
326        let extracted = format!("{}{sig}{body}\n\n", " ".repeat(indent_len));
327        let start = Position::new(range.start.line, 0);
328        let edit1 = TextEdit::new(Range::new(start, start), extracted);
329        let edit2 = TextEdit::new(Range::new(range.start, range.end), expanded);
330        let mut changes = HashMap::new();
331        changes.insert(uri.raw(), vec![edit1, edit2]);
332        action.edit = Some(WorkspaceEdit::new(changes));
333        Ok(action)
334    }
335
336    fn resolve_inline_variable_action(&self, mut action: CodeAction) -> ELSResult<CodeAction> {
337        let params = action
338            .data
339            .take()
340            .and_then(|v| serde_json::from_value::<CodeActionParams>(v).ok())
341            .ok_or("invalid params")?;
342        let uri = NormalizedUrl::new(params.text_document.uri.clone());
343        let visitor = self.get_visitor(&uri).ok_or("get_visitor")?;
344        match visitor
345            .get_min_expr(params.range.start)
346            .ok_or("get_min_expr")?
347        {
348            Expr::Def(def) => {
349                action.edit = Some(WorkspaceEdit::new(self.inline_var_def(def)));
350            }
351            Expr::Accessor(acc) => {
352                let uri = NormalizedUrl::new(
353                    Url::from_file_path(
354                        acc.var_info()
355                            .def_loc
356                            .module
357                            .as_ref()
358                            .ok_or("def_loc.module")?,
359                    )
360                    .map_err(|_| "from_file_path")?,
361                );
362                let visitor = self.get_visitor(&uri).ok_or(format!("{uri} not found"))?;
363                let range = util::loc_to_range(acc.var_info().def_loc.loc).ok_or("loc_to_range")?;
364                if let Some(Expr::Def(def)) = visitor.get_min_expr(range.start) {
365                    action.edit = Some(WorkspaceEdit::new(self.inline_var_def(def)));
366                }
367            }
368            _ => {}
369        }
370        Ok(action)
371    }
372
373    fn inline_var_def(&self, def: &erg_compiler::hir::Def) -> HashMap<Url, Vec<TextEdit>> {
374        let mut changes = HashMap::new();
375        let mut range = util::loc_to_range(def.loc()).unwrap();
376        range.end.character = u32::MAX;
377        let delete = TextEdit::new(range, "".to_string());
378        let uri = NormalizedUrl::new(
379            Url::from_file_path(def.sig.ident().vi.def_loc.module.as_ref().unwrap()).unwrap(),
380        );
381        let range = util::loc_to_range(def.body.block.loc()).unwrap();
382        let code = self.file_cache.get_ranged(&uri, range).unwrap().unwrap();
383        let expr = def.body.block.first().unwrap();
384        let code = if expr.need_to_be_closed() {
385            format!("({code})")
386        } else {
387            code
388        };
389        changes.insert(uri.raw(), vec![delete]);
390        if let Some(index) = self.shared.index.get_refs(&def.sig.ident().vi.def_loc) {
391            for ref_ in index.referrers.iter() {
392                let Some(path) = ref_.module.as_ref() else {
393                    continue;
394                };
395                let edit = TextEdit::new(util::loc_to_range(ref_.loc).unwrap(), code.clone());
396                let uri = NormalizedUrl::new(Url::from_file_path(path).unwrap());
397                changes.entry(uri.raw()).or_insert(vec![]).push(edit);
398            }
399        }
400        changes
401    }
402}
els/code_action.rs

els/
code_action.rs
