From dc64d3f3f9e60960095f81d5fc1981fa5d54c9ed Mon Sep 17 00:00:00 2001
From: l-kent <l.kent@uq.edu.au>
Date: Fri, 13 Oct 2023 14:57:07 +1000
Subject: [PATCH] fix #100

---
 src/main/scala/ir/Program.scala |  8 ++++++--
 src/main/scala/ir/Visitor.scala | 24 ++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/main/scala/ir/Program.scala b/src/main/scala/ir/Program.scala
index 4985188d3..e45005ef1 100644
--- 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
index f2e7133e0..bd2a92647 100644
--- 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
+  }
+
+}
\ No newline at end of file