package wacc

import org.scalatest.flatspec.AnyFlatSpec

class LabelGeneratorSpec extends AnyFlatSpec {
  import microWacc._
  import assemblyIR.{LabelDef, LabelArg, Directive}
  import types.?

  "getLabel" should "return unique labels" in {
    val labelGenerator = new LabelGenerator
    val labels = (1 to 10).map(_ => labelGenerator.getLabel())
    assert(labels.distinct.length == labels.length)
  }

  "getLabelDef" should "return unique labels" in {
    assert(
      LabelDef("test") == LabelDef("test") &&
        LabelDef("test").hashCode == LabelDef("test").hashCode,
      "Sanity check: LabelDef should be case-classes"
    )

    val labelGenerator = new LabelGenerator
    val labels = (List(
      Builtin.Exit,
      Builtin.Printf,
      Ident("exit", 0)(?),
      Ident("test", 0)(?)
    ) ++ RuntimeError.all.toList).map(labelGenerator.getLabelDef(_))
    assert(labels.distinct.length == labels.length)
  }

  "getLabelArg" should "return unique labels" in {
    assert(
      LabelArg("test") == LabelArg("test") &&
        LabelArg("test").hashCode == LabelArg("test").hashCode,
      "Sanity check: LabelArg should be case-classes"
    )

    val labelGenerator = new LabelGenerator
    val labels = (List(
      Builtin.Exit,
      Builtin.Printf,
      Ident("exit", 0)(?),
      Ident("test", 0)(?),
      "test",
      "test",
      "test3"
    ) ++ RuntimeError.all.toList).map {
      case s: String                      => labelGenerator.getLabelArg(s)
      case t: (CallTarget | RuntimeError) => labelGenerator.getLabelArg(t)
    }
    assert(labels.distinct.length == labels.distinct.length)
  }

  it should "return consistent labels to getLabelDef" in {
    val labelGenerator = new LabelGenerator
    val targets = (List(
      Builtin.Exit,
      Builtin.Printf,
      Ident("exit", 0)(?),
      Ident("test", 0)(?)
    ) ++ RuntimeError.all.toList)
    val labelDefs = targets.map(labelGenerator.getLabelDef(_).toString.dropRight(1)).toSet
    val labelArgs = targets.map(labelGenerator.getLabelArg(_).toString).toSet
    assert(labelDefs == labelArgs)
  }

  "generateConstants" should "generate de-duplicated labels for strings" in {
    val labelGenerator = new LabelGenerator
    val strings = List("hello", "world", "hello\u0000world", "hello", "Hello")
    val distincts = strings.distinct.length
    val labels = strings.map(labelGenerator.getLabelArg(_).toString).toSet
    val asmLines = labelGenerator.generateConstants
    assert(
      asmLines.collect { case LabelDef(name) =>
        name
      }.length == distincts
    )
    assert(
      asmLines.collect { case Directive.Asciz(str) => str }.length == distincts
    )
    assert(asmLines.collect { case LabelDef(name) => name }.toList.toSet == labels)
  }
}
