Implement exit point highlighting

diff --git a/crates/ide/src/highlight_related.rs b/crates/ide/src/highlight_related.rs
index 1daaeb43fa2fed9166a6732b452aa490fe06cb94..5493d11d07426662fa99af713166d6f49acb272d 100644 (file)
--- a/crates/ide/src/highlight_related.rs
+++ b/crates/ide/src/highlight_related.rs
@@ -42,7 +42,91 @@ pub(crate) fn highlight_related(
     }
 }
 
-fn highlight_exit_points(_token: SyntaxToken) -> Option<Vec<DocumentHighlight>> {
+fn highlight_references(
+    sema: &Semantics<RootDatabase>,
+    syntax: &SyntaxNode,
+    FilePosition { offset, file_id }: FilePosition,
+) -> Option<Vec<DocumentHighlight>> {
+    let def = references::find_def(sema, syntax, offset)?;
+    let usages = def.usages(sema).set_scope(Some(SearchScope::single_file(file_id))).all();
+
+    let declaration = match def {
+        Definition::ModuleDef(hir::ModuleDef::Module(module)) => {
+            Some(NavigationTarget::from_module_to_decl(sema.db, module))
+        }
+        def => def.try_to_nav(sema.db),
+    }
+    .filter(|decl| decl.file_id == file_id)
+    .and_then(|decl| {
+        let range = decl.focus_range?;
+        let access = references::decl_access(&def, syntax, range);
+        Some(DocumentHighlight { range, access })
+    });
+
+    let file_refs = usages.references.get(&file_id).map_or(&[][..], Vec::as_slice);
+    let mut res = Vec::with_capacity(file_refs.len() + 1);
+    res.extend(declaration);
+    res.extend(
+        file_refs
+            .iter()
+            .map(|&FileReference { access, range, .. }| DocumentHighlight { range, access }),
+    );
+    Some(res)
+}
+
+fn highlight_exit_points(token: SyntaxToken) -> Option<Vec<DocumentHighlight>> {
+    fn hl(body: Option<ast::Expr>) -> Option<Vec<DocumentHighlight>> {
+        let mut highlights = Vec::new();
+        let body = body?;
+        walk(body.syntax(), |node| {
+            match_ast! {
+                match node {
+                    ast::ReturnExpr(expr) => if let Some(token) = expr.return_token() {
+                        highlights.push(DocumentHighlight {
+                            access: None,
+                            range: token.text_range(),
+                        });
+                    },
+                    ast::TryExpr(try_) => if let Some(token) = try_.question_mark_token() {
+                        highlights.push(DocumentHighlight {
+                            access: None,
+                            range: token.text_range(),
+                        });
+                    },
+                    ast::EffectExpr(effect) => if effect.async_token().is_some() {
+                        return true;
+                    },
+                    ast::ClosureExpr(__) => return true,
+                    ast::Item(__) => return true,
+                    ast::Path(__) => return true,
+                    _ => (),
+                }
+            }
+            false
+        });
+        let tail = match body {
+            ast::Expr::BlockExpr(b) => b.tail_expr(),
+            e => Some(e),
+        };
+        if let Some(tail) = tail {
+            highlights.push(DocumentHighlight { access: None, range: tail.syntax().text_range() });
+        }
+        Some(highlights)
+    }
+    for anc in token.ancestors() {
+        return match_ast! {
+            match anc {
+                ast::Fn(fn_) => hl(fn_.body().map(ast::Expr::BlockExpr)),
+                ast::ClosureExpr(closure) => hl(closure.body()),
+                ast::EffectExpr(effect) => if effect.async_token().is_some() {
+                    None
+                } else {
+                    continue;
+                },
+                _ => continue,
+            }
+        };
+    }
     None
 }
 
@@ -54,12 +138,7 @@ fn hl(
         let mut highlights = Vec::new();
         highlights.push(DocumentHighlight { access: None, range: async_token?.text_range() });
         if let Some(body) = body {
-            let mut preorder = body.syntax().preorder();
-            while let Some(event) = preorder.next() {
-                let node = match event {
-                    WalkEvent::Enter(node) => node,
-                    WalkEvent::Leave(_) => continue,
-                };
+            walk(body.syntax(), |node| {
                 match_ast! {
                     match node {
                         ast::AwaitExpr(expr) => if let Some(token) = expr.await_token() {
@@ -68,14 +147,17 @@ fn hl(
                                 range: token.text_range(),
                             });
                         },
-                        ast::EffectExpr(__) => preorder.skip_subtree(),
-                        ast::ClosureExpr(__) => preorder.skip_subtree(),
-                        ast::Item(__) => preorder.skip_subtree(),
-                        ast::Path(__) => preorder.skip_subtree(),
+                        ast::EffectExpr(effect) => if effect.async_token().is_some() {
+                            return true;
+                        },
+                        ast::ClosureExpr(__) => return true,
+                        ast::Item(__) => return true,
+                        ast::Path(__) => return true,
                         _ => (),
                     }
                 }
-            }
+                false
+            });
         }
         Some(highlights)
     }
