Skip to content

Commit 0012480

Browse files
CryZeCryZe
committed
feat: Highlight exit points of async blocks
Async blocks act similar to async functions in that the await keywords are related, but also act like functions where the exit points are related.
1 parent 814da15 commit 0012480

File tree

1 file changed

+93
-71
lines changed

1 file changed

+93
-71
lines changed

‎src/tools/rust-analyzer/crates/ide/src/highlight_related.rs

Lines changed: 93 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -281,99 +281,95 @@ fn highlight_references(
281281
}
282282
}
283283

284-
// If `file_id` is None,
285-
pub(crate) fn highlight_exit_points(
284+
fn hl_exit_points(
286285
sema: &Semantics<'_, RootDatabase>,
287-
token: SyntaxToken,
288-
) -> FxHashMap<EditionedFileId, Vec<HighlightedRange>> {
289-
fn hl(
290-
sema: &Semantics<'_, RootDatabase>,
291-
def_token: Option<SyntaxToken>,
292-
body: ast::Expr,
293-
) -> Option<FxHashMap<EditionedFileId, FxHashSet<HighlightedRange>>> {
294-
let mut highlights: FxHashMap<EditionedFileId, FxHashSet<_>> = FxHashMap::default();
286+
def_token: Option<SyntaxToken>,
287+
body: ast::Expr,
288+
) -> Option<FxHashMap<EditionedFileId, FxHashSet<HighlightedRange>>> {
289+
let mut highlights: FxHashMap<EditionedFileId, FxHashSet<_>> = FxHashMap::default();
290+
291+
let mut push_to_highlights = |file_id, range| {
292+
if let Some(FileRange { file_id, range }) = original_frange(sema.db, file_id, range) {
293+
let hrange = HighlightedRange { category: ReferenceCategory::empty(), range };
294+
highlights.entry(file_id).or_default().insert(hrange);
295+
}
296+
};
295297

296-
let mut push_to_highlights = |file_id, range| {
297-
if let Some(FileRange { file_id, range }) = original_frange(sema.db, file_id, range) {
298-
let hrange = HighlightedRange { category: ReferenceCategory::empty(), range };
299-
highlights.entry(file_id).or_default().insert(hrange);
298+
if let Some(tok) = def_token {
299+
let file_id = sema.hir_file_for(&tok.parent()?);
300+
let range = Some(tok.text_range());
301+
push_to_highlights(file_id, range);
302+
}
303+
304+
WalkExpandedExprCtx::new(sema).walk(&body, &mut |_, expr| {
305+
let file_id = sema.hir_file_for(expr.syntax());
306+
307+
let range = match &expr {
308+
ast::Expr::TryExpr(try_) => try_.question_mark_token().map(|token| token.text_range()),
309+
ast::Expr::MethodCallExpr(_) | ast::Expr::CallExpr(_) | ast::Expr::MacroExpr(_)
310+
if sema.type_of_expr(&expr).map_or(false, |ty| ty.original.is_never()) =>
311+
{
312+
Some(expr.syntax().text_range())
300313
}
314+
_ => None,
301315
};
302316

303-
if let Some(tok) = def_token {
304-
let file_id = sema.hir_file_for(&tok.parent()?);
305-
let range = Some(tok.text_range());
306-
push_to_highlights(file_id, range);
307-
}
317+
push_to_highlights(file_id, range);
318+
});
308319

309-
WalkExpandedExprCtx::new(sema).walk(&body, &mut |_, expr| {
320+
// We should handle `return` separately, because when it is used in a `try` block,
321+
// it will exit the outside function instead of the block itself.
322+
WalkExpandedExprCtx::new(sema)
323+
.with_check_ctx(&WalkExpandedExprCtx::is_async_const_block_or_closure)
324+
.walk(&body, &mut |_, expr| {
310325
let file_id = sema.hir_file_for(expr.syntax());
311326

312327
let range = match &expr {
313-
ast::Expr::TryExpr(try_) => {
314-
try_.question_mark_token().map(|token| token.text_range())
315-
}
316-
ast::Expr::MethodCallExpr(_) | ast::Expr::CallExpr(_) | ast::Expr::MacroExpr(_)
317-
if sema.type_of_expr(&expr).map_or(false, |ty| ty.original.is_never()) =>
318-
{
319-
Some(expr.syntax().text_range())
320-
}
328+
ast::Expr::ReturnExpr(expr) => expr.return_token().map(|token| token.text_range()),
321329
_ => None,
322330
};
323331

324332
push_to_highlights(file_id, range);
325333
});
326334

327-
// We should handle `return` separately, because when it is used in a `try` block,
328-
// it will exit the outside function instead of the block itself.
329-
WalkExpandedExprCtx::new(sema)
330-
.with_check_ctx(&WalkExpandedExprCtx::is_async_const_block_or_closure)
331-
.walk(&body, &mut |_, expr| {
332-
let file_id = sema.hir_file_for(expr.syntax());
333-
334-
let range = match &expr {
335-
ast::Expr::ReturnExpr(expr) => {
336-
expr.return_token().map(|token| token.text_range())
337-
}
338-
_ => None,
339-
};
340-
341-
push_to_highlights(file_id, range);
342-
});
343-
344-
let tail = match body {
345-
ast::Expr::BlockExpr(b) => b.tail_expr(),
346-
e => Some(e),
347-
};
335+
let tail = match body {
336+
ast::Expr::BlockExpr(b) => b.tail_expr(),
337+
e => Some(e),
338+
};
348339

349-
if let Some(tail) = tail {
350-
for_each_tail_expr(&tail, &mut |tail| {
351-
let file_id = sema.hir_file_for(tail.syntax());
352-
let range = match tail {
353-
ast::Expr::BreakExpr(b) => b
354-
.break_token()
355-
.map_or_else(|| tail.syntax().text_range(), |tok| tok.text_range()),
356-
_ => tail.syntax().text_range(),
357-
};
358-
push_to_highlights(file_id, Some(range));
359-
});
360-
}
361-
Some(highlights)
340+
if let Some(tail) = tail {
341+
for_each_tail_expr(&tail, &mut |tail| {
342+
let file_id = sema.hir_file_for(tail.syntax());
343+
let range = match tail {
344+
ast::Expr::BreakExpr(b) => b
345+
.break_token()
346+
.map_or_else(|| tail.syntax().text_range(), |tok| tok.text_range()),
347+
_ => tail.syntax().text_range(),
348+
};
349+
push_to_highlights(file_id, Some(range));
350+
});
362351
}
352+
Some(highlights)
353+
}
363354

355+
// If `file_id` is None,
356+
pub(crate) fn highlight_exit_points(
357+
sema: &Semantics<'_, RootDatabase>,
358+
token: SyntaxToken,
359+
) -> FxHashMap<EditionedFileId, Vec<HighlightedRange>> {
364360
let mut res = FxHashMap::default();
365361
for def in goto_definition::find_fn_or_blocks(sema, &token) {
366362
let new_map = match_ast! {
367363
match def {
368-
ast::Fn(fn_) => fn_.body().and_then(|body| hl(sema, fn_.fn_token(), body.into())),
364+
ast::Fn(fn_) => fn_.body().and_then(|body| hl_exit_points(sema, fn_.fn_token(), body.into())),
369365
ast::ClosureExpr(closure) => {
370366
let pipe_tok = closure.param_list().and_then(|p| p.pipe_token());
371-
closure.body().and_then(|body| hl(sema, pipe_tok, body))
367+
closure.body().and_then(|body| hl_exit_points(sema, pipe_tok, body))
372368
},
373369
ast::BlockExpr(blk) => match blk.modifier() {
374-
Some(ast::BlockModifier::Async(t)) => hl(sema, Some(t), blk.into()),
370+
Some(ast::BlockModifier::Async(t)) => hl_exit_points(sema, Some(t), blk.into()),
375371
Some(ast::BlockModifier::Try(t)) if token.kind() != T![return] => {
376-
hl(sema, Some(t), blk.into())
372+
hl_exit_points(sema, Some(t), blk.into())
377373
},
378374
_ => continue,
379375
},
@@ -520,6 +516,12 @@ pub(crate) fn highlight_yield_points(
520516
if block_expr.async_token().is_none() {
521517
continue;
522518
}
519+
520+
// Async blocks act similar to closures. So we want to
521+
// highlight their exit points too.
522+
let exit_points = hl_exit_points(sema, block_expr.async_token(), block_expr.clone().into());
523+
merge_map(&mut res, exit_points);
524+
523525
hl(sema, block_expr.async_token(), Some(block_expr.into()))
524526
},
525527
ast::ClosureExpr(closure) => hl(sema, closure.async_token(), closure.body()),
@@ -876,6 +878,27 @@ pub async$0 fn foo() {
876878
);
877879
}
878880

881+
#[test]
882+
fn test_hl_exit_points_of_async_blocks() {
883+
check(
884+
r#"
885+
pub fn foo() {
886+
let x = async$0 {
887+
// ^^^^^
888+
0.await;
889+
// ^^^^^
890+
0?;
891+
// ^
892+
return 0;
893+
// ^^^^^^
894+
0
895+
// ^
896+
};
897+
}
898+
"#,
899+
);
900+
}
901+
879902
#[test]
880903
fn test_hl_let_else_yield_points() {
881904
check(
@@ -925,11 +948,10 @@ async fn foo() {
925948
async fn foo() {
926949
(async {
927950
// ^^^^^
928-
(async {
929-
0.await
930-
}).await$0 }
931-
// ^^^^^
932-
).await;
951+
(async { 0.await }).await$0
952+
// ^^^^^^^^^^^^^^^^^^^^^^^^^
953+
// ^^^^^
954+
}).await;
933955
}
934956
"#,
935957
);

0 commit comments

Comments
 (0)