fix #100

UQ-PAC · Oct 13, 2023 · dc64d3f · dc64d3f
1 parent 96bc93f
commit dc64d3f
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 2 deletions.
diff --git a/src/main/scala/ir/Program.scala b/src/main/scala/ir/Program.scala
@@ -140,7 +140,10 @@ class Procedure(
             }
 
             // update stack references
-            val rhsStackRefs = l.rhs.variables.intersect(stackRefs)
+            val variableVisitor = VariablesWithoutStoresLoads()
+            variableVisitor.visitExpr(l.rhs)
+
+            val rhsStackRefs = variableVisitor.variables.toSet.intersect(stackRefs)
             if (rhsStackRefs.nonEmpty) {
               stackRefs.add(l.lhs)
             } else if (stackRefs.contains(l.lhs) && l.lhs != stackPointer) {
@@ -160,7 +163,8 @@ class Procedure(
       for (j <- b.jumps) {
         j match {
           case g: GoTo => visitBlock(g.target)
-          case _       =>
+          case d: DirectCall => d.returnTarget.foreach(visitBlock)
+          case i: IndirectCall => i.returnTarget.foreach(visitBlock)
         }
       }
     }

diff --git a/src/main/scala/ir/Visitor.scala b/src/main/scala/ir/Visitor.scala
@@ -312,3 +312,27 @@ class ExternalRemover(external: Set[String]) extends Visitor {
     super.visitProcedure(node)
   }
 }
+
+/** Gives variables that are not contained within a MemoryStore or MemoryLoad
+  * */
+class VariablesWithoutStoresLoads extends ReadOnlyVisitor {
+  val variables: mutable.Set[Variable] = mutable.Set()
+
+  override def visitRegister(node: Register): Register = {
+    variables.add(node)
+    node
+  }
+  override def visitLocalVar(node: LocalVar): LocalVar = {
+    variables.add(node)
+    node
+  }
+
+  override def visitMemoryStore(node: MemoryStore): MemoryStore = {
+    node
+  }
+
+  override def visitMemoryLoad(node: MemoryLoad): MemoryLoad = {
+    node
+  }
+
+}