TYPE: highlight all 'await' when under the caret

src/main/kotlin/org/rust/ide/highlight/RsHighlightAwaitHandlerFactory.kt

@@ -0,0 +1,59 @@
+/*
+ * Use of this source code is governed by the MIT license that can be
+ * found in the LICENSE file.
+ */
+
+package org.rust.ide.highlight
+
+import com.intellij.codeInsight.highlighting.HighlightUsagesHandlerBase
+import com.intellij.codeInsight.highlighting.HighlightUsagesHandlerFactoryBase
+import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.util.parents
+import com.intellij.util.Consumer
+import org.rust.lang.core.psi.*
+import org.rust.lang.core.psi.ext.childrenWithLeaves
+import org.rust.lang.core.psi.ext.elementType
+
+class RsHighlightAwaitHandlerFactory : HighlightUsagesHandlerFactoryBase() {
+    override fun createHighlightUsagesHandler(editor: Editor, file: PsiFile, target: PsiElement): HighlightUsagesHandlerBase<*>? {
+        if (file !is RsFile) return null
+        val parentAsyncFunctionOrBlock = when {
+            target.isAsync() -> target.parent
+            target.isAwait() -> target.parentAsyncFunctionOrBlock()
+            else -> null
+        } ?: return null
+        return RsHighlightAsyncAwaitHandler(editor, file, parentAsyncFunctionOrBlock)
+    }
+}
+
+private fun PsiElement.isAsync() = elementType == RsElementTypes.ASYNC
+
+private fun PsiElement.isAwait() = elementType == RsElementTypes.IDENTIFIER && text == "await"
+
+private fun PsiElement.parentAsyncFunctionOrBlock(): PsiElement? = parents(withSelf = false).firstOrNull {
+    (it is RsFunction || it is RsBlockExpr) && it.childrenWithLeaves.any { leaf -> leaf.isAsync() }
+}
+
+private class RsHighlightAsyncAwaitHandler(
+    editor: Editor,
+    file: PsiFile,
+    val parentAsyncFunctionOrBlock: PsiElement
+) : HighlightUsagesHandlerBase<PsiElement>(editor, file) {
+    override fun getTargets() = listOf(parentAsyncFunctionOrBlock)
+
+    override fun selectTargets(targets: MutableList<out PsiElement>, selectionConsumer: Consumer<in MutableList<out PsiElement>>) {
+        selectionConsumer.consume(targets)
+    }
+
+    override fun computeUsages(targets: MutableList<out PsiElement>) {
+        parentAsyncFunctionOrBlock.accept(object : RsRecursiveVisitor() {
+            override fun visitDotExpr(o: RsDotExpr) {
+                o.fieldLookup?.identifier
+                    ?.takeIf { it.isAwait() && it.parentAsyncFunctionOrBlock() == parentAsyncFunctionOrBlock }
+                    ?.let(::addOccurrence)
+            }
+        })
+    }
+}

src/main/resources/META-INF/rust-core.xml

@@ -66,6 +66,7 @@
         <lang.syntaxHighlighter language="Rust"
                                 implementationClass="org.rust.ide.highlight.RsHighlighter"/>
         <highlightUsagesHandlerFactory implementation="org.rust.ide.highlight.RsHighlightExitPointsHandlerFactory"/>
+        <highlightUsagesHandlerFactory implementation="org.rust.ide.highlight.RsHighlightAwaitHandlerFactory"/>
 
         <!-- Formatter -->
 

src/test/kotlin/org/rust/ide/highlight/RsHighlightAwaitHandlerFactoryTest.kt

@@ -0,0 +1,75 @@
+/*
+ * Use of this source code is governed by the MIT license that can be
+ * found in the LICENSE file.
+ */
+
+package org.rust.ide.highlight
+
+import com.intellij.codeInsight.highlighting.HighlightUsagesHandler
+import org.intellij.lang.annotations.Language
+import org.rust.RsTestBase
+
+class RsHighlightAwaitHandlerFactoryTest : RsTestBase() {
+    fun `test highlight await in async function when caret is on async`() = doTest("""
+        async fn foo() {}
+        struct Bar { bar: i32 }
+        /*caret*/async fn test(bar: Bar) {
+            foo().await;
+            bar.bar;
+            if true {
+                foo().await;
+            }
+            async {
+                foo().await;
+                foo().await;
+            };
+        }
+    """, 4, 7)
+
+    fun `test highlight await in async function when caret is on await`() = doTest("""
+        async fn foo() {}
+        async fn test(bar: Bar) {
+            foo()./*caret*/await;
+            foo().await;
+        }
+    """, 3, 4)
+
+    fun `test highlight await in async block when caret is on async`() = doTest("""
+        async fn foo() {}
+        fn test() {
+            /*caret*/async {
+                foo().await;
+                foo().await;
+            };
+            async {
+                foo().await;
+                foo().await;
+            };
+        }
+    """, 4, 5)
+
+    fun `test highlight await in async block when caret is on await`() = doTest("""
+        async fn foo() {}
+        fn test() {
+            async {
+                foo().await;
+                foo().await;
+            };
+            async {
+                foo().await/*caret*/;
+                foo().await;
+            };
+        }
+    """, 8, 9)
+
+    private fun doTest(@Language("Rust") check: String, vararg lines: Int) {
+        InlineFile(check)
+        HighlightUsagesHandler.invoke(myFixture.project, myFixture.editor, myFixture.file)
+        val highlighters = myFixture.editor.markupModel.allHighlighters
+        val file = myFixture.file
+        val actual = highlighters.map {
+            file.findElementAt(it.startOffset)!!.lineNumber to file.text.substring(it.startOffset, it.endOffset)
+        }
+        assertSameElements(actual, lines.map { it to "await" })
+    }
+}