feat: maintain program counter and assert at beginning of blocks

this makes a fairly significant change to the GTIRB loader by inserting
PC increment statements after instructions and checking the PC at the
entry of blocks. this is done by changing the loader to maintain PC and
branchtaken assignments, and inserting a default increment if there is
no branch.

this surely has bugs

Author: katrinafyi

src/main/scala/translating/GTIRBLoader.scala

@@ -26,14 +26,18 @@ class GTIRBLoader(parserMap: immutable.Map[String, List[InsnSemantics]]) {
   private var blockCount = 0
   private var loadCounter = 0
 
-  private val opcodeSize = 4
+  val opcodeSize = 4
 
-  def visitBlock(blockUUID: ByteString, blockCountIn: Int, blockAddress: Option[BigInt]): ArrayBuffer[Statement] = {
+  def visitBlock(
+    blockUUID: ByteString,
+    blockCountIn: Int,
+    blockAddress: Option[BigInt]
+  ): ArrayBuffer[immutable.Seq[Statement]] = {
     blockCount = blockCountIn
     instructionCount = 0
     val instructions = parserMap(Base64.getEncoder.encodeToString(blockUUID.toByteArray))
 
-    val statements: ArrayBuffer[Statement] = ArrayBuffer()
+    val statements: ArrayBuffer[immutable.Seq[Statement]] = ArrayBuffer()
 
     for (instsem <- instructions) {
       constMap.clear
@@ -43,17 +47,19 @@ class GTIRBLoader(parserMap: immutable.Map[String, List[InsnSemantics]]) {
         case InsnSemantics.Error(op, err) => {
           val message = s"$op ${err.replace("\n", " :: ")}"
           Logger.warn(s"Program contains lifter unsupported opcode: $message")
-          statements.append(Assert(FalseLiteral, Some(s"Lifter error: $message")))
+          statements.append(immutable.Seq(Assert(FalseLiteral, Some(s"Lifter error: $message"))))
           instructionCount += 1
         }
-        case InsnSemantics.Result(instruction) => {
-          for ((s, i) <- instruction.zipWithIndex) {
+        case InsnSemantics.Result(aslstmts) => {
+          var stmts = immutable.LinearSeq[Statement]()
+
+          for ((s, i) <- aslstmts.zipWithIndex) {
             val label = blockAddress.map { (a: BigInt) =>
               val instructionAddress = a + (opcodeSize * instructionCount)
               instructionAddress.toString + "_" + i
             }
 
-            statements.appendAll(try {
+            stmts = stmts ++ (try {
               visitStmt(s, label)
             } catch {
               case e => {
@@ -62,6 +68,8 @@ class GTIRBLoader(parserMap: immutable.Map[String, List[InsnSemantics]]) {
               }
             })
           }
+
+          statements.append(stmts)
           instructionCount += 1
         }
       }
@@ -176,7 +184,7 @@ class GTIRBLoader(parserMap: immutable.Map[String, List[InsnSemantics]]) {
     }
 
     if (condition.isDefined) {
-      Some(TempIf(condition.get, thenStmts, elseStmts, label))
+      Some(TempIf(condition.get, thenStmts.to(immutable.LinearSeq), elseStmts.to(immutable.LinearSeq), label))
     } else {
       None
     }
@@ -298,7 +306,7 @@ class GTIRBLoader(parserMap: immutable.Map[String, List[InsnSemantics]]) {
       case "FALSE" => Some(FalseLiteral)
       case "FPCR" => Some(Register("FPCR", 32))
       // ignore the following
-      case "__BranchTaken" => None
+      case "__BranchTaken" => Some(LocalVar("__BranchTaken", BoolType))
       case "BTypeNext" => None
       case "BTypeCompatible" => None
       case "TPIDR_EL0" => Some(Register(name, 64))
@@ -682,7 +690,7 @@ class GTIRBLoader(parserMap: immutable.Map[String, List[InsnSemantics]]) {
       // ignore the following
       case "TRUE" => throw Exception(s"Boolean literal $name in LExpr ${ctx.getText}")
       case "FALSE" => throw Exception(s"Boolean literal $name in LExpr ${ctx.getText}")
-      case "__BranchTaken" => None
+      case "__BranchTaken" => Some(LocalVar("__BranchTaken", BoolType))
       case "BTypeNext" => None
       case "BTypeCompatible" => None
       case "TPIDR_EL0" => Some(Register(name, 64))

src/main/scala/translating/GTIRBToIR.scala

@@ -17,6 +17,7 @@ import scala.collection.mutable.Map
 import scala.collection.mutable.ArrayBuffer
 import scala.collection.immutable
 import scala.jdk.CollectionConverters.*
+import scala.util.{Try, Success, Failure}
 import java.util.Base64
 import java.nio.charset.*
 import scala.util.boundary
@@ -25,6 +26,13 @@ import java.nio.ByteBuffer
 import util.intrusive_list.*
 import util.Logger
 
+private def assigned(x: Statement): immutable.Set[Variable] = x match {
+  case x: Assign => x.assignees
+  case x: TempIf =>
+    x.cond.variables ++ x.thenStmts.flatMap(assigned) ++ x.elseStmts.flatMap(assigned)
+  case _ => immutable.Set.empty
+}
+
 /** TempIf class, used to temporarily store information about Jumps so that multiple parse runs are not needed.
   * Specifically, this is useful in the case that the IF statment has multiple conditions( and elses) and as such many
   * extra blocks need to be created.
@@ -38,8 +46,8 @@ import util.Logger
   */
 class TempIf(
   val cond: Expr,
-  val thenStmts: mutable.Buffer[Statement],
-  val elseStmts: mutable.Buffer[Statement],
+  val thenStmts: immutable.Seq[Statement],
+  val elseStmts: immutable.Seq[Statement],
   override val label: Option[String] = None
 ) extends NOP(label)
 
@@ -177,7 +185,9 @@ class GTIRBToIR(
 
         val statements = semanticsLoader.visitBlock(blockUUID, blockCount, block.address)
         blockCount += 1
-        block.statements.addAll(statements)
+        for ((stmts, i) <- statements.zipWithIndex) {
+          block.statements.addAll(insertPCIncrement(stmts))
+        }
 
         if (block.statements.isEmpty && !blockOutgoingEdges.contains(blockUUID)) {
           // remove blocks that are just nop padding
@@ -236,11 +246,33 @@ class GTIRBToIR(
     Program(procedures, intialProc, initialMemory)
   }
 
+  private def insertPCIncrement(isnStmts: immutable.Seq[Statement]): immutable.Seq[Statement] = {
+    Try(isnStmts.last) match {
+      case Success(x: TempIf) =>
+        isnStmts.init :+ TempIf(x.cond, insertPCIncrement(x.thenStmts), insertPCIncrement(x.elseStmts))
+      case Success(_) | Failure(_) => {
+        val branchTaken = isnStmts.exists {
+          case LocalAssign(LocalVar("__BranchTaken", BoolType, _), TrueLiteral, _) => true
+          case _: TempIf => throw Exception("encountered TempIf not at end of statement list: " + isnStmts)
+          case _ => false
+        }
+        val increment =
+          if branchTaken then Seq()
+          else
+            Seq(
+              LocalAssign(Register("_PC", 64), BinaryExpr(BVADD, Register("_PC", 64), BitVecLiteral(4, 64)), None),
+              LocalAssign(LocalVar("__BranchTaken", BoolType), FalseLiteral)
+            )
+        increment ++: isnStmts
+      }
+    }
+  }
+
   private def removePCAssign(block: Block): Option[String] = {
     block.statements.last match {
       case last @ LocalAssign(lhs: Register, _, _) if lhs.name == "_PC" =>
         val label = last.label
-        block.statements.remove(last)
+        // block.statements.remove(last)
         label
       case _ => throw Exception(s"expected block ${block.label} to have a program counter assignment at its end")
     }
@@ -309,9 +341,20 @@ class GTIRBToIR(
       throw Exception(s"block ${byteStringToString(blockUUID)} is in multiple functions")
     }
     uuidToBlock += (blockUUID -> block)
-    if (blockUUID == entranceUUID) {
+    val isEntrance = blockUUID == entranceUUID
+    if (isEntrance) {
       procedure.entryBlock = block
     }
+    val checkPCStmt: Expr => Statement =
+      if isEntrance
+      then Assume(_, None, None, false)
+      else Assert(_, None, None)
+    block.address match {
+      case Some(addr) =>
+        val assertPC = checkPCStmt(BinaryExpr(BVEQ, Register("_PC", 64), BitVecLiteral(addr, 64)))
+        block.statements.append(assertPC)
+      case _ => ()
+    }
     block
   }
 
@@ -781,7 +824,7 @@ class GTIRBToIR(
 
     val newBlocks = ArrayBuffer(trueBlock, falseBlock)
     procedure.addBlocks(newBlocks)
-    block.statements.remove(tempIf)
+    // block.statements.remove(tempIf)
 
     GoTo(newBlocks)
   }