add def-use chains

Author: MForest7

usvm-ts-dataflow/src/main/kotlin/org/usvm/dataflow/ts/infer/EtsDefUseChains.kt

@@ -0,0 +1,170 @@
+package org.usvm.dataflow.ts.infer
+
+import kotlinx.collections.immutable.PersistentSet
+import kotlinx.collections.immutable.PersistentMap
+import kotlinx.collections.immutable.mutate
+import kotlinx.collections.immutable.persistentHashMapOf
+import kotlinx.collections.immutable.persistentSetOf
+import org.jacodb.ets.base.EtsArrayAccess
+import org.jacodb.ets.base.EtsAssignStmt
+import org.jacodb.ets.base.EtsBinaryExpr
+import org.jacodb.ets.base.EtsCallExpr
+import org.jacodb.ets.base.EtsCallStmt
+import org.jacodb.ets.base.EtsCastExpr
+import org.jacodb.ets.base.EtsEntity
+import org.jacodb.ets.base.EtsIfStmt
+import org.jacodb.ets.base.EtsInstanceCallExpr
+import org.jacodb.ets.base.EtsInstanceFieldRef
+import org.jacodb.ets.base.EtsInstanceOfExpr
+import org.jacodb.ets.base.EtsLocal
+import org.jacodb.ets.base.EtsReturnStmt
+import org.jacodb.ets.base.EtsStmt
+import org.jacodb.ets.base.EtsSwitchStmt
+import org.jacodb.ets.base.EtsTernaryExpr
+import org.jacodb.ets.base.EtsThrowStmt
+import org.jacodb.ets.base.EtsUnaryExpr
+import org.jacodb.ets.base.EtsValue
+import org.jacodb.ets.model.EtsMethod
+
+interface DefUseChains<Fact, Stmt> {
+    fun nextUses(fact: Fact, from: Stmt): Set<Stmt>
+}
+
+class EtsDefUseChains(
+    val method: EtsMethod,
+) : DefUseChains<ForwardTypeDomainFact, EtsStmt> {
+    private companion object {
+        fun EtsEntity.usedLocals(): List<String> = when (this) {
+            is EtsValue -> this.usedLocals()
+            is EtsUnaryExpr -> arg.usedLocals()
+            is EtsBinaryExpr -> left.usedLocals() + right.usedLocals()
+            is EtsCallExpr -> this.usedLocals()
+            is EtsCastExpr -> arg.usedLocals()
+            is EtsInstanceOfExpr -> arg.usedLocals()
+            is EtsTernaryExpr -> condition.usedLocals() + thenExpr.usedLocals() + elseExpr.usedLocals()
+            else -> emptyList()
+        }
+
+        fun EtsValue.usedLocals(): List<String> = when (this) {
+            is EtsLocal -> listOf(name)
+            is EtsInstanceFieldRef -> listOf(instance.name)
+            is EtsArrayAccess -> array.usedLocals() + index.usedLocals()
+            else -> emptyList()
+        }
+
+        fun EtsCallExpr.usedLocals(): List<String> = when (this) {
+            is EtsInstanceCallExpr -> listOf(instance.name) + args.flatMap { it.usedLocals() }
+            else -> args.flatMap { it.usedLocals() }
+        }
+
+        fun EtsStmt.usedLocals(): List<String> = when (this) {
+            is EtsAssignStmt -> lhv.usedLocals() + rhv.usedLocals()
+            is EtsCallStmt -> expr.usedLocals()
+            is EtsReturnStmt -> returnValue?.usedLocals().orEmpty()
+            is EtsIfStmt -> condition.usedLocals()
+            is EtsSwitchStmt -> arg.usedLocals()
+            is EtsThrowStmt -> arg.usedLocals()
+            else -> emptyList()
+        }
+
+        fun EtsStmt.definedLocal(): String? =
+            if (this is EtsAssignStmt && lhv is EtsLocal) (lhv as EtsLocal).name else null
+
+        fun EtsEntity.usedThisFields(): List<String> = when (this) {
+            is EtsValue -> this.usedThisFields()
+            is EtsUnaryExpr -> arg.usedThisFields()
+            is EtsBinaryExpr -> left.usedThisFields() + right.usedThisFields()
+            is EtsCallExpr -> this.usedThisFields()
+            is EtsCastExpr -> arg.usedThisFields()
+            is EtsInstanceOfExpr -> arg.usedThisFields()
+            is EtsTernaryExpr -> condition.usedThisFields() + thenExpr.usedThisFields() + elseExpr.usedThisFields()
+            else -> emptyList()
+        }
+
+        fun EtsValue.usedThisFields(): List<String> = when (this) {
+            is EtsLocal -> emptyList()
+            is EtsInstanceFieldRef -> if (instance.name == "this") listOf(field.name) else emptyList()
+            is EtsArrayAccess -> array.usedThisFields() + index.usedThisFields()
+            else -> emptyList()
+        }
+
+        fun EtsCallExpr.usedThisFields(): List<String> = when (this) {
+            is EtsInstanceCallExpr -> args.flatMap { it.usedThisFields() }
+            else -> args.flatMap { it.usedThisFields() }
+        }
+
+        fun EtsStmt.usedThisFields(): List<String> = when (this) {
+            is EtsAssignStmt -> lhv.usedThisFields() + rhv.usedThisFields()
+            is EtsCallStmt -> expr.usedThisFields()
+            is EtsReturnStmt -> returnValue?.usedThisFields().orEmpty()
+            is EtsIfStmt -> condition.usedThisFields()
+            is EtsSwitchStmt -> arg.usedThisFields()
+            is EtsThrowStmt -> arg.usedThisFields()
+            else -> emptyList()
+        }
+
+        fun EtsStmt.definedThisField(): String? =
+            if (this is EtsAssignStmt && lhv is EtsInstanceFieldRef && (lhv as EtsInstanceFieldRef).instance.name == "this")
+                (lhv as EtsInstanceFieldRef).field.name else null
+
+        fun EtsStmt.defined(): String? = definedLocal() ?: definedThisField()?.let { "this.$it" }
+
+        fun EtsStmt.used(): List<String> = usedLocals() + usedThisFields().map { "this.$it" }
+    }
+
+    private val nextUse =
+        Array<PersistentMap<String, PersistentSet<Int>>>(method.cfg.stmts.size) { persistentHashMapOf() }
+
+    private val usedFields = method.cfg.stmts.flatMapTo(hashSetOf()) { it.usedThisFields() }
+
+    init {
+        val queue = method.cfg.stmts.toHashSet()
+        while (queue.isNotEmpty()) {
+            val stmt = queue.first()
+            queue.remove(stmt)
+
+            val useArcs = method.cfg.successors(stmt).fold(persistentHashMapOf<String, PersistentSet<Int>>()) { acc, succ ->
+                val succIndex = succ.location.index
+                val succDefined = succ.defined()
+                val succUsed = succ.used().toHashSet()
+                acc.mutate { sum ->
+                    for ((name, arcs) in nextUse[succIndex]) {
+                        if (name == succDefined || name in succUsed) continue
+                        sum.merge(name, arcs) { a, b -> a.addAll(b) }
+                    }
+                    for (name in succUsed) {
+                        if (name == succDefined) continue
+                        sum.compute(name) { _, prev -> (prev ?: persistentSetOf()).add(succIndex) }
+                    }
+                }
+            }
+
+            val stmtIndex = stmt.location.index
+            if (useArcs != nextUse[stmtIndex]) {
+                nextUse[stmtIndex] = useArcs
+                if (stmt !in method.cfg.entries) queue.addAll(method.cfg.predecessors(stmt))
+            }
+        }
+    }
+
+    override fun nextUses(fact: ForwardTypeDomainFact, from: EtsStmt): Set<EtsStmt> {
+        if (from.location.index < 0) return setOf(method.cfg.stmts.first())
+        val accessPath = (fact as? ForwardTypeDomainFact.TypedVariable ?: return method.cfg.successors(from)).variable
+
+        when (val base = accessPath.base) {
+            is AccessPathBase.Local -> {
+                val name = base.name
+                val nextStmts = nextUse[from.location.index][name] ?: return method.cfg.successors(from)
+                return nextStmts.mapTo(hashSetOf()) { method.cfg.stmts[it] }
+            }
+            is AccessPathBase.This -> {
+                val field = accessPath.accesses.firstOrNull() as? FieldAccessor
+                    ?: return method.cfg.successors(from)
+                if (field.name !in usedFields) return emptySet()
+                val nextStmts = nextUse[from.location.index]["this.$field"] ?: return method.cfg.successors(from)
+                return nextStmts.mapTo(hashSetOf()) { method.cfg.stmts[it] }
+            }
+            else -> return method.cfg.successors(from)
+        }
+    }
+}