MemTestSupport.scala

package modernsar.util

import chisel3._
import chiseltest.internal.Context
import treadle.TreadleTester
import chisel3.experimental.FixedPoint
import chisel3.experimental.Interval
import chisel3.experimental.DataMirror
import chiseltest.LiteralTypeException
import chisel3.experimental.EnumType
import chisel3.experimental.BundleLiterals._

object MemTestSupport
{
  //TODO: Put on GitHub / make pull request from this?
  implicit class TestableMem[T <: Data](x: MemBase[T])
  {
    //read the private treadle-instance from the Context, using reflection. The value
    //is cached for every memory.
    //@note: This obviously fails, if Treadle is not used as a backend
    private def getTreadleInst() =
    {
      val key = (x, classOf[TreadleTester])
      Context().backend.getVar(key) match {
        case None =>
          val field1 = Context().backend.getClass().getDeclaredField("tester")
          field1.setAccessible(true)
          val inst1 = field1.get(Context().backend)
          val field2 = inst1.getClass().getDeclaredField("tester")
          field2.setAccessible(true)
          val inst2 = field2.get(inst1).asInstanceOf[TreadleTester]
          Context().backend.setVar(key, inst2)
          inst2
        case Some(value) => value.asInstanceOf[TreadleTester]
      }
    }

    private def pokeBits(index: Int, value: BigInt, subPath: String = "") = getTreadleInst().pokeMemory(x.instanceName + subPath, index, value)

    private def pokeSubElement(index: Int, value: Data, subPath: String): Unit = value match {
      //TODO: does interval have litValue?
      case _:Bool | _:UInt | _:SInt | _:FixedPoint | _:Interval => pokeBits(index, value.litValue, subPath)
      case (rec: Record) => {
        if (subPath == "") {
          require(DataMirror.checkTypeEquivalence(x.t, value), s"Record type mismatch")
        }
        
        rec.elements foreach { case (name, value) =>
          pokeSubElement(index, value, subPath + "_" + name)
        }
      }
    }

    def poke(index: Int, value: T): Unit = pokeSubElement(index, value, "")
    
    def poke(values: Seq[T], startIndex: Int = 0): Unit =
    {
      for ((v, i) <- values.zipWithIndex) {
        poke(i + startIndex, v)
      }
    }

    private def peekBits(index: Int, subPath: String = ""): BigInt = getTreadleInst().peekMemory(x.instanceName + subPath, index)

    private def peekSubElement(index: Int, subPath: String, dataType: Data): T = dataType match {
      case (x: Bool) => peekBits(index, subPath) match {
        case x: BigInt if x == 0 => false.B.asInstanceOf[T]
        case x: BigInt if x == 1 => true.B.asInstanceOf[T]
        case x => throw new LiteralTypeException(s"peeked Bool with value $x not 0 or 1")
      }
      case (x: UInt) => peekBits(index, subPath).asUInt(DataMirror.widthOf(x)).asInstanceOf[T]
      case (x: SInt) => peekBits(index, subPath).asSInt(DataMirror.widthOf(x)).asInstanceOf[T]
      case (x: FixedPoint) => {
        val multiplier = math.pow(2, x.binaryPoint.get)
        (peekBits(index, subPath).toDouble / multiplier).F(x.binaryPoint).asInstanceOf[T]
      }
      case x: Interval =>
        peekBits(index, subPath).I(x.binaryPoint).asInstanceOf[T]
      case (x: Record) => {
        val elementValueFns = x.elements.map { case (name: String, elt: Data) =>
          (y: Record) => (y.elements(name), peekSubElement(index, subPath + "_" + name, elt))
        }.toSeq
        x.Lit(elementValueFns: _*).asInstanceOf[T]
      }
      case (x: EnumType) => {
        throw new NotImplementedError(s"peeking enums ($x) not yet supported, need programmatic enum construction")
      }
      case x => throw new LiteralTypeException(s"don't know how to peek $x")
    }

    def peek(index: Int): T = peekSubElement(index, "", x.t)

    private def expectBits(index: Int, value: BigInt, subPath: String = "", message: String = "") = getTreadleInst().expectMemory(x.instanceName + subPath, index, value, message)

    private def expectSubElement(index: Int, value: Data, subPath: String, message: String = ""): Unit = value match {
      //TODO: does interval have litValue?
      case _:Bool | _:UInt | _:SInt | _:FixedPoint | _:Interval => expectBits(index, value.litValue, subPath, message)
      case (rec: Record) => {
        if (subPath == "") {
          require(DataMirror.checkTypeEquivalence(x.t, value), s"Record type mismatch")
        }
        
        rec.elements foreach { case (name, value) =>
          expectSubElement(index, value, subPath + "_" + name, message)
        }
      }
      case x => throw new LiteralTypeException(s"don't know how to peek $x")
    }

    def expect(index: Int, value: T, message: String = "") = expectSubElement(index, value, "", message)

    def expectSeq(values: Seq[T], startIndex: Int = 0, message: String = "") =
    {
      for ((v, i) <- values.zipWithIndex) {
        expect(i + startIndex, v)
      }
    }
  }
}