chore(els): improve patch support

crates/els/hir_visitor.rs

@@ -121,11 +121,11 @@ impl<'a> HIRVisitor<'a> {
             | Expr::Set(_)
             | Expr::Tuple(_)
             | Expr::Import(_) => None,
-            Expr::PatchDef(_) | Expr::Record(_) | Expr::ReDef(_) => None,
+            Expr::Record(_) | Expr::ReDef(_) => None,
             Expr::Call(call) => self.get_call_ns(cur_ns, call, pos),
             Expr::ClassDef(class_def) => self.get_class_def_ns(cur_ns, class_def, pos),
+            Expr::PatchDef(patch_def) => self.get_patch_def_ns(cur_ns, patch_def, pos),
             Expr::Def(def) => self.get_def_ns(cur_ns, def, pos),
-            // Expr::PatchDef(patch_def) => self.get_patchdef_ns(cur_ns, patch_def, pos),
             Expr::Lambda(lambda) => self.get_lambda_ns(cur_ns, lambda, pos),
             Expr::TypeAsc(type_asc) => self.get_expr_ns(cur_ns, &type_asc.expr, pos),
             Expr::Dummy(dummy) => self.get_dummy_ns(cur_ns, dummy, pos),
@@ -163,6 +163,17 @@ impl<'a> HIRVisitor<'a> {
         self.get_exprs_ns(cur_ns, class_def.methods.iter(), pos)
     }
 
+    fn get_patch_def_ns(
+        &self,
+        mut cur_ns: Vec<Str>,
+        patch_def: &PatchDef,
+        pos: Position,
+    ) -> Option<Vec<Str>> {
+        let ns = patch_def.sig.ident().to_string_notype();
+        cur_ns.push(Str::from(ns));
+        self.get_exprs_ns(cur_ns, patch_def.methods.iter(), pos)
+    }
+
     fn get_def_ns(&self, mut cur_ns: Vec<Str>, def: &Def, pos: Position) -> Option<Vec<Str>> {
         let ns = def.sig.ident().to_string_notype();
         cur_ns.push(Str::from(ns));
@@ -230,7 +241,7 @@ impl<'a> HIRVisitor<'a> {
             Expr::Call(call) => self.get_call_t(call, token),
             Expr::ClassDef(class_def) => self.get_class_def_t(class_def, token),
             Expr::Def(def) => self.get_def_t(def, token),
-            Expr::PatchDef(patch_def) => self.get_patchdef_t(patch_def, token),
+            Expr::PatchDef(patch_def) => self.get_patch_def_t(patch_def, token),
             Expr::Lambda(lambda) => self.get_block_t(&lambda.body, token),
             Expr::Array(arr) => self.get_array_t(arr, token),
             Expr::Dict(dict) => self.get_dict_t(dict, token),
@@ -331,8 +342,15 @@ impl<'a> HIRVisitor<'a> {
         None
     }
 
-    fn get_patchdef_t(&self, patch_def: &PatchDef, token: &Token) -> Option<Type> {
+    fn get_patch_def_t(&self, patch_def: &PatchDef, token: &Token) -> Option<Type> {
         self.get_expr_t(&patch_def.base, token)
+            .or_else(|| {
+                self.return_expr_t_if_same(
+                    patch_def.sig.ident(),
+                    patch_def.sig.ident().name.token(),
+                    token,
+                )
+            })
             .or_else(|| self.get_block_t(&patch_def.methods, token))
     }
 
@@ -419,8 +437,8 @@ impl<'a> HIRVisitor<'a> {
             Expr::UnaryOp(unary) => self.get_expr_info(&unary.expr, token),
             Expr::Call(call) => self.get_call_info(call, token),
             Expr::ClassDef(class_def) => self.get_class_def_info(class_def, token),
+            Expr::PatchDef(patch_def) => self.get_patch_def_info(patch_def, token),
             Expr::Def(def) => self.get_def_info(def, token),
-            Expr::PatchDef(patch_def) => self.get_patchdef_info(patch_def, token),
             Expr::Lambda(lambda) => self.get_lambda_info(lambda, token),
             Expr::Array(arr) => self.get_array_info(arr, token),
             Expr::Dict(dict) => self.get_dict_info(dict, token),
@@ -526,6 +544,12 @@ impl<'a> HIRVisitor<'a> {
             .or_else(|| self.get_block_info(&class_def.methods, token))
     }
 
+    fn get_patch_def_info(&self, patch_def: &PatchDef, token: &Token) -> Option<VarInfo> {
+        self.get_expr_info(&patch_def.base, token)
+            .or_else(|| self.get_sig_info(&patch_def.sig, token))
+            .or_else(|| self.get_block_info(&patch_def.methods, token))
+    }
+
     fn get_block_info(&self, block: &Block, token: &Token) -> Option<VarInfo> {
         for chunk in block.iter() {
             if let Some(expr) = self.get_expr_info(chunk, token) {
@@ -549,11 +573,6 @@ impl<'a> HIRVisitor<'a> {
         None
     }
 
-    fn get_patchdef_info(&self, patch_def: &PatchDef, token: &Token) -> Option<VarInfo> {
-        self.get_expr_info(&patch_def.base, token)
-            .or_else(|| self.get_block_info(&patch_def.methods, token))
-    }
-
     fn get_lambda_info(&self, lambda: &Lambda, token: &Token) -> Option<VarInfo> {
         self.get_params_info(&lambda.params, token)
             .or_else(|| self.get_block_info(&lambda.body, token))

crates/els/inlay_hint.rs

@@ -12,7 +12,7 @@ use erg_common::error::Location;
 use erg_common::traits::{Locational, Runnable, Stream};
 
 use erg_compiler::artifact::{BuildRunnable, IncompleteArtifact};
-use erg_compiler::hir::{Block, Call, ClassDef, Def, Expr, Lambda, Signature};
+use erg_compiler::hir::{Block, Call, ClassDef, Def, Expr, Lambda, Params, PatchDef, Signature};
 use lsp_types::{InlayHint, InlayHintKind, InlayHintLabel, InlayHintParams};
 
 use crate::server::{ELSResult, Server};
@@ -97,16 +97,15 @@ impl<Checker: BuildRunnable> Server<Checker> {
             Expr::Def(def) => self.get_var_def_hint(def),
             Expr::Lambda(lambda) => self.get_lambda_hint(lambda),
             Expr::ClassDef(class_def) => self.get_class_def_hint(class_def),
+            Expr::PatchDef(patch_def) => self.get_patch_def_hint(patch_def),
             Expr::Call(call) => self.get_call_hint(call),
             _ => vec![],
         }
     }
 
-    fn get_subr_def_hint(&self, def: &Def) -> Vec<InlayHint> {
+    fn get_param_hint(&self, params: &Params) -> Vec<InlayHint> {
         let mut result = vec![];
-        result.extend(self.get_block_hint(&def.body.block));
+        for nd_param in params.non_defaults.iter() {
-        let Signature::Subr(subr) = &def.sig else { unreachable!() };
-        for nd_param in subr.params.non_defaults.iter() {
             if nd_param.raw.t_spec.is_some() {
                 continue;
             }
@@ -118,7 +117,19 @@ impl<Checker: BuildRunnable> Server<Checker> {
             );
             result.push(hint);
         }
-        for d_param in subr.params.defaults.iter() {
+        if let Some(var_params) = &params.var_params {
+            if var_params.raw.t_spec.is_some() {
+                return result;
+            }
+            let hint = type_anot(
+                var_params.ln_end().unwrap(),
+                var_params.col_end().unwrap(),
+                &var_params.vi.t,
+                false,
+            );
+            result.push(hint);
+        }
+        for d_param in params.defaults.iter() {
             if d_param.sig.raw.t_spec.is_some() {
                 continue;
             }
@@ -130,6 +141,14 @@ impl<Checker: BuildRunnable> Server<Checker> {
             );
             result.push(hint);
         }
+        result
+    }
+
+    fn get_subr_def_hint(&self, def: &Def) -> Vec<InlayHint> {
+        let mut result = vec![];
+        result.extend(self.get_block_hint(&def.body.block));
+        let Signature::Subr(subr) = &def.sig else { unreachable!() };
+        result.extend(self.get_param_hint(&subr.params));
         if def.sig.t_spec().is_none() {
             let Some(return_t) = subr.ref_t().return_t() else {
                 return result;
@@ -171,30 +190,7 @@ impl<Checker: BuildRunnable> Server<Checker> {
     fn get_lambda_hint(&self, lambda: &Lambda) -> Vec<InlayHint> {
         let mut result = vec![];
         result.extend(self.get_block_hint(&lambda.body));
-        for nd_param in lambda.params.non_defaults.iter() {
+        result.extend(self.get_param_hint(&lambda.params));
-            if nd_param.raw.t_spec.is_some() {
-                continue;
-            }
-            let hint = type_anot(
-                nd_param.ln_end().unwrap(),
-                nd_param.col_end().unwrap(),
-                &nd_param.vi.t,
-                false,
-            );
-            result.push(hint);
-        }
-        for d_param in lambda.params.defaults.iter() {
-            if d_param.sig.raw.t_spec.is_some() {
-                continue;
-            }
-            let hint = type_anot(
-                d_param.sig.ln_end().unwrap(),
-                d_param.sig.col_end().unwrap(),
-                &d_param.sig.vi.t,
-                false,
-            );
-            result.push(hint);
-        }
         let return_t = lambda.ref_t().return_t().unwrap();
         let hint = type_anot(
             lambda.params.ln_end().unwrap(),
@@ -222,6 +218,14 @@ impl<Checker: BuildRunnable> Server<Checker> {
             .collect()
     }
 
+    fn get_patch_def_hint(&self, patch_def: &PatchDef) -> Vec<InlayHint> {
+        patch_def
+            .methods
+            .iter()
+            .flat_map(|expr| self.get_expr_hint(expr))
+            .collect()
+    }
+
     fn get_call_hint(&self, call: &Call) -> Vec<InlayHint> {
         let mut result = vec![];
         let Some(call_t) = call.signature_t() else {

crates/erg_compiler/hir.rs

@@ -2226,7 +2226,7 @@ impl NoTypeDisplay for PatchDef {
 }
 
 impl_display_from_nested!(PatchDef);
-impl_locational!(PatchDef, sig);
+impl_locational!(PatchDef, sig, methods);
 
 impl HasType for PatchDef {
     #[inline]