// SPDX-License-Identifier: Apache-2.0

package chiselTests.util

import chisel3._
import chisel3.stage.WarningConfigurationAnnotation
import chisel3.testers.BasicTester
import chisel3.util.{log2Up, Counter, SparseVec}
import chiselTests.{ChiselFlatSpec, Utils}
import _root_.circt.stage.ChiselStage

// This test is checking that a SparseVec behaves _exactly_ like a dynamic
// index.  By varying the mapping and size, different corner cases of behavior
// are checked.
class SparseVecTest(size: Int, tpe: UInt, mapping: Seq[(Int, UInt)]) extends BasicTester {
  val paddedSize = BigInt(log2Up(size)).pow(2).toInt

  val reference = {
    val ttable = Wire(Vec(size, tpe))
    ttable.foreach(_ := DontCare)
    mapping.foreach {
      case (index, data) =>
        ttable(index) := data
    }

    ttable
  }

  val sparseVec = SparseVec(size, tpe, mapping)

  val (index, wrap) = Counter(0 until paddedSize)

  val failed = RegInit(Bool(), false.B)
  val resultBinary = sparseVec(index, SparseVec.Binary)
  val resultOneHot = sparseVec(index, SparseVec.OneHot)
  val resultIfElse = sparseVec(index, SparseVec.IfElse)
  when(resultBinary =/= reference(index) || resultOneHot =/= reference(index) || resultIfElse =/= reference(index)) {
    failed := true.B
    assert(false.B)
  }

  when(wrap) {
    stop()
  }

}

class SparseVecSpec extends ChiselFlatSpec with Utils {
  "SparseVec" should "work for a complete user-specified mapping" in {
    assertTesterPasses(
      new SparseVecTest(
        4,
        UInt(3.W),
        Seq(
          0 -> 1.U,
          1 -> 2.U,
          2 -> 3.U,
          3 -> 4.U
        )
      ),
      // TODO: This warning suppression doesn't work for some reason.
      annotations = Seq(WarningConfigurationAnnotation("id=4:s"))
    )
  }

  it should "work for a mapping that includes default values" in {
    assertTesterPasses(
      new SparseVecTest(
        4,
        UInt(3.W),
        Seq(
          0 -> 1.U,
          1 -> 2.U,
          3 -> 4.U
        )
      ),
      annotations = Seq(WarningConfigurationAnnotation("id=4:s"))
    )
  }

  it should "work for a mapping that includes out-of-bounds values" in {
    assertTesterPasses(
      new SparseVecTest(
        3,
        UInt(3.W),
        Seq(
          0 -> 1.U,
          1 -> 2.U,
          2 -> 3.U
        )
      ),
      annotations = Seq(WarningConfigurationAnnotation("id=4:s"))
    )
  }

  "SparseVec error behavior" should "disallow non-unique mappings" in {
    val exception = intercept[ChiselException] {
      ChiselStage.convert(new Module {
        SparseVec(2, UInt(1.W), Seq(0 -> 0.U, 0 -> 1.U))
      })
    }
    exception.getMessage should include("non-unique mapping in SparseVec")
  }

  it should "disallow a SparseVec write" in {
    intercept[ChiselException] {
      ChiselStage.convert(new Module {
        val vec = SparseVec(2, UInt(1.W), Seq(0 -> 0.U, 1 -> 1.U))
        vec(0.U) := 1.U
      })
    }
  }
}
