From 8abf27898d4d57282281ab09d644d753ec52575c Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Tue, 1 Dec 2020 11:20:25 +0300
Subject: [PATCH] Simplify FirMemberDeclaration.implicitModality

---
 .../fir/analysis/checkers/FirHelpers.kt       | 24 ++++++++++---------
 .../LightTreePositioningStrategies.kt         |  3 +++
 2 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirHelpers.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirHelpers.kt
index 7c70973ee92..b214513dce0 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirHelpers.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirHelpers.kt
@@ -12,6 +12,9 @@ import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSymbolOwner
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.diagnostics.modalityModifier
+import org.jetbrains.kotlin.fir.analysis.diagnostics.overrideModifier
+import org.jetbrains.kotlin.fir.analysis.diagnostics.visibilityModifier
 import org.jetbrains.kotlin.fir.containingClass
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
@@ -250,20 +253,21 @@ fun FirMemberDeclaration.implicitModality(context: CheckerContext): Modality {
     }
 
     val klass = context.findClosestClassOrObject() ?: return Modality.FINAL
-    val modifiers = this.modifierListOrNull() ?: return Modality.FINAL
-    if (modifiers.contains(KtTokens.OVERRIDE_KEYWORD)) {
-        val klassModifiers = klass.modifierListOrNull()
-        if (klassModifiers != null && klassModifiers.run {
-                contains(KtTokens.ABSTRACT_KEYWORD) || contains(KtTokens.OPEN_KEYWORD) || contains(KtTokens.SEALED_KEYWORD)
-            }) {
+    val source = source ?: return Modality.FINAL
+    val tree = source.treeStructure
+    if (tree.overrideModifier(source.lighterASTNode) != null) {
+        val klassModalityTokenType = klass.source?.let { tree.modalityModifier(it.lighterASTNode)?.tokenType }
+        if (klassModalityTokenType == KtTokens.ABSTRACT_KEYWORD ||
+            klassModalityTokenType == KtTokens.OPEN_KEYWORD ||
+            klassModalityTokenType == KtTokens.SEALED_KEYWORD
+        ) {
             return Modality.OPEN
         }
     }
 
-    if (
-        klass is FirRegularClass
+    if (klass is FirRegularClass
         && klass.classKind == ClassKind.INTERFACE
-        && !modifiers.contains(KtTokens.PRIVATE_KEYWORD)
+        && tree.visibilityModifier(source.lighterASTNode)?.tokenType != KtTokens.PRIVATE_KEYWORD
     ) {
         return if (this.hasBody()) Modality.OPEN else Modality.ABSTRACT
     }
@@ -271,8 +275,6 @@ fun FirMemberDeclaration.implicitModality(context: CheckerContext): Modality {
     return Modality.FINAL
 }
 
-private fun FirDeclaration.modifierListOrNull() = this.source.getModifierList()?.modifiers?.map { it.token }
-
 private fun FirDeclaration.hasBody(): Boolean = when (this) {
     is FirSimpleFunction -> this.body != null && this.body !is FirEmptyExpressionBlock
     is FirProperty -> this.setter?.body !is FirEmptyExpressionBlock? || this.getter?.body !is FirEmptyExpressionBlock?
diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategies.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategies.kt
index b4418cf6948..78863350667 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategies.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategies.kt
@@ -208,6 +208,9 @@ internal fun FlyweightCapableTreeStructure<LighterASTNode>.visibilityModifier(de
 internal fun FlyweightCapableTreeStructure<LighterASTNode>.modalityModifier(declaration: LighterASTNode): LighterASTNode? =
     modifierList(declaration)?.let { findChildByType(it, MODALITY_MODIFIERS) }
 
+internal fun FlyweightCapableTreeStructure<LighterASTNode>.overrideModifier(declaration: LighterASTNode): LighterASTNode? =
+    modifierList(declaration)?.let { findChildByType(it, KtTokens.OVERRIDE_KEYWORD) }
+
 private fun FlyweightCapableTreeStructure<LighterASTNode>.accessorNamePlaceholder(node: LighterASTNode): LighterASTNode =
     findChildByType(node, KtTokens.GET_KEYWORD) ?: findChildByType(node, KtTokens.SET_KEYWORD)!!