use erg_common::consts::ERG_MODE; use erg_common::shared::MappedRwLockReadGuard; use erg_common::traits::Locational; use erg_common::Str; use erg_compiler::erg_parser::token::Token; use erg_compiler::hir::*; use erg_compiler::varinfo::VarInfo; use lsp_types::Position; use crate::file_cache::FileCache; use crate::util::{self, NormalizedUrl}; /// This struct provides: /// * namespace where the cursor is located (`get_namespace`) /// * cursor(`Token`) -> `Expr` mapping (`get_min_expr`) /// * cursor(`Token`) -> `VarInfo` mapping (`get_info`) pub struct HIRVisitor<'a> { hir: MappedRwLockReadGuard<'a, HIR>, file_cache: &'a FileCache, uri: NormalizedUrl, strict_cmp: bool, } impl<'a> HIRVisitor<'a> { pub fn new( hir: MappedRwLockReadGuard<'a, HIR>, file_cache: &'a FileCache, uri: NormalizedUrl, ) -> Self { Self { hir, file_cache, uri, strict_cmp: ERG_MODE, } } pub fn get_namespace(&self, pos: Position) -> Vec { let name = self .uri .path() .split('/') .last() .unwrap_or_default() .split('.') .next() .unwrap_or_default(); let namespace = vec![Str::rc(name)]; if let Some(ns) = self.get_exprs_ns(namespace.clone(), self.hir.module.iter(), pos) { ns } else { namespace } } fn is_new_final_line(&self, chunk: &Expr, pos: Position) -> bool { let ln_end = chunk.ln_end().unwrap_or(0); let line = self .file_cache .get_line(&self.uri, ln_end) .unwrap_or_default(); let indent_len = line.len() - line.trim_start_matches(' ').len(); let cond = ln_end == pos.line && pos.character as usize == indent_len + 1; cond && matches!( chunk, Expr::Call(_) | Expr::Lambda(_) | Expr::Def(_) | Expr::ClassDef(_) ) } fn get_expr_ns(&self, cur_ns: Vec, chunk: &Expr, pos: Position) -> Option> { if !(util::pos_in_loc(chunk, pos) || self.is_new_final_line(chunk, pos)) { return None; } match chunk { Expr::Lit(_) | Expr::Accessor(_) | Expr::BinOp(_) | Expr::UnaryOp(_) | Expr::Array(_) | Expr::Dict(_) | Expr::Set(_) | Expr::Tuple(_) | Expr::Import(_) => 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::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), Expr::Compound(block) | Expr::Code(block) => self.get_block_ns(cur_ns, block, pos), } } fn get_exprs_ns<'s, I: Iterator>( &self, cur_ns: Vec, exprs: I, pos: Position, ) -> Option> { let mut namespaces = vec![]; for expr in exprs { if let Some(ns) = self.get_expr_ns(cur_ns.clone(), expr, pos) { namespaces.push(ns); } } namespaces.into_iter().max_by(|l, r| l.len().cmp(&r.len())) } fn get_call_ns(&self, cur_ns: Vec, call: &Call, pos: Position) -> Option> { self.get_exprs_ns(cur_ns, call.args.pos_args.iter().map(|arg| &arg.expr), pos) } fn get_class_def_ns( &self, mut cur_ns: Vec, class_def: &ClassDef, pos: Position, ) -> Option> { let ns = class_def.sig.ident().to_string_notype(); cur_ns.push(Str::from(ns)); self.get_exprs_ns(cur_ns, class_def.methods.iter(), pos) } fn get_patch_def_ns( &self, mut cur_ns: Vec, patch_def: &PatchDef, pos: Position, ) -> Option> { 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, def: &Def, pos: Position) -> Option> { let ns = def.sig.ident().to_string_notype(); cur_ns.push(Str::from(ns)); match self.get_block_ns(cur_ns.clone(), &def.body.block, pos) { Some(ns) => Some(ns), None => Some(cur_ns), } } fn get_lambda_ns( &self, mut cur_ns: Vec, lambda: &Lambda, pos: Position, ) -> Option> { let ns = lambda.name_to_string(); cur_ns.push(Str::from(ns)); match self.get_block_ns(cur_ns.clone(), &lambda.body, pos) { Some(ns) => Some(ns), None => Some(cur_ns), } } fn get_block_ns(&self, cur_ns: Vec, block: &Block, pos: Position) -> Option> { self.get_exprs_ns(cur_ns, block.iter(), pos) } fn get_dummy_ns(&self, cur_ns: Vec, dummy: &Dummy, pos: Position) -> Option> { self.get_exprs_ns(cur_ns, dummy.iter(), pos) } /// Returns the smallest expression containing `token`. Literals, accessors, containers, etc. are returned. pub fn get_min_expr(&self, token: &Token) -> Option<&Expr> { for chunk in self.hir.module.iter() { if let Some(expr) = self.get_expr(chunk, token) { return Some(expr); } } None } fn return_expr_if_same<'e>(&'e self, expr: &'e Expr, l: &Token, r: &Token) -> Option<&Expr> { if self.strict_cmp { if l.deep_eq(r) { Some(expr) } else { None } } else if l == r { Some(expr) } else { None } } fn get_expr<'e>(&'e self, expr: &'e Expr, token: &Token) -> Option<&'e Expr> { if expr.ln_end() < token.ln_begin() || expr.ln_begin() > token.ln_end() { return None; } match expr { Expr::Lit(lit) => self.return_expr_if_same(expr, &lit.token, token), Expr::Accessor(acc) => self.get_expr_from_acc(expr, acc, token), Expr::BinOp(bin) => self.get_expr_from_bin(expr, bin, token), Expr::UnaryOp(unary) => self.get_expr(&unary.expr, token), Expr::Call(call) => self.get_expr_from_call(expr, call, token), Expr::ClassDef(class_def) => self.get_expr_from_class_def(expr, class_def, token), Expr::Def(def) => self.get_expr_from_def(expr, def, token), Expr::PatchDef(patch_def) => self.get_expr_from_patch_def(expr, patch_def, token), Expr::Lambda(lambda) => self.get_expr_from_lambda(expr, lambda, token), Expr::Array(arr) => self.get_expr_from_array(expr, arr, token), Expr::Dict(dict) => self.get_expr_from_dict(expr, dict, token), Expr::Record(record) => self.get_expr_from_record(expr, record, token), Expr::Set(set) => self.get_expr_from_set(expr, set, token), Expr::Tuple(tuple) => self.get_expr_from_tuple(expr, tuple, token), Expr::TypeAsc(type_asc) => self.get_expr(&type_asc.expr, token), Expr::Dummy(dummy) => self.get_expr_from_dummy(dummy, token), Expr::Compound(block) | Expr::Code(block) => self.get_expr_from_block(block, token), Expr::ReDef(redef) => self.get_expr_from_redef(expr, redef, token), Expr::Import(_) => None, } } fn get_expr_from_acc<'e>( &'e self, expr: &'e Expr, acc: &'e Accessor, token: &Token, ) -> Option<&Expr> { match acc { Accessor::Ident(ident) => self.return_expr_if_same(expr, ident.raw.name.token(), token), Accessor::Attr(attr) => self .return_expr_if_same(expr, attr.ident.raw.name.token(), token) .or_else(|| self.get_expr(&attr.obj, token)), } } fn get_expr_from_bin<'e>( &'e self, expr: &'e Expr, bin: &'e BinOp, token: &Token, ) -> Option<&Expr> { if &bin.op == token { return Some(expr); } self.get_expr(&bin.lhs, token) .or_else(|| self.get_expr(&bin.rhs, token)) } fn get_expr_from_call<'e>( &'e self, expr: &'e Expr, call: &'e Call, token: &Token, ) -> Option<&Expr> { // args: `(1, 2)`, token: `)` call.args .paren .as_ref() .and_then(|(_, end)| self.return_expr_if_same(expr, end, token)) .or_else(|| { call.attr_name .as_ref() .and_then(|attr| self.return_expr_if_same(expr, attr.raw.name.token(), token)) }) .or_else(|| self.get_expr(&call.obj, token)) .or_else(|| self.get_expr_from_args(&call.args, token)) .or_else(|| call.loc().contains(token.loc()).then_some(expr)) } fn get_expr_from_args<'e>(&'e self, args: &'e Args, token: &Token) -> Option<&Expr> { for arg in args.pos_args.iter() { if let Some(expr) = self.get_expr(&arg.expr, token) { return Some(expr); } } if let Some(var) = &args.var_args { if let Some(expr) = self.get_expr(&var.expr, token) { return Some(expr); } } for arg in args.kw_args.iter() { if let Some(expr) = self.get_expr(&arg.expr, token) { return Some(expr); } } None } fn get_expr_from_def<'e>( &'e self, expr: &'e Expr, def: &'e Def, token: &Token, ) -> Option<&Expr> { self.return_expr_if_same(expr, def.sig.ident().raw.name.token(), token) .or_else(|| self.get_expr_from_block(&def.body.block, token)) .or_else(|| def.loc().contains(token.loc()).then_some(expr)) } fn get_expr_from_class_def<'e>( &'e self, expr: &'e Expr, class_def: &'e ClassDef, token: &Token, ) -> Option<&Expr> { class_def .require_or_sup .as_ref() .and_then(|req_sup| self.get_expr(req_sup, token)) .or_else(|| self.get_expr_from_block(&class_def.methods, token)) .or_else(|| class_def.loc().contains(token.loc()).then_some(expr)) } fn get_expr_from_block<'e>(&'e self, block: &'e Block, token: &Token) -> Option<&Expr> { for chunk in block.iter() { if let Some(expr) = self.get_expr(chunk, token) { return Some(expr); } } None } fn get_expr_from_redef<'e>( &'e self, expr: &'e Expr, redef: &'e ReDef, token: &Token, ) -> Option<&Expr> { self.get_expr_from_acc(expr, &redef.attr, token) .or_else(|| self.get_expr_from_block(&redef.block, token)) } fn get_expr_from_dummy<'e>(&'e self, dummy: &'e Dummy, token: &Token) -> Option<&Expr> { for chunk in dummy.iter() { if let Some(expr) = self.get_expr(chunk, token) { return Some(expr); } } None } fn get_expr_from_patch_def<'e>( &'e self, expr: &'e Expr, patch_def: &'e PatchDef, token: &Token, ) -> Option<&Expr> { self.return_expr_if_same(expr, patch_def.sig.name().token(), token) .or_else(|| self.get_expr(&patch_def.base, token)) .or_else(|| self.get_expr_from_block(&patch_def.methods, token)) .or_else(|| patch_def.loc().contains(token.loc()).then_some(expr)) } fn get_expr_from_lambda<'e>( &'e self, expr: &'e Expr, lambda: &'e Lambda, token: &Token, ) -> Option<&Expr> { if util::pos_in_loc(&lambda.params, util::loc_to_pos(token.loc())?) { return Some(expr); } self.get_expr_from_block(&lambda.body, token) } fn get_expr_from_array<'e>( &'e self, expr: &'e Expr, arr: &'e Array, token: &Token, ) -> Option<&Expr> { if arr.ln_end() == token.ln_end() { // arr: `[1, 2]`, token: `]` return Some(expr); } match arr { Array::Normal(arr) => self.get_expr_from_args(&arr.elems, token), _ => None, // todo!(), } } fn get_expr_from_dict<'e>( &'e self, expr: &'e Expr, dict: &'e Dict, token: &Token, ) -> Option<&Expr> { if dict.ln_end() == token.ln_end() { // arr: `{...}`, token: `}` return Some(expr); } match dict { Dict::Normal(dict) => { for kv in &dict.kvs { if let Some(expr) = self .get_expr(&kv.key, token) .or_else(|| self.get_expr(&kv.value, token)) { return Some(expr); } } None } _ => None, // todo!(), } } fn get_expr_from_record<'e>( &'e self, expr: &'e Expr, record: &'e Record, token: &Token, ) -> Option<&Expr> { if record.ln_end() == token.ln_end() { // arr: `{...}`, token: `}` return Some(expr); } for field in record.attrs.iter() { if let Some(expr) = self.get_expr_from_block(&field.body.block, token) { return Some(expr); } } None } fn get_expr_from_set<'e>( &'e self, expr: &'e Expr, set: &'e Set, token: &Token, ) -> Option<&Expr> { if set.ln_end() == token.ln_end() { // arr: `{...}`, token: `}` return Some(expr); } match set { Set::Normal(set) => self.get_expr_from_args(&set.elems, token), _ => None, // todo!(), } } fn get_expr_from_tuple<'e>( &'e self, expr: &'e Expr, tuple: &'e Tuple, token: &Token, ) -> Option<&Expr> { match tuple { Tuple::Normal(tuple) => { // arr: `(1, 2)`, token: `)` tuple .elems .paren .as_ref() .and_then(|(_, end)| self.return_expr_if_same(expr, end, token)) .or_else(|| self.get_expr_from_args(&tuple.elems, token)) } // _ => None, // todo!(), } } } impl<'a> HIRVisitor<'a> { /// Returns the smallest expression containing `token`. Literals, accessors, containers, etc. are returned. pub fn get_info(&self, token: &Token) -> Option { for chunk in self.hir.module.iter() { if let Some(expr) = self.get_expr_info(chunk, token) { return Some(expr); } } None } fn return_var_info_if_same(&self, ident: &Identifier, l: &Token, r: &Token) -> Option { if self.strict_cmp { if l.deep_eq(r) { Some(ident.vi.clone()) } else { None } } else if l == r { Some(ident.vi.clone()) } else { None } } fn get_expr_info(&self, expr: &Expr, token: &Token) -> Option { let loc = expr.loc(); if loc.ln_end() < token.ln_begin() || loc.ln_begin() > token.ln_end() { return None; } match expr { Expr::Lit(_lit) => None, Expr::Accessor(acc) => self.get_acc_info(acc, token), Expr::BinOp(bin) => self.get_bin_info(bin, token), 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::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), Expr::Record(record) => self.get_record_info(record, token), Expr::Set(set) => self.get_set_info(set, token), Expr::Tuple(tuple) => self.get_tuple_info(tuple, token), Expr::TypeAsc(type_asc) => self.get_tasc_info(type_asc, token), Expr::Dummy(dummy) => self.get_dummy_info(dummy, token), Expr::Compound(block) | Expr::Code(block) => self.get_block_info(block, token), Expr::ReDef(redef) => self.get_redef_info(redef, token), Expr::Import(_) => None, } } fn get_acc_info(&self, acc: &Accessor, token: &Token) -> Option { match acc { Accessor::Ident(ident) => { self.return_var_info_if_same(ident, ident.raw.name.token(), token) } Accessor::Attr(attr) => self .return_var_info_if_same(&attr.ident, attr.ident.raw.name.token(), token) .or_else(|| self.get_expr_info(&attr.obj, token)), } } fn get_bin_info(&self, bin: &BinOp, token: &Token) -> Option { self.get_expr_info(&bin.lhs, token) .or_else(|| self.get_expr_info(&bin.rhs, token)) } fn get_call_info(&self, call: &Call, token: &Token) -> Option { if let Some(attr) = &call.attr_name { if let Some(t) = self.return_var_info_if_same(attr, attr.raw.name.token(), token) { return Some(t); } } self.get_expr_info(&call.obj, token) .or_else(|| self.get_args_info(&call.args, token)) } fn get_args_info(&self, args: &Args, token: &Token) -> Option { for arg in args.pos_args.iter() { if let Some(expr) = self.get_expr_info(&arg.expr, token) { return Some(expr); } } if let Some(var) = &args.var_args { if let Some(expr) = self.get_expr_info(&var.expr, token) { return Some(expr); } } for arg in args.kw_args.iter() { if let Some(expr) = self.get_expr_info(&arg.expr, token) { return Some(expr); } } None } fn get_sig_info(&self, sig: &Signature, token: &Token) -> Option { match sig { Signature::Var(var) => self .return_var_info_if_same(&var.ident, var.name().token(), token) .or_else(|| { var.t_spec .as_ref() .and_then(|t_spec| self.get_expr_info(&t_spec.expr, token)) }), Signature::Subr(subr) => self .return_var_info_if_same(&subr.ident, subr.name().token(), token) .or_else(|| self.get_params_info(&subr.params, token)) .or_else(|| { subr.return_t_spec .as_ref() .and_then(|t_spec| self.get_expr_info(&t_spec.expr, token)) }), } } fn get_params_info(&self, params: &Params, token: &Token) -> Option { for param in params.non_defaults.iter() { if let Some(vi) = param .t_spec_as_expr .as_ref() .and_then(|t_spec| self.get_expr_info(t_spec, token)) { return Some(vi); } else if param.raw.pat.loc() == token.loc() { return Some(param.vi.clone()); } } if let Some(var) = ¶ms.var_params { if let Some(vi) = var .t_spec_as_expr .as_ref() .and_then(|t_spec| self.get_expr_info(t_spec, token)) { return Some(vi); } else if var.raw.pat.loc() == token.loc() { return Some(var.vi.clone()); } } for param in params.defaults.iter() { if let Some(vi) = self.get_expr_info(¶m.default_val, token) { return Some(vi); } else if let Some(vi) = param .sig .t_spec_as_expr .as_ref() .and_then(|t_spec| self.get_expr_info(t_spec, token)) { return Some(vi); } else if param.sig.raw.pat.loc() == token.loc() { return Some(param.sig.vi.clone()); } } None } fn get_def_info(&self, def: &Def, token: &Token) -> Option { self.get_sig_info(&def.sig, token) .or_else(|| self.get_block_info(&def.body.block, token)) } fn get_class_def_info(&self, class_def: &ClassDef, token: &Token) -> Option { class_def .require_or_sup .as_ref() .and_then(|req_sup| self.get_expr_info(req_sup, token)) .or_else(|| self.get_sig_info(&class_def.sig, token)) .or_else(|| self.get_block_info(&class_def.methods, token)) } fn get_patch_def_info(&self, patch_def: &PatchDef, token: &Token) -> Option { 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 { for chunk in block.iter() { if let Some(expr) = self.get_expr_info(chunk, token) { return Some(expr); } } None } fn get_redef_info(&self, redef: &ReDef, token: &Token) -> Option { self.get_acc_info(&redef.attr, token) .or_else(|| self.get_block_info(&redef.block, token)) } fn get_dummy_info(&self, dummy: &Dummy, token: &Token) -> Option { for chunk in dummy.iter() { if let Some(expr) = self.get_expr_info(chunk, token) { return Some(expr); } } None } fn get_lambda_info(&self, lambda: &Lambda, token: &Token) -> Option { self.get_params_info(&lambda.params, token) .or_else(|| self.get_block_info(&lambda.body, token)) } fn get_array_info(&self, arr: &Array, token: &Token) -> Option { match arr { Array::Normal(arr) => self.get_args_info(&arr.elems, token), _ => None, // todo!(), } } fn get_dict_info(&self, dict: &Dict, token: &Token) -> Option { match dict { Dict::Normal(dict) => { for kv in &dict.kvs { if let Some(expr) = self.get_expr_info(&kv.key, token) { return Some(expr); } else if let Some(expr) = self.get_expr_info(&kv.value, token) { return Some(expr); } } None } _ => None, // todo!(), } } fn get_record_info(&self, record: &Record, token: &Token) -> Option { for field in record.attrs.iter() { if let Some(expr) = self.get_def_info(field, token) { return Some(expr); } } None } fn get_set_info(&self, set: &Set, token: &Token) -> Option { match set { Set::Normal(set) => self.get_args_info(&set.elems, token), _ => None, // todo!(), } } fn get_tuple_info(&self, tuple: &Tuple, token: &Token) -> Option { match tuple { Tuple::Normal(tuple) => self.get_args_info(&tuple.elems, token), // _ => None, // todo!(), } } fn get_tasc_info(&self, tasc: &TypeAscription, token: &Token) -> Option { self.get_expr_info(&tasc.expr, token) .or_else(|| self.get_expr_info(&tasc.spec.expr, token)) } }