More correct keyword completion after "@"

compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt

@@ -69,6 +69,12 @@ public fun PsiElement.prevLeaf(skipEmptyElements: Boolean = false): PsiElement?
 public fun PsiElement.nextLeaf(skipEmptyElements: Boolean = false): PsiElement?
         = PsiTreeUtil.nextLeaf(this, skipEmptyElements)
 
+public val PsiElement.prevLeafs: Sequence<PsiElement>
+    get() = sequence({ prevLeaf() }, { it.prevLeaf() })
+
+public val PsiElement.nextLeafs: Sequence<PsiElement>
+    get() = sequence({ nextLeaf() }, { it.nextLeaf() })
+
 public fun PsiElement.prevLeaf(filter: (PsiElement) -> Boolean): PsiElement? {
     var leaf = prevLeaf()
     while (leaf != null && !filter(leaf)) {

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/KeywordCompletion.kt

@@ -48,9 +48,6 @@ object KeywordCompletion {
             .filter { it !in NON_ACTUAL_KEYWORDS }
             .map { it as KtKeywordToken }
 
-    private val DEFAULT_DUMMY_POSTFIX = " X"
-    private val KEYWORD_TO_DUMMY_POSTFIX = mapOf(FILE_KEYWORD to ":X")
-
     private val KEYWORDS_TO_IGNORE_PREFIX = TokenSet.create(OVERRIDE_KEYWORD /* it's needed to complete overrides that should be work by member name too */)
 
     private val COMPOUND_KEYWORDS = mapOf<KtKeywordToken, KtKeywordToken>(
@@ -192,7 +189,7 @@ object KeywordCompletion {
     private fun buildFilterByText(prefixText: String, project: Project): (KtKeywordToken) -> Boolean {
         val psiFactory = KtPsiFactory(project)
         return fun (keywordTokenType): Boolean {
-            val postfix = KEYWORD_TO_DUMMY_POSTFIX[keywordTokenType] ?: DEFAULT_DUMMY_POSTFIX
+            val postfix = if (prefixText.endsWith("@")) ":X" else " X"
             val file = psiFactory.createFile(prefixText + keywordTokenType.getValue() + postfix)
             val elementAt = file.findElementAt(prefixText.length())!!
 
@@ -201,7 +198,7 @@ object KeywordCompletion {
 
                 elementAt.getNonStrictParentOfType<PsiErrorElement>() != null -> return false
 
-                elementAt.prevLeaf { it !is PsiWhiteSpace && it !is PsiComment }?.parentsWithSelf?.any { it is PsiErrorElement } ?: false -> return false
+                isErrorElementBefore(elementAt) -> return false
 
                 keywordTokenType !is KtModifierKeywordToken -> return true
 
@@ -260,6 +257,15 @@ object KeywordCompletion {
         }
     }
 
+    private fun isErrorElementBefore(token: PsiElement): Boolean {
+        for (leaf in token.prevLeafs) {
+            if (leaf is PsiWhiteSpace || leaf is PsiComment) continue
+            if (leaf.parentsWithSelf.any { it is PsiErrorElement } ) return true
+            if (leaf.textLength != 0) break
+        }
+        return false
+    }
+
     private fun IElementType.matchesKeyword(keywordType: KtKeywordToken): Boolean {
         return when(this) {
             keywordType -> true

idea/idea-completion/testData/keywords/FileKeyword.kt

@@ -1,10 +1,11 @@
 @<caret>[]
 
 // EXIST: file
-// EXIST: class
+// EXIST: field
-// EXIST: fun
+// EXIST: get
-// EXIST: interface
+// EXIST: set
-// EXIST: object
+// EXIST: param
-// EXIST: val
+// EXIST: setparam
-// EXIST: var
+// EXIST: property
+// EXIST: receiver
 // NOTHING_ELSE