Improve error range

JetBrains · Jun 18, 2024 · a8f0c5d · a8f0c5d
1 parent 30e4dd7
commit a8f0c5d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/src/main/kotlin/org/arend/codeInsight/completion/Constants.kt b/src/main/kotlin/org/arend/codeInsight/completion/Constants.kt
@@ -4,7 +4,8 @@ import org.arend.psi.ArendElementTypes.*
 
 
 val FIXITY_KWS = listOf(INFIX_LEFT_KW, INFIX_RIGHT_KW, INFIX_NON_KW, NON_ASSOC_KW, LEFT_ASSOC_KW, RIGHT_ASSOC_KW).map { it.toString() }
-val STATEMENT_WT_KWS = listOf(FUNC_KW, SFUNC_KW, LEMMA_KW, TYPE_KW, CONS_KW, DATA_KW, CLASS_KW, RECORD_KW, INSTANCE_KW, OPEN_KW, MODULE_KW, META_KW, PLEVELS_KW, HLEVELS_KW, AXIOM_KW).map { it.toString() }
+val STATEMENT_WT_KWS_TOKENS = listOf(FUNC_KW, SFUNC_KW, LEMMA_KW, TYPE_KW, CONS_KW, DATA_KW, CLASS_KW, RECORD_KW, INSTANCE_KW, OPEN_KW, MODULE_KW, META_KW, PLEVELS_KW, HLEVELS_KW, AXIOM_KW)
+val STATEMENT_WT_KWS = STATEMENT_WT_KWS_TOKENS.map { it.toString() }
 val ACCESS_MODIFIERS = listOf(PROTECTED_KW, PRIVATE_KW).map { it.toString() }
 val CLASS_MEMBER_KWS = listOf(FIELD_KW, PROPERTY_KW, OVERRIDE_KW, DEFAULT_KW).map { it.toString() }
 val SIGMA_TELE_START_KWS = listOf(PROPERTY_KW).map { it.toString() }

diff --git a/src/main/kotlin/org/arend/highlight/BasePass.kt b/src/main/kotlin/org/arend/highlight/BasePass.kt
@@ -14,9 +14,12 @@ import com.intellij.patterns.StandardPatterns
 import com.intellij.psi.*
 import com.intellij.psi.codeStyle.CodeStyleManager
 import com.intellij.psi.impl.source.tree.LeafPsiElement
+import com.intellij.refactoring.suggested.endOffset
+import com.intellij.refactoring.suggested.startOffset
 import com.intellij.xml.util.XmlStringUtil
 import org.arend.IArendFile
 import org.arend.codeInsight.ArendCodeInsightUtils.Companion.getAllParametersForReferable
+import org.arend.codeInsight.completion.STATEMENT_WT_KWS_TOKENS
 import org.arend.codeInsight.completion.withAncestors
 import org.arend.core.context.param.DependentLink
 import org.arend.core.expr.ReferenceExpression
@@ -638,6 +641,13 @@ abstract class BasePass(protected open val file: IArendFile, editor: Editor, nam
                 }
             }
 
+            if (improvedElement is LeafPsiElement && STATEMENT_WT_KWS_TOKENS.contains(improvedElement.elementType)) {
+                val stat = improvedElement.ancestor<ArendDefinition<*>>()?.parent
+                if (stat != null) {
+                    return TextRange(stat.startOffset, improvedElement.endOffset)
+                }
+            }
+
             return improvedElement.textRange
         }