Fix ETS and TypeScript frontend contracts and simplify reviewed code

jacodb-ets/src/main/kotlin/org/jacodb/ets/dsl/BlockCfg.kt

@@ -16,9 +16,6 @@
 
 package org.jacodb.ets.dsl
 
-import org.jacodb.ets.utils.IdentityHashSet
-import java.util.IdentityHashMap
-
 data class Block(
     val id: Int,
     val statements: List<BlockStmt>,
@@ -29,41 +26,73 @@ data class BlockCfg(
     val successors: Map<Int, List<Int>>, // for IF stmt, successors are (true, false) branches
 )
 
+private class NodeOccurrence(
+    val node: Node,
+    val thenBranch: List<NodeOccurrence> = emptyList(),
+    val elseBranch: List<NodeOccurrence> = emptyList(),
+)
+
 fun Program.toBlockCfg(): BlockCfg {
-    val labelToNode: MutableMap<String, Node> = hashMapOf()
-    val targets: MutableSet<Node> = IdentityHashSet()
+    fun collectOccurrences(nodes: List<Node>): List<NodeOccurrence> = nodes.map { node ->
+        if (node is If) {
+            NodeOccurrence(
+                node = node,
+                thenBranch = collectOccurrences(node.thenBranch),
+                elseBranch = collectOccurrences(node.elseBranch),
+            )
+        } else {
+            NodeOccurrence(node)
+        }
+    }
+
+    val occurrences = collectOccurrences(nodes)
+    val labelToOccurrence: MutableMap<String, NodeOccurrence> = hashMapOf()
+    val targets: MutableSet<NodeOccurrence> = hashSetOf()
 
-    fun findLabels(nodes: List<Node>) {
-        if (nodes.lastOrNull() is Label) {
+    fun findLabels(nodes: List<NodeOccurrence>) {
+        if (nodes.lastOrNull()?.node is Label) {
             error("Label at the end of the block: $nodes")
         }
-        for ((stmt, next) in nodes.zipWithNext()) {
-            if (stmt is Label) {
-                check(next !is Label) { "Two labels in a row: $stmt, $next" }
-                check(next !is Goto) { "Label followed by goto: $stmt, $next" }
-                check(stmt.name !in labelToNode) { "Duplicate label: ${stmt.name}" }
-                labelToNode[stmt.name] = next
-            }
+        for ((occurrence, next) in nodes.zipWithNext()) {
+            val label = occurrence.node as? Label ?: continue
+
+            check(next.node !is Label) { "Two labels in a row: $label, ${next.node}" }
+            check(next.node !is Goto) { "Label followed by goto: $label, ${next.node}" }
+            check(label.name !in labelToOccurrence) { "Duplicate label: ${label.name}" }
+            labelToOccurrence[label.name] = next
         }
-        for (node in nodes) {
-            if (node is If) {
-                findLabels(node.thenBranch)
-                findLabels(node.elseBranch)
-            }
-            if (node is Goto) {
-                targets += labelToNode[node.targetLabel] ?: error("Unknown label: ${node.targetLabel}")
+        for (occurrence in nodes) {
+            findLabels(occurrence.thenBranch)
+            findLabels(occurrence.elseBranch)
+        }
+    }
+
+    findLabels(occurrences)
+
+    fun findTargets(nodes: List<NodeOccurrence>) {
+        for (occurrence in nodes) {
+            when (val node = occurrence.node) {
+                is If -> {
+                    findTargets(occurrence.thenBranch)
+                    findTargets(occurrence.elseBranch)
+                }
+
+                is Goto -> targets += labelToOccurrence[node.targetLabel]
+                    ?: error("Unknown label: ${node.targetLabel}")
+
+                else -> Unit
             }
         }
     }
 
-    findLabels(nodes)
+    findTargets(occurrences)
 
     val blocks: MutableList<Block> = mutableListOf()
     val successors: MutableMap<Int, List<Int>> = hashMapOf()
-    val stmtToBlock: MutableMap<BlockStmt, Int> = IdentityHashMap()
-    val nodeToStmt: MutableMap<Node, BlockStmt> = IdentityHashMap()
+    val occurrenceToBlock: MutableMap<NodeOccurrence, Int> = hashMapOf()
+    val pendingGotos: MutableList<Pair<Int, String>> = mutableListOf()
 
-    fun buildBlocks(nodes: List<Node>): Pair<Int, Int>? {
+    fun buildBlocks(nodes: List<NodeOccurrence>): Pair<Int, Int>? {
         if (nodes.isEmpty()) return null
 
         lateinit var currentBlock: MutableList<BlockStmt>
@@ -75,14 +104,20 @@ fun Program.toBlockCfg(): BlockCfg {
             return block
         }
 
+        fun newEmptyBranch(): Block {
+            val branch = Block(blocks.size, emptyList())
+            blocks += branch
+            return branch
+        }
+
         var block = newBlock()
         val firstBlockId = block.id
 
-        for (node in nodes) {
+        for ((index, occurrence) in nodes.withIndex()) {
+            val node = occurrence.node
             if (node is Label) continue
 
-            if (node in targets && currentBlock.isNotEmpty()) {
-                block.statements.forEach { stmtToBlock[it] = block.id }
+            if (occurrence in targets && currentBlock.isNotEmpty()) {
                 val prevBlock = block
                 block = newBlock()
                 successors[prevBlock.id] = listOf(block.id)
@@ -97,17 +132,16 @@ fun Program.toBlockCfg(): BlockCfg {
                     is CustomEts -> BlockCustomEts(node.toEts)
                     else -> error("Unexpected node: $node")
                 }
-                nodeToStmt[node] = stmt
+                occurrenceToBlock[occurrence] = block.id
                 currentBlock += stmt
             }
 
             if (node is If) {
-                block.statements.forEach { stmtToBlock[it] = block.id }
                 val ifBlock = block
                 block = newBlock()
 
-                val thenBlocks = buildBlocks(node.thenBranch)
-                val elseBlocks = buildBlocks(node.elseBranch)
+                val thenBlocks = buildBlocks(occurrence.thenBranch)
+                val elseBlocks = buildBlocks(occurrence.elseBranch)
 
                 when {
                     thenBlocks != null && elseBlocks != null -> {
@@ -147,29 +181,41 @@ fun Program.toBlockCfg(): BlockCfg {
                     }
 
                     else -> {
-                        successors[ifBlock.id] = listOf(block.id)
+                        val trueBranch = newEmptyBranch()
+                        val falseBranch = newEmptyBranch()
+                        successors[ifBlock.id] = listOf(trueBranch.id, falseBranch.id)
+                        successors[trueBranch.id] = listOf(block.id)
+                        successors[falseBranch.id] = listOf(block.id)
                     }
                 }
             } else if (node is Goto) {
-                val targetNode = labelToNode[node.targetLabel] ?: error("Unknown label: ${node.targetLabel}")
-                val target = nodeToStmt[targetNode] ?: error("No statement for $targetNode")
-                val targetBlockId = stmtToBlock[target] ?: error("No block for $target")
-                successors[block.id] = listOf(targetBlockId)
-                block.statements.forEach { stmtToBlock[it] = block.id }
+                pendingGotos += block.id to node.targetLabel
                 block = newBlock()
             } else if (node is Return) {
                 successors[block.id] = emptyList()
-                break
+                if (index < nodes.lastIndex) {
+                    block = newBlock()
+                }
             }
         }
 
-        block.statements.forEach { stmtToBlock[it] = block.id }
         val lastBlockId = block.id
 
         return Pair(firstBlockId, lastBlockId)
     }
 
-    buildBlocks(nodes)
+    buildBlocks(occurrences)
+
+    for ((sourceBlockId, label) in pendingGotos) {
+        val targetOccurrence = labelToOccurrence.getValue(label)
+        val targetBlockId = occurrenceToBlock[targetOccurrence]
+            ?: error("No block for ${targetOccurrence.node}")
+        successors[sourceBlockId] = listOf(targetBlockId)
+    }
+
+    for (block in blocks) {
+        successors.putIfAbsent(block.id, emptyList())
+    }
 
     return BlockCfg(blocks, successors)
 }

jacodb-ets/src/main/kotlin/org/jacodb/ets/dsl/Expr.kt

@@ -16,6 +16,8 @@
 
 package org.jacodb.ets.dsl
 
+import org.jacodb.ets.model.EtsClassSignature
+
 sealed interface Expr
 
 data class Local(val name: String) : LValue {
@@ -96,7 +98,14 @@ data class FieldRef(
 data class StaticFieldRef(
     val className: String,
     val fieldName: String,
+    val classSignature: EtsClassSignature? = null,
 ) : LValue {
+    init {
+        require(classSignature == null || classSignature.name == className) {
+            "Static field class name does not match its signature"
+        }
+    }
+
     override fun toString() = "$className.$fieldName"
 }
 

jacodb-ets/src/main/kotlin/org/jacodb/ets/dsl/ProgramBuilder.kt

@@ -18,6 +18,7 @@
 
 package org.jacodb.ets.dsl
 
+import org.jacodb.ets.model.EtsClassSignature
 import org.jacodb.ets.model.EtsEntity
 import org.jacodb.ets.model.EtsStmt
 import org.jacodb.ets.model.EtsStmtLocation
@@ -43,6 +44,8 @@ fun ProgramBuilder.const(value: String) = ConstantString(value)
 
 fun ProgramBuilder.fieldRef(instance: Expr, fieldName: String) = FieldRef(instance, fieldName)
 fun ProgramBuilder.staticFieldRef(className: String, fieldName: String) = StaticFieldRef(className, fieldName)
+fun ProgramBuilder.staticFieldRef(owner: EtsClassSignature, fieldName: String) =
+    StaticFieldRef(owner.name, fieldName, classSignature = owner)
 fun ProgramBuilder.arrayAccess(array: Expr, index: Expr) = ArrayAccess(array, index)
 
 fun Expr.field(name: String) = FieldRef(this, name)
@@ -84,13 +87,23 @@ class CustomBinaryExpr(
     val toEts: (left: EtsEntity, right: EtsEntity) -> EtsEntity,
 ) : Expr
 
+private fun List<Node>.snapshot(): List<Node> = map { node ->
+    when (node) {
+        is If -> node.copy(
+            thenBranch = node.thenBranch.snapshot(),
+            elseBranch = node.elseBranch.snapshot(),
+        )
+
+        else -> node
+    }
+}
+
 class ProgramBuilderImpl : ProgramBuilder {
     private val _nodes: MutableList<Node> = mutableListOf()
     val nodes: List<Node> get() = _nodes
 
-    fun build(): Program {
-        return Program(nodes)
-    }
+    /** Capture the completed node tree; later builder calls cannot change this program. */
+    fun build(): Program = Program(nodes.snapshot())
 
     override fun nop() {
         _nodes += Nop
@@ -138,10 +151,8 @@ class IfBuilder : ProgramBuilder {
         elseBuilder.apply(block)
     }
 
-    fun elseIf(condition: Expr, block: ProgramBuilder.() -> Unit): IfBuilder {
+    fun elseIf(condition: Expr, block: ProgramBuilder.() -> Unit): IfBuilder =
         elseBuilder.ifStmt(condition, block)
-        return this
-    }
 
     override fun ifStmt(condition: Expr, block: ProgramBuilder.() -> Unit): IfBuilder =
         thenBuilder.ifStmt(condition, block)