From cdfbbf61b841b8562329a9f87f4b3419a0ef0b27 Mon Sep 17 00:00:00 2001
From: Karol Stasiak <karol.m.stasiak@gmail.com>
Date: Sun, 15 Apr 2018 01:29:47 +0200
Subject: [PATCH] Fixed allocating allocating local variables twice

---
 src/main/scala/millfork/env/Environment.scala | 11 ++++++++++-
 .../millfork/output/VariableAllocator.scala   | 19 ++++++++++---------
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/src/main/scala/millfork/env/Environment.scala b/src/main/scala/millfork/env/Environment.scala
index 2a604c31..c8fdfac3 100644
--- a/src/main/scala/millfork/env/Environment.scala
+++ b/src/main/scala/millfork/env/Environment.scala
@@ -79,6 +79,15 @@ class Environment(val parent: Option[Environment], val prefix: String) {
     }.toList
   }
 
+  private def isLocalVariableName(name: String): Boolean = name.contains('$') && !name.endsWith("$") && {
+    val fname = name.substring(0, name.indexOf('$'))
+    things.get(fname) -> parent.flatMap(_.things.get(fname)) match {
+      case (Some(_: NormalFunction), _) => true
+      case (_, Some(_: NormalFunction)) => true
+      case _ => false
+    }
+  }
+
   def allocateVariables(nf: Option[NormalFunction], mem: CompiledMemory, callGraph: CallGraph, allocators: Map[String, VariableAllocator], options: CompilationOptions, onEachVariable: (String, Int) => Unit): Unit = {
     val b = get[Type]("byte")
     val p = get[Type]("pointer")
@@ -91,7 +100,7 @@ class Environment(val parent: Option[Environment], val prefix: String) {
       }
     }.toSet
     val toAdd = things.values.flatMap {
-      case m: UninitializedMemory =>
+      case m: UninitializedMemory if nf.isDefined == isLocalVariableName(m.name) =>
         val vertex = if (options.flag(CompilationFlag.VariableOverlap)) {
           nf.fold[VariableVertex](GlobalVertex) { f =>
             if (m.alloc == VariableAllocationMethod.Static) {
diff --git a/src/main/scala/millfork/output/VariableAllocator.scala b/src/main/scala/millfork/output/VariableAllocator.scala
index 60393adf..11d87f50 100644
--- a/src/main/scala/millfork/output/VariableAllocator.scala
+++ b/src/main/scala/millfork/output/VariableAllocator.scala
@@ -4,6 +4,7 @@ import millfork.error.ErrorReporting
 import millfork.node.{CallGraph, VariableVertex}
 import millfork.{CompilationFlag, CompilationOptions}
 
+import scala.annotation.tailrec
 import scala.collection.mutable
 
 /**
@@ -49,10 +50,10 @@ class AfterCodeByteAllocator(val endBefore: Int) extends ByteAllocator {
   def notifyAboutEndOfCode(org: Int): Unit = startAt = org
 }
 
-class VariableAllocator(private var pointers: List[Int], private val bytes: ByteAllocator) {
+class VariableAllocator(private val pointers: List[Int], private val bytes: ByteAllocator) {
 
-  private var pointerMap = mutable.Map[Int, Set[VariableVertex]]()
-  private var variableMap = mutable.Map[Int, mutable.Map[Int, Set[VariableVertex]]]()
+  private val pointerMap = mutable.Map[Int, Set[VariableVertex]]()
+  private val variableMap = mutable.Map[Int, mutable.Map[Int, Set[VariableVertex]]]()
 
   def allocatePointer(mem: MemoryBank, callGraph: CallGraph, p: VariableVertex): Int = {
     // TODO: search for free zeropage locations
@@ -62,16 +63,16 @@ class VariableAllocator(private var pointers: List[Int], private val bytes: Byte
         return addr
       }
     }
-    def pickFreePointer(): Int =
-      pointers match {
+    @tailrec
+    def pickFreePointer(ps: List[Int]): Int =
+      ps match {
         case Nil =>
+          ErrorReporting.trace(pointerMap.mkString(", "))
           ErrorReporting.fatal("Out of zero-page memory")
         case next :: rest =>
           if (mem.occupied(next) || mem.occupied(next + 1)) {
-            pointers = rest
-            pickFreePointer()
+            pickFreePointer(rest)
           } else {
-            pointers = rest
             mem.readable(next) = true
             mem.readable(next + 1) = true
             mem.occupied(next) = true
@@ -82,7 +83,7 @@ class VariableAllocator(private var pointers: List[Int], private val bytes: Byte
             next
           }
       }
-    pickFreePointer()
+    pickFreePointer(pointers)
   }
 
   def allocateBytes(mem: MemoryBank, callGraph: CallGraph, p: VariableVertex, options: CompilationOptions, count: Int, initialized: Boolean, writeable: Boolean): Int = {