fix: workaround the Java Psi structure when errors are present

src/main/kotlin/com/github/lppedd/highlighter/java/JavaAlwaysHighlightStrategy.kt

@@ -1,11 +1,13 @@
 package com.github.lppedd.highlighter.java
 
 import com.github.lppedd.highlighter.ReturnHighlightStrategy
-import com.intellij.psi.PsiReturnStatement
+import com.intellij.psi.PsiKeyword
 
 /**
  * @author Edoardo Luppi
  */
-object JavaAlwaysHighlightStrategy : ReturnHighlightStrategy<PsiReturnStatement> {
-  override fun isValidContext(psiElement: PsiReturnStatement) = true
+object JavaAlwaysHighlightStrategy : ReturnHighlightStrategy<PsiKeyword> {
+  // Note: it seems the Java Psi structure isn't quite right when errors are
+  //  present. Thus, we need to intercept a lower-level PsiElement
+  override fun isValidContext(psiElement: PsiKeyword) = "$psiElement" == "PsiKeyword:return"
 }

src/main/kotlin/com/github/lppedd/highlighter/java/JavaReturnAnnotator.kt

@@ -1,14 +1,14 @@
 package com.github.lppedd.highlighter.java
 
-import com.intellij.psi.PsiReturnStatement
+import com.intellij.psi.PsiKeyword
 import com.github.lppedd.highlighter.AbstractReturnAnnotator as ARA
 
 /**
  * @author Edoardo Luppi
  */
-class JavaReturnAnnotator : ARA<PsiReturnStatement>(PsiReturnStatement::class.java) {
+class JavaReturnAnnotator : ARA<PsiKeyword>(PsiKeyword::class.java) {
   private val config = JavaReturnHighlighterConfig.INSTANCE
 
-  override fun isValidContext(psiElement: PsiReturnStatement) =
+  override fun isValidContext(psiElement: PsiKeyword) =
       config.getHighlightStrategy().isValidContext(psiElement)
 }

src/main/kotlin/com/github/lppedd/highlighter/java/JavaReturnHighlighterConfig.kt

@@ -9,7 +9,7 @@ import com.intellij.openapi.components.PersistentStateComponent
 import com.intellij.openapi.components.State
 import com.intellij.openapi.components.Storage
 import com.intellij.openapi.project.ProjectManager
-import com.intellij.psi.PsiReturnStatement
+import com.intellij.psi.PsiKeyword
 
 /**
  * @author Edoardo Luppi
@@ -32,7 +32,7 @@ class JavaReturnHighlighterConfig : PersistentStateComponent<JavaConfig> {
     refreshFiles()
   }
 
-  fun getHighlightStrategy(): ReturnHighlightStrategy<PsiReturnStatement> =
+  fun getHighlightStrategy(): ReturnHighlightStrategy<PsiKeyword> =
       if (state.isOnlyTopLevelReturns) JavaTopLevelHighlightStrategy
       else JavaAlwaysHighlightStrategy
 

src/main/kotlin/com/github/lppedd/highlighter/java/JavaReturnLineMarkerProvider.kt

@@ -1,15 +1,15 @@
 package com.github.lppedd.highlighter.java
 
-import com.intellij.psi.PsiReturnStatement
+import com.intellij.psi.PsiKeyword
 import com.github.lppedd.highlighter.AbstractReturnLineMarkerProvider as ARLMP
 
 /**
  * @author Edoardo Luppi
  */
-class JavaReturnLineMarkerProvider : ARLMP<PsiReturnStatement>(PsiReturnStatement::class.java) {
+class JavaReturnLineMarkerProvider : ARLMP<PsiKeyword>(PsiKeyword::class.java) {
   private val config = JavaReturnHighlighterConfig.INSTANCE
 
   override fun getName() = "Java"
-  override fun isValidContext(psiElement: PsiReturnStatement) =
+  override fun isValidContext(psiElement: PsiKeyword) =
       config.getHighlightStrategy().isValidContext(psiElement)
 }

src/main/kotlin/com/github/lppedd/highlighter/java/JavaTopLevelHighlightStrategy.kt

@@ -9,8 +9,14 @@ import com.intellij.psi.*
 /**
  * @author Edoardo Luppi
  */
-object JavaTopLevelHighlightStrategy : ReturnHighlightStrategy<PsiReturnStatement> {
-  override fun isValidContext(psiElement: PsiReturnStatement): Boolean {
+object JavaTopLevelHighlightStrategy : ReturnHighlightStrategy<PsiKeyword> {
+  override fun isValidContext(psiElement: PsiKeyword): Boolean {
+    // Note: it seems the Java Psi structure isn't quite right when errors are
+    //  present. Thus, we need to intercept a lower-level PsiElement
+    if ("$psiElement" != "PsiKeyword:return") {
+      return false;
+    }
+
     var psi: PsiElement? = psiElement
 
     while (psi != null) {