FIR: support resolve of implicit Unit types

compiler/fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -36,6 +36,8 @@ import org.jetbrains.kotlin.types.Variance
 
 class RawFirBuilder(val session: FirSession) {
 
+    private val implicitUnitType = FirImplicitUnitType(session, null)
+
     fun buildFirFile(file: KtFile): FirFile {
         return file.accept(Visitor(), Unit) as FirFile
     }
@@ -356,6 +358,7 @@ class RawFirBuilder(val session: FirSession) {
                 // TODO: what if name is not null but we're in expression position?
                 return FirExpressionStub(session, function)
             }
+            val typeReference = function.typeReference
             val firFunction = FirMemberFunctionImpl(
                 session,
                 function,
@@ -370,7 +373,11 @@ class RawFirBuilder(val session: FirSession) {
                 function.hasModifier(KtTokens.TAILREC_KEYWORD),
                 function.hasModifier(KtTokens.EXTERNAL_KEYWORD),
                 function.receiverTypeReference.convertSafe(),
-                function.typeReference.toFirOrImplicitType(),
+                if (function.hasBlockBody()) {
+                    typeReference?.toFirOrImplicitType() ?: implicitUnitType
+                } else {
+                    typeReference.toFirOrImplicitType()
+                },
                 function.bodyExpression.toFirBody()
             )
             function.extractAnnotationsTo(firFunction)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/impl/FirTypeResolverImpl.kt

@@ -5,6 +5,8 @@
 
 package org.jetbrains.kotlin.fir.resolve.impl
 
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.resolve.FirQualifierResolver
 import org.jetbrains.kotlin.fir.resolve.FirTypeResolver
 import org.jetbrains.kotlin.fir.scopes.FirPosition
@@ -54,6 +56,8 @@ class FirTypeResolverImpl : FirTypeResolver {
         }
     }
 
+    private val implicitUnitTypeSymbols = mutableMapOf<FirSession, ConeSymbol>()
+
     override fun resolveToSymbol(
         type: FirType,
         scope: FirScope,
@@ -83,6 +87,15 @@ class FirTypeResolverImpl : FirTypeResolver {
                 // TODO: Imports
                 resolvedSymbol ?: qualifierResolver.resolveSymbol(type.qualifier)
             }
+            is FirImplicitUnitType -> implicitUnitTypeSymbols[type.session] ?: run {
+                var resolvedSymbol: ConeSymbol? = null
+                scope.processClassifiersByName(KotlinBuiltIns.FQ_NAMES.unit.shortName(), position) {
+                    resolvedSymbol = (it as ConeClassLikeSymbol)
+                    resolvedSymbol == null
+                }
+                implicitUnitTypeSymbols[type.session] = resolvedSymbol!!
+                resolvedSymbol
+            }
             else -> null
         }
     }
@@ -112,6 +125,9 @@ class FirTypeResolverImpl : FirTypeResolver {
                     type.returnType.coneTypeUnsafe()
                 )
             }
+            is FirImplicitUnitType -> {
+                resolveToSymbol(type, scope, position)!!.toConeKotlinType(emptyList())!!
+            }
             is FirDynamicType, is FirImplicitType, is FirDelegatedType -> {
                 ConeKotlinErrorType("Not supported: ${type::class.simpleName}")
             }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirRenderer.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.symbols.ConeClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.ConeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.impl.FirImplicitUnitType
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.name.SpecialNames
 import org.jetbrains.kotlin.types.Variance
@@ -410,7 +411,7 @@ class FirRenderer(builder: StringBuilder) : FirVisitorVoid() {
     }
 
     override fun visitImplicitType(implicitType: FirImplicitType) {
-        print("<implicit>")
+        print(if (implicitType is FirImplicitUnitType) "kotlin.Unit" else "<implicit>")
     }
 
     override fun visitTypeWithNullability(typeWithNullability: FirTypeWithNullability) {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirImplicitUnitType.kt

@@ -0,0 +1,19 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.types.impl
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.types.FirImplicitType
+
+class FirImplicitUnitType(
+    override val session: FirSession,
+    override val psi: PsiElement?
+) : FirImplicitType {
+    override val annotations: List<FirAnnotationCall>
+        get() = emptyList()
+}