FIR IDE: make KtExpression.getKtType() to return null for non-expressions

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtExpressionTypeProvider.kt

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtFunction
 
 public abstract class KtExpressionTypeProvider : KtAnalysisSessionComponent() {
-    public abstract fun getKtExpressionType(expression: KtExpression): KtType
+    public abstract fun getKtExpressionType(expression: KtExpression): KtType?
     public abstract fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KtType
     public abstract fun getFunctionalTypeForKtFunction(declaration: KtFunction): KtType
 
@@ -22,7 +22,15 @@ public abstract class KtExpressionTypeProvider : KtAnalysisSessionComponent() {
 }
 
 public interface KtExpressionTypeProviderMixIn : KtAnalysisSessionMixIn {
-    public fun KtExpression.getKtType(): KtType =
+    /**
+     * Get type of given expression.
+     *
+     * Return:
+     * - [KtExpression] type if given [KtExpression] is real expression;
+     * - `null` for [KtExpression] inside pacakges and import declarations;
+     * - `Unit` type for statements;
+     */
+    public fun KtExpression.getKtType(): KtType? =
         analysisSession.expressionTypeProvider.getKtExpressionType(this)
 
     /**

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirExpressionTypeProvider.kt

@@ -6,14 +6,18 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirLabel
+import org.jetbrains.kotlin.fir.FirPackageDirective
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.declarations.FirImport
 import org.jetbrains.kotlin.fir.declarations.utils.isSuspend
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.resolve.constructFunctionalType
 import org.jetbrains.kotlin.fir.typeContext
+import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
@@ -34,12 +38,13 @@ internal class KtFirExpressionTypeProvider(
     override val token: ValidityToken,
 ) : KtExpressionTypeProvider(), KtFirAnalysisSessionComponent {
 
-    override fun getKtExpressionType(expression: KtExpression): KtType = withValidityAssertion {
+    override fun getKtExpressionType(expression: KtExpression): KtType? = withValidityAssertion {
         when (val fir = expression.unwrap().getOrBuildFir(firResolveState)) {
             is FirExpression -> fir.typeRef.coneType.asKtType()
             is FirNamedReference -> fir.getReferencedElementType().asKtType()
             is FirStatement -> with(analysisSession) { builtinTypes.UNIT }
-            else -> error("Unexpected ${fir?.let { it::class }}")
+            is FirTypeRef, is FirImport, is FirPackageDirective, is FirLabel -> null
+            else -> error("Unexpected ${fir?.let { it::class }} for ${expression::class} with text `${expression.text}`")
         }
     }
 

idea/idea-frontend-fir/testData/components/expressionType/nonExpression.kt

@@ -0,0 +1 @@
+import java.<expr>util</expr>.ArrayList

idea/idea-frontend-fir/testData/components/expressionType/nonExpression.txt

@@ -0,0 +1,2 @@
+expression: util
+type: null

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/components/AbstractHLExpressionTypeTest.kt

@@ -21,7 +21,7 @@ abstract class AbstractHLExpressionTypeTest : AbstractHLApiSingleFileTest() {
     override fun doTestByFileStructure(ktFile: KtFile, module: TestModule, testServices: TestServices) {
         val expression = testServices.expressionMarkerProvider.getSelectedElement(ktFile) as KtExpression
         val type = executeOnPooledThreadInReadAction {
-            analyse(expression) { expression.getKtType().render() }
+            analyse(expression) { expression.getKtType()?.render() }
         }
         val actual = buildString {
             appendLine("expression: ${expression.text}")

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/components/HLExpressionTypeTestGenerated.java

@@ -78,6 +78,12 @@ public class HLExpressionTypeTestGenerated extends AbstractHLExpressionTypeTest
         runTest("idea/idea-frontend-fir/testData/components/expressionType/intLiteral.kt");
     }
 
+    @Test
+    @TestMetadata("nonExpression.kt")
+    public void testNonExpression() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/expressionType/nonExpression.kt");
+    }
+
     @Test
     @TestMetadata("property.kt")
     public void testProperty() throws Exception {