add procedure call inlining (#345)

Author: ailrst

src/main/scala/ir/Program.scala

@@ -610,6 +610,14 @@ class Block private (
     jump.deParent()
   }
 
+  def createBlockAfter(suffix: String): Block = {
+    val nb = Block(label + suffix)
+    parent.addBlock(nb)
+    val ojump = jump
+    replaceJump(GoTo(nb))
+    nb.replaceJump(ojump)
+  }
+
   def createBlockBetween(b2: Block, label: String = "_goto_"): Block = {
     require(nextBlocks.toSet.contains(b2))
     val b1 = this

src/main/scala/ir/cilvisitor/CILVisitor.scala

@@ -217,3 +217,4 @@ def visit_prog(v: CILVisitor, b: Program): Program = CILVisitorImpl(v).visit_pro
 def visit_stmt(v: CILVisitor, e: Statement): List[Statement] = CILVisitorImpl(v).visit_stmt(e)
 def visit_jump(v: CILVisitor, e: Jump): Jump = CILVisitorImpl(v).visit_jump(e)
 def visit_expr(v: CILVisitor, e: Expr): Expr = CILVisitorImpl(v).visit_expr(e)
+def visit_rvar(v: CILVisitor, e: Variable): Variable = CILVisitorImpl(v).visit_rvar(e)

src/main/scala/ir/transforms/Inline.scala

@@ -0,0 +1,163 @@
+package ir.transforms
+import ir.*
+import util.Logger
+import ir.cilvisitor.*
+import ir.dsl.IRToDSL.*
+import ir.dsl.*
+
+import scala.collection.mutable
+
+object Counter {
+  var id = 0
+  def next() = {
+    id += 1
+    id
+  }
+}
+
+def memoise[K, V](f: K => V): K => V = {
+  val rn = mutable.Map[K, V]()
+
+  def fun(arg: K): V = {
+    if (!rn.contains(arg)) {
+      rn(arg) = f(arg)
+    }
+
+    rn(arg)
+  }
+  fun
+}
+
+def renameBlock(s: String): String = {
+  s + "_" + (Counter.next())
+}
+
+class VarRenamer(proc: Procedure) extends CILVisitor {
+
+  def doRename(v: Variable): Variable = v match {
+    case l: LocalVar if l.name.endsWith("_in") => {
+      val name = l.name.stripSuffix("_in")
+      proc.getFreshSSAVar(name, l.getType)
+    }
+    case l: LocalVar if l.index != 0 =>
+      proc.getFreshSSAVar(l.varName, l.getType)
+    case _ => v
+  }
+  val memoed = memoise(doRename)
+
+  def rename(v: Variable) = {
+    memoed(v)
+  }
+
+  override def vlvar(v: Variable) = {
+    ChangeTo(rename(v))
+  }
+
+  override def vrvar(v: Variable) = {
+    ChangeTo(rename(v))
+  }
+
+}
+
+def convertJumpRenaming(blockName: String => String, varName: CILVisitor, x: Jump): EventuallyJump = x match {
+  case GoTo(targs, label) => EventuallyGoto(targs.map(t => DelayNameResolve(blockName(t.label))).toList, label)
+  case Unreachable(label) => EventuallyUnreachable(label)
+  case Return(label, out) =>
+    EventuallyReturn(
+      out.toList.map { case (v: Variable, e: Expr) =>
+        ((v.name, visit_expr(varName, e)))
+      }.toArray,
+      label
+    )
+}
+
+def keyToString[T](varRenamer: CILVisitor)(x: (Variable, Expr)): (String, Expr) =
+  (x(0).name, visit_expr(varRenamer, x(1)))
+
+def convertStatementRenaming(varRenamer: CILVisitor)(x: Statement): EventuallyStatement = x match {
+  case DirectCall(targ, outs, actuals, label) =>
+    directCall(
+      outs.toArray.map(v => (v(0).name, visit_rvar(varRenamer, v(1)))),
+      targ.name,
+      actuals.toArray.map(keyToString(varRenamer)): _*
+    )
+  case IndirectCall(targ, label) => indirectCall(visit_rvar(varRenamer, targ))
+  case x: NonCallStatement =>
+    CloneableStatement(visit_stmt(varRenamer, cloneStatement(x)).head.asInstanceOf[NonCallStatement])
+}
+
+def convertBlockRenaming(varRenamer: CILVisitor, blockName: String => String)(x: Block) = {
+  EventuallyBlock(
+    blockName(x.label),
+    x.statements.toArray.map(convertStatementRenaming(varRenamer)),
+    convertJumpRenaming(blockName, varRenamer, x.jump),
+    x.address
+  )
+}
+
+/**
+  * Inline a procedure call to the calling procedure;
+  *
+  * - maintains param form 
+  * - maintain dsa form if both procedures are in dsa form
+  *
+  * require invariant.SingleCallBlockEnd
+  */
+def inlineCall(prog: Program, c: DirectCall): Unit = {
+  require(c.target.entryBlock.isDefined)
+  require(c.target.returnBlock.isDefined)
+
+  val proc = c.parent.parent
+  val block = c.parent
+  val target = c.target
+  // rename ssa variables to maintain DSA in the new procedure
+  // rename block labels to avoid creating conflicts if inlining the same function multiple times
+  val varRenamer = VarRenamer(proc)
+  val blockRenamer: String => String = memoise(renameBlock)
+
+  val entry = target.entryBlock.get
+  val returnBlock = target.returnBlock.get
+
+  // the internal blocks to the target
+  val internalBlocks =
+    (target.blocks.toSet -- Seq(entry, returnBlock)).map(convertBlockRenaming(varRenamer, blockRenamer))
+
+  // clone the target procedure and apply the renaming
+  val (entryTempBlock, entryResolver) = convertBlockRenaming(varRenamer, blockRenamer)(entry).makeResolver
+  val eventuallyReturnBlock = convertBlockRenaming(varRenamer, blockRenamer)(returnBlock)
+  val afterCallBlock = block.createBlockAfter("_inlineret")
+  proc.addBlock(entryTempBlock)
+  block.replaceJump(GoTo(entryTempBlock))
+
+  // replace return in target with a jump to the aftercall block
+  val (returnTemp, resolveReturnBlock) = eventuallyReturnBlock.copy(j = unreachable).makeResolver
+  proc.addBlock(returnTemp)
+
+  // resolve internal call blocks
+  val resolvers = internalBlocks.map(_.makeResolver)
+  resolvers.foreach { case (block, _) => proc.addBlock(block) }
+  resolvers.foreach { case (_, resolve) => resolve(prog, proc) }
+
+  // remove original call statement
+  block.statements.remove(c)
+  // link the inlined blocks to the call block and the aftercall block
+  entryResolver(prog, proc)
+  block.replaceJump(GoTo(entryTempBlock))
+  resolveReturnBlock(prog, proc)
+  returnTemp.replaceJump(GoTo(afterCallBlock))
+
+  // assign the actual parameters in the caller to the renamed formal parameters in the entry block
+  val targetReturnValues: Map[String, Expr] = eventuallyReturnBlock.j match {
+    case r: EventuallyReturn => r.params.toMap
+    case _ => throw Exception("returnblock should have a return statement")
+  }
+  val outAssignments = c.outParams.map { case (formal: LocalVar, lvar: Variable) =>
+    LocalAssign(lvar, targetReturnValues(formal.name))
+  }
+  val inAssignments = c.actualParams.map { case (formal: LocalVar, actual: Expr) =>
+    LocalAssign(visit_rvar(varRenamer, formal), actual)
+  }
+  afterCallBlock.statements.prependAll(outAssignments)
+  entryTempBlock.statements.prependAll(inAssignments)
+
+}

src/main/scala/ir/transforms/InlinePLTLaunchpad.scala

@@ -0,0 +1,23 @@
+package ir.transforms
+import ir.Program
+
+def inlinePLTLaunchpad(prog: Program) = {
+  for (p <- prog.procedures) {
+
+    val candidate =
+      (p.blocks.size <= 4)
+        && p.calls.size == 1
+        && p.calls.forall(_.isExternal.contains(true))
+        && p.procName.startsWith("FUN")
+        && !p.calls.contains(p)
+
+    if (candidate) {
+      p.incomingCalls().foreach { call =>
+        inlineCall(prog, call)
+      }
+    }
+  }
+
+  applyRPO(prog)
+
+}

src/main/scala/util/RunUtils.scala

@@ -725,9 +725,10 @@ object RunUtils {
     Logger.info("[!] Simplify :: DynamicSingleAssignment")
     DebugDumpIRLogger.writeToFile(File("il-before-dsa.il"), pp_prog(program))
 
-    // transforms.DynamicSingleAssignment.applyTransform(program, liveVars)
     transforms.OnePassDSA().applyTransform(program)
 
+    transforms.inlinePLTLaunchpad(ctx.program)
+
     transforms.removeEmptyBlocks(program)
 
     AnalysisResultDotLogger.writeToFile(