@@ -92,36 +174,17 @@ fn hl(
     None
 }
 
-fn highlight_references(
-    sema: &Semantics<RootDatabase>,
-    syntax: &SyntaxNode,
-    FilePosition { offset, file_id }: FilePosition,
-) -> Option<Vec<DocumentHighlight>> {
-    let def = references::find_def(sema, syntax, offset)?;
-    let usages = def.usages(sema).set_scope(Some(SearchScope::single_file(file_id))).all();
-
-    let declaration = match def {
-        Definition::ModuleDef(hir::ModuleDef::Module(module)) => {
-            Some(NavigationTarget::from_module_to_decl(sema.db, module))
+fn walk(syntax: &SyntaxNode, mut cb: impl FnMut(SyntaxNode) -> bool) {
+    let mut preorder = syntax.preorder();
+    while let Some(event) = preorder.next() {
+        let node = match event {
+            WalkEvent::Enter(node) => node,
+            WalkEvent::Leave(_) => continue,
+        };
+        if cb(node) {
+            preorder.skip_subtree();
         }
-        def => def.try_to_nav(sema.db),
     }
-    .filter(|decl| decl.file_id == file_id)
-    .and_then(|decl| {
-        let range = decl.focus_range?;
-        let access = references::decl_access(&def, syntax, range);
-        Some(DocumentHighlight { range, access })
-    });
-
-    let file_refs = usages.references.get(&file_id).map_or(&[][..], Vec::as_slice);
-    let mut res = Vec::with_capacity(file_refs.len() + 1);
-    res.extend(declaration);
-    res.extend(
-        file_refs
-            .iter()
-            .map(|&FileReference { access, range, .. }| DocumentHighlight { range, access }),
-    );
-    Some(res)
 }
 
 #[cfg(test)]
@@ -278,6 +341,63 @@ async fn foo() {
         // ^^^^^
     ).await;
 }
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_exit_points() {
+        check(
+            r#"
+fn foo() -> u32 {
+    if true {
+        return$0 0;
+     // ^^^^^^
+    }
+
+    0?;
+  // ^
+    0xDEAD_BEEF
+ // ^^^^^^^^^^^
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_exit_points2() {
+        check(
+            r#"
+fn foo() ->$0 u32 {
+    if true {
+        return 0;
+     // ^^^^^^
+    }
+
+    0?;
+  // ^
+    0xDEAD_BEEF
+ // ^^^^^^^^^^^
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_prefer_ref_over_tail_exit() {
+        check(
+            r#"
+fn foo() -> u32 {
+// ^^^
+    if true {
+        return 0;
+    }
+
+    0?;
+
+    foo$0()
+ // ^^^
+}
 "#,
         );
     }