Skip to content

Commit

Permalink
feat: Highlight exit points of async blocks
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@CryZe
CryZe committed Sep 19, 2024
1 parent 814da15 commit 0012480
Showing 1 changed file with 93 additions and 71 deletions.
164 changes: 93 additions & 71 deletions src/tools/rust-analyzer/crates/ide/src/highlight_related.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -281,99 +281,95 @@ fn highlight_references(
}
}

// If `file_id` is None,
pub(crate) fn highlight_exit_points(
fn hl_exit_points(
sema: &Semantics<'_, RootDatabase>,
token: SyntaxToken,
) -> FxHashMap<EditionedFileId, Vec<HighlightedRange>> {
fn hl(
sema: &Semantics<'_, RootDatabase>,
def_token: Option<SyntaxToken>,
body: ast::Expr,
) -> Option<FxHashMap<EditionedFileId, FxHashSet<HighlightedRange>>> {
let mut highlights: FxHashMap<EditionedFileId, FxHashSet<_>> = FxHashMap::default();
def_token: Option<SyntaxToken>,
body: ast::Expr,
) -> Option<FxHashMap<EditionedFileId, FxHashSet<HighlightedRange>>> {
let mut highlights: FxHashMap<EditionedFileId, FxHashSet<_>> = FxHashMap::default();

let mut push_to_highlights = |file_id, range| {
if let Some(FileRange { file_id, range }) = original_frange(sema.db, file_id, range) {
let hrange = HighlightedRange { category: ReferenceCategory::empty(), range };
highlights.entry(file_id).or_default().insert(hrange);
}
};

let mut push_to_highlights = |file_id, range| {
if let Some(FileRange { file_id, range }) = original_frange(sema.db, file_id, range) {
let hrange = HighlightedRange { category: ReferenceCategory::empty(), range };
highlights.entry(file_id).or_default().insert(hrange);
if let Some(tok) = def_token {
let file_id = sema.hir_file_for(&tok.parent()?);
let range = Some(tok.text_range());
push_to_highlights(file_id, range);
}

WalkExpandedExprCtx::new(sema).walk(&body, &mut |_, expr| {
let file_id = sema.hir_file_for(expr.syntax());

let range = match &expr {
ast::Expr::TryExpr(try_) => try_.question_mark_token().map(|token| token.text_range()),
ast::Expr::MethodCallExpr(_) | ast::Expr::CallExpr(_) | ast::Expr::MacroExpr(_)
if sema.type_of_expr(&expr).map_or(false, |ty| ty.original.is_never()) =>
{
Some(expr.syntax().text_range())
}
_ => None,
};

if let Some(tok) = def_token {
let file_id = sema.hir_file_for(&tok.parent()?);
let range = Some(tok.text_range());
push_to_highlights(file_id, range);
}
push_to_highlights(file_id, range);
});

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

let range = match &expr {
ast::Expr::TryExpr(try_) => {
try_.question_mark_token().map(|token| token.text_range())
}
ast::Expr::MethodCallExpr(_) | ast::Expr::CallExpr(_) | ast::Expr::MacroExpr(_)
if sema.type_of_expr(&expr).map_or(false, |ty| ty.original.is_never()) =>
{
Some(expr.syntax().text_range())
}
ast::Expr::ReturnExpr(expr) => expr.return_token().map(|token| token.text_range()),
_ => None,
};

push_to_highlights(file_id, range);
});

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

let range = match &expr {
ast::Expr::ReturnExpr(expr) => {
expr.return_token().map(|token| token.text_range())
}
_ => None,
};

push_to_highlights(file_id, range);
});

let tail = match body {
ast::Expr::BlockExpr(b) => b.tail_expr(),
e => Some(e),
};
let tail = match body {
ast::Expr::BlockExpr(b) => b.tail_expr(),
e => Some(e),
};

if let Some(tail) = tail {
for_each_tail_expr(&tail, &mut |tail| {
let file_id = sema.hir_file_for(tail.syntax());
let range = match tail {
ast::Expr::BreakExpr(b) => b
.break_token()
.map_or_else(|| tail.syntax().text_range(), |tok| tok.text_range()),
_ => tail.syntax().text_range(),
};
push_to_highlights(file_id, Some(range));
});
}
Some(highlights)
if let Some(tail) = tail {
for_each_tail_expr(&tail, &mut |tail| {
let file_id = sema.hir_file_for(tail.syntax());
let range = match tail {
ast::Expr::BreakExpr(b) => b
.break_token()
.map_or_else(|| tail.syntax().text_range(), |tok| tok.text_range()),
_ => tail.syntax().text_range(),
};
push_to_highlights(file_id, Some(range));
});
}
Some(highlights)
}

// If `file_id` is None,
pub(crate) fn highlight_exit_points(
sema: &Semantics<'_, RootDatabase>,
token: SyntaxToken,
) -> FxHashMap<EditionedFileId, Vec<HighlightedRange>> {
let mut res = FxHashMap::default();
for def in goto_definition::find_fn_or_blocks(sema, &token) {
let new_map = match_ast! {
match def {
ast::Fn(fn_) => fn_.body().and_then(|body| hl(sema, fn_.fn_token(), body.into())),
ast::Fn(fn_) => fn_.body().and_then(|body| hl_exit_points(sema, fn_.fn_token(), body.into())),
ast::ClosureExpr(closure) => {
let pipe_tok = closure.param_list().and_then(|p| p.pipe_token());
closure.body().and_then(|body| hl(sema, pipe_tok, body))
closure.body().and_then(|body| hl_exit_points(sema, pipe_tok, body))
},
ast::BlockExpr(blk) => match blk.modifier() {
Some(ast::BlockModifier::Async(t)) => hl(sema, Some(t), blk.into()),
Some(ast::BlockModifier::Async(t)) => hl_exit_points(sema, Some(t), blk.into()),
Some(ast::BlockModifier::Try(t)) if token.kind() != T![return] => {
hl(sema, Some(t), blk.into())
hl_exit_points(sema, Some(t), blk.into())
},
_ => continue,
},
Expand Down Expand Up @@ -520,6 +516,12 @@ pub(crate) fn highlight_yield_points(
if block_expr.async_token().is_none() {
continue;
}

// Async blocks act similar to closures. So we want to
// highlight their exit points too.
let exit_points = hl_exit_points(sema, block_expr.async_token(), block_expr.clone().into());
merge_map(&mut res, exit_points);

hl(sema, block_expr.async_token(), Some(block_expr.into()))
},
ast::ClosureExpr(closure) => hl(sema, closure.async_token(), closure.body()),
Expand Down Expand Up @@ -876,6 +878,27 @@ pub async$0 fn foo() {
);
}

#[test]
fn test_hl_exit_points_of_async_blocks() {
check(
r#"
pub fn foo() {
let x = async$0 {
// ^^^^^
0.await;
// ^^^^^
0?;
// ^
return 0;
// ^^^^^^
0
// ^
};
}
"#,
);
}

#[test]
fn test_hl_let_else_yield_points() {
check(
Expand Down Expand Up @@ -925,11 +948,10 @@ async fn foo() {
async fn foo() {
(async {
// ^^^^^
(async {
0.await
}).await$0 }
// ^^^^^
).await;
(async { 0.await }).await$0
// ^^^^^^^^^^^^^^^^^^^^^^^^^
// ^^^^^
}).await;
}
"#,
);
Expand Down

0 comments on commit 0012480

Please sign in to comment.