WIP: rethink and simplify API, make DRYer

Author: kammoh

src/main/scala/chisel3/simulator/PeekPokeAPI.scala

@@ -11,6 +11,22 @@ import chisel3.internal.ExceptionHelpers
 import firrtl.options.StageUtils.dramaticMessage
 import scala.util.control.NoStackTrace
 
+private[simulator] trait Peekable[T <: Data] {
+  def peek(): T
+
+  def expect(expected: T, buildMessage: (String, T) => String)(implicit sourceInfo: SourceInfo): Unit
+
+  def expect(expected: T, message: String)(implicit sourceInfo: SourceInfo): Unit =
+    expect(expected, (_, _) => message)
+
+  def expect(expected: T)(implicit sourceInfo: SourceInfo): Unit =
+    expect(expected, (observed, expected) => s"Expectation failed: observed value $observed != $expected")
+}
+
+private[simulator] trait Pokable[T <: Data] {
+  def poke(literal: T): Unit
+}
+
 object PeekPokeAPI extends PeekPokeAPI
 
 trait PeekPokeAPI {
@@ -24,20 +40,37 @@ trait PeekPokeAPI {
         )
       )
       with NoStackTrace
+
   object FailedExpectationException {
     def apply[T](
-      observed:     T,
-      expected:     T,
-      message:      String,
-      sourceInfo:   SourceInfo,
-      extraContext: Seq[String]
+      observed:   T,
+      expected:   T,
+      message:    String,
+      sourceInfo: SourceInfo
     ): FailedExpectationException[T] = {
+      val extraContext =
+        sourceInfo match {
+          case sl: SourceLine =>
+            ExceptionHelpers.getErrorLineInFile(Seq(), sl)
+          case _ =>
+            Seq()
+        }
       val fullMessage = s"$message ${sourceInfo.makeMessage()}" +
         (if (extraContext.nonEmpty) s"\n${extraContext.mkString("\n")}" else "")
-      new FailedExpectationException(observed, expected, fullMessage)
+      new FailedExpectationException[T](observed, expected, fullMessage)
     }
   }
 
+  sealed trait AnyTestableData[T <: Data] {
+    protected def data: T
+
+    protected val simulatedModule = AnySimulatedModule.current
+
+    protected final def simulationPort = simulatedModule.port(data)
+  }
+
+  private[simulator] trait PeekPokable[T <: Data] extends Peekable[T] with Pokable[T] with AnyTestableData[T]
+
   implicit class testableClock(clock: Clock) extends AnyTestableData[Clock] {
     val data = clock
 
@@ -72,72 +105,16 @@ trait PeekPokeAPI {
     }
   }
 
