Embed type via builder

plugins/formal-verification/formver.core/src/org/jetbrains/kotlin/formver/conversion/ProgramConverter.kt

@@ -154,41 +154,22 @@ class ProgramConverter(val session: FirSession, override val config: PluginConfi
         return embedding
     }
 
-    override fun embedType(type: ConeKotlinType): TypeEmbedding = when {
-        type is ConeErrorType -> error("Encountered an erroneous type: $type")
-        type is ConeTypeParameterType -> buildType { isNullable = true; any() }
-        type.isUnit -> buildType { unit() }
-        type.isInt -> buildType { int() }
-        type.isBoolean -> buildType { boolean() }
-        type.isNothing -> buildType { nothing() }
-        type.isSomeFunctionType(session) -> {
-            val receiverType: TypeEmbedding? = type.receiverType(session)?.let { embedType(it) }
-            val paramTypes: List<TypeEmbedding> = type.valueParameterTypesWithoutReceivers(session).map(::embedType)
-            val returnType: TypeEmbedding = embedType(type.returnType(session))
-            FunctionTypeEmbedding(receiverType, paramTypes, returnType, returnsUnique = false)
+    override fun embedType(type: ConeKotlinType): TypeEmbedding = buildType { embedTypeWithBuilder(type) }
+
+    // Note: keep in mind that this function is necessary to resolve the name of the function!
+    override fun embedType(symbol: FirFunctionSymbol<*>): FunctionTypeEmbedding = buildFunctionType {
+        symbol.receiverType?.let {
+            withReceiver { embedTypeWithBuilder(it) }
         }
-        type.isNullable -> NullableTypeEmbedding(embedType(type.withNullability(ConeNullability.NOT_NULL, session.typeContext)))
-        type.isAny -> buildType { any() }
-        type is ConeClassLikeType -> {
-            val classLikeSymbol = type.toClassSymbol(session)
-            if (classLikeSymbol is FirRegularClassSymbol) {
-                embedClass(classLikeSymbol)
-            } else {
-                unimplementedTypeEmbedding(type)
+        symbol.valueParameterSymbols.forEach { param ->
+            withParam {
+                embedTypeWithBuilder(param.resolvedReturnType)
             }
         }
-        else -> unimplementedTypeEmbedding(type)
+        withReturnType { embedTypeWithBuilder(symbol.resolvedReturnType) }
+        returnsUnique = symbol.isUnique(session) || symbol is FirConstructorSymbol
     }
 
-    // Note: keep in mind that this function is necessary to resolve the name of the function!
-    override fun embedType(symbol: FirFunctionSymbol<*>): FunctionTypeEmbedding =
-        FunctionTypeEmbedding(
-            receiverType = symbol.receiverType?.let(::embedType),
-            paramTypes = symbol.valueParameterSymbols.map { embedType(it.resolvedReturnType) },
-            returnType = embedType(symbol.resolvedReturnTypeRef.coneType),
-            returnsUnique = symbol.isUnique(session) || symbol is FirConstructorSymbol,
-        )
-
     override fun embedProperty(symbol: FirPropertySymbol): PropertyEmbedding = if (symbol.isExtension) {
         embedCustomProperty(symbol)
     } else {
@@ -404,13 +385,45 @@ class ProgramConverter(val session: FirSession, override val config: PluginConfi
         return FunctionBodyEmbedding(seqnBuilder.block, returnTarget, bodyExp)
     }
 
-    private fun unimplementedTypeEmbedding(type: ConeKotlinType): TypeEmbedding =
+    private fun TypeBuilder.embedTypeWithBuilder(type: ConeKotlinType): PretypeBuilder = when {
+        type is ConeErrorType -> error("Encountered an erroneous type: $type")
+        type is ConeTypeParameterType -> {
+            isNullable = true; any()
+        }
+        type.isUnit -> unit()
+        type.isInt -> int()
+        type.isBoolean -> boolean()
+        type.isNothing -> nothing()
+        type.isSomeFunctionType(session) -> function {
+            type.receiverType(session)?.let { withReceiver { embedTypeWithBuilder(it) } }
+            type.valueParameterTypesWithoutReceivers(session).forEach { param ->
+                withParam { embedTypeWithBuilder(param) }
+            }
+            withReturnType { embedTypeWithBuilder(type.returnType(session)) }
+        }
+        type.isNullable -> {
+            isNullable = true
+            embedTypeWithBuilder(type.withNullability(ConeNullability.NOT_NULL, session.typeContext))
+        }
+        type.isAny -> any()
+        type is ConeClassLikeType -> {
+            val classLikeSymbol = type.toClassSymbol(session)
+            if (classLikeSymbol is FirRegularClassSymbol) {
+                existing(embedClass(classLikeSymbol))
+            } else {
+                unimplementedTypeEmbedding(type)
+            }
+        }
+        else -> unimplementedTypeEmbedding(type)
+    }
+
+    private fun TypeBuilder.unimplementedTypeEmbedding(type: ConeKotlinType): PretypeBuilder =
         when (config.behaviour) {
             UnsupportedFeatureBehaviour.THROW_EXCEPTION ->
                 throw NotImplementedError("The embedding for type $type is not yet implemented.")
             UnsupportedFeatureBehaviour.ASSUME_UNREACHABLE -> {
                 errorCollector.addMinorError("Requested type $type, for which we do not yet have an embedding.")
-                buildType { unit() }
+                unit()
             }
         }
 }

plugins/formal-verification/formver.core/src/org/jetbrains/kotlin/formver/embeddings/PretypeBuilder.kt

@@ -44,6 +44,7 @@ class FunctionPretypeBuilder : PretypeBuilder {
     private val paramTypes = mutableListOf<TypeEmbedding>()
     private var receiverType: TypeEmbedding? = null
     private var returnType: TypeEmbedding? = null
+    var returnsUnique: Boolean = false
 
     fun withParam(paramInit: TypeBuilder.() -> PretypeBuilder) {
         paramTypes.add(buildType { paramInit() })
@@ -61,7 +62,7 @@ class FunctionPretypeBuilder : PretypeBuilder {
 
     override fun complete(): TypeEmbedding {
         require(returnType != null) { "Return type not set" }
-        return FunctionTypeEmbedding(receiverType, paramTypes, returnType!!, returnsUnique = false)
+        return FunctionTypeEmbedding(receiverType, paramTypes, returnType!!, returnsUnique)
     }
 }
 
@@ -78,3 +79,8 @@ class ClassPretypeBuilder : PretypeBuilder {
         return ClassTypeEmbedding(className!!)
     }
 }
+
+// TODO: ensure we can build the types with the builders, without hacks like this.
+class ExistingPretypeBuilder(val embedding: TypeEmbedding) : PretypeBuilder {
+    override fun complete(): TypeEmbedding = embedding
+}

plugins/formal-verification/formver.core/src/org/jetbrains/kotlin/formver/embeddings/TypeBuilder.kt

@@ -29,6 +29,7 @@ class TypeBuilder {
     fun boolean() = BooleanPretypeBuilder
     fun function(init: FunctionPretypeBuilder.() -> Unit) = FunctionPretypeBuilder().also { it.init() }
     fun klass(init: ClassPretypeBuilder.() -> Unit) = ClassPretypeBuilder().also { it.init() }
+    fun existing(embedding: TypeEmbedding) = ExistingPretypeBuilder(embedding)
 }
 
 fun TypeBuilder.nullableAny(): AnyPretypeBuilder {