Add FirrtlCheckers and scalatest helpers for testing

1 files changed, 108 insertions, 0 deletions

diff --git a/src/test/scala/firrtlTests/FirrtlSpec.scala b/src/test/scala/firrtlTests/FirrtlSpec.scala
index 07f83142..b71e51e2 100644
--- a/src/test/scala/firrtlTests/FirrtlSpec.scala
+++ b/src/test/scala/firrtlTests/FirrtlSpec.scala
@@ -11,6 +11,7 @@ import org.scalatest.prop._
 import scala.io.Source
 
 import firrtl._
+import firrtl.ir._
 import firrtl.Parser.UseInfo
 import firrtl.annotations._
 import firrtl.transforms.{DontTouchAnnotation, NoDedupAnnotation}
@@ -119,10 +120,117 @@ trait FirrtlMatchers extends Matchers {
   }
 }
 
+object FirrtlCheckers extends FirrtlMatchers {
+  import matchers._
+  implicit class TestingFunctionsOnCircuitState(val state: CircuitState) extends AnyVal {
+    def search(pf: PartialFunction[Any, Boolean]): Boolean = state.circuit.search(pf)
+  }
+  implicit class TestingFunctionsOnCircuit(val circuit: Circuit) extends AnyVal {
+    def search(pf: PartialFunction[Any, Boolean]): Boolean = {
+      val f = pf.lift
+      def rec(node: Any): Boolean = {
+        f(node) match {
+          // If the partial function is defined on this node, return its result
+          case Some(res) => res
+          // Otherwise keep digging
+          case None =>
+            require(node.isInstanceOf[Product] || !node.isInstanceOf[FirrtlNode],
+                    "Error! Unexpected FirrtlNode that does not implement Product!")
+            val iter = node match {
+              case p: Product => p.productIterator
+              case i: Iterable[Any] => i.iterator
+              case _ => Iterator.empty
+            }
+            iter.foldLeft(false) {
+              case (res, elt) => if (res) res else rec(elt)
+            }
+        }
+      }
+      rec(circuit)
+    }
+  }
+
+  /** Checks that the emitted circuit has the expected line, both will be normalized */
+  def containLine(expectedLine: String) = new CircuitStateStringMatcher(expectedLine)
+
+  class CircuitStateStringMatcher(expectedLine: String) extends Matcher[CircuitState] {
+    override def apply(state: CircuitState): MatchResult = {
+      val emitted = state.getEmittedCircuit.value
+      MatchResult(
+        emitted.split("\n").map(normalized).contains(normalized(expectedLine)),
+        emitted + "\n did not contain \"" + expectedLine + "\"",
+        s"${state.circuit.main} contained $expectedLine"
+      )
+    }
+  }
+
+  def containTree(pf: PartialFunction[Any, Boolean]) = new CircuitStatePFMatcher(pf)
+
+  class CircuitStatePFMatcher(pf: PartialFunction[Any, Boolean]) extends Matcher[CircuitState] {
+    override def apply(state: CircuitState): MatchResult = {
+      MatchResult(
+        state.search(pf),
+        state.circuit.serialize + s"\n did not contain $pf",
+        s"${state.circuit.main} contained $pf"
+      )
+    }
+  }
+}
+
 abstract class FirrtlPropSpec extends PropSpec with PropertyChecks with FirrtlRunners with LazyLogging
 
 abstract class FirrtlFlatSpec extends FlatSpec with FirrtlRunners with FirrtlMatchers with LazyLogging
 
+// Who tests the testers?
+class TestFirrtlFlatSpec extends FirrtlFlatSpec {
+  import FirrtlCheckers._
+
+  val c = parse("""
+    |circuit Test:
+    |  module Test :
+    |    input in : UInt<8>
+    |    output out : UInt<8>
+    |    out <= in
+    |""".stripMargin)
+  val state = CircuitState(c, ChirrtlForm)
+  val compiled = (new LowFirrtlCompiler).compileAndEmit(state, List.empty)
+
+  // While useful, ScalaTest helpers should be used over search
+  behavior of "Search"
+
+  it should "be supported on Circuit" in {
+    assert(c search {
+      case Connect(_, Reference("out",_), Reference("in",_)) => true
+    })
+  }
+  it should "be supported on CircuitStates" in {
+    assert(state search {
+      case Connect(_, Reference("out",_), Reference("in",_)) => true
+    })
+  }
+  it should "be supported on the results of compilers" in {
+    assert(compiled search {
+      case Connect(_, WRef("out",_,_,_), WRef("in",_,_,_)) => true
+    })
+  }
+
+  // Use these!!!
+  behavior of "ScalaTest helpers"
+
+  they should "work for lines of emitted text" in {
+    compiled should containLine (s"input in : UInt<8>")
+    compiled should containLine (s"output out : UInt<8>")
+    compiled should containLine (s"out <= in")
+  }
+
+  they should "work for partial functions matching on subtrees" in {
+    val UInt8 = UIntType(IntWidth(8)) // BigInt unapply is weird
+    compiled should containTree { case Port(_, "in", Input, UInt8) => true }
+    compiled should containTree { case Port(_, "out", Output, UInt8) => true }
+    compiled should containTree { case Connect(_, WRef("out",_,_,_), WRef("in",_,_,_)) => true }
+  }
+}
+
 /** Super class for execution driven Firrtl tests */
 abstract class ExecutionTest(name: String, dir: String, vFiles: Seq[String] = Seq.empty) extends FirrtlPropSpec {
   property(s"$name should execute correctly") {