Basic USVM TS type system with type coercion

usvm-ts/src/main/kotlin/org/usvm/machine/TSContext.kt

@@ -27,6 +27,8 @@ import org.usvm.machine.interpreter.TSStepScope
 import org.usvm.machine.types.ExprWithTypeConstraint
 import org.usvm.machine.types.FakeType
 import org.usvm.types.single
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.contract
 
 typealias TSSizeSort = UBv32Sort
 
@@ -51,8 +53,6 @@ class TSContext(
     // TODO fix conjuncts
     fun mkTruthyExpr(expr: UExpr<out USort>, scope: TSStepScope): UBoolExpr = scope.calcOnState {
         if (expr.isFakeObject()) {
-            expr as UConcreteHeapRef
-
             val falseBranchGround = makeSymbolicPrimitive(boolSort)
 
             val conjuncts = mutableListOf<ExprWithTypeConstraint<UBoolSort>>()
@@ -117,10 +117,13 @@ class TSContext(
         }
     }
 
+    @OptIn(ExperimentalContracts::class)
     fun UExpr<out USort>.isFakeObject(): Boolean {
-        if (sort !is UAddressSort) return false
+        contract {
+            returns(true) implies (this@isFakeObject is UConcreteHeapRef)
+        }
 
-        return this is UConcreteHeapRef && address > MAGIC_OFFSET
+        return sort == addressSort && this is UConcreteHeapRef && address > MAGIC_OFFSET
     }
 
     fun mkFakeValue(

usvm-ts/src/main/kotlin/org/usvm/machine/expr/TSExprResolver.kt

@@ -74,10 +74,10 @@ import org.usvm.UExpr
 import org.usvm.USort
 import org.usvm.api.makeSymbolicPrimitive
 import org.usvm.collection.field.UFieldLValue
-import org.usvm.machine.types.FakeType
 import org.usvm.machine.TSContext
 import org.usvm.machine.interpreter.TSStepScope
 import org.usvm.machine.operator.TSBinaryOperator
+import org.usvm.machine.types.FakeType
 import org.usvm.memory.ULValue
 import org.usvm.memory.URegisterStackLValue
 

usvm-ts/src/main/kotlin/org/usvm/machine/interpreter/TSInterpreter.kt

@@ -89,11 +89,16 @@ class TSInterpreter(
     private fun visitIfStmt(scope: TSStepScope, stmt: EtsIfStmt) {
         val exprResolver = exprResolverWithScope(scope)
 
-        val boolExpr = exprResolver.resolve(stmt.condition)?.asExpr(ctx.boolSort)
-            ?: run {
-                logger.warn { "Failed to resolve condition: $stmt" }
-                return
-            }
+        val conditionExpr = exprResolver.resolve(stmt.condition) ?: run {
+            logger.warn { "Failed to resolve condition: $stmt" }
+            return
+        }
+
+        val boolExpr = if (conditionExpr.sort == ctx.boolSort) {
+            conditionExpr.asExpr(ctx.boolSort)
+        } else {
+            ctx.mkTruthyExpr(conditionExpr, scope)
+        }
 
         val (negStmt, posStmt) = applicationGraph.successors(stmt).take(2).toList()
 
@@ -102,7 +107,7 @@ class TSInterpreter(
             posStmt,
             negStmt,
             blockOnTrueState = { newStmt(posStmt) },
-            blockOnFalseState = { newStmt(negStmt) }
+            blockOnFalseState = { newStmt(negStmt) },
         )
     }
 

usvm-ts/src/main/kotlin/org/usvm/machine/operator/TSBinaryOperator.kt

@@ -4,22 +4,20 @@ import io.ksmt.sort.KFp64Sort
 import io.ksmt.utils.asExpr
 import io.ksmt.utils.cast
 import org.usvm.UAddressSort
-import org.usvm.UBoolExpr
 import org.usvm.UBoolSort
-import org.usvm.UConcreteHeapRef
 import org.usvm.UExpr
 import org.usvm.USort
 import org.usvm.api.makeSymbolicPrimitive
 import org.usvm.api.typeStreamOf
-import org.usvm.machine.types.FakeType
 import org.usvm.machine.TSContext
 import org.usvm.machine.expr.TSUndefinedSort
 import org.usvm.machine.expr.tctx
 import org.usvm.machine.interpreter.TSStepScope
 import org.usvm.machine.types.ExprWithTypeConstraint
+import org.usvm.machine.types.FakeType
 import org.usvm.machine.types.iteWriteIntoFakeObject
 import org.usvm.types.single
-import org.usvm.util.boolToFpSort
+import org.usvm.util.boolToFp
 
 sealed interface TSBinaryOperator {
 
@@ -113,9 +111,6 @@ sealed interface TSBinaryOperator {
                 val groundFalseBranch = makeSymbolicPrimitive(boolSort)
 
                 if (lhs.isFakeObject() && rhs.isFakeObject()) {
-                    lhs as UConcreteHeapRef
-                    rhs as UConcreteHeapRef
-
                     val lhsType = memory.typeStreamOf(lhs).single() as FakeType
                     val rhsType = memory.typeStreamOf(rhs).single() as FakeType
 
@@ -153,7 +148,7 @@ sealed interface TSBinaryOperator {
                     conjuncts += ExprWithTypeConstraint(
                         constraint = mkAnd(lhsType.boolTypeExpr, rhsType.fpTypeExpr),
                         expr = mkFpEqualExpr(
-                            boolToFpSort(memory.read(getIntermediateBoolLValue(lhs.address))),
+                            boolToFp(memory.read(getIntermediateBoolLValue(lhs.address))),
                             memory.read(getIntermediateFpLValue(rhs.address))
                         )
                     )
@@ -162,15 +157,14 @@ sealed interface TSBinaryOperator {
                         constraint = mkAnd(lhsType.fpTypeExpr, rhsType.boolTypeExpr),
                         expr = mkFpEqualExpr(
                             memory.read(getIntermediateFpLValue(lhs.address)),
-                            boolToFpSort(memory.read(getIntermediateBoolLValue(rhs.address)))
+                            boolToFp(memory.read(getIntermediateBoolLValue(rhs.address)))
                         )
                     )
 
                     // TODO: support objects
                 }
 
                 if (lhs.isFakeObject()) {
-                    lhs as UConcreteHeapRef
                     val lhsType = memory.typeStreamOf(lhs).single() as FakeType
 
                     scope.assert(lhsType.mkExactlyOneTypeConstraint(ctx))
@@ -189,7 +183,7 @@ sealed interface TSBinaryOperator {
                                 constraint = lhsType.fpTypeExpr,
                                 expr = mkFpEqualExpr(
                                     memory.read(getIntermediateFpLValue(lhs.address)),
-                                    boolToFpSort(rhs.asExpr(boolSort))
+                                    boolToFp(rhs.asExpr(boolSort))
                                 )
                             )
 
@@ -200,7 +194,7 @@ sealed interface TSBinaryOperator {
                             conjuncts += ExprWithTypeConstraint(
                                 constraint = lhsType.boolTypeExpr,
                                 expr = mkFpEqualExpr(
-                                    boolToFpSort(memory.read(getIntermediateBoolLValue(lhs.address))),
+                                    boolToFp(memory.read(getIntermediateBoolLValue(lhs.address))),
                                     rhs.asExpr(fp64Sort)
                                 )
                             )
@@ -235,7 +229,6 @@ sealed interface TSBinaryOperator {
                 }
 
                 if (rhs.isFakeObject()) {
-                    rhs as UConcreteHeapRef
                     val rhsType = memory.typeStreamOf(rhs).single() as FakeType
 
                     scope.assert(rhsType.mkExactlyOneTypeConstraint(ctx))
@@ -253,7 +246,7 @@ sealed interface TSBinaryOperator {
                             conjuncts += ExprWithTypeConstraint(
                                 constraint = rhsType.fpTypeExpr,
                                 expr = mkFpEqualExpr(
-                                    boolToFpSort(lhs.asExpr(boolSort)),
+                                    boolToFp(lhs.asExpr(boolSort)),
                                     memory.read(getIntermediateFpLValue(rhs.address))
                                 )
                             )
@@ -266,7 +259,7 @@ sealed interface TSBinaryOperator {
                                 constraint = rhsType.boolTypeExpr,
                                 expr = mkFpEqualExpr(
                                     lhs.asExpr(fp64Sort),
-                                    boolToFpSort(memory.read(getIntermediateBoolLValue(rhs.address)))
+                                    boolToFp(memory.read(getIntermediateBoolLValue(rhs.address)))
                                 )
                             )
 
@@ -325,11 +318,11 @@ sealed interface TSBinaryOperator {
             }
 
             if (lhs.sort is UBoolSort && rhs.sort is KFp64Sort) {
-                return mkFpEqualExpr(boolToFpSort(lhs.cast()), rhs.cast())
+                return mkFpEqualExpr(boolToFp(lhs.cast()), rhs.cast())
             }
 
             if (lhs.sort is KFp64Sort && rhs.sort is UBoolSort) {
-                return mkFpEqualExpr(lhs.cast(), boolToFpSort(rhs.cast()))
+                return mkFpEqualExpr(lhs.cast(), boolToFp(rhs.cast()))
             }
 
             // TODO: support string
@@ -399,8 +392,8 @@ sealed interface TSBinaryOperator {
         ): UExpr<out USort> {
             return mkFpAddExpr(
                 fpRoundingModeSortDefaultValue(),
-                boolToFpSort(lhs),
-                boolToFpSort(rhs)
+                boolToFp(lhs),
+                boolToFp(rhs)
             )
 
         }
@@ -444,11 +437,11 @@ sealed interface TSBinaryOperator {
 
             val fpValue = when {
                 lhs.sort is UBoolSort && rhs.sort is KFp64Sort -> {
-                    mkFpAddExpr(fpRoundingModeSortDefaultValue(), boolToFpSort(lhs.cast()), rhs.cast())
+                    mkFpAddExpr(fpRoundingModeSortDefaultValue(), boolToFp(lhs.cast()), rhs.cast())
                 }
 
                 lhs.sort is KFp64Sort && rhs.sort is UBoolSort -> {
-                    mkFpAddExpr(fpRoundingModeSortDefaultValue(), lhs.cast(), boolToFpSort(rhs.cast()))
+                    mkFpAddExpr(fpRoundingModeSortDefaultValue(), lhs.cast(), boolToFp(rhs.cast()))
                 }
 
                 else -> null

usvm-ts/src/main/kotlin/org/usvm/machine/types/FakeExprUtil.kt

@@ -21,7 +21,6 @@ fun <T : USort> TSState.extractValue(
 ): Pair<UExpr<T>?, UBoolExpr> = with(ctx) {
     when {
         value.isFakeObject() -> {
-            value as UConcreteHeapRef
             val lValue = extractIntermediateLValue(value.address)
 
             val type = memory.typeStreamOf(value).single() as FakeType

usvm-ts/src/main/kotlin/org/usvm/util/Utils.kt

@@ -6,5 +6,5 @@ import org.usvm.UExpr
 import org.usvm.machine.TSContext
 
 // Built-in KContext.bvToBool has identical implementation.
-fun TSContext.boolToFpSort(expr: UExpr<UBoolSort>): UExpr<KFp64Sort> =
+fun TSContext.boolToFp(expr: UExpr<UBoolSort>): UExpr<KFp64Sort> =
     mkIte(expr, mkFp64(1.0), mkFp64(0.0))