feat(pylyzer): pseudo-method completion

erg-lang · Sep 27, 2024 · 1baa040 · 1baa040
1 parent 1954d27
commit 1baa040
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 5 deletions.
diff --git a/crates/els/completion.rs b/crates/els/completion.rs
@@ -1,6 +1,5 @@
 use std::path::Path;
 
-use lsp_types::CompletionResponse;
 use serde_json::Value;
 
 use erg_common::config::ErgConfig;
@@ -27,13 +26,13 @@ use erg_compiler::varinfo::{AbsLocation, Mutability, VarInfo, VarKind};
 use TokenKind::*;
 
 use lsp_types::{
-    CompletionItem, CompletionItemKind, CompletionParams, Documentation, MarkedString,
-    MarkupContent, MarkupKind, Position, Range, TextEdit,
+    CompletionItem, CompletionItemKind, CompletionParams, CompletionResponse, Documentation,
+    MarkedString, MarkupContent, MarkupKind, Position, Range, TextEdit,
 };
 
 use crate::_log;
 use crate::server::{ELSResult, Flags, RedirectableStdout, Server};
-use crate::util::{self, loc_to_pos, NormalizedUrl};
+use crate::util::{self, loc_to_pos, loc_to_range, NormalizedUrl};
 
 fn comp_item_kind(t: &Type, muty: Mutability) -> CompletionItemKind {
     match t {
@@ -682,6 +681,20 @@ impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
             .set(&mut item);
             item.kind = Some(comp_item_kind(&vi.t, vi.muty));
             item.data = Some(Value::String(vi.def_loc.to_string()));
+            // s.`Function::map` => map(s)
+            if comp_kind.should_be_method() && item.label.starts_with("Function::") {
+                let receiver = self.get_receiver(&uri, pos)?;
+                if let Some(mut range) = receiver.as_ref().and_then(|expr| loc_to_range(expr.loc()))
+                {
+                    // FIXME:
+                    let s_receiver = self.file_cache.get_ranged(&uri, range)?.unwrap_or_default();
+                    range.end.character += 1;
+                    let name = item.label.trim_start_matches("Function::");
+                    let remove = TextEdit::new(range, "".to_string());
+                    item.insert_text = Some(format!("{name}({s_receiver})"));
+                    item.additional_text_edits = Some(vec![remove]);
+                }
+            }
             already_appeared.insert(item.label.clone());
             result.push(item);
         }

diff --git a/crates/els/server.rs b/crates/els/server.rs
@@ -26,7 +26,7 @@ use erg_compiler::context::{Context, ModuleContext};
 use erg_compiler::erg_parser::ast::Module;
 use erg_compiler::erg_parser::parse::{Parsable, SimpleParser};
 use erg_compiler::error::CompileWarning;
-use erg_compiler::hir::HIR;
+use erg_compiler::hir::{Expr, HIR};
 use erg_compiler::lower::ASTLowerer;
 use erg_compiler::module::{IRs, ModuleEntry, SharedCompilerResource};
 use erg_compiler::ty::{HasType, Type};
@@ -1096,6 +1096,48 @@ impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
         ctxs
     }
 
+    pub(crate) fn get_receiver(
+        &self,
+        uri: &NormalizedUrl,
+        attr_marker_pos: Position,
+    ) -> ELSResult<Option<Expr>> {
+        let Some(module) = self.get_raw_mod_ctx(uri, Duration::from_millis(100)) else {
+            return Ok(None);
+        };
+        let maybe_token = self.file_cache.get_receiver(uri, attr_marker_pos);
+        if let Some(token) = maybe_token {
+            // _log!(self, "token: {token}");
+            let mut ctxs = vec![];
+            let expr = if let Some(visitor) = self.get_visitor(uri) {
+                if let Some(expr) =
+                    loc_to_pos(token.loc()).and_then(|pos| visitor.get_min_expr(pos).cloned())
+                {
+                    let type_ctxs = module
+                        .context
+                        .get_nominal_super_type_ctxs(expr.ref_t())
+                        .unwrap_or(vec![]);
+                    ctxs.extend(type_ctxs.into_iter().map(|ctx| &ctx.ctx));
+                    if let Ok(singular_ctxs) = module
+                        .context
+                        .get_singular_ctxs_by_hir_expr(&expr, &module.context)
+                    {
+                        ctxs.extend(singular_ctxs);
+                    }
+                    Some(expr)
+                } else {
+                    _log!(self, "expr not found: {token}");
+                    None
+                }
+            } else {
+                None
+            };
+            Ok(expr)
+        } else {
+            self.send_log("token not found")?;
+            Ok(None)
+        }
+    }
+
     pub(crate) fn get_receiver_and_ctxs(
         &self,
         uri: &NormalizedUrl,

diff --git a/crates/els/util.rs b/crates/els/util.rs
@@ -91,6 +91,15 @@ pub(crate) fn loc_to_range(loc: erg_common::error::Location) -> Option<Range> {
     Some(Range::new(start, end))
 }
 
+pub(crate) fn _range_to_loc(range: Range) -> erg_common::error::Location {
+    erg_common::error::Location::range(
+        range.start.line + 1,
+        range.start.character,
+        range.end.line + 1,
+        range.end.character,
+    )
+}
+
 pub(crate) fn loc_to_pos(loc: erg_common::error::Location) -> Option<Position> {
     // FIXME: should `Position::new(loc.ln_begin()? - 1, loc.col_begin()?)`
     // but completion doesn't work (because the newline will be included)