-  trait AnyTestableData[T <: Data] {
-    protected def data: T
-
-    protected val simulatedModule = AnySimulatedModule.current
-
-    protected final def simulationPort = simulatedModule.port(data)
-  }
-
-  trait Peekable[T <: Data] extends AnyTestableData[T] {
-    def peek(): T
-
-    protected def isSigned = false
-
-    def expect(expected: T, buildMessage: (T, T) => String)(implicit sourceInfo: SourceInfo): Unit
-
-    def expect(expected: T, message: String)(implicit sourceInfo: SourceInfo): Unit =
-      expect(expected, (_, _) => message)
-
-    def expect(expected: T)(implicit sourceInfo: SourceInfo): Unit =
-      expect(expected, (observed, expected) => s"Expectation failed: observed value $observed != $expected")
-
-    def expect[U](
-      expected:     U,
-      encode:       (Simulation.Value) => U,
-      buildMessage: (U, U) => String,
-      sourceInfo:   SourceInfo
-    ): Unit = {
-      simulatedModule.willPeek()
-
-      simulationPort.check(isSigned = isSigned) { observedValue =>
-        val observed = encode(observedValue)
-        if (observed != expected) {
-          val extraContext =
-            sourceInfo match {
-              case sl: SourceLine =>
-                ExceptionHelpers.getErrorLineInFile(Seq(), sl)
-              case _ =>
-                Seq()
-            }
-          throw FailedExpectationException(
-            observed,
-            expected,
-            buildMessage(observed, expected),
-            sourceInfo,
-            extraContext
-          )
-        }
-      }
-    }
-  }
-
-  trait Pokable[T <: Data] extends AnyTestableData[T] {
-    def poke(literal: T): Unit
-  }
-
-  trait PeekPokable[T <: Data] extends Peekable[T] with Pokable[T]
-
   sealed trait TestableElement[T <: Element] extends PeekPokable[T] {
+    protected def isSigned = false
 
-    private[simulator] def encode(width: Int, value: BigInt): T
+    private[simulator] protected def encode(width: Int, value: BigInt): T
 
-    protected final def encode(value: Simulation.Value): T = {
+    private[simulator] final def encode(value: Simulation.Value): T = {
       encode(value.bitCount, value.asBigInt)
     }
 
-    protected def peekValue(): Simulation.Value = {
+    protected final def peekValue(): Simulation.Value = {
       simulatedModule.willPeek()
       simulationPort.get(isSigned = isSigned)
     }
@@ -151,51 +128,67 @@ trait PeekPokeAPI {
       simulationPort.set(value)
     }
 
+    private[simulator] protected def check[U](checkFn: Simulation.Value => Unit): Unit = {
+      simulatedModule.willPeek()
+      simulationPort.check(isSigned = isSigned)(checkFn)
+    }
+
+    private[simulator] protected final def expect[U](
+      expected:            U,
+      sameValue:           (Simulation.Value, U) => Boolean,
+      observedValToString: Simulation.Value => String,
+      buildMessage:        (String, U) => String,
+      sourceInfo:          SourceInfo
+    ): Unit = {
+      check(observedValue =>
+        if (!sameValue(observedValue, expected)) {
+          val observedStr = observedValToString(observedValue)
+          throw FailedExpectationException(
+            observedStr,
+            expected.toString,
+            buildMessage(observedStr, expected),
+            sourceInfo
+          )
+        }
+      )
+    }
+
     override def expect(expected: T)(implicit sourceInfo: SourceInfo): Unit = {
       expect(
         expected,
-        (observed, expected) => s"Expectation failed: observed value $observed != $expected"
+        (observed: String, expected: T) => s"Expectation failed: observed value $observed != $expected"
       )
     }
 
-    override def expect(expected: T, buildMessage: (T, T) => String)(implicit sourceInfo: SourceInfo): Unit = {
+    override def expect(expected: T, buildMessage: (String, T) => String)(implicit sourceInfo: SourceInfo): Unit = {
       require(expected.isLit, s"Expected value: $expected must be a literal")
 
-      simulatedModule.willPeek()
-
-      simulationPort.check(isSigned = isSigned) { observedValue =>
-        val observed = encode(observedValue)
-        if (observedValue.asBigInt != expected.litValue) {
-          val extraContext =
-            sourceInfo match {
-              case sl: SourceLine =>
-                ExceptionHelpers.getErrorLineInFile(Seq(), sl)
-              case _ =>
-                Seq()
-            }
-          throw FailedExpectationException(
-            observed,
-            expected,
-            buildMessage(observed, expected),
-            sourceInfo,
-            extraContext
-          )
-        }
-      }
-    }
-
-    final def expect(expected: BigInt)(implicit sourceInfo: SourceInfo): Unit = {
       expect(
         expected,
-        _.asBigInt,
-        (observed: BigInt, expected: BigInt) => s"Expectation failed: observed value $observed != $expected",
+        (observed: Simulation.Value, expected: T) => observed.asBigInt == expected.litValue,
+        encode(_).toString,
+        (obs: String, exp: T) => buildMessage(obs, exp),
         sourceInfo
       )
     }
 
-    final def expect(expected: BigInt, message: String)(implicit sourceInfo: SourceInfo): Unit = {
-      expect(expected, _.asBigInt, (_: BigInt, _: BigInt) => message, sourceInfo)
-    }
+    final def expect(expected: BigInt, buildMessage: (String, BigInt) => String)(
+      implicit sourceInfo: SourceInfo
+    ): Unit = expect(
+      expected,
+      (obs: Simulation.Value, exp: BigInt) => obs.asBigInt == exp,
+      _.asBigInt.toString,
+      buildMessage,
+      sourceInfo
+    )
+
+    final def expect(expected: BigInt)(implicit sourceInfo: SourceInfo): Unit = expect(
+      expected,
+      (observed: String, expected: BigInt) => s"Expectation failed: observed value $observed != $expected"
+    )
+
+    final def expect(expected: BigInt, message: String)(implicit sourceInfo: SourceInfo): Unit =
+      expect(expected, (_: String, _: BigInt) => message)
 
   }
 
@@ -242,56 +235,50 @@ trait PeekPokeAPI {
     override def encode(width: Int, value: BigInt): T = {
       data.factory.all.find(_.litValue == value).get.asInstanceOf[T]
     }
-
-    override def expect(expected: T, buildMessage: (T, T) => String)(implicit sourceInfo: SourceInfo): Unit = {
-      require(expected.isLit, "Expected value must be a literal")
-      val observedSimValue = peekValue()
-      val observedVal = observedSimValue.asBigInt
-      val expectedVal = expected.litValue
-      if (observedVal != expectedVal) {
-        throw FailedExpectationException(observedVal, expectedVal, buildMessage(peek(), expected))
-      }
-    }
   }
 
   implicit class testableRecord[T <: Record](val data: T)(implicit sourceInfo: SourceInfo) extends PeekPokable[T] {
 
     override def peek(): T = {
-      val elementValueFns = data.elements.map { case (name: String, elt: Data) =>
-        (y: Record) =>
-          (
-            y.elements(name), {
-              new testableData(elt).peek()
-            }
-          )
-      }.toSeq
-      chiselTypeOf(data).Lit(elementValueFns: _*)
+      chiselTypeOf(data).Lit(
+        data.elements.toSeq.map { case (name: String, elt: Data) =>
+          (rec: Record) => rec.elements(name) -> elt.peek()
+        }: _*
+      )
     }
 
-    override def poke(value: T): Unit = {
-      data.elements.foreach { case (name, d) =>
-        d.poke(value.elements(name))
-      }
+    override def poke(value: T): Unit = data.elements.foreach { case (name, d) =>
+      d.poke(value.elements(name))
     }
 
-    override def expect(expected: T, buildMessage: (T, T) => String)(implicit sourceInfo: SourceInfo): Unit = {
-      // TODO: not checking for isLit to allow partially specified expected record
+    def expect(expected: T, buildMessage: (String, T) => String, allowPartial: Boolean)(
+      implicit sourceInfo: SourceInfo
+    ): Unit = {
       require(DataMirror.checkTypeEquivalence(data, expected), "Type mismatch")
 
       // FIXME: I can't understand why but _not_ getting the peeked value as a `val` beforehand results in infinite recursion
-      val peekedValue = peek()
+      val peekedValue = peek().toString
 
       data.elements.foreach { case (name, d) =>
         expected.elements(name) match {
-          case DontCare => // missing fields are DontCare
-          case ve =>
+          case DontCare =>
+            if (!allowPartial) {
+              throw new Exception(
+                s"Field '$name' is not initiazlized in the expected value $expected"
+              )
+            }
+          case exp =>
             d.expect(
-              ve,
-              (obs, exp) => s"${buildMessage(peekedValue, expected)}:\n Expected value of $name to be $exp, got $obs"
+              exp,
+              (obs: String, _) =>
+                s"${buildMessage(peekedValue, expected)}:\n Expected value of field '$name' to be $exp, got $obs"
             )
         }
       }
     }
+
+    override def expect(expected: T, buildMessage: (String, T) => String)(implicit sourceInfo: SourceInfo): Unit =
+      expect(expected, buildMessage, allowPartial = false)
   }
 
   implicit class testableData[T <: Data](val data: T) extends PeekPokable[T] {
@@ -312,11 +299,11 @@ trait PeekPokeAPI {
 
     override def expect(
       expected:     T,
-      buildMessage: (T, T) => String
+      buildMessage: (String, T) => String
     )(implicit sourceInfo: SourceInfo): Unit = {
 
-      def buildMsgFn[S](observed: S, expected: S): String =
-        buildMessage(observed.asInstanceOf[T], expected.asInstanceOf[T])
+      def buildMsgFn[S](observed: String, expected: S): String =
+        buildMessage(observed, expected.asInstanceOf[T])
 
       (data, expected) match {
         case (dat: Bool, exp: Bool) =>
@@ -334,15 +321,15 @@ trait PeekPokeAPI {
             exp.length == dat.length,
             s"Vec length mismatch: Data port has ${dat.length} elements while the expected value is of length ${exp.length}"
           )
-          val peekedValue = dat.peek()
+          val peekedValue = dat.peek().toString
           dat.getElements.zip(exp.getElements).zipWithIndex.foreach { case ((datEl, valEl), index) =>
             valEl match {
               case DontCare =>
               // TODO: missing elements?
               case ve =>
                 datEl.expect(
                   ve,
-                  (o, e) => buildMessage(peekedValue.asInstanceOf[T], exp.asInstanceOf[T]) + s" at index $index"
+                  (obs: String, exp: Data) => buildMessage(peekedValue, exp.asInstanceOf[T]) + s" at index $index"
                 )
             }
           }