chore: improve loc calculation

crates/els/hir_visitor.rs

@@ -364,6 +364,11 @@ impl<'a> HIRVisitor<'a> {
         pos: Position,
     ) -> Option<&Expr> {
         self.return_expr_if_same(expr, def.sig.ident().raw.name.token(), pos)
+            .or_else(|| {
+                def.sig
+                    .t_spec_with_op()
+                    .and_then(|t_spec| self.get_expr(&t_spec.expr, pos))
+            })
             .or_else(|| self.get_expr_from_block(def.body.block.iter(), pos))
             .or_else(|| self.return_expr_if_contains(expr, pos, def))
     }

crates/erg_compiler/hir.rs

@@ -2062,7 +2062,21 @@ impl NoTypeDisplay for SubrSignature {
 }
 
 impl_display_from_nested!(SubrSignature);
-impl_locational!(SubrSignature, ident, params);
+
+impl Locational for SubrSignature {
+    // TODO: decorators
+    fn loc(&self) -> Location {
+        if let Some(return_t_spec) = &self.return_t_spec {
+            Location::concat(&self.ident, return_t_spec)
+        } else {
+            let params = self.params.loc();
+            if !params.is_unknown() {
+                return Location::concat(&self.ident, &params);
+            }
+            self.ident.loc()
+        }
+    }
+}
 
 impl HasType for SubrSignature {
     #[inline]
@@ -2141,7 +2155,7 @@ impl NoTypeDisplay for Lambda {
 }
 
 impl_display_from_nested!(Lambda);
-impl_locational!(Lambda, params, body);
+impl_locational!(Lambda, lossy params, body);
 impl_t!(Lambda);
 
 impl Lambda {
@@ -2267,6 +2281,13 @@ impl Signature {
         }
     }
 
+    pub fn params(&self) -> Option<&Params> {
+        match self {
+            Self::Var(_) => None,
+            Self::Subr(s) => Some(&s.params),
+        }
+    }
+
     pub fn captured_names(&self) -> &[Identifier] {
         match self {
             Self::Var(_) => &[],