Skip to content

Commit d6e78b0

Browse files
VeykrilVeykril
committed
feat: Handle operators like their trait functions in the IDE
1 parent cb52271 commit d6e78b0

File tree

12 files changed

+459
-68
lines changed

12 files changed

+459
-68
lines changed

‎crates/hir-expand/src/name.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -381,6 +381,7 @@ pub mod known {
381381
bitor,
382382
bitxor_assign,
383383
bitxor,
384+
branch,
384385
deref_mut,
385386
deref,
386387
div_assign,
@@ -396,6 +397,7 @@ pub mod known {
396397
not,
397398
owned_box,
398399
partial_ord,
400+
poll,
399401
r#fn,
400402
rem_assign,
401403
rem,

‎crates/hir-ty/src/infer/expr.rs

Lines changed: 7 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,13 @@ use chalk_ir::{
1010
cast::Cast, fold::Shift, DebruijnIndex, GenericArgData, Mutability, TyVariableKind,
1111
};
1212
use hir_def::{
13-
expr::{ArithOp, Array, BinaryOp, CmpOp, Expr, ExprId, Literal, Ordering, Statement, UnaryOp},
13+
expr::{ArithOp, Array, BinaryOp, CmpOp, Expr, ExprId, Literal, Statement, UnaryOp},
1414
generics::TypeOrConstParamData,
1515
path::{GenericArg, GenericArgs},
1616
resolver::resolver_for_expr,
17-
ConstParamId, FieldId, FunctionId, ItemContainerId, Lookup,
17+
ConstParamId, FieldId, ItemContainerId, Lookup,
1818
};
19-
use hir_expand::name::{name, Name};
19+
use hir_expand::name::Name;
2020
use stdx::always;
2121
use syntax::ast::RangeOp;
2222

@@ -28,7 +28,7 @@ use crate::{
2828
const_or_path_to_chalk, generic_arg_to_chalk, lower_to_chalk_mutability, ParamLoweringMode,
2929
},
3030
mapping::{from_chalk, ToChalk},
31-
method_resolution::{self, VisibleFromModule},
31+
method_resolution::{self, lang_names_for_bin_op, VisibleFromModule},
3232
primitive::{self, UintTy},
3333
static_lifetime, to_chalk_trait_id,
3434
utils::{generics, Generics},
@@ -947,7 +947,9 @@ impl<'a> InferenceContext<'a> {
947947
let lhs_ty = self.infer_expr(lhs, &lhs_expectation);
948948
let rhs_ty = self.table.new_type_var();
949949

950-
let func = self.resolve_binop_method(op);
950+
let func = lang_names_for_bin_op(op).and_then(|(name, lang_item)| {
951+
self.db.trait_data(self.resolve_lang_item(lang_item)?.as_trait()?).method_by_name(&name)
952+
});
951953
let func = match func {
952954
Some(func) => func,
953955
None => {
@@ -1473,55 +1475,4 @@ impl<'a> InferenceContext<'a> {
14731475
},
14741476
})
14751477
}
1476-
1477-
fn resolve_binop_method(&self, op: BinaryOp) -> Option<FunctionId> {
1478-
let (name, lang_item) = match op {
1479-
BinaryOp::LogicOp(_) => return None,
1480-
BinaryOp::ArithOp(aop) => match aop {
1481-
ArithOp::Add => (name!(add), name!(add)),
1482-
ArithOp::Mul => (name!(mul), name!(mul)),
1483-
ArithOp::Sub => (name!(sub), name!(sub)),
1484-
ArithOp::Div => (name!(div), name!(div)),
1485-
ArithOp::Rem => (name!(rem), name!(rem)),
1486-
ArithOp::Shl => (name!(shl), name!(shl)),
1487-
ArithOp::Shr => (name!(shr), name!(shr)),
1488-
ArithOp::BitXor => (name!(bitxor), name!(bitxor)),
1489-
ArithOp::BitOr => (name!(bitor), name!(bitor)),
1490-
ArithOp::BitAnd => (name!(bitand), name!(bitand)),
1491-
},
1492-
BinaryOp::Assignment { op: Some(aop) } => match aop {
1493-
ArithOp::Add => (name!(add_assign), name!(add_assign)),
1494-
ArithOp::Mul => (name!(mul_assign), name!(mul_assign)),
1495-
ArithOp::Sub => (name!(sub_assign), name!(sub_assign)),
1496-
ArithOp::Div => (name!(div_assign), name!(div_assign)),
1497-
ArithOp::Rem => (name!(rem_assign), name!(rem_assign)),
1498-
ArithOp::Shl => (name!(shl_assign), name!(shl_assign)),
1499-
ArithOp::Shr => (name!(shr_assign), name!(shr_assign)),
1500-
ArithOp::BitXor => (name!(bitxor_assign), name!(bitxor_assign)),
1501-
ArithOp::BitOr => (name!(bitor_assign), name!(bitor_assign)),
1502-
ArithOp::BitAnd => (name!(bitand_assign), name!(bitand_assign)),
1503-
},
1504-
BinaryOp::CmpOp(cop) => match cop {
1505-
CmpOp::Eq { negated: false } => (name!(eq), name!(eq)),
1506-
CmpOp::Eq { negated: true } => (name!(ne), name!(eq)),
1507-
CmpOp::Ord { ordering: Ordering::Less, strict: false } => {
1508-
(name!(le), name!(partial_ord))
1509-
}
1510-
CmpOp::Ord { ordering: Ordering::Less, strict: true } => {
1511-
(name!(lt), name!(partial_ord))
1512-
}
1513-
CmpOp::Ord { ordering: Ordering::Greater, strict: false } => {
1514-
(name!(ge), name!(partial_ord))
1515-
}
1516-
CmpOp::Ord { ordering: Ordering::Greater, strict: true } => {
1517-
(name!(gt), name!(partial_ord))
1518-
}
1519-
},
1520-
BinaryOp::Assignment { op: None } => return None,
1521-
};
1522-
1523-
let trait_ = self.resolve_lang_item(lang_item)?.as_trait()?;
1524-
1525-
self.db.trait_data(trait_).method_by_name(&name)
1526-
}
15271478
}

‎crates/hir-ty/src/method_resolution.rs

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -336,7 +336,7 @@ impl InherentImpls {
336336
}
337337
}
338338

