package wacc

object sizeExtensions {
  import microWacc._
  import types._
  import assemblyIR.Size

  extension (expr: Expr) {

    /** Calculate the size (bytes) of the heap required for the expression. */
    def heapSize: Int = (expr, expr.ty) match {
      case (ArrayLiter(elems), KnownType.Array(KnownType.Char)) =>
        KnownType.Int.size.toInt + elems.size.toInt * KnownType.Char.size.toInt
      case (ArrayLiter(elems), ty) =>
        KnownType.Int.size.toInt + elems.size * ty.elemSize.toInt
      case _ => expr.ty.size.toInt
    }
  }

  extension (ty: SemType) {

    /** Calculate the size (bytes) of a type in a register. */
    def size: Size = ty match {
      case KnownType.Int                                                    => Size.D32
      case KnownType.Bool | KnownType.Char                                  => Size.B8
      case KnownType.String | KnownType.Array(_) | KnownType.Pair(_, _) | ? => Size.Q64
    }

    def elemSize: Size = ty match {
      case KnownType.Array(elem) => elem.size
      case KnownType.Pair(_, _)  => Size.Q64
      case _                     => ty.size
    }
  }
}