package wacc

import org.scalatest.flatspec.AnyFlatSpec
import cats.data.Chain

class StackSpec extends AnyFlatSpec {
  import microWacc._
  import assemblyIR._
  import assemblyIR.Size._
  import assemblyIR.commonRegisters._
  import types.{KnownType, ?}
  import sizeExtensions.size

  private val RSP = Register(Q64, RegName.SP)

  "size" should "be 0 initially" in {
    val stack = new Stack
    assert(stack.size == 0)
  }

  "push" should "add an expression to the stack" in {
    val stack = new Stack
    val expr = Ident("x", 0)(?)
    val result = stack.push(expr, RAX)
    assert(stack.size == 1)
    assert(result == Push(RAX))
  }

  it should "add 2 expressions to the stack" in {
    val stack = new Stack
    val expr1 = Ident("x", 0)(?)
    val expr2 = Ident("x", 1)(?)
    val result1 = stack.push(expr1, RAX)
    val result2 = stack.push(expr2, RCX)
    assert(stack.size == 2)
    assert(result1 == Push(RAX))
    assert(result2 == Push(RCX))
  }

  it should "add a value to the stack" in {
    val stack = new Stack
    val result = stack.push(D32, RAX)
    assert(stack.size == 1)
    assert(result == Push(RAX))
  }

  "reserve" should "reserve space for an identifier" in {
    val stack = new Stack
    val ident = Ident("x", 0)(KnownType.Int)
    val result = stack.reserve(ident)
    assert(stack.size == 1)
    assert(result == Subtract(RSP, ImmediateVal(Q64.toInt)))
  }

  it should "reserve space for a register" in {
    val stack = new Stack
    val result = stack.reserve(RAX)
    assert(stack.size == 1)
    assert(result == Subtract(RSP, ImmediateVal(Q64.toInt)))
  }

  it should "reserve space for multiple values" in {
    val stack = new Stack
    val result = stack.reserve(D32, Q64, B8)
    assert(stack.size == 3)
    assert(result == Subtract(RSP, ImmediateVal(Q64.toInt * 3)))
  }

  "pop" should "remove the last value from the stack" in {
    val stack = new Stack
    stack.push(D32, RAX)
    val result = stack.pop(RAX)
    assert(stack.size == 0)
    assert(result == Pop(RAX))
  }

  "drop" should "remove the last 2 value from the stack" in {
    val stack = new Stack
    stack.push(D32, RAX)
    stack.push(Q64, RAX)
    stack.push(B8, RAX)
    val result = stack.drop(2)
    assert(stack.size == 1)
    assert(result == Add(RSP, ImmediateVal(Q64.toInt * 2)))
  }

  "withScope" should "reset stack after block" in {
    val stack = new Stack
    stack.push(D32, RAX)
    stack.push(Q64, RCX)
    stack.push(B8, RDX)
    val result = stack.withScope(() =>
      Chain(
        stack.push(Q64, RSI),
        stack.push(B8, RDI),
        stack.push(B8, RBP)
      )
    )
    assert(stack.size == 3)
    assert(
      result == Chain(
        Push(RSI),
        Push(RDI),
        Push(RBP),
        Add(RSP, ImmediateVal(Q64.toInt * 3))
      )
    )
  }

  "accessVar" should "return the correctly-sized memory location for the identifier" in {
    val stack = new Stack
    val id = Ident("x", 0)(KnownType.Int)
    stack.push(Q64, RAX)
    stack.push(id, RCX)
    stack.push(B8, RDX)
    stack.push(D32, RSI)
    val result = stack.accessVar(Ident("x", 0)(KnownType.Int))
    assert(result == MemLocation(RSP, Q64.toInt * 2, opSize = Some(KnownType.Int.size)))
  }

  "contains" should "return true if the stack contains the identifier" in {
    val stack = new Stack
    val id = Ident("x", 0)(KnownType.Int)
    stack.push(D32, RAX)
    stack.push(id, RCX)
    stack.push(B8, RDX)
    assert(stack.contains(id))
    assert(!stack.contains(Ident("x", 1)(KnownType.Int)))
  }

  "head" should "return the correct memory location for the last element" in {
    val stack = new Stack
    val id = Ident("x", 0)(KnownType.Int)
    stack.push(D32, RAX)
    stack.push(id, RCX)
    stack.push(B8, RDX)
    val result = stack.head
    assert(result == MemLocation(RSP, opSize = Some(B8)))
  }
}