339-
pub fn inherent_impl_crates_query(
339+
pub(crate) fn inherent_impl_crates_query(
340340
db: &dyn HirDatabase,
341341
krate: CrateId,
342342
fp: TyFingerprint,
@@ -419,6 +419,55 @@ pub fn def_crates(
419419
}
420420
}
421421

422+
pub fn lang_names_for_bin_op(op: syntax::ast::BinaryOp) -> Option<(Name, Name)> {
423+
use hir_expand::name;
424+
use syntax::ast::{ArithOp, BinaryOp, CmpOp, Ordering};
425+
Some(match op {
426+
BinaryOp::LogicOp(_) => return None,
427+
BinaryOp::ArithOp(aop) => match aop {
428+
ArithOp::Add => (name!(add), name!(add)),
429+
ArithOp::Mul => (name!(mul), name!(mul)),
430+
ArithOp::Sub => (name!(sub), name!(sub)),
431+
ArithOp::Div => (name!(div), name!(div)),
432+
ArithOp::Rem => (name!(rem), name!(rem)),
433+
ArithOp::Shl => (name!(shl), name!(shl)),
434+
ArithOp::Shr => (name!(shr), name!(shr)),
435+
ArithOp::BitXor => (name!(bitxor), name!(bitxor)),
436+
ArithOp::BitOr => (name!(bitor), name!(bitor)),
437+
ArithOp::BitAnd => (name!(bitand), name!(bitand)),
438+
},
439+
BinaryOp::Assignment { op: Some(aop) } => match aop {
440+
ArithOp::Add => (name!(add_assign), name!(add_assign)),
441+
ArithOp::Mul => (name!(mul_assign), name!(mul_assign)),
442+
ArithOp::Sub => (name!(sub_assign), name!(sub_assign)),
443+
ArithOp::Div => (name!(div_assign), name!(div_assign)),
444+
ArithOp::Rem => (name!(rem_assign), name!(rem_assign)),
445+
ArithOp::Shl => (name!(shl_assign), name!(shl_assign)),
446+
ArithOp::Shr => (name!(shr_assign), name!(shr_assign)),
447+
ArithOp::BitXor => (name!(bitxor_assign), name!(bitxor_assign)),
448+
ArithOp::BitOr => (name!(bitor_assign), name!(bitor_assign)),
449+
ArithOp::BitAnd => (name!(bitand_assign), name!(bitand_assign)),
450+
},
451+
BinaryOp::CmpOp(cop) => match cop {
452+
CmpOp::Eq { negated: false } => (name!(eq), name!(eq)),
453+
CmpOp::Eq { negated: true } => (name!(ne), name!(eq)),
454+
CmpOp::Ord { ordering: Ordering::Less, strict: false } => {
455+
(name!(le), name!(partial_ord))
456+
}
457+
CmpOp::Ord { ordering: Ordering::Less, strict: true } => {
458+
(name!(lt), name!(partial_ord))
459+
}
460+
CmpOp::Ord { ordering: Ordering::Greater, strict: false } => {
461+
(name!(ge), name!(partial_ord))
462+
}
463+
CmpOp::Ord { ordering: Ordering::Greater, strict: true } => {
464+
(name!(gt), name!(partial_ord))
465+
}
466+
},
467+
BinaryOp::Assignment { op: None } => return None,
468+
})
469+
}
470+
422471
/// Look up the method with the given name.
423472
pub(crate) fn lookup_method(
424473
ty: &Canonical<Ty>,

‎crates/hir/src/semantics.rs

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -357,6 +357,26 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
357357
self.imp.resolve_method_call(call).map(Function::from)
358358
}
359359

360+
pub fn resolve_await_to_poll(&self, await_expr: &ast::AwaitExpr) -> Option<Function> {
361+
self.imp.resolve_await_to_poll(await_expr).map(Function::from)
362+
}
363+
364+
pub fn resolve_prefix_expr(&self, prefix_expr: &ast::PrefixExpr) -> Option<Function> {
365+
self.imp.resolve_prefix_expr(prefix_expr).map(Function::from)
366+
}
367+
368+
pub fn resolve_index_expr(&self, index_expr: &ast::IndexExpr) -> Option<Function> {
369+
self.imp.resolve_index_expr(index_expr).map(Function::from)
370+
}
371+
372+
pub fn resolve_bin_expr(&self, bin_expr: &ast::BinExpr) -> Option<Function> {
373+
self.imp.resolve_bin_expr(bin_expr).map(Function::from)
374+
}
375+
376+
pub fn resolve_try_expr(&self, try_expr: &ast::TryExpr) -> Option<Function> {
377+
self.imp.resolve_try_expr(try_expr).map(Function::from)
378+
}
379+
360380
pub fn resolve_method_call_as_callable(&self, call: &ast::MethodCallExpr) -> Option<Callable> {
361381
self.imp.resolve_method_call_as_callable(call)
362382
}
@@ -1066,6 +1086,26 @@ impl<'db> SemanticsImpl<'db> {
10661086
self.analyze(call.syntax())?.resolve_method_call(self.db, call)
10671087
}
10681088

1089+
pub fn resolve_await_to_poll(&self, await_expr: &ast::AwaitExpr) -> Option<FunctionId> {
1090+
self.analyze(await_expr.syntax())?.resolve_await_to_poll(self.db, await_expr)
1091+
}
1092+
1093+
pub fn resolve_prefix_expr(&self, prefix_expr: &ast::PrefixExpr) -> Option<FunctionId> {
1094+
self.analyze(prefix_expr.syntax())?.resolve_prefix_expr(self.db, prefix_expr)
1095+
}
1096+
1097+
pub fn resolve_index_expr(&self, index_expr: &ast::IndexExpr) -> Option<FunctionId> {
1098+
self.analyze(index_expr.syntax())?.resolve_index_expr(self.db, index_expr)
1099+
}
1100+
1101+
pub fn resolve_bin_expr(&self, bin_expr: &ast::BinExpr) -> Option<FunctionId> {
1102+
self.analyze(bin_expr.syntax())?.resolve_bin_expr(self.db, bin_expr)
1103+
}
1104+
1105+
pub fn resolve_try_expr(&self, try_expr: &ast::TryExpr) -> Option<FunctionId> {
1106+
self.analyze(try_expr.syntax())?.resolve_try_expr(self.db, try_expr)
1107+
}
1108+
10691109
fn resolve_method_call_as_callable(&self, call: &ast::MethodCallExpr) -> Option<Callable> {
10701110
self.analyze(call.syntax())?.resolve_method_call_as_callable(self.db, call)
10711111
}

‎crates/hir/src/source_analyzer.rs

Lines changed: 111 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,17 @@ use hir_def::{
2525
Lookup, ModuleDefId, VariantId,
2626
};
2727
use hir_expand::{
28-
builtin_fn_macro::BuiltinFnLikeExpander, hygiene::Hygiene, name::AsName, HirFileId, InFile,
28+
builtin_fn_macro::BuiltinFnLikeExpander, hygiene::Hygiene, name, name::AsName, HirFileId,
29+
InFile,
2930
};
3031
use hir_ty::{
3132
diagnostics::{
3233
record_literal_missing_fields, record_pattern_missing_fields, unsafe_expressions,
3334
UnsafeExpr,
3435
},
35-
method_resolution, Adjust, Adjustment, AutoBorrow, InferenceResult, Interner, Substitution,
36-
TyExt, TyKind, TyLoweringContext,
36+
method_resolution::{self, lang_names_for_bin_op},
37+
Adjust, Adjustment, AutoBorrow, InferenceResult, Interner, Substitution, TyExt, TyKind,
38+
TyLoweringContext,
3739
};
3840
use itertools::Itertools;
3941
use smallvec::SmallVec;
@@ -255,8 +257,103 @@ impl SourceAnalyzer {
255257
) -> Option<FunctionId> {
256258
let expr_id = self.expr_id(db, &call.clone().into())?;
257259
let (f_in_trait, substs) = self.infer.as_ref()?.method_resolution(expr_id)?;
258-
let f_in_impl = self.resolve_impl_method(db, f_in_trait, &substs);
259-
f_in_impl.or(Some(f_in_trait))
260+
261+
Some(self.resolve_impl_method_or_trait_def(db, f_in_trait, &substs))
262+
}
263+
264+
pub fn resolve_await_to_poll(
265+
&self,
266+
db: &dyn HirDatabase,
267+
await_expr: &ast::AwaitExpr,
268+
) -> Option<FunctionId> {
269+
let expr_id = self.expr_id(db, &await_expr.expr()?.into())?;
270+
let ty = self.infer.as_ref()?.type_of_expr.get(expr_id)?;
271+
272+
let op_fn = db
273+
.lang_item(self.resolver.krate(), hir_expand::name![poll].to_smol_str())?
274+
.as_function()?;
275+
let substs = hir_ty::TyBuilder::subst_for_def(db, op_fn).push(ty.clone()).build();
276+
277+
Some(self.resolve_impl_method_or_trait_def(db, op_fn, &substs))
278+
}
279+
280+
pub fn resolve_prefix_expr(
281+
&self,
282+
db: &dyn HirDatabase,
283+
prefix_expr: &ast::PrefixExpr,
284+
) -> Option<FunctionId> {
285+
let lang_item_name = match prefix_expr.op_kind()? {
286+
ast::UnaryOp::Deref => name![deref],
287+
ast::UnaryOp::Not => name![not],
288+
ast::UnaryOp::Neg => name![neg],
289+
};
290+
let expr_id = self.expr_id(db, &prefix_expr.expr()?.into())?;
291+
let ty = self.infer.as_ref()?.type_of_expr.get(expr_id)?;
292+
293+
let trait_ =
294+
db.lang_item(self.resolver.krate(), lang_item_name.to_smol_str())?.as_trait()?;
295+
let op_fn = db.trait_data(trait_).method_by_name(&lang_item_name)?;
296+
let substs = hir_ty::TyBuilder::subst_for_def(db, op_fn).push(ty.clone()).build();
297+
298+
Some(self.resolve_impl_method_or_trait_def(db, op_fn, &substs))
299+
}
300+
301+
pub fn resolve_index_expr(
302+
&self,
303+
db: &dyn HirDatabase,
304+
index_expr: &ast::IndexExpr,
305+
) -> Option<FunctionId> {
306+
let base_expr_id = self.expr_id(db, &index_expr.base()?.into())?;
307+
let index_expr_id = self.expr_id(db, &index_expr.index()?.into())?;
308+
let base_ty = self.infer.as_ref()?.type_of_expr.get(base_expr_id)?;
309+
let index_ty = self.infer.as_ref()?.type_of_expr.get(index_expr_id)?;
310+
311+
let lang_item_name = name![index];
312+
313+
let trait_ =
314+
db.lang_item(self.resolver.krate(), lang_item_name.to_smol_str())?.as_trait()?;
315+
let op_fn = db.trait_data(trait_).method_by_name(&lang_item_name)?;
316+
let substs = hir_ty::TyBuilder::subst_for_def(db, op_fn)
317+
.push(base_ty.clone())
318+
.push(index_ty.clone())
319+
.build();
320+
Some(self.resolve_impl_method_or_trait_def(db, op_fn, &substs))
321+
}
322+
323+
pub fn resolve_bin_expr(
324+
&self,
325+
db: &dyn HirDatabase,
326+
binop_expr: &ast::BinExpr,
327+
) -> Option<FunctionId> {
328+
let op = binop_expr.op_kind()?;
329+
let lhs_expr_id = self.expr_id(db, &binop_expr.lhs()?.into())?;
330+
let rhs_expr_id = self.expr_id(db, &binop_expr.rhs()?.into())?;
331+
let lhs = self.infer.as_ref()?.type_of_expr.get(lhs_expr_id)?;
332+
let rhs = self.infer.as_ref()?.type_of_expr.get(rhs_expr_id)?;
333+
334+
let op_fn = lang_names_for_bin_op(op).and_then(|(name, lang_item)| {
335+
db.trait_data(db.lang_item(self.resolver.krate(), lang_item.to_smol_str())?.as_trait()?)
336+
.method_by_name(&name)
337+
})?;
338+
let substs =
339+
hir_ty::TyBuilder::subst_for_def(db, op_fn).push(lhs.clone()).push(rhs.clone()).build();
340+
341+
Some(self.resolve_impl_method_or_trait_def(db, op_fn, &substs))
342+
}
343+
344+
pub fn resolve_try_expr(
345+
&self,
346+
db: &dyn HirDatabase,
347+
try_expr: &ast::TryExpr,
348+
) -> Option<FunctionId> {
349+
let expr_id = self.expr_id(db, &try_expr.expr()?.into())?;
350+
let ty = self.infer.as_ref()?.type_of_expr.get(expr_id)?;
351+
352+
let op_fn =
353+
db.lang_item(self.resolver.krate(), name![branch].to_smol_str())?.as_function()?;
354+
let substs = hir_ty::TyBuilder::subst_for_def(db, op_fn).push(ty.clone()).build();
355+
356+
Some(self.resolve_impl_method_or_trait_def(db, op_fn, &substs))
260357
}
261358

262359
pub(crate) fn resolve_field(
@@ -666,6 +763,15 @@ impl SourceAnalyzer {
666763
let fun_data = db.function_data(func);
667764
method_resolution::lookup_impl_method(self_ty, db, trait_env, impled_trait, &fun_data.name)
668765
}
766+
767+
fn resolve_impl_method_or_trait_def(
768+
&self,
769+
db: &dyn HirDatabase,
770+
func: FunctionId,
771+
substs: &Substitution,
772+
) -> FunctionId {
773+
self.resolve_impl_method(db, func, substs).unwrap_or(func)
774+
}
669775
}
670776

671777
fn scope_for(

0 commit comments

Comments
 (0)