aftercall rename

src/main/scala/ir/IRCursor.scala

@@ -72,7 +72,7 @@ trait IntraProcIRCursor extends IRWalk[CFGPosition, CFGPosition] {
       case j: Jump => Set(j.parent.statements.lastOption.getOrElse(j.parent))
       case b: Block =>
         b.kind match {
-          case CallReturn(from) => Set(from)
+          case AfterCall(from) => Set(from)
           case Entry(proc)      => Set(proc)
           case _                => b.incomingJumps.asInstanceOf[Set[CFGPosition]]
         }
@@ -100,7 +100,7 @@ trait IntraProcBlockIRCursor extends IRWalk[CFGPosition, Block] {
       case b: Block =>
         b.kind match {
           case Entry(_)         => Set.empty
-          case CallReturn(from) => Set(from.parent)
+          case AfterCall(from) => Set(from.parent)
           case _                => b.incomingJumps.map(_.parent)
         }
       case j: Command   => pred(j.parent)
@@ -139,7 +139,7 @@ trait InterProcIRCursor extends IRWalk[CFGPosition, CFGPosition] {
       case c: Procedure       => c.incomingCalls().toSet.asInstanceOf[Set[CFGPosition]]
       case b: Block =>
         b.kind match {
-          case CallReturn(DirectCall(target, returnTarget, _)) => target.returnBlock.map(_.jump).toSet
+          case AfterCall(DirectCall(target, returnTarget, _)) => target.returnBlock.map(_.jump).toSet
           case _ => Set()
         }
       case _ => IntraProcIRCursor.pred(pos)
@@ -160,7 +160,7 @@ trait InterProcBlockIRCursor extends IRWalk[CFGPosition, Block] {
     pos match {
       case b: Block =>
         b.kind match {
-          case CallReturn(from) => from.parent.parent.returnBlock.toSet
+          case AfterCall(from) => from.parent.parent.returnBlock.toSet
           case Entry(of)        => of.incomingCalls().map(_.parent).toSet
           case _                => b.prevBlocks.toSet
         }

src/main/scala/ir/Parent.scala

@@ -1,4 +1,5 @@
-package ir 
+package ir
+import util.Logger
 
 trait HasParent[T]:
   /*
@@ -8,7 +9,14 @@ trait HasParent[T]:
       All IL structures must set the parent of the child to itself, when a child is added to itself.
    */
   private var _parent: Option[T] = None
-  def parent: T = _parent.get
+  private var last_parent: Option[T] = None
+  def parent: T = {
+    if (!hasParent) {
+      Logger.error(s"Trying to get the parent of a node that is detached from the progam: $this. Node was last attached to: $last_parent")
+    }
+    _parent.get
+  }
+
   def hasParent: Boolean = _parent.isDefined
 
   def parent_=(value: T): Unit = setParent(value)
@@ -39,6 +47,7 @@ trait HasParent[T]:
    */
   final def deParent(): Unit = if _parent.isDefined then {
     unlinkParent()
+    if _parent.isDefined then (last_parent = _parent)
     _parent = None
     }
 

src/main/scala/ir/Program.scala

@@ -200,18 +200,23 @@ class Procedure private (
   def returnBlock: Option[Block] = _returnBlock
 
   def returnBlock_=(value: Block): Unit = {
-    removeBlocks(_returnBlock)
-    _returnBlock = Some(addBlocks(value))
+    if (!returnBlock.contains(value)) {
+      removeBlocks(_returnBlock)
+      _returnBlock = Some(addBlocks(value))
+    }
   }
 
   def entryBlock: Option[Block] = _entryBlock
 
   def entryBlock_=(value: Block): Unit = {
-    removeBlocks(_entryBlock)
-    _entryBlock = Some(addBlocks(value))
+    if (!entryBlock.contains(value)) {
+      removeBlocks(_entryBlock)
+      _entryBlock = Some(addBlocks(value))
+    }
   }
 
   def addBlocks(block: Block): Block = {
+    block.parent = this
     if (!_blocks.contains(block)) {
       block.parent = this
       _blocks.add(block)
@@ -296,7 +301,7 @@ sealed trait BlockKind
 case class Regular() extends BlockKind
 
 /* Block is the fallthrough / return target of a call. */
-case class CallReturn(from: Call) extends BlockKind
+case class AfterCall(from: Call) extends BlockKind
 
 /* Block is the single return point for a procedure */
 case class Return(from: Procedure) extends BlockKind
@@ -327,13 +332,24 @@ class Block private (
   def jump_=(j: Jump): Unit = {
     if (j ne _jump) {
       _jump.deParent()
+      //_jump match {
+      //  case c: IndirectCall => c.returnTarget.filter(b=>b.kind.isInstanceOf[AfterCall]).map(b => parent.removeBlocks(b))
+      //  case c: DirectCall => c.returnTarget.filter(b=>b.kind.isInstanceOf[AfterCall]).map(b => parent.removeBlocks(b))
+      //  case _ => ()
+      //}
       _jump = j
       _jump.parent = this
+      //j match {
+      //  case c: IndirectCall => c.returnTarget = Some(parent.addBlocks(Block.afterCall(c)))
+      //  case c: DirectCall => c.returnTarget = Some(parent.addBlocks(Block.afterCall(c)))
+      //  case _ => ()
+      //}
     }
   }
 
   def replaceJump(j: Jump) = {
     jump = j
+    assert(jump.parent eq this)
     this
   }
 
@@ -439,12 +455,15 @@ object Block:
     new Block(Regular(), label, address, IntrusiveList.empty, GoTo(Seq(), Some(label + "_unknown")))
   }
 
-  def callReturn(from: Call) : Block = {
-    val jump = from match 
+  def afterCall(from: Call) : Block = {
+    val jump = from match
       case d: DirectCall => GoTo(d.returnTarget.toSet)
       case c: IndirectCall => GoTo(c.returnTarget.toSet)
-
-    new Block(CallReturn(from), from.parent.label + "_basil_callreturn", None, Seq(), jump)
+
+    jump.parent = from.parent
+    val ac = Block(AfterCall(from), from.parent.label + "_basil_aftercall", None, Seq(), jump)
+    ac.parent = from.parent.parent
+    ac
   }
 
   def procedureReturn(from: Procedure): Block = {

src/main/scala/ir/Visitor.scala

@@ -428,15 +428,15 @@ class AddCallReturnBlocks extends Visitor {
    * Add a dummy call return block with no statements after each call. 
    */
   override def visitDirectCall(node: DirectCall): Jump = {
-    val b = node.parent.parent.addBlocks(Block.callReturn(node))
+    val b = node.parent.parent.addBlocks(Block.afterCall(node))
     node.returnTarget = Some(b)
     node
   }
 
   override def visitIndirectCall(node: IndirectCall): Jump = {
     // skip return nodes
     if !(node.parent.kind.isInstanceOf[Return]) then { 
-      val b = node.parent.parent.addBlocks(Block.callReturn(node))
+      val b = node.parent.parent.addBlocks(Block.afterCall(node))
       node.returnTarget = Some(b)
     } 
     node

src/main/scala/translating/BAPToIR.scala

@@ -48,6 +48,7 @@ class BAPToIR(var program: BAPProgram, mainAddress: Int) {
         val (jump, newBlocks) = translate(b.jumps, block)
         procedure.addBlocks(newBlocks)
         block.replaceJump(jump)
+        assert(jump.hasParent)
       }
 
       // Set entry block to the block with the same address as the procedure or the first in sequence

src/main/scala/util/RunUtils.scala

@@ -96,10 +96,10 @@ object RunUtils {
     val renamer = Renamer(reserved)
     val returnUnifier = ConvertToSingleProcedureReturn()
     val callReturner = AddCallReturnBlocks()
+    IRProgram = callReturner.visitProgram(IRProgram)
     IRProgram = externalRemover.visitProgram(IRProgram)
     IRProgram = renamer.visitProgram(IRProgram)
     IRProgram = returnUnifier.visitProgram(IRProgram)
-    IRProgram = callReturner.visitProgram(IRProgram)
 
 
     q.loading.dumpIL.foreach(s => writeToFile(serialiseIL(IRProgram), s"$s-before-analysis.il"))
@@ -157,12 +157,21 @@ object RunUtils {
     Logger.info(externalAddresses)
     Logger.info("Subroutine Addresses:")
     Logger.info(subroutines)
+    var unparented = IRProgram.collect {
+      case c: Command if !c.hasParent => c
+      case b: Block if !b.hasParent => b
+    }
+    assert(unparented.isEmpty)
 
     val mergedSubroutines = subroutines ++ externalAddresses
 
 
     val cfg = ProgramCfgFactory().fromIR(IRProgram)
 
+    unparented = IRProgram.collect {
+      case c: IndirectCall if (!c.hasParent) => c
+    }
+    assert(unparented.isEmpty)
     val domain = computeDomain(IntraProcIRCursor, IRProgram.procedures)
 
     Logger.info("[!] Running Constant Propagation")
@@ -214,6 +223,13 @@ object RunUtils {
 
     Logger.info("[!] Resolving CFG")
     val (newIR, modified): (Program, Boolean) = resolveCFG(cfg, vsaResult, IRProgram)
+
+    unparented = IRProgram.collect {
+      case c: Command if !c.hasParent => c
+      case b: Block if !b.hasParent => b
+    }
+    assert(unparented.isEmpty)
+
     if (modified) {
       Logger.info(s"[!] Analysing again (iter $iteration)")
       return analyse(newIR, externalFunctions, globals, globalOffsets, config, iteration + 1)