Improve statement unwrapping logic

rpgJavaInterpreter-core/src/main/kotlin/com/smeup/rpgparser/interpreter/InterpreterCore.kt

@@ -55,7 +55,7 @@ interface InterpreterCore {
     fun setIndicators(statement: WithRightIndicators, hi: BooleanValue, lo: BooleanValue, eq: BooleanValue)
     fun eval(expression: Expression): Value
     fun execute(statements: List<Statement>)
-    fun executeUnwrappedAt(unwrappedStatements: List<Statement>, offset: Int)
+    fun executeUnwrappedAt(unwrappedStatements: List<UnwrappedStatementData>, offset: Int)
     fun dbFile(name: String, statement: Statement): EnrichedDBFile
     fun toSearchValues(searchArgExpression: Expression, fileMetadata: FileMetadata): List<String>
     fun fillDataFrom(dbFile: EnrichedDBFile, record: Record)

rpgJavaInterpreter-core/src/main/kotlin/com/smeup/rpgparser/interpreter/control_flow_exceptions.kt

@@ -1,6 +1,6 @@
 package com.smeup.rpgparser.interpreter
 
-import com.smeup.rpgparser.parsing.ast.Statement
+import com.smeup.rpgparser.parsing.ast.UnwrappedStatementData
 import com.smeup.rpgparser.utils.indexOfTag
 import com.smeup.rpgparser.utils.runIfNotEmpty
 import java.io.File
@@ -24,7 +24,7 @@ class IterException : ControlFlowException()
  * @param tag The tag of the destination label
  */
 class GotoTopLevelException(val tag: String) : ControlFlowException() {
-    internal fun indexOfTaggedStatement(statements: List<Statement>) = statements.indexOfTag(tag)
+    internal fun indexOfTaggedStatement(statements: List<UnwrappedStatementData>) = statements.indexOfTag(tag)
 }
 
 // Useful to interrupt infinite cycles in tests
@@ -40,7 +40,7 @@ class GotoException private constructor(val tag: String) : ControlFlowException(
         operator fun invoke(tag: String): GotoException = GotoException(tag.uppercase(Locale.getDefault()))
     }
 
-    internal fun indexOfTaggedStatement(statements: List<Statement>) = statements.indexOfTag(tag)
+    internal fun indexOfTaggedStatement(statements: List<UnwrappedStatementData>) = statements.indexOfTag(tag)
 }
 
 class InterpreterTimeoutException(val programName: String, val elapsed: Long, val expected: Long) : ControlFlowException() {

rpgJavaInterpreter-core/src/main/kotlin/com/smeup/rpgparser/interpreter/internal_interpreter.kt

@@ -425,7 +425,7 @@ open class InternalInterpreter(
             execute(main.stmts)
         }.exceptionOrNull()
 
-        val unwrappedStatement = main.stmts.explode(true)
+        val unwrappedStatement = main.stmts.unwrap()
 
         // Recursive deal with top level goto flow
         while (throwable is GotoTopLevelException || throwable is GotoException) {
@@ -533,26 +533,12 @@ open class InternalInterpreter(
      * @param unwrappedStatements The unwrapped list of statements. It is up to the caller to ensure statements are unwrapped.
      * @param offset Offset to start the execution from.
      */
-    override fun executeUnwrappedAt(unwrappedStatements: List<Statement>, offset: Int) {
-        /**
-         * As we execute composite statements recursively, trying to sequentially execute
-         * the unwrapped statement list would result in executing every statement multiple times.
-         *
-         * The following instruction associates to each statement the offset to add in order to find
-         * the real next statement to execute in the list.
-         */
-        val offsetAwareStatements = unwrappedStatements.map {
-            WithOffset(
-                data = it,
-                offset = if (it is CompositeStatement) it.body.size else 0
-            )
-        }
-
+    override fun executeUnwrappedAt(unwrappedStatements: List<UnwrappedStatementData>, offset: Int) {
         var index = offset
-        while (index < offsetAwareStatements.size) {
-            val offsetStatement = offsetAwareStatements[index]
-            executeWithMute(offsetStatement.data)
-            index += offsetStatement.offset + 1
+        while (index < unwrappedStatements.size) {
+            val data = unwrappedStatements[index]
+            executeWithMute(data.statement)
+            index += data.nextOperationOffset + 1
         }
     }
 

rpgJavaInterpreter-core/src/main/kotlin/com/smeup/rpgparser/parsing/ast/statements.kt

@@ -111,6 +111,50 @@ interface CompositeStatement {
     val body: List<Statement>
 }
 
+data class UnwrappedStatementData(
+    var statement: Statement,
+    var nextOperationOffset: Int,
+    var parent: UnwrappedStatementData?
+)
+
+interface CustomStatementUnwrap {
+    fun unwrap(parent: UnwrappedStatementData? = null): List<UnwrappedStatementData>
+}
+
+/**
+ * Unwrap statements while keeping their structural data
+ * @see [InterpreterCore.executeUnwrappedAt]
+ */
+fun List<Statement>.unwrap(parent: UnwrappedStatementData? = null): List<UnwrappedStatementData> {
+    val result = mutableListOf<UnwrappedStatementData>()
+    forEach {
+        when (it) {
+            is CustomStatementUnwrap -> {
+                val unwrapped = it.unwrap(parent)
+                result.addAll(unwrapped)
+            }
+            is CompositeStatement -> {
+                val current = UnwrappedStatementData(it, it.body.size, parent)
+                val body = it.body.unwrap(current)
+
+                /**
+                 * body might contain CompositeStatements recursively
+                 * we need to update the offset with the actual body length after unwrapping
+                 */
+                current.nextOperationOffset = body.size
+
+                result.add(current)
+                result.addAll(body)
+            }
+            else -> {
+                val unwrapped = UnwrappedStatementData(it, 0, parent)
+                result.add(unwrapped)
+            }
+        }
+    }
+    return result
+}
+
 fun List<Statement>.explode(preserveCompositeStatement: Boolean = false): List<Statement> {
     val result = mutableListOf<Statement>()
     forEach {
@@ -145,8 +189,8 @@ data class ExecuteSubroutine(var subroutine: ReferenceByName<Subroutine>, overri
                 interpreter.execute(body)
             }.exceptionOrNull()
 
-            val unwrappedStatements = body.explode(true)
-            val containingCU = unwrappedStatements.firstOrNull()?.getContainingCompilationUnit()
+            val unwrappedStatements = body.unwrap()
+            val containingCU = unwrappedStatements.firstOrNull()?.statement?.getContainingCompilationUnit()
 
             // Recursive deal with goto
             while (throwable is GotoException) {
@@ -157,7 +201,7 @@ data class ExecuteSubroutine(var subroutine: ReferenceByName<Subroutine>, overri
                 }
 
                 // If tag is in top level we need to quit subroutine and rollback context to top level
-                val topLevelStatements = containingCU?.main?.stmts?.explode(true)
+                val topLevelStatements = containingCU?.main?.stmts?.unwrap()
                 topLevelStatements?.let {
                     val goto = throwable as GotoException
                     val topLevelIndex = goto.indexOfTaggedStatement(it)
@@ -212,7 +256,7 @@ data class SelectStmt(
     var other: SelectOtherClause? = null,
     @Derived val dataDefinition: InStatementDataDefinition? = null,
     override val position: Position? = null
-) : Statement(position), CompositeStatement, StatementThatCanDefineData {
+) : Statement(position), CompositeStatement, StatementThatCanDefineData, CustomStatementUnwrap {
     override val loggableEntityName: String
         get() = "SELECT"
 
@@ -265,6 +309,47 @@ data class SelectStmt(
             if (other?.body != null) result.addAll(other!!.body.explode(preserveCompositeStatement = true))
             return result
         }
+
+    override fun unwrap(parent: UnwrappedStatementData?): List<UnwrappedStatementData> {
+        val switchStmt = UnwrappedStatementData(this, this.body.size, parent)
+        val casesBodies = this.cases.map { it.body.unwrap(switchStmt) }
+        val otherBody = this.other?.body?.unwrap(switchStmt) ?: emptyList()
+
+        /**
+         * Unwrapped order is like following:
+         * - SELECT statement
+         * - 1 or more WHEN cases body instructions
+         * - optional OTHER with body instructions
+         * - next operation
+         */
+        val totalCasesStatementsCount = casesBodies.sumOf { it.size }
+        val nextOperationAt = totalCasesStatementsCount + otherBody.size
+        switchStmt.nextOperationOffset = nextOperationAt
+
+        /**
+         * Each branch last instruction must point to the end of the SWITCH statement
+         * after unwrapping we need to update the offset of each 'last' statement
+         */
+
+        // Update last pointer in when bodies
+        var alreadyProcessedCasesStatements = 0
+        for (case in casesBodies) {
+            val offsetOfLastStatement = alreadyProcessedCasesStatements + case.size
+            val lastStatementMustRedirectTo = nextOperationAt - offsetOfLastStatement
+            if (case.isNotEmpty()) {
+                case.last().nextOperationOffset = lastStatementMustRedirectTo
+            }
+            alreadyProcessedCasesStatements += case.size
+        }
+
+        // Update last pointer in else body
+        if (otherBody.isNotEmpty()) {
+            val lastStatementMustRedirectTo = nextOperationAt - totalCasesStatementsCount
+            otherBody.last().nextOperationOffset = lastStatementMustRedirectTo
+        }
+
+        return listOf(switchStmt) + casesBodies.flatten() + otherBody
+    }
 }
 
 @Serializable
@@ -1059,7 +1144,7 @@ data class IfStmt(
     val elseIfClauses: List<ElseIfClause> = emptyList(),
     val elseClause: ElseClause? = null,
     override val position: Position? = null
-) : Statement(position), CompositeStatement {
+) : Statement(position), CompositeStatement, CustomStatementUnwrap {
     override val loggableEntityName: String
         get() = "IF"
 
@@ -1122,6 +1207,57 @@ data class IfStmt(
         }
     }
 
+    override fun unwrap(parent: UnwrappedStatementData?): List<UnwrappedStatementData> {
+        val ifStmt = UnwrappedStatementData(this, this.body.size, parent)
+        val thenBody = this.thenBody.unwrap(ifStmt)
+        val elseIfBodies = this.elseIfClauses.map { it.body.unwrap(ifStmt) }
+        val elseBody = this.elseClause?.body?.unwrap(ifStmt) ?: emptyList()
+
+        /**
+         * Unwrapped order is like following:
+         * - IF statement
+         * - THEN body instructions
+         * - 0 or more ELSE IF with body instructions
+         * - optional ELSE with body instructions
+         * - next operation
+         */
+        val totalElseIfStatementsCount = elseIfBodies.sumOf { it.size }
+        val nextOperationAt = thenBody.size + totalElseIfStatementsCount + elseBody.size
+        ifStmt.nextOperationOffset = nextOperationAt
+
+        /**
+         * Each branch last instruction must point to the end of the IF statement
+         * after unwrapping we need to update the offset of each 'last' statement
+         */
+
+        // Update last pointer in then body
+        if (thenBody.isNotEmpty()) {
+            val offsetOfLastStatement = thenBody.size
+            val lastStatementMustRedirectTo = nextOperationAt - offsetOfLastStatement
+            thenBody.last().nextOperationOffset = lastStatementMustRedirectTo
+        }
+
+        // Update last pointer in else if bodies
+        var alreadyProcessedElifStatements = 0
+        for (elif in elseIfBodies) {
+            val offsetOfLastStatement = alreadyProcessedElifStatements + elif.size
+            val lastStatementMustRedirectTo = nextOperationAt - offsetOfLastStatement
+            if (elif.isNotEmpty()) {
+                elif.last().nextOperationOffset = lastStatementMustRedirectTo
+            }
+            alreadyProcessedElifStatements += elif.size
+        }
+
+        // Update last pointer in else body
+        if (elseBody.isNotEmpty()) {
+            val offsetOfLastStatement = thenBody.size + totalElseIfStatementsCount
+            val lastStatementMustRedirectTo = nextOperationAt - offsetOfLastStatement
+            elseBody.last().nextOperationOffset = lastStatementMustRedirectTo
+        }
+
+        return listOf(ifStmt) + thenBody + elseIfBodies.flatten() + elseBody
+    }
+
     override fun getStatementLogRenderer(source: LogSourceProvider, action: String): LazyLogEntry {
         val entry = LogEntry(source, LogChannel.STATEMENT.getPropertyName(), action)
         return LazyLogEntry(entry) {

rpgJavaInterpreter-core/src/main/kotlin/com/smeup/rpgparser/utils/misc.kt

@@ -18,8 +18,8 @@ package com.smeup.rpgparser.utils
 
 import com.smeup.rpgparser.execution.ParsingProgram
 import com.smeup.rpgparser.parsing.ast.CompilationUnit
-import com.smeup.rpgparser.parsing.ast.Statement
 import com.smeup.rpgparser.parsing.ast.TagStmt
+import com.smeup.rpgparser.parsing.ast.UnwrappedStatementData
 import com.strumenta.kolasu.model.Node
 import com.strumenta.kolasu.model.ancestor
 import java.math.BigDecimal
@@ -176,6 +176,6 @@ internal fun <T> Stack<T>.peekOrNull(): T? = if (isNotEmpty()) {
  */
 internal fun Node.getContainingCompilationUnit() = ancestor(CompilationUnit::class.java)
 
-internal fun List<Statement>.indexOfTag(tag: String) = indexOfFirst {
-    it is TagStmt && it.tag.lowercase() == tag.lowercase()
+internal fun List<UnwrappedStatementData>.indexOfTag(tag: String) = indexOfFirst {
+    it.statement is TagStmt && (it.statement as TagStmt).tag.lowercase() == tag.lowercase()
 }