fix: variable-sized values, heap-allocated arrays (and printCharArray)

src/main/wacc/backend/Stack.scala

@@ -1,37 +1,48 @@
 package wacc
 
 import scala.collection.mutable.LinkedHashMap
+import cats.data.Chain
 
 class Stack {
   import assemblyIR._
+  import assemblyIR.Size._
   import sizeExtensions.size
   import microWacc as mw
 
-  private val RSP = Register(Size.Q64, RegName.SP)
+  private val RSP = Register(Q64, RegName.SP)
   private class StackValue(val size: Size, val offset: Int) {
-    def bottom: Int = offset + size.toInt
+    def bottom: Int = offset + elemBytes
   }
   private val stack = LinkedHashMap[mw.Expr | Int, StackValue]()
 
+  private val elemBytes: Int = Q64.toInt
+  private def sizeBytes: Int = stack.size * elemBytes
+
   /** The stack's size in bytes. */
-  def size: Int = if stack.isEmpty then 0 else stack.last._2.bottom
+  def size: Int = stack.size
 
   /** Push an expression onto the stack. */
   def push(expr: mw.Expr, src: Register): AsmLine = {
-    stack += expr -> StackValue(src.size, size)
+    stack += expr -> StackValue(src.size, sizeBytes)
     Push(src)
   }
 
-  /** Push an arbitrary register onto the stack. */
-  def push(src: Register): AsmLine = {
-    stack += stack.size -> StackValue(src.size, size)
-    Push(src)
+  /** Push a value onto the stack. */
+  def push(itemSize: Size, addr: Src): AsmLine = {
+    stack += stack.size -> StackValue(itemSize, sizeBytes)
+    Push(addr)
   }
 
   /** Reserve space for a variable on the stack. */
   def reserve(ident: mw.Ident): AsmLine = {
-    stack += ident -> StackValue(ident.ty.size, size)
-    Subtract(RSP, ImmediateVal(ident.ty.size.toInt))
+    stack += ident -> StackValue(ident.ty.size, sizeBytes)
+    Subtract(RSP, ImmediateVal(elemBytes))
+  }
+
+  /** Reserve space for a register on the stack. */
+  def reserve(src: Register): AsmLine = {
+    stack += stack.size -> StackValue(src.size, sizeBytes)
+    Subtract(RSP, ImmediateVal(src.size.toInt))
   }
 
   /** Reserve space for values on the stack.
@@ -40,45 +51,40 @@ class Stack {
     *   The sizes of the values to reserve space for.
     */
   def reserve(sizes: List[Size]): AsmLine = {
-    val totalSize = sizes
-      .map(itemSize =>
-        stack += stack.size -> StackValue(itemSize, size)
-        itemSize.toInt
-      )
-      .sum
-    Subtract(RSP, ImmediateVal(totalSize))
+    sizes.foreach { itemSize =>
+      stack += stack.size -> StackValue(itemSize, sizeBytes)
+    }
+    Subtract(RSP, ImmediateVal(elemBytes * sizes.size))
   }
 
   /** Pop a value from the stack into a register. Sizes MUST match. */
   def pop(dest: Register): AsmLine = {
-    if (dest.size != stack.last._2.size) {
-      throw new IllegalArgumentException(
-        s"Cannot pop ${stack.last._2.size} bytes into $dest (${dest.size} bytes) register"
-      )
-    }
     stack.remove(stack.last._1)
     Pop(dest)
   }
 
   /** Drop the top n values from the stack. */
   def drop(n: Int = 1): AsmLine = {
-    val totalSize = (1 to n)
-      .map(_ =>
-        val itemSize = stack.last._2.size.toInt
-        stack.remove(stack.last._1)
-        itemSize
-      )
-      .sum
-    Add(RSP, ImmediateVal(totalSize))
+    (1 to n).foreach { _ =>
+      stack.remove(stack.last._1)
+    }
+    Add(RSP, ImmediateVal(n * elemBytes))
   }
 
-  /** Get a lazy IndexAddress for a variable in the stack. */
-  def accessVar(ident: mw.Ident): () => IndexAddress = () => {
-    IndexAddress(RSP, stack.size - stack(ident).bottom)
+  /** Generate AsmLines within a scope, which is reset after the block. */
+  def withScope(block: () => Chain[AsmLine]): Chain[AsmLine] = {
+    val resetToSize = stack.size
+    var lines = block()
+    lines :+= drop(stack.size - resetToSize)
+    lines
   }
+
+  /** Get an IndexAddress for a variable in the stack. */
+  def accessVar(ident: mw.Ident): IndexAddress =
+    IndexAddress(RSP, sizeBytes - stack(ident).bottom)
+
   def contains(ident: mw.Ident): Boolean = stack.contains(ident)
-  def head: MemLocation = MemLocation(RSP)
-  def head(offset: Size): MemLocation = MemLocation(RSP, Some(offset))
-  // TODO: Might want to actually properly handle this with the LinkedHashMap too
-  def align(): AsmLine = And(RSP, ImmediateVal(-16))
+  def head: MemLocation = MemLocation(RSP, stack.last._2.size)
+
+  override def toString(): String = stack.toString
 }